python - 在 Numba 中设置结构化数组字段

Question

我想在 Numba 编译的 nopython 函数中设置 NumPy 结构化标量的整个字段。下面代码中的 desired_fn 是我想做的一个简单示例，而 working_fn 是我目前如何完成此任务的示例。

import numpy as np
import numba as nb


test_numpy_dtype = np.dtype([("blah", np.int64)])
test_numba_dtype = nb.from_dtype(test_numpy_dtype)


@nb.njit
def working_fn(thing):
    for j in range(len(thing)):
        thing[j]['blah'] += j

@nb.njit
def desired_fn(thing):
    thing['blah'] += np.arange(len(thing))


a = np.zeros(3,test_numpy_dtype)
print(a)
working_fn(a)
print(a)
desired_fn(a)

运行 desired_fn(a) 产生的错误是:

numba.errors.InternalError: unsupported array index type const('blah') in [const('blah')]
[1] During: typing of staticsetitem at /home/sam/PycharmProjects/ChessAI/playground.py (938)

这对于极其性能关键代码来说是必需的，并且将运行数十亿次，因此消除对这些类型循环的需求似乎至关重要。

Answer 1

以下作品(numba 0.37):

@nb.njit
def desired_fn(thing):
    thing.blah[:] += np.arange(len(thing))
    # or
    # thing['blah'][:] += np.arange(len(thing))

如果您主要对数据的列而不是行进行操作，则可以考虑使用不同的数据容器。 numpy 结构化数组的布局类似于结构向量而不是数组结构。这意味着当您想要更新 blah 时，您将在遍历数组时穿过非连续的内存空间。

此外，对于任何代码优化，都值得使用 timeit 或其他一些计时工具(消除 jit 代码所需的时间)来查看实际性能。您可能会发现 numba 的显式循环虽然更冗长，但实际上可能比矢量化代码更快。