Merge pull request natcap#1521 from phargogh/bugfix/1431-sdr-and-prob…

…ably-other-routed-models-have-int32-limitations

Two fixes for SDR - convert processing stack to `long` and reduce rasterizations.

HISTORY.rst

@@ -55,6 +55,12 @@ Unreleased Changes
     * Fixed an issue where certain nodata values were not being handled
       correctly, leading to pixel values of +/- infinity in the urban nature
       balance output raster.  https://github.com/natcap/invest/issues/1519
+* SDR
+    * Fixed an issue encountered in the sediment deposition function where
+      rasters with more than 2^32 pixels would raise a cryptic error relating
+      to negative dimensions. https://github.com/natcap/invest/issues/1431
+    * Optimized the creation of the summary vector by minimizing the number of
+      times the target vector needs to be rasterized.
 
 3.14.1 (2023-12-18)
 -------------------

src/natcap/invest/sdr/sdr.py

@@ -19,6 +19,7 @@
 
 from .. import gettext
 from .. import spec_utils
+from .. import urban_nature_access
 from .. import utils
 from .. import validation
 from ..model_metadata import MODEL_METADATA
@@ -609,7 +610,9 @@ def execute(args):
     target_pixel_size = (min_pixel_size, -min_pixel_size)
 
     target_sr_wkt = dem_raster_info['projection_wkt']
-    vector_mask_options = {'mask_vector_path': args['watersheds_path']}
+    vector_mask_options = {
+        'mask_vector_path': args['watersheds_path'],
+    }
     align_task = task_graph.add_task(
         func=pygeoprocessing.align_and_resize_raster_stack,
         args=(
@@ -1491,18 +1494,26 @@ def _generate_report(
     target_layer = target_vector.GetLayer()
     target_layer.SyncToDisk()
 
+    # It's worth it to check if the geometries don't significantly overlap.
+    # On large rasters, this can save a TON of time rasterizing even a
+    # relatively simple vector.
+    geometries_might_overlap = urban_nature_access._geometries_overlap(
+        watershed_results_sdr_path)
+    fields_and_rasters = [
+        ('usle_tot', usle_path), ('sed_export', sed_export_path),
+        ('sed_dep', sed_deposition_path), ('avoid_exp', avoided_export_path),
+        ('avoid_eros', avoided_erosion_path)]
+
+    # Using the list option for raster path bands so that we can reduce
+    # rasterizations, which are costly on large datasets.
+    zonal_stats_results = pygeoprocessing.zonal_statistics(
+        [(raster_path, 1) for (_, raster_path) in fields_and_rasters],
+        watershed_results_sdr_path,
+        polygons_might_overlap=geometries_might_overlap)
+
     field_summaries = {
-        'usle_tot': pygeoprocessing.zonal_statistics(
-            (usle_path, 1), watershed_results_sdr_path),
-        'sed_export': pygeoprocessing.zonal_statistics(
-            (sed_export_path, 1), watershed_results_sdr_path),
-        'sed_dep': pygeoprocessing.zonal_statistics(
-            (sed_deposition_path, 1), watershed_results_sdr_path),
-        'avoid_exp': pygeoprocessing.zonal_statistics(
-            (avoided_export_path, 1), watershed_results_sdr_path),
-        'avoid_eros': pygeoprocessing.zonal_statistics(
-            (avoided_erosion_path, 1), watershed_results_sdr_path),
-    }
+        field: stats for ((field, _), stats) in
+        zip(fields_and_rasters, zonal_stats_results)}
 
     for field_name in field_summaries:
         field_def = ogr.FieldDefn(field_name, ogr.OFTReal)

src/natcap/invest/sdr/sdr_core.pyx

@@ -436,17 +436,17 @@ def calculate_sediment_deposition(
     flow_dir_info = pygeoprocessing.get_raster_info(mfd_flow_direction_path)
     n_cols, n_rows = flow_dir_info['raster_size']
     cdef int mfd_nodata = 0
-    cdef stack[int] processing_stack
+    cdef stack[long] processing_stack
     cdef float sdr_nodata = pygeoprocessing.get_raster_info(
         sdr_path)['nodata'][0]
     cdef float e_prime_nodata = pygeoprocessing.get_raster_info(
         e_prime_path)['nodata'][0]
     cdef long col_index, row_index, win_xsize, win_ysize, xoff, yoff
     cdef long global_col, global_row, j, k
-    cdef unsigned long flat_index
+    cdef long flat_index
     cdef long seed_col = 0
     cdef long seed_row = 0
-    cdef long neighbor_row, neighbor_col
+    cdef long neighbor_row, neighbor_col, ds_neighbor_row, ds_neighbor_col
     cdef int flow_val, neighbor_flow_val, ds_neighbor_flow_val
     cdef int flow_weight, neighbor_flow_weight
     cdef float flow_sum, neighbor_flow_sum