Skip to content

disease

Climate suitability for disease outbreaks (e.g. malaria). Entry point: DiseaseRiskUseCase.

from climate_change import DiseaseRiskUseCase

Use case

use_case

DiseaseRiskUseCase

DiseaseRiskUseCase(dask_engine: DaskEngine)

Bases: BaseUseCase

Entry point for the climate-driven disease surveillance domain.

Minimal config (all optional — defaults match Kisumu County, Kenya 2021–2023):

{ "aoi_geojson": {"type": "Polygon", "coordinates": [...]}, "gee_project": "your-gee-project-id", "start_date": "2021-01-01", "end_date": "2023-12-31", "model_type": "gbm", # "gbm" | "xgboost" | "ensemble" "n_pixels": 3000, "scale": 1000, # metres — MODIS / CHIRPS native "output_dir": "outputs", "prefix": "disease", }

Source code in disease/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 disease/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 training pixels (see _preprocess_raw).

Source code in disease/use_case.py
def preprocess(self, raw_data: dict, config: AnalysisConfig) -> dict:
    """Build the feature stack and sample training 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 disease/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", "gbm")
    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: DiseaseCogResult | None = None
    try:
        raster_result = export_disease_cog(
            gbm_model=features["_gbm"],
            xgb_model=features["_xgb"],
            scaler=features["_scaler"],
            datasets=features["datasets"],
            output_dir=dict_config.get("output_dir", "outputs"),
            prefix=dict_config.get("prefix", "disease"),
            model_type=model_type,
            aoi_geojson=config.aoi_geojson,
        )
        raster_paths = {"disease_risk": raster_result["disease_risk"]}
        _apply_raster_risk_pct(result, raster_result["risk_pct"])
    except Exception as exc:
        raster_error = str(exc)

    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.get("charts", {}).get("riskDist"), total_area_ha)
    _attach_population_stats(result, raster_result)

    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 {},
    }
    if raster_error:
        metadata["raster_error"] = raster_error

    return AnalysisOutput(
        module="disease",
        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 disease/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)

    cog_paths: dict[str, str] | None = None
    cog_error: str | None = None
    raster_result: DiseaseCogResult | None = None
    try:
        raster_result = export_disease_cog(
            gbm_model=features["_gbm"],
            xgb_model=features["_xgb"],
            scaler=features["_scaler"],
            datasets=features["datasets"],
            output_dir=config.get("output_dir", "outputs"),
            prefix=config.get("prefix", "disease"),
            model_type=config.get("model_type", "gbm"),
            aoi_geojson=config.get("aoi_geojson"),
        )
        cog_paths = {"disease_risk": raster_result["disease_risk"]}
        _apply_raster_risk_pct(result, raster_result["risk_pct"])
    except Exception as exc:
        _log.warning("Disease COG export failed: %s", exc, exc_info=True)
        cog_error = str(exc)
    result["raster"] = cog_paths or {}
    if cog_error:
        result["raster_error"] = cog_error

    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.get("charts", {}).get("riskDist"), total_area_ha)
    _attach_population_stats(result, raster_result)

    result["stats"].update(
        {
            "bbox": features["bbox"],
            "prefix": config.get("prefix", "disease"),
        }
    )
    return result

run_date_ranges

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

Run the same AOI over multiple date-range configurations in parallel.

Source code in disease/use_case.py
def run_date_ranges(self, config: dict, date_ranges: list[dict]) -> list[dict]:
    """Run the same AOI over multiple date-range configurations in parallel."""
    merged = [{**config, **dr} for dr in date_ranges]
    return self._run_parallel(merged)

run_multi_regions

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

Run multiple AOI configs (same timeframe, different regions) in parallel.

Source code in disease/use_case.py
def run_multi_regions(self, configs: list[dict]) -> list[dict]:
    """Run multiple AOI configs (same timeframe, different regions) in parallel."""
    return self._run_parallel(configs)

Features

features

fetch_rainfall_4w

fetch_rainfall_4w(aoi: Geometry, end_date: str, scale: int = 1000) -> xr.Dataset

CHIRPS 28-day cumulative rainfall ending on end_date. Returns Dataset with variable 'rainfall_4w'.

Source code in disease/features.py
def fetch_rainfall_4w(
    aoi: ee.Geometry,
    end_date: str,
    scale: int = 1000,
) -> xr.Dataset:
    """
    CHIRPS 28-day cumulative rainfall ending on end_date.
    Returns Dataset with variable 'rainfall_4w'.
    """
    end_dt = ee.Date(end_date)
    start_4w = end_dt.advance(-28, "day")
    collection = _require_images(
        ee.ImageCollection("UCSB-CHG/CHIRPS/DAILY")
        .filterBounds(aoi)
        .filterDate(start_4w, end_dt)
        .select("precipitation"),
        "CHIRPS rainfall",
        cast(str, start_4w.format("YYYY-MM-dd").getInfo()),
        end_date,
    )
    chirps_4w = collection.sum().rename("rainfall_4w").clip(aoi)
    url = chirps_4w.getDownloadURL(
        {"region": aoi, "scale": scale, "crs": "EPSG:4326", "format": "GEO_TIFF"}
    )
    da = _download_band(url).squeeze()
    return xr.Dataset({"rainfall_4w": da.rename({"x": "lon", "y": "lat"})})

fetch_lst_mean

fetch_lst_mean(aoi: Geometry, start: str, end: str, scale: int = 1000) -> xr.Dataset

MODIS Terra MOD11A2 daytime LST mean over [start, end] in °C. Returns Dataset with variable 'temp_mean'.

Source code in disease/features.py
def fetch_lst_mean(
    aoi: ee.Geometry,
    start: str,
    end: str,
    scale: int = 1000,
) -> xr.Dataset:
    """
    MODIS Terra MOD11A2 daytime LST mean over [start, end] in °C.
    Returns Dataset with variable 'temp_mean'.
    """

    def _scale_lst(img: ee.Image) -> ee.Image:
        return img.multiply(0.02).subtract(273.15).copyProperties(img, ["system:time_start"])

    lst_col = _require_images(
        ee.ImageCollection("MODIS/061/MOD11A2")
        .filterBounds(aoi)
        .filterDate(start, end)
        .select("LST_Day_1km")
        .map(_scale_lst),
        "MODIS land-surface-temperature",
        start,
        end,
    )
    temp_mean = lst_col.mean().rename("temp_mean").clip(aoi)
    url = temp_mean.getDownloadURL(
        {"region": aoi, "scale": scale, "crs": "EPSG:4326", "format": "GEO_TIFF"}
    )
    da = _download_band(url).squeeze()
    return xr.Dataset({"temp_mean": da.rename({"x": "lon", "y": "lat"})})

fetch_ndwi

fetch_ndwi(aoi: Geometry, start: str, end: str, scale: int = 1000) -> xr.Dataset

Sentinel-2 MNDWI = (Green − SWIR1) / (Green + SWIR1) median composite. Returns Dataset with variable 'ndwi'.

Source code in disease/features.py
def fetch_ndwi(
    aoi: ee.Geometry,
    start: str,
    end: str,
    scale: int = 1000,
) -> xr.Dataset:
    """
    Sentinel-2 MNDWI = (Green − SWIR1) / (Green + SWIR1) median composite.
    Returns Dataset with variable 'ndwi'.
    """

    s2 = _sentinel2_mndwi_collection(aoi, start, end)
    ndwi_img = s2.median().clip(aoi)
    url = ndwi_img.getDownloadURL(
        {"region": aoi, "scale": scale, "crs": "EPSG:4326", "format": "GEO_TIFF"}
    )
    da = _download_band(url).squeeze()
    return xr.Dataset({"ndwi": da.rename({"x": "lon", "y": "lat"})})

fetch_elevation

fetch_elevation(aoi: Geometry, scale: int = 1000) -> xr.Dataset

USGS SRTM 30 m elevation. Returns Dataset with variable 'elevation'.

Source code in disease/features.py
def fetch_elevation(aoi: ee.Geometry, scale: int = 1000) -> xr.Dataset:
    """USGS SRTM 30 m elevation. Returns Dataset with variable 'elevation'."""
    elev = ee.Image("USGS/SRTMGL1_003").select("elevation").clip(aoi)
    url = elev.getDownloadURL(
        {"region": aoi, "scale": scale, "crs": "EPSG:4326", "format": "GEO_TIFF"}
    )
    da = _download_band(url).squeeze()
    return xr.Dataset({"elevation": da.rename({"x": "lon", "y": "lat"})})

fetch_pop_density

fetch_pop_density(aoi: Geometry, year: int = 2020, scale: int = 1000) -> xr.Dataset

WorldPop GP 100 m population density, log-transformed: log(1 + pop). Returns Dataset with variable 'pop_density'.

Source code in disease/features.py
def fetch_pop_density(aoi: ee.Geometry, year: int = 2020, scale: int = 1000) -> xr.Dataset:
    """
    WorldPop GP 100 m population density, log-transformed: log(1 + pop).
    Returns Dataset with variable 'pop_density'.
    """
    pop_collection = (
        ee.ImageCollection("WorldPop/GP/100m/pop")
        .filterBounds(aoi)
        .filter(ee.Filter.eq("year", year))
    )
    pop_raw = _require_population_image(pop_collection, year).select("population").clip(aoi)
    pop_log = pop_raw.add(1).log().rename("pop_density")
    url = pop_log.getDownloadURL(
        {"region": aoi, "scale": scale, "crs": "EPSG:4326", "format": "GEO_TIFF"}
    )
    da = _download_band(url).squeeze()
    return xr.Dataset({"pop_density": da.rename({"x": "lon", "y": "lat"})})

fetch_ndvi_mean

fetch_ndvi_mean(aoi: Geometry, start: str, end: str, scale: int = 1000) -> xr.Dataset

MODIS MOD13A3 monthly NDVI mean over [start, end] (scale factor 0.0001). Returns Dataset with variable 'ndvi'.

Source code in disease/features.py
def fetch_ndvi_mean(
    aoi: ee.Geometry,
    start: str,
    end: str,
    scale: int = 1000,
) -> xr.Dataset:
    """
    MODIS MOD13A3 monthly NDVI mean over [start, end] (scale factor 0.0001).
    Returns Dataset with variable 'ndvi'.
    """

    def _scale_ndvi(img: ee.Image) -> ee.Image:
        return img.multiply(0.0001).copyProperties(img, ["system:time_start"])

    ndvi_col = _require_images(
        ee.ImageCollection("MODIS/061/MOD13A3")
        .filterBounds(aoi)
        .filterDate(start, end)
        .select("NDVI")
        .map(_scale_ndvi),
        "MODIS vegetation",
        start,
        end,
    )
    ndvi_img = ndvi_col.mean().rename("ndvi").clip(aoi)
    url = ndvi_img.getDownloadURL(
        {"region": aoi, "scale": scale, "crs": "EPSG:4326", "format": "GEO_TIFF"}
    )
    da = _download_band(url).squeeze()
    return xr.Dataset({"ndvi": da.rename({"x": "lon", "y": "lat"})})

fetch_landcover

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

ESA WorldCover v200 normalised to [0, 1]. Returns Dataset with variable 'land_cover'.

Source code in disease/features.py
def fetch_landcover(aoi: ee.Geometry, scale: int = 1000) -> xr.Dataset:
    """
    ESA WorldCover v200 normalised to [0, 1].
    Returns Dataset with variable 'land_cover'.
    """
    lc = (
        ee.ImageCollection("ESA/WorldCover/v200")
        .first()
        .select("Map")
        .divide(100)
        .rename("land_cover")
        .clip(aoi)
    )
    url = lc.getDownloadURL(
        {"region": aoi, "scale": scale, "crs": "EPSG:4326", "format": "GEO_TIFF"}
    )
    da = _download_band(url).squeeze()
    return xr.Dataset({"land_cover": da.rename({"x": "lon", "y": "lat"})})

build_feature_datasets

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

Download all seven disease 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. Returns a dict keyed by band group name, consumed by cog_export.

Source code in disease/features.py
def build_feature_datasets(aoi: ee.Geometry, config: dict) -> dict[str, xr.Dataset]:
    """
    Download all seven disease 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.
    Returns a dict keyed by band group name, consumed by cog_export.
    """
    from climate_change.core.dask_engine import DaskEngine
    from climate_change.core.population import fetch_population_count_safe

    scale = config.get("scale", 1000)
    start = config.get("start_date", "2021-01-01")
    end = config.get("end_date", "2023-12-31")
    start, end = _normalise_date_window(start, end)

    return DaskEngine.run_io_parallel(
        {
            "rainfall": lambda: fetch_rainfall_4w(aoi, end_date=end, scale=scale),
            "temperature": lambda: fetch_lst_mean(aoi, start, end, scale=scale),
            "ndwi": lambda: fetch_ndwi(aoi, start, end, scale=scale),
            "elevation": lambda: fetch_elevation(aoi, scale=scale),
            "population": lambda: fetch_pop_density(aoi, scale=scale),
            "ndvi": lambda: fetch_ndvi_mean(aoi, start, end, scale=scale),
            "landcover": lambda: fetch_landcover(aoi, scale=scale),
            # Raw population COUNT for exposure reporting (population within
            # medium/high risk zones) — distinct from "population" above,
            # which is log-transformed for the ML feature and unsuitable for
            # summing. See core.population.fetch_population_count.
            "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 7-band GEE image used for stratified pixel sampling. Band order matches FEATURE_COLS. geometries=True preserved in sampling call so centroids are available for DBSCAN hotspot detection.

Source code in disease/features.py
def build_gee_feature_stack(aoi: ee.Geometry, config: dict) -> ee.Image:
    """
    Assemble the 7-band GEE image used for stratified pixel sampling.
    Band order matches FEATURE_COLS. geometries=True preserved in sampling
    call so centroids are available for DBSCAN hotspot detection.
    """
    start = config.get("start_date", "2021-01-01")
    end = config.get("end_date", "2023-12-31")
    start, end = _normalise_date_window(start, end)
    scale = config.get("scale", 1000)

    # CHIRPS 28-day rainfall
    end_dt = ee.Date(end)
    start_4w = end_dt.advance(-28, "day")
    rain_collection = _require_images(
        ee.ImageCollection("UCSB-CHG/CHIRPS/DAILY")
        .filterBounds(aoi)
        .filterDate(start_4w, end_dt)
        .select("precipitation"),
        "CHIRPS rainfall",
        cast(str, start_4w.format("YYYY-MM-dd").getInfo()),
        end,
    )
    rain_4w = rain_collection.sum().rename("rainfall_4w").clip(aoi)

    # MODIS LST → °C
    def _scale_lst(img: ee.Image) -> ee.Image:
        return img.multiply(0.02).subtract(273.15).copyProperties(img, ["system:time_start"])

    temp_collection = _require_images(
        ee.ImageCollection("MODIS/061/MOD11A2")
        .filterBounds(aoi)
        .filterDate(start, end)
        .select("LST_Day_1km")
        .map(_scale_lst),
        "MODIS land-surface-temperature",
        start,
        end,
    )
    temp_mean = temp_collection.mean().rename("temp_mean").clip(aoi)

    # Sentinel-2 MNDWI
    def _add_mndwi(img: ee.Image) -> ee.Image:
        return (
            img.normalizedDifference(["B3", "B11"])
            .rename("ndwi")
            .copyProperties(img, ["system:time_start"])
        )

    ndwi = _sentinel2_mndwi_collection(aoi, start, end).median().clip(aoi)

    # SRTM elevation
    elevation = ee.Image("USGS/SRTMGL1_003").select("elevation").clip(aoi)

    # WorldPop log(1 + pop)
    pop_collection = (
        ee.ImageCollection("WorldPop/GP/100m/pop")
        .filterBounds(aoi)
        .filter(ee.Filter.eq("year", 2020))
    )
    pop_density = (
        _require_population_image(pop_collection, 2020)
        .select("population")
        .add(1)
        .log()
        .rename("pop_density")
        .clip(aoi)
    )

    # MODIS NDVI mean
    def _scale_ndvi(img: ee.Image) -> ee.Image:
        return img.multiply(0.0001).copyProperties(img, ["system:time_start"])

    ndvi_collection = _require_images(
        ee.ImageCollection("MODIS/061/MOD13A3")
        .filterBounds(aoi)
        .filterDate(start, end)
        .select("NDVI")
        .map(_scale_ndvi),
        "MODIS vegetation",
        start,
        end,
    )
    ndvi = ndvi_collection.mean().rename("ndvi").clip(aoi)

    # ESA WorldCover normalised
    land_cover = (
        ee.ImageCollection("ESA/WorldCover/v200")
        .first()
        .select("Map")
        .divide(100)
        .rename("land_cover")
        .clip(aoi)
        .reproject("EPSG:4326", None, scale)
    )

    stack = ee.Image.cat([rain_4w, temp_mean, ndwi, elevation, pop_density, ndvi, land_cover])
    # Exclude permanent water bodies from feature preparation/sampling — disease
    # risk is a property of inhabited/inhabitable land, not open water.
    valid_land = exclusion_mask(aoi, [WATER_CLASS])
    return stack.updateMask(valid_land)

fetch_monthly_timeseries

fetch_monthly_timeseries(aoi: Geometry, start: str, end: str) -> dict[str, pd.DataFrame]

Fetch monthly area-mean NDVI, rainfall, and LST time series in parallel. The three GEE queries are independent and run concurrently via DaskEngine.run_io_parallel (ThreadPoolExecutor) sharing the GEE session. Returns dict keyed by variable name with a DatetimeIndex DataFrame.

Source code in disease/features.py
def fetch_monthly_timeseries(
    aoi: ee.Geometry,
    start: str,
    end: str,
) -> dict[str, pd.DataFrame]:
    """
    Fetch monthly area-mean NDVI, rainfall, and LST time series in parallel.
    The three GEE queries are independent and run concurrently via
    DaskEngine.run_io_parallel (ThreadPoolExecutor) sharing the GEE session.
    Returns dict keyed by variable name with a DatetimeIndex DataFrame.
    """
    from climate_change.core.dask_engine import DaskEngine

    return DaskEngine.run_io_parallel(
        {
            "ndvi": lambda: _fetch_ndvi_monthly(aoi, start, end),
            "rain": lambda: _fetch_rain_monthly(aoi, start, end),
            "lst": lambda: _fetch_lst_monthly(aoi, start, end),
        }
    )

sample_training_data

sample_training_data(feature_stack: Image, aoi: Geometry, n_pixels: int = 3000, scale: int = 1000, seed: int = 42) -> pd.DataFrame

Sample n_pixels from the GEE feature stack and assign 3-class disease risk labels. Labels: fixed absolute thresholds on the composite risk score (0-100 scale). 0 = Low Risk (score < 33) 1 = Medium Risk (33 <= score < 66) 2 = High Risk (score >= 66) Returns DataFrame with FEATURE_COLS + ['lon', 'lat', 'risk_score', 'label'].

Source code in disease/features.py
def sample_training_data(
    feature_stack: ee.Image,
    aoi: ee.Geometry,
    n_pixels: int = 3000,
    scale: int = 1000,
    seed: int = 42,
) -> pd.DataFrame:
    """
    Sample n_pixels from the GEE feature stack and assign 3-class disease risk labels.
    Labels: fixed absolute thresholds on the composite risk score (0-100 scale).
      0 = Low Risk    (score < 33)
      1 = Medium Risk (33 <= score < 66)
      2 = High Risk   (score >= 66)
    Returns DataFrame with FEATURE_COLS + ['lon', 'lat', 'risk_score', 'label'].
    """
    samples = feature_stack.sample(
        region=aoi,
        scale=scale,
        numPixels=n_pixels,
        seed=seed,
        geometries=True,
        dropNulls=False,
    )
    sample_list = cast(dict, samples.getInfo())["features"]
    if not sample_list:
        raise ValueError(
            "No pixels could be sampled for this area — it may be too small, or "
            "entirely water/built-up land that gets excluded from disease risk sampling."
        )

    raw_df = pd.DataFrame(
        [
            {
                **f["properties"],
                "lon": f["geometry"]["coordinates"][0] if f.get("geometry") else None,
                "lat": f["geometry"]["coordinates"][1] if f.get("geometry") else None,
            }
            for f in sample_list
        ]
    )
    # A band whose pixels are entirely masked over this AOI (e.g. a data-source
    # gap) never appears as a property on any sampled point, so the column is
    # absent altogether rather than merely NaN — dropna(subset=...) would raise
    # an opaque KeyError. Surface it as a clear, actionable message instead.
    missing_cols = [c for c in FEATURE_COLS if c not in raw_df.columns]
    if missing_cols:
        raise ValueError(
            f"No data was available for: {', '.join(missing_cols)} over this area "
            "and date range. Try a different area or a shorter/different date range."
        )

    df = raw_df.dropna(subset=FEATURE_COLS + ["lon", "lat"])[
        FEATURE_COLS + ["lon", "lat"]
    ].reset_index(drop=True)
    if df.empty:
        raise ValueError(
            "All sampled pixels had incomplete feature data for this area and date "
            "range. Try a different area or a shorter/different date range."
        )

    scores = _compute_risk_score(df)
    labels = np.zeros(len(df), dtype=np.intp)
    labels[scores >= RISK_SCORE_THRESHOLDS[0]] = 1
    labels[scores >= RISK_SCORE_THRESHOLDS[1]] = 2

    df["risk_score"] = scores
    df["label"] = labels
    return df

align_datasets

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

Interpolate all datasets onto the reference grid. Land cover is treated as categorical (nearest-neighbour). population_count is passed through entirely unchanged, at its own native ~100 m resolution — it is never interpolated onto this (coarser) 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 every 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 disease/features.py
def align_datasets(
    datasets: dict[str, xr.Dataset],
    ref_key: str = "elevation",
    method_continuous: InterpOptions = "linear",
    method_categorical: InterpOptions = "nearest",
) -> dict[str, xr.Dataset]:
    """
    Interpolate all datasets onto the reference grid. Land cover is treated
    as categorical (nearest-neighbour). population_count is passed through
    entirely unchanged, at its own native ~100 m resolution — it is never
    interpolated onto this (coarser) 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 every 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

DiseaseModel

DiseaseModel()

Orchestrates the full ML pipeline for a single disease surveillance analysis. Always trains Gradient Boosting + XGBoost. config['model_type'] selects which predictions drive the primary risk distribution: "gbm" — Gradient Boosting (default, highest accuracy per lab) "xgboost" — XGBoost "ensemble" — mean softmax probabilities of GBM + XGBoost Trained models and scaler are stored on self for use by cog_export.

Source code in disease/model.py
def __init__(self) -> None:
    self.gbm: GradientBoostingClassifier | None = None
    self.xgb: Booster | None = None
    self.scaler: StandardScaler | None = None

predict

predict(df: DataFrame, timeseries: dict[str, DataFrame] | None = None, config: dict | None = None) -> dict
Parameters

df : DataFrame with FEATURE_COLS + ['lon', 'lat', 'risk_score', 'label'] timeseries : dict of monthly DataFrames (ndvi, rain, lst) from fetch_monthly_timeseries config : optional dict; reads 'model_type' (default 'gbm')

Returns

dict with keys: stats, charts

Source code in disease/model.py
def predict(
    self,
    df: pd.DataFrame,
    timeseries: dict[str, pd.DataFrame] | None = None,
    config: dict | None = None,
) -> dict:
    """
    Parameters
    ----------
    df          : DataFrame with FEATURE_COLS + ['lon', 'lat', 'risk_score', 'label']
    timeseries  : dict of monthly DataFrames (ndvi, rain, lst) from fetch_monthly_timeseries
    config      : optional dict; reads 'model_type' (default 'gbm')

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

    ts = timeseries or {}
    X = df[FEATURE_COLS].fillna(df[FEATURE_COLS].median()).to_numpy(dtype=np.float64)
    y = df["label"].to_numpy(dtype=np.intp)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.20, random_state=42, stratify=_safe_stratify(y)
    )
    self.scaler = StandardScaler()
    X_train_s = self.scaler.fit_transform(X_train)
    X_test_s = self.scaler.transform(X_test)

    # Train GBM 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_gbm = client.submit(train_gbm, X_train_s, y_train, pure=False)
        f_xgb = client.submit(train_xgb, X_train_s, y_train, pure=False)
        (self.gbm, gbm_meta), (self.xgb, xgb_meta) = cast(list, client.gather([f_gbm, f_xgb]))
    else:
        self.gbm, gbm_meta = train_gbm(X_train_s, y_train)
        self.xgb, xgb_meta = train_xgb(X_train_s, y_train)

    assert self.gbm is not None
    assert self.xgb is not None

    eval_result = evaluate_models(self.gbm, self.xgb, X_test_s, y_test)
    shap_payload = compute_shap_importance(self.xgb, X_test_s)

    # Hotspot detection — applied to all pixels using the selected model
    X_all_s = self.scaler.transform(
        df[FEATURE_COLS].fillna(df[FEATURE_COLS].median()).to_numpy(dtype=np.float64)
    )
    _KEY = {"gbm": "gbm", "xgboost": "xgb", "ensemble": "ensemble"}
    result_key = _KEY.get(model_type, "gbm")

    if model_type == "gbm":
        all_preds = self.gbm.predict(X_all_s).astype(int)
    elif model_type == "xgboost":
        all_preds = np.argmax(self.xgb.predict(DMatrix(X_all_s)), axis=1).astype(int)
    else:
        proba_gbm = pad_gbm_proba(self.gbm, self.gbm.predict_proba(X_all_s))
        proba = (proba_gbm + self.xgb.predict(DMatrix(X_all_s))) / 2.0
        all_preds = np.argmax(proba, axis=1).astype(int)

    hotspots = detect_hotspots(df, all_preds)
    charts = build_disease_charts(eval_result, shap_payload, ts, hotspots, model_type)

    # Risk distribution on all pixels
    n_total = len(all_preds)
    counts = np.array([(all_preds == i).sum() for i in range(3)], dtype=np.float64)
    risk_pct = (counts / n_total * 100).round(1)
    high_risk_pct = float(risk_pct[2])

    stats = {
        "model_type": model_type,
        "n_pixels_sampled": int(len(df)),
        "gbm_cv_f1": round(gbm_meta["cv_f1_mean"], 4)
        if gbm_meta["cv_f1_mean"] is not None
        else None,
        "gbm_f1": eval_result["gbm"]["f1"],
        "gbm_accuracy": eval_result["gbm"]["accuracy"],
        "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_accuracy": eval_result["xgb"]["accuracy"],
        "ensemble_f1": eval_result["ensemble"]["f1"],
        "selected_f1": eval_result[result_key]["f1"],
        "high_risk_pct": round(high_risk_pct, 1),
        "n_hotspot_clusters": len(hotspots),
        "top_driver": shap_payload["features"][0],
    }

    _DISEASE_CLASS_NAMES = ["Low Risk", "Medium Risk", "High Risk"]
    if "lon" in df.columns and "lat" in df.columns:
        _sample_points = [
            {
                "lon": round(cast(float, df["lon"].iat[i]), 5),
                "lat": round(cast(float, df["lat"].iat[i]), 5),
                "risk_class": _DISEASE_CLASS_NAMES[int(all_preds[i])],
            }
            for i in range(len(df))
        ]
    else:
        _sample_points = []
    return {"stats": stats, "charts": charts, "_sample_points": _sample_points}

train_gbm

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

Fit Gradient Boosting Classifier with sample weights. Returns (model, metadata).

Source code in disease/model.py
def train_gbm(
    X_train: np.ndarray,
    y_train: np.ndarray,
    cv_folds: int = 5,
) -> tuple[GradientBoostingClassifier, dict]:
    """Fit Gradient Boosting Classifier with sample weights. Returns (model, metadata)."""
    sw = compute_sample_weight("balanced", y_train)
    clf = GradientBoostingClassifier(
        n_estimators=GBM_N_ESTIMATORS,
        learning_rate=GBM_LR,
        max_depth=GBM_MAX_DEPTH,
        subsample=GBM_SUBSAMPLE,
        random_state=42,
    )
    clf.fit(X_train, y_train, sample_weight=sw)
    folds = _safe_cv_folds(y_train, cv_folds)
    if folds is None:
        return clf, {"cv_f1_mean": None, "cv_f1_std": None}
    cv = StratifiedKFold(n_splits=folds, shuffle=True, random_state=42)
    cv_f1 = cross_val_score(clf, X_train, y_train, cv=cv, scoring="f1_macro")
    return clf, {"cv_f1_mean": float(cv_f1.mean()), "cv_f1_std": float(cv_f1.std())}

train_xgb

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

Fit XGBoost multi-class classifier via the low-level Booster API with sample weights.

Uses xgboost.train() rather than the XGBClassifier sklearn wrapper because the wrapper requires every class in DISEASE_CLASSES to appear in y_train (it infers num_class from np.unique(y_train) and rejects gaps). Some AOIs/time windows genuinely have zero samples of a given risk class, which is a legitimate state, not invalid input, so class count is fixed via XGB_PARAMS instead.

Returns (booster, metadata).

Source code in disease/model.py
def train_xgb(
    X_train: np.ndarray,
    y_train: np.ndarray,
    cv_folds: int = 5,
) -> tuple[Booster, dict]:
    """
    Fit XGBoost multi-class classifier via the low-level Booster API with sample weights.

    Uses xgboost.train() rather than the XGBClassifier sklearn wrapper because the
    wrapper requires every class in DISEASE_CLASSES to appear in y_train (it infers
    num_class from np.unique(y_train) and rejects gaps). Some AOIs/time windows
    genuinely have zero samples of a given risk class, which is a legitimate state,
    not invalid input, so class count is fixed via XGB_PARAMS instead.

    Returns (booster, metadata).
    """
    sw = compute_sample_weight("balanced", y_train)
    dtrain = DMatrix(X_train, label=y_train, weight=sw)
    booster = xgb_train(XGB_PARAMS, dtrain, num_boost_round=XGB_N_ESTIMATORS)

    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, shuffle=True, random_state=42)
    cv_f1_scores = []
    for train_idx, val_idx in cv.split(X_train, y_train):
        fold_booster = xgb_train(
            XGB_PARAMS,
            DMatrix(X_train[train_idx], label=y_train[train_idx]),
            num_boost_round=XGB_N_ESTIMATORS,
        )
        fold_pred = np.argmax(fold_booster.predict(DMatrix(X_train[val_idx])), axis=1)
        cv_f1_scores.append(f1_score(y_train[val_idx], fold_pred, average="macro"))

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

pad_gbm_proba

pad_gbm_proba(gbm: GradientBoostingClassifier, proba: ndarray) -> np.ndarray

Expand GBM's predict_proba output to the full DISEASE_CLASSES width.

GradientBoostingClassifier only emits a column per class it actually saw during fit (via gbm.classes_), whereas train_xgb's Booster always outputs len(DISEASE_CLASSES) columns regardless of what the AOI's data contained. Without padding, an AOI missing one risk class produces mismatched shapes (e.g. (n, 2) vs (n, 3)) the moment GBM and XGBoost probabilities are combined for the ensemble.

Source code in disease/model.py
def pad_gbm_proba(gbm: GradientBoostingClassifier, proba: np.ndarray) -> np.ndarray:
    """
    Expand GBM's predict_proba output to the full DISEASE_CLASSES width.

    GradientBoostingClassifier only emits a column per class it actually saw
    during fit (via gbm.classes_), whereas train_xgb's Booster always outputs
    len(DISEASE_CLASSES) columns regardless of what the AOI's data contained.
    Without padding, an AOI missing one risk class produces mismatched shapes
    (e.g. (n, 2) vs (n, 3)) the moment GBM and XGBoost probabilities are
    combined for the ensemble.
    """
    n_classes = len(DISEASE_CLASSES)
    if proba.shape[1] == n_classes:
        return proba
    padded = np.zeros((proba.shape[0], n_classes), dtype=proba.dtype)
    padded[:, gbm.classes_.astype(int)] = proba
    return padded

evaluate_models

evaluate_models(gbm: GradientBoostingClassifier, xgb: Booster, X_test: ndarray, y_test: ndarray) -> dict

Evaluate GBM, XGBoost, and mean-proba ensemble on the held-out test set.

Source code in disease/model.py
def evaluate_models(
    gbm: GradientBoostingClassifier,
    xgb: Booster,
    X_test: np.ndarray,
    y_test: np.ndarray,
) -> dict:
    """Evaluate GBM, XGBoost, and mean-proba ensemble on the held-out test set."""
    gbm_pred = gbm.predict(X_test).astype(int)
    proba_gbm = pad_gbm_proba(gbm, gbm.predict_proba(X_test))
    proba_xgb = xgb.predict(DMatrix(X_test))
    xgb_pred = np.argmax(proba_xgb, axis=1).astype(int)
    ens_pred = np.argmax((proba_gbm + proba_xgb) / 2.0, axis=1).astype(int)

    def _metrics(pred: np.ndarray, label: str) -> dict:
        return {
            "label": label,
            "f1": round(float(f1_score(y_test, pred, average="macro")), 4),
            "accuracy": round(float(accuracy_score(y_test, pred)), 4),
            "predictions": pred.tolist(),
        }

    return {
        "gbm": _metrics(gbm_pred, "Gradient Boosting"),
        "xgb": _metrics(xgb_pred, "XGBoost"),
        "ensemble": _metrics(ens_pred, "Ensemble (mean proba)"),
        "actuals": y_test.tolist(),
    }

compute_shap_importance

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

TreeExplainer SHAP on XGBoost (always used for SHAP regardless of model_type). Returns features sorted by descending mean |SHAP| averaged across all classes.

Source code in disease/model.py
def compute_shap_importance(xgb: Booster, X_test: np.ndarray) -> dict:
    """
    TreeExplainer SHAP on XGBoost (always used for SHAP regardless of model_type).
    Returns features sorted by descending mean |SHAP| averaged across all classes.
    """
    import shap

    X_df = pd.DataFrame(X_test, columns=FEATURE_COLS)
    explainer = shap.TreeExplainer(xgb)
    shap_vals = explainer.shap_values(X_df)

    if isinstance(shap_vals, list):
        mean_abs = np.array([np.abs(sv).mean(axis=0) for sv in shap_vals]).mean(axis=0)
    elif np.asarray(shap_vals).ndim == 3:
        mean_abs = np.abs(shap_vals).mean(axis=(0, 2))
    else:
        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(),
    }

detect_hotspots

detect_hotspots(df: DataFrame, pred_labels: ndarray, eps: float = DBSCAN_EPS, min_samples: int = DBSCAN_MIN_SAMPLES) -> list[dict]

DBSCAN spatial hotspot detection on High Risk (class 2) pixel centroids. Returns a list of cluster dicts with cluster_id, size, lon, lat. Requires df to contain 'lon' and 'lat' columns (preserved from GEE sample geometries).

Source code in disease/model.py
def detect_hotspots(
    df: pd.DataFrame,
    pred_labels: np.ndarray,
    eps: float = DBSCAN_EPS,
    min_samples: int = DBSCAN_MIN_SAMPLES,
) -> list[dict]:
    """
    DBSCAN spatial hotspot detection on High Risk (class 2) pixel centroids.
    Returns a list of cluster dicts with cluster_id, size, lon, lat.
    Requires df to contain 'lon' and 'lat' columns (preserved from GEE sample geometries).
    """
    if "lon" not in df.columns or "lat" not in df.columns:
        return []

    high_risk_mask = pred_labels == 2
    hr_df = df.loc[high_risk_mask, ["lon", "lat"]].reset_index(drop=True)
    if len(hr_df) < min_samples:
        return []

    coords = hr_df[["lon", "lat"]].to_numpy()
    cluster_labels = DBSCAN(eps=eps, min_samples=min_samples).fit_predict(coords)

    hotspots = []
    for cid in sorted(set(cluster_labels)):
        if cid == -1:
            continue
        mask = cluster_labels == cid
        hotspots.append(
            {
                "cluster_id": int(cid),
                "size": int(mask.sum()),
                "lon": round(cast(float, hr_df.loc[mask, "lon"].mean()), 4),
                "lat": round(cast(float, hr_df.loc[mask, "lat"].mean()), 4),
            }
        )
    return sorted(hotspots, key=lambda h: h["size"], reverse=True)

build_disease_charts

build_disease_charts(eval_result: dict, shap_payload: dict, timeseries: dict[str, DataFrame], hotspots: list[dict], model_type: str = 'gbm') -> dict

Assemble frontend-ready chart payloads for the disease surveillance module.

Source code in disease/model.py
def build_disease_charts(
    eval_result: dict,
    shap_payload: dict,
    timeseries: dict[str, pd.DataFrame],
    hotspots: list[dict],
    model_type: str = "gbm",
) -> dict:
    """Assemble frontend-ready chart payloads for the disease surveillance module."""
    _KEY = {"gbm": "gbm", "xgboost": "xgb", "ensemble": "ensemble"}
    result_key = _KEY.get(model_type, "gbm")
    predictions = np.array(eval_result[result_key]["predictions"])

    n_total = len(predictions)
    counts = np.array([(predictions == i).sum() for i in range(3)], dtype=np.float64)
    risk_pct = (counts / n_total * 100).round(1)

    risk_dist = {
        "labels": DISEASE_CLASSES,
        "data": risk_pct.tolist(),
        "colors": DISEASE_COLORS,
    }

    # Build time series datasets aligned to the NDVI index
    ndvi_df = timeseries.get("ndvi", pd.DataFrame())
    rain_df = timeseries.get("rain", pd.DataFrame())
    lst_df = timeseries.get("lst", pd.DataFrame())

    ts_labels = ndvi_df.index.tolist() if not ndvi_df.empty else []
    ts_datasets = []
    if not ndvi_df.empty:
        ts_datasets.append(
            {
                "label": "NDVI",
                "data": ndvi_df["ndvi"].round(4).tolist(),
                "color": "#27AE60",
            }
        )
    if not rain_df.empty:
        aligned_rain = rain_df["rain_mm"].reindex(ts_labels).round(1).tolist()
        ts_datasets.append(
            {
                "label": "Monthly rain (mm)",
                "data": aligned_rain,
                "color": "#2980B9",
            }
        )
    if not lst_df.empty:
        aligned_lst = lst_df["lst"].reindex(ts_labels).round(2).tolist()
        ts_datasets.append(
            {
                "label": "LST (°C)",
                "data": aligned_lst,
                "color": "#E74C3C",
            }
        )

    return {
        "riskDist": risk_dist,
        "timeSeries": {"labels": ts_labels, "datasets": ts_datasets},
        "shap": shap_payload,
        "hotspots": hotspots,
        "model_performance": {
            "gbm": {
                "f1": eval_result["gbm"]["f1"],
                "accuracy": eval_result["gbm"]["accuracy"],
            },
            "xgb": {
                "f1": eval_result["xgb"]["f1"],
                "accuracy": eval_result["xgb"]["accuracy"],
            },
            "ensemble": {
                "f1": eval_result["ensemble"]["f1"],
                "accuracy": eval_result["ensemble"]["accuracy"],
            },
            "selected": model_type,
        },
    }

COG export

cog_export

export_disease_cog

export_disease_cog(gbm_model: GradientBoostingClassifier, xgb_model: Booster, scaler: StandardScaler, datasets: dict[str, Dataset], output_dir: str, prefix: str, model_type: str = 'gbm', aoi_geojson: dict | None = None) -> DiseaseCogResult

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

Parameters

gbm_model : trained GradientBoostingClassifier xgb_model : trained XGBoost Booster scaler : fitted StandardScaler (from DiseaseModel.scaler) datasets : dict returned by build_feature_datasets output_dir : local directory for COG output prefix : file prefix model_type : "gbm" | "xgboost" | "ensemble"

Returns

dict with keys: 'disease_risk' → path string 'risk_pct' → [low_pct, medium_pct, high_pct] over the AOI's valid (non-nodata) pixels, matching DISEASE_CLASSES order. This reflects every pixel actually painted on the map, unlike the training-sample-based stats DiseaseModel.predict() computes, which only covers the few hundred sampled points and can diverge sharply from the full-AOI distribution.

Source code in disease/cog_export.py
def export_disease_cog(
    gbm_model: GradientBoostingClassifier,
    xgb_model: Booster,
    scaler: StandardScaler,
    datasets: dict[str, xr.Dataset],
    output_dir: str,
    prefix: str,
    model_type: str = "gbm",
    aoi_geojson: dict | None = None,
) -> DiseaseCogResult:
    """
    Build the full-AOI feature matrix from in-memory xarray Datasets,
    run inference with the selected model, classify into 3 risk classes,
    and write a Cloud Optimised GeoTIFF.

    Parameters
    ----------
    gbm_model   : trained GradientBoostingClassifier
    xgb_model   : trained XGBoost Booster
    scaler      : fitted StandardScaler (from DiseaseModel.scaler)
    datasets    : dict returned by build_feature_datasets
    output_dir  : local directory for COG output
    prefix      : file prefix
    model_type  : "gbm" | "xgboost" | "ensemble"

    Returns
    -------
    dict with keys:
      'disease_risk' → path string
      'risk_pct'     → [low_pct, medium_pct, high_pct] over the AOI's valid
                       (non-nodata) pixels, matching DISEASE_CLASSES order.
                       This reflects every pixel actually painted on the map,
                       unlike the training-sample-based stats DiseaseModel.predict()
                       computes, which only covers the few hundred sampled points
                       and can diverge sharply from the full-AOI distribution.
    """
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    cog_path = output_path / f"{prefix}_disease_risk.tif"

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

    # Extract feature arrays in FEATURE_COLS order
    rain4w = aligned["rainfall"]["rainfall_4w"].values.astype(np.float32)
    temp = aligned["temperature"]["temp_mean"].values.astype(np.float32)
    ndwi = aligned["ndwi"]["ndwi"].values.astype(np.float32)
    elev = elev_ds["elevation"].values.astype(np.float32)
    pop = aligned["population"]["pop_density"].values.astype(np.float32)
    ndvi = aligned["ndvi"]["ndvi"].values.astype(np.float32)
    lc = aligned["landcover"]["land_cover"].values.astype(np.float32)

    bands = np.stack([rain4w, temp, ndwi, elev, pop, ndvi, lc], axis=0)  # (7, H, W)
    X_full = bands.reshape(len(FEATURE_COLS), -1).T  # (n_pixels, 7)
    valid_mask = ~np.isnan(X_full).any(axis=1)
    # Permanent water bodies are excluded from the disease-risk output.
    water = np.isclose(lc.reshape(-1), _WATER_LAND_COVER)
    valid_mask &= ~water
    X_valid = scaler.transform(X_full[valid_mask])

    # Inference
    if model_type == "gbm":
        proba = gbm_model.predict_proba(X_valid)
    elif model_type == "xgboost":
        proba = xgb_model.predict(DMatrix(X_valid))
    else:  # ensemble
        proba_gbm = pad_gbm_proba(gbm_model, gbm_model.predict_proba(X_valid))
        proba = (proba_gbm + xgb_model.predict(DMatrix(X_valid))) / 2.0

    pred_classes = np.argmax(proba, axis=1).astype(np.uint8) + 1  # 1-indexed

    transform = elev_ds["elevation"].rio.transform()
    crs = elev_ds["elevation"].rio.crs or "EPSG:4326"
    risk_grid: np.ndarray = np.zeros(n_lat * n_lon, dtype=np.uint8)
    risk_grid[valid_mask] = pred_classes
    risk_grid = risk_grid.reshape(n_lat, n_lon)
    if aoi_geojson:
        inside_aoi = geometry_mask(
            [aoi_geojson],
            out_shape=(n_lat, n_lon),
            transform=transform,
            invert=True,
        )
        risk_grid[~inside_aoi] = 0

    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=[1, 2, 3])
        population_by_class = {RISK_INT[v]: by_int[str(v)] for v in (1, 2, 3)}
        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 Risk,2=Medium Risk,3=High Risk",
            MODEL_TYPE=model_type,
        )

    _log.info("COG exported: %s  (%d×%d px)", cog_path, n_lat, n_lon)

    class_counts = np.array([(risk_grid == v).sum() for v in (1, 2, 3)], dtype=np.float64)
    n_valid = class_counts.sum()
    risk_pct = (class_counts / n_valid * 100).round(1).tolist() if n_valid > 0 else [0.0, 0.0, 0.0]

    return {
        "disease_risk": str(cog_path),
        "risk_pct": risk_pct,
        "population_by_class": population_by_class,
        "total_population": total_population,
    }