Xattree update

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)