writer child binding

Author: wpbonelli

flopy4/mf6/__init__.py

@@ -1,8 +1,10 @@
 from flopy4.mf6.codec import dump, load
 from flopy4.mf6.component import Component
+from flopy4.mf6.converter import COMPONENT_CONVERTER
 from flopy4.uio import DEFAULT_REGISTRY
 
 # register io methods
-# TODO: call this "mf6" or something? since it might include binary files
-DEFAULT_REGISTRY.register_loader(Component, "ascii", lambda c: load(c.path))
-DEFAULT_REGISTRY.register_writer(Component, "ascii", lambda c: dump(c, c.path))
+# TODO: call format "mf6" or something? since it might include binary files
+DEFAULT_REGISTRY.register_writer(
+    Component, "ascii", lambda c: dump(COMPONENT_CONVERTER.unstructure(c), c.path)
+)

flopy4/mf6/context.py

@@ -1,6 +1,7 @@
 from abc import ABC
 from pathlib import Path
 
+from modflow_devtools.misc import cd
 from xattree import xattree
 
 from flopy4.mf6.component import Component
@@ -14,8 +15,20 @@ class Context(Component, ABC):
     def __attrs_post_init__(self):
         super().__attrs_post_init__()
         if self.workspace is None:
-            self.workspace = Path.cwd()
+            self.workspace = (
+                self.parent.workspace
+                if self.parent and hasattr(self.parent, "workspace")
+                else Path.cwd()
+            )
 
     @property
     def path(self) -> Path:
         return self.workspace / self.filename
+
+    def load(self, format="ascii"):
+        with cd(self.workspace):
+            super().load(format=format)
+
+    def write(self, format="ascii"):
+        with cd(self.workspace):
+            super().write(format=format)

flopy4/mf6/converter.py

@@ -1,15 +1,62 @@
+from collections.abc import MutableMapping
 from datetime import datetime
 from pathlib import Path
 from typing import Any
 
 import xarray as xr
 import xattree
+from attrs import define
 from cattrs import Converter
 
 from flopy4.mf6.component import Component
+from flopy4.mf6.context import Context
+from flopy4.mf6.exchange import Exchange
+from flopy4.mf6.model import Model
+from flopy4.mf6.package import Package
 from flopy4.mf6.spec import fields_dict, get_blocks
 
 
+@define
+class _Binding:
+    """
+    An MF6 component binding: a record representation of the
+    component for writing to a parent component's name file.
+    """
+
+    type: str
+    fname: str
+    terms: tuple[str, ...] | None = None
+
+    def to_tuple(self):
+        if self.terms and any(self.terms):
+            return (self.type, self.fname, *self.terms)
+        else:
+            return (self.type, self.fname)
+
+    @classmethod
+    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"
+            else:
+                return f"{cls_name.upper()}6"
+
+        def _get_binding_terms(component: Component) -> tuple[str, ...] | None:
+            if isinstance(component, Exchange):
+                return (component.exgmnamea, component.exgmnameb)  # type: ignore
+            elif isinstance(component, (Model, Package)):
+                return (component.name,)  # type: ignore
+            # TODO solutions
+            return None
+
+        return cls(
+            type=_get_binding_type(component),
+            fname=component.filename or component.default_filename(),
+            terms=_get_binding_terms(component),
+        )
+
+
 def _attach_field_metadata(
     dataset: xr.Dataset, component_type: type, field_names: list[str]
 ) -> None:
@@ -29,14 +76,56 @@ def _path_to_record(field_name: str, path_value: Path) -> tuple:
 
 
 def unstructure_component(value: Component) -> dict[str, Any]:
-    data = xattree.asdict(value)
     blockspec = get_blocks(value.dfn)
     blocks: dict[str, dict[str, Any]] = {}
+    xatspec = xattree.get_xatspec(type(value))
+
+    # Handle child component bindings before converting to dict
+    if isinstance(value, Context):
+        for field_name, child_spec in xatspec.children.items():
+            if hasattr(child_spec, "metadata") and "block" in child_spec.metadata:  # type: ignore
+                block_name = child_spec.metadata["block"]  # type: ignore
+                field_value = getattr(value, field_name, None)
+
+                if block_name not in blocks:
+                    blocks[block_name] = {}
+
+                if isinstance(field_value, Component):
+                    components = [_Binding.from_component(field_value).to_tuple()]
+                elif isinstance(field_value, MutableMapping):
+                    components = [
+                        _Binding.from_component(comp).to_tuple()
+                        for comp in field_value.values()
+                        if comp is not None
+                    ]
+                elif isinstance(field_value, (list, tuple)):
+                    components = [
+                        _Binding.from_component(comp).to_tuple()
+                        for comp in field_value
+                        if comp is not None
+                    ]
+                else:
+                    continue
+
+                if components:
+                    blocks[block_name][field_name] = components
+
+    data = xattree.asdict(value)
+
     for block_name, block in blockspec.items():
-        blocks[block_name] = {}
+        if block_name not in blocks:
+            blocks[block_name] = {}
         period_data = {}
         period_blocks = {}  # type: ignore
+
         for field_name in block.keys():
+            # Skip child components that have been processed as bindings
+            if isinstance(value, Context) and field_name in xatspec.children:
+                child_spec = xatspec.children[field_name]
+                if hasattr(child_spec, "metadata") and "block" in child_spec.metadata:  # type: ignore
+                    if child_spec.metadata["block"] == block_name:  # type: ignore
+                        continue
+
             field_value = data[field_name]
             # convert:
             #   - paths to records

flopy4/mf6/gwf/__init__.py

@@ -51,13 +51,13 @@ def budget(self):
                 self.parent.parent.workspace / f"{self.parent.name}.dis.grb",
             )
 
