Creating vectorized numpy.meshgrid off 2D arrays to create 3D meshes

Question

Is there a way to vectorize the creation of np.meshgrid?

#Example of normal np.meshgrid

x = [1.22749725, 2.40009184, 1.48602747, 2.83286752, 2.37426951]
y = [1.23816021, 1.69811451, 2.08692546, 2.13706377, 1.60298873]
gridrez=10

X,Y = np.meshgrid(np.linspace(min(x), max(x), endpoint=True, num=gridrez), 
                  np.linspace(min(y), max(y), endpoint=True, num=gridrez))
# X.shape = 10x10
# Y.shape = 10x10

I would like to replicate this functionality but instead of x,y being 1x5 arrays they are 1000x5 arrays, and the resulting X,Y would be 1000x10x10. Note I have found a way to vectorize the np.linspace so don't worry about it. Assume we have two (1000x10) arrays and want to create a meshgrid of 1000x10x10. When I feed the answer in I get a (10000,10000) meshgrid instead of 1000x10x10

Divakar · Accepted Answer · 2019-03-27T14:15:21.137

Here's one using vectorized-linspace : create_ranges -

# https://stackoverflow.com/a/40624614/ @Divakar
def create_ranges(start, stop, N, endpoint=True):
    if endpoint==1:
        divisor = N-1
    else:
        divisor = N
    steps = (1.0/divisor) * (stop - start)
    return steps[:,None]*np.arange(N) + start[:,None]

def linspace_nd(x,y,gridrez):
    a1 = create_ranges(x.min(1), x.max(1), N=gridrez, endpoint=True)
    a2 = create_ranges(y.min(1), y.max(1), N=gridrez, endpoint=True)
    out_shp = a1.shape + (a2.shape[1],)
    Xout = np.broadcast_to(a1[:,None,:], out_shp)
    Yout = np.broadcast_to(a2[:,:,None], out_shp)
    return Xout, Yout

The final outputs off linspace_nd would be 3D mesh views into the vectorized linspace outputs and as such would be memory-efficient and hence good on performance too.

Alternatively, if you need outputs with their own memory spaces and not the views, you can use np.repeat for the replications -

Xout = np.repeat(a1[:,None,:],a2.shape[1],axis=1)
Yout = np.repeat(a2[:,:,None],a1.shape[1],axis=2)

Timings to create such an array with views -

In [406]: np.random.seed(0)
     ...: x = np.random.rand(1000,5)
     ...: y = np.random.rand(1000,5)

In [408]: %timeit linspace_nd(x,y,gridrez=10)
1000 loops, best of 3: 221 µs per loop

Hi I can't seem to replicate the correct results using this, I'll add the data in the question. — Cr1064, Mar 27 '19 at 15:16
@ChristianRobinson Remember to add minimal sample data. So, maybe you can add a sample data of shape `(2,5)` random data. — Divakar, Mar 27 '19 at 15:18

Creating vectorized numpy.meshgrid off 2D arrays to create 3D meshes

1 Answers1