Numba-使用多维数组对循环进行并行处理

Question

我正在尝试用numba并行化for循环。我是这个库的新手，但经过一些研究，我制作了这个代码，与我研究的示例相比，它看起来是正确的：

@njit(nopython=True, parallel=True)
def tempFtemp(ftemp_pte, func_F, numPointsEval, pointsToEval):
for i in prange(0, numPointsEval):
ftemp_pte[:,i] = np.hstack(func_F(np.vstack(pointsToEval[:,i])))
return ftemp_pte
ftemp_pte= tempFtemp(ftemp_pte, func_F, numPointsEval, pointsToEval)

当我在我的程序思想中编译它时，我得到了错误"；非精确类型pyobject"；在线路CCD_ 1处。ftemp_pte和pointsToEval都是2d数组，numPointsEval是一个整数，func_F是一个随机函数，它将产生要存储在ftemp_pte[:,i]中的1d数组。

如果您能帮助我们找出是什么导致了这个错误，我们将不胜感激。

[编辑]

我最初拥有的顺序代码(有效(如下：

def func_F(x):
f= np.zeros((1,2))
f[0,0]= x[0,0]
n= max(np.size(x,0), np.size(x,1))    
g    = 1 + 9* np.sum(x[1:n])/(n-1)
h    = 1 - np.sqrt(f[0,0]/g)
f[0,1] = g * h
F= np.transpose(f)
return F
for i in range(0, numPointsEval):
ftemp_pte[:,i] = np.hstack(func_F(np.vstack(pointsToEval[:,i])))

我还应该提到，需要使用hstack和vstack，这样创建的数组的格式才能与ftemp_pte数组相匹配。删除这些说明会导致尺寸不匹配。

变量for i in prange(0, numPointsEval):0总是有2行x列。正确值的一个例子是

[[0.21875   0.21875   0.21875   0.21875   0.21875   0.21875   0.21875
0.21875   0.21875   0.21875   0.21875   0.21875   0.21875   0.21875
0.21875   0.21875   0.21875   0.21875   0.21875   0.21875   0.21875
0.21875   0.21875   0.21875   0.21875   0.21875   0.21875   0.21875
0.21875  ]
[0.5397286 0.5397286 0.5397286 0.5397286 0.5397286 0.5397286 0.5397286
0.5397286 0.5397286 0.5397286 0.5397286 0.5397286 0.5397286 0.5397286
0.5397286 0.5397286 0.5397286 0.5397286 0.5397286 0.5397286 0.5397286
0.5397286 0.5397286 0.5397286 0.5397286 0.5397286 0.5397286 0.5397286
0.5397286]]

我的代码最初的目的是将以下Matlab的parfor指令翻译成Python

parfor i=1:numPointsEval
ftemp_pte(:,i) = feval(func_F,pointsToEval(:,i));

如有任何帮助，将不胜感激

Answer 1

问题2的答案。

只要您不是真正堆叠阵列，而是重塑阵列，就应该避免使用hstack()和vstack()，并考虑使用reshape()或ravel()。例如：

ftemp_pte[:, i] = func_F(pointsToEval[:, i].reshape(1, -1)).ravel()

但是，numba不支持非连续阵列上的reshape()。

因此，我已经设法使您的代码与numba一起运行，通过转换所有内容来避免重新塑造数组。以下代码确实有效，可能会给你一些想法：

@nb.njit
def func_F(x):
f = np.zeros(2)    # Simple 1d array
f[0] = x[0]
n = max(x.shape)
g = 1 + 9 * np.sum(x[1:n]) / (n - 1)
h = 1 - np.sqrt(f[0] / g)
f[1] = g * h
return f
@nb.njit(parallel=True)
def tempFtemp(ftemp_pte, func_F, numPointsEval, pointsToEval):
for i in nb.prange(numPointsEval):
ftemp_pte[i] = func_F(pointsToEval[i])
return ftemp_pte
ftemp_pte = np.zeros((2, 5)).T
pointsToEval = np.zeros((2, 5)).T
numPointsEval = 5
ftemp_pte = tempFtemp(ftemp_pte, func_F, numPointsEval, pointsToEval)
print(ftemp_pte.T)

Answer 2

文档说一级函数对象可以是Numba cfunc编译的函数、JIT编译的函数和实现Wrapper地址协议的对象。

您可以传递一个JITted函数，如以下简化示例所示：

@nb.njit
def cos(a):
return np.cos(a)
@nb.njit(parallel=True)
def tempFtemp(ftemp_pte, func_F, numPointsEval, pointsToEval):
for i in nb.prange(numPointsEval):
ftemp_pte[:, i] = func_F(pointsToEval[:, i])
return ftemp_pte
ftemp_pte = tempFtemp(ftemp_pte, np.cos, numPointsEval, pointsToEval)  # Error
ftemp_pte = tempFtemp(ftemp_pte, cos, numPointsEval, pointsToEval)     # Works

这解决了"；非精确类型pyobject"；问题，但我已经从示例中删除了hstack和vstack操作，因为它们会产生自己的问题和效率低下。