giant/giant/data/transforms.py

import numpy as np

_EPS = 1e-8


def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
    return np.log(np.asarray(x, dtype=np.float32) + eps)


def inv_log_transform(y: np.ndarray, eps: float = _EPS) -> np.ndarray:
    return np.exp(np.asarray(y, dtype=np.float32)) - eps


def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarray:
    """Rotate post_dir into the local frame where pre_dir maps to ẑ (Rodrigues).

    Preserves the angle between pre_dir and post_dir; the result has post_dir
    expressed relative to a coordinate system in which the incoming particle
    travels along +z.
    """
    z = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)

    cos_t = np.clip((pre_dir * z).sum(axis=1, keepdims=True), -1.0, 1.0)  # (N,1)
    sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t ** 2))  # (N,1)

    axis = np.cross(pre_dir, z)  # (N,3); zero when pre_dir ∥ ẑ
    axis_norm = np.linalg.norm(axis, axis=1, keepdims=True)  # (N,1)

    # Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
    # are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
    safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
    axis = np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)

    kxv = np.cross(axis, post_dir)  # (N,3)
    kdv = (axis * post_dir).sum(axis=1, keepdims=True)  # (N,1)

    return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)


class Normalizer:
    def __init__(self) -> None:
        self.mean: np.ndarray | None = None
        self.std: np.ndarray | None = None

    def fit(self, X: np.ndarray) -> "Normalizer":
        self.mean = X.mean(axis=0).astype(np.float32)
        self.std = X.std(axis=0).astype(np.float32)
        self.std = np.where(self.std < _EPS, 1.0, self.std).astype(np.float32)
        return self

    def transform(self, X: np.ndarray) -> np.ndarray:
        return ((X - self.mean) / self.std).astype(np.float32)

    def inverse_transform(self, X: np.ndarray) -> np.ndarray:
        return (X * self.std + self.mean).astype(np.float32)

    def to_dict(self) -> dict:
        return {"mean": self.mean.tolist(), "std": self.std.tolist()}

    @classmethod
    def from_dict(cls, d: dict) -> "Normalizer":
        obj = cls()
        obj.mean = np.array(d["mean"], dtype=np.float32)
        obj.std = np.array(d["std"], dtype=np.float32)
        return obj


class _WelfordAccumulator:
    """Streaming mean/variance (Welford's online algorithm, batch update).

    Use to fit a Normalizer over data that doesn't fit in memory:
        acc = _WelfordAccumulator(n_features)
        for chunk in data:
            acc.update(chunk)
        normalizer = acc.to_normalizer()
    """

    def __init__(self, n_features: int) -> None:
        self.n = 0
        self._mean = np.zeros(n_features, dtype=np.float64)
        self._M2 = np.zeros(n_features, dtype=np.float64)

    def update(self, X: np.ndarray) -> None:
        X = np.asarray(X, dtype=np.float64)
        B = X.shape[0]
        new_n = self.n + B
        delta = X - self._mean
        self._mean += delta.sum(0) / new_n
        delta2 = X - self._mean
        self._M2 += (delta * delta2).sum(0)
        self.n = new_n

    def to_normalizer(self) -> "Normalizer":
        norm = Normalizer()
        norm.mean = self._mean.astype(np.float32)
        std = np.sqrt(self._M2 / max(self.n, 1)).astype(np.float32)
        norm.std = np.where(std < _EPS, 1.0, std).astype(np.float32)
        return norm


def inv_local_frame_rotation(pre_dir: np.ndarray, post_dir_local: np.ndarray) -> np.ndarray:
    """Inverse of local_frame_rotation: rotate from local frame back to world frame.

    Applies R^T (same axis, negative angle) to post_dir_local.
    """
    z = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)

    cos_t = np.clip((pre_dir * z).sum(axis=1, keepdims=True), -1.0, 1.0)
    sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t ** 2))

    axis = np.cross(pre_dir, z)
    axis_norm = np.linalg.norm(axis, axis=1, keepdims=True)
    safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
    axis = np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)

    kxv = np.cross(axis, post_dir_local)
    kdv = (axis * post_dir_local).sum(axis=1, keepdims=True)

    # Negative angle: sin_t → -sin_t
    return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)


def build_cond_features(
    data: dict[str, np.ndarray],
    pdg_map: dict[int, int],
    mat_map: dict[int, int],
    cond_normalizer: "Normalizer | None" = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Build conditioning arrays only — no target, no post-step variables."""
    cond_cont = np.column_stack([
        data["pre_pos"],
        log_transform(data["pre_energy"]),
        data["pre_dir"],
        data["layer_id"].astype(np.float32),
        data["n_sec"].astype(np.float32),
    ]).astype(np.float32)

    pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
    mat_idx = np.array([mat_map[int(m)] for m in data["material"]], dtype=np.int64)
    cond_cat = np.column_stack([pdg_idx, mat_idx])

    if cond_normalizer is not None:
        cond_cont = cond_normalizer.transform(cond_cont)

    return cond_cont, cond_cat


def build_features(
    data: dict[str, np.ndarray],
    pdg_map: dict[int, int],
    mat_map: dict[int, int],
    cond_normalizer: Normalizer | None = None,
    target_normalizer: Normalizer | None = None,
    fit: bool = False,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, Normalizer | None, Normalizer | None]:
    """Assemble (cond_cont, cond_cat, target) arrays ready for StepsDataset.

    When fit=True, new Normalizers are fitted on the supplied arrays.
    """
    post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])

    target = np.column_stack([
        log_transform(data["step_length"]),
        log_transform(data["delta_e"]),
        log_transform(data["edep"]),
        post_dir_local,
    ]).astype(np.float32)  # (N, 6)

    cond_cont = np.column_stack([
        data["pre_pos"],
        log_transform(data["pre_energy"]),
        data["pre_dir"],
        data["layer_id"].astype(np.float32),
        data["n_sec"].astype(np.float32),
    ]).astype(np.float32)  # (N, 9)

    pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
    mat_idx = np.array([mat_map[int(m)] for m in data["material"]], dtype=np.int64)
    cond_cat = np.column_stack([pdg_idx, mat_idx])  # (N, 2)

    if fit:
        cond_normalizer = Normalizer().fit(cond_cont)
        target_normalizer = Normalizer().fit(target)

    if cond_normalizer is not None:
        cond_cont = cond_normalizer.transform(cond_cont)
    if target_normalizer is not None:
        target = target_normalizer.transform(target)

    return cond_cont, cond_cat, target, cond_normalizer, target_normalizer