如何在cython（或numba）中迭代列表的列表？

Question

我想要一个函数，该函数接收一个列表列表，每个列表具有不同的大小，并且可以迭代每个子列表（包含整数），将它们作为广播传递给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中变得非常低效。所以我想知道如何对这个函数进行cythonize（或numbaize）处理。像这样的东西：

%%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

Answer 1

在Numba中，您可以使用Typed List在列表列表上进行迭代。 Numba不支持使用列表为NumPy数组建立索引，因此该函数还需要进行一些修改，以通过遍历内部列表的元素并索引到x来实现均值。

在调用jitted函数之前，您还需要将列表列表转换为类型列表的类型列表。