function to count extreme events

pepsico/app_calc.py

@@ -1,4 +1,5 @@
 import xarray as xr
+import operator as op
 
 
 #This is what we should need for the app
@@ -143,3 +144,56 @@ def strftimeb2int(strftimeb):
     }
     strftimebint = strftimeb_all[strftimeb]
     return strftimebint
+
+
+def number_extreme_events_within_days(
+    daily_data, operator, threshold, window, dim="T"
+):
+    """Count extreme events
+
+    Extreme events constitute cumulative `daily_data` over `window` -day
+    (or shorther) verifying the `operator` `threshold`
+
+    Parameters
+    ----------
+    daily_data : DataArray
+        Array of daily data
+    operator : str
+        "lt", "le", "gt", "ge"
+        determines the condition to meet with respect to `threshold`
+    threshold : float
+        threshold to meet in `daily_data` units
+    window : int
+        maximum length of days to cumulate over to constitute an event
+    dim : str, optional
+        name of `daily_data` time dimension
+
+    Returns
+    -------
+    DataArray
+        Array of count of extreme events
+
+    Notes
+    -----
+    Is meant to be used to produce yearly seasonal time series
+
+    Examples
+    --------
+    Number of rain events of >80mm in 2 days or less:
+    number_extreme_events_within_days(rain_daily_data_in_mm, "gt", 80, 2)    
+    """
+    count = 0
+    #Start with shortest events
+    for w in range(1, window+1):
+        for t in range(len(daily_data[dim])-(w-1)):
+            #Assert new event
+            new_event = getattr(op, operator)(
+                daily_data.isel({dim: slice(t, t+w)}).sum(dim), threshold
+            )
+            #Mask days having formed a new event so that they won't account again
+            #for following events (t loop) or longer events (w loop)
+            daily_data[{dim: slice(t, t+w)}] = (
+                daily_data[{dim: slice(t, t+w)}].where(~new_event)
+            )
+            count = count + new_event
+    return count

pepsico/test/test_app_calc.py

@@ -0,0 +1,14 @@
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+import app_calc
+
+def test_number_extreme_events_within_days():
+    t = pd.date_range(start="2000-05-01", end="2000-05-10", freq="1D")
+    values = 1 + 0*np.arange(t.size)
+    precip = xr.DataArray(values, coords={"T": t})
+    count = app_calc.number_extreme_events_within_days(precip, "gt", 2, 3)
+
+    np.testing.assert_array_equal(count, [3])
+