Как перебрать список списков в Cython (или Numba)? - PullRequest
Новая
       42

Как перебрать список списков в Cython (или Numba)?

0 голосов
/ 21 апреля 2020

Мне нужна функция, которая получает в качестве аргумента список списков, каждый подсписок с разным размером, и может выполнять итерацию по каждому из подсписков (которые содержат целые числа), чтобы передавать их в виде широковещательной передачи в массив numpy и выполнять различные операции (например, усреднение).

Позвольте мне привести простой пример ожидаемого поведения без использования cython: 

import numpy as np

mask = [[0, 1, 2, 4, 6, 7, 8, 9],
        [0, 1, 2, 4, 6, 7, 8, 9],
        [0, 1, 2, 4, 6, 9],
        [3, 5, 8],
        [0, 1, 2, 4, 6, 7, 8, 9],
        [3, 5, 7],
        [0, 1, 2, 4, 6, 9],
        [0, 1, 4, 5, 7, 8, 9],
        [0, 1, 3, 4, 7, 8, 9],
        [0, 1, 2, 4, 6, 7, 8, 9]] # This is the list of lists

x = np.array([2.0660689 , 2.08599832, 0.45032649, 1.05435649, 2.06010132,
              1.07633407, 0.43014785, 1.54286467, 1.644388  , 2.15417444])

def nocython(mask, x):
    out = np.empty(len(x), dtype=np.float64)
    for i, v in enumerate(mask):
        out[i] = x[v].mean()
    return out

>>> nocython(mask, x)
array([1.55425875, 1.55425875, 1.54113622, 1.25835952, 1.55425875,
       1.22451841, 1.54113622, 1.80427567, 1.80113602, 1.55425875])

Основная проблема заключается в том, что мне приходится много обрабатывать * numpy - массивы и списки масок, и циклы становятся чрезвычайно неэффективными в Python. Поэтому я хотел знать, как я могу цитонизировать (или нумизировать) эту функцию. Примерно так: 

%%cython

import numpy as np
cimport numpy as np

cdef np.ndarray[np.float64_t] cythonloop(int[:,:] mask, np.ndarray[np.float64_t] x):
    cdef Py_ssize_t i
    cdef Py_ssize_t N = len(x)
    cdef np.ndarray[np.float64_t] out = np.empty(N, dtype=np.float64)
    for i in range(N):
        out[i] = x[mask[i]]

cythonloop(mask, x)

Но это не сработает (Невозможно привести список к типу int [:,:] ').

Ни один, если я попробую это в numba 

import numba as nb

@nb.njit
def nocython(mask, x):
    out = np.empty(len(x), dtype=np.float64)
    for i, v in enumerate(mask):
        out[i] = x[v].mean()
    return out

Что дает следующую ошибку: 

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<built-in function getitem>) with argument(s) of type(s): (array(float64, 1d, A), reflected list(int64))
 * parameterized

1 Ответ

1 голос
/ 21 апреля 2020

В Numba вы можете использовать Типизированный список для итерации по списку списков. Numba не поддерживает индексирование массива NumPy списком, поэтому функция также нуждается в некоторой модификации для реализации среднего значения путем перебора элементов внутреннего списка и индексации в x.

. Вы также необходимо преобразовать список списков в типизированный список типизированных списков перед вызовом объединенной функции.

Объединение этого дает (в дополнение к коду из вашего вопроса): 

from numba import njit
from numba.typed import List

@njit
def jitted(mask, x): 
    out = np.empty(len(x), dtype=np.float64)
    for i in range(len(mask)):
        m_i = mask[i]
        s = 0 
        for j in range(len(m_i)):
            s += x[m_i[j]]
        out[i] = s / len(m_i)
    return out 

typed_mask = List()
for m in mask:
    typed_mask.append(List(m))

# Sanity check - Numba and nocython implementations produce the same result
np.testing.assert_allclose(nocython(mask, x),  jitted(typed_mask, x))

Обратите внимание, что также можно избежать превращения списка в типизированный список, поскольку Numba будет использовать Отраженный список при передаче встроенного типа списка - однако эта функция устарела и будет удалена из будущей версии Numba, поэтому вместо него рекомендуется использовать Typed List.

...