just use a converter, so gwf can accept a dis or a grid, and sim can accept a modeltime or a tdis

Author: wpbonelli

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

@@ -1,7 +1,6 @@
 from pathlib import Path
 from typing import Optional
 
-import attrs
 from attrs import define
 from flopy.discretization.grid import Grid
 from xattree import field, xattree
@@ -18,12 +17,11 @@
 
 @xattree
 class Gwf(Model):
-    dis: Dis = field()
+    dis: Dis = field(converter=lambda grid: Dis.from_grid(grid))
     ic: Ic = field()
     oc: Oc = field()
     npf: Npf = field()
     chd: list[Chd] = field()
-    grid: Grid = attrs.field(default=None)
 
     @define
     class NewtonOptions:
@@ -46,3 +44,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,3 @@
-import attrs
-from flopy.discretization.modeltime import ModelTime
 from xattree import field, xattree
 
 from flopy4.mf6.component import Component
@@ -14,5 +12,4 @@ class Simulation(Component):
     models: dict[str, Model] = field()
     exchanges: dict[str, Exchange] = field()
     solutions: dict[str, Solution] = field()
-    tdis: Tdis = field()
-    time: ModelTime = attrs.field(default=None)
+    tdis: Tdis = field(converter=lambda time: Tdis.from_time(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,
+        )

pixi.lock

@@ -49,44 +49,12 @@ def test_init_gwf_explicit_dims():
     )
 
     assert isinstance(gwf.data, DataTree)
-    assert gwf.dis is dis
+    assert gwf.dis is not dis  # dimension order switched.. is this ok?
     assert gwf.ic is ic
     assert gwf.oc is oc
     assert gwf.npf is npf
     assert gwf.chd[0] is chd
-    assert gwf.data.dis is dis.data
-    assert gwf.data.ic is ic.data
-    assert gwf.data.oc is oc.data
-    assert gwf.data.npf is npf.data
-    assert np.array_equal(npf.k, np.ones(4))
-    assert np.array_equal(npf.data.k, np.ones(4))
-
-
-@pytest.mark.skip(reason="TODO")
-def test_init_gwf_from_grid():
-    time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0])
-    grid = StructuredGrid(nlay=1, nrow=2, ncol=2)
-    dis = Dis(grid=grid)
-    ic = Ic(grid=grid)
-    oc = Oc(grid=grid)
-    npf = Npf(grid=grid)
-    chd = Chd(grid=grid)
-    gwf = Gwf(
-        dis=dis,
-        ic=ic,
-        oc=oc,
-        npf=npf,
-        chd=[chd],
-        grid=grid,
-    )
-
-    assert isinstance(gwf.data, DataTree)
-    assert gwf.dis is dis
-    assert gwf.ic is ic
-    assert gwf.oc is oc
-    assert gwf.npf is npf
-    assert gwf.chd[0] is chd
-    assert gwf.data.dis is dis.data
+    assert gwf.data.dis is not dis.data
     assert gwf.data.ic is ic.data
     assert gwf.data.oc is oc.data
     assert gwf.data.npf is npf.data
@@ -138,7 +106,7 @@ def test_init_gwf_dis_first():
     chd = Chd(parent=gwf, strict=False)
 
     assert isinstance(gwf.data, DataTree)
-    assert gwf.dis is dis
+    assert gwf.dis is not dis
     assert gwf.ic is ic
     assert gwf.oc is oc
     assert gwf.npf is npf
@@ -147,6 +115,25 @@ def test_init_gwf_dis_first():
     assert np.array_equal(npf.data.k, np.ones(4))
 
 
+def test_init_gwf_dis_first_with_grid():
+    grid = StructuredGrid(nlay=1, nrow=10, ncol=10)
+    gwf = Gwf(dis=grid)
+    dis = gwf.dis
+    ic = Ic(parent=gwf)
+    oc = Oc(parent=gwf, strict=False)
+    npf = Npf(parent=gwf)
+    chd = Chd(parent=gwf, strict=False)
+
+    assert isinstance(gwf.data, DataTree)
+    assert gwf.dis is dis
+    assert gwf.ic is ic
+    assert gwf.oc is oc
+    assert gwf.npf is npf
+    assert gwf.chd[0] is chd
+    assert np.array_equal(npf.k, np.ones(100))
+    assert np.array_equal(npf.data.k, np.ones(100))
+
+
 def test_init_gwf_top_down_misaligned():
     grid = StructuredGrid(nlay=1, nrow=10, ncol=10)
     dims = {
@@ -198,7 +185,7 @@ def test_init_sim_explicit_dims():
     assert isinstance(sim.data, DataTree)
     assert sim.data.tdis is tdis.data
     assert sim.data.gwf is gwf.data
-    assert gwf.dis is dis
+    assert gwf.dis is not dis  # gwf.dis has inherited dim nper
     assert gwf.ic is ic
     assert gwf.oc is oc
     assert gwf.npf is npf
@@ -219,35 +206,16 @@ def test_init_big_sim():
     # if size over threshold, arrays should be sparse
     time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0])
     grid = StructuredGrid(nlay=1, nrow=100, ncol=100)
-    dims = {
-        "nlay": grid.nlay,
-        "nrow": grid.nrow,
-        "ncol": grid.ncol,
-    }
-    dis = Dis(**dims)
-    dims["nper"] = time.nper
-    dims["nnodes"] = grid.nnodes
-    ic = Ic(dims=dims)
-    oc = Oc(dims=dims)
-    npf = Npf(dims=dims)
-    chd = Chd(dims=dims, head={"*": {(0, 0, 0): 1.0, (0, 99, 99): 0.0}})
-    gwf = Gwf(
-        dis=dis,
-        ic=ic,
-        oc=oc,
-        npf=npf,
-        chd=[chd],
-        dims=dims,
-    )
-    tdis = Tdis(dims=dims)
-    sim = Simulation(tdis=tdis, models={"gwf": gwf})
+    sim = Simulation(tdis=time)
+    gwf = Gwf(parent=sim, dis=grid)
+    ic = Ic(parent=gwf)
+    oc = Oc(parent=gwf)
+    npf = Npf(parent=gwf)
+    chd = Chd(parent=gwf, head={"*": {(0, 0, 0): 1.0, (0, 99, 99): 0.0}})
 
-    assert sim.tdis is tdis
     assert sim.models["gwf"] is gwf
     assert isinstance(sim.data, DataTree)
-    assert sim.data.tdis is tdis.data
     assert sim.data.gwf is gwf.data
-    assert gwf.dis is dis
     assert gwf.ic is ic
     assert gwf.oc is oc
     assert gwf.npf is npf