Skip to content

Commit df7d31c

Browse files
wpbonelliwpbonelli
committed
twri wip
1 parent ddc1cbe commit df7d31c

File tree

12 files changed

+385
-21
lines changed

12 files changed

+385
-21
lines changed

‎docs/examples/twri.py‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
# # TWRI
2+
#
3+
# Import dependencies.
4+
5+
from pathlib import Path
6+
7+
import numpy as np
8+
import xarray as xr
9+
10+
import flopy4
11+
12+
# Timing
13+
time = flopy4.mf6.utils.time.Time.from_timestamps(
14+
["2000-01-01", "2000-01-02", "2000-01-03", "2000-01-04"]
15+
)
16+
nper = time.nper
17+
18+
# Grid
19+
nlay = 5
20+
nrow = 15
21+
ncol = 15
22+
shape = (nlay, nrow, ncol)
23+
nodes = np.prod(shape)
24+
delr = np.ones((ncol,)) * 5000.0
25+
delc = np.ones((nrow,)) * 5000.0
26+
idomain = np.ones(shape, dtype=int)
27+
top = np.ones((nrow, ncol), dtype=float) * 200.0
28+
bottom = np.stack([np.full((nrow, ncol), val) for val in [-150.0, -200.0, -300.0, -350.0, -450.0]])
29+
grid = flopy4.mf6.utils.grid.StructuredGrid(
30+
nlay=nlay, nrow=nrow, ncol=ncol, top=top, botm=bottom, delr=delr, delc=delc, idomain=idomain
31+
)
32+
dims = {"nper": nper, **dict(grid.dataset.sizes)} # TODO: temporary
33+
34+
# Grid discretization
35+
dis = flopy4.mf6.gwf.Dis.from_grid(grid=grid)
36+
37+
# Constant head boundary on the left
38+
chd = flopy4.mf6.gwf.Chd(
39+
head={"*": {(k, i, 0): 0.0 for k in range(nlay) for i in range(nrow)}},
40+
print_input=True,
41+
print_flows=True,
42+
save_flows=True,
43+
dims=dims,
44+
)
45+
46+
# Drain in the center left of the model
47+
elevation = [0.0, 0.0, 10.0, 20.0, 30.0, 50.0, 70.0, 90.0, 100.0]
48+
conductance = 1.0
49+
drn = flopy4.mf6.gwf.Drn(
50+
elev={"*": {(0, 7, j): elevation[j] for j in range(9)}},
51+
cond={"*": {(0, 7, j): conductance for j in range(9)}},
52+
print_input=True,
53+
print_flows=True,
54+
save_flows=True,
55+
dims=dims,
56+
)
57+
58+
# Initial conditions
59+
ic = flopy4.mf6.gwf.Ic(strt=0.0, dims=dims)
60+
61+
# Node properties
62+
icelltype = np.stack([np.full((nrow, ncol), val) for val in [1, 0, 0, 0, 0]])
63+
k = np.stack([np.full((nrow, ncol), val) for val in [1.0e-3, 1.0e-8, 1.0e-4, 5.0e-7, 2.0e-4]])
64+
k33 = np.stack([np.full((nrow, ncol), val) for val in [1.0e-3, 1.0e-8, 1.0e-4, 5.0e-7, 2.0e-4]])
65+
npf = flopy4.mf6.gwf.Npf(
66+
# TODO: no need for reshaping once array structuring converter is done
67+
icelltype=icelltype.reshape((nodes,)),
68+
k=k.reshape((nodes,)),
69+
k33=k33.reshape((nodes,)),
70+
cvoptions=flopy4.mf6.gwf.Npf.CvOptions(variablecv=True, dewatered=True),
71+
perched=True,
72+
save_flows=True,
73+
dims=dims,
74+
)
75+
76+
# Storage
77+
sto = flopy4.mf6.gwf.Sto(
78+
storagecoefficient=False,
79+
ss=1.0e-5,
80+
sy=0.15,
81+
transient={"*": False},
82+
iconvert=0,
83+
dims=dims,
84+
)
85+
86+
# Uniform recharge on the top layer
87+
rch_rate = xr.full_like(grid.idomain.sel(k=1), 3.0e-8, dtype=float)
88+
rch = flopy4.mf6.gwf.Rch(recharge=rch_rate, dims=dims)
89+
90+
# Output control
91+
# TODO: show both ways to set up the Oc package, strings
92+
# and proper record types? and/or with perioddata param?
93+
oc = flopy4.mf6.gwf.Oc(save_head="all", save_budget="all", dims=dims)
94+
95+
# Wells scattered throughout the model
96+
wel_q = -5.0
97+
wel_nodes = [
98+
[3, 4, 10, -5.0],
99+
[2, 3, 5, -5.0],
100+
[2, 5, 11, -5.0],
101+
[0, 8, 7, -5.0],
102+
[0, 8, 9, -5.0],
103+
[0, 8, 11, -5.0],
104+
[0, 8, 13, -5.0],
105+
[0, 10, 7, -5.0],
106+
[0, 10, 9, -5.0],
107+
[0, 10, 11, -5.0],
108+
[0, 10, 13, -5.0],
109+
[0, 12, 7, -5.0],
110+
[0, 12, 9, -5.0],
111+
[0, 12, 11, -5.0],
112+
[0, 12, 13, -5.0],
113+
]
114+
wel = flopy4.mf6.gwf.Wel(
115+
q={"*": {(layer, row, col): wel_q for layer, row, col, wel_q in wel_nodes}},
116+
dims=dims,
117+
)
118+
119+
# Flow model
120+
gwf = flopy4.mf6.gwf.Gwf(
121+
dis=grid,
122+
chd=chd,
123+
ic=ic,
124+
npf=npf,
125+
sto=sto,
126+
oc=oc,
127+
rch=rch,
128+
wel=wel,
129+
drn=drn,
130+
dims=dims,
131+
)
132+
133+
# Solver
134+
ims = flopy4.mf6.Ims(
135+
print_option="summary",
136+
outer_dvclose=1.0e-4,
137+
outer_maximum=500,
138+
under_relaxation=None,
139+
inner_dvclose=1.0e-4,
140+
inner_rclose=0.001,
141+
inner_maximum=100,
142+
linear_acceleration="cg",
143+
scaling_method=None,
144+
reordering_method=None,
145+
relaxation_factor=0.97,
146+
)
147+
148+
# TDIS
149+
tdis = flopy4.mf6.simulation.Tdis.from_time(time)
150+
151+
# Create workspace
152+
workspace = Path(__file__).parent / "twri"
153+
workspace.mkdir(parents=True, exist_ok=True)
154+
155+
# Create simulation
156+
sim = flopy4.mf6.simulation.Simulation(
157+
name="twri",
158+
tdis=tdis,
159+
models={"gwf": gwf},
160+
solutions={"ims": ims},
161+
workspace=workspace,
162+
)
163+
164+
# Write input files and run the simulation
165+
sim.write()
166+
sim.run() # assumes the ``mf6`` executable is available on your PATH.
167+
168+
# Load head results
169+
gwf_ws = Path(workspace) / gwf.name
170+
head = flopy4.mf6.utils.open_hds(
171+
gwf_ws / f"{gwf.name}.hds",
172+
gwf_ws / f"{dis.name}.dis.grb",
173+
)
174+
175+
# Plot head results
176+
head.isel(layer=0, time=0).plot.contourf()

