basinledger 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,7 @@
1
+ """Basin Ledger: a basin's water, quantity and quality, balanced day by day across coupled models.
2
+
3
+ SWAT+ (through swatplus-bmi), MODFLOW 6 (through libmf6 and modflowapi) and water-user components (from
4
+ waterledger-models) are stepped one day at a time by one Python driver, on NextGen geography. See docs/design.md.
5
+ """
6
+
7
+ __version__ = "0.1.0"
@@ -0,0 +1,66 @@
1
+ """The command line.
2
+
3
+ python -m basinledger run basin.toml [--run NAME] [--days N] [--baseline] [--out DIR]
4
+ [--save-at DAY --save DIR] [--restart DIR]
5
+ python -m basinledger status basin.toml
6
+
7
+ ``run`` steps the example's models through its period (or ``--days``), checking the ledger every day, and writes
8
+ ledger.csv, series.csv and summary.json to the run's outputs folder. ``--baseline`` runs it without its water users.
9
+ ``--save-at`` stops after that day and saves the state to ``--save``; ``--restart`` carries on from a saved state.
10
+ ``status`` prints whether the example is calibrated or a screening run.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import argparse
16
+ import json
17
+ import sys
18
+
19
+ from . import config as C
20
+ from .driver import Basin
21
+
22
+
23
+ def main(argv=None) -> int:
24
+ ap = argparse.ArgumentParser(prog="basinledger", description="Basin Ledger: coupled basin water accounting.")
25
+ sub = ap.add_subparsers(dest="cmd", required=True)
26
+ r = sub.add_parser("run", help="run an example")
27
+ r.add_argument("config")
28
+ r.add_argument("--run", default=None, help="run name (default: base, or baseline with --baseline)")
29
+ r.add_argument("--days", type=int, default=None)
30
+ r.add_argument("--baseline", action="store_true", help="without the example's water users")
31
+ r.add_argument("--out", default=None, help="outputs folder (default: <run folder>/outputs)")
32
+ r.add_argument("--save-at", type=int, default=None, help="stop after this many days and save the state")
33
+ r.add_argument("--save", default=None, help="folder for the saved state")
34
+ r.add_argument("--restart", default=None, help="carry on from a saved state")
35
+ s = sub.add_parser("status", help="calibrated or screening run")
36
+ s.add_argument("config")
37
+ a = ap.parse_args(argv)
38
+
39
+ if a.cmd == "status":
40
+ cfg = C.load(a.config)
41
+ print(f"{cfg.name}: {cfg.calibration_status}")
42
+ return 0
43
+
44
+ if (a.save_at is None) != (a.save is None):
45
+ ap.error("--save-at and --save go together")
46
+ run = a.run or ("baseline" if a.baseline else "base")
47
+ b = Basin.from_config(a.config, run=run, users=not a.baseline, days=a.days, restart=a.restart)
48
+ try:
49
+ if a.save_at is not None:
50
+ b.run(a.save_at - b.day)
51
+ b.save(a.save)
52
+ else:
53
+ b.run()
54
+ out = b.write_outputs(a.out)
55
+ finally:
56
+ b.finalize()
57
+ summary = json.loads((out / "summary.json").read_text())
58
+ print(f"{summary['example']} ({summary['status']}): {summary['days']} days, ledger closed every day")
59
+ print(f"outputs: {out}")
60
+ if summary["status_note"]:
61
+ print("note: " + summary["status_note"])
62
+ return 0
63
+
64
+
65
+ if __name__ == "__main__":
66
+ sys.exit(main())
@@ -0,0 +1,227 @@
1
+ """Calibration: parameters, observations, the forward model, statistics, and the calibration record.
2
+
3
+ Built:
4
+ - the record and its status (no record, or a status other than "calibrated", means a screening run);
5
+ - parameters and observations as tables (CSV);
6
+ - simulated values at observations, from a run's series (channel flow, heads at named cells);
7
+ - statistics: NSE, KGE, percent bias, a low-flow NSE (on log flows), mean error and RMSE;
8
+ - the forward model: run an example and write simulated values at every observation, for PEST++ to read.
9
+
10
+ Not built yet (docs/design.md, Calibration):
11
+ - applying a parameter set to the models' input files: SWAT+ through ``calibration.cal``, MODFLOW 6 arrays by zone or
12
+ pilot point (flopy, pyEMU), users through their config. ``apply_parameters`` refuses anything it can't apply yet;
13
+ - writing PEST++ IES control and template files (pyEMU);
14
+ - fetching observations: USGS daily streamflow and groundwater levels (``dataretrieval``), matched to hydrofabric
15
+ flowpaths and nexuses and to MODFLOW 6 cells.
16
+
17
+ Tables:
18
+
19
+ parameters.csv name, model, target, value, lower, upper, transform
20
+ model: swatplus | modflow6 | user; target: what it changes (a calibration.cal parameter and
21
+ object, a MODFLOW 6 package array and zone, a user's config key); transform: none | log
22
+ observations.csv name, kind, object, date, value, weight, source
23
+ kind: flow (m3 s-1, object = channel or hydrofabric id) | head (m, object = a named head
24
+ output); date ISO; source: e.g. "USGS 04119055 dv"
25
+ """
26
+
27
+ from __future__ import annotations
28
+
29
+ import csv
30
+ import tomllib
31
+ from dataclasses import dataclass
32
+ from pathlib import Path
33
+
34
+ import numpy as np
35
+
36
+ STATUSES = ("calibrated", "screening")
37
+ SCREENING = "screening"
38
+
39
+
40
+ def model_statuses(record: Path | None, models: list[str]) -> dict[str, str]:
41
+ """Each model's status from an example's record.
42
+
43
+ A record gives either one status for the example (top-level ``status``, ``[calibration]``, ``[validation]``)
44
+ or one per model (``[models.<name>]`` with the same keys), e.g. streamflow from NextGen calibrated and MODFLOW 6
45
+ a screening model. A model the record doesn't name is screening. A calibrated model needs its calibration and
46
+ validation statistics.
47
+ """
48
+ if record is None:
49
+ return {m: SCREENING for m in models}
50
+ with open(record, "rb") as f:
51
+ rec = tomllib.load(f)
52
+ per = rec.get("models")
53
+ out = {}
54
+ for m in models:
55
+ entry = per.get(m) if per is not None else rec
56
+ status = (entry or {}).get("status", SCREENING)
57
+ if status not in STATUSES:
58
+ raise ValueError(f"{record}: status must be one of {STATUSES}, not {status!r}")
59
+ if status == "calibrated":
60
+ missing = [k for k in ("calibration", "validation") if k not in entry]
61
+ if missing:
62
+ raise ValueError(f"{record}: a calibrated record needs its {' and '.join(missing)} statistics")
63
+ out[m] = status
64
+ return out
65
+
66
+
67
+ def record_targets(record: Path | None) -> dict[str, str]:
68
+ """What each model in a record was calibrated to (its ``target``), e.g. streamflow at a gage."""
69
+ if record is None:
70
+ return {}
71
+ with open(record, "rb") as f:
72
+ rec = tomllib.load(f)
73
+ return {m: e["target"] for m, e in (rec.get("models") or {}).items() if "target" in e}
74
+
75
+
76
+ def record_status(record: Path | None, models: list[str] | None = None) -> str:
77
+ """'calibrated' only if every model the example runs is calibrated; otherwise a screening run."""
78
+ st = model_statuses(record, models or ["all"])
79
+ return "calibrated" if all(v == "calibrated" for v in st.values()) else SCREENING
80
+
81
+
82
+ @dataclass
83
+ class Parameter:
84
+ name: str
85
+ model: str
86
+ target: str
87
+ value: float
88
+ lower: float
89
+ upper: float
90
+ transform: str = "none"
91
+
92
+ def __post_init__(self):
93
+ if self.model not in ("swatplus", "modflow6", "user"):
94
+ raise ValueError(f"parameter {self.name}: model must be swatplus, modflow6 or user")
95
+ if not self.lower <= self.value <= self.upper:
96
+ raise ValueError(f"parameter {self.name}: {self.value} outside [{self.lower}, {self.upper}]")
97
+
98
+
99
+ @dataclass
100
+ class Observation:
101
+ name: str
102
+ kind: str # flow | head
103
+ object: str # channel (or hydrofabric id) for flow; named head output for head
104
+ date: str
105
+ value: float
106
+ weight: float = 1.0
107
+ source: str = ""
108
+
109
+
110
+ def read_parameters(path: Path) -> list[Parameter]:
111
+ with open(path, newline="") as f:
112
+ return [Parameter(r["name"], r["model"], r["target"], float(r["value"]), float(r["lower"]),
113
+ float(r["upper"]), r.get("transform") or "none") for r in csv.DictReader(f)]
114
+
115
+
116
+ def read_observations(path: Path) -> list[Observation]:
117
+ with open(path, newline="") as f:
118
+ obs = [Observation(r["name"], r["kind"], r["object"], r["date"], float(r["value"]),
119
+ float(r.get("weight") or 1.0), r.get("source", "")) for r in csv.DictReader(f)]
120
+ for o in obs:
121
+ if o.kind not in ("flow", "head"):
122
+ raise ValueError(f"observation {o.name}: kind must be flow or head")
123
+ return obs
124
+
125
+
126
+ def apply_parameters(model: str, folder: Path, params: list[Parameter]):
127
+ """Write a parameter set into a run's copy of one model's inputs, before it starts (the driver's ``prepare``
128
+ hook). Not built yet for any model."""
129
+ mine = [p for p in params if p.model == model]
130
+ if mine:
131
+ raise NotImplementedError(f"applying {model} parameters is planned; see docs/design.md, Calibration")
132
+
133
+
134
+ def simulated(series: dict[str, list[float]], dates: list[str], observations: list[Observation]) -> np.ndarray:
135
+ """The run's value at each observation: channel flow (``flow:<object>``) or head (``head:<object>``) on its
136
+ date."""
137
+ index = {d: i for i, d in enumerate(dates)}
138
+ out = np.empty(len(observations))
139
+ for j, o in enumerate(observations):
140
+ key = f"{o.kind}:{o.object}"
141
+ if key not in series:
142
+ raise KeyError(f"observation {o.name}: the run has no {key!r} output (add it under [outputs])")
143
+ if o.date not in index:
144
+ raise KeyError(f"observation {o.name}: {o.date} is outside the run")
145
+ out[j] = series[key][index[o.date]]
146
+ return out
147
+
148
+
149
+ def nse(sim, obs) -> float:
150
+ sim, obs = np.asarray(sim, float), np.asarray(obs, float)
151
+ return float(1 - np.sum((sim - obs) ** 2) / np.sum((obs - obs.mean()) ** 2))
152
+
153
+
154
+ def log_nse(sim, obs, floor: float = 1e-6) -> float:
155
+ """NSE on log flows: weights low flows, where a large user shows."""
156
+ return nse(np.log(np.maximum(sim, floor)), np.log(np.maximum(obs, floor)))
157
+
158
+
159
+ def kge(sim, obs) -> float:
160
+ """Kling-Gupta efficiency (Gupta et al. 2009)."""
161
+ sim, obs = np.asarray(sim, float), np.asarray(obs, float)
162
+ r = np.corrcoef(sim, obs)[0, 1]
163
+ return float(1 - np.sqrt((r - 1) ** 2 + (sim.std() / obs.std() - 1) ** 2 + (sim.mean() / obs.mean() - 1) ** 2))
164
+
165
+
166
+ def pbias(sim, obs) -> float:
167
+ """Percent bias, positive when the model overestimates."""
168
+ sim, obs = np.asarray(sim, float), np.asarray(obs, float)
169
+ return float(100 * np.sum(sim - obs) / np.sum(obs))
170
+
171
+
172
+ def mean_error(sim, obs) -> float:
173
+ return float(np.mean(np.asarray(sim, float) - np.asarray(obs, float)))
174
+
175
+
176
+ def rmse(sim, obs) -> float:
177
+ return float(np.sqrt(np.mean((np.asarray(sim, float) - np.asarray(obs, float)) ** 2)))
178
+
179
+
180
+ def statistics(observations: list[Observation], sim: np.ndarray) -> dict[str, dict[str, float]]:
181
+ """Statistics per site (kind and object): flow gets NSE, log NSE, KGE and percent bias; heads mean error, RMSE
182
+ and RMSE over the observed range."""
183
+ out = {}
184
+ for key in dict.fromkeys((o.kind, o.object) for o in observations):
185
+ idx = [i for i, o in enumerate(observations) if (o.kind, o.object) == key]
186
+ s = sim[idx]
187
+ o = np.array([observations[i].value for i in idx])
188
+ if key[0] == "flow":
189
+ out[f"flow:{key[1]}"] = {"n": len(idx), "nse": nse(s, o), "log_nse": log_nse(s, o), "kge": kge(s, o),
190
+ "pbias": pbias(s, o)}
191
+ else:
192
+ rng = float(o.max() - o.min())
193
+ out[f"head:{key[1]}"] = {"n": len(idx), "mean_error": mean_error(s, o), "rmse": rmse(s, o),
194
+ "rmse_over_range": rmse(s, o) / rng if rng > 0 else float("nan")}
195
+ return out
196
+
197
+
198
+ def run_forward_model(config: str | Path, parameters: Path | None, observations: Path, out: Path,
199
+ run: str = "calibration") -> Path:
200
+ """The forward model PEST++ calls: apply a parameter set, run the example, and write the simulated value at
201
+ every observation (name, simulated) to ``out``."""
202
+ from . import config as C
203
+ from .driver import Basin
204
+
205
+ cfg = C.load(config)
206
+ obs = read_observations(observations)
207
+ params = read_parameters(parameters) if parameters else []
208
+ b = Basin(cfg, run=run, prepare=lambda model, folder: apply_parameters(model, folder, params))
209
+ try:
210
+ b.run()
211
+ sim = simulated(b.series, b.dates(), obs)
212
+ finally:
213
+ b.finalize()
214
+ with open(out, "w", newline="") as f:
215
+ w = csv.writer(f)
216
+ w.writerow(["name", "simulated"])
217
+ for o, s in zip(obs, sim):
218
+ w.writerow([o.name, repr(float(s))])
219
+ return out
220
+
221
+
222
+ def write_pest_files(*args, **kwargs):
223
+ raise NotImplementedError("PEST++ IES control and template files (pyEMU) are planned; see docs/design.md")
224
+
225
+
226
+ def fetch_usgs(*args, **kwargs):
227
+ raise NotImplementedError("USGS streamflow and groundwater level downloads are planned; see docs/design.md")
basinledger/chain.py ADDED
@@ -0,0 +1,106 @@
1
+ """The utility chain: source -> distribution -> customers -> sewer -> wastewater plant -> outfall.
2
+
3
+ The chain is accounting, not a model. Each day it takes the volumes of the users on the system, adds the system's
4
+ existing customers, and places the water in the hydrologic models through routes:
5
+
6
+ produced = delivered / (1 - losses) withdrawn from the sources, by share
7
+ delivered = existing demand + users' withdrawals
8
+ losses = produced - delivered distribution losses (non-revenue water)
9
+ sewer = existing demand x sewered + users' returns + inflow and infiltration
10
+ outfall = sewer the plant's discharge, to its route
11
+
12
+ Config (one ``[[utilities]]`` table per system):
13
+
14
+ name = "city"
15
+ sources = [{route = "well:pw2", share = 0.6}, {route = "swat:place_of_use:intake", share = 0.4}]
16
+ losses = 0.10 # fraction of production
17
+ demand = 0.003 # m3 s-1, existing customers (a constant for now; later a daily series)
18
+ sewered = 0.8 # share of existing demand that reaches the sewer
19
+ infiltration = 0.0002 # m3 s-1 of inflow and infiltration into the sewer
20
+ outfall = "swat:point_source:wwtp"
21
+
22
+ Ledger accounts, all in m3 per day, under ``utility:<name>:``. They pair as: produced = delivered + losses;
23
+ produced = the sum of the source shares; sewer = outfall. Each route then carries its own pair of requested and
24
+ delivered.
25
+
26
+ Not yet: losses returned to the aquifer as recharge, inflow and infiltration taken from the aquifer, a daily demand
27
+ series, and the plant's own consumption and loads.
28
+ """
29
+
30
+ from __future__ import annotations
31
+
32
+ from dataclasses import dataclass, field
33
+
34
+ from .units import DAY
35
+
36
+
37
+ @dataclass
38
+ class Source:
39
+ route: str
40
+ share: float
41
+
42
+
43
+ @dataclass
44
+ class UtilityChain:
45
+ name: str
46
+ sources: list[Source]
47
+ outfall: str
48
+ losses: float = 0.0
49
+ demand: float = 0.0
50
+ sewered: float = 1.0
51
+ infiltration: float = 0.0
52
+ _withdrawals: float = field(default=0.0, repr=False)
53
+ _returns: float = field(default=0.0, repr=False)
54
+
55
+ def __post_init__(self):
56
+ total = sum(s.share for s in self.sources)
57
+ if not self.sources or abs(total - 1.0) > 1e-9:
58
+ raise ValueError(f"utility {self.name}: source shares sum to {total}, not 1")
59
+ if not 0 <= self.losses < 1 or not 0 <= self.sewered <= 1:
60
+ raise ValueError(f"utility {self.name}: losses must be in [0, 1) and sewered in [0, 1]")
61
+ for s in self.sources:
62
+ if s.route.startswith("utility:"):
63
+ raise ValueError(f"utility {self.name}: a source can't be another utility (not yet)")
64
+
65
+ @classmethod
66
+ def from_config(cls, c: dict) -> "UtilityChain":
67
+ return cls(name=c["name"], sources=[Source(s["route"], float(s["share"])) for s in c["sources"]],
68
+ outfall=c["outfall"], losses=float(c.get("losses", 0.0)), demand=float(c.get("demand", 0.0)),
69
+ sewered=float(c.get("sewered", 1.0)), infiltration=float(c.get("infiltration", 0.0)))
70
+
71
+ def add_user(self, withdrawal: float, returned: float):
72
+ """A user's day on this system: m3 withdrawn from the tap, m3 returned to the sewer."""
73
+ self._withdrawals += withdrawal
74
+ self._returns += returned
75
+
76
+ def close_day(self, ledger) -> dict[str, float]:
77
+ """The day's requests on each route (m3): sources positive withdrawals, the outfall a discharge."""
78
+ existing = self.demand * DAY
79
+ delivered = existing + self._withdrawals
80
+ produced = delivered / (1.0 - self.losses)
81
+ sewer = existing * self.sewered + self._returns + self.infiltration * DAY
82
+ k = f"utility:{self.name}:"
83
+ ledger.enter(k + "produced", produced)
84
+ ledger.enter(k + "delivered", delivered)
85
+ ledger.enter(k + "losses", produced - delivered)
86
+ ledger.enter(k + "existing_demand", existing)
87
+ ledger.enter(k + "users_withdrawal", self._withdrawals)
88
+ ledger.enter(k + "users_return", self._returns)
89
+ ledger.enter(k + "infiltration", self.infiltration * DAY)
90
+ ledger.enter(k + "consumed_by_existing", existing * (1 - self.sewered))
91
+ ledger.enter(k + "sewer", sewer)
92
+ requests = {}
93
+ for s in self.sources:
94
+ requests[s.route] = requests.get(s.route, 0.0) + produced * s.share
95
+ ledger.enter(k + f"source:{s.route}", produced * s.share)
96
+ requests[self.outfall] = requests.get(self.outfall, 0.0) + sewer
97
+ ledger.enter(k + "outfall", sewer)
98
+ self._withdrawals = self._returns = 0.0
99
+ return requests
100
+
101
+ def declare(self, ledger):
102
+ k = f"utility:{self.name}:"
103
+ ledger.pair(k + "produced", [k + "delivered", k + "losses"])
104
+ ledger.pair(k + "produced", [k + f"source:{s.route}" for s in self.sources])
105
+ ledger.pair(k + "delivered", [k + "existing_demand", k + "users_withdrawal"])
106
+ ledger.pair(k + "sewer", k + "outfall")
basinledger/config.py ADDED
@@ -0,0 +1,214 @@
1
+ """An example's configuration: one TOML file that wires real models for one basin and its water users.
2
+
3
+ Paths are relative to the config file. ``{swatplus_refdata}`` stands for the example projects that ship with SWAT+
4
+ (``refdata/`` in the source that ``swatplus-bmi build`` fetched), for the synthetic example.
5
+
6
+ name = "synthetic"
7
+ description = "..."
8
+
9
+ [run]
10
+ folder = "../../runs/synthetic" # each run copies the models here
11
+ days = 0 # 0: the SWAT+ project's whole period
12
+
13
+ [swatplus] # the surface model: SWAT+, stepped live through swatplus-bmi
14
+ project = "{swatplus_refdata}/Osu_1hru"
15
+ places_of_use = [{name = "intake", channel = "cha01"}] # added to the run's copy, set daily
16
+ point_sources = [{name = "wwtp", channel = "cha01"}]
17
+
18
+ # or [nextgen] # the surface model: a finished NGIAB run, replayed (coupling/nextgen.py)
19
+ # run = "{nextgen_michigan}/runs/scenarios/baseline"
20
+
21
+ [modflow6] # optional: the groundwater model
22
+ builder = "model:build" # model.py beside this file: build(workspace, days, start_date)
23
+ # workspace = "mf6" # or an existing simulation folder, copied into the run
24
+ model = "gwf"
25
+
26
+ [recharge] # SWAT+ aquifers -> MODFLOW 6 recharge (one-way)
27
+ package = "rch"
28
+ map = "all:aqu001" # every top-layer cell to one aquifer, or a CSV (see mapping.py)
29
+ area_rtol = 0.001 # optional: mapped area must match SWAT+'s aquifer area
30
+
31
+ [[wells]]
32
+ name = "pw1"
33
+ package = "wel"
34
+ cell = [0, 20, 30] # layer, row, col (0-based)
35
+
36
+ [exchange] # optional: two-way stream-aquifer exchange (MODFLOW 6 RIV <-> SWAT+)
37
+ package = "riv"
38
+ map = "all:cha01" # every RIV entry to one channel, or a CSV: object, layer, row, col
39
+
40
+ [[utilities]] # optional: a water system (see chain.py)
41
+ name = "city"
42
+ sources = [{route = "well:pw1", share = 1.0}]
43
+ outfall = "swat:point_source:wwtp"
44
+
45
+ [[rules]] # optional: feedback rules (see rules.py)
46
+ name = "low-flow cut"
47
+ kind = "curtailment"
48
+ user = "plant"
49
+ channel = "cha01"
50
+ threshold = 0.001
51
+
52
+ [outputs] # optional: daily series written with the ledger
53
+ channels = ["cha01"] # outflow, m3 s-1 (default: every channel)
54
+ heads = [{name = "pw1", cell = [0, 20, 28]}]
55
+
56
+ [calibration] # optional: the example's calibration record
57
+ record = "calibration/record.toml" # without one, the example is a screening run
58
+
59
+ [[users]]
60
+ name = "plant"
61
+ component = "basinledger.testing:SyntheticUser"
62
+ config = "user.toml"
63
+ withdrawal = "well:pw1" # source: a MODFLOW 6 well (later: a SWAT+ place of use, a utility)
64
+ return = "outfall:creek" # sink: an outfall account (later: a SWAT+ point source, a sewer)
65
+ # variables = {withdrawal = "...", gain = "..."} for components without the common names
66
+ # forcing = {source = "nextgen", catchment = "cat-652059"} hourly weather inputs (forcing.py)
67
+
68
+ Routes also include ``offsite:<name>``: water from or to somewhere outside the basin (a Great Lakes intake, a plant on
69
+ another river), entered in the ledger only.
70
+
71
+ ``{nextgen_michigan}`` stands for a checkout of github.com/northlinehydrology/nextgen-michigan (``$NEXTGEN_MICHIGAN``,
72
+ or ``~/repos/nextgen-michigan``).
73
+ """
74
+
75
+ from __future__ import annotations
76
+
77
+ import os
78
+ import tomllib
79
+ from dataclasses import dataclass, field
80
+ from pathlib import Path
81
+
82
+
83
+ def nextgen_michigan() -> Path:
84
+ return Path(os.environ.get("NEXTGEN_MICHIGAN", Path.home() / "repos" / "nextgen-michigan"))
85
+
86
+
87
+ def swatplus_refdata() -> Path:
88
+ if os.environ.get("SWATPLUS_REFDATA"):
89
+ return Path(os.environ["SWATPLUS_REFDATA"])
90
+ from swatplus_bmi.build import DEFAULT_TAG, home
91
+ return home() / "build" / DEFAULT_TAG / "swatplus" / "refdata"
92
+
93
+
94
+ @dataclass
95
+ class UserConfig:
96
+ name: str
97
+ component: str
98
+ config: Path | None
99
+ withdrawal: str
100
+ returns: str
101
+ variables: dict = field(default_factory=dict)
102
+ constituents: dict | None = None
103
+ forcing: dict | None = None
104
+
105
+
106
+ @dataclass
107
+ class WellConfig:
108
+ name: str
109
+ package: str
110
+ cell: tuple[int, int, int]
111
+
112
+
113
+ @dataclass
114
+ class ExampleConfig:
115
+ path: Path
116
+ name: str
117
+ description: str
118
+ run_folder: Path
119
+ days: int
120
+ swat_project: Path | None
121
+ mf6_builder: str | None
122
+ mf6_workspace: Path | None
123
+ mf6_model: str
124
+ recharge_package: str
125
+ recharge_map: str
126
+ area_rtol: float | None
127
+ wells: list[WellConfig]
128
+ users: list[UserConfig]
129
+ calibration_record: Path | None = None
130
+ places_of_use: list[tuple[str, str]] = field(default_factory=list)
131
+ point_sources: list[tuple[str, str]] = field(default_factory=list)
132
+ exchange: dict | None = None
133
+ utilities: list[dict] = field(default_factory=list)
134
+ rules: list[dict] = field(default_factory=list)
135
+ outputs: dict = field(default_factory=dict)
136
+ nextgen_run: Path | None = None
137
+
138
+ @property
139
+ def surface(self) -> str:
140
+ return "nextgen" if self.nextgen_run else "swatplus"
141
+
142
+ @property
143
+ def has_modflow6(self) -> bool:
144
+ return bool(self.mf6_builder or self.mf6_workspace)
145
+
146
+ @property
147
+ def models(self) -> list[str]:
148
+ """The hydrologic models a run of this example uses, for its calibration status."""
149
+ return [self.surface] + (["modflow6"] if self.has_modflow6 else [])
150
+
151
+ @property
152
+ def folder(self) -> Path:
153
+ return self.path.parent
154
+
155
+ @property
156
+ def calibration_status(self) -> str:
157
+ """'calibrated' or 'screening' (no record, or a record that doesn't say calibrated)."""
158
+ from .calibration import record_status
159
+ return record_status(self.calibration_record, self.models)
160
+
161
+
162
+ def _path(base: Path, p: str) -> Path:
163
+ if "{swatplus_refdata}" in p:
164
+ p = p.replace("{swatplus_refdata}", str(swatplus_refdata()))
165
+ p = p.replace("{nextgen_michigan}", str(nextgen_michigan()))
166
+ q = Path(p).expanduser()
167
+ return q if q.is_absolute() else (base / q).resolve()
168
+
169
+
170
+ def load(path: str | Path) -> ExampleConfig:
171
+ path = Path(path).resolve()
172
+ with open(path, "rb") as f:
173
+ c = tomllib.load(f)
174
+ base = path.parent
175
+ known = {"name", "description", "run", "swatplus", "modflow6", "recharge", "wells", "users", "calibration",
176
+ "exchange", "utilities", "rules", "outputs", "nextgen"}
177
+ if set(c) - known:
178
+ raise ValueError(f"{path}: unknown sections {sorted(set(c) - known)}")
179
+ mf = c.get("modflow6", {})
180
+ if "modflow6" in c and bool(mf.get("builder")) == bool(mf.get("workspace")):
181
+ raise ValueError(f"{path}: [modflow6] needs one of builder or workspace")
182
+ if ("swatplus" in c) == ("nextgen" in c):
183
+ raise ValueError(f"{path}: needs one surface model, [swatplus] or [nextgen]")
184
+ sw = c.get("swatplus", {})
185
+ rch = c.get("recharge", {})
186
+ users = []
187
+ for u in c.get("users", []):
188
+ users.append(UserConfig(
189
+ name=u["name"], component=u["component"],
190
+ config=_path(base, u["config"]) if u.get("config") else None,
191
+ withdrawal=u["withdrawal"], returns=u["return"],
192
+ variables=u.get("variables", {}), constituents=u.get("constituents"), forcing=u.get("forcing")))
193
+ return ExampleConfig(
194
+ path=path, name=c["name"], description=c.get("description", ""),
195
+ run_folder=_path(base, c.get("run", {}).get("folder", f"runs/{c['name']}")),
196
+ days=int(c.get("run", {}).get("days", 0)),
197
+ swat_project=_path(base, sw["project"]) if sw else None,
198
+ mf6_builder=mf.get("builder"),
199
+ mf6_workspace=_path(base, mf["workspace"]) if mf.get("workspace") else None,
200
+ mf6_model=mf.get("model", "gwf"),
201
+ recharge_package=rch.get("package", "rch"),
202
+ recharge_map=rch.get("map", ""),
203
+ area_rtol=rch.get("area_rtol"),
204
+ wells=[WellConfig(w["name"], w.get("package", "wel"), tuple(w["cell"])) for w in c.get("wells", [])],
205
+ users=users,
206
+ calibration_record=_path(base, c["calibration"]["record"]) if c.get("calibration", {}).get("record") else None,
207
+ places_of_use=[(x["name"], x["channel"]) for x in sw.get("places_of_use", [])],
208
+ point_sources=[(x["name"], x["channel"]) for x in sw.get("point_sources", [])],
209
+ exchange=c.get("exchange"),
210
+ utilities=c.get("utilities", []),
211
+ rules=c.get("rules", []),
212
+ outputs=c.get("outputs", {}),
213
+ nextgen_run=_path(base, c["nextgen"]["run"]) if "nextgen" in c else None,
214
+ )
File without changes