形象化展示numpy.stack

https://numpy.org/doc/stable/reference/generated/numpy.stack.html#numpy.stack

1. 一维数组

import numpy as np

a = np.arange(4)
b = np.arange(4)
c = np.stack([a,b])
d = np.stack([a,b], axis=1)
print('a -->', a.shape,':\n', a)
print('b -->', b.shape,':\n', b)
print('c -->', c.shape,':\n', c)
print('d -->', d.shape,':\n', d)

输出：

a --> (4,) :
 [0 1 2 3]
b --> (4,) :
 [0 1 2 3]
c --> (2, 4) :
 [[0 1 2 3]
 [0 1 2 3]]
d --> (4, 2) :
 [[0 0]
 [1 1]
 [2 2]
 [3 3]]

2. 二维数组

import numpy as np

a = np.arange(8).reshape(2,4)
b = np.arange(8,16).reshape(2,4)
c = np.arange(16,24).reshape(2,4)

d = np.stack([a,b,c], axis=0)
e = np.stack([a,b,c], axis=1)
f = np.stack([a,b,c], axis=2)

dd = [[[0,1,2,3],[4,5,6,7]],[[8,9,10,11],[12,13,14,15]],[[16,17,18,19],[20,21,22,23]]]
ee = [[[0,1,2,3],[8,9,10,11],[16,17,18,19]],[[4,5,6,7],[12,13,14,15],[20,21,22,23]]]
ff = [[[0,8,16],[1,9,17],[2,10,18],[3,11,19]],[[4,12,20],[5,13,21],[6,14,22],[7,15,23]]]
print(dd == d)
print(ee == e)
print(ff == f)

print('a -->', a.shape,':\n', a)
print('b -->', b.shape,':\n', b)
print('c -->', c.shape,':\n', c)

print('d -->', d.shape,':\n', d)
print('e -->', e.shape,':\n', e)
print('f -->', f.shape,':\n', f)

输出：

[[[ True  True  True  True]
  [ True  True  True  True]]

 [[ True  True  True  True]
  [ True  True  True  True]]

 [[ True  True  True  True]
  [ True  True  True  True]]]
[[[ True  True  True  True]
  [ True  True  True  True]
  [ True  True  True  True]]

 [[ True  True  True  True]
  [ True  True  True  True]
  [ True  True  True  True]]]
[[[ True  True  True]
  [ True  True  True]
  [ True  True  True]
  [ True  True  True]]

 [[ True  True  True]
  [ True  True  True]
  [ True  True  True]
  [ True  True  True]]]
a --> (2, 4) :
 [[0 1 2 3]
 [4 5 6 7]]
b --> (2, 4) :
 [[ 8  9 10 11]
 [12 13 14 15]]
c --> (2, 4) :
 [[16 17 18 19]
 [20 21 22 23]]
d --> (3, 2, 4) :
 [[[ 0  1  2  3]
  [ 4  5  6  7]]

 [[ 8  9 10 11]
  [12 13 14 15]]

 [[16 17 18 19]
  [20 21 22 23]]]
e --> (2, 3, 4) :
 [[[ 0  1  2  3]
  [ 8  9 10 11]
  [16 17 18 19]]

 [[ 4  5  6  7]
  [12 13 14 15]
  [20 21 22 23]]]
f --> (2, 4, 3) :
 [[[ 0  8 16]
  [ 1  9 17]
  [ 2 10 18]
  [ 3 11 19]]

 [[ 4 12 20]
  [ 5 13 21]
  [ 6 14 22]
  [ 7 15 23]]]

2.1 如何理解堆叠的规则呢？

Join a sequence of arrays along a new axis.，所以堆叠后的数组比原数组多一个轴（维度）。
对于axis=0特别好理解，就是把原数组每个都看成一个整体再组成一个新的数组，新数组的第一个元素就是参与堆叠的第一个数组，新数组的第二个元素就是参与堆叠的第二个数组，以此类推，新数组最后一个元素就是参与堆叠的最后一个数组。对于3个2*4的数组而言，堆叠后数组的形状是3*2*4。

对于axis=1可以把参与堆叠的每个数组拉平按行排列好，如下图所示。堆叠后的新数组比原数组也会多一个维度，新数组形状第一个维度的值是原数组的第一个维度的值，第二个维度的值是参与堆叠的数组的个数，第三个维度的值是原数组第二个维度的值。对于3个2*4的数组而言，堆叠后数组的形状是2*3*4；新数组的取值第一个元素是下图中蓝色区域，每行是一个数组；第二个元素是下图中绿色区域，每行是一个数组。

对于axis=2操作上类似axis=1，新产生的数组在形状上第一个维度相同，第二个维度与第三个维度与axis=1的情况相互交换。对于3个2*4的数组而言，堆叠后数组的形状是2*4*3。新数组的取值第一个元素是下图中蓝色区域做转置产生的4个数组；第二个元素是下图中绿色区域做转置产生的4个数组。