仅在python中使用ndarray的特定尺寸进行广播

Question

考虑TxFxM ndarray。我希望将其与共轭数相乘，仅用于M维，同时保持其他维与以下代码中呈现的相同：

import numpy as np

T=2
F=3
M=4

x=np.random.rand(T,F,M)

result=np.zeros((T,F,M,M))

for i in range(x.shape[0]):
    for j in range(x.shape[1]):
            result[i,j,:,:]=np.matmul(np.expand_dims(x[i,j,:],axis=1),np.expand_dims(x[i,j,:],axis=0).conj())

如果我像np.matmul(x,x.conj().T)那样简单地使用广播，则广播操作将继续在更高维度上进行并保持倍增。另一方面，由于两个循环，我的实现非常慢，并且对我的理解非常不合常规。

是否有实现此S.T.它会运行得更快吗？

P.S。

• 我的尺寸显然要大T=3000,F=1024,M=4，并且此操作会重复进行，因此我需要快速实现。
• 我计划对此维度T进行平均，因此，如果总体实施速度更快，我将非常感兴趣。

Answer 1

如果在x和x.conj()的两个不同位置注入单例尺寸，则可以通过广播计算所需的数组。如果x的形状为(T,F,M)，则形状为(T,F,M,1)和(T,F,1,M)的数组将按照您想要的方式广播到(T,F,M,M)。这是您的示例，其中包含复杂数据，以确保我们不会丢失某些内容：

import numpy as np 

T,F,M = 2,3,4 

x = np.random.rand(T,F,M) + np.random.rand(T,F,M)*1j 

result = np.zeros((T,F,M,M), dtype=complex) 

# loop
for i in range(x.shape[0]): 
    for j in range(x.shape[1]):
        result[i,j,:,:] = np.matmul(np.expand_dims(x[i,j,:],axis=1),
                                    np.expand_dims(x[i,j,:],axis=0).conj())

# broadcasting
result2 = x[..., None] * x[..., None, :].conj()

# proof
print(np.array_equal(result, result2))  # True

由于您提到要在T大小的维度上取平均值，因此我们必须考虑是否值得将该维度放在最后，以便该平均值尽可能使用连续的内存块。这意味着以下选项：

def summed_original(x): 
    """Assume x.shape == (T, F, M), return summed array of shape (F, M, M)""" 
    return (x[..., None] * x[..., None, :].conj()).mean(0) 

def summed_transposed(x): 
    """Assume x.shape == (F, M, T), return summed array of shape (F, M, M)""" 
    return (x[..., None, :] * x[:, None, ...].conj()).mean(-1) 

x_transposed = x.transpose(1, 2, 0).copy()  # ensure contiguous copy 
print(np.allclose(summed_original(x), summed_transposed(x_transposed)))  # True

如您所见，这两个函数计算的是相同的东西，但是它们假定输入具有不同的存储顺序。之所以如此重要，是因为将原始数组放在不同的内存布局中可能会更快（如果需要的话，会在开始时进行一次转置和复制）。

因此，让我们使用IPython的%timeit魔术和您的实际尺寸为它们计时：

T,F,M = 3000,1024,4 
x = np.random.rand(T, F, M) + np.random.rand(T, F, M)*1j 
x_transposed = x.transpose(1, 2, 0).copy() 
print(np.allclose(summed_original(x), summed_transposed(x_transposed)))  # True

%timeit summed_original(x)
# 500 ms ± 16.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit summed_transposed(x_transposed)
# 352 ms ± 2.31 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
如您所见，对于您的特定大小和类型，重新安排数组的尺寸似乎很值得，以便T尺寸对应于连续的内存块，有助于CPU进行缓存。您可以在开始时通过.transpose(...).copy()调用来执行此操作，或者甚至更好的方法是可以首先构造数组，使代码最佳。