Implement StreamObject.removeedgeeffects

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.

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])