(Advanced) Indexing

Um kompakten Code zu schreiben, bei dem die Operationen möglichst innerhalb von numpy bleiben, ohne Kontrollstrukturen in Python zu schreiben, brauchen wir einen Weg, um numpy die Intention zu vermitteln. Das Werkzeug dazu sind clevere Indizierungen.

Zunächst lassen sich einzelne ELemente genau wie bei Listen indizieren:

a = np.array([1, 2, 3, 4, 5])
a[0]  # 1
a[-1]  # 5

Bereiche lassen sich mit slices indizieren:

a = np.array([1, 2, 3, 4, 5])
a[1:3]  # array([2, 3])

Dabei werden das erste einzuschließende und das erste auszuschließende Element angegeben: also einschließlich des zweiten Elements mit Index 1 aber ohne das vierte Element mit Index 3. Wird ein Wert ausgelassen, so wird das erste oder letzte Element angenommen. Werden beide Begrenzungen weggelassen, so erfolgt eine Kopie des gesamten arrays:

a = np.array([1, 2, 3, 4, 5])
a[:]   # array([1, 2, 3, 4, 5])
a[1:]  # array([2, 3, 4, 5])

In mehreren Dimensionen werden die Begrenzungen durch Kommata getrennt:

a = np.arange(9).reshape(3,3)
a[1:3, 1:3]  # array([[4, 5],
             #        [7, 8]])

Für ein einzelnes Element gilt daher analog:

a = np.arange(9).reshape(3,3)
a[1, 1]  # 4

Diese Technik lässt sich auch für Zuweisungen verwenden:

a = np.arange(9).reshape(3,3)
a[0, 0] = 42   # array([[42,  1,  2],
               #        [ 3,  4,  5],
               #        [ 6,  7,  8]])
a[1:, :] = 0   # array([[42,  1,  2],
               #        [ 0,  0,  0],
               #        [ 0,  0,  0]])

Was ist aber, wenn ich eine Liste von Elementen habe, die ich gerne extrahieren möchte? Anders als bei Listen akzeptiert numpy auch Listen als Indizes:

a = np.arange(9)**2
a[[1, 3, 5]]  # array([ 1,  9, 25])

Auch Listen (oder arrays) von Wahrheitswerten sind möglich:

a = np.arange(3)**2
a[[False, True, True]] # array([1, 4])

Das ist besonders dann hilfreich, wenn man eine Bedingung auf das array anwenden möchte und dabei broadcasting nutzen möchte:

a = np.arange(9)
a % 2 == 0     # array([ True, False,  True, False,  True, False,  True, False,  True])
a[a % 2 == 0]  # array([0, 2, 4, 6, 8])

In mehreren Dimensionen muss man etwas aufpassen, weil nicht alle Kombinationen aus den mehreren angegebenen Indizes möglich sind:

a = np.arange(9).reshape(3,3)
a[[0, 1], [1, 2]]  # array([1, 5])

Will man stattdessen tatsächlich alle Kombinationen aus den beiden Listen, so benötigt man np.ix_:

a = np.arange(9).reshape(3,3)
a[np.ix_([0, 1], [1, 2])]  # array([[1, 2],
                           #        [4, 5]])

Eine neue Dimension lässt sich mit np.newaxis hinzufügen:

a = np.arange(3)
a[:, np.newaxis]  # array([[0],
                  #        [1],
                  #        [2]])

Das ist insbesondere für den Ersatz von Schleifen hilfreich.

Indizierung als Schleifen-Ersatz

Angenommen, wir wollen die paarweisen Abstände zwischen Teilchen in 1D berechnen. In normalem Python wäre das eine verschachtelte Schleife:

def pairwise_distances_python(positions: list[float]):
    distances = []
    for i, p1 in enumerate(positions):
        row = []
        for j, p2 in enumerate(positions):
            row.append(abs(p1 - p2))
        distances.append(row)
    return distances

pairwise_distances_python([1, 2, 3]) # [[0, 1, 2], [1, 0, 1], [2, 1, 0]]

Zwar ließe sich das Prinzip auch auf numpy übertragen:

def pairwise_distances_numpy_DONTDOTHIS(positions: np.ndarray):
    distances = np.zeros(positions.shape * 2)
    for i in range(positions.shape[0]):
        for j in range(positions.shape[0]):
            distances[i, j] = abs(positions[i] - positions[j])
    return distances
pairwise_distances_numpy_DONTDOTHIS(np.array([1,2,3])) # array([[0., 1., 2.],
                                                       #        [1., 0., 1.],
                                                       #        [2., 1., 0.]])

Mit Indizierung und Broadcasting geht das aber viel übersichtlicher:

def pairwise_distances_numpy(positions: np.ndarray):
    return np.abs(positions[:, np.newaxis] - positions[np.newaxis, :])

Der Trick hier ist dass die erste Teilausdruck positions[:, np.newaxis] ein 2D-Array erzeugt, das die Positionen in den Zeilen speichert. Das zweite positions[np.newaxis, :] erzeugt ein 2D-Array, das die Positionen in den Spalten enthält. Die Subtraktion erfolgt nun mit Broadcasting, bei dem das erste Array in alle Zeilen einer Matrix kopiert wird, während das zweite Array in allen Zeilen kopiert wird. Das Ergebnis ist eine Matrix, in der die Differenzen aller Paare von Positionen stehen. Das entspricht:

array([[0, 0, 0],         array([[0, 1, 2],
       [1, 1, 1],    -           [0, 1, 2],
       [2, 2, 2]])               [0, 1, 2]])