A normal meshgrid on two 1d vectors returns a matrix for each, containing duplicates of itself to fit the length of the other.

import numpy as np

a, b = np.meshgrid(np.arange(2), np.arange(3, 6))
a
Out[22]: 
array([[0, 1],
       [0, 1],
       [0, 1]])
b
Out[23]: 
array([[3, 3],
       [4, 4],
       [5, 5]])

I want the same, but with each element being a n-d volume, with the meshgrid only on the 1st dimension:

v1
Out[17]: 
array([[0, 1, 2, 3, 4],
       [5, 6, 7, 8, 9]])

v2
Out[18]: 
array([[11, 12, 13, 14, 15, 16],
       [17, 18, 19, 20, 21, 22],
       [23, 24, 25, 26, 27, 28]])

v1.shape
Out[19]: (2, 5)

v2.shape
Out[20]: (3, 6)

I want two meshgrids v1_mesh.shape==(2, 3, 5) and v2_mesh.shape==(2, 3, 6).

v1_mesh[i, :, :] == v1 and v2_mesh[:, j, :] == v2 for all relevant indices, just like a standard meshgrid.

That is a total of 2*3=6 == np.prod([v1.shape[0], v2.shape[0]]) combinations.

Using a, b = np.meshgrid(v1, v2) gives a.shape == b.shape == (np.prod(v1.shape), np.prod(v1.shape)) which is more combinations than I wanted. I only want combinations along the 1st axis.

CodePudding user response：

meshgrid specifies that the inputs are 1d, so in your case it is effectively ravel them first, hence the prod shape.

In [3]: v1=np.arange(10).reshape(2,5)
In [5]: v2=np.arange(11,29).reshape(3,6)

These 2 arrays should be the equivalent of meshgrid with sparse (the meshgrid code does this, except is uses reshape instead of the None indexing).

In [6]: v11, v21 = v1[:,None,:], v2[None,:,:]
In [7]: v11.shape
Out[7]: (2, 1, 5)
In [8]: v21.shape
Out[8]: (1, 3, 6)

We can flesh them out to the full shape with repeat:

In [9]: v12 = np.repeat(v11,3,1)
In [10]: v22 = np.repeat(v21,2,0)
In [11]: v12.shape
Out[11]: (2, 3, 5)
In [12]: v22.shape
Out[12]: (2, 3, 6)

meshgrid uses broadcast_arrays to expand the 'dense' dimensions, but repeat is simpler to understand.

You could also create arrays e.g. v1_mesh with the right shape, and assign values, taking advantage of broadcasting.