Copy soil layer setting csv to input folder + add a BMI MODFLOW coupling notebook

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,30 @@
 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, 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] = {var.name: var.dtype for var in VARIABLES}
 
-VARNAME_DTYPE: dict[str, str] = {
-    "respiration": "float64",
-    "soil_temperature": "float64",
-}
+VARNAME_GRID: dict[str, int] = {var.name: var.grid for var in VARIABLES}
 
-VARNAME_GRID: dict[str, int] = {
-    "respiration": 0,
-    "soil_temperature": 1,
-}
+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 +58,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 +109,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
 
 

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"],
+    ),
+)

PyStemmusScope/config_io.py

@@ -124,9 +124,11 @@ def _copy_data(input_dir: Path, config: dict) -> None:
             str(config[folder]), str(input_dir / folder), dirs_exist_ok=True
         )
 
-    # copy input_data.xlsx
+    # copy input_data.xlsx and input_soilLayThick.csv
     shutil.copy(str(config["input_data"]), str(input_dir))
-
+    soilLay_file = Path(str(config["input_data"]).replace("input_data.xlsx", "input_soilLayThick.csv"))
+    if soilLay_file.exists():
+        shutil.copy(soilLay_file, str(input_dir))
 
 def _update_config_file(
     input_dir: Path,