Merge pull request #749 from PyAutoLabs/feature/point-source-cpu-p3

perf(point): JAX PointSolver step 0 deflects only the unique lattice vertices (PyAutoArray#568 phase 3)

Author: James Nightingale

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,

test_autolens/point/triangles/test_static_lattice_jax.py

@@ -0,0 +1,299 @@
+"""
+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
+
+
+# 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)")
+
+    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]))]
+
+
+
+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
+        )
+
+
+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
+    )