numpy Vektorisierung statt für Schleifen

Question

Ich schrieb etwas Code in Python, der gut funktioniert, aber sehr langsam ist; Ich denke aufgrund der for-Schleifen. Ich hoffe, man kann die folgenden Operationen mit numpy Befehlen beschleunigen. Lass mich das Ziel definieren.

Angenommen, ich habe ein 2D-Numy-Array all_CMs mit den Abmessungen row x col. Betrachten Sie zum Beispiel ein 6 x 11 Array (siehe Zeichnung unten).

1. Ich möchte, den Mittelwert für alle Zeilen berechnen, das heißt Summe ⱼ aᵢⱼ in einem Array führt. Dies kann natürlich leicht durchgeführt werden. (I nennen diesen Wert CM_tilde)

2. Nun, für jede Reihe I den Mittelwert von einigen ausgewählten Werte berechnet werden soll, und zwar alle Werte unterhalb einer bestimmten Schwelle durch ihre Summe Rechen- und es durch die Anzahl aller Spalten Dividieren (N). Wenn der Wert über diesem definierten Schwellenwert liegt, wird der Wert CM_tilde (Mittelwert der gesamten Zeile) hinzugefügt. Dieser Wert wird CM genannt

3. Danach wird der CM Wert wird von jedem Element subtrahiert in der Reihe

Zusätzlich dazu ich eine numpy Array oder eine Liste haben wollen, wenn diese alle CM Werte aufgeführt sind .

Die Figur:

Der folgende Code funktioniert, aber sehr langsam (vor allem, wenn die Arrays immer groß)

CM_tilde = np.mean(data, axis=1) 
N = data.shape[1] 
data_cm = np.zeros((data.shape[0], data.shape[1], data.shape[2])) 
all_CMs = np.zeros((data.shape[0], data.shape[2])) 
for frame in range(data.shape[2]): 
    for row in range(data.shape[0]): 
     CM=0 
     for col in range(data.shape[1]): 
      if data[row, col, frame] < (CM_tilde[row, frame]+threshold): 
       CM += data[row, col, frame] 
      else: 
       CM += CM_tilde[row, frame] 
     CM = CM/N 
     all_CMs[row, frame] = CM 
     # calculate CM corrected value 
     for col in range(data.shape[1]): 
      data_cm[row, col, frame] = data[row, col, frame] - CM 
    print "frame: ", frame 
return data_cm, all_CMs 

Irgendwelche Ideen?

Answer 1

Es ist ganz einfach zu vektorisieren, was Sie tun:

import numpy as np 

#generate dummy data 
nrows=6 
ncols=11 
nframes=3 
threshold=0.3 
data=np.random.rand(nrows,ncols,nframes) 

CM_tilde = np.mean(data, axis=1) 
N = data.shape[1] 

all_CMs2 = np.mean(np.where(data < (CM_tilde[:,None,:]+threshold),data,CM_tilde[:,None,:]),axis=1) 
data_cm2 = data - all_CMs2[:,None,:] 

Vergleicht man dies mit Originalen:

In [684]: (data_cm==data_cm2).all() 
Out[684]: True 

In [685]: (all_CMs==all_CMs2).all() 
Out[685]: True 

Die Logik ist, dass wir mit Arrays der Größe arbeiten [nrows,ncols,nframes] gleichzeitig. Der Haupttrick besteht darin, Pythons Broadcasting zu nutzen, indem man CM_tilde der Größe [nrows,nframes] in CM_tilde[:,None,:] der Größe [nrows,1,nframes] dreht. Python verwendet dann die gleichen Werte für jede Spalte, da dies eine Singleton-Dimension dieser modifizierten CM_tilde ist.

von np.where verwenden wir wählen (basierend auf den threshold), ob wir den entsprechenden Wert von data oder wieder erhalten möchten, den Broadcast-Wert von CM_tilde. Eine neue Verwendung von np.mean ermöglicht es uns all_CMs2 zu berechnen.

Im letzten Schritt nutzten wir das Broadcasting, indem wir dieses neue all_CMs2 direkt von den entsprechenden Elementen von data abzogen.

Es kann hilfreich sein, Code auf diese Weise zu vektorisieren, indem Sie die impliziten Indizes Ihrer temporären Variablen betrachten. Was ich meine ist, dass Ihre temporäre Variable CM innerhalb einer Schleife über [nrows,nframes] lebt, und ihr Wert wird mit jeder Iteration zurückgesetzt. Dies bedeutet, dass CM in Wirklichkeit eine Menge CM[row,frame] ist (später explizit dem 2d-Array zugewiesen all_CMs), und von hier aus ist es einfach zu sehen, dass Sie es konstruieren können, indem Sie eine entsprechende CMtmp[row,col,frames] Menge entlang seiner Spaltendimension summieren. Wenn es hilft, können Sie den np.where(...) Teil als CMtmp für diesen Zweck benennen und dann np.mean(CMtmp,axis=1) daraus berechnen. Das gleiche Ergebnis natürlich, aber wahrscheinlich transparenter.

Answer 2

Hier ist meine Vektorisierung Ihrer Funktion. Ich habe von innen nach außen gearbeitet und frühere Versionen auskommentiert, als ich mitging. Also die erste Schleife, die ich vektorisiert habe, hat ### Kommentarzeichen.

Es ist nicht so sauber und gut durchdacht wie @Andras's Antwort, aber hoffentlich ist es lehrreich, geben eine Idee, wie Sie dieses Problem inkrementell beheben können.

def foo2(data, threshold): 
    CM_tilde = np.mean(data, axis=1) 
    N = data.shape[1] 
    #data_cm = np.zeros((data.shape[0], data.shape[1], data.shape[2])) 
    ##all_CMs = np.zeros((data.shape[0], data.shape[2])) 
    bmask = data < (CM_tilde[:,None,:] + threshold) 
    CM = np.zeros_like(data) 
    CM[:] = CM_tilde[:,None,:] 
    CM[bmask] = data[bmask] 
    CM = CM.sum(axis=1) 
    CM = CM/N 
    all_CMs = CM.copy() 
    """ 
    for frame in range(data.shape[2]): 
     for row in range(data.shape[0]): 
      ###print(frame, row) 
      ###mask = data[row, :, frame] < (CM_tilde[row, frame]+threshold) 
      ###print(mask) 
      ##mask = bmask[row,:,frame] 
      ##CM = data[row, mask, frame].sum() 
      ##CM += (CM_tilde[row, frame]*(~mask)).sum() 

      ##CM = CM/N 
      ##all_CMs[row, frame] = CM 
      ## calculate CM corrected value 
      #for col in range(data.shape[1]): 
      # data_cm[row, col, frame] = data[row, col, frame] - CM[row,frame] 
     print "frame: ", frame 
    """ 
    data_cm = data - CM[:,None,:] 
    return data_cm, all_CMs 

Die Ausgangsübereinstimmungen für diesen kleinen Testfall, der mir mehr als alles geholfen hat, die Dimensionen richtig zu machen.

threshold = .1 
data = np.arange(4*3*2,dtype=float).reshape(4,3,2)