Optimizing the numpy ndarray indexing operation

Question

I have a numpy operation that looks like this:

for i in range(i_max): for j in range(j_max): r[i, j, x[i, j], y[i, j]] = c[i, j] 

where x , y and c have the same shape.

Can numpy indexing be used to speed up this operation?

I tried using:

 i = numpy.arange(i_max) j = numpy.arange(j_max) r[i, j, x, y] = c 

However, I did not get the expected result.

Answer 1

Indexing arrays must be broadcasted for this to work. The only change that needs to be made is to add an axis to the first index i to fit the form with the rest. A quick way to accomplish this is by indexing with None (which is equivalent to numpy.newaxis ):

 i = numpy.arange(i_max) j = numpy.arange(j_max) r[i[:,None], j, x, y] = c 

Answer 2

Using linear indexing -

 d0,d1,d2,d3 = r.shape np.put(r,np.arange(i_max)[:,None]*d1*d2*d3 + np.arange(j_max)*d2*d3 + x*d3 +y,c) 

---

Benchmarking and verification

Define Functions -

 def linear_indx(r,x,y,c,i_max,j_max): d0,d1,d2,d3 = r.shape np.put(r,np.arange(i_max)[:,None]*d1*d2*d3 + np.arange(j_max)*d2*d3 + x*d3 +y,c) return r def org_app(r,x,y,c,i_max,j_max): for i in range(i_max): for j in range(j_max): r[i, j, x[i,j], y[i,j]] = c[i,j] return r 

Setting input arrays and reference values ​​-

 In [134]: # Setup input arrays ...: i_max = 40 ...: j_max = 50 ...: D0 = 60 ...: D1 = 70 ...: N = 80 ...: ...: r = np.zeros((D0,D1,N,N)) ...: c = np.random.rand(i_max,j_max) ...: ...: x = np.random.randint(0,N,(i_max,j_max)) ...: y = np.random.randint(0,N,(i_max,j_max)) ...: In [135]: # Make copies for testing, as both functions make in-situ changes ...: r1 = r.copy() ...: r2 = r.copy() ...: In [136]: # Verify results by comparing with original loopy approach ...: np.allclose(linear_indx(r1,x,y,c,i_max,j_max),org_app(r2,x,y,c,i_max,j_max)) Out[136]: True In [137]: # Make copies for testing, as both functions make in-situ changes ...: r1 = r.copy() ...: r2 = r.copy() ...: In [138]: %timeit linear_indx(r1,x,y,c,i_max,j_max) 10000 loops, best of 3: 115 µs per loop In [139]: %timeit org_app(r2,x,y,c,i_max,j_max) 100 loops, best of 3: 2.25 ms per loop