Example of weighted flow accumulation for NCI NOXN work using pygeoprocessing.routing.

routing-example.py

import logging
import os

import pygeoprocessing
import pygeoprocessing.routing
from osgeo import gdal

logging.basicConfig(level=logging.INFO)
LOGGER = logging.getLogger(__name__)


def doit(dem_path, flow_dir_weights, workspace):
    LOGGER.info('Starting script')
    if not os.path.exists(workspace):
        os.makedirs(workspace)

    aligned_dem_path = os.path.join(workspace, 'aligned_dem.tif')
    aligned_weights_path = os.path.join(workspace, 'aligned_weights.tif')

    # Align the DEM and weights raster
    pygeoprocessing.align_and_resize_raster_stack(
        [dem_path, flow_dir_weights],
        [aligned_dem_path, aligned_weights_path],
        ['bilinear', 'bilinear'],
        pygeoprocessing.get_raster_info(dem_path)['pixel_size'],
        bounding_box_mode='intersection')

    # Flow direction
    flow_direction_path = os.path.join(workspace, 'flow_dir.tif')
    pygeoprocessing.routing.flow_dir_d8(
        (aligned_dem_path, 1), flow_direction_path)

    # Flow accumulation with weights.
    flow_accumulation_path = os.path.join(workspace, 'flow_accumulation.tif')
    pygeoprocessing.routing.flow_accumulation_d8(
        (flow_direction_path, 1), flow_accumulation_path,
        (aligned_weights_path, 1))


if __name__ == '__main__':
    doit('globe_dem_shedsandh1k_DEM.tif', 'n_load.tif', 'workspace')