Skip to content

flood

Flood susceptibility and event-risk mapping. Entry point: FloodRiskUseCase.

from climate_change import FloodRiskUseCase

Use case

use_case

FloodRiskUseCase

FloodRiskUseCase(dask_engine: DaskEngine)

Bases: BaseUseCase

Entry point for the flood risk domain.

Minimal config (all optional — defaults match the Niger flood notebook example):

{ "aoi_geojson": {"type": "Polygon", "coordinates": [...]}, "gee_project": "your-gee-project-id", "model_type": "ensemble", # "rf" | "xgboost" | "ensemble" "output_dir": "outputs", "prefix": "flood_2022_2023", "flood_event": "Niger River flood pulse Oct 2022-Jan 2023", # Main analysis period "start_date": "2022-01-01", "end_date": "2023-01-31", # Flood event windows "pre_flood_start": "2022-05-01", "pre_flood_end": "2022-07-31", "post_flood_start": "2022-10-01", "post_flood_end": "2023-01-31", # Rainfall windows "rain_7d_start": "2022-08-18", "rain_7d_end": "2022-08-25", "rain_30d_start": "2022-08-01", "rain_30d_end": "2022-08-31", # Sentinel-2 MNDWI window "mndwi_start": "2022-10-01", "mndwi_end": "2023-01-31", # JRC flood label window "flood_label_start":"2021-01-01", "flood_label_end":"2021-12-31", # Sampling "n_pixels": 3000, "scale": 90, }

Source code in flood/use_case.py
def __init__(self, dask_engine: DaskEngine) -> None:
    super().__init__(dask_engine)

fetch_data

fetch_data(config: AnalysisConfig) -> dict

Authenticate GEE and parse the AOI geometry (see _fetch_from_dict).

Source code in flood/use_case.py
def fetch_data(self, config: AnalysisConfig) -> dict:
    """Authenticate GEE and parse the AOI geometry (see `_fetch_from_dict`)."""
    flat: dict = {
        "aoi_geojson": config.aoi_geojson,
        "start_date": config.start_date,
        "end_date": config.end_date,
        "country": config.country,
        **config.extra_params,
    }
    return self._fetch_from_dict(flat)

preprocess

preprocess(raw_data: dict, config: AnalysisConfig) -> dict

Build the feature stack and sample flooded/dry pixels (see _preprocess_raw).

Source code in flood/use_case.py
def preprocess(self, raw_data: dict, config: AnalysisConfig) -> dict:
    """Build the feature stack and sample flooded/dry pixels (see `_preprocess_raw`)."""
    return self._preprocess_raw(raw_data)

run_model

run_model(features: dict, config: AnalysisConfig) -> AnalysisOutput

Train models, export the risk COG, and package an AnalysisOutput.

Source code in flood/use_case.py
def run_model(self, features: dict, config: AnalysisConfig) -> AnalysisOutput:
    """Train models, export the risk COG, and package an `AnalysisOutput`."""
    model_type = config.extra_params.get("model_type", "ensemble")
    dict_config = {"model_type": model_type, **config.extra_params}
    result = self._run_model_dict(features, dict_config)
    raster_paths: dict[str, str] | None = None
    raster_error: str | None = None
    raster_result: FloodCogResult | None = None
    try:
        raster_result = export_flood_cog(
            rf_model=features["_rf"],
            xgb_model=features["_xgb"],
            datasets=features["datasets"],
            output_dir=dict_config.get("output_dir", "outputs"),
            prefix=dict_config.get("prefix", "flood"),
            model_type=model_type,
            aoi_geojson=config.aoi_geojson,
        )
        raster_paths = {"flood_risk": raster_result["flood_risk"]}
    except Exception as exc:
        raster_error = str(exc)

    if raster_paths and raster_paths.get("flood_risk"):
        distribution = flood_raster_distribution(raster_paths["flood_risk"])
        labels = distribution["labels"]
        percentages = distribution["percentages"]
        counts = distribution["counts"]
        result["charts"]["risk_distribution"] = {
            "labels": labels,
            "data": percentages,
            "counts": counts,
            "colors": [
                "#E74C3C",
                "#E67E22",
                "#F1C40F",
                "#2ECC71",
            ],
        }
        pct_by_label = dict(zip(labels, percentages, strict=False))
        count_by_label = dict(zip(labels, counts, strict=False))
        result["stats"].update(
            {
                "very_high_risk_pct": pct_by_label.get("Very High", 0.0),
                "high_risk_pct": pct_by_label.get("High", 0.0),
                "medium_risk_pct": pct_by_label.get("Medium", 0.0),
                "low_risk_pct": pct_by_label.get("Low", 0.0),
                "very_high_risk_pixels": count_by_label.get("Very High", 0),
                "high_risk_pixels": count_by_label.get("High", 0),
                "medium_risk_pixels": count_by_label.get("Medium", 0),
                "low_risk_pixels": count_by_label.get("Low", 0),
                "mapped_pixel_count": distribution["valid_pixel_count"],
            }
        )
        total_area_ha = _aoi_area_ha(config.aoi_geojson)
        if total_area_ha:
            result["stats"]["total_area_ha"] = round(total_area_ha, 1)
            _attach_risk_area(result["charts"]["risk_distribution"], total_area_ha)
        _attach_population_stats(result, raster_result, labels)

    geojson_features = []
    if config.aoi_geojson:
        geojson_features.append(
            {
                "type": "Feature",
                "geometry": config.aoi_geojson,
                "properties": {"type": "boundary"},
            }
        )
    for p in result.pop("_sample_points", []):
        geojson_features.append(
            {
                "type": "Feature",
                "geometry": {"type": "Point", "coordinates": [p["lon"], p["lat"]]},
                "properties": {"risk_class": p["risk_class"]},
            }
        )
    metadata = {
        "model": model_type,
        "country": config.country,
        "start_date": config.start_date,
        "end_date": config.end_date,
        "raster": raster_paths or {},
        "windows": {
            "pre_flood": {
                "start": dict_config.get("pre_flood_start") or dict_config.get("pre_sar_start"),
                "end": dict_config.get("pre_flood_end") or dict_config.get("pre_sar_end"),
            },
            "post_flood": {
                "start": dict_config.get("post_flood_start")
                or dict_config.get("flood_sar_start"),
                "end": dict_config.get("post_flood_end") or dict_config.get("flood_sar_end"),
            },
            "rainfall_7d": {
                "start": dict_config.get("rain_7d_start"),
                "end": dict_config.get("rain_7d_end"),
            },
            "rainfall_30d": {
                "start": dict_config.get("rain_30d_start"),
                "end": dict_config.get("rain_30d_end"),
            },
            "mndwi": {
                "start": dict_config.get("mndwi_start"),
                "end": dict_config.get("mndwi_end"),
            },
        },
    }
    if raster_error:
        metadata["raster_error"] = raster_error

    return AnalysisOutput(
        module="flood",
        geojson={"type": "FeatureCollection", "features": geojson_features},
        raster_path=raster_paths,
        stats={**result["stats"], "country": config.country},
        shap=result.get("charts", {}).get("shap"),
        charts=result.get("charts", {}),
        metadata=metadata,
    )

run

run(config: dict) -> dict

Single AOI analysis — full pipeline in one call.

Source code in flood/use_case.py
def run(self, config: dict) -> dict:
    """Single AOI analysis — full pipeline in one call."""
    raw_data = self._fetch_from_dict(config)
    features = self._preprocess_raw(raw_data)
    result = self._run_model_dict(features, config)

    model_type = config.get("model_type", "ensemble")
    raster_paths: dict[str, str] | None = None
    raster_error: str | None = None
    raster_result: FloodCogResult | None = None
    try:
        raster_result = export_flood_cog(
            rf_model=features["_rf"],
            xgb_model=features["_xgb"],
            datasets=features["datasets"],
            output_dir=config.get("output_dir", "outputs"),
            prefix=config.get("prefix", "flood"),
            model_type=model_type,
            aoi_geojson=config.get("aoi_geojson"),
        )
        raster_paths = {"flood_risk": raster_result["flood_risk"]}
    except Exception as exc:
        raster_error = str(exc)

    result["raster"] = raster_paths or {}
    result["stats"].update(
        {
            "bbox": features["bbox"],
            "flood_event": config.get("flood_event", ""),
            "prefix": config.get("prefix", "flood"),
        }
    )
    if raster_error:
        result["raster_error"] = raster_error

    if raster_paths and raster_paths.get("flood_risk"):
        distribution = flood_raster_distribution(raster_paths["flood_risk"])
        labels = distribution["labels"]
        percentages = distribution["percentages"]
        counts = distribution["counts"]
        result["charts"]["risk_distribution"] = {
            "labels": labels,
            "data": percentages,
            "counts": counts,
            "colors": [
                "#E74C3C",
                "#E67E22",
                "#F1C40F",
                "#2ECC71",
            ],
        }
        pct_by_label = dict(zip(labels, percentages, strict=False))
        count_by_label = dict(zip(labels, counts, strict=False))
        result["stats"].update(
            {
                "very_high_risk_pct": pct_by_label.get("Very High", 0.0),
                "high_risk_pct": pct_by_label.get("High", 0.0),
                "medium_risk_pct": pct_by_label.get("Medium", 0.0),
                "low_risk_pct": pct_by_label.get("Low", 0.0),
                "very_high_risk_pixels": count_by_label.get("Very High", 0),
                "high_risk_pixels": count_by_label.get("High", 0),
                "medium_risk_pixels": count_by_label.get("Medium", 0),
                "low_risk_pixels": count_by_label.get("Low", 0),
                "mapped_pixel_count": distribution["valid_pixel_count"],
            }
        )
        total_area_ha = _aoi_area_ha(config.get("aoi_geojson"))
        if total_area_ha:
            result["stats"]["total_area_ha"] = round(total_area_ha, 1)
            _attach_risk_area(result["charts"]["risk_distribution"], total_area_ha)
        _attach_population_stats(result, raster_result, labels)
    return result

run_date_ranges

run_date_ranges(config: dict, date_ranges: list[dict]) -> list[dict]

Not supported — flood is single-area only. See UseCaseInfo.single_area_only.

Source code in flood/use_case.py
def run_date_ranges(self, config: dict, date_ranges: list[dict]) -> list[dict]:
    """Not supported — flood is single-area only. See UseCaseInfo.single_area_only."""
    raise ValueError(
        "Flood analysis only supports a single AOI/date-range at a time. Each "
        "run requires six independent satellite date windows (pre/post-flood "
        "SAR, 7-day and 30-day rainfall, Sentinel-2 MNDWI, JRC flood label) to "
        "all resolve to non-empty imagery — running many configs in parallel "
        "multiplies the odds that one window comes up empty, so the "
        "two-date-range comparison and multi-region batch modes are disabled "
        "for this module rather than run unreliably."
    )

run_multi_regions

run_multi_regions(configs: list[dict]) -> list[dict]

Not supported — flood is single-area only. See UseCaseInfo.single_area_only.

Source code in flood/use_case.py
def run_multi_regions(self, configs: list[dict]) -> list[dict]:
    """Not supported — flood is single-area only. See UseCaseInfo.single_area_only."""
    raise ValueError(
        "Flood analysis only supports a single AOI/date-range at a time. Each "
        "run requires six independent satellite date windows (pre/post-flood "
        "SAR, 7-day and 30-day rainfall, Sentinel-2 MNDWI, JRC flood label) to "
        "all resolve to non-empty imagery — running many configs in parallel "
        "multiplies the odds that one window comes up empty, so the "
        "two-date-range comparison and multi-region batch modes are disabled "
        "for this module rather than run unreliably."
    )

Features

features

fetch_terrain

fetch_terrain(aoi: Geometry, scale: int = 90) -> xr.Dataset

Download SRTM elevation from GEE. Returns Dataset with variable 'elevation' and (lat, lon) coords.

Source code in flood/features.py
def fetch_terrain(aoi: ee.Geometry, scale: int = 90) -> xr.Dataset:
    """
    Download SRTM elevation from GEE.
    Returns Dataset with variable 'elevation' and (lat, lon) coords.
    """
    dem = ee.Image("USGS/SRTMGL1_003").select("elevation").clip(aoi)
    da = _download_image(dem, aoi, scale).squeeze()
    return xr.Dataset({"elevation": da.rename({"x": "lon", "y": "lat"})})

fetch_twi

fetch_twi(aoi: Geometry, scale: int = 500) -> xr.Dataset

Compute Topographic Wetness Index from HydroSHEDS flow accumulation + SRTM slope. Returns Dataset with variables 'twi' and 'flow_acc_log'.

Source code in flood/features.py
def fetch_twi(aoi: ee.Geometry, scale: int = 500) -> xr.Dataset:
    """
    Compute Topographic Wetness Index from HydroSHEDS flow accumulation + SRTM slope.
    Returns Dataset with variables 'twi' and 'flow_acc_log'.
    """
    dem = ee.Image("USGS/SRTMGL1_003").select("elevation").reproject("EPSG:4326", None, 500)
    slope_rad = (
        ee.Terrain.products(dem)
        .select("slope")
        .clip(aoi)
        .multiply(np.pi / 180)
        .max(ee.Image(0.001))
    )
    flow_acc = ee.Image("WWF/HydroSHEDS/30ACC").select("b1").clip(aoi)
    twi = flow_acc.multiply(810_000).divide(slope_rad.tan()).log().rename("twi")
    flow_acc_log = flow_acc.add(1).log().rename("flow_acc_log")

    raw = _download_image(
        twi.addBands(flow_acc_log), aoi, scale, band_names=["twi", "flow_acc_log"]
    )
    ds = raw.to_dataset(dim="band").rename({"x": "lon", "y": "lat"})
    return ds

fetch_sar_change

fetch_sar_change(aoi: Geometry, pre_start: str, pre_end: str, flood_start: str, flood_end: str, scale: int = 90) -> xr.Dataset

Compute Sentinel-1 VV backscatter change (pre − flood). Positive values indicate a backscatter drop → open water / flood signal. Returns Dataset with variable 'vv_change'.

Source code in flood/features.py
def fetch_sar_change(
    aoi: ee.Geometry,
    pre_start: str,
    pre_end: str,
    flood_start: str,
    flood_end: str,
    scale: int = 90,
) -> xr.Dataset:
    """
    Compute Sentinel-1 VV backscatter change (pre − flood).
    Positive values indicate a backscatter drop → open water / flood signal.
    Returns Dataset with variable 'vv_change'.
    """
    sar = (
        ee.ImageCollection("COPERNICUS/S1_GRD")
        .filterBounds(aoi)
        .filter(ee.Filter.eq("instrumentMode", "IW"))
        .filter(ee.Filter.listContains("transmitterReceiverPolarisation", "VV"))
        .select("VV")
    )
    vv_change = (
        sar.filterDate(pre_start, pre_end)
        .mean()
        .unmask(0)
        .clip(aoi)
        .subtract(sar.filterDate(flood_start, flood_end).mean().unmask(0).clip(aoi))
        .rename("vv_change")
    )
    da = _download_image(vv_change, aoi, scale).squeeze()
    return xr.Dataset({"vv_change": da.rename({"x": "lon", "y": "lat"})})

fetch_rainfall

fetch_rainfall(aoi: Geometry, start_7d: str, end_7d: str, start_30d: str, end_30d: str, scale: int = 500) -> xr.Dataset

Download CHIRPS cumulative rainfall for two windows. Returns Dataset with variables 'rainfall_7d' and 'rainfall_30d'.

Source code in flood/features.py
def fetch_rainfall(
    aoi: ee.Geometry,
    start_7d: str,
    end_7d: str,
    start_30d: str,
    end_30d: str,
    scale: int = 500,
) -> xr.Dataset:
    """
    Download CHIRPS cumulative rainfall for two windows.
    Returns Dataset with variables 'rainfall_7d' and 'rainfall_30d'.
    """
    chirps = ee.ImageCollection("UCSB-CHG/CHIRPS/DAILY").filterBounds(aoi)
    rain = (
        chirps.filterDate(start_7d, end_7d)
        .sum()
        .clip(aoi)
        .rename("rainfall_7d")
        .addBands(chirps.filterDate(start_30d, end_30d).sum().clip(aoi).rename("rainfall_30d"))
    )
    raw = _download_image(
        rain,
        aoi,
        scale,
        band_names=["rainfall_7d", "rainfall_30d"],
    )
    return raw.to_dataset(dim="band").rename({"x": "lon", "y": "lat"})

fetch_landcover

fetch_landcover(aoi: Geometry, scale: int = 100) -> xr.Dataset

Download ESA WorldCover 2021 land cover. Returns Dataset with variable 'Map' (raw class values 10–95).

Source code in flood/features.py
def fetch_landcover(aoi: ee.Geometry, scale: int = 100) -> xr.Dataset:
    """
    Download ESA WorldCover 2021 land cover.
    Returns Dataset with variable 'Map' (raw class values 10–95).
    """
    worldcover = ee.ImageCollection("ESA/WorldCover/v200").first().select("Map").clip(aoi)
    da = _download_image(worldcover, aoi, scale).squeeze()
    return xr.Dataset({"Map": da.rename({"x": "lon", "y": "lat"})})

fetch_dist_river

fetch_dist_river(aoi: Geometry, scale: int = 90) -> xr.Dataset

Compute Euclidean distance to permanent water (JRC occurrence ≥ 70 %). Returns Dataset with variable 'dist_river' in metres.

Source code in flood/features.py
def fetch_dist_river(aoi: ee.Geometry, scale: int = 90) -> xr.Dataset:
    """
    Compute Euclidean distance to permanent water (JRC occurrence ≥ 70 %).
    Returns Dataset with variable 'dist_river' in metres.
    """
    water_mask = ee.Image("JRC/GSW1_4/GlobalSurfaceWater").select("occurrence").gte(70).clip(aoi)
    dist_river = (
        water_mask.fastDistanceTransform(256, "pixels")
        .sqrt()
        .multiply(ee.Image.pixelArea().sqrt())
        .rename("dist_river")
        .clip(aoi)
    )
    da = _download_image(dist_river, aoi, scale).squeeze()
    return xr.Dataset({"dist_river": da.rename({"x": "lon", "y": "lat"})})

fetch_mndwi

fetch_mndwi(aoi: Geometry, start_date: str, end_date: str, scale: int = 90) -> xr.Dataset

Compute Sentinel-2 MNDWI = (Green − SWIR1) / (Green + SWIR1). Returns Dataset with variable 'mndwi'.

Source code in flood/features.py
def fetch_mndwi(
    aoi: ee.Geometry,
    start_date: str,
    end_date: str,
    scale: int = 90,
) -> xr.Dataset:
    """
    Compute Sentinel-2 MNDWI = (Green − SWIR1) / (Green + SWIR1).
    Returns Dataset with variable 'mndwi'.
    """
    s2 = (
        ee.ImageCollection("COPERNICUS/S2_SR_HARMONIZED")
        .filterBounds(aoi)
        .filterDate(start_date, end_date)
        .filter(ee.Filter.lt("CLOUDY_PIXEL_PERCENTAGE", 30))
        .select(["B3", "B11"])
        .median()
        .clip(aoi)
    )
    mndwi = s2.normalizedDifference(["B3", "B11"]).rename("mndwi")
    da = _download_image(mndwi, aoi, scale).squeeze()
    return xr.Dataset({"mndwi": da.rename({"x": "lon", "y": "lat"})})

build_feature_datasets

build_feature_datasets(aoi: Geometry, config: dict) -> dict[str, xr.Dataset]

Download all feature bands from GEE in parallel. The seven fetch calls are independent HTTP requests; they run concurrently via DaskEngine.run_io_parallel (ThreadPoolExecutor) sharing the GEE session. The dict is keyed by band group name and consumed by both sample_training_data and cog_export.export_flood_cog.

Source code in flood/features.py
def build_feature_datasets(aoi: ee.Geometry, config: dict) -> dict[str, xr.Dataset]:
    """
    Download all feature bands from GEE in parallel.
    The seven fetch calls are independent HTTP requests; they run concurrently
    via DaskEngine.run_io_parallel (ThreadPoolExecutor) sharing the GEE session.
    The dict is keyed by band group name and consumed by both sample_training_data
    and cog_export.export_flood_cog.
    """
    from climate_change.core.dask_engine import DaskEngine
    from climate_change.core.population import fetch_population_count_safe

    scale = config.get("scale", 90)
    pre_start, pre_end = _pre_flood_window(config)
    post_start, post_end = _post_flood_window(config)
    rain_7d_start, rain_7d_end = _rain_7d_window(config)
    rain_30d_start, rain_30d_end = _rain_30d_window(config)
    mndwi_start, mndwi_end = _mndwi_window(config)

    return DaskEngine.run_io_parallel(
        {
            "terrain": lambda: fetch_terrain(aoi, scale=scale),
            "twi": lambda: fetch_twi(aoi, scale=500),
            "sar": lambda: fetch_sar_change(
                aoi,
                pre_start,
                pre_end,
                post_start,
                post_end,
                scale=scale,
            ),
            "rainfall": lambda: fetch_rainfall(
                aoi,
                rain_7d_start,
                rain_7d_end,
                rain_30d_start,
                rain_30d_end,
                scale=500,
            ),
            "landcover": lambda: fetch_landcover(aoi, scale=100),
            "dist_river": lambda: fetch_dist_river(aoi, scale=scale),
            "mndwi": lambda: fetch_mndwi(
                aoi,
                mndwi_start,
                mndwi_end,
                scale=scale,
            ),
            # Raw population COUNT for exposure reporting (population within
            # medium/high/very-high risk zones). Best-effort — see
            # core.population.fetch_population_count_safe.
            "population_count": lambda: fetch_population_count_safe(aoi, scale=scale),
        }
    )

build_gee_feature_stack

build_gee_feature_stack(aoi: Geometry, config: dict) -> ee.Image

Assemble the 10-band GEE image used for stratified sampling. Static (event-independent) and dynamic (event-specific) bands are built concurrently in background threads, then selected into FEATURE_COLS order.

Source code in flood/features.py
def build_gee_feature_stack(aoi: ee.Geometry, config: dict) -> ee.Image:
    """
    Assemble the 10-band GEE image used for stratified sampling.
    Static (event-independent) and dynamic (event-specific) bands are built
    concurrently in background threads, then selected into FEATURE_COLS order.
    """
    from concurrent.futures import ThreadPoolExecutor

    with ThreadPoolExecutor(max_workers=2) as pool:
        f_static = pool.submit(_build_static_image, aoi)
        f_dynamic = pool.submit(_build_dynamic_image, aoi, config)
        static_img = f_static.result()
        dynamic_img = f_dynamic.result()

    stack = ee.Image.cat([static_img, dynamic_img]).select(FEATURE_COLS)
    # Exclude permanent water bodies (ESA WorldCover) — flood risk applies to
    # land that can flood, not water that already is water year-round.
    valid_land = exclusion_mask(aoi, [WATER_CLASS])
    return stack.updateMask(valid_land)

sample_training_data

sample_training_data(aoi: Geometry, config: dict, n_pixels: int = 5000, scale: int = 90, seed: int = 42) -> pd.DataFrame

Build the 10-band feature stack and derive JRC flood labels concurrently in background threads, then sample n_pixels flooded + n_pixels non-flooded pixels.

Static features (elevation, TWI, dist_river, landcover, coords), dynamic features (vv_change, rainfall, mndwi), and the JRC label image are all fetched in parallel — no external feature_stack argument required.

1 = seasonal flood water not part of the permanent baseline

0 = all other land

Returns a DataFrame with columns = FEATURE_COLS + ['is_flooded'].

Source code in flood/features.py
def sample_training_data(
    aoi: ee.Geometry,
    config: dict,
    n_pixels: int = 5000,
    scale: int = 90,
    seed: int = 42,
) -> pd.DataFrame:
    """
    Build the 10-band feature stack and derive JRC flood labels concurrently in
    background threads, then sample n_pixels flooded + n_pixels non-flooded pixels.

    Static features (elevation, TWI, dist_river, landcover, coords), dynamic
    features (vv_change, rainfall, mndwi), and the JRC label image are all
    fetched in parallel — no external feature_stack argument required.

    Labels: 1 = seasonal flood water not part of the permanent baseline
            0 = all other land
    Returns a DataFrame with columns = FEATURE_COLS + ['is_flooded'].
    """
    from concurrent.futures import ThreadPoolExecutor

    flood_start, flood_end = _label_window(config)

    with ThreadPoolExecutor(max_workers=3) as pool:
        f_static = pool.submit(_build_static_image, aoi)
        f_dynamic = pool.submit(_build_dynamic_image, aoi, config)
        f_label = pool.submit(_build_jrc_label, aoi, flood_start, flood_end)
        static_img = f_static.result()
        dynamic_img = f_dynamic.result()
        flood_label = f_label.result()

    feature_stack = ee.Image.cat([static_img, dynamic_img]).select(FEATURE_COLS)
    # Exclude permanent water bodies (ESA WorldCover) — flood risk applies to
    # land that can flood, not water that already is water year-round.
    valid_land = exclusion_mask(aoi, [WATER_CLASS])
    feature_stack = feature_stack.updateMask(valid_land)
    labeled = feature_stack.addBands(flood_label)
    records = _sample_labeled_pixels(
        labeled=labeled,
        flood_label=flood_label,
        aoi=aoi,
        scale=scale,
        n_pixels=n_pixels,
        seed=seed,
    )
    df = pd.DataFrame([f["properties"] for f in records]).dropna()
    df["is_flooded"] = df["is_flooded"].astype(int)
    return df

align_datasets

align_datasets(datasets: dict[str, Dataset], ref_key: str = 'terrain', method_continuous: InterpOptions = 'linear', method_categorical: InterpOptions = 'nearest') -> dict[str, xr.Dataset]

Interpolate all datasets onto the reference grid (terrain at 90 m by default). Landcover is treated as categorical and uses nearest-neighbour interpolation. population_count is passed through entirely unchanged, at its own native ~100 m resolution — it is never interpolated onto this reference grid. Summing population within a risk zone requires the opposite: upsampling the (categorical) classification onto population's native grid — see core.population.population_exposure, which cog_export.py uses for this instead of relying on align_datasets. Interpolating population here would silently corrupt totals (verified against live GEE data — see core.population.fetch_population_count's docstring).

