Uncertainty-Aware Early Transient Classification Framework

A complete ML pipeline for classifying astronomical transients from the PLAsTiCC dataset, featuring:

• Four classifiers: XGBoost baseline, scratch Transformer, Astromer (pretrained), Moirai/Chronos (foundation model).
• Early truncation: simulate real-world early-alert classification at 10%, 30%, 50%, and 100% of each object's light curve.
• Uncertainty quantification: MC Dropout with 50 stochastic forward passes.
• Calibration evaluation: reliability diagrams + Expected Calibration Error (ECE).
• Prioritization simulation: Top-K rare-class (kilonova, TDE) recall curves as the primary paper figure.

---

Repository structure

plasticc_pipeline/
├── config.py            All paths and hyperparameters (single source of truth)
├── data_loader.py       Data loading, normalisation, truncation, DataLoaders
├── features.py          Hand-crafted features for XGBoost
├── models.py            All four model definitions
├── uncertainty.py       MC Dropout inference (model-agnostic)
├── calibration.py       Reliability diagrams and ECE
├── prioritization.py    Follow-up budget / Top-K recall simulation
├── train.py             Training loop (PyTorch models)
├── evaluate.py          Full evaluation across models × truncation fractions
├── main.py              Pipeline orchestrator with CLI
├── colab_setup.py       One-shot Colab environment bootstrap
├── requirements.txt
├── README.md
└── RUN_IN_COLAB.md

---

Quick start (local)

1. Clone and install

git clone https://github.com/YOUR_USERNAME/plasticc_pipeline.git
cd plasticc_pipeline
pip install -r requirements.txt

2. Download data

Download the PLAsTiCC dataset from Kaggle:

kaggle competitions download -c PLAsTiCC-2018
unzip PLAsTiCC-2018.zip -d data/

Required files:

• data/training_set.csv
• data/training_set_metadata.csv

3. Run the full pipeline

# All four models, all four truncation fractions
python main.py

# Single fraction (faster, for debugging)
python main.py --truncation 0.3

# Specific models only
python main.py --models transformer xgboost

# Fewer MC samples for quick testing
python main.py --mc-samples 10 --truncation 0.5

4. Find results

---

Models

All deep models use MC Dropout (p=0.1, 50 forward passes) for uncertainty.

---

Hyperparameters

All hyperparameters are centralised in config.py:

---

Reproducing paper figures

The primary result figure (topk_recall_curves.png) is generated automatically by evaluate.py / main.py. To regenerate it from saved MC inference results:

from prioritization import run_prioritization, plot_topk_curves
import pickle, config

with open("results/mc_results.pkl", "rb") as f:
    mc = pickle.load(f)

label_map = ...   # from data_loader.build_label_map
run_prioritization(mc, label_map)

---

Colab

See RUN_IN_COLAB.md for step-by-step Colab instructions.

---

Citation

If you use this code, please cite:

@misc{plasticc_pipeline,
  title  = {Uncertainty-Aware Early Transient Classification Framework},
  year   = {2025},
  url    = {https://github.com/YOUR_USERNAME/plasticc_pipeline}
}