Xattree update (#164)

- use new xattree builtin dict to array converter, remove custom impl (todo: support tabular data too... xattree should provide a single converter that supports either dict or table input)
- get rid of star syntax for sparse dict representations of arrays. this was more trouble than it was worth to generalize, and the caller can always use a comprehension to easily generate a still-pretty-sparse dict
- accommodate get_xatspec refactor (it returns XatSpec now, not the flattened chainmap)
- miscelleanous other cleanup and notes
- todo: xattree needs to distinguish array fill value from default. can't always interpret default unambiguously for optional arrays. e.g., when an optional array should be None (i.e. not present) by default, but should have a fill value when populated from a sparse dict or table representation

Author: wpbonelli

docs/dev/sdd.md

@@ -85,7 +85,7 @@ import numpy as np
 @define
 class Ic(Package):
     """Initial conditions package"""
-    strt: NDArray[np.floating] = field(...)
+    strt: NDArray[np.float64] = field(...)
     export_array_ascii: bool = field(...)
     export_array_netcdf: bool = field(...)
 ```

docs/examples/quickstart.py

@@ -20,14 +20,14 @@
 npf = Npf(parent=gwf, save_specific_discharge=True)
 chd = Chd(
     parent=gwf,
-    head={"*": {(0, 0, 0): 1.0, (0, 9, 9): 0.0}},
+    head={0: {(0, 0, 0): 1.0, (0, 9, 9): 0.0}},
 )
 oc = Oc(
     parent=gwf,
     budget_file=f"{gwf.name}.bud",
     head_file=f"{gwf.name}.hds",
-    save_head={"*": "all"},
-    save_budget={"*": "all"},
+    save_head={0: "all"},
+    save_budget={0: "all"},
 )
 
 # sim.write()

flopy4/mf6/adapters.py

@@ -185,7 +185,7 @@ def __init__(
             self._data = package.data
         else:
             raise Exception("Input package has no data")
-        self._spec = get_xatspec(type(package))
+        self._spec = get_xatspec(type(package)).flat
         if modelgrid:
             self._grid = modelgrid
         elif model:

flopy4/mf6/codec/__init__.py

@@ -9,7 +9,6 @@
 
 from flopy4.mf6 import filters
 from flopy4.mf6.codec.converter import (
-    structure_array,
     unstructure_array,
     unstructure_chd,
     unstructure_component,
@@ -85,7 +84,6 @@ def dump(data, path: str | PathLike) -> None:
 
 
 __all__ = [
-    "structure_array",
     "unstructure_array",
     "loads",
     "load",

flopy4/mf6/codec/converter.py

@@ -1,93 +1,16 @@
-from typing import Any, Tuple
+from typing import Any
 
 import numpy as np
 import sparse
 import xattree
-from numpy.typing import NDArray
 from xarray import DataArray
-from xattree import get_xatspec
 
-from flopy4.adapters import get_cellid, get_nn
+from flopy4.adapters import get_cellid
 from flopy4.mf6.component import Component
-from flopy4.mf6.config import SPARSE_THRESHOLD
 from flopy4.mf6.constants import FILL_DNODATA
 from flopy4.mf6.spec import get_blocks, is_list_field
 
 
-# TODO: convert to a cattrs structuring hook so we don't have to
-# apply separately to all array fields?
-def structure_array(value, self_, field) -> NDArray:
-    """
-    Convert a sparse dictionary representation of an array to a
-    dense numpy array or a sparse COO array.
-    """
-
-    if not isinstance(value, dict):
-        # if not a dict, assume it's a numpy array
-        # and let xarray deal with it if it isn't
-        return value
-
-    # get spec
-    spec = get_xatspec(type(self_))
-    field = spec[field.name]
-    if not field.dims:
-        raise ValueError(f"Field {field} missing dims")
-
-    # resolve dims
-    explicit_dims = self_.__dict__.get("dims", {})
-    inherited_dims = dict(self_.parent.data.dims) if self_.parent else {}
-    dims = inherited_dims | explicit_dims
-    shape = [dims.get(d, d) for d in field.dims]
-    unresolved = [d for d in shape if isinstance(d, str)]
-    if any(unresolved):
-        raise ValueError(f"Couldn't resolve dims: {unresolved}")
-
-    if np.prod(shape) > SPARSE_THRESHOLD:
-        a: dict[Tuple[Any, ...], Any] = dict()
-
-        def set_(arr, val, *ind):
-            arr[tuple(ind)] = val
-
-        def final(arr):
-            coords = np.array(list(map(list, zip(*arr.keys()))))
-            return sparse.COO(
-                coords,
-                list(arr.values()),
-                shape=shape,
-                fill_value=field.default or FILL_DNODATA,
-            )
-    else:
-        a = np.full(shape, FILL_DNODATA, dtype=field.dtype)  # type: ignore
-
-        def set_(arr, val, *ind):
-            arr[ind] = val
-
-        def final(arr):
-            arr[arr == FILL_DNODATA] = field.default or FILL_DNODATA
-            return arr
-
-    # populate array. TODO: is there a way to do this
-    # without hardcoding awareness of kper and cellid?
-    if "nper" in dims:
-        for kper, period in value.items():
-            if kper == "*":
-                kper = 0
-            match len(shape):
-                case 1:
-                    set_(a, period, kper)
-                case _:
-                    for cellid, v in period.items():
-                        nn = get_nn(cellid, **dims)
-                        set_(a, v, kper, nn)
-            if kper == "*":
-                break
-    else:
-        for cellid, v in value.items():
-            nn = get_nn(cellid, **dims)
-            set_(a, v, nn)
-    return final(a)
-
-
 def unstructure_array(value: DataArray) -> dict:
     """
     Convert a dense numpy array or a sparse COO array to a sparse

flopy4/mf6/gwf/chd.py

@@ -2,11 +2,10 @@
 from typing import ClassVar, Optional
 
 import numpy as np
-from attrs import Converter, define
+from attrs import define
 from numpy.typing import NDArray
-from xattree import xattree
+from xattree import dict_to_array_converter, xattree
 
-from flopy4.mf6.codec import structure_array
 from flopy4.mf6.constants import FILL_DNODATA
 from flopy4.mf6.package import Package
 from flopy4.mf6.spec import array, field
@@ -34,24 +33,24 @@ class Steps:
     obs_filerecord: Optional[Path] = field(block="options", default=None)
     dev_no_newton: bool = field(default=False, metadata={"block": "options"})
     maxbound: Optional[int] = field(block="dimensions", default=None)
-    head: Optional[NDArray[np.floating]] = array(
+    head: Optional[NDArray[np.float64]] = array(
         block="period",
         dims=(
             "nper",
             "nnodes",
         ),
         default=None,
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        converter=dict_to_array_converter,
         reader="urword",
     )
-    aux: Optional[NDArray[np.floating]] = array(
+    aux: Optional[NDArray[np.float64]] = array(
         block="period",
         dims=(
             "nper",
             "nnodes",
         ),
         default=None,
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        converter=dict_to_array_converter,
         reader="urword",
     )
     boundname: Optional[NDArray[np.str_]] = array(
@@ -61,15 +60,15 @@ class Steps:
             "nnodes",
         ),
         default=None,
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        converter=dict_to_array_converter,
         reader="urword",
     )
     steps: Optional[NDArray[np.object_]] = array(
         Steps,
         block="period",
         dims=("nper", "nnodes"),
         default=None,
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        converter=dict_to_array_converter,
         reader="urword",
     )
 
@@ -79,8 +78,25 @@ def __attrs_post_init__(self):
         # in post init. but this only works when values
         # are set in the initializer, not when they are
         # set later.
-        maxhead = len(np.where(self.head != FILL_DNODATA)) if self.head is not None else 0
-        maxaux = len(np.where(self.aux != FILL_DNODATA)) if self.aux is not None else 0
-        maxboundname = len(np.where(self.boundname != "")) if self.boundname is not None else 0
+        if self.head is None:
+            maxhead = 0
+        else:
+            head = self.head if self.head.data.shape == self.head.shape else self.head.todense()
+            maxhead = len(np.where(head != FILL_DNODATA))
+        if self.aux is None:
+            maxaux = 0
+        else:
+            aux = self.aux if self.aux.data.shape == self.aux.shape else self.aux.todense()
+            maxaux = len(np.where(aux != FILL_DNODATA))
+        if self.boundname is None:
+            maxboundname = 0
+        else:
+            boundname = (
+                self.boundname
+                if self.boundname.data.shape == self.boundname.shape
+                else self.boundname.todense()
+            )
+            maxboundname = len(np.where(boundname != ""))
+
         # maxsteps = len(np.where(self.steps != None)) if self.steps is not None else 0
         self.maxbound = max(maxhead, maxaux, maxboundname)

flopy4/mf6/gwf/dis.py

@@ -1,10 +1,8 @@
 import numpy as np
-from attrs import Converter
 from flopy.discretization.structuredgrid import StructuredGrid
 from numpy.typing import NDArray
-from xattree import xattree
+from xattree import dict_to_array_converter, xattree
 
-from flopy4.mf6.codec import structure_array
 from flopy4.mf6.package import Package
 from flopy4.mf6.spec import array, dim, field
 
@@ -25,48 +23,51 @@ class Dis(Package):
         coord="lay",
         scope="gwf",
         default=1,
+        group="grid",
     )
     ncol: int = dim(
         block="dimensions",
         coord="col",
         scope="gwf",
         default=2,
+        group="grid",
     )
     nrow: int = dim(
         block="dimensions",
         coord="row",
         scope="gwf",
         default=2,
+        group="grid",
     )
-    delr: NDArray[np.floating] = array(
+    delr: NDArray[np.float64] = array(
         block="griddata",
         default=1.0,
         dims=("ncol",),
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        converter=dict_to_array_converter,
     )
-    delc: NDArray[np.floating] = array(
+    delc: NDArray[np.float64] = array(
         block="griddata",
         default=1.0,
         dims=("nrow",),
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        converter=dict_to_array_converter,
     )
-    top: NDArray[np.floating] = array(
+    top: NDArray[np.float64] = array(
         block="griddata",
         default=1.0,
         dims=("nrow", "ncol"),
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        converter=dict_to_array_converter,
     )
-    botm: NDArray[np.floating] = array(
+    botm: NDArray[np.float64] = array(
         block="griddata",
         default=0.0,
         dims=("nlay", "nrow", "ncol"),
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        converter=dict_to_array_converter,
     )
-    idomain: NDArray[np.integer] = array(
+    idomain: NDArray[np.int32] = array(
         block="griddata",
         default=1,
         dims=("nlay", "nrow", "ncol"),
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        converter=dict_to_array_converter,
     )
     nnodes: int = dim(
         coord="node",

flopy4/mf6/gwf/ic.py

@@ -1,20 +1,18 @@
 import numpy as np
-from attrs import Converter
 from numpy.typing import NDArray
-from xattree import xattree
+from xattree import dict_to_array_converter, xattree
 
-from flopy4.mf6.codec import structure_array
 from flopy4.mf6.package import Package
 from flopy4.mf6.spec import array, field
 
 
 @xattree
 class Ic(Package):
-    strt: NDArray[np.floating] = array(
+    strt: NDArray[np.float64] = array(
         block="packagedata",
         dims=("nnodes",),
         default=1.0,
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
+        converter=dict_to_array_converter,
     )
     export_array_ascii: bool = field(block="options", default=False)
     export_array_netcdf: bool = field(block="options", default=False)