feat: ShapeSolver source→image mappings, subplot_mappings, multiple-image positions (image-source-mappings P2)

[Jammy2211]

Summary

Phase 2 of the image-source-mappings epic. Closes #719.

ShapeSolver becomes the source→image engine for parametric sources, so the colour-matched "one look" mapping figure that Phase 1 (PyAutoArray#517) built for pixelized sources now works for any source, and the same solver output feeds the numerical quantities (per-image magnification, brightest image-plane coordinate of each multiple image for spectroscopic fibre pointing).

• ShapeSolver validation suite (test_autolens/point/triangles/test_shape_solver.py): SIS analytic magnification, agreement with a grid-trace oracle, refinement convergence, all shape types, multi-plane plane_redshift, and the JAX path (deferred — see the audit). The commented-out test in test_extended.py is revived.
• ShapeSolver.image_regions_from / mapping_from: kept triangles grouped into the separate multiple images and exposed as Phase 1 ImageRegion objects (triangle polygons + pixel coverage), plus a per-image find_magnification split.
• autolens/lens/mappings.py (new): mappings_from_fit (pixelized → engine A, parametric → ShapeSolver with a Circle at the light centre), multiple_image_positions_from (arcsec) and multiple_image_pixel_coordinates_from (pixels), magnifications_from, source_mapping_from, traced_region_from. No RA/Dec in the library — the workspace guide (Phase 3) shows the astropy conversion.
• subplot_mappings(fit) in autolens/imaging/plot/fit_imaging_plots.py, exported from autolens.plot, wired to the existing plots.yaml: subplot_mappings key.

Depends on PyAutoArray#518 (Phase 2a, Shape.contains / Shape.boundary); CI here is source-installed against PyAutoArray main and goes green once #518 merges. Merge order: #518 first.

ShapeSolver audit

The 2026-08-19 audit's findings, re-verified against the code as it stands, plus what the new validation suite found. Every "fixed" row is fixed in the solver, with a regression test named below; nothing was worked around in the mappings layer.

Finding	Status	Evidence
find_magnification hardcodes xp=np and never consults self._xp, so use_jax=True is silently ignored	fixed	xp=None now resolves through ShapeSolver._xp_from, which consults the flag and raises instead of ignoring it (see the next row for why raising, not switching). solve_triangles and steps likewise take xp=None rather than requiring it positionally. test_shape_solver.py::test_jax_solver_is_rejected_rather_than_silently_wrong
NEW use_jax=True is not merely ignored, it is wrong: the JAX triangle containers truncate every refinement step to ArrayTriangles.MAX_CONTAINING_SIZE = 15 triangles to keep static shapes. That is ample for a Point and meaningless for a Shape with area. Measured on an Isothermal sphere: JAX keeps 15 of the ~800 triangles NumPy keeps and reports a magnification of 0.13 where the truth is 6.86. Routing xp through self._xp without noticing this would have replaced a silently-ignored flag with a silently wrong number.	deferred (guarded)	Lifting the cap is a redesign of the JAX containers — the cap is what makes their shapes static, and an extended source has no static bound — so ShapeSolver raises NotImplementedError naming MAX_CONTAINING_SIZE. test_shape_solver.py::test_jax_and_numpy_kept_triangles_agree is xfail(strict=True) with the reason recorded, and is the test to un-xfail when the containers grow a dynamic kept set.
JAX triangle area cannot survive jax.jit (__len__ returns a traced count_nonzero)	deferred	Still reproduces on this branch: jax.jit(lambda: solver.solve_triangles(...).area) raises TracerIntegerConversionError (__index__ on a traced int64[]). Same root cause family as the cap above — deferred with it, and recorded in the ShapeSolver class docstring.
No per-image split	fixed	find_magnification(per_image=True) and image_regions_from group the kept triangles into connected components by shared vertices. On an Isothermal sphere with theta_E = 1.0 and a source at beta = 0.3 they give 4.324 and 2.330 against the analytic 1 + theta_E/beta = 4.333 and theta_E/beta - 1 = 2.333 (0.2 % and 0.13 %). test_shape_solver.py::test_per_image_magnification_matches_the_analytic_isothermal_sphere
No magnification-threshold filter — magnification_threshold was accepted by the constructor and used only by PointSolver.solve, never by ShapeSolver	fixed	Each connected component's magnification (its share of the kept area over the shape's area) is compared against the threshold in image_regions_from / per_image. Without it the singular centre of an Isothermal sphere contributes a spurious two-triangle "image" of magnification 0.004 at the lens centre. test_shape_solver.py::test_magnification_threshold_removes_the_central_image
Image area over-covers by up to one triangle edge per boundary triangle, with no convergence check or error estimate	fixed (documented + measured)	find_magnification's docstring now states the bound (perimeter * side_length, so halving pixel_scale_precision halves the error, and the area approaches the truth from above) and separates it from the physical finite-source error, quoting measured numbers: at pixel_scale_precision=0.005 a Circle of radius 0.2 / 0.1 / 0.05 gives 7.12 / 6.76 / 6.66 against an analytic 6.67. test_shape_solver.py::test_total_magnification_converges_as_the_source_shrinks asserts the error shrinks with the radius and ends below 1 %.
Sole test commented out since ~2024; sole consumer an unused workspace demo	fixed	test_extended.py::test_solver_basic (the identity-lens magnification of 1) revived, plus two siblings; 16 further tests in the new test_shape_solver.py. The solver now has three in-library consumers (autolens/lens/mappings.py, subplot_mappings, the visualizer).
NEW the kept-triangle area is not monotone over steps() — asserting monotonicity would have pinned a property the algorithm does not have. Over seven steps it runs 0.2143413, 0.2154238, 0.2126498, 0.2122439, 0.2123877, 0.2124807, 0.2124794: step 0 keeps whole coarse triangles, and each later step filters the neighbourhood of the previous kept set, which is larger than that set.	documented	The invariants that do hold are pinned instead: the search envelope step.neighbourhood.area is monotone non-increasing, and the final refinement changes the kept area by less than any earlier one. test_shape_solver.py::test_refinement_shrinks_the_search_and_converges, and the steps docstring.
NEW exact float equality does not identify coincident triangle vertices. A lattice point reached from two triangles is computed by two different sequences of floating-point operations and differs in the last ulp, so AbstractTriangles.indices (built by np.unique) splits the two images of an Isothermal sphere into 397 fragments rather than 2.	fixed	_vertex_labels_from quantises vertices to a thousandth of a triangle edge before deduplicating — far below the spacing of distinct vertices, far above the ulp.
NEW (PyAutoArray) CoordinateArrayTriangles{,Np}.for_limits_and_scale named its first pair of limits x_min/x_max and tiled them along element 0 of every vertex, while element 0 is y everywhere else. Both the keyword caller (AbstractSolver._initial_triangles) and the positional caller (AbstractTriangles.for_grid) therefore tiled the transposed rectangle. A square grid hides it; on a 24×80 grid of 0.05" pixels PointSolver.solve found one of an Isothermal's two images instead of both, the missing one lying well inside the grid.	fixed on PyAutoArray#518	Commit 67b1cb75. test_autoarray/structures/triangles/test_coordinate.py::test_for_limits_and_scale__element_0_spans_the_y_limits and ::test_for_grid__rectangular_grid_is_tiled_the_same_way_round
NEW (PyAutoArray) a stray bbbb statement in ArrayTrianglesNp.with_vertices — a NameError on every call — committed 2025-11-13 and uncaught because the solver only ever reaches the CoordinateArrayTrianglesNp override of that name	fixed on PyAutoArray#518	Commit 67b1cb75. ::test_array_triangles_with_vertices_returns_new_triangles
NEW (PyAutoArray) plot_regions drew one label per region entry, at the mean of all that region's polygons. A mapping's image regions are its multiple images, on opposite sides of the lens, so the label landed at their midpoint — on empty sky, labelling nothing, in exactly the figure the overlay exists for.	fixed on PyAutoArray#518	Commit a9a9120a. The label is repeated once per polygon at that polygon's own centre; a single-polygon region is unchanged. test_autoarray/plot/test_utils.py::test__plot_regions__multi_polygon_region_is_labelled_once_per_polygon and ::test__plot_regions__single_polygon_region_is_labelled_once
NEW subplot_mappings originally wrapped the engine in try/except Exception → logger.warning, so a source plane which could not be mapped drew four panels with no regions on them	fixed (guard removed)	A figure with no regions reads as "this source has no multiple images", which is a wrong answer rather than a missing one, so the error now propagates and names shape=. The identical guard around the visualizer's call in PlotterImaging.fit_imaging is removed for the same reason. test_fit_imaging_plots.py::test__subplot_mappings__unmappable_plane__raises_rather_than_drawing_nothing
NEW ShapeSolver(use_jax=True) on a JAX-free install raised ModuleNotFoundError from _xp before the rejection fired — _xp_from resolved xp=None through self._xp (which imports jax.numpy) and only then checked the result, so the user got the wrong error and one which blamed the environment	fixed	The use_jax case now raises before self._xp is touched, and the shared message is a plain class-level string, so nothing on the rejection path imports jax. Caught by the unittest-nojax CI leg; test_use_jax_solver_is_rejected_rather_than_silently_wrong deliberately carries no importorskip and now covers find_magnification / solve_triangles / steps, with the explicit xp=jax.numpy route split into its own guarded test. Commit c376f10

