lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Reimplement rollout-vs-truth comparison on the streaming analysis module

Browse Source 
Merged in every non-analysis change from the MoE-prototype branch (routing,
training, data pipeline, streaming rollout output), keeping this branch's
lean streaming giant/analysis.py and rebuilding the rollout-vs-truth feature
natively on it instead of resurrecting the old numpy SampleCollection path.

- Add RolloutVsTruth, accepted anywhere Tier 1-3 functions take a predict-parquet
  source: decodes a giant rollout file and a held-out truth file into
  RAW_TARGET_NAMES space via a polars port of the forward local-frame rotation,
  fully streaming (no SampleCollection, no eager materialization).
- Add compute_rollout_vs_truth_observables_pl for Tier 4, reusing
  EventObservables (now backed by independent real_table/gen_table to support
  unequal rollout/truth event counts) so every existing shower-observable plot
  function works unchanged for both one-step and full-rollout comparisons.
- Update analysis/rollout_validation.ipynb to the new API and CLAUDE.md's
  architecture description; add test coverage for the new source type.
- Fix a pre-existing return-type mismatch in giant.rollout.rollout() (found by
  `ty check`): the on_chunk summary-dict branch didn't match the declared
  dict[str, np.ndarray] return type, now expressed as a RolloutSummary TypedDict.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-17 12:32:30 +02:00
