quickstart closer, update tests

Author: wpbonelli

docs/examples/quickstart.py

@@ -1,26 +1,22 @@
-from flopy4.mf6 import Simulation, Tdis
-from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc
+import numpy as np
+
+from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf
+from flopy4.mf6.simulation import Simulation
+from flopy4.mf6.tdis import Tdis
 
 # TODO rewrite bottom up
 
 ws = "./mymodel"
 name = "mymodel"
-sim = Simulation(name=name)
-tdis = Tdis(sim)
-gwf = Gwf(sim, name=name, save_flows=True)
-dis = Dis(gwf, nrow=10, ncol=10)
-ic = Ic(gwf)
-npf = Npf(gwf, save_specific_discharge=True)
+tdis = Tdis()
+sim = Simulation(name=name, tdis=tdis)
+dis = Dis(nrow=10, ncol=10)
+gwf = Gwf(parent=sim, name=name, save_flows=True, dis=dis)
+ic = Ic(parent=gwf)
+npf = Npf(parent=gwf, save_specific_discharge=True)
 chd = Chd(
-    gwf,
-    head={(0, 0, 0): 1.0, (0, 9, 9): 0.0},
-)
-
-# list input in the mf6 paradigm, just stored in xarray.
-# this is straightforward to implement. even in flopy3?
-assert all(
-    period == Chd.Period((0, 0, 0), 1.0)
-    for period in chd.data["stress_period_data"]
+    parent=gwf,
+    head={"*": {(0, 0, 0): 1.0, (0, 9, 9): 0.0}},
 )
 
 # adopt xarray paradigm. transpose lists to arrays. this
@@ -30,15 +26,22 @@
 # we want to "align" everything to the discretization in
 # both time and space.
 
+assert chd.data["head"][0, 0] == 1.0
+assert chd.data["head"][0, 99] == 0.0
+assert np.allclose(chd.data["head"][:, 1:99], np.full(98, 1e30))
+
+# TODO: xarray index aliasing nlay/ncol/nrow to k/i/j?
 # assert chd.data["head"].loc(dict(k=0, i=0, j=0)) == 1.
 # assert chd.data["head"].loc(dict(k=0, i=9, j=9)) == 0.
 
-oc = Oc(
-    gwf,
-    budget_filerecord=f"{name}.bud",
-    head_filerecord=f"{name}.hds",
-    saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")],
-)
+# TODO OC!
+
+# oc = Oc(
+#     gwf,
+#     budget_filerecord=f"{name}.bud",
+#     head_filerecord=f"{name}.hds",
+#     saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")],
+# )
 
 # xarray style. this is how imod-python does it too.
 # assert oc.data["save_head"] == "all"

flopy4/mf6/converters.py

@@ -19,7 +19,7 @@ def convert_array(value, self_, field) -> NDArray:
 
     # resolve dims
     explicit_dims = self_.__dict__.get("dims", {})
-    inherited_dims = self_.parent.data.dims if self_.parent else {}
+    inherited_dims = dict(self_.parent.data.dims) if self_.parent else {}
     dims = inherited_dims | explicit_dims
     shape = [dims.get(d, d) for d in field.dims]
     unresolved = [d for d in shape if isinstance(d, str)]
@@ -38,13 +38,13 @@ def _get_nn(cellid):
                 return k * dims["ncpl"] + j
             case 3:
                 k, i, j = cellid
-                return k * dims["row"] * dims["col"] + i * dims["col"] + j
+                return k * dims["nrow"] * dims["ncol"] + i * dims["ncol"] + j
             case _:
                 raise ValueError(f"Invalid cellid: {cellid}")
 
     # populate array. TODO: is there a way to do this
     # without hardcoding awareness of kper and cellid?
-    if "per" in dims:
+    if "nper" in dims:
         for kper, period in value.items():
             if kper == "*":
                 kper = 0

flopy4/mf6/gwf/chd.py

