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

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

您可以先在 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

计时，使用built-in 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

Answer 2

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)

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

Altering arrays of different dimensions to be broadcasted together

python

arrays

numpy

broadcast

dimensions