From b346b6a0a092098f23964e274edb220ea9757ad0 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 24 Sep 2026 12:55:13 +0200 Subject: [PATCH 1/3] perf(point): JAX PointSolver step 0 deflects only the unique lattice vertices (PyAutoArray#568 phase 3) `AbstractSolver._initial_triangles` builds the JAX tiling with `static_vertices=True`, so step 0 ray-traces the cached geometric vertex table (11 859 rows for the 100x100 @ 0.2" grid, not the flat 69 849) and gathers back through `indices`. The lattice depends only on the solver geometry (static pytree aux data), so the table is a compile-time constant and a change of geometry retraces with the matching table. NumPy path unchanged. New test_static_lattice_jax.py (skips without jax): shape guard on the first traced deflection grid (red on main: 69 849), geometry change changes the table, positions/image counts vs the flat table (<= 1e-10, generic and near-caustic sources) and vs NumPy, log L (rel 1e-12) and non-zero jax.grad vs the flat table, vmap vs scalar. Co-Authored-By: Claude Opus 5.5 --- autolens/point/solver/shape_solver.py | 33 ++- .../triangles/test_static_lattice_jax.py | 249 ++++++++++++++++++ 2 files changed, 277 insertions(+), 5 deletions(-) create mode 100644 test_autolens/point/triangles/test_static_lattice_jax.py diff --git a/autolens/point/solver/shape_solver.py b/autolens/point/solver/shape_solver.py index d7b8f2418..1f35b01cb 100644 --- a/autolens/point/solver/shape_solver.py +++ b/autolens/point/solver/shape_solver.py @@ -108,17 +108,40 @@ def _initial_triangles(self, xp): """ Build the initial triangle tiling for the stored image-plane extent, using the JAX or NumPy triangle implementation depending on ``xp``. + + On the JAX path the tiling carries its static vertex table + (``static_vertices=True``, see + `autoarray.structures.triangles.coordinate_array.static_vertex_table`): the + step-0 lattice depends only on the solver's geometry, which is static pytree + aux data, so its geometrically unique vertices are a cached compile-time + constant. Step 0's ``triangles.vertices`` -- the grid every step-0 deflection is + evaluated on -- is then the ``(V, 2)`` unique table instead of the flat + ``(3N, 2)`` one (11 859 rather than 69 849 rows for the +-9.9" / 0.2" lattice + of a 100x100, 0.2" grid), and ``triangles.indices`` gathers the traced + vertices back to the triangles. Later steps are data-dependent and keep the + flat table. A change of geometry changes the pytree aux data, so it retraces + with the matching table. The NumPy path is unchanged (its sibling deduplicates + with dynamic shapes). """ if xp.__name__.startswith("jax"): from autoarray.structures.triangles.coordinate_array import ( - CoordinateArrayTriangles as triangle_cls, + CoordinateArrayTriangles, ) - else: - from autoarray.structures.triangles.coordinate_array_np import ( - CoordinateArrayTrianglesNp as triangle_cls, + + return CoordinateArrayTriangles.for_limits_and_scale( + y_min=self.y_min, + y_max=self.y_max, + x_min=self.x_min, + x_max=self.x_max, + scale=self.scale, + static_vertices=True, ) - return triangle_cls.for_limits_and_scale( + from autoarray.structures.triangles.coordinate_array_np import ( + CoordinateArrayTrianglesNp, + ) + + return CoordinateArrayTrianglesNp.for_limits_and_scale( y_min=self.y_min, y_max=self.y_max, x_min=self.x_min, diff --git a/test_autolens/point/triangles/test_static_lattice_jax.py b/test_autolens/point/triangles/test_static_lattice_jax.py new file mode 100644 index 000000000..7934d7112 --- /dev/null +++ b/test_autolens/point/triangles/test_static_lattice_jax.py @@ -0,0 +1,249 @@ +""" +JAX tests of the static step-0 lattice on the `PointSolver` path (point-source CPU phase 3, +PyAutoArray#568). + +On the JAX path `AbstractSolver._initial_triangles` builds its tiling with +``static_vertices=True``, so step 0 deflects only the geometrically unique lattice vertices +(11 859 rather than 69 849 rows for the +-9.9" / 0.2" lattice of a 100x100, 0.2" grid) and +gathers them back to the triangles. These tests pin that shape (the guard is red on the flat +table), that a change of geometry changes the table, and that positions, image counts, the +likelihood, `jax.grad` and `jax.vmap` agree with the flat per-triangle table the solver used +before -- including near-caustic sources. + +Every test skips on the NumPy-only (no ``[optional]`` extras) CI leg. +""" + +import importlib + +import numpy as np +import pytest + + +if importlib.util.find_spec("jax") is None: + pytestmark = pytest.mark.skip(reason="requires jax (the [optional] extras)") + + def test__placeholder_requires_jax(): # pragma: no cover + pass + +else: + import jax + import jax.numpy as jnp + + jax.config.update("jax_enable_x64", True) + + import autolens as al + from autolens.jax.registration import register_tracer_classes + from autolens.point.solver.shape_solver import AbstractSolver + from autoarray.structures.triangles.coordinate_array import ( + CoordinateArrayTriangles, + static_vertex_table, + ) + + def _solver(shape_native=(100, 100), pixel_scales=0.2): + return al.PointSolver.for_grid( + grid=al.Grid2D.uniform( + shape_native=shape_native, pixel_scales=pixel_scales + ), + pixel_scale_precision=0.001, + magnification_threshold=0.1, + use_jax=True, + ) + + def _tracer(einstein_radius=1.6, source=(0.07, 0.07)): + lens = al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.0, 0.0), + einstein_radius=einstein_radius, + ell_comps=(0.05, 0.03), + ), + ) + return al.Tracer( + galaxies=[ + lens, + al.Galaxy(redshift=1.0, point_0=al.ps.Point(centre=source)), + ] + ) + + def _flat_initial_triangles(self, xp): + """The pre-phase-3 JAX step-0 tiling: the flat (3N, 2) per-triangle vertex table.""" + return CoordinateArrayTriangles.for_limits_and_scale( + y_min=self.y_min, + y_max=self.y_max, + x_min=self.x_min, + x_max=self.x_max, + scale=self.scale, + ) + + def _traced_grid_rows(solver, tracer, source, monkeypatch): + """Trace one jitted solve and record the row count of every deflected grid.""" + rows = [] + plane_grid = AbstractSolver._plane_grid + + def recording(self, tracer, grid, xp, plane_redshift=None): + rows.append(int(np.shape(grid)[0])) + return plane_grid(self, tracer, grid, xp, plane_redshift) + + monkeypatch.setattr(AbstractSolver, "_plane_grid", recording) + jax.jit( + lambda: solver.solve(tracer, source_plane_coordinate=source).array + ).lower() + monkeypatch.setattr(AbstractSolver, "_plane_grid", plane_grid) + return rows + + def _solve(solver, einstein_radius, source): + return np.asarray( + jax.jit( + lambda e, s: solver.solve( + _tracer(e, (s[0], s[1])), source_plane_coordinate=(s[0], s[1]) + ).array + )(einstein_radius, jnp.array(source)) + ) + + def _finite_sorted(positions): + positions = positions[np.all(np.isfinite(positions), axis=1)] + return positions[np.lexsort((positions[:, 1], positions[:, 0]))] + + # Generic, near-caustic (the SIE's tangential caustic is ~0.1" across) and outer sources. + SOURCES = [ + (0.07, 0.07), + (0.0, 0.0), + (0.012, 0.0), + (0.0, 0.015), + (0.03, 0.03), + (0.3, -0.2), + (1.2, 0.1), + (1.58, 0.0), + ] + + +def test__step_0_deflects_the_unique_lattice_vertices(monkeypatch): + """ + The shape guard: the first deflected grid of a jitted solve on the 100x100, 0.2" grid is the + 11 859-row unique-vertex table, not the 69 849-row flat one. Red before phase 3. + """ + solver = _solver() + tracer = _tracer() + register_tracer_classes(tracer) + + rows = _traced_grid_rows(solver, tracer, (0.07, 0.07), monkeypatch) + + assert rows[0] == 11859 + assert rows[0] != 3 * 23283 + assert len(rows) == solver.n_steps + + +def test__changing_the_geometry_changes_the_step_0_table(monkeypatch): + solver = _solver() + tracer = _tracer() + register_tracer_classes(tracer) + + solver.scale = 0.3 + rows = _traced_grid_rows(solver, tracer, (0.07, 0.07), monkeypatch) + + expected = static_vertex_table( + float(solver.y_min), + float(solver.y_max), + float(solver.x_min), + float(solver.x_max), + 0.3, + )[0].shape[0] + assert rows[0] == expected + assert rows[0] != 11859 + + +@pytest.mark.parametrize("source", SOURCES) +def test__positions_match_the_flat_table(source, monkeypatch): + """Identical image counts and positions to <= 1e-10 against the pre-phase-3 flat table.""" + solver = _solver() + register_tracer_classes(_tracer()) + + static = _finite_sorted(_solve(solver, 1.6, source)) + + jax.clear_caches() + monkeypatch.setattr(AbstractSolver, "_initial_triangles", _flat_initial_triangles) + flat = _finite_sorted(_solve(solver, 1.6, source)) + + assert static.shape == flat.shape + assert static.shape[0] > 0 + np.testing.assert_allclose(static, flat, rtol=0.0, atol=1e-10) + + +@pytest.mark.parametrize("source", [(0.07, 0.07), (0.0, 0.015), (0.3, -0.2)]) +def test__positions_match_numpy(source): + solver = _solver() + register_tracer_classes(_tracer()) + + positions = _finite_sorted(_solve(solver, 1.6, source)) + + solver_np = al.PointSolver.for_grid( + grid=al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.2), + pixel_scale_precision=0.001, + magnification_threshold=0.1, + ) + positions_np = _finite_sorted( + np.asarray( + solver_np.solve(_tracer(1.6, source), source_plane_coordinate=source).array + ) + ) + + assert positions.shape == positions_np.shape + np.testing.assert_allclose(positions, positions_np, rtol=0.0, atol=1e-10) + + +def _log_likelihood_and_grad(solver): + positions = al.Grid2DIrregular([(1.53, 0.35), (-1.24, -0.28)]) + noise_map = al.ArrayIrregular([0.05, 0.05]) + + def log_likelihood(einstein_radius): + fit = al.FitPositionsImagePairAll( + name="point_0", + data=positions, + noise_map=noise_map, + tracer=_tracer(einstein_radius, (0.07, 0.07)), + solver=solver, + xp=jnp, + ) + return fit.log_likelihood + + value, grad = jax.jit(jax.value_and_grad(log_likelihood))(1.6) + return float(value), float(grad) + + +def test__log_likelihood_and_grad_match_the_flat_table(monkeypatch): + solver = _solver() + register_tracer_classes(_tracer()) + + value, grad = _log_likelihood_and_grad(solver) + + jax.clear_caches() + monkeypatch.setattr(AbstractSolver, "_initial_triangles", _flat_initial_triangles) + value_flat, grad_flat = _log_likelihood_and_grad(solver) + + assert np.isfinite(value) + assert grad != 0.0, "jax.grad through the solver is silently zero" + assert abs(value - value_flat) <= 1e-12 * abs(value_flat) + assert grad == pytest.approx(grad_flat, rel=1e-10) + + +def test__vmap_matches_scalar_solves(): + solver = _solver() + register_tracer_classes(_tracer()) + + sources = jnp.array([(0.07, 0.07), (0.0, 0.015), (0.3, -0.2), (0.03, 0.03)]) + + batched = np.asarray( + jax.jit( + jax.vmap( + lambda s: solver.solve( + _tracer(1.6, (s[0], s[1])), source_plane_coordinate=(s[0], s[1]) + ).array + ) + )(sources) + ) + + for i, source in enumerate(np.asarray(sources)): + scalar = _finite_sorted(_solve(solver, 1.6, tuple(source))) + np.testing.assert_allclose( + _finite_sorted(batched[i]), scalar, rtol=0.0, atol=1e-10 + ) From 972d454e55a945881ab6dde3aea11f01f344cfe6 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 24 Sep 2026 15:16:31 +0200 Subject: [PATCH 2/3] test(point): pin the source-on-step-0-vertex tie case for the static lattice (PyAutoArray#568 phase 3) A source placed bit-exactly on traced step-0 lattice vertex (-0.8, -1.99185843) of a simple SIE: the static-lattice JAX solver returns exactly the two true images (the vertex root and the counter-image ~(0.41013634, 0.89463539)); the pre-phase-3 flat table returns three (a duplicate of the vertex root). Human gate decision 2026-09-24: PASS, pinned as a test. Co-Authored-By: Claude Opus 5.5 --- .../triangles/test_static_lattice_jax.py | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/test_autolens/point/triangles/test_static_lattice_jax.py b/test_autolens/point/triangles/test_static_lattice_jax.py index 7934d7112..8d131e149 100644 --- a/test_autolens/point/triangles/test_static_lattice_jax.py +++ b/test_autolens/point/triangles/test_static_lattice_jax.py @@ -247,3 +247,49 @@ def test__vmap_matches_scalar_solves(): np.testing.assert_allclose( _finite_sorted(batched[i]), scalar, rtol=0.0, atol=1e-10 ) + + +def test__source_on_a_step_0_vertex_returns_the_two_true_images(): + """ + Pinned tie regression (human gate decision, 2026-09-24: PASS). + + The source of a simple SIE is placed bit-exactly on the image of the traced step-0 lattice + vertex v = (-0.8, -1.99185843) (flat triangle vertex 31 732 of the +-9.9" / 0.2" lattice), + and solved under ``jit`` with the source a traced input, as in a likelihood. The system has + exactly two images: the vertex itself and a counter-image at ~(0.41013634, 0.89463539) + (Newton roots of the lens equation; NumPy's 7 positions collapse onto the same two roots). + + The static step-0 lattice returns exactly those two images. The pre-phase-3 flat table + (main before PyAutoArray#568 phase 3) returns three here: it keeps two triangles either side + of the tie vertex, |p - v| ~ 9e-4 each, which both Newton-converge to the vertex root -- it + double-counts one image. That flat control is not self-consistent at such measure-zero ties + either: eager and ``jit`` solves differ in the step-0 kept set and often in the image count, + and a 1e-9 nudge of the source gives two images on both paths. The phase-3 image-count gate + therefore passes this case as the removal of a duplicate, not the loss of an image. + """ + solver = _solver() + register_tracer_classes(_tracer()) + + lattice = CoordinateArrayTriangles.for_limits_and_scale( + y_min=-9.9, y_max=9.9, x_min=-9.9, x_max=9.9, scale=0.2 + ) + vertex = np.asarray(lattice.triangles).reshape(-1, 2)[31732] + np.testing.assert_allclose(vertex, (-0.8, -1.99185843), rtol=0.0, atol=1e-8) + + lens_only = al.Tracer( + galaxies=[_tracer().galaxies[0], al.Galaxy(redshift=1.0)] + ) + deflection = np.asarray( + lens_only.deflections_yx_2d_from( + grid=al.Grid2DIrregular([tuple(vertex)]) + ).array + )[0] + source = tuple(float(value) for value in vertex - deflection) + + positions = _finite_sorted(_solve(solver, 1.6, source)) + + true_images = np.array([vertex, (0.41013634, 0.89463539)]) + assert positions.shape == (2, 2) + np.testing.assert_allclose( + positions, true_images, rtol=0.0, atol=2.0 * solver.pixel_scale_precision + ) From c9ba49fb24fff6e6c83bb8f471005844c9c034ba Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 24 Sep 2026 15:26:48 +0200 Subject: [PATCH 3/3] test: define parametrize constants outside the jax guard (fixes unittest-nojax collection, #568) The module-level @pytest.mark.parametrize decorators read these plain-data constants at collection time; defining them inside the jax-present branch raised NameError on the NumPy-only CI leg. Verified: no-jax collection skips cleanly, with-jax collected count unchanged. Co-Authored-By: Claude Opus 5.5 --- .../triangles/test_static_lattice_jax.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/test_autolens/point/triangles/test_static_lattice_jax.py b/test_autolens/point/triangles/test_static_lattice_jax.py index 8d131e149..5d766d208 100644 --- a/test_autolens/point/triangles/test_static_lattice_jax.py +++ b/test_autolens/point/triangles/test_static_lattice_jax.py @@ -19,6 +19,21 @@ import pytest +# Plain data, defined outside the jax guard: module-level @parametrize reads it at collection +# time, including on the NumPy-only matrix env. +# Generic, near-caustic (the SIE's tangential caustic is ~0.1" across) and outer sources. +SOURCES = [ + (0.07, 0.07), + (0.0, 0.0), + (0.012, 0.0), + (0.0, 0.015), + (0.03, 0.03), + (0.3, -0.2), + (1.2, 0.1), + (1.58, 0.0), +] + + if importlib.util.find_spec("jax") is None: pytestmark = pytest.mark.skip(reason="requires jax (the [optional] extras)") @@ -104,17 +119,6 @@ def _finite_sorted(positions): positions = positions[np.all(np.isfinite(positions), axis=1)] return positions[np.lexsort((positions[:, 1], positions[:, 0]))] - # Generic, near-caustic (the SIE's tangential caustic is ~0.1" across) and outer sources. - SOURCES = [ - (0.07, 0.07), - (0.0, 0.0), - (0.012, 0.0), - (0.0, 0.015), - (0.03, 0.03), - (0.3, -0.2), - (1.2, 0.1), - (1.58, 0.0), - ] def test__step_0_deflects_the_unique_lattice_vertices(monkeypatch):