在 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)

运行ning 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 时，你在遍历数组时正在移动非连续内存 space。

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

Answer 2

没有numba，访问字段值并不比访问二维数组的列慢：

In [1]: arr2 = np.zeros((10000), dtype='i,i')
In [2]: arr2.dtype
Out[2]: dtype([('f0', '<i4'), ('f1', '<i4')])

修改字段：

In [4]: %%timeit x = arr2.copy()
   ...: x['f0'] += 1
   ...: 
16.2 µs ± 13.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

如果我将字段分配给新变量，则时间相同：

In [5]: %%timeit x = arr2.copy()['f0']
   ...: x += 1
   ...: 
15.2 µs ± 14.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

如果我构造一个相同大小的一维数组会更快：

In [6]: %%timeit x = np.zeros(arr2.shape, int)
   ...: x += 1
   ...: 
8.01 µs ± 15.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

但访问二维数组列的时间相似：

In [7]: %%timeit x = np.zeros((arr2.shape[0],2), int)
   ...: x[:,0] += 1
   ...: 
17.3 µs ± 23.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

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

Setting structured array field in Numba

python

numpy

numba