Numba 兼容的 Numpy Meshgrid

Question

我有一个函数可以从 3D 网格生成扁平数组，如下所示。

我试图在我的代码中实现 Numba，但是当 Numba 遇到这个函数时它会抛出一个错误（因为 Numba 不支持 Numpy 的 meshgrid 或 mgrid 函数）。有没有其他方法可以让这个 Numba 兼容？

def meshgrid_flat(max=1.0, sampling=100):
    #
    s = np.linspace(-max,max,sampling)
    X, Y, Z = np.meshgrid(s,s,s,indexing="ij")
    x, y, z = X.ravel(), Y.ravel(), Z.ravel()
    #
    return x, y, z

Answer 1

您可以自己实现一个适合用途的网格网格，并根据需要对其进行麻木。由于您想要一个带有索引 ij 的 3D 网格，因此可以使用以下实现。

import numba

@numba.jit(nopython=True)
def meshgrid(x, y, z):
    xx = np.empty(shape=(x.size, y.size, z.size), dtype=x.dtype)
    yy = np.empty(shape=(x.size, y.size, z.size), dtype=y.dtype)
    zz = np.empty(shape=(x.size, y.size, z.size), dtype=z.dtype)
    for i in range(z.size):
        for j in range(y.size):
            for k in range(x.size):
                xx[i,j,k] = k  # change to x[k] if indexing xy
                yy[i,j,k] = j  # change to y[j] if indexing xy
                zz[i,j,k] = i  # change to z[i] if indexing xy
    return zz, yy, xx

这似乎比 numpy 快得多，因为它特定于您的需求。

 x, y, z = np.arange(100), np.arange(100), np.arange(10)

 %timeit np.meshgrid(x, y, z, indexing="ij")
177 µs ± 9.07 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

 %timeit meshgrid(x, y, z)
47.3 µs ± 3.22 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

您也可以修改上面的实现以返回平面数组。只需将 np.empty 大小更改为三个向量大小的乘积，然后更改循环内的索引：

@numba.jit(nopython=True)
def meshgrid_flat_3d(x):
    xx = np.empty(shape=(x.size * x.size * x.size), dtype=x.dtype)
    yy = np.empty_like(xx)
    zz = np.empty_like(xx)
    for i in range(x.size):
        for j in range(x.size):
            for k in range(x.size):
                xx[i*x.size**2 + j*x.size + k] = k
                yy[i*x.size**2 + j*x.size + k] = j  
                zz[i*x.size**2 + j*x.size + k] = i  
    return zz, yy, xx