Repository: natcap/invest (python)

Issue:
Create a mask raster for SDR/NDR areas that don't drain to stream
As discussed in #609, some SDR/NDR outputs now have `nodata` in areas that don't drain to a stream. @newtpatrol has suggested that:
> these pixels receive a distinctly bogus value like -9999, so they'd be really obvious to people, and they wouldn't have to pour through their inputs to see what's causing NoData in the output

This would clarify that nodata areas come from an input, while areas with `-9999` are undefined because they don't drain to a stream. 

Pull request description:
Create a mask indicating NDR, SDR drainage
This PR introduces a new raster output to the SDR and NDR models, really just a `raster_calculator` call indicating which pixels do and do not flow to the stream.

Fixes #621

This is a draft PR until I've finished the UG updates.

# Checklist
- [x] Updated HISTORY.rst (if these changes are user-facing)

- [x] Updated the user's guide (if needed)


Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index b75b28a1b..8344b3faa 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -47,6 +47,16 @@ Unreleased Changes
       of 0.5 indicate same abundance between baseline and current/future
       LULC; values 0.5 to 1 indicate less abundance in current/future LULC
       and therefore higher rarity.
+* SDR
+    * Added a new raster to the model's workspace,
+      ``intermediate_outputs/what_drains_to_stream[suffix].tif``.  This raster
+      has pixel values of 1 where DEM pixels flow to an identified stream, and
+      0 where they do not.
+* NDR
+    * Added a new raster to the model's workspace,
+      ``intermediate_outputs/what_drains_to_stream[suffix].tif``.  This raster
+      has pixel values of 1 where DEM pixels flow to an identified stream, and
+      0 where they do not.
 
 3.9.2 (2021-10-29)
 ------------------
diff --git a/src/natcap/invest/ndr/ndr.py b/src/natcap/invest/ndr/ndr.py
index 5ee5cc85f..e4b869dfb 100644
--- a/src/natcap/invest/ndr/ndr.py
+++ b/src/natcap/invest/ndr/ndr.py
@@ -7,10 +7,11 @@ import pickle
 import numpy
 import pygeoprocessing
 import pygeoprocessing.routing
-from osgeo import gdal, ogr
 import taskgraph
+from osgeo import gdal, ogr
 
-from .. import utils, validation, MODEL_METADATA
+from .. import MODEL_METADATA, utils, validation
+from ..sdr import sdr
 from . import ndr_core
 
 LOGGER = logging.getLogger(__name__)
@@ -227,6 +228,7 @@ _INTERMEDIATE_BASE_FILES = {
     'flow_direction_path': 'flow_direction.tif',
     'thresholded_slope_path': 'thresholded_slope.tif',
     'dist_to_channel_path': 'dist_to_channel.tif',
+    'drainage_mask': 'what_drains_to_stream.tif',
 }
 
 _CACHE_BASE_FILES = {
@@ -545,6 +547,15 @@ def execute(args):
         target_path_list=[f_reg['dist_to_channel_path']],
         task_name='dist to channel')
 
