From b68df8b068e34548ca9fd8460426043c9c4d1eeb Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 6 Aug 2026 13:42:20 -0400 Subject: [PATCH] fix: make PointMass and SMBH JAX-compatible PointMass.deflections_yx_2d_from routed through the legacy radial_grid_from / _cartesian_grid_via_radial_from helpers, which return an ArrayIrregular wrapper under xp=jnp on the irregular PSF-evaluation grids every imaging fit uses - jnp.multiply rejects it. Rewritten with raw xp ops (same maths, parity ~1e-18); potential_2d_from and convergence_2d_from hardened the same way (the convergence central-pixel argmin code was dead - its assignment was commented out). SMBH.__init__ used np.sqrt on the traced mass, raising TracerArrayConversionError; now tracer-safe pow. Co-Authored-By: Claude Fable 5 --- autogalaxy/profiles/mass/point/point.py | 23 +++++++++-------------- autogalaxy/profiles/mass/point/smbh.py | 2 +- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/autogalaxy/profiles/mass/point/point.py b/autogalaxy/profiles/mass/point/point.py index b78390da..2245cde0 100644 --- a/autogalaxy/profiles/mass/point/point.py +++ b/autogalaxy/profiles/mass/point/point.py @@ -50,28 +50,23 @@ def __init__( self.einstein_radius = einstein_radius def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): - squared_distances = np.square(grid[:, 0] - self.centre[0]) + np.square( - grid[:, 1] - self.centre[1] - ) - central_pixel = np.argmin(squared_distances) - - convergence = np.zeros(shape=grid.shape[0]) - # convergence[central_pixel] = np.pi * self.einstein_radius ** 2.0 - return convergence + return xp.zeros(grid.shape[0]) @aa.decorators.to_array @aa.decorators.transform def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): - r = xp.sqrt(grid.array[:, 0] ** 2 + grid.array[:, 1] ** 2 + 1e-20) - return self.einstein_radius ** 2 * xp.log(r) + y = xp.asarray(grid.array[:, 0]) + x = xp.asarray(grid.array[:, 1]) + r = xp.sqrt(y**2 + x**2 + 1e-20) + return self.einstein_radius**2 * xp.log(r) @aa.decorators.to_vector_yx @aa.decorators.transform def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): - grid_radii = self.radial_grid_from(grid=grid, xp=xp, **kwargs) - return self._cartesian_grid_via_radial_from( - grid=grid, radius=self.einstein_radius**2 / grid_radii, xp=xp - ) + y = xp.asarray(grid.array[:, 0]) + x = xp.asarray(grid.array[:, 1]) + alpha = self.einstein_radius**2 / (y**2 + x**2 + 1e-20) + return xp.stack((alpha * y, alpha * x), axis=-1) @property def is_point_mass(self): diff --git a/autogalaxy/profiles/mass/point/smbh.py b/autogalaxy/profiles/mass/point/smbh.py index e7392bcf..6ebceb40 100644 --- a/autogalaxy/profiles/mass/point/smbh.py +++ b/autogalaxy/profiles/mass/point/smbh.py @@ -60,7 +60,7 @@ def __init__( ) ) mass_angular = mass / critical_surface_density - einstein_radius = np.sqrt(mass_angular / np.pi) + einstein_radius = (mass_angular / np.pi) ** 0.5 super().__init__(centre=centre, einstein_radius=einstein_radius)