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

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 作为函数或方法复制了这一点。

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)

Answer 2

您可以先在a 上创建一个新轴（axis = 2），然后沿着这个新轴使用np.repeat：

np.repeat(a[:,:,None], 3, axis = 2)

或者另一种方法，展平数组，重复元素，然后重塑：

np.repeat(a.ravel(), 3).reshape(n,n,3)

---

结果对比：

import numpy as np
n = 4
a=np.random.random((n,n))
b=a.flatten().tolist()
a1=np.array(zip(b,b,b))
a1.shape=n,n,3
# a1 is the result from the original method

(np.repeat(a[:,:,None], 3, axis = 2) == a1).all()
# True

(np.repeat(a.ravel(), 3).reshape(4,4,3) == a1).all()
# True

---

计时，使用内置numpy.repeat也显示提速：

import numpy as np
n = 4
a=np.random.random((n,n))
​
def rep():
    b=a.flatten().tolist()
    a1=np.array(zip(b,b,b))
    a1.shape=n,n,3

%timeit rep()
# 100000 loops, best of 3: 7.11 µs per loop

%timeit np.repeat(a[:,:,None], 3, axis = 2)
# 1000000 loops, best of 3: 1.64 µs per loop

%timeit np.repeat(a.ravel(), 3).reshape(4,4,3)
# 1000000 loops, best of 3: 1.9 µs per loop