python - 改变不同维度的数组一起广播

Question

我正在寻找一种更优化的方法来将 (n,n) 或 (n,n,1) 矩阵转换为 (n,n,3) 矩阵。我从 (n,n,3) 开始，但是在我对第二个轴求和到 (n,n) 之后，我的尺寸变小了。本质上，我想保持数组的原始大小，并让第二个轴重复 3 次。我需要这个的原因是我稍后会用另一个 (n,n,3) 数组广播它，但它们需要相同的维度。

我目前的方法可行，但看起来并不优雅。

a0=np.random.random((n,n))
b=a.flatten().tolist()
a=np.array(zip(b,b,b))
a.shape=n,n,3

此设置具有预期的结果，但笨拙且难以遵循。有没有办法通过复制第二个索引直接从 (n,n) 到 (n,n,3)？或者也许是一种不缩小数组大小的方法？

Answer 1

None 或 np.newaxis 是向数组添加维度的常用方法。 reshape 与 (3,3,1) 同样有效:

In [64]: arr=np.arange(9).reshape(3,3)
In [65]: arr1 = arr[...,None]
In [66]: arr1.shape
Out[66]: (3, 3, 1)

repeat as function or method replicates this.

In [72]: arr2=arr1.repeat(3,axis=2)
In [73]: arr2.shape
Out[73]: (3, 3, 3)
In [74]: arr2[0,0,:]
Out[74]: array([0, 0, 0])

但您可能不需要这样做。广播 (3,3,1) 与 (3,3,3) 一起工作。

In [75]: (arr1+arr2).shape
Out[75]: (3, 3, 3)

实际上它会用 (3,) 广播以产生 (3,3,3)。

In [77]: arr1+np.ones(3,int)
Out[77]: 
array([[[1, 1, 1],
        [2, 2, 2],
        ...
       [[7, 7, 7],
        [8, 8, 8],
        [9, 9, 9]]])

所以 arr1+np.zeros(3,int) 是将 (3,3,1) 扩展为 (3,3,3) 的另一种方法。

广播规则是:

(3,3,1) + (3,) => (3,3,1) + (1,1,3) => (3,3,3)

广播根据需要在开始时添加维度。

当你在一个轴上求和时，你可以通过一个参数来保持原来的维数:

In [78]: arr2.sum(axis=2).shape
Out[78]: (3, 3)
In [79]: arr2.sum(axis=2, keepdims=True).shape
Out[79]: (3, 3, 1)

如果您想沿任意维度从数组中减去平均值，这会很方便:

arr2-arr2.mean(axis=2, keepdims=True)