Repository: IMMM-SFA/mosartwmpy (language: python)

Issue:
Add docstrings to all classes and methods
Add docstrings to all classes and methods


Discussion / hints:

Fix patch:
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,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/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()


Test patch:
diff --git a/mosartwmpy/tests/__init__.py b/mosartwmpy/tests/__init__.py
deleted file mode 100644
index e69de29..0000000