Overview
autolens_profiling#121 isolated a JAX autodiff defect at the exact circular default-prior mean of elliptical light profiles. An off-centre full FitImaging Sersic likelihood is finite at ell_comps=(0, 0), and gradients at the ±1e-8 and ±1e-6 Cartesian axes are finite with norm about 1.27892, but jax.grad returns [NaN, NaN] at the origin.
The former q-angle structural finding was a false positive: Sersic fits the Cartesian ell_comps_0/1 parameters, while axis ratio and angle are derived polar quantities. The actual defect is that the current image path differentiates through sqrt(e_y^2 + e_x^2) and atan2 separately at the polar origin even though their combined eccentric-radius limit is differentiable.
Evidence: PyAutoLabs/autolens_profiling#122, finding likelihood.imaging-sersic.ell-comps-origin-nonfinite-gradient.
Bounded remedy
Implement a Cartesian eccentric-radius path for the Sersic image family that avoids separately differentiating polar radius and angle at ell_comps=(0, 0). For shifted coordinates (y, x), f^2=e_x^2+e_y^2, the existing rotated eccentric radius is algebraically equivalent (before the existing high-ellipticity clamp) to:
r_ecc^2 = ((1+f^2-2e_x)x^2 - 4e_yxy + (1+f^2+2e_x)y^2) / (1-f^2)
Use this equivalence, or another demonstrably equivalent combined Cartesian formulation. Do not assign arbitrary derivatives to axis_ratio_and_angle_from; those polar outputs are individually non-differentiable at the origin and only the combined image geometry has a unique Cartesian derivative.
Acceptance
- NumPy image values for representative Sersic-family profiles agree with the existing q-angle path to floating-point noise across circular and anisotropic cases.
- JAX
jit and grad at ell_comps=(0, 0) are finite on an off-centre grid.
- The origin gradient agrees with a central finite-difference estimate and the bounded ±
1e-8/±1e-6 neighbourhood.
- The downstream full-likelihood reproducer from autolens_profiling#122 returns a finite origin gradient without changing its likelihood value.
- Preserve the existing
fac <= 0.999 behavior for traced out-of-domain ellipticities.
- Add NumPy unit coverage in PyAutoGalaxy; validate JAX through the downstream workspace/profiling reproducer per repository policy.
Boundary
No public API, prior, configuration, or physical circular-orientation convention change. Keep the fix in PyAutoGalaxy and do not import autolens.
Overview
autolens_profiling#121isolated a JAX autodiff defect at the exact circular default-prior mean of elliptical light profiles. An off-centre fullFitImagingSersic likelihood is finite atell_comps=(0, 0), and gradients at the ±1e-8and ±1e-6Cartesian axes are finite with norm about1.27892, butjax.gradreturns[NaN, NaN]at the origin.The former q-angle structural finding was a false positive: Sersic fits the Cartesian
ell_comps_0/1parameters, while axis ratio and angle are derived polar quantities. The actual defect is that the current image path differentiates throughsqrt(e_y^2 + e_x^2)andatan2separately at the polar origin even though their combined eccentric-radius limit is differentiable.Evidence: PyAutoLabs/autolens_profiling#122, finding
likelihood.imaging-sersic.ell-comps-origin-nonfinite-gradient.Bounded remedy
Implement a Cartesian eccentric-radius path for the Sersic image family that avoids separately differentiating polar radius and angle at
ell_comps=(0, 0). For shifted coordinates(y, x),f^2=e_x^2+e_y^2, the existing rotated eccentric radius is algebraically equivalent (before the existing high-ellipticity clamp) to:Use this equivalence, or another demonstrably equivalent combined Cartesian formulation. Do not assign arbitrary derivatives to
axis_ratio_and_angle_from; those polar outputs are individually non-differentiable at the origin and only the combined image geometry has a unique Cartesian derivative.Acceptance
jitandgradatell_comps=(0, 0)are finite on an off-centre grid.1e-8/±1e-6neighbourhood.fac <= 0.999behavior for traced out-of-domain ellipticities.Boundary
No public API, prior, configuration, or physical circular-orientation convention change. Keep the fix in PyAutoGalaxy and do not import
autolens.