lars
55332db67a
CI / Lint (ruff check) (push) Successful in 31s
CI / Format (ruff format) (push) Successful in 32s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 26s
CI / Type check (ty) (push) Successful in 29s
CI / Format (ruff format) (pull_request) Successful in 33s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 35s
CI / Tests (pull_request) Successful in 3m47s
CI / Tests (push) Successful in 3m55s
Rejoins lines that only wrapped because they exceeded the old 88-char limit; ruff check and the full test suite (725 passed) are unaffected.
371 lines
11 KiB
Python
371 lines
11 KiB
Python
"""Render smoke test — skipped where plotstyle / LaTeX is unavailable."""
|
||
|
||
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
|
||
|
||
|
||
def _try_render(reduced: list[Reduced], out: Path):
|
||
from giant.analysis.render import render_all
|
||
|
||
for r in reduced:
|
||
r.save(out / "reduced" / f"{r.id}.json")
|
||
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):
|
||
# render_mod.subprocess *is* the stdlib subprocess module, so a blanket
|
||
# patch of .run would also swallow the real subprocess.run calls
|
||
# matplotlib's texmanager makes to compile LaTeX during savefig — only
|
||
# intercept the "gallery generate" call itself and pass everything else
|
||
# (LaTeX included) through to the real subprocess.run.
|
||
calls = []
|
||
real_run = render_mod.subprocess.run
|
||
|
||
def fake_run(*a, **k):
|
||
if a and a[0] and a[0][0] == "gallery":
|
||
calls.append((a, k))
|
||
return None
|
||
return real_run(*a, **k)
|
||
|
||
monkeypatch.setattr(render_mod.subprocess, "run", fake_run)
|
||
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(
|
||
"m",
|
||
"marginals",
|
||
"overlay_hist",
|
||
"Overlay",
|
||
"x",
|
||
{
|
||
"edges": [0, 1, 2, 3],
|
||
"rollout": [1, 2, 3],
|
||
"reference": [3, 2, 1],
|
||
"log_y": False,
|
||
},
|
||
),
|
||
Reduced(
|
||
"g",
|
||
"marginals",
|
||
"grouped_hist",
|
||
"Grouped",
|
||
"x",
|
||
{
|
||
"edges": [0, 1, 2],
|
||
"groups": {"a": {"rollout": [1, 2], "reference": [2, 1]}},
|
||
"log_y": False,
|
||
},
|
||
),
|
||
Reduced(
|
||
"p",
|
||
"shower",
|
||
"profile",
|
||
"Profile",
|
||
"depth",
|
||
{
|
||
"edges": [0, 1, 2],
|
||
"rollout_mean": [1, 2],
|
||
"rollout_std": [0.1, 0.2],
|
||
"reference_mean": [1.1, 1.9],
|
||
"reference_std": [0.1, 0.1],
|
||
"ylabel": "e",
|
||
},
|
||
),
|
||
Reduced(
|
||
"b",
|
||
"species",
|
||
"bar",
|
||
"Bar",
|
||
"species",
|
||
{
|
||
"labels": ["e-", "gamma"],
|
||
"rollout": [0.6, 0.4],
|
||
"reference": [0.5, 0.5],
|
||
"ylabel": "frac",
|
||
},
|
||
),
|
||
Reduced(
|
||
"s",
|
||
"species",
|
||
"single_hist",
|
||
"Single",
|
||
"x",
|
||
{"edges": [0, 1, 2], "rollout": [5, 1], "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)
|
||
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
|