优化numpy ndarray索引操作

I have a numpy operation that looks like the following:

我有一个看起来如下的numpy操作:

 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.

其中x,y和c具有相同的形状。

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

是否可以使用numpy的高级索引来加速此操作?

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.

但是,我没有得到我预期的结果。

2 个解决方案

#1

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以使形状与其余索引匹配。快速实现此目的的方法是使用None进行索引(相当于numpy.newaxis):

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

#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

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

#1

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

#2