Add pdg energy/length contribution pie plots

pdg_contribution_table_pl sums real/generated total deposited energy and
total step_length per pdg species over the whole file (pure lazy polars
group_by, no post_pos reconstruction needed for these scalars). Adds
plot_pdg_energy_share/plot_pdg_length_share, each rendering two pies
(real vs generated) so the per-species breakdown can be compared directly,
plus a matching notebook section.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

giant/analysis.py

@@ -44,6 +44,16 @@ corrupted by partial events::
     plot_transverse_profile(obs)
     plot_shower_max_depth(obs)
 
+For the dataset-wide breakdown of which particle species (pdg) contributed
+how much of the total energy/length, see `pdg_contribution_table_pl`::
+
+    from giant.analysis import pdg_contribution_table_pl
+    from giant.analysis import plot_pdg_energy_share, plot_pdg_length_share
+
+    table = pdg_contribution_table_pl("path/to/steps_predicted_local.parquet")
+    plot_pdg_energy_share(table)
+    plot_pdg_length_share(table)
+
 Four tiers of checks, building on the aggregate marginal/KL check in
 `giant.validate.validate_marginals`:
 
@@ -1342,3 +1352,114 @@ def plot_shower_max_depth(observables: EventObservables, bins: int = 30):
     ax.legend(fontsize=8)
     fig.tight_layout()
     return fig
+
+
+# ---------------------------------------------------------------------------
+# Particle-species (pdg) contribution shares
+#
+# Dataset-wide (not per-event) breakdown of which pdg species contributed how
+# much of the total deposited energy / total length traveled. Unlike the
+# event-level checks above, this only needs scalar sums — no post_pos
+# reconstruction, no shower axis — so it's a single lazy polars group_by.
+# ---------------------------------------------------------------------------
+
+_PDG_NAMES = {11: "e-", -11: "e+", 22: "gamma", 2112: "n", 2212: "p"}
+
+
+def _pdg_label(pdg: int) -> str:
+    if pdg in _PDG_NAMES:
+        return _PDG_NAMES[pdg]
+    if abs(pdg) > 1_000_000_000:
+        return f"ion{pdg}"
+    return str(pdg)
+
+
+def pdg_contribution_table_pl(source: str | Path | pl.LazyFrame) -> pl.DataFrame:
+    """Total edep / step_length contributed by each pdg species, real vs generated.
+
+    One row per pdg code, sorted by pdg. Pure lazy polars `group_by` over the
+    whole file — `edep`/`step_length` are scalars unaffected by the
+    local-frame rotation, so this only needs the same `exp(...) - eps` de-log
+    transform `inv_log_transform` does, expressed directly as a polars expr.
+    """
+    lf = _scan_predicted_local(source)
+
+    def _delog(col: str) -> pl.Expr:
+        return pl.col(col).exp() - _LOG_EPS
+
+    return (
+        lf.group_by("pdg")
+        .agg(
+            _delog("true_log_edep").sum().alias("real_total_edep"),
+            _delog("pred_log_edep").sum().alias("gen_total_edep"),
+            _delog("true_log_step_length").sum().alias("real_total_length"),
+            _delog("pred_log_step_length").sum().alias("gen_total_length"),
+        )
+        .collect()
+        .sort("pdg")
+    )
+
+
+def _pdg_pie_shares(
+    table: pl.DataFrame, real_col: str, gen_col: str, max_slices: int
+) -> tuple[list[str], np.ndarray, np.ndarray]:
+    """Pie-ready (labels, real_values, gen_values), lumping small contributors into 'other'.
+
+    Ranked by combined real+gen contribution so the same species end up in
+    the same slice position in both pies, making them easier to compare.
+    """
+    pdg = table["pdg"].to_numpy()
+    real = table[real_col].to_numpy()
+    gen = table[gen_col].to_numpy()
+
+    order = np.argsort(-(real + gen))
+    pdg, real, gen = pdg[order], real[order], gen[order]
+
+    if len(pdg) > max_slices:
+        keep = max_slices - 1
+        labels = [_pdg_label(int(p)) for p in pdg[:keep]] + ["other"]
+        real = np.append(real[:keep], real[keep:].sum())
+        gen = np.append(gen[:keep], gen[keep:].sum())
+    else:
+        labels = [_pdg_label(int(p)) for p in pdg]
+
+    return labels, real, gen
+
+
+def _plot_pdg_pie(
+    table: pl.DataFrame, real_col: str, gen_col: str, suptitle: str, max_slices: int
+):
+    labels, real_vals, gen_vals = _pdg_pie_shares(table, real_col, gen_col, max_slices)
+
+    fig, axes = plt.subplots(1, 2, figsize=(9, 4.5))
+    for ax, vals, title in [
+        (axes[0], real_vals, "real"),
+        (axes[1], gen_vals, "generated"),
+    ]:
+        ax.pie(vals, labels=labels, autopct="%1.1f%%", startangle=90)
+        ax.set_title(title)
+    fig.suptitle(suptitle)
+    fig.tight_layout()
+    return fig
+
+
+def plot_pdg_energy_share(table: pl.DataFrame, max_slices: int = 6):
+    """Real-vs-generated pies of total deposited energy share by pdg species."""
+    return _plot_pdg_pie(
+        table,
+        "real_total_edep",
+        "gen_total_edep",
+        "deposited energy share by particle type",
+        max_slices,
+    )
+
+
+def plot_pdg_length_share(table: pl.DataFrame, max_slices: int = 6):
+    """Real-vs-generated pies of total length-traveled share by pdg species."""
+    return _plot_pdg_pie(
+        table,
+        "real_total_length",
+        "gen_total_length",
+        "length traveled share by particle type",
+        max_slices,
+    )

