various fixes

Author: wpbonelli

flopy4/mf6/codec/writer/filters.py

@@ -134,6 +134,7 @@ def array2const(value: xr.DataArray) -> Scalar:
         return value.max().item()
     if np.issubdtype(value.dtype, np.floating):
         return f"{value.max().item():.8f}"
+    return value.ravel()[0]
 
 
 def data2list(value: list | dict | xr.Dataset | xr.DataArray):
@@ -149,7 +150,7 @@ def data2list(value: list | dict | xr.Dataset | xr.DataArray):
         return
 
     if isinstance(value, dict):
-        for name, val in value.values():
+        for name, val in value.items():
             yield (name, val)
         return
 

flopy4/mf6/codec/writer/templates/macros.jinja

@@ -19,7 +19,7 @@
 {% endmacro %}
 
 {% macro record(value) %}
-{{ inset ~ value|join(" ")|upper -}}
+{{ inset ~ value|join(" ") -}}
 {% endmacro %}
 
 {% macro list(name, value) %}

flopy4/mf6/component.py

@@ -202,12 +202,15 @@ def write(self, format: str = MF6) -> None:
     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)
+        data.pop("nodes", None)  # TODO: find a better way to omit
         if blocks:
-            blocks_ = {}
-            for field_name, field_meta in self.data.attrs["metadata"].items():
-                block_name = field_meta["block"]
+            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] = data[field_name]
+                blocks_[block_name][field_name] = field_value
             return blocks_
         return data

flopy4/mf6/converter.py

