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() + ]