Numba Compatible Numpy Meshgrid

user_1663725 0
Viewed 526

I have a function that generates flattened arrays from a 3D meshgrid as shown below.

I'm trying to implement Numba throughout my code, but when Numba encounters this function it throws an error (since Numba doesn't support Numpy's meshgrid or mgrid functions). Is there an alternative way I can make this Numba compatible?

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
numba multidimensional-array numpy python
1 Answers

You can implement a fit-for-purpose meshgrid yourself and numbify it if you want. Since you want a 3D meshgrid with indexing ij, the following implementation would do.

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

This seems to be anecdotally faster than numpy, since it is specific to your needs.

 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)

You can also modify the implementation above to return flat arrays instead. Just change the np.empty size to the product of the three vector sizes, and change the indexing within the loop:

@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
user_35680920
Related