Skip to content

drought

Composite Drought Index (CDI) monitoring and forecasting. Entry point: DroughtUseCase.

from climate_change import DroughtUseCase

Use case

use_case

DroughtUseCase 

DroughtUseCase(dask_engine: DaskEngine)

Bases: BaseUseCase

Drought risk module — wraps the drought-monitoring package via cdi_runner.py. Implements BaseUseCase and registers itself in core.runner.MODULE_MAP.

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

fetch_data 

fetch_data(config: AnalysisConfig) -> dict

Extract bbox and year range from AnalysisConfig. GEE authentication and data fetching happen inside preprocess().

Source code in drought/use_case.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def fetch_data(self, config: AnalysisConfig) -> dict:
    """
    Extract bbox and year range from AnalysisConfig.
    GEE authentication and data fetching happen inside preprocess().
    """
    lons, lats = _lons_lats(config.aoi_geojson)
    return {
        "bbox": [min(lons), min(lats), max(lons), max(lats)],
        "aoi_geojson": config.aoi_geojson,
        "start_year": int(config.start_date[:4]),
        "end_year": int(config.end_date[:4]),
        "gee_project": config.extra_params.get("gee_project", ""),
        "country": config.country,
    }

preprocess 

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

Fetch ERA5-Land + MODIS from GEE and compute CDI time-series / annual maps. Delegates to a module-level joblib-cached function so self is never part of the cache key (joblib >= 1.4 broke ignore=["self"] on methods). Coordinates are rounded to 6 d.p. before hashing to prevent float-jitter cache misses for semantically identical AOIs.

Source code in drought/use_case.py 
77
78
79
80
81
82
83
84
85
86
def preprocess(self, raw_data: dict, config: AnalysisConfig) -> dict:
    """
    Fetch ERA5-Land + MODIS from GEE and compute CDI time-series / annual maps.
    Delegates to a module-level joblib-cached function so `self` is never
    part of the cache key (joblib >= 1.4 broke ignore=["self"] on methods).
    Coordinates are rounded to 6 d.p. before hashing to prevent float-jitter
    cache misses for semantically identical AOIs.
    """
    ensure_gee(raw_data.get("gee_project") or config.extra_params.get("gee_project", ""))
    return _run_cdi_pipeline_cached(_round_floats(raw_data))

run_model 

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

Run DroughtModel (LSTM or drought_monitoring statistical forecast), build GeoJSON from the spatial CDI map, and package as AnalysisOutput.

Source code in drought/use_case.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def run_model(self, features: dict, config: AnalysisConfig) -> AnalysisOutput:
    """
    Run DroughtModel (LSTM or drought_monitoring statistical forecast),
    build GeoJSON from the spatial CDI map, and package as AnalysisOutput.
    """
    model_type = config.extra_params.get("model_type", "lstm")
    result = DroughtModel().predict(features, {"model_type": model_type})
    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"].get("severity_distribution"), total_area_ha)
    geojson = self._cdi_to_geojson(features, config.aoi_geojson)
    raster_paths: dict[str, str] | None = None
    raster_error: str | None = None
    try:
        cog_paths = export_cdi_cog(
            features,
            output_dir=config.extra_params.get("output_dir", "outputs"),
            aoi_geojson=config.aoi_geojson,
        )
        raster_paths = {key: str(path) for key, path in cog_paths.items()}
    except Exception as exc:
        raster_error = str(exc)

    metadata = {
        "model": f"drought-monitoring CDI + {model_type}",
        "package_version": "0.1.7",
        "country": config.country,
        "start_date": config.start_date,
        "end_date": config.end_date,
        "raster": raster_paths or {},
        "spatial_resolution_source": "drought_monitoring.yearly_drought_maps",
    }
    if raster_error:
        metadata["raster_error"] = raster_error

    return AnalysisOutput(
        module="drought",
        geojson=geojson,
        raster_path=raster_paths,
        stats={**result["stats"], "country": config.country},
        shap=None,  # CDI is an ensemble index — no single SHAP decomposition
        charts=result["charts"],
        metadata=metadata,
    )

