Merge pull request #19 from IMMM-SFA/feature/docstrings-and-types

adds docstrings and type hinting to most things; fixes #5

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())

mosartwmpy/__init__.py

@@ -0,0 +1 @@
+from .model import Model

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
+    return config

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

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

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

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
+            load_reservoirs(self, config, parameters)

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
 

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(

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()
 

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

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)