python - 2D 到 3D numpy 数组的高效转换

Question

我有一个数据数组，其中包含 N 个粒子在时间步 1 到 M 内的 ndim 坐标。数组中的列通常代表每个粒子“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 捕获它。我的解决方案将 narticles 作为参数，而不是 ndim，其中 narticles * ndim == data.shape[1]。我犯这个错误的部分原因是我更改了变量 ndim 的名称。在这种情况下，我会避免使用 ndim 作为变量名，因为它与属性 data.ndim 太相似，属性是数组的维数。这是更新的解决方案，但我已将 ndim 替换为dim1`。它与您原来的功能更相似。

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