Merge pull request #80 from mdbartos/rfsm

Add profile extraction

pysheds/grid.py

@@ -1755,7 +1755,60 @@ def compute_hand(self, fdir, dem, drainage_mask, out_name='hand', dirmap=None,
         assert (np.asarray(dem.shape) == np.asarray(fdir.shape)).all()
         assert (np.asarray(dem.shape) == np.asarray(mask.shape)).all()
         if routing.lower() == 'dinf':
-            raise NotImplementedError('Only implemented for D8 routing.')
+            try:
+                # Split dinf flowdir
+                fdir_0, fdir_1, prop_0, prop_1 = self.angle_to_d8(fdir, dirmap=dirmap)
+                # Find invalid cells
+                invalid_cells = ((fdir < 0) | (fdir > (np.pi * 2)))
+                # Pad the rim
+                dirleft_0, dirright_0, dirtop_0, dirbottom_0 = self._pop_rim(fdir_0,
+                                                                            nodata=nodata_in_fdir)
+                dirleft_1, dirright_1, dirtop_1, dirbottom_1 = self._pop_rim(fdir_1,
+                                                                            nodata=nodata_in_fdir)
+                maskleft, maskright, masktop, maskbottom = self._pop_rim(mask, nodata=0)
+                mask = mask.ravel()
+                # Ensure proportion of flow is never zero
+                fdir_0.flat[prop_0 == 0] = fdir_1.flat[prop_0 == 0]
+                fdir_1.flat[prop_1 == 0] = fdir_0.flat[prop_1 == 0]
+                # Set nodata cells to zero
+                fdir_0[invalid_cells] = 0
+                fdir_1[invalid_cells] = 0
+                # Create indexing arrays for convenience
+                visited = np.zeros(fdir.size, dtype=np.bool)
+                # nvisited = np.zeros(fdir.size, dtype=int)
+                r_dirmap = np.array(dirmap)[[4, 5, 6, 7, 0, 1, 2, 3]].tolist()
+                source = np.flatnonzero(mask)
+                hand = -np.ones(fdir.size, dtype=np.int)
+                hand[source] = source
+                visited[source] = True
+                # nvisited[source] += 1
+                for _ in range(fdir.size):
+                    selection = self._select_surround_ravel(source, fdir.shape)
+                    ix = (((fdir_0.flat[selection] == r_dirmap) |
+                          (fdir_1.flat[selection] == r_dirmap)) &
+                          (hand.flat[selection] < 0) &
+                          (~visited.flat[selection])
+                     )
+                    # TODO: Not optimized (a lot of copying here)
+                    parent = np.tile(source, (len(dirmap), 1)).T[ix]
+                    child = selection[ix]
+                    if not child.size:
+                        break
+                    visited.flat[child] = True
+                    hand[child] = hand[parent]
+                    source = np.unique(child)
+                hand = hand.reshape(dem.shape)
+                hand = np.where(hand != -1, dem - dem.flat[hand], nodata_out)
+            except:
+                raise
+            finally:
+                mask = mask.reshape(dem.shape)
+                self._replace_rim(fdir_0, dirleft_0, dirright_0, dirtop_0, dirbottom_0)
+                self._replace_rim(fdir_1, dirleft_1, dirright_1, dirtop_1, dirbottom_1)
+                self._replace_rim(mask, maskleft, maskright, masktop, maskbottom)
+            return self._output_handler(data=hand, out_name=out_name, properties=properties,
+                                        inplace=inplace, metadata=metadata)
+
         elif routing.lower() == 'd8':
             try:
                 dirleft, dirright, dirtop, dirbottom = self._pop_rim(fdir, nodata=nodata_in_fdir)
@@ -2508,6 +2561,104 @@ def to_raster(self, data_name, file_name, profile=None, view=True, blockxsize=25
         with rasterio.open(file_name, 'w', **profile) as dst:
             dst.write(np.asarray(data), 1)
 
+    def extract_profiles(self, fdir, mask, dirmap=None, nodata_in=None, routing='d8',
+                         apply_mask=True, ignore_metadata=False, **kwargs):
+        """
+        Generates river profiles from flow_direction and mask arrays.
+ 
+        Parameters
+        ----------
+        fdir : str or Raster
+               Flow direction data.
+               If str: name of the dataset to be viewed.
+               If Raster: a Raster instance (see pysheds.view.Raster)
+        mask : np.ndarray or Raster
+               Boolean array indicating channelized regions
+        dirmap : list or tuple (length 8)
+                 List of integer values representing the following
+                 cardinal and intercardinal directions (in order):
+                 [N, NE, E, SE, S, SW, W, NW]
+        nodata_in : int or float
+                     Value to indicate nodata in input array.
+        routing : str
+                  Routing algorithm to use:
+                  'd8'   : D8 flow directions
+        apply_mask : bool
+               If True, "mask" the output using self.mask.
+        ignore_metadata : bool
+                          If False, require a valid affine transform and CRS.
+ 
+        Returns
+        -------
+        profiles : np.ndarray
+                   Array of channel profiles
+        """
+        if routing.lower() != 'd8':
+            raise NotImplementedError('Only implemented for D8 routing.')
+        # TODO: If two "forks" are directly connected, it can introduce a gap
+        nodata_in = self._check_nodata_in(fdir, nodata_in)
+        fdir_props = {}
+        acc_props = {}
+        fdir = self._input_handler(fdir, apply_mask=apply_mask, nodata_view=nodata_in,
+                                   properties=fdir_props,
+                                   ignore_metadata=ignore_metadata, **kwargs)
+        try:
+            assert(fdir.shape == mask.shape)
+        except:
+            raise ValueError('Flow direction and accumulation grids not aligned.')
+        dirmap = self._set_dirmap(dirmap, fdir)
+        flat_idx = np.arange(fdir.size)
+        fdir_orig_type = fdir.dtype
+        try:
+            mintype = np.min_scalar_type(fdir.size)
+            fdir = fdir.astype(mintype)
+            flat_idx = flat_idx.astype(mintype)
+            startnodes, endnodes = self._construct_matching(fdir, flat_idx,
+                                                            dirmap=dirmap)
+            start = startnodes[mask.flat[startnodes]]
+            end = fdir.flat[start]
+            # Find nodes with indegree > 1
+            indegree = (np.bincount(end)).astype(np.uint8)
+            forks_end = np.flatnonzero(indegree > 1)
+            # Find fork nodes
+            is_fork = np.in1d(end, forks_end)
+            forks = pd.Series(end[is_fork], index=start[is_fork])
+            # Cut endnode at forks
+            endnodes[start[is_fork]] = 0
+            endnodes[0] = 0
+            end = endnodes[start]
+            no_pred = ~np.in1d(start, end)
+            start = start[no_pred]
+            end = endnodes[start]
+            ixes = []
+            ixes.append(start)
+            ixes.append(end)
+            while end.any():
+                end = endnodes[end]
+                ixes.append(end)
+            ixes = np.column_stack(ixes)
+            forkorder = pd.Series(np.arange(len(ixes)), index=ixes[:, 0])
+            profiles = []
+            connections = {}
+            for row in ixes:
+                profile = row[row != 0]
+                profile_start, profile_end = profile[0], profile[-1]
+                start_num = forkorder.at[profile_start]
+                if profile_end in forks.index:
+                    profile_end = forks.at[profile_end]
+                if profile_end in forkorder.index:
+                    end_num = forkorder.at[profile_end]
+                else:
+                    end_num = -1
+                profiles.append(profile)
+                connections.update({start_num : end_num})
+        except:
+            raise
+        finally:
+            self._unflatten_fdir(fdir, flat_idx, dirmap)
+            fdir = fdir.astype(fdir_orig_type)
+        return profiles, connections
+
     def extract_river_network(self, fdir, acc, threshold=100,
                               dirmap=None, nodata_in=None, routing='d8',
                               apply_mask=True, ignore_metadata=False, **kwargs):

pysheds/rfsm.py

@@ -7,14 +7,16 @@
 from itertools import combinations, chain
 
 class RFSM:
-    def __init__(self, dem, max_levels=100, max_spills=100):
+    def __init__(self, dem, max_levels=100, max_spills=100, min_size=0, boundary=None):
         self.dem = dem
         if isinstance(dem, Raster):
             self.x, self.y = abs(dem.affine.a), abs(dem.affine.e)
         else:
             self.x, self.y = 1, 1
         self.max_levels = max_levels
         self.max_spills = max_spills
+        self.min_size = min_size
+        self.boundary = boundary
         self.shape = dem.shape
         self.size = dem.size
         self.exit_node = Node(name='exit', vol=np.inf, cumulative_vol=np.inf)
@@ -39,22 +41,34 @@ def compute_waterlevel(self, input_vol):
 
     def find_depressions(self):
         # Identify depressions in DEM
-        ix = np.arange(self.size).reshape(self.shape)
-        top = ix[0, :]
-        bottom = ix[-1, :]
-        left = ix[:, 0]
-        right = ix[:, -1]
-        # Rotate the barrier to ensure that all depressions touching edges are captured
-        # TODO: This may not capture all depressions if they are touching all edges
-        holes = np.zeros(self.shape)
-        for combination in combinations([top, bottom, left, right], 3):
+        if self.boundary is None:
+            # Rotate the barrier to ensure that all depressions touching edges are captured
+            # TODO: This may not capture all depressions if they are touching all edges
+            ix = np.arange(self.size).reshape(self.shape)
+            top = ix[0, :]
+            bottom = ix[-1, :]
+            left = ix[:, 0]
+            right = ix[:, -1]
+            holes = np.zeros(self.shape)
+            for combination in combinations([top, bottom, left, right], 3):
+                mask = np.copy(self.dem)
+                exterior = np.concatenate(combination)
+                mask.flat[exterior] = self.dem.max()
+                seed = np.copy(mask)
+                seed[1:-1, 1:-1] = self.dem.max()
+                rec = skimage.morphology.reconstruction(seed, mask, method='erosion')
+                holes[(rec - self.dem) > 0] = self.dem[(rec - self.dem) > 0]
+        else:
             mask = np.copy(self.dem)
-            exterior = np.concatenate(combination)
-            mask.flat[exterior] = self.dem.max()
+            mask[self.boundary > 0] = self.dem.max() + 1
+            mask[self.boundary < 0] = self.dem.min() - 1
             seed = np.copy(mask)
-            seed[1:-1, 1:-1] = self.dem.max()
+            seed[1:-1, 1:-1] = mask.max()
             rec = skimage.morphology.reconstruction(seed, mask, method='erosion')
-            holes[(rec - self.dem) > 0] = self.dem[(rec - self.dem) > 0]
+            holes = rec - mask
+        if self.min_size:
+            holes_mask = skimage.morphology.remove_small_objects(holes != 0, min_size=self.min_size)
+            holes = np.where(holes_mask, holes, 0)
         # Specify levels
         levels = []
         mask = np.where(holes == 0, holes.min(), holes)
@@ -68,6 +82,10 @@ def find_depressions(self):
         for _ in range(self.max_levels):
             diff = rec - mask
             holes = np.where(diff, holes, 0)
+            if self.min_size:
+                holes_mask = skimage.morphology.remove_small_objects(holes != 0,
+                                                                     min_size=self.min_size)
+                holes = np.where(holes_mask, holes, 0)
             mask = np.where(holes == 0, holes.min(), holes)
             seed = np.where(holes == 0, mask, holes.max())
             rec = skimage.morphology.reconstruction(seed, mask, method='erosion')