fix: compose ClipperPriorBoxJoint with a bijector/scaler instead of refusing it

[Jammy2211]

Overview

ClipperPriorBoxJoint is refused outright whenever a scaler or bijector is set — at construction in multi_start_gradient/search.py:365-384 and again at projection time in clipper.py:522-535. The refusal is stronger than the geometry requires: a coordinate whose bijector kind is identity with linear scale s maps a disk of radius R to a disk of radius R / s, which is still a disk and still exactly projectable. Only a genuinely non-linear pair (log, logit, or two identity coordinates with different scales, which gives an ellipse) has no closed form.

The cost is concrete: any gradient arm that wants a log/logit reparameterisation anywhere in the model must drop the joint clipper entirely, so its ell_comps lanes settle at the box corner |e| = 1.414 — outside the disk, where the axis-ratio conversion saturates and the gradient is flat. Phase 8B's log_reg arms do exactly this.

Plan

• Teach the bijector to report, per coordinate, the linear scale it applies when (and only when) that coordinate is identity-kind.
• Replace the clipper's blanket refusal with a per-pair resolution: a ball pair whose two members are both identity-kind under a common linear scale s is projected onto radius R / s; anything else raises, now naming the offending index pair and its kinds.
• Delete the construction-time refusal in MultiStartGradient and move the check to model-resolution time inside _fit, so a genuinely bad combination still dies before any likelihood evaluation but a good one is simply allowed.
• Do not round-trip through the bijector inside the clipper — the logit epsilon clamps would break bit-identity on unrelated coordinates, saturate gradients, and add traced ops.
• No new constructor arguments, so __identifier_fields__ and every stored result key are untouched.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main (f466dce)	clean

Suggested branch: feature/joint-clipper-compose-with-bijector

Implementation Steps

1. autofit/non_linear/bijector.py — add AbstractBijector.identity_scales beside kinds (~:236-240): a list whose entry i is float(self._scale[i]) when self._kind_code[i] == _IDENTITY and None otherwise, guarded by _check_resolved(). ~12 lines.

2. autofit/non_linear/clipper.py — replace the blanket raise at ~:522-535 with a _pairs_in_stepped_coordinates(pairs, scale, bijector) helper returning (i, j, radius / s) when both members of a pair share a common linear scale s (read from bijector.identity_scales, or from the raw scale array on the scaler path), and raising the existing ValueError — now naming the index pair and its kinds — otherwise. Add a "Composition with a scaler/bijector" docstring section carrying the R -> R/s argument.

3. autofit/non_linear/search/mle/multi_start_gradient/search.py — delete the construction-time refusal (:365-384); call the clipper helper on model.ball_constraint_index_pairs() inside _fit, immediately after self.bijector.from_model(model=model) (:884-885) and before _vmapped is built/compiled. Update the clipper / bijector docstring paragraphs (~:148-201). bfgs/search.py is untouched — LBFGS still refuses the joint clipper wholesale, since scipy has no ball.

Tests (NumPy only — JAX stays out of the unit suite)

• test_autofit/non_linear/test_clipper.py — rewrite test__a_scaler_or_bijector_is_refused_rather_than_applied (:700) into the composition cases: BijectorNone bit-identical to no bijector; a per-path log on a non-ell_comps path still projects the corner onto 0.999 with the angle preserved; BijectorLogit on the ball model still raises and names the pair; scale=[2,2,1] projects onto 0.999/2; scale=[2,3,1] raises; BijectorDiagonal(ScalerPriorWidth()) with equal widths is accepted. Restate the search-wiring tests (:774-784) from "refused at construction" to "refused at model resolution, before the first step", keeping the identifier assertions verbatim and adding one that a bijector leaves Identifier(...) unchanged.
• test_autofit/non_linear/test_bijector.py — F5 pin test__round_tripping_a_per_path_map_is_bit_exact_where_it_is_identity, right after the objective-composition test (~:388-409).
• test_autofit/non_linear/search/mle/test_multi_start_gradient.py — the joint clipper plus a per-path bijector constructs and is accepted at model resolution; a logit-on-ell_comps combination raises before the first step.

Key Files

• autofit/non_linear/bijector.py — the per-coordinate kind/scale cache and its accessors.
• autofit/non_linear/clipper.py — ClipperPriorBoxJoint.project, the ball projection.
• autofit/non_linear/search/mle/multi_start_gradient/search.py — where the two strategies meet.

Campaign implication (not done here)

To actually use the disk, the Phase 8B logit arm must be restated as a BijectorPerPath with logit on every path except ell_comps.*. That is a config change in autolens_profiling, not a library change, and is out of scope for this issue.

Original Prompt

Click to expand starting prompt

ClipperPriorBoxJoint is refused whenever a bijector/scaler is set (multi_start_gradient/search.py:365-384, clipper.py:522-535); Phase 8B log_reg arms therefore settle at the ell_comps box corner |e|=1.414. Compose them: resolve each ball pair against the map and refuse only genuinely non-linear pairs. Also add the in-process F5 pin test.

[Jammy2211]

Library PR created

PR: #1540 (pending-release)
Commit: 11fdd4671 on feature/joint-clipper-compose-with-bijector

Option B implemented as planned, with one naming deviation: the clipper helper is pairs_in_stepped_coordinates (public), not _pairs_in_stepped_coordinates. It is called from multi_start_gradient/search.py, so it belongs beside project / bounds_from_model in the clipper's cross-module consumer API rather than behind an underscore another module has to reach through.

