From bbed8e715aa0f252c87004f235104c2260f2148c Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Wed, 4 Jun 2025 15:08:16 -0400 Subject: [PATCH] write the quickstart simulation! not yet correctly.. --- flopy4/mf6/component.py | 28 +++++++++++++++++++++------- flopy4/mf6/context.py | 1 + flopy4/mf6/filters.py | 3 ++- flopy4/mf6/gwf/dis.py | 1 + flopy4/mf6/model.py | 3 ++- flopy4/mf6/package.py | 5 ++++- flopy4/mf6/simulation.py | 1 + flopy4/mf6/templates/blocks.jinja | 2 +- flopy4/mf6/templates/macros.jinja | 2 +- test/conftest.py | 2 ++ test/test_codec.py | 3 --- test/test_component.py | 22 +++++++++++++++++++--- 12 files changed, 55 insertions(+), 18 deletions(-) diff --git a/flopy4/mf6/component.py b/flopy4/mf6/component.py index cdf8348d..72a4b0de 100644 --- a/flopy4/mf6/component.py +++ b/flopy4/mf6/component.py @@ -38,9 +38,18 @@ class Component(ABC, MutableMapping): @property def path(self) -> Path: + """Get the path to the component's input file.""" return Path.cwd() / self.filename - def _default_filename(self) -> str: + def default_filename(self) -> str: + """ + Generate a default filename for the component. + By default, this is the component's name then + the class name in lowercase, separated by dot. + + Override this method in subclasses to provide + a custom default filename. + """ name = self.name # type: ignore cls_name = self.__class__.__name__.lower() return f"{name}.{cls_name}" @@ -50,10 +59,6 @@ def __attrs_init_subclass__(cls): COMPONENTS[cls.__name__.lower()] = cls cls.dfn = cls.get_dfn() - def __attrs_post_init__(self): - if not self.filename: - self.filename = self._default_filename() - def __getitem__(self, key): return self.children[key] # type: ignore @@ -89,12 +94,21 @@ def get_dfn(cls) -> Dfn: **blocks, ) + def _preio(self, format: str) -> None: + """Place for any pre-IO setup""" + if not self.filename: + self.filename = self.default_filename() + def load(self, format: str) -> None: + """Load the component from an input file.""" + self._preio(format=format) self._load(format=format) for child in self.children.values(): # type: ignore - child.load(format) + child.load(format=format) def write(self, format: str) -> None: + """Write the component to an input file.""" + self._preio(format=format) self._write(format=format) for child in self.children.values(): # type: ignore - child.write(format) + child.write(format=format) diff --git a/flopy4/mf6/context.py b/flopy4/mf6/context.py index 737031e6..8320b77c 100644 --- a/flopy4/mf6/context.py +++ b/flopy4/mf6/context.py @@ -12,6 +12,7 @@ class Context(Component, ABC): workspace: Path = field(default=None) def __attrs_post_init__(self): + super().__attrs_post_init__() if self.workspace is None: self.workspace = Path.cwd() diff --git a/flopy4/mf6/filters.py b/flopy4/mf6/filters.py index 5d5f0739..596a36f0 100644 --- a/flopy4/mf6/filters.py +++ b/flopy4/mf6/filters.py @@ -95,7 +95,7 @@ def array_chunks(value: xr.DataArray, chunks: Mapping[Hashable, int] | None = No } value = value.chunk(chunks) for chunk in value.data.blocks: - yield chunk.compute() + yield np.squeeze(chunk.compute()) def array2string(value: NDArray) -> str: @@ -112,6 +112,7 @@ def array2string(value: NDArray) -> str: if value.ndim == 1: # add an axis to 1d arrays so np.savetxt writes elements on 1 line value = value[None] + value = np.atleast_1d(value) format = ( "%d" if np.issubdtype(value.dtype, np.integer) diff --git a/flopy4/mf6/gwf/dis.py b/flopy4/mf6/gwf/dis.py index 8529de4e..2580f940 100644 --- a/flopy4/mf6/gwf/dis.py +++ b/flopy4/mf6/gwf/dis.py @@ -76,6 +76,7 @@ class Dis(Package): def __attrs_post_init__(self): self.nnodes = self.ncol * self.nrow * self.nlay + super().__attrs_post_init__() def to_grid(self) -> StructuredGrid: """ diff --git a/flopy4/mf6/model.py b/flopy4/mf6/model.py index da05129f..ca0367c7 100644 --- a/flopy4/mf6/model.py +++ b/flopy4/mf6/model.py @@ -7,4 +7,5 @@ @xattree class Model(Component, ABC): - pass + 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 ae5fad9c..8471c58d 100644 --- a/flopy4/mf6/package.py +++ b/flopy4/mf6/package.py @@ -7,4 +7,7 @@ @xattree class Package(Component, ABC): - pass + def default_filename(self) -> str: + name = self.parent.name if self.parent else self.name # type: ignore + cls_name = self.__class__.__name__.lower() + return f"{name}.{cls_name}" diff --git a/flopy4/mf6/simulation.py b/flopy4/mf6/simulation.py index 1da1afa0..bb1203f3 100644 --- a/flopy4/mf6/simulation.py +++ b/flopy4/mf6/simulation.py @@ -30,6 +30,7 @@ class Simulation(Context): filename: str = field(default="mfsim.nam", init=False) def __attrs_post_init__(self): + super().__attrs_post_init__() if self.filename != "mfsim.nam": warn( "Simulation filename must be 'mfsim.nam'.", diff --git a/flopy4/mf6/templates/blocks.jinja b/flopy4/mf6/templates/blocks.jinja index 07392dac..1b0ba597 100644 --- a/flopy4/mf6/templates/blocks.jinja +++ b/flopy4/mf6/templates/blocks.jinja @@ -1,7 +1,7 @@ {% import 'macros.jinja' as macros with context %} {% for block_name, block_ in (dfn|dict_blocks).items() %} BEGIN {{ block_name.upper() }} -{% for field in block_.values() -%} +{% for field in block_.values() if (field|field_value) is not none -%} {{ macros.field(field) }} {%- endfor %} END {{ block_name.upper() }} diff --git a/flopy4/mf6/templates/macros.jinja b/flopy4/mf6/templates/macros.jinja index 86bf07ff..63baf5b2 100644 --- a/flopy4/mf6/templates/macros.jinja +++ b/flopy4/mf6/templates/macros.jinja @@ -62,7 +62,7 @@ this macro receives the block definition. from that it looks up the value of the one variable with the same name as the block, which custom converter has made sure exists in a sparse dict representation of -an array. we need to spin this out into a block for +an array. we need to expand this into a block for each stress period. #} {% set dict = data[block_name] %} diff --git a/test/conftest.py b/test/conftest.py index 5a43042d..294565eb 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,5 +1,7 @@ from pathlib import Path +pytest_plugins = ["modflow_devtools.fixtures"] + PROJ_ROOT_PATH = Path(__file__).parents[1] DOCS_PATH = PROJ_ROOT_PATH / "docs" EXAMPLES_PATH = DOCS_PATH / "examples" diff --git a/test/test_codec.py b/test/test_codec.py index b365bde7..5c43840b 100644 --- a/test/test_codec.py +++ b/test/test_codec.py @@ -1,5 +1,3 @@ -import pytest - from flopy4.mf6.codec import dumps @@ -37,7 +35,6 @@ def test_dumps_oc(): assert result -@pytest.mark.xfail(reason="TODO 3D arrays") def test_dumps_dis(): from flopy4.mf6.gwf import Dis diff --git a/test/test_component.py b/test/test_component.py index 325b1ef5..80792f8d 100644 --- a/test/test_component.py +++ b/test/test_component.py @@ -272,11 +272,27 @@ def test_ims_dfn(): assert "inner_maximum" in set(dfn["linear"].keys()) -def test_write_ascii(tmp_path): +def test_write_ascii(function_tmpdir): + sim_name = "sim" time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0]) - sim = Simulation(tdis=time, workspace=tmp_path) + grid = StructuredGrid(nlay=1, nrow=10, ncol=10) + sim = Simulation(tdis=time, workspace=function_tmpdir, name=sim_name) + gwf_name = "gwf" + gwf = Gwf(parent=sim, dis=grid, name=gwf_name) + ic = Ic(parent=gwf) + oc = Oc(parent=gwf) + npf = Npf(parent=gwf) + chd = Chd(parent=gwf, head={"*": {(0, 0, 0): 1.0, (0, 9, 9): 0.0}}) + sim.write() - files = list(Path(tmp_path).glob("*")) + files = list(Path(function_tmpdir).glob("*")) file_names = [f.name for f in files] assert "mfsim.nam" in file_names + assert f"{sim_name}.tdis" in file_names + assert f"{gwf_name}.nam" in file_names + assert f"{gwf_name}.dis" in file_names + assert f"{gwf_name}.ic" in file_names + assert f"{gwf_name}.oc" in file_names + assert f"{gwf_name}.npf" in file_names + assert f"{gwf_name}.chd" in file_names