API Changes

PyAutoLens — new public API

• autolens.lens.mappings (new module, exported as al.mappings):
  • source_mapping_from(tracer, grid, shape, plane_index=-1, plane_redshift=None, pixel_scale_precision=None, magnification_threshold=0.1) -> aa.Mapping
  • traced_region_from(tracer, region_grid, plane_index=-1) -> np.ndarray — source-plane convex hull of a traced image-plane region, as a closed (N, 2) (y, x) loop.
  • mappings_from_fit(fit, plane_index=-1, shape=None, grid=None, pixel_scale_precision=None, magnification_threshold=0.1, **clump_kwargs) -> List[aa.Mapping] — engine A when the plane has(cls=aa.Pixelization), engine B otherwise.
  • multiple_image_positions_from(fit, plane_index=-1, use_centroid=False, **mapping_kwargs) -> aa.Grid2DIrregular — the brightest pixel (or flux-weighted centroid) of each multiple image, in arcsec.
  • multiple_image_pixel_coordinates_from(fit, plane_index=-1, use_centroid=False, **mapping_kwargs) -> List[Tuple[float, float]] — the same, as 1-based pixel-centre WCS pixel coordinates via Geometry2D.pixel_coordinates_wcs_2d_from. No RA/Dec in the library.
  • magnifications_from(fit, plane_index=-1, shape=None, grid=None, pixel_scale_precision=None, magnification_threshold=0.1, **clump_kwargs) -> List[float] — engine B returns the absolute per-image magnification (find_magnification(per_image=True)); engine A has no source-plane area to divide by and returns each image's share of its clump's total model flux (summing to 1 per clump). The two scales are different and the docstring says so.
