perf: static PointSolver step-0 lattice (point-source CPU phase 3)

[Jammy2211]

Overview

Phase 2 of the point-source CPU speed-up campaign (Mind prompt pointsolver_cpu_speed_phases_2_4.md; phase 1 record complete/2026/09/point-source-cpu-p1.md, autolens_profiling#297/#298). On the JAX PointSolver path, CoordinateArrayTriangles._vertices_and_indices calls jnp.unique(..., size=3N, fill_value=nan). Static shapes make the "deduplicated" table exactly 3N rows, so it saves zero deflection evaluations and costs a lexicographic sort of 3N fp64 rows, run twice per refinement step. On the phase-1 RAL baseline (PyAutoArray 22e6d608) the step-0 ray-trace prefix is 22.5 of the 24.7 ms fused call. September 17 scratch evidence (reported, monkeypatched laptop): 40.0 → 8.1 ms simple, 117 → 48 ms per cluster solve, likelihood bit-identical.

This task removes the sort in PyAutoArray only (API unchanged; NumPy sibling untouched), adds the missing JAX unit tests, and measures the change with an interleaved in-process A/B cell on a quiet RAL CPU host, recorded in the campaign ledger.

Session model: Fable 5.1 (architect); execution delegated to Opus subagents. Heart at plan time: RED (release validation FAILED (stage integrate)); PR-open and merge need the live human override.

Plan

• PyAutoArray: make the JAX _vertices_and_indices return the flat 3N vertex table and an arange index map (no jnp.unique); document why; keep vertices/indices/with_vertices signatures. NumPy class unchanged. PyAutoLens needs no change (its _plane_triangles is the only consumer).
• PyAutoArray: new JAX unit tests — vertex/triangle round-trip, NaN padding, NumPy-vs-JAX kept-triangle parity on a small lattice, and an HLO assertion that the traced containment has no sort op.
• Red controls: autolens_workspace_test point-source jax_likelihood/* vmap pins + NumPy-vs-JAX parity and jax_grad/gradient.py finite-difference checks against the worktree library; autolens_profiling image_plane_solved pin.
• autolens_profiling: new interleaved A/B cell (control = old unique implementation injected, nodedup = new, library = installed) with bit-identity, position, gradient and vmap gates, dispersion and HLO sort counts; RAL ral submit; campaign-note phase-2 section.
• Ship library-first: PyAutoArray PR, then the profiling PR linked to it.

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary, library)
• autolens_profiling (A/B cell, submit, results, note)
• read-only: PyAutoLens (autolens/point/solver/shape_solver.py:426-443 consumer), autolens_workspace_test (red controls)

Branch Survey

Repository	Current Branch	Dirty?
array/PyAutoArray	main (11b9347)	clean
lens/autolens_profiling	main (0435184)	1 untracked (dataset/abell_1201/)
lens/PyAutoLens	main (2aaa1c1a8)	clean

Suggested branch: feature/point-source-cpu-p2
Worktree: ~/Code/PyAutoLabs-wt/point-source-cpu-p2/ (PyAutoArray + autolens_profiling)
Conflict guard: PyAutoArray and autolens_profiling are claimed by certified-positive-solver (PRs #567/#299 MERGED, awaiting its close-out); file sets disjoint (autoarray/util/jax_active_set.py, settings, inversion dispatch, library_solver_injection.py) → parallel worktree, parallel-claim: recorded.

Implementation Steps

1. autoarray/structures/triangles/coordinate_array.py _vertices_and_indices (JAX): flat = self.triangles.reshape(-1, 2); indices = jnp.arange(flat.shape[0]).reshape(-1, 3); return (flat, indices). Docstrings on vertices/indices: not deduplicated on the JAX path (static shapes make dedup a no-op that costs a sort; NaN padding rows trace to NaN and every Shape.mask rejects them). with_vertices unchanged.
2. test_autoarray/structures/triangles/test_coordinate_jax.py (importorskip("jax")): round-trip (vertices.reshape(-1,3,2) == triangles, with_vertices(vertices).triangles == triangles); NaN padding after for_indexes([...,-1]) stays NaN and containing_indices(Point) returns only finite indices padded with -1; parity with CoordinateArrayTrianglesNp (for_limits_and_scale → identity with_vertices(vertices) → containing_indices(Point) → kept coordinates rounded 8 dp as sets) below the 15-triangle cap; jax.jit(...).lower(...).as_text() of the containment contains no sort.
3. Gates: pytest test_autoarray -q; downstream: autolens_workspace_test scripts/point_source/jax_likelihood/{image_plane,point,source_plane,fluxes_time_delays}.py and jax_grad/gradient.py against the worktree library (pins rtol 1e-4, parity, FD gradients); autolens_profiling scripts/point_source/likelihood_runtime/image_plane_solved.py once (pin 7.743201200876817, eager==JIT, pinned_drift empty).
4. autolens_profiling scripts/point_source/likelihood_breakdown/vertex_dedup_ab.py: injection context managers (pattern: scripts/misc/likelihood_breakdown/library_solver_injection.py) swapping CoordinateArrayTriangles._vertices_and_indices between control (old unique code, verbatim), nodedup (new code), library (installed, self-labelled by lowering a small solve and counting sort); distinct function objects + jax.clear_caches() per route; fused production likelihoods with varying parameters for the simple solved/plain model (phase-1 harness) and the two-source cluster Repeat/RepeatSolved (cluster cell model/grid); interleaved rows ≥ 20 rounds × ≥ 20 warmed calls, median/p10/min/max, ms/likelihood and vmap-4 throughput; gates: log L == (report Δ if not), positions array_equal, jax.grad allclose 1e-10, vmap parity; provenance (provenance.py), len(os.sched_getaffinity(0)), XLA flags, loadavg, HLO sort counts, cost_analysis() FLOPs, MDI from control dispersion; outputs results/breakdown/point_source/vertex_dedup_ab_<config>.{json,png}; smoke early-exit.
5. hpc/batch_cpu/submit_breakdown_point_source_vertex_dedup_ab_ral_cpu_fp64 (clone of the phase-1 submit; ral, 8 CPUs, fresh JAX cache, provenance, --config-name hpc_ral_cpu_fp64); run from a RAL worktree of the profiling branch; shared PyAutoArray on RAL stays on main (the cell injects both variants).
6. Laptop witness run, then RAL run; rsync results; validate gates.
7. results/notes/point_source_cpu_campaign.md phase-2 section (mechanism, diff, A/B tables laptop+RAL, gates, dispersion/MDI, sort counts, ulp caveat outcome, accept/reject decision, phase-3 handoff); disposition table updated; build_readme.py regen (register a label only if the A/B JSON is dashboard-tabled).
8. Ship: ship_library (PyAutoArray PR, ## API Changes behavioural note, pending-release), then ship_workspace (profiling PR, ## Upstream PR); /prm library-first.

Key Files

• array/PyAutoArray/autoarray/structures/triangles/coordinate_array.py — the fix (_vertices_and_indices, docstrings)
• array/PyAutoArray/autoarray/structures/triangles/{array.py,shape.py,coordinate_array_np.py} — read-only context (ArrayTriangles.triangles gather, Shape.mask NaN behaviour, NumPy sibling)
• array/PyAutoArray/test_autoarray/structures/triangles/test_coordinate_jax.py — new tests
• lens/PyAutoLens/autolens/point/solver/shape_solver.py:426-443 — sole consumer (unchanged)
• lens/autolens_profiling/scripts/point_source/likelihood_breakdown/vertex_dedup_ab.py — new A/B cell
• lens/autolens_profiling/hpc/batch_cpu/submit_breakdown_point_source_vertex_dedup_ab_ral_cpu_fp64 — new submit
• lens/autolens_profiling/results/notes/point_source_cpu_campaign.md — phase-2 section
• lens/autolens_workspace_test/scripts/point_source/{jax_likelihood/*,jax_grad/gradient.py} — red controls (run, not edited)

Out of scope

Phases 3–4; GPU rows; post-fix re-run of the phase-1 breakdown cells on RAL (needs the released library on the shared install — first phase-3 step); the missing solver_use_jax_parity.py referenced by PyAutoLens test_use_jax.py (follow-up).

Original Prompt

Click to expand starting prompt

Point-source CPU speed-up campaign — phases 2–4: redundant-sort removal and measured iteration

Type: research
Target: autolens_profiling
Repos:

• autolens_profiling
• PyAutoLens
• PyAutoArray
  Themes:
• point-source
• profiling
• cluster
• jax
  Difficulty: large
  Autonomy: supervised
  Priority: normal
  Status: formalised
  Consequence: judge
  Review-minutes: 20
  Unattended: ready
  Epic: cluster-strong-lensing
  Filed: 2026-09-17
  Updated: 2026-09-23
  Parent-record: complete/2026/09/point-source-cpu-p1.md

Phase 1 shipped — remainder re-filed (2026-09-23)

Phase 1 ("Reproduce and publish the CPU evidence") is COMPLETE: autolens_profiling#297 /
PR #298 merged at 9f5a3ba, record complete/2026/09/point-source-cpu-p1.md, ledger
autolens_profiling/results/notes/point_source_cpu_campaign.md. This prompt now carries
phases 2–4 only. Next bounded task = phase 2: remove the JAX-only throwaway
jnp.unique at CoordinateArrayTriangles._vertices_and_indices as used by
_plane_triangles (PyAutoArray + PyAutoLens, library-first), red control = bit-identical
likelihood + cluster solved positions + gradient/vmap parity against the frozen baseline
(PyAutoArray 22e6d608, PyAutoLens 2aaa1c1a8; RAL ral partition rows
hpc_ral_cpu_fp64: point-source fused solved 24.69 ms, cluster fused plain 128.05 ms),
interleaved A/B on RAL with distinct function objects + jax.clear_caches(), medians and
dispersion, fused control, len(os.sched_getaffinity(0)) and XLA intra-op threads recorded.
The GPU campaign takes its A100 baseline on the frozen revisions before phase 2 merges, or
from those SHAs afterwards.

Campaign contract (2026-09-19)

This execution plan supersedes the earlier deliverable/gate wording retained below.
This is one of two optimization campaigns in a three-task plan. The independent
shared breakdown task is their common prerequisite.
It is a phased campaign: at start-dev, issue only the next bounded phase (one task /
one PR per member), retaining this prompt as the campaign intent until all phases
are resolved. Do not attempt a cross-library, multi-PR campaign as a single task.
Use start-dev and the applicable library/workspace worktree and ship procedures;
obtain implementation-plan approval before source edits. No profiling or library
implementation was performed during this consolidation.

Measurement and acceptance contract

• Use @autolens_profiling for timing and versioned JSON + PNG evidence, with
  README/dashboard regeneration. Correctness evidence belongs in library tests
  and @autolens_workspace_test, not a timing-only assertion of scientific validity.
• Record exact library/profiling commits, JAX/jaxlib versions, device, precision,
  XLA flags, thread settings, model/data seed, source planes, solver grid/scale,
  precision, neighborhood degree, capacity, warm-up, repetitions and cache state.
  Historical numbers are leads, not comparable current baselines.
• Separate tracing/lowering, compilation, first execution and warmed runtime.
  Synchronize device outputs with block_until_ready; pass varying parameter
  inputs through the production likelihood so constant folding cannot fake work.
  Report repeated/interleaved A/B medians and dispersion on identical hardware,
  both ms/likelihood and batch throughput, plus memory where relevant.
• Retain a fused end-to-end production likelihood control. Prefix timing
  differences can change fusion and contain noise: report residuals/negative
  differences honestly and corroborate with a device trace rather than claiming
  independently timed steps sum to the fused runtime.
• Compare likelihood, image counts/positions, NaN padding/masks, magnification
  filtering and source-redshift handling. Cover perturbed models, doubles/quads,
  near-caustic/critical configurations and cluster multi-source/multiplane cases.
  Preserve custom_jvp, eager/JIT/vmap parity and gradient correctness; distinguish
  nondifferentiable image-topology transitions from failures in smooth regions.
• Each iteration: baseline -> one hypothesis -> bounded prototype -> correctness
  gate -> repeated full-likelihood A/B -> accept/reject -> reprofile and rank the
  remaining bottleneck. State minimum detectable improvement from observed noise;
  keep a change only for a repeatable material gain without correctness or
  unacceptable compile/memory regressions. Record negative results too.
• Stop when remaining cost is explained and no worthwhile measured lever remains,
  or a concrete external blocker prevents the next experiment. Do not promise the
  historical speedup, optimize indefinitely, or weaken correctness to hit a target.

Split request (verbatim)

break into 3, with the first build the shared likelihood breakdown first, which I guess is CPU or GPU agnostic but maybe not, then PR them into main

Original consolidation request (verbatim)

We did two reviews or assessments of the point source likelihood function recently one for CPU which was JAX and numba sparse (it concluded numba spaerse not worth it) and one for GPU. We may of made some prompts but I want you to assess all that the review put forward and ultimately end with two mind task or prompts, which could be epics, which will profile them with autolens_profiling and iteratively work on the speed up. One was focused in particular on writing an autolens_profiling likelihood_breakdown script, which it may of wrote or just planned, this would likely be the task before we go into specific CPU or GPU speed up

CPU campaign: dependencies and phase order

Start condition: task 1 shipped in
autolens_profiling#293
and is recorded in complete/2026/09/point-source-shared-breakdown.md; it owns
the shared scripts/point_source/likelihood_breakdown/ instrument and CPU
reference results under results/breakdown/point_source/.
CPU source optimization waits for that merged instrument and CPU baseline, not
for the GPU campaign. Preserve the exact unoptimized library revisions and
configuration so task 3 can reproduce an A100 baseline even if CPU fixes land
first. Avoid two branches changing the same solver at once.

1. Reproduce and publish the CPU evidence. Re-run the simple solved likelihood
   and the 13-component, two-source cluster case with the shared harness. The
   September 17 scratch note/JSONs are not committed in the inspected profiling
   tree; recover them if available, otherwise reproduce and explicitly label the
   old measurements as reported evidence. Measure unsolved and solved paths
   separately; do not attribute a likelihood-variant difference to hardware.
2. Remove redundant vertex deduplication, if reproduced. The recorded
   _vertices_and_indices sort is still present in the inspected PyAutoArray
   checkout. Test the JAX-only throwaway conversion at _plane_triangles, not
   blanket removal of all unique/remove_duplicates operations. Reproduce the
   red control on the original path, confirm masks/index semantics and full
   likelihood/gradient parity, then ship the bounded library change before
   refreshing workspace results. The recorded 40.0 -> 8.1 ms simple and
   117 -> 48 ms/source cluster gains are hypotheses to reproduce on current code.
3. Precompute static initial geometry. If deflections dominate after phase 2,
   evaluate construction-time NumPy unique vertices plus an index map for the
   initial lattice, reused as immutable JAX inputs. Count actual traced vertices,
   report setup/memory/amortization, and verify cache invalidation when solver
   geometry changes. Never cache model-dependent deflections or warm-start from
   the preceding sampler call. This lever was proposed, not measured.
4. Profile the residue and iterate. Rank the following by new evidence:
   (a) safe grid-extent guidance; (b) initial scale versus refinement count;
   (c) MAX_CONTAINING_SIZE/neighborhood fan-out; (d) cluster dPIE/NFW deflections.
   Grid/capacity changes require image-completeness evidence and must be reported
   as separate configurations, not silently substituted into the speed comparison.
   Profile components through the existing scripts/lens/deflections/ surface
   when needed; any resulting PyAutoGalaxy work is a separately scoped phase.

Disposition of every CPU assessment recommendation

• Numba/sparse rewrite: do not pursue by default. The assessment found no
  reusable kernel and existing unpadded NumPy was 96 ms/solve, about 12x slower
  than the patched JAX path. Reopen only if new post-fix measurements demonstrate
  a substantial unsolved bottleneck and preserve JVP/vmap/GPU contracts.
• Step count: seven refinement steps reportedly cost ~3 ms in total; fewer
  than four degraded the likelihood. Low priority, never a blind accuracy trade.
• Extent: reported 30x30 / 100x100 / 140x140 timings were 6.4 / 39.8 / 72 ms;
  a fiducial bit-identical likelihood does not prove coverage across a prior.
• Fit/chi-squared, beta-star and magnification filter: reported 0.34% of the
  original call. Re-rank after the main fix; don't assume the old fraction holds.
• Duplicate model_data solve: reportedly eliminated by XLA CSE; verify in
  fused execution before proposing Python caching as a performance fix.
• Unmeasured controls: vmap batches 1/4/16; post-fix grid/n_steps and capacity
  sweeps; constant-folding-pass flag A/B; direct deflections on the actual
  276,507-point cluster input (the old lower bound was extrapolated).
  Recompile independent function objects / fresh processes for monkeypatch A/B;
  JAX function-identity caching must not reuse the old executable.
• Warm starts: reject cross-call mutable solver state; maintain a pure
  likelihood for arbitrary sampler order and transformations.

Completion evidence

Commit a CPU campaign note in results/notes/ with baseline/final comparisons,
all candidate dispositions, reproducible commands, linked artifacts and shipped
phase PRs. Include a GPU regression check for shared library changes and state
any unmeasured hardware limitation. A ranked list alone is not completion: run
and decide the warranted bounded iterations, including justified no-go results.

Preserved September 17 assessment and provenance

The following is historical context; the campaign contract above governs new work.

PointSolver image-plane chi-squared CPU speed-up: is the JAX CPU path sub-optimal enough for a sparse/numba lever?

User request (verbatim):

In autolens_profiling, we have done lots of work speeding up imaging and interferometer on CCD. Now, I want us to speed up the point_source image plane chi-squared, with this issue focusing on CPU. This will likely use the JAX implementation, for other LH functions we had existing numba code to build on and there was lots of sparsity to exploit. Im not sure doing a whole numba CPU implementation is worth it. However, its worth some research, and asking if the JAX CPU implementation is sufficiently sub optimal that there are obvious low hangign fruit improvements with a clever CPU sparse approach, noting that cluster modeling will rely heavily on the PointSolver.

We may not have a likelihood_breakdown of this likelihood function yet, in which case we should make one before trying any kind of CPU or JAX code. I have another claude chat looking into that so bare that in mind and we wont do any real dev work until that is made

Gate

Gated on the point_source likelihood_breakdown cell (scripts/point_source/likelihood_breakdown/) being built by a concurrent session; no CPU or JAX implementation work starts until that breakdown exists. This session's deliverable is the research note + ranked candidate levers. Research phase COMPLETE 2026-09-17 (see Findings); the implementation phase (lever 1 first, library-first in PyAutoArray/PyAutoLens) still waits on the breakdown cell.

Baseline

Committed CPU runtime for image_plane_solved (v2026.7.23.1): eager 274 ms/call, single-JIT full pipeline 38.6 ms/call, vmap(3) 28 ms/call — results/runtime/point_source/image_plane_solved/.

Where the code lives

• autolens_profiling — profiling cells and committed runtime results.
• PyAutoLens — solver source, autolens/point/solver/.
• PyAutoArray — triangle machinery, autoarray/structures/triangles/.

Related

Sibling (not parent): draft/research/autolens_profiling/point_solver_profiling_cells.md — the PointSolver profiling-cells research prompt in the same cluster-strong-lensing epic.

Findings (2026-09-17 research run, CPU, read-only — no repo edited)

Full note + reproducible scratch scripts/JSONs: session scratchpad pointsolver/pointsolver_cpu_research_note.md (to be landed in autolens_profiling/results/notes/ with the fix PR). Machine: WSL2 laptop, 8 cores, NPROC=8, fp64, jax 0.10.2, autolens 2026.8.17.1; speed-ups are within-run interleaved ratios.

Answer: the JAX CPU path is sub-optimal, but not from sparsity or padding — ~84% of the call is a jnp.unique sort that is a no-op under jit.

• Mechanism. CoordinateArrayTriangles._vertices_and_indices (PyAutoArray autoarray/structures/triangles/coordinate_array.py) calls jnp.unique(flat_vertices, size=3*N, fill_value=nan). Under jit the static size=3N means the "deduplicated" table has exactly as many rows as the input (69 849 for the 100×100 simple grid, 59% NaN fill), so it saves zero deflection evaluations and costs one lexicographic sort of 69 849 fp64 rows. The ArrayTriangles it builds is used only for containing_indices and thrown away every step. JAX-path-only defect: the NumPy sibling's np.unique genuinely shrinks 69 849 → 28 665.
• Proof (monkeypatched in scratch, library untouched). Full AnalysisPoint.log_likelihood_function on the image_plane_solved config: 40.0 → 8.1 ms/call (4.9× median, 5.4× p10), log-likelihood bit-identical (7.743201200876812), HLO sorts 15 → 7, FLOPs 15.58M → 6.50M. Real cluster solver config (200×200 @ 0.7″, 92 169 triangles, 13 mass components, 3 planes, 2 sources): 117 → 48 ms per source (2.4×), solved positions identical; a 2-source cluster likelihood 237 → 98 ms.
• Lower bound. Bare deflections for one solve ≈ 6 ms (69 849 + 7×720 points) inside a 26–42 ms call → 77–86% bookkeeping. Post-fix the simple solve is ~1.4× its deflection bound; the cluster solve is at its (extrapolated) bound, so at cluster scale the remaining lever is the dPIE/NFW deflection code, not the solver.
• Sensitivity. n_steps: all 7 refinement steps cost ~3 ms total (~0.45 ms each) — not a lever, and LL degrades below 4 steps. Grid extent: 30×30 = 6.4 ms vs 100×100 = 39.8 ms vs 140×140 = 72 ms with bit-identical LL — the cell tiles ±9.9″ to find images at ~1.6″. Fit/χ² + β* centre + magnification filter = 0.34% of the call. The solve is invoked twice per likelihood (model_data property) but XLA CSEs it — no lever.
• NumPy path (the existing unpadded sparse implementation). 96 ms/solve doing 2.5× fewer deflection evaluations — 12× slower than the fixed JAX path.
• Numba verdict: not worth it. No existing kernel, no exploitable sparsity beyond what the NumPy path already has, headroom gone after a ~20-line fix, and a callback loop would forfeit the custom_jvp implicit gradient, vmap and the GPU path.

Ranked levers: (1) drop the throwaway jnp.unique on the JAX trace path (PyAutoArray + PyAutoLens _plane_triangles; red-control bit-identical LL + cluster positions) — 4.9× simple / 2.4× cluster; (2) precompute the initial lattice's unique vertices + index map once at solver construction in NumPy (the step-0 tiling is static), cutting step-0 deflection count ~2.4–6× with no runtime sort — matters most at cluster scale where deflections dominate post-fix (not measured); (3) solver grid extent guidance / defaults (6.2×, zero code); (4) coarser initial scale + more steps (~4× on step 0, completeness risk, not measured); (5) MAX_CONTAINING_SIZE / 240-triangle fan-out (≲1.5 ms, correctness knob); (6) mass-profile deflection code at cluster scale. Recommend against warm-starting across sampler calls (breaks the pure-function contract vmap/custom_jvp rely on).

Not measured: vmap batch 1/4/16 post-fix, MAX_CONTAINING_SIZE sweep, post-fix grid/n_steps sweep, --xla_disable_hlo_passes=constant_folding A/B, bare deflections at 276 507 points. Trap: jax caches jaxprs on function identity — a monkeypatch A/B needs a distinct function object and jax.clear_caches() before each compile.

🤖 Generated with Claude Code

[Jammy2211]

Phase 2 branches ready — 2026-09-23

Both branches are pushed; no PR is open yet because Heart is RED and PR-open needs the live human override.

PyAutoArray feature/point-source-cpu-p2 (head 25894d10, 2 commits on main 11b93476): CoordinateArrayTriangles._vertices_and_indices returns the flat 3N vertex table + arange index map, no jnp.unique; API unchanged; NumPy sibling untouched; new test_coordinate_jax.py (10 tests; the HLO no-sort guard fails on main with traced containment carries 39 sort op(s)). pytest test_autoarray: 1626 passed. Downstream red controls (autolens_workspace_test point-source jax_likelihood ×4 and jax_grad/gradient.py) all pass and print likelihoods identical to main; autolens_profiling image_plane_solved pin identical (7.743201200876817 / 7.743201200876812).

autolens_profiling feature/point-source-cpu-p2 (head 4c5141e, 5 commits): interleaved A/B cell, RAL ral submit, laptop witness + RAL job 350582 rows, campaign-note phase-2 section. RAL (AMD EPYC 7763, 8 CPUs, cold cache), median ms/call control → nodedup: simple solved 24.00 → 5.38 (4.47× [4.37, 4.57]), simple plain 23.76 → 5.39 (4.40×), vmap-4 31.31 → 11.91 (2.63×), cluster solved 155.22 → 78.32 (1.98× [1.93, 2.03]), cluster plain 153.85 → 77.71 (1.98×). All 19 gates pass with Δ = 0.0 (log-likelihoods, solved positions, gradients, vmap parity bit-identical). HLO sorts 15 → 7 / 25 → 12. Control reproduces phase 1's 24.69 ms within 3 %. Decision in the ledger: ACCEPT. Owed: the GPU regression check for the shared library change (not run in phase 2).

Heart RED at ship time — exact reason strings from pyauto-heart readiness --json (verdict red, score 45):

• release validation FAILED (stage integrate) (red)
• workspace validation not passing (4 failed, cloud#35579888156: autolens notebooks/cluster/modeling.ipynb, autolens notebooks/weak/a2744.ipynb, autolens scripts/cluster/modeling.py, +1 more) (yellow)
• manifest drift: hub organism blurb (organs present) — 7 mismatch(es) vs PyAutoMind/repos.yaml (yellow)

None is caused by or repaired by these branches. Opening the two pending-release PRs requires the human's contemporaneous authorization under AUTONOMY.md "Human override for Heart RED (development only)"; merge is library-first and a separate /prm decision.

Follow-ups found: jax.grad of an AnalysisPoint likelihood is silently all-zero without autofit.jax.register_model(model) (forward values correct) — library issue candidate; laptop witness JSON provenance strings nodedup_body / control_source.body mis-captured (timings unaffected, RAL JSON authoritative); GPU regression check owed; post-release first phase-3 step = re-run the phase-1 breakdown cells on RAL after HPCPullPyAuto.

🤖 Generated with Claude Code

https://claude.ai/code/session_01MMSQRTc7Zthxj1pxRjDVaF

[Jammy2211]

Library PR Created — 2026-09-24

• PyAutoArray: perf(triangles): drop the throwaway jnp.unique vertex dedup on the JAX PointSolver path (#568) #569 (feature/point-source-cpu-p2 @ 25894d10, label pending-release)
• autolens_profiling (linked, merge after the library PR): profiling(point_source): vertex-dedup A/B, RAL CPU + A100 rows, campaign phase 2 (PyAutoArray#568) autolens_profiling#301 (@ 680c0b7, label pending-release)

API changes: None — internal changes only. No signature changes. The contents of the JAX-only CoordinateArrayTriangles.vertices / .indices properties change: the vertex table is no longer deduplicated, and padding triangles carry NaN vertices instead of -1 indices. Solver results are bit-identical.

Workspace impact: none for autolens_workspace, which is behaviour-preserving. The only downstream PR is the linked profiling-evidence PR.

GPU regression check (RAL A100 80GB, job 350587, fp64, imported PyAutoArray 25894d10):

• simple solved: 1.695 → 0.843 ms (2.01×)
• simple plain: 1.684 → 0.822 ms (2.05×)
• vmap-4: 1.865 → 0.958 ms (1.95×)
• cluster solved: 4.206 → 2.457 ms (1.71×)
• cluster plain: 3.860 → 2.105 ms (1.83×)

Every correctness gate is bit-identical. Compile time falls about 40% for the simple cell. The protocol resolves about 4%. CPU (RAL EPYC, job 350582) speed-ups are 4.47× for simple and 1.98× for cluster; see the PR.

Next: /prm (human). Merge PyAutoArray#569 first, then autolens_profiling#301. No merge or release has been done or armed.

Heart RED development override

Shipped under the AUTONOMY.md "Human override for Heart RED (development only)".

• Authorization (live human, 2026-09-24, this session), faithfully quoted: asked to authorize the development-only override for PyAutoArray#568 after a GPU check, the user replied "GPU check and then I authorize". The GPU check passed (RAL A100 job 350587), so the authorization is in effect.
• Scope: commit, push and open the pending-release PRs (PyAutoArray, then the linked autolens_profiling PR). No merge, no release. Merge needs its own human /prm with every required check green.
• Exact current Heart RED reasons (pyauto-heart readiness, verdict RED, score 45):
  • release validation FAILED (stage integrate)
  • workspace validation not passing (4 failed, cloud#35579888156: autolens notebooks/cluster/modeling.ipynb, autolens notebooks/weak/a2744.ipynb, autolens scripts/cluster/modeling.py, +1 more)
  • manifest drift: hub organism blurb (organs present) — 7 mismatch(es) vs PyAutoMind/repos.yaml
• Passed branch gates: pytest test_autoarray 1626 passed at 25894d10 (re-run at ship); +10 new JAX tests, including an HLO sort-count guard that fails on main (39 sort ops) and passes on the branch. Downstream autolens_workspace_test point-source jax_likelihood ×4 and jax_grad/gradient.py print likelihoods and gradients identical to main. The autolens_profiling image_plane_solved pin is unchanged. A/B correctness gates are bit-identical on CPU (RAL job 350582) and GPU (RAL job 350587). Profiling build_readme.py --check and check_submits.py --check pass.
• This branch does not fix Heart. Heart stays RED for release purposes.

[Jammy2211]

Shipped — phase 2 (vertex-dedup removal)

Phase 2 of the point-source CPU campaign is merged, library first (human /prm, 2026-09-24, all CI legs green):

• PyAutoArray — perf(triangles): drop the throwaway jnp.unique vertex dedup on the JAX PointSolver path (#568) #569 → merge 681938ae616f47e1fc4b3ef7812fb6c33c0d0443 (branch head 25894d10: fix 273e152e + JAX triangle tests). Pending release — merged, not yet on PyPI.
• autolens_profiling — profiling(point_source): vertex-dedup A/B, RAL CPU + A100 rows, campaign phase 2 (PyAutoArray#568) autolens_profiling#301 → merge 93757902b512120c9ee17eaa6334a90eb6bcceb5 (branch head 680c0b77): vertex_dedup_ab.py A/B cell, RAL CPU + A100 submits/results, campaign ledger phase 2.

What changed. CoordinateArrayTriangles._vertices_and_indices on the JAX path no longer calls jnp.unique(..., size=3N) (a no-op dedup under jit that cost one fp64 lexicographic sort per refinement step); it returns the flat vertex table and an identity index map. API unchanged; NumPy sibling untouched.

Evidence (results/notes/point_source_cpu_campaign.md, phase 2):

• RAL CPU job 350582 (ral, 8 CPUs, fp64): simple likelihood 24.00 → 5.38 ms (4.47×, 90 % CI 4.37–4.57); two-source cluster 155 → 78 ms (1.98×); HLO sorts 15 → 7 / 25 → 12.
• RAL A100 job 350587 (fp64): no GPU regression — nodedup 1.71–2.05× faster than control.
• Log-likelihoods, solved positions, jax.grad and vmap-4 bit-identical across routes on both hosts; PyAutoArray suite 1626 passed; new HLO no-sort guard red on main, green on the branch.

Issue stays open. This issue carries the campaign's remaining phases. Phases 3–4 (phase 3 = precompute the static initial lattice, ranked first because deflections are 70 % of post-fix FLOPs; phase 4 = residue iteration) are re-filed as PyAutoMind/draft/research/autolens_profiling/pointsolver_cpu_speed_phases_3_4.md, and the phase-2 record is PyAutoMind/complete/2026/09/point-source-cpu-p2.md. Next, after the PyAutoArray release: HPCPullPyAuto on RAL, re-run the phase-1 breakdown cells under a new label so the dashboard rows move, then phase 3.

[Jammy2211]

Phase 3 plan — precompute the static step-0 lattice of the JAX PointSolver

This issue is reused for phase 3 of the point-source CPU campaign (phase 2 shipped in #569, merged 681938ae, pending release). Prompt: PyAutoMind/active/pointsolver_cpu_speed_phases_3_4.md (phase 4 stays open in that prompt).

Human decisions (2026-09-24): tolerance gate, geometric 11 859-vertex table, on by default for the JAX PointSolver if the A/B accepts.

Suggested branch: feature/point-source-cpu-p3 in PyAutoArray, PyAutoLens, autolens_profiling (library first).

Context

Phase 2 (PyAutoArray#569, merged 681938a, pending release) removed the throwaway
jnp.unique; the simple likelihood is now deflection-dominated (70 % of 7.05M FLOP).
Step 0 deflects all 3N = 69 849 triangle-vertex slots of a static lattice. Verified
in-session on canonical PyAutoArray (±9.9″, 0.2″, 23 283 triangles): 28 665 exact-float
unique, 11 859 geometrically unique (duplicates differ by ~1 ulp). Deflecting only
the 11 859 and gathering back cuts step-0 deflections 5.9× (~2.4M FLOP ≈ 34 % of the call).
Prompt: PyAutoMind/draft/research/autolens_profiling/pointsolver_cpu_speed_phases_3_4.md
(phase 4 stays in the draft). Issue: reuse PyAutoArray#568 (re-title to phase 3, post the
plan) rather than a fresh issue.

Decisions (user, 2026-09-24): tolerance gate, geometric 11 859 table, on-by-default for
the JAX PointSolver if the A/B accepts.

High-level

1. PyAutoArray: cached NumPy builder of the static unique-vertex table + index map; opt-in
   static_vertices path on CoordinateArrayTriangles.for_limits_and_scale.
2. PyAutoLens: JAX _initial_triangles turns it on; shape guard + parity/grad/vmap tests.
3. autolens_profiling: static_lattice_ab.py A/B (control / exact 28 665 / lattice 11 859 /
   library), constant-folding cell, thread + memory provenance, ≥20 calls/round.
4. RAL CPU (host-pinned, in-job control) + A100 regression leg via branch clones.
5. Accept/reject by the stop rule; record in the campaign note; ship library-first
   (PyAutoArray → PyAutoLens → autolens_profiling). If rejected, ship only the profiling
   negative result and drop the library branches.

Detailed

PyAutoArray autoarray/structures/triangles/coordinate_array.py

• static_vertex_table(y_min, y_max, x_min, x_max, scale) — module-level
  functools.lru_cache keyed on floats; integer lattice keys (cx+dx, 2cy+dy) → unique
  canonical-float vertices (V,2) + inverse index map (N,3); read-only NumPy arrays.
• CoordinateArrayTriangles(..., vertex_table=None); vertices/indices return the table
  when set; derived instances (for_indexes, up_sample, neighborhood) do not inherit it.
• for_limits_and_scale(..., static_vertices=False); default path and phase-2 tests untouched.
• Tests (test_autoarray/structures/triangles/test_coordinate_jax.py): V = 11 859 for the
  ±9.9/0.2 lattice; unique[inverse] ≈ triangles within a few ulp; kept triangles match NumPy;
  derived lattices drop the table; cache key changes with scale; read-only; test_no_sort kept.

PyAutoLens autolens/point/solver/shape_solver.py _initial_triangles (107-127): pass
static_vertices=True on the JAX path only. Solver pytree unchanged (aux floats), so geometry
change ⇒ new cache key. Tests under test_autolens/point/triangles/: shape guard (first
traced deflection grid has 11 859 rows, not 69 849); mutating scale changes the table;
JAX vs NumPy positions; grad + vmap parity; near-caustic / edge-tie cases. Docstring note on
the step-0 vertices shape.

autolens_profiling scripts/point_source/likelihood_breakdown/static_lattice_ab.py
(template vertex_dedup_ab.py): routes injected into AbstractSolver._initial_triangles,
fresh closures + jax.clear_caches(); library labels itself by traced row count.
Constant-folding cell in a separate process (PyAutoNerves jax_wrapper.py re-adds
--xla_disable_hlo_passes=constant_folding; use a dummy pass name containing the substring
and verify folding from HLO — no Nerves change). Record XLA thread count, affinity,
memory_analysis() + RSS, table build ms/MB, cache-hit time, FLOPs, compile time; 20 × ≥20.
Submits: CPU ral pinned (--nodelist/--constraint, in-job control) + A100 leg modelled on
hpc/batch_gpu/submit_breakdown_point_source_vertex_dedup_ab_a100_fp64. Campaign-note
phase-3 section + README line.

Gates (tolerance): identical image counts and exact containing_indices sets; logL
rel diff ≤ 1e-12; positions ≤ 1e-10; grad/vmap parity to tolerance; workspace_test
point-source pins (rtol 1e-4) + jax_grad/gradient.py unchanged.
Stop / reject if: simple gain < max(5 %, 2× MDE) with no cluster gain; any gate fails;
compile time +20 %; memory or GPU regression. Record negative results.

Setup

• Branch feature/point-source-cpu-p3 in PyAutoArray, PyAutoLens, autolens_profiling;
  worktree ~/Code/PyAutoLabs-wt/point-source-cpu-p3/.
• autolens_profiling is claimed by certified-solver-phase-b (disjoint files): record a
  parallel-claim: per the build: raise caps on scipy, astropy, scikit-image, scikit-learn #267 / phase-2 precedent.
• activate.sh currently puts the certified-solver-phase-b PyAutoArray worktree on
  PYTHONPATH — build the task env explicitly and assert autoarray.__file__ / source
  revisions in every JSON.
• Canonical PyAutoArray (3) and autolens_profiling (10) are behind origin — ff-pull first.
• Release not required to start: RAL uses a branch clone on PYTHONPATH (phase 2 recipe).

Verification

PyAutoArray + PyAutoLens full test suites; the new shape guard red on main, green on
branch; workspace_test point-source scripts; laptop A/B witness then RAL CPU + A100 jobs;
build_readme.py --check, check_submits.py --check; Heart consulted at ship.

🤖 Generated with Claude Code

[Jammy2211]

Library PR Created — phase 3 (static step-0 lattice)

• PyAutoArray: perf(triangles): cached static step-0 vertex table for CoordinateArrayTriangles (#568 phase 3) #570 (head ad0bf97b) — merge first
• PyAutoLens: perf(point): JAX PointSolver step 0 deflects only the unique lattice vertices (PyAutoArray#568 phase 3) PyAutoLens#749 (head 972d454e) — merge second
• autolens_profiling: profiling(point_source): static-lattice A/B, RAL CPU + A100 rows, campaign phase 3 (PyAutoArray#568) autolens_profiling#305 (head 629ce0c1) — merge last

All three carry pending-release. Nothing merged, nothing released.

Result: ACCEPT. RAL CPU job 350636 (Xeon 8490H, 8 CPUs, fp64, folding off), control → library: simple solved 3.54 → 1.76 ms (2.01×), vmap-4 1.43×, two-source cluster 47.4 → 9.1 ms (5.21×); folding on 2.38 / 1.39 / 2.10 / 4.54 / 3.91×. Compile +3.3 to +12.2 % (folding off), worst +16.5 % (folding on), A100 −1.6 to +13.2 % — all under +20 %. A100 job 350637: no regression, 1.01–1.07×. XLA temp memory −56 % to −88 %. 31 / 31 gates bit-identical on all three runs.

API changes: PyAutoArray adds public static_vertex_table(y_min, y_max, x_min, x_max, scale), a vertex_table=None constructor argument on CoordinateArrayTriangles, and for_limits_and_scale(..., static_vertices=False) (opt-in, defaults unchanged). PyAutoLens: behavioural change only — JAX PointSolver step-0 vertices is 11 859 rows instead of 69 849 on the simple lattice; results identical within tolerance.

Tie case: PASSED by human decision 2026-09-24, pinned as PyAutoLens test__source_on_a_step_0_vertex_returns_the_two_true_images (source bit-exactly on step-0 vertex (−0.8, −1.99185843): static lattice returns exactly the 2 true images; the flat control returns a duplicate third).

Workspace impact: additive / opt-in API, no workspace script uses the new symbols; no migration needed. autolens_workspace_test point-source scripts identical to main.

Next: human /prm (library-first, CI green) → PyAutoArray/PyAutoLens release → phase 4 (profile the residue; re-rank in the campaign note).

Heart RED development override

Shipped under the AUTONOMY.md "Human override for Heart RED (development only)".

• Authorization (live human, 2026-09-24, this session), faithfully quoted: asked "When ingestion finishes, how should I ship? Heart is RED…", with the options "Ship with RED override (Recommended): Authorize the development-only Heart RED override for PyAutoArray#568 phase 3…" or hold, the human selected "Ship with RED override (Recommended)".
• Scope: commit, push and open the pending-release PRs (PyAutoArray, PyAutoLens, then autolens_profiling). No merge, no release. Merge needs its own human /prm with every required check green.
• Exact current Heart RED reasons (pyauto-heart readiness, re-read immediately before opening these PRs, verdict RED, score 45):
  • release validation FAILED (stage integrate)
  • workspace validation not passing (4 failed, cloud#35579888156: autolens notebooks/cluster/modeling.ipynb, autolens notebooks/weak/a2744.ipynb, autolens scripts/cluster/modeling.py, +1 more)
  • manifest drift: hub organism blurb (organs present) — 7 mismatch(es) vs PyAutoMind/repos.yaml
• Passed branch gates: pytest test_autoarray 1645 passed at ad0bf97b; pytest test_autolens 756 passed, 1 xfailed at 972d454e (both re-run at ship); the step-0 shape guard (11 859 traced rows) is red on main (69 849) and green on the branch; autolens_workspace_test point-source jax_likelihood ×4 and jax_grad/gradient.py output identical to main (no rtol-1e-4 pin moved); A/B 31 / 31 gates bit-identical on RAL CPU job 350636 (folding off and on) and A100 job 350637; tie case pinned as a test; profiling build_readme.py --check and check_submits.py --check pass.
• This branch does not fix Heart. Heart stays RED for release purposes.

🤖 Generated with Claude Code

[Jammy2211]

Shipped — phase 3 (static step-0 lattice)

Merged by the human's /prm on 2026-09-24, library first, every CI leg green:

• PyAutoArray: perf(triangles): cached static step-0 vertex table for CoordinateArrayTriangles (#568 phase 3) #570 — merge 7fa8d2714f7da43fec2b23c31bff580764831670 (head d6c5e524)
• PyAutoLens: perf(point): JAX PointSolver step 0 deflects only the unique lattice vertices (PyAutoArray#568 phase 3) PyAutoLens#749 — merge 86054bbc19baa2f710a12e04a94fa27ab37cc354 (head c9ba49fb2)
• autolens_profiling: profiling(point_source): static-lattice A/B, RAL CPU + A100 rows, campaign phase 3 (PyAutoArray#568) autolens_profiling#305 — merge 9ad12abdaa44659c95b59ec024dfbbad378b415d (head 629ce0c1)

Pending release: PyAutoArray#570 and PyAutoLens#749 are merged, not yet released. Heart stays RED for release; the development-only override covered the ship.

What changed:

• PyAutoArray: static_vertex_table(...) — a cached NumPy builder of the geometric unique-vertex table for the static step-0 triangle lattice (11 859 of 69 849 slots on the simple ±9.9″/0.2″ lattice; 46 516 of 276 507 on the cluster lattice), plus opt-in static_vertices / vertex_table on CoordinateArrayTriangles.
• PyAutoLens: the JAX PointSolver _initial_triangles turns it on, so step 0 deflects only unique vertices. The tie case (source bit-exactly on a step-0 vertex: the old flat path returns a duplicate image) was passed by human decision and is pinned as test__source_on_a_step_0_vertex_returns_the_two_true_images.
• CI fix after the first ship: d6c5e524 (PyAutoArray) / c9ba49fb2 (PyAutoLens) define the parametrize constants outside the jax guard, which fixed test collection on the unittest-nojax leg.
• autolens_profiling: static_lattice_ab.py cell, RAL CPU + A100 submits, results, and the phase-3 section of results/notes/point_source_cpu_campaign.md.

Evidence: RAL CPU job 350636 (Xeon 8490H, pinned): simple 2.1×, two-source cluster 5.2×, 2.1–4.5× with constant folding on. RAL A100 job 350637: 1.01–1.07×, so no GPU regression. Compile time up at most +16.5 % (limit +20 %), XLA memory 3–4× lower. All 31 gates bit-identical on every host.

Closing this issue. Phases 2 and 3 are done. Phase 4 (profile the residue and iterate) is re-filed as the Mind backlog draft draft/research/autolens_profiling/pointsolver_cpu_speed_phase_4.md and will get its own issue at start_dev.

Mind record: complete/2026/09/point-source-cpu-p3.md.