Des boucles for plus rapides avec des tableaux en 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]

Je travaille avec de très longues boucles for; %%timeitci-dessus avec le passremplacement des rendements d'opération

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

c'est inacceptable dans le contexte (C ++ a pris 6,5 secondes). Il n'y a aucune raison pour que ci-dessus soit fait avec des objets Python; les tableaux ont des types bien définis. L'implémentation de cela en C / C ++ en tant qu'extension est une exagération du côté des développeurs et des utilisateurs; Je passe juste des tableaux à boucler et à faire de l'arithmétique.

Existe-t-il un moyen de dire à Numpy "déplacer cette logique vers C", ou une autre bibliothèque qui peut gérer des boucles imbriquées impliquant uniquement des tableaux? Je le cherche pour le cas général, pas pour des solutions de contournement pour cet exemple spécifique (mais si vous en avez un, je peux ouvrir un Q&A séparé).

Answer 1

C'est essentiellement l'idée derrière Numba . Pas aussi rapide que C, mais il peut se rapprocher ... Il utilise un compilateur jit pour compiler du code python sur machine et il est compatible avec la plupart des fonctions Numpy. (Dans la documentation, vous trouvez tous les détails)

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

Python pur:

%%timeit

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

338 ms ± 12,6 ms par boucle (moyenne ± écart type de 7 courses, 1 boucle chacune)

Avec Numba:

%%timeit

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

12 ms ± 534 µs par boucle (moyenne ± écart standard de 7 analyses, 100 boucles chacune)