Bucles for más rápidos con matrices 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]

Trabajo con bucles for muy largos; %%timeiten arriba con la passsustitución de los rendimientos de la operación

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

esto es inaceptable en contexto (C ++ tomó 6.5 segundos). No hay ninguna razón para hacer lo anterior con objetos Python; las matrices tienen tipos bien definidos. Implementar esto en C / C ++ como una extensión es una exageración tanto para el desarrollador como para el usuario; Solo estoy pasando matrices para hacer un ciclo y hacer aritmética.

¿Hay alguna manera de decirle a Numpy "mueva esta lógica a C", u otra biblioteca que pueda manejar bucles anidados que involucren solo matrices? Lo busco para el caso general, no soluciones para este ejemplo específico (pero si tiene uno, puedo abrir una sesión de preguntas y respuestas por separado).

Answer 1

Esta es básicamente la idea detrás de Numba . No tan rápido como C, pero puede acercarse ... Utiliza un compilador jit para compilar código Python en la máquina y es compatible con la mayoría de las funciones de Numpy. (En los documentos encuentras todos los detalles)

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 puro:

%%timeit

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

338 ms ± 12,6 ms por bucle (media ± desviación estándar de 7 corridas, 1 bucle cada una)

Con Numba:

%%timeit

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

12 ms ± 534 µs por bucle (media ± desviación estándar de 7 corridas, 100 bucles cada una)