run 

run(config: dict) -> dict

Single AOI pipeline — dict config for internal parallel use.

Source code in drought/use_case.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
def run(self, config: dict) -> dict:
    """Single AOI pipeline — dict config for internal parallel use."""
    raw_data = self._fetch_from_dict(config)
    features = _run_cdi_pipeline_local(raw_data)
    result = DroughtModel().predict(features, config)

    output_dir = config.get("output_dir", "outputs")
    cog_paths = export_cdi_cog(
        features,
        output_dir,
        aoi_geojson=config.get("aoi_geojson"),
    )
    result["raster"] = {k: str(v) for k, v in cog_paths.items()}

    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"].get("severity_distribution"), total_area_ha)
    return result

run_date_ranges 

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

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

Source code in drought/use_case.py 
156
157
158
159
160
161
162
163
164
165
def run_date_ranges(self, config: dict, date_ranges: list[dict]) -> list[dict]:
    """Not supported — drought is single-area only. See UseCaseInfo.single_area_only."""
    raise ValueError(
        "Drought analysis only supports a single AOI/date-range at a time. "
        "Its GEE access goes through xee (xarray-Earth-Engine), whose lazy "
        "computation graph requires a forced synchronous Dask scheduler "
        "(see _run_cdi_pipeline_local) that doesn't survive being dispatched "
        "to a remote worker — 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 — drought is single-area only. See UseCaseInfo.single_area_only.

Source code in drought/use_case.py 
167
168
169
170
171
172
173
174
175
176
def run_multi_regions(self, configs: list[dict]) -> list[dict]:
    """Not supported — drought is single-area only. See UseCaseInfo.single_area_only."""
    raise ValueError(
        "Drought analysis only supports a single AOI/date-range at a time. "
        "Its GEE access goes through xee (xarray-Earth-Engine), whose lazy "
        "computation graph requires a forced synchronous Dask scheduler "
        "(see _run_cdi_pipeline_local) that doesn't survive being dispatched "
        "to a remote worker — the two-date-range comparison and multi-region "
        "batch modes are disabled for this module rather than run unreliably."
    )

Features

features

DroughtSequenceDataset 

DroughtSequenceDataset(X: ndarray, y: ndarray, seq_len: int = SEQ_LEN, horizon: int = FORECAST_H)

Bases: Dataset

PyTorch dataset yielding (input_window, target_sequence) pairs.

Source code in drought/features.py 
49
50
51
52
53
54
55
56
57
58
59
def __init__(
    self,
    X: np.ndarray,
    y: np.ndarray,
    seq_len: int = SEQ_LEN,
    horizon: int = FORECAST_H,
):
    self.X = X
    self.y = y
    self.seq_len = seq_len
    self.horizon = horizon

build_features 

build_features(df: DataFrame, lags: tuple = (1, 3, 6, 12)) -> pd.DataFrame

Enrich CDI DataFrame with lag features, rolling statistics, and seasonal sine/cosine encodings. Rows with NaN (warm-up period) are dropped.

Source code in drought/features.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def build_features(df: pd.DataFrame, lags: tuple = (1, 3, 6, 12)) -> pd.DataFrame:
    """
    Enrich CDI DataFrame with lag features, rolling statistics, and seasonal
    sine/cosine encodings. Rows with NaN (warm-up period) are dropped.
    """
    feat = df[["PDI", "TDI", "VDI", "CDI"]].copy()
    for lag in lags:
        feat[f"CDI_lag{lag}"] = feat["CDI"].shift(lag)
    feat["CDI_roll3_mean"] = feat["CDI"].rolling(3).mean()
    feat["CDI_roll6_std"] = feat["CDI"].rolling(6).std()
    months = pd.DatetimeIndex(feat.index).month
    feat["sin_month"] = np.sin(2 * np.pi * months / 12)
    feat["cos_month"] = np.cos(2 * np.pi * months / 12)
    return feat.dropna()

prepare_datasets 

prepare_datasets(feat_df: DataFrame, seq_len: int = SEQ_LEN, horizon: int = FORECAST_H, holdout_months: int = HOLDOUT_MONTHS) -> tuple[DroughtSequenceDataset, DroughtSequenceDataset, StandardScaler, torch.Tensor]

Scale features, split train / test, and package into PyTorch datasets.

Source code in drought/features.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def prepare_datasets(
    feat_df: pd.DataFrame,
    seq_len: int = SEQ_LEN,
    horizon: int = FORECAST_H,
    holdout_months: int = HOLDOUT_MONTHS,
) -> tuple[DroughtSequenceDataset, DroughtSequenceDataset, StandardScaler, torch.Tensor]:
    """
    Scale features, split train / test, and package into PyTorch datasets.
    """
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(feat_df[INPUT_COLS])
    y = feat_df["CDI"].to_numpy(dtype=float)
    split = len(X_scaled) - holdout_months
    train_ds = DroughtSequenceDataset(X_scaled[:split], y[:split], seq_len, horizon)
    test_ds = DroughtSequenceDataset(X_scaled[split:], y[split:], seq_len, horizon)

    last_seq = torch.tensor(
        X_scaled[-seq_len:].reshape(1, seq_len, len(INPUT_COLS)),
        dtype=torch.float32,
    )
    return train_ds, test_ds, scaler, last_seq

Model

model

DroughtLSTM 

DroughtLSTM(n_features: int = len(INPUT_COLS), hidden: int = HIDDEN, layers: int = LAYERS, horizon: int = FORECAST_H)

Bases: Module

2-layer LSTM with dropout. Final hidden state → linear head → FORECAST_H outputs. All six forecast steps are predicted simultaneously (multi-output regression).

Source code in drought/model.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def __init__(
    self,
    n_features: int = len(INPUT_COLS),
    hidden: int = HIDDEN,
    layers: int = LAYERS,
    horizon: int = FORECAST_H,
):
    super().__init__()
    # PyTorch only applies recurrent dropout between stacked LSTM layers.
    # Passing a non-zero value for a single layer has no effect and emits a
    # warning, so disable it for that configuration.
    recurrent_dropout = 0.2 if layers > 1 else 0.0
    self.lstm = nn.LSTM(
        n_features,
        hidden,
        layers,
        batch_first=True,
        dropout=recurrent_dropout,
    )
    self.fc = nn.Linear(hidden, horizon)

DroughtModel

Orchestrates the drought ML pipeline.

model_type (from config, default 'lstm'): 'lstm' — train DroughtLSTM + MC Dropout forecast (existing) 'drought_monitoring' — use drought_monitoring package's built-in forecast

Designed to be called from DroughtUseCase.run_model().

train_lstm 

train_lstm(train_ds: DroughtSequenceDataset, test_ds: DroughtSequenceDataset, device: device | None = None) -> tuple[DroughtLSTM, dict]

Train with early stopping (patience = PATIENCE epochs on val MSE). Best weights are restored before returning.

Source code in drought/model.py 
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def train_lstm(
    train_ds: DroughtSequenceDataset,
    test_ds: DroughtSequenceDataset,
    device: torch.device | None = None,
) -> tuple[DroughtLSTM, dict]:
    """
    Train with early stopping (patience = PATIENCE epochs on val MSE).
    Best weights are restored before returning.
    """
    device = device or torch.device("cpu")

    model = DroughtLSTM().to(device)
    optimiser = torch.optim.Adam(model.parameters(), lr=LR)
    criterion = nn.MSELoss()

    train_loader = DataLoader(train_ds, batch_size=BATCH, shuffle=True)
    test_loader = DataLoader(test_ds, batch_size=BATCH)

    train_losses, val_losses = [], []
    best_val, patience_ctr = float("inf"), 0
    best_state = copy.deepcopy(model.state_dict())

    for epoch in range(EPOCHS):
        model.train()
        batch_losses = []
        for Xb, yb in train_loader:
            optimiser.zero_grad()
            loss = criterion(model(Xb.to(device)), yb.to(device))
            loss.backward()
            optimiser.step()
            batch_losses.append(loss.item())
        train_losses.append(float(np.mean(batch_losses)))

        model.eval()
        with torch.no_grad():
            vl = [criterion(model(Xb.to(device)), yb.to(device)).item() for Xb, yb in test_loader]
        val_losses.append(float(np.mean(vl)))

        if val_losses[-1] < best_val:
            best_val = val_losses[-1]
            best_state = copy.deepcopy(model.state_dict())
            patience_ctr = 0
        else:
            patience_ctr += 1
            if patience_ctr >= PATIENCE:
                model.load_state_dict(best_state)
                return model, {
                    "train_losses": train_losses,
                    "val_losses": val_losses,
                    "best_val_mse": best_val,
                    "stopped_epoch": epoch + 1,
                }

    model.load_state_dict(best_state)
    return model, {
        "train_losses": train_losses,
        "val_losses": val_losses,
        "best_val_mse": best_val,
        "stopped_epoch": EPOCHS,
    }

evaluate_lstm 

evaluate_lstm(model: DroughtLSTM, test_ds: DroughtSequenceDataset, device: device | None = None) -> dict

1-step-ahead evaluation on the held-out test set.

Source code in drought/model.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def evaluate_lstm(
    model: DroughtLSTM,
    test_ds: DroughtSequenceDataset,
    device: torch.device | None = None,
) -> dict:
    """
    1-step-ahead evaluation on the held-out test set.
    """
    device = device or torch.device("cpu")
    loader = DataLoader(test_ds, batch_size=BATCH)

    model.eval()
    preds_all, true_all = [], []
    with torch.no_grad():
        for Xb, yb in loader:
            preds_all.append(model(Xb.to(device)).cpu().numpy()[:, 0])
            true_all.append(yb[:, 0].numpy())

    preds = np.concatenate(preds_all)
    true = np.concatenate(true_all)

    return {
        "mae": float(mean_absolute_error(true, preds)),
        "rmse": float(np.sqrt(np.mean((preds - true) ** 2))),
        "predictions": preds.round(4).tolist(),
        "actuals": true.round(4).tolist(),
    }

forecast_with_uncertainty 

forecast_with_uncertainty(model: DroughtLSTM, last_seq: Tensor, future_dates: DatetimeIndex, n_samples: int = 500, device: device | None = None) -> dict

MC Dropout forecast: enable dropout at inference and run n_samples stochastic forward passes. The spread of the resulting distribution is the model's epistemic (model) uncertainty — it does NOT capture data or structural uncertainty.

Source code in drought/model.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def forecast_with_uncertainty(
    model: DroughtLSTM,
    last_seq: torch.Tensor,
    future_dates: pd.DatetimeIndex,
    n_samples: int = 500,
    device: torch.device | None = None,
) -> dict:
    """
    MC Dropout forecast: enable dropout at inference and run n_samples stochastic
    forward passes. The spread of the resulting distribution is the model's
    epistemic (model) uncertainty — it does NOT capture data or structural uncertainty.
    """
    device = device or torch.device("cpu")
    last_seq = last_seq.to(device)
    model.train()  # activate dropout layers for MC sampling
    mc_samples: list[np.ndarray] = []
    with torch.no_grad():
        for _ in range(n_samples):
            mc_samples.append(model(last_seq).cpu().numpy().flatten())
    model.eval()
    mc_preds: np.ndarray = np.array(mc_samples)  # (n_samples, FORECAST_H)

    ci_lo = np.percentile(mc_preds, 2.5, axis=0)
    ci_hi = np.percentile(mc_preds, 97.5, axis=0)
    ci_half = (ci_hi - ci_lo) / 2
    avg_half = float(ci_half.mean())
    return {
        "dates": future_dates.strftime("%Y-%m").tolist(),
        "mean": mc_preds.mean(axis=0).round(4).tolist(),
        "std": mc_preds.std(axis=0).round(4).tolist(),
        "ci_lower": ci_lo.round(4).tolist(),
        "ci_upper": ci_hi.round(4).tolist(),
        "ci_half_width": ci_half.round(4).tolist(),
        "ci_level": 0.95,
        "ci_summary": f"95% CI: ±{avg_half:.2f}",
        "n_samples": n_samples,
    }

run_kmeans_typology 

run_kmeans_typology(ds, n_clusters: int = 4) -> dict

Cluster pixel CDI trajectories into drought typologies using KMeans. Clusters are ranked by ascending mean CDI (most drought-prone = cluster 0).

Source code in drought/model.py 
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
def run_kmeans_typology(ds, n_clusters: int = 4) -> dict:
    """
    Cluster pixel CDI trajectories into drought typologies using KMeans.
    Clusters are ranked by ascending mean CDI (most drought-prone = cluster 0).
    """
    cdi_np = ds["CDI"].values  # (time, lon, lat)
    n_time, n_lon, n_lat = cdi_np.shape
    pixel_matrix = cdi_np.reshape(n_time, n_lon * n_lat).T  # (pixels, time)
    valid_mask = np.isfinite(pixel_matrix).any(axis=1)
    valid_pixels = pixel_matrix[valid_mask]
    if valid_pixels.size == 0:
        return {
            "label_map": np.full((n_lon, n_lat), -1, dtype=int).tolist(),
            "lons": ds.lon.values.tolist(),
            "lats": ds.lat.values.tolist(),
            "clusters": [],
            "n_clusters": 0,
            "nodata_pixel_count": int(pixel_matrix.shape[0]),
        }

    pixel_means = np.nanmean(valid_pixels, axis=1)
    global_mean = float(np.nanmean(valid_pixels))
    pixel_means = np.where(np.isfinite(pixel_means), pixel_means, global_mean)
    row_idx, col_idx = np.where(~np.isfinite(valid_pixels))
    valid_pixels = valid_pixels.copy()
    valid_pixels[row_idx, col_idx] = pixel_means[row_idx]

    effective_clusters = min(n_clusters, len(valid_pixels))
    ss = SScaler()
    pixel_scaled = ss.fit_transform(valid_pixels)

    km = KMeans(n_clusters=effective_clusters, random_state=42, n_init=10)
    labels = km.fit_predict(pixel_scaled)

    # Rank clusters: 0 = most drought-prone (lowest mean CDI)
    cluster_means = {c: float(valid_pixels[labels == c].mean()) for c in range(effective_clusters)}
    rank_map = {
        old: new for new, old in enumerate(sorted(cluster_means, key=lambda c: cluster_means[c]))
    }
    labels_r = np.vectorize(rank_map.get)(labels)
    full_labels = np.full(pixel_matrix.shape[0], -1, dtype=int)
    full_labels[valid_mask] = labels_r
    label_map_r = full_labels.reshape(n_lon, n_lat)

    years = ds.time.dt.year.values.tolist()
    clusters = [
        {
            "cluster_id": int(c),
            "mean_cdi": round(cluster_means[{v: k for k, v in rank_map.items()}[c]], 4),
            "pixel_count": int((labels_r == c).sum()),
            "pixel_pct": round(float((labels_r == c).sum()) / len(labels_r) * 100, 1),
            "trajectory": valid_pixels[labels_r == c].mean(axis=0).round(4).tolist(),
            "years": years,
        }
        for c in range(effective_clusters)
    ]

    return {
        "label_map": label_map_r.tolist(),
        "lons": ds.lon.values.tolist(),
        "lats": ds.lat.values.tolist(),
        "clusters": clusters,
        "n_clusters": effective_clusters,
        "nodata_pixel_count": int((~valid_mask).sum()),
    }

drought_severity_stats 

drought_severity_stats(latest_cdi: ndarray) -> dict

Summarise the most recent CDI raster into a mean value and the percentage of valid pixels falling in each of the 8 standard drought/wet severity bands (see _classify_cdi_value in cdi_runner.py for the same cutoffs).

Source code in drought/model.py 
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
def drought_severity_stats(latest_cdi: np.ndarray) -> dict:
    """
    Summarise the most recent CDI raster into a mean value and the percentage
    of valid pixels falling in each of the 8 standard drought/wet severity
    bands (see `_classify_cdi_value` in `cdi_runner.py` for the same cutoffs).
    """
    valid = latest_cdi[np.isfinite(latest_cdi)]
    if valid.size == 0:
        return {
            "latest_mean_cdi": 0.0,
            "aoi_valid_pixel_count": 0,
            "extreme_pct": 0.0,
            "severe_pct": 0.0,
            "moderate_pct": 0.0,
            "mild_pct": 0.0,
            "near_normal_pct": 0.0,
            "mild_wet_pct": 0.0,
            "moderately_wet_pct": 0.0,
            "very_wet_pct": 0.0,
        }
    return {
        "latest_mean_cdi": round(float(np.nanmean(valid)), 4),
        "aoi_valid_pixel_count": int(valid.size),
        "extreme_pct": round(float((valid < 0.50).mean() * 100), 1),
        "severe_pct": round(float(((valid >= 0.50) & (valid < 0.65)).mean() * 100), 1),
        "moderate_pct": round(float(((valid >= 0.65) & (valid < 0.80)).mean() * 100), 1),
        "mild_pct": round(float(((valid >= 0.80) & (valid < 0.90)).mean() * 100), 1),
        "near_normal_pct": round(float(((valid >= 0.90) & (valid < 1.10)).mean() * 100), 1),
        "mild_wet_pct": round(float(((valid >= 1.10) & (valid < 1.20)).mean() * 100), 1),
        "moderately_wet_pct": round(float(((valid >= 1.20) & (valid < 1.30)).mean() * 100), 1),
        "very_wet_pct": round(float((valid >= 1.30).mean() * 100), 1),
    }

CDI pipeline

cdi_runner

run_cdi_pipeline 

run_cdi_pipeline(raw_data: dict) -> dict

Full CDI pipeline.

Source code in drought/cdi_runner.py 
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
def run_cdi_pipeline(raw_data: dict) -> dict:
    """
    Full CDI pipeline.
    """
    bbox = raw_data["bbox"]  # [lon_min, lat_min, lon_max, lat_max]
    aoi = raw_data.get("aoi_geojson") or bbox
    start_year = raw_data["start_year"]
    end_year = raw_data["end_year"]

    # Monthly area-averaged time series
    precip = fetch_era5_precip(aoi, start_year=start_year, end_year=end_year)
    temp = fetch_era5_temp(aoi, start_year=start_year, end_year=end_year)
    ndvi = fetch_modis_ndvi(aoi, start_year=start_year, end_year=end_year)

    df = compute_all(precip, temp, ndvi, window=3, weights=(0.50, 0.25, 0.25))
    df = _temporally_fill_dataframe(df)
    df["severity"] = classify_cdi(df["CDI"]).map(_normalize_drought_class)

    # Annual pixel-wise spatial CDI maps
    # end_year - 1 because the current partial year has no complete annual map
    ds = _yearly_drought_maps(
        aoi,
        bbox,
        start_year=start_year,
        end_year=end_year - 1,
    )
    ds = _temporally_fill_dataset(ds)
    ds = _mask_dataset_water_bodies(ds, raw_data.get("aoi_geojson"))
    ds = _mask_dataset_to_aoi(ds, raw_data.get("aoi_geojson"))
    return {
        "cdi_maps": ds,
        "cdi_series": df,
        "precip": precip,
        "temp": temp,
        "ndvi": ndvi,
        "bbox": bbox,
        "aoi_geojson": raw_data.get("aoi_geojson"),
        "country": raw_data.get("country"),
        "start_year": start_year,
        "end_year": end_year,
    }

build_drought_charts 

build_drought_charts(features: dict) -> dict

Build chart data payloads for the frontend.

Source code in drought/cdi_runner.py 
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
def build_drought_charts(features: dict) -> dict:
    """
    Build chart data payloads for the frontend.
    """
    df = features["cdi_series"]

    timeseries = {
        "labels": df.index.strftime("%Y-%m").tolist(),
        "datasets": [
            {"label": "CDI", "data": df["CDI"].tolist(), "color": "#C0392B"},
            {"label": "PDI", "data": df["PDI"].tolist(), "color": "#2980B9"},
            {"label": "TDI", "data": df["TDI"].tolist(), "color": "#E67E22"},
            {"label": "VDI", "data": df["VDI"].tolist(), "color": "#27AE60"},
        ],
    }

    annual = df["CDI"].resample("YE").mean()
    anomaly = {
        "labels": annual.index.year.tolist(),
        "data": (annual - annual.mean()).round(4).tolist(),
        "mean": float(annual.mean()),
    }

    seasonal = df["CDI"].groupby(df.index.month).mean()
    seasonal_chart = {
        "labels": list(range(1, 13)),
        "data": seasonal.round(4).tolist(),
    }

    temporal_severity_dist = df["severity"].value_counts(normalize=True).mul(100).round(1)
    ds = features.get("cdi_maps")
    latest = ds["CDI"].isel(time=-1).values if ds is not None else np.array([])
    valid = latest[np.isfinite(latest)]
    total_population: float | None = None
    population_affected: float | None = None
    pop_by_class: dict[str, float] = {}
    if valid.size:
        classes = np.array([_classify_cdi_value(float(value)) for value in valid])
        counts = {label: int((classes == label).sum()) for label in DROUGHT_CLASS_ORDER}
        labels = [label for label in DROUGHT_CLASS_ORDER if counts[label] > 0]
        data = [round(counts[label] / valid.size * 100, 1) for label in labels]

        pop_ds = (
            _fetch_population(features.get("aoi_geojson"), features.get("country"))
            if ds is not None
            else None
        )
        if pop_ds is not None and ds is not None:
            # Classify the full 2D grid, explicitly transposed to (lat, lon)
            # order — ds["CDI"]'s raw array is (lon, lat) (see
            # compute_spatial_uncertainty's comments below), which would
            # silently mismatch _dataset_transform's (lat, lon) assumption
            # and warp the upsampled classification incorrectly otherwise.
            lat_dim, lon_dim = _spatial_dims(ds)
            latest_latlon = ds["CDI"].isel(time=-1).transpose(lat_dim, lon_dim).values
            # _dataset_transform always assumes row 0 = north (standard
            # raster convention), but ds's lat coordinate may be stored
            # ascending (south-to-north) rather than descending — the same
            # mismatch _mask_dataset_to_aoi already guards against above.
            lats = ds[lat_dim].values
            if len(lats) > 1 and float(lats[0]) < float(lats[-1]):
                latest_latlon = latest_latlon[::-1, :]
            class_grid = np.full(latest_latlon.shape, "", dtype=object)
            finite_mask = np.isfinite(latest_latlon)
            class_grid[finite_mask] = [
                _classify_cdi_value(float(v)) for v in latest_latlon[finite_mask]
            ]
            class_transform = _dataset_transform(ds)
            pop_by_class = population_exposure(
                class_grid, class_transform, pop_ds, classes=DROUGHT_CLASS_ORDER
            )
            total_population = round(sum(pop_by_class.values()), 1)
            population_affected = at_risk_population(pop_by_class, DROUGHT_AT_RISK_LABELS)
    else:
        labels = temporal_severity_dist.index.tolist()
        data = temporal_severity_dist.tolist()
    severity = {
        "labels": labels,
        "data": data,
        "colors": [DROUGHT_CLASS_COLORS.get(str(label), "#95A5A6") for label in labels],
        "basis": "latest_spatial_aoi",
        "valid_pixel_count": int(valid.size),
    }
    if total_population is not None:
        severity["data_population"] = [pop_by_class.get(label, 0.0) for label in labels]
        severity["total_population"] = total_population
        severity["population_affected"] = population_affected
    temporal_severity = {
        "labels": temporal_severity_dist.index.tolist(),
        "data": temporal_severity_dist.tolist(),
        "colors": [
            DROUGHT_CLASS_COLORS.get(str(label), "#95A5A6")
            for label in temporal_severity_dist.index
        ],
        "basis": "monthly_aoi_series",
    }

    return {
        "timeseries": timeseries,
        "anomaly": anomaly,
        "seasonal": seasonal_chart,
        "severity_distribution": severity,
        "temporal_severity_distribution": temporal_severity,
    }

compute_spatial_uncertainty 

compute_spatial_uncertainty(features: dict) -> dict

Two spatial uncertainty metrics derived from the annual CDI dataset.

Source code in drought/cdi_runner.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
def compute_spatial_uncertainty(features: dict) -> dict:
    """
    Two spatial uncertainty metrics derived from the annual CDI dataset.
    """
    ds = features["cdi_maps"]
    cdi_std = ds["CDI"].std("time").values  # (lon, lat)
    stacked = np.stack(
        [ds["PDI"].values, ds["TDI"].values, ds["VDI"].values], axis=0
    )  # (3, time, lon, lat)
    component_spread = np.nanmean(np.nanstd(stacked, axis=0), axis=0)  # (lon, lat)
    finite_cdi_std = cdi_std[np.isfinite(cdi_std)]
    finite_component_spread = component_spread[np.isfinite(component_spread)]

    return {
        "lons": ds.lon.values.tolist(),
        "lats": ds.lat.values.tolist(),
        "temporal_std": cdi_std.T.tolist(),  # (lat, lon)
        "component_spread": component_spread.T.tolist(),  # (lat, lon)
        "temporal_std_stats": {
            "min": float(finite_cdi_std.min()) if finite_cdi_std.size else 0.0,
            "max": float(finite_cdi_std.max()) if finite_cdi_std.size else 0.0,
            "mean": float(finite_cdi_std.mean()) if finite_cdi_std.size else 0.0,
        },
        "component_spread_stats": {
            "min": float(finite_component_spread.min()) if finite_component_spread.size else 0.0,
            "max": float(finite_component_spread.max()) if finite_component_spread.size else 0.0,
            "mean": float(finite_component_spread.mean()) if finite_component_spread.size else 0.0,
        },
    }

export_cdi_cog 

export_cdi_cog(features: dict, output_dir: str, aoi_geojson: dict | None = None) -> dict[str, Path]

Export CDI maps to Cloud Optimised GeoTIFFs. Called only at the save step.

Source code in drought/cdi_runner.py 
588
589
590
591
592
593
594
595
596
597
598
599
600
601
def export_cdi_cog(
    features: dict,
    output_dir: str,
    aoi_geojson: dict | None = None,
) -> dict[str, Path]:
    """Export CDI maps to Cloud Optimised GeoTIFFs. Called only at the save step."""
    from drought_monitoring.io import cdi_stack_to_cog

    ds = features["cdi_maps"]
    prefix = f"drought_{features['start_year']}_{features['end_year']}"
    paths = cdi_stack_to_cog(ds, output_dir=output_dir, prefix=prefix)
    if aoi_geojson:
        paths = {key: _clip_cog_to_aoi(path, aoi_geojson) for key, path in paths.items()}
    return paths