feat: joint ell_comps disk constraint the MultiStart clipper projects onto (ClipperPriorBoxJoint)

[Jammy2211]

Overview

ell_comps priors are an independent box per component, so 21.5 % of the prior area is non-physical; JAX tracer-mode validation is a no-op there, so gradient lanes settle at |e| ≥ 1 (20.1 % of MultiStart lane best points campaign-wide, 0 of 246 hit lanes; positions-on doubles it — autolens_profiling#182). PyAutoFit already has a class-declared, traceable constraint channel (autofit/mapper/prior_model/constraint.py) and PyAutoGalaxy's EllProfile already declares the disk constraint — but MultiStartGradient only counts violations. Decision (Option 1): give the declaration projectable structure and make the clipper project onto the disk. Reparameterisation (Option 2) rejected for migration cost.

Plan

• New __model_ball_constraints__ class hook (sibling of __model_constraint__, duck-typed) → [(("ell_comps",), radius)]; EllProfile declares radius = convert.ELL_COMPS_MAGNITUDE_CLAMP (0.999) — NOT 1−margin, which sits in the dead-gradient annulus [0.999, 1).
• AbstractModel.ball_constraint_index_pairs() → static (i0, i1, radius) via TuplePrior.prior_tuples + priors_ordered_by_id identity map.
• New ClipperPriorBoxJoint(ClipperPriorBox): box clip first, then jittable radial shrink (xp.where, xp.maximum(r, tiny) — no NaN at r=0, no grad poisoning), clipped_mask True on both pair members so momentum resets. Opt-in subclass ⇒ existing ClipperPriorBox identifiers unchanged; pin __identifier_fields__=("margin","strict_epsilon") on the base (no-op today).
• Guard/warn in AbstractBFGS._bounds_from (a ball is not a scipy Bounds).
• PyAutoGalaxy: widen analysis/dataset.py save_results catch to SamplesException/FitException (mirror PyAutoLens#713) — separate commit.
• Quantify on imaging/mge/hst CPU (two sweep.py arms, n_starts 16) — report fraction ≥1.0 and ≥0.999.
• validate_ell_comps untouched.

Detailed implementation plan

Work Classification

Library (PyAutoFit primary, PyAutoGalaxy).

Affected Repositories

• PyAutoFit (primary)
• PyAutoGalaxy

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean
./PyAutoGalaxy	main	clean

Suggested branch: feature/ell-comps-disk-constraint
Worktree root: ~/Code/PyAutoLabs-wt/ell-comps-disk-constraint/

Implementation Steps

1. autofit/mapper/prior_model/constraint.py — MODEL_BALL_CONSTRAINT = "__model_ball_constraints__", declares_ball_constraints(cls), ball_constraints_for(cls).
2. autofit/mapper/prior_model/abstract.py (beside constrained_model_tuples ~:836) — ball_constraint_index_pairs(self), cached.
3. autofit/non_linear/clipper.py — ClipperPriorBoxJoint; docstring "Two consumers" extended; base __identifier_fields__ pinned.
4. autofit/__init__.py (~:80-82) — export.
5. autofit/non_linear/search/mle/bfgs/search.py:117-147 — guard/warn on a joint clipper.
6. .../multi_start_gradient/search.py:1315-1319 — optional joint-clip count in the log line.
7. autogalaxy/profiles/geometry_profiles.py (~:240) — __model_ball_constraints__ = ((("ell_comps",), convert.ELL_COMPS_MAGNITUDE_CLAMP),).
8. autogalaxy/analysis/analysis/dataset.py (~:178) — exception widening + test.
9. Fix stale comment autolens_profiling/scripts/misc/searches/clipper_campaign.py:23-24 (clipper DOES enter the identifier) — follow-up PR, not this one.

Tests

• test_autofit/non_linear/test_clipper.py TestJointBall: (−1,−1)→radius r angle-preserved; interior bit-identical; (0,0) not NaN; mask on both members; batched input; xp=np only (+ optional importorskip jax jit/grad-finite).
• multi_start_gradient: lane seeded at |e|=1.4 is ≤ radius after one step.
• test_autogalaxy/profiles/test_model_constraint.py: every EllProfile subclass declares the ball; validate_ell_comps standalone unchanged.
• test_autogalaxy/analysis/analysis/test_analysis_dataset.py: mirrored save_results test.

Key Files

• PyAutoFit/autofit/non_linear/clipper.py, autofit/mapper/prior_model/{constraint,abstract}.py, autofit/non_linear/search/mle/bfgs/search.py
• PyAutoGalaxy/autogalaxy/profiles/geometry_profiles.py, autogalaxy/convert.py (CLAMP), autogalaxy/profiles/validate.py (untouched), autogalaxy/analysis/analysis/dataset.py

Original Prompt

Click to expand starting prompt

Claude Development Prompt: Joint ell_comps Disk Constraint

Type: feature
Target: PyAutoGalaxy
Difficulty: medium
Autonomy: supervised
Priority: normal
Status: formalised
Issued: 2026-08-27

Goal

Give the elliptical-components parametrisation a constraint the searches can
actually honour, so a gradient MAP optimizer cannot spend its budget — or
finish — outside the physical region.

ell_comps = (e1, e2) is physical only inside the unit disk |e| < 1. Every
profile's default priors are two INDEPENDENT box priors on e1 and e2
(typically [-1, 1] each), which is a SQUARE. The disk covers pi/4 of that
square, so 21.5% of the prior box is non-physical by area, and nothing in
the search stops a lane from sitting there.

What was measured (autolens_profiling, 2026-08-27 harvest, issue #182)

• The existing validation is a no-op where it matters. validate_ell_comps
  fires on standalone profile construction, but not on the tracer-mode path a
  MultiStart lane actually evaluates — so a lane at |e| >= 1 is never told.
• 20.1% of MultiStart lane best points end at |e| >= 1 (1,252 of 6,240
  lane best points across the recorded campaign).
• 0 of 246 lanes that HIT the target basin end there. The two populations
  are cleanly separated: ending outside the disk is a property of failed lanes,
  never of successful ones. That is what makes this a searchability defect
  rather than a cosmetic one — the wasted 20% is wasted budget, not a wasted
  answer.
• Positions-on roughly DOUBLES the out-of-disk fraction (17% -> 29% of best
  points on the Phase-4 diagnostic arms). The positions penalty pushes lanes
  into the corners the box permits and the geometry does not.

Full write-ups: autolens_profiling/results/notes/inference/ phase_03_prodigy_reliability/RESULTS.md and phase_04_positions/RESULTS.md.

Options to weigh (this prompt does not pick one)

1. A joint disk assertion the clipper honours. Express e1^2 + e2^2 < 1 as
   a constraint the search's prior-box clipper can project onto, rather than as
   a post-hoc validation that raises. The clipper already insets a box; a
   ball projection is the same operation in a different norm. Cheapest change,
   keeps the existing parameter names and every recorded identifier stable, and
   is the only option that helps a search already in flight.
2. Reparameterise. Sample in a coordinate whose whole domain is physical —
   e.g. (magnitude, angle) with a bounded magnitude, or a squashing map onto
   the disk. Removes the problem by construction, but changes the model's
   parameter set, and therefore every target_id, prior config and recorded
   posterior that names ell_comps_0/ell_comps_1. It needs a migration story
   before it needs an implementation.

Hard constraint

Never make validate_ell_comps fire on tracers. Raising from the
tracer-mode path turns a 20%-of-lanes condition into a 20%-of-lanes crash in
the middle of multi-hour GPU fits, and would make a MultiStart search die on
exactly the lanes it is supposed to clip and move on from. The fix belongs in
the search's constraint handling (option 1) or in the parametrisation (option
2) — not in a validator that converts a survivable state into an exception.

Requirements

• Decide between the two options above with the evidence in the write-ups
  named, and record the decision before implementing.
• Whatever lands, add a test that a lane placed at |e| >= 1 is handled by the
  chosen mechanism (projected, or unrepresentable) rather than merely detected.
• Quantify the effect on the measured 20.1% figure on at least one cell before
  claiming the defect is closed.
• Do not change validate_ell_comps's standalone-construction behaviour.

Out of scope

• The positions-penalty scaling question (Gate B pt 2, autolens_profiling
  issue feature/pickle #182) — the doubling is evidence FOR this task, not a task itself.
• Re-running any autolens_profiling campaign.

[Jammy2211]

Library PR Created

Option 1 implemented (a joint disk constraint the clipper honours), as
recorded before implementing. Option 2 (reparameterise) was not taken: it changes
the model's parameter set and therefore every target_id, prior config and
recorded posterior naming ell_comps_0/ell_comps_1, and it cannot help a search
already in flight.

Repo	Commit	PR
PyAutoFit	96dad5191	#1538
PyAutoGalaxy	1d8ae08b9	PyAutoLabs/PyAutoGalaxy#589

Merge order: PyAutoFit#1538 first. PyAutoGalaxy#589 declares geometry the
PyAutoFit PR introduces the machinery to read, and its tests reference
ball_constraint_index_pairs() / af.ClipperPriorBoxJoint.

What landed

• __model_ball_constraints__ — a class-declared ((path, radius), ...), beside
  the existing __model_constraint__. The measure says how far outside; the ball
  says what the region is, which is what a search needs to project.
• AbstractPriorModel.ball_constraint_index_pairs() — resolves declarations to
  (index_0, index_1, radius) indices into the physical parameter vector; static
  geometry, cached, resolved outside the step loop.
• af.ClipperPriorBoxJoint — opt-in ClipperPriorBox subclass: box clip, then
  a radial shrink onto each declared ball. Jittable (double-where so grad is
  finite at the origin; xp.maximum(r, tiny) divisor floor). Masks both
  members of a moved pair, so MultiStartGradient's momentum reset zeroes the
  outward velocity in both coordinates rather than letting the pair spiral back out.
• EllProfile.__model_ball_constraints__ with radius = ELL_COMPS_MAGNITUDE_CLAMP
  (0.999), not 1 - margin: 0.999–1.0 is a dead-gradient annulus where the
  axis-ratio conversion saturates, so projecting there would move a lane from a
  region the model rejects into one the optimizer cannot leave.
• Unsupported combinations raise: LBFGS + joint clipper on a ball-declaring
  model (a ball is not a scipy.optimize.Bounds), and joint clipper + a
  scaler/bijector (neither maps a disk to a disk) — the latter at construction, so
  a multi-hour fit does not die a minute in.
• PyAutoGalaxy AnalysisDataset.save_results catch widened to
  (AttributeError, af.exc.SamplesException, af.exc.FitException), mirroring
  PyAutoLens#713 / PyAutoFit#1535.

Requirements from the prompt

• "add a test that a lane placed at |e| >= 1 is handled rather than merely
  detected" — done, at three levels: the clipper (test_clipper.py), the search
  step contract (test_multi_start_gradient.py, including that the default
  ClipperNone/ClipperPriorBox demonstrably cannot move the same lane), and
  end-to-end on a real Isothermal lens model (test_model_constraint.py).
• "quantify the effect on the measured 20.1% figure on at least one cell" —
  200,000 uniform draws from the ell_comps prior box of an Isothermal lens
  model: 21.57% outside the disk before projection (analytic 1 - pi/4 =
  21.46%) → 0.00% after. The canonical point (0.9, 0.9), |e| = 1.2728,
  projects to |e| = 0.999000 with the angle preserved and both members masked.
• "never make validate_ell_comps fire on tracers" — honoured; the validator
  is byte-unchanged and a test class now pins its standalone-construction
  behaviour (rejects |e| > 1 and the corner (0.8, 0.8), accepts the saturating
  annulus, returns early for a non-concrete magnitude).
• "do not change validate_ell_comps's standalone-construction behaviour" —
  honoured, as above.

Tests

Suite	Result	Runtime
test_autofit/ (full)	2288 passed, 3 skipped	97s
test_autogalaxy/ (full)	1144 passed	45s
test_autolens/ (downstream, against this branch)	553 passed	89s

New tests: 17 in PyAutoFit test_clipper.py, 15 in test_model_constraint.py,
3 in test_multi_start_gradient.py; 11 in PyAutoGalaxy. Unit tests are
numpy-only — the traced jit/grad/vmap behaviour was verified
out-of-band (all finite at the origin, at the projected point and interior).
black: the set of files it would reformat is identical before and after the
branch (pre-existing, untouched).

Identifier stability, measured across checkouts. The same script on main
(cbd30065) and on the branch gives identical hashes for
MultiStartAdam+ClipperPriorBox (45e493b6…), MultiStartAdam+ClipperNone
(e1e937f9…), LBFGS+ClipperPriorBox (1caa475d…) and ClipperPriorBox()
alone (5f855685…), and ClipperPriorBox.project/bounds_from_model return
bit-identical values. ClipperPriorBoxJoint() hashes differently, as it must.

Heart

pyauto-heart readiness is RED, single red reason:
release validation FAILED (stage integrate) — the known, human-authorised
override for this task, unrelated to these branches. Shipped to PR-open only
under that override; nothing merged.

Workspace impact — option (iii), no impact

API changes are additive and opt-in. ClipperPriorBox appears in no user-facing
workspace script (autofit_workspace, autogalaxy_workspace, autolens_workspace,
autolens_workspace_test, HowToLens, HowToFit — zero matches); the only
consumer is autolens_profiling/scripts/misc/searches/, whose usages are all still
valid since nothing was removed or renamed. max_log_likelihood_galaxies call
sites in autogalaxy_workspace are unaffected — only save_results' internal
catch widened. No no_run.yaml entries touch the changed symbols.

A /smoke_test confirmation is owed at merge time rather than now: merge is
already blocked by the Heart override, and no workspace script exercises the new
opt-in clipper, so a smoke run today would not touch the changed code path.

Next

Human review of both PRs → merge PyAutoFit#1538 first, then PyAutoGalaxy#589.
A follow-up could add an autolens_workspace example configuring
af.MultiStartAdam(clipper=af.ClipperPriorBoxJoint()) — option (i) — but nothing
requires it.

🤖 Generated with Claude Code

https://claude.ai/code/session_011joies4k5TdRqezPUK8YET

[Jammy2211]

Shipped

• PyAutoFit: feat: project ell_comps onto its disk with an opt-in joint clipper (#1537) #1538 (merged)
• PyAutoGalaxy: feat: declare the ell_comps disk on EllProfile (PyAutoFit#1537) PyAutoGalaxy#589 (merged)

Option 1 of the prompt — a joint disk constraint the clipper honours — was the
one implemented; the reparameterisation (option 2) was not taken, since it re-keys
every ell_comps_0/ell_comps_1 identifier and prior config and needs a migration
story first.

What changed

PyAutoFit (#1538) — the structural half of the constraint machinery, beside the
existing __model_constraint__ violation measure:

• __model_ball_constraints__: a class-declared ((path, radius), ...) naming a
  tuple prior confined to a disk. Duck-typed like __model_constraint__, so a
  profile library declares its own geometry without inheriting from PyAutoFit.
• AbstractPriorModel.ball_constraint_index_pairs() resolves declarations to
  (index_0, index_1, radius) triples into the physical parameter vector, cached
  under an underscore-prefixed __dict__ key (the parameterization convention).
• af.ClipperPriorBoxJoint: an opt-in ClipperPriorBox subclass that clips the
  box, then radially shrinks each declared pair onto its ball. Jittable — the factor
  is a where, the radius is compared squared, and the double-where idiom keeps
  grad finite at the origin. Both members of a moved pair are masked, so
  MultiStartGradient zeroes the outward momentum in both coordinates.
• Refusals rather than silent degradation: AbstractBFGS refuses a joint clipper on
  a model declaring a ball (a ball is not a scipy.optimize.Bounds), and
  AbstractMultiStartGradient refuses one alongside a non-default scaler or
  bijector at construction (neither change of variables maps a disk to a disk).
• The default clipper is untouched: ClipperPriorBox.__identifier_fields__ is pinned
  to ("margin", "strict_epsilon"), so existing search identifiers and output
  directories stay byte-identical.

PyAutoGalaxy (#589):

• EllProfile.__model_ball_constraints__ = ((("ell_comps",), convert.ELL_COMPS_MAGNITUDE_CLAMP),)
  — declared once at EllProfile, so it reaches every elliptical light and mass
  profile. The radius is the clamp (0.999), not 1 - margin: between 0.999 and
  1.0 the axis-ratio conversion saturates and the likelihood is flat radially, so a
  lane projected into that annulus would have nothing to climb back out on.
  Spherical subclasses inherit the declaration but pin ell_comps to an instance,
  so no pair resolves and nothing is projected.
• validate_ell_comps is unchanged, per the prompt's hard constraint — the fix
  lives in the search's constraint handling, never in a validator that converts a
  survivable state into an exception.
• Also widens AnalysisDataset.save_results' catch to
  (AttributeError, af.exc.SamplesException, af.exc.FitException), mirroring
  PyAutoLens#713: writing an optional output file must never kill a completed fit
  before paths.completed() (PyAutoFit#1535).

Workspace

No workspace smoke run was needed: the clipper is opt-in and no workspace script
selects it, so every existing script keeps the default ClipperPriorBox and its
byte-identical identifiers.