Source code for hydrobricks.preprocessing.catchment_discretization

from __future__ import annotations

import itertools
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np

from hydrobricks._exceptions import (
    ConfigurationError,
    DataError,
    DependencyError,
    ModelError,
)
from hydrobricks._optional import HAS_GEOPANDAS, HAS_PYPROJ, gpd, rasterio

if TYPE_CHECKING:
    from hydrobricks.catchment import Catchment


[docs] class CatchmentDiscretization: """ Class to handle the discretization of catchments. """ def __init__(self, catchment: Catchment) -> None: """ Initialize the Discretization class. Parameters ---------- catchment The catchment object. """ self.catchment: Catchment = catchment
[docs] def create_elevation_bands( self, method: str = "equal_intervals", number: int = 100, distance: int = 50, min_elevation: int | None = None, max_elevation: int | None = None, ) -> None: """ Construction of the elevation bands based on the chosen method. Creates hydro units based on elevation bands using either equal-interval contours or quantile-based discretization. Results are stored in the catchment's map_unit_ids and hydro_units objects. Parameters ---------- method The method to build the elevation bands: 'equal_intervals' = fixed contour intervals (needs 'distance' parameter) 'quantiles' = quantiles of the catchment area (same surface; needs 'number' parameter) number Number of bands to create for the 'quantiles' method. distance Distance (m) between the contour lines for the 'equal_intervals' method. min_elevation Minimum elevation of the elevation bands (to homogenize between runs). max_elevation Maximum elevation of the elevation bands (to homogenize between runs). """ self.discretize_by( "elevation", method, number, distance, min_elevation, max_elevation )
[docs] def discretize_by( self, criteria: str, elevation_method: str = "equal_intervals", elevation_number: int = 100, elevation_distance: int = 100, min_elevation: int | None = None, max_elevation: int | None = None, slope_method: str = "equal_intervals", slope_number: int = 6, slope_distance: int = 15, min_slope: int = 0, max_slope: int = 90, radiation_method: str = "equal_intervals", radiation_number: int = 5, radiation_distance: int = 50, min_radiation: int | None = None, max_radiation: int | None = None, sub_catchments: str | Path | None = None, ) -> None: """ Construction of the elevation bands based on the chosen method. Discretizes the catchment into hydro units based on multiple criteria (elevation, slope, aspect, radiation, sub_catchments). Creates all combinations of specified criteria and populates the map_unit_ids and hydro_units objects. Parameters ---------- criteria The criteria to use to discretize the catchment (can be combined): 'elevation' = elevation bands 'aspect' = aspect according to the cardinal directions (4 classes) 'radiation' = potential radiation (Hock, 1999) 'slope' = slope in degrees 'sub_catchments' = membership to a sub-catchment polygon (keeps a hydro unit from spanning spatially-distant areas); requires the 'sub_catchments' argument elevation_method The method to build the elevation bands: 'equal_intervals' = fixed contour intervals (provide the 'elevation_distance' parameter) 'quantiles' = quantiles of the catchment area (same surface; provide the 'elevation_number' parameter) elevation_number Number of elevation bands to create for the 'quantiles' method. elevation_distance Distance (m) between the contour lines for the 'equal_intervals' method. min_elevation Minimum elevation of the elevation bands (to homogenize between runs). max_elevation Maximum elevation of the elevation bands (to homogenize between runs). slope_method The method to build the slope categories: 'equal_intervals' = fixed slope intervals (needs 'slope_distance' parameter) 'quantiles' = quantiles of the catchment area (same surface; provide the 'slope_number' parameter) slope_number Number of slope bands to create for the 'quantiles' method. slope_distance Distance (degrees) between the slope lines for the 'equal_intervals' method. min_slope Minimum slope of the slope bands (to homogenize between runs). max_slope Maximum slope of the slope bands (to homogenize between runs). radiation_method The method to build the radiation categories: 'equal_intervals' = fixed radiation intervals (provide the 'radiation_distance' parameter) 'quantiles' = quantiles of the catchment area (same surface; provide the 'radiation_number' parameter) radiation_number Number of radiation bands to create when using the 'quantiles' method. radiation_distance Distance (W/m2) in terms of radiation for the 'equal_intervals' method. min_radiation Minimum radiation of the radiation bands (to homogenize between runs). max_radiation Maximum radiation of the radiation bands (to homogenize between runs). sub_catchments Path to a vector file of sub-catchment polygons. Required when 'sub_catchments' is in ``criteria``. Each polygon defines a sub-catchment; a DEM cell is assigned to the first polygon covering it. The sub-catchment polygons must cover the whole catchment: if any catchment cell is left uncovered, a ``DataError`` is raised (rather than silently dropping that area from the hydro units). """ if not HAS_PYPROJ: raise DependencyError( "pyproj is required for catchment discretization.", package_name="pyproj", operation="CatchmentDiscretization.create_hydro_units", install_command="pip install pyproj", ) if isinstance(criteria, str): criteria = [criteria] if ( self.catchment.topography.slope is None or self.catchment.topography.aspect is None ): self.catchment.topography.calculate_slope_aspect() if ( "radiation" in criteria and self.catchment.solar_radiation.mean_annual_radiation is None ): raise ModelError("Please first compute the radiation.") self.map_unit_ids = np.zeros(self.catchment.dem_data.shape) # Create a dict to store the criteria criteria_dict = {} if "elevation" in criteria: criteria_dict["elevation"] = [] if elevation_method == "equal_intervals": dist = elevation_distance if min_elevation is None: min_elevation = np.nanmin(self.catchment.dem_data) min_elevation = min_elevation - (min_elevation % dist) if max_elevation is None: max_elevation = np.nanmax(self.catchment.dem_data) max_elevation = max_elevation + (dist - max_elevation % dist) elevations = np.arange(min_elevation, max_elevation + dist, dist) for i in range(len(elevations) - 1): criteria_dict["elevation"].append(elevations[i : i + 2]) elif elevation_method == "quantiles": elevations = np.nanquantile( self.catchment.dem_data, np.linspace(0, 1, num=elevation_number) ) for i in range(len(elevations) - 1): criteria_dict["elevation"].append(elevations[i : i + 2]) else: raise ConfigurationError( "Unknown elevation band creation method.", item_name="elevation_method", item_value=elevation_method, reason="Invalid method value", ) if "slope" in criteria: criteria_dict["slope"] = [] if slope_method == "equal_intervals": dist = slope_distance if min_slope is None: min_slope = np.nanmin(self.catchment.topography.slope) min_slope = min_slope - (min_slope % dist) if max_slope is None: max_slope = np.nanmax(self.catchment.topography.slope) max_slope = max_slope + (dist - max_slope % dist) slopes = np.arange(min_slope, max_slope + dist, dist) for i in range(len(slopes) - 1): criteria_dict["slope"].append(slopes[i : i + 2]) elif slope_method == "quantiles": slopes = np.nanquantile( self.catchment.topography.slope, np.linspace(0, 1, num=slope_number) ) for i in range(len(slopes) - 1): criteria_dict["slope"].append(slopes[i : i + 2]) else: raise ConfigurationError( "Unknown slope band creation method.", item_name="slope_method", item_value=slope_method, reason="Invalid method value", ) if "aspect" in criteria: criteria_dict["aspect"] = ["N", "E", "S", "W"] if "radiation" in criteria: if self.catchment.solar_radiation.mean_annual_radiation is None: raise ConfigurationError( "No radiation data available. Compute the radiation first.", reason="Missing required radiation data", ) criteria_dict["radiation"] = [] if radiation_method == "equal_intervals": dist = radiation_distance if min_radiation is None: min_radiation = np.nanmin( self.catchment.solar_radiation.mean_annual_radiation ) min_radiation = min_radiation - (min_radiation % dist) if max_radiation is None: max_radiation = np.nanmax( self.catchment.solar_radiation.mean_annual_radiation ) max_radiation = max_radiation + (dist - max_radiation % dist) radiations = np.arange(min_radiation, max_radiation + dist, dist) for i in range(len(radiations) - 1): criteria_dict["radiation"].append(radiations[i : i + 2]) elif radiation_method == "quantiles": radiations = np.nanquantile( self.catchment.solar_radiation.mean_annual_radiation, np.linspace(0, 1, num=radiation_number), ) for i in range(len(radiations) - 1): criteria_dict["radiation"].append(radiations[i : i + 2]) else: raise ConfigurationError( "Unknown radiation band creation method. " "Only 'equal_intervals' and 'quantiles' are supported.", item_name="radiation_method", item_value=radiation_method, reason="Invalid method value", ) sub_catchments_map = None if "sub_catchments" in criteria: if sub_catchments is None: raise ConfigurationError( "The 'sub_catchments' criterion requires the 'sub_catchments' " "argument (path to the sub-catchment polygons).", item_name="sub_catchments", item_value=None, reason="Missing sub-catchment file", ) sub_catchments_map, n_sub = self._build_sub_catchments_map(sub_catchments) criteria_dict["sub_catchments"] = list(range(1, n_sub + 1)) res_elevation = [] res_elevation_min = [] res_elevation_max = [] res_slope = [] res_slope_min = [] res_slope_max = [] res_aspect_class = [] res_radiation = [] res_radiation_min = [] res_radiation_max = [] res_sub_catchment = [] combinations = list(itertools.product(*criteria_dict.values())) combinations_keys = list(criteria_dict.keys()) # Position of each criterion within a combination tuple (computed once). criterion_positions = {name: pos for pos, name in enumerate(combinations_keys)} unit_id = 1 for combination in combinations: mask_unit = np.ones(self.catchment.dem_data.shape, dtype=bool) # Mask nan values mask_unit[np.isnan(self.catchment.dem_data)] = False for criterion_name, criterion in zip(combinations_keys, combination): if criterion_name == "elevation": mask_elev = np.logical_and( self.catchment.dem_data >= criterion[0], self.catchment.dem_data < criterion[1], ) mask_unit = np.logical_and(mask_unit, mask_elev) elif criterion_name == "slope": mask_slope = np.logical_and( self.catchment.topography.slope >= criterion[0], self.catchment.topography.slope < criterion[1], ) mask_unit = np.logical_and(mask_unit, mask_slope) elif criterion_name == "aspect": if criterion == "N": mask_aspect = np.logical_or( np.logical_and( self.catchment.topography.aspect >= 315, self.catchment.topography.aspect <= 360, ), np.logical_and( self.catchment.topography.aspect >= 0, self.catchment.topography.aspect < 45, ), ) elif criterion == "E": mask_aspect = np.logical_and( self.catchment.topography.aspect >= 45, self.catchment.topography.aspect < 135, ) elif criterion == "S": mask_aspect = np.logical_and( self.catchment.topography.aspect >= 135, self.catchment.topography.aspect < 225, ) elif criterion == "W": mask_aspect = np.logical_and( self.catchment.topography.aspect >= 225, self.catchment.topography.aspect < 315, ) else: raise ConfigurationError( "Unknown aspect value.", item_name="aspect", item_value=criterion, reason="Invalid aspect direction", ) mask_unit = np.logical_and(mask_unit, mask_aspect) elif criterion_name == "radiation": radiation = self.catchment.solar_radiation.mean_annual_radiation mask_radiation = np.logical_and( radiation >= criterion[0], radiation < criterion[1] ) mask_unit = np.logical_and(mask_unit, mask_radiation) elif criterion_name == "sub_catchments": mask_sub = sub_catchments_map == criterion mask_unit = np.logical_and(mask_unit, mask_sub) # If the unit is empty, skip it if np.count_nonzero(mask_unit) == 0: continue # Check that all cells in unit_ids are 0 assert np.count_nonzero(self.map_unit_ids[mask_unit]) == 0 # Set the unit id self.map_unit_ids[mask_unit] = unit_id # Set the mean elevation of the unit if elevation is a criterion if "elevation" in criteria_dict.keys(): elevations = combination[criterion_positions["elevation"]] res_elevation.append(round(float(np.mean(elevations)), 2)) res_elevation_min.append(round(float(elevations[0]), 2)) res_elevation_max.append(round(float(elevations[1]), 2)) # Set the mean slope of the unit if slope is a criterion if "slope" in criteria_dict.keys(): slopes = combination[criterion_positions["slope"]] res_slope.append(round(float(np.mean(slopes)), 2)) res_slope_min.append(round(float(slopes[0]), 2)) res_slope_max.append(round(float(slopes[1]), 2)) # Get the aspect class if aspect is a criterion if "aspect" in criteria_dict.keys(): res_aspect_class.append(combination[criterion_positions["aspect"]]) # Get the radiation class if radiation is a criterion if "radiation" in criteria_dict.keys(): radiations = combination[criterion_positions["radiation"]] res_radiation.append(round(float(np.mean(radiations)), 2)) res_radiation_min.append(round(float(radiations[0]), 2)) res_radiation_max.append(round(float(radiations[1]), 2)) # Get the sub-catchment id if sub_catchments is a criterion if "sub_catchments" in criteria_dict.keys(): sub_id = combination[criterion_positions["sub_catchments"]] res_sub_catchment.append(str(sub_id)) unit_id += 1 self.catchment.map_unit_ids = self.map_unit_ids.astype(rasterio.uint16) if res_elevation: self.catchment.hydro_units.add_property(("elevation", "m"), res_elevation) self.catchment.hydro_units.add_property( ("elevation_min", "m"), res_elevation_min ) self.catchment.hydro_units.add_property( ("elevation_max", "m"), res_elevation_max ) if res_slope: self.catchment.hydro_units.add_property(("slope", "deg"), res_slope) self.catchment.hydro_units.add_property(("slope_min", "deg"), res_slope_min) self.catchment.hydro_units.add_property(("slope_max", "deg"), res_slope_max) if res_aspect_class: self.catchment.hydro_units.add_property( ("aspect_class", "-"), res_aspect_class ) if res_radiation: self.catchment.hydro_units.add_property( ("radiation", "W/m2"), res_radiation ) self.catchment.hydro_units.add_property( ("radiation_min", "W/m2"), res_radiation_min ) self.catchment.hydro_units.add_property( ("radiation_max", "W/m2"), res_radiation_max ) if res_sub_catchment: self.catchment.hydro_units.add_property( ("sub_catchment", "-"), res_sub_catchment ) self.catchment.initialize_land_cover_fractions() self.catchment.get_hydro_units_attributes() self.catchment.hydro_units.populate_bounded_instance()
def _build_sub_catchments_map( self, sub_catchments: str | Path ) -> tuple[np.ndarray, int]: """Rasterize sub-catchment polygons onto the DEM grid. Reads the sub-catchment polygons, reprojects them to the catchment CRS and assigns each DEM cell to the first polygon covering it (ids 1..N in feature order). The sub-catchment polygons must cover the whole catchment: if any catchment cell is left uncovered, a ``DataError`` is raised (rather than silently dropping that area from the hydro units). Parameters ---------- sub_catchments Path to the vector file of sub-catchment polygons. Returns ------- A tuple ``(sub_catchments_map, n_sub)`` with the per-cell integer id grid over the DEM and the number of sub-catchments. Raises ------ DataError If part of the catchment is not covered by any sub-catchment polygon. """ if not HAS_GEOPANDAS: raise DependencyError( "geopandas is required to discretize by sub-catchments.", package_name="geopandas", operation="CatchmentDiscretization.discretize_by", install_command="pip install geopandas", ) gdf = gpd.read_file(sub_catchments) gdf = gdf.to_crs(self.catchment.crs) if gdf.empty: raise DataError( "The sub-catchments file contains no polygons.", data_type="sub-catchments", reason="Empty file", ) sub_map = np.zeros(self.catchment.dem_data.shape, dtype=int) assigned = np.zeros(sub_map.shape, dtype=bool) for i in range(len(gdf)): subset = gdf.iloc[[i]] masked = self.catchment.mask_dem(subset, nodata=-9999, all_touched=True) present = (masked != -9999) & ~assigned sub_map[present] = i + 1 assigned |= present # Ensure the sub-catchments cover the whole catchment. Catchment cells are # those with a valid elevation; any such cell left with id 0 would be silently # dropped from the hydro units, so fail instead. inside_catchment = ~np.isnan(self.catchment.dem_data) uncovered = inside_catchment & (sub_map == 0) n_uncovered = int(np.count_nonzero(uncovered)) if n_uncovered > 0: uncovered_area = n_uncovered * self.catchment.get_dem_pixel_area() raise DataError( f"The sub-catchment polygons do not cover the whole catchment: " f"{n_uncovered} cell(s) ({uncovered_area:.0f} m²) are left uncovered. " f"Provide sub-catchments that tile the entire catchment.", data_type="sub-catchments", reason="Incomplete catchment coverage", ) return sub_map, len(gdf)