Optimize a numpy ndarray indexing operation

Question

I have a numpy operation that looks like the following:

 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.

Is it possible to use numpy's advanced indexing to speed this operation up?

I tried using:

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

However, I didn't get the result I expected.

Divakar · Accepted Answer · 2015-09-17 06:04:20Z

6

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

Setup input arrays and benchmark -

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

edited Sep 17, 2015 at 6:04

answered Sep 17, 2015 at 5:34

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

hpaulj Over a year ago

Would np.ravel_multi_index simplify this index generation?

Divakar Over a year ago

@hpaulj That might, haven't used that function that much though. I guess I prefer the crude way of calculating this :)

user2379410 · Accepted Answer · 2015-09-17 07:19:17Z

4

The indexing arrays need to be broadcastable for this to work. The only change needed is to add an axis to the first index i to match the shape with the rest. The 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

answered Sep 17, 2015 at 7:19

user2379410

1 Comment

hpaulj Over a year ago

In other tests I've found that flattened indexing tends to be faster, even taking into account the time needed to generate the flat index. But your [:,None] is the key to both approaches. Your version is also clearer.

Collectives™ on Stack Overflow

Optimize a numpy ndarray indexing operation

2 Answers 2

2 Comments

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related