feat: project ell_comps onto its disk with an opt-in joint clipper (#1537)

[Jammy2211]

Summary

Closes the searchability defect in #1537: ell_comps = (e1, e2) is physical only
inside the unit disk, but every profile's priors are two independent boxes — a
square whose corners describe no ellipse at all. 1 - pi/4 = 21.5% of the
declared prior area is non-physical, and 20.1% of recorded MultiStart lane
best points end at |e| >= 1 (1,252 of 6,240 across the autolens_profiling
campaign, issue #182), while 0 of the 246 lanes that reach the target basin
do. That makes it wasted budget rather than a wasted answer. No box clipper can
see it: (0.8, 0.8) is inside both prior boxes and outside the disk.

This is option 1 of the two the prompt weighed — a joint constraint the
clipper honours — chosen over reparameterisation because it keeps every parameter
name, prior config and recorded identifier stable, and is the only option that
helps a search already in flight.

Three pieces:

1. __model_ball_constraints__ — a class-declared ((path, radius), ...)
   naming the TuplePrior confined to a disk. It sits beside the existing
   __model_constraint__, which measures a violation and therefore cannot say
   how to fix one; this states the geometry's structure, which is what a
   search needs to put a lane back inside it. Duck-typed exactly like
   __model_constraint__, so PyAutoGalaxy states its own geometry without
   inheriting from PyAutoFit (companion PR below).
2. AbstractPriorModel.ball_constraint_index_pairs() — resolves declarations
   to (index_0, index_1, radius) triples indexed into the physical parameter
   vector, so a consumer needs nothing but the vector and this list. Static
   geometry, resolved once outside the step loop and cached under an
   underscore-prefixed __dict__ key (the parameterization convention that the
   pytree-flattening and ModelInstance paths both skip). A declared ball is
   skipped, not raised on, when the component pinned its coordinates to an
   instance (spherical profiles) or fixed one of them.
3. af.ClipperPriorBoxJoint — an opt-in ClipperPriorBox subclass that
   clips the box first and then radially shrinks each declared pair onto its
   ball. A projection rather than a penalty deliberately: past the ellipticity
   clamp the conversion to an axis ratio saturates, so the objective is flat and
   a gradient has nothing to say — the projection does not need one, and moves
   the lane in one step to the nearest point where one exists.

Jittability. The shrink survives jit/vmap/grad: the factor is an
xp.where (never a Python branch), the radius is compared squared, the sqrt
argument is substituted before the sqrt (the double-where idiom — grad
evaluates both branches, sqrt(0) has infinite derivative, and 0 * inf is
NaN), the divisor is floored at tiny, and the factors are assembled as one
multiplicative vector applied with a single multiply, so the traced program is a
fixed handful of ops under both numpy and JAX.

Both members of a moved pair are masked. The mask's consumer is
MultiStartGradient's momentum reset: a lane pushed out of the disk carries
outward momentum in both coordinates, and zeroing one leaves the pair
spiralling straight back out.

Unsupported combinations raise, they do not degrade. AbstractBFGS._bounds_from
refuses a joint clipper on a model that declares a ball (a ball is not expressible
as a scipy.optimize.Bounds; keyed on the model so the clipper can still be
configured once for a whole pipeline), and AbstractMultiStartGradient refuses one
alongside a non-default scaler/bijector at construction, not at the first step
— neither change of variables maps a disk to a disk.

API Changes

Additive and opt-in. af.ClipperPriorBoxJoint is a new clipper; nothing selects
it by default and no existing search or model behaviour changes. EllProfile
gains a declaration in the companion PyAutoGalaxy PR. Two previously-silent
misconfigurations now raise (the joint clipper under LBFGS with a ball-declaring
model, and the joint clipper with a scaler/bijector), which is new behaviour only
for code using the new clipper.

The default clipper's identifier is byte-identical.
ClipperPriorBox.__identifier_fields__ is pinned to ("margin", "strict_epsilon")
— exactly what the identifier's __init__-argspec fallback already inferred — so
existing search identifiers and output directories are unchanged, verified below.
The pin is what keeps it that way: a subclass taking a further constructor
argument would otherwise silently re-key every stored ClipperPriorBox result.

See full details below.

Test Plan

• test_autofit/ full suite: 2288 passed, 3 skipped (97s)
• test_autofit/mapper + test_autofit/non_linear: 1597 passed, 2 skipped
• New coverage: 17 tests across test_clipper.py (projection, angle
  preservation, origin-is-not-NaN, batched lanes, box-before-ball ordering,
  mask on both members and only them, scaler/bijector refusal, identifier
  stability, search wiring), test_model_constraint.py (declaration
  normalisation, malformed-declaration failure at composition, index
  resolution through nested Collections, linked-prior dedup, fixed-coordinate
  skips, caching), and test_multi_start_gradient.py (a lane seeded outside
  the disk lands inside after one clipper call; the momentum reset zeroes
  both components; the default clippers demonstrably cannot move it).
• Unit tests stay numpy-only — no import jax anywhere in the new tests.
  The traced jit/grad/vmap behaviour was verified out-of-band (below).
• Downstream: test_autogalaxy/ 1144 passed and test_autolens/ 553 passed
  against this branch.
• black parity: the set of files black --check would reformat is
  identical before and after this branch (pre-existing, untouched).

Identifier stability, measured. The same script run against main
(cbd30065) and this branch produces identical hashes:

	main	this branch
MultiStartAdam + ClipperPriorBox	45e493b60b1b5bf12bda38f2ac31f450	45e493b60b1b5bf12bda38f2ac31f450
MultiStartAdam + ClipperNone	e1e937f9ca1a3e770b88f03ba9304313	e1e937f9ca1a3e770b88f03ba9304313
LBFGS + ClipperPriorBox	1caa475d9e84b8e1f78797b2a5182f75	1caa475d9e84b8e1f78797b2a5182f75
ClipperPriorBox() alone	5f8556851fb663e4a076182214c75e11	5f8556851fb663e4a076182214c75e11

ClipperPriorBox.project and bounds_from_model return bit-identical values on
both. ClipperPriorBoxJoint() hashes differently (af9832ec…), as it must — it
changes where a lane can sit.

Effect on the measured 20.1%, on one cell. 200,000 uniform draws from the
ell_comps prior box of an Isothermal lens model:

• outside the disk before projection: 21.57% (analytic 1 - pi/4 = 21.46%)
• outside the disk after projection: 0.00%

and the canonical failing point ell_comps = (0.9, 0.9), |e| = 1.2728, is
projected to |e| = 0.999000 with the angle preserved
(e0 == e1 == 0.7063996744), the mask set on both members and on nothing else.
The default ClipperPriorBox leaves the same point at |e| = 1.2728 with an
all-false mask.

Traced behaviour (out-of-band, since unit tests are numpy-only). Under
jax.jit the same point projects to |e| = 0.999000; the origin is finite; and
jax.grad is finite at the origin, at the projected point and at an interior
point — the case a naive radial shrink returns a correct value and a NaN
gradient for. jax.vmap over both is finite.

Heart gate. pyauto-heart readiness is RED with the single red reason
release validation FAILED (stage integrate). That is a known, human-authorised
override for this task (see #1537 and the PyAutoMind record) and is unrelated to
this branch. Shipped to PR-open only under that override — not merged.

Full API Changes (for automation & release notes)

Added

• af.ClipperPriorBoxJoint(margin=1.0e-6, strict_epsilon=1.0e-12) — ClipperPriorBox
  subclass that clips the prior box and then radially projects each pair declared
  by __model_ball_constraints__ onto its ball. Opt-in; a strict no-op on a model
  that declares no geometry. Raises ValueError when handed a scale or bijector.
• autofit.mapper.prior_model.constraint.MODEL_BALL_CONSTRAINT — the
  "__model_ball_constraints__" attribute name.
• autofit.mapper.prior_model.constraint.declares_ball_constraints(cls) -> bool
• autofit.mapper.prior_model.constraint.ball_constraints_for(cls) -> tuple —
  normalises the declaration to ((path_tuple, float_radius), ...); raises
  AssertionError on a malformed entry or a bare-string path.
• AbstractPriorModel.ball_constraint_index_pairs() -> tuple — sorted,
  de-duplicated (index_0, index_1, radius) triples into the physical parameter
  vector. Cached in __dict__["_ball_constraint_index_pairs_cache"].
• Model.has_ball_constraints -> bool
• ClipperPriorBox.__identifier_fields__ = ("margin", "strict_epsilon") —
  pinned; reproduces the previous argspec-inferred identifier exactly.

Changed Behaviour

• AbstractBFGS._bounds_from(model) (private) — raises exc.SearchException when
  the configured clipper is a ClipperPriorBoxJoint and model declares a
  ball. Unchanged for every other clipper, and for a joint clipper on a model
  with no declared geometry.
• AbstractMultiStartGradient.__init__ — raises ValueError when a
  ClipperPriorBoxJoint is passed together with a non-default scaler or
  bijector. Raised at construction, so a multi-hour fit does not die a minute in.
• A class may now declare __model_ball_constraints__; a malformed declaration
  raises AssertionError at af.Model(cls) composition time rather than when a
  search first asks for index pairs.

Migration

• None required. To adopt:
  search = af.MultiStartAdam(clipper=af.ClipperPriorBoxJoint()) — with the
  companion PyAutoGalaxy PR merged, every elliptical profile's ell_comps is then
  projected onto |e| <= 0.999 after each step, and its momentum reset in both
  components.
• Do not combine it with a scaler/bijector, or use it with af.LBFGS on a
  ball-declaring model; both now raise with a message naming the alternative.

Companion (depends on this PR): PyAutoLabs/PyAutoGalaxy#589 — declares the
disk on EllProfile.

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

https://claude.ai/code/session_011joies4k5TdRqezPUK8YET