Schnellere for-Schleifen mit Arrays in Python

Question

N, M = 1000, 4000000
a = np.random.uniform(0, 1, (N, M))
k = np.random.randint(0, N, (N, M))

out = np.zeros((N, M))
for i in range(N):
    for j in range(M):
        out[k[i, j], j] += a[i, j]

Ich arbeite mit sehr langen for-Schleifen; %%timeitweiter oben mit passErsetzen der Betriebsausbeuten

1min 19s ± 663 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Dies ist im Kontext nicht akzeptabel (C ++ dauerte 6,5 Sekunden). Es gibt keinen Grund, dies oben mit Python-Objekten zu tun. Arrays haben genau definierte Typen. Die Implementierung in C / C ++ als Erweiterung ist sowohl für Entwickler als auch für Benutzer ein Overkill. Ich übergebe nur Arrays, um sie zu schleifen und zu rechnen.

Gibt es eine Möglichkeit, Numpy mitzuteilen, dass diese Logik nach C verschoben werden soll, oder eine andere Bibliothek, die verschachtelte Schleifen verarbeiten kann, an denen nur Arrays beteiligt sind? Ich suche es für den allgemeinen Fall, nicht für dieses spezielle Beispiel (aber wenn Sie eines haben, kann ich ein separates Q & A eröffnen).

Answer 1

Dies ist im Grunde die Idee hinter Numba . Nicht so schnell wie C, aber es kann nahe kommen ... Es verwendet einen JIT-Compiler, um Python-Code für die Maschine zu kompilieren, und es ist mit den meisten Numpy-Funktionen kompatibel. (In den Dokumenten finden Sie alle Details)

import numpy as np
from numba import njit


@njit
def f(N, M):
    a = np.random.uniform(0, 1, (N, M))
    k = np.random.randint(0, N, (N, M))

    out = np.zeros((N, M))
    for i in range(N):
        for j in range(M):
            out[k[i, j], j] += a[i, j]
    return out


def f_python(N, M):
    a = np.random.uniform(0, 1, (N, M))
    k = np.random.randint(0, N, (N, M))

    out = np.zeros((N, M))
    for i in range(N):
        for j in range(M):
            out[k[i, j], j] += a[i, j]
    return out

Pure Python:

%%timeit

N, M = 100, 4000
f_python(M, N)

338 ms ± 12,6 ms pro Schleife (Mittelwert ± Standardabweichung von 7 Läufen, jeweils 1 Schleife)

Mit Numba:

%%timeit

N, M = 100, 4000
f(M, N)

12 ms ± 534 µs pro Schleife (Mittelwert ± Standardabweichung von 7 Läufen, jeweils 100 Schleifen)