try object arrays for perioddata

docs/examples/quickstart.py

@@ -26,24 +26,22 @@
     parent=gwf,
     budget_file=f"{gwf.name}.bud",
     head_file=f"{gwf.name}.hds",
-    save_head={"*": "all"},
-    save_budget={"*": "all"},
+    saverecord={"*": {"head": "all", "budget": "all"}},
 )
 
 # sim.write()
 sim.run(verbose=True)
 
 # check CHD
-assert chd.data["head"][0, 0] == 1.0
-assert chd.data.head.sel(per=0)[99] == 0.0
-assert np.allclose(chd.data.head[:, 1:99], np.full(98, 1e30))
+assert chd.data.perioddata[0, 0].head == 1.0
+assert chd.data.perioddata.sel(per=0, node=99).head == 0.0
+assert chd.data.perioddata.sel(per=0, node=98).head == 1e30
 
 # check DIS
 assert gwf.dis.data.botm.sel(lay=0, col=0, row=0) == 0.0
 
 # check OC
-assert oc.data["save_head"][0] == "all"
-assert oc.data.save_head.sel(per=0) == "all"
+oc.data.perioddata[0].saverecord
 
 # get head and budget results
 budget = gwf.output.budget.squeeze()

flopy4/mf6/context.py

@@ -9,6 +9,8 @@
 
 @xattree
 class Context(Component, ABC):
+    """Base class for MF6 components associated with a workspace directory on disk."""
+
     workspace: Path = field(default=None)
 
     def __attrs_post_init__(self):

flopy4/mf6/exchange.py

@@ -9,7 +9,9 @@
 
 @xattree
 class Exchange(Package, ABC):
-    # mypy doesn't understand that kw_only=True on the
+    """Base class for MF6 exchange packages."""
+
+    # mypy doesn't understand that kw_only=True on base
     # Component means we can have required fields here
     exgtype: type = field()  # type: ignore
     exgfile: Path = field()  # type: ignore

flopy4/mf6/gwf/chd.py

@@ -7,7 +7,6 @@
 from xattree import 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
 
@@ -17,12 +16,10 @@ class Chd(Package):
     multi_package: ClassVar[bool] = True
 
     @define(slots=False)
-    class Steps:
-        all: bool = field()
-        first: bool = field()
-        last: bool = field()
-        steps: list[int] = field()
-        frequency: int = field()
+    class PeriodData:
+        head: float = field()
+        aux: Optional[tuple[np.floating]] = field(default=None)
+        boundname: Optional[str] = field(default=None)
 
     auxiliary: Optional[list[str]] = array(block="options", default=None)
     auxmultname: Optional[str] = field(block="options", default=None)
