Add a D8 stream extraction function to pygeoprocessing.routing

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

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',

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)

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)