add to_dict method to component (#212)

supports either fields or blocks (nested). useful for the writer internally and also presumably for users.

includes a hack to omit "derived" dimensions from the dict like nodes, which is referenced in DFNs (therefore by components in their array definitions) but isn't actually a variable in any DFN. a better way would be an option on the dim field decorator, maybe. come back to it.

Author: wpbonelli

flopy4/mf6/component.py

@@ -1,12 +1,13 @@
 from abc import ABC
 from collections.abc import MutableMapping
 from pathlib import Path
-from typing import ClassVar
+from typing import Any, ClassVar
 
 import numpy as np
 from attrs import fields
 from modflow_devtools.dfn import Dfn, Field
 from packaging.version import Version
+from xattree import asdict as xattree_asdict
 from xattree import xattree
 
 from flopy4.mf6.constants import FILL_DNODATA, MF6
@@ -179,21 +180,40 @@ def get_dfn(cls) -> Dfn:
             blocks=blocks,
         )
 
-    def _preio(self, format: str = MF6) -> None:
-        # prep for io operations
-        if not self.filename:
-            self.filename = self.default_filename()
-
     def load(self, format: str = MF6) -> None:
         """Load the component and any children."""
-        self._preio(format=format)
+        # TODO: setting filename is a temp hack to get the parent's
+        # name as this component's filename stem, if it has one. an
+        # actual solution is to auto-set the filename when children
+        # are attached to parents.
+        self.filename = self.filename or self.default_filename()
         self._load(format=format)
         for child in self.children.values():  # type: ignore
             child.load(format=format)
 
     def write(self, format: str = MF6) -> None:
         """Write the component and any children."""
-        self._preio(format=format)
+        # TODO: setting filename is a temp hack to get the parent's
+        # name as this component's filename stem, if it has one. an
+        # actual solution is to auto-set the filename when children
+        # are attached to parents.
+        self.filename = self.filename or self.default_filename()
         self._write(format=format)
         for child in self.children.values():  # type: ignore
             child.write(format=format)
+
+    def to_dict(self, blocks: bool = False) -> dict[str, Any]:
+        """Convert the component to a dictionary representation."""
+        data = xattree_asdict(self)
+        data.pop("filename")
+        data.pop("workspace", None)  # might be a Context
+        data.pop("nodes", None)  # TODO: find a better way to omit
+        if blocks:
+            blocks_ = {}  # type: ignore
+            for field_name, field_value in data.items():
+                block_name = self.dfn.fields[field_name].block
+                if block_name not in blocks_:
+                    blocks_[block_name] = {}
+                blocks_[block_name][field_name] = field_value
+            return blocks_
+        return data

flopy4/spec.py

@@ -1,6 +1,8 @@
 """
 Wrap `xattree` and `attrs` specification utilities.
 These include field decorators and introspection functions.
+TODO: add `derived` option to dims? or more generic option
+to any field indicating it is not part of the formal spec?
 """
 
 from attrs import NOTHING, Attribute

test/test_component.py

@@ -416,3 +416,114 @@ def test_write_ascii(function_tmpdir):
     assert f"{gwf_name}.oc" in file_names
     assert f"{gwf_name}.npf" in file_names
     assert f"{gwf_name}.chd" in file_names
+
+
+def test_to_dict_fields():
+    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,
+        "nper": time.nper,
+        "nodes": grid.nnodes,
+    }
+
+    chd = Chd(dims=dims, head={0: {(0, 0, 0): 1.0, (0, 9, 9): 0.0}})
+    result = chd.to_dict()
+
+    assert "head" in result
+    assert result["head"][0, 0] == 1.0
+    assert result["head"][0, 99] == 0.0
+
+    npf = Npf(dims=dims, k=5.0)
+    result = npf.to_dict()
+
+    assert "filename" not in result
+    assert "k" in result
+    assert "icelltype" in result
+    assert "k33" in result
+    assert np.array_equal(result["k"], np.full(100, 5.0))
+
+
+def test_to_dict_blocks():
+    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,
+        "nper": time.nper,
+        "nodes": grid.nnodes,
+    }
+
+    chd = Chd(
+        dims=dims,
+        print_flows=True,
+        head={0: {(0, 0, 0): 1.0, (0, 9, 9): 0.0}},
+    )
+    result = chd.to_dict(blocks=True)
+
+    assert "options" in result
+    assert "period" in result
+    assert "print_flows" in result["options"]
+    assert result["options"]["print_flows"] is True
+    assert "head" in result["period"]
+    assert result["period"]["head"][0, 0] == 1.0
+    assert result["period"]["head"][0, 99] == 0.0
+
+    npf = Npf(dims=dims, save_flows=True, k=2.0)
+    result = npf.to_dict(blocks=True)
+
+    assert "options" in result
+    assert "griddata" in result
+    assert "save_flows" in result["options"]
+    assert result["options"]["save_flows"] is True
+    assert "k" in result["griddata"]
+    assert np.array_equal(result["griddata"]["k"], np.full(100, 2.0))
+
+
+def test_to_dict_on_component():
+    dims = {
+        "nper": 1,
+        "nlay": 1,
+        "nrow": 2,
+        "ncol": 2,
+        "nodes": 4,
+    }
+    dis = Dis(dims=dims)
+    result = dis.to_dict()
+
+    assert "filename" not in result
+    assert "nlay" in result
+
+
+def test_to_dict_on_context():
+    time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0])
+    ims = Ims(models=["gwf"])
+    sim = Simulation(tdis=time, solutions={"ims": ims})
+
+    result = sim.to_dict()
+
+    assert "filename" not in result
+    assert "workspace" not in result
+    assert "tdis" in result
+
+
+def test_to_dict_excludes_derived_dims():
+    # TODO eventually revise to test exclusion of all derived dimensions,
+    # once we have a mechanism to mark them as such
+    dims = {
+        "nper": 1,
+        "nlay": 1,
+        "nrow": 2,
+        "ncol": 2,
+        "nodes": 4,
+    }
+    dis = Dis(dims=dims)
+    result = dis.to_dict()
+
+    assert "nlay" in result
+    assert "nrow" in result
+    assert "ncol" in result
+    assert "nodes" not in result