From 567837023a97ccf9082c01f1ea5aa129274772d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Cousin?= <remic@iri.columbia.edu>
Date: Wed, 31 May 2023 14:52:40 -0400
Subject: [PATCH 01/13] rain of option for cess_date and tests rearranged

---
 enacts/{wat_bal => }/agronomy.py            |  0
 enacts/calc.py                              | 23 ++++++++++++++++-----
 enacts/{onset => }/tests/data_test_calc.py  |  0
 enacts/{wat_bal => }/tests/test_agronomy.py |  0
 enacts/{onset => }/tests/test_calc.py       |  3 ++-
 enacts/wat_bal/maproom_monit.py             |  2 +-
 6 files changed, 21 insertions(+), 7 deletions(-)
 rename enacts/{wat_bal => }/agronomy.py (100%)
 rename enacts/{onset => }/tests/data_test_calc.py (100%)
 rename enacts/{wat_bal => }/tests/test_agronomy.py (100%)
 rename enacts/{onset => }/tests/test_calc.py (99%)

diff --git a/enacts/wat_bal/agronomy.py b/enacts/agronomy.py
similarity index 100%
rename from enacts/wat_bal/agronomy.py
rename to enacts/agronomy.py
diff --git a/enacts/calc.py b/enacts/calc.py
index 7fa229422..e250bccb5 100644
--- a/enacts/calc.py
+++ b/enacts/calc.py
@@ -367,7 +367,8 @@ def cess_date_step(cess_yesterday, dry_spell_length, dry_spell_length_thresh):
 
 
 def cess_date(
-    daily_data, 
+    daily_data,
+    is_soil_moisture,
     dry_thresh, 
     dry_spell_length_thresh,
     et=None,
@@ -384,6 +385,8 @@ def cess_date(
     ----------
     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
@@ -417,13 +420,22 @@ def cess_date(
     --------
     """
     # Initializing
-    spell_length = (
-        daily_data.isel({time_dim: 0}) < dry_thresh
-    ).astype("timedelta64[D]")
+    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:
+            sm = water_balance_step(sminit, daily_data.isel({time_dim: 0}), et, taw)
+    spell_length = (sm < dry_thresh).astype("timedelta64[D]")
     cess_delta = xr.DataArray(np.timedelta64("NaT", "D"))
     # Loop
     for t in daily_data[time_dim][1:]:
-        dry_day = daily_data.sel({time_dim: t}) < dry_thresh
+        if is_soil_moisture:
+            sm = daily_data.sel({time_dim: t})
+        else:
+            sm = water_balance_step(sm, daily_data.isel({time_dim: t}), et, taw)
+        dry_day = sm < dry_thresh
         spell_length = (spell_length + dry_day.astype("timedelta64[D]")) * dry_day
         cess_delta = cess_date_step(
             cess_delta,
@@ -805,6 +817,7 @@ def seasonal_cess_date(
         .groupby(grouped_daily_data["seasons_starts"])
         .map(
             cess_date,
+            is_soil_moisture=True,
             dry_thresh=dry_thresh,
             dry_spell_length_thresh=dry_spell_length_thresh,
         )
diff --git a/enacts/onset/tests/data_test_calc.py b/enacts/tests/data_test_calc.py
similarity index 100%
rename from enacts/onset/tests/data_test_calc.py
rename to enacts/tests/data_test_calc.py
diff --git a/enacts/wat_bal/tests/test_agronomy.py b/enacts/tests/test_agronomy.py
similarity index 100%
rename from enacts/wat_bal/tests/test_agronomy.py
rename to enacts/tests/test_agronomy.py
diff --git a/enacts/onset/tests/test_calc.py b/enacts/tests/test_calc.py
similarity index 99%
rename from enacts/onset/tests/test_calc.py
rename to enacts/tests/test_calc.py
index 71ac18c4e..55fa21c5a 100644
--- a/enacts/onset/tests/test_calc.py
+++ b/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, 1, 3)
+    cess_delta = calc.cess_date(daily_sm, True, 1, 3)
     expected = xr.DataArray(
         [np.nan, 0, np.nan, np.nan, 2, 1, 0, np.nan]
     ).astype("timedelta64[D]")
@@ -558,6 +558,7 @@ def test_cess_date():
 def call_cess_date(data):
     cessations = calc.cess_date(
         daily_data=data,
+        is_soil_moisture=True,
         dry_thresh=5,
         dry_spell_length_thresh=3,
     )
diff --git a/enacts/wat_bal/maproom_monit.py b/enacts/wat_bal/maproom_monit.py
index 731980c43..781d58958 100644
--- a/enacts/wat_bal/maproom_monit.py
+++ b/enacts/wat_bal/maproom_monit.py
@@ -17,7 +17,7 @@
 import datetime
 
 import xarray as xr
-from . import agronomy as ag
+import agronomy as ag
 
 import psycopg2
 from psycopg2 import sql

From 0dfcf363c28ab04f9fb1e533c088209881834f76 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Cousin?= <remic@iri.columbia.edu>
Date: Thu, 1 Jun 2023 10:28:37 -0400
Subject: [PATCH 02/13] test for cess_date driven by rain

---
 enacts/calc.py            |  2 +-
 enacts/tests/test_calc.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/enacts/calc.py b/enacts/calc.py
index e250bccb5..465c2c45c 100644
--- a/enacts/calc.py
+++ b/enacts/calc.py
@@ -434,7 +434,7 @@ def cess_date(
         if is_soil_moisture:
             sm = daily_data.sel({time_dim: t})
         else:
-            sm = water_balance_step(sm, daily_data.isel({time_dim: t}), et, taw)
+            sm = water_balance_step(sm, daily_data.sel({time_dim: t}), et, taw)
         dry_day = sm < dry_thresh
         spell_length = (spell_length + dry_day.astype("timedelta64[D]")) * dry_day
         cess_delta = cess_date_step(
diff --git a/enacts/tests/test_calc.py b/enacts/tests/test_calc.py
index 55fa21c5a..367642793 100644
--- a/enacts/tests/test_calc.py
+++ b/enacts/tests/test_calc.py
@@ -555,6 +555,37 @@ def test_cess_date():
     assert np.array_equal(cess_delta.squeeze("T"), expected, equal_nan=True)
 
 
+def test_cess_date_rain():
+
+    t = pd.date_range(start="2000-05-01", end="2000-05-05", freq="1D")
+    daily_rain = xr.DataArray(
+        [[7, 5, 5, 5, 5],
+         [5, 5, 5, 5, 7],
+         [7, 5, 5, 5, 3],
+         [7, 5 ,5 ,3 ,5],
+         [7, 5, 3, 5 ,5],
+         [7, 3, 5, 5, 5],
+         [5, 5, 5, 5, 5],
+         [5, 5, 7, 3, 5]],
+        dims=["X", "T"], coords={"T": t}
+    )
+    cess_delta = calc.cess_date(
+        daily_rain,
+        False,
+        1,
+        3,
+        et=xr.DataArray(5),
+        taw=xr.DataArray(10),
+        sminit=xr.DataArray(0),
+    )
+    expected = xr.DataArray(
+        [np.nan, 0, np.nan, np.nan, 2, 1, 0, np.nan]
+    ).astype("timedelta64[D]")
+
+    assert cess_delta["T"] == daily_rain["T"][0]
+    assert np.array_equal(cess_delta.squeeze("T"), expected, equal_nan=True)
+
+
 def call_cess_date(data):
     cessations = calc.cess_date(
         daily_data=data,

From d480c31d8c6ef065330fd0b499e20955326ec68b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Cousin?= <remic@iri.columbia.edu>
Date: Thu, 1 Jun 2023 10:40:19 -0400
Subject: [PATCH 03/13] user can input numbers

---
 enacts/calc.py            | 3 +++
 enacts/tests/test_calc.py | 2 +-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/enacts/calc.py b/enacts/calc.py
index 465c2c45c..b12dcd451 100644
--- a/enacts/calc.py
+++ b/enacts/calc.py
@@ -426,6 +426,9 @@ def cess_date(
         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]")
     cess_delta = xr.DataArray(np.timedelta64("NaT", "D"))
diff --git a/enacts/tests/test_calc.py b/enacts/tests/test_calc.py
index 367642793..638efa9ab 100644
--- a/enacts/tests/test_calc.py
+++ b/enacts/tests/test_calc.py
@@ -576,7 +576,7 @@ def test_cess_date_rain():
         3,
         et=xr.DataArray(5),
         taw=xr.DataArray(10),
-        sminit=xr.DataArray(0),
+        sminit=0,
     )
     expected = xr.DataArray(
         [np.nan, 0, np.nan, np.nan, 2, 1, 0, np.nan]

From 2f8ac462be4bfdc064045e9c714bff6258ed703c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Cousin?= <remic@iri.columbia.edu>
Date: Thu, 13 Jul 2023 17:45:35 -0400
Subject: [PATCH 04/13] more readable tests

---
 enacts/tests/test_calc.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/enacts/tests/test_calc.py b/enacts/tests/test_calc.py
index 638efa9ab..fb6ee0627 100644
--- a/enacts/tests/test_calc.py
+++ b/enacts/tests/test_calc.py
@@ -571,11 +571,11 @@ def test_cess_date_rain():
     )
     cess_delta = calc.cess_date(
         daily_rain,
-        False,
-        1,
-        3,
-        et=xr.DataArray(5),
-        taw=xr.DataArray(10),
+        is_soil_moisture=False,
+        dry_thresh=1,
+        dry_spell_length_thresh=3,
+        et=5,
+        taw=10,
         sminit=0,
     )
     expected = xr.DataArray(

From 95e6cba34eff5c283fcf835f7eb88c8509e47221 Mon Sep 17 00:00:00 2001
From: Aaron Kaplan <aaron@iri.columbia.edu>
Date: Thu, 27 Jul 2023 15:07:20 -0400
Subject: [PATCH 05/13] Remove unnecessary nesting

---
 enacts/calc.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/enacts/calc.py b/enacts/calc.py
index b12dcd451..7f0ba7efa 100644
--- a/enacts/calc.py
+++ b/enacts/calc.py
@@ -425,11 +425,10 @@ def cess_date(
     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)
+        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]")
     cess_delta = xr.DataArray(np.timedelta64("NaT", "D"))
     # Loop

From 9a2fb67b90af1898072999d3d128c10f5a13a2eb Mon Sep 17 00:00:00 2001
From: Aaron Kaplan <aaron@iri.columbia.edu>
Date: Thu, 27 Jul 2023 15:35:01 -0400
Subject: [PATCH 06/13] Split cess_date_from_rain and cess_date_from_sm

---
 enacts/calc.py            | 100 ++++++++++++++------------------------
 enacts/tests/test_calc.py |   8 ++-
 2 files changed, 39 insertions(+), 69 deletions(-)

diff --git a/enacts/calc.py b/enacts/calc.py
index 7f0ba7efa..c7a7600e5 100644
--- a/enacts/calc.py
+++ b/enacts/calc.py
@@ -366,77 +366,50 @@ def cess_date_step(cess_yesterday, dry_spell_length, dry_spell_length_thresh):
     ) - np.timedelta64(1, "D")
 
 
-def cess_date(
+def cess_date_from_sm(
     daily_data,
-    is_soil_moisture,
     dry_thresh, 
     dry_spell_length_thresh,
-    et=None,
-    taw=None,
-    sminit=None,
     time_dim="T",
 ):
-    """Calculate cessation date.
-
-    Find first day of the first dry spell where: 
-        Soil moisture falls below `dry_thresh` for `min_dry_days` days.
+    spell_length = np.timedelta64(0)
+    cess_delta = xr.DataArray(np.timedelta64("NaT", "D"))
+    # Loop
+    for t in daily_data[time_dim]:
+        sm = daily_data.sel({time_dim: t})
+        dry_day = sm < dry_thresh
+        spell_length = (spell_length + dry_day.astype("timedelta64[D]")) * dry_day
+        cess_delta = cess_date_step(
+            cess_delta,
+            spell_length,
+            np.timedelta64(dry_spell_length_thresh, "D"),
+        )
+    # Delta reference (and coordinate) back to first time point of daily_data
+    cess_delta = (
+        daily_data[time_dim][-1]
+        + cess_delta
+        - daily_data[time_dim][0].expand_dims(dim=time_dim)
+    )
+    return cess_delta
 
-    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
+def cess_date_from_rain(
+    daily_data,
+    dry_thresh, 
+    dry_spell_length_thresh,
+    et,
+    taw,
+    sminit,
+    time_dim="T",
+):
+    sm = xr.DataArray(sminit)
+    et = xr.DataArray(et)
+    taw = xr.DataArray(taw)
 
-    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")
-        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]")
+    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)
+    for t in daily_data[time_dim]:
+        sm = water_balance_step(sm, daily_data.sel({time_dim: t}), et, taw)
         dry_day = sm < dry_thresh
         spell_length = (spell_length + dry_day.astype("timedelta64[D]")) * dry_day
         cess_delta = cess_date_step(
@@ -818,8 +791,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,
         )
diff --git a/enacts/tests/test_calc.py b/enacts/tests/test_calc.py
index fb6ee0627..734c65735 100644
--- a/enacts/tests/test_calc.py
+++ b/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(
+    cessations = calc.cess_date_from_sm(
         daily_data=data,
-        is_soil_moisture=True,
         dry_thresh=5,
         dry_spell_length_thresh=3,
     )

From e6104a0c7f1bd69451bfb319fea062a89d84af8e Mon Sep 17 00:00:00 2001
From: Aaron Kaplan <aaron@iri.columbia.edu>
Date: Thu, 27 Jul 2023 16:27:56 -0400
Subject: [PATCH 07/13] Factor out code duplicated between cess_date_from_*

---
 enacts/calc.py            | 51 ++++++++++++++++++---------------------
 enacts/tests/test_calc.py |  2 +-
 2 files changed, 25 insertions(+), 28 deletions(-)

diff --git a/enacts/calc.py b/enacts/calc.py
index c7a7600e5..e290ec380 100644
--- a/enacts/calc.py
+++ b/enacts/calc.py
@@ -367,33 +367,19 @@ def cess_date_step(cess_yesterday, dry_spell_length, dry_spell_length_thresh):
 
 
 def cess_date_from_sm(
-    daily_data,
+    daily_sm,
     dry_thresh, 
     dry_spell_length_thresh,
     time_dim="T",
 ):
-    spell_length = np.timedelta64(0)
-    cess_delta = xr.DataArray(np.timedelta64("NaT", "D"))
-    # Loop
-    for t in daily_data[time_dim]:
-        sm = daily_data.sel({time_dim: t})
-        dry_day = sm < dry_thresh
-        spell_length = (spell_length + dry_day.astype("timedelta64[D]")) * dry_day
-        cess_delta = cess_date_step(
-            cess_delta,
-            spell_length,
-            np.timedelta64(dry_spell_length_thresh, "D"),
-        )
-    # Delta reference (and coordinate) back to first time point of daily_data
-    cess_delta = (
-        daily_data[time_dim][-1]
-        + cess_delta
-        - daily_data[time_dim][0].expand_dims(dim=time_dim)
-    )
-    return cess_delta
+    def sm_func(_, t):
+        return daily_sm.sel({time_dim: t})
+
+    return _cess_date(dry_thresh, dry_spell_length_thresh, sm_func, daily_sm[time_dim])
+
 
 def cess_date_from_rain(
-    daily_data,
+    daily_rain,
     dry_thresh, 
     dry_spell_length_thresh,
     et,
@@ -401,15 +387,26 @@ def cess_date_from_rain(
     sminit,
     time_dim="T",
 ):
-    sm = xr.DataArray(sminit)
+    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]:
-        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(
@@ -419,9 +416,9 @@ def cess_date_from_rain(
         )
     # 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
 
diff --git a/enacts/tests/test_calc.py b/enacts/tests/test_calc.py
index 734c65735..8cb0dc377 100644
--- a/enacts/tests/test_calc.py
+++ b/enacts/tests/test_calc.py
@@ -587,7 +587,7 @@ def test_cess_date_rain():
 
 def call_cess_date(data):
     cessations = calc.cess_date_from_sm(
-        daily_data=data,
+        daily_sm=data,
         dry_thresh=5,
         dry_spell_length_thresh=3,
     )

From 07766f44f83604838ca06554b8f20bd6b3a8a4ff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Cousin?= <remic@iri.columbia.edu>
Date: Wed, 2 Aug 2023 11:09:41 -0400
Subject: [PATCH 08/13] docs for the 2 new cess_date functions

---
 enacts/calc.py | 72 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 72 insertions(+)

diff --git a/enacts/calc.py b/enacts/calc.py
index e290ec380..405db3713 100644
--- a/enacts/calc.py
+++ b/enacts/calc.py
@@ -372,6 +372,32 @@ def cess_date_from_sm(
     dry_spell_length_thresh,
     time_dim="T",
 ):
+    """Calculate cessation date from daily soil moisture.
+
+    A cessation date is found at the first day of the first dry spell
+    where a dry spell is defined as at least `dry_spell_length_thresh` consecutive days
+    of soil moisture `daily_sm` less than `dry_thresh` .
+    
+    Parameters
+    ----------
+    daily_sm : DataArray
+        Array of daily soil moisture values.
+    dry_thresh : float
+        A dry day is when soil mositure `daily_sm` is lesser than `dry_thresh` .
+    dry_spell_length_thresh : int
+        A dry spell is at least `dry_spell_length_thresh` dry days.
+    time_dim : str, optional
+        Time coordinate in `daily_sm` (default `time_dim` ="T").   
+
+    Returns
+    -------
+    cess_delta : DataArray[np.timedelta64]
+        Difference between first day of `daily_sm` and cessation date.
+
+    See Also
+    --------
+    cess_date_step, cess_date_from_rain
+    """
     def sm_func(_, t):
         return daily_sm.sel({time_dim: t})
 
@@ -387,6 +413,52 @@ def cess_date_from_rain(
     sminit,
     time_dim="T",
 ):
+    """Calculate cessation date from daily rainfall.
+
+    A cessation date is found at the first day of the first dry spell
+    where a dry spell is defined as at least `dry_spell_length_thresh` consecutive days
+    of soil moisture less than `dry_thresh` .
+    Soil moisture is estimated through a simple water balance model
+    driven by rainfall `daily_rain` , evapotranspiration `et`, soil capacity `taw`
+    and initislized at soil moisture value `sminit` .
+    
+    Parameters
+    ----------
+    daily_rain : DataArray
+        Array of daily rainfall.
+    dry_thresh : float
+        A dry day is when soil mositure is lesser than `dry_thresh` .
+    dry_spell_length_thresh : int
+        A dry spell is at least `dry_spell_length_thresh` dry days.
+    et : DataArray
+        Evapotranspiration.
+    taw : DataArray
+        Total available water.
+    sminit : DataArray
+        Soil moisture initialization. If DataArray, must not have `time_dim` dim.
+    time_dim : str, optional
+        Time coordinate in `daily_rain` (default `time_dim` ="T").   
+
+    Returns
+    -------
+    cess_delta : DataArray[np.timedelta64]
+        Difference between first day of `daily_rain` and cessation date.
+
+    See Also
+    --------
+    cess_date_step, water_balance_step, cess_date_from_sm
+
+    Notes
+    -----
+    Cessation date is typically derived from soil moisture.
+    However, soil moisture data is rare, while rainfall data is plenty.
+    This function is identical to cess_date_from_sm except that at each time step,
+    soil moisture is estimated by water_balance_step.
+    The result would be the same as applying water_balance to `daily_rain`
+    to estimate soil moisture and use that in cess_date_from_sm.
+    Using cess_date_from_rain instead would be more efficient
+    if the user wishes not to keep the intermediary result of estimated soil moisture.
+    """
     sminit = xr.DataArray(sminit)
     et = xr.DataArray(et)
     taw = xr.DataArray(taw)

From 27a936035974686ac2f81116500b356e70a7ab26 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Cousin?= <remic@iri.columbia.edu>
Date: Wed, 2 Aug 2023 11:41:43 -0400
Subject: [PATCH 09/13] seasonal function for new cess_date functions and their
 tests

---
 enacts/calc.py            | 59 ++++++++++++++++++++++++++++++-
 enacts/tests/test_calc.py | 74 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 130 insertions(+), 3 deletions(-)

diff --git a/enacts/calc.py b/enacts/calc.py
index 405db3713..9458569bd 100644
--- a/enacts/calc.py
+++ b/enacts/calc.py
@@ -792,7 +792,7 @@ def seasonal_onset_date(
     # + seasonal_onset_date.onset_delta
     return seasonal_onset_date
 
-def seasonal_cess_date(
+def seasonal_cess_date_from_sm(
     soil_moisture,
     search_start_day,
     search_start_month,
@@ -874,6 +874,63 @@ def seasonal_cess_date(
     # + seasonal_onset_date.onset_delta
     return seasonal_cess_date
 
+
+def seasonal_cess_date_from_rain(
+    daily_rain,
+    search_start_day,
+    search_start_month,
+    search_days,
+    dry_thresh,
+    dry_spell_length_thresh,
+    et,
+    taw,
+    sminit,
+    time_dim="T"
+):
+    # Deal with leap year cases
+    if search_start_day == 29 and search_start_month == 2:
+        search_start_day = 1
+        search_start_month = 3
+
+    # Find an acceptable end_day/_month
+    first_end_date = sel_day_and_month(
+        daily_rain[time_dim], search_start_day, search_start_month
+    )[0] + np.timedelta64(
+        search_days,
+        "D",
+    )
+
+    end_day = first_end_date.dt.day.values
+
+    end_month = first_end_date.dt.month.values
+
+    # Apply daily grouping by season
+    grouped_daily_data = daily_tobegroupedby_season(
+        daily_rain, search_start_day, search_start_month, end_day, end_month
+    )
+    # Apply cess_date
+    seasonal_data = (
+        grouped_daily_data[daily_rain.name]
+        .groupby(grouped_daily_data["seasons_starts"])
+        .map(
+            cess_date_from_rain,
+            dry_thresh=dry_thresh,
+            dry_spell_length_thresh=dry_spell_length_thresh,
+            et=et,
+            taw=taw,
+            sminit=sminit,
+        )
+    ).rename("cess_delta")
+    # Get the seasons ends
+    seasons_ends = grouped_daily_data["seasons_ends"].rename({"group": time_dim})
+    seasonal_cess_date = xr.merge([seasonal_data, seasons_ends])
+
+    # Tip to get dates from timedelta search_start_day
+    # seasonal_onset_date = seasonal_onset_date[time_dim]
+    # + seasonal_onset_date.onset_delta
+    return seasonal_cess_date
+
+
 def seasonal_sum(
     daily_data,
     start_day,
diff --git a/enacts/tests/test_calc.py b/enacts/tests/test_calc.py
index 8cb0dc377..70dea58a2 100644
--- a/enacts/tests/test_calc.py
+++ b/enacts/tests/test_calc.py
@@ -351,7 +351,8 @@ def test_seasonal_cess_date_keeps_returning_same_outputs():
         sminit=0,
         time_dim="T"
     ).to_array(name="soil moisture").squeeze("variable", drop=True)
-    cessds = calc.seasonal_cess_date(
+    print(wb.where((wb["T"].dt.day == 31) & (wb["T"].dt.month == 8), drop=True))
+    cessds = calc.seasonal_cess_date_from_sm(
         soil_moisture=wb,
         search_start_day=1,
         search_start_month=9,
@@ -374,6 +375,24 @@ def test_seasonal_cess_date_keeps_returning_same_outputs():
         equal_nan=True,  
     )
 
+def test_seasonal_cess_date_from_rain_keeps_returning_same_outputs():
+
+    precip = data_test_calc.multi_year_data_sample()
+    cessds = calc.seasonal_cess_date_from_rain(
+        daily_rain=precip,
+        search_start_day=1,
+        search_start_month=9,
+        search_days=90,
+        dry_thresh=5,
+        dry_spell_length_thresh=3,
+        et=5,
+        taw=60,
+        sminit=33.57026932, # from previous test sm output on 8/31/2000
+    )
+    cess = (cessds.cess_delta + cessds["T"]).squeeze()
+
+    assert cess[0] == pd.to_datetime("2000-09-21T00:00:00.000000000")
+
 def test_seasonal_onset_date():
     t = pd.date_range(start="2000-01-01", end="2005-02-28", freq="1D")
     # this is rr_mrg.sel(T=slice("2000", "2005-02-28")).isel(X=150, Y=150).precip
@@ -459,7 +478,7 @@ def test_seasonal_cess_date():
         sminit=0,
         time_dim="T"
     ).to_array(name="soil moisture")
-    cessds = calc.seasonal_cess_date(
+    cessds = calc.seasonal_cess_date_from_sm(
         soil_moisture=wb,
         search_start_day=1,
         search_start_month=9,
@@ -484,6 +503,57 @@ def test_seasonal_cess_date():
     ).all()
 
 
+def test_seasonal_cess_date_from_rain():
+    t = pd.date_range(start="2000-01-01", end="2005-02-28", freq="1D")
+    synthetic_precip = xr.DataArray(np.zeros(t.size), dims=["T"], coords={"T": t}) + 1.1
+    synthetic_precip = xr.where(
+        (synthetic_precip["T"] == pd.to_datetime("2000-03-29"))
+        | (synthetic_precip["T"] == pd.to_datetime("2000-03-30"))
+        | (synthetic_precip["T"] == pd.to_datetime("2000-03-31"))
+        | (synthetic_precip["T"] == pd.to_datetime("2001-04-30"))
+        | (synthetic_precip["T"] == pd.to_datetime("2001-05-01"))
+        | (synthetic_precip["T"] == pd.to_datetime("2001-05-02"))
+        | (synthetic_precip["T"] == pd.to_datetime("2002-04-01"))
+        | (synthetic_precip["T"] == pd.to_datetime("2002-04-02"))
+        | (synthetic_precip["T"] == pd.to_datetime("2002-04-03"))
+        | (synthetic_precip["T"] == pd.to_datetime("2003-05-16"))
+        | (synthetic_precip["T"] == pd.to_datetime("2003-05-17"))
+        | (synthetic_precip["T"] == pd.to_datetime("2003-05-18"))
+        | (synthetic_precip["T"] == pd.to_datetime("2004-03-01"))
+        | (synthetic_precip["T"] == pd.to_datetime("2004-03-02"))
+        | (synthetic_precip["T"] == pd.to_datetime("2004-03-03")),
+        7,
+        synthetic_precip,
+    ).rename("synthetic_precip")    
+    cessds = calc.seasonal_cess_date_from_rain(
+        daily_rain=synthetic_precip,
+        search_start_day=1,
+        search_start_month=9,
+        search_days=90,
+        dry_thresh=5,
+        dry_spell_length_thresh=3,
+        et=5,
+        taw=60,
+        sminit=0,
+    )
+    cess = (cessds.cess_delta + cessds["T"]).squeeze()
+    assert (
+        cess
+        == pd.to_datetime(
+            xr.DataArray(
+                [
+                    "2000-09-01T00:00:00.000000000",
+                    "2001-09-01T00:00:00.000000000",
+                    "2002-09-01T00:00:00.000000000",
+                    "2003-09-01T00:00:00.000000000",
+                    "2004-09-01T00:00:00.000000000",
+                ],dims=["T"],coords={"T": cess["T"]},
+            )
+        )
+    ).all()
+
+
+
 def precip_sample():
 
     t = pd.date_range(start="2000-05-01", end="2000-06-30", freq="1D")

From 927f2139b00b704dad32d607c45ca704ad46e883 Mon Sep 17 00:00:00 2001
From: remic <remic@iri.columbia.edu>
Date: Thu, 3 Aug 2023 09:38:01 -0400
Subject: [PATCH 10/13] adatpted to new cess date functions

---
 enacts/onset/maproom.py | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/enacts/onset/maproom.py b/enacts/onset/maproom.py
index e29f8b711..126727e58 100644
--- a/enacts/onset/maproom.py
+++ b/enacts/onset/maproom.py
@@ -554,14 +554,16 @@ def cess_plots(
             return error_fig, error_fig, tab_style
         precip.load()
         try:
-            soil_moisture = calc.water_balance(precip, 5, 60, 0)["soil_moisture"]
-            cess_delta = calc.seasonal_cess_date(
-                soil_moisture,
+            cess_delta = calc.seasonal_cess_date_from_rain(
+                precip,
                 int(cess_start_day),
                 calc.strftimeb2int(cess_start_month),
                 int(cess_search_days),
                 int(cess_soil_moisture),
                 int(cess_dry_spell),
+                5,
+                60,
+                60./3.,
                 time_dim="T",
             )
             isnan = np.isnan(cess_delta["cess_delta"]).all()
@@ -698,14 +700,16 @@ def length_plots(
             )
             return error_fig, error_fig, tab_style
         try:
-            soil_moisture = calc.water_balance(precip, 5, 60, 0)["soil_moisture"]
-            cess_delta = calc.seasonal_cess_date(
-                soil_moisture,
+            cess_delta = calc.seasonal_cess_date_from_rain(
+                precip,
                 int(cess_start_day),
                 calc.strftimeb2int(cess_start_month),
                 int(cess_search_days),
                 int(cess_soil_moisture),
                 int(cess_dry_spell),
+                5,
+                60,
+                60./3.,
                 time_dim="T",
             )
             isnan = np.isnan(cess_delta["cess_delta"]).all()
@@ -850,16 +854,16 @@ def onset_tile(tz, tx, ty):
             dry_spell_search,
         )
         if ("length" in map_choice) | ("total" in map_choice):
-            soil_moisture = calc.water_balance(
-                precip_tile, 5, 60, 0
-            )["soil_moisture"]
-            cess_dates = calc.seasonal_cess_date(
-                soil_moisture,
+            cess_dates = calc.seasonal_cess_date_from_rain(
+                precip_tile,
                 cess_start_day,
                 cess_start_month1,
                 cess_search_days,
                 cess_soil_moisture,
                 cess_dry_spell,
+                5,
+                60,
+                60./3.,
             )
             if "length" in map_choice:
                 if cess_dates["T"][0] < onset_dates["T"][0]:

From 215e03e718973407a10fd1fd0a7930bb540e3d5e Mon Sep 17 00:00:00 2001
From: remic <remic@iri.columbia.edu>
Date: Tue, 12 Sep 2023 14:27:23 -0400
Subject: [PATCH 11/13] fixing typo, better names

---
 enacts/calc.py | 34 +++++++++++++++++-----------------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/enacts/calc.py b/enacts/calc.py
index 9458569bd..8cbbebc7a 100644
--- a/enacts/calc.py
+++ b/enacts/calc.py
@@ -420,7 +420,7 @@ def cess_date_from_rain(
     of soil moisture less than `dry_thresh` .
     Soil moisture is estimated through a simple water balance model
     driven by rainfall `daily_rain` , evapotranspiration `et`, soil capacity `taw`
-    and initislized at soil moisture value `sminit` .
+    and initialized at soil moisture value `sminit` .
     
     Parameters
     ----------
@@ -763,13 +763,13 @@ def seasonal_onset_date(
     end_month = first_end_date.dt.month.values
 
     # Apply daily grouping by season
-    grouped_daily_data = daily_tobegroupedby_season(
+    seasonally_labeled_daily_data = daily_tobegroupedby_season(
         daily_rain, search_start_day, search_start_month, end_day, end_month
     )
     # Apply onset_date
     seasonal_data = (
-        grouped_daily_data[daily_rain.name]
-        .groupby(grouped_daily_data["seasons_starts"])
+        seasonally_labeled_daily_data[daily_rain.name]
+        .groupby(seasonally_labeled_daily_data["seasons_starts"])
         .map(
             onset_date,
             wet_thresh=wet_thresh,
@@ -784,7 +784,7 @@ def seasonal_onset_date(
         .rename({"seasons_starts": time_dim})
     )
     # Get the seasons ends
-    seasons_ends = grouped_daily_data["seasons_ends"].rename({"group": time_dim})
+    seasons_ends = seasonally_labeled_daily_data["seasons_ends"].rename({"group": time_dim})
     seasonal_onset_date = xr.merge([seasonal_data, seasons_ends])
 
     # Tip to get dates from timedelta search_start_day
@@ -852,13 +852,13 @@ def seasonal_cess_date_from_sm(
     end_month = first_end_date.dt.month.values
 
     # Apply daily grouping by season
-    grouped_daily_data = daily_tobegroupedby_season(
+    seasonally_labeled_daily_data = daily_tobegroupedby_season(
         soil_moisture, search_start_day, search_start_month, end_day, end_month
     )
     # Apply cess_date
     seasonal_data = (
-        grouped_daily_data[soil_moisture.name]
-        .groupby(grouped_daily_data["seasons_starts"])
+        seasonally_labeled_daily_data[soil_moisture.name]
+        .groupby(seasonally_labeled_daily_data["seasons_starts"])
         .map(
             cess_date_from_sm,
             dry_thresh=dry_thresh,
@@ -866,7 +866,7 @@ def seasonal_cess_date_from_sm(
         )
     ).rename("cess_delta")
     # Get the seasons ends
-    seasons_ends = grouped_daily_data["seasons_ends"].rename({"group": time_dim})
+    seasons_ends = seasonally_labeled_daily_data["seasons_ends"].rename({"group": time_dim})
     seasonal_cess_date = xr.merge([seasonal_data, seasons_ends])
 
     # Tip to get dates from timedelta search_start_day
@@ -905,13 +905,13 @@ def seasonal_cess_date_from_rain(
     end_month = first_end_date.dt.month.values
 
     # Apply daily grouping by season
-    grouped_daily_data = daily_tobegroupedby_season(
+    seasonally_labeled_daily_data = daily_tobegroupedby_season(
         daily_rain, search_start_day, search_start_month, end_day, end_month
     )
     # Apply cess_date
     seasonal_data = (
-        grouped_daily_data[daily_rain.name]
-        .groupby(grouped_daily_data["seasons_starts"])
+        seasonally_labeled_daily_data[daily_rain.name]
+        .groupby(seasonally_labeled_daily_data["seasons_starts"])
         .map(
             cess_date_from_rain,
             dry_thresh=dry_thresh,
@@ -922,7 +922,7 @@ def seasonal_cess_date_from_rain(
         )
     ).rename("cess_delta")
     # Get the seasons ends
-    seasons_ends = grouped_daily_data["seasons_ends"].rename({"group": time_dim})
+    seasons_ends = seasonally_labeled_daily_data["seasons_ends"].rename({"group": time_dim})
     seasonal_cess_date = xr.merge([seasonal_data, seasons_ends])
 
     # Tip to get dates from timedelta search_start_day
@@ -973,16 +973,16 @@ def seasonal_sum(
     Examples
     --------
     """
-    grouped_daily_data = daily_tobegroupedby_season(
+    seasonally_labeled_daily_data = daily_tobegroupedby_season(
         daily_data, start_day, start_month, end_day, end_month
     )
     seasonal_data = (
-        grouped_daily_data[daily_data.name]
-        .groupby(grouped_daily_data["seasons_starts"])
+        seasonally_labeled_daily_data[daily_data.name]
+        .groupby(seasonally_labeled_daily_data["seasons_starts"])
         .sum(dim=time_dim, skipna=True, min_count=min_count)
         #        .rename({"seasons_starts": time_dim})
     )
-    seasons_ends = grouped_daily_data["seasons_ends"].rename({"group": time_dim})
+    seasons_ends = seasonally_labeled_daily_data["seasons_ends"].rename({"group": time_dim})
     summed_seasons = xr.merge([seasonal_data, seasons_ends])
     return summed_seasons
 

From d2ecc451419f4125e94db7ee202efde3123023d8 Mon Sep 17 00:00:00 2001
From: remic <remic@iri.columbia.edu>
Date: Thu, 26 Oct 2023 04:06:30 -0400
Subject: [PATCH 12/13] removing a print

---
 enacts/tests/test_calc.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/enacts/tests/test_calc.py b/enacts/tests/test_calc.py
index 70dea58a2..10396aec9 100644
--- a/enacts/tests/test_calc.py
+++ b/enacts/tests/test_calc.py
@@ -351,7 +351,6 @@ def test_seasonal_cess_date_keeps_returning_same_outputs():
         sminit=0,
         time_dim="T"
     ).to_array(name="soil moisture").squeeze("variable", drop=True)
-    print(wb.where((wb["T"].dt.day == 31) & (wb["T"].dt.month == 8), drop=True))
     cessds = calc.seasonal_cess_date_from_sm(
         soil_moisture=wb,
         search_start_day=1,

From b3819ef9bef31587c98869af275f8f687415b02c Mon Sep 17 00:00:00 2001
From: remic <remic@iri.columbia.edu>
Date: Mon, 15 Apr 2024 11:14:55 -0400
Subject: [PATCH 13/13] easier to read but not allowed

---
 enacts/wat_bal/maproom_monit.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/enacts/wat_bal/maproom_monit.py b/enacts/wat_bal/maproom_monit.py
index 781d58958..b5dc2bfdf 100644
--- a/enacts/wat_bal/maproom_monit.py
+++ b/enacts/wat_bal/maproom_monit.py
@@ -512,7 +512,7 @@ def wat_bal_plots(
     ts = wat_bal_ts(
         precip,
         map_choice,
-        et=5,
+        5,
         taw,
         int(planting_day),
         calc.strftimeb2int(planting_month),
@@ -534,7 +534,7 @@ def wat_bal_plots(
     ts2 = wat_bal_ts(
         precip,
         map_choice,
-        et=5,
+        5,
         taw,
         int(planting2_day),
         calc.strftimeb2int(planting2_month),
@@ -653,8 +653,8 @@ def wat_bal_tile(tz, tx, ty):
 
     sm, drainage, et_crop, et_crop_red, planting_date = wat_bal(
         precip_tile,
-        et=5,
-        taw=taw_tile,
+        5,
+        taw_tile,
         planting_day,
         planting_month1,
         kc_init_length,