如何在python中生成n维网格

Question

I want to generate an n-dimensional grid.

For a 3D grid, I have the following working code (which creates a grid of 5X5X5 between (-1,1 )

import numpy as np
subdivision = 5
step = 1.0/subdivision
grid= np.mgrid[ step-1 : 1.0-step: complex(0, subdivision),
                step-1 : 1.0-step: complex(0, subdivision),
                step-1 : 1.0-step: complex(0, subdivision)]

I want to generalize this to n dimensions so something like

grid = np.mgrid[step-1 : 1.0-step: complex(0,subdivision) for i in range(n)]

But this obviously doesnt work. I also tried

temp = [np.linspace(step-1 , 1.0-step, subdivision) for i in range(D)]
grid = np.mgrid[temp]

But this doesn't work either since np.mgrid accepts slices

Answer 1

Instead of using complex you can define the step size explicitly using real numbers. In my opinion this is more concise:

grid= np.mgrid[ step-1 : 1.0: step * 2,
                step-1 : 1.0: step * 2,
                step-1 : 1.0: step * 2]

Dissecting above snippet, we see that step-1 : 1.0: step * 2 defines a slice, and separating them by , creates a tuple of three slices, which is passed to np.mgrid.__getitem__ .

We can generalize this to n dimensions by constructing a tuple of n slices:

n = 3
grid= np.mgrid[tuple(slice(step - 1, 1, step * 2) for _ in range(n))]

Answer 2

As suggested by kazemakase , you should replace the "short hand" slicing notations step-1 : 1.0-step: complex(0,subdivision) with an explicit call to slice , and then combine it in a " tuple generator":

D = 6
grid = np.mgrid[tuple(slice(step-1, 1.0-step, complex(0,subdivision)) for i in range(D))]

Results with a 6D grid.

Answer 3

You can use meshgrid and linspace to do what you want.

import numpy as np
X1, X2, X3 = np.meshgrid(*[np.linspace(-1,1,5),
                           np.linspace(-1,1,5),
                           np.linspace(-1,1,5)])

For many dimensions, you can do

D = 4
subdivision = 5
temp = [np.linspace(-1.0 , 1.0, subdivision) for i in range(D)]
res_to_unpack = np.meshgrid(*temp)
assert(len(res_to_unpack)==D)