高效转换2D到3D numpy数组

Question

我有一个数据数组，包含n个粒子的ndim坐标，时间跨度为1到M.数组中的列通常表示每个粒子'p'的（x，y，z），数组中的每一行代表另一个时间点' 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

我想将数组转换为3D格式，使得每个粒子都在numpy数组的不同（M x ndim）'切片'中。 我目前正在做以下事情：

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 

我对NumPy有基本的了解，但是希望提高我对数组操作的熟练程度。 如何重写上述函数以消除循环并使用更多'NumPythonic'结构？

谢谢。

-C

Answer 1

原始解决方案，与您的功能略有不同。

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

更新：我已经更新了我原来的答案，让我的错误更清楚，谢谢unutbu抓住它。 我的解决方案使用nparticles作为参数而不是ndim ，其中nparticles * ndim == data.shape[1] 。 我犯了错误，部分原因是我更改了变量ndim的名称。 在这种情况下，我会避免使用ndim作为变量名，因为它太过类似于data.ndim属性，它是数组的维数。 这是更新的解决方案，但我已经用ndim by替换ndim by ndim。 它更类似于您的原始功能。

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

Answer 2

怎么样：

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

例如，如果

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

然后alt_3D(data)产生：

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

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

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

（这与Bago的答案不同。）