• ShapeSolver.image_regions_from(tracer, shape, plane_redshift=None, grid=None, xp=None) -> List[aa.ImageRegion] (new)
• ShapeSolver.mapping_from(tracer, shape, plane_redshift=None, grid=None) -> aa.Mapping (new)
• autolens.imaging.plot.fit_imaging_plots.subplot_mappings(fit, plane_index=-1, shape=None, output_path=None, output_filename="mappings", output_format=None, colormap=None, threshold=None, min_pixels=None, total_clumps=None, pix_indexes=None, weight_threshold=0.0, pixel_scale_precision=None, magnification_threshold=0.1, region_alpha=0.25, image_plane_lines=None, image_plane_line_colors=None, source_plane_lines=None, source_plane_line_colors=None, title_prefix=None) (new), exported as aplt.subplot_fit_imaging_mappings. Writes mappings_{plane_index}.png.

PyAutoLens — changed

• ShapeSolver.find_magnification(tracer, shape, xp=np, plane_redshift=None) → (tracer, shape, xp=None, plane_redshift=None, per_image=False). xp now defaults to None (resolving to self._xp) rather than np; per_image=True returns List[float] instead of float. A ShapeSolver built with use_jax=True now raises NotImplementedError where it previously returned a NumPy answer — deliberate, see the audit.
• AbstractSolver.solve_triangles(tracer, shape, xp, plane_redshift=None) and AbstractSolver.steps(...): xp is now optional (xp=None → self._xp). Back-compatible for every existing caller.
• AbstractSolver.__init__ gains a non-pytree self.grid attribute (None unless set by for_grid), which for_grid now populates so image_regions_from knows which pixels to report on. It is deliberately excluded from tree_flatten.
• _plot_source_plane (private) gains regions / region_alpha / region_labels.
• config/visualize/plots.yaml: the inversion: subplot_mappings key — dead since the mapping visuals were removed in 2025 — is now read by PlotterImaging.fit_imaging, which writes the fit-level subplot for every source plane. Its comment is updated to say it now covers parametric sources too. The default stays false.

