giant/giant/analysis/sources.py

"""Canonical world-frame LazyFrame builders for the two kinds of comparison input.

The analysis compares one or more autoregressive ``giant rollout`` runs (the
*generated* side — one named series each, see ``RolloutSpec``) against a single
raw miniCaloSim steps file shared by all of them (the *reference* / real side).
Every rollout is the same *kind* of file regardless of how many there are, so
``Side`` stays binary: it describes a file's schema (rollout column layout +
synthetic-termination rows + per-track secondary view, vs. reference
``sec_*_list`` columns), not series identity. Both kinds carry a
**shared world-frame physical column subset** under identical names, so no
renaming or coordinate decode is needed — everything is already in world-frame
mm / MeV:

    event_id, track_id, step_no, pdg,
    pre_x, pre_y, pre_z, pre_E, pre_dx, pre_dy, pre_dz,
    post_x, post_y, post_z, post_E, post_dx, post_dy, post_dz,
    edep, step_length, material, layer_id

(rollout: ``rollout.py:_RECORD_KEYS``; reference: minicalosim ``RunAction.cc``
Steps ntuple passed through by ``dwarf convert``.)

The two files differ in their *extra* columns — the rollout adds ``parent_id``,
``generation``, ``n_sec_pred``, ``termination_reason``; the reference adds
``process``, field columns, ``child_track_ids`` and the ``sec_*_list`` secondary
birth-state lists. Those are only touched by the side-specific helpers here
(synthetic-row filtering, the secondary view).

Nothing in this module (or ``reduce.py``) imports plotstyle — compute runs on
HTCondor workers that have no LaTeX toolchain.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from pathlib import Path

import polars as pl
import pyarrow.parquet as pq

from giant.constants import (
    PREDICT_COORD_METADATA_KEY,
    ROLLOUT_COORD_VALUE,
    TERM_ENERGY_CUTOFF,
    TERM_ESCAPED,
    TERM_MAX_STEPS,
    TERM_UNKNOWN_PDG,
)
from giant.data.loader import event_id_offset, find_parquet_files

# Helper column name for the per-shard offset join in open_side; dropped before
# the LazyFrame is returned, so it never leaks into a caller's schema.
_SOURCE_PATH_COL = "__source_path"

# The world-frame physical columns both sides share under identical names.
PHYS_COLS: tuple[str, ...] = (
    "event_id",
    "pdg",
    "pre_x",
    "pre_y",
    "pre_z",
    "pre_E",
    "pre_dx",
    "pre_dy",
    "pre_dz",
    "post_x",
    "post_y",
    "post_z",
    "post_E",
    "post_dx",
    "post_dy",
    "post_dz",
    "edep",
    "step_length",
    "material",
    "layer_id",
)

# Rollout rows written purely for bookkeeping (a track's forced stop): they carry
# step_length=0, post_pos=pre_pos, and — for every reason but escaped — the
# track's whole remaining pre_E dumped into edep so the shower still conserves
# energy. They are not physical steps (the reference has no equivalent), so a
# per-step marginal comparison must drop them; a per-event energy total must keep
# them. See rollout.py's terminal-row handling.
SYNTHETIC_TERMINATION_REASONS: frozenset[str] = frozenset(
    {TERM_ESCAPED, TERM_UNKNOWN_PDG, TERM_ENERGY_CUTOFF, TERM_MAX_STEPS}
)


class Side(str, Enum):
    """Which of the two comparison-input *kinds* a file is."""

    rollout = "rollout"
    reference = "reference"


@dataclass
class RolloutSpec:
    """One named rollout input, as fed to ``build_context``/``Bundle.open``.

    ``name`` is the series' identity throughout the rest of the pipeline (a
    plot's ``payload["series"]`` key, a figure's legend label, its color) —
    resolved once in ``condor.load_rollout_yamls`` from ``--label`` or the
    YAML stem, then threaded through unchanged. ``checkpoint`` /
    ``type_embedding_l1_dist`` are only used by the router/type-embedding
    diagnostics (``catalog.py``'s ``chunkable=False`` specs).
    """

    name: str
    source: str | Path | pl.LazyFrame
    checkpoint: str | None = None
    type_embedding_l1_dist: dict | None = None


@dataclass
class RolloutSide:
    """One rollout's opened frames + per-checkpoint diagnostic inputs (``catalog.Bundle.rollouts`` value)."""

    all: pl.LazyFrame  # rollout, all rows (incl. synthetic termination rows)
    phys: pl.LazyFrame  # rollout, physical steps only
    checkpoint: str | None = None  # from the rollout YAML; router_gating only
    # Diagnostic pre-aggregated at rollout time (giant.rollout.
    # L1DistCollector.summary()) — from the rollout YAML, type_embedding_l1_distance
    # only. Unlike checkpoint, this needs no live model: it's already a
    # finished histogram, just passed through.
    type_embedding_l1_dist: dict | None = None


def _check_rollout_metadata(path: Path) -> None:
    """Raise if ``path`` carries coord metadata that isn't the rollout tag.

    A missing tag (older rollout output, predating tagging) is allowed through,
    matching ``giant rollout``'s own leniency; a tag that is present but wrong is
    a real mismatch and worth failing on before the column layout is trusted.
    """
    metadata = pq.read_schema(path).metadata or {}
    coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
    if coord is not None and coord.decode() != ROLLOUT_COORD_VALUE:
        raise ValueError(f"{path} is not a rollout file (coord={coord.decode()!r}); expected `giant rollout` output")


def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
    """Lazily scan one side's file, verifying the rollout tag when applicable.

    Returns the *full* lazy scan (no column projection) so downstream reductions
    can push their own narrow projection into the parquet read — the single
    biggest lever on a larger-than-RAM file. ``pl.LazyFrame`` inputs pass straight
    through (used by tests).

    ``pdg`` is cast to a canonical ``Int64`` here: the rollout writer and the
    reference file's upstream ROOT→parquet conversion don't agree on integer
    width, and an uncast mismatch only surfaces later as a ``pl.concat``
    ``SchemaError`` (e.g. in ``build_context``'s pdg-count merge).

    The reference (a rollout's seed ``dataset``) may be a directory of parquet
    shards, or a ``.manifest`` naming a subset, rather than a single file — each
    such shard is a separate Geant4 job whose own ``event_id`` numbering
    restarts from 0, so a multi-shard load offsets every shard's ids by
    ``giant.data.loader.event_id_offset(file_index)`` to keep them globally
    unique, exactly as the training/rollout data pipeline already does
    (``giant/data/loader.py``). ``file_index`` comes from
    ``find_parquet_files``'s deterministic ordering — the same list and
    ordering ``giant rollout`` used (via ``_seed_from_data``) to offset the
    rollout side's own ``event_id``s, so both sides agree on what an
    ``event_id`` means. There is no overflow guard here (unlike
    ``loader._offset_event_id``): checking it would cost an eager
    ``event_id``-column read per shard in every condor compute job, and
    ``giant rollout`` already ran that check over this exact file list when it
    produced the seed.
    """
    if isinstance(source, pl.LazyFrame):
        return source.with_columns(pl.col("pdg").cast(pl.Int64))
    path = Path(source)
    if side is Side.rollout:
        _check_rollout_metadata(path)
        lf = pl.scan_parquet(path)
    else:
        files = find_parquet_files(path)
        if len(files) == 1:
            lf = pl.scan_parquet(files[0])
        else:
            offsets = {str(p): event_id_offset(i) for i, p in enumerate(files)}
            lf = (
                pl.scan_parquet(files, include_file_paths=_SOURCE_PATH_COL)
                .with_columns(
                    pl.col("event_id") + pl.col(_SOURCE_PATH_COL).replace_strict(offsets, return_dtype=pl.Int64)
                )
                .drop(_SOURCE_PATH_COL)
            )
    return lf.with_columns(pl.col("pdg").cast(pl.Int64))


def physical_steps(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
    """Real, physical steps only — drops the rollout's synthetic termination rows.

    The predicate is pushed down so the dropped rows are never decoded. The
    reference has no such rows, so it is returned unchanged.
    """
    if side is Side.reference:
        return lf
    return lf.filter(~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS)))


