"""Run the ViT-BL dRICH particle-ID model on a raw Geant4/EDM4hep simulation output file, end to end: DRICHHits.cellID -> resolved (x,y) -> rasterized image -> model -> predicted species. Requires: torch, uproot, awkward, numpy, huggingface_hub pip install torch uproot awkward numpy huggingface_hub Usage: python predict_from_root.py sim_2212_30.0_2.0_3.14.root The model architecture (RingViT) is copied inline below so this script has no dependency beyond pip-installable packages -- it does not require cloning the training repo. """ import re import sys import awkward as ak import numpy as np import torch import torch.nn as nn import uproot from huggingface_hub import hf_hub_download MODEL_REPO = "deepaksamuel-cuk/drich-vit-baseline" DATA_REPO = "deepaksamuel-cuk/simhits" NAMES = ["electron", "pion", "kaon", "proton"] PID2LABEL = {11: 0, 211: 1, 321: 2, 2212: 3} FULL_IMG = 384 FULL_WINDOW = 3540.0 # mm -- full detector extent, matches the real dRICH sensor plane MOM_SCALE = 60.0 ETA_MID, ETA_HALF = 2.5, 1.0 # ---------------------------------------------------------------- model --- class TransformerBlock(nn.Module): def __init__(self, dim, heads, mlp_ratio=4.0, dropout=0.0): super().__init__() self.norm1 = nn.LayerNorm(dim) self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True) self.norm2 = nn.LayerNorm(dim) hidden = int(dim * mlp_ratio) self.mlp = nn.Sequential( nn.Linear(dim, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, dim), nn.Dropout(dropout), ) def forward(self, x): h = self.norm1(x) x = x + self.attn(h, h, h, need_weights=False)[0] x = x + self.mlp(self.norm2(x)) return x class RingViT(nn.Module): """Same architecture as the training repo's vit/model.py -- 8-layer, 256-dim ViT with an extra learned 'kinematics token' (p, eta, cos phi, sin phi) alongside the cls token and 24x24=576 image patch tokens.""" def __init__(self, img_size=384, patch_size=16, in_chans=1, num_classes=4, dim=256, depth=8, heads=8, mlp_ratio=4.0, dropout=0.1, n_kin=4): super().__init__() self.n_patches = (img_size // patch_size) ** 2 self.patch_embed = nn.Conv2d(in_chans, dim, kernel_size=patch_size, stride=patch_size) self.cls_token = nn.Parameter(torch.zeros(1, 1, dim)) self.kin_embed = nn.Sequential(nn.Linear(n_kin, dim), nn.GELU(), nn.Linear(dim, dim)) self.pos_embed = nn.Parameter(torch.zeros(1, 2 + self.n_patches, dim)) self.pos_drop = nn.Dropout(dropout) self.blocks = nn.ModuleList( [TransformerBlock(dim, heads, mlp_ratio, dropout) for _ in range(depth)]) self.norm = nn.LayerNorm(dim) self.head = nn.Linear(dim, num_classes) def forward(self, img, kin): B = img.shape[0] x = self.patch_embed(img).flatten(2).transpose(1, 2) cls = self.cls_token.expand(B, -1, -1) k = self.kin_embed(kin).unsqueeze(1) x = torch.cat([cls, k, x], dim=1) + self.pos_embed x = self.pos_drop(x) for blk in self.blocks: x = blk(x) x = self.norm(x) return self.head(x[:, 0]) # ------------------------------------------------------------- rasterize -- def rasterize(pts, img_size=FULL_IMG, window=FULL_WINDOW): """pixel value = log1p(hit count) -- multiple hits in the same cell are summed (np.add.at), never overwritten or capped. Same convention as every image this model was trained on.""" img = np.zeros((img_size, img_size), dtype=np.float32) if len(pts) == 0: return img rel = (pts - (-window / 2)) * (img_size / window) ij = np.floor(rel).astype(np.int64) ok = (ij[:, 0] >= 0) & (ij[:, 0] < img_size) & (ij[:, 1] >= 0) & (ij[:, 1] < img_size) ij = ij[ok] np.add.at(img, (ij[:, 1], ij[:, 0]), 1.0) return np.log1p(img) def parse_kinematics_from_filename(path): """Sim files follow the {PID}_{MOM}_{ETA}_{PHI}[_suffix].root naming convention (see claude.md) -- pulls out the first four underscore- separated numeric tokens, tolerant of an optional trailing suffix.""" stem = path.split("/")[-1] stem = re.sub(r"\.root$", "", stem) stem = re.sub(r"^(sim|rec|ana)_", "", stem) m = re.match(r"^(-?\d+)_([\d.]+)_([\d.]+)_([\d.]+)", stem) if not m: raise ValueError(f"could not parse PID_MOM_ETA_PHI from filename: {path}") pid_s, mom_s, eta_s, phi_s = m.groups() return int(pid_s), float(mom_s), float(eta_s), float(phi_s) def main(): if len(sys.argv) != 2: print(f"usage: python {sys.argv[0]} ") sys.exit(1) root_path = sys.argv[1] pid, mom, eta, phi = parse_kinematics_from_filename(root_path) print(f"from filename: PID={pid}, p={mom:.2f} GeV/c, eta={eta:.2f}, phi={phi:.2f} rad") # -- extract raw photon hits from the sim file -- f = uproot.open(root_path) t = f["events"] cellid = ak.to_numpy(ak.flatten(t["DRICHHits/DRICHHits.cellID"].array())) print(f"{len(cellid)} raw photon hits") # -- resolve cellID -> (x, y) mm via the precomputed sensor geometry lookup -- lookup_path = hf_hub_download(DATA_REPO, "cellid_positions.npz", repo_type="dataset") lk = np.load(lookup_path) cid_sorted, xyz = lk["cellids"], lk["xyz"] idx = np.searchsorted(cid_sorted, cellid) idx = np.clip(idx, 0, len(cid_sorted) - 1) valid = cid_sorted[idx] == cellid if not valid.all(): print(f"warning: dropping {(~valid).sum()} hits with unknown cellID") pts = xyz[idx[valid], :2].astype(np.float32) img = rasterize(pts) kin = np.array([mom / MOM_SCALE, (eta - ETA_MID) / ETA_HALF, np.cos(phi), np.sin(phi)], dtype=np.float32) # -- load the model -- ckpt_path = hf_hub_download(MODEL_REPO, "best.pt") model = RingViT(num_classes=4) model.load_state_dict(torch.load(ckpt_path, map_location="cpu")) model.eval() timg = torch.from_numpy(img).unsqueeze(0).unsqueeze(0) # (1, 1, 384, 384) tkin = torch.from_numpy(kin).unsqueeze(0) # (1, 4) with torch.no_grad(): logits = model(timg, tkin) probs = torch.softmax(logits, dim=1)[0] pred = int(probs.argmax()) print(f"\npredicted: {NAMES[pred]} (p={probs[pred]:.3f})") for i, name in enumerate(NAMES): print(f" {name:10s} {probs[i]:.4f}") if pid in PID2LABEL: true_c = PID2LABEL[abs(pid)] verdict = "CORRECT" if pred == true_c else "WRONG" print(f"\ntrue (from filename): {NAMES[true_c]} -- {verdict}") if __name__ == "__main__": main()