Optimieren eines numpy ndarray Indexierungsvorgang

Question

Ich habe eine numpy Operation, die wie folgt aussieht:

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

wo x, y und c die gleiche Form haben.

Ist es möglich, numpy der Indexierung zu verwenden, um diese Operation zu beschleunigen?

Ich habe versucht, mit: Allerdings

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

, ich das Ergebnis nicht bekommen ich erwartet hatte.

Answer 1

Die Index-Arrays müssen broadcastable sein, damit dies funktioniert. Die einzige Änderung erforderlich ist, eine Achse zu dem ersten Index hinzuzufügen i die Form mit dem Rest zu passen. Der schnelle Weg dies zu tun ist durch Indizierung mit None (die numpy.newaxis entspricht):

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

Answer 2

Verwendung 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 und Verifikation

Funktionen definieren -

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-Eingabearrays und 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