Add basic OBS package support

imod/mf6/obs_gwf.py

@@ -0,0 +1,145 @@
+from pathlib import Path
+from typing import Any
+import os
+
+from imod.mf6.package import Package
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+import xugrid as xu
+from imod.mf6.write_context import WriteContext
+
+
+class GroundwaterFlowObservations(Package):
+    """
+    Observation Utility for GroundwaterFlowModel.
+
+    Supports "head", "drawdown", and "flow-ja-face".
+
+    """
+    _pkg_id = "obs"
+    _template = Package._initialize_template(_pkg_id)
+    _init_schemata = {}
+    _write_schemata = {}
+
+    def __init__(
+        self,
+        obs_name: xr.DataArray,
+        obs_type: xr.DataArray,
+        obs_id: xr.DataArray,
+        obs_id2: xr.DataArray|None=None,
+        obs_file=None,
+        digits=None,
+        print_input=False,
+        binary=True,
+    ):
+        if obs_id2 is None:
+            obs_id2 = xr.full_like(obs_id, np.nan, dtype=float)
+
+        dict_dataset: dict[str, Any] = {
+            "obs_name": obs_name,
+            "obs_type": obs_type,
+            "obs_id": obs_id,
+            "obs_id2": obs_id2,
+            "obs_file": obs_file,
+            "digits": digits,
+            "print_input": print_input,
+            "binary": binary
+        }
+        super().__init__(dict_dataset)
+
+    def _get_output_filepath(self, directory: Path, pkgname: str):
+        binary = self.dataset["binary"].values[()]
+        if binary:
+            ext = "bsv"
+        else:
+            ext = "csv"
+
+        filepath = self.dataset["obs_file"].values[()]
+        if filepath is None:
+            filepath = directory / f"{pkgname}.{ext}"
+        else:
+            if not isinstance(filepath, str | Path):
+                raise ValueError(
+                    f"{filepath} should be of type str or Path. However it is of type {type(filepath)}"
+                )
+            filepath = Path(filepath)
+
+        if filepath.is_absolute():
+            path = filepath
+        else:
+            # Get path relative to the simulation name file.
+            sim_directory = directory.parent
+            path = Path(os.path.relpath(filepath, sim_directory))
+        return path
+
+    def _render(self, directory, pkgname, globaltimes, binary):
+        d: dict[str, Any] = {}
+        for varname in ("print_input", "digits", "binary"):
+            value = self.dataset[varname].values[()]
+            if self._valid(value):
+                d[varname] = value
+        d["obs_file"] = self._get_output_filepath(directory, pkgname)
+        return self._template.render(d)
+
+    def _obs_rows_dataframe(self) -> pd.DataFrame:
+        columns = [
+            self.dataset[["obs_name", "obs_type"]].to_dataframe(),
+            self.dataset["obs_id"].to_pandas(),
+            # astype("Int64") maps NaN to pd.NA, and to_csv will write pd.NA as an empty field
+            self.dataset["obs_id2"].to_pandas().astype("Int64")
+        ]
+        df = pd.concat(columns, axis=1)
+        # Indent the rows by two spaces.
+        # We cannot rely on spaces, as they will be quoted.
+        df.insert(0, "__indent0", "")
+        df.insert(0, "__indent1", "")
+        return df
+
+    def _append_obs_rows(self, filename: Path | str) -> None:
+        with open(filename, "a") as f:
+            self._obs_rows_dataframe().to_csv(
+                f, header=False, index=False, sep=" ", lineterminator="\n"
+            )
+            f.write("end continuous\n\n")
+        return
+
+    def _write_blockfile(self, pkgname, globaltimes, write_context: WriteContext):
+        super()._write_blockfile(pkgname, globaltimes, write_context)
+        filename = write_context.write_directory / f"{pkgname}.{self._pkg_id}"
+        self._append_obs_rows(filename)
+        return
+
+    @classmethod
+    def from_boolean_grid(cls, mask, obs_type: str, obs_file: str | Path = None, binary=True):
+        # TODO: ugly method name, mask as name also dubious
+        if isinstance(mask, xr.DataArray):
+            if not mask.dims == ("layer", "y", "x"):
+                raise ValueError()
+        elif isinstance(mask, xu.UgridDataArray):
+            # TODO:
+            if not mask.dims == ("layer", "mesh2d_nFace"):
+                raise ValueError()
+        else:
+            raise TypeError(f"Expected DataArray or UgridDataArray, received: {type(mask)}")
+
+        if not np.issubdtype(mask.dtype, np.bool_):
+            raise TypeError()
+
+        # NOTE: modflow indices are 1-based.
+        obs_id = xr.DataArray(
+            data=np.column_stack(np.nonzero(mask.to_numpy())) + 1,
+            dims=("observation", "dimension"),
+        )
+        n_obs, _ = obs_id.shape
+        return cls(
+            obs_name=xr.DataArray(data=np.full(n_obs, "0"), dims=("observation",)),
+            obs_type=xr.DataArray(data=np.full(n_obs, obs_type), dims=("observation",)),
+            obs_id=obs_id,
+            obs_id2=None,
+            obs_file=obs_file,
+            digits=None,
+            print_input=False,
+            binary=binary,
+        )

imod/mf6/out/obs.py

