File size: 4,793 Bytes
aa05499 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """Self-contained input preprocessing for ParticleViT (PyTorch only).
The model was trained on inputs passed through a frozen, parametric
per-feature transform that maps each of the four continuous kinematic features
(delta eta, delta phi, log pT, log E) to an approximately standard-normal
distribution. The transform constants live in
`omnilearned_parametric_normalization.json` and MUST be applied at inference;
feeding raw features yields meaningless predictions.
Feature layout per particle (9 channels), matching the OmniLearned corpus:
0:4 continuous kinematics (delta eta, delta phi, log pT, log E) -> normalized
4 categorical particle-ID code (dense integer id) -> passthrough
5:9 continuous vertex / tracking features -> passthrough
A particle slot is "real" iff its log pT channel (index 2) is non-zero; padded
slots are all-zero. Normalization is applied only to real particles.
"""
from __future__ import annotations
import json
import math
from pathlib import Path
import torch
PAD_FEATURE_IDX = 2 # log pT; zero for padded slots
def build_attn_mask(X_raw: torch.Tensor) -> torch.Tensor:
"""Real-particle mask (B, L) from raw, un-normalized inputs."""
return X_raw[:, :, PAD_FEATURE_IDX] != 0
def _normal_icdf(probs: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
clipped = probs.clamp(eps, 1.0 - eps)
return math.sqrt(2.0) * torch.erfinv(2.0 * clipped - 1.0)
def _laplace_cdf(values: torch.Tensor, loc: float, scale: float) -> torch.Tensor:
centered = values - loc
return torch.where(
centered < 0.0,
0.5 * torch.exp(centered / scale),
1.0 - 0.5 * torch.exp(-centered / scale),
)
def _yeo_johnson(values: torch.Tensor, lmbda: float) -> torch.Tensor:
positive = values >= 0.0
if abs(lmbda) < 1e-8:
pos = torch.log1p(values)
else:
pos = (torch.pow(values + 1.0, lmbda) - 1.0) / lmbda
if abs(lmbda - 2.0) < 1e-8:
neg = -torch.log1p(-values)
else:
neg = -(torch.pow(1.0 - values, 2.0 - lmbda) - 1.0) / (2.0 - lmbda)
return torch.where(positive, pos, neg)
def _transform_feature(values: torch.Tensor, params: dict) -> torch.Tensor:
transform = str(params["transform"])
if transform == "laplace_mixture_cdf_to_normal":
w = float(params["weight"])
probs = w * _laplace_cdf(values, float(params["loc"]), float(params["core_scale"]))
probs = probs + (1.0 - w) * _laplace_cdf(
values, float(params["loc"]), float(params["tail_scale"])
)
return _normal_icdf(probs, eps=1e-4)
if transform == "symmetric_halfnormal_mixture_angle_cdf_to_normal":
centered = values - float(params["loc"])
abs_centered = torch.abs(centered)[:, None]
scales = torch.tensor(params["scales"], dtype=values.dtype, device=values.device)
weights = torch.tensor(params["weights"], dtype=values.dtype, device=values.device)
abs_cdf = torch.sum(weights * torch.erf(abs_centered / (scales * math.sqrt(2.0))), dim=1)
probs = torch.where(centered >= 0.0, 0.5 + 0.5 * abs_cdf, 0.5 - 0.5 * abs_cdf)
return _normal_icdf(probs)
if transform == "yeo_johnson_standardized":
t = _yeo_johnson(values, float(params["lambda"]))
return (t - float(params["mean"])) / float(params["std"])
raise ValueError(f"Unsupported normalization transform: {transform}")
def load_normalization(path: str | Path) -> list[dict]:
"""Load the per-feature normalization parameters from the JSON file."""
with Path(path).open() as f:
return json.load(f)["features"]
def normalize(
X_raw: torch.Tensor,
normalization: str | Path | list[dict],
attn_mask: torch.Tensor | None = None,
) -> torch.Tensor:
"""Apply the frozen parametric normalization to a raw input batch.
Args:
X_raw: (B, L, 9) raw particle features (OmniLearned units).
normalization: path to omnilearned_parametric_normalization.json, or the
loaded list of per-feature params.
attn_mask: optional (B, L) real-particle mask; if None it is derived
from the log pT channel of X_raw.
Returns:
(B, L, 9) tensor with features 0:4 normalized; other channels untouched.
"""
params = load_normalization(normalization) if not isinstance(normalization, list) else normalization
if attn_mask is None:
attn_mask = build_attn_mask(X_raw)
attn_mask = attn_mask.bool()
out = X_raw.float().clone()
for feat in params:
idx = int(feat["feature_idx"])
values = out[:, :, idx]
values[attn_mask] = _transform_feature(values[attn_mask], feat)
out[:, :, idx] = values
return out
|