Implement groundwater coupling changes in BMI (#101)

* Implement groundwater coupling changes in BMI

* Implement new groundwater coupling variables

* Programatically get/set values

* Remove outdated test

* Pin dependencies to working versions

* Fix issues related to variables containing all timesteps

* Add notebook demonstrating the groundwater coupling through BMI

* Load docker image from test config file instead of hardcoded

* Add checks for set_value to ensure arrays are right size

* Add tests for new set_value input checks

* stemmus_scope:1.6.0 image is now available

* Skip notebooks in linting for now

Author: BSchilperoort

Dockerfile

@@ -1,4 +1,4 @@
-FROM ghcr.io/ecoextreml/stemmus_scope:1.5.0
+FROM ghcr.io/ecoextreml/stemmus_scope:1.6.0
 
 LABEL maintainer="Bart Schilperoort <b.schilperoort@esciencecenter.nl>"
 LABEL org.opencontainers.image.source = "https://github.com/EcoExtreML/STEMMUS_SCOPE_Processing"

PyStemmusScope/bmi/docker_process.py

@@ -83,7 +83,7 @@ class StemmusScopeDocker:
     """Communicate with a STEMMUS_SCOPE Docker container."""
 
     # Default image, can be overridden with config:
-    compatible_tags = ("1.5.0",)
+    compatible_tags = ("1.6.0",)
 
     _process_ready_phrase = b"Select BMI mode:"
     _process_finalized_phrase = b"Finished clean up."

PyStemmusScope/bmi/implementation.py

@@ -8,31 +8,31 @@
 import numpy as np
 from bmipy.bmi import Bmi
 from PyStemmusScope.bmi.utils import InapplicableBmiMethods
+from PyStemmusScope.bmi.utils import nested_set
+from PyStemmusScope.bmi.variable_reference import VARIABLES
+from PyStemmusScope.bmi.variable_reference import BmiVariable
 from PyStemmusScope.config_io import read_config
 
 
-MODEL_INPUT_VARNAMES: tuple[str, ...] = ("soil_temperature",)
-
-MODEL_OUTPUT_VARNAMES: tuple[str, ...] = (
-    "soil_temperature",
-    "respiration",
+MODEL_INPUT_VARNAMES: tuple[str, ...] = tuple(
+    var.name for var in VARIABLES if var.input
 )
 
-MODEL_VARNAMES: tuple[str, ...] = tuple(
-    set(MODEL_INPUT_VARNAMES + MODEL_OUTPUT_VARNAMES)
+MODEL_OUTPUT_VARNAMES: tuple[str, ...] = tuple(
+    var.name for var in VARIABLES if var.output
 )
 
-VARNAME_UNITS: dict[str, str] = {"respiration": "unknown", "soil_temperature": "degC"}
+MODEL_VARS: dict[str, BmiVariable] = {var.name: var for var in VARIABLES}
+
+MODEL_VARNAMES: tuple[str, ...] = tuple(var.name for var in VARIABLES)
+
+VARNAME_UNITS: dict[str, str] = {var.name: var.units for var in VARIABLES}
 
-VARNAME_DTYPE: dict[str, str] = {
-    "respiration": "float64",
-    "soil_temperature": "float64",
-}
+VARNAME_DTYPE: dict[str, str] = {var.name: var.dtype for var in VARIABLES}
 
-VARNAME_GRID: dict[str, int] = {
-    "respiration": 0,
-    "soil_temperature": 1,
-}
+VARNAME_GRID: dict[str, int] = {var.name: var.grid for var in VARIABLES}
+
+VARNAME_LOC: dict[str, list[str]] = {var.name: var.keys for var in VARIABLES}
 
 NO_STATE_MSG = (
     "The model state is not available. Please run `.update()` before requesting "
@@ -59,23 +59,32 @@ def load_state(config: dict) -> h5py.File:
     return h5py.File(matfile, mode="a")
 
 
-def get_variable(state: h5py.File, varname: str) -> np.ndarray:
+def get_variable(
+    state: h5py.File, varname: str
+) -> np.ndarray:  # noqa: PLR0911 PLR0912 C901
     """Get a variable from the model state.
 
     Args:
         state: STEMMUS_SCOPE model state
         varname: Variable name
     """
-    if varname == "respiration":
-        return state["fluxes"]["Resp"][0]
+    if varname not in MODEL_VARNAMES:
+        msg = "Unknown variable name"
+        raise ValueError(msg)
+
+    # deviating implemetation:
     elif varname == "soil_temperature":
         return state["TT"][0, :-1]
+
+    # default implementation:
+    _s = state
+    for _loc in VARNAME_LOC[varname]:
+        _s = _s.get(_loc)
+
+    if MODEL_VARS[varname].all_timesteps:
+        return _s[0].astype(VARNAME_DTYPE[varname])[[int(state["KT"][0])]]
     else:
-        if varname in MODEL_VARNAMES:
-            msg = "Varname is missing in get_variable! Contact devs."
-        else:
-            msg = "Unknown variable name"
-        raise ValueError(msg)
+        return _s[0].astype(VARNAME_DTYPE[varname])
 
 
 def set_variable(
@@ -101,16 +110,21 @@ def set_variable(
     else:
         vals = value
 
+    if varname in MODEL_OUTPUT_VARNAMES and varname not in MODEL_INPUT_VARNAMES:
+        msg = "This variable is a model output variable only. You cannot set it."
+        raise ValueError(msg)
+    elif varname not in MODEL_INPUT_VARNAMES:
+        msg = "Uknown variable name"
+        raise ValueError(msg)
+
+    # deviating implementations:
     if varname == "soil_temperature":
         state["TT"][0, :-1] = vals
+    elif varname == "groundwater_coupling_enabled":
+        state["GroundwaterSettings"]["GroundwaterCoupling"][0] = vals.astype("float")
+    # default:
     else:
-        if varname in MODEL_OUTPUT_VARNAMES and varname not in MODEL_INPUT_VARNAMES:
-            msg = "This variable is a model output variable only. You cannot set it."
-        elif varname in MODEL_VARNAMES:
-            msg = "Varname is missing in set_variable! Contact devs."
-        else:
-            msg = "Uknown variable name"
-        raise ValueError(msg)
+        nested_set(state, VARNAME_LOC[varname] + [0], vals)
     return state
 
 
@@ -401,6 +415,9 @@ def set_value(self, name: str, src: np.ndarray) -> None:
         """
         if self.state is None:
             raise ValueError(NO_STATE_MSG)
+        if src.size != self.get_grid_size(self.get_var_grid(name)):
+            msg = f"Size of `src` and variable '{name}' grid size are not equal!"
+            raise ValueError(msg)
         self.state = set_variable(self.state, name, src)
 
     def set_value_at_indices(
@@ -419,6 +436,9 @@ def set_value_at_indices(
         """
         if self.state is None:
             raise ValueError(NO_STATE_MSG)
+        if inds.size != src.size:
+            msg = "Sizes of `inds` and `src` are not equal!"
+            raise ValueError(msg)
         self.state = set_variable(self.state, name, src, inds)
 
     ### GRID INFO ###

PyStemmusScope/bmi/utils.py

@@ -1,4 +1,6 @@
 """Utilities for the STEMMUS_SCOPE Basic Model Interface."""
+from typing import Any
+from typing import Union
 import numpy as np
 
 
@@ -64,3 +66,18 @@ def get_grid_nodes_per_face(
     ) -> np.ndarray:
         """Get the number of nodes for each face."""
         raise NotImplementedError(INAPPLICABLE_GRID_METHOD_MSG)
+
+
+def nested_set(dic: dict, keys: Union[list, tuple], value: Any) -> None:
+    """Set a value in a nested dictionary programatically.
+
+    E.g.: dict[keys[0]][keys[1]] = value
+
+    Args:
+        dic: Dictionary to be modified.
+        keys: Iterable of keys that are used to find the right value.
+        value: The new value.
+    """
+    for key in keys[:-1]:
+        dic = dic.setdefault(key, {})
+    dic[keys[-1]] = value

PyStemmusScope/bmi/variable_reference.py

@@ -0,0 +1,154 @@
+"""Variable reference to inform the BMI implementation."""
+from dataclasses import dataclass
+
+
+@dataclass
+class BmiVariable:
+    """Holds all info to inform the BMI implementation."""
+
+    name: str
+    dtype: str
+    input: bool
+    output: bool
+    units: str
+    grid: int
+    keys: list[str]
+    all_timesteps: bool = False
+
+
+VARIABLES: tuple[BmiVariable, ...] = (
+    # atmospheric vars:
+    BmiVariable(
+        name="respiration",
+        dtype="float64",
+        input=False,
+        output=True,
+        units="cm s-1",
+        grid=0,
+        keys=["fluxes", "Resp"],
+    ),
+    BmiVariable(
+        name="evaporation_total",
+        dtype="float64",
+        input=False,
+        output=True,
+        units="cm s-1",
+        grid=0,
+        keys=["EVAP"],
+        all_timesteps=True,
+    ),
+    # soil vars:
+    BmiVariable(
+        name="soil_temperature",
+        dtype="float64",
+        input=True,
+        output=True,
+        units="degC",
+        grid=1,
+        keys=["TT"],
+    ),
+    BmiVariable(
+        name="soil_moisture",
+        dtype="float64",
+        input=True,
+        output=True,
+        units="m3 m-3",
+        grid=1,
+        keys=["SoilVariables", "Theta_U"],
+    ),
+    BmiVariable(
+        name="soil_root_water_uptake",
+        dtype="float64",
+        input=False,
+        output=True,
+        units="cm s-1",
+        grid=0,
+        keys=["RWUs"],
+    ),
+    # surface runoff
+    BmiVariable(
+        name="surface_runoff_total",
+        dtype="float64",
+        input=False,
+        output=True,
+        units="cm s-1",
+        grid=0,
+        keys=["RS"],
+    ),
+    BmiVariable(
+        name="surface_runoff_hortonian",
+        dtype="float64",
+        input=False,
+        output=True,
+        units="cm s-1",
+        grid=0,
+        keys=["ForcingData", "R_Dunn"],
+        all_timesteps=True,
+    ),
+    BmiVariable(
+        name="surface_runoff_dunnian",
+        dtype="float64",
+        input=False,
+        output=True,
+        units="cm s-1",
+        grid=0,
+        keys=["ForcingData", "R_Hort"],
+        all_timesteps=True,
+    ),
+    # groundwater vars (STEMMUS_SCOPE)
+    BmiVariable(
+        name="groundwater_root_water_uptake",
+        dtype="float64",
+        input=False,
+        output=True,
+        units="cm s-1",
+        grid=0,
+        keys=["RWUg"],
+    ),
+    BmiVariable(
+        name="groundwater_recharge",
+        dtype="float64",
+        input=False,
+        output=True,
+        units="cm s-1",
+        grid=0,
+        keys=["gwfluxes", "recharge"],
+    ),
+    # groundwater (coupling) vars
+    BmiVariable(
+        name="groundwater_coupling_enabled",
+        dtype="bool",
+        input=True,
+        output=False,
+        units="-",
+        grid=0,
+        keys=["GroundwaterSettings", "GroundwaterCoupling"],
+    ),
+    BmiVariable(
+        name="groundwater_head_bottom_layer",
+        dtype="float64",
+        input=True,
+        output=False,
+        units="cm",
+        grid=0,
+        keys=["GroundwaterSettings", "headBotmLayer"],
+    ),
+    BmiVariable(
+        name="groundwater_temperature",
+        dtype="float64",
+        input=True,
+        output=False,
+        units="degC",
+        grid=0,
+        keys=["GroundwaterSettings", "tempBotm"],
+    ),
+    BmiVariable(
+        name="groundwater_elevation_top_aquifer",
+        dtype="float64",
+        input=True,
+        output=False,
+        units="cm",
+        grid=0,
+        keys=["GroundwaterSettings", "topLevel"],
+    ),
+)

docs/bmi.md

@@ -19,10 +19,10 @@ To use the Docker image, use the `DockerImage` setting in the configuration file
 ```sh
 WorkDir=/home/username/tmp/stemmus_scope
 ...
-DockerImage=ghcr.io/ecoextreml/stemmus_scope:1.5.0
+DockerImage=ghcr.io/ecoextreml/stemmus_scope:1.6.0
 ```
 
-It is best to add the version tag here too (`:1.5.0`), this way the BMI will warn you if the version might be incompatible.
+It is best to add the version tag here too (`:1.6.0`), this way the BMI will warn you if the version might be incompatible.
 
 Note that the `docker` package for python is required here. Install this with `pip install PyStemmusScope[docker]`.
 Additionally, [Docker](https://docs.docker.com/get-docker/) itself has to be installed.