Numba 对没有输入的函数进行矢量化

Question

我想使用 numba.vectorize 并行化一个函数，但我的函数不接受任何输入。目前，我为从未使用过的函数使用虚拟数组和虚拟输入。

有没有更elegant/fast的方式（可能不用numba.vectorize）？

代码示例（不是我的实际代码，仅用于演示我如何丢弃输入）：

import numpy as np
from numba import vectorize

@vectorize(["int32(int32)"], nopython=True)
def particle_path(discard_me):
    x = 0
    for _ in range(10):
        x += np.random.uniform(0, 1)
    return np.int32(x)

arr = particle_path(np.empty(1024, dtype=np.int32))
print(arr)

Answer 1

如果您只是处理一维数组，那么您可以使用以下内容，其中数组必须在函数外实例化。似乎没有任何理由在这里使用 vectorize，您可以简单地使用 jit 实现目标，尽管您必须使用它显式地在数组元素上编写循环。如果你的数组总是一维的，那么你可以使用：

import numpy as np
from numba import jit

@jit(nopython=True)
def particle_path(out):
    for i in range(len(out)):
        x = 0
        for _ in range(10):
            x += np.random.uniform(0, 1)
        out[i] = x

arr = np.empty(1024, dtype=np.int32)
particle_path(arr)

您可以使用 flat 属性类似地处理任意维数组（并确保使用 .size 获取数组中元素的总数）：

@jit(nopython=True)
def particle_path(out):
    for i in range(out.size):
        x = 0
        for _ in range(10):
            x += np.random.uniform(0, 1)
        out.flat[i] = x

arr = np.empty(1024, dtype=np.int32)
particle_path(arr)

最后，如果您每次运行函数都需要一个新数组，则可以在函数内部创建数组（如果您将重复调用该函数并想覆盖同一个数组，因此节省了一遍又一遍地重新分配同一个数组的时间。

@jit(nopython=True)
def particle_path(num):
    out = np.empty(shape=num, dtype=np.int32)
    for i in range(num):
        x = 0
        for _ in range(10):
            x += np.random.uniform(0, 1)
        out[i] = x
    return out

arr = particle_path(1024)

Numba 对没有输入的函数进行矢量化

Numba vectorize for function with no input

python

numpy

vectorization

numba