Index function first occurrence

climix/etc/metadata.yml

@@ -20,6 +20,20 @@ index_functions:
       condition:
         kind: operator
 
+  first_occurrence:
+    description: |
+      Calculates the first time 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, locate the first occurrence when condition is met.
+    parameters:
+      threshold:
+        kind: quantity
+      condition:
+        kind: operator
+
   spell_length:
     description: |
       Calculates statistics on lengths of spells.

climix/index_functions.py

 # -*- coding: utf-8 -*-
 
+from datetime import datetime
+
 from cf_units import Unit
 import dask.array as da
 import numpy as np
@@ -82,6 +84,51 @@ class CountOccurrences:
         return res.astype('float32')
 
 
+class FirstOccurrence:
+    def __init__(self, threshold, condition):
+        self.threshold = threshold
+        self.condition = NUMPY_OPERATORS[condition]
+        self.lazy_condition = DASK_OPERATORS[condition]
+        self.standard_name = None
+        self.units = Unit('days')
+        self.NO_OCCURRENCE = np.inf
+
+    def prepare(self, input_cube):
+        change_units(self.threshold, input_cube.units, input_cube.standard_name)
+
+    def call_func(self, data, axis, **kwargs):
+        axis = normalize_axis(axis, data.ndim)
+        cond = self.condition(data, self.threshold.data)
+        res = np.ma.where(cond.any(axis=axis),
+                          cond.argmax(axis=axis),
+                          self.NO_OCCURRENCE)
+        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)
+        cond = self.lazy_condition(data, self.threshold.data)
+        res = da.where(cond.any(axis=axis),
+                       cond.argmax(axis=axis),
+                       self.NO_OCCURRENCE)
+        res = da.ma.masked_array(da.ma.getdata(res), mask)
+        return res.astype('float32')
+
+    def post_process(self, collapsed_cube, data_result, coords,
+                     period, **kwargs):
+        time = collapsed_cube.coord('time')
+        calendar = time.units.calendar
+        offsets = np.empty_like(time.points, dtype=data_result.dtype)
+        for i, representative_date in enumerate(time.cells()):
+            year = representative_date.point.year
+            start_date = datetime(year, period.first_month_number, 1)
+            units = Unit('days since {}-01-01'.format(year), calendar=calendar)
+            offsets[i] = units.date2num(start_date)
+        collapsed_cube.data = (collapsed_cube.core_data()
+                               + offsets[:, None, None])
+        return collapsed_cube
+
+
 class SpellLength:
     def __init__(self, threshold, condition, reducer):
         self.threshold = threshold

setup.py

@@ -50,6 +50,7 @@ setuptools.setup(
         ],
         'climix.index_functions': [
             'count_occurrences=climix.index_functions:CountOccurrences',
+            'first_occurrence=climix.index_functions:FirstOccurrence',
             'spell_length=climix.index_functions:SpellLength',
             'thresholded_statistics=climix.index_functions:ThresholdedStatistics',
         ],