Home > other >  Creating 3D numpy array of vectors with components from set
Creating 3D numpy array of vectors with components from set

Time:11-15

I have a list of floats x1,...,xN. I want to create 3D numpy array of dimension NxNxN such that element i,j,k is the 3-element tuple x_i,x_j,x_k. So it's a cube with a vector defined at each grid site.

Example: if x1,...,xN=0,1 the 2x2x2 grid I wish to create is

layer 0: 
(0,0,0), (1,0,0)
(0,1,0), (1,1,0)

layer 1: 
(0,0,1), (1,0,1)
(0,1,1), (1,1,1)

If x1,...,xN is 0,...,N-1 then this is easily accomplished:

np.indices((N,N,N)).swapaxes(0,3).swapaxes(0,1).swapaxes(1,2)

For a generic list x1,...,xN, what is the fastest way to create the desired array?

I am particularly interested in the perhaps not so generic case where N=2k and x1,...,xN is 0,1,...,k-1,-k,-k 1,...,-1 This ordering comes from how numpy.fft returns frequencies for a FFT.

For that particular x1,...,xN you could use the line of code above, and then do the appropriate shifts vector component by component, but I'm not sure this is the fastest way, and how to do those shifts in the fastest way.

CodePudding user response:

Creating indices (np.indices((N,N,N))) is taking most of the time here and I found no faster way to crete indices (e.g. meshgrid is much slower).

What I propose is just very slightly faster way of transposing (proposing this more for the sake of clean code):

np.indices((N,N,N)).transpose(1, 2, 3, 0)

If you create such a grid of indices very often you can always cache your results and re-use them later.

  • Related