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 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)

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的答案不同。）

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

2 个答案: