3D 阵列上的 Numba 和 NumPy nanmean - 简单的替代方案？

Question

我花了几个小时将庞大的代码库转换为与 Numba 兼容，显然还有 1 个问题需要修复（好吧，在弹出其他 Numba 错误之前，尽管这是在函数末尾附近） .基本上我最后留下了一个 3D NumPy 数组，其中 NaN 填充了实际上不存在的任何空间（这些是日期和相关价格，每月不同的工作日是一个例子，其中数组的一部分将有几个月的 NaN天数较少，例如二月与一月相比工作日较少）。

现在这个形状的简单解决方案：

3darray.shape = (6, 8192, 22)

例如，因为它包含 NaN，所以使用：

averages = numpy.nanmean(3darray, axis=2)

然后我得到了我想要的。每 22 天，减去 NaN，对整个数组进行平均，所以我得到一个 (6, 8192) 形状的二维矩阵。

但 Numba 不喜欢这样，因为看起来他们只在单个值上实现此功能（？或一维或二维数组，不知道：https://github.com/numba/numba/blob/master/numba/np/arraymath.py#L1092-L1109）所以就像我的其余代码一样，我必须编写一些疯狂的索引才能让 Numba 开心吗？尽管我在这里使用“nanmasks”进行数学运算时遇到了一些困难，尽管它只是在不存在 NaN 的情况下进行平均...

nanmask = ~np.isnan(3darray.reshape(-1)) #flatten the array to 1D, get all indexes that aren't NaN

太好了，现在我知道哪些值不是 NaN（在平面一维矩阵中）...但是接下来呢？ Numba 中必须有某种方法可以使我忽略的这个简单数学变得更容易。感谢任何帮助。

Answer 1

我为 3D 矩阵 (nanmean) 编写了这个“有趣”的 Numba 兼容函数...在 myrtlecat 指出 Numba 没有为整个矩阵以外的任何东西实现 nanmean 之后。在我的 PC 上的 (6,8192,22) 形状 3D 矩阵上以 0.5 毫秒左右的速度运行，并匹配 np.nanmean(3Darray, axis=2):[=11= 的 np.nanmean 函数结果]

import numpy as np
import numba as nb

@nb.jit(cache=True, parallel=True, nogil=True)
def nanmean3D(array):
    output = np.empty((array.shape[0],array.shape[1]))
    for i in nb.prange(array.shape[0]):
        for j in range(array.shape[1]):
            output[i,j]=np.nanmean(array[i,:,:][j,:])
    return output

希望对 nanmean 和 Numba 有疑问的其他人可以在 3D 矩阵上使用上述函数或修改它以供他们使用。

3D 阵列上的 Numba 和 NumPy nanmean - 简单的替代方案？

Numba and NumPy nanmean on a 3D array - easy alternative?

python

arrays

numpy

nan

numba