Skip to content

Commit

Permalink
simplified, working?
Browse files Browse the repository at this point in the history
  • Loading branch information
@wpbonelli
wpbonelli committed Dec 11, 2024
1 parent d25f351 commit c27a5bb
Showing 1 changed file with 110 additions and 185 deletions.
295 changes: 110 additions & 185 deletions docs/examples/attrs_demo.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,164 +3,135 @@
# This example demonstrates a tentative `attrs`-based object model.

from pathlib import Path
from typing import List, Literal, Optional, Union
from typing import Any, List, Literal, NamedTuple, Optional

# ## GWF IC
import numpy as np
from attr import asdict, define, field, fields_dict
from cattr import Converter
from flopy.discretization import StructuredGrid
from numpy.typing import NDArray

# We can define block classes where variable descriptions become
# the variable's docstring. Ideally we can come up with a Python
# input specification that is equivalent to (and convertible to)
# the original MF6 input specification, while knowing as little
# as possible about the MF6 input format; but anything we can't
# get rid of can go in field `metadata`.
from xarray import Dataset, DataTree


@define
class Options:
class GwfIc:
strt: NDArray[np.float64] = field(
metadata={"block": "packagedata", "shape": "(nodes)"}
)
export_array_ascii: bool = field(
default=False,
metadata={"longname": "export array variables to netcdf output files"},
default=False, metadata={"block": "options"}
)
"""
keyword that specifies input griddata arrays should be
written to layered ascii output files.
"""

export_array_netcdf: bool = field(
default=False, metadata={"longname": "starting head"}
default=False,
metadata={"block": "options"},
)
"""
keyword that specifies input griddata arrays should be
written to the model output netcdf file.
"""


# Eventually we may be able to take advantage of NumPy
# support for shape parameters:
# https://github.com/numpy/numpy/issues/16544
#
# We can still take advantage of type parameters.
def __attrs_post_init__(self):
# TODO: setup attributes for blocks?
self.data = DataTree(Dataset({"strt": self.strt}), name="ic")