parent 6744d05948
commit 81ec225b86
5 changed files with 1423 additions and 544 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+1 -1
View File
@@ -50,7 +50,7 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
**Samplers** (`giant/sample.py`): DDPM, DDIM, and flow matching (ODE integration, ~10 steps). Flow matching is the primary mode.
**Validation** (`giant/validate.py`): step-level marginal comparisons. Shower-level (rollout) observables live in `giant/analysis.py` (`compute_rollout_observables` + `plot_rollout_*`), fed by `giant rollout` output.
**Validation** (`giant/validate.py`): step-level marginal comparisons. `giant/analysis.py` is a fully-streaming (lazy polars) diagnostics module, sized for predict/rollout files larger than RAM, with no in-memory `SampleCollection` and no full-array materialization. It covers one-step-ahead `giant predict --coord local` output (`compute_event_observables_pl` + `plot_total_energy`/`plot_longitudinal_profile`/etc. for shower-level observables, plus the marginal/correlation/constraint tiers) and, via the `RolloutVsTruth` source type, a full autoregressive `giant rollout` shower compared against held-out truth data (`compute_rollout_vs_truth_observables_pl` for shower-level observables, reusing the same plot functions) — see the module docstring.
**Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
analysis/rollout_validation.ipynb
+123 -257
View File
File diff suppressed because one or more lines are too long
giant/analysis.py
+846 -269
View File
File diff suppressed because it is too large Load Diff
giant/rollout.py
+9 -2
View File
@@ -17,7 +17,7 @@ treated as detector leakage and not deposited.
from __future__ import annotations
from collections import Counter
from typing import Callable
from typing import Callable, TypedDict
import numpy as np
import torch
@@ -127,6 +127,13 @@ _RECORD_DTYPES: dict[str, type] = {
}
class RolloutSummary(TypedDict):
"""`rollout()`'s return shape when streaming to `on_chunk` instead of materializing rows."""
n_rows: int
termination_reason_counts: dict[str, int]
class _Recorder:
"""Accumulates per-step rows into column lists, materialised at the end —
or, when `sink` is given, streams each non-empty chunk to it immediately
@@ -276,7 +283,7 @@ def rollout(
max_tracks_per_event: int | None = None,
escape_threshold: float | None = None,
on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
) -> dict[str, np.ndarray]:
) -> dict[str, np.ndarray] | RolloutSummary:
"""Run showers to completion.
By default, returns a step-record dict (see _RECORD_KEYS) with the whole
tests/test_analysis.py
+444 -15
View File
@@ -11,7 +11,9 @@ import pytest
from giant.analysis import (
RAW_TARGET_NAMES,
RolloutVsTruth,
compute_event_observables_pl,
compute_rollout_vs_truth_observables_pl,
constraint_report_pl,
correlation_matrices_pl,
marginal_table_pl,
@@ -22,6 +24,8 @@ from giant.analysis import (
plot_kl_bars_pl,
plot_longitudinal_profile,
plot_marginals,
plot_mean_energy_per_step,
plot_mean_length_per_step,
plot_pairwise,
plot_pdg_energy_share,
plot_pdg_length_share,
@@ -35,12 +39,17 @@ from giant.constants import (
PREDICT_COORD_METADATA_KEY,
PREDICT_SCHEMA_VERSION,
PREDICT_SCHEMA_VERSION_KEY,
ROLLOUT_COORD_VALUE,
TERM_ENERGY_CUTOFF,
TERM_NATURAL_END,
)
from giant.data.transforms import (
energy_simplex_decode,
inv_log_transform,
local_frame_rotation,
log_transform,
reconstruct_post_pos,
travel_direction,
)
# ---------------------------------------------------------------------------
@@ -541,9 +550,10 @@ def test_compute_event_observables_pl_matches_manual_reconstruction(tmp_path):
)
obs = compute_event_observables_pl(path, depth_bins=5, transverse_bins=5)
table = obs.event_table.sort("event_id")
real_table = obs.real_table.sort("event_id")
gen_table = obs.gen_table.sort("event_id")
for i, eid in enumerate(table["event_id"].to_list()):
for i, eid in enumerate(real_table["event_id"].to_list()):
(
real_total_edep,
gen_total_edep,
@@ -555,28 +565,28 @@ def test_compute_event_observables_pl_matches_manual_reconstruction(tmp_path):
gen_rms,
) = expected[eid]
np.testing.assert_allclose(
table["real_total_edep"][i], real_total_edep, rtol=1e-4
real_table["total_edep"][i], real_total_edep, rtol=1e-4
)
np.testing.assert_allclose(
table["gen_total_edep"][i], gen_total_edep, rtol=1e-4
gen_table["total_edep"][i], gen_total_edep, rtol=1e-4
)
np.testing.assert_allclose(
table["real_total_length"][i], real_total_length, rtol=1e-4
real_table["total_length"][i], real_total_length, rtol=1e-4
)
np.testing.assert_allclose(
table["gen_total_length"][i], gen_total_length, rtol=1e-4
gen_table["total_length"][i], gen_total_length, rtol=1e-4
)
np.testing.assert_allclose(
table["real_centroid_depth"][i], real_centroid, rtol=1e-3, atol=1e-4
real_table["centroid_depth"][i], real_centroid, rtol=1e-3, atol=1e-4
)
np.testing.assert_allclose(
table["gen_centroid_depth"][i], gen_centroid, rtol=1e-3, atol=1e-4
gen_table["centroid_depth"][i], gen_centroid, rtol=1e-3, atol=1e-4
)
np.testing.assert_allclose(
table["real_transverse_rms"][i], real_rms, rtol=1e-3, atol=1e-4
real_table["transverse_rms"][i], real_rms, rtol=1e-3, atol=1e-4
)
np.testing.assert_allclose(
table["gen_transverse_rms"][i], gen_rms, rtol=1e-3, atol=1e-4
gen_table["transverse_rms"][i], gen_rms, rtol=1e-3, atol=1e-4
)
@@ -591,7 +601,8 @@ def test_compute_event_observables_pl_profile_shapes(tmp_path):
assert obs.gen_depth_profile.shape == (7,)
assert obs.real_transverse_profile.shape == (4,)
assert obs.gen_transverse_profile.shape == (4,)
assert len(obs.event_table) == 3
assert len(obs.real_table) == 3
assert len(obs.gen_table) == 3
def test_compute_event_observables_pl_accepts_lazyframe(tmp_path):
@@ -604,8 +615,8 @@ def test_compute_event_observables_pl_accepts_lazyframe(tmp_path):
)
np.testing.assert_allclose(
obs_from_lf.event_table.sort("event_id")["real_total_edep"].to_numpy(),
obs_from_path.event_table.sort("event_id")["real_total_edep"].to_numpy(),
obs_from_lf.real_table.sort("event_id")["total_edep"].to_numpy(),
obs_from_path.real_table.sort("event_id")["total_edep"].to_numpy(),
)
@@ -655,8 +666,8 @@ def test_compute_event_observables_pl_approx_median(tmp_path):
exact_length_median = np.median(true_length)
obs = compute_event_observables_pl(path)
approx_edep = obs.event_table["real_median_edep"][0]
approx_length = obs.event_table["real_median_length"][0]
approx_edep = obs.real_table["median_edep"][0]
approx_length = obs.real_table["median_length"][0]
# log-bin interpolation: expect within a few percent of the true median.
np.testing.assert_allclose(approx_edep, exact_edep_median, rtol=0.05)
@@ -744,3 +755,421 @@ def test_plot_pdg_energy_share_caps_slices():
fig = plot_pdg_energy_share(table, max_slices=4)
for ax in fig.axes:
assert len(ax.patches) == 4
# ---------------------------------------------------------------------------
# Rollout vs. held-out truth: unpaired `RolloutVsTruth` source (Tier 1-3) and
# `compute_rollout_vs_truth_observables_pl` (Tier 4)
# ---------------------------------------------------------------------------
_WORLD_FRAME_COLS = [
"event_id",
"pdg",
"material",
"layer_id",
"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",
]
def _write_world_frame_parquet(
path, n=200, rng=None, rollout=True, termination_reasons=None
):
"""A `giant rollout` output or truth-schema steps file, in raw world-frame units.
`rollout=True` adds `track_id`/`termination_reason`/`n_sec_pred` (rollout-
only columns `RolloutVsTruth` reads) and tags `ROLLOUT_COORD_VALUE`
metadata; `rollout=False` is a plain truth-schema file (no metadata tag,
matching a real held-out/val parquet). `termination_reasons`, if given,
overrides the (rollout-only) `termination_reason` column — pass a fixed
array to test synthetic-row dropping deterministically.
"""
rng = rng or np.random.default_rng(0)
event_id = rng.integers(0, 5, n)
pre_pos = rng.standard_normal((n, 3)).astype(np.float32) * 5
pre_dir = _unit_vectors(rng, n)
pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
post_dir = _unit_vectors(rng, n)
step_length = rng.uniform(0.1, 5.0, n).astype(np.float32)
travel_dir_world = _unit_vectors(rng, n)
post_pos = pre_pos + step_length[:, None] * travel_dir_world
post_E = (pre_E * rng.uniform(0.3, 0.99, n)).astype(np.float32)
edep = rng.uniform(0, 1, n).astype(np.float32)
columns = {
"event_id": event_id,
"pdg": rng.choice([11, -11, 22], n),
"material": rng.choice(["W", "Pb"], n),
"layer_id": rng.integers(0, 10, n).astype(np.int32),
"pre_x": pre_pos[:, 0],
"pre_y": pre_pos[:, 1],
"pre_z": pre_pos[:, 2],
"pre_E": pre_E,
"pre_dx": pre_dir[:, 0],
"pre_dy": pre_dir[:, 1],
"pre_dz": pre_dir[:, 2],
"post_x": post_pos[:, 0],
"post_y": post_pos[:, 1],
"post_z": post_pos[:, 2],
"post_E": post_E,
"post_dx": post_dir[:, 0],
"post_dy": post_dir[:, 1],
"post_dz": post_dir[:, 2],
"edep": edep,
"step_length": step_length,
}
metadata = None
if rollout:
columns["track_id"] = rng.integers(0, 3, n).astype(np.int64)
columns["termination_reason"] = (
termination_reasons
if termination_reasons is not None
else rng.choice(["", TERM_NATURAL_END], n)
)
columns["n_sec_pred"] = rng.integers(0, 3, n).astype(np.int32)
metadata = {PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}
table = pa.table(columns)
if metadata is not None:
table = table.replace_schema_metadata(metadata)
pq.write_table(table, path)
return columns
def _raw_from_world_frame(columns):
"""(N, 9) `RAW_TARGET_NAMES` oracle for a world-frame steps dict, via numpy.
Independent re-derivation of `_world_frame_local_exprs`'s forward Rodrigues
rotation, using `giant.data.transforms.local_frame_rotation`/
`travel_direction` directly rather than the module's polars expressions.
"""
pre_pos = np.column_stack(
[columns["pre_x"], columns["pre_y"], columns["pre_z"]]
).astype(np.float32)
pre_dir = np.column_stack(
[columns["pre_dx"], columns["pre_dy"], columns["pre_dz"]]
).astype(np.float32)
post_pos = np.column_stack(
[columns["post_x"], columns["post_y"], columns["post_z"]]
).astype(np.float32)
post_dir = np.column_stack(
[columns["post_dx"], columns["post_dy"], columns["post_dz"]]
).astype(np.float32)
pre_E = np.asarray(columns["pre_E"], dtype=np.float32)
post_E = np.asarray(columns["post_E"], dtype=np.float32)
post_dir_local = local_frame_rotation(pre_dir, post_dir)
travel_dir_local = local_frame_rotation(
pre_dir, travel_direction(pre_pos, post_pos)
)
return np.column_stack(
[
np.asarray(columns["step_length"], dtype=np.float32),
pre_E - post_E,
np.asarray(columns["edep"], dtype=np.float32),
post_dir_local,
travel_dir_local,
]
).astype(np.float32)
def test_rollout_vs_truth_decodes_raw_targets_correctly(tmp_path):
rng = np.random.default_rng(1)
rollout_path = tmp_path / "rollout.parquet"
truth_path = tmp_path / "truth.parquet"
rollout_cols = _write_world_frame_parquet(
rollout_path, n=150, rng=rng, rollout=True
)
truth_cols = _write_world_frame_parquet(truth_path, n=200, rng=rng, rollout=False)
expected_gen = _raw_from_world_frame(rollout_cols)
expected_real = _raw_from_world_frame(truth_cols)
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
table = marginal_table_pl(source).to_pandas().set_index("dim")
assert (table["n"] == len(truth_cols["pdg"])).all()
assert (table["n_gen"] == len(rollout_cols["pdg"])).all()
for j, name in enumerate(RAW_TARGET_NAMES):
np.testing.assert_allclose(
table.loc[name, "real_mean"], expected_real[:, j].mean(), atol=1e-3
)
np.testing.assert_allclose(
table.loc[name, "gen_mean"], expected_gen[:, j].mean(), atol=1e-3
)
def test_rollout_vs_truth_allows_unpaired_lengths(tmp_path):
rng = np.random.default_rng(2)
rollout_path = tmp_path / "rollout.parquet"
truth_path = tmp_path / "truth.parquet"
_write_world_frame_parquet(rollout_path, n=30, rng=rng, rollout=True)
_write_world_frame_parquet(truth_path, n=500, rng=rng, rollout=False)
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
table = marginal_table_pl(source)
assert (table["n"] == 500).all()
assert (table["n_gen"] == 30).all()
def test_rollout_vs_truth_drops_synthetic_termination_rows(tmp_path):
rng = np.random.default_rng(3)
rollout_path = tmp_path / "rollout.parquet"
truth_path = tmp_path / "truth.parquet"
n = 100
# Half the rollout rows are synthetic bookkeeping (dropped), half are real.
reasons = np.where(np.arange(n) % 2 == 0, TERM_ENERGY_CUTOFF, "")
rollout_cols = _write_world_frame_parquet(
rollout_path, n=n, rng=rng, rollout=True, termination_reasons=reasons
)
_write_world_frame_parquet(truth_path, n=50, rng=rng, rollout=False)
kept = reasons != TERM_ENERGY_CUTOFF
expected_gen = _raw_from_world_frame(
{k: np.asarray(v)[kept] for k, v in rollout_cols.items()}
)
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
table = marginal_table_pl(source)
assert (table["n_gen"] == kept.sum()).all()
step_length_row = table.filter(pl.col("dim") == "step_length")
np.testing.assert_allclose(
step_length_row["gen_mean"][0], expected_gen[:, 0].mean(), atol=1e-3
)
def test_rollout_vs_truth_rejects_wrong_coord_metadata(tmp_path):
rollout_path = tmp_path / "rollout.parquet"
truth_path = tmp_path / "truth.parquet"
_write_world_frame_parquet(rollout_path, rollout=True)
_write_world_frame_parquet(truth_path, rollout=False)
# Overwrite with a mismatched coord tag (predict-schema "local").
table = pq.read_table(rollout_path).replace_schema_metadata(
{PREDICT_COORD_METADATA_KEY: "local"}
)
pq.write_table(table, rollout_path)
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
with pytest.raises(ValueError, match="not a rollout file"):
marginal_table_pl(source)
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
def test_rollout_vs_truth_downstream_plots_run_without_error(tmp_path, group_by):
rollout_path = tmp_path / "rollout.parquet"
truth_path = tmp_path / "truth.parquet"
_write_world_frame_parquet(rollout_path, n=150, rollout=True)
_write_world_frame_parquet(truth_path, n=150, rollout=False)
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
assert plot_marginals(source, group_by=group_by) is not None
assert plot_kl_bars_pl(source, group_by=group_by) is not None
def test_rollout_vs_truth_joint_and_constraint_checks_run(tmp_path):
rollout_path = tmp_path / "rollout.parquet"
truth_path = tmp_path / "truth.parquet"
_write_world_frame_parquet(rollout_path, n=150, rollout=True)
_write_world_frame_parquet(truth_path, n=150, rollout=False)
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
assert plot_correlation_matrices(source) is not None
assert plot_pairwise(source, n_sample=50) is not None
assert plot_direction_alignment(source) is not None
assert plot_constraint_violations(source) is not None
report = constraint_report_pl(source).to_pandas()
assert set(report["check"]) == {
"post_dir unit norm",
"travel_dir unit norm",
"step_length >= 0",
"delta_e >= 0",
"edep >= 0",
}
def _write_truth_shower(path, event_id, n_steps_per_event, rng):
"""One-track-per-event world-frame shower: highest-`pre_E` row fixes the axis."""
rows_per_event = n_steps_per_event
rows = []
for i, eid in enumerate(event_id):
n = rows_per_event
pre_pos = rng.normal(size=(n, 3)).astype(np.float32) * 5
pre_E = rng.uniform(1.0, 50.0, n).astype(np.float32)
pre_E[0] = 1000.0 # entry step
axis = _unit_vectors(rng, 1)[0]
pre_dir = np.tile(axis, (n, 1)).astype(np.float32)
step_length = rng.uniform(0.1, 5.0, n).astype(np.float32)
post_pos = pre_pos + step_length[:, None] * axis
edep = rng.uniform(0, 1, n).astype(np.float32)
for j in range(n):
rows.append(
{
"event_id": eid,
"pdg": 11,
"material": "Pb",
"layer_id": 0,
"pre_x": pre_pos[j, 0],
"pre_y": pre_pos[j, 1],
"pre_z": pre_pos[j, 2],
"pre_E": pre_E[j],
"pre_dx": pre_dir[j, 0],
"pre_dy": pre_dir[j, 1],
"pre_dz": pre_dir[j, 2],
"post_x": post_pos[j, 0],
"post_y": post_pos[j, 1],
"post_z": post_pos[j, 2],
"edep": edep[j],
"step_length": step_length[j],
"track_id": i,
"termination_reason": TERM_NATURAL_END if j == n - 1 else "",
"n_sec_pred": 0,
}
)
return rows
def test_compute_rollout_vs_truth_observables_pl_matches_manual_reconstruction(
tmp_path,
):
rng = np.random.default_rng(4)
rollout_rows = _write_truth_shower(
tmp_path / "_r", event_id=[0, 1, 2], n_steps_per_event=15, rng=rng
)
truth_rows = _write_truth_shower(
tmp_path / "_t", event_id=[0, 1, 2], n_steps_per_event=20, rng=rng
)
rollout_path = tmp_path / "rollout.parquet"
truth_path = tmp_path / "truth.parquet"
pq.write_table(
pa.table(
{k: [r[k] for r in rollout_rows] for k in rollout_rows[0]}
).replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}),
rollout_path,
)
truth_cols = [
"event_id",
"pdg",
"material",
"layer_id",
"pre_x",
"pre_y",
"pre_z",
"pre_E",
"pre_dx",
"pre_dy",
"pre_dz",
"post_x",
"post_y",
"post_z",
"edep",
"step_length",
]
pq.write_table(
pa.table({k: [r[k] for r in truth_rows] for k in truth_cols}), truth_path
)
obs = compute_rollout_vs_truth_observables_pl(
rollout_path, truth_path, depth_bins=5, transverse_bins=5
)
rollout_df = pd.DataFrame(rollout_rows)
truth_df = pd.DataFrame(truth_rows)
expected_gen_total = rollout_df.groupby("event_id")["edep"].sum().sort_index()
expected_real_total = truth_df.groupby("event_id")["edep"].sum().sort_index()
gen_table = obs.gen_table.sort("event_id")
real_table = obs.real_table.sort("event_id")
np.testing.assert_allclose(
gen_table["total_edep"].to_numpy(), expected_gen_total.to_numpy(), rtol=1e-4
)
np.testing.assert_allclose(
real_table["total_edep"].to_numpy(), expected_real_total.to_numpy(), rtol=1e-4
)
assert (gen_table["n_tracks"].to_numpy() == 1).all()
assert (gen_table["leaked_E"].to_numpy() == 0).all()
def test_compute_rollout_vs_truth_observables_pl_profile_shapes(tmp_path):
rng = np.random.default_rng(5)
rollout_rows = _write_truth_shower(
tmp_path / "_r", event_id=[0, 1], n_steps_per_event=10, rng=rng
)
truth_rows = _write_truth_shower(
tmp_path / "_t", event_id=[0, 1, 2], n_steps_per_event=10, rng=rng
)
rollout_path = tmp_path / "rollout.parquet"
truth_path = tmp_path / "truth.parquet"
pq.write_table(
pa.table(
{k: [r[k] for r in rollout_rows] for k in rollout_rows[0]}
).replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}),
rollout_path,
)
truth_cols = [
"event_id",
"pdg",
"material",
"layer_id",
"pre_x",
"pre_y",
"pre_z",
"pre_E",
"pre_dx",
"pre_dy",
"pre_dz",
"post_x",
"post_y",
"post_z",
"edep",
"step_length",
]
pq.write_table(
pa.table({k: [r[k] for r in truth_rows] for k in truth_cols}), truth_path
)
obs = compute_rollout_vs_truth_observables_pl(
rollout_path, truth_path, depth_bins=6, transverse_bins=4
)
assert obs.depth_edges.shape == (7,)
assert obs.transverse_edges.shape == (5,)
assert len(obs.gen_table) == 2
assert len(obs.real_table) == 3
assert plot_total_energy(obs) is not None
assert plot_total_length(obs) is not None
assert plot_mean_energy_per_step(obs) is not None
assert plot_mean_length_per_step(obs) is not None
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_compute_rollout_vs_truth_observables_pl_rejects_wrong_coord_metadata(
tmp_path,
):
rollout_path = tmp_path / "rollout.parquet"
truth_path = tmp_path / "truth.parquet"
_write_world_frame_parquet(rollout_path, rollout=True)
_write_world_frame_parquet(truth_path, rollout=False)
table = pq.read_table(rollout_path).replace_schema_metadata(
{PREDICT_COORD_METADATA_KEY: "local"}
)
pq.write_table(table, rollout_path)
with pytest.raises(ValueError, match="not a rollout file"):
compute_rollout_vs_truth_observables_pl(rollout_path, truth_path)