diff --git a/data/shared/match_station_basin.py b/data/shared/match_station_basin.py new file mode 100644 index 00000000..f99747a7 --- /dev/null +++ b/data/shared/match_station_basin.py @@ -0,0 +1,254 @@ +"""Find unique HydroBASINS basins for GloFAS stations. + +For each GloFAS station in a country, finds the highest-level (largest) +HydroBASINS basin that contains only that station among all stations +from the same country. + +The main function `match_station_basin()` should be called once per country +with all stations from that country. + +Workflow: + 1. Load stations for a specific country + 2. Spatial join stations to basin polygons at each level (4-12) + 3. For each level, count how many stations from the same country fall in each basin + 4. For each station, find the coarsest level where it's the only one in the basin + 5. Return GeoDataFrame with stationCode and corresponding basin geometries + +Key outputs: + - GeoDataFrame with stationCode, basin geometry, basin level, and HYBAS_ID + - Basins larger than max_basin_area_km2 are excluded from consideration +""" + +import glob +import os + +import geopandas as gpd + +MIN_LEVEL = 4 +MAX_LEVEL = 12 +MAX_BASIN_AREA_KM2 = 50_000 # discard basins larger than this + + +def match_station_basin( + stations_gdf: gpd.GeoDataFrame, + basins_dir: str, + country_code: str, + min_level: int = MIN_LEVEL, + max_level: int = MAX_LEVEL, + max_basin_area_km2: float = MAX_BASIN_AREA_KM2, +) -> gpd.GeoDataFrame: + """ + Find the highest-level (largest) HydroBASINS basin for all GloFAS stations in a country. + + Parameters + ---------- + stations_gdf : gpd.GeoDataFrame + GeoDataFrame with GloFAS stations from a specific country. Must include: + - 'stationCode': unique station identifier + - geometry: Point geometries + basins_dir : str + Path to directory containing HydroBASINS shapefiles named like + 'hybas_*_lev{level:02d}_*.shp' + country_code : str + Country identifier code (e.g., 'UGA', 'KEN', 'TZA') + min_level : int, default 4 + Minimum (coarsest) basin level to consider + max_level : int, default 12 + Maximum (finest) basin level to consider + max_basin_area_km2 : float, default 50000 + Maximum basin area in km² to consider (larger basins are filtered out) + + Returns + ------- + gpd.GeoDataFrame + GeoDataFrame with columns: + - 'stationCode': station identifier + - geometry: basin polygon geometry + - 'basin_level': the level at which the basin was found + - 'HYBAS_ID': HydroBASINS basin ID + """ + # Validate input + required_cols = ["stationCode", "geometry"] + missing = [col for col in required_cols if col not in stations_gdf.columns] + if missing: + raise ValueError(f"stations_gdf missing required columns: {missing}") + + # Add country as a temporary column + stations_with_country = stations_gdf.copy() + stations_with_country["_country"] = country_code + + # Load basins excluding levels outside the specified range + basins = _load_basins_from_path(basins_dir, min_level, max_level) + + # Reproject all basin levels to match the stations CRS so that spatial + # joins and the output geometries are all in a consistent coordinate system + stations_crs = stations_with_country.crs + if stations_crs is None: + raise ValueError( + "stations_gdf must have a CRS set. " + "Assign a CRS with stations_gdf.set_crs(...) before calling this function." + ) + reprojected_basins = {} + for level, gdf in basins.items(): + if gdf.crs is None: + pattern = os.path.join(basins_dir, f"*lev{level:02d}*.shp") + matching_files = sorted(glob.glob(pattern)) + filenames = ", ".join(os.path.basename(p) for p in matching_files) + raise ValueError( + f"Basin shapefile for level {level} is missing CRS metadata" + + (f": {filenames}" if filenames else f" (pattern: {pattern})") + ) + reprojected_basins[level] = ( + gdf.to_crs(stations_crs) if gdf.crs != stations_crs else gdf + ) + basins = reprojected_basins + + # Find best basin for each station + best_basin_per_station = _find_highest_unique_basin( + stations_with_country, basins, max_basin_area_km2 + ) + + # Build result GeoDataFrame with basin geometries + result_rows = [] + for _, row in stations_gdf.iterrows(): + code = row["stationCode"] + if code in best_basin_per_station: + hybas_id, level = best_basin_per_station[code] + # Get the basin geometry + basin_gdf = basins[level] + basin_geom = basin_gdf[basin_gdf["HYBAS_ID"] == hybas_id].geometry.iloc[0] + result_rows.append( + { + "stationCode": code, + "HYBAS_ID": hybas_id, + "basin_level": level, + "geometry": basin_geom, + } + ) + + if not result_rows: + # Return empty GeoDataFrame with correct schema + return gpd.GeoDataFrame( + columns=["stationCode", "HYBAS_ID", "basin_level", "geometry"], + crs=stations_gdf.crs, + ) + + result_gdf = gpd.GeoDataFrame(result_rows, crs=stations_gdf.crs) + + # Print summary + print(f"Assigned {len(result_gdf)} out of {len(stations_gdf)} stations to basins") + for lvl in range(min_level, max_level + 1): + n = (result_gdf["basin_level"] == lvl).sum() + if n > 0: + print(f" Level {lvl:2d}: {n:4d} stations") + + return result_gdf + + +def _load_basins_from_path(basins_dir: str, min_level: int, max_level: int) -> dict: + """Load HydroBASINS shapefiles from directory. + + Searches for shapefiles matching the pattern '*lev{level:02d}*.shp' + for each level from min_level to max_level. + + Parameters + ---------- + basins_dir : str + Path to directory containing HydroBASINS shapefiles + min_level : int + Minimum basin level to load + max_level : int + Maximum basin level to load + + Returns + ------- + dict + Dictionary mapping level (int) to GeoDataFrame with basin polygons. + Each GeoDataFrame must contain 'HYBAS_ID', 'SUB_AREA', and 'geometry'. + """ + basins = {} + for level in range(min_level, max_level + 1): + # Try to find shapefile matching common naming patterns + pattern = os.path.join(basins_dir, f"*lev{level:02d}*.shp") + matches = glob.glob(pattern) + if matches: + print(f" Loading level {level} from {matches[0]} …") + basins[level] = gpd.read_file(matches[0]) + else: + print(f" Level {level}: shapefile not found, skipping") + return basins + + +def _find_highest_unique_basin( + stations: gpd.GeoDataFrame, + basins: dict, + max_basin_area_km2: float, +) -> dict: + """Find the coarsest basin level where each station is unique. + + For each station, identifies the largest (coarsest level) basin where + the station is the only one from its country. Iterates from coarsest + to finest levels and stops at the first level where the station is alone. + + Algorithm: + 1. Spatial join stations to basins at each level + 2. Count same-country stations per basin at each level + 3. For each station, find coarsest level where count == 1 + 4. Return mapping of stationCode to (HYBAS_ID, level) + + Parameters + ---------- + stations : gpd.GeoDataFrame + Stations with '_country' column and Point geometries + basins : dict + Dictionary mapping level to basin GeoDataFrames (already filtered + by min/max level during loading) + max_basin_area_km2 : float + Maximum basin area to consider (larger basins filtered out) + + Returns + ------- + dict + Mapping from stationCode (str) to tuple of (HYBAS_ID, level) + """ + # Pre-compute: spatial join all stations to each level once + print(" Spatial-joining stations to basins at each level …") + joins = {} + for level, basin_gdf in basins.items(): + # filter out basins that exceed the maximum area + small_enough = basin_gdf[basin_gdf["SUB_AREA"] <= max_basin_area_km2] + joined = gpd.sjoin( + stations, + small_enough[["HYBAS_ID", "geometry"]], + how="left", + predicate="within", + ) + # keep station code, country, and basin ID + joins[level] = joined[["stationCode", "_country", "HYBAS_ID"]].copy() + joins[level]["level"] = level + + # For each level, count same-country stations per basin + print(" Counting same-country stations per basin at each level …") + for level in joins: + df = joins[level] + counts = ( + df.groupby(["_country", "HYBAS_ID"]) + .size() + .reset_index(name="n_same_country") + ) + joins[level] = df.merge(counts, on=["_country", "HYBAS_ID"], how="left") + + # For each station, find the coarsest level where it is alone + print(" Assigning best level per station …") + best = {} + + for level in sorted(joins.keys()): + df = joins[level] + # stations that are alone in their country+basin at this level + alone = df[df["n_same_country"] == 1] + for _, row in alone.iterrows(): + code = row["stationCode"] + if code not in best: # only keep the coarsest (first found) + best[code] = (row["HYBAS_ID"], level) + + return best