From 18a1c3a09858e7101fb24e9d5c115d0f66e1b8c2 Mon Sep 17 00:00:00 2001 From: travis Date: Thu, 11 Feb 2021 10:40:13 -0800 Subject: [PATCH] adds docstrings and type hinting to most things; fixes #5 --- __init__.py | 0 launch.py | 5 +- mosartwmpy/__init__.py | 1 + mosartwmpy/config/__init__.py | 0 mosartwmpy/config/config.py | 68 +-- mosartwmpy/config/parameters.py | 57 +++ mosartwmpy/direct_to_ocean/__init__.py | 0 mosartwmpy/direct_to_ocean/direct_to_ocean.py | 19 +- mosartwmpy/flood/__init__.py | 0 mosartwmpy/flood/flood.py | 15 +- mosartwmpy/grid/__init__.py | 0 mosartwmpy/grid/grid.py | 22 +- mosartwmpy/hillslope/__init__.py | 0 mosartwmpy/hillslope/routing.py | 15 +- mosartwmpy/hillslope/state.py | 12 +- mosartwmpy/input/__init__.py | 0 mosartwmpy/input/demand.py | 14 +- mosartwmpy/input/runoff.py | 14 +- mosartwmpy/main_channel/__init__.py | 0 mosartwmpy/main_channel/irrigation.py | 13 +- mosartwmpy/main_channel/kinematic_wave.py | 16 +- mosartwmpy/main_channel/routing.py | 18 +- mosartwmpy/main_channel/state.py | 15 +- mosartwmpy/{mosartwmpy.py => model.py} | 112 ++--- mosartwmpy/output/__init__.py | 0 mosartwmpy/output/output.py | 41 +- mosartwmpy/reservoirs/grid.py | 137 ++++++ mosartwmpy/reservoirs/regulation.py | 42 +- mosartwmpy/reservoirs/reservoirs.py | 344 +------------- mosartwmpy/reservoirs/state.py | 205 ++++++++ mosartwmpy/state/__init__.py | 0 mosartwmpy/state/state.py | 448 +++++++++--------- mosartwmpy/subnetwork/__init__.py | 0 mosartwmpy/subnetwork/irrigation.py | 13 +- mosartwmpy/subnetwork/routing.py | 18 +- mosartwmpy/subnetwork/state.py | 15 +- mosartwmpy/tests/__init__.py | 0 mosartwmpy/update/__init__.py | 0 mosartwmpy/update/update.py | 20 +- mosartwmpy/utilities/__init__.py | 0 mosartwmpy/utilities/inherit_docs.py | 20 + mosartwmpy/utilities/pretty_timer.py | 11 +- mosartwmpy/utilities/timing.py | 11 +- 43 files changed, 1001 insertions(+), 740 deletions(-) delete mode 100644 __init__.py delete mode 100644 mosartwmpy/config/__init__.py create mode 100644 mosartwmpy/config/parameters.py delete mode 100644 mosartwmpy/direct_to_ocean/__init__.py delete mode 100644 mosartwmpy/flood/__init__.py delete mode 100644 mosartwmpy/grid/__init__.py delete mode 100644 mosartwmpy/hillslope/__init__.py delete mode 100644 mosartwmpy/input/__init__.py delete mode 100644 mosartwmpy/main_channel/__init__.py rename mosartwmpy/{mosartwmpy.py => model.py} (76%) delete mode 100644 mosartwmpy/output/__init__.py create mode 100644 mosartwmpy/reservoirs/grid.py create mode 100644 mosartwmpy/reservoirs/state.py delete mode 100644 mosartwmpy/state/__init__.py delete mode 100644 mosartwmpy/subnetwork/__init__.py delete mode 100644 mosartwmpy/tests/__init__.py delete mode 100644 mosartwmpy/update/__init__.py delete mode 100644 mosartwmpy/utilities/__init__.py create mode 100644 mosartwmpy/utilities/inherit_docs.py diff --git a/__init__.py b/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/launch.py b/launch.py index d0c3df5..07f2a2c 100644 --- a/launch.py +++ b/launch.py @@ -1,13 +1,12 @@ if __name__ == '__main__': - import logging import matplotlib.pyplot as plt import pandas as pd import numpy as np - from mosartwmpy.mosartwmpy import Model + from mosartwmpy import Model # launch simulation mosart_wm = Model() mosart_wm.initialize('./config.yaml') - mosart_wm.update_until(mosart_wm.get_end_time()) + mosart_wm.update_until(mosart_wm.get_end_time()) \ No newline at end of file diff --git a/mosartwmpy/__init__.py b/mosartwmpy/__init__.py index e69de29..7040c63 100644 --- a/mosartwmpy/__init__.py +++ b/mosartwmpy/__init__.py @@ -0,0 +1 @@ +from .model import Model \ No newline at end of file diff --git a/mosartwmpy/config/__init__.py b/mosartwmpy/config/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/config/config.py b/mosartwmpy/config/config.py index 50438fa..f0f56e8 100644 --- a/mosartwmpy/config/config.py +++ b/mosartwmpy/config/config.py @@ -1,72 +1,18 @@ -import logging -import numpy as np from benedict import benedict +from benedict.dicts import benedict as Benedict -def get_config(config_file_path): +def get_config(config_file_path: str) -> Benedict: """Configuration object for the model, using the Benedict type. + Args: config_file_path (string): path to the user defined configuration yaml file + + Returns: + Benedict: A Benedict instance containing the merged configuration """ config = benedict('./config_defaults.yaml', format='yaml') if config_file_path and config_file_path != '': config.merge(benedict(config_file_path, format='yaml'), overwrite=True) - return config - - -class Parameters: - """Constant parameters used in the model.""" - - def __init__(self): - """Initialize the constants.""" - - # TODO better document what these are used for and what they should be and maybe they should be part of config? - - # TINYVALUE - self.tiny_value = 1.0e-14 - # new and improved even tinier value, MYTINYVALUE - self.tinier_value = 1.0e-50 - # small value, for less precise arithmatic - self.small_value = 1.0e-10 - # radius of the earth [m] - self.radius_earth = 6.37122e6 - # a small value in order to avoid abrupt change of hydraulic radius - self.slope_1_def = 0.1 - self.inverse_sin_atan_slope_1_def = 1.0 / (np.sin(np.arctan(self.slope_1_def))) - # flood threshold - excess water will be sent back to ocean - self.flood_threshold = 1.0e36 # [m3]? - # liquid/ice effective velocity # TODO is this used anywhere - self.effective_tracer_velocity = 10.0 # [m/s]? - # minimum river depth - self.river_depth_minimum = 1.0e-4 # [m]? - # coefficient to adjust the width of the subnetwork channel - self.subnetwork_width_parameter = 1.0 - # minimum hillslope (replaces 0s from grid file) - self.hillslope_minimum = 0.005 - # minimum subnetwork slope (replaces 0s from grid file) - self.subnetwork_slope_minimum = 0.0001 - # minimum main channel slope (replaces 0s from grid file) - self.channel_slope_minimum = 0.0001 - # kinematic wave condition # TODO what is it? - self.kinematic_wave_condition = 1.0e6 - - # reservoir parameters # TODO better describe - self.reservoir_minimum_flow_condition = 0.05 - self.reservoir_flood_control_condition = 1.0 - self.reservoir_small_magnitude_difference = 0.01 - self.reservoir_regulation_release_parameter = 0.85 - self.reservoir_runoff_capacity_condition = 0.1 - self.reservoir_flow_volume_ratio = 0.9 - - # number of supply iterations - self.reservoir_supply_iterations = 3 - - # minimum depth to perform irrigation extraction [m] - self.irrigation_extraction_condition = 0.1 - # maximum fraction of flow that can be extracted from main channel - self.irrigation_extraction_maximum_fraction = 0.5 - - # just a string... probably can dispose of these if we never do ICE separately - self.LIQUID_TRACER = 0 - self.ICE_TRACER = 1 \ No newline at end of file + return config \ No newline at end of file diff --git a/mosartwmpy/config/parameters.py b/mosartwmpy/config/parameters.py new file mode 100644 index 0000000..9873a1e --- /dev/null +++ b/mosartwmpy/config/parameters.py @@ -0,0 +1,57 @@ +import numpy as np + +class Parameters: + """Constant parameters used in the model.""" + + def __init__(self): + """Initialize the constants.""" + + # TODO better document what these are used for and what they should be and maybe they should be part of config? + + # TINYVALUE + self.tiny_value = 1.0e-14 + # new and improved even tinier value, MYTINYVALUE + self.tinier_value = 1.0e-50 + # small value, for less precise arithmatic + self.small_value = 1.0e-10 + # radius of the earth [m] + self.radius_earth = 6.37122e6 + # a small value in order to avoid abrupt change of hydraulic radius + self.slope_1_def = 0.1 + self.inverse_sin_atan_slope_1_def = 1.0 / (np.sin(np.arctan(self.slope_1_def))) + # flood threshold - excess water will be sent back to ocean + self.flood_threshold = 1.0e36 # [m3]? + # liquid/ice effective velocity # TODO is this used anywhere + self.effective_tracer_velocity = 10.0 # [m/s]? + # minimum river depth + self.river_depth_minimum = 1.0e-4 # [m]? + # coefficient to adjust the width of the subnetwork channel + self.subnetwork_width_parameter = 1.0 + # minimum hillslope (replaces 0s from grid file) + self.hillslope_minimum = 0.005 + # minimum subnetwork slope (replaces 0s from grid file) + self.subnetwork_slope_minimum = 0.0001 + # minimum main channel slope (replaces 0s from grid file) + self.channel_slope_minimum = 0.0001 + # kinematic wave condition # TODO what is it? + self.kinematic_wave_condition = 1.0e6 + + # reservoir parameters # TODO better describe + self.reservoir_minimum_flow_condition = 0.05 + self.reservoir_flood_control_condition = 1.0 + self.reservoir_small_magnitude_difference = 0.01 + self.reservoir_regulation_release_parameter = 0.85 + self.reservoir_runoff_capacity_condition = 0.1 + self.reservoir_flow_volume_ratio = 0.9 + + # number of supply iterations + self.reservoir_supply_iterations = 3 + + # minimum depth to perform irrigation extraction [m] + self.irrigation_extraction_condition = 0.1 + # maximum fraction of flow that can be extracted from main channel + self.irrigation_extraction_maximum_fraction = 0.5 + + # TODO probably can dispose of these if we never do ICE separately + self.LIQUID_TRACER = 0 + self.ICE_TRACER = 1 \ No newline at end of file diff --git a/mosartwmpy/direct_to_ocean/__init__.py b/mosartwmpy/direct_to_ocean/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/direct_to_ocean/direct_to_ocean.py b/mosartwmpy/direct_to_ocean/direct_to_ocean.py index 80752a2..70dc3c3 100644 --- a/mosartwmpy/direct_to_ocean/direct_to_ocean.py +++ b/mosartwmpy/direct_to_ocean/direct_to_ocean.py @@ -2,10 +2,21 @@ import numexpr as ne import pandas as pd -def direct_to_ocean(state, grid, parameters, config): - ### - ### Direct transfer to outlet point - ### +from benedict.dicts import benedict as Benedict + +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State + +def direct_to_ocean(state: State, grid: Grid, parameters: Parameters, config: Benedict) -> None: + """Direct transfer to outlet point; mutates state. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + config (Benedict): the model configuration + """ # direct to ocean # note - in fortran mosart this direct_to_ocean forcing could be provided from LND component, but we don't seem to be using it diff --git a/mosartwmpy/flood/__init__.py b/mosartwmpy/flood/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/flood/flood.py b/mosartwmpy/flood/flood.py index 2f50cd7..c80d17a 100644 --- a/mosartwmpy/flood/flood.py +++ b/mosartwmpy/flood/flood.py @@ -1,6 +1,19 @@ import numpy as np -def flood(state, grid, parameters, config): +from benedict.dicts import benedict as Benedict +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State + +def flood(state: State, grid: Grid, parameters: Parameters, config: Benedict) -> None: + """Excess runoff is removed from the available groundwater; mutates state. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + config (Benedict): the model configuration + """ ### ### Compute Flood diff --git a/mosartwmpy/grid/__init__.py b/mosartwmpy/grid/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/grid/grid.py b/mosartwmpy/grid/grid.py index f2a65ff..da908f8 100644 --- a/mosartwmpy/grid/grid.py +++ b/mosartwmpy/grid/grid.py @@ -4,27 +4,28 @@ import xarray as xr from xarray import open_dataset +from benedict.dicts import benedict as Benedict -from mosartwmpy.reservoirs.reservoirs import load_reservoirs +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.reservoirs.grid import load_reservoirs class Grid(): """Class to store grid related values that are constant throughout a simulation.""" - def __init__(self, config=None, parameters=None, cores=1, empty=False): + def __init__(self, config: Benedict = None, parameters: Parameters = None, empty: bool = False): """Initialize the Grid class. Args: - config (Benedict): Model configuration benedict instance - parameters (Parameters): Model parameters instance - cores (int): How many CPUs are in use - empty (bool): Set to True to return an empty instance + config (Benedict): the model configuration + parameters (Parameters): the model parameters + empty (bool): if true will return an empty instance """ # shortcut to get an empty grid instance if empty: return - # initialize all properties that need to be shared across cores + # initialize all properties self.drainage_fraction = np.empty(0) self.local_drainage_area = np.empty(0) self.total_drainage_area_multi = np.empty(0) @@ -315,9 +316,4 @@ def __init__(self, config=None, parameters=None, cores=1, empty=False): # note that reservoir grid is assumed to be the same as the domain grid if config.get('water_management.enabled', False): logging.debug(' - reservoirs') - load_reservoirs(self, config, parameters) - - # label each grid cell with a processor number it belongs to - self.process = pd.cut(self.outlet_id, bins=cores, labels=False) - # remove non interesting grid cells from process list - self.process[np.logical_not(self.mosart_mask > 0)] = -1 \ No newline at end of file + load_reservoirs(self, config, parameters) \ No newline at end of file diff --git a/mosartwmpy/hillslope/__init__.py b/mosartwmpy/hillslope/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/hillslope/routing.py b/mosartwmpy/hillslope/routing.py index 8d1d969..45b94e4 100644 --- a/mosartwmpy/hillslope/routing.py +++ b/mosartwmpy/hillslope/routing.py @@ -1,11 +1,20 @@ import numpy as np import numexpr as ne +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State from mosartwmpy.hillslope.state import update_hillslope_state -def hillslope_routing(state, grid, parameters, delta_t): - # perform the hillslope routing for the whole grid - # TODO describe what is happening heres +def hillslope_routing(state: State, grid: Grid, parameters: Parameters, delta_t: float) -> None: + """Tracks the storage of runoff water in the hillslope and the flow of runoff water from the hillslope into the channels. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + delta_t (float): the timestep for this subcycle (overall timestep / subcycles) + """ base_condition = (grid.mosart_mask > 0) & state.euler_mask diff --git a/mosartwmpy/hillslope/state.py b/mosartwmpy/hillslope/state.py index 93eaa24..cca37b5 100644 --- a/mosartwmpy/hillslope/state.py +++ b/mosartwmpy/hillslope/state.py @@ -1,8 +1,16 @@ import numpy as np import numexpr as ne -def update_hillslope_state(state, base_condition): - # update hillslope water depth +from mosartwmpy.state.state import State + +def update_hillslope_state(state: State, base_condition: np.ndarray) -> None: + """Updates the depth of water remaining in the hillslope. + + Args: + state (State): the current model state; will be mutated + base_condition (np.ndarray): a boolean array representing where the update should occur in the state + """ + state.hillslope_depth = calculate_hillslope_depth(base_condition, state.hillslope_storage, state.hillslope_depth) calculate_hillslope_depth = ne.NumExpr( diff --git a/mosartwmpy/input/__init__.py b/mosartwmpy/input/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/input/demand.py b/mosartwmpy/input/demand.py index adfc2fa..e8e2147 100644 --- a/mosartwmpy/input/demand.py +++ b/mosartwmpy/input/demand.py @@ -1,14 +1,24 @@ import numpy as np import regex as re +from datetime import datetime from xarray import open_dataset +from benedict.dicts import benedict as Benedict +from mosartwmpy.state.state import State from mosartwmpy.utilities.timing import timing # TODO this currently can only act on the entire domain... need to make more robust so can be handled in parallel # @timing -def load_demand(state, config, current_time): +def load_demand(state: State, config: Benedict, current_time: datetime) -> None: + """Loads water demand from file into the state for each grid cell. + + Args: + state (State): the current model state; will be mutated + config (Benedict): the model configuration + current_time (datetime): the current time of the simulation + """ # demand path can have placeholders for year and month and day, so check for those and replace if needed path = config.get('water_management.demand.path') @@ -17,6 +27,8 @@ def load_demand(state, config, current_time): path = re.sub('\{d[^}]*}', current_time.strftime('%d'), path) demand = open_dataset(path) + + # TODO still won't work with data on Constance, since there is no time axis state.reservoir_monthly_demand = np.array(demand[config.get('water_management.demand.demand')].sel({config.get('water_management.demand.time'): current_time}, method='pad')).flatten() diff --git a/mosartwmpy/input/runoff.py b/mosartwmpy/input/runoff.py index 7ead493..c0c2b02 100644 --- a/mosartwmpy/input/runoff.py +++ b/mosartwmpy/input/runoff.py @@ -2,10 +2,22 @@ from datetime import datetime, time, timedelta from xarray import open_dataset +from benedict.dicts import benedict as Benedict +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State from mosartwmpy.utilities.timing import timing # @timing -def load_runoff(state, grid, config, current_time): +def load_runoff(state: State, grid: Grid, config: Benedict, current_time: datetime) -> None: + """Loads runoff from file into the state for each grid cell. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + config (Benedict): the model configuration + current_time (datetime): the current time of the simulation + """ + # note that the forcing is provided in mm/s # the flood section needs m3/s, but the routing needs m/s, so be aware of the conversions # method="pad" means the closest time in the past is selected from the file diff --git a/mosartwmpy/main_channel/__init__.py b/mosartwmpy/main_channel/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/main_channel/irrigation.py b/mosartwmpy/main_channel/irrigation.py index 209987c..f4d582c 100644 --- a/mosartwmpy/main_channel/irrigation.py +++ b/mosartwmpy/main_channel/irrigation.py @@ -1,12 +1,21 @@ import numpy as np import numexpr as ne +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State from mosartwmpy.main_channel.state import update_main_channel_state from mosartwmpy.utilities.timing import timing # @timing -def main_channel_irrigation(state, grid, parameters): - # main channel routing irrigation extraction +def main_channel_irrigation(state: State, grid: Grid, parameters: Parameters) -> None: + """Tracks the supply of water from the main river channel extracted into the grid cells. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + """ depth_condition = calculate_depth_condition(grid.mosart_mask, state.euler_mask, state.tracer, parameters.LIQUID_TRACER, state.channel_depth, parameters.irrigation_extraction_condition) diff --git a/mosartwmpy/main_channel/kinematic_wave.py b/mosartwmpy/main_channel/kinematic_wave.py index 6afa453..b86397b 100644 --- a/mosartwmpy/main_channel/kinematic_wave.py +++ b/mosartwmpy/main_channel/kinematic_wave.py @@ -1,8 +1,20 @@ import numpy as np import numexpr as ne -def kinematic_wave_routing(state, grid, parameters, delta_t, base_condition): - # classic kinematic wave routing method +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State + +def kinematic_wave_routing(state: State, grid: Grid, parameters: Parameters, delta_t: float, base_condition: np.ndarray) -> None: + """Tracks the storage and flow of water in the main channel using the kinematic wave routing method. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + delta_t (float): the timestep for this subcycle (overall timestep / subcycles) + base_condition (np.ndarray): a boolean array representing where the update should occur in the state + """ # estimation of inflow state.channel_inflow_upstream = calculate_channel_inflow_upstream(base_condition, state.channel_outflow_sum_upstream_instant, state.channel_inflow_upstream) diff --git a/mosartwmpy/main_channel/routing.py b/mosartwmpy/main_channel/routing.py index 4a07c52..3b13635 100644 --- a/mosartwmpy/main_channel/routing.py +++ b/mosartwmpy/main_channel/routing.py @@ -1,14 +1,26 @@ import numpy as np import numexpr as ne +from benedict.dicts import benedict as Benedict + +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State from mosartwmpy.main_channel.kinematic_wave import kinematic_wave_routing from mosartwmpy.main_channel.state import update_main_channel_state from mosartwmpy.utilities.timing import timing # @timing -def main_channel_routing(state, grid, parameters, config, delta_t): - # perform the main channel routing - # TODO describe what is happening here +def main_channel_routing(state: State, grid: Grid, parameters: Parameters, config: Benedict, delta_t: float) -> None: + """Tracks the storage and flow of water in the main river channels. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + config (Benedict): the model configuration + delta_t (float): the timestep for this subcycle (overall timestep / subcycles) + """ tmp_outflow_downstream = 0.0 * state.zeros local_delta_t = (delta_t / config.get('simulation.routing_iterations') / grid.iterations_main_channel) diff --git a/mosartwmpy/main_channel/state.py b/mosartwmpy/main_channel/state.py index 2fd94bd..488508f 100644 --- a/mosartwmpy/main_channel/state.py +++ b/mosartwmpy/main_channel/state.py @@ -1,8 +1,19 @@ import numpy as np import numexpr as ne -def update_main_channel_state(state, grid, parameters, base_condition): - # update the physical properties of the main channel +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State + +def update_main_channel_state(state: State, grid: Grid, parameters: Parameters, base_condition: np.ndarray) -> None: + """Updates the physical properties of the main river channel based on current state. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + base_condition (np.ndarray): a boolean array representing where the update should occur in the state + """ condition = calculate_storage_condition(grid.channel_length, state.channel_storage) state.channel_cross_section_area = calculate_channel_cross_section_area(base_condition, condition, state.channel_storage, grid.channel_length, state.channel_cross_section_area) diff --git a/mosartwmpy/mosartwmpy.py b/mosartwmpy/model.py similarity index 76% rename from mosartwmpy/mosartwmpy.py rename to mosartwmpy/model.py index 6f05af3..228faee 100644 --- a/mosartwmpy/mosartwmpy.py +++ b/mosartwmpy/model.py @@ -14,9 +14,11 @@ from pathlib import Path from pathvalidate import sanitize_filename from timeit import default_timer as timer +from typing import Tuple from xarray import open_dataset -from mosartwmpy.config.config import get_config, Parameters +from mosartwmpy.config.config import get_config +from mosartwmpy.config.parameters import Parameters from mosartwmpy.grid.grid import Grid from mosartwmpy.input.runoff import load_runoff from mosartwmpy.input.demand import load_demand @@ -25,9 +27,11 @@ from mosartwmpy.state.state import State from mosartwmpy.update.update import update from mosartwmpy.utilities.pretty_timer import pretty_timer +from mosartwmpy.utilities.inherit_docs import inherit_docs +@inherit_docs class Model(Bmi): - """[summary] + """The mosartwmpy basic model interface. Args: Bmi (Bmi): The Basic Model Interface class @@ -37,7 +41,6 @@ class Model(Bmi): """ def __init__(self): - """Initialize the mosartwmpy BMI""" self.name = None self.config = benedict() self.grid = None @@ -52,13 +55,10 @@ def __init__(self): self.reservoir_streamflow_schedule = None self.reservoir_demand_schedule = None self.reservoir_prerelease_schedule = None + self.git_hash = None + self.git_untracked = None - def initialize(self, config_file_path: str = None): - """Prepare the mosartwmpy model prior to running a simulation - - Args: - config_file_path (str, optional): Path to a yaml file for overriding default settings. Defaults to None. - """ + def initialize(self, config_file_path: str) -> None: t = timer() @@ -79,12 +79,12 @@ def initialize(self, config_file_path: str = None): logging.info('Initalizing model.') logging.info(self.config.dump()) try: - githash = subprocess.check_output(['git', 'describe', '--always']).strip().decode('utf-8') - untracked = subprocess.check_output(['git', 'diff', '--name-only']).strip().decode('utf-8').split('\n') - logging.info(f'Version: {githash}') - if len(untracked) > 0: + self.git_hash = subprocess.check_output(['git', 'describe', '--always']).strip().decode('utf-8') + self.git_untracked = subprocess.check_output(['git', 'diff', '--name-only']).strip().decode('utf-8').split('\n') + logging.info(f'Version: {self.git_hash}') + if len(self.git_untracked) > 0: logging.info(f'Uncommitted changes:') - for u in untracked: + for u in self.git_untracked: logging.info(f' * {u}') except: pass @@ -98,7 +98,7 @@ def initialize(self, config_file_path: str = None): # load grid try: - self.grid = Grid(config=self.config, parameters=self.parameters, cores=self.cores) + self.grid = Grid(config=self.config, parameters=self.parameters) except Exception as e: logging.exception('Failed to load grid file; see below for stacktrace.') raise e @@ -139,7 +139,7 @@ def initialize(self, config_file_path: str = None): logging.info(f'Initialization completed in {pretty_timer(timer() - t)}.') - def update(self): + def update(self) -> None: t = timer() step = datetime.fromtimestamp(self.get_current_time()).isoformat(" ") # perform one timestep @@ -184,153 +184,153 @@ def update(self): logging.exception('Failed to write output or restart file; see below for stacktrace.') raise e - def update_until(self, time: float): + def update_until(self, time: float) -> None: # perform timesteps until time t = timer() while self.get_current_time() < time: self.update() logging.info(f'Simulation completed in {pretty_timer(timer() - t)}.') - def finalize(self): + def finalize(self) -> None: # simulation is over so free memory, write data, etc return - def get_component_name(self): + def get_component_name(self) -> str: # TODO include version/hash info? - return 'Mosart' + return f'mosartwmpy ({self.git_hash})' - def get_input_item_count(self): + def get_input_item_count(self) -> int: # TODO return 0 - def get_output_item_count(self): + def get_output_item_count(self) -> int: # TODO return 0 - def get_input_var_names(self): + def get_input_var_names(self) -> Tuple[str]: # TODO return [] - def get_output_var_names(self): + def get_output_var_names(self) -> Tuple[str]: # TODO return [] - def get_var_grid(self, name: str): + def get_var_grid(self, name: str) -> int: # only one grid used in mosart, so it is the 0th grid return 0 - def get_var_type(self, name: str): + def get_var_type(self, name: str) -> str: # TODO return 'TODO' - def get_var_units(self, name: str): + def get_var_units(self, name: str) -> str: # TODO return 'TODO' - def get_var_itemsize(self, name: str): + def get_var_itemsize(self, name: str) -> int: # TODO return 0 - def get_var_nbytes(self, name: str): + def get_var_nbytes(self, name: str) -> int: # TODO return 0 - def get_var_location(self, name: str): + def get_var_location(self, name: str) -> str: # node, edge, face return 'node' - def get_current_time(self): + def get_current_time(self) -> float: return self.current_time.timestamp() - def get_start_time(self): + def get_start_time(self) -> float: return datetime.combine(self.config.get('simulation.start_date'), time.min).timestamp() - def get_end_time(self): + def get_end_time(self) -> float: return datetime.combine(self.config.get('simulation.end_date'), time.max).timestamp() - def get_time_units(self): + def get_time_units(self) -> str: return 's' - def get_time_step(self): + def get_time_step(self) -> float: return float(self.config.get('simulation.timestep')) - def get_value(self, name: str, array): + def get_value(self, name: str, dest: np.ndarray) -> np.ndarray: # TODO copy values into array return - def get_value_ptr(self, name, array): + def get_value_ptr(self, name: str) -> np.ndarray: # TODO set array to current array pointer return - def get_value_at_indices(self, array, indices): + def get_value_at_indices(self, name: str, dest: np.ndarray, inds: np.ndarray) -> np.ndarray: # TODO copy values from indices into array return - def set_value(self, name: str, array): + def set_value(self, name: str, src: np.ndarray) -> None: # TODO set values of name from array return - def set_value_at_indices(self, name: str, indices, array): + def set_value_at_indices(self, name: str, inds: np.ndarray, src: np.ndarray) -> None: # TODO set values of name at indices from array return - def get_grid_type(self, grid: int = 0): + def get_grid_type(self, grid: int = 0) -> str: return 'uniform_rectilinear' - def get_grid_rank(self, grid: int = 0): + def get_grid_rank(self, grid: int = 0) -> int: return 2 - def get_grid_size(self, grid: int = 0): + def get_grid_size(self, grid: int = 0) -> int: return self.grid.cell_count - def get_grid_shape(self, grid: int = 0, shape = np.empty(2, dtype=int)): + def get_grid_shape(self, grid: int = 0, shape: np.ndarray = np.empty(2, dtype=int)) -> np.ndarray: shape[0] = self.grid.unique_latitudes.size shape[1] = self.grid.unique_longitudes.size return shape - def get_grid_spacing(self, grid: int = 0, spacing = np.empty(2)): + def get_grid_spacing(self, grid: int = 0, spacing: np.ndarray = np.empty(2)) -> np.ndarray: # assumes uniform grid spacing[0] = self.grid.latitude_spacing spacing[1] = self.grid.longitude_spacing return spacing - def get_grid_origin(self, grid: int = 0, origin = np.empty(2)): + def get_grid_origin(self, grid: int = 0, origin: np.ndarray = np.empty(2)) -> np.ndarray: origin[0] = self.grid.unique_latitudes[0] origin[1] = self.grid.unique_longitudes[0] return origin - def get_grid_x(self, grid: int = 0, x = None): + def get_grid_x(self, grid: int = 0, x: np.ndarray = None) -> np.ndarray: if not x: x = np.empty(self.get_grid_shape()[0]) x[:] = self.grid.unique_latitudes return x - def get_grid_y(self, grid: int = 0, y = None): + def get_grid_y(self, grid: int = 0, y: np.ndarray = None) -> np.ndarray: if not y: y = np.empty(self.get_grid_shape()[1]) y[:] = self.grid.unique_longitudes return y - def get_grid_z(self, grid: int = 0, z = None): + def get_grid_z(self, grid: int = 0, z: np.ndarray = None) -> np.ndarray: raise NotImplementedError - def get_grid_node_count(self, grid: int = 0): + def get_grid_node_count(self, grid: int = 0) -> int: raise NotImplementedError - def get_grid_edge_count(self, grid: int = 0): + def get_grid_edge_count(self, grid: int = 0) -> int: raise NotImplementedError - def get_grid_face_count(self, grid: int = 0): + def get_grid_face_count(self, grid: int = 0) -> int: raise NotImplementedError - def get_grid_edge_nodes(self, grid: int = 0, edge_nodes = None): + def get_grid_edge_nodes(self, grid: int = 0, edge_nodes: np.ndarray = None) -> np.ndarray: raise NotImplementedError - def get_grid_face_edges(self, grid: int = 0, face_edges = None): + def get_grid_face_edges(self, grid: int = 0, face_edges: np.ndarray = None) -> np.ndarray: raise NotImplementedError - def get_grid_face_nodes(self, grid: int = 0, face_nodes = None): + def get_grid_face_nodes(self, grid: int = 0, face_nodes: np.ndarray = None) -> np.ndarray: raise NotImplementedError - def get_grid_nodes_per_face(self, grid: int = 0, nodes_per_face = None): + def get_grid_nodes_per_face(self, grid: int = 0, nodes_per_face: np.ndarray = None) -> np.ndarray: raise NotImplementedError \ No newline at end of file diff --git a/mosartwmpy/output/__init__.py b/mosartwmpy/output/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/output/output.py b/mosartwmpy/output/output.py index f287d74..86dd6e3 100644 --- a/mosartwmpy/output/output.py +++ b/mosartwmpy/output/output.py @@ -8,12 +8,13 @@ from mosartwmpy.utilities.timing import timing def initialize_output(self): - # setup output buffer and averaging + """Initializes the output buffer.""" + logging.info('Initalizing output buffer.') if self.config.get('simulation.output_resolution') % self.config.get('simulation.timestep') != 0 or self.config.get('simulation.output_resolution') < self.config.get('simulation.timestep'): raise Exception('The `simulation.output_resolution` must be greater than or equal to and evenly divisible by the `simulation.timestep`.') for output in self.config.get('simulation.output'): - if getattr(self.state, output.get('variable'), None) is not None: + if getattr(self.state, output.get('variable'), None) is not None and len(getattr(self.state, output.get('variable'))) > 0: if self.output_buffer is None: self.output_buffer = pd.DataFrame(self.state.zeros, columns=[output.get('name')]) else: @@ -21,12 +22,12 @@ def initialize_output(self): # @timing def update_output(self): - # handle updating output buffer and writing to file when appropriate + """Updates the output buffer based on the current state, and averages and writes to file when appropriate.""" # update buffer self.output_n += 1 for output in self.config.get('simulation.output'): - if getattr(self.state, output.get('variable'), None) is not None: + if getattr(self.state, output.get('variable'), None) is not None and len(getattr(self.state, output.get('variable'))) > 0: self.output_buffer.loc[:, output.get('name')] += getattr(self.state, output.get('variable')) # if a new period has begun: average output buffer, write to file, and zero output buffer @@ -36,14 +37,14 @@ def update_output(self): write_output(self) self.output_n = 0 for output in self.config.get('simulation.output'): - if getattr(self.state, output.get('variable'), None) is not None: + if getattr(self.state, output.get('variable'), None) is not None and len(getattr(self.state, output.get('variable'))) > 0: self.output_buffer.loc[:, output.get('name')] = 0.0 * self.state.zeros # check if restart file if need check_restart(self) def write_output(self): - # handle writing output to file + """Writes the output buffer and requested grid variables to a netcdf file.""" # TODO only daily resolution is currently supported - need to support arbitrary resolutions # check the write frequency to see if writing to new file or appending to existing file @@ -93,7 +94,7 @@ def write_output(self): frame.lat.attrs['units'] = 'degrees_north' frame.lon.attrs['units'] = 'degrees_east' for output in self.config.get('simulation.output'): - if getattr(self.state, output.get('variable'), None) is not None: + if getattr(self.state, output.get('variable'), None) is not None and len(getattr(self.state, output.get('variable'))) > 0: if output.get('long_name'): frame[output.get('name')].attrs['long_name'] = output.get('long_name') if output.get('units'): @@ -110,7 +111,8 @@ def write_output(self): if len(self.config.get('simulation.grid_output', [])) > 0: grid_frame = pd.DataFrame(self.grid.latitude, columns=['latitude']).join(pd.DataFrame(self.grid.longitude, columns=['longitude'])) for grid_output in self.config.get('simulation.grid_output'): - grid_frame = grid_frame.join(pd.DataFrame(getattr(self.grid, grid_output.get('variable')), columns=[grid_output.get('variable')])) + if getattr(self.grid, grid_output.get('variable'), None) is not None: + grid_frame = grid_frame.join(pd.DataFrame(getattr(self.grid, grid_output.get('variable')), columns=[grid_output.get('variable')])) grid_frame = grid_frame.rename(columns={ 'latitude': 'lat', 'longitude': 'lon' @@ -119,18 +121,20 @@ def write_output(self): lat=grid_frame.lat.astype(np.float32), lon=grid_frame.lon.astype(np.float32) ) - for grid_output in self.config.get('simulation.grid_output', []): - frame = frame.assign({ - grid_output.get('name'): grid_frame[grid_output.get('variable')] - }) - if grid_output.get('long_name'): - frame[grid_output.get('name')].attrs['long_name'] = grid_output.get('long_name') - if grid_output.get('units'): - frame[grid_output.get('name')].attrs['units'] = grid_output.get('units') + for grid_output in self.config.get('simulation.grid_output'): + if getattr(self.grid, grid_output.get('variable'), None) is not None: + frame = frame.assign({ + grid_output.get('name'): grid_frame[grid_output.get('variable')] + }) + if grid_output.get('long_name'): + frame[grid_output.get('name')].attrs['long_name'] = grid_output.get('long_name') + if grid_output.get('units'): + frame[grid_output.get('name')].attrs['units'] = grid_output.get('units') frame.to_netcdf(filename, unlimited_dims=['time']) def check_restart(self): - # check if new restart file is desired + """Checks if a restart file is needed based on the current simulation time.""" + frequency = self.config.get('simulation.restart_file_frequency') is_needed = False if frequency == 'daily': @@ -149,7 +153,8 @@ def check_restart(self): write_restart(self) def write_restart(self): - # save all state values to file + """Writes the state to a netcdf file, with the current simulation time in the file name.""" + logging.info('Writing restart file.') x = self.state.to_dataframe().to_xarray() filename = f'./output/{self.name}/{self.name}_restart_{self.current_time.year}_{self.current_time.strftime("%m")}_{self.current_time.strftime("%d")}.nc' diff --git a/mosartwmpy/reservoirs/grid.py b/mosartwmpy/reservoirs/grid.py new file mode 100644 index 0000000..e7c7dc3 --- /dev/null +++ b/mosartwmpy/reservoirs/grid.py @@ -0,0 +1,137 @@ +import logging +import numpy as np +import pandas as pd + +from xarray import concat, open_dataset, Dataset +from xarray.ufuncs import logical_not +from benedict.dicts import benedict as Benedict + +from mosartwmpy.config.parameters import Parameters + +def load_reservoirs(self, config: Benedict, parameters: Parameters) -> None: + """Loads the reservoir information from file onto the grid. + + Args: + config (Benedict): the model configuration + parameters (Parameters): the model parameters + """ + + logging.info('Loading reservoir file.') + + # reservoir parameter file + reservoirs = open_dataset(config.get('water_management.reservoirs.path')) + + # load reservoir variables + for key, value in config.get('water_management.reservoirs.variables').items(): + setattr(self, key, np.array(reservoirs[value]).flatten()) + + # correct the fields with different units + # surface area from km^2 to m^2 + self.reservoir_surface_area = self.reservoir_surface_area * 1.0e6 + # capacity from millions m^3 to m^3 + self.reservoir_storage_capacity = self.reservoir_storage_capacity * 1.0e6 + + # map dams to all their dependent grid cells + # this will be a table of many to many relationship of grid cell ids to reservoir ids + self.reservoir_to_grid_mapping = reservoirs[ + config.get('water_management.reservoirs.grid_to_reservoir') + ].to_dataframe().reset_index()[[ + config.get('water_management.reservoirs.grid_to_reservoir_reservoir_dimension'), + config.get('water_management.reservoirs.grid_to_reservoir') + ]].rename(columns={ + config.get('water_management.reservoirs.grid_to_reservoir_reservoir_dimension'): 'reservoir_id', + config.get('water_management.reservoirs.grid_to_reservoir'): 'grid_cell_id' + }) + # drop nan grid ids + self.reservoir_to_grid_mapping = self.reservoir_to_grid_mapping[self.reservoir_to_grid_mapping.grid_cell_id.notna()] + # correct to zero-based grid indexing for grid cell + self.reservoir_to_grid_mapping.loc[:, self.reservoir_to_grid_mapping.grid_cell_id.name] = self.reservoir_to_grid_mapping.grid_cell_id.values - 1 + # set to integer + self.reservoir_to_grid_mapping = self.reservoir_to_grid_mapping.astype(int) + + # count of the number of reservoirs that can supply each grid cell + self.reservoir_count = np.array(pd.DataFrame(self.id).join( + self.reservoir_to_grid_mapping.groupby('grid_cell_id').count().rename(columns={'reservoir_id': 'reservoir_count'}), + how='left' + ).reservoir_count) + + # index by grid cell + self.reservoir_to_grid_mapping = self.reservoir_to_grid_mapping.set_index('grid_cell_id') + + # prepare the month or epiweek based reservoir schedules mapped to the domain + prepare_reservoir_schedule(self, config, parameters, reservoirs) + + reservoirs.close() + + +def prepare_reservoir_schedule(self, config: Benedict, parameters: Parameters, reservoirs: Dataset) -> None: + """Establishes the reservoir schedule and flow. + + Args: + config (Benedict): the model configuration + parameters (Parameters): the model parameters + reservoirs (Dataset): the reservoir dataset loaded from file + """ + + # the reservoir streamflow and demand are specified by the time resolution and reservoir id + # so let's remap those to the actual mosart domain for ease of use + + # TODO i had wanted to convert these all to epiweeks no matter what format provided, but we don't know what year all the data came from + + # streamflow flux + streamflow_time_name = config.get('water_management.reservoirs.streamflow_time_resolution') + streamflow = reservoirs[config.get('water_management.reservoirs.streamflow')] + schedule = None + for t in np.arange(streamflow.shape[0]): + flow = streamflow[t, :].to_pandas().to_frame('streamflow') + sched = pd.DataFrame(self.reservoir_id, columns=['reservoir_id']).merge(flow, how='left', left_on='reservoir_id', right_index=True)[['streamflow']].to_xarray().expand_dims( + {streamflow_time_name: 1}, + axis=0 + ) + if schedule is None: + schedule = sched + else: + schedule = concat([schedule, sched], dim=streamflow_time_name) + self.reservoir_streamflow_schedule = schedule.assign_coords( + # if monthly, convert to 1 based month index (instead of starting from 0) + {streamflow_time_name: (streamflow_time_name, schedule[streamflow_time_name].values + (1 if streamflow_time_name == 'month' else 0))} + ).streamflow + + # demand volume + demand_time_name = config.get('water_management.reservoirs.demand_time_resolution') + demand = reservoirs[config.get('water_management.reservoirs.demand')] + schedule = None + for t in np.arange(demand.shape[0]): + dem = demand[t, :].to_pandas().to_frame('demand') + sched = pd.DataFrame(self.reservoir_id, columns=['reservoir_id']).merge(dem, how='left', left_on='reservoir_id', right_index=True)[['demand']].to_xarray().expand_dims( + {demand_time_name: 1}, axis=0 + ) + if schedule is None: + schedule = sched + else: + schedule = concat([schedule, sched], dim=demand_time_name) + self.reservoir_demand_schedule = schedule.assign_coords( + # if monthly, convert to 1 based month index (instead of starting from 0) + {demand_time_name: (demand_time_name, schedule[demand_time_name].values + (1 if demand_time_name == 'month' else 0))} + ).demand + + # initialize prerelease based on long term mean flow and demand (Biemans 2011) + # TODO this assumes demand and flow use the same timescale :( + flow_avg = self.reservoir_streamflow_schedule.mean(dim=streamflow_time_name) + demand_avg = self.reservoir_demand_schedule.mean(dim=demand_time_name) + prerelease = (1.0 * self.reservoir_streamflow_schedule) + prerelease[:,:] = flow_avg + # note that xarray `where` modifies the false values + condition = (demand_avg >= (0.5 * flow_avg)) & (flow_avg > 0) + prerelease = prerelease.where( + logical_not(condition), + demand_avg/ 10 + 9 / 10 * flow_avg * self.reservoir_demand_schedule / demand_avg + ) + prerelease = prerelease.where( + condition, + prerelease.where( + logical_not((flow_avg + self.reservoir_demand_schedule - demand_avg) > 0), + flow_avg + self.reservoir_demand_schedule - demand_avg + ) + ) + self.reservoir_prerelease_schedule = prerelease \ No newline at end of file diff --git a/mosartwmpy/reservoirs/regulation.py b/mosartwmpy/reservoirs/regulation.py index 6a80b8b..9f72415 100644 --- a/mosartwmpy/reservoirs/regulation.py +++ b/mosartwmpy/reservoirs/regulation.py @@ -2,13 +2,24 @@ import numexpr as ne import pandas as pd +from benedict.dicts import benedict as Benedict from numba import jit, prange +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State from mosartwmpy.utilities.timing import timing # @timing -def regulation(state, grid, parameters, delta_t): - # regulation of the flow from the reservoirs, applied to flow entering the grid cell, i.e. subnetwork storage downstream of reservoir +def regulation(state: State, grid: Grid, parameters: Parameters, delta_t: float) -> None: + """Regulates the flow across the reservoirs. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + delta_t (float): the timestep for this routing iteration (overall timestep / subcycles / routing iterations) + """ base_condition = ( (grid.mosart_mask > 0) & @@ -79,9 +90,16 @@ def regulation(state, grid, parameters, delta_t): ) # @timing -def extraction_regulated_flow(state, grid, parameters, config, delta_t): - # extract water from the reservoir release - # the extraction needs to be distributed across the dependent cells demand +def extraction_regulated_flow(state: State, grid: Grid, parameters: Parameters, config: Benedict, delta_t: float) -> None: + """Tracks the supply of water extracted from the reservoirs to fulfill demand from dependent grid cells. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + config (Benedict): the model configuration + delta_t (float): the timestep for this subcycle (overall timestep / subcycles) + """ # notes from fortran mosart: # This is an iterative algorithm that converts main channel flow @@ -126,9 +144,6 @@ def extraction_regulated_flow(state, grid, parameters, config, delta_t): # 3. if any sum of dam fraction < 1.0 for a gridcell, then provide fraction of # demand to gridcell prorated by the dam fraction of each dam. # - # dam_uptake is the amount of water removed from the dam, it's a global array. - # Gridcells from different tasks will accumluate the amount of water removed - # from each dam in this array. has_reservoir = np.isfinite(grid.reservoir_id) @@ -206,7 +221,16 @@ def extraction_regulated_flow(state, grid, parameters, config, delta_t): state.channel_outflow_downstream[:] -= pd.DataFrame(grid.reservoir_id, columns=['reservoir_id']).merge(reservoir_demand_flow.flow_volume, how='left', left_on='reservoir_id', right_index=True).flow_volume.fillna(0).values / delta_t @jit(nopython=True, parallel=True, fastmath=True, cache=True) -def parallel_is_in(a, b): +def parallel_is_in(a: np.ndarray, b: np.ndarray) -> np.ndarray: + """Just-in-time compiled parallel C code for determing the boolean mask on array a values that exist in array b. This is slightly faster than pandas/numpy. + + Args: + a (np.ndarray): the array of values to mask + b (np.ndarray): the array of values for comparison + + Returns: + np.ndarray: a boolean mask of a in b. + """ s = set(b) l = len(a) result = np.full(l, False) diff --git a/mosartwmpy/reservoirs/reservoirs.py b/mosartwmpy/reservoirs/reservoirs.py index 5b02b4c..c348368 100644 --- a/mosartwmpy/reservoirs/reservoirs.py +++ b/mosartwmpy/reservoirs/reservoirs.py @@ -1,324 +1,14 @@ -import logging import numpy as np -import pandas as pd +from datetime import datetime from epiweeks import Week -from xarray import concat, open_dataset -from xarray.ufuncs import logical_not - -# TODO in fortran mosart there is a StorCalibFlag that affects how storage targets are calculated -- code so far is written assuming that it == 0 - -def load_reservoirs(grid, config, parameters): - """[summary] - - Args: - grid (Grid): Grid instance for this simulation - config (Benedict): Model configuration benedict instance - parameters (Parameters): Model parameters instance - """ - - logging.info('Loading reservoir file.') - - # reservoir parameter file - reservoirs = open_dataset(config.get('water_management.reservoirs.path')) - - # load reservoir variables - for key, value in config.get('water_management.reservoirs.variables').items(): - setattr(grid, key, np.array(reservoirs[value]).flatten()) - - # correct the fields with different units - # surface area from km^2 to m^2 - grid.reservoir_surface_area = grid.reservoir_surface_area * 1.0e6 - # capacity from millions m^3 to m^3 - grid.reservoir_storage_capacity = grid.reservoir_storage_capacity * 1.0e6 - - # map dams to all their dependent grid cells - # this will be a table of many to many relationship of grid cell ids to reservoir ids - grid.reservoir_to_grid_mapping = reservoirs[ - config.get('water_management.reservoirs.grid_to_reservoir') - ].to_dataframe().reset_index()[[ - config.get('water_management.reservoirs.grid_to_reservoir_reservoir_dimension'), - config.get('water_management.reservoirs.grid_to_reservoir') - ]].rename(columns={ - config.get('water_management.reservoirs.grid_to_reservoir_reservoir_dimension'): 'reservoir_id', - config.get('water_management.reservoirs.grid_to_reservoir'): 'grid_cell_id' - }) - # drop nan grid ids - grid.reservoir_to_grid_mapping = grid.reservoir_to_grid_mapping[grid.reservoir_to_grid_mapping.grid_cell_id.notna()] - # correct to zero-based grid indexing for grid cell - grid.reservoir_to_grid_mapping.loc[:, grid.reservoir_to_grid_mapping.grid_cell_id.name] = grid.reservoir_to_grid_mapping.grid_cell_id.values - 1 - # set to integer - grid.reservoir_to_grid_mapping = grid.reservoir_to_grid_mapping.astype(int) - - # count of the number of reservoirs that can supply each grid cell - grid.reservoir_count = np.array(pd.DataFrame(grid.id).join( - grid.reservoir_to_grid_mapping.groupby('grid_cell_id').count().rename(columns={'reservoir_id': 'reservoir_count'}), - how='left' - ).reservoir_count) - - # index by grid cell - grid.reservoir_to_grid_mapping = grid.reservoir_to_grid_mapping.set_index('grid_cell_id') - - # prepare the month or epiweek based reservoir schedules mapped to the domain - prepare_reservoir_schedule(grid, config, parameters, reservoirs) - - reservoirs.close() +from benedict.dicts import benedict as Benedict +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.state.state import State +from mosartwmpy.grid.grid import Grid -def prepare_reservoir_schedule(grid, config, parameters, reservoirs): - # the reservoir streamflow and demand are specified by the time resolution and reservoir id - # so let's remap those to the actual mosart domain for ease of use - - # TODO i had wanted to convert these all to epiweeks no matter what format provided, but we don't know what year all the data came from - - # streamflow flux - streamflow_time_name = config.get('water_management.reservoirs.streamflow_time_resolution') - streamflow = reservoirs[config.get('water_management.reservoirs.streamflow')] - schedule = None - for t in np.arange(streamflow.shape[0]): - flow = streamflow[t, :].to_pandas().to_frame('streamflow') - sched = pd.DataFrame(grid.reservoir_id, columns=['reservoir_id']).merge(flow, how='left', left_on='reservoir_id', right_index=True)[['streamflow']].to_xarray().expand_dims( - {streamflow_time_name: 1}, - axis=0 - ) - if schedule is None: - schedule = sched - else: - schedule = concat([schedule, sched], dim=streamflow_time_name) - grid.reservoir_streamflow_schedule = schedule.assign_coords( - # if monthly, convert to 1 based month index (instead of starting from 0) - {streamflow_time_name: (streamflow_time_name, schedule[streamflow_time_name].values + (1 if streamflow_time_name == 'month' else 0))} - ).streamflow - - # demand volume - demand_time_name = config.get('water_management.reservoirs.demand_time_resolution') - demand = reservoirs[config.get('water_management.reservoirs.demand')] - schedule = None - for t in np.arange(demand.shape[0]): - dem = demand[t, :].to_pandas().to_frame('demand') - sched = pd.DataFrame(grid.reservoir_id, columns=['reservoir_id']).merge(dem, how='left', left_on='reservoir_id', right_index=True)[['demand']].to_xarray().expand_dims( - {demand_time_name: 1}, axis=0 - ) - if schedule is None: - schedule = sched - else: - schedule = concat([schedule, sched], dim=demand_time_name) - grid.reservoir_demand_schedule = schedule.assign_coords( - # if monthly, convert to 1 based month index (instead of starting from 0) - {demand_time_name: (demand_time_name, schedule[demand_time_name].values + (1 if demand_time_name == 'month' else 0))} - ).demand - - # initialize prerelease based on long term mean flow and demand (Biemans 2011) - # TODO this assumes demand and flow use the same timescale :( - flow_avg = grid.reservoir_streamflow_schedule.mean(dim=streamflow_time_name) - demand_avg = grid.reservoir_demand_schedule.mean(dim=demand_time_name) - prerelease = (1.0 * grid.reservoir_streamflow_schedule) - prerelease[:,:] = flow_avg - # note that xarray `where` modifies the false values - condition = (demand_avg >= (0.5 * flow_avg)) & (flow_avg > 0) - prerelease = prerelease.where( - logical_not(condition), - demand_avg/ 10 + 9 / 10 * flow_avg * grid.reservoir_demand_schedule / demand_avg - ) - prerelease = prerelease.where( - condition, - prerelease.where( - logical_not((flow_avg + grid.reservoir_demand_schedule - demand_avg) > 0), - flow_avg + grid.reservoir_demand_schedule - demand_avg - ) - ) - grid.reservoir_prerelease_schedule = prerelease - - -def initialize_reservoir_state(state, grid, config, parameters): - """[summary] - - Args: - state (State): Model state instance - grid (Grid): Model grid instance - config (Config): Model config instance - parameters (Parameters): Model parameters instance - """ - - # reservoir storage at the start of the operation year - state.reservoir_storage_operation_year_start = 0.85 * grid.reservoir_storage_capacity - - # initial storage in each reservoir - state.reservoir_storage = 0.9 * grid.reservoir_storage_capacity - - initialize_reservoir_start_of_operation_year(state, grid, config, parameters) - - -def initialize_reservoir_start_of_operation_year(state, grid, config, parameters): - # define the start of the operation - define irrigation releases pattern - # multiple hydrograph - 1 peak, 2 peaks, multiple small peaks - # TODO this all depends on the schedules being monthly :( - - streamflow_time_name = config.get('water_management.reservoirs.streamflow_time_resolution') - - # find the peak flow and peak flow month for each reservoir - peak = np.max(grid.reservoir_streamflow_schedule.values, axis=0) - month_start_operations = grid.reservoir_streamflow_schedule.idxmax(dim=streamflow_time_name).values - - # correct the month start for reservoirs where average flow is greater than a small value and magnitude of peak flow difference from average is greater than smaller value - # TODO a little hard to follow the logic here but it seem to be related to number of peaks/troughs - flow_avg = grid.reservoir_streamflow_schedule.mean(dim=streamflow_time_name).values - condition = flow_avg > parameters.reservoir_minimum_flow_condition - number_of_sign_changes = 0 * flow_avg - count = 1 + number_of_sign_changes - count_max = 1 + number_of_sign_changes - month = 1 * month_start_operations - sign = np.where( - np.abs(peak - flow_avg) > parameters.reservoir_small_magnitude_difference, - np.where( - peak - flow_avg > 0, - 1, - -1 - ), - 1 - ) - current_sign = 1 * sign - for t in grid.reservoir_streamflow_schedule[streamflow_time_name].values: - # if not an xarray object with coords, the sel doesn't work, so that why the strange definition here - i = grid.reservoir_streamflow_schedule.idxmax(dim=streamflow_time_name) - i = i.where( - i + t > 12, - i + t - ).where( - i + t <= 12, - i + t - 12 - ) - flow = grid.reservoir_streamflow_schedule.sel({ - streamflow_time_name: i.fillna(1) - }) - flow = np.where(np.isfinite(i), flow, np.nan) - current_sign = np.where( - np.abs(flow - flow_avg) > parameters.reservoir_small_magnitude_difference, - np.where( - flow - flow_avg > 0, - 1, - -1 - ), - sign - ) - number_of_sign_changes = np.where( - current_sign != sign, - number_of_sign_changes + 1, - number_of_sign_changes - ) - change_condition = (current_sign != sign) & (current_sign > 0) & (number_of_sign_changes > 0) & (count > count_max) - count_max = np.where( - change_condition, - count, - count_max - ) - month_start_operations = np.where( - condition & change_condition, - month, - month_start_operations - ) - month = np.where( - current_sign != sign, - i, - month - ) - count = np.where( - current_sign != sign, - 1, - count + 1 - ) - sign = 1 * current_sign - - # setup flood control for reservoirs with average flow greater than a larger value - # TODO this is also hard to follow, but seems related to months in a row with high or low flow - month_flood_control_start = 0 * month_start_operations - month_flood_control_end = 0 * month_start_operations - condition = flow_avg > parameters.reservoir_flood_control_condition - match = 0 * month - keep_going = np.where( - np.isfinite(month_start_operations), - True, - False - ) - # TODO why 8? - for j in np.arange(8): - t = j+1 - # if not an xarray object with coords, the sel doesn't work, so that why the strange definitions here - month = grid.reservoir_streamflow_schedule.idxmax(dim=streamflow_time_name) - month = month.where( - month_start_operations - t < 1, - month_start_operations - t - ).where( - month_start_operations - t >= 1, - month_start_operations - t + 12 - ) - month_1 = month.where( - month_start_operations - t + 1 < 1, - month_start_operations - t + 1 - ).where( - month_start_operations - t + 1 >= 1, - month_start_operations - t + 1 + 12 - ) - month_2 = month.where( - month_start_operations - t - 1 < 1, - month_start_operations - t - 1 - ).where( - month_start_operations - t - 1 >= 1, - month_start_operations - t - 1 + 12 - ) - flow = grid.reservoir_streamflow_schedule.sel({ - streamflow_time_name: month.fillna(1) - }) - flow = np.where(np.isfinite(month), flow, np.nan) - flow_1 = grid.reservoir_streamflow_schedule.sel({ - streamflow_time_name: month_1.fillna(1) - }) - flow_1 = np.where(np.isfinite(month_1), flow_1, np.nan) - flow_2 = grid.reservoir_streamflow_schedule.sel({ - streamflow_time_name: month_2.fillna(1) - }) - flow_2 = np.where(np.isfinite(month_2), flow_2, np.nan) - end_condition = (flow >= flow_avg) & (flow_2 <= flow_avg) & (match == 0) - month_flood_control_end = np.where( - condition & end_condition & keep_going, - month, - month_flood_control_end - ) - match = np.where( - condition & end_condition & keep_going, - 1, - match - ) - start_condition = (flow <= flow_1) & (flow <= flow_2) & (flow <= flow_avg) - month_flood_control_start = np.where( - condition & start_condition & keep_going, - month, - month_flood_control_start - ) - keep_going = np.where( - condition & start_condition & keep_going, - False, - keep_going - ) - # note: in fortran mosart there's a further condition concerning hydropower, but it doesn't seem to be used - - # if flood control is active, enforce the flood control targets - flood_control_condition = (grid.reservoir_use_flood_control > 0) & (month_flood_control_start == 0) - month = np.where( - month_flood_control_end - 2 < 0, - month_flood_control_end - 2 + 12, - month_flood_control_end - 2 - ) - month_flood_control_start = np.where( - condition & flood_control_condition, - month, - month_flood_control_start - ) - - state.reservoir_month_start_operations = month_start_operations - state.reservoir_month_flood_control_start = month_flood_control_start - state.reservoir_month_flood_control_end = month_flood_control_end - +# TODO in fortran mosart there is a StorCalibFlag that affects how storage targets are calculated -- code so far is written assuming that it == 0 def reservoir_release(state, grid, config, parameters, current_time): # compute release from reservoirs @@ -338,7 +28,6 @@ def reservoir_release(state, grid, config, parameters, current_time): storage_targets(state, grid, config, parameters, current_time) - def regulation_release(state, grid, config, parameters, current_time): # compute the expected monthly release based on Biemans (2011) @@ -375,17 +64,24 @@ def regulation_release(state, grid, config, parameters, current_time): ) ) -def storage_targets(state, grid, config, parameters, current_time): - # define the necessary drop in storage based on storage targets at the start of the month - # should not be run during the euler solve # TODO is that because it's expensive? - +def storage_targets(state: State, grid: Grid, config: Benedict, parameters: Parameters, current_time: datetime) -> None: + """Define the necessary drop in storage based on the reservoir storage targets at the start of the month. + + Args: + state (State): the model state + grid (Grid): the model grid + config (Config): the model configuration + parameters (Parameters): the model parameters + current_time (datetime): the current simulation time + """ + # TODO the logic here is really hard to follow... can it be simplified or made more readable? - + # TODO this is still written assuming monthly, but here's the epiweek for when that is relevant epiweek = Week.fromdate(current_time).week month = current_time.month streamflow_time_name = config.get('water_management.reservoirs.streamflow_time_resolution') - + # if flood control active and has a flood control start flood_control_condition = (grid.reservoir_use_flood_control > 0) & (state.reservoir_month_flood_control_start > 0) # modify release in order to maintain a certain storage level @@ -455,4 +151,4 @@ def storage_targets(state, grid, config, parameters, current_time): (state.reservoir_release> grid.reservoir_streamflow_schedule.mean(dim=streamflow_time_name).values) & (first_condition | second_condition), grid.reservoir_streamflow_schedule.mean(dim=streamflow_time_name).values, state.reservoir_release - ) \ No newline at end of file + ) diff --git a/mosartwmpy/reservoirs/state.py b/mosartwmpy/reservoirs/state.py new file mode 100644 index 0000000..6d305a9 --- /dev/null +++ b/mosartwmpy/reservoirs/state.py @@ -0,0 +1,205 @@ +import numpy as np + +from benedict.dicts import benedict as Benedict + +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid + +# TODO in fortran mosart there is a StorCalibFlag that affects how storage targets are calculated -- code so far is written assuming that it == 0 + +def initialize_reservoir_state(self, grid: Grid, config: Benedict, parameters: Parameters) -> None: + """Initializes the reservoir state. + + Args: + grid (Grid): the model grid + config (Config): the model configuration + parameters (Parameters): the model parameters + """ + + # reservoir storage at the start of the operation year + self.reservoir_storage_operation_year_start = 0.85 * grid.reservoir_storage_capacity + + # initial storage in each reservoir + self.reservoir_storage = 0.9 * grid.reservoir_storage_capacity + + initialize_reservoir_start_of_operation_year(self, grid, config, parameters) + + +def initialize_reservoir_start_of_operation_year(self, grid: Grid, config: Benedict, parameters: Parameters) -> None: + """Determines the start of operation for each reservoir, which influences the irrigation release patterns + + Args: + grid (Grid): the model grid + config (Config): the model configuration + parameters (Parameters): the model parameters + """ + + # Note from fortran mosart: + # multiple hydrograph - 1 peak, 2 peaks, multiple small peaks + + # TODO this all depends on the schedules being monthly :( + + streamflow_time_name = config.get('water_management.reservoirs.streamflow_time_resolution') + + # find the peak flow and peak flow month for each reservoir + peak = np.max(grid.reservoir_streamflow_schedule.values, axis=0) + month_start_operations = grid.reservoir_streamflow_schedule.idxmax(dim=streamflow_time_name).values + + # correct the month start for reservoirs where average flow is greater than a small value and magnitude of peak flow difference from average is greater than smaller value + # TODO a little hard to follow the logic here but it seem to be related to number of peaks/troughs + flow_avg = grid.reservoir_streamflow_schedule.mean(dim=streamflow_time_name).values + condition = flow_avg > parameters.reservoir_minimum_flow_condition + number_of_sign_changes = 0 * flow_avg + count = 1 + number_of_sign_changes + count_max = 1 + number_of_sign_changes + month = 1 * month_start_operations + sign = np.where( + np.abs(peak - flow_avg) > parameters.reservoir_small_magnitude_difference, + np.where( + peak - flow_avg > 0, + 1, + -1 + ), + 1 + ) + current_sign = 1 * sign + for t in grid.reservoir_streamflow_schedule[streamflow_time_name].values: + # if not an xarray object with coords, the sel doesn't work, so that why the strange definition here + i = grid.reservoir_streamflow_schedule.idxmax(dim=streamflow_time_name) + i = i.where( + i + t > 12, + i + t + ).where( + i + t <= 12, + i + t - 12 + ) + flow = grid.reservoir_streamflow_schedule.sel({ + streamflow_time_name: i.fillna(1) + }) + flow = np.where(np.isfinite(i), flow, np.nan) + current_sign = np.where( + np.abs(flow - flow_avg) > parameters.reservoir_small_magnitude_difference, + np.where( + flow - flow_avg > 0, + 1, + -1 + ), + sign + ) + number_of_sign_changes = np.where( + current_sign != sign, + number_of_sign_changes + 1, + number_of_sign_changes + ) + change_condition = (current_sign != sign) & (current_sign > 0) & (number_of_sign_changes > 0) & (count > count_max) + count_max = np.where( + change_condition, + count, + count_max + ) + month_start_operations = np.where( + condition & change_condition, + month, + month_start_operations + ) + month = np.where( + current_sign != sign, + i, + month + ) + count = np.where( + current_sign != sign, + 1, + count + 1 + ) + sign = 1 * current_sign + + # setup flood control for reservoirs with average flow greater than a larger value + # TODO this is also hard to follow, but seems related to months in a row with high or low flow + month_flood_control_start = 0 * month_start_operations + month_flood_control_end = 0 * month_start_operations + condition = flow_avg > parameters.reservoir_flood_control_condition + match = 0 * month + keep_going = np.where( + np.isfinite(month_start_operations), + True, + False + ) + # TODO why 8? + for j in np.arange(8): + t = j+1 + # if not an xarray object with coords, the sel doesn't work, so that why the strange definitions here + month = grid.reservoir_streamflow_schedule.idxmax(dim=streamflow_time_name) + month = month.where( + month_start_operations - t < 1, + month_start_operations - t + ).where( + month_start_operations - t >= 1, + month_start_operations - t + 12 + ) + month_1 = month.where( + month_start_operations - t + 1 < 1, + month_start_operations - t + 1 + ).where( + month_start_operations - t + 1 >= 1, + month_start_operations - t + 1 + 12 + ) + month_2 = month.where( + month_start_operations - t - 1 < 1, + month_start_operations - t - 1 + ).where( + month_start_operations - t - 1 >= 1, + month_start_operations - t - 1 + 12 + ) + flow = grid.reservoir_streamflow_schedule.sel({ + streamflow_time_name: month.fillna(1) + }) + flow = np.where(np.isfinite(month), flow, np.nan) + flow_1 = grid.reservoir_streamflow_schedule.sel({ + streamflow_time_name: month_1.fillna(1) + }) + flow_1 = np.where(np.isfinite(month_1), flow_1, np.nan) + flow_2 = grid.reservoir_streamflow_schedule.sel({ + streamflow_time_name: month_2.fillna(1) + }) + flow_2 = np.where(np.isfinite(month_2), flow_2, np.nan) + end_condition = (flow >= flow_avg) & (flow_2 <= flow_avg) & (match == 0) + month_flood_control_end = np.where( + condition & end_condition & keep_going, + month, + month_flood_control_end + ) + match = np.where( + condition & end_condition & keep_going, + 1, + match + ) + start_condition = (flow <= flow_1) & (flow <= flow_2) & (flow <= flow_avg) + month_flood_control_start = np.where( + condition & start_condition & keep_going, + month, + month_flood_control_start + ) + keep_going = np.where( + condition & start_condition & keep_going, + False, + keep_going + ) + # note: in fortran mosart there's a further condition concerning hydropower, but it doesn't seem to be used + + # if flood control is active, enforce the flood control targets + flood_control_condition = (grid.reservoir_use_flood_control > 0) & (month_flood_control_start == 0) + month = np.where( + month_flood_control_end - 2 < 0, + month_flood_control_end - 2 + 12, + month_flood_control_end - 2 + ) + month_flood_control_start = np.where( + condition & flood_control_condition, + month, + month_flood_control_start + ) + + self.reservoir_month_start_operations = month_start_operations + self.reservoir_month_flood_control_start = month_flood_control_start + self.reservoir_month_flood_control_end = month_flood_control_end \ No newline at end of file diff --git a/mosartwmpy/state/__init__.py b/mosartwmpy/state/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/state/state.py b/mosartwmpy/state/state.py index 84e1542..ee68648 100644 --- a/mosartwmpy/state/state.py +++ b/mosartwmpy/state/state.py @@ -4,19 +4,233 @@ from datetime import datetime, time from xarray import open_dataset -from mosartwmpy.reservoirs.reservoirs import initialize_reservoir_state +from benedict.dicts import benedict as Benedict + +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.reservoirs.state import initialize_reservoir_state class State(): - def __init__(self, grid=None, config=None, parameters=None, grid_size=None, empty=False): - """Initialize the state. + # flow [m3/s] + # flow + flow: np.ndarray = np.empty(0) + # outflow into downstream links from previous timestep [m3/s] + # eroup_lagi + outflow_downstream_previous_timestep: np.ndarray = np.empty(0) + # outflow into downstream links from current timestep [m3/s] + # eroup_lagf + outflow_downstream_current_timestep: np.ndarray = np.empty(0) + # initial outflow before dam regulation at current timestep [m3/s] + # erowm_regi + outflow_before_regulation: np.ndarray = np.empty(0) + # final outflow after dam regulation at current timestep [m3/s] + # erowm_regf + outflow_after_regulation: np.ndarray = np.empty(0) + # outflow sum of upstream gridcells, average [m3/s] + # eroutUp_avg + outflow_sum_upstream_average: np.ndarray = np.empty(0) + # lateral flow from hillslope, including surface and subsurface runoff generation components, average [m3/s] + # erlat_avg + lateral_flow_hillslope_average: np.ndarray = np.empty(0) + # routing storage [m3] + # volr + storage: np.ndarray = np.empty(0) + # routing change in storage [m3/s] + # dvolrdt + delta_storage: np.ndarray = np.empty(0) + # routing change in storage masked for land [m3/s] + # dvolrdtlnd + delta_storage_land: np.ndarray = np.empty(0) + # routing change in storage masked for ocean [m3/s] + # dvolrdtocn + delta_storage_ocean: np.ndarray = np.empty(0) + # basin derived flow [m3/s] + # runoff + runoff: np.ndarray = np.empty(0) + # return direct flow [m3/s] + # runofftot + runoff_total: np.ndarray = np.empty(0) + # runoff masked for land [m3/s] + # runofflnd + runoff_land: np.ndarray = np.empty(0) + # runoff masked for ocean [m3/s] + # runoffocn + runoff_ocean: np.ndarray = np.empty(0) + # direct flow [m3/s] + # direct + direct: np.ndarray = np.empty(0) + # direct-to-ocean forcing [m3/s] + # qdto + direct_to_ocean: np.ndarray = np.empty(0) + # flood water [m3/s] + # flood + flood: np.ndarray = np.empty(0) + # hillslope surface water storage [m] + # wh + hillslope_storage: np.ndarray = np.empty(0) + # change of hillslope water storage [m/s] + # dwh + hillslope_delta_storage: np.ndarray = np.empty(0) + # depth of hillslope surface water [m] + # yh + hillslope_depth: np.ndarray = np.empty(0) + # surface runoff from hillslope [m/s] + # qsur + hillslope_surface_runoff: np.ndarray = np.empty(0) + # subsurface runoff from hillslope [m/s] + # qsub + hillslope_subsurface_runoff: np.ndarray = np.empty(0) + # runoff from glacier, wetlands, and lakes [m/s] + # qgwl + hillslope_wetland_runoff: np.ndarray = np.empty(0) + # overland flor from hillslope into subchannel (outflow is negative) [m/s] + # ehout + hillslope_overland_flow: np.ndarray = np.empty(0) + # subnetwork water storage [m3] + # wt + subnetwork_storage: np.ndarray = np.empty(0) + # subnetwork water storage at previous timestep [m3] + # wt_last + subnetwork_storage_previous_timestep: np.ndarray = np.empty(0) + # change of subnetwork water storage [m3] + # dwt + subnetwork_delta_storage: np.ndarray = np.empty(0) + # depth of subnetwork water [m] + # yt + subnetwork_depth: np.ndarray = np.empty(0) + # cross section area of subnetwork [m2] + # mt + subnetwork_cross_section_area: np.ndarray = np.empty(0) + # hydraulic radii of subnetwork [m] + # rt + subnetwork_hydraulic_radii: np.ndarray = np.empty(0) + # wetness perimeter of subnetwork [m] + # pt + subnetwork_wetness_perimeter: np.ndarray = np.empty(0) + # subnetwork flow velocity [m/s] + # vt + subnetwork_flow_velocity: np.ndarray = np.empty(0) + # subnetwork mean travel time of water within travel [s] + # tt + subnetwork_mean_travel_time: np.ndarray = np.empty(0) + # subnetwork evaporation [m/s] + # tevap + subnetwork_evaporation: np.ndarray = np.empty(0) + # subnetwork lateral inflow from hillslope [m3/s] + # etin + subnetwork_lateral_inflow: np.ndarray = np.empty(0) + # subnetwork discharge into main channel (outflow is negative) [m3/s] + # etout + subnetwork_discharge: np.ndarray = np.empty(0) + # main channel storage [m3] + # wr + channel_storage: np.ndarray = np.empty(0) + # change in main channel storage [m3] + # dwr + channel_delta_storage: np.ndarray = np.empty(0) + # main channel storage at last timestep [m3] + # wr_last + channel_storage_previous_timestep: np.ndarray = np.empty(0) + # main channel water depth [m] + # yr + channel_depth: np.ndarray = np.empty(0) + # cross section area of main channel [m2] + # mr + channel_cross_section_area: np.ndarray = np.empty(0) + # hydraulic radii of main channel [m] + # rr + channel_hydraulic_radii: np.ndarray = np.empty(0) + # wetness perimeter of main channel[m] + # pr + channel_wetness_perimeter: np.ndarray = np.empty(0) + # main channel flow velocity [m/s] + # vr + channel_flow_velocity: np.ndarray = np.empty(0) + # main channel evaporation [m/s] + # erlg + channel_evaporation: np.ndarray = np.empty(0) + # lateral flow from hillslope [m3/s] + # erlateral + channel_lateral_flow_hillslope: np.ndarray = np.empty(0) + # inflow from upstream links [m3/s] + # erin + channel_inflow_upstream: np.ndarray = np.empty(0) + # outflow into downstream links [m3/s] + # erout + channel_outflow_downstream: np.ndarray = np.empty(0) + # outflow into downstream links from previous timestep [m3/s] + # TRunoff%eroup_lagi + channel_outflow_downstream_previous_timestep: np.ndarray = np.empty(0) + # outflow into downstream links from current timestep [m3/s] + # TRunoff%eroup_lagf + channel_outflow_downstream_current_timestep: np.ndarray = np.empty(0) + # initial outflow before dam regulation at current timestep [m3/s] + # TRunoff%erowm_regi + channel_outflow_before_regulation: np.ndarray = np.empty(0) + # final outflow after dam regulation at current timestep [m3/s] + # TRunoff%erowm_regf + channel_outflow_after_regulation: np.ndarray = np.empty(0) + # outflow sum of upstream gridcells, instantaneous [m3/s] + # eroutUp + channel_outflow_sum_upstream_instant: np.ndarray = np.empty(0) + # outflow sum of upstream gridcells, average [m3/s] + # TRunoff%eroutUp_avg + channel_outflow_sum_upstream_average: np.ndarray = np.empty(0) + # lateral flow from hillslope, including surface and subsurface runoff generation components, average [m3/s] + # TRunoff%erlat_avg + channel_lateral_flow_hillslope_average: np.ndarray = np.empty(0) + # flux for adjustment of water balance residual in glacier, wetlands, and lakes [m3/s] + # ergwl + channel_wetland_flux: np.ndarray = np.empty(0) + # streamflow from outlet, positive is out [m3/s] + # flow + channel_flow: np.ndarray = np.empty(0) + # tracer, i.e. liquid or ice - TODO ice not implemented yet + tracer: np.ndarray = np.empty(0) + # euler mask - which cells to perform the euler calculation on + euler_mask: np.ndarray = np.empty(0) + # a column of always all zeros, to use as a utility + zeros: np.ndarray = np.empty(0) + + + ## Reservoir related state variables + + # reservoir streamflow schedule + reservoir_streamflow: np.ndarray = np.empty(0) + # StorMthStOp + reservoir_storage_operation_year_start: np.ndarray = np.empty(0) + # storage [m3] + reservoir_storage: np.ndarray = np.empty(0) + # MthStOp, + reservoir_month_start_operations: np.ndarray = np.empty(0) + # MthStFC + reservoir_month_flood_control_start: np.ndarray = np.empty(0) + # MthNdFC + reservoir_month_flood_control_end: np.ndarray = np.empty(0) + # release [m3/s] + reservoir_release: np.ndarray = np.empty(0) + # supply [m3/s] + reservoir_supply: np.ndarray = np.empty(0) + # monthly demand [m3/s] (demand0) + reservoir_monthly_demand: np.ndarray = np.empty(0) + # unmet demand volume within sub timestep (deficit) [m3] + reservoir_demand: np.ndarray = np.empty(0) + # unmet demand over whole timestep [m3] + reservoir_deficit: np.ndarray = np.empty(0) + # potential evaporation [mm/s] # TODO this doesn't appear to be initialized anywhere currently + reservoir_potential_evaporation: np.ndarray = np.empty(0) + + + def __init__(self, grid: Grid = None, config: Benedict = None, parameters: Parameters = None, grid_size: int = None, empty: bool = False): + """Initialize the model state. Args: - grid (Grid): Model grid instance - config (Config): Model configuration instance - parameters (Parameters): Model parameters instance - grid_size (int): Flattened grid size - empty (bool): Set to True to return an empty instance + grid (Grid): the model grid + config (Config): the model configuration + parameters (Parameters): the model parameters + grid_size (int): size of the flattened grid + empty (bool): if True, return an empty instance """ # shortcut to get an empty state instance @@ -25,189 +239,8 @@ def __init__(self, grid=None, config=None, parameters=None, grid_size=None, empt # initialize all the state variables logging.info('Initializing state variables.') - for var in [ - # flow [m3/s] - # flow - 'flow', - # outflow into downstream links from previous timestep [m3/s] - # eroup_lagi - 'outflow_downstream_previous_timestep', - # outflow into downstream links from current timestep [m3/s] - # eroup_lagf - 'outflow_downstream_current_timestep', - # initial outflow before dam regulation at current timestep [m3/s] - # erowm_regi - 'outflow_before_regulation', - # final outflow after dam regulation at current timestep [m3/s] - # erowm_regf - 'outflow_after_regulation', - # outflow sum of upstream gridcells, average [m3/s] - # eroutUp_avg - 'outflow_sum_upstream_average', - # lateral flow from hillslope, including surface and subsurface runoff generation components, average [m3/s] - # erlat_avg - 'lateral_flow_hillslope_average', - # routing storage [m3] - # volr - 'storage', - # routing change in storage [m3/s] - # dvolrdt - 'delta_storage', - # routing change in storage masked for land [m3/s] - # dvolrdtlnd - 'delta_storage_land', - # routing change in storage masked for ocean [m3/s] - # dvolrdtocn - 'delta_storage_ocean', - # basin derived flow [m3/s] - # runoff - 'runoff', - # return direct flow [m3/s] - # runofftot - 'runoff_total', - # runoff masked for land [m3/s] - # runofflnd - 'runoff_land', - # runoff masked for ocean [m3/s] - # runoffocn - 'runoff_ocean', - # direct flow [m3/s] - # direct - 'direct', - # direct-to-ocean forcing [m3/s] - # qdto - 'direct_to_ocean', - # flood water [m3/s] - # flood - 'flood', - # hillslope surface water storage [m] - # wh - 'hillslope_storage', - # change of hillslope water storage [m/s] - # dwh - 'hillslope_delta_storage', - # depth of hillslope surface water [m] - # yh - 'hillslope_depth', - # surface runoff from hillslope [m/s] - # qsur - 'hillslope_surface_runoff', - # subsurface runoff from hillslope [m/s] - # qsub - 'hillslope_subsurface_runoff', - # runoff from glacier, wetlands, and lakes [m/s] - # qgwl - 'hillslope_wetland_runoff', - # overland flor from hillslope into subchannel (outflow is negative) [m/s] - # ehout - 'hillslope_overland_flow', - # subnetwork water storage [m3] - # wt - 'subnetwork_storage', - # subnetwork water storage at previous timestep [m3] - # wt_last - 'subnetwork_storage_previous_timestep', - # change of subnetwork water storage [m3] - # dwt - 'subnetwork_delta_storage', - # depth of subnetwork water [m] - # yt - 'subnetwork_depth', - # cross section area of subnetwork [m2] - # mt - 'subnetwork_cross_section_area', - # hydraulic radii of subnetwork [m] - # rt - 'subnetwork_hydraulic_radii', - # wetness perimeter of subnetwork [m] - # pt - 'subnetwork_wetness_perimeter', - # subnetwork flow velocity [m/s] - # vt - 'subnetwork_flow_velocity', - # subnetwork mean travel time of water within travel [s] - # tt - 'subnetwork_mean_travel_time', - # subnetwork evaporation [m/s] - # tevap - 'subnetwork_evaporation', - # subnetwork lateral inflow from hillslope [m3/s] - # etin - 'subnetwork_lateral_inflow', - # subnetwork discharge into main channel (outflow is negative) [m3/s] - # etout - 'subnetwork_discharge', - # main channel storage [m3] - # wr - 'channel_storage', - # change in main channel storage [m3] - # dwr - 'channel_delta_storage', - # main channel storage at last timestep [m3] - # wr_last - 'channel_storage_previous_timestep', - # main channel water depth [m] - # yr - 'channel_depth', - # cross section area of main channel [m2] - # mr - 'channel_cross_section_area', - # hydraulic radii of main channel [m] - # rr - 'channel_hydraulic_radii', - # wetness perimeter of main channel[m] - # pr - 'channel_wetness_perimeter', - # main channel flow velocity [m/s] - # vr - 'channel_flow_velocity', - # main channel evaporation [m/s] - # erlg - 'channel_evaporation', - # lateral flow from hillslope [m3/s] - # erlateral - 'channel_lateral_flow_hillslope', - # inflow from upstream links [m3/s] - # erin - 'channel_inflow_upstream', - # outflow into downstream links [m3/s] - # erout - 'channel_outflow_downstream', - # outflow into downstream links from previous timestep [m3/s] - # TRunoff%eroup_lagi - 'channel_outflow_downstream_previous_timestep', - # outflow into downstream links from current timestep [m3/s] - # TRunoff%eroup_lagf - 'channel_outflow_downstream_current_timestep', - # initial outflow before dam regulation at current timestep [m3/s] - # TRunoff%erowm_regi - 'channel_outflow_before_regulation', - # final outflow after dam regulation at current timestep [m3/s] - # TRunoff%erowm_regf - 'channel_outflow_after_regulation', - # outflow sum of upstream gridcells, instantaneous [m3/s] - # eroutUp - 'channel_outflow_sum_upstream_instant', - # outflow sum of upstream gridcells, average [m3/s] - # TRunoff%eroutUp_avg - 'channel_outflow_sum_upstream_average', - # lateral flow from hillslope, including surface and subsurface runoff generation components, average [m3/s] - # TRunoff%erlat_avg - 'channel_lateral_flow_hillslope_average', - # flux for adjustment of water balance residual in glacier, wetlands, and lakes [m3/s] - # ergwl - 'channel_wetland_flux', - # streamflow from outlet, positive is out [m3/s] - # flow - 'channel_flow', - # tracer, i.e. liquid or ice - TODO ice not implemented yet - 'tracer', - # euler mask - which cells to perform the euler calculation on - 'euler_mask', - # a column of always all zeros, to use as a utility - 'zeros' - ]: - setattr(self, var, np.zeros(grid_size)) + for key in [key for key in dir(self) if isinstance(getattr(self, key), np.ndarray)]: + setattr(self, key, np.zeros(grid_size)) # set tracer to liquid everywhere... TODO ice is not implemented self.tracer = np.full(grid_size, parameters.LIQUID_TRACER) @@ -221,36 +254,9 @@ def __init__(self, grid=None, config=None, parameters=None, grid_size=None, empt if config.get('water_management.enabled', False): logging.debug(' - reservoirs') - for var in [ - # reservoir streamflow schedule - 'reservoir_streamflow', - # StorMthStOp - 'reservoir_storage_operation_year_start', - # storage [m3] - 'reservoir_storage', - # MthStOp, - 'reservoir_month_start_operations', - # MthStFC - 'reservoir_month_flood_control_start', - # MthNdFC - 'reservoir_month_flood_control_end', - # release [m3/s] - 'reservoir_release', - # supply [m3/s] - 'reservoir_supply', - # monthly demand [m3/s] (demand0) - 'reservoir_monthly_demand', - # unmet demand volume within sub timestep (deficit) [m3] - 'reservoir_demand', - # unmet demand over whole timestep [m3] - 'reservoir_deficit', - # potential evaporation [mm/s] # TODO this doesn't appear to be initialized anywhere currently - 'reservoir_potential_evaporation' - ]: - setattr(self, var, np.zeros(grid_size)) initialize_reservoir_state(self, grid, config, parameters) - def to_dataframe(self): + def to_dataframe(self) -> pd.DataFrame: """Builds a dataframe from all the state values. Returns: @@ -263,8 +269,8 @@ def to_dataframe(self): return df @staticmethod - def from_dataframe(df): - """Creates a State instance from a dataframe. + def from_dataframe(df: pd.DataFrame) -> 'State': + """Creates a State instance from columns in a dataframe. Args: df (pd.DataFrame): the dataframe from which to build the state diff --git a/mosartwmpy/subnetwork/__init__.py b/mosartwmpy/subnetwork/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/subnetwork/irrigation.py b/mosartwmpy/subnetwork/irrigation.py index 93ed0f1..ad39be9 100644 --- a/mosartwmpy/subnetwork/irrigation.py +++ b/mosartwmpy/subnetwork/irrigation.py @@ -1,12 +1,21 @@ import numpy as np import numexpr as ne +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State from mosartwmpy.subnetwork.state import update_subnetwork_state from mosartwmpy.utilities.timing import timing # @timing -def subnetwork_irrigation(state, grid, parameters): - # subnetwork channel routing irrigation extraction +def subnetwork_irrigation(state: State, grid: Grid, parameters: Parameters) -> None: + """Tracks the supply of water from the subnetwork river channels extracted into the grid cells. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + """ depth_condition = calculate_depth_condition(grid.mosart_mask, state.euler_mask, state.tracer, parameters.LIQUID_TRACER, state.subnetwork_depth, parameters.irrigation_extraction_condition) diff --git a/mosartwmpy/subnetwork/routing.py b/mosartwmpy/subnetwork/routing.py index cf769d8..b28f5b9 100644 --- a/mosartwmpy/subnetwork/routing.py +++ b/mosartwmpy/subnetwork/routing.py @@ -1,13 +1,25 @@ import numpy as np import numexpr as ne +from benedict.dicts import benedict as Benedict + +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State from mosartwmpy.subnetwork.state import update_subnetwork_state from mosartwmpy.utilities.timing import timing # @timing -def subnetwork_routing(state, grid, parameters, config, delta_t): - # perform the subnetwork (tributary) routing - # TODO describe what is happening here +def subnetwork_routing(state: State, grid: Grid, parameters: Parameters, config: Benedict, delta_t: float) -> None: + """Tracks the storage and flow of water in the subnetwork river channels. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + config (Benedict): the model configuration + delta_t (float): the timestep for this subcycle (overall timestep / subcycles) + """ state.channel_lateral_flow_hillslope[:] = 0 local_delta_t = (delta_t / config.get('simulation.routing_iterations') / grid.iterations_subnetwork) diff --git a/mosartwmpy/subnetwork/state.py b/mosartwmpy/subnetwork/state.py index 13886da..cbff5bb 100644 --- a/mosartwmpy/subnetwork/state.py +++ b/mosartwmpy/subnetwork/state.py @@ -1,8 +1,19 @@ import numpy as np import numexpr as ne -def update_subnetwork_state(state, grid, parameters, base_condition): - # update the physical properties of the subnetwork +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State + +def update_subnetwork_state(state: State, grid: Grid, parameters: Parameters, base_condition: np.ndarray) -> None: + """Updates the physical properties of the subnetwork river channels based on current state. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + base_condition (np.ndarray): a boolean array representing where the update should occur in the state + """ # update state variables length_condition = calculate_length_condition(grid.subnetwork_length, state.subnetwork_storage) diff --git a/mosartwmpy/tests/__init__.py b/mosartwmpy/tests/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/update/__init__.py b/mosartwmpy/update/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/update/update.py b/mosartwmpy/update/update.py index 2f5c3db..45f3d20 100644 --- a/mosartwmpy/update/update.py +++ b/mosartwmpy/update/update.py @@ -2,6 +2,12 @@ import pandas as pd import logging +from benedict.dicts import benedict as Benedict + +from mosartwmpy.config.parameters import Parameters +from mosartwmpy.grid.grid import Grid +from mosartwmpy.state.state import State + from mosartwmpy.direct_to_ocean.direct_to_ocean import direct_to_ocean from mosartwmpy.flood.flood import flood from mosartwmpy.hillslope.routing import hillslope_routing @@ -18,7 +24,15 @@ pd.options.mode.chained_assignment = None # TODO add docstrings to each method/class -def update(state, grid, parameters, config): +def update(state: State, grid: Grid, parameters: Parameters, config: Benedict) -> None: + """Advance the simulation one timestamp. + + Args: + state (State): the current model state; will be mutated + grid (Grid): the model grid + parameters (Parameters): the model parameters + config (Benedict): the model configuration + """ # perform one timestep # ignore nan, overflow, underflow, and div by 0 warnings, since they are handled correctly @@ -213,6 +227,4 @@ def update(state, grid, parameters, config): grid.land_mask >= 2, state.delta_storage, 0 - ) - - #return state + ) \ No newline at end of file diff --git a/mosartwmpy/utilities/__init__.py b/mosartwmpy/utilities/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/mosartwmpy/utilities/inherit_docs.py b/mosartwmpy/utilities/inherit_docs.py new file mode 100644 index 0000000..de769ca --- /dev/null +++ b/mosartwmpy/utilities/inherit_docs.py @@ -0,0 +1,20 @@ +import types + +def inherit_docs(cls: object) -> object: + """A class decorator to automatically inherit docstrings from parents. + + Args: + cls (object): The class that needs docstrings + + Returns: + object: The class with docstrings inherited from parents + """ + + for name, func in vars(cls).items(): + if isinstance(func, types.FunctionType) and not func.__doc__: + for parent in cls.__bases__: + parfunc = getattr(parent, name, None) + if parfunc and getattr(parfunc, '__doc__', None): + func.__doc__ = parfunc.__doc__ + break + return cls \ No newline at end of file diff --git a/mosartwmpy/utilities/pretty_timer.py b/mosartwmpy/utilities/pretty_timer.py index e9af4c5..21759b4 100644 --- a/mosartwmpy/utilities/pretty_timer.py +++ b/mosartwmpy/utilities/pretty_timer.py @@ -1,5 +1,12 @@ -def pretty_timer(seconds): - # format elapsed times in a human friendly way +def pretty_timer(seconds: float) -> str: + """Formats an elapsed time in a human friendly way. + + Args: + seconds (float): a duration of time in seconds + + Returns: + str: human friendly string representing the duration + """ if seconds < 1: return f'{round(seconds * 1.0e3, 0)} milliseconds' elif seconds < 60: diff --git a/mosartwmpy/utilities/timing.py b/mosartwmpy/utilities/timing.py index 894faa4..04434d1 100644 --- a/mosartwmpy/utilities/timing.py +++ b/mosartwmpy/utilities/timing.py @@ -1,10 +1,19 @@ import logging from functools import wraps from timeit import default_timer as timer +from typing import Callable from mosartwmpy.utilities.pretty_timer import pretty_timer -def timing(f): +def timing(f: Callable) -> Callable: + """Decorator for timing a method and outputting to log. + + Args: + f (Callable): the method to time + + Returns: + Callable: the same method that now reports its timing + """ @wraps(f) def wrap(*args, **kw): t = timer()