cess_date refactoring

enacts/calc.py

@@ -366,78 +366,47 @@ def cess_date_step(cess_yesterday, dry_spell_length, dry_spell_length_thresh):
     ) - np.timedelta64(1, "D")
 
 
-def cess_date(
-    daily_data,
-    is_soil_moisture,
+def cess_date_from_sm(
+    daily_sm,
     dry_thresh, 
     dry_spell_length_thresh,
-    et=None,
-    taw=None,
-    sminit=None,
     time_dim="T",
 ):
-    """Calculate cessation date.
+    def sm_func(_, t):
+        return daily_sm.sel({time_dim: t})
 
-    Find first day of the first dry spell where: 
-        Soil moisture falls below `dry_thresh` for `min_dry_days` days.
+    return _cess_date(dry_thresh, dry_spell_length_thresh, sm_func, daily_sm[time_dim])
 
-    Parameters
-    ----------
-    daily_data : DataArray
-        Array of daily soil moisture or daily rainfall.
-    is_soil_moisture : boolean
-        True if `daily_data` is soil moisture, False if it is rainfall.
-    dry_thresh : float
-        Soil moisture threshold to determine a dry day.
-    dry_spell_length_thresh : int
-        Minimum number of dry days in a row to be considered a dry spell.
-    et : DataArray, optional
-        Evapotranspiration used by `water_balance` if `daily_data` is rainfall.
-        Can be a single value with no dimensions or axes.
-    taw : DataArray, optional
-        Total available water used by `water_balance` if `daily_data` is rainfall.
-        Can be a single value with no dimensions or axes.
-    sminit : DataArray, optional
-        Soil moisture initialization used by `water_balance` if `daily_data` is rainfall.
-        If DataArray, must not have `time_dim` dim.
-        Can be a single value with no dimensions or axes.
-    time_dim : str, optional
-        Time coordinate in `soil_moisture` (default `time_dim`="T"). 
-
-    Returns
-    -------
-    cess_delta : DataArray[np.timedelta64] 
-        Difference between first day of `daily_data` 
-        and cessation date.
-
-    See Also
-    --------
-    water_balance
 
-    Notes
-    -----
-    Examples
-    --------
-    """
-    # Initializing
-    if is_soil_moisture:
-        sm = daily_data.isel({time_dim: 0})
-    else:
-        if sminit is None or et is None or taw is None:
-            raise Exception("sminit, et and taw can not be None")
-        else:
-            sminit = xr.DataArray(sminit)
-            et = xr.DataArray(et)
-            taw = xr.DataArray(taw)
-            sm = water_balance_step(sminit, daily_data.isel({time_dim: 0}), et, taw)
-    spell_length = (sm < dry_thresh).astype("timedelta64[D]")
+def cess_date_from_rain(
+    daily_rain,
+    dry_thresh, 
+    dry_spell_length_thresh,
+    et,
+    taw,
+    sminit,
+    time_dim="T",
+):
+    sminit = xr.DataArray(sminit)
+    et = xr.DataArray(et)
+    taw = xr.DataArray(taw)
+
+    def sm_func(sm, t):
+        if sm is None:
+            sm = sminit
+        return water_balance_step(
+            sm, daily_rain.sel({time_dim: t}), et, taw
+        )
+
+    return _cess_date(dry_thresh, dry_spell_length_thresh, sm_func, daily_rain[time_dim])
+
+
+def _cess_date(dry_thresh, dry_spell_length_thresh, sm_func, time_coord):
+    spell_length = np.timedelta64(0)
     cess_delta = xr.DataArray(np.timedelta64("NaT", "D"))
-    # Loop
-    for t in daily_data[time_dim][1:]:
-        if is_soil_moisture:
-            sm = daily_data.sel({time_dim: t})
-        else:
-            sm = water_balance_step(sm, daily_data.sel({time_dim: t}), et, taw)
+    sm = None
+    for t in time_coord:
+        sm = sm_func(sm, t)
         dry_day = sm < dry_thresh
         spell_length = (spell_length + dry_day.astype("timedelta64[D]")) * dry_day
         cess_delta = cess_date_step(
@@ -447,9 +416,9 @@ def cess_date(
         )
     # Delta reference (and coordinate) back to first time point of daily_data
     cess_delta = (
-        daily_data[time_dim][-1]
+        time_coord[-1]
         + cess_delta
-        - daily_data[time_dim][0].expand_dims(dim=time_dim)
+        - time_coord[0].expand_dims(dim=time_coord.name)
     )
     return cess_delta
 
@@ -819,8 +788,7 @@ def seasonal_cess_date(
         grouped_daily_data[soil_moisture.name]
         .groupby(grouped_daily_data["seasons_starts"])
         .map(
-            cess_date,
-            is_soil_moisture=True,
+            cess_date_from_sm,
             dry_thresh=dry_thresh,
             dry_spell_length_thresh=dry_spell_length_thresh,
         )

enacts/tests/test_calc.py

@@ -546,7 +546,7 @@ def test_cess_date():
          [0, 0, 2, 0, 0]],
         dims=["X", "T"], coords={"T": t}
     )
-    cess_delta = calc.cess_date(daily_sm, True, 1, 3)
+    cess_delta = calc.cess_date_from_sm(daily_sm, 1, 3)
     expected = xr.DataArray(
         [np.nan, 0, np.nan, np.nan, 2, 1, 0, np.nan]
     ).astype("timedelta64[D]")
@@ -569,9 +569,8 @@ def test_cess_date_rain():
          [5, 5, 7, 3, 5]],
         dims=["X", "T"], coords={"T": t}
     )
-    cess_delta = calc.cess_date(
+    cess_delta = calc.cess_date_from_rain(
         daily_rain,
-        is_soil_moisture=False,
         dry_thresh=1,
         dry_spell_length_thresh=3,
         et=5,
@@ -587,9 +586,8 @@ def test_cess_date_rain():
 
 
 def call_cess_date(data):
-    cessations = calc.cess_date(
-        daily_data=data,
-        is_soil_moisture=True,
+    cessations = calc.cess_date_from_sm(
+        daily_sm=data,
         dry_thresh=5,
         dry_spell_length_thresh=3,
     )