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::
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