adapt time utilities from imod-python (#244)

wpbonelli · web-flow · commit c10d4d2993b3 · 2025-10-28T20:41:28.000-04:00
Begin to implement the time side of #226 and also to address #243 with a Time subclass of flopy3's ModelTime with extra capabilities. Adapt some time-related utilities from imod-python supporting the ability to generate a TDIS from timestamps - https://github.com/Deltares/imod-python/blob/87bbed86a6e36b1b65644625aabc6c3d5b6fda63/imod/util/time.py#L98 - https://github.com/Deltares/imod-python/blob/87bbed86a6e36b1b65644625aabc6c3d5b6fda63/imod/mf6/simulation.py#L177 This can be the foundation for a similar method to create a TDIS from simulation packages later on.
diff --git a/docs/examples/quickstart.py b/docs/examples/quickstart.py
@@ -2,16 +2,16 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-from flopy.discretization.modeltime import ModelTime
 from flopy.discretization.structuredgrid import StructuredGrid
 
 from flopy4.mf6.gwf import Chd, Gwf, Ic, Npf, Oc
 from flopy4.mf6.ims import Ims
 from flopy4.mf6.simulation import Simulation
+from flopy4.mf6.utils.time import Time
 
 name = "quickstart"
 workspace = Path(__file__).parent / name
-time = ModelTime(perlen=[1.0], nstp=[1])
+time = Time(perlen=[1.0], nstp=[1])
 grid = StructuredGrid(
     nlay=1,
     nrow=10,
diff --git a/flopy4/mf6/component.py b/flopy4/mf6/component.py
@@ -11,7 +11,7 @@
 
 from flopy4.mf6.constants import MF6
 from flopy4.mf6.spec import field, fields_dict, to_field
-from flopy4.mf6.utils.grid_utils import update_maxbound
+from flopy4.mf6.utils.grid import update_maxbound
 from flopy4.uio import IO, Loader, Writer
 
 COMPONENTS = {}
diff --git a/flopy4/mf6/gwf/chd.py b/flopy4/mf6/gwf/chd.py
@@ -10,7 +10,7 @@
 from flopy4.mf6.converter import structure_array
 from flopy4.mf6.package import Package
 from flopy4.mf6.spec import array, field, path
-from flopy4.mf6.utils.grid_utils import update_maxbound
+from flopy4.mf6.utils.grid import update_maxbound
 from flopy4.utils import to_path
 
 
diff --git a/flopy4/mf6/gwf/drn.py b/flopy4/mf6/gwf/drn.py
@@ -10,7 +10,7 @@
 from flopy4.mf6.converter import structure_array
 from flopy4.mf6.package import Package
 from flopy4.mf6.spec import array, field, path
-from flopy4.mf6.utils.grid_utils import update_maxbound
+from flopy4.mf6.utils.grid import update_maxbound
 from flopy4.utils import to_path
 
 
diff --git a/flopy4/mf6/gwf/rch.py b/flopy4/mf6/gwf/rch.py
@@ -10,7 +10,7 @@
 from flopy4.mf6.converter import structure_array
 from flopy4.mf6.package import Package
 from flopy4.mf6.spec import array, field, path
-from flopy4.mf6.utils.grid_utils import update_maxbound
+from flopy4.mf6.utils.grid import update_maxbound
 from flopy4.utils import to_path
 
 
diff --git a/flopy4/mf6/gwf/wel.py b/flopy4/mf6/gwf/wel.py
@@ -10,7 +10,7 @@
 from flopy4.mf6.converter import structure_array
 from flopy4.mf6.package import Package
 from flopy4.mf6.spec import array, field, path
-from flopy4.mf6.utils.grid_utils import update_maxbound
+from flopy4.mf6.utils.grid import update_maxbound
 from flopy4.utils import to_path
 
 
diff --git a/flopy4/mf6/simulation.py b/flopy4/mf6/simulation.py
@@ -1,7 +1,6 @@
 from os import PathLike
 from warnings import warn
 
-from flopy.discretization.modeltime import ModelTime
 from modflow_devtools.misc import cd, run_cmd
 from xattree import xattree
 
@@ -11,14 +10,15 @@
 from flopy4.mf6.solution import Solution
 from flopy4.mf6.spec import field
 from flopy4.mf6.tdis import Tdis
+from flopy4.mf6.utils.time import Time
 
 
 def convert_time(value):
-    if isinstance(value, ModelTime):
+    if isinstance(value, Time):
         return Tdis.from_time(value)
     if isinstance(value, Tdis):
         return value
-    raise TypeError(f"Expected ModelTime or Tdis, got {type(value)}")
+    raise TypeError(f"Expected Time or Tdis, got {type(value)}")
 
 
 @xattree
@@ -41,8 +41,8 @@ def __attrs_post_init__(self):
             self.filename = "mfsim.nam"
 
     @property
-    def time(self) -> ModelTime:
-        """Return the simulation time discretization."""
+    def time(self) -> Time:
+        """Return a `Time` object describing the simulation's time discretization."""
         return self.tdis.to_time()
 
     def run(self, exe: str | PathLike = "mf6", verbose: bool = False) -> None:
diff --git a/flopy4/mf6/tdis.py b/flopy4/mf6/tdis.py
@@ -3,13 +3,13 @@
 
 import numpy as np
 from attrs import Converter, define
-from flopy.discretization.modeltime import ModelTime
-from numpy.typing import NDArray
+from numpy.typing import ArrayLike, NDArray
 from xattree import ROOT, xattree
 
 from flopy4.mf6.converter import structure_array
 from flopy4.mf6.package import Package
 from flopy4.mf6.spec import array, dim, field
+from flopy4.mf6.utils.time import Time
 
 
 @xattree
@@ -42,9 +42,9 @@ class PeriodData:
         converter=Converter(structure_array, takes_self=True, takes_field=True),
     )
 
-    def to_time(self) -> ModelTime:
-        """Convert to a `ModelTime` object."""
-        return ModelTime(
+    def to_time(self) -> Time:
+        """Convert to a `Time` object."""
+        return Time(
             nper=self.nper,
             time_units=self.time_units,
             start_date_time=self.start_date_time,
@@ -54,8 +54,8 @@ def to_time(self) -> ModelTime:
         )
 
     @classmethod
-    def from_time(cls, time: ModelTime) -> "Tdis":
-        """Create a time discretization from a `ModelTime`."""
+    def from_time(cls, time: Time) -> "Tdis":
+        """Create a time discretization from a `Time` object."""
         return cls(
             nper=time.nper,
             time_units=None if time.time_units in [None, "unknown"] else time.time_units,
@@ -64,3 +64,33 @@ def from_time(cls, time: ModelTime) -> "Tdis":
             nstp=time.nstp,
             tsmult=time.tsmult,
         )
+
+    @classmethod
+    def from_timestamps(
+        cls,
+        timestamps: ArrayLike,
+        nstp: Optional[ArrayLike] = None,
+        tsmult: Optional[ArrayLike] = None,
+    ) -> "Tdis":
+        """
+        Create a time discretization from timestamps.
+
+        Parameters
+        ----------
+        timestamps : sequence of datetime-likes
+            Stress period start times
+        nstp : int or sequence of int, optional
+            Number of timesteps per stress period. If scalar, applied to all periods.
+            If None, defaults to 1 for all periods.
+        tsmult : float or sequence of float, optional
+            Timestep multiplier per stress period. If scalar, applied to all periods.
+            If None, defaults to 1.0 for all periods.
+
+        Returns
+        -------
+        Tdis
+            Time discretization object
+        """
+
+        time = Time.from_timestamps(timestamps, nstp=nstp, tsmult=tsmult)
+        return cls.from_time(time)
diff --git a/flopy4/mf6/utils/cbc_reader.py b/flopy4/mf6/utils/cbc_reader.py
@@ -13,7 +13,7 @@
 from flopy.discretization import StructuredGrid
 
 from flopy4.adapters import StructuredGridWrapper
-from flopy4.mf6.utils.grid_utils import get_coords
+from flopy4.mf6.utils.grid import get_coords
 
 
 @define
diff --git a/flopy4/mf6/utils/grid.py b/flopy4/mf6/utils/grid.py
diff --git a/flopy4/mf6/utils/heads_reader.py b/flopy4/mf6/utils/heads_reader.py
@@ -8,7 +8,7 @@
 import xarray as xr
 from flopy.discretization import StructuredGrid
 
-from .grid_utils import get_coords
+from .grid import get_coords
 
 
 def open_hds(
diff --git a/flopy4/mf6/utils/time.py b/flopy4/mf6/utils/time.py
@@ -0,0 +1,99 @@
+import numpy as np
+import pandas as pd
+from flopy.discretization.modeltime import ModelTime
+from numpy.typing import ArrayLike
+
+
+class Time(ModelTime):
+    """Extend flopy3's ModelTime"""
+
+    @classmethod
+    def from_timestamps(
+        cls,
+        timestamps: ArrayLike,
+        nstp: ArrayLike | None = None,
+        tsmult: ArrayLike | None = None,
+    ) -> "Time":
+        """
+        Create a time discretization from timestamps.
+
+        Parameters
+        ----------
+        timestamps : sequence of datetime-likes
+            Stress period start times
+        nstp : int or sequence of int, optional
+            Number of timesteps per stress period. If scalar, applied to all periods.
+            If None, defaults to 1 for all periods.
+        tsmult : float or sequence of float, optional
+            Timestep multiplier per stress period. If scalar, applied to all periods.
+            If None, defaults to 1.0 for all periods.
+
+        Returns
+        -------
+        Time
+            Time discretization object
+        """
+
+        unique_times = np.unique(np.hstack(pd.to_datetime(timestamps)))  # np.unique also sorts
+        if len(unique_times) < 2:
+            raise ValueError("Need at least two timestamps to create time discretization")
+
+        timedeltas = to_timedeltas(unique_times)
+        perlen = np.array(
+            [delta.total_seconds() / 86400.0 for delta in timedeltas], dtype=np.float64
+        )
+        nper = len(perlen)
+
+        if nstp is None:
+            nstp_array = np.ones(nper, dtype=np.int64)
+        elif np.isscalar(nstp):
+            nstp_array = np.full(nper, nstp, dtype=np.int64)
+        else:
+            nstp_array = np.array(nstp, dtype=np.int64)
+            if len(nstp_array) != nper:
+                raise ValueError(
+                    f"nstp length ({len(nstp_array)}) must match number of periods ({nper})"
+                )
+
+        if tsmult is None:
+            tsmult_array = np.ones(nper, dtype=np.float64)
+        elif np.isscalar(tsmult):
+            tsmult_array = np.full(nper, tsmult, dtype=np.float64)
+        else:
+            tsmult_array = np.array(tsmult, dtype=np.float64)
+            if len(tsmult_array) != nper:
+                raise ValueError(
+                    f"tsmult length ({len(tsmult_array)}) must match number of periods ({nper})"
+                )
+
+        return cls(
+            perlen=perlen,
+            nstp=nstp_array,
+            tsmult=tsmult_array,
+            time_units="days",
+            start_datetime=pd.to_datetime(unique_times[0]).to_pydatetime(),
+        )
+
+
+def to_timedeltas(
+    timestamps: ArrayLike,
+) -> list[pd.Timedelta]:
+    """
+    Converts a sequence of datetime-like objects to a list of timedelta
+    objects representing the durations between consecutive time points.
+
+    Parameters
+    ----------
+    timestamps : sequence of datetime-likes
+        A sequence of datetime-like objects representing time points.
+
+    Returns
+    -------
+    timedeltas : list of pd.Timedelta
+        A list of durations between consecutive time points.
+    """
+    if len(np.atleast_1d(timestamps)) < 2:
+        return []
+
+    timestamps = pd.to_datetime(timestamps)
+    return [pd.Timedelta(end - start) for start, end in zip(timestamps[:-1], timestamps[1:])]  # type: ignore
diff --git a/test/test_codec.py b/test/test_codec.py
@@ -118,11 +118,10 @@ def test_dumps_dis():
 
 
 def test_dumps_tdis():
-    from flopy.discretization.modeltime import ModelTime
-
     from flopy4.mf6.tdis import Tdis
+    from flopy4.mf6.utils.time import Time
 
-    tdis = Tdis.from_time(ModelTime(perlen=[1.0, 2.0], nstp=[1, 2]))
+    tdis = Tdis.from_time(Time(perlen=[1.0, 2.0], nstp=[1, 2]))
     tdis.time_units = "days"
 
     dumped = dumps(COMPONENT_CONVERTER.unstructure(tdis))
@@ -404,11 +403,10 @@ def test_dumps_gwf():
 
 
 def test_dumps_simulation():
-    from flopy.discretization.modeltime import ModelTime
-
     from flopy4.mf6.gwf import Dis, Gwf, Ic, Npf, Oc
     from flopy4.mf6.simulation import Simulation
     from flopy4.mf6.tdis import Tdis
+    from flopy4.mf6.utils.time import Time
 
     # Create model components
     dis = Dis(nlay=1, nrow=5, ncol=5, delr=100.0, delc=100.0)
@@ -427,7 +425,7 @@ def test_dumps_simulation():
     )
 
     # Create time discretization
-    time = ModelTime(perlen=[1.0], nstp=[1])
+    time = Time(perlen=[1.0], nstp=[1])
     tdis = Tdis.from_time(time)
 
     # Create simulation
diff --git a/test/test_component.py b/test/test_component.py
diff --git a/test/test_time_utils.py b/test/test_time_utils.py
