Adopt imod code (#125)

* Add output budget and head to Gwf model

Still use the imod functions.

* Copy imod code for reading heads lazily

* Copy imod code for reading info.

Make use of StructuredGrid instead of imod grid info.

* update envs, default sim_ws to none

* appease mypy

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 = Path("./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
 
@@ -10,17 +13,40 @@
 from flopy4.mf6.gwf.npf import Npf
 from flopy4.mf6.gwf.oc import Oc
 from flopy4.mf6.model import Model
+from flopy4.mf6.utils import open_hds
 
 __all__ = ["Gwf", "Chd", "Dis", "Ic", "Npf", "Oc"]
 
 
 @xattree
 class Gwf(Model):
+    @define
+    class Output:
+        parent: "Gwf" = attrs.field(repr=False)
+
+        @property
+        def head(self) -> xr.DataArray:
+            return open_hds(
+                self.parent.parent.sim_ws / f"{self.parent.name}.hds",  # type: ignore
+                self.parent.parent.sim_ws / f"{self.parent.name}.dis.grb",  # type: ignore
+            )
+
+        @property
+        def budget(self):
+            return imod.mf6.open_cbc(
+                self.parent.parent.sim_ws / "mymodel.bud",
+                self.parent.parent.sim_ws / "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
@@ -13,3 +15,4 @@ class Simulation(Component):
     exchanges: dict[str, Exchange] = field()
     solutions: dict[str, Solution] = field()
     tdis: Tdis = field()
+    sim_ws: Path = field(default=None)

flopy4/mf6/utils/__init__.py

@@ -0,0 +1,3 @@
+from .heads_reader import open_hds
+
+__all__ = ["open_hds"]

flopy4/mf6/utils/heads_reader.py

@@ -0,0 +1,206 @@
+import collections
+import os
+import struct
+from pathlib import Path
+from typing import Any
+
+import dask
+import numpy as np
+import pandas as pd
+import xarray as xr
+from flopy.discretization import StructuredGrid
+
+
+def open_hds(
+    hds_path: Path,
+    grb_path: Path,
+    dry_nan: bool = False,
+    simulation_start_time: np.datetime64 | None = None,
+    time_unit: str | None = "d",
+) -> xr.DataArray:
+    """
+    Open modflow6 heads (.hds) file.
+
+    The data is lazily read per timestep and automatically converted into
+    (dense) xr.DataArrays for DIS.
+    The conversion is done via the information stored in the Binary Grid file
+    (GRB).
+
+
+    Parameters
+    ----------
+    hds_path: pathlib.Path
+    grb_path: pathlib.Path
+    dry_nan: bool, default value: False.
+        Whether to convert dry values to NaN.
+    simulation_start_time : Optional datetime
+        The time and date corresponding to the beginning of the simulation.
+        Use this to convert the time coordinates of the output array to
+        calendar time/dates.
+        Time_unit must also be present if this argument is present.
+    time_unit: Optional str
+        The time unit MF6 is working in, in string representation.
+        Only used if simulation_start_time was provided.
+        Admissible values are:
+        ns -> nanosecond
+        ms -> microsecond
+        s -> second
+        m -> minute
+        h -> hour
+        d -> day
+        w -> week
+        Units "month" or "year" are not supported,
+        as they do not represent unambiguous timedelta values durations.
+
+    Returns
+    -------
+    head: xr.DataArray
+    """
+    grid = StructuredGrid.from_binary_grid_file(grb_path)
+    return _open_hds_dis(
+        hds_path, grid, dry_nan, simulation_start_time, time_unit
+    )
+
+
+def _open_hds_dis(
+    path: Path,
+    grid: StructuredGrid,
+    dry_nan: bool,
+    simulation_start_time: np.datetime64 | None = None,
+    time_unit: str | None = "d",
+) -> xr.DataArray:
+    nlayer, nrow, ncol = (
+        grid.nlay,
+        grid.nrow,
+        grid.ncol,
+    )
+    filesize = os.path.getsize(path)
+    ntime = filesize // (nlayer * (52 + (nrow * ncol * 8)))
+    coords = get_coords(grid)
+    coords["time"] = read_times(path, ntime, nlayer, nrow, ncol)
+
+    dask_list = []
+    # loop over times and add delayed arrays
+    for i in range(ntime):
+        # TODO verify dimension order
+        pos = i * (nlayer * (52 + nrow * ncol * 8))
+        a = dask.delayed(read_hds_timestep)(
+            path, nlayer, nrow, ncol, dry_nan, pos
+        )
+        x = dask.array.from_delayed(
+            a, shape=(nlayer, nrow, ncol), dtype=np.float64
+        )
+        dask_list.append(x)
+
+    daskarr = dask.array.stack(dask_list, axis=0)
+    data_array = xr.DataArray(
+        daskarr, coords, ("time", "layer", "y", "x"), name="head"
+    )
+    if simulation_start_time is not None:
+        data_array = assign_datetime_coords(
+            data_array, simulation_start_time, time_unit
+        )
+    return data_array
+
+
+def get_coords(grid: StructuredGrid) -> dict[str, Any]:
+    # unpack tuples
+    xmin, xmax, ymin, ymax = grid.extent
+    dx, dy = (grid.delr, -grid.delc)
+    coords: collections.OrderedDict[str, Any] = collections.OrderedDict()
+    # from cell size to x and y coordinates
+    if isinstance(dx, (int, float, np.int_)):  # equidistant
+        coords["x"] = np.arange(xmin + dx / 2.0, xmax, dx)
+        coords["y"] = np.arange(ymax + dy / 2.0, ymin, dy)
+        coords["dx"] = np.array(float(dx))
+        coords["dy"] = np.array(float(dy))
+    else:  # nonequidistant
+        # even though IDF may store them as float32,
+        # we always convert them to float64
+        dx = dx.astype(np.float64)
+        dy = dy.astype(np.float64)
+        coords["x"] = xmin + np.cumsum(dx) - 0.5 * dx
+        coords["y"] = ymax + np.cumsum(dy) - 0.5 * dy
+        if np.allclose(dx, dx[0]) and np.allclose(dy, dy[0]):
+            coords["dx"] = np.array(float(dx[0]))
+            coords["dy"] = np.array(float(dy[0]))
+        else:
+            coords["dx"] = ("x", dx)
+            coords["dy"] = ("y", dy)
+    coords["layer"] = np.arange(1, grid.nlay + 1)
+    return coords
+
+
+def read_times(
+    path: Path, ntime: int, nlayer: int, nrow: int, ncol: int
+) -> np.ndarray:
+    """
+    Reads all total simulation times.
+    """
+    times = np.empty(ntime, dtype=np.float64)
+
+    # Compute how much to skip to the next timestamp
+    start_of_header = 16
+    rest_of_header = 28
+    data_single_layer = nrow * ncol * 8
+    header = 52
+    nskip = (
+        rest_of_header
+        + data_single_layer
+        + (nlayer - 1) * (header + data_single_layer)
+        + start_of_header
+    )
+
+    with open(path, "rb") as f:
+        f.seek(start_of_header)
+        for i in range(ntime):
+            times[i] = struct.unpack("d", f.read(8))[
+                0
+            ]  # total simulation time
+            f.seek(nskip, 1)
+    return times
+
+
+def read_hds_timestep(
+    path: Path, nlayer: int, nrow: int, ncol: int, dry_nan: bool, pos: int
+) -> np.ndarray:
+    """
+    Reads all values of one timestep.
+    """
+    ncell_per_layer = nrow * ncol
+    with open(path, "rb") as f:
+        f.seek(pos)
+        a1d = np.empty(nlayer * nrow * ncol, dtype=np.float64)
+        for k in range(nlayer):
+            f.seek(52, 1)  # skip kstp, kper, pertime
+            a1d[k * ncell_per_layer : (k + 1) * ncell_per_layer] = np.fromfile(
+                f, np.float64, nrow * ncol
+            )
+
+    a3d = a1d.reshape((nlayer, nrow, ncol))
+    return _to_nan(a3d, dry_nan)
+
+
+def assign_datetime_coords(
+    da: xr.DataArray,
+    simulation_start_time: np.datetime64,
+    time_unit: str | None = "d",
+) -> xr.DataArray:
+    if "time" not in da.coords:
+        raise ValueError(
+            "cannot convert time column,"
+            " because a time column could not be found"
+        )
+
+    time = pd.Timestamp(simulation_start_time) + pd.to_timedelta(
+        da["time"], unit=time_unit
+    )
+    return da.assign_coords(time=time)
+
+
+def _to_nan(a: np.ndarray, dry_nan: bool) -> np.ndarray:
+    # TODO: this could really use a docstring?
+    a[a == 1e30] = np.nan
+    if dry_nan:
+        a[a == -1e30] = np.nan
+    return a