[Jammy2211]

Summary

The JAX PointSolver now tests step-0 containment on the static initial lattice without building the (N, 3, 2) triangle array. That gather was the largest single cost of the single-source point likelihood: about 0.9 ms of ~2.1 ms on RAL CPU in phase 4a. This is phase 4b of the point-source CPU speed-up (#579). The solver geometry and the defaults are unchanged. The kept triangle set, and so every result, is bit-identical.

• static_lattice_layout (lru-cached, ints only) derives the closed-form layout of static_vertex_table's index map, checks it element by element against the real map, and falls back (grid=None) if any parity class breaks the strided pattern.
• ArrayTriangles.containing_indices sends an exact Point (not the Circle/Triangle/Polygon/Square subclasses, which override mask) through _step0_point_mask when the triangles carry a Step0Layout. That only happens on the step-0 lattice from for_limits_and_scale(static_vertices=True) → with_vertices; derived lattices (for_indexes, up_sample, neighborhood) drop it, so refinement is untouched.
• The private switch array._STEP0_CONTAINMENT selects the route: gather (the old path), nopad, components (six 1-D gathers) or structured (the default: strided slices of the traced vertex table, with no gather at all). It is read at trace time and exists for the profiling A/B.

API Changes

None — internal changes only. The new step0_layout=None keyword on ArrayTriangles / CoordinateArrayTriangles is optional and set internally. Step0Layout and static_lattice_layout are new internal helpers.
See full details below.

Test Plan

• test_autoarray/structures/triangles/test_coordinate_jax.py: 54 passed.
  • A bit-identity fuzz of each route against gather under jit + vmap over 5 geometries, including alternating key-row widths, a padded grid and an SIS-deflected lattice. It uses about 5.7e4 points per route; PYAUTO_STEP0_FUZZ_SCALE=4 gives about 1.5e5 and passed.
  • A check that every geometry actually gets a structured layout.
  • A refinement-path-unchanged test, including a Circle.
  • An HLO guard that the default route never builds the full triangle array. It fails on main 14d6336 and on the gather/nopad routes, and passes on components/structured.
• PyAutoArray full suite: 1732 passed.
• PyAutoLens full suite against this branch: 758 passed, 1 xfailed. test_static_lattice_jax.py, including the step-0 vertex tie test, passes 16/16 on all four routes.
• Interleaved A/B (solver_config_sweep.py --step0-route; 20 rounds × 20 calls; bootstrap 90 % CI). Every route is bit-identical in log L (fiducial 7.743201200876812), positions, image counts, grad (with register_model, non-zero) and vmap 1/4/16. Completeness is 200/200 prior and 200/200 stress.

run	gather	components	structured	vmap-16 structured ×gather
RAL EPYC 7702, CPU-only (job 357335)	3.184 ms	1.979 (1.61×)	1.971 (1.62× [1.60, 1.63])	1.87×
laptop, 8 threads	2.221 ms	1.555 (1.43×)	1.613 (1.38×)	1.91×
RAL A100 (job 357322)	0.837 ms	0.839	0.838 (1.00×)	no regression

Step-0 containment falls by 66 %. Compile time is +3–5 %. structured uses 47 % less temporary memory. The quotable Xeon 8490H row (job 357321, queued on a saturated ral partition) will land with the autolens_profiling data PR.

Full API Changes (for automation & release notes)

Added

• autoarray.structures.triangles.array.Step0Layout(n_rows, n_cols, grid=None): frozen dataclass (internal)
• autoarray.structures.triangles.coordinate_array.static_lattice_layout(y_min, y_max, x_min, x_max, scale): lru-cached (internal)
• ArrayTriangles(..., step0_layout=None), CoordinateArrayTriangles(..., step0_layout=None): optional keyword, set internally

Changed Behaviour

• ArrayTriangles.containing_indices(Point) on the static step-0 lattice now uses strided slicing instead of the (N, 3, 2) gather. The output is bit-identical.

Part of #579 (library half; the profiling data PR follows).

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

[Jammy2211]

Workspace PR: PyAutoLabs/autolens_profiling#330