For-loops mais rápidos com arrays em 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]

Eu trabalho com loops for muito longos; %%timeitacima com a passsubstituição dos rendimentos da operação

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

isso é inaceitável no contexto (C ++ levou 6,5 segundos). Não há razão para que acima seja feito com objetos Python; arrays têm tipos bem definidos. Implementar isso em C / C ++ como uma extensão é um exagero tanto para o desenvolvedor quanto para o usuário; Estou apenas passando matrizes para loop e fazer aritmética.

Existe uma maneira de dizer ao Numpy "mover esta lógica para C", ou outra biblioteca que pode lidar com loops aninhados envolvendo apenas matrizes? Eu procuro para o caso geral, não para soluções alternativas para este exemplo específico (mas se você tiver uma, posso abrir uma sessão de perguntas e respostas separada).

Answer 1

Essa é basicamente a ideia por trás do Numba . Não tão rápido quanto C, mas pode chegar perto ... Ele usa um compilador jit para compilar o código Python para a máquina e é compatível com a maioria das funções Numpy. (Nos documentos você encontra todos os detalhes)

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 loop (média ± desvio padrão de 7 execuções, 1 loop cada)

Com Numba:

%%timeit

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

12 ms ± 534 µs por loop (média ± desvio padrão de 7 execuções, 100 loops cada)