add rolling_sum

hdc/algo/accessors.py

@@ -1,5 +1,5 @@
 """Xarray Accesor classes."""
-from typing import List, Optional, Union
+from typing import Iterable, List, Optional, Union
 from warnings import warn
 
 from dask import is_dask_collection
@@ -444,7 +444,7 @@ def spi(
         self,
         calibration_start: Optional[str] = None,
         calibration_stop: Optional[str] = None,
-        groups: Optional[List[int]] = None,
+        groups: Optional[Iterable[int]] = None,
     ):
         """Calculate the SPI along the time dimension.
 
@@ -460,10 +460,11 @@ def spi(
         if not self._check_for_timedim():
             raise MissingTimeError("SPI requires a time dimension!")
 
+        # pylint: disable=import-outside-toplevel
         from .ops.stats import (
             gammastd_yxt,
             gammastd_grp,
-        )  # pylint: disable=import-outside-toplevel
+        )
 
         tix = self._obj.get_index("time")
 

hdc/algo/ops/stats.py

@@ -240,8 +240,8 @@ def gammastd_grp(xx, groups, num_groups, cal_start, cal_stop, yy):
 
     This calculates gammastd across xx for indivual groups
     defined in `groups`. These need to be in ascending order from
-    0 to num_groups - 1."""
-
+    0 to num_groups - 1.
+    """
     for grp in range(num_groups):
         grp_ix = groups == grp
         pix = xx[grp_ix]
@@ -492,3 +492,17 @@ def mean_grp(xx, groups, num_groups, nodata, yy):
             avg = avg / n
 
         yy[grp_ix] = avg
+
+
+@guvectorize(["(float32[:], float64, float32[:])"], "(n),() -> (n)")
+def rolling_sum(xx, window_size, yy):
+    """Calculate moving window sum over specified size."""
+    n = xx.size
+    yy[:] = 0.0
+    for ii in range(n):
+        if ii - window_size + 1 < 0:
+            yy[ii] = 0.0
+            continue
+
+        for jj in range(ii - window_size + 1, ii + 1):
+            yy[ii] += xx[jj]

tests/test_algos.py

@@ -30,6 +30,7 @@
     mk_variance_s,
     mk_sens_slope,
     mann_kendall_trend_1d,
+    rolling_sum,
 )
 
 
@@ -189,6 +190,16 @@ def test_mean_grp(ts):
     assert (res == 1.02697).all()
 
 
+def test_rolling_sum():
+    res = rolling_sum(np.arange(10).astype("float32"), 3)
+    np.testing.assert_almost_equal(
+        res,
+        np.array(
+            [0.0, 0.0, 3.0, 6.0, 9.0, 12.0, 15.0, 18.0, 21.0, 24.0], dtype="float32"
+        ),
+    )
+
+
 def test_ws2dgu(ts):
     _ts = ts * 10
     z = ws2dgu(_ts, 10, 0)