codegen baseline (#310)

Author: mjreno

docs/dev/dfn-schema-plan.md

@@ -345,7 +345,7 @@ Type `integer`.
 
 ###### `time_series`
 
-`boolean (default: false)`. Marks fields where the parser accepts either a numeric literal or a time-series name (referencing a `utl-ts` object). Not inferrable from structural type. Also appears on array fields (where it references a `utl-tas` object instead). Note that `utl-tas` currently only works with layered arrays, not full-grid arrays, though generalizing has been considered.
+`boolean (default: false)`. Marks fields where the parser accepts either a numeric literal or a time-series name (referencing a `utl-ts` object). Not inferable from structural type. Also appears on array fields (where it references a `utl-tas` object instead). Note that `utl-tas` currently only works with layered arrays, not full-grid arrays, though generalizing has been considered.
 
 ###### `pk`
 
@@ -371,7 +371,7 @@ Type `double`.
 
 ###### `time_series`
 
-`boolean (default: false)`. Marks fields where the parser accepts either a numeric literal or a time-series name (referencing a `utl-ts` object). Not inferrable from structural type. Also appears on array fields (where it references a `utl-tas` object instead). Note that `utl-tas` currently only works with layered arrays, not full-grid arrays, though generalizing has been considered.
+`boolean (default: false)`. Marks fields where the parser accepts either a numeric literal or a time-series name (referencing a `utl-ts` object). Not inferable from structural type. Also appears on array fields (where it references a `utl-tas` object instead). Note that `utl-tas` currently only works with layered arrays, not full-grid arrays, though generalizing has been considered.
 
 #### Path
 

docs/examples/circle.py

@@ -248,7 +248,7 @@ def plot_head_ugrid(head, cbc, workspace):
     outer_maximum=500,
     under_relaxation=None,
     inner_dvclose=1.0e-4,
-    inner_rclose=0.001,
+    rclose=flopy4.mf6.Ims.Rclose(inner_rclose=0.001),
     inner_maximum=100,
     linear_acceleration="cg",
     reordering_method=None,

docs/examples/frenchman-flat.py

@@ -656,38 +656,13 @@ def plot_head_ugrid(head, cbc, grid, workspace):
     dims=dims,
 )
 
-# Save heads at every time step; save budget only at selected steps to keep
-# output file size manageable for this large model.
+# Save heads at every time step; save budget at the last step of each period.
 oc = flopy4.mf6.gwf.Oc(
     budget_file=Path("ff.cbc"),
     head_file=Path("ff.hds"),
-    perioddata={
-        0: flopy4.mf6.gwf.Oc.PrintSaveSetting(
-            printrecord=[
-                flopy4.mf6.gwf.Oc.PrintRecord(
-                    "budget",
-                    flopy4.mf6.gwf.Oc.Steps(
-                        steps=(
-                            0,
-                            99,
-                        )
-                    ),
-                ),
-            ],
-            saverecord=[
-                flopy4.mf6.gwf.Oc.SaveRecord("head", flopy4.mf6.gwf.Oc.Steps(all=True)),
-                flopy4.mf6.gwf.Oc.SaveRecord("budget", flopy4.mf6.gwf.Oc.Steps(steps=(0,))),
-            ],
-        ),
-        1: flopy4.mf6.gwf.Oc.PrintSaveSetting(
-            printrecord=[
-                flopy4.mf6.gwf.Oc.PrintRecord("budget", flopy4.mf6.gwf.Oc.Steps(last=True)),
-            ],
-            saverecord=[
-                flopy4.mf6.gwf.Oc.SaveRecord("head", flopy4.mf6.gwf.Oc.Steps(all=True)),
-            ],
-        ),
-    },
+    save_head={"0": "all", 1: "all"},
+    save_budget={"0": "STEPS 1"},
+    print_budget={"0": "STEPS 1 15", 1: "last"},
     dims=dims,
 )
 
@@ -719,7 +694,7 @@ def plot_head_ugrid(head, cbc, grid, workspace):
     under_relaxation_gamma=0.000000,
     under_relaxation_momentum=0.000000,
     inner_dvclose=0.00001,
-    inner_rclose=0.1,
+    rclose=flopy4.mf6.Ims.Rclose(inner_rclose=0.1),
     inner_maximum=100,
     linear_acceleration="bicgstab",
     number_orthogonalizations=0,

