deepaksamuel-cuk/drich-vit-baseline

ViT-BL: dRICH particle-ID baseline (Vision Transformer)

A Vision Transformer (Dosovitskiy et al. 2020) trained from scratch (no ImageNet pretraining) to identify the species of a charged particle — electron / pion / kaon / proton — from the Cherenkov ring pattern it produces in the EPIC dRICH detector (dual radiator: aerogel + gas).

Training data, full provenance, and example-image viewers are published separately at deepaksamuel-cuk/simhits. No injected noise, no synthetic mixing, no QE hit loss, no pixel-gap simulation, no augmentation anywhere in training or evaluation — every training image is a real simulated photon-hit pattern.

Files

file	contents
best.pt	model weights (state_dict), the checkpoint with the best validation accuracy across training
log.json	full per-epoch training history (train/val loss, val accuracy, timing)

Architecture

• Patch embedding: Conv2d(1, 256, kernel_size=16, stride=16) on a 384x384 single-channel input image → 24x24 = 576 patch tokens
• An extra learned kinematics token: a small MLP embeds (p/60, (eta-2.5)/1.0, cos(phi), sin(phi)) into one 256-dim token, concatenated alongside the cls token and the 576 patch tokens (578 tokens total) — added because ring size alone is degenerate across species (a kaon and proton can share a ring radius at different momenta), so the model needs momentum/direction to disambiguate
• 8 pre-norm transformer blocks, dim=256, 8 heads, MLP ratio 4x, dropout 0.1 — standard self-attention (encoder-only, no decoder)
• Classification head: Linear(256, 4) on the final cls token output
• ~9M parameters

Training recipe

• Data: train_bl_ath (527,697 events — above each species' own aerogel Cherenkov threshold, n=1.026), full-detector-extent images (384x384 px, pixel = log1p(hit count))
• 100 epochs, batch 256, AdamW (lr 3e-4, weight decay 0.05), cosine schedule with 1-epoch warmup, label smoothing 0.05, gradient clip 1.0
• No augmentation at any point — train, val, and test all see only real, unmodified simulated hits

Performance (ViT-BL-ATH: tested on above-threshold events)

Confusion matrix diagonal (per-species efficiency), test_bl_ath (58,670 events):

species	efficiency
electron	93.9%
pion	85.5%
kaon	85.9%
proton	87.4%

Known failure mode: efficiency drops sharply at the forward edge of the generated acceptance (eta > ~3.3), where "kaon" acts as a systematic attractor class for electron/pion/proton misidentifications (kaon itself collapses toward "proton" instead) — see the companion analysis code in the training repo for details.

Inputs the model expects

Two tensors per event:

• img: (B, 1, 384, 384) float32 — log1p(hit count) per pixel, full detector extent (3540mm window, matching the dRICH sensor plane's physical size), no crop/centering
• kin: (B, 4) float32 — [momentum/60.0, (eta-2.5)/1.0, cos(phi), sin(phi)]

Output: (B, 4) logits, order [electron, pion, kaon, proton].

Quick start: random-input sanity check

import torch
from huggingface_hub import hf_hub_download
# RingViT class definition -- see predict_from_root.py in this repo, or
# the training repo's vit/model.py

ckpt_path = hf_hub_download("deepaksamuel-cuk/drich-vit-baseline", "best.pt")
model = RingViT(num_classes=4)
model.load_state_dict(torch.load(ckpt_path, map_location="cpu"))
model.eval()

img = torch.zeros(1, 1, 384, 384)   # replace with a real rasterized event
kin = torch.zeros(1, 4)              # replace with real [p, eta, phi] features
with torch.no_grad():
    probs = torch.softmax(model(img, kin), dim=1)
print(probs)  # [P(electron), P(pion), P(kaon), P(proton)]

Running inference on a raw simulation file

See predict_from_root.py in this repo for a complete, self-contained script (includes the model class inline, no need to clone the training repo) that takes a raw Geant4 simulation .root output file, extracts DRICHHits/DRICHHits.cellID, resolves it to physical hit positions via the sensor geometry lookup table (cellid_positions.npz, hosted in the simhits data repo), rasterizes it exactly as in training, and prints the predicted species with per-class probabilities:

pip install torch uproot awkward numpy huggingface_hub
python predict_from_root.py sim_2212_30.0_2.0_3.14.root

Verified end-to-end against a real local sim file before publishing this script (correctly predicted "pion" at 96% confidence on a true pion event not seen during training).

Related

• Training/eval data + provenance: deepaksamuel-cuk/simhits