From bb7428c780358a042f26c620d5158131f98eb641 Mon Sep 17 00:00:00 2001 From: w-bonelli Date: Tue, 21 Oct 2025 17:41:46 -0400 Subject: [PATCH] add to_dict method to component --- flopy4/mf6/component.py | 36 ++++++++++--- flopy4/spec.py | 2 + test/test_component.py | 111 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 141 insertions(+), 8 deletions(-) diff --git a/flopy4/mf6/component.py b/flopy4/mf6/component.py index 93951569..8a748965 100644 --- a/flopy4/mf6/component.py +++ b/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 diff --git a/flopy4/spec.py b/flopy4/spec.py index 3d9b5ddb..0721fdca 100644 --- a/flopy4/spec.py +++ b/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 diff --git a/test/test_component.py b/test/test_component.py index a62c94c1..05d43dcd 100644 --- a/test/test_component.py +++ b/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