Skip to content

Commit

Permalink
Fix issues related to variables containing all timesteps
Browse files Browse the repository at this point in the history
  • Loading branch information
BSchilperoort committed Jul 15, 2024
1 parent 18bfe6d commit c0a75a2
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 23 deletions.
2 changes: 1 addition & 1 deletion PyStemmusScope/bmi/docker_process.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
Expand Down
14 changes: 7 additions & 7 deletions PyStemmusScope/bmi/implementation.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
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 VARIABLES, BmiVariable
from PyStemmusScope.config_io import read_config


Expand All @@ -21,6 +21,8 @@
var.name for var in VARIABLES if var.output
)

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}
Expand All @@ -29,7 +31,7 @@

VARNAME_GRID: dict[str, int] = {var.name: var.grid for var in VARIABLES}

VARNAME_LOC: dict[str, list[str]] = {var.name: var.loc 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 "
Expand Down Expand Up @@ -78,13 +80,11 @@ def get_variable(
for _loc in VARNAME_LOC[varname]:
_s = _s.get(_loc)

if VARNAME_GRID[varname] == 0:
if MODEL_VARS[varname].all_timesteps:
return _s[0].astype(VARNAME_DTYPE[varname])[[int(state["KT"][0])]]
else:
return _s[0].astype(VARNAME_DTYPE[varname])

# something's gone wrong:
msg = "Varname is missing in get_variable! Contact devs."
raise ValueError(msg)


def set_variable(
state: h5py.File,
Expand Down
34 changes: 19 additions & 15 deletions PyStemmusScope/bmi/variable_reference.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@ class BmiVariable:
output: bool
units: str
grid: int
loc: list[str]
keys: list[str]
all_timesteps: bool = False


VARIABLES: tuple[BmiVariable, ...] = (
Expand All @@ -24,7 +25,7 @@ class BmiVariable:
output=True,
units="cm s-1",
grid=0,
loc=["fluxes", "Resp"],
keys=["fluxes", "Resp"],
),
BmiVariable(
name="evaporation_total",
Expand All @@ -33,7 +34,8 @@ class BmiVariable:
output=True,
units="cm s-1",
grid=0,
loc=["EVAP"],
keys=["EVAP"],
all_timesteps=True,
),
# soil vars:
BmiVariable(
Expand All @@ -43,7 +45,7 @@ class BmiVariable:
output=True,
units="degC",
grid=1,
loc=["TT"],
keys=["TT"],
),
BmiVariable(
name="soil_moisture",
Expand All @@ -52,7 +54,7 @@ class BmiVariable:
output=True,
units="m3 m-3",
grid=1,
loc=["SoilVariables", "Theta_U"],
keys=["SoilVariables", "Theta_U"],
),
BmiVariable(
name="soil_root_water_uptake",
Expand All @@ -61,7 +63,7 @@ class BmiVariable:
output=True,
units="cm s-1",
grid=0,
loc=["RWUs"],
keys=["RWUs"],
),
# surface runoff
BmiVariable(
Expand All @@ -71,7 +73,7 @@ class BmiVariable:
output=True,
units="cm s-1",
grid=0,
loc=["RS"],
keys=["RS"],
),
BmiVariable(
name="surface_runoff_hortonian",
Expand All @@ -80,7 +82,8 @@ class BmiVariable:
output=True,
units="cm s-1",
grid=0,
loc=["ForcingData", "R_Dunn"],
keys=["ForcingData", "R_Dunn"],
all_timesteps=True,
),
BmiVariable(
name="surface_runoff_dunnian",
Expand All @@ -89,7 +92,8 @@ class BmiVariable:
output=True,
units="cm s-1",
grid=0,
loc=["ForcingData", "R_Hort"],
keys=["ForcingData", "R_Hort"],
all_timesteps=True,
),
# groundwater vars (STEMMUS_SCOPE)
BmiVariable(
Expand All @@ -99,7 +103,7 @@ class BmiVariable:
output=True,
units="cm s-1",
grid=0,
loc=["RWUg"],
keys=["RWUg"],
),
BmiVariable(
name="groundwater_recharge",
Expand All @@ -108,7 +112,7 @@ class BmiVariable:
output=True,
units="cm s-1",
grid=0,
loc=["gwfluxes", "recharge"],
keys=["gwfluxes", "recharge"],
),
# groundwater (coupling) vars
BmiVariable(
Expand All @@ -118,7 +122,7 @@ class BmiVariable:
output=False,
units="-",
grid=0,
loc=["GroundwaterSettings", "GroundwaterCoupling"],
keys=["GroundwaterSettings", "GroundwaterCoupling"],
),
BmiVariable(
name="groundwater_head_bottom_layer",
Expand All @@ -127,7 +131,7 @@ class BmiVariable:
output=False,
units="cm",
grid=0,
loc=["GroundwaterSettings", "headBotmLayer"],
keys=["GroundwaterSettings", "headBotmLayer"],
),
BmiVariable(
name="groundwater_temperature",
Expand All @@ -136,7 +140,7 @@ class BmiVariable:
output=False,
units="degC",
grid=0,
loc=["GroundwaterSettings", "tempBotm"],
keys=["GroundwaterSettings", "tempBotm"],
),
BmiVariable(
name="groundwater_elevation_top_aquifer",
Expand All @@ -145,6 +149,6 @@ class BmiVariable:
output=False,
units="cm",
grid=0,
loc=["GroundwaterSettings", "toplevel"],
keys=["GroundwaterSettings", "topLevel"],
),
)

0 comments on commit c0a75a2

Please sign in to comment.