Wählen Sie Werte aus einer Reihe von Arrays nach einem Array von Permutationen

Question

Ich habe 3 numpy Arrays der Form 2xN (mit N groß, ein paar Millionen), nennen sie a1, a2, a3. Dann habe ich ein weiteres Array von Form Nx3, deren Zeilenwerte sich auf eines der Arrays a1, a2, a3 beziehen, nennen es Permutationen. Diese Permutationen Array wie folgt aussieht: [[0, 1, 2], [1,2,0], [1,0,2], ... bis zu N Zeilen]

Ich möchte Erstellen Sie weitere 3 numply Arrays b1, b2, b3 der Form 2xN, die den Inhalt der ursprünglichen a1, a2, a3 haben, aber ihre Spalten wurden entsprechend den Zeilen des Permutationsarrays permutiert.

Ich habe versucht, Indexierung stacking die 3 Arrays und numpy.choose, aber ich kann es nicht zum arbeiten. Ich suche nach einer Lösung ohne Python-Schleifen. Jede Hilfe würde sehr geschätzt werden!

EDIT

Nur um zu klären ich die Python-Schleife Umsetzung zeigen, was ich versuche zu tun: b1, b2, b3 von bb

aa = np.dstack((a1, a2, a3)) 
bb = np.empty_like(aa) 
for i, o in enumerate(permutations): 
    bb[:,i, np.arange(3)] = aa[:, i, o] 

Dann würde ich abrufen.

Answer 1

Mit fancy-indexing, könnten Sie tun -

bb = aa[:,np.arange(N),permutations.T] 

Bitte beachten Sie, dass diese (2,3,N) der Form sein würde. Also, um zu wählen b1, b2, b3, würden Sie tun:

b1,b2,b3 = bb[:,0,:], bb[:,1,:], bb[:,2,:] 

Oder wenn Sie bb bestehen aus der gleichen Form zu sein, wie mit dem Code geschrieben, könnte man hinzufügen:

bb = bb.swapaxes(1,2) 

---

Hier ist ein weiterer Ansatz lineare Indizierung, Schneiden und natürlich NumPy broadcasting -

idx = permutations + 3*np.arange(N)[:,None]  
bb = aa.reshape(2,-1)[:,idx].reshape(2,N,3) 

Dies würde eine bb der gleichen Form wie mit dem Posted-Loopy-Code erstellen.

---

Runtime Test

In [189]: def original_app(aa,permutations): 
    ...:  bb = np.empty_like(aa) 
    ...:  for i, o in enumerate(permutations): 
    ...:   bb[:,i, np.arange(3)] = aa[:, i, o] 
    ...:  return bb 
    ...: 
    ...: 
    ...: def linear_index_app(aa,permutations): 
    ...:  idx = permutations + 3*np.arange(N)[:,None]  
    ...:  return aa.reshape(2,-1)[:,idx].reshape(2,N,3) 
    ...: 

In [190]: # Setup input arrays 
    ...: N = 10000 
    ...: a1 = np.random.rand(2,N) 
    ...: a2 = np.random.rand(2,N) 
    ...: a3 = np.random.rand(2,N) 
    ...: 
    ...: permutations = np.random.randint(0,3,(N,3)) 
    ...: aa = np.dstack((a1, a2, a3)) 


In [191]: %timeit original_app(aa,permutations) 
10 loops, best of 3: 128 ms per loop 

In [192]: %timeit aa[:,np.arange(N),permutations.T] 
1000 loops, best of 3: 972 µs per loop 

In [193]: %timeit linear_index_app(aa,permutations) 
1000 loops, best of 3: 1.02 ms per loop 

So scheint wie fancy-indexing die beste der Partie!