Index function thresholded percentile (closes #85)

climix/etc/metadata.yml

@@ -130,6 +130,23 @@ index_functions:
       percentiles:
         kind: quantity
 
+  thresholded_percentile:
+    description: |
+      Calculates percentiles of data for which some condition is met.
+      First, the threshold is transformed to the same standard_name and units as
+      the input data.
+      Then the thresholding is performed as condition(data, threshold), ie
+      if condition is <, data < threshold.
+      Finally, the percentiles is calculated for those data values that fulfill
+      the condition. The default percentile method is "linear" option of numpy.
+    parameters:
+      threshold:
+        kind: quantity
+      condition:
+        kind: operator
+      percentiles:
+        kind: quantity
+
   temperature_sum:
     description: |
       Calculates the temperature sum above/below a threshold. First, the threshold

climix/index_functions/__init__.py

@@ -12,5 +12,6 @@ from .index_functions import (  # noqa: F401
     SpellLength,
     Statistics,
     ThresholdedStatistics,
+    ThresholdedPercentile,    
     TemperatureSum,
 )

climix/index_functions/index_functions.py

@@ -319,6 +319,48 @@ class Percentile(IndexFunction):
         return res.astype('float32')
 
 
+class ThresholdedPercentile(ThresholdMixin, IndexFunction):
+    def __init__(self, threshold, condition,
+                 percentiles, interpolation='linear'):
+        super().__init__(threshold, condition)
+        points = percentiles.points
+        assert np.all(points > 0)
+        assert np.all(points < 100)
+        self.percentiles = percentiles
+        self.interpolation = interpolation
+
+    def prepare(self, input_cubes):
+        super().prepare(input_cubes)
+        ref_cube = next(iter(input_cubes.values()))
+        self.standard_name = ref_cube.standard_name
+        self.units = ref_cube.units
+
+    def call_func(self, data, axis, **kwargs):
+        axis = normalize_axis(axis, data.ndim)
+        mask = np.ma.getmaskarray(data).any(axis=axis)
+        comb = self.condition(data, self.threshold.points)
+        res = np.percentile(np.ma.masked_where(~comb, data),
+                            q=self.percentiles.points, axis=axis,
+                            interpolation=self.interpolation)
+        res = np.ma.masked_array(da.ma.getdata(res), mask)
+        return res.astype('float32')
+
+    def lazy_func(self, data, axis, **kwargs):
+        axis = normalize_axis(axis, data.ndim)
+        mask = da.ma.getmaskarray(data).any(axis=axis)
+        comb = self.condition(data, self.threshold.points)
+
+        def percentile(arr):
+            return np.percentile(arr,
+                                 q=self.percentiles.points,
+                                 interpolation=self.interpolation)
+        res = da.apply_along_axis(
+            percentile, axis=axis,
+            arr=np.ma.masked_where(~comb, data)).squeeze()
+        res = da.ma.masked_array(da.ma.getdata(res), mask)
+        return res.astype('float32')
+
+
 class TemperatureSum(ThresholdMixin, IndexFunction):
     def __init__(self, threshold, condition):
         super().__init__(threshold, condition, units=Unit('days'))

setup.py

@@ -60,6 +60,7 @@ setuptools.setup(
             'statistics=climix.index_functions:Statistics',
             'percentile=climix.index_functions:Percentile',
             'thresholded_statistics=climix.index_functions:ThresholdedStatistics',
+            'thresholded_percentile=climix.index_functions:ThresholdedPercentile',            
             'temperature_sum=climix.index_functions:TemperatureSum',
             'diurnal_temperature_range=climix.index_functions:DiurnalTemperatureRange',
             'interday_diurnal_temperature_range=climix.index_functions:InterdayDiurnalTemperatureRange',