Each dataset is chunked before interpolation so that xarray produces Dask-backed lazy arrays; dask.compute() then materialises them all in a single parallel scheduler pass.

Source code in flood/features.py
def align_datasets(
    datasets: dict[str, xr.Dataset],
    ref_key: str = "terrain",
    method_continuous: InterpOptions = "linear",
    method_categorical: InterpOptions = "nearest",
) -> dict[str, xr.Dataset]:
    """
    Interpolate all datasets onto the reference grid (terrain at 90 m by default).
    Landcover is treated as categorical and uses nearest-neighbour interpolation.
    population_count is passed through entirely unchanged, at its own native
    ~100 m resolution — it is never interpolated onto this reference grid.
    Summing population within a risk zone requires the opposite: upsampling
    the (categorical) classification onto population's native grid — see
    core.population.population_exposure, which cog_export.py uses for this
    instead of relying on align_datasets. Interpolating population here would
    silently corrupt totals (verified against live GEE data — see
    core.population.fetch_population_count's docstring).

    Each dataset is chunked before interpolation so that xarray produces
    Dask-backed lazy arrays; dask.compute() then materialises them all
    in a single parallel scheduler pass.
    """
    from dask.base import compute

    _CHUNK = {"lat": 256, "lon": 256}
    ref = datasets[ref_key]

    lazy: dict[str, xr.Dataset] = {ref_key: ref}
    for key, ds in datasets.items():
        if key == ref_key or ds is None or key == "population_count":
            continue
        method = method_categorical if key == "landcover" else method_continuous
        lazy[key] = ds.chunk(_CHUNK).interp(lat=ref.lat, lon=ref.lon, method=method)

    keys = list(lazy)
    computed = compute(*[lazy[k] for k in keys])
    result = dict(zip(keys, computed))
    pop_ds = datasets.get("population_count")
    if pop_ds is not None:
        result["population_count"] = pop_ds
    return result