Files

File	Change
autofit/non_linear/bijector.py	+30 / -1 — new identity_scales property
autofit/non_linear/clipper.py	+139 / -13 — pairs_in_stepped_coordinates, project composes, "Composition with a scaler/bijector" docstring
autofit/non_linear/search/mle/multi_start_gradient/search.py	+49 / -20 — construction refusal deleted, check moved into _fit, docstrings updated
test_autofit/non_linear/test_clipper.py	+206 / -20
test_autofit/non_linear/test_bijector.py	+44 / -0 (F5 pin)
test_autofit/non_linear/search/mle/test_multi_start_gradient.py	+75 / -0

Validation

• test_autofit/ full: 2301 passed, 3 skipped (base f466dce1a: 2288/3 → +13 tests), 80s.
• Downstream on the worktree autofit: test_autogalaxy/profiles 716 passed; test_autolens/analysis 62 passed.
• Traced CPU integration (JAX, outside the unit suite): 8/8 lanes end at |e| = 0.999 under jit/vmap, 874 clipped lane-steps (so the assertion is not vacuous), 0 NaN value/grad lane-steps. The logit-on-ell_comps counterfactual raises before the first step and names the pair.

Note on formatting

The base files were not black-25.1.0-clean, so a plain black run produced unrelated re-wrap churn in all five touched files. That churn was reverse-patched out — the diff is change-only.

Follow-up (not here)

The Phase 8B logit arm must be restated as a BijectorPerPath with logit on every path except ell_comps.* to actually use the disk — an autolens_profiling config change.

Shipped under the human-authorised Heart RED override; the sole red reason is "release validation FAILED (stage integrate)", unrelated to this change. Not merged — awaiting CI and review.

[Jammy2211]

Shipped

Merged in PR #1540 (merge commit 54aa087 → main).

ClipperPriorBoxJoint was refused outright whenever a scaler or bijector was set — at construction in MultiStartGradient.__init__ and again at projection time in ClipperPriorBoxJoint.project. That refusal was stronger than the geometry requires: a ball survives a common linear rescale of both its members, since theta_i**2 + theta_j**2 <= R**2 is exactly phi_i**2 + phi_j**2 <= (R/s)**2. Only a genuinely non-linear pair has no closed form — a log/logit kind, or two identity coordinates with different scales (an ellipse, whose nearest-point projection is the root of a quartic).

What changed

• autofit/non_linear/bijector.py — new AbstractBijector.identity_scales: the per-coordinate linear scale where the kind is identity, None where it is not. None rather than 1.0 for log/logit deliberately, so a caller cannot compose with it silently and wrongly.
• autofit/non_linear/clipper.py — new ClipperPriorBoxJoint.pairs_in_stepped_coordinates(pairs, scale, bijector) resolves each declared ball pair against the map, returning (i, j, radius/s) for a pair sharing one common positive linear scale and raising the existing ValueError — now naming the offending index pair and its kinds — otherwise. project calls it in place of the blanket raise.
• .../multi_start_gradient/search.py — the construction-time refusal is deleted; the check runs once in _fit, right after bijector.from_model(model=model) and before _vmapped is built. Whether the two compose is a question about the model, and no model exists at construction — but this is still before any likelihood evaluation, so a bad combination does not die a minute into a multi-hour fit. bfgs/search.py is untouched (scipy has no ball).

The projection is deliberately not round-tripped through the bijector: that would drag the logit epsilon clamps across coordinates the ball has nothing to do with, breaking the interior-point path's bit-identity, saturating gradients, and adding traced ops per step. Dividing the radius costs nothing and is exact.

API changes — additive only. New identity_scales and pairs_in_stepped_coordinates; the MultiStartAdam(clipper=ClipperPriorBoxJoint(), bijector=...) combination no longer raises at construction — it either works or raises the same ValueError class at model resolution. __identifier_fields__ is unchanged, so no stored result is re-keyed (pinned by a test). Nothing in PyAutoGalaxy or PyAutoLens calls the changed methods.

Validation: test_autofit/ 2301 passed, 3 skipped (base f466dce1a: 2288 — +13 new tests) · test_autogalaxy/profiles 716 passed · test_autolens/analysis 62 passed. New NumPy-only tests cover composition, refusal-with-named-pair, common vs mixed scales, BijectorDiagonal(ScalerPriorWidth()), the search wiring restated to "refused at model resolution", and the F5 pin (test_bijector.py — identity coordinates round-trip with ==, bit-exact; the log one to rel 1e-12), which is what the radius-division shortcut rests on. This is also the sound, in-process F5 owed by autolens_profiling#185.

A traced JAX/CPU integration run outside the unit suite confirmed the behaviour end-to-end: a BijectorPerPath({"einstein_radius": "log"}) fit whose optimum sits at |e| = 2.0 held every lane at |e| = 0.999 over 874 clipped lane-steps, 0 NaN; the BijectorLogit counterfactual raised before the first step, naming the pair.

Campaign implication (not in this PR): to actually use the disk, the Phase 8B logit arm must be restated as a BijectorPerPath carrying logit on every path except ell_comps.* — a config change in autolens_profiling, not a library change.

Shipped under the human-authorised Heart RED override (sole red reason: release validation FAILED (stage integrate)).

Closed via PyAutoMind close-out: complete/2026/08/joint-clipper-compose-with-bijector.md.