"""Delineate catchments and basin maps for mHM/mRM setups.
The module builds hydrologic grids from flow direction or DEM data, supports
global and local domains, single- and multi-gauge delineation, outlet matching
by area or shape, and writes basin ID, idgauges, mask, gauge metadata, and
basin shape outputs.
Authors
-------
- Simon Lüdke
"""
import csv
import logging
from pathlib import Path
from typing import Optional
import numpy as np
import pyflwdir
import xarray as xr
from joblib import Parallel, delayed
from scipy.ndimage import binary_dilation
from mhm_tools.common.constants import NC_ENCODE_MASK
from mhm_tools.common.file_handler import (
get_coord_values,
get_xarray_ds_from_file,
write_xarray_to_file,
)
from mhm_tools.common.logger import ErrorLogger, log_arguments
from mhm_tools.common.netcdf import generate_bounds
from mhm_tools.common.provenance import apply_output_provenance
from mhm_tools.common.resolution_handler import Resolution
from mhm_tools.common.utils import (
coord_to_index,
cut_to_filled_area,
distance_100m_units,
find_best_gauge_location_by_area,
get_upscaling_factor,
)
from mhm_tools.common.xarray_utils import get_dtype
from mhm_tools.pre.create_id_gauges import write_gauge_id
logger = logging.getLogger(__name__)
# GLOBAL VARIABLES
FDIR_FILLVALUE = {"d8": 247, "ldd": 255}
FDIR_SINKVALUE = {"d8": 0, "ldd": 5}
FACC_FILLVALUE = 0
FILLVALUE = -9999
OUTPUT_VARIABLES = ("flwdir", "basin", "uparea_grid", "upgrid", "grdare", "elevtn")
GAUGE_INFO_COLUMNS = (
"id",
"lon",
"lat",
"lon_old",
"lat_old",
"distance",
"area",
"old_area",
"area_error",
"score",
"shape_error",
"method",
)
# use d8 for basinex, ldd for mRM version in Ulysses
OUTPUT_FTYPE = "ldd"
CUTOFF_THRESHOLD = 175
def _shape_crs(is_latlon):
"""Return a CRS string for shapefile operations based on lat/lon usage."""
crs = "EPSG:4326" if is_latlon else None
logger.debug(f"Resolved shape CRS to {crs}")
return crs
def _find_shape_file(shape_folder, gauge_id):
"""Find a shapefile in a folder that contains the gauge_id in its name."""
if not shape_folder or gauge_id is None:
return None
shape_dir = Path(shape_folder)
if not shape_dir.is_dir():
return None
matching_shapes = sorted(shape_dir.glob(f"*{gauge_id}*.shp"))
if not matching_shapes:
return None
if len(matching_shapes) > 1:
logger.warning(
f"Multiple shapefiles matched gauge_id {gauge_id} in {shape_dir}. "
f"Using {matching_shapes[0].name}"
)
logger.debug(f"Using shapefile {matching_shapes[0]} for gauge_id {gauge_id}")
return matching_shapes[0]
def _read_reference_shape(shape_folder, gauge_id, latlon):
"""Read a reference catchment shape and normalize its CRS."""
shape_path = _find_shape_file(shape_folder, gauge_id)
if shape_path is None:
return None, None
try:
import geopandas as gpd
except Exception as exc:
error_msg = "geopandas is required for shape-based gauge correction."
with ErrorLogger(logger):
raise ImportError(error_msg) from exc
reference_shape = gpd.read_file(shape_path)
crs = _shape_crs(latlon)
if crs is not None:
if reference_shape.crs is None:
reference_shape = reference_shape.set_crs(crs)
elif str(reference_shape.crs) != str(crs):
reference_shape = reference_shape.to_crs(crs)
logger.debug(f"Loaded reference shape {shape_path}")
return reference_shape, shape_path
def _as_affine(affine_transform):
try:
from rasterio.transform import Affine
except Exception as exc:
error_msg = "rasterio is required for basin shapefile operations."
with ErrorLogger(logger):
raise ImportError(error_msg) from exc
if isinstance(affine_transform, Affine):
return affine_transform
try:
return Affine(*affine_transform)
except Exception:
return Affine.from_gdal(*affine_transform)
def _rasterize_shape_to_mask(reference_shape, grid_shape, affine_transform):
"""Rasterize a reference catchment shape to the active grid."""
try:
from rasterio import features
except Exception as exc:
error_msg = "rasterio is required for shape-based gauge correction."
with ErrorLogger(logger):
raise ImportError(error_msg) from exc
if reference_shape is None or reference_shape.empty:
return None
geometries = [
geom
for geom in reference_shape.geometry
if geom is not None and not getattr(geom, "is_empty", False)
]
if not geometries:
return None
mask = features.rasterize(
((geom, 1) for geom in geometries),
out_shape=grid_shape,
transform=_as_affine(affine_transform),
fill=0,
dtype="uint8",
all_touched=True,
)
return mask.astype(bool)
def _cell_polygon(affine_transform, row_idx, col_idx):
"""Create a polygon for one raster cell from an affine transform."""
try:
from shapely.geometry import Polygon
except Exception as exc:
error_msg = "shapely is required for shape/grid overlap calculation."
with ErrorLogger(logger):
raise ImportError(error_msg) from exc
corners = [
affine_transform * (col_idx, row_idx),
affine_transform * (col_idx + 1, row_idx),
affine_transform * (col_idx + 1, row_idx + 1),
affine_transform * (col_idx, row_idx + 1),
]
return Polygon(corners)
def _calculate_shape_area_on_grid(
reference_shape, cell_area, grid_shape, affine_transform
):
"""Calculate shape-covered area on a grid in km2.
Fractional cell overlap is used when shapely operations are available. If that
fails, fall back to summing cells from the existing rasterized shape mask.
"""
if reference_shape is None or reference_shape.empty or cell_area is None:
return None
cell_area = np.asarray(cell_area)
if cell_area.shape != tuple(grid_shape):
msg = (
f"cell_area shape {cell_area.shape} does not match grid shape "
f"{tuple(grid_shape)} for shape area calculation."
)
with ErrorLogger(logger):
raise ValueError(msg)
affine_transform = _as_affine(affine_transform)
shape_mask = _rasterize_shape_to_mask(
reference_shape,
grid_shape,
affine_transform,
)
if shape_mask is None or not np.any(shape_mask):
return None
try:
reference_geom = _geometry_union(reference_shape.geometry)
if reference_geom.is_empty:
return None
area = 0.0
for row_idx, col_idx in zip(*np.where(shape_mask)):
cell_geom = _cell_polygon(affine_transform, int(row_idx), int(col_idx))
cell_geom_area = cell_geom.area
if cell_geom_area <= 0:
continue
overlap_area = reference_geom.intersection(cell_geom).area
overlap_fraction = np.clip(overlap_area / cell_geom_area, 0.0, 1.0)
area += float(cell_area[row_idx, col_idx]) * float(overlap_fraction)
if area > 0:
return float(area)
except Exception as exc:
logger.debug(
"Could not calculate fractional shape/grid overlap; falling back to "
f"rasterized-cell area sum: {exc}"
)
area = float(np.sum(cell_area[shape_mask]))
return area if area > 0 else None
def _vectorize_mask_to_gdf(basin_mask, affine_transform, crs, value_name="basin"):
"""Vectorize a basin mask into a GeoDataFrame."""
try:
import geopandas as gpd
from rasterio import features
except Exception as exc:
error_msg = (
"geopandas and rasterio are required for basin shapefile operations."
"Please install them using `conda install geopandas`"
)
with ErrorLogger(logger):
raise ImportError(error_msg) from exc
affine_transform = _as_affine(affine_transform)
data = basin_mask.astype(np.uint8)
# Extract polygons for non-zero cells
feats_gen = features.shapes(
data,
mask=data.astype(bool),
transform=affine_transform,
connectivity=8,
)
features_list = [
{"geometry": geom, "properties": {value_name: 1}} for geom, _ in feats_gen
]
if not features_list:
logger.debug("No vectorizable features found in basin mask.")
return gpd.GeoDataFrame(
columns=[value_name, "geometry"], geometry="geometry", crs=crs
)
gdf = gpd.GeoDataFrame.from_features(features_list, crs=crs)
gdf[value_name] = gdf[value_name].astype(np.uint8)
logger.debug(f"Vectorized basin mask into {len(gdf)} features.")
return gdf
def _shape_iou(reference_gdf, candidate_gdf):
"""Compute intersection-over-union for two GeoDataFrames."""
if (
reference_gdf is None
or candidate_gdf is None
or reference_gdf.empty
or candidate_gdf.empty
):
return 0.0
reference_geom = _geometry_union(reference_gdf.geometry)
candidate_geom = _geometry_union(candidate_gdf.geometry)
if reference_geom.is_empty or candidate_geom.is_empty:
return 0.0
union_area = reference_geom.union(candidate_geom).area
if union_area == 0:
return 0.0
intersection_area = reference_geom.intersection(candidate_geom).area
iou = float(intersection_area / union_area)
logger.debug(f"Computed shape IoU: {iou:.4f}")
return iou
def _geometry_union(geometry):
if hasattr(geometry, "union_all"):
return geometry.union_all()
return geometry.unary_union
def _coords_from_transform(affine_transform, grid_shape):
"""Derive coordinate arrays from an affine transform and grid shape."""
try:
from rasterio.transform import Affine, xy
except Exception as exc:
error_msg = "rasterio is required to derive coordinates from transform."
with ErrorLogger(logger):
raise ImportError(error_msg) from exc
if not isinstance(affine_transform, Affine):
try:
affine_transform = Affine(*affine_transform)
except Exception:
affine_transform = Affine.from_gdal(*affine_transform)
n_rows, n_cols = grid_shape
col_indices = np.arange(n_cols)
row_indices = np.arange(n_rows)
x_coords = np.array(
[xy(affine_transform, 0, col, offset="center")[0] for col in col_indices]
)
y_coords = np.array(
[xy(affine_transform, row, 0, offset="center")[1] for row in row_indices]
)
logger.debug(
f"Derived coordinate arrays from transform with lengths {len(x_coords)} "
f"(x) and {len(y_coords)} (y)."
)
return x_coords, y_coords
def _combine_shape_bounds(bounds_list):
"""Combine multiple (minx, miny, maxx, maxy) tuples into a single bounding box."""
if not bounds_list:
return None
min_x = min(b[0] for b in bounds_list)
min_y = min(b[1] for b in bounds_list)
max_x = max(b[2] for b in bounds_list)
max_y = max(b[3] for b in bounds_list)
return min_x, min_y, max_x, max_y
def _buffer_bounds(bounds, latlon):
"""Apply a buffer to bounds (degrees if latlon, otherwise meters)."""
if bounds is None:
return None
buffer_value = 1.0 if latlon else 100_000.0
min_x, min_y, max_x, max_y = bounds
return (
min_x - buffer_value,
min_y - buffer_value,
max_x + buffer_value,
max_y + buffer_value,
)
def _slices_from_bounds(bounds, lat_values):
"""Create lon/lat slices from bounds, respecting latitude ordering."""
if bounds is None:
return None
min_x, min_y, max_x, max_y = bounds
if lat_values[0] > lat_values[-1]:
lat_slice = slice(max_y, min_y)
else:
lat_slice = slice(min_y, max_y)
lon_slice = slice(min_x, max_x)
return {"lat": lat_slice, "lon": lon_slice}
def _shape_bounds_from_folder(shape_folder, gauge_ids, latlon):
"""Compute a buffered bounding box from available gauge shapefiles."""
if not shape_folder or gauge_ids is None:
logger.warning(
f"No shape_folder or gauge_ids provided; cannot compute shape bounds. {shape_folder}, {gauge_ids}"
)
return None
try:
import geopandas as gpd
except Exception as exc:
logger.warning(f"geopandas missing; cannot use shapefile bounds: {exc}")
return None
gauge_id_list = [gauge_ids] if not isinstance(gauge_ids, list) else gauge_ids
bounds_list = []
logger.debug(
f"Looking for shapefiles in {shape_folder} for gauge_ids: {gauge_id_list}"
)
for gauge_id in gauge_id_list:
logger.info(f"Looking for shapefile for gauge_id {gauge_id} in {shape_folder}")
shape_path = _find_shape_file(shape_folder, gauge_id)
logger.debug(f"Found shapefile {shape_path} for gauge_id {gauge_id}")
if shape_path is None:
continue
try:
gdf = gpd.read_file(shape_path)
except Exception as exc:
logger.warning(f"Failed to read shapefile {shape_path}: {exc}")
continue
bounds_list.append(tuple(gdf.total_bounds))
combined_bounds = _combine_shape_bounds(bounds_list)
if combined_bounds is None:
return None
buffered_bounds = _buffer_bounds(combined_bounds, latlon)
logger.info(f"Using shapefile bounds with buffer for slicing: {buffered_bounds}")
return buffered_bounds
[docs]
def create_cell_area(ds, lat_name="lat", lon_name="lon"):
"""Create a cell area data array in km2."""
logger.info("Create cell area data array.")
lat = ds[lat_name].data
lon = ds[lon_name].data
# calculate cellsize in kilometers
R = 6371 # radius of the earth in kilometers
lat_rad = np.deg2rad(lat)
lon_rad = np.deg2rad(lon)
dlat = np.abs(np.gradient(lat_rad))
dlon = np.abs(np.gradient(lon_rad))
# create 2D arrays for lat and lon
dlat_2d, dlon_2d = np.meshgrid(dlat, dlon, indexing="ij")
lat_2d = np.tile(lat_rad[:, np.newaxis], (1, len(lon)))
# calculate area
cell_areas = R**2 * dlat_2d * dlon_2d * np.cos(lat_2d)
return xr.DataArray(
cell_areas,
coords={lat_name: lat, lon_name: lon},
dims=[lat_name, lon_name],
name="cell_area",
attrs={
"title": "cell area",
"units": "km2",
"creator": "Department of Computational Hydrosystems",
"institution": "Helmholtz Centre for Environmental Research - UFZ",
},
)
def _normalize_output_vars(output_vars):
if output_vars is None:
return set(OUTPUT_VARIABLES)
if isinstance(output_vars, str):
output_vars = [val.strip() for val in output_vars.split(",") if val.strip()]
selected = {val.strip() for val in output_vars if str(val).strip()}
if not selected:
msg = "output_vars must contain at least one variable name."
with ErrorLogger(logger):
raise ValueError(msg)
unknown = selected - set(OUTPUT_VARIABLES)
if unknown:
msg = (
"Unknown output vars: "
f"{sorted(unknown)}. Valid options: {', '.join(OUTPUT_VARIABLES)}."
)
with ErrorLogger(logger):
raise ValueError(msg)
return selected
def _to_numeric_gauge_ids(raw_ids, context="output"):
"""Convert gauge IDs to integers, generating surrogate IDs when needed."""
if isinstance(raw_ids, (str, bytes)):
raw_ids = [raw_ids]
elif not isinstance(raw_ids, list):
raw_ids = list(raw_ids)
key_to_numeric = {}
used_numeric = {}
surrogate_candidates = []
for raw in raw_ids:
key = "" if raw is None else str(raw).strip()
if key in key_to_numeric:
continue
try:
numeric_id = int(key)
except (TypeError, ValueError):
surrogate_candidates.append(key)
continue
if numeric_id in used_numeric and used_numeric[numeric_id] != key:
logger.warning(
f"Gauge id '{key}' collides with '{used_numeric[numeric_id]}' "
f"after int conversion for {context}; using surrogate id."
)
surrogate_candidates.append(key)
continue
key_to_numeric[key] = numeric_id
used_numeric[numeric_id] = key
next_surrogate = max(used_numeric) + 1 if used_numeric else 1
for key in surrogate_candidates:
if key in key_to_numeric:
continue
while next_surrogate in used_numeric:
next_surrogate += 1
key_to_numeric[key] = next_surrogate
used_numeric[next_surrogate] = key
logger.warning(
f"Gauge id '{key if key else '<empty>'}' is non-numeric for {context}; "
f"using surrogate id {next_surrogate}."
)
next_surrogate += 1
return [key_to_numeric["" if raw is None else str(raw).strip()] for raw in raw_ids]
[docs]
def write_gauges_out(gauges, output_target):
"""Write gauge information to CSV and NetCDF files."""
write_gauges_to_csv(gauges, output_target.with_suffix(".csv"))
write_gauges_to_nc(gauges, output_target.with_suffix(".nc"))
[docs]
def write_gauges_to_csv(gauges, output_target, filename="gauges_info.csv"):
"""Write gauge information to a CSV file.
Parameters
----------
gauges : Gauge | list[Gauge] | dict | list[dict]
Gauge objects or dict rows.
output_target : str | Path
Either a target file path or a target directory.
filename : str
Filename used when output_target is a directory.
"""
if not gauges:
logger.warning("No gauges to write to CSV.")
return
if isinstance(gauges, (Gauge, dict)):
gauges = [gauges]
output_target = Path(output_target)
if output_target.is_dir() or not output_target.suffix:
output_path = output_target / str(filename)
else:
output_path = output_target
output_path.parent.mkdir(parents=True, exist_ok=True)
rows = []
for gauge in gauges:
if isinstance(gauge, dict):
row = {
"id": gauge.get("id", gauge.get("gauge_id")),
"lon": gauge.get("lon"),
"lat": gauge.get("lat"),
"lon_old": gauge.get("lon_old"),
"lat_old": gauge.get("lat_old"),
"distance": gauge.get("distance", gauge.get("distance_error", np.nan)),
"area": gauge.get("area"),
"old_area": gauge.get("old_area", gauge.get("area_old")),
"area_error": gauge.get("area_error", gauge.get("error")),
"score": gauge.get("score"),
"shape_error": gauge.get("shape_error"),
"method": gauge.get("method", gauge.get("used_method")),
}
else:
row = {
"id": gauge.gauge_id,
"lon": gauge.lon,
"lat": gauge.lat,
"lon_old": gauge.lon_old,
"lat_old": gauge.lat_old,
"distance": gauge.distance_error,
"area": gauge.area,
"old_area": gauge.area_old,
"area_error": gauge.area_error,
"score": gauge.score,
"shape_error": gauge.shape_error,
"method": gauge.method,
}
rows.append(row)
with output_path.open("w", newline="", encoding="utf-8") as fp:
writer = csv.DictWriter(fp, fieldnames=GAUGE_INFO_COLUMNS)
writer.writeheader()
writer.writerows(rows)
logger.info(f"Wrote gauge information for {len(rows)} gauges to {output_path}")
[docs]
def write_gauges_to_nc(gauges, output_target, filename="gauges_info.nc"):
"""Write gauge information to a NetCDF file with CF-style station metadata.
Parameters
----------
gauges : Gauge | list[Gauge] | dict | list[dict]
Gauge objects or dict rows.
output_target : str | Path
Either a target file path or a target directory.
filename : str
Filename used when output_target is a directory.
"""
if not gauges:
logger.warning("No gauges to write to NetCDF.")
return
if isinstance(gauges, (Gauge, dict)):
gauges = [gauges]
output_target = Path(output_target)
if output_target.is_dir() or not output_target.suffix:
output_path = output_target / str(filename)
else:
output_path = output_target
output_path.parent.mkdir(parents=True, exist_ok=True)
station_ids_raw = []
lons = []
lats = []
areas = []
for gauge in gauges:
if isinstance(gauge, dict):
station_id = gauge.get("gauge_id", gauge.get("id"))
lon = gauge.get("lon")
lat = gauge.get("lat")
area = gauge.get("area")
else:
station_id = gauge.gauge_id
lon = gauge.lon
lat = gauge.lat
area = gauge.area
station_ids_raw.append(station_id)
lons.append(np.nan if lon is None else float(lon))
lats.append(np.nan if lat is None else float(lat))
areas.append(np.nan if area is None else float(area))
if not station_ids_raw:
logger.warning("No gauge ids available for NetCDF output.")
return
station_ids = _to_numeric_gauge_ids(station_ids_raw, context="gauges_info.nc")
station = np.asarray(station_ids, dtype=np.int64)
ds = xr.Dataset(
data_vars={
"lon": ("station", np.asarray(lons, dtype=np.float64)),
"lat": ("station", np.asarray(lats, dtype=np.float64)),
"area": ("station", np.asarray(areas, dtype=np.float64)),
},
coords={"station": ("station", station)},
attrs={
"title": "SCC gauges specification",
"Conventions": "CF-1.12",
"institution": "UFZ",
"source": "Catchment delineation with mhm-tools",
},
)
ds["station"].attrs.update(
{
"long_name": "Station ID",
"coordinates": "lon lat",
}
)
ds["lon"].attrs.update(
{
"long_name": "longitude",
"standard_name": "longitude",
"units": "degrees_east",
}
)
ds["lat"].attrs.update(
{
"long_name": "latitude",
"standard_name": "latitude",
"units": "degrees_north",
}
)
ds["area"].attrs.update(
{
"standard_name": "catchment_area",
"long_name": (
"catchment area based on correction to Merit Hydro by Peter Burek et al. (2023)"
),
"units": "km2",
}
)
encoding = {
"station": {"dtype": "int64"},
"lon": {"dtype": "float64", "_FillValue": np.nan},
"lat": {"dtype": "float64", "_FillValue": np.nan},
"area": {"dtype": "float64", "_FillValue": np.nan},
}
ds = apply_output_provenance(ds)
ds.to_netcdf(output_path, engine="netcdf4", format="NETCDF4", encoding=encoding)
logger.info(
f"Wrote gauge information for {len(station_ids)} gauges to {output_path}"
)
# CLASSES
[docs]
class Gauge:
"""Class to hold gauge information."""
def __init__(self, gauge_id=None, lat=None, lon=None, area=None, id=None):
if gauge_id is None:
gauge_id = id
self.gauge_id = gauge_id
# Backward-compatible alias for older call sites using `gauge.id`.
self.id = gauge_id
self.lat = lat
self.lon = lon
self.area = area
self.area_old = None
self.lat_old = None
self.lon_old = None
self.distance_error = None
self.area_error = None
self.score = None
self.shape_error = None
self.method = None
[docs]
def update(
self,
area=None,
lat=None,
lon=None,
distance_error=None,
area_error=None,
score=None,
shape_error=None,
method=None,
):
"""Update gauge information, preserving old values."""
if area is not None:
self.area_old = self.area
self.area = area
if lat is not None:
self.lat_old = self.lat
self.lat = lat
if lon is not None:
self.lon_old = self.lon
self.lon = lon
if distance_error is not None:
self.distance_error = distance_error
if area_error is not None:
self.area_error = area_error
if score is not None:
self.score = score
if shape_error is not None:
self.shape_error = shape_error
if method is not None:
self.method = method
[docs]
class Catchment:
"""Catchment class deliniating catchmetns with pyflowdir."""
def __init__(
self,
ds,
var_name,
var="data",
ftype=None,
transform=None,
out_var_name=None,
do_shift=False,
resolutions: Resolution = None,
upscale=False,
latlon=True,
l0_precision: int = 9,
):
self.flwdir = None
self.basin = None
self.upgrid = None
self.uparea_grid = None
self.grdare = None
self.elevtn = None
self.cell_area = None
self._fdir = None
self.gauge_ids = []
self.gauge_lats = []
self.gauge_lons = []
self.ftype = ftype
self.catchment_mask = None
self.resolutions = resolutions if resolutions is not None else Resolution()
if self.resolutions.l0 is None:
self.resolutions.l0 = round(
abs(ds.lon.data[1] - ds.lon.data[0]), l0_precision
)
self.upscaled_resolution = self.resolutions.l0
self.do_upscale = upscale
self.is_upscaled = False
self.out_var_name = (
out_var_name if out_var_name is not None else f"{var_name}.nc"
)
self.VARIABLES = {
"flwdir": {
"title": f"flow direction ({self.ftype.upper()})",
"_FillValue": FDIR_FILLVALUE[self.ftype],
"units": "-",
},
"basin": {
"title": "basin Id",
"_FillValue": 0,
"units": "-",
},
"uparea_grid": {
"title": "accumulated data values along the flow directions",
"_FillValue": FACC_FILLVALUE,
"units": "-",
},
"upgrid": {
"title": "upstream area",
"_FillValue": FACC_FILLVALUE,
"units": "km2",
},
"grdare": {
"title": "rectangular grid area",
"_FillValue": FILLVALUE,
"units": "m2",
},
"elevtn": {
"title": "outlet pixel elevation",
"_FillValue": float(FILLVALUE),
"units": "m",
},
}
if not isinstance(self.out_var_name, str):
self.out_var_name = f"{var_name}.nc"
self.do_shift = do_shift
self.latlon = latlon
self.latlon = latlon
self.ds = ds
logger.debug(f"self.ds: {self.ds}")
self.transform = transform
data = self._modify_data(self.ds[var_name])
if self.do_shift:
transform = list(self.transform)
transform[2] = 0
self.transform = tuple(transform)
self.input_da = data
self.input_da = data
if var == "fdir":
if "nodata_value" in self.input_da.attrs:
old_no_data_val = self.input_da.attrs["nodata_value"]
elif "_FillValue" in self.input_da.attrs:
old_no_data_val = self.input_da.attrs["_FillValue"]
elif "missing_value" in self.input_da.attrs:
old_no_data_val = self.input_da.attrs["missing_value"]
else:
old_no_data_val = np.nan
self.input_da.attrs["_FillValue"] = FDIR_FILLVALUE[ftype]
self.input_da.attrs["nodata_value"] = FDIR_FILLVALUE[ftype]
self.input_da = self.input_da.where(
(data != old_no_data_val) & ~np.isnan(data),
FDIR_FILLVALUE[ftype],
)
logger.debug(self.input_da)
self.add_fdir(latlon=latlon)
elif var == "dem":
self.add_dem(latlon=latlon)
else:
with ErrorLogger(logger):
raise NotImplementedError
def _modify_data(self, data):
# correct circumspanning data
if self.do_shift:
shift = int(data.shape[1] / 2)
if isinstance(data, xr.DataArray):
return data.roll({data.dims[1]: shift}, roll_coords=False)
return np.roll(data, shift, axis=1)
return data
def _revert_data(self, data):
# correct circumspanning data
if self.do_shift:
return np.roll(data, int(data.shape[1] / 2), axis=1)
return data
[docs]
def add_dem(self, latlon):
"""Init the FlwdirRaster class from dem."""
# perform checks
# self.input_ds = fill_nan_with_neighbors(self.input_ds)
self.elevtn = self.input_da.data
if self._fdir is None:
# Create a flow direction object
logger.info("add_dem")
self._fdir = pyflwdir.from_dem(
data=self.elevtn,
nodata=np.nan,
transform=self.transform,
latlon=latlon,
)
self.get_fdir()
[docs]
def add_fdir(self, latlon):
"""Init the FlwdirRaster class from fdir."""
# perform check
data = self.input_da.data
if self._fdir is None:
data = data.astype(np.uint8)
self._fdir = pyflwdir.from_array(
data=data, ftype=self.ftype, transform=self.transform, latlon=latlon
)
self.get_fdir()
[docs]
def get_current_coordinates(self):
"""Build L1 coordinate arrays based on the dataset extent and L1 resolution."""
lon_coords = self.ds.lon.data
lat_coords = self.ds.lat.data
input_resolution = self.resolutions.l0
if (
self.resolutions.l1 is not None
and input_resolution != self.resolutions.l1
and self.do_upscale
and self.is_upscaled
):
# Rebuild coordinates for the coarser L1 grid
lon_coords = np.arange(
lon_coords.min() - input_resolution / 2 + self.resolutions.l1 / 2,
lon_coords.max() + self.resolutions.l1 / 2,
self.resolutions.l1,
)
lat_coords = np.arange(
lat_coords.max() + input_resolution / 2 - self.resolutions.l1 / 2,
lat_coords.min() - self.resolutions.l1 / 2,
-self.resolutions.l1,
)
logger.debug(
f"Computed L1 coords: lon={len(lon_coords)}, lat={len(lat_coords)}"
)
return lon_coords, lat_coords
[docs]
def compute_cell_area(self):
"""Create a cell area data array in km2."""
logger.info("Create cell area data array.")
lon, lat = self.get_current_coordinates()
# calculate cellsize in kilometers
R = 6371 # radius of the earth in kilometers
lat_rad = np.deg2rad(lat)
lon_rad = np.deg2rad(lon)
dlat = np.abs(np.gradient(lat_rad))
dlon = np.abs(np.gradient(lon_rad))
# create 2D arrays for lat and lon
dlat_2d, dlon_2d = np.meshgrid(dlat, dlon, indexing="ij")
lat_2d = np.tile(lat_rad[:, np.newaxis], (1, len(lon)))
# calculate area
cell_areas = R**2 * dlat_2d * dlon_2d * np.cos(lat_2d)
self.cell_area = cell_areas
[docs]
def calc_upstream_area(self):
"""Use pyflwdir to calculate the upstream area from flow direction by providing cell areas."""
if self._fdir is None:
logger.error("Flow direction is not initialized.")
return None
if self.cell_area is None:
self.compute_cell_area()
if self.cell_area.shape != self._fdir.shape:
msg = (
f"cell_area shape {self.cell_area.shape} does not match "
f"flow-direction shape {self._fdir.shape}."
)
with ErrorLogger(logger):
raise ValueError(msg)
return self._fdir.accuflux(self.cell_area, nodata=-9999)
def _coord_to_index(self, lat, lon, lat_vals=None, lon_vals=None):
"""Map latitude/longitude or indices to integer grid indices."""
if "lat" not in self.ds.coords or "lon" not in self.ds.coords:
msg = "Dataset is missing latitude/longitude coordinates."
with ErrorLogger(logger):
raise ValueError(msg)
lat_vals = self.ds.lat.data if lat_vals is None else lat_vals
lon_vals = self.ds.lon.data if lon_vals is None else lon_vals
if isinstance(lat, (int, np.integer)):
i = int(lat)
logger.debug(
f"Was given latitude index {i} directly. Corresponding lat_value {lat_vals[i]}"
)
elif lat < min(lat_vals) or lat > max(lat_vals):
logger.error(
f"Given latitude {lat} is outside dataset bounds ({min(lat_vals)}, {max(lat_vals)}). Clipping to bounds."
)
i = None
else:
i = int(np.abs(lat_vals - float(lat)).argmin())
logger.debug(
f"Mapped latitude {float(lat)} to index {i} with lat_value {lat_vals[i]}"
)
if isinstance(lon, (int, np.integer)):
j = int(lon)
logger.debug(
f"Was given longitude index {j} directly. Corresponding lon_value {lon_vals[j]}"
)
elif lon < min(lon_vals) or lon > max(lon_vals):
logger.error(
f"Given longitude {lon} is outside dataset bounds ({min(lon_vals)}, {max(lon_vals)}). Clipping to bounds."
)
j = None
else:
j = int(np.abs(lon_vals - float(lon)).argmin())
logger.debug(
f"Mapped longitude {float(lon)} to index {j} with lon_value {lon_vals[j]}"
)
if i is None or j is None:
msg = (
"Could not map given coordinates to valid indices within "
"dataset bounds."
)
with ErrorLogger(logger):
raise ValueError(msg)
i = int(np.clip(i, 0, len(lat_vals) - 1))
j = int(np.clip(j, 0, len(lon_vals) - 1))
return i, j
def _coords_in_domain(self, gauge_coords, lat_vals=None, lon_vals=None):
"""Check whether gauge coordinates can be mapped to the active domain."""
if gauge_coords is None:
return False
lat_vals = self.ds.lat.data if lat_vals is None else lat_vals
lon_vals = self.ds.lon.data if lon_vals is None else lon_vals
lat, lon = gauge_coords
lat_valid = (
0 <= lat < len(lat_vals)
if isinstance(lat, (int, np.integer))
else min(lat_vals) <= lat <= max(lat_vals)
)
lon_valid = (
0 <= lon < len(lon_vals)
if isinstance(lon, (int, np.integer))
else min(lon_vals) <= lon <= max(lon_vals)
)
return lat_valid and lon_valid
def _coords_l1(self):
"""Build L1 coordinate arrays based on the dataset extent and L1 resolution."""
lon_coords = self.ds.lon.data
lat_coords = self.ds.lat.data
input_resolution = self.resolutions.l0
if (
self.resolutions.l1 is not None
and input_resolution != self.resolutions.l1
and self.do_upscale
):
# Rebuild coordinates for the coarser L1 grid
lon_coords = np.arange(
lon_coords.min() - input_resolution / 2 + self.resolutions.l1 / 2,
lon_coords.max() + self.resolutions.l1 / 2,
self.resolutions.l1,
)
lat_coords = np.arange(
lat_coords.max() + input_resolution / 2 - self.resolutions.l1 / 2,
lat_coords.min() - self.resolutions.l1 / 2,
-self.resolutions.l1,
)
logger.debug(
f"Computed L1 coords: lon={len(lon_coords)}, lat={len(lat_coords)}"
)
return lon_coords, lat_coords
[docs]
def update_gauge_coords(self, gauge_id, gauge_lat, gauge_lon):
"""Update stored gauge coordinates for a gauge_id or add it if missing."""
if gauge_id is None:
return
if gauge_id in self.gauge_ids:
gauge_index = self.gauge_ids.index(gauge_id)
self.gauge_lats[gauge_index] = gauge_lat
self.gauge_lons[gauge_index] = gauge_lon
logger.debug(f"Updated gauge {gauge_id} to {gauge_lat}/{gauge_lon}")
else:
self.gauge_ids.append(gauge_id)
self.gauge_lats.append(gauge_lat)
self.gauge_lons.append(gauge_lon)
logger.debug(f"Added gauge {gauge_id} at {gauge_lat}/{gauge_lon}")
[docs]
def correct_gauge_location_l1_from_shape(
self,
l0_shape_gdf,
gauge_coords,
max_distance_cells=5,
max_error=0.5,
ref_catchment_area=None,
reference_upstream_area=None,
):
"""Correct gauge coordinates at L1 using L0 shape similarity and upstream area."""
if ref_catchment_area is None and reference_upstream_area is not None:
ref_catchment_area = reference_upstream_area
if l0_shape_gdf is None or l0_shape_gdf.empty:
return None
if self._fdir is None:
return None
lon_coords, lat_coords = self._coords_l1()
if (
len(lat_coords) != self._fdir.shape[0]
or len(lon_coords) != self._fdir.shape[1]
):
lon_coords, lat_coords = _coords_from_transform(
getattr(self._fdir, "transform", self.transform), self._fdir.shape
)
l1_upstream_area = self.uparea_grid
if l1_upstream_area is None:
logger.warning(
"L1 upstream area grid missing; cannot apply shape correction."
)
return None
l1_result = self.find_best_gauge_location_shape(
upstream_area=l1_upstream_area,
gauge_coords=gauge_coords,
ref_catchment_area=ref_catchment_area,
shape_folder=None,
gauge_id=None,
max_distance_cells=max_distance_cells,
max_error=max_error,
reference_shape_gdf=l0_shape_gdf,
lat_values=lat_coords,
lon_values=lon_coords,
)
if l1_result is None:
logger.warning("No L1 candidate basins matched the L0 shape.")
return None
def _iou_for_index(row_idx, col_idx):
try:
linear_idx = np.ravel_multi_index((row_idx, col_idx), self._fdir.shape)
basin = self._fdir.basins(idxs=np.array([linear_idx], dtype=np.int64))
candidate_mask = basin > 0
gdf = _vectorize_mask_to_gdf(
candidate_mask,
getattr(self._fdir, "transform", self.transform),
_shape_crs(self.latlon),
)
return _shape_iou(l0_shape_gdf, gdf)
except Exception as exc:
logger.debug(
f"Could not compute L1 shape IoU for outlet candidate: {exc}"
)
return np.nan
best_candidate_index = l1_result[0]
candidate_iou = _iou_for_index(
int(best_candidate_index[0]), int(best_candidate_index[1])
)
gauge_row = int(np.abs(lat_coords - float(gauge_coords[0])).argmin())
gauge_col = int(np.abs(lon_coords - float(gauge_coords[1])).argmin())
naive_iou = _iou_for_index(gauge_row, gauge_col)
if (
np.isfinite(candidate_iou)
and np.isfinite(naive_iou)
and candidate_iou < naive_iou
):
logger.info(
f"L1 shape candidate IoU {candidate_iou:.4f} is worse than naive "
f"IoU {naive_iou:.4f}; keeping naive L1 gauge location."
)
return float(lat_coords[gauge_row]), float(lon_coords[gauge_col])
new_lat = float(lat_coords[int(best_candidate_index[0])])
new_lon = float(lon_coords[int(best_candidate_index[1])])
logger.info(
f"L1 shape-based correction selected {new_lat:.4f}/{new_lon:.4f} "
"using L0 shape reference."
)
return new_lat, new_lon
[docs]
def derive_gauge_coords_from_shape(
self,
upstream_area,
reference_shape,
lat_values=None,
lon_values=None,
shape_label="reference shape",
):
"""Use the highest upstream area inside a shape as initial gauge location."""
if upstream_area is None:
logger.warning(
"Upstream area grid missing; cannot derive gauge from shape."
)
return None
if lat_values is None or lon_values is None:
lat_values = self.ds.lat.data
lon_values = self.ds.lon.data
shape_mask = _rasterize_shape_to_mask(
reference_shape,
upstream_area.shape,
getattr(self._fdir, "transform", self.transform),
)
if shape_mask is None or not np.any(shape_mask):
logger.warning(f"Could not rasterize {shape_label} onto the active grid.")
return None
candidates = shape_mask & np.isfinite(upstream_area)
if not np.any(candidates):
logger.warning(f"No finite upstream-area cells inside {shape_label}.")
return None
values = np.where(candidates, upstream_area, -np.inf)
row_idx, col_idx = np.unravel_index(np.nanargmax(values), values.shape)
gauge_lat = float(lat_values[int(row_idx)])
gauge_lon = float(lon_values[int(col_idx)])
logger.info(
f"Derived gauge location {gauge_lat:.6f}/{gauge_lon:.6f} from highest "
f"upstream area inside {shape_label}."
)
return gauge_lat, gauge_lon
[docs]
def calculate_shape_area_on_current_grid(
self,
reference_shape,
shape_label="reference shape",
):
"""Calculate the area covered by a shape on the current grid in km2."""
if reference_shape is None or getattr(reference_shape, "empty", False):
return None
if self._fdir is None:
logger.warning(
"Flow direction is not initialized; cannot calculate shape area."
)
return None
if self.cell_area is None:
self.compute_cell_area()
shape_area = _calculate_shape_area_on_grid(
reference_shape,
self.cell_area,
self._fdir.shape,
getattr(self._fdir, "transform", self.transform),
)
if shape_area is None:
logger.warning(
f"Could not calculate area covered by {shape_label} on the current "
"grid."
)
return None
logger.info(
f"Calculated {shape_area:.2f} km2 covered by {shape_label} on the "
"current grid."
)
return shape_area
[docs]
def find_best_gauge_location_shape( # noqa: PLR0911, PLR0912, PLR0915
self,
upstream_area,
gauge_coords,
ref_catchment_area,
shape_folder,
gauge_id,
max_distance_cells=2,
max_error=0.25,
reference_shape_gdf=None,
lat_values=None,
lon_values=None,
limit_by_error=True,
started_from_shape=False,
):
"""Find best gauge location using shape similarity."""
if upstream_area is None:
logger.warning("Upstream area grid missing for shape-based correction.")
return None
if reference_shape_gdf is None:
reference_shape, shape_path = _read_reference_shape(
shape_folder, gauge_id, self.latlon
)
if reference_shape is None:
logger.debug(f"No reference shapefile found for gauge_id {gauge_id}")
return None
shape_label = shape_path.name
crs = _shape_crs(self.latlon)
else:
reference_shape = reference_shape_gdf
shape_label = "reference shape"
crs = reference_shape.crs if hasattr(reference_shape, "crs") else None
if lat_values is None or lon_values is None:
lat_values = self.ds.lat.data
lon_values = self.ds.lon.data
has_reference_area = ref_catchment_area is not None
if ref_catchment_area is None:
ref_catchment_area = self.calculate_shape_area_on_current_grid(
reference_shape,
shape_label=shape_label,
)
if gauge_coords is None:
started_from_shape = True
gauge_coords = self.derive_gauge_coords_from_shape(
upstream_area,
reference_shape,
lat_values=lat_values,
lon_values=lon_values,
shape_label=shape_label,
)
if gauge_coords is None:
return None
try:
gauge_row, gauge_col = self._coord_to_index(
gauge_coords[0],
gauge_coords[1],
lat_vals=lat_values,
lon_vals=lon_values,
)
except ValueError:
logger.warning(
f"Gauge coordinates {gauge_coords} are outside the active domain; "
f"deriving gauge from {shape_label}."
)
started_from_shape = True
gauge_coords = self.derive_gauge_coords_from_shape(
upstream_area,
reference_shape,
lat_values=lat_values,
lon_values=lon_values,
shape_label=shape_label,
)
if gauge_coords is None:
return None
gauge_row, gauge_col = self._coord_to_index(
gauge_coords[0],
gauge_coords[1],
lat_vals=lat_values,
lon_vals=lon_values,
)
max_cells = (
int(max(0, round(max_distance_cells)))
if max_distance_cells is not None
else 0
)
row_min = max(0, gauge_row - max_cells)
row_max = min(len(lat_values) - 1, gauge_row + max_cells)
col_min = max(0, gauge_col - max_cells)
col_max = min(len(lon_values) - 1, gauge_col + max_cells)
if row_min > row_max:
row_min, row_max = row_max, row_min
if col_min > col_max:
col_min, col_max = col_max, col_min
# Limit candidate search to a local neighborhood around the gauge
sub = upstream_area[row_min : row_max + 1, col_min : col_max + 1]
if ref_catchment_area is not None and limit_by_error:
candidate_indices = np.where(
(sub >= ref_catchment_area * (1 - max_error))
& (sub <= ref_catchment_area * (1 + max_error))
)
if candidate_indices[0].size == 0:
logger.warning(
"No candidates within area error bounds; expanding search to all finite upstream area values."
)
candidate_indices = np.where(np.isfinite(sub))
else:
candidate_indices = np.where(np.isfinite(sub))
logger.debug(
f"Shape-based candidate count: {len(candidate_indices[0])}",
)
best_candidate_index = None
best_candidate_score = np.inf
best_candidate_shape_iou = 0.0
best_candidate_upstream_area = None
best_candidate_distance_100m = None
logger.debug(
f"Using resolution {self.upscaled_resolution} for shape-based distance scoring."
)
lat_deg = float(lat_values[gauge_row]) if self.latlon else None
for cand_row, cand_col in zip(candidate_indices[0], candidate_indices[1]):
row_idx = int(cand_row + row_min)
col_idx = int(cand_col + col_min)
linear = np.ravel_multi_index((row_idx, col_idx), self._fdir.shape)
basin = self._fdir.basins(
idxs=np.array([linear], dtype=np.int64),
# streams=streams_mask,
)
basin_mask = basin > 0
if not np.any(basin_mask):
continue
basin_transform = getattr(self._fdir, "transform", self.transform)
gdf = _vectorize_mask_to_gdf(
basin_mask, basin_transform, crs, value_name="basin"
)
shape_overlap_ratio = _shape_iou(reference_shape, gdf)
upstream_value = upstream_area[row_idx, col_idx]
distance_100m = distance_100m_units(
row_idx - gauge_row,
col_idx - gauge_col,
l0_resolution=self.upscaled_resolution,
lat_deg=lat_deg,
)
if ref_catchment_area:
upstream_area_ratio = (
min(upstream_value, ref_catchment_area)
/ max(upstream_value, ref_catchment_area)
if upstream_value and ref_catchment_area
else 0.0
)
score = np.hypot(
1 - upstream_area_ratio, 1 - shape_overlap_ratio
) / np.sqrt(
2
) # sqrt(x1**2 + x2**2)
else:
score = 1 - shape_overlap_ratio
if (
best_candidate_index is None
or score < best_candidate_score
or (
np.isclose(score, best_candidate_score)
and best_candidate_distance_100m is not None
and distance_100m < best_candidate_distance_100m
)
):
best_candidate_index = (row_idx, col_idx)
best_candidate_score = score
best_candidate_shape_iou = shape_overlap_ratio
best_candidate_upstream_area = upstream_value
best_candidate_distance_100m = distance_100m
if (
ref_catchment_area is None
and best_candidate_shape_iou <= 0
and not started_from_shape
):
derived_coords = self.derive_gauge_coords_from_shape(
upstream_area,
reference_shape,
lat_values=lat_values,
lon_values=lon_values,
shape_label=shape_label,
)
if derived_coords is None:
logger.warning(
f"No shape-overlap candidate found around {gauge_coords} and no "
"shape-derived start could be computed."
)
return None
logger.warning(
f"No shape-overlap candidate found around gauge coordinates {gauge_coords}; "
f"retrying from highest upstream-area cell inside {shape_label}.",
)
return self.find_best_gauge_location_shape(
upstream_area,
derived_coords,
ref_catchment_area,
shape_folder,
gauge_id,
max_distance_cells=max_distance_cells,
max_error=max_error,
reference_shape_gdf=reference_shape,
lat_values=lat_values,
lon_values=lon_values,
limit_by_error=limit_by_error,
started_from_shape=True,
)
if best_candidate_index is None or best_candidate_score > max_error:
logger.warning("No suitable candidate found for shape-based correction.")
logger.warning(
f"Score {best_candidate_score:.3f} > {max_error} from IoU {best_candidate_shape_iou:.3f}; and area {best_candidate_upstream_area:.0f}km^2 / {ref_catchment_area:.0f}km^2."
)
return None
area_error = (
abs(1 - best_candidate_upstream_area / ref_catchment_area)
if ref_catchment_area and best_candidate_upstream_area
else -9999
)
logger.info(
f"Shape-based gauge correction used {shape_label} with IoU {best_candidate_shape_iou:.3f}; area error {area_error:.3f} and distance {best_candidate_distance_100m/10:.2f}km.",
)
used_method = "shape-area" if has_reference_area else "shape"
return (
best_candidate_index,
area_error,
best_candidate_distance_100m,
best_candidate_score,
1 - best_candidate_shape_iou,
used_method,
)
[docs]
def get_best_gauge_coordinate(
self,
upstream_area,
gauge_coords,
ref_catchment_area,
max_distance_cells,
max_error,
method,
shape_folder=None,
gauge_id=None,
raise_on_fallback=True,
):
"""Get best gauge coordinates given target catchment area."""
shape_result = None
score = np.nan
shape_error = np.nan
used_method = None
if shape_folder and gauge_id is not None:
try:
shape_result = self.find_best_gauge_location_shape(
upstream_area,
gauge_coords,
ref_catchment_area,
shape_folder,
gauge_id,
max_distance_cells=max_distance_cells,
max_error=max_error,
)
except Exception as exc:
logger.warning(
f"Shape-based gauge correction failed for {gauge_id}: {exc}"
)
shape_result = None
if gauge_coords is not None and not self._coords_in_domain(gauge_coords):
logger.warning(
f"Gauge coordinates {gauge_coords} are outside the active domain; "
"skipping area-only correction."
)
gauge_coords = None
if shape_result is not None:
(
outlet_idx,
error,
distance_error,
score,
shape_error,
used_method,
) = shape_result
new_lat = float(self.ds.lat.data[outlet_idx[0]])
new_lon = float(self.ds.lon.data[outlet_idx[1]])
gauge_lat = new_lat
gauge_lon = new_lon
logger.info(
f"Moved outlet to {new_lat}/{new_lon} using shape similarity "
f"(distance {distance_error / 10:.2f}km)."
)
elif ref_catchment_area is not None and gauge_coords is not None:
if method != "all":
outlet_idx, error, distance_error = find_best_gauge_location_by_area(
ds=self.ds,
upstream_area=upstream_area,
gauge_coords=gauge_coords,
ref_catchment_area=ref_catchment_area,
resolutions=self.resolutions,
max_distance_cells=max_distance_cells,
max_error=max_error,
method=method,
raise_on_fallback=raise_on_fallback,
)
used_method = f"area-{method}"
score = error * 100 + 2 * distance_error if method == "burek" else error
else:
outlet_idx_bx, error_bx, distance_error_bx = (
find_best_gauge_location_by_area(
ds=self.ds,
upstream_area=upstream_area,
gauge_coords=gauge_coords,
ref_catchment_area=ref_catchment_area,
resolutions=self.resolutions,
max_distance_cells=max_distance_cells,
max_error=max_error,
method="basinex",
raise_on_fallback=True,
)
)
outlet_idx_bu, error_bu, distance_error_bu = (
find_best_gauge_location_by_area(
ds=self.ds,
upstream_area=upstream_area,
gauge_coords=gauge_coords,
ref_catchment_area=ref_catchment_area,
resolutions=self.resolutions,
max_distance_cells=max_distance_cells,
max_error=max_error,
method="burek",
raise_on_fallback=True,
)
)
logger.info("Results of basin correction:")
logger.info(f"Burek: lat: {float(self.ds.lat.data[outlet_idx_bu[0]])}")
logger.info(f"Burek: lon: {float(self.ds.lon.data[outlet_idx_bu[1]])}")
logger.info(f"Burek: error {error_bu}")
logger.info(f"Burek: distance change {distance_error_bu/10}km")
if method in ("all", "basinex"):
if method == "basinex":
outlet_idx_bx = outlet_idx
error_bx = error
distance_error_bx = distance_error
logger.info(
f"BasinEx: lat: {float(self.ds.lat.data[outlet_idx_bx[0]])}"
)
logger.info(
f"BasinEx: lon: {float(self.ds.lon.data[outlet_idx_bx[1]])}"
)
logger.info(f"BasinEx: error {error_bx}")
logger.info(f"BasinEx: distance change {distance_error_bx/10}km")
if method == "all":
score_bx = error_bx
score_bu = error_bu * 100 + 2 * distance_error_bu
if error_bx < error_bu:
logger.info("using BasinEx location")
outlet_idx = outlet_idx_bx
error = error_bx
distance_error = distance_error_bx
score = score_bx
used_method = "area-basinex"
else:
logger.info("Using Burek location")
outlet_idx = outlet_idx_bu
error = error_bu
distance_error = distance_error_bu
score = score_bu
used_method = "area-burek"
new_lat = float(self.ds.lat.data[outlet_idx[0]])
new_lon = float(self.ds.lon.data[outlet_idx[1]])
gauge_lat = new_lat
gauge_lon = new_lon
logger.info(
f"Moved outlet {distance_error/10:.2}km shifting latitude {float(gauge_coords[0])} to {new_lat} and longitude {float(gauge_coords[1])} to {new_lon}."
)
elif shape_result is None and shape_folder and gauge_id is not None:
error_msg = f"Shape-based correction failed for gauge_id {gauge_id}, and no area-based correction applied."
with ErrorLogger(logger):
raise ValueError(error_msg)
else:
logger.warning(
"No catchment area or shape provided; falling back to original gauge coordinates."
)
if gauge_coords is None:
msg = "Gauge coordinates are required when neither shape-based nor area-based correction succeeds."
with ErrorLogger(logger):
raise ValueError(msg)
outlet_idx = coord_to_index(self.ds, gauge_coords[0], gauge_coords[1])
gauge_lat = float(gauge_coords[0])
gauge_lon = float(gauge_coords[1])
error = None
distance_error = 0.0
return (
outlet_idx,
error,
gauge_lat,
gauge_lon,
distance_error / 10,
score,
shape_error,
used_method,
)
[docs]
def delineation_sanity_check(
self,
catchment_mask,
basin,
uparea_at_outlet,
ref_catchment_area,
max_error,
raise_on_sanity_check,
gauge_id,
):
"""Perform sanity checks on the delineated basin."""
try:
mean_cell_area = (
float(np.mean(self.cell_area[catchment_mask]))
if self.cell_area is not None
else np.nan
)
unique_vals = np.unique(basin[catchment_mask])
cell_count = int(np.sum(catchment_mask))
delineated_area = (
float(np.sum(self.cell_area[catchment_mask]))
if self.cell_area is not None
else np.nan
)
logger.info(
f"Basin unique values: {unique_vals} | cells in basin: {cell_count} "
f"| mean cell area: {mean_cell_area:.6f} km2"
)
area_difference = delineated_area - uparea_at_outlet
area_difference_percent = (
area_difference / uparea_at_outlet * 100.0
if uparea_at_outlet != 0
else np.nan
)
logger.info(
f"Upstream area reported at selected outlet cell = "
f"{uparea_at_outlet:.2f} km2. Difference (sum_cells - "
f"upstream_at_outlet) = {area_difference:.2f} km2 "
f"({area_difference_percent:.2f}%)."
)
if ref_catchment_area is not None:
area_error = (delineated_area - ref_catchment_area) / ref_catchment_area
logger.info(
f"Delineated basin area (sum of cell_area[basin>0]) = "
f"{delineated_area:.2f} km2; reference area = "
f"{ref_catchment_area:.2f} km2; error = {area_error * 100.0:.2f}%"
)
if abs(area_error) > max_error * 2:
with ErrorLogger(logger):
msg = f"Delineated basin area ({delineated_area:2f} km2) differs from reference area ({ref_catchment_area:2f} km2) by more than twice the max error {max_error*100:.2f}%. Adjust max_error or max_distance_cells."
if raise_on_sanity_check:
raise ValueError(msg)
logger.warning(msg)
return True
# warn if the two area measures disagree substantially
else:
logger.warning(
"No reference catchment area provided; skipping area consistency check."
)
if not np.isclose(delineated_area, uparea_at_outlet, rtol=0.02, atol=1e-6):
with ErrorLogger(logger):
msg = f"Gauge ID {gauge_id}: " if gauge_id is not None else ""
msg += f"Sum of cell areas inside the basin ({delineated_area:2f} km2) differs from "
msg += f"upstream area at outlet ({uparea_at_outlet:2f} km2). Investigate flow-direction "
msg += "masking, nodata handling or area units."
if raise_on_sanity_check:
raise ValueError(msg)
logger.warning(msg)
return True
return False
except Exception as e:
logger.exception(f"Sanity check failed: {e}")
if raise_on_sanity_check:
raise e
return True
[docs]
def delineate_basin(
self,
gauge,
max_distance_cells=5,
max_error=0.25,
raise_on_sanity_check=True,
upstream_area=None,
mask_catchment: Optional[bool] = True,
save_coords: Optional[bool] = True,
gauge_opti_method: Optional[str] = "basinex",
shape_folder: Optional[str] = None,
raise_on_fallback: Optional[bool] = True,
):
"""Delineate the basin for a given lat and lon."""
# Target area in km2 we want to match (can be adjusted/replaced by caller later)
ref_catchment_area = gauge.area
gauge_coords = (
(gauge.lat, gauge.lon)
if gauge.lat is not None and gauge.lon is not None
else None
)
gauge_id = getattr(gauge, "gauge_id", getattr(gauge, "id", None))
# Compute upstream area (in km2) using accuflux and cell areas
if self.cell_area is None:
self.compute_cell_area()
if upstream_area is None:
try:
upstream_area = self.calc_upstream_area()
except Exception:
logger.exception("Failed to compute upstream area (accuflux).")
if upstream_area is None:
msg = "Could not calculate upstream area. Flow direction may be uninitialized."
with ErrorLogger(logger):
raise ValueError(msg)
gauge_str = f" (ID: {gauge_id})" if gauge_id is not None else ""
logger.info(
f"Delineating basin {gauge_str} for gauge coordinates {gauge_coords} and reference catchment area {ref_catchment_area} km2."
)
(
outlet_idx,
error,
gauge_lat,
gauge_lon,
distance_error,
score,
shape_error,
used_method,
) = self.get_best_gauge_coordinate(
upstream_area=upstream_area,
gauge_coords=gauge_coords,
ref_catchment_area=ref_catchment_area,
max_distance_cells=max_distance_cells,
max_error=max_error,
method=gauge_opti_method,
shape_folder=shape_folder,
gauge_id=gauge_id,
raise_on_fallback=raise_on_fallback,
)
outlet_linear_idx = np.ravel_multi_index(outlet_idx, self._fdir.shape)
try:
basin = self._fdir.basins(
idxs=np.array([outlet_linear_idx], dtype=np.int64),
)
except Exception as e:
logger.exception(f"pyflwdir.basins(idxs=...) failed for {outlet_idx}: {e}")
try:
all_basins = self._fdir.basins()
basin_id = int(all_basins[outlet_idx])
basin = np.where(all_basins == basin_id, basin_id, 0)
except Exception as e2:
logger.exception(f"Fallback basins() also failed: {e2}")
return gauge
catchment_mask = basin > 0
logger.debug(
f"mask statistics: min={basin.min()}, max={basin.max()}, all true? {np.all(catchment_mask)}"
)
uparea_at_outlet = (
upstream_area[outlet_idx] if upstream_area is not None else np.nan
)
failed_sanity_check = self.delineation_sanity_check(
catchment_mask,
basin,
uparea_at_outlet,
ref_catchment_area,
max_error,
raise_on_sanity_check,
gauge_id,
)
if np.all(catchment_mask):
logger.error("No catchment found for the given coordinates")
return gauge
if mask_catchment:
self.catchment_mask = catchment_mask
self.basin = basin
try:
fillv = self.VARIABLES["basin"]["_FillValue"]
self.basin = np.where(self.catchment_mask, self.basin, fillv)
except Exception:
logger.debug("Could not set basin fill values")
if not failed_sanity_check:
if save_coords:
self.save_coords(gauge_id, gauge_lat, gauge_lon)
gauge.update(
lat=gauge_lat,
lon=gauge_lon,
area=uparea_at_outlet,
distance_error=distance_error,
area_error=error,
score=score,
shape_error=shape_error,
method=used_method,
)
return gauge
[docs]
def save_coords(self, gauge_id, gauge_lat, gauge_lon):
"""Save gauge coordinates."""
if gauge_id is not None:
self.gauge_ids.append(gauge_id)
self.gauge_lats.append(gauge_lat)
self.gauge_lons.append(gauge_lon)
[docs]
def write_basin_shape(self, out_dir, gauge_id=None, basin_mask=None):
"""Write a basin shapefile from a basin mask."""
if basin_mask is None:
basin_mask = self.catchment_mask
if basin_mask is None or not np.any(basin_mask):
return
try:
gdf = _vectorize_mask_to_gdf(
basin_mask, self.transform, _shape_crs(self.latlon), value_name="basin"
)
except Exception as exc:
logger.warning(f"Could not write basin shapefile: {exc}")
return
output_dir = Path(out_dir)
output_dir.mkdir(parents=True, exist_ok=True)
suffix = f"_{gauge_id}" if gauge_id is not None else ""
out_path = output_dir / f"basin{suffix}.shp"
try:
gdf.to_file(out_path)
logger.info(f"Wrote basin shapefile to {out_path}")
except Exception as exc:
logger.warning(f"Failed to write basin shapefile {out_path}: {exc}")
[docs]
def get_upscaling_factor(self, max_resolution=False, l1=False, l2=True):
"""Create upscaling factor."""
return get_upscaling_factor(
self.resolutions, max_resolution=max_resolution, l1=l1, l2=l2
)
[docs]
def upscale(self, var):
"""Upscale flow direction to l1_resolution if that is int multipe of data resolution."""
factor, upscaled_resolution = get_upscaling_factor(self.resolutions, l1=True)
if factor == 1:
self.get_facc()
return
# if we upscale the do_upscale flag should be true
self.do_upscale = True
logger.info(
f"Upscaling flow direction to {upscaled_resolution} with the fator {factor}."
)
fdir_upscaled, upscaling_indices = self._fdir.upscale(factor, method="ihu")
subareas = self._fdir.ucat_area(idxs_out=upscaling_indices, unit="km2")[1]
uparea1 = fdir_upscaled.accuflux(subareas)
flwerr = self._fdir.upscale_error(fdir_upscaled, upscaling_indices)
percentage_error = np.sum(flwerr == 0) / np.sum(flwerr != 255) * 100
logger.info(f"upscaling error in {percentage_error:.2f}% of cells")
logger.debug(f"Upscaled form {self._fdir.shape} to {fdir_upscaled.shape}")
self._fdir = fdir_upscaled
self.get_fdir()
self.uparea_grid = uparea1 # replaces self.get_facc
self.cell_area = (
None # reset cell area to be recalculated at new resolution when needed
)
self.is_upscaled = True
self.upscaled_resolution = upscaled_resolution
if var == "dem":
lat_size, lon_size = self.input_da.shape
# Ensure the dimensions are evenly divisible by the factor
if lat_size % factor != 0 or lon_size % factor != 0:
msg = f"Data dimensions must be divisible by the upscaling factor of {factor}. Lat ({lat_size}/{factor})={lat_size / factor:.2f}; Lon ({lon_size}/{factor})={lon_size / factor:.2f}"
with ErrorLogger(logger):
raise ValueError(msg)
# Reshape and aggregate data
reshaped = self.input_da.values.reshape(
lat_size // factor, factor, lon_size // factor, factor
)
aggregated = reshaped.mean(axis=(1, 3)) # Conservative mean over each block
# Create new DataArray
self.elevtn = aggregated
[docs]
def get_basins(self):
"""Perform the calculation of the catchment ids."""
self.basin = self._fdir.basins()
[docs]
def get_fdir(self):
"""Perform the calculation of the flow direction."""
logger.debug("Get flwdir as array.")
self.flwdir = self._fdir.to_array(ftype=self.ftype or OUTPUT_FTYPE)
[docs]
def get_upstream_area(self):
"""Perform the calculation of the upstream catchment area."""
# upgrid = self._fdir.upstream_area(unit="km2").astype(int)
self.upgrid = self.calc_upstream_area().astype(int)
[docs]
def get_grid_area(self):
"""Perform the calculation of the catchment area."""
self.get_upstream_area()
self.grdare = self._fdir.area.astype(int)
[docs]
def get_facc(self):
"""Get the flow accumulation area."""
logger.info("Calculate flow accumulation...")
data = np.ones_like(self.flwdir).astype(np.uint32)
data[~self._fdir.mask.reshape(data.shape)] = 0
self.uparea_grid = self._fdir.accuflux(data, nodata=0)
[docs]
@staticmethod
def create_frame(ds, frame=0, frame_value=0):
"""If a frame is used this frame is set to no data values as a frame."""
logger.info(f"Creating a frame of {frame} cells around the domain.")
if frame > 0:
for var in ds.data_vars:
data = ds.variables[var].data[:]
# set bounds to -9999.
data[:frame, :] = frame_value
data[-frame:, :] = frame_value
data[:, :frame] = frame_value
data[:, -frame:] = frame_value
ds.variables[var].data[:] = data
return ds
[docs]
def fill_adjacent_missing_with_sink(self, da, fill_value, sink_value):
"""Replace all missing values adjacent to non-missing values with 0 in an xarray Dataset.
Parameters
----------
da (xr.Dataset): Input dataset.
Returns
-------
xr.Dataset: Dataset with adjacent missing values replaced with 0.
"""
# Mask of missing values
missing_mask = da == fill_value
# Mask of non-missing values
non_missing_mask = ~missing_mask
# Dilate the non-missing mask to include adjacent cells
adjacent_mask = binary_dilation(
non_missing_mask, structure=np.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
)
# Identify adjacent missing values
adjacent_missing = adjacent_mask & missing_mask
# Replace adjacent missing values with 0
return xr.where(adjacent_missing, sink_value, da)
[docs]
@log_arguments()
def write(
self,
out_path,
single_file=True,
format="nc",
cellsize=None,
cut_by_basin=False,
mask_file=None,
frame=1,
buffer=0,
variables=None,
id_gauges_out_path=None,
):
"""Write the produced data to one or multiple files."""
data_vars = {}
out_path = Path(out_path)
if not out_path.is_dir():
out_path.mkdir(parents=True, exist_ok=True)
gauges_out_path = (
Path(id_gauges_out_path) if id_gauges_out_path is not None else out_path
)
gauges_out_path.mkdir(parents=True, exist_ok=True)
selected_vars = _normalize_output_vars(variables)
lat_slice_idx, lon_slice_idx = None, None
lat_slice, lon_slice = None, None
if cut_by_basin:
lat_slice_idx, lon_slice_idx = cut_to_filled_area(
ds=self.ds,
resolutions=self.resolutions,
catchment_mask=self.catchment_mask,
buffer=buffer,
)
else:
lat_slice, lon_slice = slice(84, -56), slice(None)
for var_name in (v for v in self.VARIABLES if v in selected_vars):
data_var = self.processing_data_variable(
var_name,
cut_by_basin,
lat_slice,
lon_slice,
lat_slice_idx,
lon_slice_idx,
)
if data_var is None:
continue
if single_file:
data_vars[var_name] = data_var
else:
self.write_single_variable_file(
data_var, var_name, out_path, cellsize, format
)
if single_file:
if not data_vars:
msg = "No data variables available to write."
with ErrorLogger(logger):
raise ValueError(msg)
ds = self.write_basin_id_file(data_vars, frame, out_path)
# use basin_id to create a mask file
if "basin" in ds.data_vars:
self.write_mask_file(ds, mask_file)
if self.gauge_ids:
logger.info("Writing gauges file.")
# create empty ds with mask l0 extend and fill the data_var called data with -9999 values
id_da = xr.DataArray(
np.full(ds.basin.shape, -9999, dtype=int),
coords={"lat": ds.lat, "lon": ds.lon},
dims=["lat", "lon"],
)
id_ds = id_da.to_dataset(name="idgauges")
gauge_ids_numeric = _to_numeric_gauge_ids(
self.gauge_ids,
context="idgauges",
)
id_ds = write_gauge_id(
ds=id_ds,
id=gauge_ids_numeric,
lat=self.gauge_lats,
lon=self.gauge_lons,
data_var="idgauges",
)
write_xarray_to_file(
id_ds, gauges_out_path / "idgauges.nc", "idgauges"
)
write_xarray_to_file(
id_ds,
gauges_out_path / "idgauges.asc",
"idgauges",
resolution=self.upscaled_resolution,
)
else:
logger.info("No gauges to write, skipping gauges file.")
else:
logger.info("No basin variable written; skipping mask/gauge files.")
[docs]
def write_single_variable_file(
self, data_var, var_name, out_path, cellsize, format
):
"""Write a single data variable to a specified file path."""
# set some attributes
for coord in data_var.coords:
data_var[coord].attrs = self.ds[coord].attrs
data_var.attrs = {
"title": self.VARIABLES[var_name]["title"],
"units": self.VARIABLES[var_name]["units"],
"creator": "Department of Computational Hydrosystems",
"institution": "Helmholtz Centre for Environmental Research - UFZ",
}
fname = out_path / f"{var_name}.{format}"
if format == "nc":
write_xarray_to_file(
data_var,
fname,
encoding={
var_name: {
"dtype": get_dtype(data_var[var_name]),
"_FillValue": self.VARIABLES[var_name]["_FillValue"],
}
},
)
elif format == "asc":
cellsize = cellsize or abs(float(data_var["lon"][1] - data_var["lon"][0]))
is_ascending = bool(data_var["lat"][0] < data_var["lat"][-1])
with fname.open("w") as file_object:
file_object.write(f"ncols {data_var[var_name].shape[1]}\n")
file_object.write(f"nrows {data_var[var_name].shape[0]}\n")
file_object.write(
f"xllcorner {float(data_var['lon'][0] - cellsize / 2)}\n"
)
if is_ascending:
file_object.write(
f"yllcorner {float(data_var['lat'][0] - cellsize / 2)}\n"
)
else:
file_object.write(
f"yllcorner {float(data_var['lat'][-1] - cellsize / 2)}\n"
)
file_object.write(f"cellsize {cellsize}\n")
file_object.write(
f"nodata_value {self.VARIABLES[var_name]['_FillValue']}\n"
)
if is_ascending:
vals = data_var[var_name].values[::-1, :]
else:
vals = data_var[var_name].values
np.savetxt(file_object, vals, delimiter=" ", fmt="%s")
else:
with ErrorLogger(logger):
msg = f'Format "{format}" unknown, use one of ["nc", "asc"]'
raise Exception(msg)
[docs]
def processing_data_variable(
self,
var_name,
cut_by_basin,
lat_slice=None,
lon_slice=None,
lat_slice_idx=None,
lon_slice_idx=None,
):
"""Process data variable, masking it and croping it spatial dimensions."""
logger.info(f"Processing {var_name}")
data = getattr(self, var_name)
if data is None:
logger.warning(f"No data for {var_name}")
return None
if cut_by_basin:
data[~self.catchment_mask] = self.VARIABLES[var_name]["_FillValue"]
if data is None:
logger.warning(f"No data for {var_name}")
return None
lon, lat = self._coords_l1()
logger.debug(
f"lon_min {np.min(lon):.3f}, lon_max {np.max(lon):.3f}, resulution: {self.resolutions.l1}"
)
logger.debug(f"{var_name} - mean {np.nanmean(data)}, max {np.nanmax(data)}")
logger.debug(f"Shape {data.shape}, lon {len(lon)}, lat {len(lat)}")
data_var = xr.Dataset(
{var_name: (["lat", "lon"], self._revert_data(data))},
coords={
"lon": lon, # [slice(3555, 3565)],
"lat": lat, # [slice(860, 870)],
},
)
if lat_slice is not None and lon_slice is not None:
logger.info(f"Cutting {var_name} data to correct spatial dimensions")
data_var = data_var.sel(lat=lat_slice, lon=lon_slice)
elif lat_slice_idx is not None and lon_slice_idx is not None:
logger.info(f"Cutting {var_name} data to correct spatial dimensions")
data_var = data_var.isel(lat=lat_slice_idx, lon=lon_slice_idx)
logger.debug(data_var)
return data_var
[docs]
def write_basin_id_file(self, data_vars, frame, out_path):
"""Write the basin_id file to specified path and set a sink value frame if specified."""
logger.info("Write to single file.")
logger.debug(f"data_vars: {data_vars}")
ds = xr.merge(data_vars.values())
# set some attributes
for coord in ds.coords:
ds[coord].attrs = self.ds[coord].attrs
ds.attrs = {
"title": "Hydrologic information",
"creator": "Department of Computational Hydrosystems",
"institution": "Helmholtz Centre for Environmental Research - UFZ",
}
for var_name in ds.data_vars:
ds[var_name].attrs = {
"long_name": self.VARIABLES[var_name]["title"],
"standard_name": self.VARIABLES[var_name]["title"],
"units": self.VARIABLES[var_name]["units"],
}
# logger.debug(f"lat_slice: {lat_slice}, lon_slice: {lon_slice}")
logger.debug(f"ds: {ds}")
ds = self.create_frame(ds, frame, FDIR_SINKVALUE[self.ftype])
# For the flow dir map fill masked cells adjecent to filled cells with sink instead of missing value
# fdir_filled = self.fill_adjacent_missing_with_sink(
# ds["flwdir"], FDIR_FILLVALUE[self.ftype], FDIR_SINKVALUE[self.ftype]
# )
# ds["flwdir"].data[:] = fdir_filled.data[:]
write_xarray_to_file(
ds,
out_path / self.out_var_name,
encoding={
var_name: {
"dtype": get_dtype(ds[var_name]),
"_FillValue": self.VARIABLES[var_name]["_FillValue"],
}
for var_name in ds.data_vars
},
)
logger.info(f"Basin Id has been written to {out_path / self.out_var_name}")
return ds
def _cell_edges(self, centers: np.ndarray) -> np.ndarray:
"""Compute edges (len=N+1) from center coords (len=N) on a regular grid."""
c = np.asarray(centers)
d = np.diff(c)
left = c[0] - 0.5 * d[0]
right = c[-1] + 0.5 * d[-1]
mids = (c[:-1] + c[1:]) / 2.0
return np.concatenate(([left], mids, [right]))
def _coarse_centers_from_edges(
self, edges: np.ndarray, k: int, n_blocks: int, ascending: bool
) -> np.ndarray:
"""
Given fine-grid edges, build coarse-grid centers for block size k.
Ensures coarse edges == fine edges over the cropped window.
"""
# we assume you've cropped L0 so len(fine_centers) is divisible by k
# The window's left edge and right edge are edges[0] and edges[k*n_blocks]
left_edge = edges[0]
dx_coarse = edges[k] - edges[0] # = k * dx_fine (works for asc/desc)
# centers are midpoints of each coarse cell
n = np.arange(n_blocks)
centers = left_edge + (n + 0.5) * dx_coarse
if not ascending and centers[0] < centers[-1]:
centers = centers[::-1]
return centers
[docs]
def upscale_mask_with_correct_coords(
self,
da: xr.DataArray,
factor: Optional[int] = None,
lon_name: str = "lon",
lat_name: str = "lat",
) -> xr.DataArray:
"""
Coarsen a 2D mask-like field by integer factor and assign correct coarse coords.
so that coarse *edges* equal fine *edges* of the cropped window.
"""
if factor is None:
factor, upscaled_resolution = get_upscaling_factor(
self.resolutions, l2=True
)
if factor < 1:
msg = "factor must be >= 1"
with ErrorLogger(logger):
raise ValueError(msg)
logger.info(f"Upscaling mask with factor {factor} to {upscaled_resolution}.")
# 1) coarsen over lon/lat windows
kx = ky = int(factor)
coarsen_map = {}
if lon_name in da.dims:
coarsen_map[lon_name] = kx
if lat_name in da.dims:
coarsen_map[lat_name] = ky
# Treat only explicit mask=1 (or True) as land; ignore fill values and NaNs.
if da.dtype == bool:
cond = da
else:
fillv = da.attrs.get("_FillValue", da.encoding.get("_FillValue"))
da_clean = da
if fillv is not None:
da_clean = da.where(da != fillv)
cond = da_clean == 1
out = cond.coarsen(coarsen_map, boundary="trim").any().astype("int8")
# 2) compute correct coarse coordinates from fine edges
lon_f = da[lon_name].values
lat_f = da[lat_name].values
lon_edges = self._cell_edges(lon_f)
lat_edges = self._cell_edges(lat_f)
asc_lon = lon_f[0] < lon_f[-1]
asc_lat = lat_f[0] < lat_f[-1]
logger.debug(f"asc_lon: {asc_lon}, asc_lat: {asc_lat}")
n_lon_blocks = out.sizes.get(lon_name, 1)
n_lat_blocks = out.sizes.get(lat_name, 1)
# figure out which portion of edges we used after boundary="trim":
# Since you cropped L0 to a multiple of factor, the coarsen starts at index 0
# and uses exactly n_blocks*k cells. So we can take edges[0 : n_blocks*k + 1].
lon_edges_win = (
lon_edges[: n_lon_blocks * kx + 1]
if asc_lon
else lon_edges[-(n_lon_blocks * kx + 1) :]
)
lat_edges_win = (
lat_edges[: n_lat_blocks * ky + 1]
if asc_lat
else lat_edges[-(n_lat_blocks * ky + 1) :]
)
lon_coarse = self._coarse_centers_from_edges(
lon_edges_win, kx, n_lon_blocks, asc_lon
)
lat_coarse = self._coarse_centers_from_edges(
lat_edges_win, ky, n_lat_blocks, asc_lat
)
asc_lat_coarse = lat_coarse[0] < lat_coarse[-1]
if asc_lat != asc_lat_coarse:
logger.warning(
"Coarse lat coordinate ascending order does not match fine grid; check calculations."
)
out = out.assign_coords({lon_name: lon_coarse, lat_name: lat_coarse})
out.name = "mask_L2"
# 3) (optional) log edges for verification
try:
lon_edges_coarse = self._cell_edges(out[lon_name].values)
lat_edges_coarse = self._cell_edges(out[lat_name].values)
logger.info(
f"Coarse lon edges: {lon_edges_coarse[0]:.6f} .. {lon_edges_coarse[-1]:.6f} "
f"(should equal fine window edges: {lon_edges_win[0]:.6f} .. {lon_edges_win[-1]:.6f})"
)
logger.info(
f"Coarse lat edges: {lat_edges_coarse[0]:.6f} .. {lat_edges_coarse[-1]:.6f} "
f"(should equal fine window edges: {lat_edges_win[0]:.6f} .. {lat_edges_win[-1]:.6f})"
)
except IndexError:
logger.debug("Could not log coarse edges for verification.")
logger.debug(f"lon_coarse: {out[lon_name].values}")
logger.debug(f"lat_coarse: {out[lat_name].values}")
return out
[docs]
def write_mask_file(self, ds, mask_file):
"""Write basin mask to specified path."""
if mask_file is not None:
logger.info("Writing mask file")
# name the variable mask
mask = np.where(
ds.basin > 0, 1, 0
) # if self.catchment_mask is None else self.catchment_mask
mask_file = Path(mask_file)
mask_da = xr.DataArray(
mask, coords={"lat": ds.lat, "lon": ds.lon}, dims=["lat", "lon"]
)
mask_da["lat"].attrs.update(
{
"units": "degrees_north",
"long_name": "latitude",
"standard_name": "latitude",
"axis": "Y",
}
)
mask_da["lon"].attrs.update(
{
"units": "degrees_east",
"long_name": "longitude",
"standard_name": "longitude",
"axis": "X",
}
)
logger.debug(
f"Created mask dataarray with shape {mask_da.shape} and stats min {mask_da.min().item()}, max {mask_da.max().item()}"
)
mask_ds = xr.Dataset({"land_mask": mask_da, "mask": mask_da})
mask_upscaled = None
if self.do_upscale:
mask_upscaled = mask_da
elif self.resolutions.l2 is not None:
mask_upscaled = self.upscale_mask_with_correct_coords(mask_da)
if mask_upscaled is not None:
mask_upscaled = mask_upscaled.rename({"lat": "lat_l2", "lon": "lon_l2"})
mask_ds["mask_l2"] = mask_upscaled
dims = set(mask_ds.dims)
all_coords = set(mask_ds.coords)
dim_coords = all_coords & dims # intersection
for var in dim_coords:
bounds_name = f"{var}_bnds"
try:
mask_ds.coords[bounds_name] = generate_bounds(mask_ds[var])
mask_ds[var].attrs["bounds"] = bounds_name
except IndexError:
logger.info(f"Could not generate bounds for coord {var}")
encoding = {
v: {"zlib": True, "complevel": 4, "shuffle": True, **NC_ENCODE_MASK}
for v in mask_ds.data_vars
}
write_xarray_to_file(mask_ds, mask_file, encoding=encoding)
logger.info(f"Mask file has been written to {mask_file}")
else:
logger.info("No mask file path specified.")
[docs]
def cut_to_filled_area(
self, buffer=0, repeat=False, raise_on_l2_alignment_mismatch=False
):
"""Create lat and lon slices to cut the data to the filled area."""
return cut_to_filled_area(
ds=self.ds,
resolutions=self.resolutions,
catchment_mask=self.catchment_mask,
buffer=buffer,
repeat=repeat,
raise_on_l2_alignment_mismatch=raise_on_l2_alignment_mismatch,
)
[docs]
def merge_catchment(path1, path2, out_path):
"""Merge the rolled and non-rolled file."""
# read the rolled and non-rolled files
ds1 = get_xarray_ds_from_file(path1, engine="netcdf4")
ds2 = get_xarray_ds_from_file(path2, engine="netcdf4")
# select all the basins in the border area
border_cells = (ds1.lon > CUTOFF_THRESHOLD) | (ds1.lon < (CUTOFF_THRESHOLD * -1))
mask_ids = np.unique(ds1["basin"].where(border_cells))
mask_ids = mask_ids[np.isfinite(mask_ids)]
mask_ids = mask_ids[mask_ids != ds1["basin"].attrs.get("_FillValue", 0)]
mask_ids = mask_ids[mask_ids != 0]
# get a mask of all the border area basins
mask = ds1["basin"].isin(mask_ids)
# modify the ids to avoid overlaps
ds2["basin"] = ds2["basin"] + ds1["basin"].max().item() + 1
# in the border area, use the rolled data, else the original
merged = xr.where(mask, ds2.reindex_like(ds1, method="nearest"), ds1)
write_xarray_to_file(merged, out_path)
[docs]
def is_data_global(ds, coordinate_slice):
"""Check if the longitude data is global."""
if coordinate_slice is not None:
ds_sliced = ds.sel(lon=coordinate_slice["lon"])
else:
ds_sliced = ds
try:
return (
"lon" in ds_sliced.coords
and ds_sliced.lon.min() < (CUTOFF_THRESHOLD * -1)
and ds_sliced.lon.max() > CUTOFF_THRESHOLD
)
except Exception as e:
logger.warning(e)
return False
def _is_list_of_float_tuples(a):
return isinstance(a, list) and all(
isinstance(x, tuple) and len(x) == 2 and all(isinstance(v, float) for v in x)
for x in a
)
[docs]
@log_arguments()
def create_catchment( # noqa: PLR0913, PLR0912, PLR0915
input_file,
output_path,
var_name,
var,
ftype,
gauge_coords=None,
coordinate_slices=None,
mask_file=None,
resolutions: Resolution = None,
frame=1,
upscale=False,
latlon=True,
available_mem=None,
ref_catchment_area=None,
max_distance_cells=5,
max_error=0.1,
gauge_ids=None,
ncpus=1,
output_vars=None,
gauge_opti_method="basinex",
shape_folder=None,
gauge_info_file="gauges_info",
id_gauges_out_path=None,
raise_on_fallback=True,
):
"""Delineate global or local catchments from flow direction/DEM data.
Supports full-domain basin mapping, single- or multi-gauge delineation from
coordinates and area, and optional shape-based outlet matching. Writes
basin IDs, gauge ID files, masks, and optional gauge metadata/shapes.
"""
logger.info(
f"Creating catchment file for {var_name} using {var} and {ftype} from {input_file}"
)
output_path = Path(output_path)
id_gauges_out_path = (
output_path if id_gauges_out_path is None else Path(id_gauges_out_path)
)
id_gauges_out_path.mkdir(parents=True, exist_ok=True)
if _is_list_of_float_tuples(gauge_coords) and len(gauge_coords) == 1:
gauge_coords = gauge_coords[0]
if isinstance(ref_catchment_area, list) and len(ref_catchment_area) == 1:
ref_catchment_area = ref_catchment_area[0]
if isinstance(ref_catchment_area, list):
if len(ref_catchment_area) != 1:
msg = "If gauge_coords is a list of one tuple, ref_catchment_area (if provided) must be a single value or a list of one value."
raise ValueError(msg)
ref_catchment_area = ref_catchment_area[0]
if resolutions is None:
resolutions = Resolution()
if var not in {"fdir", "dem"}:
with ErrorLogger(logger):
msg = f"Unexpected value for var={var}, must be 'fdir' or 'dem'"
raise ValueError(msg)
output_vars = _normalize_output_vars(output_vars)
chunking = available_mem is not None
with get_xarray_ds_from_file(
input_file,
var_name,
normalize_latlon_coords=True,
force_decending_y=True,
available_mem_gib=available_mem,
chunking=chunking,
) as input_ds:
coord_slices_default = False
if coordinate_slices is None:
coordinate_slices = {"lat": slice(None, None), "lon": slice(None, None)}
coord_slices_default = True
if coord_slices_default and shape_folder:
bounds = _shape_bounds_from_folder(shape_folder, gauge_ids, latlon=latlon)
if bounds is not None:
slices = _slices_from_bounds(bounds, input_ds.lat.data)
if slices is not None:
coordinate_slices = slices
# transform
transform = get_transformation_matrix_nc(input_ds, var_name)
logger.info(transform)
if coordinate_slices is None:
if shape_folder:
bounds = _shape_bounds_from_folder(
shape_folder, gauge_ids, latlon=latlon
)
logger.debug(f"Extracted bounds from shape_folder: {bounds}")
if bounds is not None:
slices = _slices_from_bounds(bounds, input_ds.lat.data)
logger.debug(f"Calculated slices from bounds: {slices}")
if slices is not None:
coordinate_slices = slices
logger.info(
f"Updated coordinate_slices based on shape_folder: {coordinate_slices}"
)
else:
logger.debug("No shape_folder provided, using no coordinate slices.")
coordinate_slices = {"lat": slice(None, None), "lon": slice(None, None)}
logger.debug(f"Using coordinate_slices: {coordinate_slices}")
def _compute_requested_outputs(catchment):
needs_uparea_grid = "uparea_grid" in output_vars
needs_upgrid = "upgrid" in output_vars
needs_grdare = "grdare" in output_vars
needs_basin = "basin" in output_vars
if resolutions.l1 is not None and upscale:
catchment.upscale(var)
elif needs_uparea_grid:
catchment.get_facc()
if needs_basin and catchment.basin is None:
catchment.get_basins()
if needs_grdare:
catchment.get_grid_area()
elif needs_upgrid:
catchment.get_upstream_area()
def _compute_requested_outputs(catchment):
needs_uparea_grid = "uparea_grid" in output_vars
needs_upgrid = "upgrid" in output_vars
needs_grdare = "grdare" in output_vars
needs_basin = "basin" in output_vars
if resolutions.l1 is not None and upscale:
catchment.upscale(var)
elif needs_uparea_grid:
catchment.get_facc()
if needs_basin and catchment.basin is None:
catchment.get_basins()
if needs_grdare:
catchment.get_grid_area()
elif needs_upgrid:
catchment.get_upstream_area()
has_shape_gauges = shape_folder is not None and gauge_ids is not None
if (
gauge_coords is None
and not has_shape_gauges
and is_data_global(input_ds, coordinate_slices)
):
logger.info("Creating global basin id file...")
if "basin" in output_vars:
temp_file1 = "hydro1.nc"
global_catchments = Catchment(
ds=input_ds,
var_name=var_name,
var=var,
ftype=ftype,
transform=transform,
latlon=latlon,
out_var_name=temp_file1,
do_shift=False,
resolutions=resolutions,
upscale=upscale,
)
# create a shifted version of the catchment to avoid border effects
temp_file2 = "hydro2.nc"
global_catchments_shifted = Catchment(
ds=input_ds,
var_name=var_name,
var=var,
ftype=ftype,
transform=transform,
latlon=latlon,
out_var_name=temp_file2,
do_shift=True,
resolutions=resolutions,
upscale=upscale,
)
catchments = [global_catchments, global_catchments_shifted]
for c in catchments:
_compute_requested_outputs(c)
c.write(
output_path,
single_file=True,
frame=frame,
mask_file=mask_file,
variables=output_vars,
id_gauges_out_path=id_gauges_out_path,
)
# add paths to the temp files
temp_file1 = Path(output_path, "hydro1.nc")
temp_file2 = Path(output_path, "hydro2.nc")
logger.info("Merging catchment files")
merge_catchment(
temp_file1,
temp_file2,
Path(output_path, "basin_ids.nc"),
)
# remove the temporary files
temp_file1.unlink()
temp_file2.unlink()
else:
global_catchments = Catchment(
ds=input_ds,
var_name=var_name,
var=var,
ftype=ftype,
transform=transform,
latlon=latlon,
out_var_name="basin_ids.nc",
do_shift=False,
resolutions=resolutions,
upscale=upscale,
)
_compute_requested_outputs(global_catchments)
global_catchments.write(
output_path,
single_file=True,
frame=frame,
mask_file=mask_file,
variables=output_vars,
id_gauges_out_path=id_gauges_out_path,
)
return
input_ds_sliced = input_ds.sel(
lat=coordinate_slices["lat"], lon=coordinate_slices["lon"]
)
logger.info("Cropped input dataset:")
logger.info(
f" lat {input_ds_sliced.lat.data[0]}, {input_ds_sliced.lat.data[-1]}"
)
logger.info(
f" lon {input_ds_sliced.lon.data[0]}, {input_ds_sliced.lon.data[-1]}"
)
sliced_transform = get_transformation_matrix_nc(input_ds_sliced, var_name)
single_shape_gauge = (
gauge_coords is None
and shape_folder is not None
and gauge_ids is not None
and not isinstance(gauge_ids, list)
)
if (
gauge_coords is not None and isinstance(gauge_coords, tuple)
) or single_shape_gauge:
logger.info(f"Creating catchment for gauge coordinates {gauge_coords}")
c = Catchment(
ds=input_ds_sliced,
var_name=var_name,
var=var,
ftype=ftype,
transform=sliced_transform,
latlon=latlon,
out_var_name="basin_ids.nc",
do_shift=False,
resolutions=resolutions,
upscale=upscale,
)
single_ref_area = (
ref_catchment_area[0]
if isinstance(ref_catchment_area, list) and ref_catchment_area
else ref_catchment_area
)
gauge = Gauge(
gauge_id=gauge_ids if not isinstance(gauge_ids, list) else gauge_ids[0],
lat=gauge_coords[0] if gauge_coords is not None else None,
lon=gauge_coords[1] if gauge_coords is not None else None,
area=single_ref_area,
)
gauge = c.delineate_basin(
gauge,
max_distance_cells=max_distance_cells,
max_error=max_error,
gauge_opti_method=gauge_opti_method,
shape_folder=shape_folder,
raise_on_fallback=raise_on_fallback,
)
l0_shape_gdf = None
if upscale and c.catchment_mask is not None:
write_gauges_out(
gauge, id_gauges_out_path / f"{gauge_info_file}_{resolutions.l0}"
)
try:
l0_shape_gdf = _vectorize_mask_to_gdf(
c.catchment_mask,
c.transform,
_shape_crs(c.latlon),
value_name="basin",
)
except Exception as exc:
logger.warning(f"Could not build L0 basin shape: {exc}")
else:
write_gauges_out(gauge, id_gauges_out_path / gauge_info_file)
c.write_basin_shape(
output_path / "shapes",
gauge_id=gauge_ids if not isinstance(gauge_ids, list) else gauge_ids[0],
)
_compute_requested_outputs(c)
if upscale and c.is_upscaled and l0_shape_gdf is not None and c.gauge_lats:
logger.info("Applying L1 correction using L0 basin shape.")
new_coords = c.correct_gauge_location_l1_from_shape(
l0_shape_gdf,
(c.gauge_lats[-1], c.gauge_lons[-1]),
max_distance_cells=max_distance_cells,
max_error=max_error,
ref_catchment_area=ref_catchment_area,
)
if new_coords is not None:
c.update_gauge_coords(
gauge_ids if not isinstance(gauge_ids, list) else gauge_ids[0],
new_coords[0],
new_coords[1],
)
gauge.update(lat=new_coords[0], lon=new_coords[1])
write_gauges_out(
gauge,
id_gauges_out_path / f"{gauge_info_file}_{c.upscaled_resolution}",
)
c.write(
output_path,
single_file=True,
cut_by_basin=True,
mask_file=mask_file,
frame=frame,
buffer=frame,
variables=output_vars,
id_gauges_out_path=id_gauges_out_path,
)
return
logger.info("Creating basin id file for region.")
c = Catchment(
ds=input_ds_sliced,
var_name=var_name,
var=var,
ftype=ftype,
transform=sliced_transform,
latlon=latlon,
out_var_name="basin_ids.nc",
do_shift=False,
resolutions=resolutions,
upscale=upscale,
)
gauge_infos = None
gauges = []
shape_only_multi_gauges = (
gauge_coords is None
and shape_folder is not None
and isinstance(gauge_ids, list)
)
if (
_is_list_of_float_tuples(gauge_coords)
and isinstance(gauge_ids, list)
and len(gauge_coords) == len(gauge_ids)
) or shape_only_multi_gauges:
logger.info(f"Creating catchments for gauge coordinates {gauge_coords}")
upstream_area = c.calc_upstream_area()
lon = get_coord_values(input_ds_sliced, lon=True)
lat = get_coord_values(input_ds_sliced, lat=True)
def _process_gauge(i, gc, lon, lat):
ref_area = (
ref_catchment_area[i]
if isinstance(ref_catchment_area, list)
else ref_catchment_area
)
if gc is not None and not (
lon.min() <= gc[1] <= lon.max() and lat.min() <= gc[0] <= lat.max()
):
logger.warning(
f"Gauge coordinate {gc} is outside the domain lon: [{lon.min()}, {lon.max()}], lat: [{lat.min()}, {lat.max()}]"
)
if shape_folder is None:
return None
gc = None
(
outlet_idx,
error,
gauge_lat,
gauge_lon,
distance_error,
score,
shape_error,
used_method,
) = c.get_best_gauge_coordinate(
upstream_area=upstream_area,
gauge_coords=gc,
ref_catchment_area=ref_area,
max_distance_cells=max_distance_cells,
max_error=max_error,
method=gauge_opti_method,
shape_folder=shape_folder,
gauge_id=gauge_ids[i],
raise_on_fallback=raise_on_fallback,
)
return {
"gauge_id": gauge_ids[i],
"gauge_lat": gauge_lat,
"gauge_lon": gauge_lon,
"lat_old": gc[0] if gc is not None else np.nan,
"lon_old": gc[1] if gc is not None else np.nan,
"area_old": ref_area,
"outlet_idx": outlet_idx,
"error": error,
"distance_error": distance_error,
"score": score,
"shape_error": shape_error,
"method": used_method,
"ref_area": ref_area,
}
gauge_infos = Parallel(n_jobs=ncpus, prefer="threads")(
delayed(_process_gauge)(
i,
None if gauge_coords is None else gc,
lon,
lat,
)
for i, gc in enumerate(
gauge_ids if gauge_coords is None else gauge_coords
)
)
gauge_infos = [gi for gi in gauge_infos if gi is not None]
if not gauge_infos:
logger.warning("No valid gauge coordinates found inside domain.")
else:
logger.info(
f"Found {len(gauge_infos)} valid gauge coordinates inside domain."
)
shape_dir = Path(output_path) / "shapes"
outlet_idxs = [gi["outlet_idx"] for gi in gauge_infos]
outlet_linear = np.array(
[np.ravel_multi_index(idx, c._fdir.shape) for idx in outlet_idxs],
dtype=np.int64,
)
# combine all river masks
streams_mask = None
logger.info("Delineating basins for all gauges.")
# if streams_mask is None or not np.any(streams_mask):
# logger.warning("No river mask found for any gauge.")
# streams_mask = np.asarray(c._fdir.stream_order() >= 4, dtype=bool)
try:
basins = c._fdir.basins(
idxs=outlet_linear,
streams=streams_mask,
)
except Exception as exc:
logger.exception(f"pyflwdir.basins(idxs=...) failed: {exc}")
with ErrorLogger(logger):
raise exc
logger.info("Performing sanity checks for all delineated basins.")
gauges = []
for gi in gauge_infos:
outlet_idx = gi["outlet_idx"]
basin_id = int(basins[outlet_idx])
if basin_id == 0:
logger.warning(
f"No basin id found for gauge_id {gi['gauge_id']} at "
f"outlet {outlet_idx}"
)
continue
catchment_mask = basins == basin_id
uparea_at_outlet = (
upstream_area[outlet_idx]
if upstream_area is not None
else np.nan
)
failed = c.delineation_sanity_check(
catchment_mask,
basins,
uparea_at_outlet,
gi["ref_area"],
max_error,
False,
gi["gauge_id"],
)
if failed:
continue
if upscale:
try:
gi["l0_shape_gdf"] = _vectorize_mask_to_gdf(
catchment_mask,
c.transform,
_shape_crs(c.latlon),
value_name="basin",
)
except Exception as exc:
logger.warning(
f"Could not build L0 basin shape for gauge "
f"{gi['gauge_id']}: {exc}"
)
gi["l0_shape_gdf"] = None
c.save_coords(gi["gauge_id"], gi["gauge_lat"], gi["gauge_lon"])
gauge = Gauge(
gauge_id=gi["gauge_id"],
lat=gi["lat_old"],
lon=gi["lon_old"],
area=gi["area_old"],
)
gauge.update(
lat=gi["gauge_lat"],
lon=gi["gauge_lon"],
area=uparea_at_outlet,
distance_error=gi["distance_error"],
area_error=gi["error"],
score=gi["score"],
shape_error=gi["shape_error"],
method=gi["method"],
)
gauges.append(gauge)
c.write_basin_shape(
shape_dir,
gauge_id=gi["gauge_id"],
basin_mask=catchment_mask,
)
if upscale:
write_gauges_out(
gauges,
id_gauges_out_path / f"{gauge_info_file}_{resolutions.l0}",
)
else:
write_gauges_out(gauges, id_gauges_out_path / gauge_info_file)
_compute_requested_outputs(c)
if upscale and c.is_upscaled and gauges:
logger.info(f"Applying L1 correction for {len(gauge_infos)} gauges.")
for gauge in gauges:
gi = next(
(
item
for item in gauge_infos
if item["gauge_id"] == gauge.gauge_id
),
None,
)
if gi is None:
continue
l0_shape_gdf = gi.get("l0_shape_gdf")
if l0_shape_gdf is None:
continue
new_coords = c.correct_gauge_location_l1_from_shape(
l0_shape_gdf,
(gauge.lat, gauge.lon),
max_distance_cells=max_distance_cells,
max_error=max_error,
ref_catchment_area=gi.get("ref_area"),
)
if new_coords is not None:
c.update_gauge_coords(
gi["gauge_id"],
new_coords[0],
new_coords[1],
)
gauge.update(lat=new_coords[0], lon=new_coords[1])
write_gauges_out(
gauges,
id_gauges_out_path / f"{gauge_info_file}_{c.upscaled_resolution}",
)
c.write(
output_path,
single_file=True,
mask_file=mask_file,
frame=frame,
variables=output_vars,
id_gauges_out_path=id_gauges_out_path,
)