-    dis: Dis = field(converter=convert_grid)
-    ic: Ic = field()
-    oc: Oc = field()
-    npf: Npf = field()
-    chd: list[Chd] = field()
-    wel: list[Wel] = field()
-    drn: list[Drn] = field()
+    dis: Dis = field(converter=convert_grid, block="packages")
+    ic: Ic = field(block="packages")
+    oc: Oc = field(block="packages")
+    npf: Npf = field(block="packages")
+    chd: list[Chd] = field(block="packages")
+    wel: list[Wel] = field(block="packages")
+    drn: list[Drn] = field(block="packages")
     output: Output = attrs.field(
         default=attrs.Factory(lambda self: Gwf.Output(self), takes_self=True)
     )

flopy4/mf6/model.py

@@ -2,10 +2,10 @@
 
 from xattree import xattree
 
-from flopy4.mf6.component import Component
+from flopy4.mf6.context import Context
 
 
 @xattree
-class Model(Component, ABC):
+class Model(Context, ABC):
     def default_filename(self) -> str:
         return f"{self.name}.nam"  # type: ignore

flopy4/mf6/simulation.py

@@ -23,10 +23,10 @@ def convert_time(value):
 
 @xattree
 class Simulation(Context):
-    models: dict[str, Model] = field()
-    exchanges: dict[str, Exchange] = field()
-    solutions: dict[str, Solution] = field()
-    tdis: Tdis = field(converter=convert_time)
+    models: dict[str, Model] = field(block="models")
+    exchanges: dict[str, Exchange] = field(block="exchanges")
+    solutions: dict[str, Solution] = field(block="solutiongroup")
+    tdis: Tdis = field(converter=convert_time, block="timing")
     filename: str = field(default="mfsim.nam", init=False)
 
     def __attrs_post_init__(self):
@@ -52,13 +52,3 @@ def run(self, exe: str | PathLike = "mf6", verbose: bool = False) -> None:
                     f"Simulation {self.name}: {exe} failed with "  # type: ignore
                     f"return code {ret}, output:\n\n{out + err} "
                 )
-
-    def load(self, format="ascii"):
-        """Load the simulation."""
-        with cd(self.workspace):
-            super().load(format=format)
-
-    def write(self, format="ascii"):
-        """Write the simulation."""
-        with cd(self.workspace):
-            super().write(format=format)

flopy4/mf6/spec.py

@@ -233,7 +233,7 @@ def get_dfn_field_type(attribute: Attribute) -> FieldType:
     kind = xatmeta["kind"]
     match kind:
         case "child":
-            raise ValueError(f"Top-level field should not be a child: {attribute.name}")
+            return "recarray"  # Child components become tabular bindings
         case "array":
             return "recarray"
         case "coord":

test/test_codec.py

@@ -264,3 +264,83 @@ def test_dumps_wel_with_auxiliary():
     assert "45 -25.0 2.0" in result  # (1,3,4) -> node 45, q=-25.0, aux=2.0
     assert "1e+30" not in result
     assert "1.0e+30" not in result
+
+
+def test_dumps_gwf():
+    from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc
+
+    dis = Dis(nlay=1, nrow=10, ncol=10, delr=100.0, delc=100.0)
+    gwf = Gwf(name="test_model", dis=dis)
+    ic = Ic(parent=gwf, strt=1.0)
+    npf = Npf(parent=gwf, k=1.0)
+    oc = Oc(parent=gwf, head_file="test.hds", budget_file="test.bud", dims={"nper": 1})
+    chd = Chd(parent=gwf, head={0: {(0, 0, 0): 10.0}}, dims={"nper": 1})
+
+    gwf = Gwf(
+        name="test_model",
+        dis=dis,
+        ic=ic,
+        npf=npf,
+        oc=oc,
+        chd=[chd],
+    )
+
+    result = dumps(COMPONENT_CONVERTER.unstructure(gwf))
+    print("GWF model result:")
+    print(result)
+
+    # Check that child component bindings are included
+    assert "DIS6" in result
+    assert "IC6" in result
+    assert "NPF6" in result
+    assert "OC6" in result
+    assert "test_model.dis" in result
+    assert "test_model.ic" in result
+    assert "test_model.npf" in result
+    assert "test_model.oc" in result
+
+
+def test_dumps_simulation():
+    from flopy.discretization.modeltime import ModelTime
+
+    from flopy4.mf6.gwf import Dis, Gwf, Ic, Npf, Oc
+    from flopy4.mf6.simulation import Simulation
+    from flopy4.mf6.tdis import Tdis
+
+    # Create model components
+    dis = Dis(nlay=1, nrow=5, ncol=5, delr=100.0, delc=100.0)
+    gwf = Gwf(name="model1", dis=dis)
+    ic = Ic(parent=gwf, strt=1.0)
+    npf = Npf(parent=gwf, k=1.0)
+    oc = Oc(parent=gwf, head_file="model1.hds", budget_file="model1.bud", dims={"nper": 1})
+
+    # Create model
+    gwf = Gwf(
+        name="model1",
+        dis=dis,
+        ic=ic,
+        npf=npf,
+        oc=oc,
+    )
+
+    # Create time discretization
+    time = ModelTime(perlen=[1.0], nstp=[1])
+    tdis = Tdis.from_time(time)
+
+    # Create simulation
+    sim = Simulation(
+        name="test_sim",
+        models={"model1": gwf},
+        exchanges={},
+        solutions={},
+        tdis=tdis,
+    )
+
+    result = dumps(COMPONENT_CONVERTER.unstructure(sim))
+    print("Simulation result:")
+    print(result)
+
+    # Check that model bindings are included
+    assert "GWF6" in result
+    assert "model1" in result
+    assert "TDIS6" in result