passing minimal test

Author: wpbonelli

docs/dev/sdd.md

@@ -42,7 +42,21 @@ Component types include:
 - **model**: a simulated hydrological process
 - **package**: a subcomponent of a model or simulation
 
-Certain subsets of packages have distinguishing characteristics. A **stress package** represents a forcing. A **basic package** contains only input variables applying statically to the entire simulation. An **advanced package** contains time-variable (i.e. transient) input data. In some cases, only a single instance of a package is expected — in other cases, arbitrarily many. Packages for which the latter is true are called **multi-packages**.
+Certain subsets of packages have distinguishing characteristics. A **stress package** represents a forcing. A **basic package** contains only input variables applying statically to the entire simulation. An **advanced package** contains time-variable (i.e. transient) input data. Usually only a single instance of a package is expected — when arbitrarily many are permitted, the package is called a **multi-package**. A **subpackage** is a concept only recognized by the product, not by MODFLOW 6 — a package linked to its parent not by a separate input file, but directly (i.e., subpackage data provided to the parent's initializer method). Subpackages may be attached to packages, models, or simulations.
+
+```mermaid
+classDiagram
+    Simulation *-- "1+" Package
+    Simulation *-- "1+" Model
+    Simulation *-- "1+" Variable
+    Simulation *-- "1+" Subpackage
+    Model *-- "1+" Package
+    Model *-- "1+" Subpackage
+    Model *-- "1+" Variable
+    Package *-- "1+" Subpackage
+    Package *-- "1+" Variable
+    Subpackage *-- "1+" Variable
+```
 
 Components are specified by **definition files**. A **definition** specifies input variables for a single MF6 component. A **block** is a named collection of input variables. A definition file specifies exactly one component. A component may contain zero or more blocks. Each block must contain at least one variable.
 

docs/examples/quickstart.py

@@ -1,64 +1,55 @@
-import flopy4
-ws = './mymodel'
-name = 'mymodel'
-sim = flopy4.mf6.MFSimulation(sim_name=name, sim_ws=ws, exe_name='mf6')
-tdis = flopy4.mf6.ModflowTdis(sim)
-ims = flopy4.mf6.ModflowIms(sim)
-gwf = flopy4.mf6.ModflowGwf(sim, modelname=name, save_flows=True)
-dis = flopy4.mf6.ModflowGwfdis(gwf, nrow=10, ncol=10)
-ic = flopy4.mf6.ModflowGwfic(gwf)
-npf = flopy4.mf6.ModflowGwfnpf(gwf, save_specific_discharge=True)
-chd = flopy4.mf6.ModflowGwfchd(gwf) # self.stress_period_data = sparse.DOK((ncol, nrow, nlay))
-chd.stress_period_data[0,0,0] = 1.
-chd.stress_period_data[0,9,9] = 0.
-
-# chd = flopy4.mf6.ModflowGwfchd(gwf, stress_period_data=sparse.COO([[0,0], [0,9], [0,9]], [1., 0.]))
-# chd = flopy4.mf6.ModflowGwfchd(gwf, stress_period_data=[Chd.PeriodData((0, 0, 0), 1.),
-#                                                        Chd.PeriodData[(0, 9, 9), 0.]])
-budget_file = name + '.bud'
-head_file = name + '.hds'
-oc = flopy4.mf6.ModflowGwfoc(gwf,
-                            budget_filerecord=budget_file,
-                            head_filerecord=head_file,
-                            save={tis.get_period("10-20-2020"): {"head": "ALL", "budget": "ALL"}}
-                            print={"budget": np.ones()}
-                            # save_head="ALL",
-                            # save_budget=np.ones(nstp),
-                            # print_head="",
-                            # print_budget="",
-oc.saverecord["HEAD"]="ALL"
-oc.saverecord["BUDGET"]=np.ones((nstp))
-
-# Normally you would write and run the simulation
+from flopy4.mf6 import Simulation, Tdis
+from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc
+
+ws = "./mymodel"
+name = "mymodel"
+sim = Simulation(name=name, path=ws, exe="mf6")
+tdis = Tdis(sim)
+gwf = Gwf(sim, name=name, save_flows=True)
+dis = Dis(gwf, nrow=10, ncol=10)
+ic = Ic(gwf)
+npf = Npf(gwf, save_specific_discharge=True)
+chd = Chd(
+    gwf,
+    stress_period_data=[
+        # TODO? accept raw tuples for
+        # parity with flopy3 example?
+        Chd.StressPeriodData(cellid=(0, 0, 0), head=1.0),
+        Chd.StressPeriodData(cellid=(0, 9, 9), head=1.0),
+    ],
+)
+
+# TODO? xarray alternative
+# chd.data["stress_period_data"].loc(dict(i=0, j=0, k=0)) = 1.
+# chd.data["stress_period_data"].loc(dict(i=0, j=9, k=9)) = 0.
+
+# TODO? sparse array alternative
+# spd = sparse.COO([[0,0], [0,9], [0,9]], [1., 0.])
+# chd = flopy4.mf6.ModflowGwfchd(gwf, stress_period_data=spd)
+
+budget_file = name + ".bud"
+head_file = name + ".hds"
+oc = Oc(
+    gwf,
+    budget_filerecord=budget_file,
+    head_filerecord=head_file,
+    perioddata=[("HEAD", "ALL"), ("BUDGET", "ALL")],
+    # TODO: format we want to support
+    # save={"head": {0: "ALL"}, "budget": {0: "ALL"}},
+    # print={"budget": {0: np.ones((tdis.nstp[0]))}, "budget": {0, "ALL"}},
+)
+
+# TODO? mock some output
 # sim.write_simulation()
 # sim.run_simulation()
 
-
-head = gwf.output.head().get_data()
-bud = gwf.output.budget()
-
-spdis = bud.get_data(text='DATA-SPDIS')[0]
-qx, qy, qz = flopy.utils.postprocessing.get_specific_discharge(spdis, gwf)
-pmv = flopy.plot.PlotMapView(gwf)
-pmv.plot_array(head)
-pmv.plot_grid(colors='white')
-pmv.contour_array(head, levels=[.2, .4, .6, .8], linewidths=3.)
-pmv.plot_vector(qx, qy, normalize=True, color="white")
-
-
-
-
-# from demo:
-
-# # Create a simulation.
-# 
-# sim = Simulation()
-# tdis = Tdis(sim=sim, nper=1, perioddata=[Tdis.PeriodData()])
-# gwf = Gwf(sim=sim)
-# dis = Dis(model=gwf)
-# ic = Ic(model=gwf, strt=1.0)
-# oc = Oc(model=gwf, perioddata=[Oc.Steps()])
-# npf = Npf(model=gwf, icelltype=0, k=1.0)
-# 
-# # View the data tree.
-# sim.data
+# TODO? try to reproduce plots with xarray
+# head = gwf.output.head().get_data()
+# bud = gwf.output.budget()
+# spdis = bud.get_data(text='DATA-SPDIS')[0]
+# qx, qy, qz = flopy.utils.postprocessing.get_specific_discharge(spdis, gwf)
+# pmv = flopy.plot.PlotMapView(gwf)
+# pmv.plot_array(head)
+# pmv.plot_grid(colors='white')
+# pmv.contour_array(head, levels=[.2, .4, .6, .8], linewidths=3.)
+# pmv.plot_vector(qx, qy, normalize=True, color="white")

flopy4/__init__.py

@@ -0,0 +1,195 @@
+from typing import Any, Optional, get_origin
+
+import numpy as np
+from attr import Attribute, fields_dict
+from numpy.typing import ArrayLike, NDArray
+from xarray import Dataset, DataTree
+
+
+def _parse_dim_names(shape: str) -> tuple[str, ...]:
+    return tuple(
+        [
+            dim.strip()
+            for dim in shape.strip()
+            .replace("(", "")
+            .replace(")", "")
+            .split(",")
+            if any(dim)
+        ]
+    )
+
+
+def _try_resolve_dim(data: Optional[DataTree], name: str) -> int | str:
+    name = name.strip()
+    if data is None:
+        return name
+    value = data.get(name, None)
+    if value is not None:
+        return value.item()
+    root = data.root
+    paths = [
+        "tdis",
+        "dis",
+        "gwf/dis",
+    ]
+    for path in paths:
+        try:
+            key = f"{path}/{name}"
+            return root[key].item()
+        except:
+            try:
+                return root[path].dims[name]
+            except:
+                pass
+    return name
+
+
+def _try_resolve_shape(data: DataTree, attr: Attribute) -> tuple[int | str]:
+    shape = attr.metadata.get("shape", None)
+    if shape is None:
+        raise ValueError(f"Array {attr.name} missing shape metadata")
+    shape = [_try_resolve_dim(data, dim) for dim in _parse_dim_names(shape)]
+    return shape
+
+
+def _reshape_array(value: ArrayLike, shape: tuple[int]) -> Optional[NDArray]:
+    value = np.array(value)
+    if value.shape == ():
+        return np.full(shape, value.item())
+    elif value.shape != shape:
+        raise ValueError(
+            f"Shape mismatch, got {value.shape}, expected {shape}"
+        )
+    return value
+
+
+def resolve_array(
+    self, attr: Attribute, value: Optional[ArrayLike]
+) -> Optional[NDArray]:
+    """
+    Resolve an array-like value to a numpy array with the correct shape.
+    The shape is determined by the shape metadata of the attribute, and
+    the dimensions are resolved by looking up the corresponding values
+    in the data tree.
+    """
+    if value is None:
+        return None
+    shape = _try_resolve_shape(self.data, attr)
+    unresolved = [dim for dim in shape if not isinstance(dim, int)]
+    if any(unresolved):
+        raise ValueError(
+            f"Class '{type(self).__name__}' "
+            f"failed to resolve dims: {', '.join(unresolved)}"
+        )
+    return _reshape_array(value, shape)
+
+
+def _bind_tree(self, parent):
+    parent.data = parent.data.assign({self.data.name: self.data})
+    self.data = parent.data[self.data.name]
+    grandparent = getattr(parent, "parent", None)
+    if grandparent is not None:
+        _bind_tree(parent, grandparent)
+
+
+def init_tree(self, parent=None, **kwargs):
+    """
+    Initialize a data tree for a component instance.
+    """
+    cls = type(self)
+    cls_name = cls.__name__.lower()
+    spec = fields_dict(cls)
+    data = Dataset()
+    dims = set()
+
+    # add arrays
+    for name, attr in spec.items():
+        value = kwargs.get(name, attr.default)
+        shape = attr.metadata.get("shape", None)
+        if shape is not None:
+            dim_names = _parse_dim_names(shape)
+            shape = [
+                _try_resolve_dim(parent.data.root if parent else None, dim)
+                for dim in dim_names
+            ]
+            shape = tuple(
+                [
+                    (dim if isinstance(dim, int) else kwargs.get(dim, dim))
+                    for dim in shape
+                ]
+            )
+            unresolved = [dim for dim in shape if not isinstance(dim, int)]
+            if any(unresolved):
+                raise ValueError(
+                    f"Class '{cls_name}' "
+                    f"failed to resolve dims: {', '.join(unresolved)}"
+                )
+            dims.update(dim_names)
+            value = _reshape_array(value, shape)
+            if value.shape == ():
+                raise ValueError(
+                    f"Failed to resolve array '{name}', "
+                    f"make sure these dimensions exist: "
+                    f"{','.join(dims)}"
+                )
+            data[name] = (dim_names, value)
+
+    # add scalars
+    for name, value in spec.items():
+        if name in data or name in dims:
+            continue
+        value = kwargs.get(name, attr.default)
+        data[name] = value
+
+    self.data = DataTree(data, name=cls_name)
+    if parent is not None:
+        self.parent = parent
+        _bind_tree(self, parent)
+
+
+def getattribute(self, name: str) -> Any:
+    """Override `__getattribute__` to proxy attribute access."""
+    cls = type(self)
+    spec = fields_dict(cls)
+    if name in spec:
+        value = self.data.get(name, None)
+        if value is not None:
+            return value
+        value = self.data.dims.get(name, None)
+        if value is not None:
+            return value
+        value = self.data.attrs.get(name, None)
+        if value is not None:
+            return value
+    return super(cls, self).__getattribute__(name)
+
+
+def setattribute(self, attr: Attribute, value: Any):
+    """Hook for setting attribute values."""
+    cls = type(self)
+    spec = fields_dict(cls)
+    if attr.name not in spec:
+        raise AttributeError(f"{cls.__name__} has no attribute {attr.name}")
+    if value is None:
+        return
+    self.data[attr.name] = (
+        (
+            _parse_dim_names(attr.metadata["shape"]),
+            resolve_array(self, attr, value),
+        )
+        if get_origin(attr.type) in [list, np.ndarray]
+        else value
+    )
+    # TODO run validation?
+
+
+def component(cls):
+    """
+    Decorator for component classes.
+
+    This decorator adds a data tree to the class instance, on
+    top of the existing attributes. The data tree is used to
+    store both scalars and arrays. Attributes proxy the tree.
+    """
+    cls.__getattribute__ = getattribute
+    return cls