Skip to content

Commit 709f271

Browse files
wpbonelliwpbonelli
authored
quickstart closer, update tests (#104)
1 parent 19fc3cd commit 709f271

File tree

11 files changed

+842
-719
lines changed

11 files changed

+842
-719
lines changed

‎docs/examples/quickstart.py‎

Lines changed: 26 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,45 +1,38 @@
1-
from flopy4.mf6 import Simulation, Tdis
2-
from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc
1+
import numpy as np
32

4-
# TODO rewrite bottom up
3+
from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc
4+
from flopy4.mf6.simulation import Simulation
5+
from flopy4.mf6.tdis import Tdis
56

67
ws = "./mymodel"
78
name = "mymodel"
8-
sim = Simulation(name=name)
9-
tdis = Tdis(sim)
10-
gwf = Gwf(sim, name=name, save_flows=True)
11-
dis = Dis(gwf, nrow=10, ncol=10)
12-
ic = Ic(gwf)
13-
npf = Npf(gwf, save_specific_discharge=True)
9+
tdis = Tdis()
10+
sim = Simulation(name=name, tdis=tdis)
11+
dis = Dis(nrow=10, ncol=10)
12+
gwf = Gwf(parent=sim, name=name, save_flows=True, dis=dis)
13+
ic = Ic(parent=gwf)
14+
npf = Npf(parent=gwf, save_specific_discharge=True)
1415
chd = Chd(
15-
gwf,
16-
head={(0, 0, 0): 1.0, (0, 9, 9): 0.0},
16+
parent=gwf,
17+
head={"*": {(0, 0, 0): 1.0, (0, 9, 9): 0.0}},
1718
)
18-
19-
# list input in the mf6 paradigm, just stored in xarray.
20-
# this is straightforward to implement. even in flopy3?
21-
assert all(
22-
period == Chd.Period((0, 0, 0), 1.0)
23-
for period in chd.data["stress_period_data"]
19+
oc = Oc(
20+
parent=gwf,
21+
budget_file=f"{name}.bud",
22+
head_file=f"{name}.hds",
23+
save_head={"*": "all"},
24+
save_budget={"*": "all"},
2425
)
2526

26-
# adopt xarray paradigm. transpose lists to arrays. this
27-
# is no longer "sparse": in the sense that lists specify
28-
# <maxbound> features of interest, where arrays specify
29-
# the entire domain. but to maximize the value of xarray
30-
# we want to "align" everything to the discretization in
31-
# both time and space.
27+
# check CHD
28+
assert chd.data["head"][0, 0].item() == 1.0
29+
assert chd.data["head"][0, 99].item() == 0.0
30+
assert np.allclose(chd.data["head"][:, 1:99], np.full(98, 1e30))
3231

32+
# TODO: xarray index aliasing nlay/ncol/nrow to k/i/j?
3333
# assert chd.data["head"].loc(dict(k=0, i=0, j=0)) == 1.
3434
# assert chd.data["head"].loc(dict(k=0, i=9, j=9)) == 0.
3535

36-
oc = Oc(
37-
gwf,
38-
budget_filerecord=f"{name}.bud",
39-
head_filerecord=f"{name}.hds",
40-
saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")],
41-
)
42-
43-
# xarray style. this is how imod-python does it too.
44-
# assert oc.data["save_head"] == "all"
45-
# assert oc.data["save_budget"] == "all"
36+
# check OC
37+
assert oc.data["save_head"][0].item() == "all"
38+
assert oc.data["save_budget"][0].item() == "all"

‎flopy4/mf6/converters.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,15 +19,18 @@ def convert_array(value, self_, field) -> NDArray:
1919

2020
# resolve dims
2121
explicit_dims = self_.__dict__.get("dims", {})
22-
inherited_dims = self_.parent.data.dims if self_.parent else {}
22+
inherited_dims = dict(self_.parent.data.dims) if self_.parent else {}
2323
dims = inherited_dims | explicit_dims
2424
shape = [dims.get(d, d) for d in field.dims]
2525
unresolved = [d for d in shape if isinstance(d, str)]
2626
if any(unresolved):
2727
raise ValueError(f"Couldn't resolve dims: {unresolved}")
2828

2929
# create array
30-
a = np.full(shape, fill_value=FILL_DNODATA, dtype=field.dtype)
30+
# TDOD: support other fill values, configurable by field?
31+
a = np.full(
32+
shape, fill_value=field.default or FILL_DNODATA
33+
) # , dtype=field.dtype)
3134

3235
def _get_nn(cellid):
3336
match len(cellid):
@@ -38,19 +41,19 @@ def _get_nn(cellid):
3841
return k * dims["ncpl"] + j
3942
case 3:
4043
k, i, j = cellid
41-
return k * dims["row"] * dims["col"] + i * dims["col"] + j
44+
return k * dims["nrow"] * dims["ncol"] + i * dims["ncol"] + j
4245
case _:
4346
raise ValueError(f"Invalid cellid: {cellid}")
4447

4548
# populate array. TODO: is there a way to do this
4649
# without hardcoding awareness of kper and cellid?
47-
if "per" in dims:
50+
if "nper" in dims:
4851
for kper, period in value.items():
4952
if kper == "*":
5053
kper = 0
5154
match len(shape):
5255
case 1:
53-
a[kper] = value
56+
a[kper] = period
5457
case _:
5558
for cellid, v in period.items():
5659
nn = _get_nn(cellid)

‎flopy4/mf6/gwf/chd.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -42,34 +42,34 @@ class Steps:
4242
)
4343
head: Optional[NDArray[np.floating]] = array(
4444
dims=(
45-
"per",
46-
"node",
45+
"nper",
46+
"nnodes",
4747
),
4848
default=None,
4949
metadata={"block": "period"},
5050
converter=Converter(convert_array, takes_self=True, takes_field=True),
5151
)
5252
aux: Optional[NDArray[np.floating]] = array(
5353
dims=(
54-
"per",
55-
"node",
54+
"nper",
55+
"nnodes",
5656
),
5757
default=None,
5858
metadata={"block": "period"},
5959
converter=Converter(convert_array, takes_self=True, takes_field=True),
6060
)
6161
boundname: Optional[NDArray[np.str_]] = array(
6262
dims=(
63-
"per",
64-
"node",
63+
"nper",
64+
"nnodes",
6565
),
6666
default=None,
6767
metadata={"block": "period"},
6868
converter=Converter(convert_array, takes_self=True, takes_field=True),
6969
)
7070
steps: Optional[NDArray[np.object_]] = array(
7171
Steps,
72-
dims=("per", "node"),
72+
dims=("nper", "nnodes"),
7373
default=None,
7474
metadata={"block": "period"},
7575
converter=Converter(convert_array, takes_self=True, takes_field=True),

‎flopy4/mf6/gwf/dis.py‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -21,60 +21,57 @@ class Dis(Package):
2121
default=False, metadata={"block": "options"}
2222
)
2323
nlay: int = dim(
24-
name="lay",
2524
scope="gwf",
2625
default=1,
2726
metadata={
2827
"block": "dimensions",
2928
},
3029
)
3130
ncol: int = dim(
32-
name="col",
3331
scope="gwf",
3432
default=2,
3533
metadata={
3634
"block": "dimensions",
3735
},
3836
)
3937
nrow: int = dim(
40-
name="row",
4138
scope="gwf",
4239
default=2,
4340
metadata={
4441
"block": "dimensions",
4542
},
4643
)
4744
delr: NDArray[np.floating] = array(
48-
dims=("col",),
45+
dims=("ncol",),
4946
default=1.0,
5047
metadata={"block": "griddata"},
5148
converter=Converter(convert_array, takes_self=True, takes_field=True),
5249
)
5350
delc: NDArray[np.floating] = array(
54-
dims=("row",),
51+
dims=("nrow",),
5552
default=1.0,
5653
metadata={"block": "griddata"},
5754
converter=Converter(convert_array, takes_self=True, takes_field=True),
5855
)
5956
top: NDArray[np.floating] = array(
60-
dims=("col", "row"),
57+
dims=("ncol", "nrow"),
6158
default=1.0,
6259
metadata={"block": "griddata"},
6360
converter=Converter(convert_array, takes_self=True, takes_field=True),
6461
)
6562
botm: NDArray[np.floating] = array(
66-
dims=("col", "row", "lay"),
63+
dims=("ncol", "nrow", "nlay"),
6764
default=0.0,
6865
metadata={"block": "griddata"},
6966
converter=Converter(convert_array, takes_self=True, takes_field=True),
7067
)
7168
idomain: NDArray[np.integer] = array(
72-
dims=("col", "row", "lay"),
69+
dims=("ncol", "nrow", "nlay"),
7370
default=1,
7471
metadata={"block": "griddata"},
7572
converter=Converter(convert_array, takes_self=True, takes_field=True),
7673
)
77-
nnodes: int = dim(name="node", scope="gwf", init=False)
74+
nnodes: int = dim(scope="gwf", init=False)
7875

7976
def __attrs_post_init__(self):
8077
self.nnodes = self.ncol * self.nrow * self.nlay

‎flopy4/mf6/gwf/ic.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
@xattree
1111
class Ic(Package):
1212
strt: NDArray[np.floating] = array(
13-
dims=("node",),
13+
dims=("nnodes",),
1414
default=1.0,
1515
metadata={"block": "packagedata"},
1616
converter=Converter(convert_array, takes_self=True, takes_field=True),

‎flopy4/mf6/gwf/npf.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -65,49 +65,49 @@ class Xt3dOptions:
6565
default=None, metadata={"block": "options"}
6666
)
6767
icelltype: NDArray[np.integer] = array(
68-
dims=("node",),
68+
dims=("nnodes",),
6969
default=0,
7070
metadata={"block": "griddata"},
7171
converter=Converter(convert_array, takes_self=True, takes_field=True),
7272
)
7373
k: NDArray[np.floating] = array(
74-
dims=("node",),
74+
dims=("nnodes",),
7575
default=1.0,
7676
metadata={"block": "griddata"},
7777
converter=Converter(convert_array, takes_self=True, takes_field=True),
7878
)
7979
k22: Optional[NDArray[np.floating]] = array(
80-
dims=("node",),
80+
dims=("nnodes",),
8181
default=None,
8282
metadata={"block": "griddata"},
8383
converter=Converter(convert_array, takes_self=True, takes_field=True),
8484
)
8585
k33: Optional[NDArray[np.floating]] = array(
86-
dims=("node",),
86+
dims=("nnodes",),
8787
default=None,
8888
metadata={"block": "griddata"},
8989
converter=Converter(convert_array, takes_self=True, takes_field=True),
9090
)
9191
angle1: Optional[NDArray[np.floating]] = array(
92-
dims=("node",),
92+
dims=("nnodes",),
9393
default=None,
9494
metadata={"block": "griddata"},
9595
converter=Converter(convert_array, takes_self=True, takes_field=True),
9696
)
9797
angle2: Optional[NDArray[np.floating]] = array(
98-
dims=("node",),
98+
dims=("nnodes",),
9999
default=None,
100100
metadata={"block": "griddata"},
101101
converter=Converter(convert_array, takes_self=True, takes_field=True),
102102
)
103103
angle3: Optional[NDArray[np.floating]] = array(
104-
dims=("node",),
104+
dims=("nnodes",),
105105
default=None,
106106
metadata={"block": "griddata"},
107107
converter=Converter(convert_array, takes_self=True, takes_field=True),
108108
)
109109
wetdry: Optional[NDArray[np.floating]] = array(
110-
dims=("node",),
110+
dims=("nnodes",),
111111
default=None,
112112
metadata={"block": "griddata"},
113113
converter=Converter(convert_array, takes_self=True, takes_field=True),

‎flopy4/mf6/gwf/oc.py‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -53,17 +53,27 @@ class Period:
5353
format: Optional[Format] = field(
5454
default=None, init=False, metadata={"block": "options"}
5555
)
56-
saverecord: Optional[NDArray[np.object_]] = array(
57-
Period,
58-
dims=("per",),
59-
default=None,
56+
save_head: Optional[NDArray[np.str_ | np.integer]] = array(
57+
dims=("nper",),
58+
default="all",
6059
metadata={"block": "perioddata"},
6160
converter=Converter(convert_array, takes_self=True, takes_field=True),
6261
)
63-
printrecord: Optional[NDArray[np.object_]] = array(
64-
Period,
65-
dims=("per",),
66-
default=None,
62+
save_budget: Optional[NDArray[np.str_ | np.integer]] = array(
63+
dims=("nper",),
64+
default="all",
65+
metadata={"block": "perioddata"},
66+
converter=Converter(convert_array, takes_self=True, takes_field=True),
67+
)
68+
print_head: Optional[NDArray[np.str_ | np.integer]] = array(
69+
dims=("nper",),
70+
default="all",
71+
metadata={"block": "perioddata"},
72+
converter=Converter(convert_array, takes_self=True, takes_field=True),
73+
)
74+
print_budget: Optional[NDArray[np.str_ | np.integer]] = array(
75+
dims=("nper",),
76+
default="all",
6777
metadata={"block": "perioddata"},
6878
converter=Converter(convert_array, takes_self=True, takes_field=True),
6979
)

‎flopy4/mf6/tdis.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ class PeriodData:
1919
tsmult: float
2020

2121
nper: int = dim(
22-
name="per",
2322
default=1,
2423
scope=ROOT,
2524
metadata={"block": "dimensions"},
@@ -32,19 +31,19 @@ class PeriodData:
3231
)
3332
perlen: NDArray[np.floating] = array(
3433
default=1.0,
35-
dims=("per",),
34+
dims=("nper",),
3635
metadata={"block": "perioddata"},
3736
converter=Converter(convert_array, takes_self=True, takes_field=True),
3837
)
3938
nstp: NDArray[np.integer] = array(
4039
default=1,
41-
dims=("per",),
40+
dims=("nper",),
4241
metadata={"block": "perioddata"},
4342
converter=Converter(convert_array, takes_self=True, takes_field=True),
4443
)
4544
tsmult: NDArray[np.floating] = array(
4645
default=1.0,
47-
dims=("per",),
46+
dims=("nper",),
4847
metadata={"block": "perioddata"},
4948
converter=Converter(convert_array, takes_self=True, takes_field=True),
5049
)

0 commit comments

Comments
 (0)