Model

model

FloodModel

FloodModel()

Orchestrates the full ML pipeline for a single flood event. Always trains both RF and XGBoost. config['model_type'] selects which probabilities drive the primary stats, risk map, and COG export: "rf" — Random Forest "xgboost" — XGBoost "ensemble" — mean of RF + XGBoost (default) Trained models are stored on self.rf / self.xgb so the use case can pass them directly to cog_export.export_flood_cog.

Source code in flood/model.py
def __init__(self) -> None:
    self.rf: RandomForestClassifier | None = None
    self.xgb: Booster | None = None

predict

predict(df: DataFrame, config: dict | None = None) -> dict
Parameters

df : DataFrame with columns = FEATURE_COLS + ['is_flooded'] config : optional dict; reads config['model_type'] (default 'ensemble')

Returns

dict with keys: stats, charts

Source code in flood/model.py
def predict(self, df: pd.DataFrame, config: dict | None = None) -> dict:
    """
    Parameters
    ----------
    df     : DataFrame with columns = FEATURE_COLS + ['is_flooded']
    config : optional dict; reads config['model_type'] (default 'ensemble')

    Returns
    -------
    dict with keys: stats, charts
    """
    model_type = (config or {}).get("model_type", "ensemble")
    if model_type not in VALID_MODEL_TYPES:
        raise ValueError(f"model_type must be one of {VALID_MODEL_TYPES}, got '{model_type}'")

    X = df[FEATURE_COLS].to_numpy(dtype=np.float64)
    y = df["is_flooded"].to_numpy(dtype=np.intp)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=_safe_stratify(y)
    )

    # Train RF and XGBoost concurrently when the Dask cluster is running,
    # otherwise fall back to sequential execution (e.g. during unit tests).
    from climate_change.core.dask_engine import DaskEngine

    client = DaskEngine.get_client_if_running()
    if client is not None:
        f_rf = client.submit(train_rf, X_train, y_train, pure=False)
        f_xgb = client.submit(train_xgb, X_train, y_train, X_test, y_test, pure=False)
        (self.rf, rf_meta), (self.xgb, xgb_meta) = cast(list, client.gather([f_rf, f_xgb]))
    else:
        self.rf, rf_meta = train_rf(X_train, y_train)
        self.xgb, xgb_meta = train_xgb(X_train, y_train, X_test, y_test)

    assert self.rf is not None
    assert self.xgb is not None

    eval_result = evaluate_models(self.rf, self.xgb, X_test, y_test)
    shap_payload = compute_shap_importance(self.xgb, X_test)
    uncertainty = compute_uncertainty(
        np.array(eval_result["rf"]["probabilities"]),
        np.array(eval_result["xgb"]["probabilities"]),
    )
    charts = build_flood_charts(eval_result, shap_payload, uncertainty, model_type)
    # Risk percentages from the selected model
    result_key = _MODEL_TYPE_KEY.get(model_type, "ensemble")
    primary_probs = np.array(eval_result[result_key]["probabilities"])
    risk_labels = classify_flood_risk(primary_probs)

    _ORDER = ["Very High", "High", "Medium", "Low"]
    pct_arr = (
        pd.Series(risk_labels)
        .value_counts(normalize=True)
        .mul(100)
        .round(1)
        .reindex(_ORDER)
        .fillna(0.0)
        .to_numpy(dtype=np.float64)
    )
    very_high_pct, high_pct, medium_pct, low_pct = (float(v) for v in pct_arr)

    stats = {
        "model_type": model_type,
        "n_pixels_sampled": int(len(df)),
        "flooded_pct": round(float(y.mean() * 100), 1),
        "rf_cv_f1": round(rf_meta["cv_f1_mean"], 4)
        if rf_meta["cv_f1_mean"] is not None
        else None,
        "rf_f1": eval_result["rf"]["f1"],
        "rf_auc": eval_result["rf"]["auc"],
        "xgb_cv_f1": round(xgb_meta["cv_f1_mean"], 4)
        if xgb_meta["cv_f1_mean"] is not None
        else None,
        "xgb_f1": eval_result["xgb"]["f1"],
        "xgb_auc": eval_result["xgb"]["auc"],
        "ensemble_f1": eval_result["ensemble"]["f1"],
        "ensemble_auc": eval_result["ensemble"]["auc"],
        "selected_f1": eval_result[result_key]["f1"],
        "selected_auc": eval_result[result_key]["auc"],
        "selected_threshold": eval_result[result_key]["threshold"],
        "top_flood_driver": shap_payload["features"][0],
        "very_high_risk_pct": very_high_pct,
        "high_risk_pct": high_pct,
        "medium_risk_pct": medium_pct,
        "low_risk_pct": low_pct,
        **uncertainty,
    }

    lon_idx = FEATURE_COLS.index("longitude")
    lat_idx = FEATURE_COLS.index("latitude")
    _sample_points = [
        {
            "lon": round(float(X_test[i, lon_idx]), 5),
            "lat": round(float(X_test[i, lat_idx]), 5),
            "risk_class": str(risk_labels[i]),
        }
        for i in range(len(X_test))
    ]
    return {"stats": stats, "charts": charts, "_sample_points": _sample_points}

