Skip to content

Declare the ell_comps saturation constraint on EllProfile - #572

Merged
Jammy2211 merged 1 commit into
mainfrom
claude/jax-sampling-flat-gradients-ptmqnl
Aug 15, 2026
Merged

Declare the ell_comps saturation constraint on EllProfile#572
Jammy2211 merged 1 commit into
mainfrom
claude/jax-sampling-flat-gradients-ptmqnl

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

What this adds

validate_ell_comps already owns the ell_comps geometry, but it signals by raising, which works only for concrete scalars. Under a JAX trace the condition is a tracer and a raise is impossible, so it returns early (validate.py:153-154) — that escape hatch is load-bearing, or every jitted likelihood would crash instead of sampling.

The consequence is that gradient searches see nothing when a lane walks past the clamp into the region where the axis ratio saturates and the radial gradient dies. Such a lane is finite, differentiable, and can never come back.

EllProfile.__model_constraint__ states the same geometry as a traced, non-negative distance beyond the clamp. PyAutoFit reads it on the traced path to count trapped multi-start lanes (PyAutoLabs/PyAutoFit#1475).

Placement

Declared once on EllProfile — the single base every elliptical light and mass profile inherits, and the one site where ell_comps is assigned. No profile opts in individually.

Spherical profiles subclass their elliptical counterpart (IsothermalSphIsothermal → … → EllProfile), so they inherit it too. With ell_comps pinned at (0, 0) it is always satisfied — correct, if a few wasted ops. Pinned by test rather than left implicit.

The threshold is the clamp's, not the guard's

Deliberately 0.999, not 1.0. They answer different questions — the clamp is where the gradient dies, the guard is where the geometry stops meaning anything — and the annulus between them is reachable:

magnitude radial derivative validate_ell_comps
0.9985 alive valid
0.9995 exactly zero valid
1.2 exactly zero rejected

Keying the constraint to the guard's threshold would miss that band entirely.

One definition for the clamp

Adds ELL_COMPS_MAGNITUDE_CLAMP. The clamp was a bare 0.999 literal at three sites — convert.py's JAX and NumPy branches, and the Sersic Cartesian eccentric-radius path from #571 — with the guard's separate 1.0 in a fourth file and nothing relating them. That is exactly the drift that produced the reachable annulus above; the constant states the relationship in one place.

Value unchanged at every site. This is a literal-to-constant swap, not a numerical change.

No behaviour change

The guard still rejects magnitude >= 1.0 exactly as before. The constraint is read only by callers that ask for it, and nothing in this repo does.

Verification

  • Full suite: 1101 passed, 1 skipped, 0 failures — including the Sersic image path the constant swap touches.
  • 12 new tests covering the declaration's reach across elliptical profiles, the spherical inheritance case, the guard/clamp threshold split, and that the conversion still saturates at the constant.
  • Integration against the PyAutoFit branch, on a real ag.mp.Isothermal (no stand-in class), driven by af.MultiStartProdigy:
prodigy step 100/300 | best log_post -916.8119 | alive 27/32 | constrained  7/32
prodigy step 300/300 | best log_post -431.7808 | alive 22/32 | constrained  6/32

n_value_nan_lane_steps      2344
n_grad_nan_lane_steps          0
n_constrained_lane_steps    1057

best lane einstein_radius = 1.5968   (truth 1.6)

The constraint is discovered through inheritance alone — nothing in that model declares anything.

Merge order

PyAutoFit#1475 defines the protocol and should merge first. This declaration is inert without it — __model_constraint__ is a plain method that nothing calls until PyAutoFit looks for it — so there is no hard import dependency and this cannot break on an older PyAutoFit.


Generated by Claude Code

`validate_ell_comps` already owns this geometry, but it signals by raising,
which works only for concrete scalars — under a JAX trace the condition is a
tracer and a `raise` is impossible, so it returns early (validate.py:153-154).
Gradient searches therefore see nothing when a lane walks past the clamp into
the region where the axis ratio saturates and the radial gradient dies.

`EllProfile.__model_constraint__` states the same geometry as a traced,
non-negative distance beyond the clamp, which PyAutoFit consumes on the traced
path to count trapped multi-start lanes (PyAutoFit #1475). It is declared once
on `EllProfile`, the single base every elliptical light and mass profile
inherits and the one site where `ell_comps` is assigned, so no profile opts in
individually. Spherical profiles subclass their elliptical counterpart and so
inherit it too, with `ell_comps` pinned at (0, 0) it is always satisfied.

The threshold is deliberately the clamp's 0.999, not the guard's 1.0. They
answer different questions and the annulus between them is reachable: at
magnitude 0.9995 the radial derivative is already exactly zero while
validate_ell_comps still calls the point valid.

Also introduces ELL_COMPS_MAGNITUDE_CLAMP so the clamp has one definition. It
was a bare literal at three sites — convert.py's JAX and NumPy branches and the
Sersic Cartesian eccentric-radius path — with the guard's separate 1.0 in a
fourth file and nothing relating them. Value unchanged at every site.

No behaviour change: the guard still rejects magnitude >= 1.0 exactly as before,
and the constraint is read only by callers that ask for it.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_012GAFoogitLceTsgA7bfB4k
@Jammy2211
Jammy2211 merged commit 695b27c into main Aug 15, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants