Metadata-Version: 2.4
Name: da-stdk
Version: 1.0.1
Summary: DA-STDK: Data-adaptive spatio-temporal distributional prediction (cluster-adaptive bases, conformal calibration)
Author: Wen-Ting Wang
Author-email: egpivo@gmail.com
Requires-Python: >=3.10.13,<4.0.0
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Requires-Dist: joblib (>=1.3.0,<2.0.0)
Requires-Dist: k-means-constrained (>=0.7.3,<0.8.0)
Requires-Dist: matplotlib (>=3.7.0,<4.0.0)
Requires-Dist: numpy (>=1.24.0,<2.0.0)
Requires-Dist: pandas (>=2.0.0,<3.0.0)
Requires-Dist: pyyaml (>=6.0.1,<7.0.0)
Requires-Dist: scikit-learn (>=1.3.0,<2.0.0)
Requires-Dist: scipy (>=1.12.0,<2.0.0)
Requires-Dist: seaborn (>=0.13.0,<0.14.0)
Requires-Dist: torch (>=2.0.0,<3.0.0)
Requires-Dist: tqdm (>=4.65.0,<5.0.0)
Description-Content-Type: text/markdown

# DA-STDK

[![PyPI](https://img.shields.io/pypi/v/da-stdk)](https://pypi.org/project/da-stdk/)
[![Python](https://img.shields.io/pypi/pyversions/da-stdk)](https://pypi.org/project/da-stdk/)
[![Tests](https://github.com/STLABTW/da-stdk/actions/workflows/test.yaml/badge.svg)](https://github.com/STLABTW/da-stdk/actions/workflows/test.yaml)
[![Code Quality](https://github.com/STLABTW/da-stdk/actions/workflows/code-quality.yaml/badge.svg)](https://github.com/STLABTW/da-stdk/actions/workflows/code-quality.yaml)

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

![Model backbone](https://raw.githubusercontent.com/STLABTW/da-stdk/main/artifacts/backbone.png)

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](https://python-poetry.org/) are enough for most use:

```bash
poetry install --with dev
```

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

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

```python
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**

```bash
poetry run python scripts/train_default.py
```

**KAUST benchmark (multiple scenarios / models)**

```bash
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`](scripts/README.md).

## Layout

| Path | Contents |
|------|----------|
| `da_stdk/` | Models, training, data I/O, conformal utils, viz |
| `scripts/` | Training and experiment drivers |
| `configs/` | YAML configs |
| `data/` | KAUST CSVs (large; not on PyPI) |

## Dev

```bash
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 (link forthcoming)

```bibtex
@misc{kim2026cluster,
  title        = {Cluster-Aware Conformal Calibration for Spatio-Temporal Distributional Prediction},
  author       = {Kim, Gooyoung and Lim, Chae Young and Wang, Wen-Ting and Huang, Hao-Yun and Wu, Wei-Ying},
  year         = {2026},
  note         = {arXiv preprint, forthcoming},
}
```

When the arXiv entry is available, add `eprint`, `archivePrefix`, and `primaryClass` (or `url`) to the BibTeX above.