‎flopy4/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from flopy4 import mf6
2+
3+
__all__ = ["mf6"]

‎flopy4/mf6/__init__.py‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,25 @@
55
from tomli import load as load_toml
66
from tomli_w import dump as dump_toml
77

8+
# Import submodules to make them accessible via flopy4.mf6.*
9+
from flopy4.mf6 import gwf, simulation, solution, utils
810
from flopy4.mf6.codec import dump as dump_mf6
911
from flopy4.mf6.codec import load as load_mf6
1012
from flopy4.mf6.component import Component
1113
from flopy4.mf6.converter import structure, unstructure
14+
from flopy4.mf6.ims import Ims
15+
from flopy4.mf6.simulation import Simulation
16+
from flopy4.mf6.tdis import Tdis
1217
from flopy4.uio import DEFAULT_REGISTRY
1318

19+
__all__ = ["gwf", "simulation", "solution", "utils", "Ims", "Tdis", "Simulation"]
20+
21+
22+
class WriteError(Exception):
23+
"""An error occurred while writing a component."""
24+
25+
pass
26+
1427

1528
def _load_mf6(path: Path) -> Component:
1629
with open(path, "r") as fp:
@@ -29,17 +42,26 @@ def _load_toml(path: Path) -> Component:
2942

3043
def _write_mf6(component: Component) -> None:
3144
with open(component.path, "w") as fp:
32-
dump_mf6(unstructure(component), fp)
45+
data = unstructure(component)
46+
try:
47+
dump_mf6(data, fp)
48+
except Exception as e:
49+
raise WriteError(
50+
f"Failed to write MF6 format file for component '{component.name}' " # type: ignore
51+
f"of type {component.__class__.__name__}"
52+
) from e
3353

3454

3555
def _write_json(component: Component) -> None:
3656
with open(component.path, "w") as fp:
37-
dump_json(unstructure(component), fp, indent=4)
57+
data = unstructure(component)
58+
dump_json(data, fp, indent=4)
3859

3960

4061
def _write_toml(component: Component) -> None:
4162
with open(component.path, "wb") as fp:
42-
dump_toml(unstructure(component), fp)
63+
data = unstructure(component)
64+
dump_toml(data, fp)
4365

4466

4567
DEFAULT_REGISTRY.register_loader(Component, "mf6", _load_mf6)

‎flopy4/mf6/codec/writer/filters.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
from flopy4.mf6.constants import FILL_DNODATA
1212

13-
ArrayHow = Literal["constant", "internal", "external"]
13+
ArrayHow = Literal["constant", "internal", "external", "layered internal"]
1414

1515

