From facb2b307fd141eef3e3476c068ee2649ce21ad9 Mon Sep 17 00:00:00 2001 From: Bart Schilperoort Date: Wed, 22 Nov 2023 16:03:43 +0100 Subject: [PATCH] Use STEMMUS_SCOPE's interactive mode, h5py for IO --- PyStemmusScope/bmi.py | 98 ++++++++++++++++++++++++++----------------- 1 file changed, 59 insertions(+), 39 deletions(-) diff --git a/PyStemmusScope/bmi.py b/PyStemmusScope/bmi.py index fb0e367a..a29124ad 100644 --- a/PyStemmusScope/bmi.py +++ b/PyStemmusScope/bmi.py @@ -2,12 +2,11 @@ import os import sys from pathlib import Path -from subprocess import run -from typing import Any -import hdf5storage import numpy as np from bmipy.bmi import Bmi from PyStemmusScope.config_io import read_config +import subprocess +import h5py MODEL_INPUT_VARNAMES: tuple[str] = ("soil_temperature",) @@ -42,20 +41,7 @@ def ipython_info(): return ip -def run_cmd(args): - """Run a command (i.e. run the matlab model executable). - - Args: - args: Arguments to run. - """ - if ipython_info() == "notebook": - result = run(args, shell=True) - else: - result = run(args) - result.check_returncode() - - -def load_state(config: dict) -> dict[str, Any]: +def load_state(config: dict) -> h5py.File: """Load the STEMMUS_SCOPE model state. Args: @@ -65,10 +51,10 @@ def load_state(config: dict) -> dict[str, Any]: Model state, as a dict. """ matfile = Path(config["OutputPath"]) / "STEMMUS_SCOPE_state.mat" - return hdf5storage.loadmat(matfile, appendmat=False) + return h5py.File(matfile, mode="a") -def get_variable(state: dict, varname: str) -> None: +def get_variable(state: h5py.File, varname: str) -> None: """Get a variable from the model state. Args: @@ -77,9 +63,9 @@ def get_variable(state: dict, varname: str) -> None: """ match varname: case "respiration": - return state["fluxes"]["Resp"].flatten() + return state["fluxes"]["Resp"][0] case "soil_temperature": - return state["TT"].flatten()[:-1] + return state["TT"][0,:-1] case _: if varname in MODEL_VARNAMES: msg = "Varname is missing in get_variable! Contact devs." @@ -88,7 +74,7 @@ def get_variable(state: dict, varname: str) -> None: raise ValueError(msg) -def set_variable(state: dict, varname: str, value: np.ndarray) -> dict: +def set_variable(state: h5py.File, varname: str, value: np.ndarray) -> dict: """Set a variable in the model state. Args: @@ -101,9 +87,9 @@ def set_variable(state: dict, varname: str, value: np.ndarray) -> dict: """ match varname: case "respiration": - state["fluxes"]["Resp"][0][0][0] = value + state["fluxes"]["Resp"][0] = value case "soil_temperature": - state["TT"][:-1, 0] = value + state["TT"][0,:-1] = value case _: if varname in MODEL_VARNAMES: msg = "Varname is missing in get_variable! Contact devs." @@ -113,13 +99,30 @@ def set_variable(state: dict, varname: str, value: np.ndarray) -> dict: return state +def is_alive(process: subprocess.Popen) -> bool: + """Check if the process is alive, and raise an exception if it is not.""" + if process.poll() is not None: + msg = f"Model terminated with return code {process.poll()}" + raise ConnectionError(msg) + return True + + +def wait_for_model(process: subprocess.Popen, phrase=b"Select run mode:") -> None: + """Wait for model to be ready for interaction.""" + output = b'' + while is_alive(process) and phrase not in output: + output+=process.stdout.read(1) + + class StemmusScopeBmi(Bmi): """STEMMUS_SCOPE Basic Model Interface.""" config_file: str = "" config: dict = {} - state: dict = {} - state_file: Path + state: h5py.File | None = None + state_file: Path | None = None + + matlab_process: subprocess.Popen | None = None def initialize(self, config_file: str) -> None: """Perform startup tasks for the model. @@ -132,19 +135,35 @@ def initialize(self, config_file: str) -> None: self.exe_file = self.config["ExeFilePath"] self.state_file = Path(self.config["OutputPath"]) / "STEMMUS_SCOPE_state.mat" - args = f"{self.exe_file} {self.config_file} initialize" + args = [self.exe_file, self.config_file, "interactive"] # set matlab log dirc os.environ["MATLAB_LOG_DIR"] = str(self.config["InputPath"]) - run_cmd(args) + self.matlab_process = subprocess.Popen( + args, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + bufsize=0, + ) + + wait_for_model(self.matlab_process) + is_alive(self.matlab_process) + self.matlab_process.stdin.write(b"initialize\n") + wait_for_model(self.matlab_process) def update(self) -> None: """Advance the model state by one time step.""" - if self.state != {}: - hdf5storage.savemat(self.state_file, self.state, appendmat=False) - args = f"{self.exe_file} {self.config_file} update" - run_cmd(args) + if self.state is not None: + self.state = self.state.close() # Close file to allow matlab to write + + if self.matlab_process is None: + msg = "Run initialize before trying to update the model." + raise AttributeError(msg) + + is_alive(self.matlab_process) + self.matlab_process.stdin.write(b"update\n") + wait_for_model(self.matlab_process) self.state = load_state(self.config) @@ -158,8 +177,9 @@ def update_until(self, time: float) -> None: def finalize(self) -> None: """Finalize the STEMMUS_SCOPE model.""" - args = f"{self.exe_file} {self.config_file} finalize" - run_cmd(args) + is_alive(self.matlab_process) + self.matlab_process.stdin.write(b"finalize\n") + wait_for_model(self.matlab_process, phrase=b"Finished clean up.") def get_component_name(self) -> str: """Name of the component. @@ -232,7 +252,7 @@ def get_time_units(self) -> str: def get_time_step(self) -> float: """Return the current time step of the model.""" - return self.state["KT"][0][0].flatten().astype("float") + return float(self.state["KT"][0]) def get_value(self, name: str, dest: np.ndarray) -> np.ndarray: """Get a copy of values of the given variable. @@ -321,7 +341,7 @@ def get_grid_size(self, grid: int) -> int: if grid == 0: return 1 if grid == 1: - return int(self.state["ModelSettings"]["mN"].flatten()[0]) - 1 + return int(self.state["ModelSettings"]["mN"][0]) - 1 def get_grid_type(self, grid: int) -> str: """Get the grid type as a string. @@ -345,7 +365,7 @@ def get_grid_x(self, grid: int, x: np.ndarray) -> np.ndarray: Returns: The input numpy array that holds the grid's column x-coordinates. """ - x[:] = self.state["SiteProperties"]["latitude"][0][0].flatten().astype("float") + x[:] = self.state["SiteProperties"]["latitude"][0] return x def get_grid_y(self, grid: int, y: np.ndarray) -> np.ndarray: @@ -358,7 +378,7 @@ def get_grid_y(self, grid: int, y: np.ndarray) -> np.ndarray: Returns: The input numpy array that holds the grid's column y-coordinates. """ - y[:] = self.state["SiteProperties"]["latitude"].flatten().astype("float") + y[:] = self.state["SiteProperties"]["latitude"][0] return y def get_grid_z(self, grid: int, z: np.ndarray) -> np.ndarray: @@ -375,8 +395,8 @@ def get_grid_z(self, grid: int, z: np.ndarray) -> np.ndarray: z[:] = ( -np.hstack( ( + self.state["ModelSettings"]["DeltZ_R"][:,0].cumsum()[::-1], np.array([0.0]), - self.state["ModelSettings"]["DeltZ_R"].flatten().cumsum(), ) ) / 100