python - 循环中的 numpy column_stack

Question

我不知道这个标题是否合适，但让我告诉你我想做什么，

In [56]: import numpy as np

In [57]: a= np.random.rand(2,2,2); a
Out[57]: 
array([[[0.4300565 , 0.82251319],
        [0.56113378, 0.83284255]],

       [[0.00822414, 0.28256243],
        [0.16648411, 0.33381438]]])

In [58]: b=np.random.rand(2); b
Out[58]: array([0.8035224 , 0.09884653])

In [59]: np.stack(( np.column_stack((b,a[:,i,:])) for i in range(a.shape[1])))
Out[59]: 
array([[[0.8035224 , 0.4300565 , 0.82251319],
        [0.09884653, 0.00822414, 0.28256243]],

       [[0.8035224 , 0.56113378, 0.83284255],
        [0.09884653, 0.16648411, 0.33381438]]])

所以，我想将一个数组作为列堆叠到一个内轴。是否可以在 numpy 中更高效、更简洁地执行循环结构？我尝试使用 numpy insert 但无法做到。

编辑:

另一个例子

In [110]: a= np.random.rand(5,3,3); a
Out[110]: 
array([[[0.27506756, 0.82334411, 0.7004287 ],
        [0.6834928 , 0.28457133, 0.6275462 ],
        [0.49744358, 0.25131814, 0.56422852]],

       [[0.82591597, 0.92367306, 0.04652992],
        [0.98545051, 0.92813944, 0.14360307],
        [0.85454081, 0.8254149 , 0.5637401 ]],

       [[0.59545519, 0.41563571, 0.41937218],
        [0.90980491, 0.30169504, 0.96630809],
        [0.06713389, 0.64357544, 0.12901734]],

       [[0.47566444, 0.33476802, 0.26635363],
        [0.4678913 , 0.53028241, 0.03112231],
        [0.68445959, 0.07113376, 0.86651669]],

       [[0.66951982, 0.01827502, 0.43831829],
        [0.02798567, 0.36880876, 0.55029074],
        [0.40127051, 0.6311474 , 0.51015882]]])

In [111]:  b= np.random.rand(5,2); b
Out[111]: 
array([[0.01659589, 0.15320541],
       [0.79025065, 0.28041334],
       [0.56024173, 0.49317082],
       [0.28229119, 0.46010724],
       [0.72239851, 0.62075004]])

In [112]: np.stack(( np.column_stack((b,a[:,i,:])) for i in range(a.shape[1])))
Out[112]: 
array([[[0.01659589, 0.15320541, 0.27506756, 0.82334411, 0.7004287 ],
        [0.79025065, 0.28041334, 0.82591597, 0.92367306, 0.04652992],
        [0.56024173, 0.49317082, 0.59545519, 0.41563571, 0.41937218],
        [0.28229119, 0.46010724, 0.47566444, 0.33476802, 0.26635363],
        [0.72239851, 0.62075004, 0.66951982, 0.01827502, 0.43831829]],

       [[0.01659589, 0.15320541, 0.6834928 , 0.28457133, 0.6275462 ],
        [0.79025065, 0.28041334, 0.98545051, 0.92813944, 0.14360307],
        [0.56024173, 0.49317082, 0.90980491, 0.30169504, 0.96630809],
        [0.28229119, 0.46010724, 0.4678913 , 0.53028241, 0.03112231],
        [0.72239851, 0.62075004, 0.02798567, 0.36880876, 0.55029074]],

       [[0.01659589, 0.15320541, 0.49744358, 0.25131814, 0.56422852],
        [0.79025065, 0.28041334, 0.85454081, 0.8254149 , 0.5637401 ],
        [0.56024173, 0.49317082, 0.06713389, 0.64357544, 0.12901734],
        [0.28229119, 0.46010724, 0.68445959, 0.07113376, 0.86651669],
        [0.72239851, 0.62075004, 0.40127051, 0.6311474 , 0.51015882]]])

Answer 1

连接的一种变体是索引赋值:

对于第一个例子:

In [245]: a=np.arange(8).reshape(2,2,2); b=np.array([100,200])
In [246]: c = np.zeros((2,2,3), a.dtype)
In [247]: c[:,:,0]=b
In [248]: c[:,:,1:]=a.transpose(1,0,2)
In [249]: c
Out[249]: 
array([[[100,   0,   1],
        [200,   4,   5]],

       [[100,   2,   3],
        [200,   6,   7]]])

第二个:

In [250]: a1 = np.arange(5*3*3).reshape(5,3,3)
In [251]: b1 = np.arange(10).reshape(5,2)
In [252]: c1 = np.zeros((3,5,5),a.dtype)
In [253]: c1[:,:,:2]=b1
In [254]: c1[:,:,2:]=a1.transpose(1,0,2)
In [255]: c1
Out[255]: 
array([[[ 0,  1,  0,  1,  2],
        [ 2,  3,  9, 10, 11],
        [ 4,  5, 18, 19, 20],
        [ 6,  7, 27, 28, 29],
        [ 8,  9, 36, 37, 38]],

       [[ 0,  1,  3,  4,  5],
        [ 2,  3, 12, 13, 14],
        [ 4,  5, 21, 22, 23],
        [ 6,  7, 30, 31, 32],
        [ 8,  9, 39, 40, 41]],

       [[ 0,  1,  6,  7,  8],
        [ 2,  3, 15, 16, 17],
        [ 4,  5, 24, 25, 26],
        [ 6,  7, 33, 34, 35],
        [ 8,  9, 42, 43, 44]]])

从a 和b 导出c 的形状留给读者作为练习。 :)

np.stack(或 np.array)在第 2 轴上的迭代实际上是部分转置(或前 2 个轴的互换):

In [261]: np.stack([a[:,i,:] for i in range(a.shape[1])])
Out[261]: 
array([[[0, 1],
        [4, 5]],

       [[2, 3],
        [6, 7]]])
In [262]: a.transpose(1,0,2)
Out[262]: 
array([[[0, 1],
        [4, 5]],

       [[2, 3],
        [6, 7]]])

我们也可以在第一个轴上迭代，然后在第二个轴上加入:

In [263]: np.stack(a, axis=1)
Out[263]: 
array([[[0, 1],
        [4, 5]],

       [[2, 3],
        [6, 7]]])

对 Ankit's answer 的改进使用 concatenate 是:

np.concatenate([np.repeat(b[None,:,None], 2, axis=0), a.transpose(1,0,2)], axis=2)
np.concatenate([np.repeat(b1[None,:,:], 3, axis=0), a1.transpose(1,0,2)], axis=2)