tests/test_analysis.py

@@ -19,6 +19,7 @@ from giant.analysis import (
     load_predicted_local,
     marginal_table,
     marginal_table_pl,
+    pdg_contribution_table_pl,
     plot_constraint_violations,
     plot_correlation_matrices,
     plot_direction_alignment,
@@ -27,6 +28,8 @@ from giant.analysis import (
     plot_longitudinal_profile,
     plot_marginals,
     plot_pairwise,
+    plot_pdg_energy_share,
+    plot_pdg_length_share,
     plot_shower_max_depth,
     plot_total_energy,
     plot_total_length,
@@ -400,10 +403,11 @@ def _write_event_level_parquet(path, rng=None):
         _make_event_level_arrays(rng)
     )
     n = len(event_id)
+    pdg = rng.choice([11, -11, 22], n)
     table = pa.table(
         {
             "event_id": event_id,
-            "pdg": rng.choice([11, -11, 22], n),
+            "pdg": pdg,
             "pre_x": pre_pos[:, 0],
             "pre_y": pre_pos[:, 1],
             "pre_z": pre_pos[:, 2],
@@ -431,7 +435,7 @@ def _write_event_level_parquet(path, rng=None):
         }
     )
     pq.write_table(table, path)
-    return event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local
+    return event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local, pdg
 
 
 def _expected_event_table(
@@ -480,7 +484,7 @@ def _expected_event_table(
 
 def test_compute_event_observables_pl_matches_manual_reconstruction(tmp_path):
     path = tmp_path / "event_level.parquet"
-    event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local = (
+    event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local, _pdg = (
         _write_event_level_parquet(path)
     )
     expected = _expected_event_table(
@@ -566,3 +570,67 @@ def test_event_level_plots_run_without_error(tmp_path):
     assert plot_longitudinal_profile(obs) is not None
     assert plot_transverse_profile(obs) is not None
     assert plot_shower_max_depth(obs) is not None
+
+
+def test_pdg_contribution_table_pl_matches_manual_sums(tmp_path):
+    path = tmp_path / "event_level.parquet"
+    _, _, _, _, true_log_local, pred_log_local, pdg = _write_event_level_parquet(path)
+
+    real_edep = inv_log_transform(true_log_local[:, 2])
+    gen_edep = inv_log_transform(pred_log_local[:, 2])
+    real_length = inv_log_transform(true_log_local[:, 0])
+    gen_length = inv_log_transform(pred_log_local[:, 0])
+
+    expected = {}
+    for p in np.unique(pdg):
+        mask = pdg == p
+        expected[int(p)] = (
+            float(real_edep[mask].sum()),
+            float(gen_edep[mask].sum()),
+            float(real_length[mask].sum()),
+            float(gen_length[mask].sum()),
+        )
+
+    table = pdg_contribution_table_pl(path).sort("pdg")
+    for i, p in enumerate(table["pdg"].to_list()):
+        real_e, gen_e, real_l, gen_l = expected[int(p)]
+        np.testing.assert_allclose(table["real_total_edep"][i], real_e, rtol=1e-4)
+        np.testing.assert_allclose(table["gen_total_edep"][i], gen_e, rtol=1e-4)
+        np.testing.assert_allclose(table["real_total_length"][i], real_l, rtol=1e-4)
+        np.testing.assert_allclose(table["gen_total_length"][i], gen_l, rtol=1e-4)
+
+
+def test_pdg_contribution_table_pl_accepts_lazyframe(tmp_path):
+    path = tmp_path / "event_level.parquet"
+    _write_event_level_parquet(path)
+
+    from_path = pdg_contribution_table_pl(path).sort("pdg")
+    from_lf = pdg_contribution_table_pl(pl.scan_parquet(path)).sort("pdg")
+
+    np.testing.assert_allclose(
+        from_lf["real_total_edep"].to_numpy(), from_path["real_total_edep"].to_numpy()
+    )
+
+
+def test_pdg_pie_plots_run_without_error(tmp_path):
+    path = tmp_path / "event_level.parquet"
+    _write_event_level_parquet(path)
+    table = pdg_contribution_table_pl(path)
+
+    assert plot_pdg_energy_share(table) is not None
+    assert plot_pdg_length_share(table) is not None
+
+
+def test_plot_pdg_energy_share_caps_slices():
+    table = pl.DataFrame(
+        {
+            "pdg": list(range(10)),
+            "real_total_edep": [float(10 - i) for i in range(10)],
+            "gen_total_edep": [float(10 - i) for i in range(10)],
+            "real_total_length": [float(10 - i) for i in range(10)],
+            "gen_total_length": [float(10 - i) for i in range(10)],
+        }
+    )
+    fig = plot_pdg_energy_share(table, max_slices=4)
+    for ax in fig.axes:
+        assert len(ax.patches) == 4