1

I want to generate a four dimensional array with dimensions (dim,N,N,N). The first component ndim =3 and N corresponds to the grid length. How can one elegantly generate such an array using python ?

here is my 'ugly' implementation:

qvec=np.zeros([ndim,N,N,N])  

freq   = np.arange(-(N-1)/2.,+(N+1)/2.)
x, y, z = np.meshgrid(freq[range(N)], freq[range(N)], freq[range(N)],indexing='ij')

qvec[0,:,:,:]=x
qvec[1,:,:,:]=y
qvec[2,:,:,:]=z
kmario23
  • 57,311
  • 13
  • 161
  • 150
Dr. Thanos
  • 13
  • 2

1 Answers1

0

Your implementation looks good enough to me. However, here are some improvements to make it prettier:

qvec=np.empty([ndim,N,N,N])  

freq   = np.arange(-(N-1)/2.,+(N+1)/2.)
x, y, z = np.meshgrid(*[freq]*ndim, indexing='ij')

qvec[0,...]=x   # qvec[0] = x
qvec[1,...]=y   # qvec[1] = y
qvec[2,...]=z   # qvec[2] = z

The improvements are:

  • Using numpy.empty() instead of numpy.zeros()
  • Getting rid of the range(N) indexing since that would give the same freq array
  • Using iterable unpacking and utilizing ndim
  • Using the ellipsis notation for dimensions (this is also not needed)

So, after incorporating all of the above points, the below piece of code would suffice:

qvec=np.empty([ndim,N,N,N])  
freq   = np.arange(-(N-1)/2.,+(N+1)/2.)
x, y, z = np.meshgrid(*[freq]*ndim, indexing='ij')    
qvec[0:ndim] = x, y, z

Note: I'm assuming N is same since you used same variable name.

kmario23
  • 57,311
  • 13
  • 161
  • 150
  • 1
    thank you for the suggestion. It does look pretty right now :). Yes N is the same. – Dr. Thanos May 09 '19 at 05:50
  • how exactly does *[freq]*ndim work ?I have usually seen *a on the left-hand side but not sure how this works here. – Dr. Thanos May 09 '19 at 06:30
  • `[freq]*3` is equivalent to `[freq, freq, freq]`. It just repeats the inside element that many times, here 3. Using a `*` at the front unpacks these elements (i.e. it removes the `list` since that's what `numpy.meshgrid` needs. HTH :) – kmario23 May 09 '19 at 07:49