From 44a5df1fccb736f18c2b429a575aa1259139e78a Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Sat, 7 Jun 2025 19:48:17 -0400 Subject: [PATCH 1/2] wip --- docs/examples/quickstart.py | 11 ++-- flopy4/mf6/context.py | 2 + flopy4/mf6/exchange.py | 4 +- flopy4/mf6/gwf/chd.py | 48 ++------------ flopy4/mf6/gwf/oc.py | 127 +++++++----------------------------- flopy4/mf6/model.py | 2 + flopy4/mf6/package.py | 2 + flopy4/mf6/simulation.py | 2 + flopy4/mf6/solution.py | 2 + flopy4/mf6/tdis.py | 42 +++++------- 10 files changed, 66 insertions(+), 176 deletions(-) diff --git a/docs/examples/quickstart.py b/docs/examples/quickstart.py index cba8044e..47927dce 100644 --- a/docs/examples/quickstart.py +++ b/docs/examples/quickstart.py @@ -26,23 +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["saverecord"][0] == "all" assert oc.data.save_head.sel(per=0) == "all" # get head and budget results diff --git a/flopy4/mf6/context.py b/flopy4/mf6/context.py index 8320b77c..bb5b3316 100644 --- a/flopy4/mf6/context.py +++ b/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): diff --git a/flopy4/mf6/exchange.py b/flopy4/mf6/exchange.py index ca0ccb57..3452cc1e 100644 --- a/flopy4/mf6/exchange.py +++ b/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 diff --git a/flopy4/mf6/gwf/chd.py b/flopy4/mf6/gwf/chd.py index d61a4307..81d42f34 100644 --- a/flopy4/mf6/gwf/chd.py +++ b/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,34 +42,6 @@ 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 @@ -79,8 +49,4 @@ 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 - # 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 diff --git a/flopy4/mf6/gwf/oc.py b/flopy4/mf6/gwf/oc.py index 6bcae76a..0de15641 100644 --- a/flopy4/mf6/gwf/oc.py +++ b/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,37 @@ 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 OCSettingOption: + first: bool = field(default=True) + last: bool = field(default=False) + all: bool = field(default=False) + steps: Optional[list[int]] = field(default=None) + frequency: Optional[int] = field(default=None) @define(slots=False) - class Period: + class OCSetting: + ocsetting: Optional[list["Oc.OCSettingOption"]] = field(default=None) + + @define(slots=False) + class SaveRecord: + rtype: str = field() + ocsetting: "Oc.OCSetting" = field() + + @define(slots=False) + class PrintRecord: rtype: str = field() - steps: "Oc.Steps" = field() + ocsetting: "Oc.OCSetting" = field() budget_file: Optional[Path] = field( block="options", @@ -108,47 +58,20 @@ 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) + saverecord: Optional[NDArray[np.object_]] = array( + SaveRecord, 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", + default=None, dims=("nper",), converter=Converter(structure_array, takes_self=True, takes_field=True), reader="urword", ) - print_budget: Optional[NDArray[np.object_]] = array( - Steps, + printrecord: Optional[NDArray[np.object_]] = array( + PrintRecord, 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 diff --git a/flopy4/mf6/model.py b/flopy4/mf6/model.py index ca0367c7..2124c776 100644 --- a/flopy4/mf6/model.py +++ b/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 diff --git a/flopy4/mf6/package.py b/flopy4/mf6/package.py index 8471c58d..22183563 100644 --- a/flopy4/mf6/package.py +++ b/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() diff --git a/flopy4/mf6/simulation.py b/flopy4/mf6/simulation.py index 46484f01..e072782e 100644 --- a/flopy4/mf6/simulation.py +++ b/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() diff --git a/flopy4/mf6/solution.py b/flopy4/mf6/solution.py index dfcafd73..7d721d40 100644 --- a/flopy4/mf6/solution.py +++ b/flopy4/mf6/solution.py @@ -7,4 +7,6 @@ @xattree class Solution(Package, ABC): + """Base class for MF6 solution packages.""" + pass diff --git a/flopy4/mf6/tdis.py b/flopy4/mf6/tdis.py index 79e8396b..b11631e3 100644 --- a/flopy4/mf6/tdis.py +++ b/flopy4/mf6/tdis.py @@ -1,6 +1,7 @@ from datetime import datetime from typing import Optional +import attrs import numpy as np from attrs import Converter, define from flopy.discretization.modeltime import ModelTime @@ -14,11 +15,11 @@ @xattree class Tdis(Package): - @define + @define(slots=False) class PeriodData: - perlen: float - nstp: int - tsmult: float + perlen: float = 1.0 + nstp: int = 1 + tsmult: float = 1.0 nper: int = dim( block="dimensions", @@ -28,23 +29,9 @@ class PeriodData: ) time_units: Optional[str] = field(block="options", default=None) start_date_time: Optional[datetime] = field(block="options", default=None) - perlen: NDArray[np.floating] = array( + perioddata: NDArray[np.object_] = array( + PeriodData, block="perioddata", - default=1.0, - dims=("nper",), - converter=Converter(structure_array, takes_self=True, takes_field=True), - reader="urword", - ) - nstp: NDArray[np.integer] = array( - block="perioddata", - default=1, - dims=("nper",), - converter=Converter(structure_array, takes_self=True, takes_field=True), - reader="urword", - ) - tsmult: NDArray[np.floating] = array( - block="perioddata", - default=1.0, dims=("nper",), converter=Converter(structure_array, takes_self=True, takes_field=True), reader="urword", @@ -52,23 +39,26 @@ class PeriodData: def to_time(self) -> ModelTime: """Convert to a `ModelTime` object.""" + perioddata = np.rec.fromrecords([attrs.astuple(pd) for pd in self.perioddata.to_numpy()]) return ModelTime( nper=self.nper, time_units=self.time_units, start_date_time=self.start_date_time, - perlen=self.perlen, - nstp=self.nstp, - tsmult=self.tsmult, + perlen=perioddata.perlen, + nstp=perioddata.nstp, + tsmult=perioddata.tsmult, ) @classmethod def from_time(cls, time: ModelTime) -> "Tdis": """Create a time discretization from a `ModelTime`.""" + perioddata = np.rec.fromrecords( + *zip(time.perlen, time.nstp, time.tsmult), + dtype=[("perlen", float), ("nstp", int), ("tsmult", float)], + ) return cls( nper=time.nper, time_units=time.time_units, start_date_time=time.start_datetime, - perlen=time.perlen, - nstp=time.nstp, - tsmult=time.tsmult, + perioddata=perioddata, ) From 325729406b737cb8ec2d9450d1fd2d049b335e09 Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Sun, 8 Jun 2025 22:57:41 -0400 Subject: [PATCH 2/2] rework oc --- docs/examples/quickstart.py | 3 +-- flopy4/mf6/gwf/oc.py | 25 +++++++++---------------- 2 files changed, 10 insertions(+), 18 deletions(-) diff --git a/docs/examples/quickstart.py b/docs/examples/quickstart.py index 47927dce..1d73e81a 100644 --- a/docs/examples/quickstart.py +++ b/docs/examples/quickstart.py @@ -41,8 +41,7 @@ assert gwf.dis.data.botm.sel(lay=0, col=0, row=0) == 0.0 # check OC -assert oc.data["saverecord"][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() diff --git a/flopy4/mf6/gwf/oc.py b/flopy4/mf6/gwf/oc.py index 0de15641..5d6b4e7d 100644 --- a/flopy4/mf6/gwf/oc.py +++ b/flopy4/mf6/gwf/oc.py @@ -22,17 +22,13 @@ class FormatRecord: format: Literal["exponential", "fixed", "general", "scientific"] = field(default="general") @define(slots=False) - class OCSettingOption: + class OCSetting: first: bool = field(default=True) last: bool = field(default=False) all: bool = field(default=False) - steps: Optional[list[int]] = field(default=None) + steps: Optional[tuple[int]] = field(default=None) frequency: Optional[int] = field(default=None) - @define(slots=False) - class OCSetting: - ocsetting: Optional[list["Oc.OCSettingOption"]] = field(default=None) - @define(slots=False) class SaveRecord: rtype: str = field() @@ -43,6 +39,11 @@ 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", converter=to_path, @@ -59,16 +60,8 @@ class PrintRecord: default=None, ) headprintrecord: Optional[FormatRecord] = field(block="options", default=None, init=False) - saverecord: Optional[NDArray[np.object_]] = array( - SaveRecord, - block="period", - default=None, - dims=("nper",), - converter=Converter(structure_array, takes_self=True, takes_field=True), - reader="urword", - ) - printrecord: Optional[NDArray[np.object_]] = array( - PrintRecord, + perioddata: Optional[NDArray[np.object_]] = array( + PeriodData, block="period", default=None, dims=("nper",),