fix: make PointMass and SMBH JAX-compatible

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):

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)