diff --git a/HISTORY.rst b/HISTORY.rst index 460307c5..97ee691a 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -3,6 +3,11 @@ Release History Unreleased Changes ------------------ +* Added a D8 stream extraction function at + ``pygeoprocessing.routing.extract_streams_d8`` which takes a D8 flow + accumulation raster and a flow accumulation threshold, setting all pixels + with accumulation above that threshold to 1 and all other valid pixels to 0. + https://github.com/natcap/pygeoprocessing/issues/272 * Adding a new function, ``pygeoprocessing.create_raster_from_bounding_box``, that enables the creation of a new raster from a bounding box. https://github.com/natcap/pygeoprocessing/issues/276 diff --git a/src/pygeoprocessing/routing/__init__.py b/src/pygeoprocessing/routing/__init__.py index 39f7d8e4..7f85b828 100644 --- a/src/pygeoprocessing/routing/__init__.py +++ b/src/pygeoprocessing/routing/__init__.py @@ -1,20 +1,17 @@ -from pygeoprocessing.routing.routing import ( - fill_pits, - fill_pits, - flow_dir_d8, - flow_accumulation_d8, - flow_dir_mfd, - flow_accumulation_mfd, - distance_to_channel_d8, - distance_to_channel_mfd, - extract_streams_mfd, - detect_outlets, - extract_strahler_streams_d8, - calculate_subwatershed_boundary, - detect_lowest_drain_and_sink, - ) -from pygeoprocessing.routing.watershed import ( - delineate_watersheds_d8) +from pygeoprocessing.routing.helper_functions import extract_streams_d8 +from pygeoprocessing.routing.routing import calculate_subwatershed_boundary +from pygeoprocessing.routing.routing import detect_lowest_drain_and_sink +from pygeoprocessing.routing.routing import detect_outlets +from pygeoprocessing.routing.routing import distance_to_channel_d8 +from pygeoprocessing.routing.routing import distance_to_channel_mfd +from pygeoprocessing.routing.routing import extract_strahler_streams_d8 +from pygeoprocessing.routing.routing import extract_streams_mfd +from pygeoprocessing.routing.routing import fill_pits +from pygeoprocessing.routing.routing import flow_accumulation_d8 +from pygeoprocessing.routing.routing import flow_accumulation_mfd +from pygeoprocessing.routing.routing import flow_dir_d8 +from pygeoprocessing.routing.routing import flow_dir_mfd +from pygeoprocessing.routing.watershed import delineate_watersheds_d8 __all__ = ( 'fill_pits', @@ -25,6 +22,7 @@ 'flow_accumulation_mfd', 'distance_to_channel_d8', 'distance_to_channel_mfd', + 'extract_streams_d8', 'extract_streams_mfd', 'delineate_watersheds_d8', 'detect_outlets', diff --git a/src/pygeoprocessing/routing/helper_functions.py b/src/pygeoprocessing/routing/helper_functions.py new file mode 100644 index 00000000..a1e68124 --- /dev/null +++ b/src/pygeoprocessing/routing/helper_functions.py @@ -0,0 +1,56 @@ +import numpy +from osgeo import gdal + +from ..geoprocessing import get_raster_info +from ..geoprocessing import raster_calculator +from ..geoprocessing_core import DEFAULT_GTIFF_CREATION_TUPLE_OPTIONS + + +def extract_streams_d8( + flow_accum_raster_path_band, flow_threshold, target_stream_raster_path, + raster_driver_creation_tuple=DEFAULT_GTIFF_CREATION_TUPLE_OPTIONS): + """Extract D8 streams based on a flow accumulation threshold. + + Creates an unsigned byte raster where pixel values of 1 indicate the + presence of a stream and pixel values of 0 indicate the absence of a + stream. Any flow accumulation pixels greater than ``flow_threshold`` are + considered stream pixels. Nodata values found in the input flow + accumulation raster propagate through to the target stream raster. + + Args: + flow_accum_raster_path_band (tuple): A (path, band) tuple indicating + the path to a D8 flow accumulation raster and the band index to + use. + flow_threshold (number): The flow threshold. Flow accumulation values + greater than this threshold are considered stream pixels, values + less than this threshold are non-stream pixels. + target_stream_raster_path (string): Where the target streams raster + should be written. + raster_driver_creation_tuple (tuple): A tuple where the first element + is the GDAL driver name of the target raster and the second element + is an iterable of raster creation options for the selected driver. + + Returns: + ``None`` + """ + flow_accum_nodata = get_raster_info( + flow_accum_raster_path_band[0])['nodata'][ + flow_accum_raster_path_band[1]-1] + target_nodata = 255 + + def _threshold_streams(flow_accum): + out_matrix = numpy.full(flow_accum.shape, target_nodata, + dtype=numpy.uint8) + + valid_pixels = numpy.ones(flow_accum.shape, dtype=bool) + if flow_accum_nodata is not None: + valid_pixels = ~numpy.isclose(flow_accum, flow_accum_nodata) + + over_threshold = flow_accum > flow_threshold + out_matrix[valid_pixels & over_threshold] = 1 + out_matrix[valid_pixels & ~over_threshold] = 0 + return out_matrix + + raster_calculator( + [flow_accum_raster_path_band], _threshold_streams, + target_stream_raster_path, gdal.GDT_Byte, target_nodata) diff --git a/tests/test_routing.py b/tests/test_routing.py index 5be070f8..c429cde4 100644 --- a/tests/test_routing.py +++ b/tests/test_routing.py @@ -4,13 +4,12 @@ import tempfile import unittest -from osgeo import gdal import numpy import numpy.testing -import scipy.interpolate - import pygeoprocessing import pygeoprocessing.routing +import scipy.interpolate +from osgeo import gdal from test_geoprocessing import _array_to_raster @@ -1311,3 +1310,46 @@ def test_channel_not_exist_distance(self): numpy.testing.assert_almost_equal( pygeoprocessing.raster_to_numpy_array(distance_path), expected_result) + + def test_extract_streams_d8(self): + """PGP.routing: test d8 stream thresholding.""" + # test without nodata + flow_accum_no_nodata_path = os.path.join( + self.workspace_dir, 'flow_accum_no_nodata.tif') + flow_accum = numpy.array([ + [2, 3, 4, 4, 4, 4], + [3, 5, 6, 2, 3, 4], + [2, 3, 4, 1, 2, 3], + [0, 0, 0, 4, 4, 0]], dtype=numpy.int32) + _array_to_raster( + flow_accum, target_nodata=None, + target_path=flow_accum_no_nodata_path) + + expected_array = numpy.array([ + [0, 0, 1, 1, 1, 1], + [0, 1, 1, 0, 0, 1], + [0, 0, 1, 0, 0, 0], + [0, 0, 0, 1, 1, 0]], dtype=numpy.uint8) + target_streams_path = os.path.join( + self.workspace_dir, 'streams.tif') + pygeoprocessing.routing.extract_streams_d8( + (flow_accum_no_nodata_path, 1), 3.5, target_streams_path) + numpy.testing.assert_array_equal( + pygeoprocessing.raster_to_numpy_array(target_streams_path), + expected_array) + + # test with nodata + n = 255 # nodata value + expected_array = numpy.array([ + [0, 0, 1, 1, 1, 1], + [0, 1, 1, 0, 0, 1], + [0, 0, 1, 0, 0, 0], + [n, n, n, 1, 1, n]], dtype=numpy.uint8) + _array_to_raster( + flow_accum, target_nodata=0, + target_path=flow_accum_no_nodata_path) + pygeoprocessing.routing.extract_streams_d8( + (flow_accum_no_nodata_path, 1), 3.5, target_streams_path) + numpy.testing.assert_array_equal( + pygeoprocessing.raster_to_numpy_array(target_streams_path), + expected_array)