From 6eea5ce4e89c8ea179bb2dc93e47603d5bebb718 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 22 May 2026 09:19:22 +0100 Subject: [PATCH] test(viz): Visualization Sanity rollout across visualization_jax*.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase D.2.a — adds __Visualization Sanity__ blocks to each existing visualization_jax*.py script in autolens_workspace_test (3 scripts). Companion to Phase D.1 (PRs #113 / #54). These scripts are the single-shot JAX-backed visualization path — no Nautilus search, the Sanity block goes at the end of the script after the existing PILOT SUCCEEDED print. Scripts updated: - scripts/imaging/visualization_jax.py — SIE Sanity (correctness + perf) - scripts/interferometer/visualization_jax.py — SIE Sanity + interferometer fit.model_data (complex Visibilities) finite + non-zero via analysis.fit_from(instance=instance). - scripts/point_source/visualization_jax.py — SIE Sanity only (no FoM gate; prior-median position can legitimately give chi²=-inf). Locally validated end-to-end on all three scripts: 1 tangential CC, einstein_radius=1.1938, warm-call latencies 88-98 ms (all under the 100 ms threshold), interferometer model_data sum 11118. Phase D.2.b (new modeling_visualization_jit + viz_jax for ellipse, weak lensing, quantity viz_jit) follows once this lands. Closes #114. Co-Authored-By: Claude Opus 4.7 (1M context) --- scripts/imaging/visualization_jax.py | 49 +++++++++++++++++ scripts/interferometer/visualization_jax.py | 61 +++++++++++++++++++++ scripts/point_source/visualization_jax.py | 47 ++++++++++++++++ 3 files changed, 157 insertions(+) diff --git a/scripts/imaging/visualization_jax.py b/scripts/imaging/visualization_jax.py index 940461a2..e226243a 100644 --- a/scripts/imaging/visualization_jax.py +++ b/scripts/imaging/visualization_jax.py @@ -148,3 +148,52 @@ image_path / "fit.png" ).exists(), "fit.png was not produced" print("PILOT SUCCEEDED — JAX-backed visualization produced fit.png/tracer.png.") + + +""" +__Visualization Sanity__ + +Phase D.2.a rollout of the Sanity-block pattern from PR #111 / #113. +Same imaging-template SIE assertions as the modeling_visualization_jit +variant — catches the silent-zero / cache-busting failure class on the +single-shot JAX-backed visualization path. Uses a deterministic SIE +tracer so assertions are independent of the script's specific model. +""" +import time as _sanity_time +import numpy as _sanity_np +from autogalaxy.operate.lens_calc import LensCalc as _SanityLensCalc + +_sanity_lens = al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.0, 0.0), einstein_radius=1.2, ell_comps=(0.1, 0.0) + ), +) +_sanity_source = al.Galaxy(redshift=1.0) +_sanity_tracer = al.Tracer(galaxies=[_sanity_lens, _sanity_source]) +_sanity_od = _SanityLensCalc.from_tracer(_sanity_tracer) + +_tc_list = _sanity_od.tangential_critical_curve_list_via_zero_contour_from() +assert len(_tc_list) > 0, ( + "no tangential critical curves returned by zero_contour — algorithmic " + "regression (PyAutoGalaxy abd7b717 / PyAutoFit #1280 family)" +) +_er_sanity = _sanity_od.einstein_radius_via_zero_contour_from() +assert _sanity_np.isfinite(float(_er_sanity)) and float(_er_sanity) > 0.0, ( + f"Einstein radius via zero_contour returned {_er_sanity} — should be " + "finite and positive for the SIE sanity tracer (einstein_radius=1.2)" +) +print( + f" PASS Visualization Sanity (correctness): " + f"{len(_tc_list)} tangential CC, einstein_radius={float(_er_sanity):.4f}" +) + +_sanity_od.tangential_critical_curve_list_via_zero_contour_from() # warm cache +_t0 = _sanity_time.perf_counter() +_sanity_od.tangential_critical_curve_list_via_zero_contour_from() +_warm_dt = _sanity_time.perf_counter() - _t0 +assert _warm_dt < 0.1, ( + f"zero_contour warm call took {_warm_dt * 1000:.1f} ms (> 100 ms) — " + "closure cache-busting bug from PyAutoGalaxy #433 may have regressed" +) +print(f" PASS Visualization Sanity (perf): warm call {_warm_dt * 1000:.1f} ms") diff --git a/scripts/interferometer/visualization_jax.py b/scripts/interferometer/visualization_jax.py index bc42cb77..7bdff78a 100644 --- a/scripts/interferometer/visualization_jax.py +++ b/scripts/interferometer/visualization_jax.py @@ -144,3 +144,64 @@ print( "PILOT SUCCEEDED — JAX-backed interferometer visualization produced fit.png/tracer.png." ) + + +""" +__Visualization Sanity__ + +Phase D.2.a rollout. Lensing-side SIE Sanity (silent-zero / cache-busting +regression class) plus interferometer-specific `fit.model_data` (complex +visibilities) finite + non-zero on the script's actual fit. +""" +import time as _sanity_time +import numpy as _sanity_np +from autogalaxy.operate.lens_calc import LensCalc as _SanityLensCalc + +_sanity_lens = al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.0, 0.0), einstein_radius=1.2, ell_comps=(0.1, 0.0) + ), +) +_sanity_source = al.Galaxy(redshift=1.0) +_sanity_tracer = al.Tracer(galaxies=[_sanity_lens, _sanity_source]) +_sanity_od = _SanityLensCalc.from_tracer(_sanity_tracer) + +_tc_list = _sanity_od.tangential_critical_curve_list_via_zero_contour_from() +assert len(_tc_list) > 0, ( + "no tangential critical curves returned by zero_contour — algorithmic " + "regression (PyAutoGalaxy abd7b717 / PyAutoFit #1280 family)" +) +_er_sanity = _sanity_od.einstein_radius_via_zero_contour_from() +assert _sanity_np.isfinite(float(_er_sanity)) and float(_er_sanity) > 0.0, ( + f"Einstein radius via zero_contour returned {_er_sanity} — should be " + "finite and positive for the SIE sanity tracer (einstein_radius=1.2)" +) +print( + f" PASS Visualization Sanity (lensing): " + f"{len(_tc_list)} tangential CC, einstein_radius={float(_er_sanity):.4f}" +) + +_sanity_od.tangential_critical_curve_list_via_zero_contour_from() # warm cache +_t0 = _sanity_time.perf_counter() +_sanity_od.tangential_critical_curve_list_via_zero_contour_from() +_warm_dt = _sanity_time.perf_counter() - _t0 +assert _warm_dt < 0.1, ( + f"zero_contour warm call took {_warm_dt * 1000:.1f} ms (> 100 ms) — " + "closure cache-busting bug from PyAutoGalaxy #433 may have regressed" +) +print(f" PASS Visualization Sanity (perf): warm call {_warm_dt * 1000:.1f} ms") + +# Interferometer-specific: model_data (complex Visibilities) finite + non-zero. +_fit_for_vis = analysis.fit_from(instance=instance) +_mv = _sanity_np.asarray(_fit_for_vis.model_data) +assert _sanity_np.isfinite(_mv).all(), ( + "fit.model_data (visibilities) have nan/inf — NUFFT/inversion collapse" +) +assert float(_sanity_np.abs(_mv).sum()) > 0.0, ( + "fit.model_data (visibilities) all-zero — NUFFT/inversion collapse" +) +print( + f" PASS Visualization Sanity (interferometer): " + f"|model_data|.sum() = {float(_sanity_np.abs(_mv).sum()):.4f}" +) diff --git a/scripts/point_source/visualization_jax.py b/scripts/point_source/visualization_jax.py index 235fdd40..c35e6c91 100644 --- a/scripts/point_source/visualization_jax.py +++ b/scripts/point_source/visualization_jax.py @@ -137,3 +137,50 @@ image_path / "fit.png" ).exists(), f"fit.png was not produced. Files present: {list(image_path.iterdir())}" print("PILOT SUCCEEDED — JAX-backed point-source visualization produced fit.png.") + + +""" +__Visualization Sanity__ + +Phase D.2.a rollout. SIE-tracer lensing-side check only — no +point-source-specific FoM assertion (prior-median position can +legitimately give chi² = -inf if outside the image-pair basin). +""" +import time as _sanity_time +import numpy as _sanity_np +from autogalaxy.operate.lens_calc import LensCalc as _SanityLensCalc + +_sanity_lens = al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.0, 0.0), einstein_radius=1.2, ell_comps=(0.1, 0.0) + ), +) +_sanity_source = al.Galaxy(redshift=1.0) +_sanity_tracer = al.Tracer(galaxies=[_sanity_lens, _sanity_source]) +_sanity_od = _SanityLensCalc.from_tracer(_sanity_tracer) + +_tc_list = _sanity_od.tangential_critical_curve_list_via_zero_contour_from() +assert len(_tc_list) > 0, ( + "no tangential critical curves returned by zero_contour — algorithmic " + "regression (PyAutoGalaxy abd7b717 / PyAutoFit #1280 family)" +) +_er_sanity = _sanity_od.einstein_radius_via_zero_contour_from() +assert _sanity_np.isfinite(float(_er_sanity)) and float(_er_sanity) > 0.0, ( + f"Einstein radius via zero_contour returned {_er_sanity} — should be " + "finite and positive for the SIE sanity tracer (einstein_radius=1.2)" +) +print( + f" PASS Visualization Sanity (lensing): " + f"{len(_tc_list)} tangential CC, einstein_radius={float(_er_sanity):.4f}" +) + +_sanity_od.tangential_critical_curve_list_via_zero_contour_from() # warm cache +_t0 = _sanity_time.perf_counter() +_sanity_od.tangential_critical_curve_list_via_zero_contour_from() +_warm_dt = _sanity_time.perf_counter() - _t0 +assert _warm_dt < 0.1, ( + f"zero_contour warm call took {_warm_dt * 1000:.1f} ms (> 100 ms) — " + "closure cache-busting bug from PyAutoGalaxy #433 may have regressed" +) +print(f" PASS Visualization Sanity (perf): warm call {_warm_dt * 1000:.1f} ms")