Add output budget and head to Gwf model

Still use the imod functions.

Author: deltamarnix

docs/examples/quickstart.py

@@ -1,6 +1,5 @@
 from pathlib import Path
 
-import imod.mf6
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -9,11 +8,11 @@
 from flopy4.mf6.simulation import Simulation
 from flopy4.mf6.tdis import Tdis
 
-ws = "./mymodel"
+ws = "./quickstart_data"
 name = "mymodel"
 tdis = Tdis()
 ims = Ims()
-sim = Simulation(name=name, tdis=tdis, solutions={"ims": ims})
+sim = Simulation(name=name, tdis=tdis, solutions={"ims": ims}, sim_ws=ws)
 dis = Dis(nrow=10, ncol=10)
 gwf = Gwf(parent=sim, name=name, save_flows=True, dis=dis)
 ic = Ic(parent=gwf)
@@ -43,17 +42,12 @@
 assert oc.data.save_head.sel(per=0) == "all"
 
 # PLOT
-# create budget reader
-bpth = Path("./quickstart_data/mymodel.bud")
-grbpth = Path("./quickstart_data/mymodel.dis.grb")
 
 # set specific discharge
-spdis = imod.mf6.open_cbc(bpth, grbpth, merge_to_dataset=True)
-
-# create head reader
-hpth = Path("./quickstart_data/mymodel.hds")
-heads = imod.mf6.open_hds(hpth, grbpth)
+spdis = gwf.output.budget
+heads = gwf.output.head
 sq = heads.squeeze()
+
 fig, ax = plt.subplots()
 ax.tick_params()
 ax.set_xticks(np.arange(0, 11, 2), minor=False)

flopy4/mf6/gwf/__init__.py

@@ -1,6 +1,9 @@
 from pathlib import Path
 from typing import Optional
 
+import attrs
+import imod
+import xarray as xr
 from attrs import define
 from xattree import field, xattree
 
@@ -16,11 +19,33 @@
 
 @xattree
 class Gwf(Model):
+    @define
+    class Output:
+        parent: "Gwf" = field(repr=False)
+
+        @property
+        def head(self) -> xr.DataArray:
+            return imod.mf6.open_hds(
+                Path("quickstart_data", f"{self.parent.name}.hds"),
+                Path("quickstart_data", f"{self.parent.name}.dis.grb"),
+            )
+
+        @property
+        def budget(self):
+            return imod.mf6.open_cbc(
+                Path("./quickstart_data/mymodel.bud"),
+                Path("./quickstart_data/mymodel.dis.grb"),
+                merge_to_dataset=True,
+            )
+
     dis: Dis = field()
     ic: Ic = field()
     oc: Oc = field()
     npf: Npf = field()
     chd: list[Chd] = field()
+    output: Output = attrs.field(
+        default=attrs.Factory(lambda self: Gwf.Output(self), takes_self=True)
+    )
 
     @define
     class NewtonOptions:

flopy4/mf6/simulation.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 from xattree import field, xattree
 
 from flopy4.mf6.component import Component
@@ -9,6 +11,7 @@
 
 @xattree
 class Simulation(Component):
+    sim_ws: Path = field()
     models: dict[str, Model] = field()
     exchanges: dict[str, Exchange] = field()
     solutions: dict[str, Solution] = field()