将 1D 数组填充到 3D 数组中

Filling 1D arrays into 3D arrays vectorisation

我有一个函数init_tensor(),它将维度(N,N)的2d矩阵广播到dim(M,N,N)的3d块矩阵中,这样就有M个维度NXN的矩阵:

def init_tensor(input_state, sample_size):
    return np.broadcast_to(input_state, (sample_size,)+input_state.shape)

例如,如果我想创建 3 (4x4) 个矩阵,那么我可以这样做:

init_tensor(np.eye(4, dtype=complex), 3)
Out[462]: 
array([[[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
        [0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j],
        [0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j],
        [0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j]],

       [[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
        [0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j],
        [0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j],
        [0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j]],

       [[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
        [0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j],
        [0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j],
        [0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j]]])

我遇到的问题是我有一些带有 dim (1,M) 的数组,我想将其作为元素填充到 3D 数组中。对于一个简单的例子,如果 M 是 3 并且我有:

lambda1 = [l11,l12,l13]
lambda2 = [l21,l22,l23]
lambda3 = [l31,l32,l33]
tau1 = [t11,t12,t13]
tau2 = [t21,t22,t23]
tau3 = [t31,t32,t33]

我想要一种矢量化的方式,我可以将它们填充到张量中,使其变为:

array([[[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
        [   t11,    l11, 0.+0.j, 0.+0.j],
        [   t21, 0.+0.j,    l21, 0.+0.j],
        [   t31, 0.+0.j, 0.+0.j,    l31]],

       [[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
        [   t12,    l12, 0.+0.j, 0.+0.j],
        [   t22, 0.+0.j,    l22, 0.+0.j],
        [   t32, 0.+0.j, 0.+0.j,    l32]],

       [[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
        [   t13,    l13, 0.+0.j, 0.+0.j],
        [   t23, 0.+0.j,    l23, 0.+0.j],
        [   t33, 0.+0.j, 0.+0.j,    l33]]])

张量矩阵的深度将始终与一维数组的长度相同,M 的值可以在 1 到 100 之间变化。

您可以使用 numpys advanced indexing 得到这个结果:

import numpy as np

m = 3
n = 4
arr = np.ones((m, n, n), dtype=int)

lmbda = 10 * np.arange(1, (n-1)**2+1).reshape(n-1, n-1)
# lmbda = [[l11,l12,l13]
#          [l21,l22,l23]
#          [l31,l32,l33]]
# array([[10, 20, 30],
#        [40, 50, 60],
#        [70, 80, 90]])
tau = -10 * np.arange(1, (n-1)**2+1).reshape(n-1, n-1)
# array([[-10, -20, -30],
#        [-40, -50, -60],
#        [-70, -80, -90]])

arr[:, range(1, n), range(1, n)] = lmbda.T
arr[:, range(1, n), 0] = tau.T  # arr[:, 1:, 0] = tau.T  # alternative
# array([[[  1,   1,   1,   1],
#         [-10,  10,   1,   1],
#         [-40,   1,  40,   1],
#         [-70,   1,   1,  70]],
#        [[  1,   1,   1,   1],
#         [-20,  20,   1,   1],
#         [-50,   1,  50,   1],
#         [-80,   1,   1,  80]],
#        [[  1,   1,   1,   1],
#         [-30,  30,   1,   1],
#         [-60,   1,  60,   1],
#         [-90,   1,   1,  90]]])