How to get meshgrid of vectors of matrices in numpy?

Question

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 nd 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.

Answer 1

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.