1616
def array_how(value: xr.DataArray) -> ArrayHow:
@@ -26,7 +26,11 @@ def array_how(value: xr.DataArray) -> ArrayHow:
2626
return "external"
2727
if value.max() == value.min():
2828
return "constant"
29-
return "internal"
29+
if value.ndim <= 2:
30+
return "internal"
31+
if value.ndim == 3:
32+
return "layered internal"
33+
raise ValueError(f"Arrays with ndim > 3 are not supported, got ndim={value.ndim}")
3034

3135

3236
def array2const(value: xr.DataArray) -> Scalar:
@@ -95,7 +99,7 @@ def array2string(value: NDArray) -> str:
9599
buffer = StringIO()
96100
value = np.asarray(value)
97101
if value.ndim > 2:
98-
raise ValueError("Only 1D and 2D arrays are supported.")
102+
raise ValueError(f"Only 1D and 2D arrays are supported, got ndim={value.ndim}")
99103
if value.ndim == 1:
100104
# add an axis to 1d arrays so np.savetxt writes elements on 1 line
101105
value = value[None]

‎flopy4/mf6/codec/writer/templates/macros.jinja‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,16 @@
3737
{% for layer in value -%}
3838
{{ inset }}CONSTANT {{ layer|array2const }}
3939
{%- endfor %}
40+
{% elif how == "layered internal" %}
41+
{% for layer in value %}
42+
43+
{{ inset }}INTERNAL
44+
{% for chunk in layer|array2chunks -%}
45+
{{ (2 * inset) ~ chunk|array2string }}
46+
{%- endfor %}
47+
{%- endfor %}
4048
{% elif how == "internal" %}
41-
INTERNAL
49+
{{ inset }}INTERNAL
4250
{% for chunk in value|array2chunks -%}
4351
{{ (2 * inset) ~ chunk|array2string }}
4452
{%- endfor %}

‎flopy4/mf6/config.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
# TODO use https://environ-config.readthedocs.io/en/stable/?
22

3-
SPARSE_THRESHOLD = 1000
3+
SPARSE_THRESHOLD = 100000

‎flopy4/mf6/context.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,12 @@
77
from flopy4.mf6.component import Component
88
from flopy4.mf6.constants import MF6
99
from flopy4.mf6.spec import field
10+
from flopy4.utils import to_path
1011

1112

1213
@xattree
1314
class Context(Component, ABC):
14-
workspace: Path = field(default=None)
15+
workspace: Path = field(default=None, converter=to_path)
1516

1617
def __attrs_post_init__(self):
1718
super().__attrs_post_init__()

‎flopy4/mf6/gwf/__init__.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,15 @@
1313
from flopy4.mf6.gwf.ic import Ic
1414
from flopy4.mf6.gwf.npf import Npf
1515
from flopy4.mf6.gwf.oc import Oc
16+
from flopy4.mf6.gwf.rch import Rch
1617
from flopy4.mf6.gwf.sto import Sto
1718
from flopy4.mf6.gwf.wel import Wel
1819
from flopy4.mf6.model import Model
1920
from flopy4.mf6.spec import field, path
2021
from flopy4.mf6.utils import open_cbc, open_hds
2122
from flopy4.utils import to_path
2223

23-
__all__ = ["Gwf", "Chd", "Dis", "Drn", "Ic", "Npf", "Oc", "Sto", "Wel"]
24+
__all__ = ["Gwf", "Chd", "Dis", "Drn", "Ic", "Npf", "Oc", "Sto", "Wel", "Rch"]
2425

2526

2627
def convert_grid(value):
@@ -80,6 +81,7 @@ def budget(self):
8081
chd: list[Chd] = field(block="packages")
8182
wel: list[Wel] = field(block="packages")
8283
drn: list[Drn] = field(block="packages")
84+
rch: list[Rch] = field(block="packages")
8385
output: Output = attrs.field(
8486
default=attrs.Factory(lambda self: Gwf.Output(self), takes_self=True)
8587
)

‎flopy4/mf6/simulation.py‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,11 +34,14 @@ def default_filename(self) -> str:
3434
def __attrs_post_init__(self):
3535
super().__attrs_post_init__()
3636
if self.filename != "mfsim.nam":
37-
warn(
38-
"Simulation filename must be 'mfsim.nam'.",
39-
UserWarning,
40-
)
37+
if self.filename is not None:
38+
warn(
39+
"Simulation filename must be 'mfsim.nam'.",
40+
UserWarning,
41+
)
4142
self.filename = "mfsim.nam"
43+
for model in self.models.values():
44+
model.workspace = self.workspace
4245

4346
@property
4447
def time(self) -> Time:

‎flopy4/mf6/utils/__init__.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
1+
# Import submodules to make them accessible via flopy4.mf6.utils.*
2+
from flopy4.mf6.utils import grid, time
3+
14
from .cbc_reader import open_cbc
25
from .heads_reader import open_hds
36

4-
__all__ = ["open_hds", "open_cbc"]
7+
__all__ = ["open_hds", "open_cbc", "grid", "time"]

0 commit comments

Comments
 (0)