def reset_elevations(seg):
"""Reset segment elevations above (upsegs) and below (outseg) a node.
"""
oseg = graph[seg]
all_upsegs = np.array(list(get_upsegs(graph_r, seg)) +
[seg]) # all segments upstream of node
elevmin_s = np.min([elevations[s] for s in all_upsegs
]) # minimum current elevation upstream of node
oldmin_s = elevations[seg]
elevs = [elevmin_s, oldmin_s]
if oseg > 0: # if segment is not an outlet,
if start_elevations is not None:
elevs.append(
elevmax[oseg]) # outseg start elevation (already updated)
# set segment end elevation as min of
# upstream elevations, current elevation, outseg start elevation
elevations[seg] = np.min(elevs)
# if the node is not an outlet, reset the outseg max if the current min is lower
if oseg > 0:
if start_elevations is not None:
next_reach_elev = elevmax[oseg]
elevmax[graph[seg]] = np.min([elevmin_s, next_reach_elev])
else:
next_reach_elev = elevations[oseg]
elevations[graph[seg]] = np.min([elevmin_s, next_reach_elev])
def arbolate_sum(segment, lengths, routing):
"""Compute the total length of all tributaries upstream from
segment, including that segment, using the supplied lengths and
routing connections.
Parameters
----------
segment : int or list of ints
Segment or node number that is also a key in the lengths
and routing dictionaries.
lengths : dict
Dictionary of lengths keyed by segment (node) numbers,
including those in segment.
routing : dict
Dictionary describing routing connections between
segments (nodes); values represent downstream connections.
Returns
-------
asum : float or dict
Arbolate sums for each segment.
"""
scalar = False
if np.isscalar(segment):
scalar = True
segment = [segment]
graph_r = make_graph(list(routing.values()), list(routing.keys()))
asum = {}
for s in segment:
upsegs = get_upsegs(graph_r, s)
lnupsegs = [lengths[s] for s in upsegs]
asum[s] = np.sum(lnupsegs) + lengths[s]
return asum
def test_get_upsegs(sfr_test_numbering):
rd, sd = sfr_test_numbering
graph = dict(zip(sd.nseg, sd.outseg))
graph_r = make_graph(list(graph.values()), list(graph.keys()))
upsegs = []
for s in sd.nseg:
upsegs.append(get_upsegs(graph_r, s))
assert upsegs[1] == {1, 3, 4, 5, 6, 7, 8, 9}
def arbolate_sum(segment, lengths, routing, starting_asums=None):
"""Compute the total length of all tributaries upstream from
segment, including that segment, using the supplied lengths and
routing connections.
Parameters
----------
segment : int or list of ints
Segment or node number that is also a key in the lengths
and routing dictionaries.
lengths : dict
Dictionary of lengths keyed by segment (node) numbers,
including those in segment.
routing : dict
Dictionary describing routing connections between
segments (nodes); values represent downstream connections.
starting_asums : dict
Option to supply starting arbolate sum values for any of
the segments. By default, None.
Returns
-------
asum : float or dict
Arbolate sums for each segment.
"""
if np.isscalar(segment):
segment = [segment]
graph_r = make_graph(list(routing.values()), list(routing.keys()))
asum = {}
for s in segment:
upsegs = get_upsegs(graph_r, s)
lnupsegs = [lengths[us] for us in upsegs]
upstream_starting_asums = [0.]
segment_starting_asum = 0.
if starting_asums is not None:
upstream_starting_asums = [
starting_asums.get(us, 0.) for us in upsegs
]
segment_starting_asum = starting_asums.get(s, 0.)
asum[s] = np.sum(lnupsegs) + lengths[s] + np.sum(
upstream_starting_asums) + segment_starting_asum
return asum