classify_flood_risk

classify_flood_risk(prob: ndarray) -> np.ndarray

Map flood probability array to 4-class string labels.

Source code in flood/model.py
def classify_flood_risk(prob: np.ndarray) -> np.ndarray:
    """Map flood probability array to 4-class string labels."""
    risk = np.full(prob.shape, "Low", dtype=object)
    risk[(prob >= 0.25) & (prob < 0.50)] = "Medium"
    risk[(prob >= 0.50) & (prob < 0.75)] = "High"
    risk[prob >= 0.75] = "Very High"
    return risk

train_rf

train_rf(X_train: ndarray, y_train: ndarray, cv_folds: int = 5) -> tuple[RandomForestClassifier, dict]

Fit a balanced Random Forest and report CV F1. Returns (model, metadata).

Source code in flood/model.py
def train_rf(
    X_train: np.ndarray,
    y_train: np.ndarray,
    cv_folds: int = 5,
) -> tuple[RandomForestClassifier, dict]:
    """Fit a balanced Random Forest and report CV F1. Returns (model, metadata)."""
    rf = RandomForestClassifier(
        n_estimators=RF_N_ESTIMATORS,
        max_depth=RF_MAX_DEPTH,
        min_samples_leaf=RF_MIN_SAMPLES_LEAF,
        class_weight="balanced",
        random_state=42,
        n_jobs=-1,
    )
    rf.fit(X_train, y_train)
    folds = _safe_cv_folds(y_train, cv_folds)
    if folds is None:
        return rf, {"cv_f1_mean": None, "cv_f1_std": None}
    cv_f1 = cross_val_score(rf, X_train, y_train, cv=folds, scoring="f1")
    return rf, {"cv_f1_mean": float(cv_f1.mean()), "cv_f1_std": float(cv_f1.std())}