@@ -42,34 +42,34 @@ class Steps:
     )
     head: Optional[NDArray[np.floating]] = array(
         dims=(
-            "per",
-            "node",
+            "nper",
+            "nnodes",
         ),
         default=None,
         metadata={"block": "period"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     aux: Optional[NDArray[np.floating]] = array(
         dims=(
-            "per",
-            "node",
+            "nper",
+            "nnodes",
         ),
         default=None,
         metadata={"block": "period"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     boundname: Optional[NDArray[np.str_]] = array(
         dims=(
-            "per",
-            "node",
+            "nper",
+            "nnodes",
         ),
         default=None,
         metadata={"block": "period"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     steps: Optional[NDArray[np.object_]] = array(
         Steps,
-        dims=("per", "node"),
+        dims=("nper", "nnodes"),
         default=None,
         metadata={"block": "period"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),

flopy4/mf6/gwf/dis.py

@@ -21,60 +21,57 @@ class Dis(Package):
         default=False, metadata={"block": "options"}
     )
     nlay: int = dim(
-        name="lay",
         scope="gwf",
         default=1,
         metadata={
             "block": "dimensions",
         },
     )
     ncol: int = dim(
-        name="col",
         scope="gwf",
         default=2,
         metadata={
             "block": "dimensions",
         },
     )
     nrow: int = dim(
-        name="row",
         scope="gwf",
         default=2,
         metadata={
             "block": "dimensions",
         },
     )
     delr: NDArray[np.floating] = array(
-        dims=("col",),
+        dims=("ncol",),
         default=1.0,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     delc: NDArray[np.floating] = array(
-        dims=("row",),
+        dims=("nrow",),
         default=1.0,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     top: NDArray[np.floating] = array(
-        dims=("col", "row"),
+        dims=("ncol", "nrow"),
         default=1.0,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     botm: NDArray[np.floating] = array(
-        dims=("col", "row", "lay"),
+        dims=("ncol", "nrow", "nlay"),
         default=0.0,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     idomain: NDArray[np.integer] = array(
-        dims=("col", "row", "lay"),
+        dims=("ncol", "nrow", "nlay"),
         default=1,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
-    nnodes: int = dim(name="node", scope="gwf", init=False)
+    nnodes: int = dim(scope="gwf", init=False)
 
     def __attrs_post_init__(self):
         self.nnodes = self.ncol * self.nrow * self.nlay

flopy4/mf6/gwf/ic.py

@@ -10,7 +10,7 @@
 @xattree
 class Ic(Package):
     strt: NDArray[np.floating] = array(
-        dims=("node",),
+        dims=("nnodes",),
         default=1.0,
         metadata={"block": "packagedata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),

flopy4/mf6/gwf/npf.py

@@ -65,49 +65,49 @@ class Xt3dOptions:
         default=None, metadata={"block": "options"}
     )
     icelltype: NDArray[np.integer] = array(
-        dims=("node",),
+        dims=("nnodes",),
         default=0,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     k: NDArray[np.floating] = array(
-        dims=("node",),
+        dims=("nnodes",),
         default=1.0,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     k22: Optional[NDArray[np.floating]] = array(
-        dims=("node",),
+        dims=("nnodes",),
         default=None,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     k33: Optional[NDArray[np.floating]] = array(
-        dims=("node",),
+        dims=("nnodes",),
         default=None,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     angle1: Optional[NDArray[np.floating]] = array(
-        dims=("node",),
+        dims=("nnodes",),
         default=None,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     angle2: Optional[NDArray[np.floating]] = array(
-        dims=("node",),
+        dims=("nnodes",),
         default=None,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     angle3: Optional[NDArray[np.floating]] = array(
-        dims=("node",),
+        dims=("nnodes",),
         default=None,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     wetdry: Optional[NDArray[np.floating]] = array(
-        dims=("node",),
+        dims=("nnodes",),
         default=None,
         metadata={"block": "griddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),

flopy4/mf6/gwf/oc.py

@@ -55,14 +55,14 @@ class Period:
     )
     saverecord: Optional[NDArray[np.object_]] = array(
         Period,
-        dims=("per",),
+        dims=("nper",),
         default=None,
         metadata={"block": "perioddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     printrecord: Optional[NDArray[np.object_]] = array(
         Period,
-        dims=("per",),
+        dims=("nper",),
         default=None,
         metadata={"block": "perioddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),

flopy4/mf6/tdis.py

@@ -19,7 +19,6 @@ class PeriodData:
         tsmult: float
 
     nper: int = dim(
-        name="per",
         default=1,
         scope=ROOT,
         metadata={"block": "dimensions"},
@@ -32,19 +31,19 @@ class PeriodData:
     )
     perlen: NDArray[np.floating] = array(
         default=1.0,
-        dims=("per",),
+        dims=("nper",),
         metadata={"block": "perioddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     nstp: NDArray[np.integer] = array(
         default=1,
-        dims=("per",),
+        dims=("nper",),
         metadata={"block": "perioddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )
     tsmult: NDArray[np.floating] = array(
         default=1.0,
-        dims=("per",),
+        dims=("nper",),
         metadata={"block": "perioddata"},
         converter=Converter(convert_array, takes_self=True, takes_field=True),
     )

test/test_component.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 from flopy.discretization import StructuredGrid
 from flopy.discretization.modeltime import ModelTime
 from xarray import DataTree
@@ -27,13 +28,12 @@ def test_init_gwf_explicit_dims():
     time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0])
     grid = StructuredGrid(nlay=1, nrow=2, ncol=2)
     dims = {
-        "per": time.nper,
-        "lay": grid.nlay,
-        "row": grid.nrow,
-        "col": grid.ncol,
-        "node": grid.nnodes,
+        "nper": time.nper,
+        "nlay": grid.nlay,
+        "nrow": grid.nrow,
+        "ncol": grid.ncol,
+        "nnodes": grid.nnodes,
     }
-
     dis = Dis(dims=dims)
     ic = Ic(dims=dims)
     oc = Oc(dims=dims)
@@ -80,18 +80,37 @@ def test_init_gwf_dis_first():
     assert np.array_equal(npf.data.k, np.ones(4))
 
 
-def test_init_sim():
-    time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0])
+def test_init_gwf_top_down_misaligned():
     grid = StructuredGrid(nlay=1, nrow=10, ncol=10)
     dims = {
-        "per": time.nper,
-        "lay": grid.nlay,
-        "row": grid.nrow,
-        "col": grid.ncol,
-        "node": grid.nnodes,
+        "nrow": grid.nrow,
+        "ncol": grid.ncol,
     }
+    gwf = Gwf()
+    with pytest.raises(
+        ValueError, match=r"group '/dis' is not aligned with its parents"
+    ):
+        Dis(parent=gwf, **dims)
 
-    dis = Dis(dims=dims)
+    # passing dims explicitly to gwf doesn't work either.
+    # one MUST create the component declaring dims first.
+    with pytest.raises(
+        ValueError, match=r"group '/dis' is not aligned with its parents"
+    ):
+        Gwf(dims=dims)
+
+
+def test_init_sim_explicit_dims():
+    time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0])
+    grid = StructuredGrid(nlay=1, nrow=10, ncol=10)
+    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)