DTW- Python implementation

DTW.py

import rasterio
import numpy as np
from skimage.graph import MCP

""" Adopted from Schönauer, M., Maack, J., 2021. R-code for calculating depth-to-water (DTW) maps using GRASS GIS. Zenodo. doi: 10.5281/zenodo.5638517 
adopted to Python by Ardo.J, 2024"""

# Setting Paths, Assuming that these layers are calculated outside of this code.
dem_path = "breached.tif"
slope_path = "slope.tif"
fd8_path = "fd8.tif"
output_path_template = "DTW_FIA_{fia}_ha.tif"
#to-do include water body polygons


# Run the function for the desired FIAs. 
# For calculating the  flow paths, it is necessary to define a threshold (t) for the minimal flow initiation  area  (FIA)
# FIA numbers should be higher for well drained and rocky soils, and lower for loamy and clayey soils. 
# Extended droughts and frost periods would also increase FIA
fia_values = [1, 2, 4, 6, 8]

# Read DEM metadata (for transform information)
with rasterio.open(dem_path) as src:
    dem_meta = src.meta

# Read flow accumulation
with rasterio.open(fd8_path) as src:
    accumulation = src.read(1)

# Read slope
with rasterio.open(slope_path) as src:
    slope = src.read(1)

# Function to calculate DTW
def calc_dtw(fia):
    # Calculate flow initiation area threshold
    t = fia * 10000 / (dem_meta['transform'][0] ** 2)
    print(f'Treshold is:{t}')

    # Create binary flow lines
    flow_lines = np.where(accumulation >= t, 1, 0)

    # Prepare a mask for flow lines
    flow_lines_coords = np.argwhere(flow_lines == 1)

    # Initialize cost array with high values 
    cost = np.full_like(slope, np.inf)

    # Ensure all coordinates are within bounds
    valid_coords = [tuple(coord) for coord in flow_lines_coords if 0 <= coord[0] < slope.shape[0] and 0 <= coord[1] < slope.shape[1]]

    # Set cost of flow line cells to 0
    for coord in valid_coords:
        cost[coord[0], coord[1]] = 0

    # Calculate least-cost paths using MCP
    mcp = MCP(slope)
    costs, _ = mcp.find_costs(valid_coords)

    # Correct cost grid to achieve [m]
    dtw_xha = costs * dem_meta['transform'][0] / 100

    # Update metadata for compression
    dem_meta.update(compress='lzw', dtype=rasterio.float32)

    # Save the raster
    dtw_path = output_path_template.format(fia=fia)
    with rasterio.open(dtw_path, 'w', **dem_meta) as dst:
        dst.write(dtw_xha, 1)


for fia in fia_values:
    calc_dtw(fia)