Index function thresholded running statistics (closes #157)

climix/etc/metadata.yml

@@ -159,6 +159,22 @@ index_functions:
       reducer:
         kind: reducer
 
+  thresholded_running_statistics:
+    description: |
+      First calculate a statistic within a moving window, then calculate a
+      different statistic across the moving windows.
+    parameters:
+      threshold:
+        kind: quantity
+      condition:
+        kind: operator
+      rolling_aggregator:
+        kind: reducer
+      window_size:
+        kind: quantity
+      reducer:
+        kind: reducer
+
   temperature_sum:
     description: |
       Calculates the temperature sum above/below a threshold. First, the threshold

climix/index_functions/__init__.py

@@ -13,6 +13,7 @@ from .index_functions import (  # noqa: F401
     ThresholdedPercentile,
     ThresholdedStatistics,
     RunningStatistics,
+    ThresholdedRunningStatistics,
     TemperatureSum,
 )
 

climix/index_functions/index_functions.py

@@ -432,6 +432,21 @@ class RunningStatistics(RollingWindowMixin, IndexFunction):
         return cube, res_data
 
 
+class ThresholdedRunningStatistics(ThresholdMixin, RunningStatistics):
+    def __init__(self, threshold, condition, rolling_aggregator, window_size, reducer):
+        super().__init__(threshold, condition, rolling_aggregator, window_size, reducer)
+
+    def call_func(self, data, axis, **kwargs):
+        comb = self.condition(data, self.threshold.points)
+        thresholded_data = np.where(comb, data, 0.0)
+        return super().call_func(thresholded_data, axis, **kwargs)
+
+    def lazy_func(self, data, axis, **kwargs):
+        comb = self.condition(data, self.threshold.points)
+        thresholded_data = da.where(comb, data, 0.0)
+        return super().call_func(thresholded_data, axis, **kwargs)
+
+
 class TemperatureSum(ThresholdMixin, IndexFunction):
     def __init__(self, threshold, condition):
         super().__init__(threshold, condition, units=Unit("days"))

setup.py

@@ -64,6 +64,7 @@ setuptools.setup(
             "statistics=climix.index_functions:Statistics",
             "temperature_sum=climix.index_functions:TemperatureSum",
             "thresholded_percentile=climix.index_functions:ThresholdedPercentile",
+            "thresholded_running_statistics=climix.index_functions:ThresholdedRunningStatistics",  # noqa: E501
             "thresholded_statistics=climix.index_functions:ThresholdedStatistics",
         ],
     },