docs/examples/quickstart.py

@@ -64,7 +64,16 @@
 
 sim = Simulation(name=name, workspace=workspace, tdis=time)
 gwf_name = "mymodel"
-ims = Ims(parent=sim, models=[gwf_name])  # registered with sim; references the model by name
+ims = Ims(
+    parent=sim,
+    models=[gwf_name],
+    outer_dvclose=1e-3,
+    outer_maximum=25,
+    inner_maximum=50,
+    inner_dvclose=1e-3,
+    rclose=Ims.Rclose(inner_rclose=0.1),
+    linear_acceleration="cg",
+)  # registered with sim; references the model by name
 gwf = Gwf(parent=sim, name=gwf_name, save_flows=True, dis=grid)
 
 # Node-property flow: isotropic conductivity; saves specific-discharge for quiver plots.

docs/examples/quickstart_expanded.py

@@ -91,8 +91,8 @@
 head_file = name + ".hds"
 oc = Oc(
     parent=gwf,
-    budget_filerecord=budget_file,
-    head_filerecord=head_file,
+    budget_file=budget_file,
+    head_file=head_file,
     # 1) tuples (like flopy3)
     perioddata=[("HEAD", "ALL"), ("BUDGET", "ALL")],
     # 2) typed records

docs/examples/twri.py

@@ -118,7 +118,7 @@ def plot_head(head, workspace):
     icelltype=icelltype,
     k=k,
     k33=k33,
-    cvoptions=flopy4.mf6.gwf.Npf.CvOptions(dewatered=True),
+    cvoptions=flopy4.mf6.gwf.Npf.Cvoptions(dewatered=True),
     perched=True,
     save_flows=True,
     dims=dims,
@@ -199,7 +199,7 @@ def plot_head(head, workspace):
     outer_maximum=500,
     under_relaxation=None,
     inner_dvclose=1.0e-4,
-    inner_rclose=0.001,
+    rclose=flopy4.mf6.Ims.Rclose(inner_rclose=0.001),
     inner_maximum=100,
     linear_acceleration="cg",
     scaling_method=None,

flopy4/mf6/__init__.py

@@ -6,18 +6,35 @@
 from tomli_w import dump as dump_toml
 
 # Import submodules to make them accessible via flopy4.mf6.*
-from flopy4.mf6 import gwf, simulation, solution, utils
+from flopy4.mf6 import gwe, gwf, gwt, prt, simulation, solution, utils
 from flopy4.mf6.codec import dump as dump_mf6
 from flopy4.mf6.codec import load as load_mf6
 from flopy4.mf6.component import Component
 from flopy4.mf6.converter import structure, unstructure
+from flopy4.mf6.ems import Ems
+from flopy4.mf6.exchange import GwfGwe, GwfGwt
 from flopy4.mf6.ims import Ims
 from flopy4.mf6.netcdf import NetCDFModel
 from flopy4.mf6.simulation import Simulation
 from flopy4.mf6.tdis import Tdis
 from flopy4.uio import DEFAULT_REGISTRY
 
-__all__ = ["gwf", "simulation", "solution", "utils", "Ims", "NetCDFModel", "Tdis", "Simulation"]
+__all__ = [
+    "gwf",
+    "gwt",
+    "gwe",
+    "prt",
+    "simulation",
+    "solution",
+    "utils",
+    "Ems",
+    "GwfGwe",
+    "GwfGwt",
+    "Ims",
+    "NetCDFModel",
+    "Tdis",
+    "Simulation",
+]
 
 
 class WriteError(Exception):

flopy4/mf6/binding.py

@@ -29,7 +29,7 @@ def from_component(cls, component: Component) -> "Binding":
         def _get_binding_type(component: Component) -> str:
             cls_name = component.__class__.__name__
             if isinstance(component, Exchange):
-                return f"{'-'.join([cls_name[:2], cls_name[3:]]).upper()}6"
+                return f"{cls_name[:3].upper()}6-{cls_name[3:].upper()}6"
             elif isinstance(component, Solution):
                 return f"{component.slntype}6"
             else:

flopy4/mf6/codec/writer/filters.py

@@ -243,38 +243,64 @@ def dataset2list(value: xr.Dataset):
                 yield (*name.split("_"), val)  # type: ignore
 
         else:
-            vals = []
-            for name in value.data_vars.keys():
-                val = value[name]
-                val = val.item() if val.shape == () else val
-                vals.append(val)
-            yield tuple(vals)
+            row: list[Any] = []
+            for name, da in value.data_vars.items():
+                val = da.item() if da.shape == () else da
+                if kw := da.attrs.get("row_keyword", False):
+                    if val:
+                        row.append(kw if isinstance(kw, str) else str(name).upper())
+                else:
+                    row.extend(da.attrs.get("prefix", ()))
+                    row.append(val)
+            yield tuple(row)
         return
 
     combined_mask: Any = None
+    spatial_da = None  # a non-aux DataArray for deriving spatial dims/mask
     for name, first in value.data_vars.items():
-        mask = nonempty(first)
-        combined_mask = mask if combined_mask is None else combined_mask | mask
+        if "naux" in first.dims:
+            # nonempty gives (spatial..., naux) bool array; collapse naux with any()
+            mask = nonempty(first).any(axis=-1)
+        else:
+            mask = nonempty(first)
+            spatial_da = first
+        combined_mask = mask if combined_mask is None else (combined_mask | mask)
     if combined_mask is None or not np.any(combined_mask):
         return
 
-    spatial_dims = [d for d in first.dims if d in ("nlay", "nrow", "ncol", "nodes")]
+    if spatial_da is None:
+        spatial_da = first
+    spatial_dims = [d for d in spatial_da.dims if d in ("nlay", "nrow", "ncol", "nodes")]
     has_spatial_dims = len(spatial_dims) > 0
     indices = np.where(combined_mask)
     for i in range(len(indices[0])):
         if is_oc:
             for name in value.data_vars.keys():
                 val = value[name][tuple(idx[i] for idx in indices)]
                 val = val.item() if val.shape == () else val
-                yield (*name.split("_"), val)  # type: ignore
+                yield (*str(name).split("_"), val)  # type: ignore
         else:
-            vals = []
-            for name in value.data_vars.keys():
-                val = value[name][tuple(idx[i] for idx in indices)]
+            row2: list[Any] = []
+            for name, da in value.data_vars.items():
+                val = da[tuple(idx[i] for idx in indices)]
                 val = val.item() if val.shape == () else val
-                vals.append(val)
+                if kw := da.attrs.get("row_keyword", False):
+                    if val:
+                        row2.append(kw if isinstance(kw, str) else str(name).upper())
+                else:
+                    row2.extend(da.attrs.get("prefix", ()))
+                    if da.attrs.get("cellid"):
+                        if isinstance(val, tuple):
+                            row2.extend(c + 1 for c in val)
+                        else:
+                            row2.append(val + 1)
+                    else:
+                        if hasattr(val, "ndim") and val.ndim > 0:
+                            row2.extend(float(v) for v in np.asarray(val).flat)
+                        else:
+                            row2.append(val)
             if has_spatial_dims:
                 cellid = tuple(idx[i] + 1 for idx in indices)
-                yield cellid + tuple(vals)
+                yield tuple(cellid) + tuple(row2)
             else:
-                yield tuple(vals)
+                yield tuple(row2)

flopy4/mf6/component.py

@@ -15,7 +15,7 @@
 from flopy4.mf6.write_context import WriteContext
 from flopy4.uio import IO, Loader, Writer
 
-COMPONENTS = {}
+COMPONENTS: dict[str, type] = {}
 """MF6 component registry."""
 
 
@@ -95,7 +95,14 @@ def _update_maxbound_if_needed(self):
 
     @classmethod
     def __attrs_init_subclass__(cls):
-        COMPONENTS[cls.__name__.lower()] = cls
+        key = cls.__name__.lower()
+        COMPONENTS[key] = cls
+        # Also register a model-qualified key (e.g. "gwf-ic") for classes in a
+        # model subpackage (flopy4.mf6.<model>.<pkg>), giving deterministic
+        # per-model lookup when multiple models share a class name like "ic".
+        parts = cls.__module__.split(".")
+        if len(parts) >= 4 and parts[0] == "flopy4" and parts[1] == "mf6":
+            COMPONENTS[f"{parts[2]}-{key}"] = cls
 
     def __getitem__(self, key):
         # We use `children` from `xattree` to implement MutableMapping.