first draft of array converters

Author: wpbonelli

docs/examples/quickstart.py

@@ -1,19 +1,19 @@
-from flopy4.mf6 import Sim, Tdis
+from flopy4.mf6 import Simulation, Tdis
 from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc
 
 # TODO rewrite bottom up
 
 ws = "./mymodel"
 name = "mymodel"
-sim = Sim(name=name, path=ws, exe="mf6")
+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)
 chd = Chd(
     gwf,
-    stress_period_data=[[(0, 0, 0), 1.0], [(0, 9, 9), 0.0]],
+    head={(0, 0, 0): 1.0, (0, 9, 9): 0.0},
 )
 
 # list input in the mf6 paradigm, just stored in xarray.

docs/examples/quickstart_expanded.py

@@ -26,12 +26,12 @@
 changes to support.
 """
 
-from flopy4.mf6 import Sim, Tdis
+from flopy4.mf6 import Simulation, Tdis
 from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc
 
 ws = "./mymodel"
 name = "mymodel"
-sim = Sim(name=name, path=ws, exe="mf6")
+sim = Simulation(name=name, path=ws, exe="mf6")
 tdis = Tdis(sim)
 gwf = Gwf(sim, name=name, save_flows=True)
 dis = Dis(gwf, nrow=10, ncol=10)

flopy4/mf6/gwf/chd.py

@@ -1,13 +1,42 @@
 from pathlib import Path
 from typing import Optional
 
+import attrs
 import numpy as np
 from attrs import define
 from numpy.typing import NDArray
-from xattree import array, field, xattree
+from xattree import _get_xatspec, array, field, xattree
 
 from flopy4.mf6 import Package
 
+dnodata = 1e30
+
+
+def _get_nn(ncol, nrow, k, i, j):
+    return k * nrow * ncol + i * ncol + j
+
+
+def _convert_array(value, self_, field):
+    if not isinstance(value, dict):
+        return value
+
+    inherited_dims = self_.__dict__.get("dims", {})
+    spec = _get_xatspec(type(self_))
+    field = spec.arrays["head"]
+    shape = field.dims
+    if not shape:
+        raise ValueError()
+    dims = [inherited_dims.get(d, d) for d in shape]
+    # TODO pull out dtype from annotation
+    a = np.full(dims, fill_value=dnodata, dtype=np.float64)
+    for kper, period in value.items():
+        if kper == "*":
+            kper = 0
+        for cellid, v in period.items():
+            nn = _get_nn(inherited_dims["col"], inherited_dims["row"], *cellid)
+            a[kper, nn] = v
+    return a
+
 
 @xattree(multi="list")
 class Chd(Package):
@@ -46,6 +75,9 @@ class Steps:
         ),
         default=None,
         metadata={"block": "period"},
+        converter=attrs.Converter(
+            _convert_array, takes_self=True, takes_field=True
+        ),
     )
     aux: Optional[NDArray[np.floating]] = array(
         dims=(

test/test_component.py

@@ -79,7 +79,7 @@ def test_init_gwf_dis_first():
 
 def test_init_sim():
     time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0])
-    grid = StructuredGrid(nlay=1, nrow=2, ncol=3)
+    grid = StructuredGrid(nlay=1, nrow=10, ncol=10)
     dims = {
         "per": time.nper,
         "lay": grid.nlay,
@@ -92,7 +92,7 @@ def test_init_sim():
     ic = Ic(dims=dims)
     oc = Oc(dims=dims)
     npf = Npf(dims=dims)
-    chd = Chd(dims=dims)
+    chd = Chd(dims=dims, head={"*": {(0, 0, 0): 1.0, (0, 9, 9): 0.0}})
     gwf = Gwf(
         dis=dis,
         ic=ic,
@@ -114,5 +114,5 @@ def test_init_sim():
     assert gwf.oc is oc
     assert gwf.npf is npf
     assert gwf.chd[0] is chd
-    assert np.array_equal(sim.models["gwf"].npf.k, np.ones(6))
-    assert np.array_equal(sim.models["gwf"].npf.data.k, np.ones(6))
+    assert np.array_equal(sim.models["gwf"].npf.k, np.ones(100))
+    assert np.array_equal(sim.models["gwf"].npf.data.k, np.ones(100))