bottom up working

Author: wpbonelli

flopy4/__init__.py

@@ -1,4 +1,5 @@
-from collections.abc import Iterable, Mapping
+from collections.abc import Mapping
+from itertools import chain
 from pathlib import Path
 from typing import Annotated, Any, Optional, get_origin
 
@@ -215,6 +216,7 @@ def init_tree(
     name: Optional[str] = None,
     parent: Optional[_HasTree] = None,
     children: Optional[Mapping[str, _HasTree]] = None,
+    **kwargs,
 ):
     """
     Initialize a data tree for a component class instance.
@@ -271,7 +273,7 @@ def _yield_arrays(spec, vals):
                 var,
                 value=vals.pop(var.name, var.default),
                 tree=parent.data.root if parent else None,
-                **scalar_vals,
+                **{**scalar_vals, **kwargs},
             )
             if val is not None:
                 yield (var.name, (dims, val))
@@ -347,11 +349,7 @@ def setattribute(self: _Component, attr: Attribute, value: Any):
     self.data.update({attr.name: value})
 
 
-def component(
-    maybe_cls: Optional[type[_IsAttrs]] = None,
-    *,
-    align: Optional[Iterable[str]] = None,
-) -> type[_Component]:
+def component(maybe_cls: Optional[type[_IsAttrs]] = None) -> type[_Component]:
     """
     Attach a data tree to an `attrs` class instance, and use
     the data tree for attribute storage: intercept gets/sets
