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')
path = re.sub('\{y[^}]*}', current_time.strftime('%Y'), path)
path = re.sub('\{m[^}]*}', current_time.strftime('%m'), path)
path = re.sub('\{d[^}]*}', current_time.strftime('%d'), path)
demand = open_dataset(path)
# if the demand file has a time axis, use it; otherwise assume data is just 2d
if config.get('water_management.demand.time', None) in demand:
state.grid_cell_demand_rate = np.array(demand[config.get('water_management.demand.demand')].sel({config.get('water_management.demand.time'): current_time}, method='pad')).flatten()
else:
state.grid_cell_demand_rate = np.array(demand[config.get('water_management.demand.demand')]).flatten()
# fill missing values with 0
state.grid_cell_demand_rate = np.where(
np.logical_not(np.isfinite(state.grid_cell_demand_rate)),
0,
state.grid_cell_demand_rate
)
demand.close()
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
runoff = open_dataset(config.get('runoff.path'))
sel = {config.get('runoff.time'): current_time}
if config.get('runoff.variables.surface_runoff', None) is not None:
state.hillslope_surface_runoff = 0.001 * grid.land_fraction * grid.area * np.array(
runoff[config.get('runoff.variables.surface_runoff')].sel(
sel, method='pad')).flatten()
if config.get('runoff.variables.subsurface_runoff', None) is not None:
state.hillslope_subsurface_runoff = 0.001 * grid.land_fraction * grid.area * np.array(
runoff[config.get('runoff.variables.subsurface_runoff')].sel(
sel, method='pad')).flatten()
if config.get('runoff.variables.wetland_runoff', None) is not None:
state.hillslope_wetland_runoff = 0.001 * grid.land_fraction * grid.area * np.array(
runoff[config.get('runoff.variables.wetland_runoff')].sel(
sel, method='pad')).flatten()
runoff.close()
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
### Remove excess liquid water from land
###
### TODO tcraig leaves a comment here concerning surface_runoff in fortran mosart:
### "This seems like an odd approach, you
### might create negative forcing. why not take it out of
### the volr directly? it's also odd to compute this
### at the initial time of the time loop. why not do
### it at the end or even during the run loop as the
### new volume is computed. fluxout depends on volr, so
### how this is implemented does impact the solution."
###
# flux sent back to land
state.flood = np.where(
(grid.land_mask == 1) & (state.storage > parameters.flood_threshold) &
(state.tracer == parameters.LIQUID_TRACER),
(state.storage - parameters.flood_threshold) /
config.get('simulation.timestep'), 0)
# remove this flux from the input runoff from land
state.hillslope_surface_runoff = np.where(
state.tracer == parameters.LIQUID_TRACER,
state.hillslope_surface_runoff - state.flood,
state.hillslope_surface_runoff)
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): 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
"""
# flow [m3/s]
# flow
self.flow: np.ndarray = np.empty(0)
# outflow into downstream links from previous timestep [m3/s]
# eroup_lagi
self.outflow_downstream_previous_timestep: np.ndarray = np.empty(0)
# outflow into downstream links from current timestep [m3/s]
# eroup_lagf
self.outflow_downstream_current_timestep: np.ndarray = np.empty(0)
# initial outflow before dam regulation at current timestep [m3/s]
# erowm_regi
self.outflow_before_regulation: np.ndarray = np.empty(0)
# final outflow after dam regulation at current timestep [m3/s]
# erowm_regf
self.outflow_after_regulation: np.ndarray = np.empty(0)
# outflow sum of upstream gridcells, average [m3/s]
# eroutUp_avg
self.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
self.lateral_flow_hillslope_average: np.ndarray = np.empty(0)
# routing storage [m3]
# volr
self.storage: np.ndarray = np.empty(0)
# routing change in storage [m3/s]
# dvolrdt
self.delta_storage: np.ndarray = np.empty(0)
# routing change in storage masked for land [m3/s]
# dvolrdtlnd
self.delta_storage_land: np.ndarray = np.empty(0)
# routing change in storage masked for ocean [m3/s]
# dvolrdtocn
self.delta_storage_ocean: np.ndarray = np.empty(0)
# basin derived flow [m3/s]
# runoff
self.runoff: np.ndarray = np.empty(0)
# return direct flow [m3/s]
# runofftot
self.runoff_total: np.ndarray = np.empty(0)
# runoff masked for land [m3/s]
# runofflnd
self.runoff_land: np.ndarray = np.empty(0)
# runoff masked for ocean [m3/s]
# runoffocn
self.runoff_ocean: np.ndarray = np.empty(0)
# direct flow [m3/s]
# direct
self.direct: np.ndarray = np.empty(0)
# direct-to-ocean forcing [m3/s]
# qdto
self.direct_to_ocean: np.ndarray = np.empty(0)
# flood water [m3/s]
# flood
self.flood: np.ndarray = np.empty(0)
# hillslope surface water storage [m]
# wh
self.hillslope_storage: np.ndarray = np.empty(0)
# change of hillslope water storage [m/s]
# dwh
self.hillslope_delta_storage: np.ndarray = np.empty(0)
# depth of hillslope surface water [m]
# yh
self.hillslope_depth: np.ndarray = np.empty(0)
# surface runoff from hillslope [m/s]
# qsur
self.hillslope_surface_runoff: np.ndarray = np.empty(0)
# subsurface runoff from hillslope [m/s]
# qsub
self.hillslope_subsurface_runoff: np.ndarray = np.empty(0)
# runoff from glacier, wetlands, and lakes [m/s]
# qgwl
self.hillslope_wetland_runoff: np.ndarray = np.empty(0)
# overland flor from hillslope into subchannel (outflow is negative) [m/s]
# ehout
self.hillslope_overland_flow: np.ndarray = np.empty(0)
# subnetwork water storage [m3]
# wt
self.subnetwork_storage: np.ndarray = np.empty(0)
# subnetwork water storage at previous timestep [m3]
# wt_last
self.subnetwork_storage_previous_timestep: np.ndarray = np.empty(0)
# change of subnetwork water storage [m3]
# dwt
self.subnetwork_delta_storage: np.ndarray = np.empty(0)
# depth of subnetwork water [m]
# yt
self.subnetwork_depth: np.ndarray = np.empty(0)
# cross section area of subnetwork [m2]
# mt
self.subnetwork_cross_section_area: np.ndarray = np.empty(0)
# hydraulic radii of subnetwork [m]
# rt
self.subnetwork_hydraulic_radii: np.ndarray = np.empty(0)
# wetness perimeter of subnetwork [m]
# pt
self.subnetwork_wetness_perimeter: np.ndarray = np.empty(0)
# subnetwork flow velocity [m/s]
# vt
self.subnetwork_flow_velocity: np.ndarray = np.empty(0)
# subnetwork mean travel time of water within travel [s]
# tt
self.subnetwork_mean_travel_time: np.ndarray = np.empty(0)
# subnetwork evaporation [m/s]
# tevap
self.subnetwork_evaporation: np.ndarray = np.empty(0)
# subnetwork lateral inflow from hillslope [m3/s]
# etin
self.subnetwork_lateral_inflow: np.ndarray = np.empty(0)
# subnetwork discharge into main channel (outflow is negative) [m3/s]
# etout
self.subnetwork_discharge: np.ndarray = np.empty(0)
# main channel storage [m3]
# wr
self.channel_storage: np.ndarray = np.empty(0)
# change in main channel storage [m3]
# dwr
self.channel_delta_storage: np.ndarray = np.empty(0)
# main channel storage at last timestep [m3]
# wr_last
self.channel_storage_previous_timestep: np.ndarray = np.empty(0)
# main channel water depth [m]
# yr
self.channel_depth: np.ndarray = np.empty(0)
# cross section area of main channel [m2]
# mr
self.channel_cross_section_area: np.ndarray = np.empty(0)
# hydraulic radii of main channel [m]
# rr
self.channel_hydraulic_radii: np.ndarray = np.empty(0)
# wetness perimeter of main channel[m]
# pr
self.channel_wetness_perimeter: np.ndarray = np.empty(0)
# main channel flow velocity [m/s]
# vr
self.channel_flow_velocity: np.ndarray = np.empty(0)
# main channel evaporation [m/s]
# erlg
self.channel_evaporation: np.ndarray = np.empty(0)
# lateral flow from hillslope [m3/s]
# erlateral
self.channel_lateral_flow_hillslope: np.ndarray = np.empty(0)
# inflow from upstream links [m3/s]
# erin
self.channel_inflow_upstream: np.ndarray = np.empty(0)
# outflow into downstream links [m3/s]
# erout
self.channel_outflow_downstream: np.ndarray = np.empty(0)
# outflow into downstream links from previous timestep [m3/s]
# TRunoff%eroup_lagi
self.channel_outflow_downstream_previous_timestep: np.ndarray = np.empty(
0)
# outflow into downstream links from current timestep [m3/s]
# TRunoff%eroup_lagf
self.channel_outflow_downstream_current_timestep: np.ndarray = np.empty(
0)
# initial outflow before dam regulation at current timestep [m3/s]
# TRunoff%erowm_regi
self.channel_outflow_before_regulation: np.ndarray = np.empty(0)
# final outflow after dam regulation at current timestep [m3/s]
# TRunoff%erowm_regf
self.channel_outflow_after_regulation: np.ndarray = np.empty(0)
# outflow sum of upstream gridcells, instantaneous [m3/s]
# eroutUp
self.channel_outflow_sum_upstream_instant: np.ndarray = np.empty(0)
# outflow sum of upstream gridcells, average [m3/s]
# TRunoff%eroutUp_avg
self.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
self.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
self.channel_wetland_flux: np.ndarray = np.empty(0)
# streamflow from outlet, positive is out [m3/s]
# flow
self.channel_flow: np.ndarray = np.empty(0)
# tracer, i.e. liquid or ice - TODO ice not implemented yet
self.tracer: np.ndarray = np.empty(0)
# euler mask - which cells to perform the euler calculation on
self.euler_mask: np.ndarray = np.empty(0)
# a column of always all zeros, to use as a utility
self.zeros: np.ndarray = np.empty(0)
## Reservoir related state variables
# reservoir streamflow schedule
self.reservoir_streamflow: np.ndarray = np.empty(0)
# StorMthStOp
self.reservoir_storage_operation_year_start: np.ndarray = np.empty(0)
# storage [m3]
self.reservoir_storage: np.ndarray = np.empty(0)
# MthStOp,
self.reservoir_month_start_operations: np.ndarray = np.empty(0)
# MthStFC
self.reservoir_month_flood_control_start: np.ndarray = np.empty(0)
# MthNdFC
self.reservoir_month_flood_control_end: np.ndarray = np.empty(0)
# release [m3/s]
self.reservoir_release: np.ndarray = np.empty(0)
# supply [m3/s]
self.grid_cell_supply: np.ndarray = np.empty(0)
# demand rate [m3/s] (demand0)
self.grid_cell_demand_rate: np.ndarray = np.empty(0)
# unmet demand volume within sub timestep [m3]
self.grid_cell_unmet_demand: np.ndarray = np.empty(0)
# unmet demand over whole timestep [m3]
self.grid_cell_deficit: np.ndarray = np.empty(0)
# potential evaporation [mm/s] # TODO this doesn't appear to be initialized anywhere currently
self.reservoir_potential_evaporation: np.ndarray = np.empty(0)
# shortcut to get an empty state instance
if empty:
return
# initialize all the state variables
logging.info('Initializing state variables.')
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)
# set euler_mask to 1 everywhere... TODO not really necessary without ice implemented
self.euler_mask = np.where(self.tracer == parameters.LIQUID_TRACER,
True, False)
if config.get('water_management.enabled', False):
logging.debug(' - reservoirs')
initialize_reservoir_state(self, grid, config, parameters)
def __init__(self,
config: Benedict = None,
parameters: Parameters = None,
empty: bool = False):
"""Initialize the Grid class.
Args:
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
logging.info('Loading grid file.')
# open dataset
grid_dataset = open_dataset(config.get('grid.path'))
# create grid from longitude and latitude dimensions
self.unique_longitudes = np.array(
grid_dataset[config.get('grid.longitude')])
self.unique_latitudes = np.array(
grid_dataset[config.get('grid.latitude')])
self.cell_count = self.unique_longitudes.size * self.unique_latitudes.size
self.longitude_spacing = abs(self.unique_longitudes[1] -
self.unique_longitudes[0])
self.latitude_spacing = abs(self.unique_latitudes[1] -
self.unique_latitudes[0])
self.longitude, self.latitude = np.meshgrid(
grid_dataset[config.get('grid.longitude')],
grid_dataset[config.get('grid.latitude')])
self.longitude = self.longitude.flatten()
self.latitude = self.latitude.flatten()
for key, value in config.get('grid.variables').items():
setattr(self, key, np.array(grid_dataset[value]).flatten())
# free memory
grid_dataset.close()
# use ID and dnID field to calculate masks, upstream, downstream, and outlet indices, as well as count of upstream cells
logging.debug(
' - masks, downstream, upstream, and outlet cell indices')
# ocean/land mask
# 1 == land
# 2 == ocean
# 3 == ocean outlet from land
# rtmCTL%mask
self.land_mask = np.where(
self.downstream_id >= 0, 1,
np.where(np.isin(self.id, self.downstream_id), 3, 2))
# TODO this is basically the same as the above... should consolidate code to just use one of these masks
# mosart ocean/land mask
# 0 == ocean
# 1 == land
# 2 == outlet
# TUnit%mask
self.mosart_mask = np.where(
np.array(self.flow_direction) < 0, 0,
np.where(
np.array(self.flow_direction) == 0, 2,
np.where(np.array(self.flow_direction) > 0, 1, 0)))
# determine final downstream outlet of each cell
# this essentially slices up the grid into discrete basins
# first remap cell ids into cell indices for the (reshaped) 1d grid
id_hashmap = {}
for i, _id in enumerate(self.id):
id_hashmap[int(_id)] = int(i)
# convert downstream ids into downstream indices
self.downstream_id = np.array([
id_hashmap[int(i)] if int(i) in id_hashmap else -1
for i in self.downstream_id
],
dtype=int)
# update the id to be zero-indexed (note this makes them one less than fortran mosart ids)
self.id = np.arange(self.id.size)
# follow each cell downstream to compute outlet id
size = self.downstream_id.size
self.outlet_id = np.full(size, -1)
self.upstream_id = np.full(size, -1)
self.upstream_cell_count = np.full(size, 0)
for i in np.arange(size):
if self.downstream_id[i] >= 0:
# mark as upstream cell of downstream cell
self.upstream_id[self.downstream_id[i]] = i
if self.land_mask[i] == 1:
# land
j = i
while self.land_mask[j] == 1:
self.upstream_cell_count[j] += 1
j = int(self.downstream_id[j])
if self.land_mask[j] == 3:
# found the ocean outlet
self.upstream_cell_count[j] += 1
self.outlet_id[i] = j
else:
# ocean
self.upstream_cell_count[i] += 1
self.outlet_id[i] = i
# recalculate area to fill in missing values
# assumes grid spacing is in degrees and uniform
logging.debug(' - area')
deg2rad = np.pi / 180.0
self.area = np.where(
self.local_drainage_area <= 0,
np.absolute(
parameters.radius_earth**2 * deg2rad * self.longitude_spacing *
(np.sin(deg2rad *
(self.latitude + 0.5 * self.latitude_spacing)) -
np.sin(deg2rad *
(self.latitude - 0.5 * self.latitude_spacing)))),
self.local_drainage_area,
)
# update zero slopes to a small number
self.hillslope = np.where(self.hillslope <= 0,
parameters.hillslope_minimum, self.hillslope)
self.subnetwork_slope = np.where(self.subnetwork_slope <= 0,
parameters.subnetwork_slope_minimum,
self.subnetwork_slope)
self.channel_slope = np.where(self.channel_slope <= 0,
parameters.channel_slope_minimum,
self.channel_slope)
# load the land grid to get the land fraction; if it's not there, default to 1
# TODO need to just add this field to mosart grid file
try:
land = open_dataset(config.get('grid.land.path'))
self.land_fraction = np.array(
land[config.get('grid.land.land_fraction')]).flatten()
land.close()
except:
self.land_fraction = np.full(self.id.size, 1.0)
logging.debug(' - main channel iterations')
# parameter for calculating number of main channel iterations needed
# phi_r
phi_main = np.where(
(self.mosart_mask > 0) & (self.channel_length > 0),
self.total_drainage_area_single * np.sqrt(self.channel_slope) /
(self.channel_length * self.channel_width), 0)
# sub timesteps needed for main channel
# numDT_r
self.iterations_main_channel = np.where(
phi_main >= 10,
np.maximum(
1,
np.floor(1 + config.get('simulation.subcycles') *
np.log10(phi_main))),
np.where((self.mosart_mask > 0) & (self.channel_length > 0),
1 + config.get('simulation.subcycles'), 1))
logging.debug(' - subnetwork substeps')
# total main channel length [m]
# rlenTotal
self.total_channel_length = self.area * self.drainage_density
self.total_channel_length = np.where(
self.channel_length > self.total_channel_length,
self.channel_length, self.total_channel_length)
# hillslope length [m]
# constrain hillslope length
# there is a TODO in fortran mosart that says: "allievate the outlier in drainage density estimation."
# hlen
channel_length_minimum = np.sqrt(self.area)
hillslope_max_length = np.maximum(channel_length_minimum, 1000)
self.hillslope_length = np.where(
self.channel_length > 0,
self.area / self.total_channel_length / 2.0, 0)
self.hillslope_length = np.where(
self.hillslope_length > hillslope_max_length, hillslope_max_length,
self.hillslope_length)
# subnetwork channel length [m]
# tlen
self.subnetwork_length = np.where(
self.channel_length > 0,
np.where(
self.channel_length >= channel_length_minimum,
self.area / self.channel_length / 2.0 - self.hillslope_length,
self.area / channel_length_minimum / 2.0 -
self.hillslope_length), 0)
self.subnetwork_length = np.where(self.subnetwork_length < 0, 0,
self.subnetwork_length)
# subnetwork channel width (adjusted from input file) [m]
# twidth
self.subnetwork_width = np.where(
(self.channel_length > 0) & (self.subnetwork_width >= 0),
np.where(
(self.subnetwork_length > 0) &
((self.total_channel_length - self.channel_length) /
self.subnetwork_length > 1),
parameters.subnetwork_width_parameter * self.subnetwork_width *
((self.total_channel_length - self.channel_length) /
self.subnetwork_length), self.subnetwork_width), 0)
self.subnetwork_width = np.where(
(self.subnetwork_length > 0) & (self.subnetwork_width <= 0), 0,
self.subnetwork_width)
# parameter for calculating number of subnetwork iterations needed
# phi_t
phi_sub = np.where(
self.subnetwork_length > 0,
(self.area * np.sqrt(self.subnetwork_slope)) /
(self.subnetwork_length * self.subnetwork_width),
0,
)
# sub timesteps needed for subnetwork
# numDT_t
self.iterations_subnetwork = np.where(
self.subnetwork_length > 0,
np.where(
phi_sub >= 10,
np.maximum(
np.floor(1 + config.get('simulation.subcycles') *
np.log10(phi_sub)), 1),
np.where(self.subnetwork_length > 0,
1 + config.get('simulation.subcycles'), 1)), 1)
# if water management is enabled, load the reservoir parameters and build the grid cell mapping
# 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)
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
month_condition = state.reservoir_month_flood_control_start <= state.reservoir_month_flood_control_end
total_condition = flood_control_condition & (
(month_condition &
(month >= state.reservoir_month_flood_control_start) &
(month < state.reservoir_month_flood_control_end)) |
(np.logical_not(month_condition) &
(month >= state.reservoir_month_flood_control_start) |
(month < state.reservoir_month_flood_control_end)))
drop = 0 * state.reservoir_month_flood_control_start
n_month = 0 * drop
for m in np.arange(1, 13): # TODO assumes monthly
m_and_condition = (m >= state.reservoir_month_flood_control_start) & (
m < state.reservoir_month_flood_control_end)
m_or_condition = (m >= state.reservoir_month_flood_control_start) | (
m < state.reservoir_month_flood_control_end)
drop = np.where(
(month_condition & m_and_condition) |
(np.logical_not(month_condition) & m_or_condition),
np.where(
grid.reservoir_streamflow_schedule.sel({
streamflow_time_name: m
}).values >= grid.reservoir_streamflow_schedule.mean(
dim=streamflow_time_name).values, drop + 0, drop + np.abs(
grid.reservoir_streamflow_schedule.mean(
dim=streamflow_time_name).values -
grid.reservoir_streamflow_schedule.sel({
streamflow_time_name:
m
}).values)), drop)
n_month = np.where((month_condition & m_and_condition) |
(np.logical_not(month_condition) & m_or_condition),
n_month + 1, n_month)
state.reservoir_release = np.where(
total_condition & (n_month > 0),
state.reservoir_release + drop / n_month, state.reservoir_release)
# now need to make sure it will fill up but issue with spilling in certain hydro-climate conditions
month_condition = state.reservoir_month_flood_control_end <= state.reservoir_month_start_operations
first_condition = flood_control_condition & month_condition & (
(month >= state.reservoir_month_flood_control_end) &
(month < state.reservoir_month_start_operations))
second_condition = flood_control_condition & np.logical_not(
month_condition) & (
(month >= state.reservoir_month_flood_control_end) |
(month < state.reservoir_month_start_operations))
# TODO this logic exists in fortran mosart but isn't used...
# fill = 0 * drop
# n_month = 0 * drop
# for m in np.arange(1,13): # TODO assumes monthly
# m_condition = (m >= self.state.reservoir_month_flood_control_end.values) &
# (self.reservoir_streamflow_schedule.sel({streamflow_time_name: m}).values > self.reservoir_streamflow_schedule.mean(dim=streamflow_time_name).values) & (
# (first_condition & (m <= self.state.reservoir_month_start_operations)) |
# (second_condition & (m <= 12))
# )
# fill = np.where(
# m_condition,
# fill + np.abs(self.reservoir_streamflow_schedule.mean(dim=streamflow_time_name).values - self.reservoir_streamflow_schedule.sel({streamflow_time_name: m}).values),
# fill
# )
# n_month = np.where(
# m_condition,
# n_month + 1,
# n_month
# )
state.reservoir_release = np.where(
(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)
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(
~condition, demand_avg / 10 +
9 / 10 * flow_avg * self.reservoir_demand_schedule / demand_avg)
prerelease = prerelease.where(
condition,
prerelease.where(
~((flow_avg + self.reservoir_demand_schedule - demand_avg) > 0),
flow_avg + self.reservoir_demand_schedule - demand_avg))
self.reservoir_prerelease_schedule = prerelease
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 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
source_direct = 1.0 * state.direct_to_ocean
# wetland runoff
wetland_runoff_volume = state.hillslope_wetland_runoff * config.get('simulation.timestep') / config.get('simulation.subcycles')
river_volume_minimum = parameters.river_depth_minimum * grid.area
# if wetland runoff is negative and it would bring main channel storage below the minimum, send it directly to ocean
condition = ((state.channel_storage + wetland_runoff_volume) < river_volume_minimum) & (state.hillslope_wetland_runoff < 0)
source_direct = np.where(
condition,
source_direct + state.hillslope_wetland_runoff,
source_direct
)
state.hillslope_wetland_runoff = np.where(
condition,
0,
state.hillslope_wetland_runoff
)
# remove remaining wetland runoff (negative and positive)
source_direct = source_direct + state.hillslope_wetland_runoff
state.hillslope_wetland_runoff = 0.0 * state.zeros
# runoff from hillslope
# remove negative subsurface water
condition = state.hillslope_subsurface_runoff < 0
source_direct = np.where(
condition,
source_direct + state.hillslope_subsurface_runoff,
source_direct
)
state.hillslope_subsurface_runoff = np.where(
condition,
0,
state.hillslope_subsurface_runoff
)
# remove negative surface water
condition = state.hillslope_surface_runoff < 0
source_direct = np.where(
condition,
source_direct + state.hillslope_surface_runoff,
source_direct
)
state.hillslope_surface_runoff = np.where(
condition,
0,
state.hillslope_surface_runoff
)
# if ocean cell or ice tracer, remove the rest of the sub and surface water
# other cells will be handled by mosart euler
condition = (grid.mosart_mask == 0) | (state.tracer == parameters.ICE_TRACER)
source_direct = np.where(
condition,
source_direct + state.hillslope_subsurface_runoff + state.hillslope_surface_runoff,
source_direct
)
state.hillslope_subsurface_runoff = np.where(
condition,
0,
state.hillslope_subsurface_runoff
)
state.hillslope_surface_runoff = np.where(
condition,
0,
state.hillslope_surface_runoff
)
state.direct[:] = source_direct
# send the direct water to outlet for each tracer
state.direct[:] = pd.DataFrame(
grid.id, columns=['id']
).merge(
pd.DataFrame(
state.direct, columns=['direct']
).join(
pd.DataFrame(
grid.outlet_id, columns=['outlet_id']
)
).groupby('outlet_id').sum(),
how='left',
left_on='id',
right_index=True
).direct.fillna(0.0).values
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)
# step through max iterations, masking out the unnecessary cells each time
base_condition = (grid.mosart_mask > 0) & state.euler_mask
sub_condition = grid.subnetwork_length > grid.hillslope_length # has tributaries
for _ in np.arange(np.nanmax(grid.iterations_subnetwork)):
iteration_condition = base_condition & (grid.iterations_subnetwork > _)
state.subnetwork_flow_velocity = calculate_subnetwork_flow_velocity(
iteration_condition, sub_condition,
state.subnetwork_hydraulic_radii, grid.subnetwork_slope,
grid.subnetwork_manning, state.subnetwork_flow_velocity)
state.subnetwork_discharge = calculate_subnetwork_discharge(
iteration_condition, sub_condition, state.subnetwork_flow_velocity,
state.subnetwork_cross_section_area,
state.subnetwork_lateral_inflow, state.subnetwork_discharge)
discharge_condition = calculate_discharge_condition(
iteration_condition, sub_condition, state.subnetwork_storage,
state.subnetwork_lateral_inflow, state.subnetwork_discharge,
local_delta_t, parameters.tiny_value)
state.subnetwork_discharge = update_subnetwork_discharge(
discharge_condition, state.subnetwork_lateral_inflow,
state.subnetwork_storage, local_delta_t,
state.subnetwork_discharge)
state.subnetwork_flow_velocity = update_flow_velocity(
discharge_condition, state.subnetwork_cross_section_area,
state.subnetwork_discharge, state.subnetwork_flow_velocity)
state.subnetwork_delta_storage = calculate_subnetwork_delta_storage(
iteration_condition, state.subnetwork_lateral_inflow,
state.subnetwork_discharge, state.subnetwork_delta_storage)
# update storage
state.subnetwork_storage_previous_timestep = calculate_subnetwork_storage_previous_timestep(
iteration_condition, state.subnetwork_storage,
state.subnetwork_storage_previous_timestep)
state.subnetwork_storage = calculate_subnetwork_storage(
iteration_condition, state.subnetwork_storage,
state.subnetwork_delta_storage, local_delta_t)
update_subnetwork_state(state, grid, parameters, iteration_condition)
state.channel_lateral_flow_hillslope = calculate_channel_lateral_flow_hillslope(
iteration_condition, state.channel_lateral_flow_hillslope,
state.subnetwork_discharge)
# average lateral flow over substeps
state.channel_lateral_flow_hillslope = average_channel_lateral_flow_hillslope(
base_condition, state.channel_lateral_flow_hillslope,
grid.iterations_subnetwork)
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