Improve time discretization

nlmod/dims/time.py

@@ -1,14 +1,16 @@
 import datetime as dt
 import logging
+import warnings
 
 import numpy as np
 import pandas as pd
+import xarray as xr
 from xarray import IndexVariable
 
 logger = logging.getLogger(__name__)
 
 
-def set_ds_time(
+def set_ds_time_deprecated(
     ds,
     time=None,
     steady_state=False,
@@ -71,6 +73,13 @@ def set_ds_time(
     ds : xarray.Dataset
         dataset with time variant model data
     """
+
+    warnings.warn(
+        "this function is deprecated and will eventually be removed, "
+        "please use nlmod.time.set_ds_time() in the future.",
+        DeprecationWarning,
+    )
+
     # checks
     if time_units.lower() != "days":
         raise NotImplementedError()
@@ -127,6 +136,99 @@ def set_ds_time(
     return ds
 
 
+def set_ds_time(ds, time, start, steady=True, time_units="DAYS", nstp=1, tsmult=1.0):
+    """Set time discretisation for model dataset.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        model dataset
+    time : array-like
+        array-like of floats (indicating elapsed time) or timestamps corresponding to
+        the end of each stress period in the model.
+    start : str or pandas.Timestamp, optional
+        model start datetime as string or pandas Timestamp, if None, defaults to
+        1 january 2000.
+    steady : arraylike or bool, optional
+        arraylike indicating which stress periods are steady-state, by default True,
+        which sets all stress periods to steady-state.
+    time_units : str, optional
+        time units, by default "DAYS"
+    nstp : int or array-like, optional
+        number of steps per stress period, stored in ds.attrs, default is 1
+    tsmult : float, optional
+        timestep multiplier within stress periods, stored in ds.attrs, default is 1.0
+
+    Returns
+    -------
+    ds : xarray.Dataset
+        model dataset with added time coordinate
+
+    """
+    logger.info(
+        "This is the new version of set_ds_time()."
+        " If you're looking for the old behavior,"
+        "use `nlmod.time.set_ds_time_deprecated()`."
+    )
+
+    # parse start datetime
+    if isinstance(start, str):
+        start = pd.Timestamp(start)
+    elif isinstance(start, (pd.Timestamp, np.datetime64)):
+        pass
+    else:
+        raise TypeError("Cannot parse start datetime.")
+
+    # convert time to Timestamps
+    # calculate time idx
+    if isinstance(time[0], (int, np.integer, float)):
+        time = pd.Timestamp(start) + pd.to_timedelta(time, time_units)
+    elif isinstance(time[0], str):
+        time = pd.to_datetime(time)
+    elif isinstance(time[0], (pd.Timestamp, np.datetime64, xr.core.variable.Variable)):
+        pass
+    else:
+        raise TypeError("Cannot process 'time' argument. Datatype not understood.")
+
+    # set steady
+    if isinstance(steady, bool):
+        steady = steady * np.ones(len(time))
+
+    ds = ds.assign_coords(coords={"time": time})
+    if time_units == "D":
+        time_units = "DAYS"
+    ds.time.attrs["time_units"] = time_units
+    ds.time.attrs["start"] = str(start)
+    ds.time.attrs["steady"] = steady
+    ds.time.attrs["nstp"] = nstp
+    ds.time.attrs["tsmult"] = tsmult
+    return ds
+
+
+def ds_time_idx_from_tdis_settings(start, perlen, nstp=1, tsmult=1.0, time_units="D"):
+    deltlist = []
+    for kper, delt in enumerate(perlen):
+        if not isinstance(nstp, int):
+            kstpkper = nstp[kper]
+        else:
+            kstpkper = nstp
+
+        if not isinstance(tsmult, float):
+            tsm = tsmult[kper]
+        else:
+            tsm = tsmult
+
+        if tsm > 1.0:
+            delt0 = delt * (tsm - 1) / (tsm**kstpkper - 1)
+            delt = delt0 * tsm ** np.arange(kstpkper)
+        else:
+            delt = np.ones(kstpkper) * delt / kstpkper
+        deltlist.append(delt)
+
+    dt_arr = np.cumsum(np.concatenate(deltlist))
+    return ds_time_idx(dt_arr, start_datetime=start, time_units=time_units)
+
+
 def estimate_nstp(
     forcing, perlen=1, tsmult=1.1, nstp_min=1, nstp_max=25, return_dt_arr=False
 ):
@@ -214,6 +316,15 @@ def estimate_nstp(
 
 
 def ds_time_from_model(gwf):
+    warnings.warn(
+        "this function was renamed to `ds_time_idx_from_model`. "
+        "Please use the new function name.",
+        DeprecationWarning,
+    )
+    return ds_time_idx_from_model(gwf)
+
+
+def ds_time_idx_from_model(gwf):
     """Get time index variable from model (gwf or gwt).
 
     Parameters
@@ -227,10 +338,19 @@ def ds_time_from_model(gwf):
         time coordinate for xarray data-array or dataset
     """
 
-    return ds_time_from_modeltime(gwf.modeltime)
+    return ds_time_idx_from_modeltime(gwf.modeltime)
 
 
 def ds_time_from_modeltime(modeltime):
+    warnings.warn(
+        "this function was renamed to `ds_time_idx_from_model`. "
+        "Please use the new function name.",
+        DeprecationWarning,
+    )
+    return ds_time_idx_from_modeltime(modeltime)
+
+
+def ds_time_idx_from_modeltime(modeltime):
     """Get time index variable from modeltime object.
 
     Parameters

nlmod/gwf/gwf.py

@@ -578,18 +578,15 @@ def sto(
     """
     logger.info("creating mf6 STO")
 
-    if ds.time.steady_state:
+    if ds.time.steady_state.all():
+        logger.warning("Model is steady-state, no STO package created.")
         return None
     else:
-        if ds.time.steady_start:
-            sts_spd = {0: True}
-            trn_spd = {1: True}
-        else:
-            sts_spd = None
-            trn_spd = {0: True}
+        sts_spd = {iper: bool(b) for iper, b in enumerate(ds.time.steady)}
+        trn_spd = {iper: not bool(b) for iper, b in enumerate(ds.time.steady)}
 
         sy = _get_value_from_ds_datavar(ds, "sy", sy, default=0.2)
-        ss = _get_value_from_ds_datavar(ds, "ss", ss, default=0.000001)
+        ss = _get_value_from_ds_datavar(ds, "ss", ss, default=1e-5)
 
         sto = flopy.mf6.ModflowGwfsto(
             gwf,

nlmod/gwf/recharge.py

@@ -35,7 +35,7 @@ def ds_to_rch(gwf, ds, mask=None, pname="rch", **kwargs):
         raise ValueError("please remove nan values in recharge data array")
 
     # get stress period data
-    if ds.time.steady_state:
+    if ds.time.steady.all():
         recharge = "recharge"
         if "time" in ds["recharge"].dims:
             mask = ds["recharge"].isel(time=0) != 0
@@ -69,7 +69,7 @@ def ds_to_rch(gwf, ds, mask=None, pname="rch", **kwargs):
         **kwargs,
     )
 
-    if ds.time.steady_state:
+    if ds.time.steady.all():
         return rch
 
     # create timeseries packages
@@ -128,7 +128,7 @@ def ds_to_evt(gwf, ds, pname="evt", nseg=1, surface=None, depth=None, **kwargs):
         raise ValueError("please remove nan values in evaporation data array")
 
     # get stress period data
-    if ds.time.steady_state:
+    if ds.time.steady.all():
         if "time" in ds["evaporation"].dims:
             mask = ds["evaporation"].isel(time=0) != 0
         else:
@@ -163,7 +163,7 @@ def ds_to_evt(gwf, ds, pname="evt", nseg=1, surface=None, depth=None, **kwargs):
         **kwargs,
     )
 
-    if ds.time.steady_state:
+    if ds.time.steady.all():
         return evt
 
     # create timeseries packages

nlmod/plot/plotutil.py

@@ -6,6 +6,7 @@
 from ..dims.resample import get_affine_mod_to_world
 from ..epsg28992 import EPSG_28992
 
+
 def get_patches(ds, rotated=False):
     """Get the matplotlib patches for a vertex grid."""
     assert "icell2d" in ds.dims

nlmod/sim/sim.py

@@ -64,7 +64,7 @@ def write_and_run(sim, ds, write_ds=True, script_path=None, silent=False):
     ds.attrs["model_ran_on"] = dt.datetime.now().strftime("%Y%m%d_%H:%M:%S")
 
 
-def get_tdis_perioddata(ds):
+def get_tdis_perioddata(ds, nstp=None, tsmult=None):
     """Get tdis_perioddata from ds.
 
     Parameters
@@ -92,15 +92,21 @@ def get_tdis_perioddata(ds):
     if len(ds["time"]) > 1:
         perlen.extend(np.diff(ds["time"]) / deltat)
 
-    if "nstp" in ds:
-        nstp = ds["nstp"].values
-    else:
-        nstp = [ds.time.nstp] * len(perlen)
+    nstp = util._get_value_from_ds_datavar(
+        ds, "nstp", nstp, return_da=False, warn=False
+    )
+    nstp = util._get_value_from_ds_attr(ds.time, "nstp", nstp)
+
+    if isinstance(nstp, (int, np.integer)):
+        nstp = [nstp] * len(perlen)
+
+    nstp = util._get_value_from_ds_datavar(
+        ds, "nstp", nstp, return_da=False, warn=False
+    )
+    tsmult = util._get_value_from_ds_attr(ds.time, "tsmult", value=tsmult)
 
-    if "tsmult" in ds:
-        tsmult = ds["tsmult"].values
-    else:
-        tsmult = [ds.time.tsmult] * len(perlen)
+    if isinstance(tsmult, float):
+        tsmult = [tsmult] * len(perlen)
 
     tdis_perioddata = list(zip(perlen, nstp, tsmult))