CODEWITHSUNDEEP
pythonnumpy
itzik Ben Shabat
itzik Ben Shabat

Reputation: 927

How to generate an n-dimensional grid in python

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

Upvotes: 5

Views: 4007

Answers (3)

MB-F
MB-F

Reputation: 23647

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))]

Upvotes: 5

Shai
Shai

Reputation: 114926

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.

Upvotes: 2

Guillaume Jacquenot
Guillaume Jacquenot

Reputation: 11717

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)

Upvotes: 0

Related Questions