Issue #1244 reading dependent variable output

docs/api/changelog.rst

@@ -6,6 +6,13 @@ All notable changes to this project will be documented in this file.
 The format is based on `Keep a Changelog`_, and this project adheres to
 `Semantic Versioning`_.
 
+Unreleased
+------
+
+Added
+~~~~~
+- Added function :func:`imod.mf6.open_dvs` to read dependent variable output files like the water content file of :class:`imod.mf6.UnsaturatedZoneFlow`
+
 0.17.2
 ------
 

imod/mf6/__init__.py

@@ -39,7 +39,14 @@
 from imod.mf6.mst import MobileStorageTransfer
 from imod.mf6.npf import NodePropertyFlow
 from imod.mf6.oc import OutputControl
-from imod.mf6.out import open_cbc, open_conc, open_hds, read_cbc_headers, read_grb
+from imod.mf6.out import (
+    open_cbc,
+    open_conc,
+    open_hds,
+    read_cbc_headers,
+    read_grb,
+    open_dvs,
+)
 from imod.mf6.rch import Recharge
 from imod.mf6.riv import River
 from imod.mf6.simulation import Modflow6Simulation

imod/mf6/out/__init__.py

@@ -69,6 +69,12 @@ def open_cbc(
     "disu": disu.open_cbc,
 }
 
+_OPEN_DVS = {
+    "dis": dis.open_dvs,
+    "disv": disv.open_dvs,
+    "disu": disu.open_dvs,
+}
+
 
 def _get_function(d: Dict[str, Callable], key: str) -> Callable:
     try:
