Rolling mean over one axis in 3D-array

user_12221937 0
Viewed 520

Is there a simple way to calculate the rolling mean over one axis in a 3D-array? Let's say I have an array with x,y and time axis and I want the rolling mean over the time axis for all x and y. For 1D-arrays I use pandas:

import pandas as pd 

rolling_array = pd.Series(array).rolling(window=window).mean()

But this does not work for multidimensional data.

EDIT:

My array looks like this:

import numpy as np 

array = np.random.rand(100,100,200)

And I want the rolling mean over axis = 2

pandas numpy arrays python
1 Answers

We can use uniform_filter1d that accepts axis arg and we will make it generic to accept any n-dim array along a generic axis -

from scipy.ndimage import uniform_filter1d

def rolling_mean_along_axis(a, W, axis=-1):
    # a : Input ndarray
    # W : Window size
    # axis : Axis along which we will apply rolling/sliding mean
    hW = W//2
    L = a.shape[axis]-W+1   
    indexer = [slice(None) for _ in range(a.ndim)]
    indexer[axis] = slice(hW,hW+L)
    return uniform_filter1d(a,W,axis=axis)[tuple(indexer)]

Sample run to verify shapes :

In [70]: a = np.random.rand(10,10,10)

In [72]: rolling_mean_along_axis(a, W=5, axis=0).shape
Out[72]: (6, 10, 10)

In [73]: rolling_mean_along_axis(a, W=5, axis=1).shape
Out[73]: (10, 6, 10)

In [74]: rolling_mean_along_axis(a, W=5, axis=2).shape
Out[74]: (10, 10, 6)
user_32938810
Related