"""Digital elevation model (DEM) download and preprocessing.
Role: Provide clipped DEM and derived slope for the study area.
"""
import requests as r
import os
import numpy as np
from bakaano.utils import Utils
import zipfile
import matplotlib.pyplot as plt
import rasterio
from scipy.ndimage import convolve
[docs]
class DEM:
def __init__(self, working_dir, study_area, local_data=False, local_data_path=None):
"""
Role: Download/prepare DEM inputs for the basin.
Initialize a DEM (Digital Elevation Model) object.
Args:
working_dir (str): The parent working directory where files and outputs will be stored.
study_area (str): The path to the shapefile of the river basin or watershed.
local_data (bool, optional): Flag indicating whether to use local data instead of downloading new data. Defaults to False.
local_data_path (str, optional): Path to the local DEM geotiff tile if `local_data` is True. Defaults to None. Local DEM provided should be in the GCS WGS84 or EPSG:4326 coordinate system
Methods
-------
__init__(working_dir, study_area, local_data=False, local_data_path=None):
Initializes the DEM object with project details.
get_dem_data():
Download DEM data.
preprocess():
Preprocess downloaded data.
plot_dem():
Plot DEM data
Returns:
A DEM geotiff clipped to the study area extent to be stored in "{working_dir}/elevation" directory
"""
self.study_area = study_area
self.working_dir = working_dir
os.makedirs(f'{self.working_dir}/elevation', exist_ok=True)
self.uw = Utils(self.working_dir, self.study_area)
self.out_path = f'{self.working_dir}/elevation/dem_clipped.tif'
self.local_data = local_data
self.local_data_path = local_data_path
[docs]
def get_dem_data(self):
"""Download or ingest DEM data and clip to the study area.
If ``local_data`` is False, downloads HydroSHEDS 30s DEM and clips it.
If ``local_data`` is True, clips the provided local GeoTIFF.
Returns:
None. Writes ``dem_clipped.tif`` to ``{working_dir}/elevation``.
"""
if self.local_data is False:
if not os.path.exists(self.out_path):
url = 'https://data.hydrosheds.org/file/hydrosheds-v1-dem/hyd_glo_dem_30s.zip'
local_filename = f'{self.working_dir}/elevation/hyd_glo_dem_30s.zip'
uw = Utils(self.working_dir, self.study_area)
uw.get_bbox('EPSG:4326')
response = r.get(url, stream=True)
if response.status_code == 200:
with open(local_filename, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
print(f"File downloaded successfully and saved as '{local_filename}'")
else:
print(f"Failed to download the file. HTTP status code: {response.status_code}")
extraction_path = f'{self.working_dir}/elevation' # Directory where files will be extracted
# Open and extract the zip file
with zipfile.ZipFile(local_filename, 'r') as zip_ref:
zip_ref.extractall(extraction_path)
print(f"Files extracted to '{extraction_path}'")
self.preprocess()
else:
print(f" - DEM data already exists in {self.working_dir}/elevation; skipping download.")
else:
#print(f" - Local DEM data already provided")
try:
if not self.local_data_path:
raise ValueError("Local data path must be provided when 'local_data' is set to True.")
if not os.path.exists(self.local_data_path):
raise FileNotFoundError(f"The specified local DEM file '{self.local_data_path}' does not exist.")
if not self.local_data_path.endswith('.tif'):
raise ValueError("The local DEM file must be a GeoTIFF (.tif) file.")
self.uw.clip(raster_path=self.local_data_path, out_path=self.out_path, save_output=True)
except (ValueError, FileNotFoundError) as e:
print(f"Error: {e}")
[docs]
def preprocess(self):
"""Clip the DEM and derive slope output if missing.
Returns:
None. Writes ``dem_clipped.tif`` and ``slope_clipped.tif``.
"""
dem = f'{self.working_dir}/elevation/hyd_glo_dem_30s.tif'
self.uw.clip(raster_path=dem, out_path=self.out_path, save_output=True, crop_type=False)
#self.uw.clip(raster_path=dem, out_path=self.out_path_uncropped, save_output=True, crop_type=False)
slope_name = f'{self.working_dir}/elevation/slope_clipped.tif'
if not os.path.exists(slope_name):
self.compute_slope_percent_riserun()
[docs]
def compute_slope_percent_riserun(self):
"""Compute slope (percent rise/run) from the clipped DEM.
Returns:
None. Writes ``slope_clipped.tif`` to ``{working_dir}/elevation``.
"""
with rasterio.open(self.out_path) as src:
elevation = src.read(1).astype(float)
profile = src.profile.copy()
res_x, res_y = src.res
cellsize = np.mean([abs(res_x), abs(res_y)]) *100000
nodata = src.nodata
# Handle NoData
if nodata is not None:
elevation[elevation == nodata] = np.nan
# Fill NaNs for convolution
elevation_filled = np.where(np.isnan(elevation), np.nanmean(elevation), elevation)
# Horn kernels (adjusted for cellsize in meters)
kernel_x = np.array([[-1, 0, 1],
[-2, 0, 2],
[-1, 0, 1]]) / (8 * cellsize)
kernel_y = np.array([[1, 2, 1],
[0, 0, 0],
[-1, -2, -1]]) / (8 * cellsize)
# Compute gradients
dzdx = convolve(elevation_filled, kernel_x, mode='nearest')
dzdy = convolve(elevation_filled, kernel_y, mode='nearest')
# Slope in percent rise/run
slope_percent = np.sqrt(dzdx**2 + dzdy**2) * 100
slope_percent[np.isnan(elevation)] = np.nan
# Update metadata
profile.update(dtype=rasterio.float32, count=1, nodata=np.nan)
# Save to GeoTIFF
output_path = f'{self.working_dir}/elevation/slope_clipped.tif'
with rasterio.open(output_path, 'w', **profile) as dst:
dst.write(slope_percent.astype(np.float32), 1)
print(f"Slope saved to: {output_path}")
[docs]
def plot_dem(self):
"""Plot the clipped DEM within the study area.
Returns:
None. Displays a matplotlib plot.
"""
dem_data = self.uw.clip(raster_path=self.out_path, out_path=None, save_output=False, crop_type=True)[0]
dem_data = np.where(dem_data > 0, dem_data, np.nan)
dem_data = np.where(dem_data < 32000, dem_data, np.nan)
plt.imshow(dem_data, cmap='terrain')
plt.colorbar()