Numpy:计算大数组的协方差

Numpy: Calculate Covariance of large array

我有一个形状为 (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 内置的索引和循环似乎效率低下。

那么有没有更优雅/更快的方法呢?

从 numpy 1.20.0+ 开始,sliding_window_view 允许您提取滑动 windows。然后就可以进行协变计算了:

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 秒。