arrays - 通过减去维度来折叠立方体

Question

我有一个立方体，其中包含基于形状 [t1,t2,gamma] 的 3 个变量的概率密度，每个方向有 80 个值。我需要绘制 T，gamma 的分布。其中 T 是 t2-t1。

有没有一些聪明的方法可以把这个立方体折叠成想要的结果？我一直在努力寻找它，但我找不到。

标签： arraysnumpy

Answer 1

您可以用零填充，然后通过将维度递增或递减一来剪切数组：

# make small diagnostic example
nt1, nt2, ngm = 4, 5, 2
data = sum(np.ogrid[1:nt1+1,-1:-nt2-1:-1,100:100*ngm+100:100])
# by construction values are equal if coordinates (T,gamma) are equal, no matter how T = t2-t1 decomposes.
# Fixing gamma, for example at 1, we can see that T is constant along the diagonals
data[..., 1]
# array([[200, 199, 198, 197, 196],
#        [201, 200, 199, 198, 197],
#        [202, 201, 200, 199, 198],
#        [203, 202, 201, 200, 199]])

# now let's transform the example, first recover dimensions
nt1, nt2, ngm = data.shape
# next, zero pad
aux = np.zeros((nt1+2, nt1+nt2-2, ngm), data.dtype)
aux[1:-1, :nt2] = data
# and shear, in this case by incrementing dimension 1
sheared = aux.reshape(-1, ngm)[nt2-1:3-nt1-nt2].reshape(nt1, nt1+nt2-1, ngm)

# check result, for example at gamma = 1
sheared[..., 1]
# array([[  0,   0,   0, 200, 199, 198, 197, 196],
#        [  0,   0, 201, 200, 199, 198, 197,   0],
#        [  0, 202, 201, 200, 199, 198,   0,   0],
#        [203, 202, 201, 200, 199,   0,   0,   0]])


# corresponding values of T are now aligned and ready for further processing.