Numpy：计算大数组的协方差

Question

我有一个形状为(32,2048,2048) 的大型numpy 数组，其中每个[i,:,:] 是一组二维像素，它们是来自空间相关统计分布的样本。每个像素都有i=32 样本。

我现在需要为 2D 图像上的每一个 2x2 ROI（重叠）计算 covariance matrix，从而得到一组总形状为 (4,4,2047,2047) 的 4x4 协方差矩阵。

循环遍历所有 ROI 是可能的，在我的机器上大约需要 4 分钟：

import numpy as np
arr = np.random.normal(1000,10,(32,2048,2048))
shape = arr.shape
result = np.zeros((4,4,shape[1]-1,shape[2]-1))
for i in range(shape[1]-1):
    for j in range(shape[2]-1):
         result[:,:,i,j] = np.cov(arr[:,i:i+2,j:j+2].reshape(32,4), rowvar=False, bias=True)

但是不使用 numpy 内置的索引和循环似乎效率低下。

那么有没有更优雅/更快的方法来做到这一点？

Answer 1

从 numpy 1.20.0+ 开始，sliding_window_view 允许您提取滑动窗口。然后就可以进行协变计算了：

from numpy.lib.stride_tricks import sliding_window_view
a = sliding_window_view(arr, (2,2), axis=(1,2)).reshape(32,2047,2047,-1)

X = a - a.mean(axis=0)

out = np.einsum('ijlk,ijlm,...->kmjl', X,X,1/32)

这在我的系统上花了大约 20 秒。