2

I am currently trying to vectorize some code I had written using a large for loop in Python. The vectorized code is as follows:

rho[pi,pj] += (rho_coeff*dt)*i_frac*j_frac
rho[pi+1,pj] += (rho_coeff*dt)*ip1_frac*j_frac
rho[pi,pj+1] += (rho_coeff*dt)*i_frac*jp1_frac
rho[pi+1,pj+1] += (rho_coeff*dt)*ip1_frac*jp1_frac

Each of pi, pj, dt, i_frac, j_frac, ip1_frac, jp1_frac is a numpy array of one dimension and all of the same length. rho is a two dimensional numpy array. pi and pj make up a list of coordinates (pi,pj) which indicate which element of the matrix rho is modified. The modification involves the addition of the (rho_coeff*dt)*i_frac*j_frac term to the (pi,pj) element as well as addition of similar terms to neighbouring elements: (pi+1,pj), (pi,pj+1) and (pi+1,pj+1). Each coordinate in the list (pi, pj) has a unique dt, i_frac, j_frac, ip1_frac and jp1_frac associated with it.

The problem is that the list can have (and always will have) repeating coordinates. So instead of successively adding to rho each time the same coordinate is encountered in the list, it only adds the term corresponding to the last repeating coordinate. This problem is described briefly with an example in the Tentative Numpy Tutorial under fancy indexing with arrays of indices (see the last three examples before boolean indexing). Unfortunately they did not provide a solution to this.

Is there a way of doing this operation without resorting to a for loop? I am trying to optimize for performance and want to do away with a loop if possible.

FYI: this code forms part of a 2D particle tracking algorithm where the charge from each particle is added to the four adjacent nodes of a mesh surrounding the particle's position based on volume fractions.

Glorfindel
  • 21,988
  • 13
  • 81
  • 109
nicholls
  • 289
  • 1
  • 3
  • 8

1 Answers1

1

You are going to have to figure out the repeated items and add them together before updating your array. The following code shows a way of doing that for your first update:

rows, cols = 100, 100
items = 1000

rho = np.zeros((rows, cols))
rho_coeff, dt, i_frac, j_frac = np.random.rand(4, items)
pi = np.random.randint(1, rows-1, size=(items,))
pj = np.random.randint(1, cols-1, size=(items,))

# The following code assumes pi and pj have the same dtype
pij = np.column_stack((pi, pj)).view((np.void,
                                      2*pi.dtype.itemsize)).ravel()

unique_coords, indices = np.unique(pij, return_inverse=True)
unique_coords = unique_coords.view(pi.dtype).reshape(-1, 2)
data = rho_coeff*dt*i_frac*j_frac
binned_data = np.bincount(indices, weights=data)
rho[tuple(unique_coords.T)] += binned_data

I think you can reuse all of the unique coordinate finding above for the other updates, so the following would work:

ip1_frac, jp1_frac = np.random.rand(2, items)

unique_coords[:, 0] += 1
data =  rho_coeff*dt*ip1_frac*j_frac
binned_data = np.bincount(indices, weights=data)
rho[tuple(unique_coords.T)] += binned_data

unique_coords[:, 1] += 1
data =  rho_coeff*dt*ip1_frac*jp1_frac
binned_data = np.bincount(indices, weights=data)
rho[tuple(unique_coords.T)] += binned_data

unique_coords[:, 0] -= 1
data =  rho_coeff*dt*i_frac*jp1_frac
binned_data = np.bincount(indices, weights=data)
rho[tuple(unique_coords.T)] += binned_data
Jaime
  • 65,696
  • 17
  • 124
  • 159
  • Beautiful. I had seen bincount used in 1D and was wondering at one point if there was a similar tool for my case, guess not. Also, could you clarify what is happening when you use `.view()` on `pij`? Would there be a better way to implement `pij` in the first place so I can avoid this line on every iteration (as this code is contained in a loop)? – nicholls Jul 22 '13 at 22:15
  • 1
    If instead of having `pi` and `pj` be independent arrays, they were the columns of a `pij` contiguous array of shape `(n, 2)`, i.e. `pi = pij[:, 0]` and `pj = pij[:, 1]`, you could skip the `np.column_stack` call. You would still need to take a flattened view of the void dtype into this array, but that should be a pretty inexpensive operation, as it doesn't involve any data copying. `view` is basically taking your two coordinates as a single variable that can be used to find the unique coordinates. – Jaime Jul 22 '13 at 22:23