rough draft of routing wrapper fn and use it in ndr

src/natcap/invest/managed_raster/managed_raster.pxd

@@ -42,6 +42,7 @@ cdef class _ManagedRaster:
     cdef bytes raster_path
     cdef int band_id
     cdef int closed
+    cdef float nodata
 
     cdef inline void set(_ManagedRaster self, long xi, long yi, double value)
     cdef inline double get(_ManagedRaster self, long xi, long yi)
@@ -70,3 +71,5 @@ cdef inline int is_close(double x, double y):
     if isnan(x) and isnan(y):
         return 1
     return abs(x - y) <= (1e-8 + 1e-05 * abs(y))
+
+cdef route(flow_dir_path, seed_fn, route_fn)

src/natcap/invest/managed_raster/managed_raster.pyx

@@ -18,6 +18,55 @@ from libcpp.list cimport list as clist
 from libcpp.stack cimport stack
 from libcpp.vector cimport vector
 
+
+cdef route(flow_dir_path, seed_fn, route_fn):
+    """
+    Args:
+        seed_fn (callable): function that accepts an (x, y) coordinate
+            and returns a bool indicating if the pixel is a seed
+        route_fn (callable): function that accepts an (x, y) coordinate
+            and performs whatever routing operation is needed on that pixel.
+
+    Returns:
+        None
+    """
+
+    cdef long win_xsize, win_ysize, xoff, yoff
+    cdef long col_index, row_index, global_col, global_row
+    cdef stack[long] processing_stack
+    cdef long n_cols, n_rows
+
+    flow_dir_info = pygeoprocessing.get_raster_info(flow_dir_path)
+    n_cols, n_rows = flow_dir_info['raster_size']
+
+    for offset_dict in pygeoprocessing.iterblocks(
+            (flow_dir_path, 1), offset_only=True, largest_block=0):
+        # use cython variables to avoid python overhead of dict values
+        win_xsize = offset_dict['win_xsize']
+        win_ysize = offset_dict['win_ysize']
+        xoff = offset_dict['xoff']
+        yoff = offset_dict['yoff']
+        for row_index in range(win_ysize):
+            global_row = yoff + row_index
+            for col_index in range(win_xsize):
+                global_col = xoff + col_index
+
+                if seed_fn(global_col, global_row):
+                    processing_stack.push(global_row * n_cols + global_col)
+
+        while processing_stack.size() > 0:
+            # loop invariant, we don't push a cell on the stack that
+            # hasn't already been set for processing.
+            flat_index = processing_stack.top()
+            processing_stack.pop()
+            global_row = flat_index / n_cols
+            global_col = flat_index % n_cols
+
+            to_push = route_fn(global_col, global_row)
+            for index in to_push:
+                processing_stack.push(index)
+
+
 # this ctype is used to store the block ID and the block buffer as one object
 # inside Managed Raster
 ctypedef pair[int, double*] BlockBufferPair

src/natcap/invest/ndr/ndr_core.pyx

@@ -17,6 +17,7 @@ from ..managed_raster.managed_raster cimport _ManagedRaster
 from ..managed_raster.managed_raster cimport ManagedFlowDirRaster
 from ..managed_raster.managed_raster cimport is_close
 from ..managed_raster.managed_raster cimport INFLOW_OFFSETS
+from ..managed_raster.managed_raster cimport route
 
 cdef extern from "time.h" nogil:
     ctypedef int time_t