PyAutoArray (on #518, this branch's dependency)

• plot_regions now draws a region's label once per polygon rather than once per region (behaviour change for any multi-polygon region; single-polygon regions are unchanged).
• CoordinateArrayTriangles.for_limits_and_scale and CoordinateArrayTrianglesNp.for_limits_and_scale: parameter order changed from (x_min, x_max, y_min, y_max, scale) to (y_min, y_max, x_min, x_max, scale), matching AbstractTriangles.for_limits_and_scale and ArrayTrianglesNp. Keyword callers are unaffected in meaning but now tile the rectangle they asked for; the one positional caller (AbstractTriangles.for_grid) is fixed by the reorder.

Tests

Run with the worktree PyAutoArray on PYTHONPATH (Shape.contains / Shape.boundary from #518).

Targeted set — pytest test_autolens/lens/test_mappings.py test_autolens/point/triangles test_autolens/imaging/plot -q:

87 passed, 1 xfailed

File	Result
test_autolens/point/triangles/test_shape_solver.py (new, 18 tests)	17 passed, 1 xfailed
test_autolens/point/triangles/test_extended.py (revived)	3 passed
test_autolens/lens/test_mappings.py (new)	12 passed
test_autolens/imaging/plot/test_fit_imaging_plots.py (+3)	18 passed
the rest of test_autolens/point/triangles	unchanged, passing

Full suite — pytest test_autolens -q:

610 passed, 1 xfailed

and with JAX made unimportable (the unittest-nojax leg), test_autolens/point/triangles/test_shape_solver.py is 15 passed, 2 skipped — the two skips being the tests which genuinely need a jax.numpy module to pass in.

PyAutoArray (the phase-2a branch, with the two triangle fixes) — pytest test_autoarray -q:

1410 passed

including 3 new tests in test_autoarray/structures/triangles/test_coordinate.py and 2 in test_autoarray/plot/test_utils.py. pytest test_autoarray/plot test_autoarray/inversion/plot -q (Phase 1's aa subplot_mappings included): 42 passed.

xfails (one, strict=True):

• test_shape_solver.py::test_jax_and_numpy_kept_triangles_agree — the MAX_CONTAINING_SIZE = 15 cap on the JAX triangle containers, deferred; see the audit. Its reason string names the condition for removing it.

No JAX in the unit tests except that one case, which is importorskip-guarded. black was applied to the four files this PR adds or rewrites; the pre-existing files it edits were not black-clean on main, so they are left in their own style rather than reformatted wholesale.

Verification

(i) Renders — a simulated Isothermal (axis ratio 0.7 / 0.9, theta_E = 1.6) + Sersic source, 100×100 at 0.1", fitted at the true model, then the same dataset fitted with a RectangularUniform(30, 30) + Constant pixelization:

• parametric (engine B): /tmp/claude-1000/-home-jammy-Code-PyAutoLabs/7ac5cc85-dfd9-41e2-a4bd-725f33f0f24f/scratchpad/p2_verify/mappings_1.png
• pixelized (engine A): /tmp/claude-1000/-home-jammy-Code-PyAutoLabs/7ac5cc85-dfd9-41e2-a4bd-725f33f0f24f/scratchpad/p2_verify/mappings_pix_1.png

Both figures draw the source region filled and outlined in the source-plane panels, its multiple images in the same colour and with the same label on each image (the PyAutoArray#518 label fix above) in the two image-plane panels, and the critical curve / caustic. The two engines return the same multiple-image positions on the same data — (1.65, 1.15) and (-0.95, -0.65) — which is an independent check that the ShapeSolver engine agrees with the mapping-matrix engine.

(ii) Multiple-image positions against PointSolver — a full-resolution simulated quad (PYAUTO_SMALL_DATASETS unset, so PointSolver.solve is not short-circuited): Isothermal(axis_ratio=0.7, angle=45, theta_E=1.6) with a compact SersicSph (effective_radius=0.05) at (0.05, 0.03), inside the inner caustic; 120×120 at 0.1".

multiple_image_positions_from (y, x)	PointSolver.solve (y, x)	separation (")
(−0.950, +1.350)	(−0.940, +1.348)	0.0103
(+1.050, +1.150)	(+1.034, +1.120)	0.0343
(+1.350, −0.950)	(+1.356, −0.975)	0.0259
(−1.050, −0.950)	(−1.027, −0.967)	0.0290

Four positions each, every pair matched within 0.034" — about a third of a pixel (0.1"), and well inside the "within one pixel" requirement. The mapped positions are exact pixel centres by construction (they are the brightest pixel of each image region), which is why the residual is of order half a pixel.

(iii) Grid-trace oracle — for a Circle of radius 0.1 on an Isothermal sphere over a 100×100 / 0.05" mask, the oracle (trace the image grid, shape.contains, image_regions_from_slim_mask) finds 2 regions covering 86 pixels and the solver finds 2 regions covering 124. Every oracle pixel is in the solver's regions (the oracle is a strict subset) and no solver pixel lies outside a one-pixel dilation of the oracle — 0 pixels outside. The difference is the expected one: the solver marks every pixel containing a point which traces inside the shape, the oracle only pixels whose centre does.

Epic

image-source-mappings phase 2 — ledger PyAutoMind/draft/feature/autoarray/image_source_mappings_epic.md. Phase 3 (workspace guide + tutorial_2_mappers rewrite) opens after this and #518 release. Label pending-release.

🤖 Generated with Claude Code

https://claude.ai/code/session_011EVnQ4sEYrRM7GCvFyveUA