perf(triangles): raise PointSolver MAX_CONTAINING_SIZE headroom (15 → ~20, measured)

[Jammy2211]

Overview

The JAX PointSolver keeps at most MAX_CONTAINING_SIZE (15) containing triangles per refinement step and silently truncates anything beyond. The phase-4a sweep (RAL 356367) found an uncapped step-0 maximum of 17 (prior draw 12), and step 1 reaches exactly 15 under ±2.5/0.4, so the headroom is zero. Phase 4c of epic point-source-cpu-speed measures MCS 18/20/24 against the MCS-15 control and raises the constant to the smallest value that clears the observed maximum with margin (expected ~20), at measured cost.

Plan

1. Measure first, with no library change yet. Add MCS 18, 20 and 24 rows to the phase-4a/4b sweep harness and run an interleaved A/B against the MCS-15 control, on the laptop and on the quotable RAL 8490H. Report cost, completeness and uncapped counts.
2. Checkpoint with you: pick N, the smallest value that clears max 17 with margin (expected 20), from the cost table.
3. PyAutoArray: change the constant 15 → N and add a test that a 16–17-triangle containing set is no longer truncated.
4. PyAutoLens: make the docstrings and the ShapeSolver rejection message read the constant rather than a hard-coded "(15)", and update the test that pins that message.
5. Pins: re-check the fiducial log L and the workspace_test point-source and grad pins. Re-pin deliberately only if the padded shape changes a value, and explain why.
6. Ship library-first: PyAutoArray → PyAutoLens → autolens_profiling data PR, with an A100 no-regression row.

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary)
• PyAutoLens
• autolens_profiling

Branch Survey