@@ -65,7 +66,6 @@ def ndr_eff_calculation(
     flow_dir_info = pygeoprocessing.get_raster_info(mfd_flow_direction_path)
     n_cols, n_rows = flow_dir_info['raster_size']
 
-    cdef stack[long] processing_stack
     stream_info = pygeoprocessing.get_raster_info(stream_path)
     # cell sizes must be square, so no reason to test at this point.
     cdef float cell_size = abs(stream_info['pixel_size'][0])
@@ -102,141 +102,124 @@ def ndr_eff_calculation(
     cdef _ManagedRaster to_process_flow_directions_raster = _ManagedRaster(
         to_process_flow_directions_path, 1, True)
 
-    cdef long col_index, row_index, win_xsize, win_ysize, xoff, yoff
-    cdef long global_col, global_row
-    cdef unsigned long flat_index
-    cdef long outflow_weight, flow_dir
+    cdef long flow_dir
     cdef long ds_col, ds_row, i
     cdef float current_step_factor, step_size, crit_len
     cdef long neighbor_row, neighbor_col
     cdef int neighbor_outflow_dir, neighbor_outflow_dir_mask, neighbor_process_flow_dir
     cdef int outflow_dirs, dir_mask
 
-    for offset_dict in pygeoprocessing.iterblocks(
-            (mfd_flow_direction_path, 1), offset_only=True, largest_block=0):
-        # use cython variables to avoid python overhead of dict values
-        win_xsize = offset_dict['win_xsize']
-        win_ysize = offset_dict['win_ysize']
-        xoff = offset_dict['xoff']
-        yoff = offset_dict['yoff']
-        for row_index in range(win_ysize):
-            global_row = yoff + row_index
-            for col_index in range(win_xsize):
-                global_col = xoff + col_index
-                outflow_dirs = <int>to_process_flow_directions_raster.get(
-                    global_col, global_row)
-                should_seed = 0
-                # see if this pixel drains to nodata or the edge, if so it's
-                # a drain
-                for neighbor in (
-                        mfd_flow_direction_raster.get_downslope_neighbors(
-                            global_col, global_row, skip_oob=False)):
-                    if (neighbor.x < 0 or neighbor.x >= n_cols or
-                        neighbor.y < 0 or neighbor.y >= n_rows or
-                        to_process_flow_directions_raster.get(
-                            neighbor.x, neighbor.y) == 0):
-                        should_seed = 1
-                        outflow_dirs &= ~(1 << neighbor.direction)
-
-                if should_seed:
-                    # mark all outflow directions processed
-                    to_process_flow_directions_raster.set(
-                        global_col, global_row, outflow_dirs)
-                    processing_stack.push(global_row * n_cols + global_col)
-
-        while processing_stack.size() > 0:
-            # loop invariant, we don't push a cell on the stack that
-            # hasn't already been set for processing.
-            flat_index = processing_stack.top()
-            processing_stack.pop()
-            global_row = flat_index / n_cols
-            global_col = flat_index % n_cols
-
-            crit_len = <float>crit_len_raster.get(global_col, global_row)
-            retention_eff_lulc = retention_eff_lulc_raster.get(
-                global_col, global_row)
-            flow_dir = <int>mfd_flow_direction_raster.get(
-                    global_col, global_row)
-            if stream_raster.get(global_col, global_row) == 1:
-                # if it's a stream effective retention is 0.
-                effective_retention_raster.set(global_col, global_row, STREAM_EFFECTIVE_RETENTION)
-            elif (is_close(crit_len, crit_len_nodata) or
-                  is_close(retention_eff_lulc, retention_eff_nodata) or
-                  flow_dir == 0):
-                # if it's nodata, effective retention is nodata.
-                effective_retention_raster.set(
-                    global_col, global_row, effective_retention_nodata)
-            else:
-                working_retention_eff = 0.0
-                has_outflow = False
-                for neighbor in mfd_flow_direction_raster.get_downslope_neighbors(
-                        global_col, global_row, skip_oob=False):
-                    has_outflow = True
-                    if (neighbor.x < 0 or neighbor.x >= n_cols or
-                        neighbor.y < 0 or neighbor.y >= n_rows):
-                        continue
-
-                    if neighbor.direction % 2 == 1:
-                        step_size = <float>(cell_size * 1.41421356237)
-                    else:
-                        step_size = cell_size
-                    # guard against a critical length factor that's 0
-                    if crit_len > 0:
-                        current_step_factor = <float>(
-                            exp(-5 * step_size / crit_len))
-                    else:
-                        current_step_factor = 0.0
-
-                    neighbor_effective_retention = (
-                        effective_retention_raster.get(
-                            neighbor.x, neighbor.y))
-
-                    # Case 1: downslope neighbor is a stream pixel
-                    if neighbor_effective_retention == STREAM_EFFECTIVE_RETENTION:
-                        intermediate_retention = (
-                            retention_eff_lulc * (1 - current_step_factor))
-
-                    # Case 2: the current LULC's retention exceeds the neighbor's retention.
-                    elif retention_eff_lulc > neighbor_effective_retention:
-                        intermediate_retention = (
-                            (neighbor_effective_retention * current_step_factor) +
-                            (retention_eff_lulc * (1 - current_step_factor)))
-
-                    # Case 3: the other 2 cases have not been hit.
-                    else:
-                        intermediate_retention = neighbor_effective_retention
-
-                    working_retention_eff += (
-                        intermediate_retention * neighbor.flow_proportion)
-
-                if has_outflow:
-                    effective_retention_raster.set(
-                        global_col, global_row, working_retention_eff)
+    def ndr_seed_fn(col, row):
+        should_seed = False
+        outflow_dirs = <int>to_process_flow_directions_raster.get(col, row)
+        # see if this pixel drains to nodata or the edge, if so it's
+        # a drain
+        for neighbor in (
+                mfd_flow_direction_raster.get_downslope_neighbors(
+                    col, row, skip_oob=False)):
+            if (neighbor.x < 0 or neighbor.x >= n_cols or
+                neighbor.y < 0 or neighbor.y >= n_rows or
+                to_process_flow_directions_raster.get(
+                    neighbor.x, neighbor.y) == 0):
+                should_seed = True
+                outflow_dirs &= ~(1 << neighbor.direction)
+
+        if should_seed:
+            # mark all outflow directions processed
+            to_process_flow_directions_raster.set(
+                col, row, outflow_dirs)
+
+        return should_seed
+
+    def ndr_route_fn(col, row):
+        to_push = []
+        crit_len = <float>crit_len_raster.get(col, row)
+        retention_eff_lulc = retention_eff_lulc_raster.get(col, row)
+        flow_dir = <int>mfd_flow_direction_raster.get(col, row)
+        if stream_raster.get(col, row) == 1:
+            # if it's a stream effective retention is 0.
+            effective_retention_raster.set(col, row, STREAM_EFFECTIVE_RETENTION)
+        elif (is_close(crit_len, crit_len_nodata) or
+              is_close(retention_eff_lulc, retention_eff_nodata) or
+              flow_dir == 0):
+            # if it's nodata, effective retention is nodata.
+            effective_retention_raster.set(
+                col, row, effective_retention_nodata)
+        else:
+            working_retention_eff = 0.0
+            has_outflow = False
+            for neighbor in mfd_flow_direction_raster.get_downslope_neighbors(
+                    col, row, skip_oob=False):
+                has_outflow = True
+                if (neighbor.x < 0 or neighbor.x >= n_cols or
+                    neighbor.y < 0 or neighbor.y >= n_rows):
+                    continue
+
+                if neighbor.direction % 2 == 1:
+                    step_size = <float>(cell_size * 1.41421356237)
+                else:
+                    step_size = cell_size
+                # guard against a critical length factor that's 0
+                if crit_len > 0:
+                    current_step_factor = <float>(
+                        exp(-5 * step_size / crit_len))
                 else:
-                    raise Exception("got to a cell that has no outflow!")
-            # search upslope to see if we need to push a cell on the stack
-            # for i in range(8):
-            for neighbor in mfd_flow_direction_raster.get_upslope_neighbors(
-                    global_col, global_row):
-                neighbor_outflow_dir = INFLOW_OFFSETS[neighbor.direction]
-                neighbor_outflow_dir_mask = 1 << neighbor_outflow_dir
-                neighbor_process_flow_dir = <int>(
-                    to_process_flow_directions_raster.get(
+                    current_step_factor = 0.0
+
+                neighbor_effective_retention = (
+                    effective_retention_raster.get(
                         neighbor.x, neighbor.y))
-                if neighbor_process_flow_dir == 0:
-                    # skip, due to loop invariant this must be a nodata pixel
-                    continue
-                if neighbor_process_flow_dir & neighbor_outflow_dir_mask == 0:
-                    # no outflow
-                    continue
-                # mask out the outflow dir that this iteration processed
-                neighbor_process_flow_dir &= ~neighbor_outflow_dir_mask
-                to_process_flow_directions_raster.set(
-                    neighbor.x, neighbor.y, neighbor_process_flow_dir)
-                if neighbor_process_flow_dir == 0:
-                    # if 0 then all downslope have been processed,
-                    # push on stack, otherwise another downslope pixel will
-                    # pick it up
-                    processing_stack.push(neighbor.y * n_cols + neighbor.x)
+
+                # Case 1: downslope neighbor is a stream pixel
+                if neighbor_effective_retention == STREAM_EFFECTIVE_RETENTION:
+                    intermediate_retention = (
+                        retention_eff_lulc * (1 - current_step_factor))
+
+                # Case 2: the current LULC's retention exceeds the neighbor's retention.
+                elif retention_eff_lulc > neighbor_effective_retention:
+                    intermediate_retention = (
+                        (neighbor_effective_retention * current_step_factor) +
+                        (retention_eff_lulc * (1 - current_step_factor)))
+
+                # Case 3: the other 2 cases have not been hit.
+                else:
+                    intermediate_retention = neighbor_effective_retention
+
+                working_retention_eff += (
+                    intermediate_retention * neighbor.flow_proportion)
+
+            if has_outflow:
+                effective_retention_raster.set(
+                    col, row, working_retention_eff)
+            else:
+                raise Exception("got to a cell that has no outflow!")
+        # search upslope to see if we need to push a cell on the stack
+        # for i in range(8):
+        for neighbor in mfd_flow_direction_raster.get_upslope_neighbors(col, row):
+            neighbor_outflow_dir = INFLOW_OFFSETS[neighbor.direction]
+            neighbor_outflow_dir_mask = 1 << neighbor_outflow_dir
+            neighbor_process_flow_dir = <int>(
+                to_process_flow_directions_raster.get(
+                    neighbor.x, neighbor.y))
+            if neighbor_process_flow_dir == 0:
+                # skip, due to loop invariant this must be a nodata pixel
+                continue
+            if neighbor_process_flow_dir & neighbor_outflow_dir_mask == 0:
+                # no outflow
+                continue
+            # mask out the outflow dir that this iteration processed
+            neighbor_process_flow_dir &= ~neighbor_outflow_dir_mask
+            to_process_flow_directions_raster.set(
+                neighbor.x, neighbor.y, neighbor_process_flow_dir)
+            if neighbor_process_flow_dir == 0:
+                # if 0 then all downslope have been processed,
+                # push on stack, otherwise another downslope pixel will
+                # pick it up
+                to_push.append(neighbor.y * n_cols + neighbor.x)
+        return to_push
+
+
+    route(mfd_flow_direction_path, ndr_seed_fn, ndr_route_fn)
+
     to_process_flow_directions_raster.close()
     os.remove(to_process_flow_directions_path)