I would like to create a meshgrid of variable dimensions by specifying the dimensions with a variable i.e. specifying dim=2, rather than manually changing the expression as in the example below to set a 2D mesh grid.
How would I implement a wrapper function for this?
The problem stems from not being familiar with the syntax that mgrid uses (index_tricks).
import numpy as np
mgrid = np.mgrid[
-5:5:5j,
-5:5:5j,
]
Observed documentation for mgrid, but there seems to be no info on setting the number of dimensions with a variable.
CodePudding user response:
You can create a tuple containing the slice manually, and repeat it some number of times:
import numpy as np
num_dims = 2
mgrid = np.mgrid[(slice(-5, 5, 5j),) * num_dims]