Numba- 使用多维数组并行化 For 循环

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)

当我在我的程序思想中编译它时，我在for i in prange(0, numPointsEval): 行收到错误“non-precise type pyobject”。 ftemp_pte和pointsToEval都是二维数组，numPointsEval是一个整数，func_F是一个随机函数，它将产生一维数组存储在ftemp_pte[:,i]中。

任何有关如何找出导致此错误的原因的帮助将不胜感激。

[编辑]

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

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_ptearray。删除这些说明会导致尺寸不匹配。

变量ftemp_pte始终有 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

docs 说一等函数对象可以是 Numba cfunc 编译函数、JIT 编译函数和实现包装地址协议的对象。

您可以传递 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 操作，因为它们产生their own problems and inefficiencies。