DA-STDK

Reference code for cluster-aware conformal calibration in spatio-temporal distributional prediction.

• Cluster-adaptive spatial bases — centers and scales initialized from sampling density, so capacity follows heterogeneous observation patterns instead of a fixed grid.
• Cluster-aware conformal calibration — interval widths are calibrated within spatial clusters, with a global fallback when local samples are scarce.

Benchmarks in this repo use the KAUST spatio-temporal datasets (scenarios 2a/2b).

Architecture

Cluster-adaptive spatial basis, temporal basis, and covariates are concatenated and passed through a shared MLP trunk. Quantile heads predict multiple levels; cluster-aware CQR produces calibrated prediction intervals.

Install

Python 3.10+ and Poetry are enough for most use:

poetry install --with dev

Optional: Conda env via bash envs/conda/build_conda_env.sh then conda activate st-dadk.

pip install da-stdk   # after a PyPI release
# or locally:
pip install -e .

import da_stdk
from da_stdk.models import STDKMLP, create_model
from da_stdk.data.kaust_loader import load_kaust_csv_single

Run

Single training run

poetry run python scripts/train_default.py

KAUST benchmark (multiple scenarios / models)

make kaust
# or (train only, then analyze manually):
poetry run python scripts/run_kaust_data.py --config configs/config_default.yaml
poetry run python scripts/analyze_kaust_results.py --results_dir results/kaust_data_<timestamp>

make kaust runs all scenario×model combos and calls analyze_kaust_results.py when finished (--analyze). Use make kaust-dry to preview commands.

More scripts and flags: scripts/README.md.

Layout

Dev

make test          # pytest
make lint          # black, isort, mypy
pre-commit run --all-files

Citation

If you use this code, please cite:

Cluster-Aware Conformal Calibration for Spatio-Temporal Distributional Prediction Gooyoung Kim, Chae Young Lim, Wen-Ting Wang, Hao-Yun Huang, Wei-Ying Wu arXiv preprint arXiv:2606.06753, 2026. https://arxiv.org/abs/2606.06753

@misc{kim2026clusterawareconformalcalibrationspatiotemporal,
      title={Cluster-Aware Conformal Calibration for Spatio-Temporal Distributional Prediction},
      author={Gooyoung Kim and Chae Young Lim and Wen-Ting Wang and Hao-Yun Huang and Wei-Ying Wu},
      year={2026},
      eprint={2606.06753},
      archivePrefix={arXiv},
      primaryClass={stat.ME},
      url={https://arxiv.org/abs/2606.06753},
}