Python Forum

Full Version: Apply rolling window function over time dimension of 3D data
You're currently viewing a stripped down version of our content. View the full version with proper formatting.
import xarray as xr
import numpy as np
def get_grps(s, thresh=-1, Nmin=2):
    """
    Nmin : int > 0
        Min number of consecutive values below threshold.
    """
    s = pd.Series(s)
    m = np.logical_and.reduce([s.shift(-i).le(thresh) for i in range(Nmin)])
    if Nmin > 1:
        m = pd.Series(m, index=s.index).replace({False: np.NaN}).ffill(limit=Nmin-1).fillna(False)
    else:
        m = pd.Series(m, index=s.index)
 
    # Form consecutive groups
    gps = m.ne(m.shift(1)).cumsum().where(m)
 
    # Return None if no groups, else the aggregations
    if gps.isnull().all():
        return [0]
    else:
        d = s.groupby(gps).agg([list, sum, 'size']).reset_index(drop=True)
 
        data = len((d['sum']))  ### returning frequency
 
        data = np.array(data)
 
 
 
        return (data)
 
### create some dummy data::
 
spi = xr.DataArray(dims=("time", "lon", "lat"), data=np.random.randn(324, 180, 360))
 
### Apply xarray rolling window function::
 
spi.rolling(time=59).reduce(get_grps)
 
## then I get this error::
 
Traceback (most recent call last):
......
    raise Exception("Data must be 1-dimensional")
Exception: Data must be 1-dimensional
 
How can I modify this function to accept 3D data and then apply xarray rolling window properly