@@ -32,13 +32,13 @@ def get_binding_blocks(value: Component) -> dict[str, dict[str, list[tuple[str,
     if not isinstance(value, Context):
         return {}
 
-    blocks = {}
+    blocks = {}  # type: ignore
     xatspec = xattree.get_xatspec(type(value))
 
     for child_name, child_spec in xatspec.children.items():
         if (child := getattr(value, child_name, None)) is None:
             continue
-        if (block_name := child_spec.metadata["block"]) not in blocks:
+        if (block_name := child_spec.metadata["block"]) not in blocks:  # type: ignore
             blocks[block_name] = {}
         match child:
             case Component():
@@ -80,55 +80,57 @@ def has_tdis_dims(value: xr.DataArray) -> bool:
     return "nper" in value.dims
 
 
-def _hack_grid_dims(value, field_value):
-    # terrible hack to convert flat nodes dimension to 3d structured dims.
-    # long term solution for this is to use a custom xarray index. filters
-    # should then have access to all dimensions needed.
-    dims_ = set(field_value.dims).copy()
-    parent = value.parent  # type: ignore
-    if parent is None:
-        # TODO for standalone packages
-        return field_value
+def _hack_structured_grid_dims(value: xr.DataArray, structured_grid_dims: Mapping):
+    """
+    Temporary hack to convert flat nodes dimension to 3d structured dims.
+    long term solution for this is to use a custom xarray index. filters
+    should then have access to all dimensions needed.
+    """
 
-    if "nper" in dims_:
-        dims_.remove("nper")
-        shape = (
-            field_value.sizes["nper"],
-            parent.dims["nlay"],
-            parent.dims["nrow"],
-            parent.dims["ncol"],
+    if "nper" in (old_dims := set(value.dims).copy()):
+        old_dims.remove("nper")
+        shape: tuple[int, ...] = (
+            value.sizes["nper"],
+            structured_grid_dims["nlay"],
+            structured_grid_dims["nrow"],
+            structured_grid_dims["ncol"],
         )
-        dims = ("nper", "nlay", "nrow", "ncol")
+        dims: tuple[str, ...] = ("nper", "nlay", "nrow", "ncol")
         coords = {
-            "nper": field_value.coords["nper"],
-            "nlay": range(parent.dims["nlay"]),
-            "nrow": range(parent.dims["nrow"]),
-            "ncol": range(parent.dims["ncol"]),
+            "nper": value.coords["nper"],
+            "nlay": range(structured_grid_dims["nlay"]),
+            "nrow": range(structured_grid_dims["nrow"]),
+            "ncol": range(structured_grid_dims["ncol"]),
         }
     else:
         shape = (
-            parent.dims["nlay"],
-            parent.dims["nrow"],
-            parent.dims["ncol"],
+            structured_grid_dims["nlay"],
+            structured_grid_dims["nrow"],
+            structured_grid_dims["ncol"],
         )
         dims = ("nlay", "nrow", "ncol")
         coords = {
-            "nlay": range(parent.dims["nlay"]),
-            "nrow": range(parent.dims["nrow"]),
-            "ncol": range(parent.dims["ncol"]),
+            "nlay": range(structured_grid_dims["nlay"]),
+            "nrow": range(structured_grid_dims["nrow"]),
+            "ncol": range(structured_grid_dims["ncol"]),
         }
 
-    if dims_ == {"nodes"}:
-        field_value = xr.DataArray(
-            field_value.data.reshape(shape),
+    if old_dims == {"nodes"}:
+        value = xr.DataArray(
+            value.data.reshape(shape),
             dims=dims,
             coords=coords,
         )
 
-    return field_value
+    return value
 
 
-def unstructure_field(name: str, value: Any) -> tuple[str, Any]:
+def unstructure_field(
+    name: str,
+    value: Any,
+    # TODO: temporary, remove not needed
+    structured_grid_dims: Mapping | None,
+) -> tuple[str, Any]:
     """
     Convert:
 
@@ -166,19 +168,34 @@ def unstructure_field(name: str, value: Any) -> tuple[str, Any]:
             return name, value.isoformat()
         case xr.DataArray():
             if name == "auxiliary":
-                value = tuple(value.values.tolist())
+                return name, tuple(value.values.tolist())
             if has_grid_dims(value):
-                value = _hack_grid_dims(value, value)
+                if structured_grid_dims is None:
+                    raise ValueError("Need structured grid dimension sizes")
+                value = _hack_structured_grid_dims(value, structured_grid_dims=structured_grid_dims)
             if has_tdis_dims(value):
                 value = {kper: value.isel(nper=kper) for kper in range(value.sizes["nper"])}
             return name, value
         case _:
             return name, value
 
 
-def unstructure_block(block: dict[str, Any]) -> dict[str, Any]:
+def unstructure_block(
+    block: dict[str, Any],
+    # TODO: temporary, remove not needed
+    structured_grid_dims: Mapping | None,
+) -> dict[str, Any]:
     """Unstructure a block of data, converting fields to a suitable format."""
-    return dict([unstructure_field(block.get(field_name, None)) for field_name in block.keys()])
+    return dict(
+        [
+            unstructure_field(
+                name=field_name,
+                value=block.get(field_name, None),
+                structured_grid_dims=structured_grid_dims,
+            )
+            for field_name in block.keys()
+        ]
+    )
 
 
 def _hack_field_metadata(
@@ -195,8 +212,8 @@ def _hack_field_metadata(
 
 def segment_period_data(block: dict[str, Any], cls: type[Component]) -> dict[str, dict[str, Any]]:
     """Partition period data by stress period"""
-    arrays = {}
-    blocks = {}
+    arrays = {}  # type: ignore
+    blocks = {}  # type: ignore
     period = PERIOD.upper()
 
     for arr_name, periods in block.items():
@@ -220,14 +237,26 @@ def unstructure_component(value: Component) -> dict[str, Any]:
     data = value.to_dict(blocks=True)
     blocks: dict[str, dict[str, Any]] = {}
     blocks.update(binding_blocks := get_binding_blocks(value))
+
+    # temporary hack! TODO remove once we have a structured grid index
+    if "nlay" in value.data.dims:  # type: ignore
+        structured_grid_dims = value.data.dims  # type: ignore
+    elif value.data.parent is not None and "nlay" in value.data.parent.dims:  # type: ignore
+        structured_grid_dims = value.data.parent.dims  # type: ignore
+    else:
+        structured_grid_dims = None
+
     blocks.update(
         {
-            block_name: unstructure_block(data[block_name])
+            block_name: unstructure_block(
+                data[block_name], structured_grid_dims=structured_grid_dims
+            )
             for block_name in dfn.blocks.keys()
             if block_name not in binding_blocks
         }
     )
     if period_block := blocks.pop(PERIOD, None):
+        period_block = {k: v for k, v in period_block.items() if v is not None}
         blocks.update(segment_period_data(period_block, cls))
 
     # total temporary hack! manually set solutiongroup 1.

flopy4/mf6/solution.py

@@ -1,9 +1,8 @@
 from abc import ABC
-from pathlib import Path
-from typing import ClassVar, Optional
+from typing import ClassVar
 
 import attrs
-from xattree import field, xattree
+from xattree import xattree
 
 from flopy4.mf6.package import Package
 
@@ -12,8 +11,7 @@
 class Solution(Package, ABC):
     slntype: ClassVar[str] = "sln"
 
-    slnfname: Optional[Path] = field(default=None)  # type: ignore
     models: list[str] = attrs.field(default=attrs.Factory(list))
 
     def default_filename(self) -> str:
-        return str(self.slnfname) if self.slnfname else f"solution.{self.slntype.lower()}"
+        return f"solution.{self.slntype.lower()}"

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_codec.py

@@ -253,11 +253,11 @@ def test_dumps_drn():
     assert len(period2_lines) == 3
 
     # node elev cond
-    assert "1 1 5 10.00000000 1.00000000" in dumped  # Period 1: (0,0,4)
-    assert "2 5 1 8.00000000 2.00000000" in dumped  # Period 1: (1,4,0)
-    assert "1 2 2 12.00000000 1.50000000" in dumped  # Period 2: (0,1,1)
-    assert "1 3 4 9.00000000 0.80000000" in dumped  # Period 2: (0,2,3)
-    assert "2 4 3 7.00000000 2.20000000" in dumped  # Period 2: (1,3,2)
+    assert "1 1 5 10.0 1.0" in dumped  # Period 1: (0,0,4)
+    assert "2 5 1 8.0 2.0" in dumped  # Period 1: (1,4,0)
+    assert "1 2 2 12.0 1.5" in dumped  # Period 2: (0,1,1)
+    assert "1 3 4 9.0 0.8" in dumped  # Period 2: (0,2,3)
+    assert "2 4 3 7.0 2.2" in dumped  # Period 2: (1,3,2)
     assert "1e+30" not in dumped
     assert "1.0e+30" not in dumped
 
@@ -334,8 +334,8 @@ def test_dumps_wel_with_aux():
 
     assert len(lines) == 2
     # node q aux_value
-    assert "1 2 3 -75.00000000 1.00000000" in dumped  # (0,1,2) -> node 8, q=-75.0, aux=1.0
-    assert "2 4 5 -25.00000000 2.00000000" in dumped  # (1,3,4) -> node 45, q=-25.0, aux=2.0
+    assert "1 2 3 -75.0 1.0" in dumped  # (0,1,2) -> node 8, q=-75.0, aux=1.0
+    assert "2 4 5 -25.0 2.0" in dumped  # (1,3,4) -> node 45, q=-25.0, aux=2.0
     assert "1e+30" not in dumped
     assert "1.0e+30" not in dumped
 

test/test_mf6.py

@@ -13,7 +13,7 @@ def test_gwf_chd01(function_tmpdir):
     time = ModelTime(perlen=[5.0], nstp=[1], tsmult=[1.0], time_units="days")
 
     ims = Ims(
-        slnfname="sln1.ims",
+        filename="sln1.ims",
         models=[gwf_name],
         print_option="summary",
         outer_dvclose=1.00000000e-06,
@@ -91,7 +91,7 @@ def test_gwf_npf01(function_tmpdir):
     )
 
     ims = Ims(
-        # slnfname="sln1.ims",
+        filename="sln1.ims",
         models=[gwf_name],
         print_option="summary",
         outer_dvclose=1.00000000e-06,