Szybsze pętle for z tablicami w Pythonie

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]

Pracuję z bardzo długimi pętlami for; %%timeitpowyżej z passzastąpieniem wydajności operacyjnej

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

jest to niedopuszczalne w kontekście (C ++ zajęło 6,5 sekundy). Nie ma powodu, aby robić to z obiektami Pythona; tablice mają dobrze zdefiniowane typy. Implementacja tego w C / C ++ jako rozszerzeniu jest przesadą zarówno z punktu widzenia programisty, jak i użytkownika; Po prostu przekazuję tablice do zapętlenia i wykonywania działań arytmetycznych.

Czy istnieje sposób, aby powiedzieć Numpy'emu „przenieś tę logikę do C” lub inną bibliotekę, która może obsługiwać zagnieżdżone pętle obejmujące tylko tablice? Szukam go dla ogólnego przypadku, a nie obejścia dla tego konkretnego przykładu (ale jeśli masz taki, mogę otworzyć osobne pytania i odpowiedzi).

Answer 1

To jest w zasadzie idea Numba . Nie tak szybko jak C, ale może być blisko ... Używa kompilatora jit do kompilowania kodu Pythona na maszynie i jest kompatybilny z większością funkcji Numpy. (W dokumentach znajdziesz wszystkie szczegóły)

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

Czysty Python:

%%timeit

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

338 ms ± 12,6 ms na pętlę (średnia ± odchylenie standardowe 7 przebiegów, 1 pętla każda)

Z Numba:

%%timeit

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

12 ms ± 534 µs na pętlę (średnia ± odchylenie standardowe z 7 przebiegów, po 100 pętli każda)