train_xgb

train_xgb(X_train: ndarray, y_train: ndarray, X_val: ndarray, y_val: ndarray, cv_folds: int = 5) -> tuple[Booster, dict]

Fit XGBoost binary flood classifier via the low-level Booster API with scale_pos_weight and eval-set logging. Returns (model, metadata).

Uses xgboost.train() rather than the XGBClassifier sklearn wrapper because the wrapper's fit() rejects a y whose sorted unique values aren't exactly [0, 1] — an AOI/time window where every sampled pixel is flooded (or none are) is a legitimate data state, not invalid input, and would otherwise crash training.

Source code in flood/model.py
def train_xgb(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_val: np.ndarray,
    y_val: np.ndarray,
    cv_folds: int = 5,
) -> tuple[Booster, dict]:
    """
    Fit XGBoost binary flood classifier via the low-level Booster API with
    scale_pos_weight and eval-set logging. Returns (model, metadata).

    Uses xgboost.train() rather than the XGBClassifier sklearn wrapper because
    the wrapper's fit() rejects a y whose sorted unique values aren't exactly
    [0, 1] — an AOI/time window where every sampled pixel is flooded (or none
    are) is a legitimate data state, not invalid input, and would otherwise
    crash training.
    """
    scale_pos_weight = float((y_train == 0).sum() / max((y_train == 1).sum(), 1))
    params = {**XGB_PARAMS, "scale_pos_weight": scale_pos_weight}
    dtrain = DMatrix(X_train, label=y_train)
    dval = DMatrix(X_val, label=y_val)
    booster = xgb_train(
        params,
        dtrain,
        num_boost_round=XGB_N_ESTIMATORS,
        evals=[(dval, "validation")],
        verbose_eval=False,
    )

    folds = _safe_cv_folds(y_train, cv_folds)
    if folds is None:
        return booster, {"cv_f1_mean": None, "cv_f1_std": None}
    cv = StratifiedKFold(n_splits=folds)
    cv_f1_scores = []
    for train_idx, val_idx in cv.split(X_train, y_train):
        fold_scale_pos_weight = float(
            (y_train[train_idx] == 0).sum() / max((y_train[train_idx] == 1).sum(), 1)
        )
        fold_booster = xgb_train(
            {**XGB_PARAMS, "scale_pos_weight": fold_scale_pos_weight},
            DMatrix(X_train[train_idx], label=y_train[train_idx]),
            num_boost_round=XGB_N_ESTIMATORS,
        )
        fold_pred = (fold_booster.predict(DMatrix(X_train[val_idx])) >= 0.5).astype(int)
        cv_f1_scores.append(f1_score(y_train[val_idx], fold_pred))

    cv_f1 = np.array(cv_f1_scores)
    return booster, {"cv_f1_mean": float(cv_f1.mean()), "cv_f1_std": float(cv_f1.std())}

