Neerav-Gupta/tokamark-robustness-data

tokamark-robustness-data/
├── data/
│   ├── train_X_feat.npy          # XGBoost feature vectors, train
│   ├── train_X_ts.npy            # Time series tensors, train
│   ├── train_y.npy               # Target labels, train
│   ├── val_X_feat.npy            # XGBoost feature vectors, val
│   ├── val_X_ts.npy              # Time series tensors, val
│   ├── val_y.npy                 # Target labels, val
│   ├── test_X_feat.npy           # XGBoost feature vectors, test
│   ├── test_X_ts.npy             # Time series tensors, test
│   ├── test_y.npy                # Target labels, test
│   ├── test_raw_samples.pkl      # Raw test samples with t_cut timestamps
│   └── feature_names.json        # Feature names for X_feat columns
├── checkpoints/
│   ├── xgboost_clean.pkl         # Trained XGBoost model
│   ├── lstm_clean.pt             # Trained LSTM model
│   ├── transformer_clean.pt      # Trained Transformer model
│   └── cnn_clean.pt              # Trained CNN baseline model
└── results/
    ├── xgboost_results.json
    ├── lstm_results.json
    ├── transformer_results.json
    ├── cnn_results.json
    ├── shot_level_metrics.json
    ├── alarm_under_corruption.json
    └── alarm_mitigation_proximate.json

import numpy as np
import pickle
from huggingface_hub import snapshot_download

# Download everything
snapshot_download(
    repo_id="Neerav-Gupta/tokamark-robustness-data",
    repo_type="dataset",
    local_dir="./tokamark_robustness_data"
)

# Load arrays
X_train_ts = np.load("./tokamark_robustness_data/data/train_X_ts.npy")
y_train    = np.load("./tokamark_robustness_data/data/train_y.npy")
X_test_ts  = np.load("./tokamark_robustness_data/data/test_X_ts.npy")
y_test     = np.load("./tokamark_robustness_data/data/test_y.npy")

print(f"Train: {X_train_ts.shape}")  # (9950, 600, 18)
print(f"Test:  {X_test_ts.shape}")   # (2420, 600, 18)

# Load raw test samples (includes t_cut disruption timestamps)
with open("./tokamark_robustness_data/data/test_raw_samples.pkl", "rb") as f:
    test_samples = pickle.load(f)

# Each sample has: shot_id, window_index, input, actuator, output, t_cut
print(f"Sample keys: {list(test_samples[0].keys())}")
print(f"t_cut (disruption time): {test_samples[0]['t_cut']:.4f}s")

# Load trained LSTM checkpoint
import torch
ckpt = torch.load(
    "./tokamark_robustness_data/checkpoints/lstm_clean.pt",
    map_location="cpu"
)
print(f"LSTM n_features: {ckpt['n_features']}")

# Load trained CNN checkpoint
ckpt_cnn = torch.load(
    "./tokamark_robustness_data/checkpoints/cnn_clean.pt",
    map_location="cpu"
)
print(f"CNN n_channels: {ckpt_cnn['n_channels']}, "
      f"input_len: {ckpt_cnn['input_len']}, "
      f"backbone_hidden: {ckpt_cnn['backbone_hidden']}")

# Clone the code repository
git clone https://github.com/Neerav-Gupta/tokamark-robustness.git
cd tokamark-robustness

# Download this dataset
python -c "
from huggingface_hub import snapshot_download
snapshot_download(
    repo_id='Neerav-Gupta/tokamark-robustness-data',
    repo_type='dataset',
    local_dir='fusion_research/data'
)
"

# Train all four models and run robustness evaluation
python scripts/train_xgboost.py
python scripts/train_lstm.py
python scripts/train_transformer.py
python scripts/train_cnn_baseline.py

# Compute shot-level alarm metrics
python scripts/compute_alarm_metrics.py

# Generate all 9 figures
python scripts/analyze_results.py

@misc{gupta2026tokamark_robustness,
  title         = {Benchmarking Sensor Robustness in Plasma Diagnostic
                   Models: A Systematic Evaluation on TokaMark},
  author        = {Gupta, Neerav},
  year          = {2026},
  eprint        = {2607.11915},
  archivePrefix = {arXiv},
  primaryClass  = {physics.plasm-ph}
}

@article{rousseau2026tokamark,
  title   = {TokaMark: A Comprehensive Benchmark for MAST Tokamak
             Plasma Models},
  author  = {Rousseau, C{\'e}cile and Jackson, Samuel and
             Ordonez-Hurtado, Rodrigo H. and Amorisco, Nicola C. and
             Boschi, Tobia and Holt, George K. and Loreti, Andrea and
             Sz{\'e}kely, Eszter and Whittle, Alexander and
             Agnello, Adriano and Pamela, Stanislas and
             Pascale, Alessandra and Akers, Robert and
             Bernabe Moreno, Juan and Thorne, Sue and
             Zayats, Mykhaylo},
  journal = {arXiv preprint arXiv:2602.10132},
  year    = {2026}
}