From 68a2a00d08b8906d4937fef97aeaae5a0bf4546b Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Tue, 14 Jul 2026 13:34:42 +0100 Subject: [PATCH] =?UTF-8?q?fix:=20JAX-trace=20S=C3=A9rsic=20stellar-mass?= =?UTF-8?q?=20CSE=20deflection=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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: Claude Opus 4.8 --- autogalaxy/profiles/mass/abstract/cse.py | 55 ++++++------ autogalaxy/profiles/mass/dark/nfw.py | 11 ++- autogalaxy/profiles/mass/stellar/sersic.py | 85 +++++++++++++------ .../profiles/mass/stellar/sersic_gradient.py | 12 ++- 4 files changed, 104 insertions(+), 59 deletions(-) diff --git a/autogalaxy/profiles/mass/abstract/cse.py b/autogalaxy/profiles/mass/abstract/cse.py index 8d880dc4..60ef21fa 100644 --- a/autogalaxy/profiles/mass/abstract/cse.py +++ b/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] diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py index f129da9c..0b4edc03 100644 --- a/autogalaxy/profiles/mass/dark/nfw.py +++ b/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 diff --git a/autogalaxy/profiles/mass/stellar/sersic.py b/autogalaxy/profiles/mass/stellar/sersic.py index 8bd34424..2cd82c82 100644 --- a/autogalaxy/profiles/mass/stellar/sersic.py +++ b/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,9 +302,12 @@ 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 @@ -285,7 +315,7 @@ 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 diff --git a/autogalaxy/profiles/mass/stellar/sersic_gradient.py b/autogalaxy/profiles/mass/stellar/sersic_gradient.py index 7156fed8..5b488466 100644 --- a/autogalaxy/profiles/mass/stellar/sersic_gradient.py +++ b/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, )