Count multi-start lanes trapped outside a declared model constraint

[Jammy2211]

What this adds

_nan_lane_counts splits failing multi-start lanes into value-NaN (likelihood undefined) and gradient-NaN (defined but not differentiable). A third mode escapes both: a lane whose value and gradient are finite, but which sits on a saturating plateau with no gradient along the saturated direction, so it can never leave.

The reference case is the ell_comps magnitude clamp (jax.lax.min(fac, 0.999) in PyAutoGalaxy's convert.py). Past |ell_comps| >= 0.999 the axis ratio pins and the likelihood is exactly constant radially — measured identical to 10 significant figures at magnitude 1.0, 1.2 and 5.0. Such a lane is finite, differentiable, permanently wrong, and its flat figure of merit is indistinguishable from convergence.

This adds a third disjoint counter for it, plus the mechanism it needs.

Why not detect it from the gradient

Because that does not work, and it is worth recording why. The clamp kills only the radial derivative; the angle still comes from an unclamped arctan2, so in general position both ell_comps components carry a large non-zero gradient while their radial projection is zero only to floating-point residue. Scored against 17 genuinely trapped lanes in a JAX reproduction:

candidate detector	caught
gradient component exactly zero	0/17
all gradient components exactly zero	0/17
radial derivative exactly zero	6/17
per-parameter prior-limit escape	17/17
declared model constraint	17/17 + the corner region

Per-parameter prior limits score well only because the overshoot is gross; they provably cannot see the corner region (both components inside (-1, 1), magnitude above 1 — the whole 4 - pi area between the unit disc and unit square). The declared constraint states the valid region exactly rather than inferring it.

The mechanism

A class declares its own valid region via __model_constraint__, returning a traced non-negative violation measure:

def __model_constraint__(self, xp=np):
    magnitude = xp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2)
    return xp.maximum(magnitude - 0.999, 0.0)

This is the differentiable sibling of add_assertion, with two changes: declared on the class rather than attached per model, and evaluated as a traced value rather than raised. The second is forced — a raise on a traced condition gives TracerBoolConversionError, which is exactly why guards like validate_ell_comps return early for non-concrete scalars rather than raising.

A measure rather than a bool is deliberate: a bool suffices for counting, but a magnitude is what a penalty term would need later and carries a usable gradient back into the valid region.

Cost

Nothing flows up through the likelihood. The constraint is a pure function of the parameter vector, evaluated beside fitness.call, which is untouched. It rides the fused value_and_grad as a fourth output, on the device→host sync that already happens each step — the same pattern and rationale as the existing grad_finite reduction.

Measured: a flat +20 HLO lines at every scale tested (grid 31→256, i.e. ~1.2 ms→93 ms per call; starts 32→128), plus 128 bytes of transfer. The cost is fixed, not proportional to likelihood cost. Wall-clock overhead was inside the noise band on a shared CPU and is not reported as a number.

Models declaring no constraint short-circuit and are bit-identical.

Scope

Measurement only. It never gates a redraw and never changes stepping, matching the existing gradient-NaN counter — the policy question waits on what the counts show. No penalty term. No change to resurrect, apply_if_finite, the convergence check, or the clamp.

Nothing in PyAutoGalaxy declares a constraint yet, so on real lens models this reads zero until EllProfile opts in.

Verification

End to end on a real MultiStartProdigy run over a likelihood carrying the actual clamp:

prodigy step 400/400 | best log_post -431.9519 | alive 30/32 | constrained 15/32

n_value_nan_lane_steps       561
n_grad_nan_lane_steps          0
n_constrained_lane_steps    5689

15 of 32 lanes permanently trapped — which the existing instrumentation reports as alive 30/32 with zero gradient-NaN. The surviving lanes still recovered the truth (|ell_comps| 0.9001 against 0.90), which is the honest severity: wasted budget rather than a corrupted result, until n_starts is small or the good basin is rare.

Suite: 1745 passed, +15 new tests. Five failures are pre-existing — verified identical on a stashed clean tree (missing astropy, and a nautilus pool test).

Review note

This touches Model.__init__ and AbstractPriorModel, which every model composition path runs through, so the blast radius is wider than the feature. The short-circuit for constraint-free models is the thing most worth a careful look.

Formatting was deliberately left alone: black wants to reformat all three modified files at main too, so running it would bury the change in unrelated churn. The two new files are black-clean.

---

Generated by Claude Code