fix: JAX-trace Sérsic stellar-mass CSE deflection path (#500)

The cored-steep-ellipsoid (CSE) decomposition used by Sérsic stellar-mass
deflections was only partially ported to JAX, so a jitted likelihood over a
Sérsic mass profile crashed with TracerArrayConversionError (surfaced by
group/features/advanced/mass_stellar_dark/chaining.py on the release stack).
This is an incomplete-JAX-port bug, not a jax-0.10.2 regression — isothermal/
NFW deflect analytically and never exercised the CSE path under jit.

P1 (unblock): thread `xp` through the full CSE deflection path — radial_grid_from,
decompose_convergence_via_cse, axis_ratio, and _decompose_convergence_via_cse_from
(vectorised coefficient matrix + jnp.linalg.lstsq on the JAX path; scipy retained
for NumPy). cse_settings_from is rewritten branch-free (xp.where) so a tracer
sersic_index traces, with total_cses/sample_points frozen to a static 50/80
(conservative max) for static jit shapes and the per-index dex ranges preserved
exactly. NFW/SersicGradient decompose signatures accept `xp` for compatibility.

Validated: mass_stellar_dark/chaining.py completes (0 tracebacks); 432
test_autogalaxy mass-profile tests pass; jit round-trip nojit==jit to ~1e-11;
numpy-vs-jax parity <=2e-5 for n in [1,4]. NumPy Sérsic-mass deflections shift
~1e-3 (finer CSE count). See PyAutoGalaxy#499.

Co-authored-by: Jammy2211 <JNightingale2211@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: 3 people

autogalaxy/profiles/mass/abstract/cse.py

@@ -79,7 +79,7 @@ def deflections_via_cse_from(
         return xp.vstack((defl_y, defl_x))
 
     @abstractmethod
-    def decompose_convergence_via_cse(self, grid_radii: np.ndarray):
+    def decompose_convergence_via_cse(self, grid_radii: np.ndarray, xp=np):
         pass
 
     def _decompose_convergence_via_cse_from(
@@ -89,6 +89,7 @@ def _decompose_convergence_via_cse_from(
         radii_max: float,
         total_cses: int = 25,
         sample_points: int = 100,
+        xp=np,
     ) -> Tuple[List, List]:
         """
         Decompose the convergence of a mass profile into cored steep elliptical (cse) profiles.
@@ -115,34 +116,40 @@ def _decompose_convergence_via_cse_from(
             A list of amplitudes and core radii of every cored steep elliptical (cse) the mass profile is decomposed
             into.
         """
-        from scipy.linalg import lstsq
-
         error_sigma = 0.1  # error spread. Could be any value.
 
-        r_samples = np.logspace(np.log10(radii_min), np.log10(radii_max), sample_points)
-        y_samples = np.ones_like(r_samples) / error_sigma
+        r_samples = xp.logspace(
+            xp.log10(radii_min), xp.log10(radii_max), sample_points
+        )
+        y_samples = xp.ones_like(r_samples) / error_sigma
         y_samples_func = func(r_samples)
 
-        core_radius_list = np.logspace(
-            np.log10(radii_min), np.log10(radii_max), total_cses
+        core_radius_list = xp.logspace(
+            xp.log10(radii_min), xp.log10(radii_max), total_cses
         )
 
         # Different from Masamune's (2106.11464) method, I set S to a series fixed values. So that
         # the decomposition can be solved linearly.
-
-        coefficient_matrix = np.zeros((sample_points, total_cses))
-
-        for j in range(total_cses):
-            coefficient_matrix[:, j] = self.convergence_cse_1d_from(
-                r_samples, core_radius_list[j]
-            )
-
-        for k in range(sample_points):
-            coefficient_matrix[k] /= y_samples_func[k] * error_sigma
-
-        results = lstsq(coefficient_matrix, y_samples.T)
-
-        amplitude_list = results[0]
+        #
+        # The coefficient matrix is built vectorised (rather than the original
+        # per-column / per-row Python loops that assigned into an ``np.zeros``
+        # buffer) so the decomposition traces cleanly under ``jax.jit`` when
+        # ``xp`` is ``jax.numpy`` — a Python loop writing into a NumPy array
+        # cannot hold tracers. Element ``[k, j]`` matches the original:
+        # ``convergence_cse_1d_from(r_samples[k], core_radius_list[j]) /
+        # (y_samples_func[k] * error_sigma)``. Shape ``(sample_points, total_cses)``.
+        coefficient_matrix = self.convergence_cse_1d_from(
+            r_samples[:, None], core_radius_list[None, :], xp=xp
+        ) / (y_samples_func[:, None] * error_sigma)
+
+        if xp is np:
+            from scipy.linalg import lstsq
+
+            amplitude_list = lstsq(coefficient_matrix, y_samples.T)[0]
+        else:
+            amplitude_list = xp.linalg.lstsq(
+                coefficient_matrix, y_samples.T, rcond=None
+            )[0]
 
         return amplitude_list, core_radius_list
 
@@ -168,7 +175,7 @@ def _convergence_2d_via_cse_from(
         """
 
         amplitude_list, core_radius_list = self.decompose_convergence_via_cse(
-            grid_radii=grid_radii
+            grid_radii=grid_radii, xp=xp
         )
 
         return sum(
@@ -194,9 +201,9 @@ def _deflections_2d_via_cse_from(self, grid: np.ndarray, xp=np, **kwargs) -> np.
         """
 
         amplitude_list, core_radius_list = self.decompose_convergence_via_cse(
-            grid_radii=self.radial_grid_from(grid=grid, **kwargs)
+            grid_radii=self.radial_grid_from(grid=grid, xp=xp, **kwargs), xp=xp
         )
-        q = self.axis_ratio()
+        q = self.axis_ratio(xp)
         q2 = q**2.0
         grid_y = grid.array[:, 0]
         grid_x = grid.array[:, 1]

autogalaxy/profiles/mass/dark/nfw.py

@@ -175,7 +175,7 @@ def convergence_func(self, grid_radius: float, xp=np) -> float:
         )
 
     def decompose_convergence_via_cse(
-        self, grid_radii: np.ndarray, total_cses=30, sample_points=60
+        self, grid_radii: np.ndarray, total_cses=30, sample_points=60, xp=np
     ):
         """
         Decompose the convergence of the elliptical NFW mass profile into cored steep elliptical (cse) profiles.
@@ -203,12 +203,14 @@ def decompose_convergence_via_cse(
             into.
         """
         radii_min = 0.005
-        radii_max = max(7.5, np.max(grid_radii))
+        radii_max = xp.maximum(7.5, xp.max(grid_radii))
 
         def nfw_2d(r):
             grid_radius = (1.0 / self.scale_radius) * r + 0j
-            return np.real(
-                2.0 * self.kappa_s * self.coord_func_g(grid_radius=grid_radius)
+            return xp.real(
+                2.0
+                * self.kappa_s
+                * self.coord_func_g(grid_radius=grid_radius, xp=xp)
             )
 
         return self._decompose_convergence_via_cse_from(
@@ -217,6 +219,7 @@ def nfw_2d(r):
             radii_max=radii_max,
             total_cses=total_cses,
             sample_points=sample_points,
+            xp=xp,
         )
 
     @aa.decorators.to_vector_yx

autogalaxy/profiles/mass/stellar/sersic.py

@@ -15,8 +15,26 @@
 
 
 def cse_settings_from(
-    effective_radius, sersic_index, sersic_constant, mass_to_light_gradient
+    effective_radius, sersic_index, sersic_constant, mass_to_light_gradient, xp=np
 ):
+    """
+    Return the radial fitting range (``upper_dex`` / ``lower_dex``) and the CSE
+    decomposition resolution (``total_cses`` / ``sample_points``) used to decompose
+    a Sersic convergence profile into cored steep ellipsoids.
+
+    The standard (``mass_to_light_gradient <= 0.5``) path is written branch-free
+    using ``xp`` (``jnp.where`` / ``jnp.log10`` under JAX) so it traces cleanly
+    inside ``jax.jit`` when ``sersic_index`` / ``effective_radius`` are tracers:
+    Python ``if`` statements cannot branch on a tracer value, and ``jax.jit``
+    requires static array shapes, so ``total_cses`` / ``sample_points`` are frozen
+    to the conservative maximum previously used (50 / 80) rather than tuned per
+    Sersic index. The per-index ``upper_dex`` / ``lower_dex`` *ranges* are preserved
+    exactly via ``xp.where`` so accuracy is unchanged (only the CSE count can rise).
+
+    The ``mass_to_light_gradient > 0.5`` path (radial-gradient Sersic profiles) is
+    left as the original NumPy branching — it is only reached with a concrete
+    gradient value and is not part of the standard JAX-traced Sersic path.
+    """
     if mass_to_light_gradient > 0.5:
         if effective_radius > 0.2:
             lower_dex = 6.0
@@ -53,31 +71,40 @@ def cse_settings_from(
                 total_cses = 30
                 sample_points = 50
     else:
-        upper_dex = np.min(
-            [
-                np.log10((23.0 / sersic_constant) ** sersic_index),
-                0.85 - np.log10(effective_radius),
-            ]
+        # Static shapes for jax.jit: frozen to the former per-index maximum.
+        total_cses = 50
+        sample_points = 80
+
+        log_effective_radius = xp.log10(effective_radius)
+        base_upper_dex = xp.minimum(
+            xp.log10((23.0 / sersic_constant) ** sersic_index),
+            0.85 - log_effective_radius,
         )
 
-        if (sersic_index <= 0.9) and (sersic_index > 0.8):
-            total_cses = 50
-            sample_points = 80
-            upper_dex = np.log10((18.0 / sersic_constant) ** sersic_index)
-            lower_dex = 4.3 + np.log10(effective_radius)
-        elif sersic_index <= 0.8:
-            total_cses = 50
-            sample_points = 80
-            upper_dex = np.log10((16.0 / sersic_constant) ** sersic_index)
-            lower_dex = 4.0 + np.log10(effective_radius)
-        elif sersic_index > 3.8:
-            total_cses = 40
-            sample_points = 50
-            lower_dex = 4.5 + np.log10(effective_radius)
-        else:
-            lower_dex = 3.5 + np.log10(effective_radius)
-            total_cses = 30
-            sample_points = 50
+        # upper_dex / lower_dex select the same per-index ranges as the original
+        # if/elif ladder, expressed branch-free so a tracer sersic_index traces.
+        upper_dex = xp.where(
+            sersic_index <= 0.8,
+            xp.log10((16.0 / sersic_constant) ** sersic_index),
+            xp.where(
+                sersic_index <= 0.9,
+                xp.log10((18.0 / sersic_constant) ** sersic_index),
+                base_upper_dex,
+            ),
+        )
+        lower_dex = xp.where(
+            sersic_index <= 0.8,
+            4.0 + log_effective_radius,
+            xp.where(
+                sersic_index <= 0.9,
+                4.3 + log_effective_radius,
+                xp.where(
+                    sersic_index > 3.8,
+                    4.5 + log_effective_radius,
+                    3.5 + log_effective_radius,
+                ),
+            ),
+        )
 
     return upper_dex, lower_dex, total_cses, sample_points
 
@@ -243,7 +270,7 @@ def image_2d_via_radii_from(self, radius: np.ndarray, xp=np):
         )
 
     def decompose_convergence_via_cse(
-        self, grid_radii: np.ndarray
+        self, grid_radii: np.ndarray, xp=np
     ) -> Tuple[List, List]:
         """
         Decompose the convergence of the Sersic profile into cored steep elliptical (cse) profiles.
@@ -275,17 +302,20 @@ def decompose_convergence_via_cse(
             sersic_index=self.sersic_index,
             sersic_constant=self.sersic_constant,
             mass_to_light_gradient=0.0,
+            xp=xp,
         )
 
-        scaled_effective_radius = self.effective_radius / np.sqrt(self.axis_ratio())
+        scaled_effective_radius = self.effective_radius / xp.sqrt(
+            self.axis_ratio(xp)
+        )
         radii_min = scaled_effective_radius / 10.0**lower_dex
         radii_max = scaled_effective_radius * 10.0**upper_dex
 
         def sersic_2d(r):
             return (
                 self.mass_to_light_ratio
                 * self.intensity
-                * np.exp(
+                * xp.exp(
                     -self.sersic_constant
                     * (
                         ((r / scaled_effective_radius) ** (1.0 / self.sersic_index))
@@ -300,6 +330,7 @@ def sersic_2d(r):
             radii_max=radii_max,
             total_cses=total_cses,
             sample_points=sample_points,
+            xp=xp,
         )
 
     @property

autogalaxy/profiles/mass/stellar/sersic_gradient.py

@@ -98,7 +98,7 @@ def sersic_gradient_2D(r):
         )
 
     def decompose_convergence_via_cse(
-        self, grid_radii: np.ndarray
+        self, grid_radii: np.ndarray, xp=np
     ) -> Tuple[List, List]:
         """
         Decompose the convergence of the Sersic profile into singular isothermal elliptical (sie) profiles.
@@ -130,9 +130,12 @@ def decompose_convergence_via_cse(
             sersic_index=self.sersic_index,
             sersic_constant=self.sersic_constant,
             mass_to_light_gradient=self.mass_to_light_gradient,
+            xp=xp,
         )
 
-        scaled_effective_radius = self.effective_radius / np.sqrt(self.axis_ratio())
+        scaled_effective_radius = self.effective_radius / xp.sqrt(
+            self.axis_ratio(xp)
+        )
         radii_min = scaled_effective_radius / 10.0**lower_dex
         radii_max = scaled_effective_radius * 10.0**upper_dex
 
@@ -141,10 +144,10 @@ def sersic_gradient_2D(r):
                 self.mass_to_light_ratio
                 * self.intensity
                 * (
-                    ((self.axis_ratio() * r) / scaled_effective_radius)
+                    ((self.axis_ratio(xp) * r) / scaled_effective_radius)
                     ** -self.mass_to_light_gradient
                 )
-                * np.exp(
+                * xp.exp(
                     -self.sersic_constant
                     * (
                         ((r / scaled_effective_radius) ** (1.0 / self.sersic_index))
@@ -159,6 +162,7 @@ def sersic_gradient_2D(r):
             radii_max=radii_max,
             total_cses=total_cses,
             sample_points=sample_points,
+            xp=xp,
         )