+    _ = task_graph.add_task(
+        func=sdr._calculate_what_drains_to_stream,
+        args=(f_reg['flow_direction_path'],
+              f_reg['dist_to_channel_path'],
+              f_reg['drainage_mask']),
+        target_path_list=[f_reg['drainage_mask']],
+        dependent_task_list=[flow_dir_task, dist_to_channel_task],
+        task_name='write mask of what drains to stream')
+
     ic_task = task_graph.add_task(
         func=calculate_ic,
         args=(
diff --git a/src/natcap/invest/sdr/sdr.py b/src/natcap/invest/sdr/sdr.py
index 5e324e5d9..3eb03a4a3 100644
--- a/src/natcap/invest/sdr/sdr.py
+++ b/src/natcap/invest/sdr/sdr.py
@@ -7,17 +7,16 @@ The SDR method in this model is based on:
     large watersheds." Journal of Soil and Water Conservation 63.3 (2008):
     105-111.
 """
-import os
 import logging
+import os
 
-from osgeo import gdal, ogr
 import numpy
 import pygeoprocessing
 import pygeoprocessing.routing
 import taskgraph
-from .. import utils
-from .. import validation
-from .. import MODEL_METADATA
+from osgeo import gdal, ogr
+
+from .. import MODEL_METADATA, utils, validation
 from . import sdr_core
 
 LOGGER = logging.getLogger(__name__)
@@ -220,7 +219,8 @@ _INTERMEDIATE_BASE_FILES = {
     'w_path': 'w.tif',
     'ws_inverse_path': 'ws_inverse.tif',
     'e_prime_path': 'e_prime.tif',
-    'weighted_avg_aspect_path': 'weighted_avg_aspect.tif'
+    'weighted_avg_aspect_path': 'weighted_avg_aspect.tif',
+    'drainage_mask': 'what_drains_to_stream.tif',
 }
 
 _TMP_BASE_FILES = {
@@ -234,6 +234,7 @@ _TMP_BASE_FILES = {
 # Target nodata is for general rasters that are positive, and _IC_NODATA are
 # for rasters that are any range
 _TARGET_NODATA = -1.0
+_BYTE_NODATA = 255
 _IC_NODATA = float(numpy.finfo('float32').min)
 
 
@@ -661,6 +662,14 @@ def execute(args):
             sdr_bare_task],
         task_name='calculate sediment retention')
 
+    _ = task_graph.add_task(
+        func=_calculate_what_drains_to_stream,
+        args=(f_reg['flow_direction_path'], f_reg['d_dn_path'],
+              f_reg['drainage_mask']),
+        target_path_list=[f_reg['drainage_mask']],
+        dependent_task_list=[flow_dir_task, d_dn_task],
+        task_name='write mask of what drains to stream')
+
     _ = task_graph.add_task(
         func=_generate_report,
         args=(
@@ -676,6 +685,71 @@ def execute(args):
     task_graph.join()
 
 
+def _calculate_what_drains_to_stream(
+        flow_dir_mfd_path, dist_to_channel_mfd_path, target_mask_path):
+    """Create a mask indicating regions that do or do not drain to a stream.
+
+    This is useful because ``pygeoprocessing.distance_to_stream_mfd`` may leave
+    some unexpected regions as nodata if they do not drain to a stream.  This
+    may be confusing behavior, so this mask is intended to locate what drains
+    to a stream and what does not. A pixel doesn't drain to a stream if it has
+    a defined flow direction but undefined distance to stream.
+
+    Args:
+        flow_dir_mfd_path (string): The path to an MFD flow direction raster.
+            This raster must have a nodata value defined.
+        dist_to_channel_mfd_path (string): The path to an MFD
+            distance-to-channel raster.  This raster must have a nodata value
+            defined.
+        target_mask_path (string): The path to where the mask raster should be
+            written.
+
+    Returns:
+        ``None``
+    """
+    flow_dir_mfd_nodata = pygeoprocessing.get_raster_info(
+        flow_dir_mfd_path)['nodata'][0]
+    dist_to_channel_nodata = pygeoprocessing.get_raster_info(
+        dist_to_channel_mfd_path)['nodata'][0]
+
+    def _what_drains_to_stream(flow_dir_mfd, dist_to_channel):
+        """Determine which pixels do and do not drain to a stream.
+
+        Args:
+            flow_dir_mfd (numpy.array): A numpy array of MFD flow direction
+                values.
+            dist_to_channel (numpy.array): A numpy array of calculated
+                distances to the nearest channel.
+
+        Returns:
+            A ``numpy.array`` of dtype ``numpy.uint8`` with pixels where:
+
+                * ``255`` where ``flow_dir_mfd`` is nodata (and thus 
+                  ``dist_to_channel`` is also nodata).
+                * ``0`` where ``flow_dir_mfd`` has data and ``dist_to_channel``
+                  does not
+                * ``1`` where ``flow_dir_mfd`` has data, and
+                  ``dist_to_channel`` also has data.
+        """
+        drains_to_stream = numpy.full(
+            flow_dir_mfd.shape, _BYTE_NODATA, dtype=numpy.uint8)
+        valid_flow_dir = ~numpy.isclose(flow_dir_mfd, flow_dir_mfd_nodata)
+        valid_dist_to_channel = (
+            ~numpy.isclose(dist_to_channel, dist_to_channel_nodata) &
+            valid_flow_dir)
+
+        # Nodata where both flow_dir and dist_to_channel are nodata
+        # 1 where flow_dir and dist_to_channel have values (drains to stream)
+        # 0 where flow_dir has data and dist_to_channel doesn't (doesn't drain)
+        drains_to_stream[valid_flow_dir & valid_dist_to_channel] = 1
+        drains_to_stream[valid_flow_dir & ~valid_dist_to_channel] = 0
+        return drains_to_stream
+
+    pygeoprocessing.raster_calculator(
+        [(flow_dir_mfd_path, 1), (dist_to_channel_mfd_path, 1)],
+        _what_drains_to_stream, target_mask_path, gdal.GDT_Byte, _BYTE_NODATA)
+
+
 def _calculate_ls_factor(
         flow_accumulation_path, slope_path, avg_aspect_path, l_max,
         target_ls_prime_factor_path):


Test patch:
diff --git a/tests/test_ndr.py b/tests/test_ndr.py
index 2850f6fc2..bcb3660da 100644
--- a/tests/test_ndr.py
+++ b/tests/test_ndr.py
@@ -7,8 +7,7 @@ import unittest
 
 import numpy
 import pygeoprocessing
-from osgeo import gdal
-from osgeo import ogr
+from osgeo import gdal, ogr
 
 REGRESSION_DATA = os.path.join(
     os.path.dirname(__file__), '..', 'data', 'invest-test-data', 'ndr')
@@ -246,10 +245,18 @@ class NDRTests(unittest.TestCase):
         if mismatch_list:
             raise RuntimeError("results not expected: %s" % mismatch_list)
 
+        # We only need to test that the drainage mask exists.  Functionality
+        # for that raster is tested in SDR.
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(
+                    args['workspace_dir'], 'intermediate_outputs',
+                    'what_drains_to_stream.tif')))
+
     def test_validation(self):
         """NDR test argument validation."""
-        from natcap.invest.ndr import ndr
         from natcap.invest import validation
+        from natcap.invest.ndr import ndr
 
         # use predefined directory so test can clean up files during teardown
         args = NDRTests.generate_base_args(self.workspace_dir)
diff --git a/tests/test_sdr.py b/tests/test_sdr.py
index 52e8217d8..f0d79ab35 100644
--- a/tests/test_sdr.py
+++ b/tests/test_sdr.py
@@ -1,14 +1,12 @@
 """InVEST SDR model tests."""
-import unittest
-import tempfile
-import shutil
 import os
+import shutil
+import tempfile
+import unittest
 
-import pygeoprocessing
 import numpy
-from osgeo import ogr
-from osgeo import osr
-from osgeo import gdal
+import pygeoprocessing
+from osgeo import gdal, ogr, osr
 
 REGRESSION_DATA = os.path.join(
     os.path.dirname(__file__), '..', 'data', 'invest-test-data', 'sdr')
@@ -219,6 +217,14 @@ class SDRTests(unittest.TestCase):
             args['workspace_dir'], 'watershed_results_sdr.shp')
         assert_expected_results_in_vector(expected_results, vector_path)
 
+        # We only need to test that the drainage mask exists.  Functionality
+        # for that raster is tested elsewhere
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(
+                    args['workspace_dir'], 'intermediate_outputs',
+                    'what_drains_to_stream.tif')))
+
     def test_regression_with_undefined_nodata(self):
         """SDR base regression test with undefined nodata values.
 
@@ -346,8 +352,8 @@ class SDRTests(unittest.TestCase):
 
     def test_missing_lulc_value(self):
         """SDR test for ValueError when LULC value not found in table."""
-        from natcap.invest.sdr import sdr
         import pandas
+        from natcap.invest.sdr import sdr
 
         # use predefined directory so test can clean up files during teardown
         args = SDRTests.generate_base_args(self.workspace_dir)
@@ -369,6 +375,48 @@ class SDRTests(unittest.TestCase):
             " are: [2.]" in str(context.exception),
             f'error does not match {str(context.exception)}')
 
+    def test_what_drains_to_stream(self):
+        """SDR test for what pixels drain to a stream."""
+        from natcap.invest.sdr import sdr
+
+        srs = osr.SpatialReference()
+        srs.ImportFromEPSG(26910)  # NAD83 / UTM zone 11N
+        srs_wkt = srs.ExportToWkt()
+        origin = (463250, 4929700)
+        pixel_size = (30, -30)
+
+        flow_dir_mfd = numpy.array([
+            [0, 1],
+            [1, 1]], dtype=numpy.float64)
+        flow_dir_mfd_nodata = 0  # Matches pygeoprocessing output
+        flow_dir_mfd_path = os.path.join(self.workspace_dir, 'flow_dir.tif')
+        pygeoprocessing.numpy_array_to_raster(
+            flow_dir_mfd, flow_dir_mfd_nodata, pixel_size, origin, srs_wkt,
+            flow_dir_mfd_path)
+
+        dist_to_channel = numpy.array([
+            [10, 5],
+            [-1, 6]], dtype=numpy.float64)
+        dist_to_channel_nodata = -1  # Matches pygeoprocessing output
+        dist_to_channel_path = os.path.join(
+            self.workspace_dir, 'dist_to_channel.tif')
+        pygeoprocessing.numpy_array_to_raster(
+            dist_to_channel, dist_to_channel_nodata, pixel_size, origin,
+            srs_wkt, dist_to_channel_path)
+
+        target_what_drains_path = os.path.join(
+            self.workspace_dir, 'what_drains.tif')
+        sdr._calculate_what_drains_to_stream(
+            flow_dir_mfd_path, dist_to_channel_path, target_what_drains_path)
+
+        # 255 is the byte nodata value assigned
+        expected_drainage = numpy.array([
+            [255, 1],
+            [0, 1]], dtype=numpy.uint8)
+        what_drains = pygeoprocessing.raster_to_numpy_array(
+            target_what_drains_path)
+        numpy.testing.assert_allclose(what_drains, expected_drainage)
+
     @staticmethod
     def _assert_regression_results_equal(
             workspace_dir, file_list_path, result_vector_path,