jaluus/ParticleViT-L

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