@@ -0,0 +1,184 @@
+import os
+import struct
+from typing import BinaryIO, Optional
+
+import dask
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+from imod.mf6.out.common import FilePath, _to_nan
+from imod.mf6.utilities.dataset import assign_datetime_coords
+
+LEN_RECORD1 = 100
+LEN_RECORD2 = 4
+TIMESTAMP_SIZE = 8
+
+
+def read_times(
+    f: BinaryIO, ntime: int, nobs: int
+) -> np.ndarray:
+    times = np.empty(ntime, dtype=np.float64)
+    for i in range(ntime):
+        times[i] = struct.unpack("d", f.read(8))[0]
+        f.seek(nobs * 8, 1)
+    return times
+
+
+def read_timestep(
+    path: FilePath, dry_nan: bool, nobs: int, pos: int,
+) -> np.ndarray:
+    with open(path, "rb") as f:
+        f.seek(pos)
+        a = np.fromfile(f, np.float64, nobs)
+    return _to_nan(a, dry_nan)
+
+
+def read_record1(f: BinaryIO) -> int:
+    recordtype = f.read(5).decode("utf-8").strip()
+    precision = f.read(6).decode("utf-8").strip()
+    lenobsname = f.read(4).decode("utf-8").strip()
+    blanks = f.read(85).decode("utf-8").strip()
+    if not recordtype == "cont":
+        raise ValueError(f'recordtype is not "cont", but: "{recordtype}"')
+    if not precision == "double":
+        raise ValueError(f'precision is not "double", but: "{recordtype}"')
+    if not blanks == "":
+        raise ValueError(f'Expected 85 blanks, but received: "{blanks}"')
+    return int(lenobsname.strip())
+
+
+def read_record2(f: BinaryIO) -> int:
+    return struct.unpack("i", f.read(4))[0]
+
+
+def read_record3(f: BinaryIO, len_obsname: int, nobs: int) -> np.ndarray:
+    rawnames = np.fromfile(
+        file=f,
+        dtype=f"<S{len_obsname}",
+        count=nobs,
+    )
+    return np.char.strip(rawnames).astype(str)
+
+
+def open_obs_bsv(
+    path: FilePath,
+    dry_nan: bool=False,
+    simulation_start_time: Optional[np.datetime64] = None,
+    time_unit: Optional[str] = "d"
+) -> xr.DataArray:
+    """
+    Open modflow6 observation package binary simulated values ("bsv") file.
+
+    The data is lazily read per timestep.
+
+    Parameters
+    ----------
+    path: Union[str, pathlib.Path]
+    dry_nan: bool, default value: False.
+        Whether to convert dry values to NaN.
+    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
+    -------
+    observations: xr.DataArray
+    """
+    filesize = os.path.getsize(path)
+    with open(path, "rb") as f:
+        len_obsname = read_record1(f)
+        nobs = read_record2(f)
+        obsnames = read_record3(f, len_obsname, nobs)
+        len_record3 = len_obsname * nobs
+        # For a timestep, the timestamp is stored and all observations, as doubles.
+        len_timestep = (nobs + 1) * 8
+        ntime = (filesize - (LEN_RECORD1 + LEN_RECORD2 + len_record3)) // len_timestep
+        times = read_times(f, ntime, nobs)
+
+    dask_list = []
+    initial_skip = LEN_RECORD1 + LEN_RECORD2 + len_record3 + TIMESTAMP_SIZE
+    for i in range(ntime):
+        pos = initial_skip + (i * len_timestep)
+        a = dask.delayed(read_timestep)(path, dry_nan, nobs, pos)
+        x = dask.array.from_delayed(a, shape=(nobs,), dtype=np.float64)
+        dask_list.append(x)
+
+    daskarr = dask.array.stack(dask_list, axis=0)
+    data_array = xr.DataArray(
+        data=daskarr,
+        dims=("time", "observation"),
+        coords={"time": times, "observation": obsnames}
+    )
+    if simulation_start_time is not None:
+        data_array = assign_datetime_coords(
+            data_array, simulation_start_time, time_unit
+        )
+    return data_array
+
+
+def read_obs_csv(
+    path: FilePath,
+    dry_nan: bool=False,
+    simulation_start_time: Optional[np.datetime64] = None,
+    time_unit: Optional[str] = "d"
+) -> xr.DataArray:
+    """
+    Read modflow6 observation package CSV file.
+
+    Unlike the ``open_obs_bsv`` function, data is directly read into memory.
+
+    Parameters
+    ----------
+    path: Union[str, pathlib.Path]
+    dry_nan: bool, default value: False.
+        Whether to convert dry values to NaN.
+    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
+    -------
+    observations: xr.DataArray
+    """
+
+    df = pd.read_csv(path, index_col=0)
+    times = df.index
+    obsnames = df.columns
+    data = _to_nan(df.to_numpy(), dry_nan)
+    data_array = xr.DataArray(
+        data=data,
+        dims=("time", "observation"),
+        coords={"time": times, "observation": obsnames}
+    )
+    if simulation_start_time is not None:
+        data_array = assign_datetime_coords(
+            data_array, simulation_start_time, time_unit
+        )
+    return data_array

imod/templates/mf6/gwf-obs.j2

@@ -0,0 +1,8 @@
+begin options
+{% if digits is defined %}  digits {{digits}}
+{% endif %}
+{%- if print_input is defined %}  print_input
+{% endif -%}
+end options
+
+begin continuous fileout {{obs_file}} {% if binary %}binary{% endif %}