如何在字典中使用 numba njit 并将数组列表作为值

How to use numba njit on dictionaried with list of arrays as values

我正在尝试学习如何将 python 词典与 numba njit 结合使用。在这方面,我们必须将 pythonic 字典项转换为 numba 兼容项。在处理值是数组列表的字典时,我陷入了这一步。以下示例可以重现该问题:

sample_dict = {'size_0': [np.array([0.021, 0.022]), np.array([0.03, 0.008])],
               'size_1': [np.array([0.031]), np.array([0.036, 0.003])],
               'size_2': [np.array([], dtype=np.float64), np.array([0.043])]}
num_ = 2



String = nb.types.unicode_type
ValueArray = nb.float64[::1]

ValueList = nb.types.ListType(ValueArray)
ValueDict = nb.types.DictType(String, ValueList)

# DictValue = nb.types.Tuple([ValueArray])
# ValueDict = nb.types.DictType(String, DictValue)

# DictValue = nb.types.Tuple([nb.types.List(ValueArray, reflected=True)])
# ValueDict = nb.types.DictType(String, DictValue)

sample_nb = nb.typed.typeddict.Dict.empty(String, ValueArray)
for key, value in sample_dict.items():
    sample_nb[key] = value.copy()
    # sample_nb[key] = nb.types.List(value)


@nb.njit(ValueDict, nb.int_)
def nbTest(sample_dict, num_):
    ii = 1
    for i in range(num_):
        for j in sample_dict.values():
            ii += len(j[i])


nbTest(sample_nb, num_)

自从我尝试了各种方法以来,我遇到了不同的错误,但最常见的错误是:

No implementation of function Function() found for signature:

setitem(DictType[unicode_type,array(float64, 1d, C)]<iv=None>, unicode_type, reflected list(array(float64, 1d, C))<iv=None>)

如何准备这样的字典供 numba njit 使用?如果其中一个字典值只有一个数组(不是本例中所有字典值都包含 2 个数组),是否有任何限制?

问题来自两点:类型在字典中无效,因为它是 ValueArray 而它应该是 ValueList 以及无法使用 a 将反射列表复制到字典中基本任务。对于第二点 nb.typed.typedlist.List 可以用来解决这个问题。请注意 nb.types.List 是一种类型,而 nb.typed.typedlist.List 用于创建实例。这是结果代码:

String = nb.types.unicode_type
ValueArray = nb.float64[::1]

ValueList = nb.types.ListType(ValueArray)
ValueDict = nb.types.DictType(String, ValueList)

# DictValue = nb.types.Tuple([ValueArray])
# ValueDict = nb.types.DictType(String, DictValue)

# DictValue = nb.types.Tuple([nb.types.List(ValueArray, reflected=True)])
# ValueDict = nb.types.DictType(String, DictValue)

sample_nb = nb.typed.typeddict.Dict.empty(String, ValueList)
for key, value in sample_dict.items():
    sample_nb[key] = nb.typed.typedlist.List(value.copy())
    # sample_nb[key] = nb.types.List(value)