@define
class PackageData:
strt: NDArray[np.float64] = field(
metadata={"longname": "starting head", "shape": ("nodes")}
class GwfOc:
@define
class Format:
columns: int
width: int
digits: int
format: Literal["exponential", "fixed", "general", "scientific"]

periods: List[List[tuple]] = field(
metadata={"block": "perioddata"}
)
budget_file: Optional[Path] = field(
default=None, metadata={"block": "options"}
)
budget_csv_file: Optional[Path] = field(
default=None, metadata={"block": "options"}
)
head_file: Optional[Path] = field(
default=None, metadata={"block": "options"}
)
printhead: Optional[Format] = field(
default=None, metadata={"block": "options"}
)
"""
is the initial (starting) head---that is, head at the
beginning of the GWF Model simulation. STRT must be specified for
all simulations, including steady-state simulations. One value is
read for every model cell. For simulations in which the first stress
period is steady state, the values used for STRT generally do not
affect the simulation (exceptions may occur if cells go dry and (or)
rewet). The execution time, however, will be less if STRT includes
hydraulic heads that are close to the steady-state solution. A head
value lower than the cell bottom can be provided if a cell should
start as dry.
"""


# Putting it all together:


@define
class GwfIc:
options: Options = field()
packagedata: PackageData = field()


# ## GWF OC
#
# The output control package has a more complicated variable structure.
# Below docstrings/descriptions are omitted for space-saving.


@define
class Format:
columns: int = field()
width: int = field()
digits: int = field()
format: Literal["exponential", "fixed", "general", "scientific"] = field()


@define
class Options:
budget_file: Optional[Path] = field(default=None)
budget_csv_file: Optional[Path] = field(default=None)
head_file: Optional[Path] = field(default=None)
printhead: Optional[Format] = field(default=None)


# It's awkward to have single-parameter classes, but
# it's the only way I got `cattrs` to distinguish a
# number of choices with the same shape in a union
# like `OCSetting`. There may be a better way.


@define
class All:
all: bool = field()


@define
class First:
first: bool = field()


@define
class Last:
last: bool = field()

class GwfDis:
nlay: int = field(metadata={"block": "dimensions"})
ncol: int = field(metadata={"block": "dimensions"})
nrow: int = field(metadata={"block": "dimensions"})
delr: NDArray[np.float64] = field(
metadata={"block": "griddata", "shape": "(ncol,)"}
)
delc: NDArray[np.float64] = field(
metadata={"block": "griddata", "shape": "(nrow,)"}
)
top: NDArray[np.float64] = field(
metadata={"block": "griddata", "shape": "(ncol, nrow)"}
)
botm: NDArray[np.float64] = field(
metadata={"block": "griddata", "shape": "(ncol, nrow, nlay)"}
)
idomain: NDArray[np.float64] = field(
metadata={"block": "griddata", "shape": "(ncol, nrow, nlay)"}
)
length_units: str = field(default=None, metadata={"block": "options"})
nogrb: bool = field(default=False, metadata={"block": "options"})
xorigin: float = field(default=None, metadata={"block": "options"})
yorigin: float = field(default=None, metadata={"block": "options"})
angrot: float = field(default=None, metadata={"block": "options"})
export_array_netcdf: bool = field(
default=False, metadata={"block": "options"}
)

@define
class Steps:
steps: List[int] = field()
def __attrs_post_init__(self):
self.data = DataTree(
Dataset(
{
"nlay": self.nlay,
"ncol": self.ncol,
"nrow": self.nrow,
"delr": self.delr,
"delc": self.delc,
"top": self.top,
"botm": self.botm,
"idomain": self.idomain,
}
),
name="dis",
)
# TODO: check for parent and update dimensions
# then try to realign any existing packages?


@define
class Frequency:
frequency: int = field()
class Gwf:
dis: GwfDis = field()
ic: GwfIc = field()

def __attrs_post_init__(self):
self.data = DataTree.from_dict(
{"/dis": self.dis, "/ic": self.ic}, name="gwf"
)
self.grid = StructuredGrid(**asdict(self.dis))

PrintSave = Literal["print", "save"]
RType = Literal["budget", "head"]
OCSetting = Union[All, First, Last, Steps, Frequency]
@ic.validator
def _check_dims(self, attribute, value):
assert value.strt.shape == (
self.dis.nlay * self.dis.nrow * self.dis.ncol
)


@define
class OutputControlData:
printsave: PrintSave = field()
rtype: RType = field()
ocsetting: OCSetting = field()

@classmethod
def from_tuple(cls, t):
t = list(t)
printsave = t.pop(0)
rtype = t.pop(0)
ocsetting = {
"all": All,
"first": First,
"last": Last,
"steps": Steps,
"frequency": Frequency,
}[t.pop(0).lower()](t)
return cls(printsave, rtype, ocsetting)


Period = List[OutputControlData]
Periods = List[Period]
# We can define a package with some data.


@define
class GwfOc:
options: Options = field()
periods: Periods = field()
oc = GwfOc(
budget_file="some/file/path.cbc",
periods=[[("print", "budget", "steps", 1, 3, 5)]]
)
assert isinstance(oc.budget_file, str) # TODO path


# We now set up a `cattrs` converter to convert an unstructured
Expand All @@ -169,63 +140,19 @@ class GwfOc:
converter = Converter()


# Register a hook for the `OutputControlData.from_tuple` method.
# MODFLOW 6 defines records as tuples, from which we'll need to
# instantiate objects.


def output_control_data_hook(value, _) -> OutputControlData:
return OutputControlData.from_tuple(value)


converter.register_structure_hook(OutputControlData, output_control_data_hook)


# We can inspect the input specification with `attrs` machinery.


spec = fields_dict(OutputControlData)
assert len(spec) == 3

ocsetting = spec["ocsetting"]
assert ocsetting.type is OCSetting


# We can define a block with some data.


options = Options(
budget_file="some/file/path.cbc",
)
assert isinstance(options.budget_file, str) # TODO path
assert len(asdict(options)) == 4


# We can load a record from a tuple.


ocdata = OutputControlData.from_tuple(("print", "budget", "steps", 1, 3, 5))
assert ocdata.printsave == "print"
assert ocdata.rtype == "budget"
assert ocdata.ocsetting == Steps([1, 3, 5])


# We can load the full package from an unstructured dictionary,
# as would be returned by a separate IO layer in the future.
# (Either hand-written or using e.g. lark.)


gwfoc = converter.structure(
{
"options": {
"budget_file": "some/file/path.cbc",
"head_file": "some/file/path.hds",
"printhead": {
"columns": 1,
"width": 10,
"digits": 8,
"format": "scientific",
},
"budget_file": "some/file/path.cbc",
"head_file": "some/file/path.hds",
"printhead": {
"columns": 1,
"width": 10,
"digits": 8,
"format": "scientific",
},
"periods": [
[
Expand All @@ -236,11 +163,9 @@ def output_control_data_hook(value, _) -> OutputControlData:
},
GwfOc,
)
assert gwfoc.options.budget_file == Path("some/file/path.cbc")
assert gwfoc.options.printhead.width == 10
assert gwfoc.options.printhead.format == "scientific"
assert gwfoc.budget_file == Path("some/file/path.cbc")
assert gwfoc.printhead.width == 10
assert gwfoc.printhead.format == "scientific"
period = gwfoc.periods[0]
assert len(period) == 2
assert period[0] == OutputControlData.from_tuple(
("print", "budget", "steps", 1, 3, 5)
)
assert period[0] == ("print", "budget", "steps", 1, 3, 5)

0 comments on commit c27a5bb

Please sign in to comment.