NumPy: Evaulate Indexarray während vektorisierten Zuordnung

Question

Ich mag würde dieses NumPy Betrieb vektorisieren:

for j in range(yt): 
    for i in range(xt): 
     y[j, i] = x[idx[j, i], j, i] 

wo idx achsen 0 Index zu einer x slice enthält. Gibt es eine einfache Möglichkeit, dies zu tun?

Answer 1

können Sie verwenden:

J, I = np.ogrid[:yt, :xt] 
x[idx, J, I] 

Hier ist der Test:

import numpy as np 

yt, xt = 3, 5 
x = np.random.rand(10, 6, 7) 
y = np.zeros((yt, xt)) 
idx = np.random.randint(0, 10, (yt, xt)) 

for j in range(yt): 
    for i in range(xt): 
     y[j, i] = x[idx[j, i], j, i] 

J, I = np.ogrid[:yt, :xt] 
np.all(x[idx, J, I] == y) 

Answer 2

Hier ist ein Ansatz mit linear indexing -

zt,yt,xt = x.shape 
out = x.reshape(zt,-1)[idx.ravel(),np.arange(yt*xt)].reshape(-1,xt) 

---

Runtime-Tests & verifizieren Ausgang

Dieser Abschnitt vergleicht den vorgeschlagenen Ansatz in diesem Beitrag und die other orgid based solution auf Leistung und überprüft auch die Ausgänge.

Funktionsdefinitionen -

def original_app(x,idx): 
    _,yt,xt = x.shape 
    y = np.zeros((yt,xt)) 
    for j in range(yt): 
     for i in range(xt): 
      y[j, i] = x[idx[j, i], j, i] 
    return y 

def ogrid_based(x,idx): 
    _,yt,xt = x.shape 
    J, I = np.ogrid[:yt, :xt] 
    return x[idx, J, I] 

def reshape_based(x,idx):        
    zt,yt,xt = x.shape 
    return x.reshape(zt,-1)[idx.ravel(),np.arange(yt*xt)].reshape(-1,xt) 

Einrichtungseingänge -

In [56]: # Inputs 
    ...: zt,yt,xt = 100,100,100 
    ...: x = np.random.rand(zt,yt,xt) 
    ...: idx = np.random.randint(0,zt,(yt,xt)) 
...: 

Verify Ausgänge -

In [57]: np.allclose(original_app(x,idx),ogrid_based(x,idx)) 
Out[57]: True 

In [58]: np.allclose(original_app(x,idx),reshape_based(x,idx)) 
Out[58]: True 

Timings -

In [68]: %timeit original_app(x,idx) 
100 loops, best of 3: 6.97 ms per loop 

In [69]: %timeit ogrid_based(x,idx) 
1000 loops, best of 3: 391 µs per loop 

In [70]: %timeit reshape_based(x,idx) 
1000 loops, best of 3: 230 µs per loop