def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
    """Per-secondary birth state, one row per produced secondary.

    Canonical columns: ``event_id, energy, pdg, sdx, sdy, sdz`` (birth energy in
    MeV, PDG code, birth unit direction in the world frame). The two sides encode
    secondaries differently:

    - rollout: each secondary is its own track, so its birth state is the row with
      ``generation > 0`` and ``step_no == 0`` (``pre_E`` / ``pre_dir`` there).
    - reference: secondaries live in per-parent-step ``sec_*_list`` columns; the
      lists are exploded together and empty (no-secondary) steps drop out.
    """
    if side is Side.rollout:
        return lf.filter((pl.col("generation") > 0) & (pl.col("step_no") == 0)).select(
            "event_id",
            pl.col("pre_E").alias("energy"),
            "pdg",
            pl.col("pre_dx").alias("sdx"),
            pl.col("pre_dy").alias("sdy"),
            pl.col("pre_dz").alias("sdz"),
        )
    lists = ["sec_E_list", "sec_pdg_list", "sec_dx_list", "sec_dy_list", "sec_dz_list"]
    return (
        lf.select("event_id", *lists)
        .explode(lists)
        .drop_nulls("sec_E_list")
        .select(
            "event_id",
            pl.col("sec_E_list").alias("energy"),
            pl.col("sec_pdg_list").alias("pdg"),
            pl.col("sec_dx_list").alias("sdx"),
            pl.col("sec_dy_list").alias("sdy"),
            pl.col("sec_dz_list").alias("sdz"),
        )
    )


def secondaries_by_step(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
    """One row per produced secondary, tagged with the step that produced it.

    Canonical columns: ``step_key`` (an opaque struct identifying the emitting
    step) and ``pdg``. ``secondaries`` deliberately drops that link; the
    per-step multiplicity plots need it, so this is a separate view rather than
    extra columns every other consumer would pay for.

    - rollout: a secondary's birth row carries ``parent_id`` and a birth
      position copied verbatim from the parent step's ``post_pos``, so
      ``(event_id, parent_id, pre_pos)`` identifies the emitting step exactly —
      no join against the (large) step frame is needed.
    - reference: secondaries already live on their parent step's row, so the
      row index *is* the step key. It is only ever used as a group key inside
      one chunk's own aggregation, so indices repeating across chunks is
      harmless.
    """
    if side is Side.rollout:
        return lf.filter((pl.col("generation") > 0) & (pl.col("step_no") == 0)).select(
            pl.struct("event_id", "parent_id", "pre_x", "pre_y", "pre_z").alias("step_key"),
            "pdg",
        )
    return (
        lf.select("sec_pdg_list")
        .with_row_index("_row")
        .explode("sec_pdg_list")
        .drop_nulls("sec_pdg_list")
        .select(pl.struct("_row").alias("step_key"), pl.col("sec_pdg_list").cast(pl.Int64).alias("pdg"))
    )