diff --git a/README.md b/README.md index a495edd..871091d 100644 --- a/README.md +++ b/README.md @@ -376,6 +376,7 @@ ax.set_title('Soil types (raster)', size=14) - `stream_order` : Compute the (strahler) stream order. - `extract_river_network` : Extract river segments from a catchment and return a geojson object. + - `extract_profiles` : Extract river segments and return a list of channel indices along with a dictionary describing connectivity. - `cell_dh` : Compute the drop in elevation from each cell to its downstream neighbor. - `cell_distances` : Compute the distance from each cell to its downstream neighbor. - `cell_slopes` : Compute the slope between each cell and its downstream neighbor. @@ -471,6 +472,7 @@ Performance benchmarks on a 2015 MacBook Pro (M: million, K: thousand): | `hand` | MFD | 36M total, 770K channel | 29.8 [s] | | `stream_order` | D8 | 36M total, 1M channel | 3.99 [s] | | `extract_river_network` | D8 | 36M total, 345K channel | 4.07 [s] | +| `extract_profiles` | D8 | 36M total, 345K channel | 2.89 [s] | | `detect_pits` | N/A | 36M | 1.80 [s] | | `detect_flats` | N/A | 36M | 1.84 [s] | | `fill_pits` | N/A | 36M | 2.52 [s] | diff --git a/data/cogeo.tiff b/data/cogeo.tiff new file mode 100644 index 0000000..7d40bfb Binary files /dev/null and b/data/cogeo.tiff differ diff --git a/docs/index.md b/docs/index.md index 5db33ca..150bad8 100644 --- a/docs/index.md +++ b/docs/index.md @@ -380,6 +380,7 @@ ax.set_title('Soil types (raster)', size=14) - `stream_order` : Compute the (strahler) stream order. - `extract_river_network` : Extract river segments from a catchment and return a geojson object. + - `extract_profiles` : Extract river segments and return a list of channel indices along with a dictionary describing connectivity. - `cell_dh` : Compute the drop in elevation from each cell to its downstream neighbor. - `cell_distances` : Compute the distance from each cell to its downstream neighbor. - `cell_slopes` : Compute the slope between each cell and its downstream neighbor. @@ -475,6 +476,7 @@ Performance benchmarks on a 2015 MacBook Pro (M: million, K: thousand): | `hand` | MFD | 36M total, 770K channel | 29.8 [s] | | `stream_order` | D8 | 36M total, 1M channel | 3.99 [s] | | `extract_river_network` | D8 | 36M total, 345K channel | 4.07 [s] | +| `extract_profiles` | D8 | 36M total, 345K channel | 2.89 [s] | | `detect_pits` | N/A | 36M | 1.80 [s] | | `detect_flats` | N/A | 36M | 1.84 [s] | | `fill_pits` | N/A | 36M | 2.52 [s] | diff --git a/pysheds/__init__.py b/pysheds/__init__.py index f9aa3e1..e19434e 100644 --- a/pysheds/__init__.py +++ b/pysheds/__init__.py @@ -1 +1 @@ -__version__ = "0.3.2" +__version__ = "0.3.3" diff --git a/pysheds/_sgrid.py b/pysheds/_sgrid.py index 1fe424f..5401085 100644 --- a/pysheds/_sgrid.py +++ b/pysheds/_sgrid.py @@ -1,7 +1,7 @@ from heapq import heappop, heappush import numpy as np from numba import njit, prange -from numba.types import float64, int64, uint32, uint16, uint8, boolean, UniTuple, Tuple, List, void +from numba.types import float64, int64, uint32, uint16, uint8, boolean, UniTuple, Tuple, List, DictType, void # Functions for 'flowdir' @@ -1466,6 +1466,36 @@ def _d8_stream_network_iter_numba(fdir, indegree, orig_indegree, startnodes): endnode = fdir.flat[startnode] return profiles +@njit(Tuple((List(List(int64)), DictType(int64, int64)))(int64[:,:], uint8[:], + uint8[:], int64[:], boolean), + cache=True) +def _d8_stream_connection_iter_numba(fdir, indegree, orig_indegree, startnodes, + include_endpoint): + n = startnodes.size + profiles = [[0]] + connections = {0 : 0} + _ = profiles.pop() + _ = connections.pop(0) + for k in range(n): + startnode = startnodes.flat[k] + endnode = fdir.flat[startnode] + profile = [startnode] + while (indegree.flat[startnode] == 0): + profile.append(endnode) + indegree.flat[endnode] -= 1 + if (orig_indegree.flat[endnode] > 1): + chain_start = profile[0] + chain_end = profile[-1] + connections[chain_start] = chain_end + if not include_endpoint: + _ = profile.pop() + profiles.append(profile) + if (indegree.flat[endnode] == 0): + profile = [endnode] + startnode = endnode + endnode = fdir.flat[startnode] + return profiles, connections + @njit(float64[:,:](int64[:,:], int64[:,:], float64[:,:]), parallel=True, cache=True) diff --git a/pysheds/io.py b/pysheds/io.py index 8d477ea..b91b96a 100644 --- a/pysheds/io.py +++ b/pysheds/io.py @@ -101,7 +101,7 @@ def read_raster(data, band=1, window=None, window_crs=None, mask_geometry=False, if window is None: shape = f.shape if len(f.indexes) > 1: - data = np.ma.filled(f.read_band(band)) + data = np.ma.filled(f.read(band)) else: data = np.ma.filled(f.read()) affine = f.transform @@ -121,7 +121,7 @@ def read_raster(data, band=1, window=None, window_crs=None, mask_geometry=False, # If window crs not specified, assume it is in raster crs ix_window = f.window(*window) if len(f.indexes) > 1: - data = np.ma.filled(f.read_band(band, window=ix_window)) + data = np.ma.filled(f.read(band, window=ix_window)) else: data = np.ma.filled(f.read(window=ix_window)) affine = f.window_transform(ix_window) diff --git a/pysheds/sgrid.py b/pysheds/sgrid.py index cf1ffc4..f0521aa 100644 --- a/pysheds/sgrid.py +++ b/pysheds/sgrid.py @@ -1385,7 +1385,7 @@ def _mfd_compute_hand(self, fdir, mask, dirmap=(64, 128, 1, 2, 4, 8, 16, 32), def extract_river_network(self, fdir, mask, dirmap=(64, 128, 1, 2, 4, 8, 16, 32), routing='d8', algorithm='iterative', **kwargs): """ - Generates river segments from accumulation and flow_direction arrays. + Generates river segments from flow direction and mask. Parameters ---------- @@ -1453,6 +1453,81 @@ def extract_river_network(self, fdir, mask, dirmap=(64, 128, 1, 2, 4, 8, 16, 32) geo = geojson.FeatureCollection(featurelist) return geo + def extract_profiles(self, fdir, mask, dirmap=(64, 128, 1, 2, 4, 8, 16, 32), + include_endpoint=True, routing='d8', algorithm='iterative', + **kwargs): + """ + Extracts river segments and connectivity of river segments from flow direction and mask. + + Parameters + ---------- + fdir : Raster + Flow direction data. + mask : Raster + Boolean raster 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] + include_endpoint : bool + If True, include last cell in each river segment. + If False, do not include last cell (such that + cell indices in each profile are unique). + routing : str + Routing algorithm to use: + 'd8' : D8 flow directions + algorithm : str + Algorithm type to use: + 'iterative' : Use an iterative algorithm (recommended). + 'recursive' : Use a recursive algorithm. + + Additional keyword arguments (**kwargs) are passed to self.view. + + Returns + ------- + profiles : list of lists of ints + A list containing a collection of river profiles. Each river profile + is a list containing the indices of the grid cells inside the + river segment. Indices correspond to the flattened index of river segment + cells. + connections : dict (int : int) + A dictionary describing the connectivity of the profiles. For each + key-value pair, the key represents index of the upstream profile and + the value represents the index of the downstream profile that it drains to. + Indices correspond to the ordered elements of the `profiles` object. + """ + if routing.lower() == 'd8': + fdir_overrides = {'dtype' : np.int64, 'nodata' : fdir.nodata} + else: + raise NotImplementedError('Only implemented for `d8` routing.') + mask_overrides = {'dtype' : np.bool8, 'nodata' : False} + kwargs.update(fdir_overrides) + fdir = self._input_handler(fdir, **kwargs) + kwargs.update(mask_overrides) + mask = self._input_handler(mask, **kwargs) + # Find nodata cells and invalid cells + nodata_cells = self._get_nodata_cells(fdir) + invalid_cells = ~np.in1d(fdir.ravel(), dirmap).reshape(fdir.shape) + # Set nodata cells to zero + fdir[nodata_cells] = 0 + fdir[invalid_cells] = 0 + maskleft, maskright, masktop, maskbottom = self._pop_rim(mask, nodata=False) + masked_fdir = np.where(mask, fdir, 0).astype(np.int64) + startnodes = np.arange(fdir.size, dtype=np.int64) + endnodes = _self._flatten_fdir_numba(masked_fdir, dirmap).reshape(fdir.shape) + indegree = np.bincount(endnodes.ravel(), minlength=fdir.size).astype(np.uint8) + orig_indegree = np.copy(indegree) + startnodes = startnodes[(indegree == 0)] + profiles, connections = _self._d8_stream_connection_iter_numba(endnodes, indegree, + orig_indegree, + startnodes, + include_endpoint) + connections = dict(connections) + indices = {profile[0] : index for index, profile in enumerate(profiles)} + connections = {indices[key] : indices.setdefault(value, indices[key]) + for key, value in connections.items()} + return profiles, connections + def stream_order(self, fdir, mask, dirmap=(64, 128, 1, 2, 4, 8, 16, 32), nodata_out=0, routing='d8', algorithm='iterative', **kwargs): """ diff --git a/setup.py b/setup.py index 5f62d67..f45487f 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ from setuptools import setup setup(name='pysheds', - version='0.3.2', + version='0.3.3', description='🌎 Simple and fast watershed delineation in python.', author='Matt Bartos', author_email='mdbartos@umich.edu', diff --git a/tests/test_grid.py b/tests/test_grid.py index 28a3e6d..23bc342 100644 --- a/tests/test_grid.py +++ b/tests/test_grid.py @@ -11,6 +11,7 @@ dir_path = os.path.join(data_dir, 'dir.asc') dem_path = os.path.join(data_dir, 'dem.tif') roi_path = os.path.join(data_dir, 'roi.tif') +multiband_path = os.path.join(data_dir, 'cogeo.tiff') feature_geometry = [{'type': 'Polygon', 'coordinates': (((-97.29749977660477, 32.74000135435936), (-97.29083107907053, 32.74000328969928), @@ -331,14 +332,10 @@ def test_cell_slopes(): slopes_dinf = grid.cell_slopes(dem, fdir_dinf, routing='dinf') slopes_mfd = grid.cell_slopes(dem, fdir_mfd, routing='mfd') -# def test_set_nodata(): -# grid.set_nodata('dir', 0) - def test_to_ascii(): catch = d.catch fdir = d.fdir grid.clip_to(catch) - # np.float is depreciated grid.to_ascii(fdir, 'test_dir.asc', target_view=fdir.viewfinder, dtype=np.float64) fdir_out = grid.read_ascii('test_dir.asc', dtype=np.uint8) assert((fdir_out == fdir).all()) @@ -346,6 +343,11 @@ def test_to_ascii(): fdir_out = grid.read_ascii('test_dir.asc', dtype=np.uint8) assert((fdir_out == grid.view(fdir)).all()) +def test_read_raster(): + band_1 = grid.read_raster(multiband_path, band=1) + band_2 = grid.read_raster(multiband_path, band=2) + band_3 = grid.read_raster(multiband_path, band=3) + def test_to_raster(): catch = d.catch fdir = d.fdir @@ -406,6 +408,13 @@ def test_extract_river_network(): grid.extract_river_network(catch, acc > 20, algorithm='recursive') # TODO: Need more checks here. Check if endnodes equals next startnode +def test_extract_profiles(): + fdir = d.fdir + catch = d.catch + acc = d.acc + grid.clip_to(catch) + profiles, connections = grid.extract_profiles(catch, acc > 20) + def test_view_methods(): dem = d.dem catch = d.catch