efficient conversion of 2D to 3D numpy array

Question

I have a data array containing ndim coordinates of N particles over timesteps 1 to M. The columns in the array typically represent the (x,y,z) of each particle 'p', and each row in the array represents another time point 't':

x_t1p1  y_t1p1  z_t1p1  x_t1p2  y_t1p2  z_t1p2  ...  x_t1pN  y_t1pN  z_t1pN
x_t2p1  y_t2p1  z_t2p1  x_t2p2  y_t2p2  z_t2p2  ...  x_t2pN  y_t2pN  z_t2pN
...
x_tMp1  y_tMp1  z_tMp1  x_tMp2  y_tMp2  z_t1p2  ...  x_tMpN  y_tMpN  z_tMpN

I would like to convert the array to a 3D format such that each particle is in a different (M x ndim) 'slice' of the numpy array. I am currently doing the following:

import numpy as np
def datarray_to_3D(data, ndim=3):
    (nr,nc) = data.shape
    nparticles = nc/ndim
    dat_3D = np.zeros([nr,ndim,nparticles])
    for i in range(nparticles):
        dat_3D[:,:,i] = data[:,i*ndim:(i+1)*ndim]
    return dat_3D

I have a basic knowledge of NumPy, but would like to increase my proficiency in array manipulation. How can the above function be rewritten to eliminate the loop and use a more 'NumPythonic' structure?

Thank you.

-c

Bi Rico · Accepted Answer · 2012-08-07 22:31:06Z

2

Original solution, slightly different than your function.

def datarray_to_3D(data, nparticles=3):
    nr, nc = data.shape
    data = data.reshape(nr, nparticles, nc/nparticles)
    return np.rollaxis(data, 2, 1)

Update: I've updated my original answer to make my mistake more clear, Thanks unutbu for catching it. My solution took nparticles as an argument instead of ndim where nparticles * ndim == data.shape[1]. I made the mistake partly becuase I changed the name of your variable ndim. I would avoid using ndim as a variable name in this case because it is too similar to the attribute data.ndim which is the number of dimensions of the array. Here is the updated solution, but I've replaced ndim bydim1`. It is more similar to your original function.

def datarray_to_3D(data, dim1=3):
    nr, nc = data.shape
    data = data.reshape(nr, nc/dim1, dim1)
    return np.rollaxis(data, 2, 1)

edited Aug 7, 2012 at 22:31

answered Aug 7, 2012 at 17:44

Bi Rico

25.9k3 gold badges57 silver badges75 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

cytochrome Over a year ago

Thank you! This is just what I was looking for and allows me the chance to explore the reshape and rollaxis methods.

unutbu · Accepted Answer · 2012-08-07 18:50:45Z

1

How about:

def alt_3D(data, ndim=3):
    nr, nc = data.shape
    result = data.reshape(nr,-1,3).transpose(0,2,1)
    return result

For example, if

data = np.arange(18).reshape((-1,6))

then alt_3D(data) yields:

[[[ 0  3]
  [ 1  4]
  [ 2  5]]

 [[ 6  9]
  [ 7 10]
  [ 8 11]]

 [[12 15]
  [13 16]
  [14 17]]]

(This is a different result than Bago's answer.)

answered Aug 7, 2012 at 18:50

unutbu

886k197 gold badges1.9k silver badges1.7k bronze badges

Collectives™ on Stack Overflow

efficient conversion of 2D to 3D numpy array

2 Answers 2

1 Comment

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

Comments

Your Answer

Sign up or log in

Post as a guest

Related