diff --git a/src/topotoolbox/stream_object.py b/src/topotoolbox/stream_object.py index 44c6135..fe08b7e 100644 --- a/src/topotoolbox/stream_object.py +++ b/src/topotoolbox/stream_object.py @@ -1204,6 +1204,44 @@ def downstreamto(self, nodes) -> 'StreamObject': return self.subgraph(nal) + def removeedgeeffects(self, fd, dem=None): + """Remove segments of the stream network downstream of DEM boundaries + + Parameters + ---------- + fd: FlowObject + + dem: GridObject + If a GridObject is supplied, NaNs in the GridObject are + additionally considered as boundary pixels. + + Raises + ------ + ValueError + The supplied FlowObject is not aligned to the StreamObject. + + """ + + if not validate_alignment(self, fd): + err = "The provided StreamObject and FlowObject are not aligned" + raise ValueError(err) + + # Create a mask of the boundary pixels + mask = np.ones(self.shape, dtype=bool, order=fd.order) + mask[1:-1, 1:-1] = False + + # Add NaNs to the mask. We have to dilate to find pixels that + # border NaNs. + if dem is not None: + nans = dem.duplicate_with_new_data(np.isnan(dem)) + nans = nans.dilate(size=(3, 3)) + mask = np.logical_or(mask, nans) + + # Propagate the mask downstream + i = fd.influencemap(mask) + + return self.upstreamto(np.logical_not(i)) + def gradient(self, dem, impose=False) -> 'np.ndarray': """Calculates the stream slope for each node in the stream network S based on the associated digital elevation model DEM. diff --git a/tests/test_stream_object.py b/tests/test_stream_object.py index 856d506..c02b628 100644 --- a/tests/test_stream_object.py +++ b/tests/test_stream_object.py @@ -958,3 +958,37 @@ def test_distance(cs): d = cs.distance('node_to_node') assert np.array_equal(d[cs.source], cs.node_to_node_distance()) + + +def test_removeedgeeffects(cs, cfd): + cs2 = cs.removeedgeeffects(cfd) + + mask = np.ones(cfd.shape, dtype=bool) + mask[1:-1, 1:-1] = False + + cz = cfd.influencemap(mask) + assert np.any(cz[cs.node_indices]) + assert not np.any(cz[cs2.node_indices]) + + +def test_removeedgeeffects_order(cs, cfd, fs, ffd): + cs2 = cs.removeedgeeffects(cfd) + fs2 = fs.removeedgeeffects(ffd) + + assert isequivalent(cs2, fs2) + +def test_removeedgeeffects_nans(wide_dem): + dem = wide_dem.duplicate_with_new_data(np.array(wide_dem, copy=True)) + + dem.z[dem.z == np.max(dem)] = np.nan + + fd = topo.FlowObject(dem) + s = topo.StreamObject(fd, threshold=1) + s2 = s.removeedgeeffects(fd, dem) + + mask = np.isnan(dem) + i = fd.influencemap(mask) + + assert np.any(i) + assert np.any(i[s.node_indices]) + assert not np.any(i[s2.node_indices])