@@ -153,6 +159,55 @@ def open_hds(
     return _open(hds_path, grb_content, dry_nan, simulation_start_time, time_unit)
 
 
+def open_dvs(
+    dvs_path: FilePath,
+    grb_path: FilePath,
+    indices: np.ndarray,
+    simulation_start_time: Optional[np.datetime64] = None,
+    time_unit: Optional[str] = "d",
+) -> GridDataArray:
+    """
+    Open modflow6 dependent variable output for complex packages (like the watercontent file for the UZF-package).
+
+    The data is lazily read per timestep and automatically converted into
+    (dense) xr.DataArrays or xu.UgridDataArrays, for DIS and DISV respectively.
+    The conversion is done via the information stored in the Binary Grid file
+    (GRB).
+
+
+    Parameters
+    ----------
+    dsv_path: Union[str, pathlib.Path]
+    grb_path: Union[str, pathlib.Path]
+    indices: np.ndarray
+        The indices to map dvs variable to model nodes
+    simulation_start_time : Optional datetime
+        The time and date correpsonding 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
+    -------
+    dvs : Union[xr.DataArray, xu.UgridDataArray]
+    """
+    grb_content = read_grb(grb_path)
+    distype = grb_content["distype"]
+    _open = _get_function(_OPEN_DVS, distype)
+    return _open(dvs_path, grb_content, indices, simulation_start_time, time_unit)
+
+
 def open_conc(
     ucn_path: FilePath,
     grb_path: FilePath,

imod/mf6/out/common.py

@@ -2,6 +2,7 @@
 from typing import Any, BinaryIO, Dict, List, Union
 
 import numpy as np
+import struct
 
 # Type annotations
 IntArray = np.ndarray
@@ -29,3 +30,33 @@ def get_first_header_advanced_package(
         if "flow-ja-face" not in key and "gwf_" in key:
             return header_list[0]
     return None
+
+
+def read_name_dvs(path: FilePath) -> str:
+    """
+    Reads variable name from first header in dependent variable file.
+    """
+    with open(path, "rb") as f:
+        f.seek(24)
+        name = struct.unpack("16s", f.read(16))[0]
+    return name.decode().strip()
+
+
+def read_times_dvs(path: FilePath, ntime: int, indices: np.ndarray) -> FloatArray:
+    """
+    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 = indices.size * 8
+    nskip = rest_of_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

imod/mf6/out/dis.py

@@ -17,6 +17,8 @@
     IntArray,
     _to_nan,
     get_first_header_advanced_package,
+    read_name_dvs,
+    read_times_dvs,
 )
 
 
@@ -121,6 +123,20 @@ def read_hds_timestep(
     return _to_nan(a3d, dry_nan)
 
 
+def read_dvs_timestep(
+    path: FilePath, nlayer: int, nrow: int, ncol: int, pos: int, indices: np.ndarray
+) -> FloatArray:
+    """
+    Reads all values of one timestep.
+    """
+    with open(path, "rb") as f:
+        f.seek(pos)
+        a1d = np.full(nlayer * nrow * ncol, dtype=np.float64, fill_value=np.nan)
+        f.seek(52, 1)  # skip kstp, kper, pertime
+        a1d[indices] = np.fromfile(f, np.float64, indices.size)
+    return a1d.reshape((nlayer, nrow, ncol))
+
+
 def open_hds(
     path: FilePath,
     grid_info: Dict[str, Any],
@@ -155,6 +171,42 @@ def open_hds(
     return data_array
 
 
+def open_dvs(
+    path: FilePath,
+    grid_info: Dict[str, Any],
+    indices: np.ndarray,
+    simulation_start_time: Optional[np.datetime64] = None,
+    time_unit: Optional[str] = "d",
+) -> xr.DataArray:
+    nlayer, nrow, ncol = grid_info["nlayer"], grid_info["nrow"], grid_info["ncol"]
+    filesize = os.path.getsize(path)
+    ntime = filesize // (52 + (indices.size * 8))
+    times = read_times_dvs(path, ntime, indices)
+    dv_name = read_name_dvs(path)
+
+    coords = grid_info["coords"]
+    coords["time"] = times
+
+    dask_list = []
+    # loop over times and add delayed arrays
+    for i in range(ntime):
+        # TODO verify dimension order
+        pos = i * (52 + indices.size * 8)
+        a = dask.delayed(read_dvs_timestep)(path, nlayer, nrow, ncol, pos, indices)
+        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=dv_name
+    )
+    if simulation_start_time is not None:
+        data_array = assign_datetime_coords(
+            data_array, simulation_start_time, time_unit
+        )
+    return data_array
+
+
 def open_imeth1_budgets(
     cbc_path: FilePath, grb_content: dict, header_list: List[cbc.Imeth1Header]
 ) -> xr.DataArray:

imod/mf6/out/disu.py

@@ -77,6 +77,16 @@ def open_hds(
     raise NotImplementedError
 
 
+def open_dvs(
+    path: FilePath,
+    grid_info: Dict[str, Any],
+    indices: np.ndarray,
+    simulation_start_time: Optional[np.datetime64] = None,
+    time_unit: Optional[str] = "d",
+) -> xr.DataArray:
+    raise NotImplementedError
+
+
 def open_imeth1_budgets(
     cbc_path: FilePath, grb_content: dict, header_list: List[cbc.Imeth1Header]
 ) -> xr.DataArray:

imod/mf6/out/disv.py

@@ -18,6 +18,8 @@
     IntArray,
     _to_nan,
     get_first_header_advanced_package,
+    read_name_dvs,
+    read_times_dvs,
 )
 
 
@@ -147,6 +149,20 @@ def read_hds_timestep(
     return _to_nan(a2d, dry_nan)
 
 
+def read_dvs_timestep(
+    path: FilePath, nlayer: int, ncells_per_layer: int, pos: int, indices: np.ndarray
+) -> FloatArray:
+    """
+    Reads all values of one timestep.
+    """
+    with open(path, "rb") as f:
+        f.seek(pos)
+        a1d = np.full(nlayer * ncells_per_layer, dtype=np.float64, fill_value=np.nan)
+        f.seek(52, 1)  # skip kstp, kper, pertime
+        a1d[indices] = np.fromfile(f, np.float64, indices.size)
+    return a1d.reshape((nlayer, ncells_per_layer))
+
+
 def open_hds(
     path: FilePath,
     grid_info: Dict[str, Any],
@@ -185,6 +201,45 @@ def open_hds(
     return xu.UgridDataArray(da, grid)
 
 
+def open_dvs(
+    path: FilePath,
+    grid_info: Dict[str, Any],
+    indices: np.ndarray,
+    simulation_start_time: Optional[np.datetime64] = None,
+    time_unit: Optional[str] = "d",
+) -> xu.UgridDataArray:
+    grid = grid_info["grid"]
+    nlayer, ncells_per_layer = grid_info["nlayer"], grid_info["ncells_per_layer"]
+    filesize = os.path.getsize(path)
+    ntime = filesize // (52 + (indices.size * 8))
+    times = read_times_dvs(path, ntime, indices)
+    dv_name = read_name_dvs(path)
+
+    coords = grid_info["coords"]
+    coords["time"] = times
+
+    dask_list = []
+    # loop over times and add delayed arrays
+    for i in range(ntime):
+        # TODO verify dimension order
+        pos = i * (52 + indices.size * 8)
+        a = dask.delayed(read_dvs_timestep)(
+            path, nlayer, ncells_per_layer, pos, indices
+        )
+        x = dask.array.from_delayed(
+            a, shape=(nlayer, ncells_per_layer), dtype=np.float64
+        )
+        dask_list.append(x)
+
+    daskarr = dask.array.stack(dask_list, axis=0)
+    da = xr.DataArray(
+        daskarr, coords, ("time", "layer", grid.face_dimension), name=dv_name
+    )
+    if simulation_start_time is not None:
+        da = assign_datetime_coords(da, simulation_start_time, time_unit)
+    return xu.UgridDataArray(da, grid)
+
+
 def open_imeth1_budgets(
     cbc_path: FilePath, grb_content: dict, header_list: List[cbc.Imeth1Header]
 ) -> xu.UgridDataArray:

imod/mf6/uzf.py

@@ -1,5 +1,6 @@
 import numpy as np
 import xarray as xr
+import xugrid as xu
 
 from imod.logging import init_log_decorator
 from imod.mf6.boundary_condition import AdvancedBoundaryCondition, BoundaryCondition
@@ -97,6 +98,8 @@ class UnsaturatedZoneFlow(AdvancedBoundaryCondition):
         path to output cbc-file for UZF budgets
     budgetcsv_fileout: ({"str"}, optional)
         path to output csv-file for summed budgets
+    water_content_file: ({"str"}, optional)
+        path to output file for unsaturated zone water content
     observations: [Not yet supported.]
         Default is None.
     water_mover: [Not yet supported.]
@@ -233,6 +236,7 @@ def __init__(
         save_flows=False,
         budget_fileout=None,
         budgetcsv_fileout=None,
+        water_content_file=None,
         observations=None,
         water_mover=None,
         timeseries=None,
@@ -271,6 +275,7 @@ def __init__(
             "save_flows": save_flows,
             "budget_fileout": budget_fileout,
             "budgetcsv_fileout": budgetcsv_fileout,
+            "water_content_file": water_content_file,
             "observations": observations,
             "water_mover": water_mover,
             "timeseries": timeseries,
@@ -301,7 +306,10 @@ def fill_stress_perioddata(self):
         for var in self._period_data:
             if self.dataset[var].size == 1:  # Prevent loading large arrays in memory
                 if self.dataset[var].values[()] is None:
-                    self.dataset[var] = xr.full_like(self["infiltration_rate"], 0.0)
+                    if isinstance(self["infiltration_rate"], xu.UgridDataArray):
+                        self.dataset[var] = xu.full_like(self["infiltration_rate"], 0)
+                    else:
+                        self.dataset[var] = xr.full_like(self["infiltration_rate"], 0.0)
                 else:
                     raise ValueError("{} cannot be a scalar".format(var))
 
@@ -379,9 +387,12 @@ def _package_data_to_sparse(self):
 
         field_spec = self._get_field_spec_from_dtype(recarr)
         field_names = [i[0] for i in field_spec]
-        index_spec = [("iuzno", np.int32)] + field_spec[:3]
+        n = 3
+        if isinstance(self.dataset, xu.UgridDataset):
+            n = 2
+        index_spec = [("iuzno", np.int32)] + field_spec[:n]
         field_spec = (
-            [("landflag", np.int32)] + [("ivertcon", np.int32)] + field_spec[3:]
+            [("landflag", np.int32)] + [("ivertcon", np.int32)] + field_spec[n:]
         )
         sparse_dtype = np.dtype(index_spec + field_spec)
 

imod/templates/mf6/gwf-uzf.j2

@@ -15,6 +15,8 @@ begin options
   budget fileout {{budget_fileout}}{% endif %}
 {%- if budgetcsv_fileout is defined %}
   budgetcsv fileout {{budgetcsv_fileout}}{% endif %}
+{%- if water_content_file is defined %}
+  water_content fileout {{water_content_file}}{% endif %}
 {%- if ts_filerecord is defined %}
   ts6 filein {{ts6_filename}}{% endif %}
 {%- if obs_filerecord is defined %}