CACD-Flood: Context-Aware Concept Distillation for Trustworthy Flood Prediction

Reference implementation for the IJCAI 2026 (AI for Good track) paper:

Context-Aware Concept Distillation for Trustworthy Flood Prediction
Eli Levinkopf, Efrat Morin, Claudia Goldman.
Proceedings of the 35th International Joint Conference on Artificial Intelligence (IJCAI 2026), AI for Good Track.

Paper: paper/Context_Aware_Concept_Distillation_for_Trustworthy_Flood_Prediction.pdf

---

What this repo does

CACD distills a black-box hydrological LSTM into an interpretable surrogate that explains flood predictions through a small "Hydrological Language" of human-readable concepts (e.g., Convective Flash Flood, Delayed Snowmelt).

The framework has three stages:

1. Unsupervised Concept Discovery — SHAP attributions of the LSTM are aggregated over log-temporal windows using four robust statistics (Net Influence, Magnitude, Soft Peak Lag, rescaled HHI Concentration), then hierarchically clustered into K = 6 concepts.
2. Concept Encoder — an MLP that maps LSTM hidden states h ∈ ℝ²⁵⁶ to concept probabilities p ∈ ℝ⁶.
3. Context-Aware Decoder — a Residual Hypernetwork that uses static basin attributes to modulate a base linear law ŷ = W(xₛ)·p + b(xₛ), keeping the prediction linear in concepts (interpretable) while letting weights adapt per basin.

On 5,203 globally-distributed basins, the interpretable surrogate achieves median fidelity NSE 0.70, beating a black-box Deep MLP (0.60) and a naive linear decoder (0.25).

Repository layout

cacd-flood/
├── cacd/                            # Python package (pip install -e .)
│   ├── concept_bottleneck/          # Encoder, decoders, trainers, fidelity analyzer
│   │   ├── data/                    # LSTM output extraction & dataset builder
│   │   ├── models/                  # ConceptEncoder, SurrogateDecoder, ContextAwareDecoder
│   │   ├── training/                # Encoder + decoder trainers
│   │   └── evaluation/              # FidelityAnalyzer
│   ├── explainability/              # SHAP base + BasinPreprocessor
│   ├── model/                       # Loader for pre-trained neuralhydrology LSTM
│   ├── post_analysis/               # SHAP-based concept discovery (clustering)
│   └── utils/                       # Basin filters
├── scripts/                         # End-to-end pipeline (cbm_01 ... cbm_06)
├── configs/                         # Decoder configs (linear / mlp / context-aware)
├── paper/                           # Camera-ready PDF
├── tests/                           # Smoke tests
├── setup.py
├── LICENSE                          # MIT
└── CITATION.cff

Installation

git clone https://github.com/elilevinkopf/cacd-flood.git
cd cacd-flood

# Create a fresh environment (Python ≥ 3.9)
python -m venv .venv && source .venv/bin/activate

# Install the package in editable mode
pip install -e .

The teacher LSTM is trained with neuralhydrology, which is pulled in by setup.py.

Data

The framework operates on the Caravan global hydrological dataset (Kratzert et al., 2023). After downloading Caravan, the directory structure expected by neuralhydrology is:

data/Caravan/
├── attributes/      # static basin attributes (51 features)
├── timeseries/      # daily forcings + observed streamflow
└── shapefiles/

Set the data path in your neuralhydrology run config (used to train the teacher LSTM) and in the SHAP / dataset-building steps below.

Reproducing the paper

Pre-trained teacher LSTM weights are published as a release asset (teacher_lstm_v1.tar.gz). Place them under runs/teacher_lstm/ so the path matches runs/teacher_lstm/train/model_epoch025/ plus the standard neuralhydrology files (config.yml, model_epoch025.pt, train_data/).

The full pipeline runs end-to-end with six scripts. Each step writes its outputs under the run directory and is consumed by the next step.

# 0. (Optional) Train the teacher LSTM yourself with neuralhydrology
#    See: https://github.com/neuralhydrology/neuralhydrology
#    Or download our pre-trained weights from the GitHub release.

# 1. Extract LSTM hidden states & predictions for train / val / test
python scripts/cbm_01_extract_lstm.py \
    --run_dir runs/teacher_lstm \
    --epoch 25 \
    --output_dir runs/teacher_lstm/train/model_epoch025/cbm \
    --periods train validation test

# 2. Run SHAP analysis + log-temporal feature engineering + hierarchical
#    K-Means clustering to discover the K=6 "Hydrological Language" concepts.
python -m cacd.post_analysis.extreme_yearly_shap_analysis_subclusters \
    --run_dir runs/teacher_lstm --epoch 25

# 3. Build the master CBM dataset (encoder + decoder splits)
python scripts/cbm_02_prepare_data.py \
    --run_dir runs/teacher_lstm --epoch 25 --mode encoder
python scripts/cbm_02_prepare_data.py \
    --run_dir runs/teacher_lstm --epoch 25 --mode decoder

# 4. Train the Concept Encoder (Stage 1) — frozen after this step.
python scripts/cbm_03_train_encoder.py \
    --run_dir runs/teacher_lstm/train/model_epoch025/cbm \
    --n_epochs 100

# 5. Train the decoder (Stage 2). Pick the architecture via --config:
#       configs/cbm_decoder_linear.yml          (Naive Linear baseline)
#       configs/cbm_decoder_mlp.yml             (Deep MLP baseline)
#       configs/cbm_decoder_context_aware.yml   (Ours: Context-Aware)
python scripts/cbm_04_train_decoder.py \
    --config configs/cbm_decoder_context_aware.yml

# 6. Evaluate fidelity (Global / Median NSE / Success Rate)
python scripts/cbm_05_analyze_fidelity.py \
    --config configs/cbm_decoder_context_aware.yml

# 7. (Optional) Semantic validation of the Context-Aware Decoder
#    correlates learned per-basin weights with static attributes.
python scripts/cbm_06_semantic_validation.py \
    --config configs/cbm_decoder_context_aware.yml

Run the same step-5/6 commands with the linear and MLP configs to reproduce all three rows of Table 1 in the paper.

Expected results (Table 1, paper)

Decoder	Global NSE	Median NSE	Success Rate (NSE > 0.7)
Naive Linear (Global)	0.67	0.25	18.8 %
Deep MLP (Black Box)	0.83	0.60	39.9 %
Context-Aware (Ours)	0.88	0.70	50.3 %

License

MIT — see LICENSE.

Citation

@inproceedings{levinkopf2026cacd,
  title     = {Context-Aware Concept Distillation for Trustworthy Flood Prediction},
  author    = {Levinkopf, Eli and Morin, Efrat and Goldman, Claudia},
  booktitle = {Proceedings of the 35th International Joint Conference on Artificial Intelligence (IJCAI), AI for Good Track},
  year      = {2026},
}

Acknowledgments

The teacher LSTM is trained with neuralhydrology. We use the Caravan global hydrological dataset.