Skip to content

Instantly share code, notes, and snippets.

@phargogh
Last active August 29, 2023 05:41
Show Gist options
  • Save phargogh/f673a1c1c62d930cbb30cc62f5d7a92d to your computer and use it in GitHub Desktop.
Save phargogh/f673a1c1c62d930cbb30cc62f5d7a92d to your computer and use it in GitHub Desktop.
Example of pygeoprocessing-based D8 routing, including watershed delineation.
import os
import shutil
import numpy
import pygeoprocessing
import pygeoprocessing.routing
from osgeo import gdal
def doit(dem_path, flow_dir_weights, watershed_source_vector, workspace):
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))
# Create streams layer
flow_accumulation_nodata = pygeoprocessing.get_raster_info(
flow_accumulation_path)['nodata'][0]
streams_nodata = 255
flow_threshold = 1000 # This is the TFA
def _filter_streams(flow_accumulation):
out_array = numpy.full(flow_accumulation.shape, streams_nodata,
dtype=numpy.uint8)
valid_pixels = (flow_accumulation != flow_accumulation_nodata)
out_array[valid_pixels & (flow_accumulation >= flow_threshold)] = 1
out_array[valid_pixels & (flow_accumulation < flow_threshold)] = 0
return out_array
streams_raster_path = os.path.join(workspace, 'streams.tif')
pygeoprocessing.raster_calculator(
[(flow_accumulation_path, 1)],
_filter_streams,
streams_raster_path,
gdal.GDT_Byte,
streams_nodata)
# Delineate watersheds
target_watersheds_path = os.path.join(workspace, 'watersheds.gpkg')
pygeoprocessing.routing.delineate_watersheds_d8(
(flow_direction_path, 1),
watershed_source_vector,
target_watersheds_path)
if __name__ == '__main__':
# This function call assumes that everything is in the current working directory:
# ./
# ./my_dem.tif
# ./weights.tif
# ./points.gpkg
# ./workspace/
doit('my_dem.tif', 'weights.tif', 'points.gpkg', 'workspace')
# You could also use absolute paths for any of these.
# Relative and absolute paths can also be mixed and matched.
# For example:
doit(
'/Users/my_username/Documents/my_dem.tif',
'/mnt/data/my_project/weights.tif',
'/home/points.gpkg',
'./workspace'
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment