numpy: Matrix anhand von index-Vektor filtern

[Tyrax]

Hallo Gemeinde,

ich habe ein vermutlich recht einfaches Problem, doch ich komme einfach nicht auf eine zufriedenstellende Lösung.

Sei also eine (3xn)-Matrix M gegeben und dazu ein Vektor A der Länge n mit Einträgen 0,1 oder 2. Jetzt will ich A auf M werfen, so dass ein Vektor B rauskommt, in dem zu jeder Zeile aus M das Element steht dessen Index in A angegeben ist.

Beispiel:

Code: Alles auswählen

M = array([[ 0,  1,  2],
         [ 3,  4,  5],
         [ 6,  7,  8],
         [ 9, 10, 11]])
A = array([0, 2, 2,1])
B = array([0, 5, 8, 10])

Das sollte doch eigentlich ein Einzeiler sein. Meine Versuche gingen allerdings bisher ins Leere...

Danke und beste Grüße, Tyrax

[Tyrax]

So, eine Möglichkeit habe ich gefunden:

Code: Alles auswählen

np.choose(A, M.T)

Hat noch jemand, 'ne bessere Idee?

[schaeffkoch]

vielleicht marginal schneller

Code: Alles auswählen

b=zeros(5)
for n in enumerate(a):
   b[n[0]]=m[n[0],n[1]]

[cofi]

Das sieht ja schlimm aus. Mit Sequence Unpacking:

Code: Alles auswählen

b = zeros(5)
for i, n in enumerate(a):
    b[i] = m[i, n]

[snafu]

Nur aus Interesse: Hat das vom Threadersteller entdeckte `numpy.choose()` irgendwelche markanten Schwächen? Der ausformulierte Python-Code ist zwar nicht übelst schwer zu lesen oder so, aber man kann `choose()` doch ruhig nutzen, oder etwa nicht?

[schaeffkoch]

choose hat natürlich keine schwächen, ist eben nur anscheinend ein bisschen langsamer.

[cofi]

Da es hier um (potentiellen) Maschinencode gegen Python-Code geht und ich nicht viel Ahnung von Numpy/Scipy habe: Hast du das auch geprofiled?

[schaeffkoch]

hat mich auch gewundert, aber vielleicht liegt es am transponieren der matrix.

Code: Alles auswählen

In [57]: timeit b=choose(a,m.T)
10000 loops, best of 3: 20.5 us per loop
In [59]: timeit b=choose_enum(a, m)
10000 loops, best of 3: 19.5 us per loop
....
In [61]: m2=m.T
In [64]: timeit b=choose(a, m2)
10000 loops, best of 3: 19.8 us per loop

hm nö..vielleicht einer eine erklärung?

[snafu]

Man kann schon mit einem direkten Aufruf von `A.choose(M.T)` ein bißchen was rausholen, wenn man denn Wert darauf legt.