Skip to content

Commit 49714ef

Browse files
committed
fix: make the PointSolver SMALL_DATASETS short-circuit announce itself
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JJSb9kydhAgMgCCKh26ai2
1 parent 491e5e4 commit 49714ef

2 files changed

Lines changed: 117 additions & 0 deletions

File tree

autolens/point/solver/point_solver.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,12 @@
3131

3232
logger = logging.getLogger(__name__)
3333

34+
# One-shot latch for the ``PYAUTO_SMALL_DATASETS`` short-circuit warning in
35+
# ``PointSolver.solve``. Module-level rather than per-instance: a vmap batch calls
36+
# ``solve`` once per sampled parameter set, and the message is about the process
37+
# environment, not about any one solver.
38+
_SMALL_DATASETS_WARNED = False
39+
3440

3541
class PointSolver(AbstractSolver):
3642

@@ -101,6 +107,13 @@ def solve(
101107
normally. ``PYAUTO_SMALL_DATASETS`` is a smoke-test-only flag and is never set
102108
inside a ``jax.jit`` trace, so a plain numpy-backed ``Grid2DIrregular`` is safe
103109
here even when the surrounding analysis uses ``xp=jnp``.
110+
111+
The short-circuit announces itself with a ``logger.warning`` the first time it
112+
fires in a process. It returns the same two coordinates for every lens model, so
113+
anything derived from them — a likelihood, a chi-squared, a position pairing — is
114+
model-independent, and a parity script that compares such a value against a pinned
115+
literal is measuring nothing. That failure mode was silent until it cost a real
116+
investigation (PyAutoLens#710), hence the warning rather than a bare return.
104117
"""
105118
if xp is None:
106119
xp = self._xp
@@ -117,6 +130,19 @@ def solve(
117130
# JIT-it-yourself pattern.
118131

119132
if os.environ.get("PYAUTO_SMALL_DATASETS") == "1":
133+
global _SMALL_DATASETS_WARNED
134+
if not _SMALL_DATASETS_WARNED:
135+
_SMALL_DATASETS_WARNED = True
136+
logger.warning(
137+
"PointSolver.solve is short-circuited: PYAUTO_SMALL_DATASETS=1 is set, "
138+
"so the triangle-tiling solve is skipped and the fixed pair "
139+
"[(1.0, 0.0), (0.0, 1.0)] is returned for EVERY lens model. Any "
140+
"likelihood, chi-squared or position-pairing value computed from these "
141+
"positions is independent of the model and must not be compared against "
142+
"a pinned literal. Scripts that need a real solve should declare "
143+
"`ENV: full_datasets` (workspace test-harness) or unset the flag. This "
144+
"warning is issued once per process."
145+
)
120146
return aa.Grid2DIrregular(values=[(1.0, 0.0), (0.0, 1.0)])
121147

122148
if xp is not np:

test_autolens/weak/test_simulator_small_datasets.py

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import pytest
2+
13
import autolens as al
24

35

@@ -41,3 +43,92 @@ def test__explicit_grid__never_capped(monkeypatch):
4143
)
4244

4345
assert dataset.n_galaxies == 60
46+
47+
48+
def _solver():
49+
import autolens as al
50+
51+
grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.2)
52+
return al.PointSolver.for_grid(
53+
grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1
54+
)
55+
56+
57+
def _tracer_with_einstein_radius(einstein_radius):
58+
lens = al.Galaxy(
59+
redshift=0.5,
60+
mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=einstein_radius),
61+
)
62+
return al.Tracer(galaxies=[lens, al.Galaxy(redshift=1.0)])
63+
64+
65+
def test__point_solver__short_circuits_to_a_model_independent_pair(monkeypatch):
66+
"""Under the cap the solve is skipped entirely, so every lens model yields the same
67+
two positions. Anything derived from them is model-independent — the reason a pinned
68+
parity literal cannot be compared against a capped run (PyAutoLens#710)."""
69+
import numpy as np
70+
71+
monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1")
72+
73+
solver = _solver()
74+
75+
solved = [
76+
np.asarray(
77+
solver.solve(
78+
tracer=_tracer_with_einstein_radius(einstein_radius),
79+
source_plane_coordinate=(0.07, 0.07),
80+
).array
81+
)
82+
for einstein_radius in (1.0, 1.6, 2.5)
83+
]
84+
85+
for positions in solved:
86+
assert positions == pytest.approx(np.array([[1.0, 0.0], [0.0, 1.0]]))
87+
88+
89+
def test__point_solver__short_circuit_warns_once_per_process(monkeypatch, caplog):
90+
"""The short-circuit must not be silent, and must not flood a vmap batch."""
91+
import logging
92+
93+
from autolens.point.solver import point_solver
94+
95+
monkeypatch.setenv("PYAUTO_SMALL_DATASETS", "1")
96+
monkeypatch.setattr(point_solver, "_SMALL_DATASETS_WARNED", False)
97+
98+
solver = _solver()
99+
tracer = _tracer_with_einstein_radius(1.6)
100+
101+
with caplog.at_level(logging.WARNING, logger=point_solver.__name__):
102+
for _ in range(3):
103+
solver.solve(tracer=tracer, source_plane_coordinate=(0.07, 0.07))
104+
105+
warnings = [
106+
record
107+
for record in caplog.records
108+
if record.levelno == logging.WARNING
109+
and "PYAUTO_SMALL_DATASETS" in record.getMessage()
110+
]
111+
112+
assert len(warnings) == 1
113+
assert "EVERY lens model" in warnings[0].getMessage()
114+
115+
116+
def test__point_solver__no_short_circuit_warning_without_the_env_var(monkeypatch, caplog):
117+
import logging
118+
119+
from autolens.point.solver import point_solver
120+
121+
monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False)
122+
monkeypatch.setattr(point_solver, "_SMALL_DATASETS_WARNED", False)
123+
124+
with caplog.at_level(logging.WARNING, logger=point_solver.__name__):
125+
_solver().solve(
126+
tracer=_tracer_with_einstein_radius(1.6),
127+
source_plane_coordinate=(0.07, 0.07),
128+
)
129+
130+
assert not [
131+
record
132+
for record in caplog.records
133+
if "PYAUTO_SMALL_DATASETS" in record.getMessage()
134+
]

0 commit comments

Comments
 (0)