def test_component_output_fields(Comp):
"""Check that required output fields exist with correct dtypes and locations"""
if Comp.name in _EXCLUDE_COMPONENTS:
pytest.skip("component explicitly excluded")
component_name = Comp._name
grid = RasterModelGrid((10, 10))
_add_input_fields_to_grid(Comp, grid)
Comp(grid)
for name, meta in Comp._info.items():
if meta["intent"].endswith("out") and not meta["optional"]:
at = meta["mapping"]
if name not in grid[at]:
raise ValueError(
f"{component_name} is missing output variable: {name} at {at}"
)
expected_dtype = meta["dtype"]
actual_dtype = grid[at][name].dtype
if actual_dtype != expected_dtype:
raise FieldError(
f"{component_name} output required variable: {name} at {at} has "
f"incorrect dtype. dtype must be {expected_dtype} and is "
f"{actual_dtype}"
)
def _wrapped(grid, vals, *args, **kwds):
"""Convert the second argument to an array."""
if isinstance(vals, str):
if vals in grid[self._at]:
vals = grid[self._at][vals]
else:
raise FieldError(vals)
else:
vals = np.asarray(vals).flatten()
return func(grid, vals, *args, **kwds)
def __init__(self, grid, channel__mask=None, **kwds):
"""Initialize the DrainageDensity component.
Parameters
----------
grid : ModelGrid
Landlab ModelGrid object
channel__mask : array, optional (default is None)
Array that holds 1's where channels exist and 0's elsewhere
"""
if channel__mask is not None:
if grid.number_of_nodes != len(channel__mask):
raise ValueError('Length of channel mask is not equal to '
'number of grid nodes')
if 'channel__mask' in grid.at_node:
warn("Existing channel__mask grid field was overwritten.")
grid.at_node['channel__mask'] = channel__mask
required = ('flow__receiver_node', 'flow__link_to_receiver_node',
'channel__mask')
for name in required:
if name not in grid.at_node:
raise FieldError(
'{name}: missing required field'.format(name=name))
# Store grid
self._grid = grid
# for this component to work with Cython acceleration,
# the channel_network must be uint8, not bool...
self.channel_network = (grid.at_node['channel__mask'].view(
dtype=np.uint8))
# Flow receivers
self.flow_receivers = grid.at_node['flow__receiver_node']
# Links to receiver nodes
self.stack_links = grid.at_node['flow__link_to_receiver_node']
# Distance to channel
try:
self.distance_to_channel = grid.at_node[
'surface_to_channel__minimum_distance']
except KeyError:
self.distance_to_channel = grid.add_zeros(
'surface_to_channel__minimum_distance', at='node', dtype=float)
def _wrapped(grid, vals, *args, **kwds):
"""Convert the second argument to an array."""
if isinstance(vals, six.string_types):
if vals in grid[self._at]:
vals = grid[self._at][vals]
else:
raise FieldError(vals)
else:
expected_size = grid.size(self._at)
vals = np.asarray(vals).flatten()
if vals.size == 1:
vals = np.broadcast_to(vals, (expected_size, ))
if vals.size != expected_size:
raise ValueError(
("Array passed to function decorated with "
"use_field_name_array_or_value is not "
"the size of fields at " + self._at))
return func(grid, vals, *args, **kwds)
def __init__(self,
grid,
channel__mask=None,
area_coefficient=None,
slope_coefficient=None,
area_exponent=None,
slope_exponent=None,
channelization_threshold=None,
**kwds):
"""Initialize the DrainageDensity component.
Parameters
----------
grid : ModelGrid
channel__mask : Array that holds 1's where
channels exist and 0's elsewhere
area_coefficient : coefficient to multiply drainage area by,
for calculating channelization threshold
slope_coefficient : coefficient to multiply slope by,
for calculating channelization threshold
area_exponent : exponent to raise drainage area to,
for calculating channelization threshold
slope_exponent : exponent to raise slope to,
for calculating channelization threshold
channelization_threshold : threshold value above
which channels exist
"""
# Store grid
self._grid = grid
for name in _REQUIRED_FIELDS:
if name not in grid.at_node:
raise FieldError(
"{name}: missing required field".format(name=name))
if grid.at_node["flow__receiver_node"].size != grid.size("node"):
msg = ("A route-to-multiple flow director has been "
"run on this grid. The landlab development team has not "
"verified that DrainageDensity is compatible with "
"route-to-multiple methods. Please open a GitHub Issue "
"to start this process.")
raise NotImplementedError(msg)
if channel__mask is not None:
if area_coefficient is not None:
warn("Channel mask and area "
"coefficient supplied. Defaulting "
"to channel mask, ignoring area "
"coefficient.")
if slope_coefficient is not None:
warn("Channel mask and slope "
"coefficient supplied. Defaulting "
"to channel mask, ignoring slope "
"coefficient.")
if area_exponent is not None:
warn("Channel mask and area "
"exponent supplied. Defaulting "
"to channel mask, ignoring area "
"exponent.")
if slope_exponent is not None:
warn("Channel mask and slope "
"exponent supplied. Defaulting "
"to channel mask, ignoring slope "
"exponent.")
if channelization_threshold is not None:
warn("Channel mask and channelization "
"threshold supplied. Defaulting "
"to channel mask, ignoring "
"threshold.")
if grid.number_of_nodes != len(channel__mask):
raise ValueError("Length of channel mask is not equal to "
"number of grid nodes")
if "channel__mask" in grid.at_node:
warn("Existing channel__mask grid field was overwritten.")
if channel__mask.dtype.type is not np.uint8:
raise ValueError("mask must by np.uint8")
self._mask_as_array = True
self._update_channel_mask = self._update_channel_mask_array
grid.at_node["channel__mask"] = channel__mask
if channel__mask is None:
if area_coefficient is None:
raise ValueError("No channel mask and no area "
"coefficient supplied. Either "
"a channel mask or all 5 threshold "
"parameters are needed.")
if slope_coefficient is None:
raise ValueError("No channel mask and no slope "
"coefficient supplied. Either "
"a channel mask or all 5 threshold "
"parameters are needed.")
if area_exponent is None:
raise ValueError("No channel mask and no area "
"exponent supplied. Either "
"a channel mask or all 5 threshold "
"parameters are needed.")
if slope_exponent is None:
raise ValueError("No channel mask and no slope "
"exponent supplied. Either "
"a channel mask or all 5 threshold "
"parameters are needed.")
if channelization_threshold is None:
raise ValueError("No channel mask and no channelization "
"threshold supplied. Either "
"a channel mask or all 5 threshold "
"parameters are needed.")
self._mask_as_array = False
self._update_channel_mask = self._update_channel_mask_values
self._area_coefficient = area_coefficient
self._slope_coefficient = slope_coefficient
self._area_exponent = area_exponent
self._slope_exponent = slope_exponent
self._channelization_threshold = channelization_threshold
self._update_channel_mask()
# for this component to work with Cython acceleration,
# the channel_network must be uint8, not bool...
self._channel_network = grid.at_node["channel__mask"]
# Flow receivers
self._flow_receivers = grid.at_node["flow__receiver_node"]
# Links to receiver nodes
self._stack_links = grid.at_node["flow__link_to_receiver_node"]
# Upstream node order
self._upstream_order = grid.at_node["flow__upstream_node_order"]
# Distance to channel
if "surface_to_channel__minimum_distance" in grid.at_node:
self.distance_to_channel = grid.at_node[
"surface_to_channel__minimum_distance"]
else:
self.distance_to_channel = grid.add_zeros(
"surface_to_channel__minimum_distance", at="node", dtype=float)
def __init__(self,
grid,
channel__mask=None,
area_coefficient=None,
slope_coefficient=None,
area_exponent=None,
slope_exponent=None,
channelization_threshold=None,
**kwds):
"""Initialize the DrainageDensity component.
Parameters
----------
grid : ModelGrid
channel__mask : Array that holds 1's where
channels exist and 0's elsewhere
area_coefficient : coefficient to multiply drainage area by,
for calculating channelization threshold
slope_coefficient : coefficient to multiply slope by,
for calculating channelization threshold
area_exponent : exponent to raise drainage area to,
for calculating channelization threshold
slope_exponent : exponent to raise slope to,
for calculating channelization threshold
channelization_threshold : threshold value above
which channels exist
"""
if (grid.at_node['flow__receiver_node'].size != grid.size('node')):
msg = ('A route-to-multiple flow director has been '
'run on this grid. The landlab development team has not '
'verified that DrainageDensity is compatible with '
'route-to-multiple methods. Please open a GitHub Issue '
'to start this process.')
raise NotImplementedError(msg)
if channel__mask is not None:
if area_coefficient is not None:
warn('Channel mask and area '
'coefficient supplied. Defaulting '
'to channel mask, ignoring area '
'coefficient.')
if slope_coefficient is not None:
warn('Channel mask and slope '
'coefficient supplied. Defaulting '
'to channel mask, ignoring slope '
'coefficient.')
if area_exponent is not None:
warn('Channel mask and area '
'exponent supplied. Defaulting '
'to channel mask, ignoring area '
'exponent.')
if slope_exponent is not None:
warn('Channel mask and slope '
'exponent supplied. Defaulting '
'to channel mask, ignoring slope '
'exponent.')
if channelization_threshold is not None:
warn('Channel mask and channelization '
'threshold supplied. Defaulting '
'to channel mask, ignoring '
'threshold.')
if grid.number_of_nodes != len(channel__mask):
raise ValueError('Length of channel mask is not equal to '
'number of grid nodes')
if 'channel__mask' in grid.at_node:
warn("Existing channel__mask grid field was overwritten.")
grid.at_node['channel__mask'] = channel__mask
if channel__mask is None:
if area_coefficient is None:
raise FieldError('No channel mask and no area '
'coefficient supplied. Either '
'a channel mask or all 5 threshold '
'parameters are needed.')
if slope_coefficient is None:
raise FieldError('No channel mask and no slope '
'coefficient supplied. Either '
'a channel mask or all 5 threshold '
'parameters are needed.')
if area_exponent is None:
raise FieldError('No channel mask and no area '
'exponent supplied. Either '
'a channel mask or all 5 threshold '
'parameters are needed.')
if slope_exponent is None:
raise FieldError('No channel mask and no slope '
'exponent supplied. Either '
'a channel mask or all 5 threshold '
'parameters are needed.')
if channelization_threshold is None:
raise FieldError('No channel mask and no channelization '
'threshold supplied. Either '
'a channel mask or all 5 threshold '
'parameters are needed.')
channel__mask = (area_coefficient * \
np.power(grid.at_node['drainage_area'], area_exponent) \
* slope_coefficient * \
np.power(grid.at_node['topographic__steepest_slope'], \
slope_exponent)) > channelization_threshold
grid.at_node['channel__mask'] = channel__mask
required = ('flow__receiver_node', 'flow__link_to_receiver_node',
'topographic__steepest_slope')
for name in required:
if name not in grid.at_node:
raise FieldError(
'{name}: missing required field'.format(name=name))
# Store grid
self._grid = grid
# for this component to work with Cython acceleration,
# the channel_network must be uint8, not bool...
self.channel_network = (grid.at_node['channel__mask'].view(
dtype=np.uint8))
# Flow receivers
self.flow_receivers = grid.at_node['flow__receiver_node']
# Links to receiver nodes
self.stack_links = grid.at_node['flow__link_to_receiver_node']
# Distance to channel
try:
self.distance_to_channel = grid.at_node[
'surface_to_channel__minimum_distance']
except KeyError:
self.distance_to_channel = grid.add_zeros(
'surface_to_channel__minimum_distance', at='node', dtype=float)
def __init__(
self,
grid,
surface="topographic__elevation",
flow_metric="D8",
runoff_rate=None,
update_flow_depressions=True,
depression_handler="fill",
exponent=1,
epsilon=True,
accumulate_flow=True,
accumulate_flow_hill=False,
separate_hill_flow=False,
update_hill_depressions=False,
update_hill_flow_instantaneous=True,
hill_flow_metric="Quinn",
hill_exponent=1,
suppress_out=True,
):
"""Initialize the FlowAccumulator component.
Saves the grid, tests grid type, tests input types and
compatibility for the flow_metric and depression_finder
keyword arguments, tests the argument of runoff, and
initializes new fields.
"""
super(PriorityFloodFlowRouter, self).__init__(grid)
# Keep a local reference to the grid
self._suppress_output = partial(
suppress_output, out=suppress_out, err=suppress_out
)
# STEP 1: Testing of input values, supplied either in function call or
# as part of the grid.
self._test_water_inputs(grid, runoff_rate)
# Grid type testing
if not isinstance(self._grid, RasterModelGrid):
raise FieldError(
"Flow Accumulator Priority flood only works with regular raster grids, "
"use default Landlab flow accumulator instead"
)
node_cell_area = self._grid.cell_area_at_node.copy()
node_cell_area[self._grid.closed_boundary_nodes] = 0.0
self._node_cell_area = node_cell_area
self._runoff_rate = runoff_rate
if (flow_metric in PSINGLE_FMs) or (flow_metric in PMULTIPLE_FMs):
self._flow_metric = flow_metric
else:
raise ValueError(
f"flow metric should be one of these single flow directors : {', '.join(PSINGLE_FMs)} or multiple flow directors: {', '.join(PMULTIPLE_FMs)}"
)
if (hill_flow_metric in PSINGLE_FMs) or (hill_flow_metric in PMULTIPLE_FMs):
self._hill_flow_metric = hill_flow_metric
else:
raise ValueError(
f"flow metric should be one of these single flow directors : {', '.join(PSINGLE_FMs)} or multiple flow directors: {', '.join(PMULTIPLE_FMs)}"
)
if depression_handler == "fill":
self._depression_handler = partial(
rd.FillDepressions, epsilon=epsilon, in_place=True
)
elif depression_handler == "breach":
self._depression_handler = partial(rd.BreachDepressions, in_place=True)
else:
raise ValueError("depression_handler should be one of 'fill' or 'breach'")
self._exponent = exponent
self._separate_hill_flow = separate_hill_flow
self._update_hill_flow_instantaneous = update_hill_flow_instantaneous
self._update_flow_depressions = update_flow_depressions
self._update_hill_depressions = update_hill_depressions
self._hill_exponent = hill_exponent
if self._separate_hill_flow:
# Adjust dict
self._info["hill_drainage_area"]["optional"] = False
self._info["hill_surface_water__discharge"]["optional"] = False
self._info["hill_flow__upstream_node_order"]["optional"] = False
self._info["hill_flow__receiver_node"]["optional"] = False
self._info["hill_topographic__steepest_slope"]["optional"] = False
self._info["hill_flow__receiver_proportions"]["optional"] = False
else:
self._info["hill_drainage_area"]["optional"] = True
self._info["hill_surface_water__discharge"]["optional"] = True
self._info["hill_flow__upstream_node_order"]["optional"] = True
self._info["hill_flow__receiver_node"]["optional"] = True
self._info["hill_topographic__steepest_slope"]["optional"] = True
self._info["hill_flow__receiver_proportions"]["optional"] = True
self._accumulate_flow = accumulate_flow
self._accumulate_flow_hill = accumulate_flow_hill
if not self._accumulate_flow:
self._info["drainage_area"]["optional"] = True
self._info["surface_water__discharge"]["optional"] = True
else:
self._info["drainage_area"]["optional"] = False
self._info["surface_water__discharge"]["optional"] = False
if not self._accumulate_flow_hill:
self._info["hill_drainage_area"]["optional"] = True
self._info["hill_surface_water__discharge"]["optional"] = True
else:
self._info["hill_drainage_area"]["optional"] = False
self._info["hill_surface_water__discharge"]["optional"] = False
self.initialize_output_fields()
# Make aliases
if self._accumulate_flow:
self._drainage_area = self.grid.at_node["drainage_area"]
self._discharges = self.grid.at_node["surface_water__discharge"]
self._sort = self.grid.at_node["flow__upstream_node_order"]
# if multiple flow algorithm is made, the dimensions of the slope and receiver fields change (8 colums for all neightbors)
if flow_metric in PMULTIPLE_FMs:
self.grid.at_node["topographic__steepest_slope"] = np.zeros(
(self.grid.number_of_nodes, 8)
)
self.grid.at_node["flow__receiver_node"] = np.zeros(
(self.grid.number_of_nodes, 8), dtype=int
)
self.grid.at_node["flow__receiver_proportions"] = np.zeros(
(self.grid.number_of_nodes, 8)
)
self.grid.at_node["flow__link_to_receiver_node"] = np.zeros(
(self.grid.number_of_nodes, 8)
)
self._slope = self.grid.at_node["topographic__steepest_slope"]
self._rcvs = self.grid.at_node["flow__receiver_node"]
self._prps = self.grid.at_node["flow__receiver_proportions"]
self._recvr_link = self.grid.at_node["flow__link_to_receiver_node"]
if self._separate_hill_flow:
if self._accumulate_flow_hill:
self._hill_drainage_area = self.grid.at_node["hill_drainage_area"]
self._hill_discharges = self.grid.at_node[
"hill_surface_water__discharge"
]
if hill_flow_metric in PMULTIPLE_FMs:
self.grid.at_node["hill_topographic__steepest_slope"] = np.zeros(
(self.grid.number_of_nodes, 8)
)
self.grid.at_node["hill_flow__receiver_node"] = np.zeros(
(self.grid.number_of_nodes, 8), dtype=int
)
self.grid.at_node["hill_flow__receiver_proportions"] = np.zeros(
(self.grid.number_of_nodes, 8)
)
self._hill_slope = self.grid.at_node["hill_topographic__steepest_slope"]
self._hill_rcvs = self.grid.at_node["hill_flow__receiver_node"]
self._hill_prps = self.grid.at_node["hill_flow__receiver_proportions"]
# Create properties specific to RichDEM
self._create_richdem_properties()
def __init__(
self,
grid,
surface="topographic__elevation",
flow_director="FlowDirectorSteepest",
runoff_rate=None,
depression_finder=None,
**kwargs
):
"""Initialize the FlowAccumulator component.
Saves the grid, tests grid type, tests imput types and
compatability for the flow_director and depression_finder
keyword arguments, tests the argument of runoff_rate, and
initializes new fields.
"""
super(FlowAccumulator, self).__init__(grid)
# Keep a local reference to the grid
# Grid type testing
self._is_raster = isinstance(self._grid, RasterModelGrid)
self._is_Voroni = isinstance(self._grid, VoronoiDelaunayGrid)
self._is_Network = isinstance(self._grid, NetworkModelGrid)
self._kwargs = kwargs
# STEP 1: Testing of input values, supplied either in function call or
# as part of the grid.
self._test_water_inputs(grid, runoff_rate)
# save elevations and node_cell_area to class properites.
self._surface = surface
self._surface_values = return_array_at_node(grid, surface)
if self._is_Network:
try:
node_cell_area = self._grid.at_node["cell_area_at_node"]
except FieldError:
raise FieldError(
"In order for the FlowAccumulator to work, the "
"grid must have an at-node field called "
"cell_area_at_node."
)
else:
node_cell_area = self._grid.cell_area_at_node.copy()
node_cell_area[self._grid.closed_boundary_nodes] = 0.0
self._node_cell_area = node_cell_area
# STEP 2:
# identify Flow Director method, save name, import and initialize the correct
# flow director component if necessary
self._add_director(flow_director)
self._add_depression_finder(depression_finder)
# This component will track of the following variables.
# Attempt to create each, if they already exist, assign the existing
# version to the local copy.
# - drainage area at each node
# - receiver of each node
# - delta array
self.initialize_output_fields()
self._drainage_area = grid.at_node["drainage_area"]
self._discharges = grid.at_node["surface_water__discharge"]
self._upstream_ordered_nodes = grid.at_node["flow__upstream_node_order"]
if np.all(self._upstream_ordered_nodes == 0):
self._upstream_ordered_nodes.fill(self._grid.BAD_INDEX)
self._delta_structure = grid.at_node["flow__data_structure_delta"]
if np.all(self._delta_structure == 0):
self._delta_structure[:] = self._grid.BAD_INDEX
self._D_structure = self._grid.BAD_INDEX * grid.ones(at="link", dtype=int)
self._nodes_not_in_stack = True
if len(self._kwargs) > 0:
kwdstr = " ".join(list(self._kwargs.keys()))
raise ValueError(
"Extra kwargs passed to FlowAccumulator:{kwds}".format(kwds=kwdstr)
)
def __init__(self,
grid,
surface="topographic__elevation",
flow_director="FlowDirectorSteepest",
runoff_rate=None,
depression_finder=None,
**kwargs):
"""Initialize the FlowAccumulator component.
Saves the grid, tests grid type, tests imput types and
compatability for the flow_director and depression_finder
keyword arguments, tests the argument of runoff_rate, and
initializes new fields.
"""
super(FlowAccumulator, self).__init__(grid)
# Keep a local reference to the grid
self._grid = grid
# Grid type testing
self._is_raster = isinstance(self._grid, RasterModelGrid)
self._is_Voroni = isinstance(self._grid, VoronoiDelaunayGrid)
self._is_Network = isinstance(self._grid, NetworkModelGrid)
self.kwargs = kwargs
# STEP 1: Testing of input values, supplied either in function call or
# as part of the grid.
self._test_water_inputs(grid, runoff_rate)
# save elevations and node_cell_area to class properites.
self.surface = surface
self.surface_values = return_array_at_node(grid, surface)
if self._is_Network:
try:
node_cell_area = self._grid.at_node["cell_area_at_node"]
except FieldError:
raise FieldError(
"In order for the FlowAccumulator to work, the "
"grid must have an at-node field called "
"cell_area_at_node.")
else:
node_cell_area = self._grid.cell_area_at_node.copy()
node_cell_area[self._grid.closed_boundary_nodes] = 0.
self.node_cell_area = node_cell_area
# STEP 2:
# identify Flow Director method, save name, import and initialize the correct
# flow director component if necessary
self._add_director(flow_director)
self._add_depression_finder(depression_finder)
# This component will track of the following variables.
# Attempt to create each, if they already exist, assign the existing
# version to the local copy.
# - drainage area at each node
# - receiver of each node
# - D array
# - delta array
# - missing nodes in stack.
if "drainage_area" not in grid.at_node:
self.drainage_area = grid.add_zeros("drainage_area",
at="node",
dtype=float)
else:
self.drainage_area = grid.at_node["drainage_area"]
if "surface_water__discharge" not in grid.at_node:
self.discharges = grid.add_zeros("surface_water__discharge",
at="node",
dtype=float)
else:
self.discharges = grid.at_node["surface_water__discharge"]
if "flow__upstream_node_order" not in grid.at_node:
self.upstream_ordered_nodes = grid.add_field(
"flow__upstream_node_order",
BAD_INDEX_VALUE * grid.ones(at="node", dtype=int),
at="node",
dtype=int,
)
else:
self.upstream_ordered_nodes = grid.at_node[
"flow__upstream_node_order"]
if "flow__data_structure_delta" not in grid.at_node:
self.delta_structure = grid.add_field(
"flow__data_structure_delta",
BAD_INDEX_VALUE * grid.ones(at="node", dtype=int),
at="node",
dtype=int,
)
else:
self.delta_structure = grid.at_node["flow__data_structure_delta"]
try:
D = BAD_INDEX_VALUE * grid.ones(at="link", dtype=int)
D_structure = np.array([D], dtype=object)
self.D_structure = grid.add_field(
"flow__data_structure_D",
D_structure,
at="grid",
dtype=object,
noclobber=False,
)
except FieldError:
self.D_structure = grid.at_grid["flow__data_structure_D"]
self.nodes_not_in_stack = True
def calculate_flow__distance(grid, add_to_grid=False, noclobber=True):
"""
Calculate the along flow distance from node to outlet.
This utility calculates the along flow distance based on the results of
running flow accumulation on the grid. It will use the connectivity
used by the FlowAccumulator (e.g. D4, D8, Dinf).
Parameters
----------
grid : ModelGrid
add_to_grid : boolean, optional
Flag to indicate if the stream length field should be added to the
grid. Default is False. The field name used is ``flow__distance``.
noclobber : boolean, optional
Flag to indicate if adding the field to the grid should not clobber an
existing field with the same name. Default is True.
Returns
-------
flow__distance : float ndarray
The distance that has to be covered from an imaginary flow, located in
each node of the grid, to reach the watershed's outlet.
Examples
--------
>>> from landlab import RasterModelGrid
>>> from landlab.components import FlowAccumulator
>>> from landlab.utils.flow__distance import calculate_flow__distance
>>> mg = RasterModelGrid((5, 4), spacing=(1, 1))
>>> elev = np.array([0., 0., 0., 0.,
... 0., 21., 10., 0.,
... 0., 31., 20., 0.,
... 0., 32., 30., 0.,
... 0., 0., 0., 0.])
>>> _ = mg.add_field('node','topographic__elevation', elev)
>>> mg.set_closed_boundaries_at_grid_edges(bottom_is_closed=True,
... left_is_closed=True,
... right_is_closed=True,
... top_is_closed=True)
>>> fr = FlowAccumulator(mg, flow_director = 'D8')
>>> fr.run_one_step()
>>> flow__distance = calculate_flow__distance(mg, add_to_grid=True, noclobber=False)
>>> mg.at_node['flow__distance']
array([ 0. , 0. , 0. , 0. ,
0. , 1. , 0. , 0. ,
0. , 1.41421356, 1. , 0. ,
0. , 2.41421356, 2. , 0. ,
0. , 0. , 0. , 0. ])
Now, let's change to D4 the flow_director method, which does not
consider diagonal links bewtween nodes.
>>> from landlab import RasterModelGrid
>>> from landlab.components import FlowAccumulator
>>> from landlab.utils.flow__distance import calculate_flow__distance
>>> mg = RasterModelGrid((5, 4), spacing=(1, 1))
>>> elev = np.array([0., 0., 0., 0.,
... 0., 21., 10., 0.,
... 0., 31., 20., 0.,
... 0., 32., 30., 0.,
... 0., 0., 0., 0.])
>>> _ = mg.add_field('node','topographic__elevation', elev)
>>> mg.set_closed_boundaries_at_grid_edges(bottom_is_closed=True,
... left_is_closed=True,
... right_is_closed=True,
... top_is_closed=True)
>>> fr = FlowAccumulator(mg, flow_director = 'D4')
>>> fr.run_one_step()
>>> flow__distance = calculate_flow__distance(mg, add_to_grid=True,
... noclobber=False)
>>> mg.at_node['flow__distance']
array([ 0., 0., 0., 0.,
0., 1., 0., 0.,
0., 2., 1., 0.,
0., 3., 2., 0.,
0., 0., 0., 0.])
The flow__distance utility can also work on irregular grids. For the example we
will use a Hexagonal Model Grid, a special type of Voroni Grid that has
regularly spaced hexagonal cells.
>>> from landlab import HexModelGrid, FIXED_VALUE_BOUNDARY, CLOSED_BOUNDARY
>>> from landlab.components import FlowAccumulator
>>> from landlab.utils.flow__distance import calculate_flow__distance
>>> dx = 1
>>> hmg = HexModelGrid(5,3, dx)
>>> _ = hmg.add_field('topographic__elevation',
... hmg.node_x + np.round(hmg.node_y),
... at = 'node')
>>> hmg.status_at_node[hmg.boundary_nodes] = CLOSED_BOUNDARY
>>> hmg.status_at_node[0] = FIXED_VALUE_BOUNDARY
>>> fr = FlowAccumulator(hmg, flow_director = 'D4')
>>> fr.run_one_step()
>>> flow__distance = calculate_flow__distance(hmg,
... add_to_grid=True,
... noclobber=False)
>>> hmg.at_node['flow__distance']
array([ 0., 0., 0.,
0., 1., 2., 0.,
0., 2., 2., 3., 0.,
0., 3., 3., 0.,
0., 0., 0.])
"""
# check that flow__receiver nodes exists
if "flow__receiver_node" not in grid.at_node:
raise FieldError("A 'flow__receiver_node' field is required at the "
"nodes of the input grid.")
if "flow__upstream_node_order" not in grid.at_node:
raise FieldError(
"A 'flow__upstream_node_order' field is required at the "
"nodes of the input grid.")
# get the reciever nodes, depending on if this is to-one, or to-multiple,
# we'll need to get a different at-node field.
if grid.at_node["flow__receiver_node"].size != grid.size("node"):
to_one = False
else:
to_one = True
flow__receiver_node = grid.at_node["flow__receiver_node"]
# get the upstream node order
flow__upstream_node_order = grid.at_node["flow__upstream_node_order"]
# get downstream flow link lengths, result depends on type of grid.
if isinstance(grid, RasterModelGrid):
flow_link_lengths = grid.length_of_d8[
grid.at_node["flow__link_to_receiver_node"]]
else:
flow_link_lengths = grid.length_of_link[
grid.at_node["flow__link_to_receiver_node"]]
# create an array that representes the outlet lengths.
flow__distance = np.zeros(grid.nodes.size)
# iterate through the flow__upstream_node_order, this will already have
# identified the locations of the outlet nodes and have
for node in flow__upstream_node_order:
# get flow recievers
reciever = flow__receiver_node[node]
# assess if this is a to one (D8/D4) or to multiple (Dinf, MFD)
# flow directing method.
if to_one:
potential_outlet = reciever
else:
# if this is an outlet, the first element of the recievers will be
# the nodes ID.
potential_outlet = reciever[0]
# assess if this is an outlet or not.
if potential_outlet == node:
not_outlet = False
else:
not_outlet = True
# if not an outlet
if not_outlet:
# deal with the two cases of route to one and route to multiple.
if to_one:
# get the stream length of the downstream node
downstream_stream_length = flow__distance[reciever]
# get the stream segment length from this node to its downstream
# neigbor
stream_increment_length = flow_link_lengths[node]
else:
# non-existant links are coded with -1
useable_recievers = np.where(reciever != BAD_INDEX_VALUE)[0]
# we will have the stream flow to the downstream node with the
# shortest distance to the outlet.
# in the event of a tie, we will choose the shorter link length.
# get the flow distances of the downstream nodes
potential_downstream_stream_lengths = flow__distance[
flow__receiver_node[node]][useable_recievers]
# get the stream segment lengths from this node to its downstream
# neighbor
potential_stream_increment_lengths = flow_link_lengths[node][
useable_recievers]
# get the lowest downstream stream length.
downstream_stream_length = np.min(
potential_downstream_stream_lengths)
# determine which of the stream increments flowed to this
# downstream neighbor.
which_link = np.where(potential_downstream_stream_lengths ==
downstream_stream_length)[0]
# and choose the smallest of these links.
stream_increment_length = np.min(
potential_stream_increment_lengths[which_link])
# set the total stream length of this node
flow__distance[
node] = downstream_stream_length + stream_increment_length
# store on the grid
if add_to_grid:
grid.add_field("node",
"flow__distance",
flow__distance,
noclobber=noclobber)
return flow__distance
def get_watershed_mask(grid, outlet_id):
"""
Get the watershed of an outlet returned as a boolean array.
Parameters
----------
grid : RasterModelGrid
A landlab RasterModelGrid.
outlet_id : integer
The id of the outlet node.
Returns
-------
watershed_mask : boolean ndarray
True elements of this array correspond to nodes with flow that is
received by the outlet. The length of the array is equal to the grid
number of nodes.
Examples
--------
>>> import numpy as np
>>> from landlab import RasterModelGrid
>>> from landlab.components import FlowAccumulator
>>> from landlab.utils import get_watershed_mask
>>> rmg = RasterModelGrid((7, 7), 1)
>>> z = np.array([
... -9999., -9999., -9999., -9999., -9999., -9999., -9999.,
... -9999., 26., 0., 30., 32., 34., -9999.,
... -9999., 28., 1., 25., 28., 32., -9999.,
... -9999., 30., 3., 3., 11., 34., -9999.,
... -9999., 32., 11., 25., 18., 38., -9999.,
... -9999., 34., 32., 34., 36., 40., -9999.,
... -9999., -9999., -9999., -9999., -9999., -9999., -9999.])
>>> rmg.at_node['topographic__elevation'] = z
Only the bottom boundary is set to open.
>>> rmg.set_closed_boundaries_at_grid_edges(True, True, True, False)
>>> rmg.set_fixed_value_boundaries_at_grid_edges(False, False, False, True)
Route flow.
>>> fr = FlowAccumulator(rmg, flow_director='D8')
>>> fr.run_one_step()
>>> get_watershed_mask(rmg, 2)
array([False, False, True, False, False, False, False, False, False,
True, False, False, False, False, False, True, True, True,
True, True, False, False, True, True, True, True, True,
False, False, True, True, True, True, True, False, False,
True, True, True, True, True, False, False, False, False,
False, False, False, False], dtype=bool)
"""
if 'flow__receiver_node' not in grid.at_node:
raise FieldError("A 'flow__receiver_node' field is required at the "
"nodes of the input grid.")
if (grid.at_node['flow__receiver_node'].size != grid.size('node')):
msg = ('A route-to-multiple flow director has been '
'run on this grid. The landlab development team has not '
'verified that get_watershed_mask is compatible with '
'route-to-multiple methods. Please open a GitHub Issue '
'to start this process.')
raise NotImplementedError(msg)
grid_nodes = grid.nodes.flatten()
receiver_at_node = grid.at_node['flow__receiver_node']
# Prepare output.
watershed_mask = np.zeros(grid.number_of_nodes, dtype=bool)
for node in grid_nodes:
# Follow flow path of each node.
receiver_node = receiver_at_node[node]
outlet_not_found = True
while outlet_not_found:
node_flows_to_outlet = any([receiver_node == outlet_id,
receiver_at_node[receiver_node] ==
outlet_id])
node_is_outlet = node == outlet_id
if node_flows_to_outlet or node_is_outlet:
watershed_mask[node] = True
outlet_not_found = False
else:
receiver_node = receiver_at_node[receiver_node]
if receiver_node == receiver_at_node[receiver_node]:
# Receiver_node is a pit.
outlet_not_found = False
return watershed_mask
def get_watershed_outlet(grid, source_node_id):
"""
Get the downstream-most node (the outlet) of the source node.
Parameters
----------
grid : RasterModelGrid
A landlab RasterModelGrid.
source_node_id : integer
The id of the node in which to identify its outlet.
Returns
-------
outlet_node : integer
The id of the node that is the downstream-most node (the outlet) of the
source node.
Examples
--------
>>> import numpy as np
>>> from landlab import RasterModelGrid
>>> from landlab.components import FlowAccumulator
>>> from landlab.utils import get_watershed_outlet
>>> rmg = RasterModelGrid((7, 7), 1)
>>> z = np.array([
... -9999., -9999., -9999., -9999., -9999., -9999., -9999.,
... -9999., 26., 0., 30., 32., 34., -9999.,
... -9999., 28., 1., 25., 28., 32., -9999.,
... -9999., 30., 3., 3., 11., 34., -9999.,
... -9999., 32., 11., 25., 18., 38., -9999.,
... -9999., 34., 32., 34., 36., 40., -9999.,
... -9999., -9999., -9999., -9999., -9999., -9999., -9999.])
>>> rmg.at_node['topographic__elevation'] = z
>>> imposed_outlet = 2
>>> rmg.set_watershed_boundary_condition_outlet_id(imposed_outlet, z,
... nodata_value=-9999.)
Route flow.
>>> fr = FlowAccumulator(rmg, flow_director='D8')
>>> fr.run_one_step()
Get the grid watershed outlet.
>>> determined_outlet = get_watershed_outlet(rmg, 40)
>>> determined_outlet == imposed_outlet
True
"""
if 'flow__receiver_node' not in grid.at_node:
raise FieldError("A 'flow__receiver_node' field is required at the "
"nodes of the input grid.")
if (grid.at_node['flow__receiver_node'].size != grid.size('node')):
msg = ('A route-to-multiple flow director has been '
'run on this grid. The landlab development team has not '
'verified that get_watershed_outlet is compatible with '
'route-to-multiple methods. Please open a GitHub Issue '
'to start this process.')
raise NotImplementedError(msg)
receiver_at_node = grid.at_node['flow__receiver_node']
receiver_node = receiver_at_node[source_node_id]
outlet_not_found = True
while outlet_not_found:
node_is_outlet = receiver_node == source_node_id
node_is_boundary = grid.node_is_boundary(receiver_node)
node_is_pit = receiver_node == receiver_at_node[receiver_node]
if node_is_outlet or node_is_boundary or node_is_pit:
outlet_not_found = False
outlet_node = receiver_node
else:
receiver_node = receiver_at_node[receiver_node]
return outlet_node
def calculate_distance_to_divide(
grid, longest_path=True, add_to_grid=False, noclobber=True
):
"""
Calculate the along flow distance from drainage divide to point.
This utility calculates the along flow distance based on the results of
running flow accumulation on the grid. It will use the connectivity
used by the FlowAccumulator (e.g. D4, D8, Dinf).
Parameters
----------
grid : ModelGrid
longest_path : bool, optional
Take the longest (or shortest) path to a drainage divide. Default is
true.
add_to_grid : boolean, optional
Flag to indicate if the stream length field should be added to the
grid. Default is False. The field name used is ``distance_to_divide``.
noclobber : boolean, optional
Flag to indicate if adding the field to the grid should not clobber an
existing field with the same name. Default is True.
Returns
-------
distance_to_divide : float ndarray
The distance that has to be covered from an imaginary flow, located in
each node of the grid, to reach the watershed's outlet.
Examples
--------
>>> import numpy as np
>>> from landlab import RasterModelGrid
>>> from landlab.components import FlowAccumulator
>>> from landlab.utils.distance_to_divide import (
... calculate_distance_to_divide)
>>> mg = RasterModelGrid((5, 4))
>>> elev = np.array([0., 0., 0., 0.,
... 0., 10., 10., 0.,
... 0., 20., 20., 0.,
... 0., 30., 30., 0.,
... 0., 0., 0., 0.])
>>> _ = mg.add_field('node','topographic__elevation', elev)
>>> mg.set_closed_boundaries_at_grid_edges(
... bottom_is_closed=False,
... left_is_closed=True,
... right_is_closed=True,
... top_is_closed=True)
>>> fr = FlowAccumulator(mg, flow_director = 'D8')
>>> fr.run_one_step()
>>> distance_to_divide = calculate_distance_to_divide(
... mg,
... add_to_grid=True,
... noclobber=False)
>>> mg.at_node['distance_to_divide']
array([ 0., 3., 3., 0.,
0., 2., 2., 0.,
0., 1., 1., 0.,
0., 0., 0., 0.,
0., 0., 0., 0.])
Now, let's change to MFD the flow_director method, which routes flow to
multiple nodes.
>>> from landlab import RasterModelGrid
>>> from landlab.components import FlowAccumulator
>>> from landlab.utils.distance_to_divide import (
... calculate_distance_to_divide)
>>> mg = RasterModelGrid((5, 4), xy_spacing=(1, 1))
>>> elev = np.array([0., 0., 0., 0.,
... 0., 10., 10., 0.,
... 0., 20., 20., 0.,
... 0., 30., 30., 0.,
... 0., 0., 0., 0.])
>>> _ = mg.add_field('node','topographic__elevation', elev)
>>> mg.set_closed_boundaries_at_grid_edges(
... bottom_is_closed=False,
... left_is_closed=True,
... right_is_closed=True,
... top_is_closed=True)
>>> fr = FlowAccumulator(mg, flow_director = 'MFD')
>>> fr.run_one_step()
>>> distance_to_divide = calculate_distance_to_divide(
... mg,
... add_to_grid=True,
... noclobber=False)
>>> mg.at_node['distance_to_divide']
array([ 0., 3., 3., 0.,
0., 2., 2., 0.,
0., 1., 1., 0.,
0., 0., 0., 0.,
0., 0., 0., 0.])
The distance_to_divide utility can also work on irregular grids. For the
example we will use a Hexagonal Model Grid, a special type of Voroni Grid
that has regularly spaced hexagonal cells.
>>> from landlab import HexModelGrid, FIXED_VALUE_BOUNDARY, CLOSED_BOUNDARY
>>> from landlab.components import FlowAccumulator
>>> from landlab.utils.distance_to_divide import (
... calculate_distance_to_divide)
>>> dx = 1
>>> hmg = HexModelGrid(5,3, dx)
>>> _ = hmg.add_field('topographic__elevation',
... hmg.node_x + np.round(hmg.node_y),
... at = 'node')
>>> hmg.status_at_node[hmg.boundary_nodes] = CLOSED_BOUNDARY
>>> hmg.status_at_node[0] = FIXED_VALUE_BOUNDARY
>>> fr = FlowAccumulator(hmg, flow_director = 'D4')
>>> fr.run_one_step()
>>> distance_to_divide = calculate_distance_to_divide(
... hmg,
... add_to_grid=True,
... noclobber=False)
>>> hmg.at_node['distance_to_divide']
array([ 3., 0., 0.,
0., 2., 1., 0.,
0., 1., 1., 0., 0.,
0., 0., 0., 0.,
0., 0., 0.])
"""
# check that flow__receiver nodes exists
if "flow__receiver_node" not in grid.at_node:
raise FieldError(
"A 'flow__receiver_node' field is required at the "
"nodes of the input grid."
)
if "flow__upstream_node_order" not in grid.at_node:
raise FieldError(
"A 'flow__upstream_node_order' field is required at the "
"nodes of the input grid."
)
if "drainage_area" not in grid.at_node:
raise FieldError(
"A 'flow__upstream_node_order' field is required at the "
"nodes of the input grid."
)
# get the reciever nodes, depending on if this is to-one, or to-multiple,
# we'll need to get a different at-node field.
if grid.at_node["flow__receiver_node"].size != grid.size("node"):
to_one = False
else:
to_one = True
flow__receiver_node = grid.at_node["flow__receiver_node"]
drainage_area = grid.at_node["drainage_area"]
# get the upstream node order
flow__upstream_node_order = grid.at_node["flow__upstream_node_order"]
# get downstream flow link lengths, result depends on type of grid.
if isinstance(grid, RasterModelGrid):
flow_link_lengths = grid.length_of_d8[
grid.at_node["flow__link_to_receiver_node"]
]
else:
flow_link_lengths = grid.length_of_link[
grid.at_node["flow__link_to_receiver_node"]
]
# create an array that representes the distance to the divide.
distance_to_divide = np.zeros(grid.nodes.size)
if not longest_path:
distance_to_divide[:] = 2 * grid.size("node") * np.max(flow_link_lengths)
# iterate through the flow__upstream_node_order backwards.
for node in reversed(flow__upstream_node_order):
# if drainage are is equal to node cell area, set distance to zeros
# this should handle the drainage divide cells as boundary cells have
# their area set to zero.
if drainage_area[node] == grid.cell_area_at_node[node]:
distance_to_divide[node] = 0
# get flow recievers
reciever = flow__receiver_node[node]
if to_one:
# if not processing an outlet node.
if reciever != node:
if longest_path:
cond = (
distance_to_divide[reciever]
< distance_to_divide[node] + flow_link_lengths[node]
)
else:
cond = (
distance_to_divide[reciever]
> distance_to_divide[node] + flow_link_lengths[node]
)
if cond:
distance_to_divide[reciever] = (
distance_to_divide[node] + flow_link_lengths[node]
)
else:
# non-existant links are coded with -1
useable_recievers = np.where(reciever != BAD_INDEX_VALUE)[0]
for idx in range(len(useable_recievers)):
r = reciever[useable_recievers][idx]
fll = flow_link_lengths[node][useable_recievers][idx]
# if not processing an outlet node.
if r != node:
if longest_path:
cond = distance_to_divide[r] < distance_to_divide[node] + fll
else:
cond = distance_to_divide[r] > distance_to_divide[node] + fll
if cond:
distance_to_divide[r] = distance_to_divide[node] + fll
# store on the grid
if add_to_grid:
grid.add_field(
"node", "distance_to_divide", distance_to_divide, noclobber=noclobber
)
return distance_to_divide
def test_component_metadata(Comp):
if Comp.name not in (
"ChannelProfiler",
"DrainageDensity",
"gFlex",
"HackCalculator",
"Lithology",
"LithoLayers",
"Profiler",
"SoilMoisture",
"Vegetation",
):
print(Comp.name)
grid = RasterModelGrid((10, 10))
# verify that we can create it
for name in Comp._info.keys():
if "in" in Comp._info[name]["intent"]:
at = Comp.var_loc(name)
dtype = Comp.var_type(name)
if at == "grid":
grid.at_grid[name] = 0
else:
grid.add_zeros(at, name, dtype=dtype)
_ = Comp(grid)
# verify that all output fields are made
for name in Comp._info.keys():
if "out" in Comp._info[name]["intent"]:
if not Comp._info[name]["optional"]:
at = Comp._info[name]["mapping"]
if name not in grid[at]:
raise ValueError(
"{component} is missing output variable: {name} at {at}".format(
component=Comp._name, name=name, at=at
)
)
field = grid[at][name]
dtype = Comp._info[name]["dtype"]
try:
assert field.dtype == dtype
except AssertionError:
raise FieldError(
"{component} output required variable: {name} at {at} has incorrect dtype. dtype must be {dtype} and is {actual}".format(
component=Comp._name,
name=name,
at=at,
dtype=dtype,
actual=field.dtype,
)
)
# verify all info exist:
for name in Comp._info.keys():
info = Comp._info[name].copy()
at = Comp._info[name]["mapping"]
for attribute in _REQ_ATTRS:
if attribute in info:
info.pop(attribute)
else:
raise ValueError(
"{component} is missing attribute {attribute} about variable: {name} at {at}".format(
component=Comp._name, name=name, at=at, attribute=attribute
)
)
if len(info) > 0:
raise ValueError(
"{component} has an extra attribute {attribute} about variable: {name} at {at}".format(
component=Comp._name, name=name, at=at, attribute=attribute
)
)
# TODO: Verify that all units are UDUNITS compatible.
# TODO: Verify that all dtypes are valid.
# TODO: Verify that all mappings are valid grid locations.
if Comp._info[name]["mapping"] not in _VALID_LOCS:
raise ValueError(
"{component} mapping for variable: {name} is invalid: {at}".format(
component=Comp._name, name=name, at=at, attribute=attribute
)
)
def get_watershed_outlet(grid, source_node_id):
"""
Get the downstream-most node (the outlet) of the source node.
Parameters
----------
grid : RasterModelGrid
A landlab RasterModelGrid.
source_node_id : integer
The id of the node in which to identify its outlet.
Returns
-------
outlet_node : integer
The id of the node that is the downstream-most node (the outlet) of the
source node.
Examples
--------
>>> import numpy as np
>>> from landlab import RasterModelGrid
>>> from landlab.components import FlowRouter
>>> from landlab.utils import get_watershed_outlet
>>> rmg = RasterModelGrid((7, 7), 1)
>>> z = np.array([
... -9999., -9999., -9999., -9999., -9999., -9999., -9999.,
... -9999., 26., 0., 30., 32., 34., -9999.,
... -9999., 28., 1., 25., 28., 32., -9999.,
... -9999., 30., 3., 3., 11., 34., -9999.,
... -9999., 32., 11., 25., 18., 38., -9999.,
... -9999., 34., 32., 34., 36., 40., -9999.,
... -9999., -9999., -9999., -9999., -9999., -9999., -9999.])
>>> rmg.at_node['topographic__elevation'] = z
>>> imposed_outlet = 2
>>> rmg.set_watershed_boundary_condition_outlet_id(imposed_outlet, z,
... nodata_value=-9999.)
Route flow.
>>> fr = FlowRouter(rmg)
>>> fr.run_one_step()
Get the grid watershed outlet.
>>> determined_outlet = get_watershed_outlet(rmg, 40)
>>> determined_outlet == imposed_outlet
True
"""
if 'flow__receiver_node' not in grid.at_node:
raise FieldError("A 'flow__receiver_node' field is required at the "
"nodes of the input grid.")
receiver_at_node = grid.at_node['flow__receiver_node']
receiver_node = receiver_at_node[source_node_id]
outlet_not_found = True
while outlet_not_found:
node_is_outlet = receiver_node == source_node_id
node_is_boundary = grid.node_is_boundary(receiver_node)
node_is_pit = receiver_node == receiver_at_node[receiver_node]
if node_is_outlet or node_is_boundary or node_is_pit:
outlet_not_found = False
outlet_node = receiver_node
else:
receiver_node = receiver_at_node[receiver_node]
return outlet_node
def get_watershed_mask(grid, outlet_id):
"""Get the watershed of an outlet returned as a boolean array.
Parameters
----------
grid : RasterModelGrid
A landlab RasterModelGrid.
outlet_id : integer
The id of the outlet node.
Returns
-------
watershed_mask : boolean ndarray
True elements of this array correspond to nodes with flow that is
received by the outlet. The length of the array is equal to the grid
number of nodes.
Examples
--------
>>> import numpy as np
>>> from landlab import RasterModelGrid
>>> from landlab.components import FlowAccumulator
>>> from landlab.utils import get_watershed_mask
>>> rmg = RasterModelGrid((7, 7))
>>> z = np.array([
... -9999., -9999., -9999., -9999., -9999., -9999., -9999.,
... -9999., 26., 0., 30., 32., 34., -9999.,
... -9999., 28., 1., 25., 28., 32., -9999.,
... -9999., 30., 3., 3., 11., 34., -9999.,
... -9999., 32., 11., 25., 18., 38., -9999.,
... -9999., 34., 32., 34., 36., 40., -9999.,
... -9999., -9999., -9999., -9999., -9999., -9999., -9999.])
>>> rmg.at_node['topographic__elevation'] = z
Only the bottom boundary is set to open.
>>> rmg.set_closed_boundaries_at_grid_edges(True, True, True, False)
>>> rmg.set_fixed_value_boundaries_at_grid_edges(False, False, False, True)
Route flow.
>>> fr = FlowAccumulator(rmg, flow_director='D8')
>>> fr.run_one_step()
>>> get_watershed_mask(rmg, 2).reshape(rmg.shape)
array([[False, False, True, False, False, False, False],
[False, False, True, False, False, False, False],
[False, True, True, True, True, True, False],
[False, True, True, True, True, True, False],
[False, True, True, True, True, True, False],
[False, True, True, True, True, True, False],
[False, False, False, False, False, False, False]], dtype=bool)
"""
if "flow__receiver_node" not in grid.at_node:
raise FieldError("A 'flow__receiver_node' field is required at the "
"nodes of the input grid.")
if grid.at_node["flow__receiver_node"].size != grid.size("node"):
msg = ("A route-to-multiple flow director has been "
"run on this grid. The landlab development team has not "
"verified that get_watershed_mask is compatible with "
"route-to-multiple methods. Please open a GitHub Issue "
"to start this process.")
raise NotImplementedError(msg)
receiver_at_node = grid.at_node["flow__receiver_node"]
upstream_node_order = grid.at_node["flow__upstream_node_order"]
# Prepare output.
watershed_mask = np.zeros(grid.number_of_nodes, dtype=bool)
# loop through all nodes once based on upstream node order. This means we
# only need to loop through the nodes once.
for node in upstream_node_order:
# when the outlet_id is encountered, mark it as true, and set
# outlet_found to True.
if node == outlet_id:
watershed_mask[node] = True
# once the outlet is found, set the watershed mask to the value of
# the reciever at node, this will paint the watershed in as we move
# upstream.
if watershed_mask[receiver_at_node[node]]:
watershed_mask[node] = True
return watershed_mask
