add grid and time fields to components (#124)

Initial step towards #120, add grid property to Gwf and time property to Simulation, and use a converter so the user can pass grid/modeltime to gwf/sim instead of dis/tdis, respectively

Author: wpbonelli

docs/examples/quickstart.py

@@ -55,7 +55,5 @@
 ax.grid(which="both", color="white")
 head.plot.imshow(ax=ax)
 head.plot.contour(ax=ax, levels=[0.2, 0.4, 0.6, 0.8], linewidths=3.0)
-budget.plot.quiver(
-    x="x", y="y", u="npf-qx", v="npf-qy", ax=ax, color="white"
-)
+budget.plot.quiver(x="x", y="y", u="npf-qx", v="npf-qy", ax=ax, color="white")
 fig.savefig(ws / "quickstart.png")

flopy4/mf6/converters.py

@@ -3,7 +3,7 @@
 import numpy as np
 import sparse
 from numpy.typing import NDArray
-from xattree import _get_xatspec
+from xattree import get_xatspec
 
 from flopy4.mf6.config import SPARSE_THRESHOLD
 from flopy4.mf6.constants import FILL_DNODATA
@@ -16,8 +16,8 @@ def convert_array(value, self_, field) -> NDArray:
         return value
 
     # get spec
-    spec = _get_xatspec(type(self_))
-    field = spec.arrays[field.name]
+    spec = get_xatspec(type(self_))
+    field = spec[field.name]
     if not field.dims:
         raise ValueError(f"Field {field} missing dims")
 

flopy4/mf6/gwf/__init__.py

@@ -5,6 +5,7 @@
 import imod
 import xarray as xr
 from attrs import define
+from flopy.discretization.grid import Grid
 from xattree import field, xattree
 
 from flopy4.mf6.gwf.chd import Chd
@@ -18,6 +19,14 @@
 __all__ = ["Gwf", "Chd", "Dis", "Ic", "Npf", "Oc"]
 
 
+def convert_grid(value):
+    if isinstance(value, Grid):
+        return Dis.from_grid(value)
+    if isinstance(value, Dis):
+        return value
+    raise TypeError(f"Expected Grid or Dis, got {type(value)}")
+
+
 @xattree
 class Gwf(Model):
     @define
@@ -39,7 +48,7 @@ def budget(self):
                 merge_to_dataset=True,
             )
 
-    dis: Dis = field()
+    dis: Dis = field(converter=convert_grid)
     ic: Ic = field()
     oc: Oc = field()
     npf: Npf = field()
@@ -69,3 +78,7 @@ class NewtonOptions:
     nc_filerecord: Optional[Path] = field(
         default=None, metadata={"block": "options"}
     )
+
+    @property
+    def grid(self) -> Grid:
+        return self.dis.to_grid()

flopy4/mf6/gwf/dis.py

@@ -1,5 +1,6 @@
 import numpy as np
 from attrs import Converter
+from flopy.discretization.structuredgrid import StructuredGrid
 from numpy.typing import NDArray
 from xattree import array, dim, field, xattree
 
@@ -82,3 +83,49 @@ class Dis(Package):
 
     def __attrs_post_init__(self):
         self.nnodes = self.ncol * self.nrow * self.nlay
+
+    def to_grid(self) -> StructuredGrid:
+        """
+        Convert the discretization to a `StructuredGrid`.
+
+        Returns
+        -------
+        StructuredGrid
+            A `StructuredGrid` with the same dimensions and data as the `Dis`.
+        """
+        return StructuredGrid(
+            nlay=self.nlay,
+            nrow=self.nrow,
+            ncol=self.ncol,
+            delr=self.delr,
+            delc=self.delc,
+            top=self.top,
+            botm=self.botm,
+            idomain=self.idomain,
+        )
+
+    @classmethod
+    def from_grid(cls, grid: StructuredGrid) -> "Dis":
+        """
+        Create a discretization from a `StructuredGrid`.
+
+        Parameters
+        ----------
+        grid : StructuredGrid
+            A structured grid.
+
+        Returns
+        -------
+        Dis
+            A discretization with the same dimensions and data as the grid.
+        """
+        return Dis(
+            nlay=grid.nlay,
+            nrow=grid.nrow,
+            ncol=grid.ncol,
+            delr=grid.delr,
+            delc=grid.delc,
+            top=grid.top,
+            botm=grid.botm,
+            idomain=grid.idomain,
+        )

flopy4/mf6/simulation.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 
+from flopy.discretization.modeltime import ModelTime
 from xattree import field, xattree
 
 from flopy4.mf6.component import Component
@@ -9,10 +10,22 @@
 from flopy4.mf6.tdis import Tdis
 
 
+def convert_time(value):
+    if isinstance(value, ModelTime):
+        return Tdis.from_time(value)
+    if isinstance(value, Tdis):
+        return value
+    raise TypeError(f"Expected ModelTime or Tdis, got {type(value)}")
+
+
 @xattree
 class Simulation(Component):
     models: dict[str, Model] = field()
     exchanges: dict[str, Exchange] = field()
     solutions: dict[str, Solution] = field()
-    tdis: Tdis = field()
+    tdis: Tdis = field(converter=convert_time)
     sim_ws: Path = field(default=None)
+
+    @property
+    def time(self) -> ModelTime:
+        return self.tdis.to_time()

flopy4/mf6/tdis.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from attrs import Converter, define
+from flopy.discretization.modeltime import ModelTime
 from numpy.typing import NDArray
 from xattree import ROOT, array, dim, field, xattree
 
@@ -48,3 +49,26 @@ class PeriodData:
         metadata={"block": "perioddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
+
+    def to_time(self) -> ModelTime:
+        """Convert to a `ModelTime` object."""
+        return ModelTime(
+            nper=self.nper,
+            time_units=self.time_units,
+            start_date_time=self.start_date_time,
+            perlen=self.perlen,
+            nstp=self.nstp,
+            tsmult=self.tsmult,
+        )
+
+    @classmethod
+    def from_time(cls, time: ModelTime) -> "Tdis":
+        """Create a time discretization from a `ModelTime`."""
+        return cls(
+            nper=time.nper,
+            time_units=time.time_units,
+            start_date_time=time.start_datetime,
+            perlen=time.perlen,
+            nstp=time.nstp,
+            tsmult=time.tsmult,
+        )