Skip to content

Commit 56a0267

Browse files
authored
first draft of array converters (#102)
1 parent 9225da3 commit 56a0267

File tree

6 files changed

+213
-174
lines changed

6 files changed

+213
-174
lines changed

docs/examples/quickstart.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,19 @@
1-
from flopy4.mf6 import Sim, Tdis
1+
from flopy4.mf6 import Simulation, Tdis
22
from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc
33

44
# TODO rewrite bottom up
55

66
ws = "./mymodel"
77
name = "mymodel"
8-
sim = Sim(name=name, path=ws, exe="mf6")
8+
sim = Simulation(name=name)
99
tdis = Tdis(sim)
1010
gwf = Gwf(sim, name=name, save_flows=True)
1111
dis = Dis(gwf, nrow=10, ncol=10)
1212
ic = Ic(gwf)
1313
npf = Npf(gwf, save_specific_discharge=True)
1414
chd = Chd(
1515
gwf,
16-
stress_period_data=[[(0, 0, 0), 1.0], [(0, 9, 9), 0.0]],
16+
head={(0, 0, 0): 1.0, (0, 9, 9): 0.0},
1717
)
1818

1919
# list input in the mf6 paradigm, just stored in xarray.

docs/examples/quickstart_expanded.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,12 @@
2626
changes to support.
2727
"""
2828

29-
from flopy4.mf6 import Sim, Tdis
29+
from flopy4.mf6 import Simulation, Tdis
3030
from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc
3131

3232
ws = "./mymodel"
3333
name = "mymodel"
34-
sim = Sim(name=name, path=ws, exe="mf6")
34+
sim = Simulation(name=name, path=ws, exe="mf6")
3535
tdis = Tdis(sim)
3636
gwf = Gwf(sim, name=name, save_flows=True)
3737
dis = Dis(gwf, nrow=10, ncol=10)

flopy4/mf6/gwf/chd.py

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,42 @@
11
from pathlib import Path
22
from typing import Optional
33

4+
import attrs
45
import numpy as np
56
from attrs import define
67
from numpy.typing import NDArray
7-
from xattree import array, field, xattree
8+
from xattree import _get_xatspec, array, field, xattree
89

910
from flopy4.mf6 import Package
1011

12+
dnodata = 1e30
13+
14+
15+
def _get_nn(ncol, nrow, k, i, j):
16+
return k * nrow * ncol + i * ncol + j
17+
18+
19+
def _convert_array(value, self_, field):
20+
if not isinstance(value, dict):
21+
return value
22+
23+
inherited_dims = self_.__dict__.get("dims", {})
24+
spec = _get_xatspec(type(self_))
25+
field = spec.arrays["head"]
26+
shape = field.dims
27+
if not shape:
28+
raise ValueError()
29+
dims = [inherited_dims.get(d, d) for d in shape]
30+
# TODO pull out dtype from annotation
31+
a = np.full(dims, fill_value=dnodata, dtype=np.float64)
32+
for kper, period in value.items():
33+
if kper == "*":
34+
kper = 0
35+
for cellid, v in period.items():
36+
nn = _get_nn(inherited_dims["col"], inherited_dims["row"], *cellid)
37+
a[kper, nn] = v
38+
return a
39+
1140

1241
@xattree(multi="list")
1342
class Chd(Package):
@@ -46,6 +75,9 @@ class Steps:
4675
),
4776
default=None,
4877
metadata={"block": "period"},
78+
converter=attrs.Converter(
79+
_convert_array, takes_self=True, takes_field=True
80+
),
4981
)
5082
aux: Optional[NDArray[np.floating]] = array(
5183
dims=(

0 commit comments

Comments
 (0)