positive_class_proba

positive_class_proba(clf: RandomForestClassifier, proba: ndarray) -> np.ndarray

Return P(y=1) from a binary sklearn classifier's predict_proba output.

predict_proba only has one column when the classifier saw a single class during training (e.g. every sampled pixel in this AOI/time window was flooded, or none were) — proba[:, 1] would then raise IndexError. Falls back to reading clf.classes_ to know whether that lone column represents class 0 or class 1.

Source code in flood/model.py
def positive_class_proba(clf: RandomForestClassifier, proba: np.ndarray) -> np.ndarray:
    """
    Return P(y=1) from a binary sklearn classifier's predict_proba output.

    predict_proba only has one column when the classifier saw a single class
    during training (e.g. every sampled pixel in this AOI/time window was
    flooded, or none were) — proba[:, 1] would then raise IndexError. Falls
    back to reading clf.classes_ to know whether that lone column represents
    class 0 or class 1.
    """
    if proba.shape[1] == 2:
        return proba[:, 1]
    return proba[:, 0] if clf.classes_[0] == 1 else np.zeros(proba.shape[0])

find_best_threshold

find_best_threshold(probs: ndarray, y_true: ndarray) -> tuple[float, float]

Return (threshold, best_f1) that maximises F1 on the precision-recall curve.

Source code in flood/model.py
def find_best_threshold(probs: np.ndarray, y_true: np.ndarray) -> tuple[float, float]:
    """Return (threshold, best_f1) that maximises F1 on the precision-recall curve."""
    precisions, recalls, thresholds = precision_recall_curve(y_true, probs)
    f1s = 2 * precisions * recalls / (precisions + recalls + 1e-9)
    best_idx = int(f1s.argmax())
    return float(thresholds[best_idx]), float(f1s[best_idx])

