giant/giant/analysis/sources.py

"""Canonical world-frame LazyFrame builders for the two sides of a comparison.

The analysis compares one autoregressive ``giant rollout`` (the *generated* side)
against a raw miniCaloSim steps file (the *reference* / real side). Both 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 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 inputs a file is."""

    rollout = "rollout"
    reference = "reference"


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