@@ -34,7 +31,8 @@ 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(
+    perioddata: Optional[NDArray[np.object_]] = array(
+        PeriodData,
         block="period",
         dims=(
             "nper",
@@ -44,43 +42,11 @@ class Steps:
         converter=Converter(structure_array, takes_self=True, takes_field=True),
         reader="urword",
     )
-    aux: Optional[NDArray[np.floating]] = array(
-        block="period",
-        dims=(
-            "nper",
-            "nnodes",
-        ),
-        default=None,
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
-        reader="urword",
-    )
-    boundname: Optional[NDArray[np.str_]] = array(
-        block="period",
-        dims=(
-            "nper",
-            "nnodes",
-        ),
-        default=None,
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
-        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),
-        reader="urword",
-    )
 
     def __attrs_post_init__(self):
         # TODO set up on_setattr hooks for period block
         # arrays to update maxbound? for now do it here
         # 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
-        # maxsteps = len(np.where(self.steps != None)) if self.steps is not None else 0
-        self.maxbound = max(maxhead, maxaux, maxboundname)
+        self.maxbound = len(np.where(self.perioddata != None)) if self.perioddata is not None else 0  # noqa: E711

flopy4/mf6/gwf/oc.py

@@ -3,7 +3,6 @@
 
 import numpy as np
 from attrs import Converter, define
-from modflow_devtools.dfn import Field
 from numpy.typing import NDArray
 from xattree import xattree
 
@@ -12,86 +11,38 @@
 from flopy4.mf6.spec import array, field
 from flopy4.utils import to_path
 
-_OCSETTING = Field(
-    name="ocsetting",
-    type="keystring",
-    reader="urword",
-    children={
-        "all": Field(
-            name="all",
-            type="keyword",
-            reader="urword",
-        ),
-        "first": Field(
-            name="first",
-            type="keyword",
-            reader="urword",
-        ),
-        "last": Field(
-            name="last",
-            type="keyword",
-            reader="urword",
-        ),
-        "steps": Field(
-            name="steps",
-            type="integer",
-            reader="urword",
-        ),
-        "frequency": Field(
-            name="frequency",
-            type="integer",
-            reader="urword",
-        ),
-    },
-)
-
-_RTYPE = Field(
-    name="rtype",
-    type="string",
-    reader="urword",
-)
-
-
-def _oc_action_field(action: str) -> Field:
-    return Field(
-        name=f"{action}record",
-        type="recarray",
-        dims=("nper",),
-        block="perioddata",
-        reader="urword",
-        children={
-            action: Field(
-                name=action,
-                type="keyword",
-                reader="urword",
-            ),
-            "rtype": _RTYPE,
-            "ocsetting": _OCSETTING,
-        },
-    )
-
 
 @xattree
 class Oc(Package):
     @define(slots=False)
-    class Format:
+    class FormatRecord:
         columns: int = field(default=10)
         width: int = field(default=11)
         digits: int = field(default=4)
         format: Literal["exponential", "fixed", "general", "scientific"] = field(default="general")
 
     @define(slots=False)
-    class Steps:
-        all: bool = field()
-        first: bool = field()
-        last: bool = field()
-        steps: list[int] = field()
-        frequency: int = field()
+    class OCSetting:
+        first: bool = field(default=True)
+        last: bool = field(default=False)
+        all: bool = field(default=False)
+        steps: Optional[tuple[int]] = field(default=None)
+        frequency: Optional[int] = field(default=None)
 
     @define(slots=False)
-    class Period:
+    class SaveRecord:
         rtype: str = field()
-        steps: "Oc.Steps" = field()
+        ocsetting: "Oc.OCSetting" = field()
+
+    @define(slots=False)
+    class PrintRecord:
+        rtype: str = field()
+        ocsetting: "Oc.OCSetting" = field()
+
+    @define(slots=False)
+    class PeriodData:
+        saverecord: Optional[tuple["Oc.SaveRecord"]] = field(default=None)
+        printrecord: Optional[tuple["Oc.PrintRecord"]] = field(default=None)
 
     budget_file: Optional[Path] = field(
         block="options",
@@ -108,47 +59,12 @@ class Period:
         converter=to_path,
         default=None,
     )
-    format: Optional[Format] = field(block="options", default=None, init=False)
-    save_head: Optional[NDArray[np.object_]] = array(
-        Steps,
+    headprintrecord: Optional[FormatRecord] = field(block="options", default=None, init=False)
+    perioddata: Optional[NDArray[np.object_]] = array(
+        PeriodData,
         block="period",
-        default="all",
-        dims=("nper",),
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
-        reader="urword",
-    )
-    save_budget: Optional[NDArray[np.object_]] = array(
-        Steps,
-        block="period",
-        default="all",
-        dims=("nper",),
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
-        reader="urword",
-    )
-    print_head: Optional[NDArray[np.object_]] = array(
-        Steps,
-        block="period",
-        default="all",
-        dims=("nper",),
-        converter=Converter(structure_array, takes_self=True, takes_field=True),
-        reader="urword",
-    )
-    print_budget: Optional[NDArray[np.object_]] = array(
-        Steps,
-        block="period",
-        default="all",
+        default=None,
         dims=("nper",),
         converter=Converter(structure_array, takes_self=True, takes_field=True),
         reader="urword",
     )
-
-    # original DFN
-    # @classmethod
-    # def get_dfn(cls) -> Dfn:
-    #     """Generate the component's MODFLOW 6 definition."""
-    #     dfn = super().get_dfn()
-    #     for field_name in list(dfn["perioddata"].keys()):
-    #         dfn["perioddata"].pop(field_name)
-    #     dfn["perioddata"]["saverecord"] = _oc_action_field("save")
-    #     dfn["perioddata"]["printrecord"] = _oc_action_field("print")
-    #     return dfn

flopy4/mf6/model.py

@@ -7,5 +7,7 @@
 
 @xattree
 class Model(Component, ABC):
+    """Base class for MF6 models."""
+
     def default_filename(self) -> str:
         return f"{self.name}.nam"  # type: ignore

flopy4/mf6/package.py

@@ -7,6 +7,8 @@
 
 @xattree
 class Package(Component, ABC):
+    """Base class for MF6 packages."""
+
     def default_filename(self) -> str:
         name = self.parent.name if self.parent else self.name  # type: ignore
         cls_name = self.__class__.__name__.lower()

flopy4/mf6/simulation.py

@@ -23,6 +23,8 @@ def convert_time(value):
 
 @xattree
 class Simulation(Context):
+    """MF6 simulation."""
+
     models: dict[str, Model] = field()
     exchanges: dict[str, Exchange] = field()
     solutions: dict[str, Solution] = field()

flopy4/mf6/solution.py

@@ -7,4 +7,6 @@
 
 @xattree
 class Solution(Package, ABC):
+    """Base class for MF6 solution packages."""
+
     pass