evaluate_models

evaluate_models(rf: RandomForestClassifier, xgb: Booster, X_test: ndarray, y_test: ndarray) -> dict

Evaluate RF, XGBoost, and their ensemble on the held-out test set. Thresholds are maximised per-model via the precision-recall curve.

Source code in flood/model.py
def evaluate_models(
    rf: RandomForestClassifier,
    xgb: Booster,
    X_test: np.ndarray,
    y_test: np.ndarray,
) -> dict:
    """
    Evaluate RF, XGBoost, and their ensemble on the held-out test set.
    Thresholds are maximised per-model via the precision-recall curve.
    """
    rf_prob = positive_class_proba(rf, rf.predict_proba(X_test))
    xgb_prob = xgb.predict(DMatrix(X_test))
    ens_prob = (rf_prob + xgb_prob) / 2.0

    def _metrics(prob: np.ndarray, label: str) -> dict:
        thresh, _ = find_best_threshold(prob, y_test)
        pred = (prob >= thresh).astype(int)
        # roc_auc_score is undefined (returns NaN, which isn't valid JSON) when
        # y_test ends up with only one class — a real possibility for a scarce
        # class routed through the non-stratified split fallback in predict().
        has_both_classes = len(np.unique(y_test)) > 1
        auc = round(float(roc_auc_score(y_test, prob)), 4) if has_both_classes else None
        return {
            "label": label,
            "f1": round(float(f1_score(y_test, pred)), 4),
            "auc": auc,
            "threshold": round(thresh, 4),
            "predictions": pred.tolist(),
            "probabilities": prob.round(4).tolist(),
        }

    return {
        "rf": _metrics(rf_prob, "Random Forest"),
        "xgb": _metrics(xgb_prob, "XGBoost"),
        "ensemble": _metrics(ens_prob, "Ensemble (mean prob)"),
        "actuals": y_test.tolist(),
    }

compute_shap_importance

compute_shap_importance(xgb: Booster, X_test: ndarray) -> dict

TreeExplainer SHAP values for XGBoost, sorted by descending mean |SHAP|. XGBoost is always used for SHAP regardless of model_type — it provides the most interpretable tree-based explanations.

Source code in flood/model.py
def compute_shap_importance(xgb: Booster, X_test: np.ndarray) -> dict:
    """
    TreeExplainer SHAP values for XGBoost, sorted by descending mean |SHAP|.
    XGBoost is always used for SHAP regardless of model_type — it provides
    the most interpretable tree-based explanations.
    """
    import shap

    explainer = shap.TreeExplainer(xgb)
    shap_vals = explainer.shap_values(X_test)
    mean_abs = np.abs(shap_vals).mean(axis=0)
    rank_idx = np.argsort(mean_abs)[::-1]
    return {
        "features": [FEATURE_COLS[i] for i in rank_idx],
        "mean_abs_shap": mean_abs[rank_idx].round(4).tolist(),
    }

compute_uncertainty

compute_uncertainty(rf_prob: ndarray, xgb_prob: ndarray) -> dict

Epistemic uncertainty from RF–XGBoost probability spread. Pixels with spread > 0.20 are flagged for field validation. Always computed regardless of model_type so the UI can display it.

Source code in flood/model.py
def compute_uncertainty(rf_prob: np.ndarray, xgb_prob: np.ndarray) -> dict:
    """
    Epistemic uncertainty from RF–XGBoost probability spread.
    Pixels with spread > 0.20 are flagged for field validation.
    Always computed regardless of model_type so the UI can display it.
    """
    spread = np.abs(rf_prob - xgb_prob)
    return {
        "mean_spread": round(float(spread.mean()), 4),
        "high_uncertainty_pct": round(float((spread > 0.20).mean() * 100), 1),
        "spread_stats": {
            "min": round(float(spread.min()), 4),
            "p25": round(float(np.percentile(spread, 25)), 4),
            "p75": round(float(np.percentile(spread, 75)), 4),
            "max": round(float(spread.max()), 4),
            "mean": round(float(spread.mean()), 4),
        },
    }

build_flood_charts

build_flood_charts(eval_result: dict, shap_payload: dict, uncertainty_payload: dict, model_type: str = 'ensemble') -> dict

Assemble frontend-ready chart payloads.

Risk distribution is derived from the selected model_type's probabilities. model_performance always includes all three models for comparison.

Source code in flood/model.py
def build_flood_charts(
    eval_result: dict,
    shap_payload: dict,
    uncertainty_payload: dict,
    model_type: str = "ensemble",
) -> dict:
    """
    Assemble frontend-ready chart payloads.

    Risk distribution is derived from the selected model_type's probabilities.
    model_performance always includes all three models for comparison.
    """
    result_key = _MODEL_TYPE_KEY.get(model_type, "ensemble")
    selected_probs = np.array(eval_result[result_key]["probabilities"])

    _RISK_ORDER = ["Very High", "High", "Medium", "Low"]

    risk_labels = classify_flood_risk(selected_probs)
    counts = (
        pd.Series(risk_labels)
        .value_counts()
        .reindex(_RISK_ORDER)
        .fillna(0.0)
        .to_numpy(dtype=np.float64)
    )
    risk_pct = (counts / counts.sum() * 100).round(1)

    risk_chart = {
        "labels": _RISK_ORDER,
        "data": risk_pct.tolist(),
        "colors": [RISK_COLORS[c] for c in _RISK_ORDER],
    }

    # All three models always shown for comparison
    model_performance = {
        "rf": {"f1": eval_result["rf"]["f1"], "auc": eval_result["rf"]["auc"]},
        "xgb": {"f1": eval_result["xgb"]["f1"], "auc": eval_result["xgb"]["auc"]},
        "ensemble": {
            "f1": eval_result["ensemble"]["f1"],
            "auc": eval_result["ensemble"]["auc"],
        },
        "selected": model_type,
    }

    return {
        "risk_distribution": risk_chart,
        "shap": shap_payload,
        "uncertainty": uncertainty_payload,
        "model_performance": model_performance,
    }

COG export

cog_export

flood_raster_distribution

flood_raster_distribution(path: str | Path) -> dict

Read an already-exported flood-risk COG and tally pixel counts/percentages per risk class (Very High/High/Medium/Low), skipping nodata (<= 0) pixels.

