For loop più veloci con array 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]

Lavoro con cicli for molto lunghi; %%timeitsopra con la passsostituzione dei rendimenti dell'operazione

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

questo è inaccettabile nel contesto (C ++ ha impiegato 6,5 secondi). Non c'è motivo per cui sopra sia stato fatto con oggetti Python; gli array hanno tipi ben definiti. L'implementazione di questo in C / C ++ come estensione è eccessivo sia per gli sviluppatori che per gli utenti; Sto solo passando array su cui eseguire il loop e eseguire operazioni aritmetiche.

C'è un modo per dire a Numpy "sposta questa logica in C", o un'altra libreria in grado di gestire loop annidati che coinvolgono solo array? Lo cerco per il caso generale, non per soluzioni alternative per questo esempio specifico (ma se ne hai uno posso aprire una sessione di domande e risposte separata).

Answer 1

Questa è fondamentalmente l'idea alla base di Numba . Non veloce come il C, ma può avvicinarsi ... Utilizza un compilatore jit per compilare il codice Python sulla macchina ed è compatibile con la maggior parte delle funzioni Numpy. (Nella documentazione trovi tutti i dettagli)

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

Pitone puro:

%%timeit

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

338 ms ± 12,6 ms per loop (media ± dev. Std. Di 7 analisi, 1 loop ciascuna)

Con Numba:

%%timeit

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

12 ms ± 534 µs per loop (media ± dev. Std. Di 7 analisi, 100 loop ciascuna)