preprocessing.py

"""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