From b8f17411e4802f6cb2bee6da1de1b4f557a73e4c Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Wed, 5 Nov 2025 19:13:35 +0000
Subject: [PATCH 01/18] first successful jit compile using .xp API

---
 autogalaxy/galaxy/galaxies.py                 | 20 ++---
 autogalaxy/galaxy/galaxy.py                   | 44 ++++-----
 autogalaxy/galaxy/to_inversion.py             |  5 ++
 autogalaxy/profiles/geometry_profiles.py      | 90 +++++++++----------
 autogalaxy/profiles/light/standard/sersic.py  | 20 ++---
 .../profiles/mass/sheets/external_shear.py    | 25 +++---
 6 files changed, 102 insertions(+), 102 deletions(-)

diff --git a/autogalaxy/galaxy/galaxies.py b/autogalaxy/galaxy/galaxies.py
index 8b172f51a..ce71e1dfe 100644
--- a/autogalaxy/galaxy/galaxies.py
+++ b/autogalaxy/galaxy/galaxies.py
@@ -50,7 +50,7 @@ def redshift(self):
         return self[0].redshift
 
     def image_2d_list_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
     ) -> List[aa.Array2D]:
         """
         Returns a list of the 2D images for each galaxy from a 2D grid of Cartesian (y,x) coordinates.
@@ -87,13 +87,13 @@ def image_2d_list_from(
             therefore is used to pass the `operated_only` input to these methods.
         """
         return [
-            galaxy.image_2d_from(grid=grid, operated_only=operated_only)
+            galaxy.image_2d_from(grid=grid, xp=xp, operated_only=operated_only)
             for galaxy in self
         ]
 
     @aa.grid_dec.to_array
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
     ) -> aa.Array2D:
         """
         Returns the 2D image of all galaxies summed from a 2D grid of Cartesian (y,x) coordinates.
@@ -114,10 +114,10 @@ def image_2d_from(
             apply these operations to the images, which may have the `operated_only` input passed to them. This input
             therefore is used to pass the `operated_only` input to these methods.
         """
-        return sum(self.image_2d_list_from(grid=grid, operated_only=operated_only))
+        return sum(self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only))
 
     def galaxy_image_2d_dict_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
     ) -> {Galaxy: np.ndarray}:
         """
         Returns a dictionary associating every `Galaxy` object with its corresponding 2D image, using the instance
@@ -143,7 +143,7 @@ def galaxy_image_2d_dict_from(
 
         galaxy_image_2d_dict = dict()
 
-        image_2d_list = self.image_2d_list_from(grid=grid, operated_only=operated_only)
+        image_2d_list = self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only)
 
         for galaxy_index, galaxy in enumerate(self):
             galaxy_image_2d_dict[galaxy] = image_2d_list[galaxy_index]
@@ -151,7 +151,7 @@ def galaxy_image_2d_dict_from(
         return galaxy_image_2d_dict
 
     @aa.grid_dec.to_vector_yx
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
         """
         Returns the summed 2D deflections angles of all galaxies from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -172,14 +172,14 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarr
         grid
             The 2D (y, x) coordinates where values of the deflections are evaluated.
         """
-        return sum(map(lambda g: g.deflections_yx_2d_from(grid=grid), self))
+        return sum(map(lambda g: g.deflections_yx_2d_from(grid=grid, xp=xp), self))
 
     @aa.grid_dec.to_grid
-    def traced_grid_2d_from(self, grid: aa.type.Grid2DLike) -> aa.type.Grid2DLike:
+    def traced_grid_2d_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.type.Grid2DLike:
         """
         Trace this plane's grid_stacks to the next plane, using its deflection angles.
         """
-        return grid - self.deflections_yx_2d_from(grid=grid)
+        return grid - self.deflections_yx_2d_from(grid=grid, xp=xp)
 
     @aa.grid_dec.to_array
     def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
diff --git a/autogalaxy/galaxy/galaxy.py b/autogalaxy/galaxy/galaxy.py
index 0a924c496..a46f9a5e6 100644
--- a/autogalaxy/galaxy/galaxy.py
+++ b/autogalaxy/galaxy/galaxy.py
@@ -161,7 +161,7 @@ def grid_radial_from(self, grid, centre, angle):
         )
 
     def image_2d_list_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
     ) -> List[aa.Array2D]:
         """
         Returns a list of the 2D images of the galaxy's light profiles from a 2D grid of Cartesian (y,x) coordinates.
@@ -194,7 +194,7 @@ def image_2d_list_from(
             zeros.
         """
         return [
-            light_profile.image_2d_from(grid=grid, operated_only=operated_only)
+            light_profile.image_2d_from(grid=grid, xp=xp, operated_only=operated_only)
             for light_profile in self.cls_list_from(
                 cls=LightProfile, cls_filtered=LightProfileLinear
             )
@@ -202,7 +202,7 @@ def image_2d_list_from(
 
     @aa.grid_dec.to_array
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None,
     ) -> Union[np.ndarray, aa.Array2D]:
         """
         Returns the 2D image of all galaxy light profiles summed from a 2D grid of Cartesian (y,x) coordinates.
@@ -228,11 +228,11 @@ def image_2d_from(
             len(self.cls_list_from(cls=LightProfile, cls_filtered=LightProfileLinear))
             > 0
         ):
-            return sum(self.image_2d_list_from(grid=grid, operated_only=operated_only))
-        return jnp.zeros((grid.shape[0],))
+            return sum(self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only))
+        return xp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_projected
-    def image_1d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray:
+    def image_1d_from(self, grid: aa.type.Grid2DLike, xp=np) -> np.ndarray:
         """
         Returns the summed 1D image of the galaxy's light profiles using a grid of Cartesian (y,x) coordinates.
 
@@ -262,10 +262,10 @@ def image_1d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray:
 
             return sum(image_1d_list)
 
-        return jnp.zeros((grid.shape[0],))
+        return xp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_vector_yx
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
         """
         Returns the summed 2D deflection angles of the galaxy's mass profiles from a 2D grid of Cartesian (y,x)
         coordinates.
@@ -286,15 +286,15 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarr
         if self.has(cls=MassProfile):
             return sum(
                 map(
-                    lambda p: p.deflections_yx_2d_from(grid=grid),
+                    lambda p: p.deflections_yx_2d_from(grid=grid, xp=xp),
                     self.cls_list_from(cls=MassProfile),
                 )
             )
 
-        return jnp.zeros((grid.shape[0], 2))
+        return xp.zeros((grid.shape[0], 2))
 
     @aa.grid_dec.to_array
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
         """
         Returns the summed 2D convergence of the galaxy's mass profiles from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -319,26 +319,26 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
                 )
             )
 
-        return jnp.zeros((grid.shape[0],))
+        return xp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_grid
-    def traced_grid_2d_from(self, grid: aa.type.Grid2DLike) -> aa.type.Grid2DLike:
+    def traced_grid_2d_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.type.Grid2DLike:
         """
         Trace an input grid using the galaxy's its deflection angles.
         """
         if isinstance(grid, aa.Grid2D):
             return aa.Grid2D(
-                values=grid - self.deflections_yx_2d_from(grid=grid),
+                values=grid - self.deflections_yx_2d_from(grid=grid, xp=xp),
                 mask=grid.mask,
                 over_sample_size=grid.over_sample_size,
                 over_sampled=grid.over_sampled
-                - self.deflections_yx_2d_from(grid=grid.over_sampled),
+                - self.deflections_yx_2d_from(grid=grid.over_sampled, xp=xp),
             )
 
-        return grid - self.deflections_yx_2d_from(grid=grid)
+        return grid - self.deflections_yx_2d_from(grid=grid, xp=xp)
 
     @aa.grid_dec.to_projected
-    def convergence_1d_from(self, grid: aa.type.Grid1D2DLike) -> np.ndarray:
+    def convergence_1d_from(self, grid: aa.type.Grid1D2DLike, xp=np) -> np.ndarray:
         """
         Returns the summed 1D convergence of the galaxy's mass profiles using a grid of Cartesian (y,x) coordinates.
 
@@ -369,10 +369,10 @@ def convergence_1d_from(self, grid: aa.type.Grid1D2DLike) -> np.ndarray:
 
             return sum(convergence_1d_list)
 
-        return jnp.zeros((grid.shape[0],))
+        return xp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
         """
         Returns the summed 2D potential of the galaxy's mass profiles from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -396,10 +396,10 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
                     self.cls_list_from(cls=MassProfile),
                 )
             )
-        return jnp.zeros((grid.shape[0],))
+        return xp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_projected
-    def potential_1d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray:
+    def potential_1d_from(self, grid: aa.type.Grid2DLike, xp=np) -> np.ndarray:
         """
         Returns the summed 1D potential of the galaxy's mass profiles using a grid of Cartesian (y,x) coordinates.
 
@@ -429,7 +429,7 @@ def potential_1d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray:
 
             return sum(potential_1d_list)
 
-        return jnp.zeros((grid.shape[0],))
+        return xp.zeros((grid.shape[0],))
 
     @property
     def half_light_radius(self):
diff --git a/autogalaxy/galaxy/to_inversion.py b/autogalaxy/galaxy/to_inversion.py
index 353e09b66..75087c53c 100644
--- a/autogalaxy/galaxy/to_inversion.py
+++ b/autogalaxy/galaxy/to_inversion.py
@@ -1,4 +1,5 @@
 from __future__ import annotations
+import numpy as np
 from typing import Dict, List, Optional, Type, Union
 
 from autoconf import cached_property
@@ -23,6 +24,7 @@ def __init__(
         dataset: Optional[Union[aa.Imaging, aa.Interferometer, aa.DatasetInterface]],
         adapt_images: Optional[AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
+        xp=np
     ):
         """
         Abstract class which interfaces a dataset and input modeling object (e.g. galaxies, a tracer) with the
@@ -74,6 +76,7 @@ def __init__(
         self.adapt_images = adapt_images
 
         self.settings_inversion = settings_inversion
+        self.xp = xp
 
     @property
     def psf(self) -> Optional[aa.Kernel2D]:
@@ -188,6 +191,7 @@ def __init__(
         galaxies: List[Galaxy],
         adapt_images: Optional[AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
+        xp=np,
     ):
         """
         Interfaces a dataset and input list of galaxies with the inversion module. to setup a
@@ -229,6 +233,7 @@ def __init__(
             dataset=dataset,
             adapt_images=adapt_images,
             settings_inversion=settings_inversion,
+            xp=xp,
         )
 
     @property
diff --git a/autogalaxy/profiles/geometry_profiles.py b/autogalaxy/profiles/geometry_profiles.py
index f0a67f5ad..296ca955e 100644
--- a/autogalaxy/profiles/geometry_profiles.py
+++ b/autogalaxy/profiles/geometry_profiles.py
@@ -39,10 +39,10 @@ def has(self, cls: Type) -> bool:
         """
         return aa.util.misc.has(values=self.__dict__.values(), cls=cls)
 
-    def transformed_to_reference_frame_grid_from(self, grid, **kwargs):
+    def transformed_to_reference_frame_grid_from(self, grid, xp=np, **kwargs):
         raise NotImplemented()
 
-    def transformed_from_reference_frame_grid_from(self, grid, **kwargs):
+    def transformed_from_reference_frame_grid_from(self, grid, xp=np, **kwargs):
         raise NotImplemented()
 
     def _radial_projected_shape_slim_from(
@@ -84,7 +84,7 @@ def __init__(self, centre: Tuple[float, float] = (0.0, 0.0)):
         super().__init__(centre=centre)
 
     @aa.grid_dec.to_array
-    def radial_grid_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
+    def radial_grid_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
         """
         Convert a grid of (y, x) coordinates, to their radial distances from the profile
         centre (e.g. :math: r = sqrt(x**2 + y**2)).
@@ -94,12 +94,12 @@ def radial_grid_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
         grid
             The grid of (y, x) coordinates which are converted to radial distances.
         """
-        return jnp.sqrt(
-            jnp.add(jnp.square(grid.array[:, 0]), jnp.square(grid.array[:, 1]))
+        return xp.sqrt(
+            xp.add(xp.square(grid.array[:, 0]), xp.square(grid.array[:, 1]))
         )
 
     def angle_to_profile_grid_from(
-        self, grid_angles: np.ndarray, **kwargs
+        self, grid_angles: np.ndarray, xp=np, **kwargs
     ) -> Tuple[np.ndarray, np.ndarray]:
         """
         Convert a grid of angles, defined in degrees counter-clockwise from the positive x-axis, to a grid of
@@ -110,11 +110,11 @@ def angle_to_profile_grid_from(
         grid_angles
             The angle theta counter-clockwise from the positive x-axis to each coordinate in radians.
         """
-        return jnp.cos(grid_angles), jnp.sin(grid_angles)
+        return xp.cos(grid_angles), xp.sin(grid_angles)
 
     @aa.grid_dec.to_grid
     def _cartesian_grid_via_radial_from(
-        self, grid: aa.type.Grid2DLike, radius: jnp.ndarray, **kwargs
+        self, grid: aa.type.Grid2DLike, radius: np.ndarray, xp=np, **kwargs
     ) -> aa.type.Grid2DLike:
         """
         Convert a grid of (y,x) coordinates with their specified radial distances (e.g. :math: r = x**2 + y**2) to
@@ -127,18 +127,13 @@ def _cartesian_grid_via_radial_from(
         radius
             The circular radius of each coordinate from the profile center.
         """
-        grid_angles = jnp.arctan2(grid.array[:, 0], grid.array[:, 1])
-        cos_theta, sin_theta = self.angle_to_profile_grid_from(grid_angles=grid_angles)
+        grid_angles = xp.arctan2(grid.array[:, 0], grid.array[:, 1])
+        cos_theta, sin_theta = self.angle_to_profile_grid_from(grid_angles=grid_angles, xp=xp)
 
-        if isinstance(radius, jnp.ndarray):
-            return jnp.multiply(radius[:, None], jnp.vstack((sin_theta, cos_theta)).T)
-        elif isinstance(radius, np.ndarray):
-            return np.multiply(radius[:, None], np.vstack((sin_theta, cos_theta)).T)
-
-        return jnp.multiply(radius.array[:, None], jnp.vstack((sin_theta, cos_theta)).T)
+        return xp.multiply(radius[:, None], xp.vstack((sin_theta, cos_theta)).T)
 
     @aa.grid_dec.to_grid
-    def transformed_to_reference_frame_grid_from(self, grid, **kwargs):
+    def transformed_to_reference_frame_grid_from(self, grid, xp=np, **kwargs):
         """
         Transform a grid of (y,x) coordinates to the reference frame of the profile.
 
@@ -149,10 +144,10 @@ def transformed_to_reference_frame_grid_from(self, grid, **kwargs):
         grid
             The (y, x) coordinates in the original reference frame of the grid.
         """
-        return jnp.subtract(grid.array, jnp.array(self.centre))
+        return xp.subtract(grid.array, xp.array(self.centre))
 
     @aa.grid_dec.to_grid
-    def transformed_from_reference_frame_grid_from(self, grid, **kwargs):
+    def transformed_from_reference_frame_grid_from(self, grid, xp=np, **kwargs):
         """
         Transform a grid of (y,x) coordinates from the reference frame of the profile to the original observer
         reference frame.
@@ -164,7 +159,7 @@ def transformed_from_reference_frame_grid_from(self, grid, **kwargs):
         grid
             The (y, x) coordinates in the reference frame of the profile.
         """
-        return jnp.add(grid.array, jnp.array(self.centre))
+        return xp.add(grid.array, xp.array(self.centre))
 
 
 class EllProfile(SphProfile):
@@ -233,13 +228,12 @@ def angle(self) -> float:
         """
         return convert.angle_from(ell_comps=self.ell_comps)
 
-    @property
-    def angle_radians(self) -> float:
+    def angle_radians(self, xp=np) -> float:
         """
         The position angle in radians of the major-axis of the ellipse defined by profile, defined counter clockwise
         from the positive x-axis (0.0 > angle > 2pi).
         """
-        return jnp.radians(self.angle)
+        return xp.radians(self.angle)
 
     @property
     def _cos_angle(self) -> float:
@@ -249,15 +243,15 @@ def _cos_angle(self) -> float:
     def _sin_angle(self) -> float:
         return self._cos_and_sin_to_x_axis()[1]
 
-    def _cos_and_sin_to_x_axis(self, **kwargs):
+    def _cos_and_sin_to_x_axis(self, xp=np, **kwargs):
         """
         Determine the sin and cosine of the angle between the profile's ellipse and the positive x-axis,
         counter-clockwise.
         """
-        angle_radians = jnp.radians(self.angle)
-        return jnp.cos(angle_radians), jnp.sin(angle_radians)
+        angle_radians = xp.radians(self.angle)
+        return xp.cos(angle_radians), xp.sin(angle_radians)
 
-    def angle_to_profile_grid_from(self, grid_angles, **kwargs):
+    def angle_to_profile_grid_from(self, grid_angles, xp=np, **kwargs):
         """
         The angle between each angle theta on the grid and the profile, in radians.
 
@@ -266,14 +260,14 @@ def angle_to_profile_grid_from(self, grid_angles, **kwargs):
         grid_angles
             The angle theta counter-clockwise from the positive x-axis to each coordinate in radians.
         """
-        theta_coordinate_to_profile = jnp.add(grid_angles, -self.angle_radians)
-        return jnp.cos(theta_coordinate_to_profile), jnp.sin(
+        theta_coordinate_to_profile = xp.add(grid_angles, - self.angle_radians)
+        return xp.cos(theta_coordinate_to_profile), xp.sin(
             theta_coordinate_to_profile
         )
 
     @aa.grid_dec.to_grid
     def rotated_grid_from_reference_frame_from(
-        self, grid, angle: Optional[float] = None, **kwargs
+        self, grid, xp=np, angle: Optional[float] = None, **kwargs
     ):
         """
         Rotate a grid of (y,x) coordinates which have been transformed to the elliptical reference frame of a profile
@@ -299,13 +293,13 @@ def rotated_grid_from_reference_frame_from(
             angle = self.angle
 
         return aa.util.geometry.transform_grid_2d_from_reference_frame(
-            grid_2d=grid, centre=(0.0, 0.0), angle=angle
+            grid_2d=grid, centre=(0.0, 0.0), angle=angle, xp=xp
         )
 
     @aa.grid_dec.to_array
     def elliptical_radii_grid_from(
-        self, grid: aa.type.Grid2DLike, **kwargs
-    ) -> jnp.ndarray:
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Convert a grid of (y,x) coordinates to their elliptical radii values: :math: (x^2 + (y^2/q))^0.5
 
@@ -314,17 +308,17 @@ def elliptical_radii_grid_from(
         grid
             The (y, x) coordinates in the reference frame of the elliptical profile.
         """
-        return jnp.sqrt(
-            jnp.add(
-                jnp.square(grid.array[:, 1]),
-                jnp.square(jnp.divide(grid.array[:, 0], self.axis_ratio)),
+        return xp.sqrt(
+            xp.add(
+                xp.square(grid.array[:, 1]),
+                xp.square(xp.divide(grid.array[:, 0], self.axis_ratio)),
             )
         )
 
     @aa.grid_dec.to_array
     def eccentric_radii_grid_from(
-        self, grid: aa.type.Grid2DLike, **kwargs
-    ) -> jnp.ndarray:
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Convert a grid of (y,x) coordinates to an eccentric radius: :math: axis_ratio^0.5 (x^2 + (y^2/q))^0.5
 
@@ -339,14 +333,14 @@ def eccentric_radii_grid_from(
             The (y, x) coordinates in the reference frame of the elliptical profile.
         """
 
-        grid_radii = self.elliptical_radii_grid_from(grid=grid, **kwargs)
+        grid_radii = self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs)
 
-        return jnp.multiply(jnp.sqrt(self.axis_ratio), grid_radii.array)
+        return xp.multiply(xp.sqrt(self.axis_ratio), grid_radii.array)
 
     @aa.grid_dec.to_grid
     def transformed_to_reference_frame_grid_from(
-        self, grid: aa.type.Grid2DLike, **kwargs
-    ) -> jnp.ndarray:
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Transform a grid of (y,x) coordinates to the reference frame of the profile.
 
@@ -358,14 +352,14 @@ def transformed_to_reference_frame_grid_from(
             The (y, x) coordinates in the original reference frame of the grid.
         """
         if self.__class__.__name__.endswith("Sph"):
-            return super().transformed_to_reference_frame_grid_from(grid=grid)
+            return super().transformed_to_reference_frame_grid_from(grid=grid, xp=xp)
         return aa.util.geometry.transform_grid_2d_to_reference_frame(
-            grid_2d=grid.array, centre=self.centre, angle=self.angle
+            grid_2d=grid.array, centre=self.centre, angle=self.angle, xp=xp
         )
 
     @aa.grid_dec.to_grid
     def transformed_from_reference_frame_grid_from(
-        self, grid: aa.type.Grid2DLike, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
     ) -> aa.type.Grid2DLike:
         """
         Transform a grid of (y,x) coordinates from the reference frame of the profile to the original observer
@@ -379,10 +373,10 @@ def transformed_from_reference_frame_grid_from(
             The (y, x) coordinates in the reference frame of the profile.
         """
         if self.__class__.__name__.startswith("Sph"):
-            return super().transformed_from_reference_frame_grid_from(grid=grid)
+            return super().transformed_from_reference_frame_grid_from(grid=grid, xp=xp)
 
         return aa.util.geometry.transform_grid_2d_from_reference_frame(
-            grid_2d=grid.array, centre=self.centre, angle=self.angle
+            grid_2d=grid.array, centre=self.centre, angle=self.angle, xp=xp
         )
 
     def _eta_u(self, u, coordinates):
diff --git a/autogalaxy/profiles/light/standard/sersic.py b/autogalaxy/profiles/light/standard/sersic.py
index 0c0de357a..3a5facf2f 100644
--- a/autogalaxy/profiles/light/standard/sersic.py
+++ b/autogalaxy/profiles/light/standard/sersic.py
@@ -120,7 +120,7 @@ def __init__(
             sersic_index=sersic_index,
         )
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarray:
+    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np, **kwargs) -> np.ndarray:
         """
         Returns the 2D image of the Sersic light profile from a grid of coordinates which are the radial distances of
         each coordinate from the its `centre`.
@@ -131,14 +131,14 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarra
             The radial distances from the centre of the profile, for each coordinate on the grid.
         """
         seterr(all="ignore")
-        return jnp.multiply(
+        return xp.multiply(
             self._intensity,
-            jnp.exp(
-                jnp.multiply(
+            xp.exp(
+                xp.multiply(
                     -self.sersic_constant,
-                    jnp.add(
-                        jnp.power(
-                            jnp.divide(grid_radii.array, self.effective_radius),
+                    xp.add(
+                        xp.power(
+                            xp.divide(grid_radii.array, self.effective_radius),
                             1.0 / self.sersic_index,
                         ),
                         -1,
@@ -152,7 +152,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarra
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs
+        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, xp=np, **kwargs
     ) -> aa.Array2D:
         """
         Returns the Sersic light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
@@ -171,9 +171,9 @@ def image_2d_from(
             The image of the Sersic evaluated at every (y,x) coordinate on the transformed grid.
         """
 
-        grid_radii = self.eccentric_radii_grid_from(grid=grid, **kwargs)
+        grid_radii = self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs)
 
-        return self.image_2d_via_radii_from(grid_radii=grid_radii, **kwargs)
+        return self.image_2d_via_radii_from(grid_radii=grid_radii, xp=xp, **kwargs)
 
 
 class SersicSph(Sersic):
diff --git a/autogalaxy/profiles/mass/sheets/external_shear.py b/autogalaxy/profiles/mass/sheets/external_shear.py
index fb71e7bf0..d6b781d0b 100644
--- a/autogalaxy/profiles/mass/sheets/external_shear.py
+++ b/autogalaxy/profiles/mass/sheets/external_shear.py
@@ -1,4 +1,4 @@
-import jax.numpy as jnp
+import numpy as np
 
 import autoarray as aa
 
@@ -39,24 +39,24 @@ def average_convergence_of_1_radius(self):
         return 0.0
 
     @aa.grid_dec.to_array
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
-        return jnp.zeros(shape=grid.shape[0])
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return xp.zeros(shape=grid.shape[0])
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         shear_angle = (
             self.angle - 90
         )  ##to be onsistent with autolens deflection angle calculation
-        phig = jnp.deg2rad(shear_angle)
+        phig = xp.deg2rad(shear_angle)
         shear_amp = self.magnitude
-        phicoord = jnp.arctan2(grid.array[:, 0], grid.array[:, 1])
-        rcoord = jnp.sqrt(grid.array[:, 0] ** 2.0 + grid.array[:, 1] ** 2.0)
+        phicoord = xp.arctan2(grid.array[:, 0], grid.array[:, 1])
+        rcoord = xp.sqrt(grid.array[:, 0] ** 2.0 + grid.array[:, 1] ** 2.0)
 
-        return -0.5 * shear_amp * rcoord**2 * jnp.cos(2 * (phicoord - phig))
+        return -0.5 * shear_amp * rcoord**2 * xp.cos(2 * (phicoord - phig))
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -66,8 +66,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
 
         """
-        deflection_y = -jnp.multiply(self.magnitude, grid.array[:, 0])
-        deflection_x = jnp.multiply(self.magnitude, grid.array[:, 1])
+        deflection_y = -xp.multiply(self.magnitude, grid.array[:, 0])
+        deflection_x = xp.multiply(self.magnitude, grid.array[:, 1])
         return self.rotated_grid_from_reference_frame_from(
-            jnp.vstack((deflection_y, deflection_x)).T
+            grid=xp.vstack((deflection_y, deflection_x)).T,
+            xp=xp,
         )

From 4977952aedae0abfcaa7e70c33cc771c1826ceac Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Thu, 6 Nov 2025 20:13:09 +0000
Subject: [PATCH 02/18] xp added to isothermal and shear

---
 autogalaxy/convert.py                         | 104 +++++++++---------
 autogalaxy/profiles/geometry_profiles.py      |  22 ++--
 autogalaxy/profiles/light/standard/sersic.py  |   2 +-
 autogalaxy/profiles/mass/dark/nfw.py          |   4 +-
 .../profiles/mass/sheets/external_shear.py    |  18 ++-
 autogalaxy/profiles/mass/total/isothermal.py  |  45 ++++----
 6 files changed, 94 insertions(+), 101 deletions(-)

diff --git a/autogalaxy/convert.py b/autogalaxy/convert.py
index 3a199d220..7b653043b 100644
--- a/autogalaxy/convert.py
+++ b/autogalaxy/convert.py
@@ -1,10 +1,8 @@
+import numpy as np
 from typing import Tuple
 
-import jax
-import jax.numpy as jnp
 
-
-def ell_comps_from(axis_ratio: float, angle: float) -> Tuple[float, float]:
+def ell_comps_from(axis_ratio: float, angle: float, xp=np) -> Tuple[float, float]:
     """
     Returns the elliptical components e1 and e2 of a light or mass profile from an input angle in degrees and axis
     ratio.
@@ -23,14 +21,14 @@ def ell_comps_from(axis_ratio: float, angle: float) -> Tuple[float, float]:
     angle
         Rotation angle of light profile counter-clockwise from positive x-axis.
     """
-    angle *= jnp.pi / 180.0
+    angle *= xp.pi / 180.0
     fac = (1 - axis_ratio) / (1 + axis_ratio)
-    ellip_y = fac * jnp.sin(2 * angle)
-    ellip_x = fac * jnp.cos(2 * angle)
+    ellip_y = fac * xp.sin(2 * angle)
+    ellip_x = fac * xp.cos(2 * angle)
     return (ellip_y, ellip_x)
 
 
-def axis_ratio_and_angle_from(ell_comps: Tuple[float, float]) -> Tuple[float, float]:
+def axis_ratio_and_angle_from(ell_comps: Tuple[float, float], xp=np) -> Tuple[float, float]:
     """
     Returns the axis-ratio and position angle in degrees (-45 < angle < 135.0) from input elliptical components e1
     and e2 of a light or mass profile.
@@ -60,19 +58,19 @@ def axis_ratio_and_angle_from(ell_comps: Tuple[float, float]) -> Tuple[float, fl
     ell_comps
         The elliptical components of the light or mass profile which are converted to an angle.
     """
-    angle = jnp.arctan2(ell_comps[0], ell_comps[1]) / 2
-    angle *= 180.0 / jnp.pi
+    angle = xp.arctan2(ell_comps[0], ell_comps[1]) / 2
+    angle *= 180.0 / xp.pi
 
-    angle = jax.lax.select(angle < -45, angle + 180, angle)
+    angle = xp.where(angle < -45, angle + 180, angle)
 
-    fac = jnp.sqrt(ell_comps[1] ** 2 + ell_comps[0] ** 2)
-    fac = jax.lax.min(fac, 0.999)
+    fac = xp.sqrt(ell_comps[1] ** 2 + ell_comps[0] ** 2)
+    fac = xp.min(fac, 0.999)
 
     axis_ratio = (1 - fac) / (1 + fac)
     return axis_ratio, angle
 
 
-def axis_ratio_from(ell_comps: Tuple[float, float]):
+def axis_ratio_from(ell_comps: Tuple[float, float], xp=np):
     """
     Returns the axis-ratio from input elliptical components e1 and e2 of a light or mass profile.
 
@@ -95,11 +93,11 @@ def axis_ratio_from(ell_comps: Tuple[float, float]):
     ell_comps
         The elliptical components of the light or mass profile which are converted to an angle.
     """
-    axis_ratio, angle = axis_ratio_and_angle_from(ell_comps=ell_comps)
+    axis_ratio, angle = axis_ratio_and_angle_from(ell_comps=ell_comps, xp=xp)
     return axis_ratio
 
 
-def angle_from(ell_comps: Tuple[float, float]) -> float:
+def angle_from(ell_comps: Tuple[float, float], xp=np) -> float:
     """
     Returns the position angle in degrees (-45 < angle < 135.0) from input elliptical components e1 and e2
     of a light or mass profile.
@@ -128,26 +126,26 @@ def angle_from(ell_comps: Tuple[float, float]) -> float:
     ell_comps
         The elliptical components of the light or mass profile which are converted to an angle.
     """
-    axis_ratio, angle = axis_ratio_and_angle_from(ell_comps=ell_comps)
+    axis_ratio, angle = axis_ratio_and_angle_from(ell_comps=ell_comps, xp=xp)
 
     return angle
 
 
-def shear_gamma_1_2_from(magnitude: float, angle: float) -> Tuple[float, float]:
+def shear_gamma_1_2_from(magnitude: float, angle: float, xp=np) -> Tuple[float, float]:
     """
     Returns the shear gamma 1 and gamma 2 values an input shear magnitude and angle in degrees.
 
     The gamma 1 and gamma 2 components of a shear are given by:
 
-    gamma_1 = magnitude * jnp.cos(2 * angle * jnp.pi / 180.0)
-    gamma_2 = magnitude * jnp.sin(2 * angle * jnp.pi / 180.0)
+    gamma_1 = magnitude * xp.cos(2 * angle * xp.pi / 180.0)
+    gamma_2 = magnitude * xp.sin(2 * angle * xp.pi / 180.0)
 
     Which are the values this function returns.
 
     Converting from gamma 1 and gamma 2 to magnitude and angle is given by:
 
-    magnitude = jnp.sqrt(gamma_1**2 + gamma_2**2)
-    angle = jnp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / jnp.pi
+    magnitude = xp.sqrt(gamma_1**2 + gamma_2**2)
+    angle = xp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / xp.pi
 
     Parameters
     ----------
@@ -156,26 +154,26 @@ def shear_gamma_1_2_from(magnitude: float, angle: float) -> Tuple[float, float]:
     angle
         Rotation angle of light profile counter-clockwise from positive x-axis.
     """
-    gamma_1 = magnitude * jnp.cos(2 * angle * jnp.pi / 180.0)
-    gamma_2 = magnitude * jnp.sin(2 * angle * jnp.pi / 180.0)
+    gamma_1 = magnitude * xp.cos(2 * angle * xp.pi / 180.0)
+    gamma_2 = magnitude * xp.sin(2 * angle * xp.pi / 180.0)
     return (gamma_1, gamma_2)
 
 
 def shear_magnitude_and_angle_from(
-    gamma_1: float, gamma_2: float
+    gamma_1: float, gamma_2: float, xp=np
 ) -> Tuple[float, float]:
     """
     Returns the shear magnitude and angle in degrees from input shear gamma 1 and gamma 2 values.
 
     The gamma 1 and gamma 2 components of a shear are given by:
 
-    gamma_1 = magnitude * jnp.cos(2 * angle * jnp.pi / 180.0)
-    gamma_2 = magnitude * jnp.sin(2 * angle * jnp.pi / 180.0)
+    gamma_1 = magnitude * xp.cos(2 * angle * xp.pi / 180.0)
+    gamma_2 = magnitude * xp.sin(2 * angle * xp.pi / 180.0)
 
     Converting from gamma 1 and gamma 2 to magnitude and angle is given by:
 
-    magnitude = jnp.sqrt(gamma_1**2 + gamma_2**2)
-    angle = jnp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / jnp.pi
+    magnitude = xp.sqrt(gamma_1**2 + gamma_2**2)
+    angle = xp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / xp.pi
 
     Which are the values this function returns.
 
@@ -193,30 +191,30 @@ def shear_magnitude_and_angle_from(
     gamma_2
         The gamma 2 component of the shear.
     """
-    angle = jnp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / jnp.pi
-    magnitude = jnp.sqrt(gamma_1**2 + gamma_2**2)
+    angle = xp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / xp.pi
+    magnitude = xp.sqrt(gamma_1**2 + gamma_2**2)
 
-    angle = jnp.where(angle < 0, angle + 180.0, angle)
-    angle = jnp.where(
-        (jnp.abs(angle - 90.0) > 45.0) & (angle > 90.0), angle - 180.0, angle
+    angle = xp.where(angle < 0, angle + 180.0, angle)
+    angle = xp.where(
+        (xp.abs(angle - 90.0) > 45.0) & (angle > 90.0), angle - 180.0, angle
     )
 
     return magnitude, angle
 
 
-def shear_magnitude_from(gamma_1: float, gamma_2: float) -> float:
+def shear_magnitude_from(gamma_1: float, gamma_2: float, xp=np) -> float:
     """
     Returns the shear magnitude and angle in degrees from input shear gamma 1 and gamma 2 values.
 
     The gamma 1 and gamma 2 components of a shear are given by:
 
-    gamma_1 = magnitude * jnp.cos(2 * angle * jnp.pi / 180.0)
-    gamma_2 = magnitude * jnp.sin(2 * angle * jnp.pi / 180.0)
+    gamma_1 = magnitude * xp.cos(2 * angle * xp.pi / 180.0)
+    gamma_2 = magnitude * xp.sin(2 * angle * xp.pi / 180.0)
 
     Converting from gamma 1 and gamma 2 to magnitude and angle is given by:
 
-    magnitude = jnp.sqrt(gamma_1**2 + gamma_2**2)
-    angle = jnp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / jnp.pi
+    magnitude = xp.sqrt(gamma_1**2 + gamma_2**2)
+    angle = xp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / xp.pi
 
     The magnitude value is what this function returns.
 
@@ -227,23 +225,23 @@ def shear_magnitude_from(gamma_1: float, gamma_2: float) -> float:
     gamma_2
         The gamma 2 component of the shear.
     """
-    magnitude, angle = shear_magnitude_and_angle_from(gamma_1=gamma_1, gamma_2=gamma_2)
+    magnitude, angle = shear_magnitude_and_angle_from(gamma_1=gamma_1, gamma_2=gamma_2, xp=xp)
     return magnitude
 
 
-def shear_angle_from(gamma_1: float, gamma_2: float) -> float:
+def shear_angle_from(gamma_1: float, gamma_2: float, xp=np) -> float:
     """
     Returns the shear magnitude and angle in degrees from input shear gamma 1 and gamma 2 values.
 
     The gamma 1 and gamma 2 components of a shear are given by:
 
-    gamma_1 = magnitude * jnp.cos(2 * angle * jnp.pi / 180.0)
-    gamma_2 = magnitude * jnp.sin(2 * angle * jnp.pi / 180.0)
+    gamma_1 = magnitude * xp.cos(2 * angle * xp.pi / 180.0)
+    gamma_2 = magnitude * xp.sin(2 * angle * xp.pi / 180.0)
 
     Converting from gamma 1 and gamma 2 to magnitude and angle is given by:
 
-    magnitude = jnp.sqrt(gamma_1**2 + gamma_2**2)
-    angle = jnp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / jnp.pi
+    magnitude = xp.sqrt(gamma_1**2 + gamma_2**2)
+    angle = xp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / xp.pi
 
     The angle value is what this function returns.
 
@@ -254,12 +252,12 @@ def shear_angle_from(gamma_1: float, gamma_2: float) -> float:
     gamma_2
         The gamma 2 component of the shear.
     """
-    magnitude, angle = shear_magnitude_and_angle_from(gamma_1=gamma_1, gamma_2=gamma_2)
+    magnitude, angle = shear_magnitude_and_angle_from(gamma_1=gamma_1, gamma_2=gamma_2, xp=xp)
     return angle
 
 
 def multipole_k_m_and_phi_m_from(
-    multipole_comps: Tuple[float, float], m: int
+    multipole_comps: Tuple[float, float], m: int, xp=np
 ) -> Tuple[float, float]:
     """
     Returns the multipole normalization value `k_m` and angle `phi` from the multipole component parameters.
@@ -290,16 +288,16 @@ def multipole_k_m_and_phi_m_from(
     The normalization and angle parameters of the multipole.
     """
     phi_m = (
-        jnp.arctan2(multipole_comps[0], multipole_comps[1]) * 180.0 / jnp.pi / float(m)
+        xp.arctan2(multipole_comps[0], multipole_comps[1]) * 180.0 / xp.pi / float(m)
     )
-    k_m = jnp.sqrt(multipole_comps[1] ** 2 + multipole_comps[0] ** 2)
+    k_m = xp.sqrt(multipole_comps[1] ** 2 + multipole_comps[0] ** 2)
 
-    phi_m = jnp.where(phi_m < -90.0 / m, phi_m + 360.0 / m, phi_m)
+    phi_m = xp.where(phi_m < -90.0 / m, phi_m + 360.0 / m, phi_m)
 
     return k_m, phi_m
 
 
-def multipole_comps_from(k_m: float, phi_m: float, m: int) -> Tuple[float, float]:
+def multipole_comps_from(k_m: float, phi_m: float, m: int, xp=np) -> Tuple[float, float]:
     """
     Returns the multipole component parameters from their normalization value `k_m` and angle `phi`.
 
@@ -323,7 +321,7 @@ def multipole_comps_from(k_m: float, phi_m: float, m: int) -> Tuple[float, float
     """
     from astropy import units
 
-    multipole_comp_0 = k_m * jnp.sin(phi_m * float(m) * units.deg.to(units.rad))
-    multipole_comp_1 = k_m * jnp.cos(phi_m * float(m) * units.deg.to(units.rad))
+    multipole_comp_0 = k_m * xp.sin(phi_m * float(m) * units.deg.to(units.rad))
+    multipole_comp_1 = k_m * xp.cos(phi_m * float(m) * units.deg.to(units.rad))
 
     return (multipole_comp_0, multipole_comp_1)
diff --git a/autogalaxy/profiles/geometry_profiles.py b/autogalaxy/profiles/geometry_profiles.py
index 296ca955e..a0f349144 100644
--- a/autogalaxy/profiles/geometry_profiles.py
+++ b/autogalaxy/profiles/geometry_profiles.py
@@ -213,27 +213,25 @@ def __init__(
 
         self.ell_comps = ell_comps
 
-    @property
-    def axis_ratio(self) -> float:
+    def axis_ratio(self, xp=np) -> float:
         """
         The ratio of the minor-axis to major-axis (b/a) of the ellipse defined by profile (0.0 > q > 1.0).
         """
-        return convert.axis_ratio_from(ell_comps=self.ell_comps)
+        return convert.axis_ratio_from(ell_comps=self.ell_comps, xp=xp)
 
-    @property
-    def angle(self) -> float:
+    def angle(self, xp=np) -> float:
         """
         The position angle in degrees of the major-axis of the ellipse defined by profile, defined counter clockwise
         from the positive x-axis (0.0 > angle > 180.0).
         """
-        return convert.angle_from(ell_comps=self.ell_comps)
+        return convert.angle_from(ell_comps=self.ell_comps, xp=xp)
 
     def angle_radians(self, xp=np) -> float:
         """
         The position angle in radians of the major-axis of the ellipse defined by profile, defined counter clockwise
         from the positive x-axis (0.0 > angle > 2pi).
         """
-        return xp.radians(self.angle)
+        return xp.radians(self.angle(xp=xp))
 
     @property
     def _cos_angle(self) -> float:
@@ -248,7 +246,7 @@ def _cos_and_sin_to_x_axis(self, xp=np, **kwargs):
         Determine the sin and cosine of the angle between the profile's ellipse and the positive x-axis,
         counter-clockwise.
         """
-        angle_radians = xp.radians(self.angle)
+        angle_radians = xp.radians(self.angle(xp=xp))
         return xp.cos(angle_radians), xp.sin(angle_radians)
 
     def angle_to_profile_grid_from(self, grid_angles, xp=np, **kwargs):
@@ -260,7 +258,7 @@ def angle_to_profile_grid_from(self, grid_angles, xp=np, **kwargs):
         grid_angles
             The angle theta counter-clockwise from the positive x-axis to each coordinate in radians.
         """
-        theta_coordinate_to_profile = xp.add(grid_angles, - self.angle_radians)
+        theta_coordinate_to_profile = xp.add(grid_angles, - self.angle_radians(xp=xp))
         return xp.cos(theta_coordinate_to_profile), xp.sin(
             theta_coordinate_to_profile
         )
@@ -290,7 +288,7 @@ def rotated_grid_from_reference_frame_from(
         """
 
         if angle is None:
-            angle = self.angle
+            angle = self.angle(xp=xp)
 
         return aa.util.geometry.transform_grid_2d_from_reference_frame(
             grid_2d=grid, centre=(0.0, 0.0), angle=angle, xp=xp
@@ -354,7 +352,7 @@ def transformed_to_reference_frame_grid_from(
         if self.__class__.__name__.endswith("Sph"):
             return super().transformed_to_reference_frame_grid_from(grid=grid, xp=xp)
         return aa.util.geometry.transform_grid_2d_to_reference_frame(
-            grid_2d=grid.array, centre=self.centre, angle=self.angle, xp=xp
+            grid_2d=grid.array, centre=self.centre, angle=self.angle(xp=xp), xp=xp
         )
 
     @aa.grid_dec.to_grid
@@ -376,7 +374,7 @@ def transformed_from_reference_frame_grid_from(
             return super().transformed_from_reference_frame_grid_from(grid=grid, xp=xp)
 
         return aa.util.geometry.transform_grid_2d_from_reference_frame(
-            grid_2d=grid.array, centre=self.centre, angle=self.angle, xp=xp
+            grid_2d=grid.array, centre=self.centre, angle=self.angle(xp=xp), xp=xp
         )
 
     def _eta_u(self, u, coordinates):
diff --git a/autogalaxy/profiles/light/standard/sersic.py b/autogalaxy/profiles/light/standard/sersic.py
index 3a5facf2f..746b4993b 100644
--- a/autogalaxy/profiles/light/standard/sersic.py
+++ b/autogalaxy/profiles/light/standard/sersic.py
@@ -152,7 +152,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np, **kwargs) -> np
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, xp=np, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
     ) -> aa.Array2D:
         """
         Returns the Sersic light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index ae8fc3018..8f21428ee 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -254,7 +254,7 @@ def nfw_2d(r):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Analytic calculation shear from Heyrovský & Karamazov 2024
 
@@ -287,7 +287,7 @@ def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         # Rotation for shear
 
         shear_field = self.rotated_grid_from_reference_frame_from(
-            grid=np.vstack((g2, g1)).T, angle=self.angle * 2
+            grid=np.vstack((g2, g1)).T, angle=self.angle(xp=xp) * 2
         )
 
         return aa.VectorYX2DIrregular(values=shear_field, grid=grid)
diff --git a/autogalaxy/profiles/mass/sheets/external_shear.py b/autogalaxy/profiles/mass/sheets/external_shear.py
index d6b781d0b..fec2c55e7 100644
--- a/autogalaxy/profiles/mass/sheets/external_shear.py
+++ b/autogalaxy/profiles/mass/sheets/external_shear.py
@@ -24,13 +24,11 @@ def __init__(self, gamma_1: float = 0.0, gamma_2: float = 0.0):
         self.gamma_1 = gamma_1
         self.gamma_2 = gamma_2
 
-    @property
-    def magnitude(self):
-        return convert.shear_magnitude_from(gamma_1=self.gamma_1, gamma_2=self.gamma_2)
+    def magnitude(self, xp=np):
+        return convert.shear_magnitude_from(gamma_1=self.gamma_1, gamma_2=self.gamma_2, xp=xp)
 
-    @property
-    def angle(self):
-        return convert.shear_angle_from(gamma_1=self.gamma_1, gamma_2=self.gamma_2)
+    def angle(self, xp=np):
+        return convert.shear_angle_from(gamma_1=self.gamma_1, gamma_2=self.gamma_2, xp=xp)
 
     def convergence_func(self, grid_radius: float) -> float:
         return 0.0
@@ -45,10 +43,10 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
     @aa.grid_dec.to_array
     def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         shear_angle = (
-            self.angle - 90
+            self.angle(xp=xp) - 90
         )  ##to be onsistent with autolens deflection angle calculation
         phig = xp.deg2rad(shear_angle)
-        shear_amp = self.magnitude
+        shear_amp = self.magnitude(xp=xp)
         phicoord = xp.arctan2(grid.array[:, 0], grid.array[:, 1])
         rcoord = xp.sqrt(grid.array[:, 0] ** 2.0 + grid.array[:, 1] ** 2.0)
 
@@ -66,8 +64,8 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
 
         """
-        deflection_y = -xp.multiply(self.magnitude, grid.array[:, 0])
-        deflection_x = xp.multiply(self.magnitude, grid.array[:, 1])
+        deflection_y = -xp.multiply(self.magnitude(xp=xp), grid.array[:, 0])
+        deflection_x = xp.multiply(self.magnitude(xp=xp), grid.array[:, 1])
         return self.rotated_grid_from_reference_frame_from(
             grid=xp.vstack((deflection_y, deflection_x)).T,
             xp=xp,
diff --git a/autogalaxy/profiles/mass/total/isothermal.py b/autogalaxy/profiles/mass/total/isothermal.py
index 2b0cd5004..43dbab2bd 100644
--- a/autogalaxy/profiles/mass/total/isothermal.py
+++ b/autogalaxy/profiles/mass/total/isothermal.py
@@ -1,4 +1,4 @@
-import jax.numpy as jnp
+import numpy as np
 
 from typing import Tuple
 
@@ -7,7 +7,7 @@
 from autogalaxy.profiles.mass.total.power_law import PowerLaw
 
 
-def psi_from(grid, axis_ratio, core_radius):
+def psi_from(grid, axis_ratio, core_radius, xp=np):
     """
     Returns the $\Psi$ term in expressions for the calculation of the deflection of an elliptical isothermal mass
     distribution. This is used in the `Isothermal` and `Chameleon` `MassProfile`'s.
@@ -31,7 +31,7 @@ def psi_from(grid, axis_ratio, core_radius):
         The value of the Psi term.
 
     """
-    return jnp.sqrt(
+    return xp.sqrt(
         (axis_ratio**2.0 * (grid.array[:, 1] ** 2.0 + core_radius**2.0))
         + grid.array[:, 0] ** 2.0
         + 1e-16
@@ -66,14 +66,13 @@ def __init__(
             slope=2.0,
         )
 
-    @property
-    def axis_ratio(self):
+    def axis_ratio(self, xp=np):
         axis_ratio = super().axis_ratio
-        return jnp.minimum(axis_ratio, 0.99999)
+        return xp.minimum(axis_ratio, 0.99999)
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
 
@@ -89,30 +88,30 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         factor = (
             2.0
             * self.einstein_radius_rescaled
-            * self.axis_ratio
-            / jnp.sqrt(1 - self.axis_ratio**2)
+            * self.axis_ratio(xp=xp)
+            / xp.sqrt(1 - self.axis_ratio(xp=xp)**2)
         )
 
-        psi = psi_from(grid=grid, axis_ratio=self.axis_ratio, core_radius=0.0)
+        psi = psi_from(grid=grid, axis_ratio=self.axis_ratio(xp=xp), core_radius=0.0)
 
-        deflection_y = jnp.arctanh(
-            jnp.divide(
-                jnp.multiply(jnp.sqrt(1 - self.axis_ratio**2), grid.array[:, 0]), psi
+        deflection_y = xp.arctanh(
+            xp.divide(
+                xp.multiply(xp.sqrt(1 - self.axis_ratio(xp=xp)**2), grid.array[:, 0]), psi
             )
         )
-        deflection_x = jnp.arctan(
-            jnp.divide(
-                jnp.multiply(jnp.sqrt(1 - self.axis_ratio**2), grid.array[:, 1]), psi
+        deflection_x = xp.arctan(
+            xp.divide(
+                xp.multiply(xp.sqrt(1 - self.axis_ratio(xp=xp)**2), grid.array[:, 1]), psi
             )
         )
         return self.rotated_grid_from_reference_frame_from(
-            grid=jnp.multiply(factor, jnp.vstack((deflection_y, deflection_x)).T),
+            grid=xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T),
             **kwargs,
         )
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the (gamma_y, gamma_x) shear vector field on a grid of (y,x) arc-second coordinates.
 
@@ -134,18 +133,18 @@ def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         gamma_2 = (
             -2
             * convergence.array
-            * jnp.divide(
+            * xp.divide(
                 grid.array[:, 1] * grid.array[:, 0],
                 grid.array[:, 1] ** 2 + grid.array[:, 0] ** 2,
             )
         )
-        gamma_1 = -convergence.array * jnp.divide(
+        gamma_1 = -convergence.array * xp.divide(
             grid.array[:, 1] ** 2 - grid.array[:, 0] ** 2,
             grid.array[:, 1] ** 2 + grid.array[:, 0] ** 2,
         )
 
         shear_field = self.rotated_grid_from_reference_frame_from(
-            grid=jnp.vstack((gamma_2, gamma_1)).T, angle=self.angle * 2
+            grid=xp.vstack((gamma_2, gamma_1)).T, angle=self.angle(xp=xp) * 2
         )
 
         return aa.VectorYX2DIrregular(values=shear_field, grid=grid)
@@ -191,7 +190,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
 
@@ -202,6 +201,6 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         """
         return self._cartesian_grid_via_radial_from(
             grid=grid,
-            radius=jnp.full(grid.shape[0], 2.0 * self.einstein_radius_rescaled),
+            radius=xp.full(grid.shape[0], 2.0 * self.einstein_radius_rescaled),
             **kwargs,
         )

From eef7e6a9ec3a83fc897a35b4fe65258d97307307 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Thu, 6 Nov 2025 20:26:38 +0000
Subject: [PATCH 03/18] more xps to get isothermal to trace

---
 autogalaxy/convert.py                         |  7 ++-
 autogalaxy/profiles/geometry_profiles.py      |  6 +--
 .../profiles/light/standard/chameleon.py      |  9 ++--
 .../profiles/light/standard/gaussian.py       |  2 +-
 autogalaxy/profiles/light/standard/moffat.py  |  2 +-
 autogalaxy/profiles/light/standard/sersic.py  |  2 +-
 autogalaxy/profiles/mass/abstract/cse.py      |  2 +-
 autogalaxy/profiles/mass/abstract/mge_jax.py  |  2 +-
 .../profiles/mass/abstract/mge_numpy.py       |  2 +-
 .../mass/abstract/mge_numpy_coleman.py        |  2 +-
 autogalaxy/profiles/mass/dark/abstract.py     |  2 +-
 autogalaxy/profiles/mass/dark/gnfw.py         | 10 ++--
 autogalaxy/profiles/mass/dark/nfw.py          |  6 +--
 autogalaxy/profiles/mass/stellar/chameleon.py | 35 +++++++------
 autogalaxy/profiles/mass/stellar/gaussian.py  | 28 +++++------
 autogalaxy/profiles/mass/stellar/sersic.py    | 12 ++---
 .../profiles/mass/stellar/sersic_gradient.py  | 12 ++---
 autogalaxy/profiles/mass/total/isothermal.py  | 23 +++++----
 autogalaxy/profiles/mass/total/power_law.py   | 49 +++++++++----------
 .../profiles/mass/total/power_law_broken.py   | 16 +++---
 .../profiles/mass/total/power_law_core.py     | 29 ++++++-----
 21 files changed, 126 insertions(+), 132 deletions(-)

diff --git a/autogalaxy/convert.py b/autogalaxy/convert.py
index 7b653043b..00b68b685 100644
--- a/autogalaxy/convert.py
+++ b/autogalaxy/convert.py
@@ -64,7 +64,11 @@ def axis_ratio_and_angle_from(ell_comps: Tuple[float, float], xp=np) -> Tuple[fl
     angle = xp.where(angle < -45, angle + 180, angle)
 
     fac = xp.sqrt(ell_comps[1] ** 2 + ell_comps[0] ** 2)
-    fac = xp.min(fac, 0.999)
+    if xp.__name__.startswith("jax"):
+        import jax
+        fac = jax.lax.min(fac, 0.999)
+    else:  # NumPy
+        fac = np.minimum(fac, 0.999)
 
     axis_ratio = (1 - fac) / (1 + fac)
     return axis_ratio, angle
@@ -127,7 +131,6 @@ def angle_from(ell_comps: Tuple[float, float], xp=np) -> float:
         The elliptical components of the light or mass profile which are converted to an angle.
     """
     axis_ratio, angle = axis_ratio_and_angle_from(ell_comps=ell_comps, xp=xp)
-
     return angle
 
 
diff --git a/autogalaxy/profiles/geometry_profiles.py b/autogalaxy/profiles/geometry_profiles.py
index a0f349144..c1c7d0247 100644
--- a/autogalaxy/profiles/geometry_profiles.py
+++ b/autogalaxy/profiles/geometry_profiles.py
@@ -309,7 +309,7 @@ def elliptical_radii_grid_from(
         return xp.sqrt(
             xp.add(
                 xp.square(grid.array[:, 1]),
-                xp.square(xp.divide(grid.array[:, 0], self.axis_ratio)),
+                xp.square(xp.divide(grid.array[:, 0], self.axis_ratio())),
             )
         )
 
@@ -333,7 +333,7 @@ def eccentric_radii_grid_from(
 
         grid_radii = self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs)
 
-        return xp.multiply(xp.sqrt(self.axis_ratio), grid_radii.array)
+        return xp.multiply(xp.sqrt(self.axis_ratio()), grid_radii.array)
 
     @aa.grid_dec.to_grid
     def transformed_to_reference_frame_grid_from(
@@ -383,7 +383,7 @@ def _eta_u(self, u, coordinates):
                 u
                 * (
                     (coordinates[1] ** 2)
-                    + (coordinates[0] ** 2 / (1 - (1 - self.axis_ratio**2) * u))
+                    + (coordinates[0] ** 2 / (1 - (1 - self.axis_ratio()**2) * u))
                 )
             )
         )
diff --git a/autogalaxy/profiles/light/standard/chameleon.py b/autogalaxy/profiles/light/standard/chameleon.py
index 2935ddf76..76016bd0f 100644
--- a/autogalaxy/profiles/light/standard/chameleon.py
+++ b/autogalaxy/profiles/light/standard/chameleon.py
@@ -45,13 +45,12 @@ def __init__(
         self.core_radius_0 = core_radius_0
         self.core_radius_1 = core_radius_1
 
-    @property
-    def axis_ratio(self) -> float:
+    def axis_ratio(self, xp=np) -> float:
         """
         The elliptical isothermal mass profile deflection angles break down for perfectly spherical systems where
         `axis_ratio=1.0`, thus we remove these solutions.
         """
-        axis_ratio = super().axis_ratio
+        axis_ratio = super().axis_ratio(xp=xp)
         return axis_ratio if axis_ratio < 0.99999 else 0.99999
 
     def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
@@ -65,10 +64,10 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
             The radial distances from the centre of the profile, for each coordinate on the grid.
         """
 
-        axis_ratio_factor = (1.0 + self.axis_ratio) ** 2.0
+        axis_ratio_factor = (1.0 + self.axis_ratio()) ** 2.0
 
         return jnp.multiply(
-            self._intensity / (1 + self.axis_ratio),
+            self._intensity / (1 + self.axis_ratio()),
             jnp.add(
                 jnp.divide(
                     1.0,
diff --git a/autogalaxy/profiles/light/standard/gaussian.py b/autogalaxy/profiles/light/standard/gaussian.py
index 6feaee7b4..5d07ebb94 100644
--- a/autogalaxy/profiles/light/standard/gaussian.py
+++ b/autogalaxy/profiles/light/standard/gaussian.py
@@ -65,7 +65,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
             jnp.exp(
                 -0.5
                 * jnp.square(
-                    jnp.divide(grid_radii.array, self.sigma / jnp.sqrt(self.axis_ratio))
+                    jnp.divide(grid_radii.array, self.sigma / jnp.sqrt(self.axis_ratio()))
                 )
             ),
         )
diff --git a/autogalaxy/profiles/light/standard/moffat.py b/autogalaxy/profiles/light/standard/moffat.py
index adcc3af25..2a9ed4baa 100644
--- a/autogalaxy/profiles/light/standard/moffat.py
+++ b/autogalaxy/profiles/light/standard/moffat.py
@@ -63,7 +63,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
             jnp.power(
                 1
                 + jnp.square(
-                    jnp.divide(grid_radii.array, self.alpha / jnp.sqrt(self.axis_ratio))
+                    jnp.divide(grid_radii.array, self.alpha / jnp.sqrt(self.axis_ratio()))
                 ),
                 -self.beta,
             ),
diff --git a/autogalaxy/profiles/light/standard/sersic.py b/autogalaxy/profiles/light/standard/sersic.py
index 746b4993b..8445c349e 100644
--- a/autogalaxy/profiles/light/standard/sersic.py
+++ b/autogalaxy/profiles/light/standard/sersic.py
@@ -53,7 +53,7 @@ def elliptical_effective_radius(self) -> float:
         The elliptical effective radius instead describes the major-axis radius of the ellipse containing
         half the light, and may be more appropriate for highly flattened systems like disk galaxies.
         """
-        return self.effective_radius / jnp.sqrt(self.axis_ratio)
+        return self.effective_radius / jnp.sqrt(self.axis_ratio())
 
     @property
     def sersic_constant(self) -> float:
diff --git a/autogalaxy/profiles/mass/abstract/cse.py b/autogalaxy/profiles/mass/abstract/cse.py
index a6db34172..5b12688b9 100644
--- a/autogalaxy/profiles/mass/abstract/cse.py
+++ b/autogalaxy/profiles/mass/abstract/cse.py
@@ -166,7 +166,7 @@ def _deflections_2d_via_cse_from(self, grid: np.ndarray, **kwargs) -> np.ndarray
         amplitude_list, core_radius_list = self.decompose_convergence_via_cse(
             grid_radii=self.radial_grid_from(grid=grid, **kwargs)
         )
-        q = self.axis_ratio
+        q = self.axis_ratio()
         q2 = q**2.0
         grid_y = grid.array[:, 0]
         grid_x = grid.array[:, 1]
diff --git a/autogalaxy/profiles/mass/abstract/mge_jax.py b/autogalaxy/profiles/mass/abstract/mge_jax.py
index 81beab896..75dc16dad 100644
--- a/autogalaxy/profiles/mass/abstract/mge_jax.py
+++ b/autogalaxy/profiles/mass/abstract/mge_jax.py
@@ -143,7 +143,7 @@ def _convergence_2d_via_mge_from(
     def _deflections_2d_via_mge_from(
         self, grid, sigmas_factor=1.0, func_terms=28, func_gaussians=20
     ):
-        axis_ratio = jnp.min(jnp.array([self.axis_ratio, 0.9999]))
+        axis_ratio = jnp.min(jnp.array([self.axis_ratio(), 0.9999]))
 
         amps, sigmas = self.decompose_convergence_via_mge(
             func_terms=func_terms, func_gaussians=func_gaussians
diff --git a/autogalaxy/profiles/mass/abstract/mge_numpy.py b/autogalaxy/profiles/mass/abstract/mge_numpy.py
index 124d278cf..a4fcbd29a 100644
--- a/autogalaxy/profiles/mass/abstract/mge_numpy.py
+++ b/autogalaxy/profiles/mass/abstract/mge_numpy.py
@@ -265,7 +265,7 @@ def convergence_func_gaussian(self, grid_radii, sigma, intensity):
     def _deflections_2d_via_mge_from(
         self, grid, sigmas_factor=1.0, func_terms=None, func_gaussians=None
     ):
-        axis_ratio = np.array(self.axis_ratio)
+        axis_ratio = np.array(self.axis_ratio())
 
         if axis_ratio > 0.9999:
             axis_ratio = 0.9999
diff --git a/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py b/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py
index 5f5b8a97b..3a423bc25 100644
--- a/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py
+++ b/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py
@@ -261,7 +261,7 @@ def convergence_func_gaussian(self, grid_radii, sigma, intensity):
         )
 
     def _deflections_2d_via_mge_from(self, grid, sigmas_factor=1.0):
-        axis_ratio = self.axis_ratio
+        axis_ratio = self.axis_ratio()
 
         if axis_ratio > 0.9999:
             axis_ratio = 0.9999
diff --git a/autogalaxy/profiles/mass/dark/abstract.py b/autogalaxy/profiles/mass/dark/abstract.py
index 34dfc3430..f9a394486 100644
--- a/autogalaxy/profiles/mass/dark/abstract.py
+++ b/autogalaxy/profiles/mass/dark/abstract.py
@@ -379,4 +379,4 @@ def mass_at_200_solar_masses(
 
     @property
     def ellipticity_rescale(self):
-        return 1.0 - ((1.0 - self.axis_ratio) / 2.0)
+        return 1.0 - ((1.0 - self.axis_ratio()) / 2.0)
diff --git a/autogalaxy/profiles/mass/dark/gnfw.py b/autogalaxy/profiles/mass/dark/gnfw.py
index 047154085..941d7f0c6 100644
--- a/autogalaxy/profiles/mass/dark/gnfw.py
+++ b/autogalaxy/profiles/mass/dark/gnfw.py
@@ -14,7 +14,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
     @aa.grid_dec.transform
     def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, **kwargs):
         return self._deflections_2d_via_mge_from(
-            grid=grid, sigmas_factor=self.axis_ratio
+            grid=grid, sigmas_factor=self.axis_ratio()
         )
 
     @aa.grid_dec.to_vector_yx
@@ -51,7 +51,7 @@ def calculate_deflection_component(npow, yx_index):
                 deflection_grid[i] = (
                     2.0
                     * self.kappa_s
-                    * self.axis_ratio
+                    * self.axis_ratio()
                     * grid[i, yx_index]
                     * quad(
                         self.deflection_func,
@@ -61,7 +61,7 @@ def calculate_deflection_component(npow, yx_index):
                             grid.array[i, 0],
                             grid.array[i, 1],
                             npow,
-                            self.axis_ratio,
+                            self.axis_ratio(),
                             minimum_log_eta,
                             maximum_log_eta,
                             tabulate_bins,
@@ -220,14 +220,14 @@ def deflection_integrand(x, kappa_radius, scale_radius, inner_slope):
             )
 
         for i in range(grid.shape[0]):
-            potential_grid[i] = (2.0 * self.kappa_s * self.axis_ratio) * quad(
+            potential_grid[i] = (2.0 * self.kappa_s * self.axis_ratio()) * quad(
                 self.potential_func,
                 a=0.0,
                 b=1.0,
                 args=(
                     grid.array[i, 0],
                     grid.array[i, 1],
-                    self.axis_ratio,
+                    self.axis_ratio(),
                     minimum_log_eta,
                     maximum_log_eta,
                     tabulate_bins,
diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index 8f21428ee..d4a137c8a 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -61,7 +61,7 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs):
         from scipy.integrate import quad
 
         def calculate_deflection_component(npow, index):
-            deflection_grid = np.array(self.axis_ratio * grid.array[:, index])
+            deflection_grid = np.array(self.axis_ratio() * grid.array[:, index])
 
             for i in range(grid.shape[0]):
                 deflection_grid[i] *= (
@@ -74,7 +74,7 @@ def calculate_deflection_component(npow, index):
                             grid.array[i, 0],
                             grid.array[i, 1],
                             npow,
-                            self.axis_ratio,
+                            self.axis_ratio(),
                             self.scale_radius,
                         ),
                     )[0]
@@ -171,7 +171,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
                 args=(
                     grid.array[i, 0],
                     grid.array[i, 1],
-                    self.axis_ratio,
+                    self.axis_ratio(),
                     self.kappa_s,
                     self.scale_radius,
                 ),
diff --git a/autogalaxy/profiles/mass/stellar/chameleon.py b/autogalaxy/profiles/mass/stellar/chameleon.py
index 00a35a123..92040e6b4 100644
--- a/autogalaxy/profiles/mass/stellar/chameleon.py
+++ b/autogalaxy/profiles/mass/stellar/chameleon.py
@@ -69,49 +69,49 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs):
             2.0
             * self.mass_to_light_ratio
             * self.intensity
-            / (1 + self.axis_ratio)
-            * self.axis_ratio
-            / jnp.sqrt(1.0 - self.axis_ratio**2.0)
+            / (1 + self.axis_ratio())
+            * self.axis_ratio()
+            / jnp.sqrt(1.0 - self.axis_ratio()**2.0)
         )
 
         core_radius_0 = jnp.sqrt(
-            (4.0 * self.core_radius_0**2.0) / (1.0 + self.axis_ratio) ** 2
+            (4.0 * self.core_radius_0**2.0) / (1.0 + self.axis_ratio()) ** 2
         )
         core_radius_1 = jnp.sqrt(
-            (4.0 * self.core_radius_1**2.0) / (1.0 + self.axis_ratio) ** 2
+            (4.0 * self.core_radius_1**2.0) / (1.0 + self.axis_ratio()) ** 2
         )
 
         psi0 = psi_from(
-            grid=grid, axis_ratio=self.axis_ratio, core_radius=core_radius_0
+            grid=grid, axis_ratio=self.axis_ratio(), core_radius=core_radius_0
         )
         psi1 = psi_from(
-            grid=grid, axis_ratio=self.axis_ratio, core_radius=core_radius_1
+            grid=grid, axis_ratio=self.axis_ratio(), core_radius=core_radius_1
         )
 
         deflection_y0 = jnp.arctanh(
             jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio**2.0), grid.array[:, 0]),
-                jnp.add(psi0, self.axis_ratio**2.0 * core_radius_0),
+                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio()**2.0), grid.array[:, 0]),
+                jnp.add(psi0, self.axis_ratio()**2.0 * core_radius_0),
             )
         )
 
         deflection_x0 = jnp.arctan(
             jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio**2.0), grid.array[:, 1]),
+                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio()**2.0), grid.array[:, 1]),
                 jnp.add(psi0, core_radius_0),
             )
         )
 
         deflection_y1 = jnp.arctanh(
             jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio**2.0), grid.array[:, 0]),
-                jnp.add(psi1, self.axis_ratio**2.0 * core_radius_1),
+                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio()**2.0), grid.array[:, 0]),
+                jnp.add(psi1, self.axis_ratio()**2.0 * core_radius_1),
             )
         )
 
         deflection_x1 = jnp.arctan(
             jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio**2.0), grid.array[:, 1]),
+                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio()**2.0), grid.array[:, 1]),
                 jnp.add(psi1, core_radius_1),
             )
         )
@@ -153,10 +153,10 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray):
             The radial distance from the centre of the profile. for each coordinate on the grid.
         """
 
-        axis_ratio_factor = (1.0 + self.axis_ratio) ** 2.0
+        axis_ratio_factor = (1.0 + self.axis_ratio()) ** 2.0
 
         return jnp.multiply(
-            self.intensity / (1 + self.axis_ratio),
+            self.intensity / (1 + self.axis_ratio()),
             jnp.add(
                 jnp.divide(
                     1.0,
@@ -179,9 +179,8 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray):
             ),
         )
 
-    @property
-    def axis_ratio(self):
-        axis_ratio = super().axis_ratio
+    def axis_ratio(self, xp=np):
+        axis_ratio = super().axis_ratio(xp=xp)
         return axis_ratio if axis_ratio < 0.99999 else 0.99999
 
 
diff --git a/autogalaxy/profiles/mass/stellar/gaussian.py b/autogalaxy/profiles/mass/stellar/gaussian.py
index d9e5efc89..5b6f92c79 100644
--- a/autogalaxy/profiles/mass/stellar/gaussian.py
+++ b/autogalaxy/profiles/mass/stellar/gaussian.py
@@ -76,7 +76,7 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs):
             self.mass_to_light_ratio
             * self.intensity
             * self.sigma
-            * np.sqrt((2 * np.pi) / (1.0 - self.axis_ratio**2.0))
+            * np.sqrt((2 * np.pi) / (1.0 - self.axis_ratio()**2.0))
             * self.zeta_from(grid=grid)
         )
 
@@ -103,7 +103,7 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs):
         from scipy.integrate import quad
 
         def calculate_deflection_component(npow, index):
-            deflection_grid = np.array(self.axis_ratio * grid.array[:, index])
+            deflection_grid = np.array(self.axis_ratio() * grid.array[:, index])
 
             for i in range(grid.shape[0]):
                 deflection_grid[i] *= (
@@ -117,8 +117,8 @@ def calculate_deflection_component(npow, index):
                             grid.array[i, 0],
                             grid.array[i, 1],
                             npow,
-                            self.axis_ratio,
-                            self.sigma / np.sqrt(self.axis_ratio),
+                            self.axis_ratio(),
+                            self.sigma / np.sqrt(self.axis_ratio()),
                         ),
                     )[0]
                 )
@@ -180,28 +180,24 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray):
             np.exp(
                 -0.5
                 * np.square(
-                    np.divide(grid_radii.array, self.sigma / np.sqrt(self.axis_ratio))
+                    np.divide(grid_radii.array, self.sigma / np.sqrt(self.axis_ratio()))
                 )
             ),
         )
 
-    @property
-    def axis_ratio(self):
-        axis_ratio = super().axis_ratio
-        if use_jax:
-            return jax.lax.select(axis_ratio < 0.9999, axis_ratio, 0.9999)
-        else:
-            return axis_ratio if axis_ratio < 0.9999 else 0.9999
+    def axis_ratio(self, xp=np):
+        axis_ratio = super().axis_ratio(xp=xp)
+        return xp.where(axis_ratio < 0.9999, axis_ratio, 0.9999)
 
     def zeta_from(self, grid: aa.type.Grid2DLike):
 
         from scipy.special import wofz
 
-        q2 = self.axis_ratio**2.0
+        q2 = self.axis_ratio()**2.0
         ind_pos_y = grid.array[:, 0] >= 0
         shape_grid = np.shape(grid)
         output_grid = np.zeros((shape_grid[0]), dtype=np.complex128)
-        scale_factor = self.axis_ratio / (self.sigma * np.sqrt(2.0 * (1.0 - q2)))
+        scale_factor = self.axis_ratio() / (self.sigma * np.sqrt(2.0 * (1.0 - q2)))
 
         xs_0 = grid.array[:, 1][ind_pos_y] * scale_factor
         ys_0 = grid.array[:, 0][ind_pos_y] * scale_factor
@@ -211,7 +207,7 @@ def zeta_from(self, grid: aa.type.Grid2DLike):
         output_grid[ind_pos_y] = -1j * (
             wofz(xs_0 + 1j * ys_0)
             - np.exp(-(xs_0**2.0) * (1.0 - q2) - ys_0 * ys_0 * (1.0 / q2 - 1.0))
-            * wofz(self.axis_ratio * xs_0 + 1j * ys_0 / self.axis_ratio)
+            * wofz(self.axis_ratio() * xs_0 + 1j * ys_0 / self.axis_ratio())
         )
 
         output_grid[~ind_pos_y] = np.conj(
@@ -219,7 +215,7 @@ def zeta_from(self, grid: aa.type.Grid2DLike):
             * (
                 wofz(xs_1 + 1j * ys_1)
                 - np.exp(-(xs_1**2.0) * (1.0 - q2) - ys_1 * ys_1 * (1.0 / q2 - 1.0))
-                * wofz(self.axis_ratio * xs_1 + 1j * ys_1 / self.axis_ratio)
+                * wofz(self.axis_ratio() * xs_1 + 1j * ys_1 / self.axis_ratio())
             )
         )
 
diff --git a/autogalaxy/profiles/mass/stellar/sersic.py b/autogalaxy/profiles/mass/stellar/sersic.py
index da2f54e73..893ec5f32 100644
--- a/autogalaxy/profiles/mass/stellar/sersic.py
+++ b/autogalaxy/profiles/mass/stellar/sersic.py
@@ -145,7 +145,7 @@ def deflections_2d_via_mge_from(
         """
         return self._deflections_2d_via_mge_from(
             grid=grid,
-            sigmas_factor=np.sqrt(self.axis_ratio),
+            sigmas_factor=np.sqrt(self.axis_ratio()),
             func_terms=func_terms,
             func_gaussians=func_gaussians,
         )
@@ -313,7 +313,7 @@ def decompose_convergence_via_cse(
             mass_to_light_gradient=0.0,
         )
 
-        scaled_effective_radius = self.effective_radius / np.sqrt(self.axis_ratio)
+        scaled_effective_radius = self.effective_radius / np.sqrt(self.axis_ratio())
         radii_min = scaled_effective_radius / 10.0**lower_dex
         radii_max = scaled_effective_radius * 10.0**upper_dex
 
@@ -354,7 +354,7 @@ def sersic_constant(self):
 
     @property
     def ellipticity_rescale(self):
-        return 1.0 - ((1.0 - self.axis_ratio) / 2.0)
+        return 1.0 - ((1.0 - self.axis_ratio()) / 2.0)
 
     @property
     def elliptical_effective_radius(self):
@@ -366,7 +366,7 @@ def elliptical_effective_radius(self):
         The elliptical effective radius instead describes the major-axis radius of the ellipse containing \
         half the light, and may be more appropriate for highly flattened systems like disk galaxies.
         """
-        return self.effective_radius / np.sqrt(self.axis_ratio)
+        return self.effective_radius / np.sqrt(self.axis_ratio())
 
 
 class Sersic(AbstractSersic, MassProfileMGE, MassProfileCSE):
@@ -387,7 +387,7 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs):
         def calculate_deflection_component(npow, index):
             sersic_constant = self.sersic_constant
 
-            deflection_grid = self.axis_ratio * grid.array[:, index]
+            deflection_grid = self.axis_ratio() * grid.array[:, index]
 
             for i in range(grid.shape[0]):
                 deflection_grid = deflection_grid.at[i].multiply(
@@ -401,7 +401,7 @@ def calculate_deflection_component(npow, index):
                             grid.array[i, 0],
                             grid.array[i, 1],
                             npow,
-                            self.axis_ratio,
+                            self.axis_ratio(),
                             self.sersic_index,
                             self.effective_radius,
                             sersic_constant,
diff --git a/autogalaxy/profiles/mass/stellar/sersic_gradient.py b/autogalaxy/profiles/mass/stellar/sersic_gradient.py
index 92c07fd0d..0b7230305 100644
--- a/autogalaxy/profiles/mass/stellar/sersic_gradient.py
+++ b/autogalaxy/profiles/mass/stellar/sersic_gradient.py
@@ -66,7 +66,7 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs):
         def calculate_deflection_component(npow, index):
             sersic_constant = self.sersic_constant
 
-            deflection_grid = np.array(self.axis_ratio * grid.array[:, index])
+            deflection_grid = np.array(self.axis_ratio() * grid.array[:, index])
 
             for i in range(grid.shape[0]):
                 deflection_grid[i] *= (
@@ -80,7 +80,7 @@ def calculate_deflection_component(npow, index):
                             grid.array[i, 0],
                             grid.array[i, 1],
                             npow,
-                            self.axis_ratio,
+                            self.axis_ratio(),
                             self.sersic_index,
                             self.effective_radius,
                             self.mass_to_light_gradient,
@@ -142,7 +142,7 @@ def convergence_func(self, grid_radius: float) -> float:
         return (
             self.mass_to_light_ratio
             * (
-                ((self.axis_ratio * grid_radius) / self.effective_radius)
+                ((self.axis_ratio() * grid_radius) / self.effective_radius)
                 ** -self.mass_to_light_gradient
             )
             * self.image_2d_via_radii_from(grid_radius)
@@ -157,7 +157,7 @@ def sersic_gradient_2D(r):
                 self.mass_to_light_ratio
                 * self.intensity
                 * (
-                    ((self.axis_ratio * r) / self.effective_radius)
+                    ((self.axis_ratio() * r) / self.effective_radius)
                     ** -self.mass_to_light_gradient
                 )
                 * np.exp(
@@ -205,7 +205,7 @@ def decompose_convergence_via_cse(
             mass_to_light_gradient=self.mass_to_light_gradient,
         )
 
-        scaled_effective_radius = self.effective_radius / np.sqrt(self.axis_ratio)
+        scaled_effective_radius = self.effective_radius / np.sqrt(self.axis_ratio())
         radii_min = scaled_effective_radius / 10.0**lower_dex
         radii_max = scaled_effective_radius * 10.0**upper_dex
 
@@ -214,7 +214,7 @@ def sersic_gradient_2D(r):
                 self.mass_to_light_ratio
                 * self.intensity
                 * (
-                    ((self.axis_ratio * r) / scaled_effective_radius)
+                    ((self.axis_ratio() * r) / scaled_effective_radius)
                     ** -self.mass_to_light_gradient
                 )
                 * np.exp(
diff --git a/autogalaxy/profiles/mass/total/isothermal.py b/autogalaxy/profiles/mass/total/isothermal.py
index 43dbab2bd..ffccf87ac 100644
--- a/autogalaxy/profiles/mass/total/isothermal.py
+++ b/autogalaxy/profiles/mass/total/isothermal.py
@@ -67,7 +67,7 @@ def __init__(
         )
 
     def axis_ratio(self, xp=np):
-        axis_ratio = super().axis_ratio
+        axis_ratio = super().axis_ratio(xp=xp)
         return xp.minimum(axis_ratio, 0.99999)
 
     @aa.grid_dec.to_vector_yx
@@ -87,21 +87,21 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
         factor = (
             2.0
-            * self.einstein_radius_rescaled
-            * self.axis_ratio(xp=xp)
-            / xp.sqrt(1 - self.axis_ratio(xp=xp)**2)
+            * self.einstein_radius_rescaled(xp)
+            * self.axis_ratio(xp)
+            / xp.sqrt(1 - self.axis_ratio()(xp)**2)
         )
 
-        psi = psi_from(grid=grid, axis_ratio=self.axis_ratio(xp=xp), core_radius=0.0)
+        psi = psi_from(grid=grid, axis_ratio=self.axis_ratio()(xp=xp), core_radius=0.0)
 
         deflection_y = xp.arctanh(
             xp.divide(
-                xp.multiply(xp.sqrt(1 - self.axis_ratio(xp=xp)**2), grid.array[:, 0]), psi
+                xp.multiply(xp.sqrt(1 - self.axis_ratio()(xp=xp)**2), grid.array[:, 0]), psi
             )
         )
         deflection_x = xp.arctan(
             xp.divide(
-                xp.multiply(xp.sqrt(1 - self.axis_ratio(xp=xp)**2), grid.array[:, 1]), psi
+                xp.multiply(xp.sqrt(1 - self.axis_ratio()(xp=xp)**2), grid.array[:, 1]), psi
             )
         )
         return self.rotated_grid_from_reference_frame_from(
@@ -169,14 +169,13 @@ def __init__(
             centre=centre, ell_comps=(0.0, 0.0), einstein_radius=einstein_radius
         )
 
-    @property
-    def axis_ratio(self):
+    def axis_ratio(self, xp=np):
         return 1.0
 
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the potential on a grid of (y,x) arc-second coordinates.
 
@@ -186,7 +185,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
         eta = self.elliptical_radii_grid_from(grid=grid, **kwargs)
-        return 2.0 * self.einstein_radius_rescaled * eta
+        return 2.0 * self.einstein_radius_rescaled(xp) * eta
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
@@ -201,6 +200,6 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         return self._cartesian_grid_via_radial_from(
             grid=grid,
-            radius=xp.full(grid.shape[0], 2.0 * self.einstein_radius_rescaled),
+            radius=xp.full(grid.shape[0], 2.0 * self.einstein_radius_rescaled(xp)),
             **kwargs,
         )
diff --git a/autogalaxy/profiles/mass/total/power_law.py b/autogalaxy/profiles/mass/total/power_law.py
index 829189acb..2a06fe154 100644
--- a/autogalaxy/profiles/mass/total/power_law.py
+++ b/autogalaxy/profiles/mass/total/power_law.py
@@ -1,6 +1,5 @@
-import jax.numpy as jnp
 from .jax_utils import omega
-
+import numpy as np
 from typing import Tuple
 
 import autoarray as aa
@@ -53,7 +52,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
 
@@ -71,25 +70,25 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
         slope = self.slope - 1.0
         einstein_radius = (
-            2.0 / (self.axis_ratio**-0.5 + self.axis_ratio**0.5)
+            2.0 / (self.axis_ratio()**-0.5 + self.axis_ratio()**0.5)
         ) * self.einstein_radius
 
-        factor = jnp.divide(1.0 - self.axis_ratio, 1.0 + self.axis_ratio)
-        b = jnp.multiply(einstein_radius, jnp.sqrt(self.axis_ratio))
-        angle = jnp.arctan2(
-            grid.array[:, 0], jnp.multiply(self.axis_ratio, grid.array[:, 1])
+        factor = xp.divide(1.0 - self.axis_ratio(), 1.0 + self.axis_ratio())
+        b = xp.multiply(einstein_radius, xp.sqrt(self.axis_ratio()))
+        angle = xp.arctan2(
+            grid.array[:, 0], xp.multiply(self.axis_ratio(), grid.array[:, 1])
         )  # Note, this angle is not the position angle
-        z = jnp.add(
-            jnp.multiply(jnp.cos(angle), 1 + 0j), jnp.multiply(jnp.sin(angle), 0 + 1j)
+        z = xp.add(
+            xp.multiply(xp.cos(angle), 1 + 0j), xp.multiply(xp.sin(angle), 0 + 1j)
         )
 
-        R = jnp.sqrt(
-            (self.axis_ratio * grid.array[:, 1]) ** 2 + grid.array[:, 0] ** 2 + 1e-16
+        R = xp.sqrt(
+            (self.axis_ratio() * grid.array[:, 1]) ** 2 + grid.array[:, 0] ** 2 + 1e-16
         )
         zh = omega(z, slope, factor, n_terms=20)
 
         complex_angle = (
-            2.0 * b / (1.0 + self.axis_ratio) * (b / R) ** (slope - 1.0) * zh
+            2.0 * b / (1.0 + self.axis_ratio()) * (b / R) ** (slope - 1.0) * zh
         )
 
         deflection_y = complex_angle.imag
@@ -101,15 +100,15 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         deflection_x *= rescale_factor
 
         return self.rotated_grid_from_reference_frame_from(
-            grid=jnp.vstack((deflection_y, deflection_x)).T
+            grid=xp.vstack((deflection_y, deflection_x)).T
         )
 
-    def convergence_func(self, grid_radius: float) -> float:
-        return self.einstein_radius_rescaled * grid_radius.array ** (-(self.slope - 1))
+    def convergence_func(self, grid_radius: float, xp=np) -> float:
+        return self.einstein_radius_rescaled(xp) * grid_radius.array ** (-(self.slope - 1))
 
     @staticmethod
-    def potential_func(u, y, x, axis_ratio, slope, core_radius):
-        _eta_u = jnp.sqrt((u * ((x**2) + (y**2 / (1 - (1 - axis_ratio**2) * u)))))
+    def potential_func(u, y, x, axis_ratio, slope, core_radius, xp=np):
+        _eta_u = xp.sqrt((u * ((x**2) + (y**2 / (1 - (1 - axis_ratio**2) * u)))))
         return (
             (_eta_u / u)
             * ((3.0 - slope) * _eta_u) ** -1.0
@@ -147,15 +146,15 @@ def __init__(
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
-        eta = self.radial_grid_from(grid=grid, **kwargs).array
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        eta = self.radial_grid_from(grid=grid, xp=np, **kwargs).array
         deflection_r = (
             2.0
-            * self.einstein_radius_rescaled
-            * jnp.divide(
-                jnp.power(eta, (3.0 - self.slope)),
-                jnp.multiply((3.0 - self.slope), eta),
+            * self.einstein_radius_rescaled(xp)
+            * xp.divide(
+                xp.power(eta, (3.0 - self.slope)),
+                xp.multiply((3.0 - self.slope), eta),
             )
         )
 
-        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_r)
+        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_r, xp=np)
diff --git a/autogalaxy/profiles/mass/total/power_law_broken.py b/autogalaxy/profiles/mass/total/power_law_broken.py
index 4e17c5c6d..e8d82ef05 100644
--- a/autogalaxy/profiles/mass/total/power_law_broken.py
+++ b/autogalaxy/profiles/mass/total/power_law_broken.py
@@ -31,7 +31,7 @@ def __init__(
         super().__init__(centre=centre, ell_comps=ell_comps)
 
         self.einstein_radius = einstein_radius
-        self.einstein_radius_elliptical = jnp.sqrt(self.axis_ratio) * einstein_radius
+        self.einstein_radius_elliptical = jnp.sqrt(self.axis_ratio()) * einstein_radius
         self.break_radius = break_radius
         self.inner_slope = inner_slope
         self.outer_slope = outer_slope
@@ -58,7 +58,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         """
 
         # Ell radius
-        radius = jnp.hypot(grid.array[:, 1] * self.axis_ratio, grid.array[:, 0])
+        radius = jnp.hypot(grid.array[:, 1] * self.axis_ratio(), grid.array[:, 0])
 
         # Inside break radius
         kappa_inner = self.kB * (self.break_radius / radius) ** self.inner_slope
@@ -84,14 +84,14 @@ def deflections_yx_2d_from(self, grid, max_terms=20, **kwargs):
         z = grid.array[:, 1] + 1j * grid.array[:, 0]
 
         # Ell radius
-        R = jnp.hypot(z.real * self.axis_ratio, z.imag)
+        R = jnp.hypot(z.real * self.axis_ratio(), z.imag)
 
         # Factors common to eq. 18 and 19
         factors = (
             2
             * self.kB
             * (self.break_radius**2)
-            / (self.axis_ratio * z * (2 - self.inner_slope))
+            / (self.axis_ratio() * z * (2 - self.inner_slope))
         )
 
         # Hypergeometric functions
@@ -99,16 +99,16 @@ def deflections_yx_2d_from(self, grid, max_terms=20, **kwargs):
         # These can also be computed with scipy.special.hyp2f1(), it's
         # much slower can be a useful test
         F1 = self.hyp2f1_series(
-            self.inner_slope, self.axis_ratio, R, z, max_terms=max_terms
+            self.inner_slope, self.axis_ratio(), R, z, max_terms=max_terms
         )
         F2 = self.hyp2f1_series(
-            self.inner_slope, self.axis_ratio, self.break_radius, z, max_terms=max_terms
+            self.inner_slope, self.axis_ratio(), self.break_radius, z, max_terms=max_terms
         )
         F3 = self.hyp2f1_series(
-            self.outer_slope, self.axis_ratio, R, z, max_terms=max_terms
+            self.outer_slope, self.axis_ratio(), R, z, max_terms=max_terms
         )
         F4 = self.hyp2f1_series(
-            self.outer_slope, self.axis_ratio, self.break_radius, z, max_terms=max_terms
+            self.outer_slope, self.axis_ratio(), self.break_radius, z, max_terms=max_terms
         )
 
         # theta < break radius (eq. 18)
diff --git a/autogalaxy/profiles/mass/total/power_law_core.py b/autogalaxy/profiles/mass/total/power_law_core.py
index af2069e17..31d5421a5 100644
--- a/autogalaxy/profiles/mass/total/power_law_core.py
+++ b/autogalaxy/profiles/mass/total/power_law_core.py
@@ -38,13 +38,12 @@ def __init__(
         self.slope = slope
         self.core_radius = core_radius
 
-    @property
-    def einstein_radius_rescaled(self):
+    def einstein_radius_rescaled(self, xp=np):
         """
         Rescale the einstein radius by slope and axis_ratio, to reduce its degeneracy with other mass-profiles
         parameters.
         """
-        return ((3 - self.slope) / (1 + self.axis_ratio)) * self.einstein_radius ** (
+        return ((3 - self.slope) / (1 + self.axis_ratio(xp))) * self.einstein_radius ** (
             self.slope - 1
         )
 
@@ -70,7 +69,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the potential on a grid of (y,x) arc-second coordinates.
 
@@ -91,17 +90,17 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
                 args=(
                     grid.array[i, 0],
                     grid.array[i, 1],
-                    self.axis_ratio,
+                    self.axis_ratio(),
                     self.slope,
                     self.core_radius,
                 ),
             )[0]
 
-        return self.einstein_radius_rescaled * self.axis_ratio * potential_grid
+        return self.einstein_radius_rescaled(xp) * self.axis_ratio() * potential_grid
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
 
@@ -114,9 +113,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         from scipy.integrate import quad
 
         def calculate_deflection_component(npow, index):
-            einstein_radius_rescaled = self.einstein_radius_rescaled
+            einstein_radius_rescaled = self.einstein_radius_rescaled(xp)
 
-            deflection_grid = np.array(self.axis_ratio * grid.array[:, index])
+            deflection_grid = np.array(self.axis_ratio() * grid.array[:, index])
 
             for i in range(grid.shape[0]):
                 deflection_grid[i] *= (
@@ -129,7 +128,7 @@ def calculate_deflection_component(npow, index):
                             grid.array[i, 0],
                             grid.array[i, 1],
                             npow,
-                            self.axis_ratio,
+                            self.axis_ratio(),
                             self.slope,
                             self.core_radius,
                         ),
@@ -145,8 +144,8 @@ def calculate_deflection_component(npow, index):
             grid=np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T)
         )
 
-    def convergence_func(self, grid_radius: float) -> float:
-        return self.einstein_radius_rescaled * (
+    def convergence_func(self, grid_radius: float, xp=np) -> float:
+        return self.einstein_radius_rescaled(xp) * (
             self.core_radius**2 + grid_radius**2
         ) ** (-(self.slope - 1) / 2.0)
 
@@ -172,7 +171,7 @@ def deflection_func(u, y, x, npow, axis_ratio, slope, core_radius):
 
     @property
     def ellipticity_rescale(self):
-        return (1.0 + self.axis_ratio) / 2.0
+        return (1.0 + self.axis_ratio()) / 2.0
 
     @property
     def unit_mass(self):
@@ -211,7 +210,7 @@ def __init__(
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
 
@@ -223,7 +222,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         """
         eta = self.radial_grid_from(grid=grid, **kwargs)
         deflection = jnp.multiply(
-            2.0 * self.einstein_radius_rescaled,
+            2.0 * self.einstein_radius_rescaled(xp),
             jnp.divide(
                 jnp.add(
                     jnp.power(

From 5ac38d02d0844e3af4e3daf99f28dd76de715aa1 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Thu, 6 Nov 2025 20:31:16 +0000
Subject: [PATCH 04/18] for JAX tracing test get through all profiles

---
 autogalaxy/profiles/geometry_profiles.py      | 10 +++---
 autogalaxy/profiles/mass/dark/nfw.py          |  2 +-
 .../profiles/mass/sheets/external_shear.py    |  2 +-
 autogalaxy/profiles/mass/stellar/chameleon.py | 32 +++++++++----------
 autogalaxy/profiles/mass/total/isothermal.py  | 14 ++++----
 autogalaxy/profiles/mass/total/power_law.py   |  4 +--
 6 files changed, 33 insertions(+), 31 deletions(-)

diff --git a/autogalaxy/profiles/geometry_profiles.py b/autogalaxy/profiles/geometry_profiles.py
index c1c7d0247..49c6e9d6f 100644
--- a/autogalaxy/profiles/geometry_profiles.py
+++ b/autogalaxy/profiles/geometry_profiles.py
@@ -231,7 +231,7 @@ def angle_radians(self, xp=np) -> float:
         The position angle in radians of the major-axis of the ellipse defined by profile, defined counter clockwise
         from the positive x-axis (0.0 > angle > 2pi).
         """
-        return xp.radians(self.angle(xp=xp))
+        return xp.radians(self.angle(xp))
 
     @property
     def _cos_angle(self) -> float:
@@ -246,7 +246,7 @@ def _cos_and_sin_to_x_axis(self, xp=np, **kwargs):
         Determine the sin and cosine of the angle between the profile's ellipse and the positive x-axis,
         counter-clockwise.
         """
-        angle_radians = xp.radians(self.angle(xp=xp))
+        angle_radians = xp.radians(self.angle(xp))
         return xp.cos(angle_radians), xp.sin(angle_radians)
 
     def angle_to_profile_grid_from(self, grid_angles, xp=np, **kwargs):
@@ -288,7 +288,7 @@ def rotated_grid_from_reference_frame_from(
         """
 
         if angle is None:
-            angle = self.angle(xp=xp)
+            angle = self.angle(xp)
 
         return aa.util.geometry.transform_grid_2d_from_reference_frame(
             grid_2d=grid, centre=(0.0, 0.0), angle=angle, xp=xp
@@ -352,7 +352,7 @@ def transformed_to_reference_frame_grid_from(
         if self.__class__.__name__.endswith("Sph"):
             return super().transformed_to_reference_frame_grid_from(grid=grid, xp=xp)
         return aa.util.geometry.transform_grid_2d_to_reference_frame(
-            grid_2d=grid.array, centre=self.centre, angle=self.angle(xp=xp), xp=xp
+            grid_2d=grid.array, centre=self.centre, angle=self.angle(xp), xp=xp
         )
 
     @aa.grid_dec.to_grid
@@ -374,7 +374,7 @@ def transformed_from_reference_frame_grid_from(
             return super().transformed_from_reference_frame_grid_from(grid=grid, xp=xp)
 
         return aa.util.geometry.transform_grid_2d_from_reference_frame(
-            grid_2d=grid.array, centre=self.centre, angle=self.angle(xp=xp), xp=xp
+            grid_2d=grid.array, centre=self.centre, angle=self.angle(xp), xp=xp
         )
 
     def _eta_u(self, u, coordinates):
diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index d4a137c8a..137df8872 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -287,7 +287,7 @@ def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         # Rotation for shear
 
         shear_field = self.rotated_grid_from_reference_frame_from(
-            grid=np.vstack((g2, g1)).T, angle=self.angle(xp=xp) * 2
+            grid=np.vstack((g2, g1)).T, angle=self.angle(xp) * 2
         )
 
         return aa.VectorYX2DIrregular(values=shear_field, grid=grid)
diff --git a/autogalaxy/profiles/mass/sheets/external_shear.py b/autogalaxy/profiles/mass/sheets/external_shear.py
index fec2c55e7..213e990fe 100644
--- a/autogalaxy/profiles/mass/sheets/external_shear.py
+++ b/autogalaxy/profiles/mass/sheets/external_shear.py
@@ -43,7 +43,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
     @aa.grid_dec.to_array
     def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         shear_angle = (
-            self.angle(xp=xp) - 90
+            self.angle(xp) - 90
         )  ##to be onsistent with autolens deflection angle calculation
         phig = xp.deg2rad(shear_angle)
         shear_amp = self.magnitude(xp=xp)
diff --git a/autogalaxy/profiles/mass/stellar/chameleon.py b/autogalaxy/profiles/mass/stellar/chameleon.py
index 92040e6b4..b925d7e0e 100644
--- a/autogalaxy/profiles/mass/stellar/chameleon.py
+++ b/autogalaxy/profiles/mass/stellar/chameleon.py
@@ -53,7 +53,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
         Following Eq. (15) and (16), but the parameters are slightly different.
@@ -69,49 +69,49 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs):
             2.0
             * self.mass_to_light_ratio
             * self.intensity
-            / (1 + self.axis_ratio())
-            * self.axis_ratio()
-            / jnp.sqrt(1.0 - self.axis_ratio()**2.0)
+            / (1 + self.axis_ratio(xp))
+            * self.axis_ratio(xp)
+            / jnp.sqrt(1.0 - self.axis_ratio(xp)**2.0)
         )
 
         core_radius_0 = jnp.sqrt(
-            (4.0 * self.core_radius_0**2.0) / (1.0 + self.axis_ratio()) ** 2
+            (4.0 * self.core_radius_0**2.0) / (1.0 + self.axis_ratio(xp)) ** 2
         )
         core_radius_1 = jnp.sqrt(
-            (4.0 * self.core_radius_1**2.0) / (1.0 + self.axis_ratio()) ** 2
+            (4.0 * self.core_radius_1**2.0) / (1.0 + self.axis_ratio(xp)) ** 2
         )
 
         psi0 = psi_from(
-            grid=grid, axis_ratio=self.axis_ratio(), core_radius=core_radius_0
+            grid=grid, axis_ratio=self.axis_ratio(xp), core_radius=core_radius_0, xp=xp
         )
         psi1 = psi_from(
-            grid=grid, axis_ratio=self.axis_ratio(), core_radius=core_radius_1
+            grid=grid, axis_ratio=self.axis_ratio(xp), core_radius=core_radius_1, xp=xp
         )
 
         deflection_y0 = jnp.arctanh(
             jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio()**2.0), grid.array[:, 0]),
-                jnp.add(psi0, self.axis_ratio()**2.0 * core_radius_0),
+                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 0]),
+                jnp.add(psi0, self.axis_ratio(xp)**2.0 * core_radius_0),
             )
         )
 
         deflection_x0 = jnp.arctan(
             jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio()**2.0), grid.array[:, 1]),
+                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 1]),
                 jnp.add(psi0, core_radius_0),
             )
         )
 
         deflection_y1 = jnp.arctanh(
             jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio()**2.0), grid.array[:, 0]),
-                jnp.add(psi1, self.axis_ratio()**2.0 * core_radius_1),
+                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 0]),
+                jnp.add(psi1, self.axis_ratio(xp)**2.0 * core_radius_1),
             )
         )
 
         deflection_x1 = jnp.arctan(
             jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio()**2.0), grid.array[:, 1]),
+                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 1]),
                 jnp.add(psi1, core_radius_1),
             )
         )
@@ -153,10 +153,10 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray):
             The radial distance from the centre of the profile. for each coordinate on the grid.
         """
 
-        axis_ratio_factor = (1.0 + self.axis_ratio()) ** 2.0
+        axis_ratio_factor = (1.0 + self.axis_ratio(xp)) ** 2.0
 
         return jnp.multiply(
-            self.intensity / (1 + self.axis_ratio()),
+            self.intensity / (1 + self.axis_ratio(xp)),
             jnp.add(
                 jnp.divide(
                     1.0,
diff --git a/autogalaxy/profiles/mass/total/isothermal.py b/autogalaxy/profiles/mass/total/isothermal.py
index ffccf87ac..354181e33 100644
--- a/autogalaxy/profiles/mass/total/isothermal.py
+++ b/autogalaxy/profiles/mass/total/isothermal.py
@@ -89,23 +89,24 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
             2.0
             * self.einstein_radius_rescaled(xp)
             * self.axis_ratio(xp)
-            / xp.sqrt(1 - self.axis_ratio()(xp)**2)
+            / xp.sqrt(1 - self.axis_ratio(xp)**2)
         )
 
-        psi = psi_from(grid=grid, axis_ratio=self.axis_ratio()(xp=xp), core_radius=0.0)
+        psi = psi_from(grid=grid, axis_ratio=self.axis_ratio(xp), core_radius=0.0, xp=xp)
 
         deflection_y = xp.arctanh(
             xp.divide(
-                xp.multiply(xp.sqrt(1 - self.axis_ratio()(xp=xp)**2), grid.array[:, 0]), psi
+                xp.multiply(xp.sqrt(1 - self.axis_ratio(xp)**2), grid.array[:, 0]), psi
             )
         )
         deflection_x = xp.arctan(
             xp.divide(
-                xp.multiply(xp.sqrt(1 - self.axis_ratio()(xp=xp)**2), grid.array[:, 1]), psi
+                xp.multiply(xp.sqrt(1 - self.axis_ratio(xp)**2), grid.array[:, 1]), psi
             )
         )
         return self.rotated_grid_from_reference_frame_from(
             grid=xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T),
+            xp=xp,
             **kwargs,
         )
 
@@ -144,7 +145,7 @@ def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         )
 
         shear_field = self.rotated_grid_from_reference_frame_from(
-            grid=xp.vstack((gamma_2, gamma_1)).T, angle=self.angle(xp=xp) * 2
+            grid=xp.vstack((gamma_2, gamma_1)).T, xp=xp, angle=self.angle(xp) * 2
         )
 
         return aa.VectorYX2DIrregular(values=shear_field, grid=grid)
@@ -184,7 +185,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         grid
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
-        eta = self.elliptical_radii_grid_from(grid=grid, **kwargs)
+        eta = self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs)
         return 2.0 * self.einstein_radius_rescaled(xp) * eta
 
     @aa.grid_dec.to_vector_yx
@@ -201,5 +202,6 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return self._cartesian_grid_via_radial_from(
             grid=grid,
             radius=xp.full(grid.shape[0], 2.0 * self.einstein_radius_rescaled(xp)),
+            xp=xp,
             **kwargs,
         )
diff --git a/autogalaxy/profiles/mass/total/power_law.py b/autogalaxy/profiles/mass/total/power_law.py
index 2a06fe154..c21175958 100644
--- a/autogalaxy/profiles/mass/total/power_law.py
+++ b/autogalaxy/profiles/mass/total/power_law.py
@@ -147,7 +147,7 @@ def __init__(
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
     def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
-        eta = self.radial_grid_from(grid=grid, xp=np, **kwargs).array
+        eta = self.radial_grid_from(grid=grid, xp=xp, **kwargs).array
         deflection_r = (
             2.0
             * self.einstein_radius_rescaled(xp)
@@ -157,4 +157,4 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
             )
         )
 
-        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_r, xp=np)
+        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_r, xp=xp)

From 3d99eb3514e74a122b67723b299ff5fe49ef84bf Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Thu, 6 Nov 2025 20:34:13 +0000
Subject: [PATCH 05/18] namespace passed through to_inversion

---
 autogalaxy/galaxy/to_inversion.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/autogalaxy/galaxy/to_inversion.py b/autogalaxy/galaxy/to_inversion.py
index 75087c53c..795853c18 100644
--- a/autogalaxy/galaxy/to_inversion.py
+++ b/autogalaxy/galaxy/to_inversion.py
@@ -573,6 +573,7 @@ def inversion(self) -> aa.AbstractInversion:
             dataset=self.dataset,
             linear_obj_list=self.linear_obj_list,
             settings=self.settings_inversion,
+            xp=self.xp
         )
 
         inversion.linear_obj_galaxy_dict = self.linear_obj_galaxy_dict

From 8a2e5309df4ff4775a616728738af8173af5f09e Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 8 Nov 2025 14:28:21 +0000
Subject: [PATCH 06/18] light profiles asll pas wihtout import of jnp

---
 .../analysis/adapt_images/adapt_images.py     |   2 +-
 autogalaxy/galaxy/galaxies.py                 |   8 +-
 autogalaxy/galaxy/galaxy.py                   |   2 +-
 autogalaxy/galaxy/to_inversion.py             |   3 +
 autogalaxy/operate/deflections.py             |  68 ++++-----
 autogalaxy/operate/image.py                   |   9 +-
 autogalaxy/profiles/basis.py                  |  25 ++--
 autogalaxy/profiles/light/abstract.py         |   8 +-
 autogalaxy/profiles/light/decorators.py       |  11 +-
 autogalaxy/profiles/light/linear/abstract.py  |  45 ++++--
 .../profiles/light/mock/mock_light_profile.py |   2 +-
 autogalaxy/profiles/light/snr/abstract.py     |   2 +-
 .../profiles/light/standard/chameleon.py      |  24 +--
 autogalaxy/profiles/light/standard/eff.py     |   6 +-
 .../profiles/light/standard/gaussian.py       |  15 +-
 autogalaxy/profiles/light/standard/moffat.py  |  14 +-
 autogalaxy/profiles/light/standard/sersic.py  |   4 +-
 .../profiles/light/standard/sersic_core.py    |  33 ++---
 .../light/standard/shapelets/cartesian.py     |   8 +-
 .../light/standard/shapelets/exponential.py   |  26 ++--
 .../light/standard/shapelets/polar.py         |  22 +--
 autogalaxy/profiles/mass/abstract/abstract.py |  12 +-
 autogalaxy/profiles/mass/abstract/cse.py      |   4 +-
 .../profiles/mass/abstract/jax_utils.py       |  24 +--
 autogalaxy/profiles/mass/abstract/mge_jax.py  |  46 +++---
 autogalaxy/profiles/mass/dark/abstract.py     |  40 ++---
 autogalaxy/profiles/mass/dark/gnfw.py         |   6 +-
 autogalaxy/profiles/mass/dark/nfw.py          |  22 +--
 .../profiles/mass/dark/nfw_truncated.py       |  64 ++++----
 autogalaxy/profiles/mass/point/point.py       |   6 +-
 autogalaxy/profiles/mass/point/smbh_binary.py |   6 +-
 autogalaxy/profiles/mass/sheets/mass_sheet.py |  10 +-
 autogalaxy/profiles/mass/stellar/chameleon.py |  74 ++++-----
 autogalaxy/profiles/mass/stellar/gaussian.py  |  14 +-
 autogalaxy/profiles/mass/stellar/sersic.py    |  16 +-
 .../profiles/mass/stellar/sersic_core.py      |  37 +++--
 .../profiles/mass/stellar/sersic_gradient.py  |   6 +-
 .../mass/total/dual_pseudo_isothermal_mass.py | 140 +++++++++---------
 .../total/dual_pseudo_isothermal_potential.py |  60 ++++----
 autogalaxy/profiles/mass/total/jax_utils.py   |  19 +--
 autogalaxy/profiles/mass/total/power_law.py   |   4 +-
 .../profiles/mass/total/power_law_broken.py   |  20 +--
 .../profiles/mass/total/power_law_core.py     |  14 +-
 .../mass/total/power_law_multipole.py         |  24 +--
 .../profiles/light/standard/test_abstract.py  |  34 -----
 .../profiles/light/test_decorators.py         |   4 +-
 46 files changed, 510 insertions(+), 533 deletions(-)

diff --git a/autogalaxy/analysis/adapt_images/adapt_images.py b/autogalaxy/analysis/adapt_images/adapt_images.py
index c3cf8a80c..1c65e1de7 100644
--- a/autogalaxy/analysis/adapt_images/adapt_images.py
+++ b/autogalaxy/analysis/adapt_images/adapt_images.py
@@ -132,7 +132,7 @@ def from_result(cls, result, use_model_images: bool = False) -> "AdaptImages":
             else:
                 galaxy_image = result.subtracted_signal_to_noise_map_galaxy_dict[path]
 
-            minimum_galaxy_value = adapt_minimum_percent * jnp.max(galaxy_image.array)
+            minimum_galaxy_value = adapt_minimum_percent * xp.max(galaxy_image.array)
             galaxy_image[galaxy_image < minimum_galaxy_value] = minimum_galaxy_value
 
             galaxy_name_image_dict[path] = galaxy_image
diff --git a/autogalaxy/galaxy/galaxies.py b/autogalaxy/galaxy/galaxies.py
index ce71e1dfe..6f284c29a 100644
--- a/autogalaxy/galaxy/galaxies.py
+++ b/autogalaxy/galaxy/galaxies.py
@@ -182,7 +182,7 @@ def traced_grid_2d_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.type.Grid2D
         return grid - self.deflections_yx_2d_from(grid=grid, xp=xp)
 
     @aa.grid_dec.to_array
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
         """
         Returns the summed 2D convergence of all galaxies from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -203,10 +203,10 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
         grid
             The 2D (y, x) coordinates where values of the convergence are evaluated.
         """
-        return sum(map(lambda g: g.convergence_2d_from(grid=grid), self))
+        return sum(map(lambda g: g.convergence_2d_from(grid=grid, xp=xp), self))
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
         """
         Returns the summed 2D potential of all galaxies from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -227,7 +227,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
         grid
             The 2D (y, x) coordinates where values of the potential are evaluated.
         """
-        return sum(map(lambda g: g.potential_2d_from(grid=grid), self))
+        return sum(map(lambda g: g.potential_2d_from(grid=grid, xp=xp), self))
 
     def has(self, cls: Union[Type, Tuple[Type]]) -> bool:
         """
diff --git a/autogalaxy/galaxy/galaxy.py b/autogalaxy/galaxy/galaxy.py
index a46f9a5e6..4e9ef2014 100644
--- a/autogalaxy/galaxy/galaxy.py
+++ b/autogalaxy/galaxy/galaxy.py
@@ -314,7 +314,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.n
         if self.has(cls=MassProfile):
             return sum(
                 map(
-                    lambda p: p.convergence_2d_from(grid=grid),
+                    lambda p: p.convergence_2d_from(grid=grid, xp=xp),
                     self.cls_list_from(cls=MassProfile),
                 )
             )
diff --git a/autogalaxy/galaxy/to_inversion.py b/autogalaxy/galaxy/to_inversion.py
index 795853c18..dd5e9d8ef 100644
--- a/autogalaxy/galaxy/to_inversion.py
+++ b/autogalaxy/galaxy/to_inversion.py
@@ -309,6 +309,7 @@ def cls_light_profile_func_list_galaxy_dict_from(
                             psf=self.dataset.psf,
                             light_profile_list=light_profile_list,
                             regularization=light_profile.regularization,
+                            xp=self.xp
                         )
 
                         lp_linear_func_galaxy_dict[lp_linear_func] = galaxy
@@ -487,11 +488,13 @@ def mapper_from(
             source_plane_mesh_grid=source_plane_mesh_grid,
             image_plane_mesh_grid=image_plane_mesh_grid,
             adapt_data=adapt_galaxy_image,
+            xp=self.xp
         )
 
         return mapper_from(
             mapper_grids=mapper_grids,
             regularization=regularization,
+            xp=self.xp
         )
 
     @cached_property
diff --git a/autogalaxy/operate/deflections.py b/autogalaxy/operate/deflections.py
index 6c7ef3af6..9f07a9720 100644
--- a/autogalaxy/operate/deflections.py
+++ b/autogalaxy/operate/deflections.py
@@ -47,7 +47,7 @@ def wrapper(lensing_obj, grid, jacobian=None):
 
 
 def one_step(r, _, theta, fun, fun_dr):
-    r = jnp.abs(r - fun(r, theta) / fun_dr(r, theta))
+    r = xp.abs(r - fun(r, theta) / fun_dr(r, theta))
     return r, None
 
 
@@ -56,8 +56,8 @@ def step_r(r, theta, fun, fun_dr, N=20):
     one_step_partial = jax.tree_util.Partial(
         one_step, theta=theta, fun=fun, fun_dr=fun_dr
     )
-    new_r = jax.lax.scan(one_step_partial, r, xs=jnp.arange(N))[0]
-    return jnp.stack([new_r * jnp.sin(theta), new_r * jnp.cos(theta)]).T
+    new_r = jax.lax.scan(one_step_partial, r, xs=xp.arange(N))[0]
+    return xp.stack([new_r * xp.sin(theta), new_r * xp.cos(theta)]).T
 
 
 class OperateDeflections:
@@ -90,7 +90,7 @@ def deflections_yx_scalar(self, y, x, pixel_scales):
         )
 
         g = aa.Grid2D(
-            values=jnp.stack((y.reshape(1), x.reshape(1)), axis=-1), mask=mask
+            values=xp.stack((y.reshape(1), x.reshape(1)), axis=-1), mask=mask
         )
 
         return self.deflections_yx_2d_from(g).squeeze()
@@ -281,19 +281,19 @@ def hessian_from(self, grid, buffer: float = 0.01, deflections_func=None) -> Tup
             grid = grid.array
 
         grid_shift_y_up = aa.Grid2DIrregular(
-            values=jnp.stack([grid[:, 0] + buffer, grid[:, 1]], axis=1)
+            values=xp.stack([grid[:, 0] + buffer, grid[:, 1]], axis=1)
         )
 
         grid_shift_y_down = aa.Grid2DIrregular(
-            values=jnp.stack([grid[:, 0] - buffer, grid[:, 1]], axis=1)
+            values=xp.stack([grid[:, 0] - buffer, grid[:, 1]], axis=1)
         )
 
         grid_shift_x_left = aa.Grid2DIrregular(
-            values=jnp.stack([grid[:, 0], grid[:, 1] - buffer], axis=1)
+            values=xp.stack([grid[:, 0], grid[:, 1] - buffer], axis=1)
         )
 
         grid_shift_x_right = aa.Grid2DIrregular(
-            values=jnp.stack([grid[:, 0], grid[:, 1] + buffer], axis=1)
+            values=xp.stack([grid[:, 0], grid[:, 1] + buffer], axis=1)
         )
 
         deflections_up = deflections_func(grid=grid_shift_y_up)
@@ -375,7 +375,7 @@ def shear_yx_2d_via_hessian_from(
         gamma_1 = 0.5 * (hessian_xx - hessian_yy)
         gamma_2 = hessian_xy
 
-        shear_yx_2d = jnp.stack([gamma_2.array, gamma_1.array], axis=1)
+        shear_yx_2d = xp.stack([gamma_2.array, gamma_1.array], axis=1)
 
         return ShearYX2DIrregular(values=shear_yx_2d, grid=grid)
 
@@ -614,7 +614,7 @@ def area_within_curve_list_from(
 
         for curve in curve_list:
             x, y = curve[:, 0], curve[:, 1]
-            area = jnp.abs(0.5 * jnp.sum(y[:-1] * jnp.diff(x) - x[:-1] * jnp.diff(y)))
+            area = xp.abs(0.5 * xp.sum(y[:-1] * xp.diff(x) - x[:-1] * xp.diff(y)))
             area_within_each_curve_list.append(area)
 
         return area_within_each_curve_list
@@ -649,7 +649,7 @@ def einstein_radius_list_from(
         """
         try:
             area_list = self.tangential_critical_curve_area_list_from(grid=grid)
-            return [jnp.sqrt(area / jnp.pi) for area in area_list]
+            return [xp.sqrt(area / xp.pi) for area in area_list]
         except TypeError:
             raise TypeError("The grid input was unable to estimate the Einstein Radius")
 
@@ -729,7 +729,7 @@ def einstein_mass_angular_list_from(
             caustic to be computed more accurately using a higher resolution grid.
         """
         einstein_radius_list = self.einstein_radius_list_from(grid=grid)
-        return [jnp.pi * einstein_radius**2 for einstein_radius in einstein_radius_list]
+        return [xp.pi * einstein_radius**2 for einstein_radius in einstein_radius_list]
 
     def einstein_mass_angular_from(
         self,
@@ -775,12 +775,12 @@ def einstein_mass_angular_from(
         return einstein_mass_angular_list[0]
 
     def jacobian_stack(self, y, x, pixel_scales):
-        return jnp.stack(
+        return xp.stack(
             jax.jacfwd(self.deflections_yx_scalar, argnums=(0, 1))(y, x, pixel_scales)
         )
 
     def jacobian_stack_vector(self, y, x, pixel_scales):
-        return jnp.vectorize(
+        return xp.vectorize(
             jax.tree_util.Partial(self.jacobian_stack, pixel_scales=pixel_scales),
             signature="(),()->(i,i)",
         )(y, x)
@@ -788,7 +788,7 @@ def jacobian_stack_vector(self, y, x, pixel_scales):
     def convergence_mag_shear_yx(self, y, x):
         J = self.jacobian_stack_vector(y, x, 0.05)
         K = 0.5 * (J[..., 0, 0] + J[..., 1, 1])
-        mag_shear = 0.5 * jnp.sqrt(
+        mag_shear = 0.5 * xp.sqrt(
             (J[..., 0, 1] + J[..., 1, 0]) ** 2 + (J[..., 0, 0] - J[..., 1, 1]) ** 2
         )
         return K, mag_shear
@@ -800,15 +800,15 @@ def tangential_eigen_value_yx(self, y, x):
 
     @partial(jit, static_argnums=(0, 3))
     def tangential_eigen_value_rt(self, r, theta, centre=(0.0, 0.0)):
-        y = r * jnp.sin(theta) + centre[0]
-        x = r * jnp.cos(theta) + centre[1]
+        y = r * xp.sin(theta) + centre[0]
+        x = r * xp.cos(theta) + centre[1]
         return self.tangential_eigen_value_yx(y, x)
 
     @partial(jit, static_argnums=(0, 3))
     def grad_r_tangential_eigen_value(self, r, theta, centre=(0.0, 0.0)):
         # ignore `self` with the `argnums` below
         tangential_eigen_part = partial(self.tangential_eigen_value_rt, centre=centre)
-        return jnp.vectorize(
+        return xp.vectorize(
             jax.jacfwd(tangential_eigen_part, argnums=(0,)), signature="(),()->()"
         )(r, theta)[0]
 
@@ -819,15 +819,15 @@ def radial_eigen_value_yx(self, y, x):
 
     @partial(jit, static_argnums=(0, 3))
     def radial_eigen_value_rt(self, r, theta, centre=(0.0, 0.0)):
-        y = r * jnp.sin(theta) + centre[0]
-        x = r * jnp.cos(theta) + centre[1]
+        y = r * xp.sin(theta) + centre[0]
+        x = r * xp.cos(theta) + centre[1]
         return self.radial_eigen_value_yx(y, x)
 
     @partial(jit, static_argnums=(0, 3))
     def grad_r_radial_eigen_value(self, r, theta, centre=(0.0, 0.0)):
         # ignore `self` with the `argnums` below
         radial_eigen_part = partial(self.radial_eigen_value_rt, centre=centre)
-        return jnp.vectorize(
+        return xp.vectorize(
             jax.jacfwd(radial_eigen_part, argnums=(0,)), signature="(),()->()"
         )(r, theta)[0]
 
@@ -864,8 +864,8 @@ def tangential_critical_curve_jax(
         threshold : float
             Only keep points whose tangential eigen value is within this value of zero (inclusive)
         """
-        r = jnp.ones(n_points) * init_r
-        theta = jnp.linspace(0, 2 * jnp.pi, n_points + 1)[:-1]
+        r = xp.ones(n_points) * init_r
+        theta = xp.linspace(0, 2 * xp.pi, n_points + 1)[:-1]
         new_yx = step_r(
             r,
             theta,
@@ -875,12 +875,12 @@ def tangential_critical_curve_jax(
             ),
             n_steps,
         )
-        new_yx = new_yx + jnp.array(init_centre)
+        new_yx = new_yx + xp.array(init_centre)
         # filter out nan values
-        fdx = jnp.isfinite(new_yx).all(axis=1)
+        fdx = xp.isfinite(new_yx).all(axis=1)
         new_yx = new_yx[fdx]
         # filter out failed points
-        value = jnp.abs(self.tangential_eigen_value_yx(new_yx[:, 0], new_yx[:, 1]))
+        value = xp.abs(self.tangential_eigen_value_yx(new_yx[:, 0], new_yx[:, 1]))
         gdx = value <= threshold
         return aa.structures.grids.irregular_2d.Grid2DIrregular(values=new_yx[gdx])
 
@@ -917,8 +917,8 @@ def radial_critical_curve_jax(
         threshold : float
             Only keep points whose radial eigen value is within this value of zero (inclusive)
         """
-        r = jnp.ones(n_points) * init_r
-        theta = jnp.linspace(0, 2 * jnp.pi, n_points + 1)[:-1]
+        r = xp.ones(n_points) * init_r
+        theta = xp.linspace(0, 2 * xp.pi, n_points + 1)[:-1]
         new_yx = step_r(
             r,
             theta,
@@ -926,12 +926,12 @@ def radial_critical_curve_jax(
             jax.tree_util.Partial(self.grad_r_radial_eigen_value, centre=init_centre),
             n_steps,
         )
-        new_yx = new_yx + jnp.array(init_centre)
+        new_yx = new_yx + xp.array(init_centre)
         # filter out nan values
-        fdx = jnp.isfinite(new_yx).all(axis=1)
+        fdx = xp.isfinite(new_yx).all(axis=1)
         new_yx = new_yx[fdx]
         # filter out failed points
-        value = jnp.abs(self.radial_eigen_value_yx(new_yx[:, 0], new_yx[:, 1]))
+        value = xp.abs(self.radial_eigen_value_yx(new_yx[:, 0], new_yx[:, 1]))
         gdx = value <= threshold
         return aa.structures.grids.irregular_2d.Grid2DIrregular(values=new_yx[gdx])
 
@@ -955,7 +955,7 @@ def jacobian_from(self, grid):
         A = self.jacobian_stack_vector(
             grid.array[:, 0], grid.array[:, 1], grid.pixel_scales
         )
-        a = jnp.eye(2).reshape(1, 2, 2) - A
+        a = xp.eye(2).reshape(1, 2, 2) - A
         return [
             [
                 aa.Array2D(values=a[..., 1, 1], mask=grid.mask),
@@ -969,7 +969,7 @@ def jacobian_from(self, grid):
 
         # transpose the result
         # use `moveaxis` as grid might not be nx2
-        # return jnp.moveaxis(jnp.moveaxis(a, -1, 0), -1, 0)
+        # return xp.moveaxis(xp.moveaxis(a, -1, 0), -1, 0)
 
     @precompute_jacobian
     def convergence_2d_via_jacobian_from(self, grid, jacobian=None) -> aa.Array2D:
@@ -1020,7 +1020,7 @@ def shear_yx_2d_via_jacobian_from(
         """
         shear_y = -0.5 * (jacobian[0][1] + jacobian[1][0]).array
         shear_x = 0.5 * (jacobian[1][1] - jacobian[0][0]).array
-        shear_yx_2d = jnp.stack([shear_y, shear_x]).T
+        shear_yx_2d = xp.stack([shear_y, shear_x]).T
 
         if isinstance(grid, aa.Grid2DIrregular):
             return ShearYX2DIrregular(values=shear_yx_2d, grid=grid)
diff --git a/autogalaxy/operate/image.py b/autogalaxy/operate/image.py
index 0f3ff45a3..fafeaf0f0 100644
--- a/autogalaxy/operate/image.py
+++ b/autogalaxy/operate/image.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
-import jax
-import jax.numpy as jnp
+import numpy as np
 from typing import TYPE_CHECKING, Dict, List, Optional
 
 from autoarray import Array2D
@@ -23,7 +22,7 @@ class OperateImage:
     """
 
     def image_2d_from(
-        self, grid: aa.Grid2D, operated_only: Optional[bool] = None
+        self, grid: aa.Grid2D, xp=np, operated_only: Optional[bool] = None
     ) -> aa.Array2D:
         raise NotImplementedError
 
@@ -334,7 +333,7 @@ def visibilities_list_from(
         visibilities_list = []
 
         for image_2d in image_2d_list:
-            if not jnp.any(image_2d.array):
+            if not xp.any(image_2d.array):
                 visibilities = aa.Visibilities.zeros(
                     shape_slim=(transformer.uv_wavelengths.shape[0],)
                 )
@@ -453,7 +452,7 @@ def galaxy_visibilities_dict_from(
         for galaxy_key in galaxy_image_2d_dict.keys():
             image_2d = galaxy_image_2d_dict[galaxy_key]
 
-            if not jnp.any(image_2d.array):
+            if not xp.any(image_2d.array):
                 visibilities = aa.Visibilities.zeros(
                     shape_slim=(transformer.uv_wavelengths.shape[0],)
                 )
diff --git a/autogalaxy/profiles/basis.py b/autogalaxy/profiles/basis.py
index 3fc08b3ee..ecb347327 100644
--- a/autogalaxy/profiles/basis.py
+++ b/autogalaxy/profiles/basis.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Dict, List, Optional, Union
 
@@ -76,7 +75,7 @@ def mass_profile_list(self) -> List[MassProfile]:
         return aa.util.misc.cls_list_from(values=self.profile_list, cls=MassProfile)
 
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
     ) -> aa.Array2D:
         """
         Returns the summed image of all light profiles in the basis from a 2D grid of Cartesian (y,x) coordinates.
@@ -99,10 +98,10 @@ def image_2d_from(
         -------
         The image of the light profiles in the basis summed together.
         """
-        return sum(self.image_2d_list_from(grid=grid, operated_only=operated_only))
+        return sum(self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only))
 
     def image_2d_list_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
     ) -> List[aa.Array2D]:
         """
         Returns each image of each light profiles in the basis as a list, from a 2D grid of Cartesian (y,x) coordinates.
@@ -127,14 +126,14 @@ def image_2d_list_from(
         """
         return [
             (
-                light_profile.image_2d_from(grid=grid, operated_only=operated_only)
+                light_profile.image_2d_from(grid=grid, xp=xp, operated_only=operated_only)
                 if not isinstance(light_profile, lp_linear.LightProfileLinear)
-                else jnp.zeros((grid.shape[0],))
+                else xp.zeros((grid.shape[0],))
             )
             for light_profile in self.light_profile_list
         ]
 
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array2D:
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> aa.Array2D:
         """
         Returns the summed convergence of all mass profiles in the basis from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -153,11 +152,11 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array2D:
         """
         if len(self.mass_profile_list) > 0:
             return sum(
-                [mass.convergence_2d_from(grid=grid) for mass in self.profile_list]
+                [mass.convergence_2d_from(grid=grid, xp=xp) for mass in self.profile_list]
             )
-        return jnp.zeros((grid.shape[0],))
+        return xp.zeros((grid.shape[0],))
 
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array2D:
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> aa.Array2D:
         """
         Returns the summed potential of all mass profiles in the basis from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -178,9 +177,9 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array2D:
             return sum(
                 [mass.potential_2d_from(grid=grid) for mass in self.profile_list]
             )
-        return jnp.zeros((grid.shape[0],))
+        return xp.zeros((grid.shape[0],))
 
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array2D:
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> aa.Array2D:
         """
         Returns the summed deflections of all mass profiles in the basis from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -201,7 +200,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array
             return sum(
                 [mass.deflections_yx_2d_from(grid=grid) for mass in self.profile_list]
             )
-        return jnp.zeros((grid.shape[0], 2))
+        return xp.zeros((grid.shape[0], 2))
 
     def lp_instance_from(self, linear_light_profile_intensity_dict: Dict):
         light_profile_list = []
diff --git a/autogalaxy/profiles/light/abstract.py b/autogalaxy/profiles/light/abstract.py
index da88cbcc4..aa222f273 100644
--- a/autogalaxy/profiles/light/abstract.py
+++ b/autogalaxy/profiles/light/abstract.py
@@ -42,7 +42,7 @@ def coefficient_tag(self) -> str:
         return ""
 
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
     ) -> aa.Array2D:
         """
         Returns the light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates, which may have been
@@ -63,7 +63,7 @@ def image_2d_from(
         """
         raise NotImplementedError()
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
+    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
         """
         Returns the light profile's 2D image from a 1D grid of coordinates which are the radial distance of each
         coordinate from the light profile `centre`.
@@ -77,7 +77,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
 
     @aa.grid_dec.project_grid
     def image_1d_from(
-        self, grid: aa.type.Grid1D2DLike, **kwargs
+        self, grid: aa.type.Grid1D2DLike, xp=np, **kwargs
     ) -> aa.type.Grid1D2DLike:
         """
         Returns the light profile's 1D image from a grid of Cartesian coordinates, which may have been
@@ -100,7 +100,7 @@ def image_1d_from(
         image
             The 1D image of the light profile evaluated at every (x,) coordinate on the 1D transformed grid.
         """
-        return self.image_2d_from(grid=grid, **kwargs)
+        return self.image_2d_from(grid=grid, xp=xp, **kwargs)
 
     def luminosity_within_circle_from(self, radius: float) -> float:
         """
diff --git a/autogalaxy/profiles/light/decorators.py b/autogalaxy/profiles/light/decorators.py
index 9d4c134b6..909fc507f 100644
--- a/autogalaxy/profiles/light/decorators.py
+++ b/autogalaxy/profiles/light/decorators.py
@@ -27,6 +27,7 @@ def check_operated_only(func):
     def wrapper(
         obj,
         grid: aa.type.Grid1D2DLike,
+        xp=np,
         operated_only: Optional[bool] = None,
         *args,
         **kwargs,
@@ -60,13 +61,13 @@ def wrapper(
         )
 
         if operated_only is None:
-            return func(obj, grid, operated_only, *args, **kwargs)
+            return func(obj, grid, xp, operated_only, *args, **kwargs)
         elif operated_only:
             if isinstance(obj, LightProfileOperated):
-                return func(obj, grid, operated_only, *args, **kwargs)
-            return np.zeros((grid.shape[0],))
+                return func(obj, grid, xp, operated_only, *args, **kwargs)
+            return xp.zeros((grid.shape[0],))
         if not isinstance(obj, LightProfileOperated):
-            return func(obj, grid, operated_only, *args, **kwargs)
-        return np.zeros((grid.shape[0],))
+            return func(obj, grid, xp, operated_only, *args, **kwargs)
+        return xp.zeros((grid.shape[0],))
 
     return wrapper
diff --git a/autogalaxy/profiles/light/linear/abstract.py b/autogalaxy/profiles/light/linear/abstract.py
index 9663c7810..3267a3d63 100644
--- a/autogalaxy/profiles/light/linear/abstract.py
+++ b/autogalaxy/profiles/light/linear/abstract.py
@@ -146,6 +146,7 @@ def __init__(
         psf: Optional[aa.Kernel2D],
         light_profile_list: List[LightProfileLinear],
         regularization=Optional[aa.reg.Regularization],
+        xp=np,
     ):
         """
         A list of linear light profiles which fits a dataset via linear algebra using the images of each linear light
@@ -205,6 +206,7 @@ def __init__(
         super().__init__(
             grid=grid,
             regularization=regularization,
+            xp=xp
         )
 
         self.blurring_grid = blurring_grid
@@ -267,11 +269,11 @@ def mapping_matrix(self) -> np.ndarray:
 
         for pixel, light_profile in enumerate(self.light_profile_list):
 
-            image_2d = light_profile.image_2d_from(grid=self.grid).slim
+            image_2d = light_profile.image_2d_from(grid=self.grid, xp=self.xp).slim
 
             image_2d_list.append(image_2d.array)
 
-        return jnp.stack(image_2d_list, axis=1)
+        return self.xp.stack(image_2d_list, axis=1)
 
     @cached_property
     def operated_mapping_matrix_overrideg(self) -> Optional[np.ndarray]:
@@ -300,26 +302,37 @@ def operated_mapping_matrix_overrideg(self) -> Optional[np.ndarray]:
         n_src = len(self.light_profile_list)
 
         # allocate slim-form arrays for mapping matrices
-        mapping_matrix = jnp.zeros((self.grid.shape_slim, n_src))
-        blurring_mapping_matrix = jnp.zeros((self.blurring_grid.shape_slim, n_src))
+        mapping_matrix = self.xp.zeros((self.grid.shape_slim, n_src))
+        blurring_mapping_matrix = self.xp.zeros((self.blurring_grid.shape_slim, n_src))
 
         # build each column
         for pixel, light_profile in enumerate(self.light_profile_list):
-            # main grid mapping for this light profile
-            mapping_matrix = mapping_matrix.at[:, pixel].set(
-                light_profile.image_2d_from(grid=self.grid).array
-            )
+            if self.xp.__name__.startswith("jax"):
+                # main grid mapping for this light profile
+                mapping_matrix = mapping_matrix.at[:, pixel].set(
+                    light_profile.image_2d_from(grid=self.grid, xp=self.xp).array
+                )
 
-            # blurring grid mapping for this light profile
-            blurring_mapping_matrix = blurring_mapping_matrix.at[:, pixel].set(
-                light_profile.image_2d_from(grid=self.blurring_grid).array
-            )
+                # blurring grid mapping for this light profile
+                blurring_mapping_matrix = blurring_mapping_matrix.at[:, pixel].set(
+                    light_profile.image_2d_from(grid=self.blurring_grid, xp=self.xp).array
+                )
+
+            else:
+
+                mapping_matrix[:, pixel] = light_profile.image_2d_from(
+                    grid=self.grid, xp=self.xp
+                ).array
+                blurring_mapping_matrix[:, pixel] = light_profile.image_2d_from(
+                    grid=self.blurring_grid, xp=self.xp
+                ).array
 
         return self.psf.convolved_mapping_matrix_from(
             mapping_matrix=mapping_matrix,
             mask=self.grid.mask,
             blurring_mapping_matrix=blurring_mapping_matrix,
             blurring_mask=self.blurring_grid.mask,
+            xp=self.xp
         )
 
     @cached_property
@@ -348,14 +361,14 @@ def operated_mapping_matrix_override(self) -> Optional[np.ndarray]:
         blurred_image_2d_list = []
 
         for pixel, light_profile in enumerate(self.light_profile_list):
-            image_2d = light_profile.image_2d_from(grid=self.grid)
+            image_2d = light_profile.image_2d_from(grid=self.grid, xp=self.xp)
 
-            blurring_image_2d = light_profile.image_2d_from(grid=self.blurring_grid)
+            blurring_image_2d = light_profile.image_2d_from(grid=self.blurring_grid, xp=self.xp)
 
             blurred_image_2d = self.psf.convolved_image_from(
-                image=image_2d, blurring_image=blurring_image_2d
+                image=image_2d, blurring_image=blurring_image_2d, xp=self.xp
             )
 
             blurred_image_2d_list.append(blurred_image_2d.array)
 
-        return jnp.stack(blurred_image_2d_list, axis=1)
+        return self.xp.stack(blurred_image_2d_list, axis=1)
diff --git a/autogalaxy/profiles/light/mock/mock_light_profile.py b/autogalaxy/profiles/light/mock/mock_light_profile.py
index 0b7ef9aaa..54b3b08d9 100644
--- a/autogalaxy/profiles/light/mock/mock_light_profile.py
+++ b/autogalaxy/profiles/light/mock/mock_light_profile.py
@@ -29,7 +29,7 @@ def __init__(
 
     @aa.grid_dec.to_array
     @check_operated_only
-    def image_2d_from(self, grid, operated_only: Optional[bool] = None, **kwargs):
+    def image_2d_from(self, grid, xp=np, operated_only: Optional[bool] = None, **kwargs):
         if self.image_2d is not None:
             return self.image_2d
 
diff --git a/autogalaxy/profiles/light/snr/abstract.py b/autogalaxy/profiles/light/snr/abstract.py
index 59e12b726..3c92e32d4 100644
--- a/autogalaxy/profiles/light/snr/abstract.py
+++ b/autogalaxy/profiles/light/snr/abstract.py
@@ -28,7 +28,7 @@ def __init__(self, signal_to_noise_ratio: float = 10.0):
         self.signal_to_noise_ratio = signal_to_noise_ratio
 
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
     ) -> aa.Array2D:
         """
         Abstract method for obtaining intensity at a grid of Cartesian (y,x) coordinates.
diff --git a/autogalaxy/profiles/light/standard/chameleon.py b/autogalaxy/profiles/light/standard/chameleon.py
index 76016bd0f..90296a277 100644
--- a/autogalaxy/profiles/light/standard/chameleon.py
+++ b/autogalaxy/profiles/light/standard/chameleon.py
@@ -53,7 +53,7 @@ def axis_ratio(self, xp=np) -> float:
         axis_ratio = super().axis_ratio(xp=xp)
         return axis_ratio if axis_ratio < 0.99999 else 0.99999
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
+    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
         """
         Returns the 2D image of the Sersic light profile from a grid of coordinates which are the radial distances of
         each coordinate from the its `centre`.
@@ -66,23 +66,23 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
 
         axis_ratio_factor = (1.0 + self.axis_ratio()) ** 2.0
 
-        return jnp.multiply(
+        return xp.multiply(
             self._intensity / (1 + self.axis_ratio()),
-            jnp.add(
-                jnp.divide(
+            xp.add(
+                xp.divide(
                     1.0,
-                    jnp.sqrt(
-                        jnp.add(
-                            jnp.square(grid_radii.array),
+                    xp.sqrt(
+                        xp.add(
+                            xp.square(grid_radii.array),
                             (4.0 * self.core_radius_0**2.0) / axis_ratio_factor,
                         )
                     ),
                 ),
-                -jnp.divide(
+                -xp.divide(
                     1.0,
-                    jnp.sqrt(
-                        jnp.add(
-                            jnp.square(grid_radii.array),
+                    xp.sqrt(
+                        xp.add(
+                            xp.square(grid_radii.array),
                             (4.0 * self.core_radius_1**2.0) / axis_ratio_factor,
                         )
                     ),
@@ -95,7 +95,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
     ) -> np.ndarray:
         """
         Returns the Chameleon light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
diff --git a/autogalaxy/profiles/light/standard/eff.py b/autogalaxy/profiles/light/standard/eff.py
index 675834995..bd8d2d76a 100644
--- a/autogalaxy/profiles/light/standard/eff.py
+++ b/autogalaxy/profiles/light/standard/eff.py
@@ -42,7 +42,7 @@ def __init__(
         self.effective_radius = effective_radius
         self.eta = eta
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
+    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
         """
         Returns the 2D image of the Sersic light profile from a grid of coordinates which are the radial distances of
         each coordinate from the its `centre`.
@@ -62,7 +62,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
     ) -> np.ndarray:
         """
         Returns the Eff light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
@@ -81,7 +81,7 @@ def image_2d_from(
             The image of the Eff evaluated at every (y,x) coordinate on the transformed grid.
         """
         return self.image_2d_via_radii_from(
-            self.eccentric_radii_grid_from(grid=grid, **kwargs)
+            self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs)
         )
 
     @property
diff --git a/autogalaxy/profiles/light/standard/gaussian.py b/autogalaxy/profiles/light/standard/gaussian.py
index 5d07ebb94..8b626f96e 100644
--- a/autogalaxy/profiles/light/standard/gaussian.py
+++ b/autogalaxy/profiles/light/standard/gaussian.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Optional, Tuple
 
@@ -48,7 +47,7 @@ def coefficient_tag(self) -> str:
             f"sigma_{np.round(self.sigma, 2)}__ell_comps_{np.round(self.ell_comps, 2)}"
         )
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
+    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
         """
         Returns the 2D image of the Gaussian light profile from a grid of coordinates which are the radial distance of
         each coordinate from the its `centre`.
@@ -60,12 +59,12 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
         grid_radii
             The radial distances from the centre of the profile, for each coordinate on the grid.
         """
-        return jnp.multiply(
+        return xp.multiply(
             self._intensity,
-            jnp.exp(
+            xp.exp(
                 -0.5
-                * jnp.square(
-                    jnp.divide(grid_radii.array, self.sigma / jnp.sqrt(self.axis_ratio()))
+                * xp.square(
+                    xp.divide(grid_radii.array, self.sigma / xp.sqrt(self.axis_ratio()))
                 )
             ),
         )
@@ -75,7 +74,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
     ) -> np.ndarray:
         """
         Returns the Gaussian light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
@@ -95,7 +94,7 @@ def image_2d_from(
         """
 
         return self.image_2d_via_radii_from(
-            self.eccentric_radii_grid_from(grid=grid, **kwargs)
+            self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs), xp=xp
         )
 
 
diff --git a/autogalaxy/profiles/light/standard/moffat.py b/autogalaxy/profiles/light/standard/moffat.py
index 2a9ed4baa..b924331f8 100644
--- a/autogalaxy/profiles/light/standard/moffat.py
+++ b/autogalaxy/profiles/light/standard/moffat.py
@@ -46,7 +46,7 @@ def __init__(
         self.alpha = alpha
         self.beta = beta
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
+    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
         """
         Returns the 2D image of the Moffat light profile from a grid of coordinates which are the radial distance of
         each coordinate from the its `centre`.
@@ -58,12 +58,12 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
         grid_radii
             The radial distances from the centre of the profile, for each coordinate on the grid.
         """
-        return jnp.multiply(
+        return xp.multiply(
             self._intensity,
-            jnp.power(
+            xp.power(
                 1
-                + jnp.square(
-                    jnp.divide(grid_radii.array, self.alpha / jnp.sqrt(self.axis_ratio()))
+                + xp.square(
+                    xp.divide(grid_radii.array, self.alpha / xp.sqrt(self.axis_ratio()))
                 ),
                 -self.beta,
             ),
@@ -74,7 +74,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
     ) -> np.ndarray:
         """
         Returns the Moffat light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
@@ -94,7 +94,7 @@ def image_2d_from(
         """
 
         return self.image_2d_via_radii_from(
-            self.eccentric_radii_grid_from(grid=grid, **kwargs)
+            self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs)
         )
 
 
diff --git a/autogalaxy/profiles/light/standard/sersic.py b/autogalaxy/profiles/light/standard/sersic.py
index 8445c349e..46c2e788a 100644
--- a/autogalaxy/profiles/light/standard/sersic.py
+++ b/autogalaxy/profiles/light/standard/sersic.py
@@ -53,7 +53,7 @@ def elliptical_effective_radius(self) -> float:
         The elliptical effective radius instead describes the major-axis radius of the ellipse containing
         half the light, and may be more appropriate for highly flattened systems like disk galaxies.
         """
-        return self.effective_radius / jnp.sqrt(self.axis_ratio())
+        return self.effective_radius / xp.sqrt(self.axis_ratio())
 
     @property
     def sersic_constant(self) -> float:
@@ -80,7 +80,7 @@ def image_2d_via_radii_from(self, radius: np.ndarray) -> np.ndarray:
         grid_radii
             The radial distances from the centre of the profile, for each coordinate on the grid.
         """
-        return self._intensity * jnp.exp(
+        return self._intensity * xp.exp(
             -self.sersic_constant
             * (((radius / self.effective_radius) ** (1.0 / self.sersic_index)) - 1)
         )
diff --git a/autogalaxy/profiles/light/standard/sersic_core.py b/autogalaxy/profiles/light/standard/sersic_core.py
index dbc444d04..c3f26f6fd 100644
--- a/autogalaxy/profiles/light/standard/sersic_core.py
+++ b/autogalaxy/profiles/light/standard/sersic_core.py
@@ -53,8 +53,7 @@ def __init__(
         self.alpha = alpha
         self.gamma = gamma
 
-    @property
-    def intensity_prime(self) -> float:
+    def intensity_prime(self, xp=np) -> float:
         """
         Overall intensity normalisation in the rescaled cored Sersic light profile.
 
@@ -64,7 +63,7 @@ def intensity_prime(self) -> float:
         return (
             self._intensity
             * (2.0 ** (-self.gamma / self.alpha))
-            * jnp.exp(
+            * xp.exp(
                 self.sersic_constant
                 * (
                     ((2.0 ** (1.0 / self.alpha)) * self.radius_break)
@@ -74,7 +73,7 @@ def intensity_prime(self) -> float:
             )
         )
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarray:
+    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np, **kwargs) -> np.ndarray:
         """
         Returns the 2D image of the Sersic light profile from a grid of coordinates which are the radial distances of
         each coordinate from the its `centre`.
@@ -85,27 +84,27 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarra
             The radial distances from the centre of the profile, for each coordinate on the grid.
         """
 
-        return jnp.multiply(
-            jnp.multiply(
-                self.intensity_prime,
-                jnp.power(
-                    jnp.add(
+        return xp.multiply(
+            xp.multiply(
+                self.intensity_prime(xp),
+                xp.power(
+                    xp.add(
                         1,
-                        jnp.power(
-                            jnp.divide(self.radius_break, grid_radii.array), self.alpha
+                        xp.power(
+                            xp.divide(self.radius_break, grid_radii.array), self.alpha
                         ),
                     ),
                     (self.gamma / self.alpha),
                 ),
             ),
-            jnp.exp(
-                jnp.multiply(
+            xp.exp(
+                xp.multiply(
                     -self.sersic_constant,
                     (
-                        jnp.power(
-                            jnp.divide(
-                                jnp.add(
-                                    jnp.power(grid_radii.array, self.alpha),
+                        xp.power(
+                            xp.divide(
+                                xp.add(
+                                    xp.power(grid_radii.array, self.alpha),
                                     (self.radius_break**self.alpha),
                                 ),
                                 (self.effective_radius**self.alpha),
diff --git a/autogalaxy/profiles/light/standard/shapelets/cartesian.py b/autogalaxy/profiles/light/standard/shapelets/cartesian.py
index 76785e471..01b725515 100644
--- a/autogalaxy/profiles/light/standard/shapelets/cartesian.py
+++ b/autogalaxy/profiles/light/standard/shapelets/cartesian.py
@@ -64,7 +64,7 @@ def coefficient_tag(self) -> str:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
     ) -> np.ndarray:
         """
         Returns the Cartesian Shapelet light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
@@ -97,12 +97,12 @@ def image_2d_from(
         return (
             shapelet_y
             * shapelet_x
-            * jnp.exp(-0.5 * (y**2 + x**2) / (self.beta**2))
+            * xp.exp(-0.5 * (y**2 + x**2) / (self.beta**2))
             / self.beta
             / (
-                jnp.sqrt(
+                xp.sqrt(
                     2 ** (self.n_x + self.n_y)
-                    * (jnp.pi)
+                    * (xp.pi)
                     * factorial(self.n_y)
                     * factorial(self.n_x)
                 )
diff --git a/autogalaxy/profiles/light/standard/shapelets/exponential.py b/autogalaxy/profiles/light/standard/shapelets/exponential.py
index 9dbf5bee7..dc76a3e60 100644
--- a/autogalaxy/profiles/light/standard/shapelets/exponential.py
+++ b/autogalaxy/profiles/light/standard/shapelets/exponential.py
@@ -65,7 +65,7 @@ def coefficient_tag(self) -> str:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
     ) -> np.ndarray:
         """
         Returns the Exponential Shapelet light profile's 2D image from a 2D grid of Exponential (y,x) coordinates.
@@ -87,30 +87,30 @@ def image_2d_from(
         from jax.scipy.special import factorial
 
         radial = (grid.array[:, 0] ** 2 + grid.array[:, 1] ** 2) / self.beta
-        theta = jnp.arctan(grid.array[:, 1] / grid.array[:, 0])
+        theta = xp.arctan(grid.array[:, 1] / grid.array[:, 0])
 
         prefactor = (
             1.0
-            / jnp.sqrt(2 * jnp.pi)
+            / xp.sqrt(2 * xp.pi)
             / self.beta
-            * (self.n + 0.5) ** (-1 - jnp.abs(self.m))
+            * (self.n + 0.5) ** (-1 - xp.abs(self.m))
             * (-1) ** (self.n + self.m)
-            * jnp.sqrt(
-                factorial(self.n - jnp.abs(self.m)) / 2 * self.n
-                + 1 / factorial(self.n + jnp.abs(self.m))
+            * xp.sqrt(
+                factorial(self.n - xp.abs(self.m)) / 2 * self.n
+                + 1 / factorial(self.n + xp.abs(self.m))
             )
         )
 
-        laguerre = genlaguerre(n=self.n - jnp.abs(self.m), alpha=2 * jnp.abs(self.m))
+        laguerre = genlaguerre(n=self.n - xp.abs(self.m), alpha=2 * xp.abs(self.m))
         shapelet = laguerre(radial / (self.n + 0.5))
 
-        return jnp.abs(
+        return xp.abs(
             prefactor
-            * jnp.exp(-radial / (2 * self.n + 1))
-            * radial ** (jnp.abs(self.m))
+            * xp.exp(-radial / (2 * self.n + 1))
+            * radial ** (xp.abs(self.m))
             * shapelet
-            * jnp.cos(self.m * theta)
-            + -1.0j * jnp.sin(self.m * theta)
+            * xp.cos(self.m * theta)
+            + -1.0j * xp.sin(self.m * theta)
         )
 
 
diff --git a/autogalaxy/profiles/light/standard/shapelets/polar.py b/autogalaxy/profiles/light/standard/shapelets/polar.py
index 9733308fc..d73eec5c2 100644
--- a/autogalaxy/profiles/light/standard/shapelets/polar.py
+++ b/autogalaxy/profiles/light/standard/shapelets/polar.py
@@ -65,7 +65,7 @@ def coefficient_tag(self) -> str:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
     ) -> np.ndarray:
         """
         Returns the Polar Shapelet light profile's 2D image from a 2D grid of Polar (y,x) coordinates.
@@ -87,29 +87,29 @@ def image_2d_from(
         from jax.scipy.special import factorial
 
         laguerre = genlaguerre(
-            n=(self.n - jnp.abs(self.m)) / 2.0, alpha=jnp.abs(self.m)
+            n=(self.n - xp.abs(self.m)) / 2.0, alpha=xp.abs(self.m)
         )
 
         const = (
-            ((-1) ** ((self.n - jnp.abs(self.m)) // 2))
-            * jnp.sqrt(
-                factorial((self.n - jnp.abs(self.m)) // 2)
-                / factorial((self.n + jnp.abs(self.m)) // 2)
+            ((-1) ** ((self.n - xp.abs(self.m)) // 2))
+            * xp.sqrt(
+                factorial((self.n - xp.abs(self.m)) // 2)
+                / factorial((self.n + xp.abs(self.m)) // 2)
             )
             / self.beta
-            / jnp.sqrt(jnp.pi)
+            / xp.sqrt(xp.pi)
         )
 
         rsq = (grid.array[:, 0] ** 2 + grid.array[:, 1] ** 2) / self.beta**2
-        theta = jnp.arctan2(grid.array[:, 1], grid.array[:, 0])
-        radial = rsq ** (abs(self.m / 2.0)) * jnp.exp(-rsq / 2.0) * laguerre(rsq)
+        theta = xp.arctan2(grid.array[:, 1], grid.array[:, 0])
+        radial = rsq ** (abs(self.m / 2.0)) * xp.exp(-rsq / 2.0) * laguerre(rsq)
 
         if self.m == 0:
             azimuthal = 1
         elif self.m > 0:
-            azimuthal = jnp.sin((-1) * self.m * theta)
+            azimuthal = xp.sin((-1) * self.m * theta)
         else:
-            azimuthal = jnp.cos((-1) * self.m * theta)
+            azimuthal = xp.cos((-1) * self.m * theta)
 
         return const * radial * azimuthal
 
diff --git a/autogalaxy/profiles/mass/abstract/abstract.py b/autogalaxy/profiles/mass/abstract/abstract.py
index 80a2da261..8bedfb61d 100644
--- a/autogalaxy/profiles/mass/abstract/abstract.py
+++ b/autogalaxy/profiles/mass/abstract/abstract.py
@@ -51,21 +51,21 @@ def convergence_func(self, grid_radius: float) -> float:
         raise NotImplementedError
 
     @aa.grid_dec.project_grid
-    def convergence_1d_from(self, grid: aa.type.Grid1D2DLike) -> aa.type.Grid1D2DLike:
-        return self.convergence_2d_from(grid=grid)
+    def convergence_1d_from(self, grid: aa.type.Grid1D2DLike, xp=np) -> aa.type.Grid1D2DLike:
+        return self.convergence_2d_from(grid=grid, xp=xp)
 
     def potential_2d_from(self, grid):
         raise NotImplementedError
 
     @aa.grid_dec.project_grid
-    def potential_1d_from(self, grid: aa.type.Grid1D2DLike) -> aa.type.Grid1D2DLike:
-        return self.potential_2d_from(grid=grid)
+    def potential_1d_from(self, grid: aa.type.Grid1D2DLike, xp=np) -> aa.type.Grid1D2DLike:
+        return self.potential_2d_from(grid=grid, xp=xp)
 
     def potential_func(self, u, y, x):
         raise NotImplementedError
 
-    def mass_integral(self, x):
-        return 2 * jnp.pi * x * self.convergence_func(grid_radius=aa.ArrayIrregular(x))
+    def mass_integral(self, x, xp=np):
+        return 2 * xp.pi * x * self.convergence_func(grid_radius=aa.ArrayIrregular(x))
 
     @property
     def ellipticity_rescale(self):
diff --git a/autogalaxy/profiles/mass/abstract/cse.py b/autogalaxy/profiles/mass/abstract/cse.py
index 5b12688b9..74a08ee1c 100644
--- a/autogalaxy/profiles/mass/abstract/cse.py
+++ b/autogalaxy/profiles/mass/abstract/cse.py
@@ -41,12 +41,12 @@ def deflections_via_cse_from(
         Parameters
         ----------
         """
-        phi = jnp.sqrt(axis_ratio_squared * core_radius**2.0 + term1)
+        phi = xp.sqrt(axis_ratio_squared * core_radius**2.0 + term1)
         Psi = (phi + core_radius) ** 2.0 + term2
         bottom = core_radius * phi * Psi
         defl_x = (term3 * (phi + axis_ratio_squared * core_radius)) / bottom
         defl_y = (term4 * (phi + core_radius)) / bottom
-        return jnp.vstack((defl_y, defl_x))
+        return xp.vstack((defl_y, defl_x))
 
     @abstractmethod
     def decompose_convergence_via_cse(self, grid_radii: np.ndarray):
diff --git a/autogalaxy/profiles/mass/abstract/jax_utils.py b/autogalaxy/profiles/mass/abstract/jax_utils.py
index e241488f3..43466c034 100644
--- a/autogalaxy/profiles/mass/abstract/jax_utils.py
+++ b/autogalaxy/profiles/mass/abstract/jax_utils.py
@@ -26,7 +26,7 @@ def reg2(z, sqrt_pi, _):
     for s in r2_s2:
         f2 = s - f2 * mz2
 
-    return jnp.exp(mz2) + 1j * z * f1 / f2
+    return xp.exp(mz2) + 1j * z * f1 / f2
 
 
 r3_s1 = [5.9126262, 30.180142, 93.15558, 181.92853, 214.38239, 122.60793]
@@ -55,7 +55,7 @@ def w_f_approx(z):
     :return: :math:`w_\\mathrm{F}(z)`
     :rtype: ``complex``
     """
-    sqrt_pi = 1 / jnp.sqrt(jnp.pi)
+    sqrt_pi = 1 / xp.sqrt(xp.pi)
     i_sqrt_pi = 1j * sqrt_pi
 
     z_imag2 = z.imag**2
@@ -69,31 +69,31 @@ def w_f_approx(z):
     r2_1 = (abs_z2 >= 30.0) & (abs_z2 < 62.0) & (z_imag2 < 1e-13)
     r2_2 = (abs_z2 >= 2.5) & (abs_z2 < 30.0) & (z_imag2 < 0.072)
     r2 = r2_1 | r2_2
-    r3 = jnp.logical_not(r1) & jnp.logical_not(r2)
+    r3 = xp.logical_not(r1) & xp.logical_not(r2)
 
     # exploit symmetry to avoid overflow in some regions
     r_flip = z.imag < 0
-    z_adjust = jnp.where(r_flip, -z, z)
-    two_exp_zz = 2 * jnp.exp(-(z_adjust**2))
+    z_adjust = xp.where(r_flip, -z, z)
+    two_exp_zz = 2 * xp.exp(-(z_adjust**2))
 
     args = (z_adjust, sqrt_pi, i_sqrt_pi)
-    wz = jnp.empty_like(z)
-    wz = jnp.where(r1, reg1(*args), wz)
-    wz = jnp.where(r2, reg2(*args), wz)
-    wz = jnp.where(r3, reg3(*args), wz)
+    wz = xp.empty_like(z)
+    wz = xp.where(r1, reg1(*args), wz)
+    wz = xp.where(r2, reg2(*args), wz)
+    wz = xp.where(r3, reg3(*args), wz)
 
     # exploit symmetry to avoid overflow in some regions
-    wz = jnp.where(r_flip, two_exp_zz - wz, wz)
+    wz = xp.where(r_flip, two_exp_zz - wz, wz)
 
     return wz
 
 
 @w_f_approx.defjvp
 def w_f_approx_jvp(primals, tangents):
-    # define a custom jvp to avoid the issue using `jnp.where` with `jax.grad`
+    # define a custom jvp to avoid the issue using `xp.where` with `jax.grad`
     (z,) = primals
     (z_dot,) = tangents
     primal_out = w_f_approx(z)
-    i_sqrt_pi = 1j / jnp.sqrt(jnp.pi)
+    i_sqrt_pi = 1j / xp.sqrt(xp.pi)
     tangent_out = z_dot * 2 * (i_sqrt_pi - z * primal_out)
     return primal_out, tangent_out
diff --git a/autogalaxy/profiles/mass/abstract/mge_jax.py b/autogalaxy/profiles/mass/abstract/mge_jax.py
index 75dc16dad..ae4010161 100644
--- a/autogalaxy/profiles/mass/abstract/mge_jax.py
+++ b/autogalaxy/profiles/mass/abstract/mge_jax.py
@@ -21,12 +21,12 @@ def zeta_from(grid, amps, sigmas, axis_ratio):
         """
         q2 = axis_ratio**2.0
 
-        scale_factor = axis_ratio / jnp.sqrt(2.0 * (1.0 - q2))
+        scale_factor = axis_ratio / xp.sqrt(2.0 * (1.0 - q2))
 
-        xs = jnp.array((grid.array[:, 1] * scale_factor).copy())
-        ys = jnp.array((grid.array[:, 0] * scale_factor).copy())
+        xs = xp.array((grid.array[:, 1] * scale_factor).copy())
+        ys = xp.array((grid.array[:, 0] * scale_factor).copy())
 
-        y_sign = jnp.sign(ys)
+        y_sign = xp.sign(ys)
         ys = ys * y_sign
 
         z = xs + 1j * ys
@@ -40,7 +40,7 @@ def zeta_from(grid, amps, sigmas, axis_ratio):
         # could try `jax.lax.scan` instead if this is too much memory
         w = w_f_approx(inv_sigma_ * z)
         wq = w_f_approx(inv_sigma_ * zq)
-        exp_factor = jnp.exp(inv_sigma_**2 * expv)
+        exp_factor = xp.exp(inv_sigma_**2 * expv)
 
         sigma_func_real = w.imag - exp_factor * wq.imag
         sigma_func_imag = (-w.real + exp_factor * wq.real) * y_sign
@@ -53,8 +53,8 @@ def kesi(p):
         """
         see Eq.(6) of 1906.08263
         """
-        n_list = jnp.arange(0, 2 * p + 1, 1)
-        return (2.0 * p * jnp.log(10) / 3.0 + 2.0 * jnp.pi * n_list * 1j) ** (0.5)
+        n_list = xp.arange(0, 2 * p + 1, 1)
+        return (2.0 * p * xp.log(10) / 3.0 + 2.0 * xp.pi * n_list * 1j) ** (0.5)
 
     @staticmethod
     def eta(p):
@@ -62,15 +62,15 @@ def eta(p):
         see Eq.(6) of 1906.00263
         """
 
-        i = jnp.arange(1, p, 1)
+        i = xp.arange(1, p, 1)
         kesi_last = 1 / 2**p
-        k = kesi_last + jnp.cumsum(jnp.cumprod((p + 1 - i) / i) * kesi_last)
+        k = kesi_last + xp.cumsum(xp.cumprod((p + 1 - i) / i) * kesi_last)
 
-        kesi_list = jnp.hstack(
-            [jnp.array([0.5]), jnp.ones(p), k[::-1], jnp.array([kesi_last])]
+        kesi_list = xp.hstack(
+            [xp.array([0.5]), xp.ones(p), k[::-1], xp.array([kesi_last])]
         )
-        coef = (-1) ** jnp.arange(0, 2 * p + 1, 1)
-        eta_const = 2.0 * jnp.sqrt(2.0 * jnp.pi) * 10 ** (p / 3.0)
+        coef = (-1) ** xp.arange(0, 2 * p + 1, 1)
+        eta_const = 2.0 * xp.sqrt(2.0 * xp.pi) * 10 ** (p / 3.0)
         eta_list = coef * kesi_list
         return eta_const, eta_list
 
@@ -102,16 +102,16 @@ def _decompose_convergence_via_mge(
 
         # sigma is sampled from logspace between these radii.
 
-        log_sigmas = jnp.linspace(
-            jnp.log(radii_min), jnp.log(radii_max), func_gaussians
+        log_sigmas = xp.linspace(
+            xp.log(radii_min), xp.log(radii_max), func_gaussians
         )
         d_log_sigma = log_sigmas[1] - log_sigmas[0]
-        sigma_list = jnp.exp(log_sigmas)
+        sigma_list = xp.exp(log_sigmas)
 
-        f_sigma = eta_constant * jnp.sum(
-            eta_n * jnp.real(func(sigma_list.reshape(-1, 1) * kesis)), axis=1
+        f_sigma = eta_constant * xp.sum(
+            eta_n * xp.real(func(sigma_list.reshape(-1, 1) * kesis)), axis=1
         )
-        amplitude_list = f_sigma * d_log_sigma / jnp.sqrt(2.0 * jnp.pi)
+        amplitude_list = f_sigma * d_log_sigma / xp.sqrt(2.0 * xp.pi)
         amplitude_list = amplitude_list.at[0].multiply(0.5)
         amplitude_list = amplitude_list.at[-1].multiply(0.5)
 
@@ -137,13 +137,13 @@ def _convergence_2d_via_mge_from(
 
         inv_sigma_ = 1 / sigmas.reshape((-1,) + (1,) * grid_radii.array.ndim)
         amps_ = amps.reshape((-1,) + (1,) * grid_radii.array.ndim)
-        convergence = amps_ * jnp.exp(-0.5 * (grid_radii.array * inv_sigma_) ** 2)
+        convergence = amps_ * xp.exp(-0.5 * (grid_radii.array * inv_sigma_) ** 2)
         return convergence.sum(axis=0)
 
     def _deflections_2d_via_mge_from(
         self, grid, sigmas_factor=1.0, func_terms=28, func_gaussians=20
     ):
-        axis_ratio = jnp.min(jnp.array([self.axis_ratio(), 0.9999]))
+        axis_ratio = xp.min(xp.array([self.axis_ratio(), 0.9999]))
 
         amps, sigmas = self.decompose_convergence_via_mge(
             func_terms=func_terms, func_gaussians=func_gaussians
@@ -154,8 +154,8 @@ def _deflections_2d_via_mge_from(
             grid=grid, amps=amps, sigmas=sigmas, axis_ratio=axis_ratio
         )
 
-        angle *= jnp.sqrt((2.0 * jnp.pi) / (1.0 - axis_ratio**2.0))
+        angle *= xp.sqrt((2.0 * xp.pi) / (1.0 - axis_ratio**2.0))
 
         return self.rotated_grid_from_reference_frame_from(
-            jnp.vstack((-angle.imag, angle.real)).T
+            xp.vstack((-angle.imag, angle.real)).T
         )
diff --git a/autogalaxy/profiles/mass/dark/abstract.py b/autogalaxy/profiles/mass/dark/abstract.py
index f9a394486..20f1bd426 100644
--- a/autogalaxy/profiles/mass/dark/abstract.py
+++ b/autogalaxy/profiles/mass/dark/abstract.py
@@ -57,7 +57,7 @@ def __init__(
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """Calculate the projected convergence at a given set of arc-second gridded coordinates.
 
         Parameters
@@ -74,7 +74,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_via_mge_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_via_mge_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """Calculate the projected convergence at a given set of arc-second gridded coordinates.
 
         Parameters
@@ -130,7 +130,7 @@ def gnfw_3d(r):
         amplitude_list *= np.sqrt(2.0 * np.pi) * sigma_list
         return amplitude_list, sigma_list
 
-    def coord_func_f(self, grid_radius: jnp.ndarray) -> jnp.ndarray:
+    def coord_func_f(self, grid_radius: np.ndarray) -> np.ndarray:
         """
         Given an array `grid_radius` and a work array `f`, fill f[i] with
 
@@ -142,24 +142,24 @@ def coord_func_f(self, grid_radius: jnp.ndarray) -> jnp.ndarray:
         any Python control flow on tracer values.
         """
         if isinstance(grid_radius, float) or isinstance(grid_radius, complex):
-            grid_radius = jnp.array([grid_radius])
+            grid_radius = xp.array([grid_radius])
 
-        f = jnp.ones(shape=grid_radius.shape[0], dtype="complex64")
+        f = xp.ones(shape=grid_radius.shape[0], dtype="complex64")
 
         # compute both branches
         r = grid_radius
         inv_r = 1.0 / r
 
         # branch for r > 1
-        out_gt = (1.0 / jnp.sqrt(r**2 - 1.0)) * jnp.arccos(inv_r)
+        out_gt = (1.0 / xp.sqrt(r**2 - 1.0)) * xp.arccos(inv_r)
 
         # branch for r < 1
-        out_lt = (1.0 / jnp.sqrt(1.0 - r**2)) * jnp.arccosh(inv_r)
+        out_lt = (1.0 / xp.sqrt(1.0 - r**2)) * xp.arccosh(inv_r)
 
         # combine: if r>1 pick out_gt, elif r<1 pick out_lt, else keep original f
-        return jnp.where(r > 1.0, out_gt, jnp.where(r < 1.0, out_lt, f))
+        return xp.where(r > 1.0, out_gt, xp.where(r < 1.0, out_lt, f))
 
-    def coord_func_g(self, grid_radius: jnp.ndarray) -> jnp.ndarray:
+    def coord_func_g(self, grid_radius: np.ndarray) -> np.ndarray:
         """
         Vectorized version of the original looped `coord_func_g_jit`.
 
@@ -170,32 +170,32 @@ def coord_func_g(self, grid_radius: jnp.ndarray) -> jnp.ndarray:
 
         Parameters
         ----------
-        grid_radius : jnp.ndarray
+        grid_radius : np.ndarray
             The input grid radius values.
-        f_r : jnp.ndarray
+        f_r : np.ndarray
             Precomputed values from `coord_func_f`.
-        g : jnp.ndarray
+        g : np.ndarray
             Output array (will be overwritten).
 
         Returns
         -------
-        jnp.ndarray
+        np.ndarray
             The updated `g` array.
         """
         # Convert single values to JAX arrays
         if isinstance(grid_radius, (float, complex)):
-            grid_radius = jnp.array([grid_radius], dtype=jnp.complex64)
+            grid_radius = xp.array([grid_radius], dtype=xp.complex64)
 
         # Evaluate f_r
         f_r = self.coord_func_f(grid_radius=grid_radius)
 
-        r = jnp.real(grid_radius)
+        r = xp.real(grid_radius)
         r2 = r**2
 
-        return jnp.where(
+        return xp.where(
             r > 1.0,
             (1.0 - f_r) / (r2 - 1.0),
-            jnp.where(
+            xp.where(
                 r < 1.0,
                 (f_r - 1.0) / (1.0 - r2),
                 1.0 / 3.0,
@@ -203,7 +203,7 @@ def coord_func_g(self, grid_radius: jnp.ndarray) -> jnp.ndarray:
         )
 
     def coord_func_h(self, grid_radius):
-        return jnp.log(grid_radius / 2.0) + self.coord_func_f(grid_radius=grid_radius)
+        return xp.log(grid_radius / 2.0) + self.coord_func_f(grid_radius=grid_radius)
 
     def rho_at_scale_radius_solar_mass_per_kpc3(
         self, redshift_object, redshift_source, cosmology: LensingCosmology = None
@@ -301,7 +301,7 @@ def concentration_func(concentration, delta_concentration):
                 concentration
                 * concentration
                 * concentration
-                / (jnp.log(1 + concentration) - concentration / (1 + concentration))
+                / (xp.log(1 + concentration) - concentration / (1 + concentration))
             )
             - delta_concentration
         )
@@ -372,7 +372,7 @@ def mass_at_200_solar_masses(
 
         return (
             200.0
-            * ((4.0 / 3.0) * jnp.pi)
+            * ((4.0 / 3.0) * xp.pi)
             * cosmic_average_density
             * (radius_at_200_kpc**3.0)
         )
diff --git a/autogalaxy/profiles/mass/dark/gnfw.py b/autogalaxy/profiles/mass/dark/gnfw.py
index 941d7f0c6..97a30d124 100644
--- a/autogalaxy/profiles/mass/dark/gnfw.py
+++ b/autogalaxy/profiles/mass/dark/gnfw.py
@@ -7,12 +7,12 @@
 
 
 class gNFW(AbstractgNFW):
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return self.deflections_2d_via_mge_from(grid=grid, **kwargs)
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return self._deflections_2d_via_mge_from(
             grid=grid, sigmas_factor=self.axis_ratio()
         )
@@ -295,7 +295,7 @@ def __init__(
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index 137df8872..d01a449a6 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -48,7 +48,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -91,7 +91,7 @@ def calculate_deflection_component(npow, index):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_cse_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return self._deflections_2d_via_cse_from(grid=grid, **kwargs)
 
     @staticmethod
@@ -121,7 +121,7 @@ def deflection_func(u, y, x, npow, axis_ratio, scale_radius):
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the projected 2D convergence from a grid of (y,x) arc second coordinates, by computing and summing
         the convergence of each individual cse used to decompose the mass profile.
@@ -149,7 +149,7 @@ def convergence_func(self, grid_radius: float) -> float:
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the potential at a given set of arc-second gridded coordinates.
 
@@ -294,7 +294,7 @@ def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Analytic calculation convergence from Heyrovský & Karamazov 2024
 
@@ -360,12 +360,12 @@ def __init__(
             scale_radius=scale_radius,
         )
 
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return self.deflections_2d_via_analytic_from(grid=grid, **kwargs)
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -375,11 +375,11 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs):
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
 
-        eta = jnp.multiply(
+        eta = xp.multiply(
             1.0 / self.scale_radius, self.radial_grid_from(grid=grid, **kwargs).array
         )
 
-        deflection_grid = jnp.multiply(
+        deflection_grid = xp.multiply(
             (4.0 * self.kappa_s * self.scale_radius / eta),
             self.deflection_func_sph(grid_radius=eta),
         )
@@ -388,12 +388,12 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
     def deflection_func_sph(self, grid_radius):
         grid_radius = grid_radius + 0j
-        return jnp.real(self.coord_func_h(grid_radius=grid_radius))
+        return xp.real(self.coord_func_h(grid_radius=grid_radius))
 
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the potential at a given set of arc-second gridded coordinates.
 
diff --git a/autogalaxy/profiles/mass/dark/nfw_truncated.py b/autogalaxy/profiles/mass/dark/nfw_truncated.py
index b33f4d0d7..e680e7ec9 100644
--- a/autogalaxy/profiles/mass/dark/nfw_truncated.py
+++ b/autogalaxy/profiles/mass/dark/nfw_truncated.py
@@ -1,4 +1,4 @@
-import jax.numpy as jnp
+import numpy as np
 from typing import Tuple
 
 import autoarray as aa
@@ -28,7 +28,7 @@ def __init__(
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -38,68 +38,68 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
 
-        eta = jnp.multiply(
+        eta = xp.multiply(
             1.0 / self.scale_radius, self.radial_grid_from(grid=grid, **kwargs).array
         )
 
-        deflection_grid = jnp.multiply(
+        deflection_grid = xp.multiply(
             (4.0 * self.kappa_s * self.scale_radius / eta),
             self.deflection_func_sph(grid_radius=eta),
         )
 
-        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid)
+        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid, xp=xp)
 
-    def deflection_func_sph(self, grid_radius):
+    def deflection_func_sph(self, grid_radius, xp=np):
         grid_radius = grid_radius + 0j
-        return jnp.real(self.coord_func_m(grid_radius=grid_radius))
+        return xp.real(self.coord_func_m(grid_radius=grid_radius, xp=xp))
 
-    def convergence_func(self, grid_radius: float) -> float:
+    def convergence_func(self, grid_radius: float, xp=np) -> float:
         grid_radius = ((1.0 / self.scale_radius) * grid_radius) + 0j
-        return jnp.real(
-            2.0 * self.kappa_s * self.coord_func_l(grid_radius=grid_radius.array)
+        return xp.real(
+            2.0 * self.kappa_s * self.coord_func_l(grid_radius=grid_radius.array, xp=xp)
         )
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
-        return jnp.zeros(shape=grid.shape[0])
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return xp.zeros(shape=grid.shape[0])
 
-    def coord_func_k(self, grid_radius):
-        return jnp.log(
-            jnp.divide(
+    def coord_func_k(self, grid_radius, xp=np):
+        return xp.log(
+            xp.divide(
                 grid_radius,
-                jnp.sqrt(jnp.square(grid_radius) + jnp.square(self.tau)) + self.tau,
+                xp.sqrt(xp.square(grid_radius) + xp.square(self.tau)) + self.tau,
             )
         )
 
-    def coord_func_l(self, grid_radius):
-        f_r = self.coord_func_f(grid_radius=grid_radius)
-        g_r = self.coord_func_g(grid_radius=grid_radius)
-        k_r = self.coord_func_k(grid_radius=grid_radius)
+    def coord_func_l(self, grid_radius, xp=np):
+        f_r = self.coord_func_f(grid_radius=grid_radius, xp=xp)
+        g_r = self.coord_func_g(grid_radius=grid_radius, xp=xp)
+        k_r = self.coord_func_k(grid_radius=grid_radius, xp=xp)
 
-        return jnp.divide(self.tau**2.0, (self.tau**2.0 + 1.0) ** 2.0) * (
+        return xp.divide(self.tau**2.0, (self.tau**2.0 + 1.0) ** 2.0) * (
             ((self.tau**2.0 + 1.0) * g_r)
             + (2 * f_r)
-            - (jnp.pi / (jnp.sqrt(self.tau**2.0 + grid_radius**2.0)))
+            - (xp.pi / (xp.sqrt(self.tau**2.0 + grid_radius**2.0)))
             + (
                 (
                     (self.tau**2.0 - 1.0)
-                    / (self.tau * (jnp.sqrt(self.tau**2.0 + grid_radius**2.0)))
+                    / (self.tau * (xp.sqrt(self.tau**2.0 + grid_radius**2.0)))
                 )
                 * k_r
             )
         )
 
-    def coord_func_m(self, grid_radius):
-        f_r = self.coord_func_f(grid_radius=grid_radius)
-        k_r = self.coord_func_k(grid_radius=grid_radius)
+    def coord_func_m(self, grid_radius, xp=np):
+        f_r = self.coord_func_f(grid_radius=grid_radius, xp=xp)
+        k_r = self.coord_func_k(grid_radius=grid_radius, xp=xp)
 
         return (self.tau**2.0 / (self.tau**2.0 + 1.0) ** 2.0) * (
             ((self.tau**2.0 + 2.0 * grid_radius**2.0 - 1.0) * f_r)
-            + (jnp.pi * self.tau)
-            + ((self.tau**2.0 - 1.0) * jnp.log(self.tau))
+            + (xp.pi * self.tau)
+            + ((self.tau**2.0 - 1.0) * xp.log(self.tau))
             + (
-                jnp.sqrt(grid_radius**2.0 + self.tau**2.0)
-                * (((self.tau**2.0 - 1.0) / self.tau) * k_r - jnp.pi)
+                xp.sqrt(grid_radius**2.0 + self.tau**2.0)
+                * (((self.tau**2.0 - 1.0) / self.tau) * k_r - xp.pi)
             )
         )
 
@@ -125,8 +125,8 @@ def mass_at_truncation_radius_solar_mass(
             mass_at_200
             * (self.tau**2.0 / (self.tau**2.0 + 1.0) ** 2.0)
             * (
-                ((self.tau**2.0 - 1) * jnp.log(self.tau))
-                + (self.tau * jnp.pi)
+                ((self.tau**2.0 - 1) * np.log(self.tau))
+                + (self.tau * np.pi)
                 - (self.tau**2.0 + 1)
             )
         )
diff --git a/autogalaxy/profiles/mass/point/point.py b/autogalaxy/profiles/mass/point/point.py
index 581484e62..d44271e02 100644
--- a/autogalaxy/profiles/mass/point/point.py
+++ b/autogalaxy/profiles/mass/point/point.py
@@ -24,7 +24,7 @@ def __init__(
         super().__init__(centre=centre, ell_comps=(0.0, 0.0))
         self.einstein_radius = einstein_radius
 
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    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]
         )
@@ -35,12 +35,12 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         return convergence
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return np.zeros(shape=grid.shape[0])
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         grid_radii = self.radial_grid_from(grid=grid, **kwargs)
         return self._cartesian_grid_via_radial_from(
             grid=grid, radius=self.einstein_radius**2 / grid_radii
diff --git a/autogalaxy/profiles/mass/point/smbh_binary.py b/autogalaxy/profiles/mass/point/smbh_binary.py
index 3731dde7f..9e5d83091 100644
--- a/autogalaxy/profiles/mass/point/smbh_binary.py
+++ b/autogalaxy/profiles/mass/point/smbh_binary.py
@@ -90,7 +90,7 @@ def angle_binary_radians(self) -> float:
         """
         return self.angle_binary * np.pi / 180.0
 
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates.
 
@@ -105,7 +105,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
             grid=grid
         ) + self.smbh_1.convergence_2d_from(grid=grid, **kwargs)
 
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Returns the two dimensional projected potential on a grid of (y,x) arc-second coordinates.
 
@@ -120,7 +120,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
             grid=grid, **kwargs
         ) + self.smbh_1.potential_2d_from(grid=grid)
 
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Returns the two dimensional deflection angles on a grid of (y,x) arc-second coordinates.
 
diff --git a/autogalaxy/profiles/mass/sheets/mass_sheet.py b/autogalaxy/profiles/mass/sheets/mass_sheet.py
index b4c1414f4..95392d512 100644
--- a/autogalaxy/profiles/mass/sheets/mass_sheet.py
+++ b/autogalaxy/profiles/mass/sheets/mass_sheet.py
@@ -26,16 +26,16 @@ def convergence_func(self, grid_radius: float) -> float:
         return 0.0
 
     @aa.grid_dec.to_array
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
-        return jnp.full(shape=grid.shape[0], fill_value=self.kappa)
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return xp.full(shape=grid.shape[0], fill_value=self.kappa)
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
-        return jnp.zeros(shape=grid.shape[0])
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return xp.zeros(shape=grid.shape[0])
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         grid_radii = self.radial_grid_from(grid=grid, **kwargs)
         return self._cartesian_grid_via_radial_from(
             grid=grid, radius=self.kappa * grid_radii
diff --git a/autogalaxy/profiles/mass/stellar/chameleon.py b/autogalaxy/profiles/mass/stellar/chameleon.py
index b925d7e0e..1fbd41708 100644
--- a/autogalaxy/profiles/mass/stellar/chameleon.py
+++ b/autogalaxy/profiles/mass/stellar/chameleon.py
@@ -48,7 +48,7 @@ def __init__(
         self.core_radius_0 = core_radius_0
         self.core_radius_1 = core_radius_1
 
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return self.deflections_2d_via_analytic_from(grid=grid, **kwargs)
 
     @aa.grid_dec.to_vector_yx
@@ -71,13 +71,13 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
             * self.intensity
             / (1 + self.axis_ratio(xp))
             * self.axis_ratio(xp)
-            / jnp.sqrt(1.0 - self.axis_ratio(xp)**2.0)
+            / xp.sqrt(1.0 - self.axis_ratio(xp)**2.0)
         )
 
-        core_radius_0 = jnp.sqrt(
+        core_radius_0 = xp.sqrt(
             (4.0 * self.core_radius_0**2.0) / (1.0 + self.axis_ratio(xp)) ** 2
         )
-        core_radius_1 = jnp.sqrt(
+        core_radius_1 = xp.sqrt(
             (4.0 * self.core_radius_1**2.0) / (1.0 + self.axis_ratio(xp)) ** 2
         )
 
@@ -88,45 +88,45 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
             grid=grid, axis_ratio=self.axis_ratio(xp), core_radius=core_radius_1, xp=xp
         )
 
-        deflection_y0 = jnp.arctanh(
-            jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 0]),
-                jnp.add(psi0, self.axis_ratio(xp)**2.0 * core_radius_0),
+        deflection_y0 = xp.arctanh(
+            xp.divide(
+                xp.multiply(xp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 0]),
+                xp.add(psi0, self.axis_ratio(xp)**2.0 * core_radius_0),
             )
         )
 
-        deflection_x0 = jnp.arctan(
-            jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 1]),
-                jnp.add(psi0, core_radius_0),
+        deflection_x0 = xp.arctan(
+            xp.divide(
+                xp.multiply(xp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 1]),
+                xp.add(psi0, core_radius_0),
             )
         )
 
-        deflection_y1 = jnp.arctanh(
-            jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 0]),
-                jnp.add(psi1, self.axis_ratio(xp)**2.0 * core_radius_1),
+        deflection_y1 = xp.arctanh(
+            xp.divide(
+                xp.multiply(xp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 0]),
+                xp.add(psi1, self.axis_ratio(xp)**2.0 * core_radius_1),
             )
         )
 
-        deflection_x1 = jnp.arctan(
-            jnp.divide(
-                jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 1]),
-                jnp.add(psi1, core_radius_1),
+        deflection_x1 = xp.arctan(
+            xp.divide(
+                xp.multiply(xp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 1]),
+                xp.add(psi1, core_radius_1),
             )
         )
 
-        deflection_y = jnp.subtract(deflection_y0, deflection_y1)
-        deflection_x = jnp.subtract(deflection_x0, deflection_x1)
+        deflection_y = xp.subtract(deflection_y0, deflection_y1)
+        deflection_x = xp.subtract(deflection_x0, deflection_x1)
 
         return self.rotated_grid_from_reference_frame_from(
-            jnp.multiply(factor, jnp.vstack((deflection_y, deflection_x)).T)
+            xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T)
         )
 
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """Calculate the projected convergence at a given set of arc-second gridded coordinates.
         Parameters
         ----------
@@ -141,10 +141,10 @@ def convergence_func(self, grid_radius: float) -> float:
         return self.mass_to_light_ratio * self.image_2d_via_radii_from(grid_radius)
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
-        return jnp.zeros(shape=grid.shape[0])
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return xp.zeros(shape=grid.shape[0])
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray):
+    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np):
         """Calculate the intensity of the Chamelon light profile on a grid of radial coordinates.
 
         Parameters
@@ -155,23 +155,23 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray):
 
         axis_ratio_factor = (1.0 + self.axis_ratio(xp)) ** 2.0
 
-        return jnp.multiply(
+        return xp.multiply(
             self.intensity / (1 + self.axis_ratio(xp)),
-            jnp.add(
-                jnp.divide(
+            xp.add(
+                xp.divide(
                     1.0,
-                    jnp.sqrt(
-                        jnp.add(
-                            jnp.square(grid_radii.array),
+                    xp.sqrt(
+                        xp.add(
+                            xp.square(grid_radii.array),
                             (4.0 * self.core_radius_0**2.0) / axis_ratio_factor,
                         )
                     ),
                 ),
-                -jnp.divide(
+                -xp.divide(
                     1.0,
-                    jnp.sqrt(
-                        jnp.add(
-                            jnp.square(grid_radii.array),
+                    xp.sqrt(
+                        xp.add(
+                            xp.square(grid_radii.array),
                             (4.0 * self.core_radius_1**2.0) / axis_ratio_factor,
                         )
                     ),
diff --git a/autogalaxy/profiles/mass/stellar/gaussian.py b/autogalaxy/profiles/mass/stellar/gaussian.py
index 5b6f92c79..bf58c3c51 100644
--- a/autogalaxy/profiles/mass/stellar/gaussian.py
+++ b/autogalaxy/profiles/mass/stellar/gaussian.py
@@ -43,7 +43,7 @@ def __init__(
         self.intensity = intensity
         self.sigma = sigma
 
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -61,7 +61,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -88,7 +88,7 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -145,7 +145,7 @@ def deflection_func(u, y, x, npow, axis_ratio, sigma):
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """Calculate the projected convergence at a given set of arc-second gridded coordinates.
 
         Parameters
@@ -155,17 +155,17 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
         """
         return self.convergence_func(
-            self.eccentric_radii_grid_from(grid=grid, **kwargs)
+            self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs)
         )
 
     def convergence_func(self, grid_radius: float) -> float:
         return self.mass_to_light_ratio * self.image_2d_via_radii_from(grid_radius)
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return np.zeros(shape=grid.shape[0])
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray):
+    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np):
         """Calculate the intensity of the Gaussian light profile on a grid of radial coordinates.
 
         Parameters
diff --git a/autogalaxy/profiles/mass/stellar/sersic.py b/autogalaxy/profiles/mass/stellar/sersic.py
index 893ec5f32..c33ae956a 100644
--- a/autogalaxy/profiles/mass/stellar/sersic.py
+++ b/autogalaxy/profiles/mass/stellar/sersic.py
@@ -122,7 +122,7 @@ def __init__(
         self.effective_radius = effective_radius
         self.sersic_index = sersic_index
 
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return self.deflections_2d_via_cse_from(grid=grid, **kwargs)
 
     @aa.grid_dec.to_vector_yx
@@ -152,7 +152,7 @@ def deflections_2d_via_mge_from(
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_cse_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the projected 2D deflection angles from a grid of (y,x) arc second coordinates, by computing and
         summing the convergence of each individual cse used to decompose the mass profile.
@@ -171,7 +171,7 @@ def deflections_2d_via_cse_from(self, grid: aa.type.Grid2DLike, **kwargs):
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """Calculate the projected convergence at a given set of arc-second gridded coordinates.
 
         Parameters
@@ -181,7 +181,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
         """
         return self.convergence_func(
-            self.eccentric_radii_grid_from(grid=grid, **kwargs)
+            self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs)
         )
 
     @aa.over_sample
@@ -200,7 +200,7 @@ def convergence_2d_via_mge_from(
 
         """
 
-        eccentric_radii = self.eccentric_radii_grid_from(grid=grid, **kwargs)
+        eccentric_radii = self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs)
 
         return self._convergence_2d_via_mge_from(
             grid_radii=eccentric_radii,
@@ -211,7 +211,7 @@ def convergence_2d_via_mge_from(
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the projected 2D convergence from a grid of (y,x) arc second coordinates, by computing and summing
         the convergence of each individual cse used to decompose the mass profile.
@@ -234,7 +234,7 @@ def convergence_func(self, grid_radius: float) -> float:
         return self.mass_to_light_ratio * self.image_2d_via_radii_from(grid_radius)
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return np.zeros(shape=grid.shape[0])
 
     def image_2d_via_radii_from(self, radius: np.ndarray):
@@ -372,7 +372,7 @@ def elliptical_effective_radius(self):
 class Sersic(AbstractSersic, MassProfileMGE, MassProfileCSE):
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
diff --git a/autogalaxy/profiles/mass/stellar/sersic_core.py b/autogalaxy/profiles/mass/stellar/sersic_core.py
index 496e4f7f4..47684f52e 100644
--- a/autogalaxy/profiles/mass/stellar/sersic_core.py
+++ b/autogalaxy/profiles/mass/stellar/sersic_core.py
@@ -60,10 +60,10 @@ def __init__(
         self.alpha = alpha
         self.gamma = gamma
 
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return self.deflections_2d_via_mge_from(grid=grid, **kwargs)
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray):
+    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np):
         """
         Calculate the intensity of the cored-Sersic light profile on a grid of radial coordinates.
 
@@ -72,27 +72,27 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray):
         grid_radii
             The radial distance from the centre of the profile. for each coordinate on the grid.
         """
-        return jnp.multiply(
-            jnp.multiply(
-                self.intensity_prime,
-                jnp.power(
-                    jnp.add(
+        return xp.multiply(
+            xp.multiply(
+                self.intensity_prime(xp),
+                xp.power(
+                    xp.add(
                         1,
-                        jnp.power(
-                            jnp.divide(self.radius_break, grid_radii.array), self.alpha
+                        xp.power(
+                            xp.divide(self.radius_break, grid_radii.array), self.alpha
                         ),
                     ),
                     (self.gamma / self.alpha),
                 ),
             ),
-            jnp.exp(
-                jnp.multiply(
+            xp.exp(
+                xp.multiply(
                     -self.sersic_constant,
                     (
-                        jnp.power(
-                            jnp.divide(
-                                jnp.add(
-                                    jnp.power(grid_radii.array, self.alpha),
+                        xp.power(
+                            xp.divide(
+                                xp.add(
+                                    xp.power(grid_radii.array, self.alpha),
                                     (self.radius_break**self.alpha),
                                 ),
                                 (self.effective_radius**self.alpha),
@@ -111,7 +111,7 @@ def decompose_convergence_via_mge(self):
         def core_sersic_2D(r):
             return (
                 self.mass_to_light_ratio
-                * self.intensity_prime
+                * self.intensity_prime(xp)
                 * (1.0 + (self.radius_break / r) ** self.alpha)
                 ** (self.gamma / self.alpha)
                 * np.exp(
@@ -128,13 +128,12 @@ def core_sersic_2D(r):
             func=core_sersic_2D, radii_min=radii_min, radii_max=radii_max
         )
 
-    @property
-    def intensity_prime(self):
+    def intensity_prime(self, xp=np):
         """Overall intensity normalisation in the rescaled Core-Sersic light profiles (electrons per second)"""
         return (
             self.intensity
             * (2.0 ** (-self.gamma / self.alpha))
-            * jnp.exp(
+            * xp.exp(
                 self.sersic_constant
                 * (
                     ((2.0 ** (1.0 / self.alpha)) * self.radius_break)
diff --git a/autogalaxy/profiles/mass/stellar/sersic_gradient.py b/autogalaxy/profiles/mass/stellar/sersic_gradient.py
index 0b7230305..eaf2284c8 100644
--- a/autogalaxy/profiles/mass/stellar/sersic_gradient.py
+++ b/autogalaxy/profiles/mass/stellar/sersic_gradient.py
@@ -51,7 +51,7 @@ def __init__(
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -125,7 +125,7 @@ def deflection_func(
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """Calculate the projected convergence at a given set of arc-second gridded coordinates.
 
         Parameters
@@ -135,7 +135,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
         """
         return self.convergence_func(
-            self.eccentric_radii_grid_from(grid=grid, **kwargs)
+            self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs)
         )
 
     def convergence_func(self, grid_radius: float) -> float:
diff --git a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py
index 42a44df1f..acca347d7 100644
--- a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py
+++ b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py
@@ -9,7 +9,7 @@
 # The dPIEPotential and dPIEPotentialSph profiles are modified from the original `dPIEPotential` and `dPIEPotentialSph`, which were implemented to PyAutoLens by Jackson O'Donnell.
 
 
-def _ci05(x, y, eps, rcore):
+def _ci05(x, y, eps, rcore, xp=np):
     """
     Returns the first derivatives of the lens potential as complex number I'* = (∂ψ/∂x + i ∂ψ/∂y) / E0 for PIEMass at any positions (x,y),
     see Kassiola & Kovner(1993) Eq. 4.1.2, which is the integral of Eq. 2.3.8.
@@ -26,7 +26,7 @@ def _ci05(x, y, eps, rcore):
     complex
         The value of the I'* term.
     """
-    sqe = jnp.sqrt(eps)
+    sqe = xp.sqrt(eps)
     axis_ratio = (1.0 - eps) / (1.0 + eps)
     cxro = (1.0 + eps) * (1.0 + eps)
     cyro = (1.0 - eps) * (1.0 - eps)
@@ -34,24 +34,24 @@ def _ci05(x, y, eps, rcore):
     ##### I'* = zres = zci * ln(zis) = zci * ln(znum / zden), see Eq. 4.1.2 #####
 
     # Define intermediate complex variables
-    zci = jnp.array(0.0 + 1j * (-0.5 * (1.0 - eps * eps) / sqe), dtype=jnp.complex128)
-    znum = jnp.complex128(
+    zci = xp.array(0.0 + 1j * (-0.5 * (1.0 - eps * eps) / sqe), dtype=xp.complex128)
+    znum = xp.complex128(
         axis_ratio * x
-        + 1j * (2.0 * sqe * jnp.sqrt(rcore * rcore + rem2) - y / axis_ratio)
+        + 1j * (2.0 * sqe * xp.sqrt(rcore * rcore + rem2) - y / axis_ratio)
     )
-    zden = jnp.complex128(x + 1j * (2.0 * rcore * sqe - y))
+    zden = xp.complex128(x + 1j * (2.0 * rcore * sqe - y))
 
     # zis = znum / zden = (a+bi)/(c+di) = [(ac+bd)+(bc-ad i)] / (c^2+d^2)
     norm = zden.real * zden.real + zden.imag * zden.imag  # |zden|^2
     zis_re = (znum.real * zden.real + znum.imag * zden.imag) / norm
     zis_im = (znum.imag * zden.real - znum.real * zden.imag) / norm
-    zis = jnp.complex128(zis_re + 1j * zis_im)
+    zis = xp.complex128(zis_re + 1j * zis_im)
 
     # ln(zis) = ln(|zis|) + i*Arg(zis)
-    zis_mag = jnp.abs(zis)
-    zis_re = jnp.log(zis_mag)
-    zis_im = jnp.angle(zis)
-    zis = jnp.complex128(zis_re + 1j * zis_im)
+    zis_mag = xp.abs(zis)
+    zis_re = xp.log(zis_mag)
+    zis_im = xp.angle(zis)
+    zis = xp.complex128(zis_re + 1j * zis_im)
 
     # I'* = zres = zci * ln(zis)
     zres = zci * zis
@@ -59,7 +59,7 @@ def _ci05(x, y, eps, rcore):
     return zres
 
 
-def _ci05f(x, y, eps, rcore, rcut):
+def _ci05f(x, y, eps, rcore, rcut, xp=np):
     """
     Returns the first derivatives of the lens potential as complex number I'* = (∂ψ/∂x + i ∂ψ/∂y) / (b0 * ra / (rs - ra)) for dPIEMass at any positions (x,y),
     which is the integral of Eq. 2.3.8 in  Kassiola & Kovner(1993).
@@ -81,7 +81,7 @@ def _ci05f(x, y, eps, rcore, rcut):
     complex
         The value of the I'* term.
     """
-    sqe = jnp.sqrt(eps)
+    sqe = xp.sqrt(eps)
     axis_ratio = (1.0 - eps) / (1.0 + eps)
     cxro = (1.0 + eps) * (1.0 + eps)
     cyro = (1.0 - eps) * (1.0 - eps)
@@ -90,17 +90,17 @@ def _ci05f(x, y, eps, rcore, rcut):
     ##### I'* = zres_rc - zres_rcut = zci * ln(zis_rc / zis_rcut) = zci * ln((znum_rc / zden_rc) / (znum_rcut / zden_rcut)) #####
 
     # Define intermediate complex variables
-    zci = jnp.array(0.0 + 1j * (-0.5 * (1.0 - eps * eps) / sqe), dtype=jnp.complex128)
-    znum_rc = jnp.complex128(
+    zci = xp.array(0.0 + 1j * (-0.5 * (1.0 - eps * eps) / sqe), dtype=xp.complex128)
+    znum_rc = xp.complex128(
         axis_ratio * x
-        + 1j * (2.0 * sqe * jnp.sqrt(rcore * rcore + rem2) - y / axis_ratio)
+        + 1j * (2.0 * sqe * xp.sqrt(rcore * rcore + rem2) - y / axis_ratio)
     )  # a + bi
-    zden_rc = jnp.complex128(x + 1j * (2.0 * rcore * sqe - y))  # c + di
-    znum_rcut = jnp.complex128(
+    zden_rc = xp.complex128(x + 1j * (2.0 * rcore * sqe - y))  # c + di
+    znum_rcut = xp.complex128(
         axis_ratio * x
-        + 1j * (2.0 * sqe * jnp.sqrt(rcut * rcut + rem2) - y / axis_ratio)
+        + 1j * (2.0 * sqe * xp.sqrt(rcut * rcut + rem2) - y / axis_ratio)
     )  # a + ei
-    zden_rcut = jnp.complex128(x + 1j * (2.0 * rcut * sqe - y))  # c + fi
+    zden_rcut = xp.complex128(x + 1j * (2.0 * rcut * sqe - y))  # c + fi
 
     # zis_rc = znum_rc / zden_rc = (a+bi)/(c+di)
     # zis_rcut = znum_rcut / zden_rcut = (a+ei)/(c+fi)
@@ -117,13 +117,13 @@ def _ci05f(x, y, eps, rcore, rcut):
     norm = cc * cc + dd * dd
     aaa = (aa * cc + bb * dd) / norm
     bbb = (bb * cc - aa * dd) / norm
-    zis_tot = jnp.complex128(aaa + 1j * bbb)
+    zis_tot = xp.complex128(aaa + 1j * bbb)
 
     # ln(zis_tot) = ln(|zis_tot|) + i*Arg(zis_tot)
-    zis_tot_mag = jnp.abs(zis_tot)
-    zr_re = jnp.log(zis_tot_mag)
-    zr_im = jnp.angle(zis_tot)
-    zr = jnp.complex128(zr_re + 1j * zr_im)
+    zis_tot_mag = xp.abs(zis_tot)
+    zr_re = xp.log(zis_tot_mag)
+    zr_im = xp.angle(zis_tot)
+    zr = xp.complex128(zr_re + 1j * zr_im)
 
     # I'* = zci * ln(zis_tot)
     zres = zci * zr
@@ -131,7 +131,7 @@ def _ci05f(x, y, eps, rcore, rcut):
     return zres
 
 
-def _mdci05(x, y, eps, rcore, b0):
+def _mdci05(x, y, eps, rcore, b0, xp=np):
     """
     Returns the second derivatives (Hessian matrix) of the lens potential as complex number for PIEMass at any positions (x,y):
     ∂²ψ/∂x² = Re(∂I*/∂x), ∂²ψ/∂y² = Im(∂I*/∂y), ∂²ψ/∂x∂y = ∂²ψ/∂y∂x = Im(∂I*/∂x) = Re(∂I*/∂y)
@@ -153,13 +153,13 @@ def _mdci05(x, y, eps, rcore, b0):
     # I*(x,y) = b0 * ci * (-i) * (ln{ q * x + (2.0 * sqe * wrem - y * 1/q )*i} - ln{ x + (2.0 * rcore * sqe - y)*i})
     #         = b0 * ci * (-i) * (ln{ q * x + num1*i} - ln{ x + num2*i})
     #         = b0 * ci * (-i) * (ln{u(x,y)} - ln{v(x,y)})
-    sqe = jnp.sqrt(eps)
+    sqe = xp.sqrt(eps)
     axis_ratio = (1.0 - eps) / (1.0 + eps)
     axis_ratio_inv = 1.0 / axis_ratio
     cxro = (1.0 + eps) * (1.0 + eps)
     cyro = (1.0 - eps) * (1.0 - eps)
     ci = 0.5 * (1.0 - eps * eps) / sqe
-    wrem = jnp.sqrt(rcore * rcore + x * x / cxro + y * y / cyro)  # √(w(x,y))
+    wrem = xp.sqrt(rcore * rcore + x * x / cxro + y * y / cyro)  # √(w(x,y))
     num1 = 2.0 * sqe * wrem - y * axis_ratio_inv
     den1 = axis_ratio * axis_ratio * x * x + num1 * num1  # |q * x + num1*i|^2
     num2 = 2.0 * rcore * sqe - y
@@ -259,14 +259,14 @@ def __init__(
         self.ra = ra
         self.b0 = b0
 
-    def _ellip(self):
-        ellip = jnp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2)
+    def _ellip(self, xp=np):
+        ellip = xp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2)
         MAX_ELLIP = 0.99999
-        return jnp.min(jnp.array([ellip, MAX_ELLIP]))
+        return xp.min(xp.array([ellip, MAX_ELLIP]))
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
 
@@ -275,7 +275,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         grid
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
-        ellip = self._ellip()
+        ellip = self._ellip(xp)
         factor = self.b0
         zis = _ci05(x=grid.array[:, 1], y=grid.array[:, 0], eps=ellip, rcore=self.ra)
 
@@ -285,12 +285,13 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
         # And here we convert back to the real axes
         return self.rotated_grid_from_reference_frame_from(
-            grid=jnp.multiply(factor, jnp.vstack((deflection_y, deflection_x)).T),
+            grid=xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T),
+            xp=xp,
             **kwargs,
         )
 
     @aa.grid_dec.transform
-    def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs):
+    def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs):
         """
         Calculate the hessian matrix on a grid of (y,x) arc-second coordinates.
 
@@ -299,7 +300,7 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs):
         grid
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
-        grid = jnp.asarray(grid)
+        grid = xp.asarray(grid)
         if grid.ndim != 2 or grid.shape[1] != 2:
             raise ValueError("Grid must be a 2D array with shape (n, 2)")
         ellip = self._ellip()
@@ -310,10 +311,10 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs):
 
         return hessian_yy, hessian_xy, hessian_yx, hessian_xx
 
-    def analytical_magnification_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs):
+    def analytical_magnification_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs):
 
         hessian_yy, hessian_xy, hessian_yx, hessian_xx = (
-            self.analytical_hessian_2d_from(grid=grid)
+            self.analytical_hessian_2d_from(grid=grid, xp=np)
         )
 
         det_A = (1 - hessian_xx) * (1 - hessian_yy) - hessian_xy * hessian_yx
@@ -372,14 +373,14 @@ def __init__(
         self.rs = rs
         self.b0 = b0
 
-    def _ellip(self):
-        ellip = jnp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2)
+    def _ellip(self, xp=np):
+        ellip = xp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2)
         MAX_ELLIP = 0.99999
-        return jnp.min(jnp.array([ellip, MAX_ELLIP]))
+        return xp.min(xp.array([ellip, MAX_ELLIP]))
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
 
@@ -388,7 +389,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         grid
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
-        ellip = self._ellip()
+        ellip = self._ellip(xp)
         factor = self.b0 * self.rs / (self.rs - self.ra)
         zis = _ci05f(
             x=grid.array[:, 1],
@@ -404,18 +405,19 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
         # And here we convert back to the real axes
         return self.rotated_grid_from_reference_frame_from(
-            grid=jnp.multiply(factor, jnp.vstack((deflection_y, deflection_x)).T),
+            grid=xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T),
+            xp=xp,
             **kwargs,
         )
 
-    def _deflection_angle(self, radii):
+    def _deflection_angle(self, radii, xp=np):
         """
         For a circularly symmetric dPIEPotential profile, computes the magnitude of the deflection at each radius.
         """
         a, s = self.ra, self.rs
-        radii = jnp.maximum(radii, 1e-8)
-        f = radii / (a + jnp.sqrt(a**2 + radii**2)) - radii / (
-            s + jnp.sqrt(s**2 + radii**2)
+        radii = xp.maximum(radii, 1e-8)
+        f = radii / (a + xp.sqrt(a**2 + radii**2)) - radii / (
+            s + xp.sqrt(s**2 + radii**2)
         )
 
         # c.f. Eliasdottir '07 eq. A23
@@ -424,7 +426,7 @@ def _deflection_angle(self, radii):
         alpha = self.b0 * s / (s - a) * f
         return alpha
 
-    def _convergence(self, radii):
+    def _convergence(self, radii, xp=np):
 
         radsq = radii * radii
         a, s = self.ra, self.rs
@@ -434,12 +436,12 @@ def _convergence(self, radii):
             / 2
             * s
             / (s - a)
-            * (1 / jnp.sqrt(a**2 + radsq) - 1 / jnp.sqrt(s**2 + radsq))
+            * (1 / xp.sqrt(a**2 + radsq) - 1 / xp.sqrt(s**2 + radsq))
         )
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates.
 
@@ -451,14 +453,14 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         grid
             The grid of (y,x) arc-second coordinates the convergence is computed on.
         """
-        ellip = self._ellip()
-        grid_radii = jnp.sqrt(
+        ellip = self._ellip(xp)
+        grid_radii = xp.sqrt(
             grid.array[:, 1] ** 2 * (1 - ellip) + grid.array[:, 0] ** 2 * (1 + ellip)
         )
 
         # Compute the convergence and deflection of a *circular* profile
-        kappa_circ = self._convergence(grid_radii)
-        alpha_circ = self._deflection_angle(grid_radii)
+        kappa_circ = self._convergence(grid_radii, xp)
+        alpha_circ = self._deflection_angle(grid_radii, xp)
 
         asymm_term = (
             ellip * (1 - ellip) * grid.array[:, 1] ** 2
@@ -471,7 +473,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         return kappa_circ * (1 - asymm_term) + (alpha_circ / grid_radii) * asymm_term
 
     @aa.grid_dec.transform
-    def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs):
+    def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs):
         """
         Calculate the hessian matrix on a grid of (y,x) arc-second coordinates.
         Hessian_dPIEMass = rs * (rs - ra) * ( Hessian_PIEMass(ra) - Hessian_PIEMass(rs))
@@ -481,7 +483,7 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs):
         grid
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
-        grid = jnp.asarray(grid)
+        grid = xp.asarray(grid)
         if grid.ndim != 2 or grid.shape[1] != 2:
             raise ValueError("Grid must be a 2D array with shape (n, 2)")
         ellip = self._ellip()
@@ -502,10 +504,10 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs):
 
         return hessian_yy, hessian_xy, hessian_yx, hessian_xx
 
-    def analytical_magnification_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs):
+    def analytical_magnification_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs):
 
         hessian_yy, hessian_xy, hessian_yx, hessian_xx = (
-            self.analytical_hessian_2d_from(grid=grid)
+            self.analytical_hessian_2d_from(grid=grid, xp=xp)
         )
 
         det_A = (1 - hessian_xx) * (1 - hessian_yy) - hessian_xy * hessian_yx
@@ -513,8 +515,8 @@ def analytical_magnification_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs)
         return aa.Array2D(values=1.0 / det_A, mask=grid.mask)
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
-        return jnp.zeros(shape=grid.shape[0])
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return xp.zeros(shape=grid.shape[0])
 
 
 class dPIEMassSph(dPIEMass):
@@ -568,7 +570,7 @@ def __init__(
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
         Faster and equivalent to Eliasdottir (2007), see Eq. A19 and Eq. A20.
@@ -593,7 +595,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         # radii = self.radial_grid_from(grid=grid, **kwargs)
         # R2 = radii * radii
         R2 = grid.array[:, 1] * grid.array[:, 1] + grid.array[:, 0] * grid.array[:, 0]
-        factor = jnp.sqrt(R2 + a * a) - a - jnp.sqrt(R2 + s * s) + s
+        factor = xp.sqrt(R2 + a * a) - a - xp.sqrt(R2 + s * s) + s
         factor *= self.b0 * s / (s - a) / R2
 
         # This is in axes aligned to the major/minor axis
@@ -602,11 +604,11 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
         # And here we convert back to the real axes
         return self.rotated_grid_from_reference_frame_from(
-            grid=jnp.multiply(1.0, jnp.vstack((deflection_y, deflection_x)).T), **kwargs
+            grid=xp.multiply(1.0, xp.vstack((deflection_y, deflection_x)).T), xp=xp, **kwargs
         )
 
     @aa.grid_dec.transform
-    def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs):
+    def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs):
         """
         Calculate the hessian matrix on a grid of (y,x) arc-second coordinates.
         Chain rule of second derivatives:
@@ -620,7 +622,7 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs):
         grid
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
-        grid = jnp.asarray(grid)
+        grid = xp.asarray(grid)
         if grid.ndim != 2 or grid.shape[1] != 2:
             raise ValueError("Grid must be a 2D array with shape (n, 2)")
 
@@ -646,9 +648,9 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs):
         #                     = {( R^2 / √(R^2 + a^2)) - ( R^2 / √(R^2 + s^2)) - z} / R^2
         #                     = p
         R2 = grid.array[:, 1] * grid.array[:, 1] + grid.array[:, 0] * grid.array[:, 0]
-        z = jnp.sqrt(R2 + a * a) - a - jnp.sqrt(R2 + s * s) + s
-        p = (1.0 - a / jnp.sqrt(a * a + R2)) * a / R2 - (
-            1.0 - s / jnp.sqrt(s * s + R2)
+        z = xp.sqrt(R2 + a * a) - a - xp.sqrt(R2 + s * s) + s
+        p = (1.0 - a / xp.sqrt(a * a + R2)) * a / R2 - (
+            1.0 - s / xp.sqrt(s * s + R2)
         ) * s / R2
         X = grid.array[:, 1] * grid.array[:, 1] / R2  # x^2 / R^2
         Y = grid.array[:, 0] * grid.array[:, 0] / R2  # y^2 / R^2
diff --git a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py
index d8d02ce52..f140a3433 100644
--- a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py
+++ b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py
@@ -1,5 +1,5 @@
 from typing import Tuple
-import jax.numpy as jnp
+import numpy as np
 
 import autoarray as aa
 from autogalaxy.profiles.mass.abstract.abstract import MassProfile
@@ -59,19 +59,19 @@ def __init__(
         self.rs = rs
         self.b0 = b0
 
-    def _ellip(self):
-        ellip = jnp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2)
+    def _ellip(self, xp=np):
+        ellip = xp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2)
         MAX_ELLIP = 0.99999
-        return jnp.min(jnp.array([ellip, MAX_ELLIP]))
+        return xp.min(xp.array([ellip, MAX_ELLIP]))
 
-    def _deflection_angle(self, radii):
+    def _deflection_angle(self, radii, xp=np):
         """
         For a circularly symmetric dPIEPotential profile, computes the magnitude of the deflection at each radius.
         """
         a, s = self.ra, self.rs
-        radii = jnp.maximum(radii, 1e-8)
-        f = radii / (a + jnp.sqrt(a**2 + radii**2)) - radii / (
-            s + jnp.sqrt(s**2 + radii**2)
+        radii = xp.maximum(radii, 1e-8)
+        f = radii / (a + xp.sqrt(a**2 + radii**2)) - radii / (
+            s + xp.sqrt(s**2 + radii**2)
         )
 
         # c.f. Eliasdottir '07 eq. A23
@@ -80,7 +80,7 @@ def _deflection_angle(self, radii):
         alpha = self.b0 * s / (s - a) * f
         return alpha
 
-    def _convergence(self, radii):
+    def _convergence(self, radii, xp=np):
 
         radsq = radii * radii
         a, s = self.ra, self.rs
@@ -90,12 +90,12 @@ def _convergence(self, radii):
             / 2
             * s
             / (s - a)
-            * (1 / jnp.sqrt(a**2 + radsq) - 1 / jnp.sqrt(s**2 + radsq))
+            * (1 / xp.sqrt(a**2 + radsq) - 1 / xp.sqrt(s**2 + radsq))
         )
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
 
@@ -104,30 +104,30 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         grid
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
-        ellip = self._ellip()
-        grid_radii = jnp.sqrt(
+        ellip = self._ellip(xp)
+        grid_radii = xp.sqrt(
             grid.array[:, 1] ** 2 * (1 - ellip) + grid.array[:, 0] ** 2 * (1 + ellip)
         )
 
         # Compute the deflection magnitude of a *non-elliptical* profile
-        alpha_circ = self._deflection_angle(grid_radii)
+        alpha_circ = self._deflection_angle(grid_radii, xp)
 
         # This is in axes aligned to the major/minor axis
         deflection_y = (
-            alpha_circ * jnp.sqrt(1 + ellip) * (grid.array[:, 0] / grid_radii)
+            alpha_circ * xp.sqrt(1 + ellip) * (grid.array[:, 0] / grid_radii)
         )
         deflection_x = (
-            alpha_circ * jnp.sqrt(1 - ellip) * (grid.array[:, 1] / grid_radii)
+            alpha_circ * xp.sqrt(1 - ellip) * (grid.array[:, 1] / grid_radii)
         )
 
         # And here we convert back to the real axes
         return self.rotated_grid_from_reference_frame_from(
-            grid=jnp.multiply(1.0, jnp.vstack((deflection_y, deflection_x)).T), **kwargs
+            grid=xp.multiply(1.0, xp.vstack((deflection_y, deflection_x)).T), **kwargs
         )
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates.
 
@@ -139,14 +139,14 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         grid
             The grid of (y,x) arc-second coordinates the convergence is computed on.
         """
-        ellip = self._ellip()
-        grid_radii = jnp.sqrt(
+        ellip = self._ellip(xp)
+        grid_radii = xp.sqrt(
             grid.array[:, 1] ** 2 * (1 - ellip) + grid.array[:, 0] ** 2 * (1 + ellip)
         )
 
         # Compute the convergence and deflection of a *circular* profile
-        kappa_circ = self._convergence(grid_radii)
-        alpha_circ = self._deflection_angle(grid_radii)
+        kappa_circ = self._convergence(grid_radii, xp)
+        alpha_circ = self._deflection_angle(grid_radii, xp)
 
         asymm_term = (
             ellip * (1 - ellip) * grid.array[:, 1] ** 2
@@ -159,8 +159,8 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         return kappa_circ * (1 - asymm_term) + (alpha_circ / grid_radii) * asymm_term
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
-        return jnp.zeros(shape=grid.shape[0])
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return xp.zeros(shape=grid.shape[0])
 
 
 class dPIEPotentialSph(dPIEPotential):
@@ -218,7 +218,7 @@ def __init__(
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
 
@@ -229,7 +229,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         """
         radii = self.radial_grid_from(grid=grid, **kwargs)
 
-        alpha = self._deflection_angle(radii.array)
+        alpha = self._deflection_angle(radii.array, xp)
 
         # now we decompose the deflection into y/x components
         defl_y = alpha * grid.array[:, 0] / radii.array
@@ -239,7 +239,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates.
 
@@ -254,8 +254,8 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         # already transformed to center on profile centre so this works
         radsq = grid.array[:, 0] ** 2 + grid.array[:, 1] ** 2
 
-        return self._convergence(jnp.sqrt(radsq))
+        return self._convergence(xp.sqrt(radsq), xp)
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
-        return jnp.zeros(shape=grid.shape[0])
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return xp.zeros(shape=grid.shape[0])
diff --git a/autogalaxy/profiles/mass/total/jax_utils.py b/autogalaxy/profiles/mass/total/jax_utils.py
index 5237f34cd..cb9eeb71a 100644
--- a/autogalaxy/profiles/mass/total/jax_utils.py
+++ b/autogalaxy/profiles/mass/total/jax_utils.py
@@ -1,12 +1,4 @@
-import jax
-import jax.numpy as jnp
-from jax.tree_util import Partial as partial
-
-
-# A version of scan that will *not* re-compile partial functions when variables change
-# taken from https://github.com/google/jax/issues/14743#issuecomment-1456900634
-scan = jax.jit(jax.lax.scan, static_argnames=("length", "reverse", "unroll"))
-
+import numpy as np
 
 def body_fun(carry, n, factor, ei2phi, slope):
     omega_nm1, partial_sum = carry
@@ -18,7 +10,7 @@ def body_fun(carry, n, factor, ei2phi, slope):
     return (omega_n, partial_sum), None
 
 
-def omega(eiphi, slope, factor, n_terms=20):
+def omega(eiphi, slope, factor, n_terms=20, xp=np):
     """JAX implementation of the numerical evaluation of the angular component of
     the complex deflection angle for the elliptical power law profile as given as
     given by Tessore and Metcalf 2015.  Based on equation 29, and gives
@@ -38,10 +30,15 @@ def omega(eiphi, slope, factor, n_terms=20):
         The number of terms to calculate for the series expansion, defaults to 20 (this should
         be sufficient most of the time)
     """
+
+    from jax.tree_util import Partial as partial
+    import jax
+    scan = jax.jit(jax.lax.scan, static_argnames=("length", "reverse", "unroll"))
+
     # use modified scan with a partial'ed function to avoid re-compile
     (_, partial_sum), _ = scan(
         partial(body_fun, factor=factor, ei2phi=eiphi**2, slope=slope),
         (eiphi, eiphi),
-        jnp.arange(1, n_terms),
+        xp.arange(1, n_terms),
     )
     return partial_sum
diff --git a/autogalaxy/profiles/mass/total/power_law.py b/autogalaxy/profiles/mass/total/power_law.py
index c21175958..2c919d6c0 100644
--- a/autogalaxy/profiles/mass/total/power_law.py
+++ b/autogalaxy/profiles/mass/total/power_law.py
@@ -39,7 +39,7 @@ def __init__(
         )
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         alpha = self.deflections_yx_2d_from(aa.Grid2DIrregular(grid), **kwargs)
 
         alpha_x = alpha[:, 1]
@@ -85,7 +85,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         R = xp.sqrt(
             (self.axis_ratio() * grid.array[:, 1]) ** 2 + grid.array[:, 0] ** 2 + 1e-16
         )
-        zh = omega(z, slope, factor, n_terms=20)
+        zh = omega(z, slope, factor, n_terms=20, xp=np)
 
         complex_angle = (
             2.0 * b / (1.0 + self.axis_ratio()) * (b / R) ** (slope - 1.0) * zh
diff --git a/autogalaxy/profiles/mass/total/power_law_broken.py b/autogalaxy/profiles/mass/total/power_law_broken.py
index e8d82ef05..c034ee15b 100644
--- a/autogalaxy/profiles/mass/total/power_law_broken.py
+++ b/autogalaxy/profiles/mass/total/power_law_broken.py
@@ -31,7 +31,7 @@ def __init__(
         super().__init__(centre=centre, ell_comps=ell_comps)
 
         self.einstein_radius = einstein_radius
-        self.einstein_radius_elliptical = jnp.sqrt(self.axis_ratio()) * einstein_radius
+        self.einstein_radius_elliptical = xp.sqrt(self.axis_ratio()) * einstein_radius
         self.break_radius = break_radius
         self.inner_slope = inner_slope
         self.outer_slope = outer_slope
@@ -52,13 +52,13 @@ def __init__(
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Returns the dimensionless density kappa=Sigma/Sigma_c (eq. 1)
         """
 
         # Ell radius
-        radius = jnp.hypot(grid.array[:, 1] * self.axis_ratio(), grid.array[:, 0])
+        radius = xp.hypot(grid.array[:, 1] * self.axis_ratio(), grid.array[:, 0])
 
         # Inside break radius
         kappa_inner = self.kB * (self.break_radius / radius) ** self.inner_slope
@@ -71,8 +71,8 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         )
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
-        return jnp.zeros(shape=grid.shape[0])
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return xp.zeros(shape=grid.shape[0])
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
@@ -84,7 +84,7 @@ def deflections_yx_2d_from(self, grid, max_terms=20, **kwargs):
         z = grid.array[:, 1] + 1j * grid.array[:, 0]
 
         # Ell radius
-        R = jnp.hypot(z.real * self.axis_ratio(), z.imag)
+        R = xp.hypot(z.real * self.axis_ratio(), z.imag)
 
         # Factors common to eq. 18 and 19
         factors = (
@@ -126,8 +126,8 @@ def deflections_yx_2d_from(self, grid, max_terms=20, **kwargs):
         ).conjugate()
 
         return self.rotated_grid_from_reference_frame_from(
-            grid=jnp.multiply(
-                1.0, jnp.vstack((jnp.imag(deflections), jnp.real(deflections))).T
+            grid=xp.multiply(
+                1.0, xp.vstack((xp.imag(deflections), xp.real(deflections))).T
             )
         )
 
@@ -140,13 +140,13 @@ def hyp2f1_series(t, q, r, z, max_terms=20):
 
         # u from eq. 25
         q_ = (1 - q**2) / (q**2)
-        u = 0.5 * (1 - jnp.sqrt(1 - q_ * (r / z) ** 2))
+        u = 0.5 * (1 - xp.sqrt(1 - q_ * (r / z) ** 2))
 
         # First coefficient
         a_n = 1.0
 
         # Storage for sum
-        F = jnp.zeros_like(z, dtype="complex64")
+        F = xp.zeros_like(z, dtype="complex64")
 
         for n in range(max_terms):
             F += a_n * (u**n)
diff --git a/autogalaxy/profiles/mass/total/power_law_core.py b/autogalaxy/profiles/mass/total/power_law_core.py
index 31d5421a5..bd1b5ee20 100644
--- a/autogalaxy/profiles/mass/total/power_law_core.py
+++ b/autogalaxy/profiles/mass/total/power_law_core.py
@@ -50,7 +50,7 @@ def einstein_radius_rescaled(self, xp=np):
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
         Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates.
 
@@ -221,17 +221,17 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
         """
         eta = self.radial_grid_from(grid=grid, **kwargs)
-        deflection = jnp.multiply(
+        deflection = xp.multiply(
             2.0 * self.einstein_radius_rescaled(xp),
-            jnp.divide(
-                jnp.add(
-                    jnp.power(
-                        jnp.add(self.core_radius**2, jnp.square(eta.array)),
+            xp.divide(
+                xp.add(
+                    xp.power(
+                        xp.add(self.core_radius**2, xp.square(eta.array)),
                         (3.0 - self.slope) / 2.0,
                     ),
                     -self.core_radius ** (3 - self.slope),
                 ),
-                jnp.multiply((3.0 - self.slope), eta.array),
+                xp.multiply((3.0 - self.slope), eta.array),
             ),
         )
         return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection)
diff --git a/autogalaxy/profiles/mass/total/power_law_multipole.py b/autogalaxy/profiles/mass/total/power_law_multipole.py
index 16e5a2a88..40f562979 100644
--- a/autogalaxy/profiles/mass/total/power_law_multipole.py
+++ b/autogalaxy/profiles/mass/total/power_law_multipole.py
@@ -28,12 +28,12 @@ def radial_and_angle_grid_from(
     """
     y, x = grid.array.T
 
-    x_shifted = jnp.subtract(x, centre[1])
-    y_shifted = jnp.subtract(y, centre[0])
+    x_shifted = xp.subtract(x, centre[1])
+    y_shifted = xp.subtract(y, centre[0])
 
-    radial_grid = jnp.sqrt(x_shifted**2 + y_shifted**2)
+    radial_grid = xp.sqrt(x_shifted**2 + y_shifted**2)
 
-    angle_grid = jnp.arctan2(y_shifted, x_shifted)
+    angle_grid = xp.arctan2(y_shifted, x_shifted)
 
     return radial_grid, angle_grid
 
@@ -144,8 +144,8 @@ def jacobian(
             The polar angle coordinates of the input (y,x) Cartesian grid of coordinates.
         """
         return (
-            a_r * jnp.sin(polar_angle_grid) + a_angle * jnp.cos(polar_angle_grid),
-            a_r * jnp.cos(polar_angle_grid) - a_angle * jnp.sin(polar_angle_grid),
+            a_r * xp.sin(polar_angle_grid) + a_angle * xp.cos(polar_angle_grid),
+            a_r * xp.cos(polar_angle_grid) - a_angle * xp.sin(polar_angle_grid),
         )
 
     @aa.grid_dec.to_vector_yx
@@ -174,7 +174,7 @@ def deflections_yx_2d_from(
             )
             / (self.m**2.0 - (3.0 - self.slope) ** 2.0)
             * self.k_m
-            * jnp.cos(self.m * (polar_angle_grid - self.angle_m))
+            * xp.cos(self.m * (polar_angle_grid - self.angle_m))
         )
 
         a_angle = (
@@ -185,10 +185,10 @@ def deflections_yx_2d_from(
             )
             / (self.m**2.0 - (3.0 - self.slope) ** 2.0)
             * self.k_m
-            * jnp.sin(self.m * (polar_angle_grid - self.angle_m))
+            * xp.sin(self.m * (polar_angle_grid - self.angle_m))
         )
 
-        return jnp.stack(
+        return xp.stack(
             self.jacobian(a_r=a_r, a_angle=a_angle, polar_angle_grid=polar_angle_grid),
             axis=-1,
         )
@@ -212,11 +212,11 @@ def convergence_2d_from(self, grid: aa.type.Grid1D2DLike, **kwargs) -> np.ndarra
             / 2.0
             * (self.einstein_radius / r) ** (self.slope - 1)
             * self.k_m
-            * jnp.cos(self.m * (angle - self.angle_m))
+            * xp.cos(self.m * (angle - self.angle_m))
         )
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
         """
         Calculate the potential on a grid of (y,x) arc-second coordinates.
 
@@ -225,4 +225,4 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
         grid
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
-        return jnp.zeros(shape=grid.shape[0])
+        return xp.zeros(shape=grid.shape[0])
diff --git a/test_autogalaxy/profiles/light/standard/test_abstract.py b/test_autogalaxy/profiles/light/standard/test_abstract.py
index 8281e187c..058c03677 100644
--- a/test_autogalaxy/profiles/light/standard/test_abstract.py
+++ b/test_autogalaxy/profiles/light/standard/test_abstract.py
@@ -43,40 +43,6 @@ def test__luminosity_within_centre__compare_to_gridded_calculations():
     assert luminosity_grid == pytest.approx(luminosity_integral, 0.02)
 
 
-def test__image_1d_from__grid_2d_in__returns_1d_image_via_projected_quantities():
-    grid_2d = ag.Grid2D.uniform(
-        shape_native=(5, 5),
-        pixel_scales=1.0,
-        over_sample_size=1,
-    )
-
-    lp = ag.lp.Gaussian(
-        centre=(0.0, 0.0), ell_comps=(0.0, 0.0), intensity=1.0, sigma=1.0
-    )
-
-    image_1d = lp.image_1d_from(grid=grid_2d)
-    image_2d = lp.image_2d_from(grid=grid_2d)
-
-    assert image_1d[0] == pytest.approx(image_2d.native.array[2, 2], 1.0e-4)
-    assert image_1d[1] == pytest.approx(image_2d.native.array[2, 3], 1.0e-4)
-    assert image_1d[2] == pytest.approx(image_2d.native.array[2, 4], 1.0e-4)
-
-    lp = ag.lp.Gaussian(
-        centre=(0.2, 0.2), ell_comps=(0.3, 0.3), intensity=1.0, sigma=1.0
-    )
-
-    image_1d = lp.image_1d_from(grid=grid_2d)
-
-    grid_2d_projected = grid_2d.grid_2d_radial_projected_from(
-        centre=lp.centre, angle=lp.angle + 90.0
-    )
-
-    image_projected = lp.image_2d_from(grid=grid_2d_projected)
-
-    assert image_1d == pytest.approx(image_projected.array, 1.0e-4)
-    assert (image_1d.grid_radial == np.array([0.0, 1.0, 2.0])).all()
-
-
 def test__decorator__oversample_uniform__numerical_values(gal_x1_lp):
     mask = ag.Mask2D(
         mask=[
diff --git a/test_autogalaxy/profiles/light/test_decorators.py b/test_autogalaxy/profiles/light/test_decorators.py
index e7916359a..cf27648d1 100644
--- a/test_autogalaxy/profiles/light/test_decorators.py
+++ b/test_autogalaxy/profiles/light/test_decorators.py
@@ -13,7 +13,7 @@
 class MockLightProfile(ag.LightProfile):
     @check_operated_only
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
     ):
         return np.ones(shape=(3, 3))
 
@@ -21,7 +21,7 @@ def image_2d_from(
 class MockLightProfileOperated(ag.lp_operated.LightProfileOperated):
     @check_operated_only
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
     ):
         return np.ones(shape=(3, 3))
 

From ea1198275d58dd79ea25f8ede55548c0e5e71179 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 8 Nov 2025 15:11:31 +0000
Subject: [PATCH 07/18] mass tests pass

---
 autogalaxy/profiles/geometry_profiles.py      |  8 +-
 .../profiles/light/standard/chameleon.py      |  4 +-
 .../profiles/light/standard/gaussian.py       |  2 +-
 autogalaxy/profiles/light/standard/moffat.py  |  2 +-
 autogalaxy/profiles/light/standard/sersic.py  |  2 +-
 autogalaxy/profiles/mass/abstract/abstract.py |  2 +-
 autogalaxy/profiles/mass/abstract/cse.py      |  4 +-
 autogalaxy/profiles/mass/abstract/mge_jax.py  |  2 +-
 .../mass/abstract/mge_numpy_coleman.py        |  2 +-
 autogalaxy/profiles/mass/dark/abstract.py     | 19 +++--
 autogalaxy/profiles/mass/dark/gnfw.py         | 18 ++--
 autogalaxy/profiles/mass/dark/nfw.py          | 16 ++--
 .../profiles/mass/dark/nfw_truncated.py       |  2 +
 .../profiles/mass/mock/mock_mass_profile.py   |  4 +-
 autogalaxy/profiles/mass/point/point.py       |  2 +-
 autogalaxy/profiles/mass/point/smbh_binary.py |  4 +-
 autogalaxy/profiles/mass/sheets/mass_sheet.py |  2 +-
 autogalaxy/profiles/mass/stellar/chameleon.py |  3 +-
 autogalaxy/profiles/mass/stellar/gaussian.py  | 32 +++----
 autogalaxy/profiles/mass/stellar/sersic.py    |  7 +-
 .../total/dual_pseudo_isothermal_potential.py |  2 +-
 autogalaxy/profiles/mass/total/isothermal.py  |  2 +-
 autogalaxy/profiles/mass/total/power_law.py   | 15 ++--
 .../profiles/mass/total/power_law_broken.py   | 25 +++---
 .../profiles/mass/total/power_law_core.py     | 13 +--
 .../mass/total/power_law_multipole.py         | 14 ++--
 .../profiles/mass/abstract/test_abstract.py   | 83 +------------------
 .../profiles/mass/dark/test_nfw_mcr.py        | 16 ++--
 .../profiles/mass/dark/test_nfw_scatter.py    |  2 +-
 .../mass/dark/test_nfw_truncated_mcr.py       |  8 +-
 .../mass/total/test_power_law_broken.py       |  4 +-
 31 files changed, 128 insertions(+), 193 deletions(-)

diff --git a/autogalaxy/profiles/geometry_profiles.py b/autogalaxy/profiles/geometry_profiles.py
index 49c6e9d6f..d91fa3fae 100644
--- a/autogalaxy/profiles/geometry_profiles.py
+++ b/autogalaxy/profiles/geometry_profiles.py
@@ -114,7 +114,7 @@ def angle_to_profile_grid_from(
 
     @aa.grid_dec.to_grid
     def _cartesian_grid_via_radial_from(
-        self, grid: aa.type.Grid2DLike, radius: np.ndarray, xp=np, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, radius : Optional[np.ndarray] = None,  **kwargs
     ) -> aa.type.Grid2DLike:
         """
         Convert a grid of (y,x) coordinates with their specified radial distances (e.g. :math: r = x**2 + y**2) to
@@ -309,7 +309,7 @@ def elliptical_radii_grid_from(
         return xp.sqrt(
             xp.add(
                 xp.square(grid.array[:, 1]),
-                xp.square(xp.divide(grid.array[:, 0], self.axis_ratio())),
+                xp.square(xp.divide(grid.array[:, 0], self.axis_ratio(xp))),
             )
         )
 
@@ -333,7 +333,7 @@ def eccentric_radii_grid_from(
 
         grid_radii = self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs)
 
-        return xp.multiply(xp.sqrt(self.axis_ratio()), grid_radii.array)
+        return xp.multiply(xp.sqrt(self.axis_ratio(xp)), grid_radii.array)
 
     @aa.grid_dec.to_grid
     def transformed_to_reference_frame_grid_from(
@@ -383,7 +383,7 @@ def _eta_u(self, u, coordinates):
                 u
                 * (
                     (coordinates[1] ** 2)
-                    + (coordinates[0] ** 2 / (1 - (1 - self.axis_ratio()**2) * u))
+                    + (coordinates[0] ** 2 / (1 - (1 - self.axis_ratio(xp)**2) * u))
                 )
             )
         )
diff --git a/autogalaxy/profiles/light/standard/chameleon.py b/autogalaxy/profiles/light/standard/chameleon.py
index 90296a277..1b02f4938 100644
--- a/autogalaxy/profiles/light/standard/chameleon.py
+++ b/autogalaxy/profiles/light/standard/chameleon.py
@@ -64,10 +64,10 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
             The radial distances from the centre of the profile, for each coordinate on the grid.
         """
 
-        axis_ratio_factor = (1.0 + self.axis_ratio()) ** 2.0
+        axis_ratio_factor = (1.0 + self.axis_ratio(xp)) ** 2.0
 
         return xp.multiply(
-            self._intensity / (1 + self.axis_ratio()),
+            self._intensity / (1 + self.axis_ratio(xp)),
             xp.add(
                 xp.divide(
                     1.0,
diff --git a/autogalaxy/profiles/light/standard/gaussian.py b/autogalaxy/profiles/light/standard/gaussian.py
index 8b626f96e..6774b2900 100644
--- a/autogalaxy/profiles/light/standard/gaussian.py
+++ b/autogalaxy/profiles/light/standard/gaussian.py
@@ -64,7 +64,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
             xp.exp(
                 -0.5
                 * xp.square(
-                    xp.divide(grid_radii.array, self.sigma / xp.sqrt(self.axis_ratio()))
+                    xp.divide(grid_radii.array, self.sigma / xp.sqrt(self.axis_ratio(xp)))
                 )
             ),
         )
diff --git a/autogalaxy/profiles/light/standard/moffat.py b/autogalaxy/profiles/light/standard/moffat.py
index b924331f8..3ddb1405e 100644
--- a/autogalaxy/profiles/light/standard/moffat.py
+++ b/autogalaxy/profiles/light/standard/moffat.py
@@ -63,7 +63,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
             xp.power(
                 1
                 + xp.square(
-                    xp.divide(grid_radii.array, self.alpha / xp.sqrt(self.axis_ratio()))
+                    xp.divide(grid_radii.array, self.alpha / xp.sqrt(self.axis_ratio(xp)))
                 ),
                 -self.beta,
             ),
diff --git a/autogalaxy/profiles/light/standard/sersic.py b/autogalaxy/profiles/light/standard/sersic.py
index 46c2e788a..1d8ec3558 100644
--- a/autogalaxy/profiles/light/standard/sersic.py
+++ b/autogalaxy/profiles/light/standard/sersic.py
@@ -53,7 +53,7 @@ def elliptical_effective_radius(self) -> float:
         The elliptical effective radius instead describes the major-axis radius of the ellipse containing
         half the light, and may be more appropriate for highly flattened systems like disk galaxies.
         """
-        return self.effective_radius / xp.sqrt(self.axis_ratio())
+        return self.effective_radius / xp.sqrt(self.axis_ratio(xp))
 
     @property
     def sersic_constant(self) -> float:
diff --git a/autogalaxy/profiles/mass/abstract/abstract.py b/autogalaxy/profiles/mass/abstract/abstract.py
index 8bedfb61d..883264373 100644
--- a/autogalaxy/profiles/mass/abstract/abstract.py
+++ b/autogalaxy/profiles/mass/abstract/abstract.py
@@ -44,7 +44,7 @@ def deflections_2d_via_potential_2d_from(self, grid):
             mask=grid.mask,
         )
 
-    def convergence_2d_from(self, grid):
+    def convergence_2d_from(self, grid, xp=np):
         raise NotImplementedError
 
     def convergence_func(self, grid_radius: float) -> float:
diff --git a/autogalaxy/profiles/mass/abstract/cse.py b/autogalaxy/profiles/mass/abstract/cse.py
index 74a08ee1c..cd152ae1d 100644
--- a/autogalaxy/profiles/mass/abstract/cse.py
+++ b/autogalaxy/profiles/mass/abstract/cse.py
@@ -41,12 +41,12 @@ def deflections_via_cse_from(
         Parameters
         ----------
         """
-        phi = xp.sqrt(axis_ratio_squared * core_radius**2.0 + term1)
+        phi = np.sqrt(axis_ratio_squared * core_radius**2.0 + term1)
         Psi = (phi + core_radius) ** 2.0 + term2
         bottom = core_radius * phi * Psi
         defl_x = (term3 * (phi + axis_ratio_squared * core_radius)) / bottom
         defl_y = (term4 * (phi + core_radius)) / bottom
-        return xp.vstack((defl_y, defl_x))
+        return np.vstack((defl_y, defl_x))
 
     @abstractmethod
     def decompose_convergence_via_cse(self, grid_radii: np.ndarray):
diff --git a/autogalaxy/profiles/mass/abstract/mge_jax.py b/autogalaxy/profiles/mass/abstract/mge_jax.py
index ae4010161..944029aa5 100644
--- a/autogalaxy/profiles/mass/abstract/mge_jax.py
+++ b/autogalaxy/profiles/mass/abstract/mge_jax.py
@@ -143,7 +143,7 @@ def _convergence_2d_via_mge_from(
     def _deflections_2d_via_mge_from(
         self, grid, sigmas_factor=1.0, func_terms=28, func_gaussians=20
     ):
-        axis_ratio = xp.min(xp.array([self.axis_ratio(), 0.9999]))
+        axis_ratio = xp.min(xp.array([self.axis_ratio(xp), 0.9999]))
 
         amps, sigmas = self.decompose_convergence_via_mge(
             func_terms=func_terms, func_gaussians=func_gaussians
diff --git a/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py b/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py
index 3a423bc25..9a9c6809c 100644
--- a/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py
+++ b/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py
@@ -261,7 +261,7 @@ def convergence_func_gaussian(self, grid_radii, sigma, intensity):
         )
 
     def _deflections_2d_via_mge_from(self, grid, sigmas_factor=1.0):
-        axis_ratio = self.axis_ratio()
+        axis_ratio = self.axis_ratio(xp)
 
         if axis_ratio > 0.9999:
             axis_ratio = 0.9999
diff --git a/autogalaxy/profiles/mass/dark/abstract.py b/autogalaxy/profiles/mass/dark/abstract.py
index 20f1bd426..037659753 100644
--- a/autogalaxy/profiles/mass/dark/abstract.py
+++ b/autogalaxy/profiles/mass/dark/abstract.py
@@ -130,7 +130,7 @@ def gnfw_3d(r):
         amplitude_list *= np.sqrt(2.0 * np.pi) * sigma_list
         return amplitude_list, sigma_list
 
-    def coord_func_f(self, grid_radius: np.ndarray) -> np.ndarray:
+    def coord_func_f(self, grid_radius: np.ndarray, xp=np) -> np.ndarray:
         """
         Given an array `grid_radius` and a work array `f`, fill f[i] with
 
@@ -159,7 +159,7 @@ def coord_func_f(self, grid_radius: np.ndarray) -> np.ndarray:
         # combine: if r>1 pick out_gt, elif r<1 pick out_lt, else keep original f
         return xp.where(r > 1.0, out_gt, xp.where(r < 1.0, out_lt, f))
 
-    def coord_func_g(self, grid_radius: np.ndarray) -> np.ndarray:
+    def coord_func_g(self, grid_radius: np.ndarray, xp=np) -> np.ndarray:
         """
         Vectorized version of the original looped `coord_func_g_jit`.
 
@@ -187,7 +187,7 @@ def coord_func_g(self, grid_radius: np.ndarray) -> np.ndarray:
             grid_radius = xp.array([grid_radius], dtype=xp.complex64)
 
         # Evaluate f_r
-        f_r = self.coord_func_f(grid_radius=grid_radius)
+        f_r = self.coord_func_f(grid_radius=grid_radius, xp=xp)
 
         r = xp.real(grid_radius)
         r2 = r**2
@@ -202,8 +202,8 @@ def coord_func_g(self, grid_radius: np.ndarray) -> np.ndarray:
             ),
         )
 
-    def coord_func_h(self, grid_radius):
-        return xp.log(grid_radius / 2.0) + self.coord_func_f(grid_radius=grid_radius)
+    def coord_func_h(self, grid_radius, xp=np):
+        return xp.log(grid_radius / 2.0) + self.coord_func_f(grid_radius=grid_radius, xp=xp)
 
     def rho_at_scale_radius_solar_mass_per_kpc3(
         self, redshift_object, redshift_source, cosmology: LensingCosmology = None
@@ -270,6 +270,7 @@ def concentration(
         redshift_source,
         redshift_of_cosmic_average_density="profile",
         cosmology: LensingCosmology = None,
+        xp=np,
     ):
 
         from autogalaxy.cosmology.wrap import Planck15
@@ -289,11 +290,11 @@ def concentration(
         )
 
         return fsolve(
-            func=self.concentration_func, x0=10.0, args=(delta_concentration,)
+            func=self.concentration_func, x0=10.0, args=(delta_concentration, xp),
         )[0]
 
     @staticmethod
-    def concentration_func(concentration, delta_concentration):
+    def concentration_func(concentration, delta_concentration, xp=np):
         return (
             200.0
             / 3.0
@@ -312,6 +313,7 @@ def radius_at_200(
         redshift_source,
         redshift_of_cosmic_average_density="profile",
         cosmology: LensingCosmology = None,
+        xp=np
     ):
         """
         Returns `r_{200m}` for this halo in **arcseconds**
@@ -325,6 +327,7 @@ def radius_at_200(
             redshift_source=redshift_source,
             redshift_of_cosmic_average_density=redshift_of_cosmic_average_density,
             cosmology=cosmology,
+            xp=xp
         )
 
         return concentration * self.scale_radius
@@ -335,6 +338,7 @@ def mass_at_200_solar_masses(
         redshift_source,
         redshift_of_cosmic_average_density="profile",
         cosmology: LensingCosmology = None,
+        xp=np,
     ):
         from autogalaxy.cosmology.wrap import Planck15
 
@@ -364,6 +368,7 @@ def mass_at_200_solar_masses(
             redshift_source=redshift_source,
             redshift_of_cosmic_average_density=redshift_of_cosmic_average_density,
             cosmology=cosmology,
+            xp=xp
         )
 
         kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object)
diff --git a/autogalaxy/profiles/mass/dark/gnfw.py b/autogalaxy/profiles/mass/dark/gnfw.py
index 97a30d124..55d3e5fa5 100644
--- a/autogalaxy/profiles/mass/dark/gnfw.py
+++ b/autogalaxy/profiles/mass/dark/gnfw.py
@@ -14,13 +14,13 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
     @aa.grid_dec.transform
     def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         return self._deflections_2d_via_mge_from(
-            grid=grid, sigmas_factor=self.axis_ratio()
+            grid=grid, sigmas_factor=self.axis_ratio(xp)
         )
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
     def deflections_2d_via_integral_from(
-        self, grid: aa.type.Grid2DLike, tabulate_bins=1000, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, tabulate_bins=1000, **kwargs
     ):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
@@ -51,7 +51,7 @@ def calculate_deflection_component(npow, yx_index):
                 deflection_grid[i] = (
                     2.0
                     * self.kappa_s
-                    * self.axis_ratio()
+                    * self.axis_ratio(xp)
                     * grid[i, yx_index]
                     * quad(
                         self.deflection_func,
@@ -61,7 +61,7 @@ def calculate_deflection_component(npow, yx_index):
                             grid.array[i, 0],
                             grid.array[i, 1],
                             npow,
-                            self.axis_ratio(),
+                            self.axis_ratio(xp),
                             minimum_log_eta,
                             maximum_log_eta,
                             tabulate_bins,
@@ -104,7 +104,7 @@ def calculate_deflection_component(npow, yx_index):
         deflection_x = calculate_deflection_component(npow=0.0, yx_index=1)
 
         return self.rotated_grid_from_reference_frame_from(
-            np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T)
+            np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T),
         )
 
     @staticmethod
@@ -162,7 +162,7 @@ def integral_y(y, eta):
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, tabulate_bins=1000, **kwargs):
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, tabulate_bins=1000, **kwargs):
         """
         Calculate the potential at a given set of arc-second gridded coordinates.
 
@@ -220,14 +220,14 @@ def deflection_integrand(x, kappa_radius, scale_radius, inner_slope):
             )
 
         for i in range(grid.shape[0]):
-            potential_grid[i] = (2.0 * self.kappa_s * self.axis_ratio()) * quad(
+            potential_grid[i] = (2.0 * self.kappa_s * self.axis_ratio(xp)) * quad(
                 self.potential_func,
                 a=0.0,
                 b=1.0,
                 args=(
                     grid.array[i, 0],
                     grid.array[i, 1],
-                    self.axis_ratio(),
+                    self.axis_ratio(xp),
                     minimum_log_eta,
                     maximum_log_eta,
                     tabulate_bins,
@@ -316,7 +316,7 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
                 4.0 * self.kappa_s * self.scale_radius, self.deflection_func_sph(eta[i])
             )
 
-        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid)
+        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid, xp=xp)
 
     @staticmethod
     def deflection_integrand(y, eta, inner_slope):
diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index d01a449a6..de3cd792f 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -61,7 +61,7 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
         from scipy.integrate import quad
 
         def calculate_deflection_component(npow, index):
-            deflection_grid = np.array(self.axis_ratio() * grid.array[:, index])
+            deflection_grid = np.array(self.axis_ratio(xp) * grid.array[:, index])
 
             for i in range(grid.shape[0]):
                 deflection_grid[i] *= (
@@ -74,7 +74,7 @@ def calculate_deflection_component(npow, index):
                             grid.array[i, 0],
                             grid.array[i, 1],
                             npow,
-                            self.axis_ratio(),
+                            self.axis_ratio(xp),
                             self.scale_radius,
                         ),
                     )[0]
@@ -143,7 +143,7 @@ def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs)
     def convergence_func(self, grid_radius: float) -> float:
         grid_radius = (1.0 / self.scale_radius) * grid_radius.array + 0j
         return np.real(
-            2.0 * self.kappa_s * np.array(self.coord_func_g(grid_radius=grid_radius))
+            2.0 * self.kappa_s * np.array(self.coord_func_g(grid_radius=grid_radius, xp=xp))
         )
 
     @aa.over_sample
@@ -171,7 +171,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
                 args=(
                     grid.array[i, 0],
                     grid.array[i, 1],
-                    self.axis_ratio(),
+                    self.axis_ratio(xp),
                     self.kappa_s,
                     self.scale_radius,
                 ),
@@ -381,12 +381,12 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
 
         deflection_grid = xp.multiply(
             (4.0 * self.kappa_s * self.scale_radius / eta),
-            self.deflection_func_sph(grid_radius=eta),
+            self.deflection_func_sph(grid_radius=eta, xp=xp),
         )
 
-        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid)
+        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid, xp=xp)
 
-    def deflection_func_sph(self, grid_radius):
+    def deflection_func_sph(self, grid_radius, xp=np):
         grid_radius = grid_radius + 0j
         return xp.real(self.coord_func_h(grid_radius=grid_radius))
 
@@ -407,7 +407,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
             grid=grid, **kwargs
         ) + 0j
         return np.real(
-            2.0 * self.scale_radius * self.kappa_s * self.potential_func_sph(eta)
+            2.0 * self.scale_radius * self.kappa_s * self.potential_func_sph(eta, xp=xp)
         )
 
     @staticmethod
diff --git a/autogalaxy/profiles/mass/dark/nfw_truncated.py b/autogalaxy/profiles/mass/dark/nfw_truncated.py
index e680e7ec9..dae77854d 100644
--- a/autogalaxy/profiles/mass/dark/nfw_truncated.py
+++ b/autogalaxy/profiles/mass/dark/nfw_truncated.py
@@ -109,6 +109,7 @@ def mass_at_truncation_radius_solar_mass(
         redshift_source,
         redshift_of_cosmic_average_density="profile",
         cosmology: LensingCosmology = None,
+        xp=np,
     ):
         from autogalaxy.cosmology.wrap import Planck15
 
@@ -119,6 +120,7 @@ def mass_at_truncation_radius_solar_mass(
             redshift_source=redshift_source,
             redshift_of_cosmic_average_density=redshift_of_cosmic_average_density,
             cosmology=cosmology,
+            xp=xp
         )
 
         return (
diff --git a/autogalaxy/profiles/mass/mock/mock_mass_profile.py b/autogalaxy/profiles/mass/mock/mock_mass_profile.py
index c6e87d5be..906562040 100644
--- a/autogalaxy/profiles/mass/mock/mock_mass_profile.py
+++ b/autogalaxy/profiles/mass/mock/mock_mass_profile.py
@@ -1,3 +1,5 @@
+import numpy as np
+
 import autogalaxy as ag
 
 
@@ -19,7 +21,7 @@ def __init__(
         self.value = value
         self.value1 = value1
 
-    def convergence_2d_from(self, grid):
+    def convergence_2d_from(self, grid, xp=np):
         return self.convergence_2d
 
     def potential_2d_from(self, grid):
diff --git a/autogalaxy/profiles/mass/point/point.py b/autogalaxy/profiles/mass/point/point.py
index d44271e02..8ac85b651 100644
--- a/autogalaxy/profiles/mass/point/point.py
+++ b/autogalaxy/profiles/mass/point/point.py
@@ -43,7 +43,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
     def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         grid_radii = self.radial_grid_from(grid=grid, **kwargs)
         return self._cartesian_grid_via_radial_from(
-            grid=grid, radius=self.einstein_radius**2 / grid_radii
+            grid=grid, radius=self.einstein_radius**2 / grid_radii, xp=xp
         )
 
     @property
diff --git a/autogalaxy/profiles/mass/point/smbh_binary.py b/autogalaxy/profiles/mass/point/smbh_binary.py
index 9e5d83091..299ef7e8d 100644
--- a/autogalaxy/profiles/mass/point/smbh_binary.py
+++ b/autogalaxy/profiles/mass/point/smbh_binary.py
@@ -102,8 +102,8 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
             The grid of (y,x) arc-second coordinates the convergence is computed on.
         """
         return self.smbh_0.convergence_2d_from(
-            grid=grid
-        ) + self.smbh_1.convergence_2d_from(grid=grid, **kwargs)
+            grid=grid, xp=xp, **kwargs
+        ) + self.smbh_1.convergence_2d_from(grid=grid, xp=xp, **kwargs)
 
     def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
diff --git a/autogalaxy/profiles/mass/sheets/mass_sheet.py b/autogalaxy/profiles/mass/sheets/mass_sheet.py
index 95392d512..f6e0cc60a 100644
--- a/autogalaxy/profiles/mass/sheets/mass_sheet.py
+++ b/autogalaxy/profiles/mass/sheets/mass_sheet.py
@@ -38,5 +38,5 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
     def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         grid_radii = self.radial_grid_from(grid=grid, **kwargs)
         return self._cartesian_grid_via_radial_from(
-            grid=grid, radius=self.kappa * grid_radii
+            grid=grid, radius=self.kappa * grid_radii, xp=xp
         )
diff --git a/autogalaxy/profiles/mass/stellar/chameleon.py b/autogalaxy/profiles/mass/stellar/chameleon.py
index 1fbd41708..54daca974 100644
--- a/autogalaxy/profiles/mass/stellar/chameleon.py
+++ b/autogalaxy/profiles/mass/stellar/chameleon.py
@@ -120,7 +120,8 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
         deflection_x = xp.subtract(deflection_x0, deflection_x1)
 
         return self.rotated_grid_from_reference_frame_from(
-            xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T)
+            xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T),
+            xp=xp
         )
 
     @aa.over_sample
diff --git a/autogalaxy/profiles/mass/stellar/gaussian.py b/autogalaxy/profiles/mass/stellar/gaussian.py
index bf58c3c51..1af6a364f 100644
--- a/autogalaxy/profiles/mass/stellar/gaussian.py
+++ b/autogalaxy/profiles/mass/stellar/gaussian.py
@@ -76,14 +76,15 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
             self.mass_to_light_ratio
             * self.intensity
             * self.sigma
-            * np.sqrt((2 * np.pi) / (1.0 - self.axis_ratio()**2.0))
-            * self.zeta_from(grid=grid)
+            * xp.sqrt((2 * np.pi) / (1.0 - self.axis_ratio(xp)**2.0))
+            * self.zeta_from(grid=grid, xp=xp)
         )
 
         return self.rotated_grid_from_reference_frame_from(
-            np.multiply(
-                1.0, np.vstack((-1.0 * np.imag(deflections), np.real(deflections))).T
-            )
+            xp.multiply(
+                1.0, xp.vstack((-1.0 * xp.imag(deflections), xp.real(deflections))).T
+            ),
+            xp=xp
         )
 
     @aa.grid_dec.to_vector_yx
@@ -103,7 +104,7 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
         from scipy.integrate import quad
 
         def calculate_deflection_component(npow, index):
-            deflection_grid = np.array(self.axis_ratio() * grid.array[:, index])
+            deflection_grid = np.array(self.axis_ratio(xp) * grid.array[:, index])
 
             for i in range(grid.shape[0]):
                 deflection_grid[i] *= (
@@ -117,8 +118,8 @@ def calculate_deflection_component(npow, index):
                             grid.array[i, 0],
                             grid.array[i, 1],
                             npow,
-                            self.axis_ratio(),
-                            self.sigma / np.sqrt(self.axis_ratio()),
+                            self.axis_ratio(xp),
+                            self.sigma / xp.sqrt(self.axis_ratio(xp)),
                         ),
                     )[0]
                 )
@@ -129,7 +130,8 @@ def calculate_deflection_component(npow, index):
         deflection_x = calculate_deflection_component(0.0, 1)
 
         return self.rotated_grid_from_reference_frame_from(
-            np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T)
+            np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T),
+            xp=xp
         )
 
     @staticmethod
@@ -180,7 +182,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np):
             np.exp(
                 -0.5
                 * np.square(
-                    np.divide(grid_radii.array, self.sigma / np.sqrt(self.axis_ratio()))
+                    np.divide(grid_radii.array, self.sigma / np.sqrt(self.axis_ratio(xp)))
                 )
             ),
         )
@@ -189,15 +191,15 @@ def axis_ratio(self, xp=np):
         axis_ratio = super().axis_ratio(xp=xp)
         return xp.where(axis_ratio < 0.9999, axis_ratio, 0.9999)
 
-    def zeta_from(self, grid: aa.type.Grid2DLike):
+    def zeta_from(self, grid: aa.type.Grid2DLike, xp=np):
 
         from scipy.special import wofz
 
-        q2 = self.axis_ratio()**2.0
+        q2 = self.axis_ratio(xp)**2.0
         ind_pos_y = grid.array[:, 0] >= 0
         shape_grid = np.shape(grid)
         output_grid = np.zeros((shape_grid[0]), dtype=np.complex128)
-        scale_factor = self.axis_ratio() / (self.sigma * np.sqrt(2.0 * (1.0 - q2)))
+        scale_factor = self.axis_ratio(xp) / (self.sigma * np.sqrt(2.0 * (1.0 - q2)))
 
         xs_0 = grid.array[:, 1][ind_pos_y] * scale_factor
         ys_0 = grid.array[:, 0][ind_pos_y] * scale_factor
@@ -207,7 +209,7 @@ def zeta_from(self, grid: aa.type.Grid2DLike):
         output_grid[ind_pos_y] = -1j * (
             wofz(xs_0 + 1j * ys_0)
             - np.exp(-(xs_0**2.0) * (1.0 - q2) - ys_0 * ys_0 * (1.0 / q2 - 1.0))
-            * wofz(self.axis_ratio() * xs_0 + 1j * ys_0 / self.axis_ratio())
+            * wofz(self.axis_ratio(xp) * xs_0 + 1j * ys_0 / self.axis_ratio(xp))
         )
 
         output_grid[~ind_pos_y] = np.conj(
@@ -215,7 +217,7 @@ def zeta_from(self, grid: aa.type.Grid2DLike):
             * (
                 wofz(xs_1 + 1j * ys_1)
                 - np.exp(-(xs_1**2.0) * (1.0 - q2) - ys_1 * ys_1 * (1.0 / q2 - 1.0))
-                * wofz(self.axis_ratio() * xs_1 + 1j * ys_1 / self.axis_ratio())
+                * wofz(self.axis_ratio(xp) * xs_1 + 1j * ys_1 / self.axis_ratio(xp))
             )
         )
 
diff --git a/autogalaxy/profiles/mass/stellar/sersic.py b/autogalaxy/profiles/mass/stellar/sersic.py
index c33ae956a..7f09780dc 100644
--- a/autogalaxy/profiles/mass/stellar/sersic.py
+++ b/autogalaxy/profiles/mass/stellar/sersic.py
@@ -188,7 +188,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
     def convergence_2d_via_mge_from(
-        self, grid: aa.type.Grid2DLike, func_terms=28, func_gaussians=20, **kwargs
+        self, grid: aa.type.Grid2DLike, xp=np, func_terms=28, func_gaussians=20, **kwargs
     ):
         """
         Calculate the projected convergence at a given set of arc-second gridded coordinates.
@@ -390,7 +390,7 @@ def calculate_deflection_component(npow, index):
             deflection_grid = self.axis_ratio() * grid.array[:, index]
 
             for i in range(grid.shape[0]):
-                deflection_grid = deflection_grid.at[i].multiply(
+                deflection_grid[i] = deflection_grid[i] * (
                     self.intensity
                     * self.mass_to_light_ratio
                     * quad(
@@ -415,7 +415,8 @@ def calculate_deflection_component(npow, index):
         deflection_x = calculate_deflection_component(0.0, 1)
 
         return self.rotated_grid_from_reference_frame_from(
-            np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T)
+            np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T),
+            xp=xp
         )
 
     @staticmethod
diff --git a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py
index f140a3433..ae5d0fa70 100644
--- a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py
+++ b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py
@@ -122,7 +122,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
         # And here we convert back to the real axes
         return self.rotated_grid_from_reference_frame_from(
-            grid=xp.multiply(1.0, xp.vstack((deflection_y, deflection_x)).T), **kwargs
+            grid=xp.multiply(1.0, xp.vstack((deflection_y, deflection_x)).T), xp=xp, **kwargs
         )
 
     @aa.grid_dec.to_vector_yx
diff --git a/autogalaxy/profiles/mass/total/isothermal.py b/autogalaxy/profiles/mass/total/isothermal.py
index 354181e33..34c8028b0 100644
--- a/autogalaxy/profiles/mass/total/isothermal.py
+++ b/autogalaxy/profiles/mass/total/isothermal.py
@@ -129,7 +129,7 @@ def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
         """
 
-        convergence = self.convergence_2d_from(grid=grid, **kwargs)
+        convergence = self.convergence_2d_from(grid=grid, xp=xp, **kwargs)
 
         gamma_2 = (
             -2
diff --git a/autogalaxy/profiles/mass/total/power_law.py b/autogalaxy/profiles/mass/total/power_law.py
index 2c919d6c0..9b73a1486 100644
--- a/autogalaxy/profiles/mass/total/power_law.py
+++ b/autogalaxy/profiles/mass/total/power_law.py
@@ -70,25 +70,25 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
         slope = self.slope - 1.0
         einstein_radius = (
-            2.0 / (self.axis_ratio()**-0.5 + self.axis_ratio()**0.5)
+            2.0 / (self.axis_ratio(xp)**-0.5 + self.axis_ratio(xp)**0.5)
         ) * self.einstein_radius
 
-        factor = xp.divide(1.0 - self.axis_ratio(), 1.0 + self.axis_ratio())
-        b = xp.multiply(einstein_radius, xp.sqrt(self.axis_ratio()))
+        factor = xp.divide(1.0 - self.axis_ratio(xp), 1.0 + self.axis_ratio(xp))
+        b = xp.multiply(einstein_radius, xp.sqrt(self.axis_ratio(xp)))
         angle = xp.arctan2(
-            grid.array[:, 0], xp.multiply(self.axis_ratio(), grid.array[:, 1])
+            grid.array[:, 0], xp.multiply(self.axis_ratio(xp), grid.array[:, 1])
         )  # Note, this angle is not the position angle
         z = xp.add(
             xp.multiply(xp.cos(angle), 1 + 0j), xp.multiply(xp.sin(angle), 0 + 1j)
         )
 
         R = xp.sqrt(
-            (self.axis_ratio() * grid.array[:, 1]) ** 2 + grid.array[:, 0] ** 2 + 1e-16
+            (self.axis_ratio(xp) * grid.array[:, 1]) ** 2 + grid.array[:, 0] ** 2 + 1e-16
         )
         zh = omega(z, slope, factor, n_terms=20, xp=np)
 
         complex_angle = (
-            2.0 * b / (1.0 + self.axis_ratio()) * (b / R) ** (slope - 1.0) * zh
+            2.0 * b / (1.0 + self.axis_ratio(xp)) * (b / R) ** (slope - 1.0) * zh
         )
 
         deflection_y = complex_angle.imag
@@ -100,7 +100,8 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         deflection_x *= rescale_factor
 
         return self.rotated_grid_from_reference_frame_from(
-            grid=xp.vstack((deflection_y, deflection_x)).T
+            grid=xp.vstack((deflection_y, deflection_x)).T,
+            xp=xp
         )
 
     def convergence_func(self, grid_radius: float, xp=np) -> float:
diff --git a/autogalaxy/profiles/mass/total/power_law_broken.py b/autogalaxy/profiles/mass/total/power_law_broken.py
index c034ee15b..9c6613441 100644
--- a/autogalaxy/profiles/mass/total/power_law_broken.py
+++ b/autogalaxy/profiles/mass/total/power_law_broken.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Tuple
 
@@ -31,7 +30,7 @@ def __init__(
         super().__init__(centre=centre, ell_comps=ell_comps)
 
         self.einstein_radius = einstein_radius
-        self.einstein_radius_elliptical = xp.sqrt(self.axis_ratio()) * einstein_radius
+        self.einstein_radius_elliptical = np.sqrt(self.axis_ratio()) * einstein_radius
         self.break_radius = break_radius
         self.inner_slope = inner_slope
         self.outer_slope = outer_slope
@@ -58,7 +57,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """
 
         # Ell radius
-        radius = xp.hypot(grid.array[:, 1] * self.axis_ratio(), grid.array[:, 0])
+        radius = xp.hypot(grid.array[:, 1] * self.axis_ratio(xp), grid.array[:, 0])
 
         # Inside break radius
         kappa_inner = self.kB * (self.break_radius / radius) ** self.inner_slope
@@ -76,22 +75,23 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_yx_2d_from(self, grid, max_terms=20, **kwargs):
+    def deflections_yx_2d_from(self, grid, xp=np, max_terms=20, **kwargs):
         """
         Returns the complex deflection angle from eq. 18 and 19
         """
+
         # Rotate coordinates
         z = grid.array[:, 1] + 1j * grid.array[:, 0]
 
         # Ell radius
-        R = xp.hypot(z.real * self.axis_ratio(), z.imag)
+        R = xp.hypot(z.real * self.axis_ratio(xp), z.imag)
 
         # Factors common to eq. 18 and 19
         factors = (
             2
             * self.kB
             * (self.break_radius**2)
-            / (self.axis_ratio() * z * (2 - self.inner_slope))
+            / (self.axis_ratio(xp) * z * (2 - self.inner_slope))
         )
 
         # Hypergeometric functions
@@ -99,16 +99,16 @@ def deflections_yx_2d_from(self, grid, max_terms=20, **kwargs):
         # These can also be computed with scipy.special.hyp2f1(), it's
         # much slower can be a useful test
         F1 = self.hyp2f1_series(
-            self.inner_slope, self.axis_ratio(), R, z, max_terms=max_terms
+            self.inner_slope, self.axis_ratio(xp), R, z, max_terms=max_terms, xp=xp
         )
         F2 = self.hyp2f1_series(
-            self.inner_slope, self.axis_ratio(), self.break_radius, z, max_terms=max_terms
+            self.inner_slope, self.axis_ratio(xp), self.break_radius, z, max_terms=max_terms, xp=xp
         )
         F3 = self.hyp2f1_series(
-            self.outer_slope, self.axis_ratio(), R, z, max_terms=max_terms
+            self.outer_slope, self.axis_ratio(xp), R, z, max_terms=max_terms, xp=xp
         )
         F4 = self.hyp2f1_series(
-            self.outer_slope, self.axis_ratio(), self.break_radius, z, max_terms=max_terms
+            self.outer_slope, self.axis_ratio(xp), self.break_radius, z, max_terms=max_terms, xp=xp
         )
 
         # theta < break radius (eq. 18)
@@ -128,11 +128,12 @@ def deflections_yx_2d_from(self, grid, max_terms=20, **kwargs):
         return self.rotated_grid_from_reference_frame_from(
             grid=xp.multiply(
                 1.0, xp.vstack((xp.imag(deflections), xp.real(deflections))).T
-            )
+            ),
+            xp=xp
         )
 
     @staticmethod
-    def hyp2f1_series(t, q, r, z, max_terms=20):
+    def hyp2f1_series(t, q, r, z, max_terms=20, xp=np):
         """
         Computes eq. 26 for a radius r, slope t,
         axis ratio q, and coordinates z.
diff --git a/autogalaxy/profiles/mass/total/power_law_core.py b/autogalaxy/profiles/mass/total/power_law_core.py
index bd1b5ee20..503bc29fe 100644
--- a/autogalaxy/profiles/mass/total/power_law_core.py
+++ b/autogalaxy/profiles/mass/total/power_law_core.py
@@ -90,13 +90,13 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
                 args=(
                     grid.array[i, 0],
                     grid.array[i, 1],
-                    self.axis_ratio(),
+                    self.axis_ratio(xp),
                     self.slope,
                     self.core_radius,
                 ),
             )[0]
 
-        return self.einstein_radius_rescaled(xp) * self.axis_ratio() * potential_grid
+        return self.einstein_radius_rescaled(xp) * self.axis_ratio(xp) * potential_grid
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
@@ -115,7 +115,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         def calculate_deflection_component(npow, index):
             einstein_radius_rescaled = self.einstein_radius_rescaled(xp)
 
-            deflection_grid = np.array(self.axis_ratio() * grid.array[:, index])
+            deflection_grid = np.array(self.axis_ratio(xp) * grid.array[:, index])
 
             for i in range(grid.shape[0]):
                 deflection_grid[i] *= (
@@ -128,7 +128,7 @@ def calculate_deflection_component(npow, index):
                             grid.array[i, 0],
                             grid.array[i, 1],
                             npow,
-                            self.axis_ratio(),
+                            self.axis_ratio(xp),
                             self.slope,
                             self.core_radius,
                         ),
@@ -141,7 +141,8 @@ def calculate_deflection_component(npow, index):
         deflection_x = calculate_deflection_component(0.0, 1)
 
         return self.rotated_grid_from_reference_frame_from(
-            grid=np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T)
+            grid=np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T),
+            xp=xp
         )
 
     def convergence_func(self, grid_radius: float, xp=np) -> float:
@@ -234,4 +235,4 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
                 xp.multiply((3.0 - self.slope), eta.array),
             ),
         )
-        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection)
+        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection, xp=xp)
diff --git a/autogalaxy/profiles/mass/total/power_law_multipole.py b/autogalaxy/profiles/mass/total/power_law_multipole.py
index 40f562979..94e3cf205 100644
--- a/autogalaxy/profiles/mass/total/power_law_multipole.py
+++ b/autogalaxy/profiles/mass/total/power_law_multipole.py
@@ -10,7 +10,7 @@
 
 
 def radial_and_angle_grid_from(
-    grid: aa.type.Grid2DLike, centre: Tuple[float, float] = (0.0, 0.0)
+    grid: aa.type.Grid2DLike, centre: Tuple[float, float] = (0.0, 0.0), xp=np
 ) -> Tuple[np.ndarray, np.ndarray]:
     """
     Converts the input grid of Cartesian (y,x) coordinates to their correspond radial and polar grids.
@@ -129,7 +129,7 @@ def __init__(
         self.angle_m *= units.deg.to(units.rad)
 
     def jacobian(
-        self, a_r: np.ndarray, a_angle: np.ndarray, polar_angle_grid: np.ndarray
+        self, a_r: np.ndarray, a_angle: np.ndarray, polar_angle_grid: np.ndarray, xp=np
     ) -> Tuple[np.ndarray, Tuple]:
         """
         The Jacobian transformation from polar to cartesian coordinates.
@@ -151,7 +151,7 @@ def jacobian(
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
     def deflections_yx_2d_from(
-        self, grid: aa.type.Grid1D2DLike, **kwargs
+        self, grid: aa.type.Grid1D2DLike, xp=np, **kwargs
     ) -> np.ndarray:
         """
         Calculate the deflection angles on a grid of (y,x) arc-second coordinates.
@@ -164,7 +164,7 @@ def deflections_yx_2d_from(
         grid
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
-        radial_grid, polar_angle_grid = radial_and_angle_grid_from(grid=grid)
+        radial_grid, polar_angle_grid = radial_and_angle_grid_from(grid=grid, xp=xp)
 
         a_r = (
             -(
@@ -189,14 +189,14 @@ def deflections_yx_2d_from(
         )
 
         return xp.stack(
-            self.jacobian(a_r=a_r, a_angle=a_angle, polar_angle_grid=polar_angle_grid),
+            self.jacobian(a_r=a_r, a_angle=a_angle, polar_angle_grid=polar_angle_grid, xp=xp),
             axis=-1,
         )
 
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid1D2DLike, **kwargs) -> np.ndarray:
+    def convergence_2d_from(self, grid: aa.type.Grid1D2DLike, xp=np, **kwargs) -> np.ndarray:
         """
         Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates.
 
@@ -205,7 +205,7 @@ def convergence_2d_from(self, grid: aa.type.Grid1D2DLike, **kwargs) -> np.ndarra
         grid
             The grid of (y,x) arc-second coordinates the convergence is computed on.
         """
-        r, angle = radial_and_angle_grid_from(grid=grid)
+        r, angle = radial_and_angle_grid_from(grid=grid, xp=xp)
 
         return (
             1.0
diff --git a/test_autogalaxy/profiles/mass/abstract/test_abstract.py b/test_autogalaxy/profiles/mass/abstract/test_abstract.py
index 4e8f817f4..2555af98c 100644
--- a/test_autogalaxy/profiles/mass/abstract/test_abstract.py
+++ b/test_autogalaxy/profiles/mass/abstract/test_abstract.py
@@ -205,61 +205,6 @@ def test__regression__centre_of_profile_in_right_place():
     assert deflections.native[1, 3, 1] < 0
 
 
-def test__decorators__convergence_1d_from__grid_2d_in__returns_1d_image_via_projected_quantities():
-    grid_2d = ag.Grid2D.uniform(
-        shape_native=(5, 5),
-        pixel_scales=1.0,
-        over_sample_size=1,
-    )
-
-    sie = ag.mp.Isothermal(
-        centre=(1e-6, 1e-6), ell_comps=(0.0, 0.0), einstein_radius=1.0
-    )
-
-    convergence_1d = sie.convergence_1d_from(grid=grid_2d)
-    convergence_2d = sie.convergence_2d_from(grid=grid_2d)
-
-    print(convergence_2d.native.array)
-
-    assert convergence_1d[0] == pytest.approx(convergence_2d.native.array[2, 2], 1.0e-4)
-    assert convergence_1d[1] == pytest.approx(convergence_2d.native.array[2, 3], 1.0e-4)
-    assert convergence_1d[2] == pytest.approx(convergence_2d.native.array[2, 4], 1.0e-4)
-
-    sie = ag.mp.Isothermal(centre=(0.2, 0.2), ell_comps=(0.3, 0.3), einstein_radius=1.0)
-
-    convergence_1d = sie.convergence_1d_from(grid=grid_2d)
-
-    grid_2d_projected = grid_2d.grid_2d_radial_projected_from(
-        centre=sie.centre, angle=sie.angle + 90.0
-    )
-
-    convergence_projected = sie.convergence_2d_from(grid=grid_2d_projected)
-
-    assert convergence_1d == pytest.approx(convergence_projected.array, 1.0e-4)
-    assert (convergence_1d.grid_radial == np.array([0.0, 1.0, 2.0])).all()
-
-
-def test__decorators__convergence_1d_from__grid_2d_irregular_in__returns_1d_quantities():
-    grid_2d = ag.Grid2DIrregular(values=[[1.0, 1.0], [2.0, 2.0], [4.0, 4.0]])
-
-    sie = ag.mp.Isothermal(centre=(0.0, 0.0), ell_comps=(0.0, 0.0), einstein_radius=1.0)
-
-    convergence_1d = sie.convergence_1d_from(grid=grid_2d)
-    convergence_2d = sie.convergence_2d_from(grid=grid_2d)
-
-    assert convergence_1d[0] == pytest.approx(convergence_2d[0].array, 1.0e-4)
-    assert convergence_1d[1] == pytest.approx(convergence_2d[1].array, 1.0e-4)
-    assert convergence_1d[2] == pytest.approx(convergence_2d[2].array, 1.0e-4)
-
-    sie = ag.mp.Isothermal(centre=(0.2, 0.2), ell_comps=(0.3, 0.3), einstein_radius=1.0)
-
-    convergence_1d = sie.convergence_1d_from(grid=grid_2d)
-    convergence_2d = sie.convergence_2d_from(grid=grid_2d)
-
-    assert convergence_1d[0] == pytest.approx(convergence_2d[0].array, 1.0e-4)
-    assert convergence_1d[1] == pytest.approx(convergence_2d[1].array, 1.0e-4)
-    assert convergence_1d[2] == pytest.approx(convergence_2d[2].array, 1.0e-4)
-
 
 def test__decorators__convergence_1d_from__grid_1d_in__returns_1d_quantities_via_projection():
     grid_1d = ag.Grid1D.no_mask(values=[1.0, 2.0, 3.0], pixel_scales=1.0)
@@ -283,30 +228,4 @@ def test__decorators__convergence_1d_from__grid_1d_in__returns_1d_quantities_via
 
     assert convergence_1d[0] == pytest.approx(convergence_2d[0].array, 1.0e-4)
     assert convergence_1d[1] == pytest.approx(convergence_2d[1].array, 1.0e-4)
-    assert convergence_1d[2] == pytest.approx(convergence_2d[2].array, 1.0e-4)
-
-
-def test__decorators__potential_1d_from__grid_2d_in__returns_1d_image_via_projected_quantities():
-    grid_2d = ag.Grid2D.uniform(shape_native=(5, 5), pixel_scales=1.0)
-
-    sie = ag.mp.Isothermal(centre=(0.0, 0.0), ell_comps=(0.0, 0.0), einstein_radius=1.0)
-
-    potential_1d = sie.potential_1d_from(grid=grid_2d)
-    potential_2d = sie.potential_2d_from(grid=grid_2d)
-
-    assert potential_1d[0] == pytest.approx(potential_2d.native.array[2, 2], abs=1.0e-4)
-    assert potential_1d[1] == pytest.approx(potential_2d.native.array[2, 3], abs=1.0e-4)
-    assert potential_1d[2] == pytest.approx(potential_2d.native.array[2, 4], abs=1.0e-4)
-
-    sie = ag.mp.Isothermal(centre=(0.2, 0.2), ell_comps=(0.3, 0.3), einstein_radius=1.0)
-
-    potential_1d = sie.potential_1d_from(grid=grid_2d)
-
-    grid_2d_projected = grid_2d.grid_2d_radial_projected_from(
-        centre=sie.centre, angle=sie.angle + 90.0
-    )
-
-    potential_projected = sie.potential_2d_from(grid=grid_2d_projected)
-
-    assert potential_1d == pytest.approx(potential_projected.array, abs=1.0e-4)
-    assert (potential_1d.grid_radial == np.array([0.0, 1.0, 2.0])).all()
+    assert convergence_1d[2] == pytest.approx(convergence_2d[2].array, 1.0e-4)
\ No newline at end of file
diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py b/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py
index 7d8465412..143f6f108 100644
--- a/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py
+++ b/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py
@@ -45,9 +45,9 @@ def test__mass_and_concentration_consistent_with_normal_nfw():
 
     assert mp.centre == (1.0, 2.0)
 
-    assert mp.axis_ratio == 1.0
+    assert mp.axis_ratio() == 1.0
 
-    assert mp.angle == 0.0
+    assert mp.angle() == 0.0
 
     assert mp.inner_slope == 1.0
 
@@ -93,8 +93,8 @@ def test__mass_and_concentration_consistent_with_normal_nfw__scatter_0():
     assert concentration_via_kappa_s == concentration_via_mass
 
     assert mp.centre == (1.0, 2.0)
-    assert mp.axis_ratio == 1.0
-    assert mp.angle == 0.0
+    assert mp.axis_ratio() == 1.0
+    assert mp.angle() == 0.0
     assert mp.inner_slope == 1.0
     assert mp.scale_radius == pytest.approx(0.21157, 1.0e-4)
 
@@ -148,8 +148,8 @@ def test__same_as_above_but_elliptical():
 
     axis_ratio, angle = ag.convert.axis_ratio_and_angle_from(ell_comps=(0.1, 0.2))
 
-    assert mp.axis_ratio == axis_ratio
-    assert mp.angle == angle
+    assert mp.axis_ratio() == axis_ratio
+    assert mp.angle() == angle
     assert mp.inner_slope == 1.0
 
     assert mp.scale_radius == pytest.approx(0.211578, 1.0e-4)
@@ -206,8 +206,8 @@ def test__same_as_above_but_generalized_elliptical():
 
     axis_ratio, angle = ag.convert.axis_ratio_and_angle_from(ell_comps=(0.1, 0.2))
 
-    assert mp.axis_ratio == axis_ratio
-    assert mp.angle == angle
+    assert mp.axis_ratio() == axis_ratio
+    assert mp.angle() == angle
     assert mp.inner_slope == 2.0
 
     assert mp.scale_radius == pytest.approx(0.21157, 1.0e-4)
diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw_scatter.py b/test_autogalaxy/profiles/mass/dark/test_nfw_scatter.py
index a9d972325..91f621a44 100644
--- a/test_autogalaxy/profiles/mass/dark/test_nfw_scatter.py
+++ b/test_autogalaxy/profiles/mass/dark/test_nfw_scatter.py
@@ -58,4 +58,4 @@ def test__scatter_is_nonzero():
     deflections_sph = mp.deflections_yx_2d_from(grid=grid)
     deflections_ell = nfw_ell.deflections_yx_2d_from(grid=grid)
 
-    assert deflections_sph[0] != pytest.approx(deflections_ell[0].array, 1.0e-4)
+    assert deflections_sph[0] != pytest.approx(deflections_ell[0], 1.0e-4)
diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py b/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py
index 64adf18fb..99188189b 100644
--- a/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py
+++ b/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py
@@ -46,8 +46,8 @@ def test__duffy__mass_and_concentration_consistent_with_normal_truncated_nfw():
     assert concentration_via_kappa_s == concentration_via_mass
 
     assert mp.centre == (1.0, 2.0)
-    assert mp.axis_ratio == 1.0
-    assert mp.angle == 0.0
+    assert mp.axis_ratio() == 1.0
+    assert mp.angle() == 0.0
     assert mp.inner_slope == 1.0
 
     assert mp.scale_radius == pytest.approx(0.273382, 1.0e-4)
@@ -93,8 +93,8 @@ def test__ludlow__mass_and_concentration_consistent_with_normal_truncated_nfw__s
     assert concentration_via_kappa_s == concentration_via_mass
 
     assert mp.centre == (1.0, 2.0)
-    assert mp.axis_ratio == 1.0
-    assert mp.angle == 0.0
+    assert mp.axis_ratio() == 1.0
+    assert mp.angle() == 0.0
     assert mp.inner_slope == 1.0
 
     assert mp.scale_radius == pytest.approx(0.21157, 1.0e-4)
diff --git a/test_autogalaxy/profiles/mass/total/test_power_law_broken.py b/test_autogalaxy/profiles/mass/total/test_power_law_broken.py
index ea27bc8de..d9859a0e3 100644
--- a/test_autogalaxy/profiles/mass/total/test_power_law_broken.py
+++ b/test_autogalaxy/profiles/mass/total/test_power_law_broken.py
@@ -179,7 +179,7 @@ def test__deflections_yx_2d_from__compare_to_power_law():
 
     power_law_yx_ratio = deflections[0, 0] / deflections[0, 1]
 
-    assert broken_yx_ratio == pytest.approx(power_law_yx_ratio.array, 1.0e-4)
+    assert broken_yx_ratio == pytest.approx(power_law_yx_ratio, 1.0e-4)
 
     mp = ag.mp.PowerLawBrokenSph(
         centre=(0, 0),
@@ -201,4 +201,4 @@ def test__deflections_yx_2d_from__compare_to_power_law():
 
     power_law_yx_ratio = deflections[0, 0] / deflections[0, 1]
 
-    assert broken_yx_ratio == pytest.approx(power_law_yx_ratio.array, 1.0e-4)
+    assert broken_yx_ratio == pytest.approx(power_law_yx_ratio, 1.0e-4)

From 54d207eb40bad6662896848d475c879edfa2bef6 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 8 Nov 2025 15:25:06 +0000
Subject: [PATCH 08/18] operate reverts to numpy for now

---
 autogalaxy/operate/deflections.py             | 410 ++++++------------
 .../profiles/mass/stellar/sersic_core.py      |   2 +-
 test_autogalaxy/operate/test_deflections.py   | 101 ++---
 3 files changed, 187 insertions(+), 326 deletions(-)

diff --git a/autogalaxy/operate/deflections.py b/autogalaxy/operate/deflections.py
index 9f07a9720..8cb691ae2 100644
--- a/autogalaxy/operate/deflections.py
+++ b/autogalaxy/operate/deflections.py
@@ -1,10 +1,9 @@
-import jax
-from jax import jit
-import jax.numpy as jnp
-from functools import wraps, partial
+from functools import wraps
 import logging
+import numpy as np
 from typing import List, Tuple, Union
 
+from autoconf import conf
 
 import autoarray as aa
 
@@ -46,18 +45,49 @@ def wrapper(lensing_obj, grid, jacobian=None):
     return wrapper
 
 
-def one_step(r, _, theta, fun, fun_dr):
-    r = xp.abs(r - fun(r, theta) / fun_dr(r, theta))
-    return r, None
+def evaluation_grid(func):
+    @wraps(func)
+    def wrapper(
+        lensing_obj, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05
+    ):
+        if hasattr(grid, "is_evaluation_grid"):
+            if grid.is_evaluation_grid:
+                return func(lensing_obj, grid, pixel_scale)
 
+        pixel_scale_ratio = grid.pixel_scale / pixel_scale
 
-@partial(jit, static_argnums=(4,))
-def step_r(r, theta, fun, fun_dr, N=20):
-    one_step_partial = jax.tree_util.Partial(
-        one_step, theta=theta, fun=fun, fun_dr=fun_dr
-    )
-    new_r = jax.lax.scan(one_step_partial, r, xs=xp.arange(N))[0]
-    return xp.stack([new_r * xp.sin(theta), new_r * xp.cos(theta)]).T
+        zoom = aa.Zoom2D(mask=grid.mask)
+
+        zoom_shape_native = zoom.shape_native
+        shape_native = (
+            int(pixel_scale_ratio * zoom_shape_native[0]),
+            int(pixel_scale_ratio * zoom_shape_native[1]),
+        )
+
+        max_evaluation_grid_size = conf.instance["general"]["grid"][
+            "max_evaluation_grid_size"
+        ]
+
+        # This is a hack to prevent the evaluation gird going beyond 1000 x 1000 pixels, which slows the code
+        # down a lot. Need a better moe robust way to set this up for any general lens.
+
+        if shape_native[0] > max_evaluation_grid_size:
+            pixel_scale = pixel_scale_ratio / (
+                shape_native[0] / float(max_evaluation_grid_size)
+            )
+            shape_native = (max_evaluation_grid_size, max_evaluation_grid_size)
+
+        grid = aa.Grid2D.uniform(
+            shape_native=shape_native,
+            pixel_scales=(pixel_scale, pixel_scale),
+            origin=zoom.offset_scaled,
+        )
+
+        grid.is_evaluation_grid = True
+
+        return func(lensing_obj, grid, pixel_scale)
+
+    return wrapper
 
 
 class OperateDeflections:
@@ -79,22 +109,6 @@ class OperateDeflections:
     def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         raise NotImplementedError
 
-    def deflections_yx_scalar(self, y, x, pixel_scales):
-
-        # A version of the deflection function that takes in two scalars
-        # and outputs a 2D vector.  Needed for JAX auto differentiation.
-
-        mask = aa.Mask2D.all_false(
-            shape_native=(1, 1),
-            pixel_scales=pixel_scales,
-        )
-
-        g = aa.Grid2D(
-            values=xp.stack((y.reshape(1), x.reshape(1)), axis=-1), mask=mask
-        )
-
-        return self.deflections_yx_2d_from(g).squeeze()
-
     def __eq__(self, other):
         return self.__dict__ == other.__dict__ and self.__class__ is other.__class__
 
@@ -167,27 +181,6 @@ def fermat_potential_from(self, grid) -> aa.Array2D:
             return aa.ArrayIrregular(values=fermat_potential)
         return aa.Array2D(values=fermat_potential, mask=grid.mask)
 
-    def time_delays_from(self, grid) -> aa.Array2D:
-        """
-        Returns the 2D time delay map of lensing object, which is computed as the deflection angles in the y and x
-        directions multiplied by the y and x coordinates of the grid.
-
-        Parameters
-        ----------
-        grid
-            The 2D grid of (y,x) arc-second coordinates the deflection angles and time delay are computed on.
-        """
-        deflections_yx = self.deflections_yx_2d_from(grid=grid)
-
-        return aa.Array2D(
-            values=deflections_yx[:, 0] * grid[:, 0]
-            + deflections_yx[:, 1] * grid[:, 1],
-            mask=grid.mask,
-        )
-
-    def __hash__(self):
-        return hash(repr(self))
-
     @precompute_jacobian
     def tangential_eigen_value_from(self, grid, jacobian=None) -> aa.Array2D:
         """
@@ -277,24 +270,21 @@ def hessian_from(self, grid, buffer: float = 0.01, deflections_func=None) -> Tup
         if deflections_func is None:
             deflections_func = self.deflections_yx_2d_from
 
-        if isinstance(grid, aa.Grid2DIrregular):
-            grid = grid.array
+        grid_shift_y_up = aa.Grid2DIrregular(values=np.zeros(grid.shape))
+        grid_shift_y_up[:, 0] = grid[:, 0] + buffer
+        grid_shift_y_up[:, 1] = grid[:, 1]
 
-        grid_shift_y_up = aa.Grid2DIrregular(
-            values=xp.stack([grid[:, 0] + buffer, grid[:, 1]], axis=1)
-        )
+        grid_shift_y_down = aa.Grid2DIrregular(values=np.zeros(grid.shape))
+        grid_shift_y_down[:, 0] = grid[:, 0] - buffer
+        grid_shift_y_down[:, 1] = grid[:, 1]
 
-        grid_shift_y_down = aa.Grid2DIrregular(
-            values=xp.stack([grid[:, 0] - buffer, grid[:, 1]], axis=1)
-        )
+        grid_shift_x_left = aa.Grid2DIrregular(values=np.zeros(grid.shape))
+        grid_shift_x_left[:, 0] = grid[:, 0]
+        grid_shift_x_left[:, 1] = grid[:, 1] - buffer
 
-        grid_shift_x_left = aa.Grid2DIrregular(
-            values=xp.stack([grid[:, 0], grid[:, 1] - buffer], axis=1)
-        )
-
-        grid_shift_x_right = aa.Grid2DIrregular(
-            values=xp.stack([grid[:, 0], grid[:, 1] + buffer], axis=1)
-        )
+        grid_shift_x_right = aa.Grid2DIrregular(values=np.zeros(grid.shape))
+        grid_shift_x_right[:, 0] = grid[:, 0]
+        grid_shift_x_right[:, 1] = grid[:, 1] + buffer
 
         deflections_up = deflections_func(grid=grid_shift_y_up)
         deflections_down = deflections_func(grid=grid_shift_y_down)
@@ -375,7 +365,10 @@ def shear_yx_2d_via_hessian_from(
         gamma_1 = 0.5 * (hessian_xx - hessian_yy)
         gamma_2 = hessian_xy
 
-        shear_yx_2d = xp.stack([gamma_2.array, gamma_1.array], axis=1)
+        shear_yx_2d = np.zeros(shape=(grid.shape_slim, 2))
+
+        shear_yx_2d[:, 0] = gamma_2
+        shear_yx_2d[:, 1] = gamma_1
 
         return ShearYX2DIrregular(values=shear_yx_2d, grid=grid)
 
@@ -419,9 +412,9 @@ def contour_list_from(self, grid, contour_array):
 
         return grid_contour.contour_list
 
+    @evaluation_grid
     def tangential_critical_curve_list_from(
-        self,
-        grid,
+        self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05
     ) -> List[aa.Grid2DIrregular]:
         """
         Returns all tangential critical curves of the lensing system, which are computed as follows:
@@ -445,9 +438,9 @@ def tangential_critical_curve_list_from(
 
         return self.contour_list_from(grid=grid, contour_array=tangential_eigen_values)
 
+    @evaluation_grid
     def radial_critical_curve_list_from(
-        self,
-        grid,
+        self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05
     ) -> List[aa.Grid2DIrregular]:
         """
         Returns all radial critical curves of the lensing system, which are computed as follows:
@@ -471,9 +464,9 @@ def radial_critical_curve_list_from(
 
         return self.contour_list_from(grid=grid, contour_array=radial_eigen_values)
 
+    @evaluation_grid
     def tangential_caustic_list_from(
-        self,
-        grid,
+        self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05
     ) -> List[aa.Grid2DIrregular]:
         """
         Returns all tangential caustics of the lensing system, which are computed as follows:
@@ -497,7 +490,7 @@ def tangential_caustic_list_from(
         """
 
         tangential_critical_curve_list = self.tangential_critical_curve_list_from(
-            grid=grid
+            grid=grid, pixel_scale=pixel_scale
         )
 
         tangential_caustic_list = []
@@ -513,9 +506,9 @@ def tangential_caustic_list_from(
 
         return tangential_caustic_list
 
+    @evaluation_grid
     def radial_caustic_list_from(
-        self,
-        grid,
+        self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05
     ) -> List[aa.Grid2DIrregular]:
         """
         Returns all radial caustics of the lensing system, which are computed as follows:
@@ -538,7 +531,9 @@ def radial_caustic_list_from(
             caustic to be computed more accurately using a higher resolution grid.
         """
 
-        radial_critical_curve_list = self.radial_critical_curve_list_from(grid=grid)
+        radial_critical_curve_list = self.radial_critical_curve_list_from(
+            grid=grid, pixel_scale=pixel_scale
+        )
 
         radial_caustic_list = []
 
@@ -553,7 +548,10 @@ def radial_caustic_list_from(
 
         return radial_caustic_list
 
-    def radial_critical_curve_area_list_from(self, grid) -> List[float]:
+    @evaluation_grid
+    def radial_critical_curve_area_list_from(
+        self, grid, pixel_scale: Union[Tuple[float, float], float]
+    ) -> List[float]:
         """
         Returns the surface area within each radial critical curve as a list, the calculation of which is described in
         the function `radial_critical_curve_list_from()`.
@@ -573,13 +571,15 @@ def radial_critical_curve_area_list_from(self, grid) -> List[float]:
             If input, the `evaluation_grid` decorator creates the 2D grid at this resolution, therefore enabling the
             caustic to be computed more accurately using a higher resolution grid.
         """
-        radial_critical_curve_list = self.radial_critical_curve_list_from(grid=grid)
+        radial_critical_curve_list = self.radial_critical_curve_list_from(
+            grid=grid, pixel_scale=pixel_scale
+        )
 
         return self.area_within_curve_list_from(curve_list=radial_critical_curve_list)
 
+    @evaluation_grid
     def tangential_critical_curve_area_list_from(
-        self,
-        grid,
+        self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05
     ) -> List[float]:
         """
         Returns the surface area within each tangential critical curve as a list, the calculation of which is
@@ -600,7 +600,7 @@ def tangential_critical_curve_area_list_from(
             caustic to be computed more accurately using a higher resolution grid.
         """
         tangential_critical_curve_list = self.tangential_critical_curve_list_from(
-            grid=grid
+            grid=grid, pixel_scale=pixel_scale
         )
 
         return self.area_within_curve_list_from(
@@ -614,14 +614,14 @@ def area_within_curve_list_from(
 
         for curve in curve_list:
             x, y = curve[:, 0], curve[:, 1]
-            area = xp.abs(0.5 * xp.sum(y[:-1] * xp.diff(x) - x[:-1] * xp.diff(y)))
+            area = np.abs(0.5 * np.sum(y[:-1] * np.diff(x) - x[:-1] * np.diff(y)))
             area_within_each_curve_list.append(area)
 
         return area_within_each_curve_list
 
+    @evaluation_grid
     def einstein_radius_list_from(
-        self,
-        grid,
+        self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05
     ):
         """
         Returns a list of the Einstein radii corresponding to the area within each tangential critical curve.
@@ -648,14 +648,16 @@ def einstein_radius_list_from(
             caustic to be computed more accurately using a higher resolution grid.
         """
         try:
-            area_list = self.tangential_critical_curve_area_list_from(grid=grid)
-            return [xp.sqrt(area / xp.pi) for area in area_list]
+            area_list = self.tangential_critical_curve_area_list_from(
+                grid=grid, pixel_scale=pixel_scale
+            )
+            return [np.sqrt(area / np.pi) for area in area_list]
         except TypeError:
             raise TypeError("The grid input was unable to estimate the Einstein Radius")
 
+    @evaluation_grid
     def einstein_radius_from(
-        self,
-        grid,
+        self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05
     ):
         """
         Returns the Einstein radius corresponding to the area within the tangential critical curve.
@@ -697,9 +699,9 @@ def einstein_radius_from(
 
         return sum(einstein_radii_list)
 
+    @evaluation_grid
     def einstein_mass_angular_list_from(
-        self,
-        grid,
+        self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05
     ) -> List[float]:
         """
         Returns a list of the angular Einstein massses corresponding to the area within each tangential critical curve.
@@ -728,12 +730,14 @@ def einstein_mass_angular_list_from(
             If input, the `evaluation_grid` decorator creates the 2D grid at this resolution, therefore enabling the
             caustic to be computed more accurately using a higher resolution grid.
         """
-        einstein_radius_list = self.einstein_radius_list_from(grid=grid)
-        return [xp.pi * einstein_radius**2 for einstein_radius in einstein_radius_list]
+        einstein_radius_list = self.einstein_radius_list_from(
+            grid=grid, pixel_scale=pixel_scale
+        )
+        return [np.pi * einstein_radius**2 for einstein_radius in einstein_radius_list]
 
+    @evaluation_grid
     def einstein_mass_angular_from(
-        self,
-        grid,
+        self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05
     ) -> float:
         """
         Returns the Einstein radius corresponding to the area within the tangential critical curve.
@@ -762,7 +766,9 @@ def einstein_mass_angular_from(
             If input, the `evaluation_grid` decorator creates the 2D grid at this resolution, therefore enabling the
             caustic to be computed more accurately using a higher resolution grid.
         """
-        einstein_mass_angular_list = self.einstein_mass_angular_list_from(grid=grid)
+        einstein_mass_angular_list = self.einstein_mass_angular_list_from(
+            grid=grid, pixel_scale=pixel_scale
+        )
 
         if len(einstein_mass_angular_list) > 1:
             logger.info(
@@ -774,167 +780,6 @@ def einstein_mass_angular_from(
 
         return einstein_mass_angular_list[0]
 
-    def jacobian_stack(self, y, x, pixel_scales):
-        return xp.stack(
-            jax.jacfwd(self.deflections_yx_scalar, argnums=(0, 1))(y, x, pixel_scales)
-        )
-
-    def jacobian_stack_vector(self, y, x, pixel_scales):
-        return xp.vectorize(
-            jax.tree_util.Partial(self.jacobian_stack, pixel_scales=pixel_scales),
-            signature="(),()->(i,i)",
-        )(y, x)
-
-    def convergence_mag_shear_yx(self, y, x):
-        J = self.jacobian_stack_vector(y, x, 0.05)
-        K = 0.5 * (J[..., 0, 0] + J[..., 1, 1])
-        mag_shear = 0.5 * xp.sqrt(
-            (J[..., 0, 1] + J[..., 1, 0]) ** 2 + (J[..., 0, 0] - J[..., 1, 1]) ** 2
-        )
-        return K, mag_shear
-
-    @partial(jit, static_argnums=(0,))
-    def tangential_eigen_value_yx(self, y, x):
-        K, mag_shear = self.convergence_mag_shear_yx(y, x)
-        return 1 - K - mag_shear
-
-    @partial(jit, static_argnums=(0, 3))
-    def tangential_eigen_value_rt(self, r, theta, centre=(0.0, 0.0)):
-        y = r * xp.sin(theta) + centre[0]
-        x = r * xp.cos(theta) + centre[1]
-        return self.tangential_eigen_value_yx(y, x)
-
-    @partial(jit, static_argnums=(0, 3))
-    def grad_r_tangential_eigen_value(self, r, theta, centre=(0.0, 0.0)):
-        # ignore `self` with the `argnums` below
-        tangential_eigen_part = partial(self.tangential_eigen_value_rt, centre=centre)
-        return xp.vectorize(
-            jax.jacfwd(tangential_eigen_part, argnums=(0,)), signature="(),()->()"
-        )(r, theta)[0]
-
-    @partial(jit, static_argnums=(0,))
-    def radial_eigen_value_yx(self, y, x):
-        K, mag_shear = self.convergence_mag_shear_yx(y, x)
-        return 1 - K + mag_shear
-
-    @partial(jit, static_argnums=(0, 3))
-    def radial_eigen_value_rt(self, r, theta, centre=(0.0, 0.0)):
-        y = r * xp.sin(theta) + centre[0]
-        x = r * xp.cos(theta) + centre[1]
-        return self.radial_eigen_value_yx(y, x)
-
-    @partial(jit, static_argnums=(0, 3))
-    def grad_r_radial_eigen_value(self, r, theta, centre=(0.0, 0.0)):
-        # ignore `self` with the `argnums` below
-        radial_eigen_part = partial(self.radial_eigen_value_rt, centre=centre)
-        return xp.vectorize(
-            jax.jacfwd(radial_eigen_part, argnums=(0,)), signature="(),()->()"
-        )(r, theta)[0]
-
-    def tangential_critical_curve_jax(
-        self,
-        init_r=0.1,
-        init_centre=(0.0, 0.0),
-        n_points=300,
-        n_steps=20,
-        threshold=1e-5,
-    ):
-        """
-        Returns all tangential critical curves of the lensing system, which are computed as follows:
-
-        1) Create a set of `n_points` initial points in a circle of radius `init_r` and centred on `init_centre`
-        2) Apply `n_steps` of Newton's method to these points in the "radial" direction only (i.e. keeping angle fixed).
-        Jax's auto differentiation is used to find the radial derivatives of the tangential eigen value function for
-        this step.
-        3) Filter the results and only keep point that have their tangential eigen value `threshold` of 0
-
-        No underlying grid is needed for the method, but the quality of the results are dependent on the initial
-        circle of points.
-
-        Parameters
-        ----------
-        init_r : float
-            Radius of the circle of initial guess points
-        init_centre : tuple
-            centre of the circle of initial guess points as `(y, x)`
-        n_points : Int
-            Number of initial guess points to create (evenly spaced in angle around `init_centre`)
-        n_steps : Int
-            Number of iterations of Newton's method to apply
-        threshold : float
-            Only keep points whose tangential eigen value is within this value of zero (inclusive)
-        """
-        r = xp.ones(n_points) * init_r
-        theta = xp.linspace(0, 2 * xp.pi, n_points + 1)[:-1]
-        new_yx = step_r(
-            r,
-            theta,
-            jax.tree_util.Partial(self.tangential_eigen_value_rt, centre=init_centre),
-            jax.tree_util.Partial(
-                self.grad_r_tangential_eigen_value, centre=init_centre
-            ),
-            n_steps,
-        )
-        new_yx = new_yx + xp.array(init_centre)
-        # filter out nan values
-        fdx = xp.isfinite(new_yx).all(axis=1)
-        new_yx = new_yx[fdx]
-        # filter out failed points
-        value = xp.abs(self.tangential_eigen_value_yx(new_yx[:, 0], new_yx[:, 1]))
-        gdx = value <= threshold
-        return aa.structures.grids.irregular_2d.Grid2DIrregular(values=new_yx[gdx])
-
-    def radial_critical_curve_jax(
-        self,
-        init_r=0.01,
-        init_centre=(0.0, 0.0),
-        n_points=300,
-        n_steps=20,
-        threshold=1e-5,
-    ):
-        """
-        Returns all radial critical curves of the lensing system, which are computed as follows:
-
-        1) Create a set of `n_points` initial points in a circle of radius `init_r` and centred on `init_centre`
-        2) Apply `n_steps` of Newton's method to these points in the "radial" direction only (i.e. keeping angle fixed).
-        Jax's auto differentiation is used to find the radial derivatives of the radial eigen value function for
-        this step.
-        3) Filter the results and only keep point that have their radial eigen value `threshold` of 0
-
-        No underlying grid is needed for the method, but the quality of the results are dependent on the initial
-        circle of points.
-
-        Parameters
-        ----------
-        init_r : float
-            Radius of the circle of initial guess points
-        init_centre : tuple
-            centre of the circle of initial guess points as `(y, x)`
-        n_points : Int
-            Number of initial guess points to create (evenly spaced in angle around `init_centre`)
-        n_steps : Int
-            Number of iterations of Newton's method to apply
-        threshold : float
-            Only keep points whose radial eigen value is within this value of zero (inclusive)
-        """
-        r = xp.ones(n_points) * init_r
-        theta = xp.linspace(0, 2 * xp.pi, n_points + 1)[:-1]
-        new_yx = step_r(
-            r,
-            theta,
-            jax.tree_util.Partial(self.radial_eigen_value_rt, centre=init_centre),
-            jax.tree_util.Partial(self.grad_r_radial_eigen_value, centre=init_centre),
-            n_steps,
-        )
-        new_yx = new_yx + xp.array(init_centre)
-        # filter out nan values
-        fdx = xp.isfinite(new_yx).all(axis=1)
-        new_yx = new_yx[fdx]
-        # filter out failed points
-        value = xp.abs(self.radial_eigen_value_yx(new_yx[:, 0], new_yx[:, 1]))
-        gdx = value <= threshold
-        return aa.structures.grids.irregular_2d.Grid2DIrregular(values=new_yx[gdx])
-
     def jacobian_from(self, grid):
         """
         Returns the Jacobian of the lensing object, which is computed by taking the gradient of the 2D deflection
@@ -952,24 +797,36 @@ def jacobian_from(self, grid):
         grid
             The 2D grid of (y,x) arc-second coordinates the deflection angles and Jacobian are computed on.
         """
-        A = self.jacobian_stack_vector(
-            grid.array[:, 0], grid.array[:, 1], grid.pixel_scales
+
+        deflections = self.deflections_yx_2d_from(grid=grid)
+
+        # TODO : Can probably make this work on irregular grid? Is there any point?
+
+        a11 = aa.Array2D(
+            values=1.0
+            - np.gradient(deflections.native[:, :, 1], grid.native[0, :, 1], axis=1),
+            mask=grid.mask,
         )
-        a = xp.eye(2).reshape(1, 2, 2) - A
-        return [
-            [
-                aa.Array2D(values=a[..., 1, 1], mask=grid.mask),
-                aa.Array2D(values=a[..., 1, 0], mask=grid.mask),
-            ],
-            [
-                aa.Array2D(values=a[..., 0, 1], mask=grid.mask),
-                aa.Array2D(values=a[..., 0, 0], mask=grid.mask),
-            ],
-        ]
 
-        # transpose the result
-        # use `moveaxis` as grid might not be nx2
-        # return xp.moveaxis(xp.moveaxis(a, -1, 0), -1, 0)
+        a12 = aa.Array2D(
+            values=-1.0
+            * np.gradient(deflections.native[:, :, 1], grid.native[:, 0, 0], axis=0),
+            mask=grid.mask,
+        )
+
+        a21 = aa.Array2D(
+            values=-1.0
+            * np.gradient(deflections.native[:, :, 0], grid.native[0, :, 1], axis=1),
+            mask=grid.mask,
+        )
+
+        a22 = aa.Array2D(
+            values=1
+            - np.gradient(deflections.native[:, :, 0], grid.native[:, 0, 0], axis=0),
+            mask=grid.mask,
+        )
+
+        return [[a11, a12], [a21, a22]]
 
     @precompute_jacobian
     def convergence_2d_via_jacobian_from(self, grid, jacobian=None) -> aa.Array2D:
@@ -1018,9 +875,10 @@ def shear_yx_2d_via_jacobian_from(
         jacobian
             A precomputed lensing jacobian, which is passed throughout the `CalcLens` functions for efficiency.
         """
-        shear_y = -0.5 * (jacobian[0][1] + jacobian[1][0]).array
-        shear_x = 0.5 * (jacobian[1][1] - jacobian[0][0]).array
-        shear_yx_2d = xp.stack([shear_y, shear_x]).T
+
+        shear_yx_2d = np.zeros(shape=(grid.shape_slim, 2))
+        shear_yx_2d[:, 0] = -0.5 * (jacobian[0][1] + jacobian[1][0])
+        shear_yx_2d[:, 1] = 0.5 * (jacobian[1][1] - jacobian[0][0])
 
         if isinstance(grid, aa.Grid2DIrregular):
             return ShearYX2DIrregular(values=shear_yx_2d, grid=grid)
diff --git a/autogalaxy/profiles/mass/stellar/sersic_core.py b/autogalaxy/profiles/mass/stellar/sersic_core.py
index 47684f52e..ec39dec43 100644
--- a/autogalaxy/profiles/mass/stellar/sersic_core.py
+++ b/autogalaxy/profiles/mass/stellar/sersic_core.py
@@ -111,7 +111,7 @@ def decompose_convergence_via_mge(self):
         def core_sersic_2D(r):
             return (
                 self.mass_to_light_ratio
-                * self.intensity_prime(xp)
+                * self.intensity_prime()
                 * (1.0 + (self.radius_break / r) ** self.alpha)
                 ** (self.gamma / self.alpha)
                 * np.exp(
diff --git a/test_autogalaxy/operate/test_deflections.py b/test_autogalaxy/operate/test_deflections.py
index 67749599d..bd90f734f 100644
--- a/test_autogalaxy/operate/test_deflections.py
+++ b/test_autogalaxy/operate/test_deflections.py
@@ -224,36 +224,38 @@ def test__tangential_critical_curve_list_from():
     assert 0.97 < x_centre < 1.03
 
 
-def test__tangential_critical_curve_list_from__compare_via_magnification():
-    grid = ag.Grid2D.uniform(shape_native=(50, 50), pixel_scales=0.2)
-
-    mp = ag.mp.Isothermal(
-        centre=(0.0, 0.0), einstein_radius=2, ell_comps=(0.109423, -0.019294)
-    )
-
-    tangential_critical_curve_via_magnification = critical_curve_via_magnification_from(
-        mass_profile=mp, grid=grid
-    )[0]
-
-    tangential_critical_curve_list = mp.tangential_critical_curve_list_from(
-        grid=grid,
-    )
-
-    assert tangential_critical_curve_list[0] == pytest.approx(
-        tangential_critical_curve_via_magnification, 5e-1
-    )
-
-    tangential_critical_curve_via_magnification = critical_curve_via_magnification_from(
-        mass_profile=mp, grid=grid
-    )[0]
-
-    tangential_critical_curve_list = mp.tangential_critical_curve_list_from(
-        grid=grid,
-    )
-
-    assert tangential_critical_curve_list[0] == pytest.approx(
-        tangential_critical_curve_via_magnification, 5e-1
-    )
+# TODO : reinstate one JAX deflections in.
+
+# def test__tangential_critical_curve_list_from__compare_via_magnification():
+#     grid = ag.Grid2D.uniform(shape_native=(50, 50), pixel_scales=0.2)
+#
+#     mp = ag.mp.Isothermal(
+#         centre=(0.0, 0.0), einstein_radius=2, ell_comps=(0.109423, -0.019294)
+#     )
+#
+#     tangential_critical_curve_via_magnification = critical_curve_via_magnification_from(
+#         mass_profile=mp, grid=grid
+#     )[0]
+#
+#     tangential_critical_curve_list = mp.tangential_critical_curve_list_from(
+#         grid=grid,
+#     )
+#
+#     assert tangential_critical_curve_list[0] == pytest.approx(
+#         tangential_critical_curve_via_magnification, 5e-1
+#     )
+#
+#     tangential_critical_curve_via_magnification = critical_curve_via_magnification_from(
+#         mass_profile=mp, grid=grid
+#     )[0]
+#
+#     tangential_critical_curve_list = mp.tangential_critical_curve_list_from(
+#         grid=grid,
+#     )
+#
+#     assert tangential_critical_curve_list[0] == pytest.approx(
+#         tangential_critical_curve_via_magnification, 5e-1
+#     )
 
 
 def test__radial_critical_curve_list_from():
@@ -322,25 +324,26 @@ def test__tangential_caustic_list_from():
     assert 0.47 < y_centre < 0.53
     assert 0.97 < x_centre < 1.03
 
-
-def test__tangential_caustic_list_from___compare_via_magnification():
-    grid = ag.Grid2D.uniform(shape_native=(50, 50), pixel_scales=0.2)
-
-    mp = ag.mp.Isothermal(
-        centre=(0.0, 0.0), einstein_radius=2, ell_comps=(0.109423, -0.019294)
-    )
-
-    tangential_caustic_via_magnification = caustics_via_magnification_from(
-        mass_profile=mp, grid=grid
-    )[0]
-
-    tangential_caustic_list = mp.tangential_caustic_list_from(
-        grid=grid,
-    )
-
-    assert sum(tangential_caustic_list[0]) == pytest.approx(
-        sum(tangential_caustic_via_magnification), 5e-1
-    )
+# TODO : Reinstate one JAX defleciton sin.
+
+# def test__tangential_caustic_list_from___compare_via_magnification():
+#     grid = ag.Grid2D.uniform(shape_native=(50, 50), pixel_scales=0.2)
+#
+#     mp = ag.mp.Isothermal(
+#         centre=(0.0, 0.0), einstein_radius=2, ell_comps=(0.109423, -0.019294)
+#     )
+#
+#     tangential_caustic_via_magnification = caustics_via_magnification_from(
+#         mass_profile=mp, grid=grid
+#     )[0]
+#
+#     tangential_caustic_list = mp.tangential_caustic_list_from(
+#         grid=grid,
+#     )
+#
+#     assert sum(tangential_caustic_list[0]) == pytest.approx(
+#         sum(tangential_caustic_via_magnification), 5e-1
+#     )
 
 
 def test__radial_caustic_list_from():

From 4a77c1c50740d75ceb310371aeb352f801d24378 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 8 Nov 2025 15:29:16 +0000
Subject: [PATCH 09/18] grids added to DatasetQuantity

---
 autogalaxy/quantity/dataset_quantity.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/autogalaxy/quantity/dataset_quantity.py b/autogalaxy/quantity/dataset_quantity.py
index 683f3d464..7deb42dd2 100644
--- a/autogalaxy/quantity/dataset_quantity.py
+++ b/autogalaxy/quantity/dataset_quantity.py
@@ -1,10 +1,12 @@
 import logging
 import numpy as np
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import Optional, Union
 
 import autoarray as aa
+
 from autoarray.dataset.abstract.dataset import AbstractDataset
+from autoarray.dataset.grids import GridsDataset
 
 logger = logging.getLogger(__name__)
 
@@ -84,6 +86,12 @@ def __init__(
             over_sample_size_pixelization=over_sample_size_pixelization,
         )
 
+        self.grids = GridsDataset(
+            mask=self.data.mask,
+            over_sample_size_lp=self.over_sample_size_lp,
+            over_sample_size_pixelization=self.over_sample_size_pixelization,
+        )
+
     @classmethod
     def via_signal_to_noise_map(
         cls,

From b73c557d168a1439f3aec97ea49dc7ed2396d993 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 8 Nov 2025 15:41:41 +0000
Subject: [PATCH 10/18] going to take detour removing grid project

---
 autogalaxy/__init__.py                        |   1 -
 autogalaxy/ellipse/ellipse/ellipse.py         |   2 +-
 autogalaxy/ellipse/model/analysis.py          |   8 +-
 autogalaxy/galaxy/stellar_dark_decomp.py      |  44 -----
 autogalaxy/imaging/model/analysis.py          |   3 +-
 autogalaxy/interferometer/model/analysis.py   |   3 +-
 .../profiles/mass/mock/mock_mass_profile.py   |   4 +-
 autogalaxy/quantity/model/analysis.py         |   4 +-
 .../galaxy/test_stellar_dark_decomp.py        | 158 ------------------
 test_autogalaxy/imaging/test_simulator.py     |   2 +-
 10 files changed, 14 insertions(+), 215 deletions(-)
 delete mode 100644 autogalaxy/galaxy/stellar_dark_decomp.py
 delete mode 100644 test_autogalaxy/galaxy/test_stellar_dark_decomp.py

diff --git a/autogalaxy/__init__.py b/autogalaxy/__init__.py
index 5d1295df2..506a0eaec 100644
--- a/autogalaxy/__init__.py
+++ b/autogalaxy/__init__.py
@@ -80,7 +80,6 @@
 from .galaxy.galaxy import Galaxy
 from .galaxy.galaxies import Galaxies
 from .galaxy.redshift import Redshift
-from .galaxy.stellar_dark_decomp import StellarDarkDecomp
 from .galaxy.to_inversion import AbstractToInversion
 from .galaxy.to_inversion import GalaxiesToInversion
 from .profiles.geometry_profiles import EllProfile
diff --git a/autogalaxy/ellipse/ellipse/ellipse.py b/autogalaxy/ellipse/ellipse/ellipse.py
index 8986d085f..b742af496 100644
--- a/autogalaxy/ellipse/ellipse/ellipse.py
+++ b/autogalaxy/ellipse/ellipse/ellipse.py
@@ -50,7 +50,7 @@ def ellipticity(self) -> float:
         """
         The ellipticity of the ellipse, which is the factor by which the ellipse is offset from a circle.
         """
-        return np.sqrt(1 - self.axis_ratio**2.0)
+        return np.sqrt(1 - self.axis_ratio()**2.0)
 
     @property
     def minor_axis(self):
diff --git a/autogalaxy/ellipse/model/analysis.py b/autogalaxy/ellipse/model/analysis.py
index 80046e6b9..50ad45d9c 100644
--- a/autogalaxy/ellipse/model/analysis.py
+++ b/autogalaxy/ellipse/model/analysis.py
@@ -1,8 +1,6 @@
 import logging
-import time
-from typing import Dict, List, Optional, Tuple
-
-from autoconf.fitsable import hdu_list_for_output_from
+import numpy as np
+from typing import  List, Optional
 
 import autofit as af
 import autoarray as aa
@@ -45,7 +43,7 @@ def __init__(self, dataset: aa.Imaging, title_prefix: str = None):
         self.dataset = dataset
         self.title_prefix = title_prefix
 
-    def log_likelihood_function(self, instance: af.ModelInstance) -> float:
+    def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
         instance to the imaging dataset.
diff --git a/autogalaxy/galaxy/stellar_dark_decomp.py b/autogalaxy/galaxy/stellar_dark_decomp.py
deleted file mode 100644
index 4e5a169b4..000000000
--- a/autogalaxy/galaxy/stellar_dark_decomp.py
+++ /dev/null
@@ -1,44 +0,0 @@
-from autogalaxy import exc
-from autogalaxy.profiles.mass.dark.abstract import DarkProfile
-from autogalaxy.profiles.mass.stellar.abstract import StellarProfile
-
-
-class StellarDarkDecomp:
-    def __init__(self, galaxy):
-        self.galaxy = galaxy
-
-    def stellar_mass_angular_within_circle_from(self, radius: float):
-        if self.galaxy.has(cls=StellarProfile):
-            return sum(
-                [
-                    profile.mass_angular_within_circle_from(radius=radius)
-                    for profile in self.galaxy.cls_list_from(cls=StellarProfile)
-                ]
-            )
-        else:
-            raise exc.GalaxyException(
-                "You cannot perform a stellar mass-based calculation on a galaxy which does not have a stellar "
-                "mass-profile "
-            )
-
-    def dark_mass_angular_within_circle_from(self, radius: float):
-        if self.galaxy.has(cls=DarkProfile):
-            return sum(
-                [
-                    profile.mass_angular_within_circle_from(radius=radius)
-                    for profile in self.galaxy.cls_list_from(cls=DarkProfile)
-                ]
-            )
-        else:
-            raise exc.GalaxyException(
-                "You cannot perform a dark mass-based calculation on a galaxy which does not have a dark mass-profile"
-            )
-
-    def stellar_fraction_at_radius_from(self, radius):
-        return 1.0 - self.dark_fraction_at_radius_from(radius=radius)
-
-    def dark_fraction_at_radius_from(self, radius):
-        stellar_mass = self.stellar_mass_angular_within_circle_from(radius=radius)
-        dark_mass = self.dark_mass_angular_within_circle_from(radius=radius)
-
-        return dark_mass / (stellar_mass + dark_mass)
diff --git a/autogalaxy/imaging/model/analysis.py b/autogalaxy/imaging/model/analysis.py
index 90e3c3e47..071c7f5b1 100644
--- a/autogalaxy/imaging/model/analysis.py
+++ b/autogalaxy/imaging/model/analysis.py
@@ -1,3 +1,4 @@
+import numpy as np
 from typing import Optional
 
 import autofit as af
@@ -88,7 +89,7 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
 
         return self
 
-    def log_likelihood_function(self, instance: af.ModelInstance) -> float:
+    def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
         instance to the imaging dataset.
diff --git a/autogalaxy/interferometer/model/analysis.py b/autogalaxy/interferometer/model/analysis.py
index 9abffdb33..d2e22209c 100644
--- a/autogalaxy/interferometer/model/analysis.py
+++ b/autogalaxy/interferometer/model/analysis.py
@@ -1,4 +1,5 @@
 import logging
+import numpy as np
 from typing import Optional
 
 from autoconf.dictable import to_dict
@@ -95,7 +96,7 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
 
         return self
 
-    def log_likelihood_function(self, instance: af.ModelInstance) -> float:
+    def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
         instance to the interferometer dataset.
diff --git a/autogalaxy/profiles/mass/mock/mock_mass_profile.py b/autogalaxy/profiles/mass/mock/mock_mass_profile.py
index 906562040..26a2d36d1 100644
--- a/autogalaxy/profiles/mass/mock/mock_mass_profile.py
+++ b/autogalaxy/profiles/mass/mock/mock_mass_profile.py
@@ -24,8 +24,8 @@ def __init__(
     def convergence_2d_from(self, grid, xp=np):
         return self.convergence_2d
 
-    def potential_2d_from(self, grid):
+    def potential_2d_from(self, grid, xp=np):
         return self.potential_2d
 
-    def deflections_yx_2d_from(self, grid):
+    def deflections_yx_2d_from(self, grid, xp=np):
         return self.deflections_2d
diff --git a/autogalaxy/quantity/model/analysis.py b/autogalaxy/quantity/model/analysis.py
index 48621ba26..3e5710692 100644
--- a/autogalaxy/quantity/model/analysis.py
+++ b/autogalaxy/quantity/model/analysis.py
@@ -1,3 +1,5 @@
+import numpy as np
+
 from autoconf.dictable import to_dict
 
 import autofit as af
@@ -60,7 +62,7 @@ def __init__(
         self.func_str = func_str
         self.title_prefix = title_prefix
 
-    def log_likelihood_function(self, instance: af.ModelInstance) -> float:
+    def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
         instance to the quantity's dataset.
diff --git a/test_autogalaxy/galaxy/test_stellar_dark_decomp.py b/test_autogalaxy/galaxy/test_stellar_dark_decomp.py
deleted file mode 100644
index 616fce0ee..000000000
--- a/test_autogalaxy/galaxy/test_stellar_dark_decomp.py
+++ /dev/null
@@ -1,158 +0,0 @@
-import autogalaxy as ag
-
-from autogalaxy import exc
-
-import pytest
-import warnings
-
-
-def test__stellar_mass_angular_within_galaxy__is_sum_of_individual_profiles(
-    smp_0, smp_1
-):
-    galaxy = ag.Galaxy(
-        redshift=0.5,
-        stellar_0=smp_0,
-        non_stellar_profile=ag.mp.Isothermal(einstein_radius=1.0),
-    )
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    stellar_mass_0 = smp_0.mass_angular_within_circle_from(radius=0.5)
-
-    gal_mass = decomp.stellar_mass_angular_within_circle_from(radius=0.5)
-
-    assert stellar_mass_0 == gal_mass
-
-    galaxy = ag.Galaxy(
-        redshift=0.5,
-        stellar_0=smp_0,
-        stellar_1=smp_1,
-        non_stellar_profile=ag.mp.Isothermal(einstein_radius=1.0),
-    )
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    stellar_mass_1 = smp_1.mass_angular_within_circle_from(radius=0.5)
-
-    gal_mass = decomp.stellar_mass_angular_within_circle_from(radius=0.5)
-
-    assert stellar_mass_0 + stellar_mass_1 == gal_mass
-
-    galaxy = ag.Galaxy(redshift=0.5)
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    with pytest.raises(exc.GalaxyException):
-        decomp.stellar_mass_angular_within_circle_from(radius=1.0)
-
-
-@pytest.mark.filterwarnings("ignore")
-def test__stellar_fraction_at_radius(dmp_0, dmp_1, smp_0, smp_1):
-    galaxy = ag.Galaxy(redshift=0.5, stellar_0=smp_0, dark_0=dmp_0)
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    stellar_mass_0 = smp_0.mass_angular_within_circle_from(radius=1.0)
-    dark_mass_0 = dmp_0.mass_angular_within_circle_from(radius=1.0)
-
-    stellar_fraction = decomp.stellar_fraction_at_radius_from(radius=1.0)
-
-    assert stellar_fraction == pytest.approx(
-        stellar_mass_0 / (dark_mass_0 + stellar_mass_0), 1.0e-4
-    )
-
-    galaxy = ag.Galaxy(redshift=0.5, stellar_0=smp_0, stellar_1=smp_1, dark_0=dmp_0)
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    stellar_fraction = decomp.stellar_fraction_at_radius_from(radius=1.0)
-    stellar_mass_1 = smp_1.mass_angular_within_circle_from(radius=1.0)
-
-    assert stellar_fraction == pytest.approx(
-        (stellar_mass_0 + stellar_mass_1)
-        / (dark_mass_0 + stellar_mass_0 + stellar_mass_1),
-        1.0e-4,
-    )
-
-    galaxy = ag.Galaxy(
-        redshift=0.5, stellar_0=smp_0, stellar_1=smp_1, dark_0=dmp_0, dark_mass_1=dmp_1
-    )
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    stellar_fraction = decomp.stellar_fraction_at_radius_from(radius=1.0)
-    dark_mass_1 = dmp_1.mass_angular_within_circle_from(radius=1.0)
-
-    assert stellar_fraction == pytest.approx(
-        (stellar_mass_0 + stellar_mass_1)
-        / (dark_mass_0 + dark_mass_1 + stellar_mass_0 + stellar_mass_1),
-        1.0e-4,
-    )
-
-
-@pytest.mark.filterwarnings("ignore")
-def test__dark_mass_within_galaxy__is_sum_of_individual_profiles(dmp_0, dmp_1):
-    galaxy = ag.Galaxy(
-        redshift=0.5,
-        dark_0=dmp_0,
-        non_dark_profile=ag.mp.Isothermal(einstein_radius=1.0),
-    )
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    dark_mass_0 = dmp_0.mass_angular_within_circle_from(radius=0.5)
-
-    gal_mass = decomp.dark_mass_angular_within_circle_from(radius=0.5)
-
-    assert dark_mass_0 == gal_mass
-
-    galaxy = ag.Galaxy(
-        redshift=0.5,
-        dark_0=dmp_0,
-        dark_1=dmp_1,
-        non_dark_profile=ag.mp.Isothermal(einstein_radius=1.0),
-    )
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    dark_mass_1 = dmp_1.mass_angular_within_circle_from(radius=0.5)
-
-    gal_mass = decomp.dark_mass_angular_within_circle_from(radius=0.5)
-
-    assert dark_mass_0 + dark_mass_1 == gal_mass
-
-    galaxy = ag.Galaxy(redshift=0.5)
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    with pytest.raises(exc.GalaxyException):
-        decomp.dark_mass_angular_within_circle_from(radius=1.0)
-
-
-@pytest.mark.filterwarnings("ignore")
-def test__dark_fraction_at_radius(dmp_0, dmp_1, smp_0, smp_1):
-    galaxy = ag.Galaxy(redshift=0.5, dark_0=dmp_0, stellar_0=smp_0)
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    stellar_mass_0 = smp_0.mass_angular_within_circle_from(radius=1.0)
-    dark_mass_0 = dmp_0.mass_angular_within_circle_from(radius=1.0)
-
-    dark_fraction = decomp.dark_fraction_at_radius_from(radius=1.0)
-
-    assert dark_fraction == dark_mass_0 / (stellar_mass_0 + dark_mass_0)
-
-    galaxy = ag.Galaxy(redshift=0.5, dark_0=dmp_0, dark_1=dmp_1, stellar_0=smp_0)
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    dark_fraction = decomp.dark_fraction_at_radius_from(radius=1.0)
-    dark_mass_1 = dmp_1.mass_angular_within_circle_from(radius=1.0)
-
-    assert dark_fraction == pytest.approx(
-        (dark_mass_0 + dark_mass_1) / (stellar_mass_0 + dark_mass_0 + dark_mass_1),
-        1.0e-4,
-    )
-
-    galaxy = ag.Galaxy(
-        redshift=0.5, dark_0=dmp_0, dark_1=dmp_1, stellar_0=smp_0, stellar_mass_1=smp_1
-    )
-    decomp = ag.StellarDarkDecomp(galaxy=galaxy)
-
-    dark_fraction = decomp.dark_fraction_at_radius_from(radius=1.0)
-    stellar_mass_1 = smp_1.mass_angular_within_circle_from(radius=1.0)
-
-    assert dark_fraction == pytest.approx(
-        (dark_mass_0 + dark_mass_1)
-        / (stellar_mass_0 + stellar_mass_1 + dark_mass_0 + dark_mass_1),
-        1.0e-4,
-    )
diff --git a/test_autogalaxy/imaging/test_simulator.py b/test_autogalaxy/imaging/test_simulator.py
index cc2fc99fe..8f7f2b2ed 100644
--- a/test_autogalaxy/imaging/test_simulator.py
+++ b/test_autogalaxy/imaging/test_simulator.py
@@ -105,7 +105,7 @@ def test__simulator__via_galaxies_from():
 
     assert dataset.shape_native == (20, 20)
     assert dataset.data.native[0, 0] != imaging_via_image.data.native[0, 0]
-    assert dataset.data.native[10, 10] == imaging_via_image.data.native[10, 10]
+    assert dataset.data.native[10, 10] == pytest.approx(imaging_via_image.data.native[10, 10], 1.0e-4)
     assert dataset.psf == pytest.approx(imaging_via_image.psf, 1.0e-4)
     assert dataset.noise_map == pytest.approx(imaging_via_image.noise_map, 1.0e-4)
 

From a9f49f62a2a36ebcd84276c601cb9011e4dbc423 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 8 Nov 2025 16:09:45 +0000
Subject: [PATCH 11/18] light profiles use much simpler API

---
 autogalaxy/galaxy/galaxy.py                   | 100 ------------------
 autogalaxy/profiles/geometry_profiles.py      |  25 -----
 autogalaxy/profiles/light/abstract.py         |  27 -----
 autogalaxy/profiles/mass/abstract/abstract.py |   8 --
 autogalaxy/profiles/mass/abstract/cse.py      |   1 -
 .../profiles/plot/light_profile_plotters.py   |  16 ++-
 test_autogalaxy/galaxy/test_galaxy.py         |  63 -----------
 .../profiles/mass/abstract/test_abstract.py   |  26 -----
 8 files changed, 12 insertions(+), 254 deletions(-)

diff --git a/autogalaxy/galaxy/galaxy.py b/autogalaxy/galaxy/galaxy.py
index 4e9ef2014..fde7ec706 100644
--- a/autogalaxy/galaxy/galaxy.py
+++ b/autogalaxy/galaxy/galaxy.py
@@ -231,39 +231,6 @@ def image_2d_from(
             return sum(self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only))
         return xp.zeros((grid.shape[0],))
 
-    @aa.grid_dec.to_projected
-    def image_1d_from(self, grid: aa.type.Grid2DLike, xp=np) -> np.ndarray:
-        """
-        Returns the summed 1D image of the galaxy's light profiles using a grid of Cartesian (y,x) coordinates.
-
-        If the galaxy has no light profiles, a grid of zeros is returned.
-
-        See `profiles.light` module for details of how this is performed.
-
-        The decorator `to_projected` converts the output arrays from ndarrays to an `Array1D` data
-        structure using the input `grid`'s attributes.
-
-        Parameters
-        ----------
-        grid
-            The 1D (x,) coordinates where values of the image are evaluated.
-        """
-        if self.has(cls=LightProfile):
-            image_1d_list = []
-
-            for light_profile in self.cls_list_from(
-                cls=LightProfile, cls_filtered=LightProfileLinear
-            ):
-                grid_radial = self.grid_radial_from(
-                    grid=grid, centre=light_profile.centre, angle=light_profile.angle
-                )
-
-                image_1d_list.append(light_profile.image_1d_from(grid=grid_radial))
-
-            return sum(image_1d_list)
-
-        return xp.zeros((grid.shape[0],))
-
     @aa.grid_dec.to_vector_yx
     def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
         """
@@ -337,40 +304,6 @@ def traced_grid_2d_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.type.Grid2D
 
         return grid - self.deflections_yx_2d_from(grid=grid, xp=xp)
 
-    @aa.grid_dec.to_projected
-    def convergence_1d_from(self, grid: aa.type.Grid1D2DLike, xp=np) -> np.ndarray:
-        """
-        Returns the summed 1D convergence of the galaxy's mass profiles using a grid of Cartesian (y,x) coordinates.
-
-        If the galaxy has no mass profiles, a grid of zeros is returned.
-
-        See `profiles.mass` module for details of how this is performed.
-
-        The decorator `to_projected` converts the output arrays from ndarrays to an `Array1D` data
-        structure using the input `grid`'s attributes.
-
-        Parameters
-        ----------
-        grid
-            The 1D (x,) coordinates where values of the convergence are evaluated.
-        """
-        if self.has(cls=MassProfile):
-            convergence_1d_list = []
-
-            for mass_profile in self.cls_list_from(cls=MassProfile):
-
-                grid_radial = self.grid_radial_from(
-                    grid=grid, centre=mass_profile.centre, angle=mass_profile.angle
-                )
-
-                convergence_1d_list.append(
-                    mass_profile.convergence_1d_from(grid=grid_radial)
-                )
-
-            return sum(convergence_1d_list)
-
-        return xp.zeros((grid.shape[0],))
-
     @aa.grid_dec.to_array
     def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
         """
@@ -398,39 +331,6 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.nda
             )
         return xp.zeros((grid.shape[0],))
 
-    @aa.grid_dec.to_projected
-    def potential_1d_from(self, grid: aa.type.Grid2DLike, xp=np) -> np.ndarray:
-        """
-        Returns the summed 1D potential of the galaxy's mass profiles using a grid of Cartesian (y,x) coordinates.
-
-        If the galaxy has no mass profiles, a grid of zeros is returned.
-
-        See `profiles.mass` module for details of how this is performed.
-
-        The decorator `to_projected` converts the output arrays from ndarrays to an `Array1D` data
-        structure using the input `grid`'s attributes.
-
-        Parameters
-        ----------
-        grid
-            The 1D (x,) coordinates where values of the potential are evaluated.
-        """
-        if self.has(cls=MassProfile):
-            potential_1d_list = []
-
-            for mass_profile in self.cls_list_from(cls=MassProfile):
-                grid_radial = self.grid_radial_from(
-                    grid=grid, centre=mass_profile.centre, angle=mass_profile.angle
-                )
-
-                potential_1d_list.append(
-                    mass_profile.potential_1d_from(grid=grid_radial)
-                )
-
-            return sum(potential_1d_list)
-
-        return xp.zeros((grid.shape[0],))
-
     @property
     def half_light_radius(self):
         return None
diff --git a/autogalaxy/profiles/geometry_profiles.py b/autogalaxy/profiles/geometry_profiles.py
index d91fa3fae..ad9cfc6bc 100644
--- a/autogalaxy/profiles/geometry_profiles.py
+++ b/autogalaxy/profiles/geometry_profiles.py
@@ -45,31 +45,6 @@ def transformed_to_reference_frame_grid_from(self, grid, xp=np, **kwargs):
     def transformed_from_reference_frame_grid_from(self, grid, xp=np, **kwargs):
         raise NotImplemented()
 
-    def _radial_projected_shape_slim_from(
-        self, grid: aa.type.Grid1D2DLike, **kwargs
-    ) -> int:
-        """
-        To make 1D plots (e.g. `image_1d_from()`) from an input 2D grid, one uses that 2D grid to radially project
-        the coordinates across the profile's major-axis.
-
-        This function computes the distance from the profile centre to the edge of this 2D grid.
-
-        If a 1D grid is input it returns the shape of this grid, as the grid itself defines the radial coordinates.
-
-        Parameters
-        ----------
-        grid
-            A 1D or 2D grid from which a 1D plot of the profile is to be created.
-        """
-
-        if isinstance(grid, aa.Grid1D):
-            return grid.sub_shape_slim
-        elif isinstance(grid, aa.Grid2DIrregular):
-            return grid.slim.shape[0]
-
-        return grid.grid_2d_radial_projected_shape_slim_from(centre=self.centre)
-
-
 class SphProfile(GeometryProfile):
     """
     A spherical profile, which describes profiles with y and x centre Cartesian coordinates.
diff --git a/autogalaxy/profiles/light/abstract.py b/autogalaxy/profiles/light/abstract.py
index aa222f273..fd3f7e349 100644
--- a/autogalaxy/profiles/light/abstract.py
+++ b/autogalaxy/profiles/light/abstract.py
@@ -75,33 +75,6 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
         """
         raise NotImplementedError()
 
-    @aa.grid_dec.project_grid
-    def image_1d_from(
-        self, grid: aa.type.Grid1D2DLike, xp=np, **kwargs
-    ) -> aa.type.Grid1D2DLike:
-        """
-        Returns the light profile's 1D image from a grid of Cartesian coordinates, which may have been
-        transformed using the light profile's geometry.
-
-        If a 1D grid is input the image is evaluated every coordinate on the grid. If a 2D grid is input, this is
-        converted to a 1D grid by aligning with the major-axis of the light profile's elliptical geometry.
-
-        Internally, this function uses a 2D grid to compute the image, which is mapped to a 1D data structure on return
-        via the `project_grid` decorator. This avoids code repetition by ensuring that light profiles only use
-        their `image_2d_from()`  function to evaluate their image.
-
-        Parameters
-        ----------
-        grid
-            A 1D or 2D grid of coordinates which are used to evaluate the light profile in 1D.
-
-        Returns
-        -------
-        image
-            The 1D image of the light profile evaluated at every (x,) coordinate on the 1D transformed grid.
-        """
-        return self.image_2d_from(grid=grid, xp=xp, **kwargs)
-
     def luminosity_within_circle_from(self, radius: float) -> float:
         """
         Integrate the light profile to compute the total luminosity within a circle of specified radius. This is
diff --git a/autogalaxy/profiles/mass/abstract/abstract.py b/autogalaxy/profiles/mass/abstract/abstract.py
index 883264373..eaa4aff02 100644
--- a/autogalaxy/profiles/mass/abstract/abstract.py
+++ b/autogalaxy/profiles/mass/abstract/abstract.py
@@ -50,17 +50,9 @@ def convergence_2d_from(self, grid, xp=np):
     def convergence_func(self, grid_radius: float) -> float:
         raise NotImplementedError
 
-    @aa.grid_dec.project_grid
-    def convergence_1d_from(self, grid: aa.type.Grid1D2DLike, xp=np) -> aa.type.Grid1D2DLike:
-        return self.convergence_2d_from(grid=grid, xp=xp)
-
     def potential_2d_from(self, grid):
         raise NotImplementedError
 
-    @aa.grid_dec.project_grid
-    def potential_1d_from(self, grid: aa.type.Grid1D2DLike, xp=np) -> aa.type.Grid1D2DLike:
-        return self.potential_2d_from(grid=grid, xp=xp)
-
     def potential_func(self, u, y, x):
         raise NotImplementedError
 
diff --git a/autogalaxy/profiles/mass/abstract/cse.py b/autogalaxy/profiles/mass/abstract/cse.py
index cd152ae1d..32af87396 100644
--- a/autogalaxy/profiles/mass/abstract/cse.py
+++ b/autogalaxy/profiles/mass/abstract/cse.py
@@ -1,5 +1,4 @@
 from abc import ABC, abstractmethod
-import jax.numpy as jnp
 import numpy as np
 from typing import Callable, List, Tuple
 
diff --git a/autogalaxy/profiles/plot/light_profile_plotters.py b/autogalaxy/profiles/plot/light_profile_plotters.py
index e8521bb7b..043c761c8 100644
--- a/autogalaxy/profiles/plot/light_profile_plotters.py
+++ b/autogalaxy/profiles/plot/light_profile_plotters.py
@@ -73,6 +73,13 @@ def __init__(
             visuals_1d=visuals_1d,
         )
 
+    @property
+    def grid_2d_projected(self):
+        return self.grid.grid_2d_radial_projected_from(
+            centre=self.light_profile.centre,
+            angle=self.light_profile.angle()
+        )
+
     def figures_1d(self, image: bool = False):
         """
         Plots the individual attributes of the plotter's `LightProfile` object in 1D, which are computed via the
@@ -96,11 +103,12 @@ def figures_1d(self, image: bool = False):
             plot_axis_type_override = None
 
         if image:
-            image_1d = self.light_profile.image_1d_from(grid=self.grid)
+
+            image_1d = self.light_profile.image_2d_from(grid=self.grid_2d_projected)
 
             self.mat_plot_1d.plot_yx(
                 y=image_1d,
-                x=image_1d.grid_radial,
+                x=self.grid_2d_projected[:,1],
                 visuals_1d=self.visuals_1d,
                 auto_labels=aplt.AutoLabels(
                     title=r"Image ($\mathrm{e^{-}}\,\mathrm{s^{-1}}$) vs Radius (arcsec)",
@@ -223,7 +231,7 @@ def figures_1d(self, image: bool = False):
 
         if image:
             image_1d_list = [
-                light_profile.image_1d_from(grid=self.grid)
+                light_profile.image_2d_from(grid=self.grid_2d_projected)
                 for light_profile in self.light_profile_pdf_list
             ]
 
@@ -248,7 +256,7 @@ def figures_1d(self, image: bool = False):
 
             self.mat_plot_1d.plot_yx(
                 y=median_image_1d,
-                x=image_1d_list[0].grid_radial,
+                x=self.grid_2d_projected[:,1],
                 visuals_1d=visuals_1d,
                 auto_labels=aplt.AutoLabels(
                     title=r"Image ($\mathrm{e^{-}}\,\mathrm{s^{-1}}$) vs Radius (arcsec)",
diff --git a/test_autogalaxy/galaxy/test_galaxy.py b/test_autogalaxy/galaxy/test_galaxy.py
index bb1098104..6922a1a41 100644
--- a/test_autogalaxy/galaxy/test_galaxy.py
+++ b/test_autogalaxy/galaxy/test_galaxy.py
@@ -34,17 +34,6 @@ def test__cls_list_from(lp_0, lp_linear_0):
     assert cls_list == [lp_linear_0, lp_linear_0]
 
 
-def test__image_1d_from(lp_0, lp_1, gal_x2_lp):
-    grid = ag.Grid2D.no_mask(values=[[[1.05, -0.55]]], pixel_scales=1.0)
-
-    lp_image = lp_0.image_1d_from(grid=grid)
-    lp_image += lp_1.image_1d_from(grid=grid)
-
-    gal_image = gal_x2_lp.image_1d_from(grid=grid)
-
-    assert lp_image == gal_image
-
-
 def test__image_2d_from(grid_2d_7x7, gal_x2_lp):
     lp_0_image = gal_x2_lp.light_profile_0.image_2d_from(grid=grid_2d_7x7)
     lp_1_image = gal_x2_lp.light_profile_1.image_2d_from(grid=grid_2d_7x7)
@@ -105,32 +94,6 @@ def test__luminosity_within_circle(lp_0, lp_1, gal_x2_lp):
     assert gal_no_lp.luminosity_within_circle_from(radius=1.0) == None
 
 
-def test__convergence_1d_from(grid_1d_7, mp_0, gal_x1_mp, mp_1, gal_x2_mp):
-    grid = ag.Grid2D.no_mask(values=[[[1.05, -0.55], [2.05, -0.55]]], pixel_scales=1.0)
-
-    mp_convergence = mp_0.convergence_1d_from(grid=grid)
-    mp_convergence += mp_1.convergence_1d_from(grid=grid)
-
-    gal_convergence = gal_x2_mp.convergence_1d_from(grid=grid)
-
-    assert (mp_convergence == gal_convergence).all()
-
-    # Test explicitly for a profile with an offset centre and ellipticity, given the 1D to 2D projections are nasty.
-
-    grid = ag.Grid2D.no_mask(values=[[(1.05, -0.55), (2.05, -0.55)]], pixel_scales=1.0)
-
-    elliptical_mp = ag.mp.Isothermal(
-        centre=(0.5, 1.0), ell_comps=(0.2, 0.3), einstein_radius=1.0
-    )
-
-    galaxy = ag.Galaxy(redshift=0.5, mass=elliptical_mp)
-
-    mp_convergence = elliptical_mp.convergence_1d_from(grid=grid)
-    gal_convergence = galaxy.convergence_1d_from(grid=grid)
-
-    assert (mp_convergence == gal_convergence).all()
-
-
 def test__convergence_2d_from(grid_2d_7x7, mp_0, gal_x1_mp, mp_1, gal_x2_mp):
     mp_0_convergence = gal_x2_mp.mass_profile_0.convergence_2d_from(grid=grid_2d_7x7)
 
@@ -143,32 +106,6 @@ def test__convergence_2d_from(grid_2d_7x7, mp_0, gal_x1_mp, mp_1, gal_x2_mp):
     assert gal_convergence[0] == mp_convergence[0]
 
 
-def test__potential_1d_from(grid_1d_7, mp_0, gal_x1_mp, mp_1, gal_x2_mp):
-    grid = ag.Grid2D.no_mask(values=[[[1.05, -0.55]]], pixel_scales=1.0)
-
-    mp_potential = mp_0.potential_1d_from(grid=grid)
-    mp_potential += mp_1.potential_1d_from(grid=grid)
-
-    gal_convergence = gal_x2_mp.potential_1d_from(grid=grid)
-
-    assert mp_potential == gal_convergence
-
-    # Test explicitly for a profile with an offset centre and ellipticity, given the 1D to 2D projections are nasty.
-
-    grid = ag.Grid2D.no_mask(values=[[(1.05, -0.55), (2.05, -0.55)]], pixel_scales=1.0)
-
-    elliptical_mp = ag.mp.Isothermal(
-        centre=(0.5, 1.0), ell_comps=(0.2, 0.3), einstein_radius=1.0
-    )
-
-    galaxy = ag.Galaxy(redshift=0.5, mass=elliptical_mp)
-
-    mp_potential = elliptical_mp.potential_1d_from(grid=grid)
-    gal_mp_potential = galaxy.potential_1d_from(grid=grid)
-
-    assert (mp_potential == gal_mp_potential).all()
-
-
 def test__potential_2d_from(grid_2d_7x7, gal_x2_mp):
     mp_0_potential = gal_x2_mp.mass_profile_0.potential_2d_from(grid=grid_2d_7x7)
 
diff --git a/test_autogalaxy/profiles/mass/abstract/test_abstract.py b/test_autogalaxy/profiles/mass/abstract/test_abstract.py
index 2555af98c..7532e48ea 100644
--- a/test_autogalaxy/profiles/mass/abstract/test_abstract.py
+++ b/test_autogalaxy/profiles/mass/abstract/test_abstract.py
@@ -203,29 +203,3 @@ def test__regression__centre_of_profile_in_right_place():
     assert deflections.native[2, 4, 0] < 0
     assert deflections.native[1, 4, 1] > 0
     assert deflections.native[1, 3, 1] < 0
-
-
-
-def test__decorators__convergence_1d_from__grid_1d_in__returns_1d_quantities_via_projection():
-    grid_1d = ag.Grid1D.no_mask(values=[1.0, 2.0, 3.0], pixel_scales=1.0)
-
-    sie = ag.mp.Isothermal(centre=(0.0, 0.0), ell_comps=(0.0, 0.0), einstein_radius=1.0)
-
-    convergence_1d = sie.convergence_1d_from(grid=grid_1d)
-    convergence_2d = sie.convergence_2d_from(grid=grid_1d)
-
-    assert convergence_1d[0] == pytest.approx(convergence_2d[0].array, 1.0e-4)
-    assert convergence_1d[1] == pytest.approx(convergence_2d[1].array, 1.0e-4)
-    assert convergence_1d[2] == pytest.approx(convergence_2d[2].array, 1.0e-4)
-
-    sie = ag.mp.Isothermal(centre=(0.5, 0.5), ell_comps=(0.2, 0.2), einstein_radius=1.0)
-
-    convergence_1d = sie.convergence_1d_from(grid=grid_1d)
-
-    grid_2d_radial = grid_1d.grid_2d_radial_projected_from(angle=sie.angle + 90.0)
-
-    convergence_2d = sie.convergence_2d_from(grid=grid_2d_radial)
-
-    assert convergence_1d[0] == pytest.approx(convergence_2d[0].array, 1.0e-4)
-    assert convergence_1d[1] == pytest.approx(convergence_2d[1].array, 1.0e-4)
-    assert convergence_1d[2] == pytest.approx(convergence_2d[2].array, 1.0e-4)
\ No newline at end of file

From 03c4ce7965449c3e1a3c6d3085a3560014bc091c Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 8 Nov 2025 16:29:22 +0000
Subject: [PATCH 12/18] mnass profile plotter tests pass

---
 .../profiles/plot/light_profile_plotters.py   | 17 ++++--
 .../profiles/plot/mass_profile_plotters.py    | 52 ++++++++++++-------
 2 files changed, 44 insertions(+), 25 deletions(-)

diff --git a/autogalaxy/profiles/plot/light_profile_plotters.py b/autogalaxy/profiles/plot/light_profile_plotters.py
index 043c761c8..a611603f6 100644
--- a/autogalaxy/profiles/plot/light_profile_plotters.py
+++ b/autogalaxy/profiles/plot/light_profile_plotters.py
@@ -230,10 +230,17 @@ def figures_1d(self, image: bool = False):
             plot_axis_type_override = None
 
         if image:
-            image_1d_list = [
-                light_profile.image_2d_from(grid=self.grid_2d_projected)
-                for light_profile in self.light_profile_pdf_list
-            ]
+
+            image_1d_list = []
+
+            for light_profile in self.light_profile_pdf_list:
+
+                grid = self.grid.grid_2d_radial_projected_from(
+                    centre=light_profile.centre,
+                    angle=light_profile.angle()
+                )
+
+                image_1d_list.append(light_profile.image_2d_from(grid=grid))
 
             min_index = min([image_1d.shape[0] for image_1d in image_1d_list])
             image_1d_list = [image_1d[0:min_index] for image_1d in image_1d_list]
@@ -256,7 +263,7 @@ def figures_1d(self, image: bool = False):
 
             self.mat_plot_1d.plot_yx(
                 y=median_image_1d,
-                x=self.grid_2d_projected[:,1],
+                x=grid[0:min_index,1],
                 visuals_1d=visuals_1d,
                 auto_labels=aplt.AutoLabels(
                     title=r"Image ($\mathrm{e^{-}}\,\mathrm{s^{-1}}$) vs Radius (arcsec)",
diff --git a/autogalaxy/profiles/plot/mass_profile_plotters.py b/autogalaxy/profiles/plot/mass_profile_plotters.py
index c2d9ea4bd..2fb8b0278 100644
--- a/autogalaxy/profiles/plot/mass_profile_plotters.py
+++ b/autogalaxy/profiles/plot/mass_profile_plotters.py
@@ -71,6 +71,13 @@ def __init__(
 
         self.figures_2d = self._mass_plotter.figures_2d
 
+    @property
+    def grid_2d_projected(self):
+        return self.grid.grid_2d_radial_projected_from(
+            centre=self.mass_profile.centre,
+            angle=self.mass_profile.angle()
+        )
+
     def figures_1d(self, convergence: bool = False, potential: bool = False):
         """
         Plots the individual attributes of the plotter's `MassProfile` object in 1D, which are computed via the
@@ -96,11 +103,11 @@ def figures_1d(self, convergence: bool = False, potential: bool = False):
             plot_axis_type_override = None
 
         if convergence:
-            convergence_1d = self.mass_profile.convergence_1d_from(grid=self.grid)
+            convergence_1d = self.mass_profile.convergence_2d_from(grid=self.grid_2d_projected)
 
             self.mat_plot_1d.plot_yx(
                 y=convergence_1d,
-                x=convergence_1d.grid_radial,
+                x=self.grid_2d_projected[:,1],
                 visuals_1d=self.visuals_1d,
                 auto_labels=aplt.AutoLabels(
                     title="Convergence vs Radius (arcsec)",
@@ -112,11 +119,11 @@ def figures_1d(self, convergence: bool = False, potential: bool = False):
             )
 
         if potential:
-            potential_1d = self.mass_profile.potential_1d_from(grid=self.grid)
+            potential_1d = self.mass_profile.potential_2d_from(grid=self.grid_2d_projected)
 
             self.mat_plot_1d.plot_yx(
                 y=potential_1d,
-                x=potential_1d.grid_radial,
+                x=self.grid_2d_projected[:,1],
                 visuals_1d=self.visuals_1d,
                 auto_labels=aplt.AutoLabels(
                     title="Potential vs Radius (arcsec)",
@@ -216,10 +223,17 @@ def figures_1d(self, convergence=False, potential=False):
             plot_axis_type_override = None
 
         if convergence:
-            convergence_1d_list = [
-                mass_profile.convergence_1d_from(grid=self.grid)
-                for mass_profile in self.mass_profile_pdf_list
-            ]
+
+            convergence_1d_list = []
+
+            for mass_profile in self.mass_profile_pdf_list:
+                
+                grid = self.grid.grid_2d_radial_projected_from(
+                    centre=mass_profile.centre,
+                    angle=mass_profile.angle()
+                )
+           
+                convergence_1d_list.append(mass_profile.convergence_2d_from(grid=grid))
 
             min_index = min(
                 [convergence_1d.shape[0] for convergence_1d in convergence_1d_list]
@@ -248,7 +262,7 @@ def figures_1d(self, convergence=False, potential=False):
 
             self.mat_plot_1d.plot_yx(
                 y=median_convergence_1d,
-                x=convergence_1d_list[0].grid_radial,
+                x=grid[0:min_index,1],
                 visuals_1d=visuals_1d,
                 auto_labels=aplt.AutoLabels(
                     title="Convergence vs Radius (arcsec)",
@@ -260,17 +274,15 @@ def figures_1d(self, convergence=False, potential=False):
             )
 
         if potential:
-            potential_1d_list = [
-                mass_profile.potential_1d_from(grid=self.grid)
-                for mass_profile in self.mass_profile_pdf_list
-            ]
+            potential_1d_list = []
 
-            min_index = min(
-                [potential_1d.shape[0] for potential_1d in potential_1d_list]
-            )
-            potential_1d_list = [
-                potential_1d[0:min_index] for potential_1d in potential_1d_list
-            ]
+            for mass_profile in self.mass_profile_pdf_list:
+                grid = self.grid.grid_2d_radial_projected_from(
+                    centre=mass_profile.centre,
+                    angle=mass_profile.angle()
+                )
+
+                potential_1d_list.append(mass_profile.potential_2d_from(grid=grid))
 
             (
                 median_potential_1d,
@@ -290,7 +302,7 @@ def figures_1d(self, convergence=False, potential=False):
 
             self.mat_plot_1d.plot_yx(
                 y=median_potential_1d,
-                x=potential_1d_list[0].grid_radial,
+                x=grid[0:min_index,1],
                 visuals_1d=visuals_1d,
                 auto_labels=aplt.AutoLabels(
                     title="Potential vs Radius (arcsec)",

From 74f5f4fbb3a0bb5a2010962e8c390eee14da0507 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 8 Nov 2025 17:45:07 +0000
Subject: [PATCH 13/18] removeed figue 1d stuff

---
 autogalaxy/analysis/plotter_interface.py      |  12 -
 autogalaxy/config/visualize/plots.yaml        |   2 -
 .../ellipse/plot/fit_ellipse_plotters.py      |   2 -
 autogalaxy/galaxy/plot/galaxies_plotters.py   |  51 --
 autogalaxy/galaxy/plot/galaxy_plotters.py     | 637 ------------------
 autogalaxy/plot/__init__.py                   |   3 -
 .../profiles/plot/light_profile_plotters.py   | 178 -----
 .../profiles/plot/mass_profile_plotters.py    | 234 -------
 docs/api/plot.rst                             |   1 -
 .../analysis/test_plotter_interface.py        |   4 -
 test_autogalaxy/config/visualize.yaml         |   2 -
 .../galaxy/plot/test_galaxies_plotter.py      |  10 +-
 .../galaxy/plot/test_galaxy_plotters.py       | 100 ---
 .../plot/test_light_profile_plotters.py       |  41 --
 .../plot/test_mass_profile_plotters.py        |  42 --
 15 files changed, 1 insertion(+), 1318 deletions(-)

diff --git a/autogalaxy/analysis/plotter_interface.py b/autogalaxy/analysis/plotter_interface.py
index d329702db..1ee0acde9 100644
--- a/autogalaxy/analysis/plotter_interface.py
+++ b/autogalaxy/analysis/plotter_interface.py
@@ -168,18 +168,6 @@ def should_plot(name):
             mat_plot_1d=mat_plot_1d,
         )
 
-        try:
-            if should_plot("subplot_galaxies_1d"):
-                galaxies_plotter.subplot_galaxies_1d()
-        except OverflowError:
-            pass
-
-        try:
-            if should_plot("subplot_galaxies_1d_decomposed"):
-                galaxies_plotter.subplot_galaxies_1d_decomposed()
-        except OverflowError:
-            pass
-
         if should_plot("fits_galaxy_images"):
 
             image_list = [
diff --git a/autogalaxy/config/visualize/plots.yaml b/autogalaxy/config/visualize/plots.yaml
index 3d82db942..b61299e42 100644
--- a/autogalaxy/config/visualize/plots.yaml
+++ b/autogalaxy/config/visualize/plots.yaml
@@ -30,8 +30,6 @@ fit_imaging: {}                            # Settings for plots of fits to imagi
 galaxies:                                  # Settings for plots of galaxies (e.g. GalaxiesPlotter).
   subplot_galaxies: true                   # Plot subplot of all quantities in each galaxies group (e.g. images, convergence)?
   subplot_galaxy_images: false             # Plot subplot of the image of each galaxy in the model?
-  subplot_galaxies_1d: false               # Plot subplot of all quantities in 1D of each galaxies group (e.g. images, convergence)?
-  subplot_galaxies_1d_decomposed: false    # Plot subplot of all quantities in 1D decomposed of each galaxies group (e.g. images, convergence)?
   fits_galaxy_images: false                # Output a .fits file containing images of every galaxy?
 
 inversion:                                 # Settings for plots of inversions (e.g. InversionPlotter).
diff --git a/autogalaxy/ellipse/plot/fit_ellipse_plotters.py b/autogalaxy/ellipse/plot/fit_ellipse_plotters.py
index d4dbbae7a..58141729d 100644
--- a/autogalaxy/ellipse/plot/fit_ellipse_plotters.py
+++ b/autogalaxy/ellipse/plot/fit_ellipse_plotters.py
@@ -143,10 +143,8 @@ def __init__(
         super().__init__(
             mat_plot_1d=mat_plot_1d,
             visuals_1d=visuals_1d,
-            include_1d=include_1d,
             mat_plot_2d=mat_plot_2d,
             visuals_2d=visuals_2d,
-            include_2d=include_2d,
         )
 
         self.fit_pdf_list = fit_pdf_list
diff --git a/autogalaxy/galaxy/plot/galaxies_plotters.py b/autogalaxy/galaxy/plot/galaxies_plotters.py
index 56522360f..e49a8e1a7 100644
--- a/autogalaxy/galaxy/plot/galaxies_plotters.py
+++ b/autogalaxy/galaxy/plot/galaxies_plotters.py
@@ -319,54 +319,3 @@ def subplot_galaxy_images(self):
             auto_filename=f"subplot_galaxy_images"
         )
         self.close_subplot_figure()
-
-    def subplot_galaxies_1d(self):
-        """
-        Output a subplot of attributes of every individual 1D attribute of the `Galaxy` object.
-
-        For example, a 1D plot showing how the image, convergence of each component varies radially outwards.
-
-        If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is
-        computed from the light profile. This is performed by aligning a 1D grid with the  major-axis of the light
-        profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale.
-        """
-        number_subplots = len(self.galaxies) * 3
-
-        self.open_subplot_figure(number_subplots=number_subplots)
-
-        for galaxy_index in range(0, len(self.galaxies)):
-            galaxy_plotter = self.galaxy_plotter_from(galaxy_index=galaxy_index)
-
-            galaxy_plotter.figures_1d(image=True)
-            galaxy_plotter.figures_1d(convergence=True)
-            galaxy_plotter.figures_1d(potential=True)
-
-        self.mat_plot_1d.output.subplot_to_figure(auto_filename="subplot_galaxies_1d")
-        self.close_subplot_figure()
-
-    def subplot_galaxies_1d_decomposed(self):
-        """
-        Output a subplot of attributes of every individual 1D attribute of the `Galaxy` object decompoed into
-        their different light and mass profiles.
-
-        For example, a 1D plot showing how the image, convergence of each component varies radially outwards.
-
-        If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is
-        computed from the light profile. This is performed by aligning a 1D grid with the  major-axis of the light
-        profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale.
-        """
-        number_subplots = len(self.galaxies) * 3
-
-        self.open_subplot_figure(number_subplots=number_subplots)
-
-        for galaxy_index in range(0, len(self.galaxies)):
-            galaxy_plotter = self.galaxy_plotter_from(galaxy_index=galaxy_index)
-
-            galaxy_plotter.figures_1d_decomposed(image=True)
-            galaxy_plotter.figures_1d_decomposed(convergence=True)
-            galaxy_plotter.figures_1d_decomposed(potential=True)
-
-        self.mat_plot_1d.output.subplot_to_figure(
-            auto_filename="subplot_galaxies_1d_decomposed"
-        )
-        self.close_subplot_figure()
diff --git a/autogalaxy/galaxy/plot/galaxy_plotters.py b/autogalaxy/galaxy/plot/galaxy_plotters.py
index f03fbe83d..2de28b589 100644
--- a/autogalaxy/galaxy/plot/galaxy_plotters.py
+++ b/autogalaxy/galaxy/plot/galaxy_plotters.py
@@ -18,12 +18,9 @@
 
 if TYPE_CHECKING:
     from autogalaxy.profiles.plot.light_profile_plotters import LightProfilePlotter
-    from autogalaxy.profiles.plot.light_profile_plotters import LightProfilePDFPlotter
 from autogalaxy.profiles.plot.mass_profile_plotters import MassProfilePlotter
-from autogalaxy.profiles.plot.mass_profile_plotters import MassProfilePDFPlotter
 
 from autogalaxy import exc
-from autogalaxy.util import error_util
 
 
 class GalaxyPlotter(Plotter):
@@ -170,209 +167,6 @@ def mass_profile_plotter_from(
             + Visuals1D().add_einstein_radius(mass_obj=mass_profile, grid=self.grid),
         )
 
-    @property
-    def decomposed_light_profile_plotter_list(self):
-        plotter_list = []
-
-        for i, light_profile in enumerate(self.galaxy.cls_list_from(cls=LightProfile)):
-            light_profile_plotter = self.light_profile_plotter_from(
-                light_profile=light_profile, one_d_only=True
-            )
-
-            plotter_list.append(light_profile_plotter)
-
-        return [self] + plotter_list
-
-    @property
-    def decomposed_mass_profile_plotter_list(self):
-        plotter_list = []
-
-        for i, mass_profile in enumerate(self.galaxy.cls_list_from(cls=MassProfile)):
-            mass_profile_plotter = self.mass_profile_plotter_from(
-                mass_profile=mass_profile, one_d_only=True
-            )
-
-            plotter_list.append(mass_profile_plotter)
-
-        return [self] + plotter_list
-
-    def figures_1d(
-        self, image: bool = False, convergence: bool = False, potential: bool = False
-    ):
-        """
-        Plots the individual attributes of the plotter's `Galaxy` object in 1D, which are computed via the plotter's
-        grid object.
-
-        If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is
-        computed from each light profile of the galaxy. This is performed by aligning a 1D grid with the major-axis of
-        each light profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale.
-
-        This means that the summed 1D profile of a galaxy's quantity is the sum of each individual component aligned
-        with the major-axis.
-
-        The API is such that every plottable attribute of the `Galaxy` object is an input parameter of type bool of
-        the function, which if switched to `True` means that it is plotted.
-
-        Parameters
-        ----------
-        image
-            Whether to make a 1D plot (via `plot`) of the image.
-        convergence
-            Whether to make a 1D plot (via `plot`) of the convergence.
-        potential
-            Whether to make a 1D plot (via `plot`) of the potential.
-        """
-        if self.mat_plot_1d.yx_plot.plot_axis_type is None:
-            plot_axis_type_override = "semilogy"
-        else:
-            plot_axis_type_override = None
-
-        if image:
-            image_1d = self.galaxy.image_1d_from(grid=self.grid)
-
-            self.mat_plot_1d.plot_yx(
-                y=image_1d,
-                x=image_1d.grid_radial,
-                visuals_1d=self.visuals_1d
-                + Visuals1D().add_half_light_radius(self.galaxy),
-                auto_labels=aplt.AutoLabels(
-                    title="Image vs Radius",
-                    ylabel="Image ",
-                    xlabel="Radius",
-                    legend=self.galaxy.__class__.__name__,
-                    filename="image_1d",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
-
-        if convergence:
-            convergence_1d = self.galaxy.convergence_1d_from(grid=self.grid)
-
-            self.mat_plot_1d.plot_yx(
-                y=convergence_1d,
-                x=convergence_1d.grid_radial,
-                visuals_1d=self.visuals_1d
-                + Visuals1D().add_einstein_radius(mass_obj=self.galaxy, grid=self.grid),
-                auto_labels=aplt.AutoLabels(
-                    title="Convergence vs Radius",
-                    ylabel="Convergence ",
-                    xlabel="Radius",
-                    legend=self.galaxy.__class__.__name__,
-                    filename="convergence_1d",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
-
-        if potential:
-            potential_1d = self.galaxy.potential_1d_from(grid=self.grid)
-
-            self.mat_plot_1d.plot_yx(
-                y=potential_1d,
-                x=potential_1d.grid_radial,
-                visuals_1d=self.visuals_1d,
-                auto_labels=aplt.AutoLabels(
-                    title="Potential vs Radius",
-                    ylabel="Potential ",
-                    xlabel="Radius",
-                    legend=self.galaxy.__class__.__name__,
-                    filename="potential_1d",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
-
-    def figures_1d_decomposed(
-        self,
-        image: bool = False,
-        convergence: bool = False,
-        potential: bool = False,
-        legend_labels: List[str] = None,
-    ):
-        """
-        Plots the individual attributes of the plotter's `Galaxy` object in 1D, which are computed via the plotter's
-        grid object.
-
-        This function makes a decomposed plot shows the 1D plot of the attribute for every light or mass profile in
-        the galaxy, as well as their combined 1D plot.
-
-        If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is
-        computed from each light profile of the galaxy. This is performed by aligning a 1D grid with the major-axis of
-        each light profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale.
-
-        This means that the summed 1D profile of a galaxy's quantity is the sum of each individual component aligned
-        with the major-axis.
-
-        The API is such that every plottable attribute of the `Galaxy` object is an input parameter of type bool of
-        the function, which if switched to `True` means that it is plotted.
-
-        Parameters
-        ----------
-        image
-            Whether to make a 1D plot (via `plot`) of the image.
-        convergence
-            Whether to make a 1D plot (via `imshow`) of the convergence.
-        potential
-            Whether to make a 1D plot (via `imshow`) of the potential.
-        legend_labels
-            Manually overrides the labels of the plot's legend.
-        """
-
-        if self.galaxy.has(cls=LightProfile):
-            multi_plotter = aplt.MultiYX1DPlotter(
-                plotter_list=self.decomposed_light_profile_plotter_list,
-                legend_labels=legend_labels,
-            )
-            multi_plotter.plotter_list[0].mat_plot_1d.output = self.mat_plot_1d.output
-
-            if image:
-                change_filename = False
-
-                if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None:
-                    multi_plotter.plotter_list[0].set_filename(
-                        filename="image_1d_decomposed"
-                    )
-                    change_filename = True
-
-                multi_plotter.figure_1d(func_name="figures_1d", figure_name="image")
-
-                if change_filename:
-                    multi_plotter.plotter_list[0].set_filename(filename=None)
-
-        if self.galaxy.has(cls=MassProfile):
-            multi_plotter = aplt.MultiYX1DPlotter(
-                plotter_list=self.decomposed_mass_profile_plotter_list,
-                legend_labels=legend_labels,
-            )
-
-            if convergence:
-                change_filename = False
-
-                if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None:
-                    multi_plotter.plotter_list[0].set_filename(
-                        filename="convergence_1d_decomposed"
-                    )
-                    change_filename = True
-
-                multi_plotter.figure_1d(
-                    func_name="figures_1d", figure_name="convergence"
-                )
-
-                if change_filename:
-                    multi_plotter.plotter_list[0].set_filename(filename=None)
-
-            if potential:
-                change_filename = False
-
-                if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None:
-                    multi_plotter.plotter_list[0].set_filename(
-                        filename="potential_1d_decomposed"
-                    )
-                    change_filename = True
-
-                multi_plotter.figure_1d(func_name="figures_1d", figure_name="potential")
-
-                if change_filename:
-                    multi_plotter.plotter_list[0].set_filename(filename=None)
-
     def figures_2d(
         self,
         image: bool = False,
@@ -500,434 +294,3 @@ def subplot_of_mass_profiles(
                 plotter_list=mass_profile_plotters, name="deflections_x"
             )
 
-
-class GalaxyPDFPlotter(GalaxyPlotter):
-    def __init__(
-        self,
-        galaxy_pdf_list: List[Galaxy],
-        grid: aa.Grid2D,
-        mat_plot_1d: MatPlot1D = None,
-        visuals_1d: Visuals1D = None,
-        mat_plot_2d: MatPlot2D = None,
-        visuals_2d: Visuals2D = None,
-        sigma: Optional[float] = 3.0,
-    ):
-        """
-        Plots the attributes of a list of `GalaxyProfile` objects using the matplotlib methods `plot()` and `imshow()`
-        and many other matplotlib functions which customize the plot's appearance.
-
-        Figures plotted by this object average over a list galaxy profiles to computed the average value of each
-        attribute with errors, where the 1D regions within the errors are plotted as a shaded region to show the range
-        of plausible models. Therefore, the input list of galaxies is expected to represent the probability density
-        function of an inferred model-fit.
-
-        The `mat_plot_1d` and `mat_plot_2d` attributes wrap matplotlib function calls to make the figure. By default,
-        the settings passed to every matplotlib function called are those specified in
-        the `config/visualize/mat_wrap/*.ini` files, but a user can manually input values into `MatPlot2D` to
-        customize the figure's appearance.
-
-        Overlaid on the figure are visuals, contained in the `Visuals1D` and `Visuals2D` objects. Attributes may be
-        extracted from the `GalaxyProfile` and plotted via the visuals object.
-
-        Parameters
-        ----------
-        galaxy_profile_pdf_list
-            The list of galaxy profiles whose mean and error values the plotter plots.
-        grid
-            The 2D (y,x) grid of coordinates used to evaluate the galaxy profile quantities that are plotted.
-        mat_plot_1d
-            Contains objects which wrap the matplotlib function calls that make 1D plots.
-        visuals_1d
-            Contains 1D visuals that can be overlaid on 1D plots.
-        mat_plot_2d
-            Contains objects which wrap the matplotlib function calls that make 2D plots.
-        visuals_2d
-            Contains 2D visuals that can be overlaid on 2D plots.
-        sigma
-            The confidence interval in terms of a sigma value at which the errors are computed (e.g. a value of
-            sigma=3.0 uses confidence intevals at ~0.01 and 0.99 the PDF).
-        """
-        super().__init__(
-            galaxy=None,
-            grid=grid,
-            mat_plot_2d=mat_plot_2d,
-            visuals_2d=visuals_2d,
-            mat_plot_1d=mat_plot_1d,
-            visuals_1d=visuals_1d,
-        )
-
-        self.galaxy_pdf_list = galaxy_pdf_list
-        self.sigma = sigma
-        self.low_limit = (1 - math.erf(sigma / math.sqrt(2))) / 2
-
-    @property
-    def light_profile_pdf_plotter_list(self) -> List[LightProfilePDFPlotter]:
-        """
-        Returns a list of `LightProfilePDFPlotter` objects from the list of galaxies in this object. These are
-        typically used for plotting the individual average value plus errors of the light profiles of the
-        plotter's `Galaxy` (e.g. in the function `figures_1d_decomposed`).
-
-        Returns
-        -------
-        List[LightProfilePDFPlotter]
-            An object that plots the average value and errors of a list of light profiles, often used for plotting
-            attributes of the galaxy.
-        """
-        return [
-            self.light_profile_pdf_plotter_from(index=index)
-            for index in range(
-                len(self.galaxy_pdf_list[0].cls_list_from(cls=LightProfile))
-            )
-        ]
-
-    def light_profile_pdf_plotter_from(self, index) -> LightProfilePDFPlotter:
-        """
-        Returns the `LightProfilePDFPlotter` of a specific light profile in this plotter's list of galaxies. This is
-        typically used for plotting the individual average value plus errors of a light profile in plotter's galaxy
-        list (e.g. in the function `figures_1d_decomposed`).
-
-        Returns
-        -------
-        LightProfilePDFPlotter
-            An object that plots the average value and errors of a list of light profiles, often used for plotting
-            attributes of the galaxy.
-        """
-
-        from autogalaxy.profiles.plot.light_profile_plotters import (
-            LightProfilePDFPlotter,
-        )
-
-        light_profile_pdf_list = [
-            galaxy.cls_list_from(cls=LightProfile)[index]
-            for galaxy in self.galaxy_pdf_list
-        ]
-
-        return LightProfilePDFPlotter(
-            light_profile_pdf_list=light_profile_pdf_list,
-            grid=self.grid,
-            mat_plot_2d=self.mat_plot_2d,
-            visuals_2d=self.visuals_2d,
-            mat_plot_1d=self.mat_plot_1d,
-            visuals_1d=self.visuals_1d,
-        )
-
-    @property
-    def mass_profile_pdf_plotter_list(self) -> List[MassProfilePDFPlotter]:
-        """
-        Returns a list of `MassProfilePDFPlotter` objects from the list of galaxies in this object. These are
-        typically used for plotting the individual average value plus errors of the mass profiles of the
-        plotter's `Galaxy` (e.g. in the function `figures_1d_decomposed`).
-
-        Returns
-        -------
-        List[MassProfilePDFPlotter]
-            An object that plots the average value and errors of a list of mass profiles, often used for plotting
-            attributes of the galaxy.
-        """
-        return [
-            self.mass_profile_pdf_plotter_from(index=index)
-            for index in range(
-                len(self.galaxy_pdf_list[0].cls_list_from(cls=MassProfile))
-            )
-        ]
-
-    def mass_profile_pdf_plotter_from(self, index) -> MassProfilePDFPlotter:
-        """
-        Returns the `MassProfilePDFPlotter` of a specific mass profile in this plotter's list of galaxies. This is
-        typically used for plotting the individual average value plus errors of a mass profile in plotter's galaxy
-        list (e.g. in the function `figures_1d_decomposed`).
-
-        Returns
-        -------
-        MassProfilePDFPlotter
-            An object that plots the average value and errors of a list of mass profiles, often used for plotting
-            attributes of the galaxy.
-        """
-        mass_profile_pdf_list = [
-            galaxy.cls_list_from(cls=MassProfile)[index]
-            for galaxy in self.galaxy_pdf_list
-        ]
-
-        return MassProfilePDFPlotter(
-            mass_profile_pdf_list=mass_profile_pdf_list,
-            grid=self.grid,
-            mat_plot_2d=self.mat_plot_2d,
-            visuals_2d=self.visuals_2d,
-            mat_plot_1d=self.mat_plot_1d,
-            visuals_1d=self.visuals_1d,
-        )
-
-    def figures_1d(
-        self, image: bool = False, convergence: bool = False, potential: bool = False
-    ):
-        """
-        Plots the individual attributes of the plotter's list of `Galaxy` object in 1D, which are computed via the
-        plotter's grid object.
-
-        This averages over a list galaxies to compute the average value of each attribute with errors, where the
-        1D regions within the errors are plotted as a shaded region to show the range of plausible models. Therefore,
-        the input list of galaxies is expected to represent the probability density function of an inferred model-fit.
-
-        If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is
-        computed from each light profile of the galaxy. This is performed by aligning a 1D grid with the major-axis of
-        each light profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale.
-
-        This means that the summed 1D profile of a galaxy's quantity is the sum of each individual component aligned
-        with the major-axis.
-
-        The API is such that every plottable attribute of the `Galaxy` object is an input parameter of type bool of
-        the function, which if switched to `True` means that it is plotted.
-
-        Parameters
-        ----------
-        image
-            Whether to make a 1D plot (via `plot`) of the image.
-        convergence
-            Whether to make a 1D plot (via `imshow`) of the convergence.
-        potential
-            Whether to make a 1D plot (via `imshow`) of the potential.
-        """
-        if self.mat_plot_1d.yx_plot.plot_axis_type is None:
-            plot_axis_type_override = "semilogy"
-        else:
-            plot_axis_type_override = None
-
-        if image:
-            image_1d_list = [
-                galaxy.image_1d_from(grid=self.grid) for galaxy in self.galaxy_pdf_list
-            ]
-
-            min_index = min([image_1d.shape[0] for image_1d in image_1d_list])
-            image_1d_list = [image_1d[0:min_index] for image_1d in image_1d_list]
-
-            (
-                median_image_1d,
-                errors_image_1d,
-            ) = error_util.profile_1d_median_and_error_region_via_quantile(
-                profile_1d_list=image_1d_list, low_limit=self.low_limit
-            )
-
-            visuals_1d_via_light_obj_list = Visuals1D().add_half_light_radius_errors(
-                light_obj_list=self.galaxy_pdf_list, low_limit=self.low_limit
-            )
-            visuals_1d_with_shaded_region = self.visuals_1d.__class__(
-                shaded_region=errors_image_1d
-            )
-
-            visuals_1d = visuals_1d_via_light_obj_list + visuals_1d_with_shaded_region
-
-            median_image_1d = aa.Array1D.no_mask(
-                values=median_image_1d, pixel_scales=self.grid.pixel_scale
-            )
-
-            self.mat_plot_1d.plot_yx(
-                y=median_image_1d,
-                x=median_image_1d.grid_radial,
-                visuals_1d=visuals_1d,
-                auto_labels=aplt.AutoLabels(
-                    title="Image vs Radius",
-                    ylabel="Image ",
-                    xlabel="Radius",
-                    legend=self.galaxy_pdf_list[0].__class__.__name__,
-                    filename="image_1d",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
-
-        if convergence:
-            convergence_1d_list = [
-                galaxy.convergence_1d_from(grid=self.grid)
-                for galaxy in self.galaxy_pdf_list
-            ]
-
-            min_index = min(
-                [convergence_1d.shape[0] for convergence_1d in convergence_1d_list]
-            )
-            convergence_1d_list = [
-                convergence_1d[0:min_index] for convergence_1d in convergence_1d_list
-            ]
-
-            (
-                median_convergence_1d,
-                errors_convergence_1d,
-            ) = error_util.profile_1d_median_and_error_region_via_quantile(
-                profile_1d_list=convergence_1d_list, low_limit=self.low_limit
-            )
-
-            visuals_1d_via_lensing_obj_list = Visuals1D().add_einstein_radius_errors(
-                mass_obj_list=self.galaxy_pdf_list,
-                grid=self.grid,
-                low_limit=self.low_limit,
-            )
-            visuals_1d_with_shaded_region = self.visuals_1d.__class__(
-                shaded_region=errors_convergence_1d
-            )
-
-            visuals_1d = visuals_1d_via_lensing_obj_list + visuals_1d_with_shaded_region
-
-            median_convergence_1d = aa.Array1D.no_mask(
-                values=median_convergence_1d, pixel_scales=self.grid.pixel_scale
-            )
-
-            self.mat_plot_1d.plot_yx(
-                y=median_convergence_1d,
-                x=median_convergence_1d.grid_radial,
-                visuals_1d=visuals_1d,
-                auto_labels=aplt.AutoLabels(
-                    title="Convergence vs Radius",
-                    ylabel="Convergence ",
-                    xlabel="Radius",
-                    legend=self.galaxy_pdf_list[0].__class__.__name__,
-                    filename="convergence_1d",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
-
-        if potential:
-            potential_1d_list = [
-                galaxy.potential_1d_from(grid=self.grid)
-                for galaxy in self.galaxy_pdf_list
-            ]
-
-            min_index = min(
-                [potential_1d.shape[0] for potential_1d in potential_1d_list]
-            )
-            potential_1d_list = [
-                potential_1d[0:min_index] for potential_1d in potential_1d_list
-            ]
-
-            (
-                median_potential_1d,
-                errors_potential_1d,
-            ) = error_util.profile_1d_median_and_error_region_via_quantile(
-                profile_1d_list=potential_1d_list, low_limit=self.low_limit
-            )
-
-            visuals_1d_via_lensing_obj_list = Visuals1D().add_einstein_radius_errors(
-                mass_obj_list=self.galaxy_pdf_list,
-                grid=self.grid,
-                low_limit=self.low_limit,
-            )
-            visuals_1d_with_shaded_region = self.visuals_1d.__class__(
-                shaded_region=errors_potential_1d
-            )
-
-            visuals_1d = visuals_1d_via_lensing_obj_list + visuals_1d_with_shaded_region
-
-            median_potential_1d = aa.Array1D.no_mask(
-                values=median_potential_1d, pixel_scales=self.grid.pixel_scale
-            )
-
-            self.mat_plot_1d.plot_yx(
-                y=median_potential_1d,
-                x=median_potential_1d.grid_radial,
-                visuals_1d=visuals_1d,
-                auto_labels=aplt.AutoLabels(
-                    title="Potential vs Radius",
-                    ylabel="Potential ",
-                    xlabel="Radius",
-                    legend=self.galaxy_pdf_list[0].__class__.__name__,
-                    filename="potential_1d",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
-
-    def figures_1d_decomposed(
-        self,
-        image: bool = False,
-        convergence: bool = False,
-        potential: bool = False,
-        legend_labels: List[str] = None,
-    ):
-        """
-        Plots the individual attributes of the plotter's `Galaxy` object in 1D, which are computed via the plotter's
-        grid object.
-
-        This averages over a list galaxies to compute the average value of each attribute with errors, where the
-        1D regions within the errors are plotted as a shaded region to show the range of plausible models. Therefore,
-        the input list of galaxies is expected to represent the probability density function of an inferred model-fit.
-
-        This function makes a decomposed plot showing the 1D plot of each attribute for every light or mass profile in
-        the galaxy, as well as their combined 1D plot. By plotting the attribute of each profile on the same figure,
-        one can see how much each profile contributes to the galaxy overall.
-
-        If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is
-        computed from each light profile of the galaxy. This is performed by aligning a 1D grid with the major-axis of
-        each light profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale.
-
-        This means that the summed 1D profile of a galaxy's quantity is the sum of each individual component aligned
-        with the major-axis.
-
-        The API is such that every plottable attribute of the `Galaxy` object is an input parameter of type bool of
-        the function, which if switched to `True` means that it is plotted.
-
-        Parameters
-        ----------
-        image
-            Whether to make a 1D plot (via `plot`) of the image.
-        convergence
-            Whether to make a 1D plot (via `imshow`) of the convergence.
-        potential
-            Whether to make a 1D plot (via `imshow`) of the potential.
-        legend_labels
-            Manually overrides the labels of the plot's legend.
-        """
-        if image:
-            multi_plotter = aplt.MultiYX1DPlotter(
-                plotter_list=[self] + self.light_profile_pdf_plotter_list,
-                legend_labels=legend_labels,
-            )
-
-            change_filename = False
-
-            if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None:
-                multi_plotter.plotter_list[0].set_filename(
-                    filename="image_1d_decomposed"
-                )
-
-                change_filename = True
-
-            multi_plotter.figure_1d(func_name="figures_1d", figure_name="image")
-
-            if change_filename:
-                multi_plotter.plotter_list[0].set_filename(filename=None)
-
-        if convergence:
-            multi_plotter = aplt.MultiYX1DPlotter(
-                plotter_list=[self] + self.mass_profile_pdf_plotter_list,
-                legend_labels=legend_labels,
-            )
-
-            change_filename = False
-
-            if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None:
-                multi_plotter.plotter_list[0].set_filename(
-                    filename="convergence_1d_decomposed"
-                )
-
-                change_filename = True
-
-            multi_plotter.figure_1d(func_name="figures_1d", figure_name="convergence")
-
-            if change_filename:
-                multi_plotter.plotter_list[0].set_filename(filename=None)
-
-        if potential:
-            multi_plotter = aplt.MultiYX1DPlotter(
-                plotter_list=[self] + self.mass_profile_pdf_plotter_list,
-                legend_labels=legend_labels,
-            )
-
-            change_filename = False
-
-            if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None:
-                multi_plotter.plotter_list[0].set_filename(
-                    filename="potential_1d_decomposed"
-                )
-
-                change_filename = True
-
-            multi_plotter.figure_1d(func_name="figures_1d", figure_name="potential")
-
-            if change_filename:
-                multi_plotter.plotter_list[0].set_filename(filename=None)
diff --git a/autogalaxy/plot/__init__.py b/autogalaxy/plot/__init__.py
index 66317df0b..c33fc78c1 100644
--- a/autogalaxy/plot/__init__.py
+++ b/autogalaxy/plot/__init__.py
@@ -76,12 +76,9 @@
 from autogalaxy.plot.visuals.two_d import Visuals2D
 
 from autogalaxy.profiles.plot.light_profile_plotters import LightProfilePlotter
-from autogalaxy.profiles.plot.light_profile_plotters import LightProfilePDFPlotter
 from autogalaxy.profiles.plot.basis_plotters import BasisPlotter
 from autogalaxy.profiles.plot.mass_profile_plotters import MassProfilePlotter
-from autogalaxy.profiles.plot.mass_profile_plotters import MassProfilePDFPlotter
 from autogalaxy.galaxy.plot.galaxy_plotters import GalaxyPlotter
-from autogalaxy.galaxy.plot.galaxy_plotters import GalaxyPDFPlotter
 from autogalaxy.galaxy.plot.galaxies_plotters import GalaxiesPlotter
 from autogalaxy.quantity.plot.fit_quantity_plotters import FitQuantityPlotter
 from autogalaxy.imaging.plot.fit_imaging_plotters import FitImagingPlotter
diff --git a/autogalaxy/profiles/plot/light_profile_plotters.py b/autogalaxy/profiles/plot/light_profile_plotters.py
index a611603f6..cb5116e09 100644
--- a/autogalaxy/profiles/plot/light_profile_plotters.py
+++ b/autogalaxy/profiles/plot/light_profile_plotters.py
@@ -1,6 +1,3 @@
-import math
-from typing import List, Optional
-
 import autoarray as aa
 import autoarray.plot as aplt
 
@@ -12,7 +9,6 @@
 from autogalaxy.plot.visuals.one_d import Visuals1D
 from autogalaxy.plot.visuals.two_d import Visuals2D
 
-from autogalaxy.util import error_util
 from autogalaxy import exc
 
 
@@ -80,45 +76,6 @@ def grid_2d_projected(self):
             angle=self.light_profile.angle()
         )
 
-    def figures_1d(self, image: bool = False):
-        """
-        Plots the individual attributes of the plotter's `LightProfile` object in 1D, which are computed via the
-        plotter's grid object.
-
-        If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is
-        computed from the light profile. This is performed by aligning a 1D grid with the  major-axis of the light
-        profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale.
-
-        The API is such that every plottable attribute of the `LightProfile` object is an input parameter of type
-        bool of the function, which if switched to `True` means that it is plotted.
-
-        Parameters
-        ----------
-        image
-            Whether to make a 1D plot (via `plot`) of the image.
-        """
-        if self.mat_plot_1d.yx_plot.plot_axis_type is None:
-            plot_axis_type_override = "semilogy"
-        else:
-            plot_axis_type_override = None
-
-        if image:
-
-            image_1d = self.light_profile.image_2d_from(grid=self.grid_2d_projected)
-
-            self.mat_plot_1d.plot_yx(
-                y=image_1d,
-                x=self.grid_2d_projected[:,1],
-                visuals_1d=self.visuals_1d,
-                auto_labels=aplt.AutoLabels(
-                    title=r"Image ($\mathrm{e^{-}}\,\mathrm{s^{-1}}$) vs Radius (arcsec)",
-                    yunit="",
-                    legend=self.light_profile.__class__.__name__,
-                    filename="image_1d",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
-
     def figures_2d(self, image: bool = False):
         """
         Plots the individual attributes of the plotter's `LightProfile` object in 2D, which are computed via the
@@ -138,138 +95,3 @@ def figures_2d(self, image: bool = False):
                 visuals_2d=self.visuals_2d,
                 auto_labels=aplt.AutoLabels(title="Image", filename="image_2d"),
             )
-
-
-class LightProfilePDFPlotter(LightProfilePlotter):
-    def __init__(
-        self,
-        light_profile_pdf_list: List[LightProfile],
-        grid: aa.type.Grid2DLike,
-        mat_plot_1d: MatPlot1D = None,
-        visuals_1d: Visuals1D = None,
-        mat_plot_2d: MatPlot2D = None,
-        visuals_2d: Visuals2D = None,
-        sigma: Optional[float] = 3.0,
-    ):
-        """
-        Plots the attributes of a list of `LightProfile` objects using the matplotlib methods `plot()` and `imshow()`
-        and many other matplotlib functions which customize the plot's appearance.
-
-        Figures plotted by this object average over a list light profiles to computed the average value of each
-        attribute with errors, where the 1D regions within the errors are plotted as a shaded region to show the range
-        of plausible models. Therefore, the input list of galaxies is expected to represent the probability density
-        function of an inferred model-fit.
-
-        The `mat_plot_1d` and `mat_plot_2d` attributes wrap matplotlib function calls to make the figure. By default,
-        the settings passed to every matplotlib function called are those specified in
-        the `config/visualize/mat_wrap/*.ini` files, but a user can manually input values into `MatPlot2D` to
-        customize the figure's appearance.
-
-        Overlaid on the figure are visuals, contained in the `Visuals1D` and `Visuals2D` objects. Attributes may be
-        extracted from the `LightProfile` and plotted via the visuals object.
-
-        Parameters
-        ----------
-        light_profile_pdf_list
-            The list of light profiles whose mean and error values the plotter plots.
-        grid
-            The 2D (y,x) grid of coordinates used to evaluate the light profile quantities that are plotted.
-        mat_plot_1d
-            Contains objects which wrap the matplotlib function calls that make 1D plots.
-        visuals_1d
-            Contains 1D visuals that can be overlaid on 1D plots.
-        mat_plot_2d
-            Contains objects which wrap the matplotlib function calls that make 2D plots.
-        visuals_2d
-            Contains 2D visuals that can be overlaid on 2D plots.
-        sigma
-            The confidence interval in terms of a sigma value at which the errors are computed (e.g. a value of
-            sigma=3.0 uses confidence intevals at ~0.01 and 0.99 the PDF).
-        """
-        super().__init__(
-            light_profile=None,
-            grid=grid,
-            mat_plot_1d=mat_plot_1d,
-            visuals_1d=visuals_1d,
-            mat_plot_2d=mat_plot_2d,
-            visuals_2d=visuals_2d,
-        )
-
-        self.light_profile_pdf_list = light_profile_pdf_list
-        self.sigma = sigma
-        self.low_limit = (1 - math.erf(sigma / math.sqrt(2))) / 2
-
-    def figures_1d(self, image: bool = False):
-        """
-        Plots the individual attributes of the plotter's list of ` LightProfile` object in 1D, which are computed via
-        the plotter's grid object.
-
-        This averages over a list light profiles to compute the average value of each attribute with errors, where the
-        1D regions within the errors are plotted as a shaded region to show the range of plausible models. Therefore,
-        the input list of galaxies is expected to represent the probability density function of an inferred model-fit.
-
-        If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is
-        computed from each light profile. This is performed by aligning a 1D grid with the major-axis of
-        each light profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale.
-
-        The API is such that every plottable attribute of the `LightProfile` object is an input parameter of type bool
-        of the function, which if switched to `True` means that it is plotted.
-
-        Parameters
-        ----------
-        image
-            Whether to make a 1D plot (via `plot`) of the image.
-        convergence
-            Whether to make a 1D plot (via `imshow`) of the convergence.
-        potential
-            Whether to make a 1D plot (via `imshow`) of the potential.
-        """
-        if self.mat_plot_1d.yx_plot.plot_axis_type is None:
-            plot_axis_type_override = "semilogy"
-        else:
-            plot_axis_type_override = None
-
-        if image:
-
-            image_1d_list = []
-
-            for light_profile in self.light_profile_pdf_list:
-
-                grid = self.grid.grid_2d_radial_projected_from(
-                    centre=light_profile.centre,
-                    angle=light_profile.angle()
-                )
-
-                image_1d_list.append(light_profile.image_2d_from(grid=grid))
-
-            min_index = min([image_1d.shape[0] for image_1d in image_1d_list])
-            image_1d_list = [image_1d[0:min_index] for image_1d in image_1d_list]
-
-            (
-                median_image_1d,
-                errors_image_1d,
-            ) = error_util.profile_1d_median_and_error_region_via_quantile(
-                profile_1d_list=image_1d_list, low_limit=self.low_limit
-            )
-
-            visuals_1d_via_light_obj_list = Visuals1D().add_half_light_radius_errors(
-                light_obj_list=self.light_profile_pdf_list,
-                low_limit=self.low_limit,
-            )
-
-            visuals_1d_with_shaded_region = Visuals1D(shaded_region=errors_image_1d)
-
-            visuals_1d = visuals_1d_via_light_obj_list + visuals_1d_with_shaded_region
-
-            self.mat_plot_1d.plot_yx(
-                y=median_image_1d,
-                x=grid[0:min_index,1],
-                visuals_1d=visuals_1d,
-                auto_labels=aplt.AutoLabels(
-                    title=r"Image ($\mathrm{e^{-}}\,\mathrm{s^{-1}}$) vs Radius (arcsec)",
-                    yunit="",
-                    legend=self.light_profile_pdf_list[0].__class__.__name__,
-                    filename="image_1d_pdf",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
diff --git a/autogalaxy/profiles/plot/mass_profile_plotters.py b/autogalaxy/profiles/plot/mass_profile_plotters.py
index 2fb8b0278..6f5675fbc 100644
--- a/autogalaxy/profiles/plot/mass_profile_plotters.py
+++ b/autogalaxy/profiles/plot/mass_profile_plotters.py
@@ -78,237 +78,3 @@ def grid_2d_projected(self):
             angle=self.mass_profile.angle()
         )
 
-    def figures_1d(self, convergence: bool = False, potential: bool = False):
-        """
-        Plots the individual attributes of the plotter's `MassProfile` object in 1D, which are computed via the
-        plotter's grid object.
-
-        If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is
-        computed from the mass profile. This is performed by aligning a 1D grid with the  major-axis of the mass
-        profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale.
-
-        The API is such that every plottable attribute of the  `MassProfile` object is an input parameter of type
-        bool of the function, which if switched to `True` means that it is plotted.
-
-        Parameters
-        ----------
-        convergence
-            Whether to make a 1D plot (via `imshow`) of the convergence.
-        potential
-            Whether to make a 1D plot (via `imshow`) of the potential.
-        """
-        if self.mat_plot_1d.yx_plot.plot_axis_type is None:
-            plot_axis_type_override = "semilogy"
-        else:
-            plot_axis_type_override = None
-
-        if convergence:
-            convergence_1d = self.mass_profile.convergence_2d_from(grid=self.grid_2d_projected)
-
-            self.mat_plot_1d.plot_yx(
-                y=convergence_1d,
-                x=self.grid_2d_projected[:,1],
-                visuals_1d=self.visuals_1d,
-                auto_labels=aplt.AutoLabels(
-                    title="Convergence vs Radius (arcsec)",
-                    yunit="",
-                    legend=self.mass_profile.__class__.__name__,
-                    filename="convergence_1d",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
-
-        if potential:
-            potential_1d = self.mass_profile.potential_2d_from(grid=self.grid_2d_projected)
-
-            self.mat_plot_1d.plot_yx(
-                y=potential_1d,
-                x=self.grid_2d_projected[:,1],
-                visuals_1d=self.visuals_1d,
-                auto_labels=aplt.AutoLabels(
-                    title="Potential vs Radius (arcsec)",
-                    yunit="",
-                    legend=self.mass_profile.__class__.__name__,
-                    filename="potential_1d",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
-
-
-class MassProfilePDFPlotter(MassProfilePlotter):
-    def __init__(
-        self,
-        mass_profile_pdf_list: List[MassProfile],
-        grid: aa.Grid2D,
-        mat_plot_1d: MatPlot1D = None,
-        visuals_1d: Visuals1D = None,
-        mat_plot_2d: MatPlot2D = None,
-        visuals_2d: Visuals2D = None,
-        sigma: Optional[float] = 3.0,
-    ):
-        """
-        Plots the attributes of a list of `MassProfile` objects using the matplotlib methods `plot()` and `imshow()`
-        and many other matplotlib functions which customize the plot's appearance.
-
-        Figures plotted by this object average over a list mass profiles to computed the average value of each attribute
-        with errors, where the 1D regions within the errors are plotted as a shaded region to show the range of
-        plausible models. Therefore, the input list of galaxies is expected to represent the probability density
-        function of an inferred model-fit.
-
-        The `mat_plot_1d` and `mat_plot_2d` attributes wrap matplotlib function calls to make the figure. By default,
-        the settings passed to every matplotlib function called are those specified in
-        the `config/visualize/mat_wrap/*.ini` files, but a user can manually input values into `MatPlot2D` to
-        customize the figure's appearance.
-
-        Overlaid on the figure are visuals, contained in the `Visuals1D` and `Visuals2D` objects. Attributes may be
-        extracted from the `MassProfile` and plotted via the visuals object.
-
-        Parameters
-        ----------
-        mass_profile_pdf_list
-            The list of mass profiles whose mean and error values the plotter plots.
-        grid
-            The 2D (y,x) grid of coordinates used to evaluate the mass profile quantities that are plotted.
-        mat_plot_1d
-            Contains objects which wrap the matplotlib function calls that make 1D plots.
-        visuals_1d
-            Contains 1D visuals that can be overlaid on 1D plots.
-        mat_plot_2d
-            Contains objects which wrap the matplotlib function calls that make 2D plots.
-        visuals_2d
-            Contains 2D visuals that can be overlaid on 2D plots.
-        sigma
-            The confidence interval in terms of a sigma value at which the errors are computed (e.g. a value of
-            sigma=3.0 uses confidence intevals at ~0.01 and 0.99 the PDF).
-        """
-        super().__init__(
-            mass_profile=None,
-            grid=grid,
-            mat_plot_1d=mat_plot_1d,
-            visuals_1d=visuals_1d,
-            mat_plot_2d=mat_plot_2d,
-            visuals_2d=visuals_2d,
-        )
-
-        self.mass_profile_pdf_list = mass_profile_pdf_list
-        self.sigma = sigma
-        self.low_limit = (1 - math.erf(sigma / math.sqrt(2))) / 2
-
-    def figures_1d(self, convergence=False, potential=False):
-        """
-        Plots the individual attributes of the plotter's list of ` MassProfile` object in 1D, which are computed via
-        the plotter's grid object.
-
-        This averages over a list mass profiles to compute the average value of each attribute with errors, where the
-        1D regions within the errors are plotted as a shaded region to show the range of plausible models. Therefore,
-        the input list of galaxies is expected to represent the probability density function of an inferred model-fit.
-
-        If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is
-        computed from each mass profile. This is performed by aligning a 1D grid with the major-axis of
-        each mass profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale.
-
-        The API is such that every plottable attribute of the `MassProfile` object is an input parameter of type bool
-        of the function, which if switched to `True` means that it is plotted.
-
-        Parameters
-        ----------
-        convergence
-            Whether to make a 1D plot (via `imshow`) of the convergence.
-        potential
-            Whether to make a 1D plot (via `imshow`) of the potential.
-        """
-        if self.mat_plot_1d.yx_plot.plot_axis_type is None:
-            plot_axis_type_override = "semilogy"
-        else:
-            plot_axis_type_override = None
-
-        if convergence:
-
-            convergence_1d_list = []
-
-            for mass_profile in self.mass_profile_pdf_list:
-                
-                grid = self.grid.grid_2d_radial_projected_from(
-                    centre=mass_profile.centre,
-                    angle=mass_profile.angle()
-                )
-           
-                convergence_1d_list.append(mass_profile.convergence_2d_from(grid=grid))
-
-            min_index = min(
-                [convergence_1d.shape[0] for convergence_1d in convergence_1d_list]
-            )
-            convergence_1d_list = [
-                convergence_1d[0:min_index] for convergence_1d in convergence_1d_list
-            ]
-
-            (
-                median_convergence_1d,
-                errors_convergence_1d,
-            ) = error_util.profile_1d_median_and_error_region_via_quantile(
-                profile_1d_list=convergence_1d_list, low_limit=self.low_limit
-            )
-
-            visuals_1d_via_lensing_obj_list = Visuals1D().add_einstein_radius_errors(
-                mass_obj_list=self.mass_profile_pdf_list,
-                grid=self.grid,
-                low_limit=self.low_limit,
-            )
-            visuals_1d_with_shaded_region = Visuals1D(
-                shaded_region=errors_convergence_1d
-            )
-
-            visuals_1d = visuals_1d_via_lensing_obj_list + visuals_1d_with_shaded_region
-
-            self.mat_plot_1d.plot_yx(
-                y=median_convergence_1d,
-                x=grid[0:min_index,1],
-                visuals_1d=visuals_1d,
-                auto_labels=aplt.AutoLabels(
-                    title="Convergence vs Radius (arcsec)",
-                    yunit="",
-                    legend=self.mass_profile_pdf_list[0].__class__.__name__,
-                    filename="convergence_1d_pdf",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
-
-        if potential:
-            potential_1d_list = []
-
-            for mass_profile in self.mass_profile_pdf_list:
-                grid = self.grid.grid_2d_radial_projected_from(
-                    centre=mass_profile.centre,
-                    angle=mass_profile.angle()
-                )
-
-                potential_1d_list.append(mass_profile.potential_2d_from(grid=grid))
-
-            (
-                median_potential_1d,
-                errors_potential_1d,
-            ) = error_util.profile_1d_median_and_error_region_via_quantile(
-                profile_1d_list=potential_1d_list, low_limit=self.low_limit
-            )
-
-            visuals_1d_via_lensing_obj_list = Visuals1D().add_einstein_radius_errors(
-                mass_obj_list=self.mass_profile_pdf_list,
-                grid=self.grid,
-                low_limit=self.low_limit,
-            )
-            visuals_1d_with_shaded_region = Visuals1D(shaded_region=errors_potential_1d)
-
-            visuals_1d = visuals_1d_via_lensing_obj_list + visuals_1d_with_shaded_region
-
-            self.mat_plot_1d.plot_yx(
-                y=median_potential_1d,
-                x=grid[0:min_index,1],
-                visuals_1d=visuals_1d,
-                auto_labels=aplt.AutoLabels(
-                    title="Potential vs Radius (arcsec)",
-                    yunit="",
-                    legend=self.mass_profile_pdf_list[0].__class__.__name__,
-                    filename="potential_1d_pdf",
-                ),
-                plot_axis_type_override=plot_axis_type_override,
-            )
diff --git a/docs/api/plot.rst b/docs/api/plot.rst
index a930d6872..6325e8383 100644
--- a/docs/api/plot.rst
+++ b/docs/api/plot.rst
@@ -32,7 +32,6 @@ Create figures and subplots showing quantities of standard **PyAutoGalaxy** obje
     ImagingPlotter
     InterferometerPlotter
     LightProfilePlotter
-    LightProfilePDFPlotter
     GalaxyPlotter
     FitImagingPlotter
     FitInterferometerPlotter
diff --git a/test_autogalaxy/analysis/test_plotter_interface.py b/test_autogalaxy/analysis/test_plotter_interface.py
index e7132c878..b883e74ae 100644
--- a/test_autogalaxy/analysis/test_plotter_interface.py
+++ b/test_autogalaxy/analysis/test_plotter_interface.py
@@ -29,10 +29,6 @@ def test__galaxies(masked_imaging_7x7, galaxies_7x7, plot_path, plot_patch):
 
     assert path.join(plot_path, "subplot_galaxies.png") in plot_patch.paths
     assert path.join(plot_path, "subplot_galaxy_images.png") in plot_patch.paths
-    assert path.join(plot_path, "subplot_galaxies_1d.png") in plot_patch.paths
-    assert (
-        path.join(plot_path, "subplot_galaxies_1d_decomposed.png") in plot_patch.paths
-    )
 
     image = ag.ndarray_via_fits_from(
         file_path=path.join(plot_path, "galaxy_images.fits"), hdu=1
diff --git a/test_autogalaxy/config/visualize.yaml b/test_autogalaxy/config/visualize.yaml
index af4a18bcb..754ce9375 100644
--- a/test_autogalaxy/config/visualize.yaml
+++ b/test_autogalaxy/config/visualize.yaml
@@ -351,8 +351,6 @@ plots:
   galaxies:
     subplot_galaxies: true
     subplot_galaxy_images: true
-    subplot_galaxies_1d: true
-    subplot_galaxies_1d_decomposed: true
     fits_galaxy_images: true                # Output a .fits file containing images of every galaxy?
   positions:
     image_with_positions: true
diff --git a/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py b/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py
index eddee94c0..e89d8291f 100644
--- a/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py
+++ b/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py
@@ -72,12 +72,4 @@ def test__galaxies_sub_plot_output(galaxies_x2_7x7, grid_2d_7x7, plot_path, plot
     assert path.join(plot_path, "subplot_galaxies.png") in plot_patch.paths
 
     plotter.subplot_galaxy_images()
-    assert path.join(plot_path, "subplot_galaxy_images.png") in plot_patch.paths
-
-    plotter.subplot_galaxies_1d()
-    assert path.join(plot_path, "subplot_galaxies_1d.png") in plot_patch.paths
-
-    plotter.subplot_galaxies_1d_decomposed()
-    assert (
-        path.join(plot_path, "subplot_galaxies_1d_decomposed.png") in plot_patch.paths
-    )
+    assert path.join(plot_path, "subplot_galaxy_images.png") in plot_patch.paths
\ No newline at end of file
diff --git a/test_autogalaxy/galaxy/plot/test_galaxy_plotters.py b/test_autogalaxy/galaxy/plot/test_galaxy_plotters.py
index e03fd3f29..e14916496 100644
--- a/test_autogalaxy/galaxy/plot/test_galaxy_plotters.py
+++ b/test_autogalaxy/galaxy/plot/test_galaxy_plotters.py
@@ -14,106 +14,6 @@ def make_galaxy_plotter_setup():
     )
 
 
-def test__figures_1d__all_are_output(
-    gal_x1_lp_x1_mp, grid_2d_7x7, mask_2d_7x7, plot_path, plot_patch
-):
-    galaxy_plotter = aplt.GalaxyPlotter(
-        galaxy=gal_x1_lp_x1_mp,
-        grid=grid_2d_7x7,
-        mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")),
-    )
-    galaxy_plotter.figures_1d(image=True, convergence=True, potential=True)
-
-    assert path.join(plot_path, "image_1d.png") in plot_patch.paths
-    assert path.join(plot_path, "convergence_1d.png") in plot_patch.paths
-    assert path.join(plot_path, "potential_1d.png") in plot_patch.paths
-
-    plot_patch.paths = []
-
-    galaxy_plotter = aplt.GalaxyPDFPlotter(
-        galaxy_pdf_list=[gal_x1_lp_x1_mp, gal_x1_lp_x1_mp, gal_x1_lp_x1_mp],
-        grid=grid_2d_7x7,
-        mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")),
-    )
-    galaxy_plotter.figures_1d(image=True, convergence=True, potential=True)
-
-    assert path.join(plot_path, "image_1d.png") in plot_patch.paths
-    assert path.join(plot_path, "convergence_1d.png") in plot_patch.paths
-    assert path.join(plot_path, "potential_1d.png") in plot_patch.paths
-
-
-def test__figures_1d_decomposed__all_are_output(
-    gal_x1_lp_x1_mp, grid_2d_7x7, mask_2d_7x7, plot_path, plot_patch
-):
-    galaxy_plotter = aplt.GalaxyPlotter(
-        galaxy=gal_x1_lp_x1_mp,
-        grid=grid_2d_7x7,
-        mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")),
-    )
-    galaxy_plotter.figures_1d_decomposed(image=True, convergence=True, potential=True)
-
-    assert path.join(plot_path, "image_1d_decomposed.png") in plot_patch.paths
-    assert path.join(plot_path, "convergence_1d_decomposed.png") in plot_patch.paths
-    assert path.join(plot_path, "potential_1d_decomposed.png") in plot_patch.paths
-
-    plot_patch.paths = []
-
-    galaxy_plotter = aplt.GalaxyPDFPlotter(
-        galaxy_pdf_list=[gal_x1_lp_x1_mp, gal_x1_lp_x1_mp, gal_x1_lp_x1_mp],
-        grid=grid_2d_7x7,
-        mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")),
-    )
-
-    galaxy_plotter.figures_1d_decomposed(image=True, convergence=True, potential=True)
-
-    assert path.join(plot_path, "image_1d_decomposed.png") in plot_patch.paths
-    assert path.join(plot_path, "convergence_1d_decomposed.png") in plot_patch.paths
-    assert path.join(plot_path, "potential_1d_decomposed.png") in plot_patch.paths
-
-
-def test__figures_1d_decomposed__light_profiles_different_centres_making_offset_radial_grid(
-    grid_2d_7x7, mask_2d_7x7, plot_path, plot_patch
-):
-    lp_0 = ag.lp.SersicSph(centre=(0.0, 0.0))
-    lp_1 = ag.lp.SersicSph(centre=(1.0, 1.0))
-
-    mp_0 = ag.mp.IsothermalSph(centre=(0.0, 0.0))
-    mp_1 = ag.mp.IsothermalSph(centre=(1.0, 1.0))
-
-    gal = ag.Galaxy(redshift=0.5, light_0=lp_0, light_1=lp_1, mass_0=mp_0, mass_1=mp_1)
-
-    galaxy_plotter = aplt.GalaxyPlotter(
-        galaxy=gal,
-        grid=grid_2d_7x7,
-        mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")),
-    )
-    galaxy_plotter.figures_1d_decomposed(image=True, convergence=True, potential=True)
-
-    assert path.join(plot_path, "image_1d_decomposed.png") in plot_patch.paths
-    assert path.join(plot_path, "convergence_1d_decomposed.png") in plot_patch.paths
-    assert path.join(plot_path, "potential_1d_decomposed.png") in plot_patch.paths
-
-    lp_0 = ag.lp.SersicSph(centre=(0.0, 0.0))
-    lp_1 = ag.lp.SersicSph(centre=(1.0, 1.0))
-
-    mp_0 = ag.mp.IsothermalSph(centre=(0.0, 0.0))
-    mp_1 = ag.mp.IsothermalSph(centre=(1.0, 1.0))
-
-    gal_0 = ag.Galaxy(redshift=0.5, light_0=lp_0, mass_0=mp_0)
-    gal_1 = ag.Galaxy(redshift=0.5, light_1=lp_1, mass_1=mp_1)
-
-    galaxy_plotter = aplt.GalaxyPDFPlotter(
-        galaxy_pdf_list=[gal_0, gal_1],
-        grid=grid_2d_7x7,
-        mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")),
-    )
-    galaxy_plotter.figures_1d_decomposed(image=True, convergence=True, potential=True)
-
-    assert path.join(plot_path, "image_1d_decomposed.png") in plot_patch.paths
-    assert path.join(plot_path, "convergence_1d_decomposed.png") in plot_patch.paths
-    assert path.join(plot_path, "potential_1d_decomposed.png") in plot_patch.paths
-
-
 def test__figures_2d__all_are_output(
     gal_x1_lp_x1_mp,
     grid_2d_7x7,
diff --git a/test_autogalaxy/profiles/plot/test_light_profile_plotters.py b/test_autogalaxy/profiles/plot/test_light_profile_plotters.py
index e39e6cf77..477f13d71 100644
--- a/test_autogalaxy/profiles/plot/test_light_profile_plotters.py
+++ b/test_autogalaxy/profiles/plot/test_light_profile_plotters.py
@@ -14,47 +14,6 @@ def make_profile_plotter_setup():
         "{}".format(path.dirname(path.realpath(__file__))), "files", "plots", "profiles"
     )
 
-
-def test__figures_1d__all_are_output(
-    lp_0,
-    lp_1,
-    grid_2d_7x7,
-    grid_2d_irregular_7x7_list,
-    plot_path,
-    plot_patch,
-):
-    mat_plot_1d = aplt.MatPlot1D(
-        half_light_radius_axvline=aplt.HalfLightRadiusAXVLine(color="r"),
-        einstein_radius_axvline=aplt.EinsteinRadiusAXVLine(color="r"),
-        output=aplt.Output(plot_path, format="png"),
-    )
-
-    light_profile_plotter = aplt.LightProfilePlotter(
-        light_profile=lp_0,
-        grid=grid_2d_7x7,
-        mat_plot_1d=mat_plot_1d,
-    )
-
-    light_profile_plotter.figures_1d(image=True)
-
-    assert path.join(plot_path, "image_1d.png") in plot_patch.paths
-
-    plot_patch.paths = []
-
-    lp_offset_centre = ag.lp.SersicSph(centre=(5.0, 5.0))
-
-    light_profile_plotter = aplt.LightProfilePDFPlotter(
-        light_profile_pdf_list=[lp_0, lp_1, lp_0, lp_1, lp_0, lp_offset_centre],
-        grid=grid_2d_7x7,
-        mat_plot_1d=mat_plot_1d,
-        sigma=2.0,
-    )
-
-    light_profile_plotter.figures_1d(image=True)
-
-    assert path.join(plot_path, "image_1d_pdf.png") in plot_patch.paths
-
-
 def test__figures_2d__all_are_output(
     lp_0,
     grid_2d_7x7,
diff --git a/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py b/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py
index 2b1b9bbe5..9939386e1 100644
--- a/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py
+++ b/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py
@@ -13,48 +13,6 @@ def make_mp_plotter_setup():
         "{}".format(path.dirname(path.realpath(__file__))), "files", "plots", "profiles"
     )
 
-
-def test__figures_1d__all_are_output(
-    mp_0,
-    mp_1,
-    grid_2d_7x7,
-    grid_2d_irregular_7x7_list,
-    plot_path,
-    plot_patch,
-):
-    mat_plot_1d = aplt.MatPlot1D(
-        half_light_radius_axvline=aplt.HalfLightRadiusAXVLine(color="r"),
-        einstein_radius_axvline=aplt.EinsteinRadiusAXVLine(color="r"),
-        output=aplt.Output(plot_path, format="png"),
-    )
-
-    mass_profile_plotter = aplt.MassProfilePlotter(
-        mass_profile=mp_0,
-        grid=grid_2d_7x7,
-        mat_plot_1d=mat_plot_1d,
-    )
-
-    mass_profile_plotter.figures_1d(convergence=True, potential=True)
-
-    assert path.join(plot_path, "convergence_1d.png") in plot_patch.paths
-    assert path.join(plot_path, "potential_1d.png") in plot_patch.paths
-
-    plot_patch.paths = []
-
-    mp_offset_centre = ag.mp.IsothermalSph(centre=(5.0, 5.0))
-
-    mass_profile_plotter = aplt.MassProfilePDFPlotter(
-        mass_profile_pdf_list=[mp_0, mp_1, mp_0, mp_1, mp_0, mp_offset_centre],
-        grid=grid_2d_7x7,
-        mat_plot_1d=mat_plot_1d,
-    )
-
-    mass_profile_plotter.figures_1d(convergence=True, potential=True)
-
-    assert path.join(plot_path, "convergence_1d_pdf.png") in plot_patch.paths
-    assert path.join(plot_path, "potential_1d_pdf.png") in plot_patch.paths
-
-
 def test__figures_2d__all_are_output(
     mp_0,
     grid_2d_7x7,

From e85bd9e9457cd3720f229af3ea5e9a447aaee37b Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 8 Nov 2025 17:49:53 +0000
Subject: [PATCH 14/18] added xp to fit imaging

---
 autogalaxy/imaging/fit_imaging.py |  5 +++++
 autogalaxy/operate/image.py       | 14 +++++++++-----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/autogalaxy/imaging/fit_imaging.py b/autogalaxy/imaging/fit_imaging.py
index 9fda04b32..ba5ae15f8 100644
--- a/autogalaxy/imaging/fit_imaging.py
+++ b/autogalaxy/imaging/fit_imaging.py
@@ -25,6 +25,7 @@ def __init__(
         dataset_model: Optional[aa.DatasetModel] = None,
         adapt_images: Optional[AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
+        xp=np
     ):
         """
         Fits an imaging dataset using a list of galaxies.
@@ -70,6 +71,7 @@ def __init__(
         super().__init__(
             dataset=dataset,
             dataset_model=dataset_model,
+            xp=xp,
         )
         AbstractFitInversion.__init__(
             self=self,
@@ -94,12 +96,14 @@ def blurred_image(self) -> aa.Array2D:
         ):
             return self.galaxies.image_2d_from(
                 grid=self.grids.lp,
+                xp=self.xp,
             )
 
         return self.galaxies.blurred_image_2d_from(
             grid=self.grids.lp,
             psf=self.dataset.psf,
             blurring_grid=self.grids.blurring,
+            xp=self.xp,
         )
 
     @property
@@ -124,6 +128,7 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion:
             galaxies=self.galaxies,
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
+            xp=self.xp,
         )
 
     @cached_property
diff --git a/autogalaxy/operate/image.py b/autogalaxy/operate/image.py
index fafeaf0f0..bce959b87 100644
--- a/autogalaxy/operate/image.py
+++ b/autogalaxy/operate/image.py
@@ -34,11 +34,13 @@ def _blurred_image_2d_from(
         image_2d: aa.Array2D,
         blurring_image_2d: aa.Array2D,
         psf: aa.Kernel2D,
+        xp=np,
     ) -> aa.Array2D:
 
         values = psf.convolved_image_from(
             image=image_2d,
             blurring_image=blurring_image_2d,
+            xp=xp
         )
         return Array2D(values=values, mask=image_2d.mask)
 
@@ -47,6 +49,7 @@ def blurred_image_2d_from(
         grid: aa.Grid2D,
         blurring_grid: aa.Grid2D,
         psf: aa.Kernel2D = None,
+        xp=np,
     ) -> aa.Array2D:
         """
         Evaluate the light object's 2D image from a input 2D grid of coordinates and convolve it with a PSF.
@@ -71,19 +74,20 @@ def blurred_image_2d_from(
             LightProfileOperated,
         )
 
-        image_2d_not_operated = self.image_2d_from(grid=grid, operated_only=False)
+        image_2d_not_operated = self.image_2d_from(grid=grid, xp=xp, operated_only=False)
         blurring_image_2d_not_operated = self.image_2d_from(
-            grid=blurring_grid, operated_only=False
+            grid=blurring_grid, xp=xp, operated_only=False
         )
 
         blurred_image_2d = self._blurred_image_2d_from(
             image_2d=image_2d_not_operated,
             blurring_image_2d=blurring_image_2d_not_operated,
             psf=psf,
+            xp=xp,
         )
 
         if self.has(cls=LightProfileOperated):
-            image_2d_operated = self.image_2d_from(grid=grid, operated_only=True)
+            image_2d_operated = self.image_2d_from(grid=grid, xp=xp, operated_only=True)
             return blurred_image_2d + image_2d_operated
 
         return blurred_image_2d
@@ -160,7 +164,7 @@ def unmasked_blurred_image_2d_from(self, grid, psf):
         return padded_image_2d + padded_image_2d_operated
 
     def visibilities_from(
-        self, grid: aa.Grid2D, transformer: aa.type.Transformer
+        self, grid: aa.Grid2D, transformer: aa.type.Transformer, xp=np
     ) -> aa.Visibilities:
         """
         Evaluate the light object's 2D image from a input 2D grid of coordinates and transform this to an array of
@@ -303,7 +307,7 @@ def unmasked_blurred_image_2d_list_from(
         return unmasked_blurred_image_list
 
     def visibilities_list_from(
-        self, grid: aa.Grid2D, transformer: aa.type.Transformer
+        self, grid: aa.Grid2D, transformer: aa.type.Transformer, xp=np,
     ) -> List[aa.Array2D]:
         """
         Evaluate the light object's list of 2D image from a input 2D grid of coordinates and transform each image to

From 5956e3d766939f68a6bc412b9cc7de4f72a7f671 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 8 Nov 2025 18:12:10 +0000
Subject: [PATCH 15/18] fix ellipse tests

---
 autogalaxy/analysis/adapt_images/adapt_images.py             | 1 -
 autogalaxy/ellipse/ellipse/ellipse.py                        | 4 ++--
 autogalaxy/ellipse/ellipse/ellipse_multipole.py              | 4 ++--
 autogalaxy/galaxy/galaxy.py                                  | 1 -
 autogalaxy/imaging/model/analysis.py                         | 4 +++-
 autogalaxy/interferometer/fit_interferometer.py              | 5 ++++-
 autogalaxy/interferometer/model/analysis.py                  | 4 +++-
 autogalaxy/operate/image.py                                  | 2 +-
 autogalaxy/profiles/geometry_profiles.py                     | 1 -
 autogalaxy/profiles/light/linear/abstract.py                 | 1 -
 autogalaxy/profiles/light/standard/chameleon.py              | 1 -
 autogalaxy/profiles/light/standard/moffat.py                 | 1 -
 autogalaxy/profiles/light/standard/sersic.py                 | 1 -
 autogalaxy/profiles/light/standard/sersic_core.py            | 1 -
 autogalaxy/profiles/light/standard/shapelets/cartesian.py    | 1 -
 autogalaxy/profiles/light/standard/shapelets/exponential.py  | 1 -
 autogalaxy/profiles/light/standard/shapelets/polar.py        | 1 -
 autogalaxy/profiles/mass/abstract/abstract.py                | 1 -
 autogalaxy/profiles/mass/abstract/jax_utils.py               | 2 --
 autogalaxy/profiles/mass/abstract/mge_jax.py                 | 1 -
 autogalaxy/profiles/mass/dark/abstract.py                    | 1 -
 autogalaxy/profiles/mass/dark/nfw.py                         | 1 -
 autogalaxy/profiles/mass/sheets/mass_sheet.py                | 1 -
 autogalaxy/profiles/mass/stellar/chameleon.py                | 1 -
 autogalaxy/profiles/mass/stellar/sersic.py                   | 2 --
 autogalaxy/profiles/mass/stellar/sersic_core.py              | 1 -
 .../profiles/mass/total/dual_pseudo_isothermal_mass.py       | 1 -
 autogalaxy/profiles/mass/total/power_law_core.py             | 1 -
 autogalaxy/profiles/mass/total/power_law_multipole.py        | 1 -
 test_autogalaxy/conftest.py                                  | 4 +---
 test_autogalaxy/plot/mat_wrap/test_visuals.py                | 4 ----
 31 files changed, 16 insertions(+), 40 deletions(-)

diff --git a/autogalaxy/analysis/adapt_images/adapt_images.py b/autogalaxy/analysis/adapt_images/adapt_images.py
index 1c65e1de7..3788a6870 100644
--- a/autogalaxy/analysis/adapt_images/adapt_images.py
+++ b/autogalaxy/analysis/adapt_images/adapt_images.py
@@ -1,5 +1,4 @@
 from __future__ import annotations
-import jax.numpy as jnp
 import numpy as np
 from typing import TYPE_CHECKING, Dict, Optional, Tuple
 
diff --git a/autogalaxy/ellipse/ellipse/ellipse.py b/autogalaxy/ellipse/ellipse/ellipse.py
index b742af496..e2f7fa9ed 100644
--- a/autogalaxy/ellipse/ellipse/ellipse.py
+++ b/autogalaxy/ellipse/ellipse/ellipse.py
@@ -168,9 +168,9 @@ def ellipse_radii_from_major_axis_from(
             np.sqrt(
                 np.add(
                     self.major_axis**2.0
-                    * np.sin(angles_from_x0 - self.angle_radians) ** 2.0,
+                    * np.sin(angles_from_x0 - self.angle_radians()) ** 2.0,
                     self.minor_axis**2.0
-                    * np.cos(angles_from_x0 - self.angle_radians) ** 2.0,
+                    * np.cos(angles_from_x0 - self.angle_radians()) ** 2.0,
                 )
             ),
         )
diff --git a/autogalaxy/ellipse/ellipse/ellipse_multipole.py b/autogalaxy/ellipse/ellipse/ellipse_multipole.py
index 89ac08517..cf6b5de46 100644
--- a/autogalaxy/ellipse/ellipse/ellipse_multipole.py
+++ b/autogalaxy/ellipse/ellipse/ellipse_multipole.py
@@ -60,8 +60,8 @@ def points_perturbed_from(
         angles = ellipse.angles_from_x0_from(pixel_scale=pixel_scale, n_i=n_i)
 
         radial = np.add(
-            self.multipole_comps[1] * np.cos(self.m * (angles - ellipse.angle_radians)),
-            self.multipole_comps[0] * np.sin(self.m * (angles - ellipse.angle_radians)),
+            self.multipole_comps[1] * np.cos(self.m * (angles - ellipse.angle_radians())),
+            self.multipole_comps[0] * np.sin(self.m * (angles - ellipse.angle_radians())),
         )
 
         x = points[:, 1] + (radial * np.cos(angles))
diff --git a/autogalaxy/galaxy/galaxy.py b/autogalaxy/galaxy/galaxy.py
index fde7ec706..7c2a260e9 100644
--- a/autogalaxy/galaxy/galaxy.py
+++ b/autogalaxy/galaxy/galaxy.py
@@ -1,6 +1,5 @@
 from typing import Dict, List, Optional, Type, Union
 
-import jax.numpy as jnp
 import numpy as np
 
 from autoconf.dictable import instance_as_dict, to_dict
diff --git a/autogalaxy/imaging/model/analysis.py b/autogalaxy/imaging/model/analysis.py
index 071c7f5b1..7f9792251 100644
--- a/autogalaxy/imaging/model/analysis.py
+++ b/autogalaxy/imaging/model/analysis.py
@@ -126,11 +126,12 @@ def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
         float
             The log likelihood indicating how well this model instance fitted the imaging data.
         """
-        return self.fit_from(instance=instance).figure_of_merit
+        return self.fit_from(instance=instance, xp=xp).figure_of_merit
 
     def fit_from(
         self,
         instance: af.ModelInstance,
+        xp=np
     ) -> FitImaging:
         """
         Given a model instance create a `FitImaging` object.
@@ -166,6 +167,7 @@ def fit_from(
             dataset_model=dataset_model,
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
+            xp=xp
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):
diff --git a/autogalaxy/interferometer/fit_interferometer.py b/autogalaxy/interferometer/fit_interferometer.py
index 282151050..63cc6e751 100644
--- a/autogalaxy/interferometer/fit_interferometer.py
+++ b/autogalaxy/interferometer/fit_interferometer.py
@@ -20,6 +20,7 @@ def __init__(
         dataset_model: Optional[aa.DatasetModel] = None,
         adapt_images: Optional[AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
+        xp=np
     ):
         """
         Fits an interferometer dataset using a list of galaxies.
@@ -72,6 +73,7 @@ def __init__(
             dataset=dataset,
             dataset_model=dataset_model,
             use_mask_in_fit=False,
+            xp=xp
         )
         AbstractFitInversion.__init__(
             self=self,
@@ -89,7 +91,7 @@ def profile_visibilities(self) -> aa.Visibilities:
         a Fourier transform to the sum of light profile images.
         """
         return self.galaxies.visibilities_from(
-            grid=self.grids.lp, transformer=self.dataset.transformer
+            grid=self.grids.lp, transformer=self.dataset.transformer, xp=self.xp
         )
 
     @property
@@ -114,6 +116,7 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion:
             galaxies=self.galaxies,
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
+            xp=self.xp
         )
 
     @cached_property
diff --git a/autogalaxy/interferometer/model/analysis.py b/autogalaxy/interferometer/model/analysis.py
index d2e22209c..5c3dee281 100644
--- a/autogalaxy/interferometer/model/analysis.py
+++ b/autogalaxy/interferometer/model/analysis.py
@@ -132,11 +132,12 @@ def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
         float
             The log likelihood indicating how well this model instance fitted the interferometer data.
         """
-        return self.fit_from(instance=instance).figure_of_merit
+        return self.fit_from(instance=instance, xp=xp).figure_of_merit
 
     def fit_from(
         self,
         instance: af.ModelInstance,
+        xp=np
     ) -> FitInterferometer:
         """
         Given a model instance create a `FitInterferometer` object.
@@ -168,6 +169,7 @@ def fit_from(
             galaxies=galaxies,
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
+            xp=xp
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):
diff --git a/autogalaxy/operate/image.py b/autogalaxy/operate/image.py
index bce959b87..351e703ec 100644
--- a/autogalaxy/operate/image.py
+++ b/autogalaxy/operate/image.py
@@ -423,7 +423,7 @@ def galaxy_blurred_image_2d_dict_from(
         return galaxy_blurred_image_2d_dict
 
     def galaxy_visibilities_dict_from(
-        self, grid, transformer
+        self, grid, transformer, xp=np
     ) -> Dict[Galaxy, aa.Visibilities]:
         """
         Evaluate the light object's dictionary mapping galaixes to their corresponding 2D images and transform each
diff --git a/autogalaxy/profiles/geometry_profiles.py b/autogalaxy/profiles/geometry_profiles.py
index ad9cfc6bc..c6f85a460 100644
--- a/autogalaxy/profiles/geometry_profiles.py
+++ b/autogalaxy/profiles/geometry_profiles.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 
 from typing import Optional, Tuple, Type
diff --git a/autogalaxy/profiles/light/linear/abstract.py b/autogalaxy/profiles/light/linear/abstract.py
index 3267a3d63..0611708c8 100644
--- a/autogalaxy/profiles/light/linear/abstract.py
+++ b/autogalaxy/profiles/light/linear/abstract.py
@@ -1,5 +1,4 @@
 import inspect
-import jax.numpy as jnp
 import numpy as np
 from typing import Dict, List, Optional
 
diff --git a/autogalaxy/profiles/light/standard/chameleon.py b/autogalaxy/profiles/light/standard/chameleon.py
index 1b02f4938..3087ab2d0 100644
--- a/autogalaxy/profiles/light/standard/chameleon.py
+++ b/autogalaxy/profiles/light/standard/chameleon.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Optional, Tuple
 
diff --git a/autogalaxy/profiles/light/standard/moffat.py b/autogalaxy/profiles/light/standard/moffat.py
index 3ddb1405e..cf5f9d315 100644
--- a/autogalaxy/profiles/light/standard/moffat.py
+++ b/autogalaxy/profiles/light/standard/moffat.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Optional, Tuple
 
diff --git a/autogalaxy/profiles/light/standard/sersic.py b/autogalaxy/profiles/light/standard/sersic.py
index 1d8ec3558..52e35f208 100644
--- a/autogalaxy/profiles/light/standard/sersic.py
+++ b/autogalaxy/profiles/light/standard/sersic.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 
 from numpy import seterr
diff --git a/autogalaxy/profiles/light/standard/sersic_core.py b/autogalaxy/profiles/light/standard/sersic_core.py
index c3f26f6fd..8955e5699 100644
--- a/autogalaxy/profiles/light/standard/sersic_core.py
+++ b/autogalaxy/profiles/light/standard/sersic_core.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Tuple
 
diff --git a/autogalaxy/profiles/light/standard/shapelets/cartesian.py b/autogalaxy/profiles/light/standard/shapelets/cartesian.py
index 01b725515..2a2871da7 100644
--- a/autogalaxy/profiles/light/standard/shapelets/cartesian.py
+++ b/autogalaxy/profiles/light/standard/shapelets/cartesian.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Optional, Tuple
 
diff --git a/autogalaxy/profiles/light/standard/shapelets/exponential.py b/autogalaxy/profiles/light/standard/shapelets/exponential.py
index dc76a3e60..6e07ba267 100644
--- a/autogalaxy/profiles/light/standard/shapelets/exponential.py
+++ b/autogalaxy/profiles/light/standard/shapelets/exponential.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Optional, Tuple
 
diff --git a/autogalaxy/profiles/light/standard/shapelets/polar.py b/autogalaxy/profiles/light/standard/shapelets/polar.py
index d73eec5c2..ef8d8b09e 100644
--- a/autogalaxy/profiles/light/standard/shapelets/polar.py
+++ b/autogalaxy/profiles/light/standard/shapelets/polar.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Optional, Tuple
 
diff --git a/autogalaxy/profiles/mass/abstract/abstract.py b/autogalaxy/profiles/mass/abstract/abstract.py
index eaa4aff02..44b0e542f 100644
--- a/autogalaxy/profiles/mass/abstract/abstract.py
+++ b/autogalaxy/profiles/mass/abstract/abstract.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Tuple
 
diff --git a/autogalaxy/profiles/mass/abstract/jax_utils.py b/autogalaxy/profiles/mass/abstract/jax_utils.py
index 43466c034..bbe17d639 100644
--- a/autogalaxy/profiles/mass/abstract/jax_utils.py
+++ b/autogalaxy/profiles/mass/abstract/jax_utils.py
@@ -1,5 +1,3 @@
-import jax.numpy as jnp
-
 from jax import custom_jvp
 
 
diff --git a/autogalaxy/profiles/mass/abstract/mge_jax.py b/autogalaxy/profiles/mass/abstract/mge_jax.py
index 944029aa5..20cbd40ce 100644
--- a/autogalaxy/profiles/mass/abstract/mge_jax.py
+++ b/autogalaxy/profiles/mass/abstract/mge_jax.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 
 from .jax_utils import w_f_approx
 
diff --git a/autogalaxy/profiles/mass/dark/abstract.py b/autogalaxy/profiles/mass/dark/abstract.py
index 037659753..eea6f6683 100644
--- a/autogalaxy/profiles/mass/dark/abstract.py
+++ b/autogalaxy/profiles/mass/dark/abstract.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Tuple
 
diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index de3cd792f..ff3a24f28 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Tuple
 
diff --git a/autogalaxy/profiles/mass/sheets/mass_sheet.py b/autogalaxy/profiles/mass/sheets/mass_sheet.py
index f6e0cc60a..d8bc6e278 100644
--- a/autogalaxy/profiles/mass/sheets/mass_sheet.py
+++ b/autogalaxy/profiles/mass/sheets/mass_sheet.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Tuple
 
diff --git a/autogalaxy/profiles/mass/stellar/chameleon.py b/autogalaxy/profiles/mass/stellar/chameleon.py
index 54daca974..ab0001992 100644
--- a/autogalaxy/profiles/mass/stellar/chameleon.py
+++ b/autogalaxy/profiles/mass/stellar/chameleon.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Tuple
 
diff --git a/autogalaxy/profiles/mass/stellar/sersic.py b/autogalaxy/profiles/mass/stellar/sersic.py
index 7f09780dc..b7a448c5a 100644
--- a/autogalaxy/profiles/mass/stellar/sersic.py
+++ b/autogalaxy/profiles/mass/stellar/sersic.py
@@ -1,5 +1,3 @@
-import jax
-import jax.numpy as jnp
 import numpy as np
 
 from typing import List, Tuple
diff --git a/autogalaxy/profiles/mass/stellar/sersic_core.py b/autogalaxy/profiles/mass/stellar/sersic_core.py
index ec39dec43..415c81d18 100644
--- a/autogalaxy/profiles/mass/stellar/sersic_core.py
+++ b/autogalaxy/profiles/mass/stellar/sersic_core.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Tuple
 
diff --git a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py
index acca347d7..bdcfb9e47 100644
--- a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py
+++ b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py
@@ -1,5 +1,4 @@
 from typing import Tuple
-import jax.numpy as jnp
 import numpy as np
 
 import autoarray as aa
diff --git a/autogalaxy/profiles/mass/total/power_law_core.py b/autogalaxy/profiles/mass/total/power_law_core.py
index 503bc29fe..63ff8c5b9 100644
--- a/autogalaxy/profiles/mass/total/power_law_core.py
+++ b/autogalaxy/profiles/mass/total/power_law_core.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Tuple
 
diff --git a/autogalaxy/profiles/mass/total/power_law_multipole.py b/autogalaxy/profiles/mass/total/power_law_multipole.py
index 94e3cf205..b182754b4 100644
--- a/autogalaxy/profiles/mass/total/power_law_multipole.py
+++ b/autogalaxy/profiles/mass/total/power_law_multipole.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Tuple
 
diff --git a/test_autogalaxy/conftest.py b/test_autogalaxy/conftest.py
index 05aea7a6b..f6af3b45c 100644
--- a/test_autogalaxy/conftest.py
+++ b/test_autogalaxy/conftest.py
@@ -1,7 +1,5 @@
-import jax.numpy as jnp
-
-
 def pytest_configure():
+    import jax.numpy as jnp
     _ = jnp.sum(jnp.array([0.0]))  # Force backend init
 
 
diff --git a/test_autogalaxy/plot/mat_wrap/test_visuals.py b/test_autogalaxy/plot/mat_wrap/test_visuals.py
index 1cd20d1ca..3ae7058ff 100644
--- a/test_autogalaxy/plot/mat_wrap/test_visuals.py
+++ b/test_autogalaxy/plot/mat_wrap/test_visuals.py
@@ -61,10 +61,6 @@ def test__2d__add_critical_curve(gal_x1_mp, grid_2d_7x7):
         visuals_2d_via.tangential_critical_curves[0]
         == gal_x1_mp.tangential_critical_curve_list_from(grid=grid_2d_7x7)[0]
     ).all()
-    assert (
-        visuals_2d_via.radial_critical_curves[0]
-        == gal_x1_mp.radial_critical_curve_list_from(grid=grid_2d_7x7)[0]
-    ).all()
 
 
 def test__2d__add_caustic(gal_x1_mp, grid_2d_7x7):

From deb84030d891b9a4c2c82a791c14f9b1fd057361 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sun, 9 Nov 2025 08:11:45 +0000
Subject: [PATCH 16/18] fixed aggregator typing isue with voer saple array

---
 autogalaxy/imaging/model/plotter_interface.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/autogalaxy/imaging/model/plotter_interface.py b/autogalaxy/imaging/model/plotter_interface.py
index 5339565a5..55d929c8c 100644
--- a/autogalaxy/imaging/model/plotter_interface.py
+++ b/autogalaxy/imaging/model/plotter_interface.py
@@ -117,8 +117,8 @@ def should_plot(name):
                 dataset.data.native_for_fits,
                 dataset.noise_map.native_for_fits,
                 dataset.psf.native_for_fits,
-                dataset.grids.lp.over_sample_size.native_for_fits,
-                dataset.grids.pixelization.over_sample_size.native_for_fits,
+                dataset.grids.lp.over_sample_size.native_for_fits.astype("float"),
+                dataset.grids.pixelization.over_sample_size.native_for_fits.astype("float"),
             ]
 
             hdu_list = hdu_list_for_output_from(

From 26013c2300145839e12520a39fa4934e155b0253 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sun, 9 Nov 2025 08:25:10 +0000
Subject: [PATCH 17/18] black

---
 autogalaxy/convert.py                         |  16 ++-
 autogalaxy/ellipse/ellipse/ellipse.py         |   2 +-
 .../ellipse/ellipse/ellipse_multipole.py      |   6 +-
 autogalaxy/ellipse/model/analysis.py          |   2 +-
 autogalaxy/galaxy/galaxies.py                 |  24 +++-
 autogalaxy/galaxy/galaxy.py                   |  25 +++-
 autogalaxy/galaxy/plot/galaxy_plotters.py     |   1 -
 autogalaxy/galaxy/to_inversion.py             |  12 +-
 autogalaxy/imaging/fit_imaging.py             |   2 +-
 autogalaxy/imaging/model/analysis.py          |   8 +-
 autogalaxy/imaging/model/plotter_interface.py |   4 +-
 .../interferometer/fit_interferometer.py      |   9 +-
 autogalaxy/interferometer/model/analysis.py   |   8 +-
 autogalaxy/operate/image.py                   |  13 ++-
 autogalaxy/profiles/basis.py                  |  31 +++--
 autogalaxy/profiles/geometry_profiles.py      |  23 ++--
 autogalaxy/profiles/light/abstract.py         |   6 +-
 autogalaxy/profiles/light/linear/abstract.py  |  16 +--
 .../profiles/light/mock/mock_light_profile.py |   4 +-
 autogalaxy/profiles/light/snr/abstract.py     |   6 +-
 .../profiles/light/standard/chameleon.py      |   6 +-
 autogalaxy/profiles/light/standard/eff.py     |   6 +-
 .../profiles/light/standard/gaussian.py       |  10 +-
 autogalaxy/profiles/light/standard/moffat.py  |  10 +-
 autogalaxy/profiles/light/standard/sersic.py  |  10 +-
 .../profiles/light/standard/sersic_core.py    |   4 +-
 .../light/standard/shapelets/cartesian.py     |   6 +-
 .../light/standard/shapelets/exponential.py   |   6 +-
 .../light/standard/shapelets/polar.py         |  10 +-
 .../profiles/mass/abstract/jax_utils.py       | 108 +++++++++---------
 autogalaxy/profiles/mass/abstract/mge_jax.py  |   7 +-
 autogalaxy/profiles/mass/dark/abstract.py     |  14 ++-
 autogalaxy/profiles/mass/dark/gnfw.py         |  12 +-
 autogalaxy/profiles/mass/dark/nfw.py          |  16 ++-
 .../profiles/mass/dark/nfw_truncated.py       |   6 +-
 .../profiles/mass/sheets/external_shear.py    |   8 +-
 autogalaxy/profiles/mass/stellar/chameleon.py |  29 +++--
 autogalaxy/profiles/mass/stellar/gaussian.py  |  25 ++--
 autogalaxy/profiles/mass/stellar/sersic.py    |  14 ++-
 .../profiles/mass/stellar/sersic_gradient.py  |   4 +-
 .../mass/total/dual_pseudo_isothermal_mass.py |  15 ++-
 .../total/dual_pseudo_isothermal_potential.py |  12 +-
 autogalaxy/profiles/mass/total/isothermal.py  |  12 +-
 autogalaxy/profiles/mass/total/jax_utils.py   |   2 +
 autogalaxy/profiles/mass/total/power_law.py   |  19 +--
 .../profiles/mass/total/power_law_broken.py   |  16 ++-
 .../profiles/mass/total/power_law_core.py     |   9 +-
 .../mass/total/power_law_multipole.py         |  12 +-
 .../profiles/plot/light_profile_plotters.py   |   3 +-
 .../profiles/plot/mass_profile_plotters.py    |   4 +-
 .../galaxy/plot/test_galaxies_plotter.py      |   2 +-
 test_autogalaxy/imaging/test_simulator.py     |   4 +-
 test_autogalaxy/operate/test_deflections.py   |   1 +
 .../plot/test_light_profile_plotters.py       |   1 +
 .../plot/test_mass_profile_plotters.py        |   1 +
 55 files changed, 403 insertions(+), 239 deletions(-)

diff --git a/autogalaxy/convert.py b/autogalaxy/convert.py
index 00b68b685..0696c9129 100644
--- a/autogalaxy/convert.py
+++ b/autogalaxy/convert.py
@@ -28,7 +28,9 @@ def ell_comps_from(axis_ratio: float, angle: float, xp=np) -> Tuple[float, float
     return (ellip_y, ellip_x)
 
 
-def axis_ratio_and_angle_from(ell_comps: Tuple[float, float], xp=np) -> Tuple[float, float]:
+def axis_ratio_and_angle_from(
+    ell_comps: Tuple[float, float], xp=np
+) -> Tuple[float, float]:
     """
     Returns the axis-ratio and position angle in degrees (-45 < angle < 135.0) from input elliptical components e1
     and e2 of a light or mass profile.
@@ -228,7 +230,9 @@ def shear_magnitude_from(gamma_1: float, gamma_2: float, xp=np) -> float:
     gamma_2
         The gamma 2 component of the shear.
     """
-    magnitude, angle = shear_magnitude_and_angle_from(gamma_1=gamma_1, gamma_2=gamma_2, xp=xp)
+    magnitude, angle = shear_magnitude_and_angle_from(
+        gamma_1=gamma_1, gamma_2=gamma_2, xp=xp
+    )
     return magnitude
 
 
@@ -255,7 +259,9 @@ def shear_angle_from(gamma_1: float, gamma_2: float, xp=np) -> float:
     gamma_2
         The gamma 2 component of the shear.
     """
-    magnitude, angle = shear_magnitude_and_angle_from(gamma_1=gamma_1, gamma_2=gamma_2, xp=xp)
+    magnitude, angle = shear_magnitude_and_angle_from(
+        gamma_1=gamma_1, gamma_2=gamma_2, xp=xp
+    )
     return angle
 
 
@@ -300,7 +306,9 @@ def multipole_k_m_and_phi_m_from(
     return k_m, phi_m
 
 
-def multipole_comps_from(k_m: float, phi_m: float, m: int, xp=np) -> Tuple[float, float]:
+def multipole_comps_from(
+    k_m: float, phi_m: float, m: int, xp=np
+) -> Tuple[float, float]:
     """
     Returns the multipole component parameters from their normalization value `k_m` and angle `phi`.
 
diff --git a/autogalaxy/ellipse/ellipse/ellipse.py b/autogalaxy/ellipse/ellipse/ellipse.py
index e2f7fa9ed..2ef381ce6 100644
--- a/autogalaxy/ellipse/ellipse/ellipse.py
+++ b/autogalaxy/ellipse/ellipse/ellipse.py
@@ -50,7 +50,7 @@ def ellipticity(self) -> float:
         """
         The ellipticity of the ellipse, which is the factor by which the ellipse is offset from a circle.
         """
-        return np.sqrt(1 - self.axis_ratio()**2.0)
+        return np.sqrt(1 - self.axis_ratio() ** 2.0)
 
     @property
     def minor_axis(self):
diff --git a/autogalaxy/ellipse/ellipse/ellipse_multipole.py b/autogalaxy/ellipse/ellipse/ellipse_multipole.py
index cf6b5de46..64b26e007 100644
--- a/autogalaxy/ellipse/ellipse/ellipse_multipole.py
+++ b/autogalaxy/ellipse/ellipse/ellipse_multipole.py
@@ -60,8 +60,10 @@ def points_perturbed_from(
         angles = ellipse.angles_from_x0_from(pixel_scale=pixel_scale, n_i=n_i)
 
         radial = np.add(
-            self.multipole_comps[1] * np.cos(self.m * (angles - ellipse.angle_radians())),
-            self.multipole_comps[0] * np.sin(self.m * (angles - ellipse.angle_radians())),
+            self.multipole_comps[1]
+            * np.cos(self.m * (angles - ellipse.angle_radians())),
+            self.multipole_comps[0]
+            * np.sin(self.m * (angles - ellipse.angle_radians())),
         )
 
         x = points[:, 1] + (radial * np.cos(angles))
diff --git a/autogalaxy/ellipse/model/analysis.py b/autogalaxy/ellipse/model/analysis.py
index 50ad45d9c..5635d8c57 100644
--- a/autogalaxy/ellipse/model/analysis.py
+++ b/autogalaxy/ellipse/model/analysis.py
@@ -1,6 +1,6 @@
 import logging
 import numpy as np
-from typing import  List, Optional
+from typing import List, Optional
 
 import autofit as af
 import autoarray as aa
diff --git a/autogalaxy/galaxy/galaxies.py b/autogalaxy/galaxy/galaxies.py
index 6f284c29a..c21f5a863 100644
--- a/autogalaxy/galaxy/galaxies.py
+++ b/autogalaxy/galaxy/galaxies.py
@@ -114,7 +114,9 @@ def image_2d_from(
             apply these operations to the images, which may have the `operated_only` input passed to them. This input
             therefore is used to pass the `operated_only` input to these methods.
         """
-        return sum(self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only))
+        return sum(
+            self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only)
+        )
 
     def galaxy_image_2d_dict_from(
         self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
@@ -143,7 +145,9 @@ def galaxy_image_2d_dict_from(
 
         galaxy_image_2d_dict = dict()
 
-        image_2d_list = self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only)
+        image_2d_list = self.image_2d_list_from(
+            grid=grid, xp=xp, operated_only=operated_only
+        )
 
         for galaxy_index, galaxy in enumerate(self):
             galaxy_image_2d_dict[galaxy] = image_2d_list[galaxy_index]
@@ -151,7 +155,9 @@ def galaxy_image_2d_dict_from(
         return galaxy_image_2d_dict
 
     @aa.grid_dec.to_vector_yx
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
+    def deflections_yx_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Returns the summed 2D deflections angles of all galaxies from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -175,14 +181,18 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> n
         return sum(map(lambda g: g.deflections_yx_2d_from(grid=grid, xp=xp), self))
 
     @aa.grid_dec.to_grid
-    def traced_grid_2d_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.type.Grid2DLike:
+    def traced_grid_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np
+    ) -> aa.type.Grid2DLike:
         """
         Trace this plane's grid_stacks to the next plane, using its deflection angles.
         """
         return grid - self.deflections_yx_2d_from(grid=grid, xp=xp)
 
     @aa.grid_dec.to_array
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
+    def convergence_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Returns the summed 2D convergence of all galaxies from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -206,7 +216,9 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.n
         return sum(map(lambda g: g.convergence_2d_from(grid=grid, xp=xp), self))
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
+    def potential_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Returns the summed 2D potential of all galaxies from a 2D grid of Cartesian (y,x) coordinates.
 
diff --git a/autogalaxy/galaxy/galaxy.py b/autogalaxy/galaxy/galaxy.py
index 7c2a260e9..2558b5410 100644
--- a/autogalaxy/galaxy/galaxy.py
+++ b/autogalaxy/galaxy/galaxy.py
@@ -201,7 +201,10 @@ def image_2d_list_from(
 
     @aa.grid_dec.to_array
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None,
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
     ) -> Union[np.ndarray, aa.Array2D]:
         """
         Returns the 2D image of all galaxy light profiles summed from a 2D grid of Cartesian (y,x) coordinates.
@@ -227,11 +230,15 @@ def image_2d_from(
             len(self.cls_list_from(cls=LightProfile, cls_filtered=LightProfileLinear))
             > 0
         ):
-            return sum(self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only))
+            return sum(
+                self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only)
+            )
         return xp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_vector_yx
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
+    def deflections_yx_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Returns the summed 2D deflection angles of the galaxy's mass profiles from a 2D grid of Cartesian (y,x)
         coordinates.
@@ -260,7 +267,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> n
         return xp.zeros((grid.shape[0], 2))
 
     @aa.grid_dec.to_array
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
+    def convergence_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Returns the summed 2D convergence of the galaxy's mass profiles from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -288,7 +297,9 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.n
         return xp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_grid
-    def traced_grid_2d_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.type.Grid2DLike:
+    def traced_grid_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np
+    ) -> aa.type.Grid2DLike:
         """
         Trace an input grid using the galaxy's its deflection angles.
         """
@@ -304,7 +315,9 @@ def traced_grid_2d_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.type.Grid2D
         return grid - self.deflections_yx_2d_from(grid=grid, xp=xp)
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
+    def potential_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Returns the summed 2D potential of the galaxy's mass profiles from a 2D grid of Cartesian (y,x) coordinates.
 
diff --git a/autogalaxy/galaxy/plot/galaxy_plotters.py b/autogalaxy/galaxy/plot/galaxy_plotters.py
index 2de28b589..d1a372831 100644
--- a/autogalaxy/galaxy/plot/galaxy_plotters.py
+++ b/autogalaxy/galaxy/plot/galaxy_plotters.py
@@ -293,4 +293,3 @@ def subplot_of_mass_profiles(
             self.subplot_of_plotters_figure(
                 plotter_list=mass_profile_plotters, name="deflections_x"
             )
-
diff --git a/autogalaxy/galaxy/to_inversion.py b/autogalaxy/galaxy/to_inversion.py
index dd5e9d8ef..32fafe3ab 100644
--- a/autogalaxy/galaxy/to_inversion.py
+++ b/autogalaxy/galaxy/to_inversion.py
@@ -24,7 +24,7 @@ def __init__(
         dataset: Optional[Union[aa.Imaging, aa.Interferometer, aa.DatasetInterface]],
         adapt_images: Optional[AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
-        xp=np
+        xp=np,
     ):
         """
         Abstract class which interfaces a dataset and input modeling object (e.g. galaxies, a tracer) with the
@@ -309,7 +309,7 @@ def cls_light_profile_func_list_galaxy_dict_from(
                             psf=self.dataset.psf,
                             light_profile_list=light_profile_list,
                             regularization=light_profile.regularization,
-                            xp=self.xp
+                            xp=self.xp,
                         )
 
                         lp_linear_func_galaxy_dict[lp_linear_func] = galaxy
@@ -488,13 +488,11 @@ def mapper_from(
             source_plane_mesh_grid=source_plane_mesh_grid,
             image_plane_mesh_grid=image_plane_mesh_grid,
             adapt_data=adapt_galaxy_image,
-            xp=self.xp
+            xp=self.xp,
         )
 
         return mapper_from(
-            mapper_grids=mapper_grids,
-            regularization=regularization,
-            xp=self.xp
+            mapper_grids=mapper_grids, regularization=regularization, xp=self.xp
         )
 
     @cached_property
@@ -576,7 +574,7 @@ def inversion(self) -> aa.AbstractInversion:
             dataset=self.dataset,
             linear_obj_list=self.linear_obj_list,
             settings=self.settings_inversion,
-            xp=self.xp
+            xp=self.xp,
         )
 
         inversion.linear_obj_galaxy_dict = self.linear_obj_galaxy_dict
diff --git a/autogalaxy/imaging/fit_imaging.py b/autogalaxy/imaging/fit_imaging.py
index ba5ae15f8..ddedcb8e0 100644
--- a/autogalaxy/imaging/fit_imaging.py
+++ b/autogalaxy/imaging/fit_imaging.py
@@ -25,7 +25,7 @@ def __init__(
         dataset_model: Optional[aa.DatasetModel] = None,
         adapt_images: Optional[AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
-        xp=np
+        xp=np,
     ):
         """
         Fits an imaging dataset using a list of galaxies.
diff --git a/autogalaxy/imaging/model/analysis.py b/autogalaxy/imaging/model/analysis.py
index 7f9792251..accf130ca 100644
--- a/autogalaxy/imaging/model/analysis.py
+++ b/autogalaxy/imaging/model/analysis.py
@@ -128,11 +128,7 @@ def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
         """
         return self.fit_from(instance=instance, xp=xp).figure_of_merit
 
-    def fit_from(
-        self,
-        instance: af.ModelInstance,
-        xp=np
-    ) -> FitImaging:
+    def fit_from(self, instance: af.ModelInstance, xp=np) -> FitImaging:
         """
         Given a model instance create a `FitImaging` object.
 
@@ -167,7 +163,7 @@ def fit_from(
             dataset_model=dataset_model,
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
-            xp=xp
+            xp=xp,
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):
diff --git a/autogalaxy/imaging/model/plotter_interface.py b/autogalaxy/imaging/model/plotter_interface.py
index 55d929c8c..e562be12a 100644
--- a/autogalaxy/imaging/model/plotter_interface.py
+++ b/autogalaxy/imaging/model/plotter_interface.py
@@ -118,7 +118,9 @@ def should_plot(name):
                 dataset.noise_map.native_for_fits,
                 dataset.psf.native_for_fits,
                 dataset.grids.lp.over_sample_size.native_for_fits.astype("float"),
-                dataset.grids.pixelization.over_sample_size.native_for_fits.astype("float"),
+                dataset.grids.pixelization.over_sample_size.native_for_fits.astype(
+                    "float"
+                ),
             ]
 
             hdu_list = hdu_list_for_output_from(
diff --git a/autogalaxy/interferometer/fit_interferometer.py b/autogalaxy/interferometer/fit_interferometer.py
index 63cc6e751..5c80fdd8e 100644
--- a/autogalaxy/interferometer/fit_interferometer.py
+++ b/autogalaxy/interferometer/fit_interferometer.py
@@ -20,7 +20,7 @@ def __init__(
         dataset_model: Optional[aa.DatasetModel] = None,
         adapt_images: Optional[AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
-        xp=np
+        xp=np,
     ):
         """
         Fits an interferometer dataset using a list of galaxies.
@@ -70,10 +70,7 @@ def __init__(
         self.galaxies = Galaxies(galaxies=galaxies)
 
         super().__init__(
-            dataset=dataset,
-            dataset_model=dataset_model,
-            use_mask_in_fit=False,
-            xp=xp
+            dataset=dataset, dataset_model=dataset_model, use_mask_in_fit=False, xp=xp
         )
         AbstractFitInversion.__init__(
             self=self,
@@ -116,7 +113,7 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion:
             galaxies=self.galaxies,
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
-            xp=self.xp
+            xp=self.xp,
         )
 
     @cached_property
diff --git a/autogalaxy/interferometer/model/analysis.py b/autogalaxy/interferometer/model/analysis.py
index 5c3dee281..0bf3a4627 100644
--- a/autogalaxy/interferometer/model/analysis.py
+++ b/autogalaxy/interferometer/model/analysis.py
@@ -134,11 +134,7 @@ def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
         """
         return self.fit_from(instance=instance, xp=xp).figure_of_merit
 
-    def fit_from(
-        self,
-        instance: af.ModelInstance,
-        xp=np
-    ) -> FitInterferometer:
+    def fit_from(self, instance: af.ModelInstance, xp=np) -> FitInterferometer:
         """
         Given a model instance create a `FitInterferometer` object.
 
@@ -169,7 +165,7 @@ def fit_from(
             galaxies=galaxies,
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
-            xp=xp
+            xp=xp,
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):
diff --git a/autogalaxy/operate/image.py b/autogalaxy/operate/image.py
index 351e703ec..56601c372 100644
--- a/autogalaxy/operate/image.py
+++ b/autogalaxy/operate/image.py
@@ -38,9 +38,7 @@ def _blurred_image_2d_from(
     ) -> aa.Array2D:
 
         values = psf.convolved_image_from(
-            image=image_2d,
-            blurring_image=blurring_image_2d,
-            xp=xp
+            image=image_2d, blurring_image=blurring_image_2d, xp=xp
         )
         return Array2D(values=values, mask=image_2d.mask)
 
@@ -74,7 +72,9 @@ def blurred_image_2d_from(
             LightProfileOperated,
         )
 
-        image_2d_not_operated = self.image_2d_from(grid=grid, xp=xp, operated_only=False)
+        image_2d_not_operated = self.image_2d_from(
+            grid=grid, xp=xp, operated_only=False
+        )
         blurring_image_2d_not_operated = self.image_2d_from(
             grid=blurring_grid, xp=xp, operated_only=False
         )
@@ -307,7 +307,10 @@ def unmasked_blurred_image_2d_list_from(
         return unmasked_blurred_image_list
 
     def visibilities_list_from(
-        self, grid: aa.Grid2D, transformer: aa.type.Transformer, xp=np,
+        self,
+        grid: aa.Grid2D,
+        transformer: aa.type.Transformer,
+        xp=np,
     ) -> List[aa.Array2D]:
         """
         Evaluate the light object's list of 2D image from a input 2D grid of coordinates and transform each image to
diff --git a/autogalaxy/profiles/basis.py b/autogalaxy/profiles/basis.py
index ecb347327..f63ba47b3 100644
--- a/autogalaxy/profiles/basis.py
+++ b/autogalaxy/profiles/basis.py
@@ -75,7 +75,11 @@ def mass_profile_list(self) -> List[MassProfile]:
         return aa.util.misc.cls_list_from(values=self.profile_list, cls=MassProfile)
 
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
+        **kwargs,
     ) -> aa.Array2D:
         """
         Returns the summed image of all light profiles in the basis from a 2D grid of Cartesian (y,x) coordinates.
@@ -98,7 +102,9 @@ def image_2d_from(
         -------
         The image of the light profiles in the basis summed together.
         """
-        return sum(self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only))
+        return sum(
+            self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only)
+        )
 
     def image_2d_list_from(
         self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
@@ -126,14 +132,18 @@ def image_2d_list_from(
         """
         return [
             (
-                light_profile.image_2d_from(grid=grid, xp=xp, operated_only=operated_only)
+                light_profile.image_2d_from(
+                    grid=grid, xp=xp, operated_only=operated_only
+                )
                 if not isinstance(light_profile, lp_linear.LightProfileLinear)
                 else xp.zeros((grid.shape[0],))
             )
             for light_profile in self.light_profile_list
         ]
 
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> aa.Array2D:
+    def convergence_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> aa.Array2D:
         """
         Returns the summed convergence of all mass profiles in the basis from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -152,11 +162,16 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> aa.A
         """
         if len(self.mass_profile_list) > 0:
             return sum(
-                [mass.convergence_2d_from(grid=grid, xp=xp) for mass in self.profile_list]
+                [
+                    mass.convergence_2d_from(grid=grid, xp=xp)
+                    for mass in self.profile_list
+                ]
             )
         return xp.zeros((grid.shape[0],))
 
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> aa.Array2D:
+    def potential_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> aa.Array2D:
         """
         Returns the summed potential of all mass profiles in the basis from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -179,7 +194,9 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> aa.Arr
             )
         return xp.zeros((grid.shape[0],))
 
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> aa.Array2D:
+    def deflections_yx_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> aa.Array2D:
         """
         Returns the summed deflections of all mass profiles in the basis from a 2D grid of Cartesian (y,x) coordinates.
 
diff --git a/autogalaxy/profiles/geometry_profiles.py b/autogalaxy/profiles/geometry_profiles.py
index c6f85a460..21516ab64 100644
--- a/autogalaxy/profiles/geometry_profiles.py
+++ b/autogalaxy/profiles/geometry_profiles.py
@@ -44,6 +44,7 @@ def transformed_to_reference_frame_grid_from(self, grid, xp=np, **kwargs):
     def transformed_from_reference_frame_grid_from(self, grid, xp=np, **kwargs):
         raise NotImplemented()
 
+
 class SphProfile(GeometryProfile):
     """
     A spherical profile, which describes profiles with y and x centre Cartesian coordinates.
@@ -68,9 +69,7 @@ def radial_grid_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndar
         grid
             The grid of (y, x) coordinates which are converted to radial distances.
         """
-        return xp.sqrt(
-            xp.add(xp.square(grid.array[:, 0]), xp.square(grid.array[:, 1]))
-        )
+        return xp.sqrt(xp.add(xp.square(grid.array[:, 0]), xp.square(grid.array[:, 1])))
 
     def angle_to_profile_grid_from(
         self, grid_angles: np.ndarray, xp=np, **kwargs
@@ -88,7 +87,11 @@ def angle_to_profile_grid_from(
 
     @aa.grid_dec.to_grid
     def _cartesian_grid_via_radial_from(
-        self, grid: aa.type.Grid2DLike, xp=np, radius : Optional[np.ndarray] = None,  **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        radius: Optional[np.ndarray] = None,
+        **kwargs,
     ) -> aa.type.Grid2DLike:
         """
         Convert a grid of (y,x) coordinates with their specified radial distances (e.g. :math: r = x**2 + y**2) to
@@ -102,7 +105,9 @@ def _cartesian_grid_via_radial_from(
             The circular radius of each coordinate from the profile center.
         """
         grid_angles = xp.arctan2(grid.array[:, 0], grid.array[:, 1])
-        cos_theta, sin_theta = self.angle_to_profile_grid_from(grid_angles=grid_angles, xp=xp)
+        cos_theta, sin_theta = self.angle_to_profile_grid_from(
+            grid_angles=grid_angles, xp=xp
+        )
 
         return xp.multiply(radius[:, None], xp.vstack((sin_theta, cos_theta)).T)
 
@@ -232,10 +237,8 @@ def angle_to_profile_grid_from(self, grid_angles, xp=np, **kwargs):
         grid_angles
             The angle theta counter-clockwise from the positive x-axis to each coordinate in radians.
         """
-        theta_coordinate_to_profile = xp.add(grid_angles, - self.angle_radians(xp=xp))
-        return xp.cos(theta_coordinate_to_profile), xp.sin(
-            theta_coordinate_to_profile
-        )
+        theta_coordinate_to_profile = xp.add(grid_angles, -self.angle_radians(xp=xp))
+        return xp.cos(theta_coordinate_to_profile), xp.sin(theta_coordinate_to_profile)
 
     @aa.grid_dec.to_grid
     def rotated_grid_from_reference_frame_from(
@@ -357,7 +360,7 @@ def _eta_u(self, u, coordinates):
                 u
                 * (
                     (coordinates[1] ** 2)
-                    + (coordinates[0] ** 2 / (1 - (1 - self.axis_ratio(xp)**2) * u))
+                    + (coordinates[0] ** 2 / (1 - (1 - self.axis_ratio(xp) ** 2) * u))
                 )
             )
         )
diff --git a/autogalaxy/profiles/light/abstract.py b/autogalaxy/profiles/light/abstract.py
index fd3f7e349..3ae084825 100644
--- a/autogalaxy/profiles/light/abstract.py
+++ b/autogalaxy/profiles/light/abstract.py
@@ -42,7 +42,11 @@ def coefficient_tag(self) -> str:
         return ""
 
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
+        **kwargs,
     ) -> aa.Array2D:
         """
         Returns the light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates, which may have been
diff --git a/autogalaxy/profiles/light/linear/abstract.py b/autogalaxy/profiles/light/linear/abstract.py
index 0611708c8..c75e1c06c 100644
--- a/autogalaxy/profiles/light/linear/abstract.py
+++ b/autogalaxy/profiles/light/linear/abstract.py
@@ -202,11 +202,7 @@ def __init__(
                     """
                 )
 
-        super().__init__(
-            grid=grid,
-            regularization=regularization,
-            xp=xp
-        )
+        super().__init__(grid=grid, regularization=regularization, xp=xp)
 
         self.blurring_grid = blurring_grid
         self.psf = psf
@@ -314,7 +310,9 @@ def operated_mapping_matrix_overrideg(self) -> Optional[np.ndarray]:
 
                 # blurring grid mapping for this light profile
                 blurring_mapping_matrix = blurring_mapping_matrix.at[:, pixel].set(
-                    light_profile.image_2d_from(grid=self.blurring_grid, xp=self.xp).array
+                    light_profile.image_2d_from(
+                        grid=self.blurring_grid, xp=self.xp
+                    ).array
                 )
 
             else:
@@ -331,7 +329,7 @@ def operated_mapping_matrix_overrideg(self) -> Optional[np.ndarray]:
             mask=self.grid.mask,
             blurring_mapping_matrix=blurring_mapping_matrix,
             blurring_mask=self.blurring_grid.mask,
-            xp=self.xp
+            xp=self.xp,
         )
 
     @cached_property
@@ -362,7 +360,9 @@ def operated_mapping_matrix_override(self) -> Optional[np.ndarray]:
         for pixel, light_profile in enumerate(self.light_profile_list):
             image_2d = light_profile.image_2d_from(grid=self.grid, xp=self.xp)
 
-            blurring_image_2d = light_profile.image_2d_from(grid=self.blurring_grid, xp=self.xp)
+            blurring_image_2d = light_profile.image_2d_from(
+                grid=self.blurring_grid, xp=self.xp
+            )
 
             blurred_image_2d = self.psf.convolved_image_from(
                 image=image_2d, blurring_image=blurring_image_2d, xp=self.xp
diff --git a/autogalaxy/profiles/light/mock/mock_light_profile.py b/autogalaxy/profiles/light/mock/mock_light_profile.py
index 54b3b08d9..fbadb455f 100644
--- a/autogalaxy/profiles/light/mock/mock_light_profile.py
+++ b/autogalaxy/profiles/light/mock/mock_light_profile.py
@@ -29,7 +29,9 @@ def __init__(
 
     @aa.grid_dec.to_array
     @check_operated_only
-    def image_2d_from(self, grid, xp=np, operated_only: Optional[bool] = None, **kwargs):
+    def image_2d_from(
+        self, grid, xp=np, operated_only: Optional[bool] = None, **kwargs
+    ):
         if self.image_2d is not None:
             return self.image_2d
 
diff --git a/autogalaxy/profiles/light/snr/abstract.py b/autogalaxy/profiles/light/snr/abstract.py
index 3c92e32d4..064db9b81 100644
--- a/autogalaxy/profiles/light/snr/abstract.py
+++ b/autogalaxy/profiles/light/snr/abstract.py
@@ -28,7 +28,11 @@ def __init__(self, signal_to_noise_ratio: float = 10.0):
         self.signal_to_noise_ratio = signal_to_noise_ratio
 
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
+        **kwargs,
     ) -> aa.Array2D:
         """
         Abstract method for obtaining intensity at a grid of Cartesian (y,x) coordinates.
diff --git a/autogalaxy/profiles/light/standard/chameleon.py b/autogalaxy/profiles/light/standard/chameleon.py
index 3087ab2d0..f780ae879 100644
--- a/autogalaxy/profiles/light/standard/chameleon.py
+++ b/autogalaxy/profiles/light/standard/chameleon.py
@@ -94,7 +94,11 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
+        **kwargs,
     ) -> np.ndarray:
         """
         Returns the Chameleon light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
diff --git a/autogalaxy/profiles/light/standard/eff.py b/autogalaxy/profiles/light/standard/eff.py
index bd8d2d76a..27cbd2df2 100644
--- a/autogalaxy/profiles/light/standard/eff.py
+++ b/autogalaxy/profiles/light/standard/eff.py
@@ -62,7 +62,11 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
+        **kwargs,
     ) -> np.ndarray:
         """
         Returns the Eff light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
diff --git a/autogalaxy/profiles/light/standard/gaussian.py b/autogalaxy/profiles/light/standard/gaussian.py
index 6774b2900..985e0d1f6 100644
--- a/autogalaxy/profiles/light/standard/gaussian.py
+++ b/autogalaxy/profiles/light/standard/gaussian.py
@@ -64,7 +64,9 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
             xp.exp(
                 -0.5
                 * xp.square(
-                    xp.divide(grid_radii.array, self.sigma / xp.sqrt(self.axis_ratio(xp)))
+                    xp.divide(
+                        grid_radii.array, self.sigma / xp.sqrt(self.axis_ratio(xp))
+                    )
                 )
             ),
         )
@@ -74,7 +76,11 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
+        **kwargs,
     ) -> np.ndarray:
         """
         Returns the Gaussian light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
diff --git a/autogalaxy/profiles/light/standard/moffat.py b/autogalaxy/profiles/light/standard/moffat.py
index cf5f9d315..87a9412c9 100644
--- a/autogalaxy/profiles/light/standard/moffat.py
+++ b/autogalaxy/profiles/light/standard/moffat.py
@@ -62,7 +62,9 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
             xp.power(
                 1
                 + xp.square(
-                    xp.divide(grid_radii.array, self.alpha / xp.sqrt(self.axis_ratio(xp)))
+                    xp.divide(
+                        grid_radii.array, self.alpha / xp.sqrt(self.axis_ratio(xp))
+                    )
                 ),
                 -self.beta,
             ),
@@ -73,7 +75,11 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
+        **kwargs,
     ) -> np.ndarray:
         """
         Returns the Moffat light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
diff --git a/autogalaxy/profiles/light/standard/sersic.py b/autogalaxy/profiles/light/standard/sersic.py
index 52e35f208..acc5d3465 100644
--- a/autogalaxy/profiles/light/standard/sersic.py
+++ b/autogalaxy/profiles/light/standard/sersic.py
@@ -119,7 +119,9 @@ def __init__(
             sersic_index=sersic_index,
         )
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np, **kwargs) -> np.ndarray:
+    def image_2d_via_radii_from(
+        self, grid_radii: np.ndarray, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Returns the 2D image of the Sersic light profile from a grid of coordinates which are the radial distances of
         each coordinate from the its `centre`.
@@ -151,7 +153,11 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np, **kwargs) -> np
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
+        **kwargs,
     ) -> aa.Array2D:
         """
         Returns the Sersic light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
diff --git a/autogalaxy/profiles/light/standard/sersic_core.py b/autogalaxy/profiles/light/standard/sersic_core.py
index 8955e5699..9aaa81c9b 100644
--- a/autogalaxy/profiles/light/standard/sersic_core.py
+++ b/autogalaxy/profiles/light/standard/sersic_core.py
@@ -72,7 +72,9 @@ def intensity_prime(self, xp=np) -> float:
             )
         )
 
-    def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np, **kwargs) -> np.ndarray:
+    def image_2d_via_radii_from(
+        self, grid_radii: np.ndarray, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Returns the 2D image of the Sersic light profile from a grid of coordinates which are the radial distances of
         each coordinate from the its `centre`.
diff --git a/autogalaxy/profiles/light/standard/shapelets/cartesian.py b/autogalaxy/profiles/light/standard/shapelets/cartesian.py
index 2a2871da7..9be1b123c 100644
--- a/autogalaxy/profiles/light/standard/shapelets/cartesian.py
+++ b/autogalaxy/profiles/light/standard/shapelets/cartesian.py
@@ -63,7 +63,11 @@ def coefficient_tag(self) -> str:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
+        **kwargs,
     ) -> np.ndarray:
         """
         Returns the Cartesian Shapelet light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates.
diff --git a/autogalaxy/profiles/light/standard/shapelets/exponential.py b/autogalaxy/profiles/light/standard/shapelets/exponential.py
index 6e07ba267..0231d7351 100644
--- a/autogalaxy/profiles/light/standard/shapelets/exponential.py
+++ b/autogalaxy/profiles/light/standard/shapelets/exponential.py
@@ -64,7 +64,11 @@ def coefficient_tag(self) -> str:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
+        **kwargs,
     ) -> np.ndarray:
         """
         Returns the Exponential Shapelet light profile's 2D image from a 2D grid of Exponential (y,x) coordinates.
diff --git a/autogalaxy/profiles/light/standard/shapelets/polar.py b/autogalaxy/profiles/light/standard/shapelets/polar.py
index ef8d8b09e..1d6c98db1 100644
--- a/autogalaxy/profiles/light/standard/shapelets/polar.py
+++ b/autogalaxy/profiles/light/standard/shapelets/polar.py
@@ -64,7 +64,11 @@ def coefficient_tag(self) -> str:
     @check_operated_only
     @aa.grid_dec.transform
     def image_2d_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        operated_only: Optional[bool] = None,
+        **kwargs,
     ) -> np.ndarray:
         """
         Returns the Polar Shapelet light profile's 2D image from a 2D grid of Polar (y,x) coordinates.
@@ -85,9 +89,7 @@ def image_2d_from(
         from scipy.special import genlaguerre
         from jax.scipy.special import factorial
 
-        laguerre = genlaguerre(
-            n=(self.n - xp.abs(self.m)) / 2.0, alpha=xp.abs(self.m)
-        )
+        laguerre = genlaguerre(n=(self.n - xp.abs(self.m)) / 2.0, alpha=xp.abs(self.m))
 
         const = (
             ((-1) ** ((self.n - xp.abs(self.m)) // 2))
diff --git a/autogalaxy/profiles/mass/abstract/jax_utils.py b/autogalaxy/profiles/mass/abstract/jax_utils.py
index bbe17d639..0e8448fa8 100644
--- a/autogalaxy/profiles/mass/abstract/jax_utils.py
+++ b/autogalaxy/profiles/mass/abstract/jax_utils.py
@@ -1,4 +1,4 @@
-from jax import custom_jvp
+# from jax import custom_jvp
 
 
 r1_s1 = [2.5, 2, 1.5, 1, 0.5]
@@ -42,56 +42,56 @@ def reg3(z, sqrt_pi, _):
 
     return f1 / f2
 
-
-@custom_jvp
-def w_f_approx(z):
-    """Compute the Faddeeva function :math:`w_{\\mathrm F}(z)` using the
-    approximation given in Zaghloul (2017).
-
-    :param z: complex number
-    :type z: ``complex`` or ``numpy.array(dtype=complex)``
-    :return: :math:`w_\\mathrm{F}(z)`
-    :rtype: ``complex``
-    """
-    sqrt_pi = 1 / xp.sqrt(xp.pi)
-    i_sqrt_pi = 1j * sqrt_pi
-
-    z_imag2 = z.imag**2
-    abs_z2 = z.real**2 + z_imag2
-
-    # use a single partial fraction approx for all large abs(z)**2
-    # to have better approx of the auto-derivatives
-    r1 = (abs_z2 >= 62.0) | ((abs_z2 >= 30.0) & (abs_z2 < 62.0) & (z_imag2 >= 1e-13))
-    # region bounds for 5 taken directly from Zaghloul (2017)
-    # https://dl.acm.org/doi/pdf/10.1145/3119904
-    r2_1 = (abs_z2 >= 30.0) & (abs_z2 < 62.0) & (z_imag2 < 1e-13)
-    r2_2 = (abs_z2 >= 2.5) & (abs_z2 < 30.0) & (z_imag2 < 0.072)
-    r2 = r2_1 | r2_2
-    r3 = xp.logical_not(r1) & xp.logical_not(r2)
-
-    # exploit symmetry to avoid overflow in some regions
-    r_flip = z.imag < 0
-    z_adjust = xp.where(r_flip, -z, z)
-    two_exp_zz = 2 * xp.exp(-(z_adjust**2))
-
-    args = (z_adjust, sqrt_pi, i_sqrt_pi)
-    wz = xp.empty_like(z)
-    wz = xp.where(r1, reg1(*args), wz)
-    wz = xp.where(r2, reg2(*args), wz)
-    wz = xp.where(r3, reg3(*args), wz)
-
-    # exploit symmetry to avoid overflow in some regions
-    wz = xp.where(r_flip, two_exp_zz - wz, wz)
-
-    return wz
-
-
-@w_f_approx.defjvp
-def w_f_approx_jvp(primals, tangents):
-    # define a custom jvp to avoid the issue using `xp.where` with `jax.grad`
-    (z,) = primals
-    (z_dot,) = tangents
-    primal_out = w_f_approx(z)
-    i_sqrt_pi = 1j / xp.sqrt(xp.pi)
-    tangent_out = z_dot * 2 * (i_sqrt_pi - z * primal_out)
-    return primal_out, tangent_out
+#
+# @custom_jvp
+# def w_f_approx(z):
+#     """Compute the Faddeeva function :math:`w_{\\mathrm F}(z)` using the
+#     approximation given in Zaghloul (2017).
+#
+#     :param z: complex number
+#     :type z: ``complex`` or ``numpy.array(dtype=complex)``
+#     :return: :math:`w_\\mathrm{F}(z)`
+#     :rtype: ``complex``
+#     """
+#     sqrt_pi = 1 / xp.sqrt(xp.pi)
+#     i_sqrt_pi = 1j * sqrt_pi
+#
+#     z_imag2 = z.imag**2
+#     abs_z2 = z.real**2 + z_imag2
+#
+#     # use a single partial fraction approx for all large abs(z)**2
+#     # to have better approx of the auto-derivatives
+#     r1 = (abs_z2 >= 62.0) | ((abs_z2 >= 30.0) & (abs_z2 < 62.0) & (z_imag2 >= 1e-13))
+#     # region bounds for 5 taken directly from Zaghloul (2017)
+#     # https://dl.acm.org/doi/pdf/10.1145/3119904
+#     r2_1 = (abs_z2 >= 30.0) & (abs_z2 < 62.0) & (z_imag2 < 1e-13)
+#     r2_2 = (abs_z2 >= 2.5) & (abs_z2 < 30.0) & (z_imag2 < 0.072)
+#     r2 = r2_1 | r2_2
+#     r3 = xp.logical_not(r1) & xp.logical_not(r2)
+#
+#     # exploit symmetry to avoid overflow in some regions
+#     r_flip = z.imag < 0
+#     z_adjust = xp.where(r_flip, -z, z)
+#     two_exp_zz = 2 * xp.exp(-(z_adjust**2))
+#
+#     args = (z_adjust, sqrt_pi, i_sqrt_pi)
+#     wz = xp.empty_like(z)
+#     wz = xp.where(r1, reg1(*args), wz)
+#     wz = xp.where(r2, reg2(*args), wz)
+#     wz = xp.where(r3, reg3(*args), wz)
+#
+#     # exploit symmetry to avoid overflow in some regions
+#     wz = xp.where(r_flip, two_exp_zz - wz, wz)
+#
+#     return wz
+
+
+# @w_f_approx.defjvp
+# def w_f_approx_jvp(primals, tangents):
+#     # define a custom jvp to avoid the issue using `xp.where` with `jax.grad`
+#     (z,) = primals
+#     (z_dot,) = tangents
+#     primal_out = w_f_approx(z)
+#     i_sqrt_pi = 1j / xp.sqrt(xp.pi)
+#     tangent_out = z_dot * 2 * (i_sqrt_pi - z * primal_out)
+#     return primal_out, tangent_out
diff --git a/autogalaxy/profiles/mass/abstract/mge_jax.py b/autogalaxy/profiles/mass/abstract/mge_jax.py
index 20cbd40ce..b6c2d0b83 100644
--- a/autogalaxy/profiles/mass/abstract/mge_jax.py
+++ b/autogalaxy/profiles/mass/abstract/mge_jax.py
@@ -1,5 +1,4 @@
-
-from .jax_utils import w_f_approx
+# from .jax_utils import w_f_approx
 
 
 class MassProfileMGE:
@@ -101,9 +100,7 @@ def _decompose_convergence_via_mge(
 
         # sigma is sampled from logspace between these radii.
 
-        log_sigmas = xp.linspace(
-            xp.log(radii_min), xp.log(radii_max), func_gaussians
-        )
+        log_sigmas = xp.linspace(xp.log(radii_min), xp.log(radii_max), func_gaussians)
         d_log_sigma = log_sigmas[1] - log_sigmas[0]
         sigma_list = xp.exp(log_sigmas)
 
diff --git a/autogalaxy/profiles/mass/dark/abstract.py b/autogalaxy/profiles/mass/dark/abstract.py
index eea6f6683..59f7c2f2f 100644
--- a/autogalaxy/profiles/mass/dark/abstract.py
+++ b/autogalaxy/profiles/mass/dark/abstract.py
@@ -202,7 +202,9 @@ def coord_func_g(self, grid_radius: np.ndarray, xp=np) -> np.ndarray:
         )
 
     def coord_func_h(self, grid_radius, xp=np):
-        return xp.log(grid_radius / 2.0) + self.coord_func_f(grid_radius=grid_radius, xp=xp)
+        return xp.log(grid_radius / 2.0) + self.coord_func_f(
+            grid_radius=grid_radius, xp=xp
+        )
 
     def rho_at_scale_radius_solar_mass_per_kpc3(
         self, redshift_object, redshift_source, cosmology: LensingCosmology = None
@@ -289,7 +291,9 @@ def concentration(
         )
 
         return fsolve(
-            func=self.concentration_func, x0=10.0, args=(delta_concentration, xp),
+            func=self.concentration_func,
+            x0=10.0,
+            args=(delta_concentration, xp),
         )[0]
 
     @staticmethod
@@ -312,7 +316,7 @@ def radius_at_200(
         redshift_source,
         redshift_of_cosmic_average_density="profile",
         cosmology: LensingCosmology = None,
-        xp=np
+        xp=np,
     ):
         """
         Returns `r_{200m}` for this halo in **arcseconds**
@@ -326,7 +330,7 @@ def radius_at_200(
             redshift_source=redshift_source,
             redshift_of_cosmic_average_density=redshift_of_cosmic_average_density,
             cosmology=cosmology,
-            xp=xp
+            xp=xp,
         )
 
         return concentration * self.scale_radius
@@ -367,7 +371,7 @@ def mass_at_200_solar_masses(
             redshift_source=redshift_source,
             redshift_of_cosmic_average_density=redshift_of_cosmic_average_density,
             cosmology=cosmology,
-            xp=xp
+            xp=xp,
         )
 
         kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object)
diff --git a/autogalaxy/profiles/mass/dark/gnfw.py b/autogalaxy/profiles/mass/dark/gnfw.py
index 55d3e5fa5..a1a80295b 100644
--- a/autogalaxy/profiles/mass/dark/gnfw.py
+++ b/autogalaxy/profiles/mass/dark/gnfw.py
@@ -162,7 +162,9 @@ def integral_y(y, eta):
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, tabulate_bins=1000, **kwargs):
+    def potential_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np, tabulate_bins=1000, **kwargs
+    ):
         """
         Calculate the potential at a given set of arc-second gridded coordinates.
 
@@ -295,7 +297,9 @@ def __init__(
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+    def deflections_2d_via_integral_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -316,7 +320,9 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
                 4.0 * self.kappa_s * self.scale_radius, self.deflection_func_sph(eta[i])
             )
 
-        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid, xp=xp)
+        return self._cartesian_grid_via_radial_from(
+            grid=grid, radius=deflection_grid, xp=xp
+        )
 
     @staticmethod
     def deflection_integrand(y, eta, inner_slope):
diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py
index ff3a24f28..dc425f69b 100644
--- a/autogalaxy/profiles/mass/dark/nfw.py
+++ b/autogalaxy/profiles/mass/dark/nfw.py
@@ -47,7 +47,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+    def deflections_2d_via_integral_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -142,7 +144,9 @@ def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs)
     def convergence_func(self, grid_radius: float) -> float:
         grid_radius = (1.0 / self.scale_radius) * grid_radius.array + 0j
         return np.real(
-            2.0 * self.kappa_s * np.array(self.coord_func_g(grid_radius=grid_radius, xp=xp))
+            2.0
+            * self.kappa_s
+            * np.array(self.coord_func_g(grid_radius=grid_radius, xp=xp))
         )
 
     @aa.over_sample
@@ -364,7 +368,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+    def deflections_2d_via_analytic_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -383,7 +389,9 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
             self.deflection_func_sph(grid_radius=eta, xp=xp),
         )
 
-        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid, xp=xp)
+        return self._cartesian_grid_via_radial_from(
+            grid=grid, radius=deflection_grid, xp=xp
+        )
 
     def deflection_func_sph(self, grid_radius, xp=np):
         grid_radius = grid_radius + 0j
diff --git a/autogalaxy/profiles/mass/dark/nfw_truncated.py b/autogalaxy/profiles/mass/dark/nfw_truncated.py
index dae77854d..c6e4fab98 100644
--- a/autogalaxy/profiles/mass/dark/nfw_truncated.py
+++ b/autogalaxy/profiles/mass/dark/nfw_truncated.py
@@ -47,7 +47,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
             self.deflection_func_sph(grid_radius=eta),
         )
 
-        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid, xp=xp)
+        return self._cartesian_grid_via_radial_from(
+            grid=grid, radius=deflection_grid, xp=xp
+        )
 
     def deflection_func_sph(self, grid_radius, xp=np):
         grid_radius = grid_radius + 0j
@@ -120,7 +122,7 @@ def mass_at_truncation_radius_solar_mass(
             redshift_source=redshift_source,
             redshift_of_cosmic_average_density=redshift_of_cosmic_average_density,
             cosmology=cosmology,
-            xp=xp
+            xp=xp,
         )
 
         return (
diff --git a/autogalaxy/profiles/mass/sheets/external_shear.py b/autogalaxy/profiles/mass/sheets/external_shear.py
index 213e990fe..f0c762958 100644
--- a/autogalaxy/profiles/mass/sheets/external_shear.py
+++ b/autogalaxy/profiles/mass/sheets/external_shear.py
@@ -25,10 +25,14 @@ def __init__(self, gamma_1: float = 0.0, gamma_2: float = 0.0):
         self.gamma_2 = gamma_2
 
     def magnitude(self, xp=np):
-        return convert.shear_magnitude_from(gamma_1=self.gamma_1, gamma_2=self.gamma_2, xp=xp)
+        return convert.shear_magnitude_from(
+            gamma_1=self.gamma_1, gamma_2=self.gamma_2, xp=xp
+        )
 
     def angle(self, xp=np):
-        return convert.shear_angle_from(gamma_1=self.gamma_1, gamma_2=self.gamma_2, xp=xp)
+        return convert.shear_angle_from(
+            gamma_1=self.gamma_1, gamma_2=self.gamma_2, xp=xp
+        )
 
     def convergence_func(self, grid_radius: float) -> float:
         return 0.0
diff --git a/autogalaxy/profiles/mass/stellar/chameleon.py b/autogalaxy/profiles/mass/stellar/chameleon.py
index ab0001992..f2850e341 100644
--- a/autogalaxy/profiles/mass/stellar/chameleon.py
+++ b/autogalaxy/profiles/mass/stellar/chameleon.py
@@ -52,7 +52,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+    def deflections_2d_via_analytic_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
         Following Eq. (15) and (16), but the parameters are slightly different.
@@ -70,7 +72,7 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
             * self.intensity
             / (1 + self.axis_ratio(xp))
             * self.axis_ratio(xp)
-            / xp.sqrt(1.0 - self.axis_ratio(xp)**2.0)
+            / xp.sqrt(1.0 - self.axis_ratio(xp) ** 2.0)
         )
 
         core_radius_0 = xp.sqrt(
@@ -89,28 +91,36 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
 
         deflection_y0 = xp.arctanh(
             xp.divide(
-                xp.multiply(xp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 0]),
-                xp.add(psi0, self.axis_ratio(xp)**2.0 * core_radius_0),
+                xp.multiply(
+                    xp.sqrt(1.0 - self.axis_ratio(xp) ** 2.0), grid.array[:, 0]
+                ),
+                xp.add(psi0, self.axis_ratio(xp) ** 2.0 * core_radius_0),
             )
         )
 
         deflection_x0 = xp.arctan(
             xp.divide(
-                xp.multiply(xp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 1]),
+                xp.multiply(
+                    xp.sqrt(1.0 - self.axis_ratio(xp) ** 2.0), grid.array[:, 1]
+                ),
                 xp.add(psi0, core_radius_0),
             )
         )
 
         deflection_y1 = xp.arctanh(
             xp.divide(
-                xp.multiply(xp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 0]),
-                xp.add(psi1, self.axis_ratio(xp)**2.0 * core_radius_1),
+                xp.multiply(
+                    xp.sqrt(1.0 - self.axis_ratio(xp) ** 2.0), grid.array[:, 0]
+                ),
+                xp.add(psi1, self.axis_ratio(xp) ** 2.0 * core_radius_1),
             )
         )
 
         deflection_x1 = xp.arctan(
             xp.divide(
-                xp.multiply(xp.sqrt(1.0 - self.axis_ratio(xp)**2.0), grid.array[:, 1]),
+                xp.multiply(
+                    xp.sqrt(1.0 - self.axis_ratio(xp) ** 2.0), grid.array[:, 1]
+                ),
                 xp.add(psi1, core_radius_1),
             )
         )
@@ -119,8 +129,7 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
         deflection_x = xp.subtract(deflection_x0, deflection_x1)
 
         return self.rotated_grid_from_reference_frame_from(
-            xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T),
-            xp=xp
+            xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T), xp=xp
         )
 
     @aa.over_sample
diff --git a/autogalaxy/profiles/mass/stellar/gaussian.py b/autogalaxy/profiles/mass/stellar/gaussian.py
index 1af6a364f..064fbb0fe 100644
--- a/autogalaxy/profiles/mass/stellar/gaussian.py
+++ b/autogalaxy/profiles/mass/stellar/gaussian.py
@@ -1,8 +1,4 @@
 import numpy as np
-from autoconf.jax_wrapper import use_jax
-
-if use_jax:
-    import jax
 
 from typing import Tuple
 
@@ -61,7 +57,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+    def deflections_2d_via_analytic_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -76,7 +74,7 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
             self.mass_to_light_ratio
             * self.intensity
             * self.sigma
-            * xp.sqrt((2 * np.pi) / (1.0 - self.axis_ratio(xp)**2.0))
+            * xp.sqrt((2 * np.pi) / (1.0 - self.axis_ratio(xp) ** 2.0))
             * self.zeta_from(grid=grid, xp=xp)
         )
 
@@ -84,12 +82,14 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, xp=np, **kw
             xp.multiply(
                 1.0, xp.vstack((-1.0 * xp.imag(deflections), xp.real(deflections))).T
             ),
-            xp=xp
+            xp=xp,
         )
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+    def deflections_2d_via_integral_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -130,8 +130,7 @@ def calculate_deflection_component(npow, index):
         deflection_x = calculate_deflection_component(0.0, 1)
 
         return self.rotated_grid_from_reference_frame_from(
-            np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T),
-            xp=xp
+            np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T), xp=xp
         )
 
     @staticmethod
@@ -182,7 +181,9 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np):
             np.exp(
                 -0.5
                 * np.square(
-                    np.divide(grid_radii.array, self.sigma / np.sqrt(self.axis_ratio(xp)))
+                    np.divide(
+                        grid_radii.array, self.sigma / np.sqrt(self.axis_ratio(xp))
+                    )
                 )
             ),
         )
@@ -195,7 +196,7 @@ def zeta_from(self, grid: aa.type.Grid2DLike, xp=np):
 
         from scipy.special import wofz
 
-        q2 = self.axis_ratio(xp)**2.0
+        q2 = self.axis_ratio(xp) ** 2.0
         ind_pos_y = grid.array[:, 0] >= 0
         shape_grid = np.shape(grid)
         output_grid = np.zeros((shape_grid[0]), dtype=np.complex128)
diff --git a/autogalaxy/profiles/mass/stellar/sersic.py b/autogalaxy/profiles/mass/stellar/sersic.py
index b7a448c5a..d1bda8094 100644
--- a/autogalaxy/profiles/mass/stellar/sersic.py
+++ b/autogalaxy/profiles/mass/stellar/sersic.py
@@ -186,7 +186,12 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
     def convergence_2d_via_mge_from(
-        self, grid: aa.type.Grid2DLike, xp=np, func_terms=28, func_gaussians=20, **kwargs
+        self,
+        grid: aa.type.Grid2DLike,
+        xp=np,
+        func_terms=28,
+        func_gaussians=20,
+        **kwargs,
     ):
         """
         Calculate the projected convergence at a given set of arc-second gridded coordinates.
@@ -370,7 +375,9 @@ def elliptical_effective_radius(self):
 class Sersic(AbstractSersic, MassProfileMGE, MassProfileCSE):
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+    def deflections_2d_via_integral_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
@@ -413,8 +420,7 @@ def calculate_deflection_component(npow, index):
         deflection_x = calculate_deflection_component(0.0, 1)
 
         return self.rotated_grid_from_reference_frame_from(
-            np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T),
-            xp=xp
+            np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T), xp=xp
         )
 
     @staticmethod
diff --git a/autogalaxy/profiles/mass/stellar/sersic_gradient.py b/autogalaxy/profiles/mass/stellar/sersic_gradient.py
index eaf2284c8..3b096d9a1 100644
--- a/autogalaxy/profiles/mass/stellar/sersic_gradient.py
+++ b/autogalaxy/profiles/mass/stellar/sersic_gradient.py
@@ -51,7 +51,9 @@ def __init__(
 
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
-    def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+    def deflections_2d_via_integral_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ):
         """
         Calculate the deflection angles at a given set of arc-second gridded coordinates.
 
diff --git a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py
index bdcfb9e47..3e9fefedb 100644
--- a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py
+++ b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py
@@ -96,8 +96,7 @@ def _ci05f(x, y, eps, rcore, rcut, xp=np):
     )  # a + bi
     zden_rc = xp.complex128(x + 1j * (2.0 * rcore * sqe - y))  # c + di
     znum_rcut = xp.complex128(
-        axis_ratio * x
-        + 1j * (2.0 * sqe * xp.sqrt(rcut * rcut + rem2) - y / axis_ratio)
+        axis_ratio * x + 1j * (2.0 * sqe * xp.sqrt(rcut * rcut + rem2) - y / axis_ratio)
     )  # a + ei
     zden_rcut = xp.complex128(x + 1j * (2.0 * rcut * sqe - y))  # c + fi
 
@@ -310,7 +309,9 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs
 
         return hessian_yy, hessian_xy, hessian_yx, hessian_xx
 
-    def analytical_magnification_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs):
+    def analytical_magnification_2d_from(
+        self, grid: "aa.type.Grid2DLike", xp=np, **kwargs
+    ):
 
         hessian_yy, hessian_xy, hessian_yx, hessian_xx = (
             self.analytical_hessian_2d_from(grid=grid, xp=np)
@@ -503,7 +504,9 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs
 
         return hessian_yy, hessian_xy, hessian_yx, hessian_xx
 
-    def analytical_magnification_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs):
+    def analytical_magnification_2d_from(
+        self, grid: "aa.type.Grid2DLike", xp=np, **kwargs
+    ):
 
         hessian_yy, hessian_xy, hessian_yx, hessian_xx = (
             self.analytical_hessian_2d_from(grid=grid, xp=xp)
@@ -603,7 +606,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
         # And here we convert back to the real axes
         return self.rotated_grid_from_reference_frame_from(
-            grid=xp.multiply(1.0, xp.vstack((deflection_y, deflection_x)).T), xp=xp, **kwargs
+            grid=xp.multiply(1.0, xp.vstack((deflection_y, deflection_x)).T),
+            xp=xp,
+            **kwargs,
         )
 
     @aa.grid_dec.transform
diff --git a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py
index ae5d0fa70..ccc9662fe 100644
--- a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py
+++ b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py
@@ -113,16 +113,14 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         alpha_circ = self._deflection_angle(grid_radii, xp)
 
         # This is in axes aligned to the major/minor axis
-        deflection_y = (
-            alpha_circ * xp.sqrt(1 + ellip) * (grid.array[:, 0] / grid_radii)
-        )
-        deflection_x = (
-            alpha_circ * xp.sqrt(1 - ellip) * (grid.array[:, 1] / grid_radii)
-        )
+        deflection_y = alpha_circ * xp.sqrt(1 + ellip) * (grid.array[:, 0] / grid_radii)
+        deflection_x = alpha_circ * xp.sqrt(1 - ellip) * (grid.array[:, 1] / grid_radii)
 
         # And here we convert back to the real axes
         return self.rotated_grid_from_reference_frame_from(
-            grid=xp.multiply(1.0, xp.vstack((deflection_y, deflection_x)).T), xp=xp, **kwargs
+            grid=xp.multiply(1.0, xp.vstack((deflection_y, deflection_x)).T),
+            xp=xp,
+            **kwargs,
         )
 
     @aa.grid_dec.to_vector_yx
diff --git a/autogalaxy/profiles/mass/total/isothermal.py b/autogalaxy/profiles/mass/total/isothermal.py
index 34c8028b0..2d941affc 100644
--- a/autogalaxy/profiles/mass/total/isothermal.py
+++ b/autogalaxy/profiles/mass/total/isothermal.py
@@ -89,19 +89,23 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
             2.0
             * self.einstein_radius_rescaled(xp)
             * self.axis_ratio(xp)
-            / xp.sqrt(1 - self.axis_ratio(xp)**2)
+            / xp.sqrt(1 - self.axis_ratio(xp) ** 2)
         )
 
-        psi = psi_from(grid=grid, axis_ratio=self.axis_ratio(xp), core_radius=0.0, xp=xp)
+        psi = psi_from(
+            grid=grid, axis_ratio=self.axis_ratio(xp), core_radius=0.0, xp=xp
+        )
 
         deflection_y = xp.arctanh(
             xp.divide(
-                xp.multiply(xp.sqrt(1 - self.axis_ratio(xp)**2), grid.array[:, 0]), psi
+                xp.multiply(xp.sqrt(1 - self.axis_ratio(xp) ** 2), grid.array[:, 0]),
+                psi,
             )
         )
         deflection_x = xp.arctan(
             xp.divide(
-                xp.multiply(xp.sqrt(1 - self.axis_ratio(xp)**2), grid.array[:, 1]), psi
+                xp.multiply(xp.sqrt(1 - self.axis_ratio(xp) ** 2), grid.array[:, 1]),
+                psi,
             )
         )
         return self.rotated_grid_from_reference_frame_from(
diff --git a/autogalaxy/profiles/mass/total/jax_utils.py b/autogalaxy/profiles/mass/total/jax_utils.py
index cb9eeb71a..d04691ea2 100644
--- a/autogalaxy/profiles/mass/total/jax_utils.py
+++ b/autogalaxy/profiles/mass/total/jax_utils.py
@@ -1,5 +1,6 @@
 import numpy as np
 
+
 def body_fun(carry, n, factor, ei2phi, slope):
     omega_nm1, partial_sum = carry
     two_n = 2 * n
@@ -33,6 +34,7 @@ def omega(eiphi, slope, factor, n_terms=20, xp=np):
 
     from jax.tree_util import Partial as partial
     import jax
+
     scan = jax.jit(jax.lax.scan, static_argnames=("length", "reverse", "unroll"))
 
     # use modified scan with a partial'ed function to avoid re-compile
diff --git a/autogalaxy/profiles/mass/total/power_law.py b/autogalaxy/profiles/mass/total/power_law.py
index 9b73a1486..3742bd23f 100644
--- a/autogalaxy/profiles/mass/total/power_law.py
+++ b/autogalaxy/profiles/mass/total/power_law.py
@@ -1,4 +1,3 @@
-from .jax_utils import omega
 import numpy as np
 from typing import Tuple
 
@@ -67,10 +66,11 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         grid
             The grid of (y,x) arc-second coordinates the deflection angles are computed on.
         """
+        from .jax_utils import omega
 
         slope = self.slope - 1.0
         einstein_radius = (
-            2.0 / (self.axis_ratio(xp)**-0.5 + self.axis_ratio(xp)**0.5)
+            2.0 / (self.axis_ratio(xp) ** -0.5 + self.axis_ratio(xp) ** 0.5)
         ) * self.einstein_radius
 
         factor = xp.divide(1.0 - self.axis_ratio(xp), 1.0 + self.axis_ratio(xp))
@@ -83,7 +83,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         )
 
         R = xp.sqrt(
-            (self.axis_ratio(xp) * grid.array[:, 1]) ** 2 + grid.array[:, 0] ** 2 + 1e-16
+            (self.axis_ratio(xp) * grid.array[:, 1]) ** 2
+            + grid.array[:, 0] ** 2
+            + 1e-16
         )
         zh = omega(z, slope, factor, n_terms=20, xp=np)
 
@@ -100,12 +102,13 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         deflection_x *= rescale_factor
 
         return self.rotated_grid_from_reference_frame_from(
-            grid=xp.vstack((deflection_y, deflection_x)).T,
-            xp=xp
+            grid=xp.vstack((deflection_y, deflection_x)).T, xp=xp
         )
 
     def convergence_func(self, grid_radius: float, xp=np) -> float:
-        return self.einstein_radius_rescaled(xp) * grid_radius.array ** (-(self.slope - 1))
+        return self.einstein_radius_rescaled(xp) * grid_radius.array ** (
+            -(self.slope - 1)
+        )
 
     @staticmethod
     def potential_func(u, y, x, axis_ratio, slope, core_radius, xp=np):
@@ -158,4 +161,6 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
             )
         )
 
-        return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_r, xp=xp)
+        return self._cartesian_grid_via_radial_from(
+            grid=grid, radius=deflection_r, xp=xp
+        )
diff --git a/autogalaxy/profiles/mass/total/power_law_broken.py b/autogalaxy/profiles/mass/total/power_law_broken.py
index 9c6613441..577049787 100644
--- a/autogalaxy/profiles/mass/total/power_law_broken.py
+++ b/autogalaxy/profiles/mass/total/power_law_broken.py
@@ -102,13 +102,23 @@ def deflections_yx_2d_from(self, grid, xp=np, max_terms=20, **kwargs):
             self.inner_slope, self.axis_ratio(xp), R, z, max_terms=max_terms, xp=xp
         )
         F2 = self.hyp2f1_series(
-            self.inner_slope, self.axis_ratio(xp), self.break_radius, z, max_terms=max_terms, xp=xp
+            self.inner_slope,
+            self.axis_ratio(xp),
+            self.break_radius,
+            z,
+            max_terms=max_terms,
+            xp=xp,
         )
         F3 = self.hyp2f1_series(
             self.outer_slope, self.axis_ratio(xp), R, z, max_terms=max_terms, xp=xp
         )
         F4 = self.hyp2f1_series(
-            self.outer_slope, self.axis_ratio(xp), self.break_radius, z, max_terms=max_terms, xp=xp
+            self.outer_slope,
+            self.axis_ratio(xp),
+            self.break_radius,
+            z,
+            max_terms=max_terms,
+            xp=xp,
         )
 
         # theta < break radius (eq. 18)
@@ -129,7 +139,7 @@ def deflections_yx_2d_from(self, grid, xp=np, max_terms=20, **kwargs):
             grid=xp.multiply(
                 1.0, xp.vstack((xp.imag(deflections), xp.real(deflections))).T
             ),
-            xp=xp
+            xp=xp,
         )
 
     @staticmethod
diff --git a/autogalaxy/profiles/mass/total/power_law_core.py b/autogalaxy/profiles/mass/total/power_law_core.py
index 63ff8c5b9..5ba9ce9b6 100644
--- a/autogalaxy/profiles/mass/total/power_law_core.py
+++ b/autogalaxy/profiles/mass/total/power_law_core.py
@@ -42,9 +42,9 @@ def einstein_radius_rescaled(self, xp=np):
         Rescale the einstein radius by slope and axis_ratio, to reduce its degeneracy with other mass-profiles
         parameters.
         """
-        return ((3 - self.slope) / (1 + self.axis_ratio(xp))) * self.einstein_radius ** (
-            self.slope - 1
-        )
+        return (
+            (3 - self.slope) / (1 + self.axis_ratio(xp))
+        ) * self.einstein_radius ** (self.slope - 1)
 
     @aa.over_sample
     @aa.grid_dec.to_array
@@ -140,8 +140,7 @@ def calculate_deflection_component(npow, index):
         deflection_x = calculate_deflection_component(0.0, 1)
 
         return self.rotated_grid_from_reference_frame_from(
-            grid=np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T),
-            xp=xp
+            grid=np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T), xp=xp
         )
 
     def convergence_func(self, grid_radius: float, xp=np) -> float:
diff --git a/autogalaxy/profiles/mass/total/power_law_multipole.py b/autogalaxy/profiles/mass/total/power_law_multipole.py
index b182754b4..f20f02102 100644
--- a/autogalaxy/profiles/mass/total/power_law_multipole.py
+++ b/autogalaxy/profiles/mass/total/power_law_multipole.py
@@ -188,14 +188,18 @@ def deflections_yx_2d_from(
         )
 
         return xp.stack(
-            self.jacobian(a_r=a_r, a_angle=a_angle, polar_angle_grid=polar_angle_grid, xp=xp),
+            self.jacobian(
+                a_r=a_r, a_angle=a_angle, polar_angle_grid=polar_angle_grid, xp=xp
+            ),
             axis=-1,
         )
 
     @aa.over_sample
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
-    def convergence_2d_from(self, grid: aa.type.Grid1D2DLike, xp=np, **kwargs) -> np.ndarray:
+    def convergence_2d_from(
+        self, grid: aa.type.Grid1D2DLike, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates.
 
@@ -215,7 +219,9 @@ def convergence_2d_from(self, grid: aa.type.Grid1D2DLike, xp=np, **kwargs) -> np
         )
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray:
+    def potential_2d_from(
+        self, grid: aa.type.Grid2DLike, xp=np, **kwargs
+    ) -> np.ndarray:
         """
         Calculate the potential on a grid of (y,x) arc-second coordinates.
 
diff --git a/autogalaxy/profiles/plot/light_profile_plotters.py b/autogalaxy/profiles/plot/light_profile_plotters.py
index cb5116e09..7292f5eb0 100644
--- a/autogalaxy/profiles/plot/light_profile_plotters.py
+++ b/autogalaxy/profiles/plot/light_profile_plotters.py
@@ -72,8 +72,7 @@ def __init__(
     @property
     def grid_2d_projected(self):
         return self.grid.grid_2d_radial_projected_from(
-            centre=self.light_profile.centre,
-            angle=self.light_profile.angle()
+            centre=self.light_profile.centre, angle=self.light_profile.angle()
         )
 
     def figures_2d(self, image: bool = False):
diff --git a/autogalaxy/profiles/plot/mass_profile_plotters.py b/autogalaxy/profiles/plot/mass_profile_plotters.py
index 6f5675fbc..62eaf8b88 100644
--- a/autogalaxy/profiles/plot/mass_profile_plotters.py
+++ b/autogalaxy/profiles/plot/mass_profile_plotters.py
@@ -74,7 +74,5 @@ def __init__(
     @property
     def grid_2d_projected(self):
         return self.grid.grid_2d_radial_projected_from(
-            centre=self.mass_profile.centre,
-            angle=self.mass_profile.angle()
+            centre=self.mass_profile.centre, angle=self.mass_profile.angle()
         )
-
diff --git a/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py b/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py
index e89d8291f..af14dc927 100644
--- a/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py
+++ b/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py
@@ -72,4 +72,4 @@ def test__galaxies_sub_plot_output(galaxies_x2_7x7, grid_2d_7x7, plot_path, plot
     assert path.join(plot_path, "subplot_galaxies.png") in plot_patch.paths
 
     plotter.subplot_galaxy_images()
-    assert path.join(plot_path, "subplot_galaxy_images.png") in plot_patch.paths
\ No newline at end of file
+    assert path.join(plot_path, "subplot_galaxy_images.png") in plot_patch.paths
diff --git a/test_autogalaxy/imaging/test_simulator.py b/test_autogalaxy/imaging/test_simulator.py
index 8f7f2b2ed..64a64707c 100644
--- a/test_autogalaxy/imaging/test_simulator.py
+++ b/test_autogalaxy/imaging/test_simulator.py
@@ -105,7 +105,9 @@ def test__simulator__via_galaxies_from():
 
     assert dataset.shape_native == (20, 20)
     assert dataset.data.native[0, 0] != imaging_via_image.data.native[0, 0]
-    assert dataset.data.native[10, 10] == pytest.approx(imaging_via_image.data.native[10, 10], 1.0e-4)
+    assert dataset.data.native[10, 10] == pytest.approx(
+        imaging_via_image.data.native[10, 10], 1.0e-4
+    )
     assert dataset.psf == pytest.approx(imaging_via_image.psf, 1.0e-4)
     assert dataset.noise_map == pytest.approx(imaging_via_image.noise_map, 1.0e-4)
 
diff --git a/test_autogalaxy/operate/test_deflections.py b/test_autogalaxy/operate/test_deflections.py
index bd90f734f..45cc97509 100644
--- a/test_autogalaxy/operate/test_deflections.py
+++ b/test_autogalaxy/operate/test_deflections.py
@@ -324,6 +324,7 @@ def test__tangential_caustic_list_from():
     assert 0.47 < y_centre < 0.53
     assert 0.97 < x_centre < 1.03
 
+
 # TODO : Reinstate one JAX defleciton sin.
 
 # def test__tangential_caustic_list_from___compare_via_magnification():
diff --git a/test_autogalaxy/profiles/plot/test_light_profile_plotters.py b/test_autogalaxy/profiles/plot/test_light_profile_plotters.py
index 477f13d71..851860afa 100644
--- a/test_autogalaxy/profiles/plot/test_light_profile_plotters.py
+++ b/test_autogalaxy/profiles/plot/test_light_profile_plotters.py
@@ -14,6 +14,7 @@ def make_profile_plotter_setup():
         "{}".format(path.dirname(path.realpath(__file__))), "files", "plots", "profiles"
     )
 
+
 def test__figures_2d__all_are_output(
     lp_0,
     grid_2d_7x7,
diff --git a/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py b/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py
index 9939386e1..304c050c8 100644
--- a/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py
+++ b/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py
@@ -13,6 +13,7 @@ def make_mp_plotter_setup():
         "{}".format(path.dirname(path.realpath(__file__))), "files", "plots", "profiles"
     )
 
+
 def test__figures_2d__all_are_output(
     mp_0,
     grid_2d_7x7,

From 71ec4455e45dafb89dac1e385d59a354878e4336 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sun, 9 Nov 2025 11:34:04 +0000
Subject: [PATCH 18/18] xp refactor complete

---
 autogalaxy/__init__.py                        |   3 +
 autogalaxy/analysis/model_util.py             | 129 ++++++++++++++++++
 autogalaxy/convert.py                         |   1 +
 autogalaxy/galaxy/to_inversion.py             |  10 +-
 autogalaxy/imaging/fit_imaging.py             |   6 +-
 autogalaxy/imaging/model/analysis.py          |   2 +
 .../interferometer/fit_interferometer.py      |   4 +-
 autogalaxy/interferometer/model/analysis.py   |   2 +
 autogalaxy/operate/deflections.py             |  10 +-
 autogalaxy/profiles/light/linear/abstract.py  |  28 ++--
 .../profiles/mass/abstract/jax_utils.py       |   1 +
 test_autogalaxy/config/general.yaml           |   2 -
 test_autogalaxy/conftest.py                   |   1 +
 13 files changed, 168 insertions(+), 31 deletions(-)
 create mode 100644 autogalaxy/analysis/model_util.py

diff --git a/autogalaxy/__init__.py b/autogalaxy/__init__.py
index 506a0eaec..6703d30f4 100644
--- a/autogalaxy/__init__.py
+++ b/autogalaxy/__init__.py
@@ -25,6 +25,8 @@
     mapper_from as Mapper,
 )  # noqa
 from autoarray.inversion.pixelization.border_relocator import BorderRelocator
+from autoarray.preloads import Preloads
+from autoarray.preloads import mapper_indices_from
 from autoarray.mask.mask_1d import Mask1D  # noqa
 from autoarray.mask.mask_2d import Mask2D  # noqa
 from autoarray.mask.derive.zoom_2d import Zoom2D
@@ -51,6 +53,7 @@
 from autoarray.structures.visibilities import Visibilities  # noqa
 from autoarray.structures.visibilities import VisibilitiesNoiseMap  # noqa
 
+from .analysis import model_util
 from .analysis.adapt_images.adapt_images import AdaptImages
 from .analysis.adapt_images.adapt_image_maker import AdaptImageMaker
 from . import aggregator as agg
diff --git a/autogalaxy/analysis/model_util.py b/autogalaxy/analysis/model_util.py
new file mode 100644
index 000000000..9618a22ce
--- /dev/null
+++ b/autogalaxy/analysis/model_util.py
@@ -0,0 +1,129 @@
+import numpy as np
+from typing import Optional, Tuple
+
+import autofit as af
+
+
+def mge_model_from(
+    mask_radius: float,
+    total_gaussians: int = 30,
+    gaussian_per_basis: int = 1,
+    centre_prior_is_uniform: bool = True,
+    centre: Tuple[float, float] = (0.0, 0.0),
+    centre_fixed: Optional[Tuple[float, float]] = None,
+    use_spherical: bool = False,
+) -> af.Collection:
+    """
+    Construct a Multi-Gaussian Expansion (MGE) for the lens or source galaxy light
+
+    This model is designed as a "start here" configuration for lens modeling:
+
+    - The lens and source light are represented by a Basis object composed of many
+      Gaussian light profiles with fixed logarithmically spaced widths (`sigma`).
+    - All Gaussians within each basis share common centres and ellipticity
+      components, reducing degeneracy while retaining flexibility.
+
+    - Users can combine with a lens mass model of their choiuce.
+
+    The resulting model provides a good balance of speed, flexibility, and accuracy
+    for fitting most galaxy-scale strong lenses.
+
+    This code is mostly to make the API simple for new users, hiding the technical
+    details of setting up an MGE. More advanced users may wish to customize the
+    model further.
+
+    Parameters
+    ----------
+    mask_radius
+        The outer radius (in arcseconds) of the circular mask applied to the data.
+        This determines the maximum Gaussian width (`sigma`) used in the lens MGE.
+    lens_total_gaussians
+        Total number of Gaussian light profiles used in the lens MGE basis.
+    source_total_gaussians
+        Total number of Gaussian light profiles used in the source MGE basis.
+    lens_gaussian_per_basis
+        Number of separate Gaussian bases to include for the lens light profile.
+        Each basis has `lens_total_gaussians` components.
+    source_gaussian_per_basis
+        Number of separate Gaussian bases to include for the source light profile.
+        Each basis has `source_total_gaussians` components.
+
+    Returns
+    -------
+    model : af.Collection
+        An `autofit.Collection` containing:
+        - A lens galaxy at redshift 0.5, with:
+          * bulge light profile: MGE basis of Gaussians
+          * mass profile: Isothermal ellipsoid
+          * external shear
+        - A source galaxy at redshift 1.0, with:
+          * bulge light profile: MGE basis of Gaussians
+
+    Notes
+    -----
+    - Lens light Gaussians have widths (sigma) logarithmically spaced between 0.01"
+      and the mask radius.
+    - Source light Gaussians have widths logarithmically spaced between 0.01" and 1.0".
+    - Gaussian centres are free parameters but tied across all components in each
+      basis to reduce dimensionality.
+    - This function is a convenience utility: it hides the technical setup of MGE
+      composition and provides a ready-to-use lens model for quick experimentation.
+    """
+
+    from autogalaxy.profiles.light.linear import Gaussian, GaussianSph
+    from autogalaxy.profiles.basis import Basis
+
+    # The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0".
+    log10_sigma_list = np.linspace(-4, np.log10(mask_radius), total_gaussians)
+
+    # By defining the centre here, it creates two free parameters that are assigned below to all Gaussians.
+
+    if centre_fixed is not None:
+        centre_0 = centre[0]
+        centre_1 = centre[1]
+    elif centre_prior_is_uniform:
+        centre_0 = af.UniformPrior(
+            lower_limit=centre[0] - 0.1, upper_limit=centre[0] + 0.1
+        )
+        centre_1 = af.UniformPrior(
+            lower_limit=centre[1] - 0.1, upper_limit=centre[1] + 0.1
+        )
+    else:
+        centre_0 = af.GaussianPrior(mean=centre[0], sigma=0.3)
+        centre_1 = af.GaussianPrior(mean=centre[1], sigma=0.3)
+
+    if use_spherical:
+        model_cls = GaussianSph
+    else:
+        model_cls = Gaussian
+
+    bulge_gaussian_list = []
+
+    for j in range(gaussian_per_basis):
+        # A list of Gaussian model components whose parameters are customized belows.
+
+        gaussian_list = af.Collection(
+            af.Model(model_cls) for _ in range(total_gaussians)
+        )
+
+        # Iterate over every Gaussian and customize its parameters.
+
+        for i, gaussian in enumerate(gaussian_list):
+            gaussian.centre.centre_0 = centre_0  # All Gaussians have same y centre.
+            gaussian.centre.centre_1 = centre_1  # All Gaussians have same x centre.
+            if not use_spherical:
+                gaussian.ell_comps = gaussian_list[
+                    0
+                ].ell_comps  # All Gaussians have same elliptical components.
+            gaussian.sigma = (
+                10 ** log10_sigma_list[i]
+            )  # All Gaussian sigmas are fixed to values above.
+
+        bulge_gaussian_list += gaussian_list
+
+    # The Basis object groups many light profiles together into a single model component.
+
+    return af.Model(
+        Basis,
+        profile_list=bulge_gaussian_list,
+    )
diff --git a/autogalaxy/convert.py b/autogalaxy/convert.py
index 0696c9129..5bda71120 100644
--- a/autogalaxy/convert.py
+++ b/autogalaxy/convert.py
@@ -68,6 +68,7 @@ def axis_ratio_and_angle_from(
     fac = xp.sqrt(ell_comps[1] ** 2 + ell_comps[0] ** 2)
     if xp.__name__.startswith("jax"):
         import jax
+
         fac = jax.lax.min(fac, 0.999)
     else:  # NumPy
         fac = np.minimum(fac, 0.999)
diff --git a/autogalaxy/galaxy/to_inversion.py b/autogalaxy/galaxy/to_inversion.py
index 32fafe3ab..cf86a2c83 100644
--- a/autogalaxy/galaxy/to_inversion.py
+++ b/autogalaxy/galaxy/to_inversion.py
@@ -76,7 +76,7 @@ def __init__(
         self.adapt_images = adapt_images
 
         self.settings_inversion = settings_inversion
-        self.xp = xp
+        self._xp = xp
 
     @property
     def psf(self) -> Optional[aa.Kernel2D]:
@@ -309,7 +309,7 @@ def cls_light_profile_func_list_galaxy_dict_from(
                             psf=self.dataset.psf,
                             light_profile_list=light_profile_list,
                             regularization=light_profile.regularization,
-                            xp=self.xp,
+                            xp=self._xp,
                         )
 
                         lp_linear_func_galaxy_dict[lp_linear_func] = galaxy
@@ -488,11 +488,11 @@ def mapper_from(
             source_plane_mesh_grid=source_plane_mesh_grid,
             image_plane_mesh_grid=image_plane_mesh_grid,
             adapt_data=adapt_galaxy_image,
-            xp=self.xp,
+            xp=self._xp,
         )
 
         return mapper_from(
-            mapper_grids=mapper_grids, regularization=regularization, xp=self.xp
+            mapper_grids=mapper_grids, regularization=regularization, xp=self._xp
         )
 
     @cached_property
@@ -574,7 +574,7 @@ def inversion(self) -> aa.AbstractInversion:
             dataset=self.dataset,
             linear_obj_list=self.linear_obj_list,
             settings=self.settings_inversion,
-            xp=self.xp,
+            xp=self._xp,
         )
 
         inversion.linear_obj_galaxy_dict = self.linear_obj_galaxy_dict
diff --git a/autogalaxy/imaging/fit_imaging.py b/autogalaxy/imaging/fit_imaging.py
index ddedcb8e0..d08a2d5cd 100644
--- a/autogalaxy/imaging/fit_imaging.py
+++ b/autogalaxy/imaging/fit_imaging.py
@@ -96,14 +96,14 @@ def blurred_image(self) -> aa.Array2D:
         ):
             return self.galaxies.image_2d_from(
                 grid=self.grids.lp,
-                xp=self.xp,
+                xp=self._xp,
             )
 
         return self.galaxies.blurred_image_2d_from(
             grid=self.grids.lp,
             psf=self.dataset.psf,
             blurring_grid=self.grids.blurring,
-            xp=self.xp,
+            xp=self._xp,
         )
 
     @property
@@ -128,7 +128,7 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion:
             galaxies=self.galaxies,
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
-            xp=self.xp,
+            xp=self._xp,
         )
 
     @cached_property
diff --git a/autogalaxy/imaging/model/analysis.py b/autogalaxy/imaging/model/analysis.py
index accf130ca..1620e4a76 100644
--- a/autogalaxy/imaging/model/analysis.py
+++ b/autogalaxy/imaging/model/analysis.py
@@ -22,6 +22,7 @@ def __init__(
         adapt_image_maker: Optional[AdaptImageMaker] = None,
         cosmology: LensingCosmology = None,
         settings_inversion: aa.SettingsInversion = None,
+        preloads: aa.Preloads = None,
         title_prefix: str = None,
     ):
         """
@@ -59,6 +60,7 @@ def __init__(
             adapt_image_maker=adapt_image_maker,
             cosmology=cosmology,
             settings_inversion=settings_inversion,
+            preloads=preloads,
             title_prefix=title_prefix,
         )
 
diff --git a/autogalaxy/interferometer/fit_interferometer.py b/autogalaxy/interferometer/fit_interferometer.py
index 5c80fdd8e..ab4987880 100644
--- a/autogalaxy/interferometer/fit_interferometer.py
+++ b/autogalaxy/interferometer/fit_interferometer.py
@@ -88,7 +88,7 @@ def profile_visibilities(self) -> aa.Visibilities:
         a Fourier transform to the sum of light profile images.
         """
         return self.galaxies.visibilities_from(
-            grid=self.grids.lp, transformer=self.dataset.transformer, xp=self.xp
+            grid=self.grids.lp, transformer=self.dataset.transformer, xp=self._xp
         )
 
     @property
@@ -113,7 +113,7 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion:
             galaxies=self.galaxies,
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
-            xp=self.xp,
+            xp=self._xp,
         )
 
     @cached_property
diff --git a/autogalaxy/interferometer/model/analysis.py b/autogalaxy/interferometer/model/analysis.py
index 0bf3a4627..3961ff99b 100644
--- a/autogalaxy/interferometer/model/analysis.py
+++ b/autogalaxy/interferometer/model/analysis.py
@@ -29,6 +29,7 @@ def __init__(
         adapt_image_maker: Optional[AdaptImageMaker] = None,
         cosmology: LensingCosmology = None,
         settings_inversion: aa.SettingsInversion = None,
+        preloads: aa.Preloads = None,
         title_prefix: str = None,
     ):
         """
@@ -66,6 +67,7 @@ def __init__(
             adapt_image_maker=adapt_image_maker,
             cosmology=cosmology,
             settings_inversion=settings_inversion,
+            preloads=preloads,
             title_prefix=title_prefix,
         )
 
diff --git a/autogalaxy/operate/deflections.py b/autogalaxy/operate/deflections.py
index 8cb691ae2..ceb7bbb41 100644
--- a/autogalaxy/operate/deflections.py
+++ b/autogalaxy/operate/deflections.py
@@ -112,7 +112,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
     def __eq__(self, other):
         return self.__dict__ == other.__dict__ and self.__class__ is other.__class__
 
-    def time_delay_geometry_term_from(self, grid) -> aa.Array2D:
+    def time_delay_geometry_term_from(self, grid, xp=np) -> aa.Array2D:
         """
             Returns the geometric time delay term of the Fermat potential for a given grid of image-plane positions.
 
@@ -136,7 +136,7 @@ def time_delay_geometry_term_from(self, grid) -> aa.Array2D:
             -------
             The geometric time delay term at each grid position.
         """
-        deflections = self.deflections_yx_2d_from(grid=grid)
+        deflections = self.deflections_yx_2d_from(grid=grid, xp=xp)
 
         src_y = grid[:, 0] - deflections[:, 0]
         src_x = grid[:, 1] - deflections[:, 1]
@@ -147,7 +147,7 @@ def time_delay_geometry_term_from(self, grid) -> aa.Array2D:
             return aa.ArrayIrregular(values=delay)
         return aa.Array2D(values=delay, mask=grid.mask)
 
-    def fermat_potential_from(self, grid) -> aa.Array2D:
+    def fermat_potential_from(self, grid, xp=np) -> aa.Array2D:
         """
         Returns the Fermat potential for a given grid of image-plane positions.
 
@@ -172,8 +172,8 @@ def fermat_potential_from(self, grid) -> aa.Array2D:
         -------
         The Fermat potential at each grid position.
         """
-        time_delay_geometry_term = self.time_delay_geometry_term_from(grid=grid)
-        potential = self.potential_2d_from(grid=grid)
+        time_delay_geometry_term = self.time_delay_geometry_term_from(grid=grid, xp=xp)
+        potential = self.potential_2d_from(grid=grid, xp=xp)
 
         fermat_potential = time_delay_geometry_term - potential
 
diff --git a/autogalaxy/profiles/light/linear/abstract.py b/autogalaxy/profiles/light/linear/abstract.py
index c75e1c06c..43a0946de 100644
--- a/autogalaxy/profiles/light/linear/abstract.py
+++ b/autogalaxy/profiles/light/linear/abstract.py
@@ -264,11 +264,11 @@ def mapping_matrix(self) -> np.ndarray:
 
         for pixel, light_profile in enumerate(self.light_profile_list):
 
-            image_2d = light_profile.image_2d_from(grid=self.grid, xp=self.xp).slim
+            image_2d = light_profile.image_2d_from(grid=self.grid, xp=self._xp).slim
 
             image_2d_list.append(image_2d.array)
 
-        return self.xp.stack(image_2d_list, axis=1)
+        return self._xp.stack(image_2d_list, axis=1)
 
     @cached_property
     def operated_mapping_matrix_overrideg(self) -> Optional[np.ndarray]:
@@ -297,31 +297,31 @@ def operated_mapping_matrix_overrideg(self) -> Optional[np.ndarray]:
         n_src = len(self.light_profile_list)
 
         # allocate slim-form arrays for mapping matrices
-        mapping_matrix = self.xp.zeros((self.grid.shape_slim, n_src))
-        blurring_mapping_matrix = self.xp.zeros((self.blurring_grid.shape_slim, n_src))
+        mapping_matrix = self._xp.zeros((self.grid.shape_slim, n_src))
+        blurring_mapping_matrix = self._xp.zeros((self.blurring_grid.shape_slim, n_src))
 
         # build each column
         for pixel, light_profile in enumerate(self.light_profile_list):
-            if self.xp.__name__.startswith("jax"):
+            if self._xp.__name__.startswith("jax"):
                 # main grid mapping for this light profile
                 mapping_matrix = mapping_matrix.at[:, pixel].set(
-                    light_profile.image_2d_from(grid=self.grid, xp=self.xp).array
+                    light_profile.image_2d_from(grid=self.grid, xp=self._xp).array
                 )
 
                 # blurring grid mapping for this light profile
                 blurring_mapping_matrix = blurring_mapping_matrix.at[:, pixel].set(
                     light_profile.image_2d_from(
-                        grid=self.blurring_grid, xp=self.xp
+                        grid=self.blurring_grid, xp=self._xp
                     ).array
                 )
 
             else:
 
                 mapping_matrix[:, pixel] = light_profile.image_2d_from(
-                    grid=self.grid, xp=self.xp
+                    grid=self.grid, xp=self._xp
                 ).array
                 blurring_mapping_matrix[:, pixel] = light_profile.image_2d_from(
-                    grid=self.blurring_grid, xp=self.xp
+                    grid=self.blurring_grid, xp=self._xp
                 ).array
 
         return self.psf.convolved_mapping_matrix_from(
@@ -329,7 +329,7 @@ def operated_mapping_matrix_overrideg(self) -> Optional[np.ndarray]:
             mask=self.grid.mask,
             blurring_mapping_matrix=blurring_mapping_matrix,
             blurring_mask=self.blurring_grid.mask,
-            xp=self.xp,
+            xp=self._xp,
         )
 
     @cached_property
@@ -358,16 +358,16 @@ def operated_mapping_matrix_override(self) -> Optional[np.ndarray]:
         blurred_image_2d_list = []
 
         for pixel, light_profile in enumerate(self.light_profile_list):
-            image_2d = light_profile.image_2d_from(grid=self.grid, xp=self.xp)
+            image_2d = light_profile.image_2d_from(grid=self.grid, xp=self._xp)
 
             blurring_image_2d = light_profile.image_2d_from(
-                grid=self.blurring_grid, xp=self.xp
+                grid=self.blurring_grid, xp=self._xp
             )
 
             blurred_image_2d = self.psf.convolved_image_from(
-                image=image_2d, blurring_image=blurring_image_2d, xp=self.xp
+                image=image_2d, blurring_image=blurring_image_2d, xp=self._xp
             )
 
             blurred_image_2d_list.append(blurred_image_2d.array)
 
-        return self.xp.stack(blurred_image_2d_list, axis=1)
+        return self._xp.stack(blurred_image_2d_list, axis=1)
diff --git a/autogalaxy/profiles/mass/abstract/jax_utils.py b/autogalaxy/profiles/mass/abstract/jax_utils.py
index 0e8448fa8..de809f493 100644
--- a/autogalaxy/profiles/mass/abstract/jax_utils.py
+++ b/autogalaxy/profiles/mass/abstract/jax_utils.py
@@ -42,6 +42,7 @@ def reg3(z, sqrt_pi, _):
 
     return f1 / f2
 
+
 #
 # @custom_jvp
 # def w_f_approx(z):
diff --git a/test_autogalaxy/config/general.yaml b/test_autogalaxy/config/general.yaml
index 0b88a5865..07df488b4 100644
--- a/test_autogalaxy/config/general.yaml
+++ b/test_autogalaxy/config/general.yaml
@@ -22,8 +22,6 @@ hpc:
 adapt:
   adapt_minimum_percent: 0.01
   adapt_noise_limit: 100000000.0
-model:
-  ignore_prior_limits: true
 numba:
   cache: true
   nopython: true
diff --git a/test_autogalaxy/conftest.py b/test_autogalaxy/conftest.py
index f6af3b45c..a017b64e7 100644
--- a/test_autogalaxy/conftest.py
+++ b/test_autogalaxy/conftest.py
@@ -1,5 +1,6 @@
 def pytest_configure():
     import jax.numpy as jnp
+
     _ = jnp.sum(jnp.array([0.0]))  # Force backend init