@@ -365,26 +363,37 @@ def component(
 
     def wrap(cls):
         init_self = cls.__init__
+        spec = fields_dict(cls)
 
         def init(self, *args, **kwargs):
             name = kwargs.pop("name", None)
             children = kwargs.pop("children", None)
             parent = args[0] if args and any(args) else None
 
             # resolve dims from grid and time discretizations
+            # get dims from spec
+            dim_kwargs = {}
+            dims_used = set(
+                chain(*[var.metadata.get("dims", []) for var in spec.values()])
+            )
             grid: Grid = kwargs.pop("grid", None)
             time: ModelTime = kwargs.pop("time", None)
-            diss = [dis for dis in [grid, time] if dis]
-            if align:
-                for dim in align:
-                    for dis in diss:
-                        attr = getattr(dis, dim, None)
-                        if attr is not None:
-                            kwargs[dim] = attr
+            if grid:
+                grid_dims = ["nlay", "nrow", "ncol", "nnodes"]
+                for dim in grid_dims:
+                    if dim in dims_used:
+                        dim_kwargs[dim] = getattr(grid, dim)
+            if time:
+                time_dims = ["nper", "ntstp"]
+                for dim in time_dims:
+                    if dim in dims_used:
+                        dim_kwargs[dim] = getattr(time, dim)
 
             # run the original __init__, then set up the tree
             init_self(self, **kwargs)
-            init_tree(self, name=name, parent=parent, children=children)
+            init_tree(
+                self, name=name, parent=parent, children=children, **dim_kwargs
+            )
 
             # override attribute access
             cls.__getattr__ = getattribute

flopy4/mf6/__init__.py

@@ -34,7 +34,7 @@ class Model(Component):
     pass
 
 
-@component(align=["nper"])
+@component
 @define(slots=False, on_setattr=setattribute)
 class Tdis(Package):
     @define(slots=False)

flopy4/mf6/gwf/chd.py

@@ -7,7 +7,7 @@
 from flopy4.mf6 import Package
 
 
-@component(align="nper")
+@component
 @define(slots=False, on_setattr=setattribute)
 class Chd(Package):
     multi = True

flopy4/mf6/gwf/dis.py

@@ -8,7 +8,7 @@
 from flopy4.mf6 import Package
 
 
-@component(align=["nlay", "ncol", "nrow"])
+@component
 @define(slots=False, on_setattr=setattribute)
 class Dis(Package):
     length_units: str = field(
@@ -45,7 +45,10 @@ class Dis(Package):
         default=1,
         metadata={"block": "griddata", "dims": ("ncol", "nrow", "nlay")},
     )
-    nodes: Optional[int] = field(default=None)
+    nnodes: Optional[int] = field(default=None)
 
     def __attrs_post_init__(self):
-        self.nodes = self.ncol * self.nrow * self.nlay
+        try:
+            self.nnodes = self.ncol * self.nrow * self.nlay
+        except:
+            pass

flopy4/mf6/gwf/ic.py

@@ -6,12 +6,12 @@
 from flopy4.mf6 import Package
 
 
-@component(align=["nodes"])
+@component
 @define(slots=False, on_setattr=setattribute)
 class Ic(Package):
     strt: NDArray[np.floating] = field(
         default=1.0,
-        metadata={"block": "packagedata", "dims": ("nodes",)},
+        metadata={"block": "packagedata", "dims": ("nnodes",)},
     )
     export_array_ascii: bool = field(
         default=False, metadata={"block": "options"}

flopy4/mf6/gwf/npf.py

@@ -9,7 +9,7 @@
 from flopy4.mf6 import Package
 
 
-@component(align=["nodes"])
+@component
 @define(slots=False, on_setattr=setattribute)
 class Npf(Package):
     @define(slots=False)
@@ -66,33 +66,33 @@ class Xt3dOptions:
     )
     icelltype: NDArray[np.integer] = field(
         default=0,
-        metadata={"block": "griddata", "dims": ("nodes",)},
+        metadata={"block": "griddata", "dims": ("nnodes",)},
     )
     k: NDArray[np.floating] = field(
         default=1.0,
-        metadata={"block": "griddata", "dims": ("nodes",)},
+        metadata={"block": "griddata", "dims": ("nnodes",)},
     )
     k22: Optional[NDArray[np.floating]] = field(
         default=None,
-        metadata={"block": "griddata", "dims": ("nodes",)},
+        metadata={"block": "griddata", "dims": ("nnodes",)},
     )
     k33: Optional[NDArray[np.floating]] = field(
         default=None,
-        metadata={"block": "griddata", "dims": ("nodes",)},
+        metadata={"block": "griddata", "dims": ("nnodes",)},
     )
     angle1: Optional[NDArray[np.floating]] = field(
         default=None,
-        metadata={"block": "griddata", "dims": ("nodes",)},
+        metadata={"block": "griddata", "dims": ("nnodes",)},
     )
     angle2: Optional[NDArray[np.floating]] = field(
         default=None,
-        metadata={"block": "griddata", "dims": ("nodes",)},
+        metadata={"block": "griddata", "dims": ("nnodes",)},
     )
     angle3: Optional[NDArray[np.floating]] = field(
         default=None,
-        metadata={"block": "griddata", "dims": ("nodes",)},
+        metadata={"block": "griddata", "dims": ("nnodes",)},
     )
     wetdry: Optional[NDArray[np.floating]] = field(
         default=None,
-        metadata={"block": "griddata", "dims": ("nodes",)},
+        metadata={"block": "griddata", "dims": ("nnodes",)},
     )

flopy4/mf6/gwf/oc.py

@@ -12,7 +12,7 @@
 )
 
 
-@component(align="nper")
+@component
 @define(slots=False, on_setattr=setattribute)
 class Oc(Package):
     @define(slots=False)

test/test_component.py

@@ -55,18 +55,19 @@ def test_init_top_down():
 
 def test_init_bottom_up():
     time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0])
-    grid = StructuredGrid()
+    grid = StructuredGrid(nlay=1, nrow=2, ncol=2)
     dis = Dis(grid=grid)
     ic = Ic(grid=grid)
     oc = Oc(grid=grid)
     npf = Npf(grid=grid)
     gwf = Gwf(
+        grid=grid,
         children={
             "dis": dis,
             "ic": ic,
             "oc": oc,
             "npf": npf,
-        }
+        },
     )
     tdis = Tdis(time=time)
     sim = Sim(children={"tdis": tdis, "gwf": gwf})