Source code in flood/cog_export.py
def flood_raster_distribution(path: str | Path) -> dict:
    """
    Read an already-exported flood-risk COG and tally pixel counts/percentages
    per risk class (Very High/High/Medium/Low), skipping nodata (<= 0) pixels.
    """
    with rasterio.open(path) as src:
        data = src.read(1)

    valid_values = data[data > 0]
    counts = {
        label: int((valid_values == RISK_VALUE_BY_LABEL[label]).sum()) for label in RISK_ORDER
    }
    total = int(valid_values.size)
    percentages = {
        label: round((count / total * 100), 1) if total else 0.0 for label, count in counts.items()
    }
    return {
        "labels": list(RISK_ORDER),
        "counts": [counts[label] for label in RISK_ORDER],
        "percentages": [percentages[label] for label in RISK_ORDER],
        "valid_pixel_count": total,
    }

export_flood_cog

export_flood_cog(rf_model: RandomForestClassifier, xgb_model: Booster, datasets: dict[str, Dataset], output_dir: str, prefix: str, model_type: str = 'ensemble', aoi_geojson: dict | None = None) -> FloodCogResult

Build the full-AOI feature matrix from in-memory xarray Datasets, run inference with the selected model, classify into 4 risk classes, and write a Cloud Optimised GeoTIFF.

Source code in flood/cog_export.py
def export_flood_cog(
    rf_model: RandomForestClassifier,
    xgb_model: Booster,
    datasets: dict[str, xr.Dataset],
    output_dir: str,
    prefix: str,
    model_type: str = "ensemble",
    aoi_geojson: dict | None = None,
) -> FloodCogResult:
    """
    Build the full-AOI feature matrix from in-memory xarray Datasets,
    run inference with the selected model, classify into 4 risk classes,
    and write a Cloud Optimised GeoTIFF.
    """
    if model_type not in VALID_MODEL_TYPES:
        raise ValueError(f"model_type must be one of {VALID_MODEL_TYPES}, got '{model_type}'")

    output_dir_path = Path(output_dir)
    output_dir_path.mkdir(parents=True, exist_ok=True)
    cog_path = output_dir_path / f"{prefix}_flood_risk.tif"

    # Align all bands to the terrain (90 m) reference grid
    aligned = align_datasets(datasets, ref_key="terrain")
    terrain = aligned["terrain"]
    ref_lat = terrain.lat
    ref_lon = terrain.lon
    n_lat, n_lon = len(ref_lat), len(ref_lon)

    # Extract feature arrays (all float32)
    elevation = terrain["elevation"].values.astype(np.float32)
    twi_ = aligned["twi"]["twi"].values.astype(np.float32)
    dist_river_ = aligned["dist_river"]["dist_river"].values.astype(np.float32)
    vv_change_ = aligned["sar"]["vv_change"].values.astype(np.float32)
    rain7_ = aligned["rainfall"]["rainfall_7d"].values.astype(np.float32)
    rain30_ = aligned["rainfall"]["rainfall_30d"].values.astype(np.float32)
    mndwi_ = aligned["mndwi"]["mndwi"].values.astype(np.float32)
    landcover_ = aligned["landcover"]["Map"].values.astype(np.float32)
    landcover_ = (landcover_ / 10.0) - 1.0  # normalise: match §8

    lon_grid = np.broadcast_to(ref_lon.values[np.newaxis, :], (n_lat, n_lon)).astype(np.float32)
    lat_grid = np.broadcast_to(ref_lat.values[:, np.newaxis], (n_lat, n_lon)).astype(np.float32)

    # Stack → (n_pixels, n_features) and mask nodata
    bands = np.stack(
        [
            elevation,
            twi_,
            dist_river_,
            vv_change_,
            rain7_,
            rain30_,
            mndwi_,
            landcover_,
            lon_grid,
            lat_grid,
        ],
        axis=0,
    )  # (10, n_lat, n_lon)
    X_full = bands.reshape(len(FEATURE_COLS), -1).T  # (n_pixels, 10)
    valid_mask = ~np.isnan(X_full).any(axis=1)
    # Permanent water bodies (ESA WorldCover) are excluded from flood-risk output.
    water = np.isclose(landcover_.reshape(-1), _WATER_LANDCOVER)
    valid_mask &= ~water

    transform = terrain["elevation"].rio.transform()
    crs = terrain["elevation"].rio.crs or "EPSG:4326"
    geometries = _aoi_geometries(aoi_geojson)
    if geometries:
        inside_aoi = geometry_mask(
            geometries,
            out_shape=(n_lat, n_lon),
            transform=transform,
            invert=True,
        ).reshape(-1)
        valid_mask = valid_mask & inside_aoi

    X_valid = X_full[valid_mask]

    # Inference — route by model_type
    risk_grid: np.ndarray = np.zeros(n_lat * n_lon, dtype=np.uint8)
    if X_valid.size:
        prob_rf = positive_class_proba(rf_model, rf_model.predict_proba(X_valid))
        prob_xgb = xgb_model.predict(DMatrix(X_valid))
        if model_type == "rf":
            selected_prob = prob_rf
        elif model_type == "xgboost":
            selected_prob = prob_xgb
        else:  # ensemble
            selected_prob = (prob_rf + prob_xgb) / 2.0

        # Classify and reconstruct spatial grid
        risk_labels = classify_flood_risk(selected_prob)
        risk_valid = np.array([RISK_INT[r] for r in risk_labels], dtype=np.uint8)
        risk_grid[valid_mask] = risk_valid
    risk_grid = risk_grid.reshape(n_lat, n_lon)

    population_by_class: dict[str, float] = {}
    total_population: float | None = None
    pop_ds = aligned.get("population_count")
    if pop_ds is not None:
        by_int = population_exposure(
            risk_grid, transform, pop_ds, classes=list(RISK_LABEL_BY_VALUE)
        )
        population_by_class = {RISK_LABEL_BY_VALUE[v]: by_int[str(v)] for v in RISK_LABEL_BY_VALUE}
        total_population = round(sum(population_by_class.values()), 1)

    with rasterio.open(
        cog_path,
        "w",
        driver="COG",
        height=n_lat,
        width=n_lon,
        count=1,
        dtype=np.uint8,
        crs=crs,
        transform=transform,
        nodata=0,
        compress="lzw",
    ) as dst:
        dst.write(risk_grid, 1)
        dst.update_tags(
            RISK_CLASSES="0=nodata,1=Low,2=Medium,3=High,4=Very High",
            MODEL_TYPE=model_type,
        )

    uncompressed_mb = risk_grid.nbytes / 1e6
    _log.info(
        "COG exported: %s  (%d×%d px, %.1f MB uncompressed)",
        cog_path,
        n_lat,
        n_lon,
        uncompressed_mb,
    )
    return {
        "flood_risk": str(cog_path),
        "population_by_class": population_by_class,
        "total_population": total_population,
    }