support overriding istarf reservoir behavior with generic reservoir o…

…perations (#98)

mosartwmpy/_version.py

@@ -1 +1 @@
-__version__ = "0.4.4"
+__version__ = "0.5.0"

mosartwmpy/config_defaults.yaml

@@ -353,6 +353,10 @@ water_management:
                 reservoir_release_min: release_min
                 reservoir_release_p_one: release_p_one
                 reservoir_release_p_two: release_p_two
+                # field name for a value that allows setting the behavior of the reservoir release per reservoir;
+                # currently, if this field holds a value of "generic" for a particular reservoir,
+                # istarf is disabled for that reservoir, any other value in the field enables istarf
+                reservoir_behavior: fit
         # database of grid cell indices that can extract from each reservoir
         dependencies:
             # path to the database dependency file; can be absolute or relative to the source code root

mosartwmpy/grid/grid.py

@@ -102,6 +102,7 @@ class Grid:
     reservoir_release_min: np.ndarray = np.empty(0)
     reservoir_release_p_one: np.ndarray = np.empty(0)
     reservoir_release_p_two: np.ndarray = np.empty(0)
+    reservoir_behavior: np.ndarray = np.empty(0)
     reservoir_dependency_database: pd.DataFrame = pd.DataFrame()
     grid_index_to_reservoirs_map: Dict = Dict.empty(
         key_type=types.int64,
@@ -171,7 +172,11 @@ def __init__(self, config: Benedict = None, parameters: Parameters = None, empty
         # TUnit%mask
         self.mosart_mask = np.where(
             np.array(self.flow_direction) < 0,
-            0,
+            np.where(
+                self.land_fraction > 0,
+                2,
+                0,
+            ),
             np.where(
                 np.array(self.flow_direction) == 0,
                 2,
@@ -376,11 +381,12 @@ def __init__(self, config: Benedict = None, parameters: Parameters = None, empty
     def __getitem__(self, item):
         return getattr(self, item)
 
-    def to_files(self, path: str) -> None:
+    def to_files(self, path: str, mask: np.ndarray) -> None:
         """Builds a dataframe from all the grid values.
         
         Args:
             path (str): the file path to save the grid zip file to
+            mask (ndarray): the mask applied by the model
         """
 
         if not path.endswith('.zip'):
@@ -402,7 +408,18 @@ def to_files(self, path: str) -> None:
         npdf = pd.DataFrame()
         for key in [key for key in keys if isinstance(getattr(self, key), np.ndarray)]:
             if getattr(self, key).size > 0:
-                npdf[key] = getattr(self, key)
+                vector = getattr(self, key)
+                unmasked = np.empty_like(mask, dtype=vector.dtype)
+                if vector.dtype == float:
+                    unmasked[:] = np.nan
+                elif vector.dtype == int:
+                    unmasked[:] = -9999
+                elif vector.dtype == bool:
+                    unmasked[:] = False
+                elif vector.dtype == np.object:
+                    unmasked[:] = np.nan
+                unmasked[mask] = vector
+                npdf[key] = unmasked
 
         # handle dataframes
         dfs = []

mosartwmpy/model.py

@@ -196,14 +196,8 @@ def update(self) -> None:
                         logging.debug(f'Reading demand rate input from file.')
                         # load the demand from file
                         load_demand(self.name, self.state, self.config, self.current_time, self.farmerABM, self.mask)
-                # only compute new release if it's the very start of simulation or new month
-                # unless ISTARF mode is enabled, in which case update the release if it's the start of a new day
-                if self.current_time == datetime.combine(self.config.get('simulation.start_date'), time.min) or \
-                    (self.config.get('water_management.reservoirs.enable_istarf') and self.current_time ==
-                     datetime(self.current_time.year, self.current_time.month, self.current_time.day, 0, 0, 0)) or \
-                        self.current_time == datetime(self.current_time.year, self.current_time.month, 1):
-                    # update reservoir release targets
-                    reservoir_release(self.state, self.grid, self.config, self.parameters, self.current_time)
+                # update reservoir release targets (logic for when to do this is inside the method)
+                reservoir_release(self.state, self.grid, self.config, self.parameters, self.current_time, self.mask)
                 # zero supply and demand
                 self.state.grid_cell_supply[:] = 0
                 self.state.grid_cell_unmet_demand[:] = 0

mosartwmpy/reservoirs/istarf.py

@@ -1,5 +1,6 @@
 import numba as nb
 import numpy as np
+import pandas as pd
 
 from datetime import datetime
 from benedict.dicts import benedict as Benedict
@@ -10,16 +11,22 @@
 from mosartwmpy.utilities.epiweek import get_epiweek_from_datetime
 
 
-def istarf_release(state: State, grid: Grid, config: Benedict, parameters: Parameters, current_time: datetime):
+def istarf_release(state: State, grid: Grid, current_time: datetime):
     # estimate reservoir release using ISTARF which is based on harmonic functions
 
     # restrict epiweek to [1, 52]
     epiweek = np.minimum(float(get_epiweek_from_datetime(current_time)), 52.0)
 
+    # boolean array indicating which cells use the istarf rules;
+    # if behavior is "generic", then just keep the monthly generic release value instead
+    uses_istarf = np.array([(x.lower() != 'generic') if pd.notna(x) else False for x in grid.reservoir_behavior])
+
+    # initialize the release array
     daily_release = np.zeros(len(grid.reservoir_id))
 
     compute_istarf_release(
         epiweek,
+        uses_istarf,
         grid.reservoir_id,
         grid.reservoir_upper_min,
         grid.reservoir_upper_max,
@@ -48,12 +55,17 @@ def istarf_release(state: State, grid: Grid, config: Benedict, parameters: Param
     )
 
     # join release back into the grid as m3/s
-    state.reservoir_release = daily_release / (24.0 * 60.0 * 60.0)
+    # except where generic behavior is requested
+    state.reservoir_release = np.where(
+        uses_istarf,
+        daily_release / (24.0 * 60.0 * 60.0),
+        state.reservoir_release
+    )
 
 
 @nb.jit(
     "void("
-        "float64, float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], "
+        "float64, boolean[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], "
         "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], "
         "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]"
     ")",
@@ -64,6 +76,7 @@ def istarf_release(state: State, grid: Grid, config: Benedict, parameters: Param
 )
 def compute_istarf_release(
     epiweek,
+    uses_istarf,
     reservoir_id,
     upper_min,
     upper_max,
@@ -95,7 +108,7 @@ def compute_istarf_release(
 
     for i in nb.prange(len(release)):
 
-        if np.isfinite(reservoir_id[i]):
+        if np.isfinite(reservoir_id[i]) and uses_istarf[i]:
 
             max_normal = np.minimum(
                 upper_max[i],

mosartwmpy/reservoirs/release.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from datetime import datetime
+from datetime import datetime, time
 from benedict.dicts import benedict as Benedict
 
 from mosartwmpy.config.parameters import Parameters
@@ -9,63 +9,76 @@
 from mosartwmpy.reservoirs.istarf import istarf_release
 
 
-def reservoir_release(state: State, grid: Grid, config: Benedict, parameters: Parameters, current_time: datetime):
+def reservoir_release(state: State, grid: Grid, config: Benedict, parameters: Parameters, current_time: datetime, mask: np.ndarray):
     # compute release from reservoirs
 
-    # if ISTARF enabled, use new module; else use old module
-    if config.get('water_management.reservoirs.enable_istarf'):
-        istarf_release(state, grid, config, parameters, current_time)
+    # to support enabling some reservoirs to use generic rules and others to use special rules,
+    # let's always compute release with the generic rules at the beginning of the month, and then let the
+    # special rules update the release afterward, where required
 
-    else:
+    # at the beginning of simulation and start of each month,
+    # apply the generic reservoir operating rules to update the release
+    if (
+        (current_time == datetime.combine(config.get('simulation.start_date'), time.min)) or
+        (current_time == datetime(current_time.year, current_time.month, 1))
+    ):
         month = current_time.month
-        # if it's the start of the operational year for the reservoir, set it's start of op year storage to the current storage
+        # if it's the start of the operational year for the reservoir,
+        # set it's start of op year storage to the current storage
         state.reservoir_storage_operation_year_start = np.where(
             state.reservoir_month_start_operations == month,
             state.reservoir_storage,
             state.reservoir_storage_operation_year_start
         )
-        regulation_release(state, grid, config, parameters, current_time)
-        storage_targets(state, grid, config, parameters, current_time)
+        regulation_release(state, grid, parameters, current_time, mask)
+        storage_targets(state, grid, current_time, mask)
 
+    # if ISTARF is enabled, and it is the start of a day, update the release targets for non-generic reservoirs
+    if (
+        config.get('water_management.reservoirs.enable_istarf') and
+        (current_time == datetime(current_time.year, current_time.month, current_time.day, 0, 0, 0))
+    ):
+        istarf_release(state, grid, current_time)
 
-def regulation_release(state, grid, config, parameters, current_time):
+
+def regulation_release(state, grid, parameters, current_time, mask):
     # compute the expected monthly release based on Biemans (2011)
 
     # initialize to the average flow
-    state.reservoir_release = grid.reservoir_streamflow_schedule.mean(dim='month').values
+    state.reservoir_release = grid.reservoir_streamflow_schedule.mean(dim='month').values[mask]
 
     # TODO what is k
     k = state.reservoir_storage_operation_year_start / (
             parameters.reservoir_regulation_release_parameter * grid.reservoir_storage_capacity)
 
     # TODO what is factor
     factor = np.where(
-        grid.reservoir_runoff_capacity > parameters.reservoir_runoff_capacity_parameter,
-        (2.0 / grid.reservoir_runoff_capacity) ** 2.0,
+        grid.reservoir_runoff_capacity.values[mask] > parameters.reservoir_runoff_capacity_parameter,
+        (2.0 / grid.reservoir_runoff_capacity.values[mask]) ** 2.0,
         0
     )
 
     # release is some combination of prerelease, average flow in the time period, and total average flow
     state.reservoir_release = np.where(
-        np.logical_or(grid.reservoir_use_electricity, grid.reservoir_use_irrigation),
+        np.logical_or(grid.reservoir_use_electricity == True, grid.reservoir_use_irrigation == True),
         np.where(
-            grid.reservoir_runoff_capacity <= 2.0,
-            k * grid.reservoir_prerelease_schedule.sel({'month': current_time.month}).values,
+            grid.reservoir_runoff_capacity.values[mask] <= 2.0,
+            k * grid.reservoir_prerelease_schedule.sel({'month': current_time.month}).values[mask],
             k * factor * grid.reservoir_prerelease_schedule.sel({
-                'month': current_time.month}).values + (1 - factor) * grid.reservoir_streamflow_schedule.sel({
-                    'month': current_time.month}).values
+                'month': current_time.month}).values[mask] + (1 - factor) * grid.reservoir_streamflow_schedule.sel({
+                    'month': current_time.month}).values[mask]
         ),
         np.where(
-            grid.reservoir_runoff_capacity <= 2.0,
-            k * grid.reservoir_streamflow_schedule.mean(dim='month').values,
+            grid.reservoir_runoff_capacity.values[mask] <= 2.0,
+            k * grid.reservoir_streamflow_schedule.mean(dim='month').values[mask],
             k * factor * grid.reservoir_streamflow_schedule.mean(
-                dim='month').values + (1 - factor) * grid.reservoir_streamflow_schedule.sel({
-                    'month': current_time.month}).values
+                dim='month').values[mask] + (1 - factor) * grid.reservoir_streamflow_schedule.sel({
+                    'month': current_time.month}).values[mask]
         )
     )
 
 
-def storage_targets(state: State, grid: Grid, config: Benedict, parameters: Parameters, current_time: datetime) -> None:
+def storage_targets(state: State, grid: Grid, current_time: datetime, mask: np.ndarray) -> None:
     """Define the necessary drop in storage based on the reservoir storage targets at the start of the month.
 
     Args:
@@ -79,7 +92,7 @@ def storage_targets(state: State, grid: Grid, config: Benedict, parameters: Para
     # TODO the logic here is really hard to follow... can it be simplified or made more readable?
 
     # 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)
+    flood_control_condition = (grid.reservoir_use_flood_control == True) & (state.reservoir_month_flood_control_start > 0)
     # modify release in order to maintain a certain storage level
     month_condition = state.reservoir_month_flood_control_start <= state.reservoir_month_flood_control_end
     total_condition = flood_control_condition & (
@@ -98,9 +111,9 @@ def storage_targets(state: State, grid: Grid, config: Benedict, parameters: Para
         drop = np.where(
             (month_condition & m_and_condition) | (np.logical_not(month_condition) & m_or_condition),
             np.where(
-                grid.reservoir_streamflow_schedule.sel({'month': m}).values >= grid.reservoir_streamflow_schedule.mean(dim='month').values,
+                grid.reservoir_streamflow_schedule.sel({'month': m}).values[mask] >= grid.reservoir_streamflow_schedule.mean(dim='month').values[mask],
                 drop + 0,
-                drop + np.abs(grid.reservoir_streamflow_schedule.mean(dim='month').values - grid.reservoir_streamflow_schedule.sel({'month': m}).values)
+                drop + np.abs(grid.reservoir_streamflow_schedule.mean(dim='month').values[mask] - grid.reservoir_streamflow_schedule.sel({'month': m}).values[mask])
             ),
             drop
         )
@@ -125,7 +138,7 @@ def storage_targets(state: State, grid: Grid, config: Benedict, parameters: Para
         (current_time.month < state.reservoir_month_start_operations)
     )
     state.reservoir_release = np.where(
-        (state.reservoir_release > grid.reservoir_streamflow_schedule.mean(dim='month').values) & (first_condition | second_condition),
-        grid.reservoir_streamflow_schedule.mean(dim='month').values,
+        (state.reservoir_release > grid.reservoir_streamflow_schedule.mean(dim='month').values[mask]) & (first_condition | second_condition),
+        grid.reservoir_streamflow_schedule.mean(dim='month').values[mask],
         state.reservoir_release
     )

mosartwmpy/tests/test_istarf.py

@@ -12,6 +12,7 @@ class IstarfTest(unittest.TestCase):
     def test_is_correct(self):
 
         epiweek = np.minimum(float(get_epiweek_from_datetime(date(1994, 10, 1))), 52.0)
+        uses_istarf = np.array([True])
         reservoir_id = np.array([0.0])
         upper_min = np.array([-np.Inf])
         upper_max = np.array([np.Inf])
@@ -40,6 +41,7 @@ def test_is_correct(self):
 
         compute_istarf_release(
             epiweek,
+            uses_istarf,
             reservoir_id,
             upper_min,
             upper_max,

setup.py

@@ -30,29 +30,29 @@ def readme():
         ]
     },
     install_requires=[
-        'bmipy==2.0',
-        'click==8.0.1',
-        'contextily==1.2.0',
-        'dask[complete]==2021.10.0',
-        'geopandas==0.10.2',
-        'h5netcdf==0.11.0',
-        'hvplot==0.7.3',
-        'matplotlib==3.4.3',
-        'nc-time-axis==1.4.0',
-        'netCDF4==1.5.7',
-        'numba==0.53.1',
-        'numpy==1.20.3',
-        'pandas==1.3.4',
-        'pathvalidate==2.5.0',
-        'psutil==5.8.0',
-        'pyarrow==6.0.0',
-        'pyomo==6.2',
-        'python-benedict==0.24.3',
-        'regex==2021.10.23',
-        'requests==2.26.0',
-        'rioxarray==0.8.0',
-        'tqdm==4.62.3',
-        'xarray==0.19.0'
+        'bmipy~=2.0',
+        'click~=8.0.1',
+        'contextily~=1.2.0',
+        'dask[complete]~=2021.10.0',
+        'geopandas~=0.10.2',
+        'h5netcdf~=0.11.0',
+        'hvplot~=0.7.3',
+        'matplotlib~=3.4.3',
+        'nc-time-axis~=1.4.0',
+        'netCDF4~=1.5.7',
+        'numba~=0.53.1',
+        'numpy~=1.20.3',
+        'pandas~=1.3.4',
+        'pathvalidate~=2.5.0',
+        'psutil~=5.8.0',
+        'pyarrow~=6.0.0',
+        'pyomo~=6.2',
+        'python-benedict~=0.24.3',
+        'regex~=2021.10.23',
+        'requests~=2.26.0',
+        'rioxarray~=0.8.0',
+        'tqdm~=4.62.3',
+        'xarray~=0.19.0'
     ],
     extras_require={
         'dev': [