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..8cbbebc7a 100644 --- a/enacts/calc.py +++ b/enacts/calc.py @@ -366,64 +366,120 @@ def cess_date_step(cess_yesterday, dry_spell_length, dry_spell_length_thresh): ) - np.timedelta64(1, "D") -def cess_date( - daily_data, +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. - - Find first day of the first dry spell where: - Soil moisture falls below `dry_thresh` for `min_dry_days` days. + """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_data : DataArray - Array of daily soil moisture or daily rainfall. + daily_sm : DataArray + Array of daily soil moisture values. dry_thresh : float - Soil moisture threshold to determine a dry day. + A dry day is when soil mositure `daily_sm` is lesser than `dry_thresh` . 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. + A dry spell is at least `dry_spell_length_thresh` dry days. time_dim : str, optional - Time coordinate in `soil_moisture` (default `time_dim`="T"). - + Time coordinate in `daily_sm` (default `time_dim` ="T"). + Returns ------- - cess_delta : DataArray[np.timedelta64] - Difference between first day of `daily_data` - and cessation date. + 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}) + + return _cess_date(dry_thresh, dry_spell_length_thresh, sm_func, daily_sm[time_dim]) + + +def cess_date_from_rain( + daily_rain, + dry_thresh, + dry_spell_length_thresh, + et, + taw, + 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 initialized 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 -------- - water_balance + cess_date_step, water_balance_step, cess_date_from_sm Notes ----- - Examples - -------- + 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. """ - # Initializing - spell_length = ( - daily_data.isel({time_dim: 0}) < dry_thresh - ).astype("timedelta64[D]") + 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:]: - dry_day = daily_data.sel({time_dim: t}) < dry_thresh + 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( cess_delta, @@ -432,9 +488,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 @@ -707,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, @@ -728,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 @@ -736,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, @@ -796,21 +852,77 @@ def seasonal_cess_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( 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, + dry_spell_length_thresh=dry_spell_length_thresh, + ) + ).rename("cess_delta") + # Get the seasons ends + 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 + # seasonal_onset_date = seasonal_onset_date[time_dim] + # + 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 + seasonally_labeled_daily_data = daily_tobegroupedby_season( + daily_rain, search_start_day, search_start_month, end_day, end_month + ) + # Apply cess_date + seasonal_data = ( + seasonally_labeled_daily_data[daily_rain.name] + .groupby(seasonally_labeled_daily_data["seasons_starts"]) .map( - cess_date, + 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}) + 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 @@ -818,6 +930,7 @@ def seasonal_cess_date( # + seasonal_onset_date.onset_delta return seasonal_cess_date + def seasonal_sum( daily_data, start_day, @@ -860,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 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]: 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 86% rename from enacts/onset/tests/test_calc.py rename to enacts/tests/test_calc.py index 71ac18c4e..10396aec9 100644 --- a/enacts/onset/tests/test_calc.py +++ b/enacts/tests/test_calc.py @@ -351,7 +351,7 @@ 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( + cessds = calc.seasonal_cess_date_from_sm( soil_moisture=wb, search_start_day=1, search_start_month=9, @@ -374,6 +374,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 +477,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 +502,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") @@ -546,7 +615,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_from_sm(daily_sm, 1, 3) expected = xr.DataArray( [np.nan, 0, np.nan, np.nan, 2, 1, 0, np.nan] ).astype("timedelta64[D]") @@ -555,9 +624,39 @@ 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_from_rain( + daily_rain, + dry_thresh=1, + dry_spell_length_thresh=3, + et=5, + taw=10, + sminit=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, + cessations = calc.cess_date_from_sm( + daily_sm=data, 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..b5dc2bfdf 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 @@ -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,