rolling statistics for numpy 2D arrays with strides

strides_2d.py

import numpy as np
from numpy.lib.stride_tricks import as_strided

#
# 1-D sample
# based on: http://www.rigtorp.se/2011/01/01/rolling-statistics-numpy.html
def strides_1d(a, r):

    ax = np.zeros(shape=(a.shape[0]+2*r))#, dtype=np.uint8) # `a` extended
    ax[:] = np.nan
    ax[r:ax.shape[0]-r] = a
    
    s = as_strided(ax, shape=(a.shape + (1+2*r,)), strides=(ax.strides[0], ax.strides[0]))
    return s

#----------
r = 1
a = np.arange(7, dtype=np.uint16)
strides_1d(a,r)


# ----------------------------
# 2-D sample


def strides_2d(a, r, linear=True):
    
    ax = np.zeros(shape=(a.shape[0] + 2*r[0], a.shape[1] + 2*r[1]))
    ax[:] = np.nan
    ax[r[0]:ax.shape[0]-r[0], r[1]:ax.shape[1]-r[1]] = a
    
    shape = a.shape + (1+2*r[0], 1+2*r[1])
    strides = ax.strides + ax.strides
    s = as_strided(ax, shape=shape, strides=strides)
    
    return s.reshape(a.shape + (shape[2]*shape[3],)) if linear else s


#----------
r = [1,3] # radii
a = np.arange(3*4).reshape(3,4)
b =  strides_2d(a, r)
b = np.ma.masked_array(b, mask=np.isnan(b))

# moving average and standard deviation with radii `r` of matrix `a`
b.std(axis=-1)
np.ma.mean(b)