I have a large numpy array of shape (32,2048,2048), where every [i,:,:] is a 2D set of pixels which are samples from a spatially correlated statistical distribution. With i=32 samples for each pixel.
I now need to calculate the covariance matrix for everey 2x2 ROI (overlapping) on the 2D image, resulting in a a set of 4x4 covariance matrices with total shape of (4,4,2047,2047).
Looping over all ROIs in a loop is of cause possible and takes about 4 minutes on my maschine:
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)
But indexing and looping without using numpy build-ins seems inefficient.
So is there a more elegant / faster way of doing this?
From numpy 1.20.0+, there is sliding_window_view that allows you to extract the sliding windows. Then you can perform the covariant calculation:
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)
This took about 20s on my system.