Repository	Current Branch	Dirty?
./array/PyAutoArray	main (origin/main e281abf, includes #580 4383ea8)	clean
./lens/PyAutoLens	main	clean
./lens/autolens_profiling	main	untracked dataset/abell_1201/ only

Suggested branch: feature/pointsolver-mcs-headroom

Step 1 — measurement (autolens_profiling worktree)

• scripts/point_source_image/likelihood_breakdown/solver_config_sweep.py: add _cfg("mcs18"|"mcs20"|"mcs24", "max_containing_size", mcs=…) beside the existing mcs8/mcs10 rows (~l.300).
  • The existing max_containing_size() patch (~l.402) already rewrites ArrayTriangles.__init__ / for_limits_and_scale __defaults__ and the module constant, and the max_containing_size_proof gate asserts the traced (mcs, 2) shape. Reuse both; add nothing new.
  • Add a --configs selection, or a --mcs-only preset, so the run covers control plus mcs15/18/20/24 and skips the full 26-config grid.
• Per row: median + bootstrap 90 % CI, the paired-round ratio, compile, FLOPs, XLA temp memory, vmap 1/4/16.
• Completeness: 200 prior + 200 stress draws against the default and the fine reference.
  • Extend uncapped_containing_counts to report, per MCS, draws_exceeding_cap (already present, ~l.1493) for every refinement step, not only step 0.
  • List each draw whose image set differs from MCS 15 and explain it (expected: draw 12 and any other truncating draws).
• Runs: laptop witness (not quotable), then the RAL 8490H euclid-ral-compute-10-4. Use the phase-4b submit pattern: branch-clone provenance asserts, and WALL-BASIS cell: point_source_image/solver_config_sweep/simple with a measured-wall source.
  • Run check_submits.py --check locally first.
  • Prefer a quiet node; record the loadavg.

Step 2 — decision checkpoint (you)

• The rule: choose the smallest N ≥ 18 whose uncapped max over all steps and draws is ≤ N − 3. It must add no new image-set change beyond the explained truncation fixes, keep compile within +20 %, and add no worse than about 5 % to the scalar median.
• I bring you the table. You choose N. No default change without your call (this is an epic rule).

Step 3 — PyAutoArray (array/PyAutoArray)

• autoarray/structures/triangles/array.py:13 MAX_CONTAINING_SIZE = N. The defaults at l.65 and l.110 bind the constant at import, so a source change propagates.
  • Check that CoordinateArrayTriangles.with_vertices → ArrayTriangles(...) inherits it; it uses the default.
  • Grep for any hard-coded 15 in autoarray/structures/triangles/.
• Tests in test_autoarray/structures/triangles/test_coordinate_jax.py:
  • Build a lattice and point where the containing set has 16–17 members; draw 12's geometry, or a vertex-on-fold construction. Assert that all of them are kept under jit, where MCS 15 would drop some.
  • Assert the padded output shape is (N,).
  • test_np.py:25 passes max_containing_size=5 explicitly and is unaffected.
• The 4b _STEP0_CONTAINMENT routes pad to self.max_containing_size, and the 4b fuzz/HLO tests stay green. The HLO guard checks the triangle-array size, not MCS.

Step 4 — PyAutoLens (lens/PyAutoLens)

• autolens/point/solver/shape_solver.py:585,600: turn the "(15)" docstring and message into text built from aa...array.MAX_CONTAINING_SIZE. It stays a module-level plain string, so the "same message whether or not JAX is installed" property holds.
• test_autolens/point/triangles/test_shape_solver.py:578,601,608: update the pinned message to use the constant.
• point_solver.py:15,61,263 and implicit_diff.py:35,159 name the constant without a number, so no change.
• Run the full PyAutoLens suite against the PyAutoArray branch.

Step 5 — pins and downstream

• Fiducial 7.743201200876812: padded rows are inf sentinels. Check whether the log L is bit-identical at N. If not, trace the reason (a log-sum-exp over padded pairs) and re-pin FIDUCIAL_SOLVED_LOG_L_BY_BACKEND deliberately, with an explanation.
• autolens_workspace_test scripts/point_source/jax_likelihood/* (the exact −83.38049778 asserts) and jax_grad/gradient.py: run them against the branch. Any changed pin goes into a workspace_test PR with its reason. Positions shape (MCS, 2): grep consumers.
• The targeted smoke is the same 31-script point-source set as perf(triangles): step-0 Point containment without the (N,3,2) gather (#579) #580.

Step 6 — A100 + ship

• An A100 no-regression row for control against N.
• Ship with /ship_library for PyAutoArray and PyAutoLens (Heart gate), then /ship_workspace for the autolens_profiling data plus the ledger "Phase 4c" section in results/notes/point_source_cpu_campaign.md. Add workspace_test only if a pin changed.

Explicitly not in scope

• A runtime overflow counter in the library. Any on-path count costs something on the JAX default, and the harness's uncapped NumPy count already keeps truncation visible for profiling. I'd file it as a draft if you want it.
• Grid-extent defaults (separate prompts).
• nopad removal (campaign remainder).

Parallel claims (human-approved 2026-09-27)

• PyAutoArray is claimed by interferometer-sparse-cache (PR perf(interferometer): cache curvature_matrix / data_vector on sparse and mapping inversions (#581) #582 awaiting merge). It touches only autoarray/inversion/inversion/interferometer*, so it is disjoint from structures/triangles/.
• PyAutoLens is claimed by workspace-config-cleanup (Lens#751 merged, awaiting release). This task touches only autolens/point/solver/shape_solver.py and its test, so it is disjoint.
• autolens_profiling is claimed by point-source-source-plane-p2c. That task touches scripts/point_source_source/… and the source-plane ledger; this one touches scripts/point_source_image/…, the new hpc submits, results/breakdown/point_source_image/ and the CPU ledger. Disjoint.

Verification

• Laptop and RAL 8490H A/B JSONs with all_gates_pass, the per-draw explanations of every image-set change, and the uncapped max per step.
• The new PyAutoArray truncation test fails at MCS 15 and passes at N.
• Full PyAutoArray and PyAutoLens suites.
• workspace_test point-source jax_likelihood + jax_grad, and the 31-script targeted smoke.
• check_submits.py --check and build_readme.py --check clean.
• A100 row.

Original Prompt

Click to expand starting prompt

Point-source CPU speed-up phase 4c — raise MAX_CONTAINING_SIZE headroom (15 → ~20, measured)

Type: feature
Target: autoarray
Repos:

• PyAutoArray
• PyAutoLens
• autolens_profiling
  Themes:
• point-source
• profiling
  Difficulty: small
  Autonomy: supervised
  Priority: normal
  Status: formalised
  Consequence: judge
  Review-minutes: 15
  Unattended: ready
  Epic: point-source-cpu-speed
  Filed: 2026-09-26
  Parent: complete/2026/09/point-source-cpu-p4.md (issue autolens_profiling#314)

Goal

Raise PyAutoArray's MAX_CONTAINING_SIZE from 15 to about 20. Measure 18, 20 and 24 and pick
the smallest that clears the observed maximum with margin. The cap is the static per-step capacity
of containing triangles in the JAX PointSolver, and it silently truncates when exceeded. Human
decision (2026-09-26): accepted as a phase-4a follow-up.

Evidence (phase 4a sweep, RAL job 356367, Xeon 8490H, fp64)

Source: lens/autolens_profiling/results/breakdown/point_source_image/solver_config_sweep_hpc_ral_cpu_fp64.json,
key uncapped_containing_counts, counted in NumPy with no cap on 200 prior draws. Ledger: campaign
note, "Phase 4a".

• Default geometry (±9.9″ / 0.2 / 1e-3): step 0 has median 9, p99 15 and max 17, on prior draw
  12.
  • 23 draws exceed 12, and 1 exceeds 15.
  • Steps 1–7 peak at 13 / 11 / 9 / 7 / 5 / 5 / 5.
  • Draw 12 still passed completeness, because the truncated entries were the spurious fold-line
    candidates. That is luck, not a guarantee.
• Under ±2.5/0.4, step 1 reaches exactly 15 (p99 15). The headroom is zero there too, and a smaller
  workspace grid (see draft/feature/autolens_workspace/pointsolver_grid_extent_per_package.md)
  pushes counts up.
• The cap is load-bearing. MCS 8 and MCS 10 lose images (prior multiplicity 62 % and 86.5 %), at
  1.11× and 1.06× speed.

Method

• Harness: lens/autolens_profiling/scripts/point_source_image/likelihood_breakdown/solver_config_sweep.py.
  Add MCS 18 / 20 / 24 routes (the max_containing_size block already exists).
• Report the cost of each as its own interleaved row: median, CI, compile, FLOPs and temp memory. The
  cap sets the refinement grid size (720 rows per step at MCS 15: 240 kept triangles × 3) and the f64[60] neighbourhood
  sorts, so expect a modest cost, and measure it.
• Prove no image-set change on the 200 prior + 200 stress draws against both the default and the
  reference, except where the old cap was truncating (draw 12 and any others the uncapped count
  flags). Those draws must be explained one by one.
• Fiducial log L 7.743201200876812 must stay bit-identical, or, if the padded shape changes it,
  explain why and re-pin deliberately. Check autolens_workspace_test point_source/jax_likelihood/*
  and jax_grad/gradient.py pins, and positions shape (MCS, 2).
• Consider a debug-mode overflow counter or warning: a non-JAX or jax.debug-gated count of how
  often the uncapped containing set exceeds the cap, so truncation is never silent again. It must not
  cost anything on the default JAX path.
• A100 no-regression row if the default changes.

Traps

• The __defaults__ binding trap. MAX_CONTAINING_SIZE is a module constant, but
  ArrayTriangles.__init__ and for_limits_and_scale bind it as a default argument at import.
  Patching the module constant alone does nothing. The 4a cell had to patch those functions'
  __defaults__ too, and restore them by value. In the library change, make sure every default
  reads the new value (or switch to a None sentinel resolved at call time), and grep PyAutoLens for
  any copy of the constant.
• PyAutoArray is currently claimed by interferometer-transform-real-scatter. Check at start_dev.
  Coordinate with phase 4b (shipped 2026-09-27, record complete/2026/09/pointsolver-step0-gather.md, PyAutoArray#580), which also touches the
  triangles code, and prefer landing 4b first.

[Jammy2211]

Decision (human, 2026-09-27): MAX_CONTAINING_SIZE 15 → 20.

Evidence: RAL 8490H job 358976 (loaded node, load around 207; ratios between rows are valid), plus a laptop witness run. All caps give identical image sets, positions and log L on every prior, stress and stream draw, and the fiducial 7.743201200876812 is bit-identical at every cap. The uncapped maximum is 17 (prior draw 12, step 0); no draw exceeds the cap at 18 or above. MCS 20 costs +6.2 % on the scalar call (CI +5.0 to +8.5 %), +8.6 % at vmap-16 and +7 % compile. It misses the ~5 % cost budget, and that was accepted for 3 triangles of headroom. Gradients differ by at most 6.1e-14 relative, with log L bit-identical (the padded shape changes the backward sum order). Harness commits (local): a40e140 d5a971e d6757ad 90afbeb 94eac23.

[Jammy2211]

Library PRs Created (2026-09-27): PyAutoArray #584 (MCS 15 → 20, with a truncation test that fails at 15 and passes at 20). PyAutoLens PyAutoLabs/PyAutoLens#753 (the ShapeSolver message reads the constant). autolens_profiling data PyAutoLabs/autolens_profiling#335.

Suites: PyAutoArray 1738 passed; PyAutoLens 759 passed, 1 xfailed. No pin changed, the targeted smoke passes 31/31, and the A100 row is 1.00×. Heart was YELLOW for manifest drift and the incomplete rehearsal, and the human acknowledged it. Next: /prm, library-first.

[Jammy2211]

Shipped

Point-source CPU speed, phase 4c: MAX_CONTAINING_SIZE raised 15 → 20.

PRs (all merged 2026-09-27)

• PyAutoArray#584 — perf(triangles): raise PointSolver MAX_CONTAINING_SIZE 15 → 20 (#583) #584 (library, pending-release)
• PyAutoLens#753 — docs(point): ShapeSolver JAX message reads MAX_CONTAINING_SIZE (PyAutoArray#583) PyAutoLens#753 (library, pending-release)
• autolens_profiling#335 — research: point-source CPU phase 4c — MAX_CONTAINING_SIZE headroom A/B (PyAutoArray#583) autolens_profiling#335

What changed

• MAX_CONTAINING_SIZE 15 → 20. The human chose 20 at +6.2 % scalar cost (a miss of the ~5 % rule, accepted) to get 3 triangles of headroom over the observed maximum of 17.
• PyAutoLens: the ShapeSolver JAX rejection message now quotes MAX_CONTAINING_SIZE from autoarray.
• New truncation test: fails at 15, passes at 20.

Evidence

• RAL 8490H job 358976 (loaded node): mcs20 at 0.941 [0.922, 0.953] × control, compile +7 %. A100 job 359102: 1.00×.
• Image sets, positions, log L and fiducial are bit-identical. Gradients differ by ≤ 6.1e-14 relative.
• Suites: PyAutoArray 1738 passed; PyAutoLens 759 passed + 1 xfail. Pins unchanged; smoke 31/31. Heart YELLOW acknowledged by the human.

Open follow-ups (recorded in Mind draft/research/autolens_profiling/pointsolver_cpu_speed_campaign_remainder.md, "Phase 4c leftovers")

• A100 vmap-4 +8 % single cell (likely noise).
• Draw-12 fold-line attribution, carried from 4a and not re-verified.

The library changes reach users with the next PyAutoArray / PyAutoLens release.