modify code to accept 3d array

[Staph]

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