diff --git a/tests/test_render.py b/tests/test_render.py index 49bb84c..f5691c1 100644 --- a/tests/test_render.py +++ b/tests/test_render.py @@ -4,10 +4,12 @@ from __future__ import annotations from pathlib import Path +import numpy as np import pytest pytest.importorskip("plotstyle") +from giant.analysis import render as render_mod # noqa: E402 from giant.analysis.reduced import Reduced # noqa: E402 @@ -19,6 +21,148 @@ def _try_render(reduced: list[Reduced], out: Path): return render_all(out / "reduced", out / "plots") +def test_render_router_diagnostics_and_edge_cases(tmp_path: Path): + reduced = [ + Reduced( + "rg", + "router", + "router_gating", + "Router gating", + "pre-step energy [MeV]", + { + "n_experts": 2, + "log_x": True, + "router_type": "energy", + "rollout": { + "centers": [1.0, 10.0, 100.0], + "means": [[0.6, 0.4], [0.5, 0.5], [0.4, 0.6]], + }, + "reference": { + "centers": [1.0, 10.0, 100.0], + "means": [[0.55, 0.45], [0.5, 0.5], [0.45, 0.55]], + }, + }, + ), + Reduced( + "rs", + "router", + "router_share", + "Router share", + "species", + { + "categories": ["e-", "gamma"], + "n_experts": 2, + "router_type": "energy", + "rollout": {"e-": [0.7, 0.3], "gamma": [0.2, 0.8]}, + "reference": {"e-": [0.6, 0.4], "gamma": [0.3, 0.7]}, + }, + ), + Reduced( + "ru", + "router", + "unavailable", + "Router unavailable", + "x", + {"note": "router diagnostics unavailable: no router in this run"}, + ), + Reduced( + "g4", + "marginals", + "grouped_hist", + "Grouped (4)", + "x", + { + "edges": [0, 1, 2], + "groups": { + lbl: {"rollout": [1, 2], "reference": [2, 1]} + for lbl in ("a", "b", "c", "d") + }, + "log_y": True, + }, + ), + Reduced( + "sl", + "species", + "single_hist", + "Single (log-x)", + "x", + {"edges": [1, 10, 100], "rollout": [5, 1], "log_x": True, "log_y": True}, + ), + ] + try: + pdfs = _try_render(reduced, tmp_path) + except RuntimeError as e: # LaTeX missing at render time + pytest.skip(f"LaTeX rendering unavailable: {e}") + assert len(pdfs) == len(reduced) + assert all(p.exists() for p in pdfs) + + +def test_render_all_run_gallery_invokes_subprocess(tmp_path: Path, monkeypatch): + calls = [] + monkeypatch.setattr( + render_mod.subprocess, "run", lambda *a, **k: calls.append((a, k)) + ) + reduced = [ + Reduced( + "s", + "species", + "single_hist", + "Single", + "x", + {"edges": [0, 1, 2], "rollout": [5, 1]}, + ) + ] + for r in reduced: + r.save(tmp_path / "reduced" / f"{r.id}.json") + try: + render_mod.render_all( + tmp_path / "reduced", tmp_path / "plots", run_gallery=True + ) + except RuntimeError as e: + pytest.skip(f"LaTeX rendering unavailable: {e}") + + assert len(calls) == 1 + args, kwargs = calls[0] + assert args[0] == ["gallery", "generate", "--source", str(tmp_path / "plots")] + assert kwargs == {"check": True} + + +def test_render_run_glues_condor_run_meta_into_render_all(tmp_path: Path, monkeypatch): + from giant.analysis import condor as condor_mod + + run_dir = tmp_path / "run" + (run_dir / "reduced").mkdir(parents=True) + + merge_calls = [] + monkeypatch.setattr( + condor_mod, "merge_all", lambda rd: merge_calls.append(Path(rd)) + ) + meta = condor_mod.RunMeta( + rollout="rollout.parquet", + reference="reference.parquet", + run_dir=str(run_dir), + title="my-run", + plot_meta={"checkpoint": "ckpt/best.pt"}, + ) + monkeypatch.setattr(condor_mod.RunMeta, "load", classmethod(lambda cls, p: meta)) + + Reduced( + "s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "rollout": [1]} + ).save(run_dir / "reduced" / "s.json") + + try: + pdfs = render_mod.render_run(run_dir) + except RuntimeError as e: + pytest.skip(f"LaTeX rendering unavailable: {e}") + + assert merge_calls == [run_dir] + assert len(pdfs) == 1 + plot_meta = (run_dir / "plots" / "species" / "s.yaml").read_text() + assert "checkpoint" in plot_meta + root_meta = (run_dir / "plots" / "metadata.yaml").read_text() + assert "my-run" in root_meta + + def test_render_one_of_each_kind(tmp_path: Path): reduced = [ Reduced( @@ -90,3 +234,135 @@ def test_render_one_of_each_kind(tmp_path: Path): assert len(pdfs) == len(reduced) assert all(p.exists() for p in pdfs) assert (tmp_path / "plots" / "metadata.yaml").exists() + + +# ── pure-function helpers: no matplotlib figure needed ────────────────── + + +def test_density_zero_total_returns_counts_unchanged(): + counts = np.array([0.0, 0.0, 0.0]) + out = render_mod._density(counts, np.array([0.0, 1.0, 2.0, 3.0])) + np.testing.assert_array_equal(out, counts) + + +def test_density_normalizes_by_total_and_bin_width(): + counts = [1, 3] + edges = np.array([0.0, 2.0, 4.0]) # bin width 2 + out = render_mod._density(counts, edges) + np.testing.assert_allclose(out, np.array([1, 3]) / (4 * 2)) + + +def test_router_summary_disabled_is_off(): + assert render_mod._router_summary({"enabled": False, "type": "energy"}) == "off" + assert render_mod._router_summary({}) == "off" + + +def test_router_summary_enabled_formats_type_and_n_experts(): + cfg = {"enabled": True, "type": "energy", "n_experts": 8} + assert render_mod._router_summary(cfg) == "energy×8" + + +def test_figure_params_v2_basics_and_router_and_epoch(): + mc = { + "stage1_model": { + "hidden_dim": 256, + "n_res_blocks": 4, + "generator": "flow", + "router": {"enabled": True, "type": "energy", "n_experts": 4}, + }, + "conditioning": {"particle": {"type": "physical"}}, + } + run_meta = {"training_epoch": 12, "best_val_loss": 0.123456, "steps": 10} + params = render_mod._figure_params(run_meta | {"model_config": mc}) + assert params == { + "hidden_dim": 256, + "n_res_blocks": 4, + "mode": "flow", + "conditioning": "physical", + "router": "energy×4", + "epoch": 12, + "best_val_loss": 0.1235, + "steps": 10, + } + + +def test_figure_params_v2_wgan_reports_noise_dim_not_steps(): + mc = { + "stage1_model": { + "generator": "wgan", + "wgan": {"noise_dim": 32}, + }, + } + run_meta = {"model_config": mc, "steps": 10} + params = render_mod._figure_params(run_meta) + assert params["mode"] == "wgan" + assert params["noise_dim"] == 32 + assert "steps" not in params + + +def test_figure_params_v2_reports_mode_s2_only_when_it_differs(): + same = { + "stage1_model": {"generator": "flow"}, + "stage2_model": {"generator": "flow"}, + } + assert "mode_s2" not in render_mod._figure_params({"model_config": same}) + + mixed = { + "stage1_model": {"generator": "flow"}, + "stage2_model": {"generator": "wgan"}, + } + params = render_mod._figure_params({"model_config": mixed}) + assert params["mode_s2"] == "wgan" + + +def test_figure_params_old_shape_basics(): + run_meta = { + "model_config": { + "hidden_dim": 128, + "n_blocks": 3, + "mode": "ddpm", + "conditioning": "embedding", + "router": {"enabled": False}, + }, + "training_epoch": 5, + "best_val_loss": 0.5, + "steps": 20, + } + params = render_mod._figure_params(run_meta) + assert params == { + "hidden_dim": 128, + "n_blocks": 3, + "mode": "ddpm", + "conditioning": "embedding", + "router": "off", + "epoch": 5, + "best_val_loss": 0.5, + "steps": 20, + } + + +def test_figure_params_old_shape_wgan_reports_noise_dim_not_steps(): + run_meta = { + "model_config": {"mode": "wgan", "noise_dim": 16}, + "steps": 20, + } + params = render_mod._figure_params(run_meta) + assert params["noise_dim"] == 16 + assert "steps" not in params + + +def test_plot_metadata_includes_note_and_run_meta_parameters(): + r = Reduced( + "u", "router", "unavailable", "Unavailable", "x", {"note": "no router data"} + ) + meta = render_mod._plot_metadata(r, {"title": "run-1", "checkpoint": "ckpt.pt"}) + assert meta["note"] == "no router data" + assert meta["parameters"] == {"checkpoint": "ckpt.pt"} + assert "title" not in meta["parameters"] + + +def test_plot_metadata_omits_parameters_when_run_meta_empty(): + r = Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1]}) + meta = render_mod._plot_metadata(r, {}) + assert "parameters" not in meta + assert "note" not in meta