Более быстрые циклы for с массивами в 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]

Я работаю с очень длинными циклами for; %%timeitна выше с passзаменой операции дает

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

это неприемлемо в контексте (C ++ занял 6,5 секунды). Нет никаких причин для выполнения вышеизложенного с объектами Python; массивы имеют четко определенные типы. Реализация этого в C / C ++ в качестве расширения является излишним как для разработчиков, так и для пользователей; Я просто передаю массивы в цикл и выполняю арифметические операции.

Есть ли способ сказать Numpy «переместить эту логику на C» или другую библиотеку, которая может обрабатывать вложенные циклы, включающие только массивы? Я ищу его для общего случая, а не для обхода этого конкретного примера (но если он у вас есть, я могу открыть отдельный раздел вопросов и ответов).

Answer 1

Это в основном идея Нумбы . Не так быстро, как C, но он может приблизиться ... Он использует jit-компилятор для компиляции кода Python на машину и совместим с большинством функций Numpy. (В документации вы найдете все подробности)

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:

%%timeit

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

338 мс ± 12,6 мс на цикл (среднее ± стандартное отклонение из 7 прогонов, по 1 циклу в каждом)

С Нумбой:

%%timeit

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

12 мс ± 534 мкс на цикл (среднее ± стандартное отклонение из 7 прогонов, по 100 циклов в каждом)