From 49714ef06334580e409b73f240d9f801345f5415 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 23 Aug 2026 14:31:24 +0000 Subject: [PATCH] fix: make the PointSolver SMALL_DATASETS short-circuit announce itself MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Under `PYAUTO_SMALL_DATASETS=1`, `PointSolver.solve` skips the triangle-tiling solve and returns the fixed pair `[(1.0, 0.0), (0.0, 1.0)]`. That is intentional — the solve is the dominant cost in smoke runs and is meaningless on downsized grids — but it was completely silent, and the returned pair is identical for every lens model. Verified here at einstein_radius 1.0 / 1.6 / 2.5: same two coordinates, no log line even at DEBUG with warnings forced on. Anything derived from those positions is therefore model-independent, so a parity script comparing such a value against a pinned literal is measuring nothing. That is exactly how it failed: `autolens_workspace_test`'s `point_source/jax_likelihood/point.py` produced a different wrong value on every run and cost a full investigation before the cause was found (#710). The short-circuit now emits one `logger.warning` per process — latched at module level rather than per instance, because a vmap batch calls `solve` once per sampled parameter set and the condition is about the process environment. Warn rather than raise: every smoke script that legitimately relies on the speedup keeps working, and only the silence goes away. Tests cover all three properties: the pair is model-independent, the warning fires exactly once across repeated calls, and nothing is emitted when the flag is unset. Refs #710 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JJSb9kydhAgMgCCKh26ai2 --- autolens/point/solver/point_solver.py | 26 ++++++ .../weak/test_simulator_small_datasets.py | 91 +++++++++++++++++++ 2 files changed, 117 insertions(+) diff --git a/autolens/point/solver/point_solver.py b/autolens/point/solver/point_solver.py index aa13a4af2..cd9e1b8bf 100644 --- a/autolens/point/solver/point_solver.py +++ b/autolens/point/solver/point_solver.py @@ -31,6 +31,12 @@ logger = logging.getLogger(__name__) +# One-shot latch for the ``PYAUTO_SMALL_DATASETS`` short-circuit warning in +# ``PointSolver.solve``. Module-level rather than per-instance: a vmap batch calls +# ``solve`` once per sampled parameter set, and the message is about the process +# environment, not about any one solver. +_SMALL_DATASETS_WARNED = False + class PointSolver(AbstractSolver): @@ -101,6 +107,13 @@ def solve( normally. ``PYAUTO_SMALL_DATASETS`` is a smoke-test-only flag and is never set inside a ``jax.jit`` trace, so a plain numpy-backed ``Grid2DIrregular`` is safe here even when the surrounding analysis uses ``xp=jnp``. + + The short-circuit announces itself with a ``logger.warning`` the first time it + fires in a process. It returns the same two coordinates for every lens model, so + anything derived from them — a likelihood, a chi-squared, a position pairing — is + model-independent, and a parity script that compares such a value against a pinned + literal is measuring nothing. That failure mode was silent until it cost a real + investigation (PyAutoLens#710), hence the warning rather than a bare return. """ if xp is None: xp = self._xp @@ -117,6 +130,19 @@ def solve( # JIT-it-yourself pattern. if os.environ.get("PYAUTO_SMALL_DATASETS") == "1": + global _SMALL_DATASETS_WARNED + if not _SMALL_DATASETS_WARNED: + _SMALL_DATASETS_WARNED = True + logger.warning( + "PointSolver.solve is short-circuited: PYAUTO_SMALL_DATASETS=1 is set, " + "so the triangle-tiling solve is skipped and the fixed pair " + "[(1.0, 0.0), (0.0, 1.0)] is returned for EVERY lens model. Any " + "likelihood, chi-squared or position-pairing value computed from these " + "positions is independent of the model and must not be compared against " + "a pinned literal. Scripts that need a real solve should declare " + "`ENV: full_datasets` (workspace test-harness) or unset the flag. This " + "warning is issued once per process." + ) return aa.Grid2DIrregular(values=[(1.0, 0.0), (0.0, 1.0)]) if xp is not np: diff --git a/test_autolens/weak/test_simulator_small_datasets.py b/test_autolens/weak/test_simulator_small_datasets.py index f013908f1..51dca944a 100644 --- a/test_autolens/weak/test_simulator_small_datasets.py +++ b/test_autolens/weak/test_simulator_small_datasets.py @@ -1,3 +1,5 @@ +import pytest + import autolens as al @@ -41,3 +43,92 @@ def test__explicit_grid__never_capped(monkeypatch): ) assert dataset.n_galaxies == 60 + + +def _solver(): + import autolens as al + + grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.2) + return al.PointSolver.for_grid( + grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1 + ) + + +def _tracer_with_einstein_radius(einstein_radius): + lens = al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=einstein_radius), + ) + return al.Tracer(galaxies=[lens, al.Galaxy(redshift=1.0)]) + + +def test__point_solver__short_circuits_to_a_model_independent_pair(monkeypatch): + """Under the cap the solve is skipped entirely, so every lens model yields the same + two positions. Anything derived from them is model-independent — the reason a pinned + parity literal cannot be compared against a capped run (PyAutoLens#710).""" + import numpy as np + + monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1") + + solver = _solver() + + solved = [ + np.asarray( + solver.solve( + tracer=_tracer_with_einstein_radius(einstein_radius), + source_plane_coordinate=(0.07, 0.07), + ).array + ) + for einstein_radius in (1.0, 1.6, 2.5) + ] + + for positions in solved: + assert positions == pytest.approx(np.array([[1.0, 0.0], [0.0, 1.0]])) + + +def test__point_solver__short_circuit_warns_once_per_process(monkeypatch, caplog): + """The short-circuit must not be silent, and must not flood a vmap batch.""" + import logging + + from autolens.point.solver import point_solver + + monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1") + monkeypatch.setattr(point_solver, "_SMALL_DATASETS_WARNED", False) + + solver = _solver() + tracer = _tracer_with_einstein_radius(1.6) + + with caplog.at_level(logging.WARNING, logger=point_solver.__name__): + for _ in range(3): + solver.solve(tracer=tracer, source_plane_coordinate=(0.07, 0.07)) + + warnings = [ + record + for record in caplog.records + if record.levelno == logging.WARNING + and "PYAUTO_SMALL_DATASETS" in record.getMessage() + ] + + assert len(warnings) == 1 + assert "EVERY lens model" in warnings[0].getMessage() + + +def test__point_solver__no_short_circuit_warning_without_the_env_var(monkeypatch, caplog): + import logging + + from autolens.point.solver import point_solver + + monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False) + monkeypatch.setattr(point_solver, "_SMALL_DATASETS_WARNED", False) + + with caplog.at_level(logging.WARNING, logger=point_solver.__name__): + _solver().solve( + tracer=_tracer_with_einstein_radius(1.6), + source_plane_coordinate=(0.07, 0.07), + ) + + assert not [ + record + for record in caplog.records + if "PYAUTO_SMALL_DATASETS" in record.getMessage() + ]