Repository: Jammy2211/PyAutoGalaxy (language: python)

Issue:
Revert "Main to JAX"
Reverts Jammy2211/PyAutoGalaxy#201


Discussion / hints:

Fix patch:
diff --git a/autogalaxy/__init__.py b/autogalaxy/__init__.py
index 08fb9d3d0..26df5df98 100644
--- a/autogalaxy/__init__.py
+++ b/autogalaxy/__init__.py
@@ -4,7 +4,6 @@
 from autoarray.dataset.interferometer.dataset import Interferometer  # noqa
 from autoarray.dataset.dataset_model import DatasetModel
 from autoarray.dataset.over_sampling import OverSamplingDataset
-from autoarray.inversion.inversion.mapper_valued import MapperValued
 from autoarray.inversion.pixelization import mesh  # noqa
 from autoarray.inversion import regularization as reg  # noqa
 from autoarray.inversion.pixelization import image_mesh
@@ -53,8 +52,6 @@
 
 from .analysis.adapt_images.adapt_images import AdaptImages
 from .analysis.adapt_images.adapt_image_maker import AdaptImageMaker
-from .analysis.maker import FitMaker
-from .analysis.preloads import Preloads
 from . import aggregator as agg
 from . import exc
 from . import plot
@@ -111,6 +108,8 @@
 from .gui.clicker import Clicker
 from .gui.scribbler import Scribbler
 
+from .analysis.clump_model import ClumpModel
+
 from autoconf import conf
 
 conf.instance.register(__file__)
diff --git a/autogalaxy/aggregator/agg_util.py b/autogalaxy/aggregator/agg_util.py
index 58b818d3e..205944d4b 100644
--- a/autogalaxy/aggregator/agg_util.py
+++ b/autogalaxy/aggregator/agg_util.py
@@ -68,102 +68,3 @@ def adapt_images_from(
         adapt_images_list.append(adapt_images)
 
     return adapt_images_list
-
-
-def mesh_grids_of_planes_list_from(
-    fit: af.Fit, total_fits: int, use_preloaded_grid: bool
-) -> List[Optional[aa.Grid2D]]:
-    """
-    Returns the image-plane pixelization grid(s) used by the fit.
-
-    A subset of image-mesh objects (e.g. `Hilbert`, `KMeans`) create the grid of points that act as the mesh
-    centres (e.g. the centers of Voronoi cells) in the image-plane. For lensing calculations this may then be
-    traced to the source-plane to form the pixelization.
-
-    This calculation can depend on the library used to determine the image-plane grid and may have a random element
-    associated with it. This means that performing an analysis on a super computer and then downloading the results
-    for inspection on your laptop may produce different image-plane grids, changing the result and quantities like
-    the `log_likelihood`.
-
-    By storing this grid as a result in the `files` folder and loading it via the database before creating a fit
-    this possible mismatch is removed, ensuring results on a local computer are identical to those computer elsewhere.
-
-    Parameters
-    ----------
-    fit
-        A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
-    total_fits
-        The total number of `Analysis` objects summed to create the fit.
-    use_preloaded_grid
-        Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This grid
-        may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is used
-        as the fit.
-
-    Returns
-    -------
-    The list of image-plane mesh centres used when creating the overall fit.
-    """
-
-    if use_preloaded_grid:
-        if not fit.children:
-            return [fit.value(name="preload_mesh_grids_of_planes")]
-        else:
-            try:
-                return fit.child_values(name="preload_mesh_grids_of_planes")
-            except AttributeError:
-                return [None] * total_fits
-    else:
-        return [None] * total_fits
-
-
-def preloads_from(
-    preloads_cls,
-    use_preloaded_grid: bool,
-    mesh_grids_of_planes: List,
-    use_w_tilde: Optional[bool] = False,
-) -> aa.Preloads:
-    """
-    Returns a `Preloads` object associated with a fit loaded via the database.
-
-    When loading results via the database, the preloads class may have certain attributes associated with it
-    in order to perform the fit. The main purpose is to use the same image-plane mesh centres for pixelization where
-    the mesh is computed in the image-plane (see `agg_util.mesh_grids_of_planes_list_from`).
-
-    The preloads may also switch off `w_tilde` so fits are computed faster locally as they do not need to recompute
-    w_tilde.
-
-    Parameters
-    ----------
-    preloads_cls
-        The `Preloads` object used to create the preloads (this varies across
-        projects like `autogalaxy` and `autolens`).
-    use_preloaded_grid
-        Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This grid
-        may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is used
-        as the fit.
-    mesh_grids_of_planes
-        The list of image-plane mesh centres used when creating the overall fit which are associated with the
-        preloads.
-    use_w_tilde
-        Whether to use the w-tilde formalism when recomputing fits.
-
-    Returns
-    -------
-    The preloads object associated with the fit loaded via the database.
-    """
-    preloads = preloads_cls()
-
-    if use_preloaded_grid:
-        if mesh_grids_of_planes is not None:
-            preloads = preloads_cls(
-                image_plane_mesh_grid_pg_list=mesh_grids_of_planes,
-                use_w_tilde=use_w_tilde,
-            )
-
-            if len(preloads.image_plane_mesh_grid_pg_list) == 2:
-                if type(preloads.image_plane_mesh_grid_pg_list[1]) != list:
-                    preloads.image_plane_mesh_grid_pg_list[1] = [
-                        preloads.image_plane_mesh_grid_pg_list[1]
-                    ]
-
-    return preloads
diff --git a/autogalaxy/aggregator/ellipse/fit_ellipse.py b/autogalaxy/aggregator/ellipse/fit_ellipse.py
index 0fe101295..031a7c24a 100644
--- a/autogalaxy/aggregator/ellipse/fit_ellipse.py
+++ b/autogalaxy/aggregator/ellipse/fit_ellipse.py
@@ -60,6 +60,7 @@ def _fit_ellipse_from(
     for dataset, ellipse_list, multipole_lists in zip(
         dataset_list, ellipse_list_list, multipole_list_list
     ):
+
         for ellipse, multipole_list in zip(ellipse_list, multipole_lists):
             fit_dataset_list.append(
                 FitEllipse(
diff --git a/autogalaxy/aggregator/galaxies.py b/autogalaxy/aggregator/galaxies.py
index ed1c3224e..a6711ac54 100644
--- a/autogalaxy/aggregator/galaxies.py
+++ b/autogalaxy/aggregator/galaxies.py
@@ -41,14 +41,14 @@ def _galaxies_from(fit: af.Fit, instance: af.ModelInstance) -> List[Galaxy]:
     if instance is not None:
         galaxies = instance.galaxies
 
-        if hasattr(instance, "extra_galaxies"):
-            galaxies = galaxies + fit.instance.extra_galaxies
+        if hasattr(instance, "clumps"):
+            galaxies = galaxies + fit.instance.clumps
 
     else:
         galaxies = fit.instance.galaxies
 
-        if hasattr(fit.instance, "extra_galaxies"):
-            galaxies = galaxies + fit.instance.extra_galaxies
+        if hasattr(fit.instance, "clumps"):
+            galaxies = galaxies + fit.instance.clumps
 
     if fit.children is not None:
         if len(fit.children) > 0:
diff --git a/autogalaxy/aggregator/imaging/fit_imaging.py b/autogalaxy/aggregator/imaging/fit_imaging.py
index e805df1ee..8ba62ae73 100644
--- a/autogalaxy/aggregator/imaging/fit_imaging.py
+++ b/autogalaxy/aggregator/imaging/fit_imaging.py
@@ -7,8 +7,6 @@
 import autofit as af
 import autoarray as aa
 
-from autogalaxy.analysis.preloads import Preloads
-
 from autogalaxy.aggregator.imaging.imaging import _imaging_from
 from autogalaxy.aggregator.galaxies import _galaxies_from
 from autogalaxy.aggregator.dataset_model import _dataset_model_from
@@ -19,7 +17,6 @@ def _fit_imaging_from(
     fit: af.Fit,
     instance: Optional[af.ModelInstance] = None,
     settings_inversion: aa.SettingsInversion = None,
-    use_preloaded_grid: bool = True,
 ) -> List[FitImaging]:
     """
     Returns a list of `FitImaging` objects from a `PyAutoFit` sqlite database `Fit` object.
@@ -51,10 +48,6 @@ def _fit_imaging_from(
         randomly from the PDF).
     settings_inversion
         Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
-    use_preloaded_grid
-        Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This grid
-        may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is used
-        as the fit.
     """
 
     from autogalaxy.imaging.fit_imaging import FitImaging
@@ -69,25 +62,14 @@ def _fit_imaging_from(
 
     settings_inversion = settings_inversion or fit.value(name="settings_inversion")
 
-    mesh_grids_of_planes_list = agg_util.mesh_grids_of_planes_list_from(
-        fit=fit, total_fits=len(dataset_list), use_preloaded_grid=use_preloaded_grid
-    )
-
     fit_dataset_list = []
 
-    for dataset, galaxies, dataset_model, adapt_images, mesh_grids_of_planes in zip(
+    for dataset, galaxies, dataset_model, adapt_images in zip(
         dataset_list,
         galaxies_list,
         dataset_model_list,
         adapt_images_list,
-        mesh_grids_of_planes_list,
     ):
-        preloads = agg_util.preloads_from(
-            preloads_cls=Preloads,
-            use_preloaded_grid=use_preloaded_grid,
-            mesh_grids_of_planes=mesh_grids_of_planes,
-            use_w_tilde=False,
-        )
 
         fit_dataset_list.append(
             FitImaging(
@@ -96,7 +78,6 @@ def _fit_imaging_from(
                 dataset_model=dataset_model,
                 adapt_images=adapt_images,
                 settings_inversion=settings_inversion,
-                preloads=preloads,
             )
         )
 
@@ -108,7 +89,6 @@ def __init__(
         self,
         aggregator: af.Aggregator,
         settings_inversion: Optional[aa.SettingsInversion] = None,
-        use_preloaded_grid: bool = True,
     ):
         """
         Interfaces with an `PyAutoFit` aggregator object to create instances of `FitImaging` objects from the results
@@ -142,15 +122,10 @@ def __init__(
             A `PyAutoFit` aggregator object which can load the results of model-fits.
         settings_inversion
             Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
-        use_preloaded_grid
-            Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This
-            grid may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is
-            used as the fit.
         """
         super().__init__(aggregator=aggregator)
 
         self.settings_inversion = settings_inversion
-        self.use_preloaded_grid = use_preloaded_grid
 
     def object_via_gen_from(
         self, fit, instance: Optional[af.ModelInstance] = None
@@ -172,5 +147,4 @@ def object_via_gen_from(
             fit=fit,
             instance=instance,
             settings_inversion=self.settings_inversion,
-            use_preloaded_grid=self.use_preloaded_grid,
         )
diff --git a/autogalaxy/aggregator/interferometer/fit_interferometer.py b/autogalaxy/aggregator/interferometer/fit_interferometer.py
index 316e3aac4..5e395927b 100644
--- a/autogalaxy/aggregator/interferometer/fit_interferometer.py
+++ b/autogalaxy/aggregator/interferometer/fit_interferometer.py
@@ -7,8 +7,6 @@
 import autofit as af
 import autoarray as aa
 
-from autogalaxy.analysis.preloads import Preloads
-
 from autogalaxy.aggregator import agg_util
 from autogalaxy.aggregator.interferometer.interferometer import _interferometer_from
 from autogalaxy.aggregator.dataset_model import _dataset_model_from
@@ -20,7 +18,6 @@ def _fit_interferometer_from(
     instance: Optional[af.ModelInstance] = None,
     real_space_mask: Optional[aa.Mask2D] = None,
     settings_inversion: aa.SettingsInversion = None,
-    use_preloaded_grid: bool = True,
 ) -> List[FitInterferometer]:
     """
     Returns a list of `FitInterferometer` objects from a `PyAutoFit` sqlite database `Fit` object.
@@ -53,10 +50,6 @@ def _fit_interferometer_from(
         randomly from the PDF).
     settings_inversion
         Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
-    use_preloaded_grid
-        Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This grid
-        may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is used
-        as the fit.
     """
     from autogalaxy.interferometer.fit_interferometer import FitInterferometer
 
@@ -73,25 +66,14 @@ def _fit_interferometer_from(
 
     settings_inversion = settings_inversion or fit.value(name="settings_inversion")
 
-    mesh_grids_of_planes_list = agg_util.mesh_grids_of_planes_list_from(
-        fit=fit, total_fits=len(dataset_list), use_preloaded_grid=use_preloaded_grid
-    )
-
     fit_dataset_list = []
 
-    for dataset, galaxies, dataset_model, adapt_images, mesh_grids_of_planes in zip(
+    for dataset, galaxies, dataset_model, adapt_images in zip(
         dataset_list,
         galaxies_list,
         dataset_model_list,
         adapt_images_list,
-        mesh_grids_of_planes_list,
     ):
-        preloads = agg_util.preloads_from(
-            preloads_cls=Preloads,
-            use_preloaded_grid=use_preloaded_grid,
-            mesh_grids_of_planes=mesh_grids_of_planes,
-            use_w_tilde=False,
-        )
 
         fit_dataset_list.append(
             FitInterferometer(
@@ -100,7 +82,6 @@ def _fit_interferometer_from(
                 dataset_model=dataset_model,
                 adapt_images=adapt_images,
                 settings_inversion=settings_inversion,
-                preloads=preloads,
             )
         )
 
@@ -112,7 +93,6 @@ def __init__(
         self,
         aggregator: af.Aggregator,
         settings_inversion: Optional[aa.SettingsInversion] = None,
-        use_preloaded_grid: bool = True,
         real_space_mask: Optional[aa.Mask2D] = None,
     ):
         """
@@ -147,15 +127,10 @@ def __init__(
             A `PyAutoFit` aggregator object which can load the results of model-fits.
         settings_inversion
             Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
-        use_preloaded_grid
-            Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This
-            grid may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is
-            used as the fit.
         """
         super().__init__(aggregator=aggregator)
 
         self.settings_inversion = settings_inversion
-        self.use_preloaded_grid = use_preloaded_grid
         self.real_space_mask = real_space_mask
 
     def object_via_gen_from(
@@ -178,5 +153,4 @@ def object_via_gen_from(
             fit=fit,
             instance=instance,
             settings_inversion=self.settings_inversion,
-            use_preloaded_grid=self.use_preloaded_grid,
         )
diff --git a/autogalaxy/analysis/analysis/analysis.py b/autogalaxy/analysis/analysis/analysis.py
index d4ea19600..15645c66b 100644
--- a/autogalaxy/analysis/analysis/analysis.py
+++ b/autogalaxy/analysis/analysis/analysis.py
@@ -42,8 +42,8 @@ def galaxies_via_instance_from(
         """
         Create a list of galaxies from a model instance, which is used to fit the dataset.
 
-        The instance may only contain galaxies, in which case this function is redundant. However, if extra galaxies
-        are included, the instance will contain both galaxies and extra galaxies, and they should be added to create
+        The instance may only contain galaxies, in which case this function is redundant. However, if the clumns
+        API is being used, the instance will contain both galaxies and clumps, and they should be added to create
         the single list of galaxies used to fit the dataset.
 
         Parameters
@@ -56,12 +56,11 @@ def galaxies_via_instance_from(
         -------
         A list of galaxies that is used to then fit the dataset.
         """
-        if hasattr(instance, "extra_galaxies"):
-            if getattr(instance, "extra_galaxies", None) is not None:
-                return Galaxies(
-                    galaxies=instance.galaxies + instance.extra_galaxies,
-                    run_time_dict=run_time_dict
-                )
+        if hasattr(instance, "clumps"):
+            return Galaxies(
+                galaxies=instance.galaxies + instance.clumps,
+                run_time_dict=run_time_dict,
+            )
 
         return Galaxies(galaxies=instance.galaxies, run_time_dict=run_time_dict)
 
@@ -139,7 +138,7 @@ def profile_log_likelihood_function(
         An `info_dict` is also created which stores information on aspects of the model and dataset that dictate
         run times, so the profiled times can be interpreted with this context.
 
-        The results of this profiling are then output to hard-disk in the `preloads` folder of the model-fit results,
+        The results of this profiling are then output to hard-disk in the `prefiling` folder of the model-fit results,
         which they can be inspected to ensure run-times are as expected.
 
         Parameters
@@ -148,7 +147,7 @@ def profile_log_likelihood_function(
             An instance of the model that is being fitted to the data by this analysis (whose parameters have been set
             via a non-linear search).
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
 
         Returns
@@ -192,33 +191,33 @@ def profile_log_likelihood_function(
 
         try:
             info_dict["image_pixels"] = self.dataset.grids.uniform.shape_slim
-            info_dict[
-                "sub_total_light_profiles"
-            ] = self.dataset.grids.uniform.over_sampler.sub_total
+            info_dict["sub_total_light_profiles"] = (
+                self.dataset.grids.uniform.over_sampler.sub_total
+            )
         except AttributeError:
             pass
 
         if fit.model_obj.has(cls=aa.Pixelization):
             info_dict["use_w_tilde"] = fit.inversion.settings.use_w_tilde
             try:
-                info_dict[
-                    "sub_total_pixelization"
-                ] = self.dataset.grids.pixelization.over_sampler.sub_total
+                info_dict["sub_total_pixelization"] = (
+                    self.dataset.grids.pixelization.over_sampler.sub_total
+                )
             except AttributeError:
                 pass
-            info_dict[
-                "use_positive_only_solver"
-            ] = fit.inversion.settings.use_positive_only_solver
-            info_dict[
-                "force_edge_pixels_to_zeros"
-            ] = fit.inversion.settings.force_edge_pixels_to_zeros
+            info_dict["use_positive_only_solver"] = (
+                fit.inversion.settings.use_positive_only_solver
+            )
+            info_dict["force_edge_pixels_to_zeros"] = (
+                fit.inversion.settings.force_edge_pixels_to_zeros
+            )
             info_dict["use_w_tilde_numpy"] = fit.inversion.settings.use_w_tilde_numpy
             info_dict["source_pixels"] = len(fit.inversion.reconstruction)
 
             if hasattr(fit.inversion, "w_tilde"):
-                info_dict[
-                    "w_tilde_curvature_preload_size"
-                ] = fit.inversion.w_tilde.curvature_preload.shape[0]
+                info_dict["w_tilde_curvature_preload_size"] = (
+                    fit.inversion.w_tilde.curvature_preload.shape[0]
+                )
 
         self.output_profiling_info(
             paths=paths, run_time_dict=run_time_dict, info_dict=info_dict
@@ -239,7 +238,7 @@ def output_profiling_info(
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
         run_time_dict
             A dictionary containing the profiling times of the functions called by the `log_likelihood_function`.
diff --git a/autogalaxy/analysis/analysis/dataset.py b/autogalaxy/analysis/analysis/dataset.py
index 28c3cdd64..ac5e957e0 100644
--- a/autogalaxy/analysis/analysis/dataset.py
+++ b/autogalaxy/analysis/analysis/dataset.py
@@ -1,17 +1,13 @@
 import copy
 import logging
 from typing import Optional, Union
-import os
 
-from autoconf import conf
 from autoconf.dictable import to_dict, output_to_json
 import autofit as af
 import autoarray as aa
 
 from autogalaxy.analysis.adapt_images.adapt_image_maker import AdaptImageMaker
 from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages
-from autogalaxy.analysis.maker import FitMaker
-from autogalaxy.analysis.preloads import Preloads
 from autogalaxy.cosmology.lensing import LensingCosmology
 from autogalaxy.cosmology.wrap import Planck15
 from autogalaxy.analysis.analysis.analysis import Analysis
@@ -62,18 +58,8 @@ def __init__(
 
         self.settings_inversion = settings_inversion or aa.SettingsInversion()
 
-        self.preloads = self.preloads_cls()
-
         self.title_prefix = title_prefix
 
-    @property
-    def preloads_cls(self):
-        return Preloads
-
-    @property
-    def fit_maker_cls(self):
-        return FitMaker
-
     @property
     def adapt_images(self):
         if self._adapt_images is not None:
@@ -82,67 +68,6 @@ def adapt_images(self):
         if self.adapt_image_maker is not None:
             return self.adapt_image_maker.adapt_images
 
-    def set_preloads(self, paths: af.DirectoryPaths, model: af.Collection):
-        """
-        It is common for the model to have components whose parameters are all fixed, and thus the way that component
-        fits the data does not change. For example, if all parameter associated with the light profiles of galaxies
-        in the model are fixed, the image generated from these galaxies will not change irrespective of the model
-        parameters chosen by the non-linear search.
-
-        Preloading exploits this to speed up the log likelihood function, by inspecting the model and storing in memory
-        quantities that do not change. For the example above, the image of all galaxies would be stored in memory and
-        to perform every fit in the `log_likelihood_funtion`.
-
-        This function sets up all preload quantities, which are described fully in the `preloads` modules. This
-        occurs directly before the non-linear search begins, to ensure the model parameterization is fixed.
-
-        Parameters
-        ----------
-        paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
-            visualization and the pickled objects used by the aggregator output by this function.
-        model
-            The model object, which includes model components representing the galaxies that are fitted to
-            the imaging data.
-        """
-
-        logger.info(
-            "PRELOADS - Setting up preloads, may take a few minutes for fits using an inversion."
-        )
-
-        self.preloads = self.preloads_cls()
-
-        settings_inversion_original = copy.copy(self.settings_inversion)
-
-        self.settings_inversion.image_mesh_min_mesh_pixels_per_pixel = None
-        self.settings_inversion.image_mesh_adapt_background_percent_threshold = None
-
-        fit_maker = self.fit_maker_cls(model=model, fit_from=self.fit_from)
-
-        fit_0 = fit_maker.fit_via_model_from(unit_value=0.45)
-        fit_1 = fit_maker.fit_via_model_from(unit_value=0.55)
-
-        if fit_0 is None or fit_1 is None:
-            self.preloads = self.preloads_cls(failed=True)
-
-            self.settings_inversion = settings_inversion_original
-
-        else:
-            self.preloads = self.preloads_cls.setup_all_via_fits(
-                fit_0=fit_0, fit_1=fit_1
-            )
-
-            if conf.instance["general"]["test"]["check_preloads"]:
-                self.preloads.check_via_fit(fit=fit_0)
-
-            self.settings_inversion = settings_inversion_original
-
-        if isinstance(paths, af.DatabasePaths):
-            return
-
-        os.makedirs(paths.profile_path, exist_ok=True)
-        self.preloads.output_info_to_summary(file_path=paths.profile_path)
-
     def save_attributes(self, paths: af.DirectoryPaths):
         """
         Before the model-fit via the non-linear search begins, this routine saves attributes of the `Analysis` object
@@ -165,7 +90,7 @@ def save_attributes(self, paths: af.DirectoryPaths):
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization, and the pickled objects used by the aggregator output by this function.
         """
         paths.save_fits(
@@ -220,7 +145,7 @@ def save_results(self, paths: af.DirectoryPaths, result: ResultDataset):
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
         result
             The result of a model fit, including the non-linear search, samples and maximum likelihood tracer.
diff --git a/autogalaxy/analysis/chaining_util.py b/autogalaxy/analysis/chaining_util.py
index 5affa0bc6..52353670d 100644
--- a/autogalaxy/analysis/chaining_util.py
+++ b/autogalaxy/analysis/chaining_util.py
@@ -166,19 +166,19 @@ def source_from(
     return source_custom_model_from(result=result, source_is_model=True)
 
 
-def extra_galaxies_from(
+def clumps_from(
     result: Result,
     light_as_model: bool = False,
     mass_as_model: bool = False,
     free_centre: bool = False,
 ) -> af.Collection:
     """
-    The extra galaxies API models the light and / or mass of additional galaxies surrouding the main galaxy or
-    strong lens system.
+    The clump API models the light and / or mass of additional galaxies surrouding the main galaxy or strong lens
+    system.
 
-    This function performs model composition of extra galaxies for fits using the search chaining API. It makes it
-    possible to pass the extra galaxy parameters from a previous search to a new search, such that the new extra
-    galaxies are either treated as an instance or model component.
+    This function performs model composition of clumps for fits using the search chaining API. It makes it possible to
+    pass the clump parameters from a previous search to a new search, such that the new clumps are either treated
+    as an instance or model component.
 
     This function currently requires that mass profiles are `IsothermalSph` objects and that light profiles are
     `Sersic` objects. This will be generalised in the future.
@@ -186,70 +186,57 @@ def extra_galaxies_from(
     Parameters
     ----------
     result
-        The result, which includes extra galaxies, of the previous search, which via prior passing are used to create
-        the new extra galaxy model.
+        The result, which includes clumps, of the previous search, which via prior passing are used to create the new
+        clump model.
     light_as_model
-        If `True`, the extra_galaxy light is passed as a model component, else it is a fixed instance.
+        If `True`, the clump light is passed as a model component, else it is a fixed instance.
     mass_as_model
-        If `True`, the extra_galaxy mass is passed as a model component, else it is a fixed instance.
+        If `True`, the clump mass is passed as a model component, else it is a fixed instance.
     free_centre
         If `True`, the requested mass and/or light model centres are passed as a model, else they are fixed.
 
     Returns
     -------
     af.Collection
-        A collection of extra galaxy `Galaxy` objects, where each is either an instance or model component.
+        A collection of clump `Galaxy` objects, where each clump is either an instance or model component.
     """
     # ideal API:
 
-    # extra_galaxies = result.instance.extra_galaxies.as_model((LightProfile, mp.MassProfile,), fixed="centre", prior_pass=True)
-
-    if result.instance.extra_galaxies is None:
-        return None
+    # clumps = result.instance.clumps.as_model((LightProfile, mp.MassProfile,), fixed="centre", prior_pass=True)
 
     if mass_as_model:
-        extra_galaxies = result.instance.extra_galaxies.as_model((MassProfile,))
-
-        for extra_galaxy_index in range(len(result.instance.extra_galaxies)):
-            if hasattr(result.instance.extra_galaxies[extra_galaxy_index], "mass"):
-                extra_galaxies[
-                    extra_galaxy_index
-                ].mass.centre = result.instance.extra_galaxies[
-                    extra_galaxy_index
+        clumps = result.instance.clumps.as_model((MassProfile,))
+
+        for clump_index in range(len(result.instance.clumps)):
+            if hasattr(result.instance.clumps[clump_index], "mass"):
+                clumps[clump_index].mass.centre = result.instance.clumps[
+                    clump_index
                 ].mass.centre
-                extra_galaxies[
-                    extra_galaxy_index
-                ].mass.einstein_radius = result.model.extra_galaxies[
-                    extra_galaxy_index
+                clumps[clump_index].mass.einstein_radius = result.model.clumps[
+                    clump_index
                 ].mass.einstein_radius
                 if free_centre:
-                    extra_galaxies[
-                        extra_galaxy_index
-                    ].mass.centre = result.model.extra_galaxies[
-                        extra_galaxy_index
+                    clumps[clump_index].mass.centre = result.model.clumps[
+                        clump_index
                     ].mass.centre
 
     elif light_as_model:
-        extra_galaxies = result.instance.extra_galaxies.as_model((LightProfile,))
-
-        for extra_galaxy_index in range(len(result.instance.extra_galaxies)):
-            if extra_galaxies[extra_galaxy_index].bulge is not None:
-                extra_galaxies[
-                    extra_galaxy_index
-                ].bulge.centre = result.instance.extra_galaxies[
-                    extra_galaxy_index
-                ].bulge.centre
+        clumps = result.instance.clumps.as_model((LightProfile,))
+
+        for clump_index in range(len(result.instance.clumps)):
+            if clumps[clump_index].light is not None:
+                clumps[clump_index].light.centre = result.instance.clumps[
+                    clump_index
+                ].light.centre
                 if free_centre:
-                    extra_galaxies[
-                        extra_galaxy_index
-                    ].bulge.centre = result.model.extra_galaxies[
-                        extra_galaxy_index
-                    ].bulge.centre
+                    clumps[clump_index].light.centre = result.model.clumps[
+                        clump_index
+                    ].light.centre
 
     else:
-        extra_galaxies = result.instance.extra_galaxies.as_model(())
+        clumps = result.instance.clumps.as_model(())
 
-    return extra_galaxies
+    return clumps
 
 
 def lp_chain_tracer_from(light_result, settings_search):
diff --git a/autogalaxy/analysis/clump_model.py b/autogalaxy/analysis/clump_model.py
new file mode 100644
index 000000000..509c40f83
--- /dev/null
+++ b/autogalaxy/analysis/clump_model.py
@@ -0,0 +1,204 @@
+from typing import List, Optional, Type
+
+import autofit as af
+import autoarray as aa
+
+from autogalaxy.galaxy.galaxy import Galaxy
+from autogalaxy.profiles.light.abstract import LightProfile
+from autogalaxy.profiles.mass.abstract.abstract import MassProfile
+
+
+class ClumpModel:
+    def __init__(
+        self,
+        redshift: float,
+        centres: aa.Grid2DIrregular,
+        light_cls: Optional[Type[LightProfile]] = None,
+        mass_cls: Optional[Type[MassProfile]] = None,
+        einstein_radius_upper_limit: Optional[float] = None,
+        unfix_centres: bool = False,
+    ):
+        """
+        The clump API allows creates model components which model the light and mass of galaxies that are nearby the
+        main galaxy(s) of interest.
+
+        The `ClumpModel` object handles the creation of these model components to streamline model composition with
+        multiple clumps.
+
+        Every galaxy which is modeled as a clump has its centre input into this object which is fixed to this value
+        for model-fitting. All clumps are created as model `Galaxy` objects with a shard input redshift.
+
+        The light and mass profiles of the clumps are input via the `light_cls` and `mass_cls` inputs. If either is
+        omitted the clumps are not assigned a light or mass model.
+
+        Parameters
+        ----------
+        redshift
+            The redshift value of all clumps, which is likely the same as the main galaxy redshift.
+        centres
+            The centre of every clump in the model, whose light and mass profile centres are fixed to this value
+            throughout the model-fit.
+        light_cls
+            The light profile given to all clumps; if omitted all clumps have no light profile.
+        mass_cls
+            The mass profile given to all clumps; if omitted all clumps have no mass profile.
+        einstein_radius_upper_limit
+            The upper limit given to any mass model's `einstein_radius` parameter (e.g. if `IsothermalSph` profiles
+            are used to model clumps).
+        unfix_centres
+            If required, change the mass and light centres from fixed values to Uniform Prior Models +/- 0.1 around
+            the input centres.
+        """
+        self.redshift = redshift
+        self.centres = centres
+
+        self.light_cls = light_cls
+        self.mass_cls = mass_cls
+
+        self.einstein_radius_upper_limit = einstein_radius_upper_limit
+        self.unfix_centres = unfix_centres
+
+        self.centre_prior_half_width = 0.1
+
+    @property
+    def total_clumps(self) -> int:
+        return len(self.centres.in_list)
+
+    def unfix_centre(obj, centre):
+        obj.centre.centre_0 = af.UniformPrior(
+            lower_limit=new_centre[0] - centre_prior_half_width,
+            upper_limit=new_centre[0] + centre_prior_half_width,
+        )
+        obj.centre.centre_1 = af.UniformPrior(
+            lower_limit=new_centre[1] - centre_prior_half_width,
+            upper_limit=new_centre[1] + centre_prior_half_width,
+        )
+        return obj
+
+    @property
+    def light_list(self) -> Optional[List[af.Model]]:
+        """
+        Returns a list of every clump's light model, where the centre of that light model is fixed to its corresponding
+        input clump's centre, unless specified to be free.
+        """
+        if self.light_cls is None:
+            return None
+
+        light_list = []
+
+        for centre in self.centres.in_list:
+            if self.unfix_centres:
+                light_list.append(unfix_centre(af.Model(self.light_cls), centre))
+            else:
+                light_list.append(af.Model(self.light_cls, centre=centre))
+
+        return light_list
+
+    @property
+    def mass_list(self) -> Optional[List[af.Model]]:
+        """
+        Returns a list of every clump's mass model, where the centre of that mass model is fixed to its corresponding
+        input clump's centre, unless specified to be free.
+        """
+        if self.mass_cls is None:
+            return None
+
+        mass_list = []
+
+        for centre in self.centres.in_list:
+            if self.unfix_centres:
+                mass = unfix_centres(af.Model(self.mass_cls), centre)
+            else:
+                mass = af.Model(self.mass_cls, centre=centre)
+
+            if (
+                hasattr(mass, "einstein_radius")
+                and self.einstein_radius_upper_limit is not None
+            ):
+                mass.einstein_radius = af.UniformPrior(
+                    lower_limit=0.0, upper_limit=self.einstein_radius_upper_limit
+                )
+            mass_list.append(mass)
+
+        return mass_list
+
+    @property
+    def clumps_light_only(self) -> af.Collection:
+        """
+        Returns all clumps as a `Collection` model object, which can be easily added to an overall model `Collection`
+        of galaxies (see the `clumps` property below for more details).
+
+        All clumps only contain their model mass profiles, which is important for certain pipelines which omit
+        the clump mass profiles.
+        """
+        clumps_light_only = {}
+
+        for i in range(self.total_clumps):
+            light = self.light_list[i] if self.light_cls is not None else None
+
+            clumps_light_only[f"clump_{i}"] = af.Model(
+                Galaxy, redshift=self.redshift, light=light
+            )
+
+        return af.Collection(**clumps_light_only)
+
+    @property
+    def clumps_mass_only(self) -> af.Collection:
+        """
+        Returns all clumps as a `Collection` model object, which can be easily added to an overall model `Collection`
+        of galaxies (see the `clumps` property below for more details).
+
+        All clumps only contain their model mass profiles, which is important for certain pipelines which omit
+        the clump light profiles.
+        """
+        clumps_mass_only = {}
+
+        for i in range(self.total_clumps):
+            mass = self.mass_list[i] if self.mass_cls is not None else None
+
+            clumps_mass_only[f"clump_{i}"] = af.Model(
+                Galaxy, redshift=self.redshift, mass=mass
+            )
+
+        return af.Collection(**clumps_mass_only)
+
+    @property
+    def clumps(self) -> af.Collection:
+        """
+        Returns all clumps as a `Collection` model object, which can be easily added to an overall model `Collection`
+        of galaxies.
+
+        To make this `Collection` every clump centre and the input `light_cls` / `mass_cls` are used to create a
+        dictionary of model `Galaxy` objects with fixed light and mass profile centres. Their redshifts use the
+        input redshift.
+
+        The keys of this dictionary are numerically ordered as `clump_0`, `clump_2` etc.
+        """
+        clumps = {}
+
+        for i in range(self.total_clumps):
+            light = self.light_list[i] if self.light_cls is not None else None
+            mass = self.mass_list[i] if self.mass_cls is not None else None
+
+            clumps[f"clump_{i}"] = af.Model(
+                Galaxy, redshift=self.redshift, light=light, mass=mass
+            )
+
+        return af.Collection(**clumps)
+
+
+class ClumpModelDisabled:
+    def __init__(self):
+        pass
+
+    @property
+    def clumps_light_only(self):
+        return af.Collection()
+
+    @property
+    def clumps_mass_only(self):
+        return af.Collection()
+
+    @property
+    def clumps(self):
+        return af.Collection()
diff --git a/autogalaxy/analysis/maker.py b/autogalaxy/analysis/maker.py
deleted file mode 100644
index 154042ee9..000000000
--- a/autogalaxy/analysis/maker.py
+++ /dev/null
@@ -1,128 +0,0 @@
-import logging
-from typing import Callable, Union
-
-import autoarray as aa
-import autofit as af
-from autoconf import conf
-from autofit.exc import PriorLimitException
-from autogalaxy.analysis.preloads import Preloads
-
-logger = logging.getLogger(__name__)
-
-logger.setLevel(level="INFO")
-
-
-class FitMaker:
-    def __init__(self, model: af.Collection, fit_from: Callable):
-        """
-        Makes fits using an input PyAutoFit `model`, where the parameters of the model are drawn from its prior. This
-        uses an input `fit_from`, which given an `instance` of the model creates the fit object.
-
-        This is used for implicit preloading in the `Analysis` classes, whereby the created fits are compared against
-        one another to determine whether certain components of the analysis can be preloaded.
-
-        This includes functionality for creating the fit via the model in different ways, so that if certain
-        models are ill-defined another is used instead.
-
-        Parameters
-        ----------
-        model
-            A **PyAutoFit** model object which via its parameters and their priors can created instances of the model.
-        fit_from
-            A function which given the instance of the model creates a `Fit` object.
-        """
-
-        self.model = model
-        self.fit_from = fit_from
-
-    @property
-    def preloads_cls(self):
-        return Preloads
-
-    def fit_via_model_from(
-        self, unit_value: float
-    ) -> Union[aa.FitImaging, aa.FitInterferometer]:
-        """
-        Create a fit via the model.
-
-        This first tries to compute the fit from the input `unit_value`, where the `unit_value` defines unit hyper
-        cube values of each parameter's prior in the model, used to map each value to physical values for the fit.
-
-        If this model fit produces an `Exception` because the parameter combination does not fit the data accurately,
-        a sequence of random fits are instead used into an exception is not returned. However, if the number
-        of `preload_attempts` defined in the configuration files is exceeded a None is returned.
-
-        Parameters
-        ----------
-        unit_value
-            The unit hyper cube values of each parameter's prior in the model, used to map each value to physical
-            values for the fit.
-
-        Returns
-        -------
-        fit
-            A fit object where an instance of the model has been fitted to the data.
-        """
-        try:
-            try:
-                return self.fit_unit_instance_from(unit_value=unit_value)
-            except IndexError as e:
-                raise Exception from e
-        except (Exception, PriorLimitException):
-            return self.fit_random_instance_from()
-
-    def fit_unit_instance_from(
-        self, unit_value: float
-    ) -> Union[aa.FitImaging, aa.FitInterferometer]:
-        """
-        Create a fit via the model using an input `unit_value`, where the `unit_value` defines unit hyper
-        cube values of each parameter's prior in the model, used to map each value to physical values for the fit.
-
-        Parameters
-        ----------
-        unit_value
-            The unit hyper cube values of each parameter's prior in the model, used to map each value to physical
-            values for the fit.
-
-        Returns
-        -------
-        fit
-            A fit object where an instance of the model has been fitted to the data.
-        """
-        instance = self.model.instance_from_unit_vector(
-            unit_vector=[unit_value] * self.model.prior_count, ignore_prior_limits=True
-        )
-
-        fit = self.fit_from(
-            instance=instance,
-        )
-        fit.figure_of_merit
-        return fit
-
-    def fit_random_instance_from(self) -> Union[aa.FitImaging, aa.FitInterferometer]:
-        """
-        Create a fit via the model by guessing a sequence of random fits until an exception is not returned. If
-        the number of `preload_attempts` defined in the configuration files is exceeded a None is returned.
-
-        Returns
-        -------
-        fit
-            A fit object where an instance of the model has been fitted to the data.
-        """
-
-        preload_attempts = conf.instance["general"]["analysis"]["preload_attempts"]
-
-        for i in range(preload_attempts):
-            try:
-                instance = self.model.random_instance(ignore_prior_limits=True)
-
-                fit = self.fit_from(
-                    instance=instance,
-                )
-
-                fit.figure_of_merit
-
-                return fit
-
-            except Exception as e:
-                continue
diff --git a/autogalaxy/analysis/plotter_interface.py b/autogalaxy/analysis/plotter_interface.py
index bc5646635..df0da17d5 100644
--- a/autogalaxy/analysis/plotter_interface.py
+++ b/autogalaxy/analysis/plotter_interface.py
@@ -379,7 +379,7 @@ def adapt_images(
         Visualizes the adapt images used by a model-fit for adaptive pixelization mesh's and regularization.
 
         Images are output to the `image` folder of the `image_path` in a subfolder called `adapt`. When
-        used with a non-linear search the `image_path` is the output folder of the non-linear search.
+        used with a non-linear search the `image_path` points to the search's results folder.
 
         Visualization includes an image of the overall adapt model image and a subplot of all galaxy images on the same
         figure.
diff --git a/autogalaxy/analysis/preloads.py b/autogalaxy/analysis/preloads.py
deleted file mode 100644
index b1d513e86..000000000
--- a/autogalaxy/analysis/preloads.py
+++ /dev/null
@@ -1,213 +0,0 @@
-import logging
-import numpy as np
-from os import path
-from typing import Dict, Optional, List
-
-import autofit as af
-import autoarray as aa
-
-logger = logging.getLogger(__name__)
-
-logger.setLevel(level="INFO")
-
-
-class Preloads(aa.Preloads):
-    def __init__(
-        self,
-        w_tilde: Optional[aa.WTildeImaging] = None,
-        use_w_tilde: Optional[bool] = None,
-        blurred_image: Optional[aa.Array2D] = None,
-        image_plane_mesh_grid_pg_list: Optional[List[List[aa.Grid2D]]] = None,
-        relocated_grid: Optional[aa.Grid2D] = None,
-        mapper_list: Optional[aa.AbstractMapper] = None,
-        mapper_galaxy_dict: Optional[Dict[aa.AbstractMapper, "Galaxy"]] = None,
-        linear_func_operated_mapping_matrix_dict=None,
-        data_linear_func_matrix_dict=None,
-        mapper_operated_mapping_matrix_dict=None,
-        operated_mapping_matrix: Optional[np.ndarray] = None,
-        curvature_matrix: Optional[np.ndarray] = None,
-        regularization_matrix: Optional[np.ndarray] = None,
-        log_det_regularization_matrix_term: Optional[float] = None,
-        traced_mesh_grids_list_of_planes=None,
-        image_plane_mesh_grid_list=None,
-        failed=False,
-    ):
-        """
-        Class which offers a concise API for settings up the preloads, which before a model-fit are set up via
-        a comparison of two fits using two different models. If a quantity in these two fits is identical, it does
-        not change thoughout the model-fit and can therefore be preloaded to avoid computation, speeding up
-        the analysis.
-
-        For example, the image-plane source-plane pixelization grid (which may be computationally expensive to compute
-        via a KMeans algorithm) does not change for the majority of model-fits, because the associated model parameters
-        are fixed. Preloading avoids rerruning the KMeans algorithm for every model fitted, by preloading it in memory
-        and using this preload in every fit.
-
-        Parameters
-        ----------
-        blurred_image
-            The preloaded array of values containing the blurred image of a model fit (e.g. that light profile of
-            every galaxy in the model). This can be preloaded when no light profiles in the model vary.
-        w_tilde
-            A class containing values that enable an inversion's linear algebra to use the w-tilde formalism. This can
-            be preloaded when no component of the model changes the noise map (e.g. galaxies are fixed).
-        use_w_tilde
-            Whether to use the w tilde formalism, which superseeds the value in `SettingsInversions` such that w tilde
-            will be disabled for model-fits it is not applicable (e.g. because the noise-map changes).
-        traced_grids_of_planes_for_inversion
-            The two dimensional grids corresponding to the traced grids in a lens fit. This can be preloaded when no
-             mass profiles in the model vary.
-        image_plane_mesh_grid_pg_list
-            The two dimensional grids corresponding to the sparse image plane grids in a lens fit, that is ray-traced to
-            the source plane to form the source pixelization. This can be preloaded when no pixelizations in the model
-            vary.
-        relocated_grid
-            The two dimensional grids corresponding to the grid that has had its border pixels relocated for a
-            pixelization in a lens fit. This can be preloaded when no mass profiles in the model vary.
-        mapper_list
-            The mapper of a fit, which preloading avoids recalculation of the mapping matrix and image to source
-            pixel mappings. This can be preloaded when no pixelizations in the model vary.
-        operated_mapping_matrix
-            A matrix containing the mappings between PSF blurred image pixels and source pixels used in the linear
-            algebra of an inversion. This can be preloaded when no mass profiles and pixelizations in the model vary.
-
-        Returns
-        -------
-        The preloads object used to skip certain calculations in the log likelihood function.
-        """
-        super().__init__(
-            w_tilde=w_tilde,
-            use_w_tilde=use_w_tilde,
-            relocated_grid=relocated_grid,
-            image_plane_mesh_grid_pg_list=image_plane_mesh_grid_pg_list,
-            mapper_list=mapper_list,
-            linear_func_operated_mapping_matrix_dict=linear_func_operated_mapping_matrix_dict,
-            data_linear_func_matrix_dict=data_linear_func_matrix_dict,
-            mapper_operated_mapping_matrix_dict=mapper_operated_mapping_matrix_dict,
-            operated_mapping_matrix=operated_mapping_matrix,
-            curvature_matrix=curvature_matrix,
-            regularization_matrix=regularization_matrix,
-            log_det_regularization_matrix_term=log_det_regularization_matrix_term,
-            traced_mesh_grids_list_of_planes=traced_mesh_grids_list_of_planes,
-            image_plane_mesh_grid_list=image_plane_mesh_grid_list,
-        )
-
-        self.mapper_galaxy_dict = mapper_galaxy_dict
-        self.blurred_image = blurred_image
-        self.failed = failed
-
-    @classmethod
-    def setup_all_via_fits(cls, fit_0, fit_1) -> "Preloads":
-        """
-        Setup the Preloads from two fits which use two different lens model of a model-fit.
-
-        Parameters
-        ----------
-        fit_0
-            The first fit corresponding to a model with a specific set of unit-values.
-        fit_1
-            The second fit corresponding to a model with a different set of unit-values.
-
-        Returns
-        -------
-        Preloads
-            Preloads which are set up based on the fit's passed in specific to a lens model.
-
-        """
-
-        preloads = cls()
-
-        if isinstance(fit_0, aa.FitImaging):
-            preloads.set_w_tilde_imaging(fit_0=fit_0, fit_1=fit_1)
-            preloads.set_blurred_image(fit_0=fit_0, fit_1=fit_1)
-
-        preloads.set_mapper_list(fit_0=fit_0, fit_1=fit_1)
-
-        if preloads.mapper_list is not None:
-            preloads.mapper_galaxy_dict = fit_0.galaxies_to_inversion.mapper_galaxy_dict
-
-        preloads.set_operated_mapping_matrix_with_preloads(fit_0=fit_0, fit_1=fit_1)
-        preloads.set_linear_func_inversion_dicts(fit_0=fit_0, fit_1=fit_1)
-        preloads.set_curvature_matrix(fit_0=fit_0, fit_1=fit_1)
-        preloads.set_regularization_matrix_and_term(fit_0=fit_0, fit_1=fit_1)
-
-        return preloads
-
-    def set_blurred_image(self, fit_0, fit_1):
-        """
-        If the `LightProfile`'s in a model are all fixed parameters their corresponding image and therefore PSF blurred
-        image do not change during the model fit and can therefore be preloaded.
-
-        This function compares the blurred image of two fit's corresponding to two model instances, and preloads
-        the blurred image if the blurred image of both fits are the same.
-
-        The preload is typically used though out search chaining pipelines, as it is common to fix the lens light for
-        the majority of model-fits.
-
-        Parameters
-        ----------
-        fit_0
-            The first fit corresponding to a model with a specific set of unit-values.
-        fit_1
-            The second fit corresponding to a model with a different set of unit-values.
-        """
-        self.blurred_image = None
-
-        precision = 1e-8
-
-        if (np.max(abs(fit_0.blurred_image - fit_1.blurred_image)) < precision) and (
-            np.sum(fit_0.blurred_image) > precision
-        ):
-            self.blurred_image = fit_0.blurred_image
-
-            logger.info(
-                "PRELOADS - Blurred image (e.g. the image of all light profiles) is preloaded for this model-fit."
-            )
-
-    def output_info_to_summary(self, file_path):
-        file_preloads = path.join(file_path, "preloads.summary")
-
-        if self.failed:
-            logger.info(
-                "PRELOADS - Preloading failed as models gave too many exceptions, preloading therefore "
-                "not used."
-            )
-
-            af.formatter.output_list_of_strings_to_file(
-                file=file_preloads, list_of_strings=["FAILED"]
-            )
-
-        else:
-            af.formatter.output_list_of_strings_to_file(
-                file=file_preloads, list_of_strings=self.info
-            )
-
-    @property
-    def info(self) -> List[str]:
-        """
-        The information on what has or has not been preloaded, which is written to the file `preloads.summary`.
-
-        Returns
-        -------
-            A list of strings containing True and False values as to whether a quantity has been preloaded.
-        """
-        line = [f"W Tilde = {self.w_tilde is not None}\n"]
-        line += [f"Use W Tilde = {self.use_w_tilde}\n\n"]
-        line += [f"Blurred Image = {np.count_nonzero(self.blurred_image) != 0}\n"]
-        line += [f"Mapper = {self.mapper_list is not None}\n"]
-        line += [
-            f"Blurred Mapping Matrix = {self.operated_mapping_matrix is not None}\n"
-        ]
-        line += [
-            f"Inversion Linear Func (Linear Light Profile) Dicts = {self.linear_func_operated_mapping_matrix_dict is not None}\n"
-        ]
-        line += [f"Curvature Matrix = {self.curvature_matrix is not None}\n"]
-        line += [
-            f"Curvature Matrix Mapper Diag = {self.curvature_matrix_mapper_diag is not None}\n"
-        ]
-        line += [f"Regularization Matrix = {self.regularization_matrix is not None}\n"]
-        line += [
-            f"Log Det Regularization Matrix Term = {self.log_det_regularization_matrix_term is not None}\n"
-        ]
-
-        return line
diff --git a/autogalaxy/config/general.yaml b/autogalaxy/config/general.yaml
index 61f8b09a0..361abe732 100644
--- a/autogalaxy/config/general.yaml
+++ b/autogalaxy/config/general.yaml
@@ -1,5 +1,3 @@
-analysis:
-  preload_attempts: 250
 fits:
   flip_for_ds9: true
 grid:
@@ -11,6 +9,4 @@ inversion:
   use_border_relocator: true          # If True, by default a pixelization's border is used to relocate all pixels outside its border to the border.
 test:
   check_likelihood_function: true   # if True, when a search is resumed the likelihood of a previous sample is recalculated to ensure it is consistent with the previous run.
-  check_preloads: false
-  exception_override: false
-  preloads_check_threshold: 1.0     # If the figure of merit of a fit with and without preloads is greater than this threshold, the check preload test fails and an exception raised for a model-fit. 
+  exception_override: false
\ No newline at end of file
diff --git a/autogalaxy/convert.py b/autogalaxy/convert.py
index 3364b1812..8ff432d7b 100644
--- a/autogalaxy/convert.py
+++ b/autogalaxy/convert.py
@@ -1,5 +1,5 @@
 from astropy import units
-import numpy as np
+from autofit.jax_wrapper import numpy as np
 from typing import Tuple
 
 
diff --git a/autogalaxy/ellipse/model/analysis.py b/autogalaxy/ellipse/model/analysis.py
index 20a904e7a..07f92f41c 100644
--- a/autogalaxy/ellipse/model/analysis.py
+++ b/autogalaxy/ellipse/model/analysis.py
@@ -181,7 +181,7 @@ def save_attributes(self, paths: af.DirectoryPaths):
          Parameters
          ----------
          paths
-             The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+             The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
              visualization, and the pickled objects used by the aggregator output by this function.
         """
         paths.save_fits(
@@ -217,7 +217,7 @@ def profile_log_likelihood_function(
         An `info_dict` is also created which stores information on aspects of the model and dataset that dictate
         run times, so the profiled times can be interpreted with this context.
 
-        The results of this profiling are then output to hard-disk in the `preloads` folder of the model-fit results,
+        The results of this profiling are then output to hard-disk in the `profiling` folder of the model-fit results,
         which they can be inspected to ensure run-times are as expected.
 
         Parameters
@@ -226,7 +226,7 @@ def profile_log_likelihood_function(
             An instance of the model that is being fitted to the data by this analysis (whose parameters have been set
             via a non-linear search).
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
 
         Returns
diff --git a/autogalaxy/ellipse/model/plotter_interface.py b/autogalaxy/ellipse/model/plotter_interface.py
index bce0ecdf6..cbbf52fc6 100644
--- a/autogalaxy/ellipse/model/plotter_interface.py
+++ b/autogalaxy/ellipse/model/plotter_interface.py
@@ -17,7 +17,7 @@ def imaging(self, dataset: aa.Imaging):
         Visualizes an `Imaging` dataset object.
 
         Images are output to the `image` folder of the `image_path` in a subfolder called `imaging`. When used with
-        a non-linear search the `image_path` is the output folder of the non-linear search.
+        a non-linear search the `image_path` points to the search's results folder.
         `.
         Visualization includes individual images of attributes of the dataset (e.g. the image, noise map, PSF) and a
         subplot of all these attributes on the same figure.
diff --git a/autogalaxy/ellipse/model/visualizer.py b/autogalaxy/ellipse/model/visualizer.py
index 42862f101..781bf6213 100644
--- a/autogalaxy/ellipse/model/visualizer.py
+++ b/autogalaxy/ellipse/model/visualizer.py
@@ -20,10 +20,10 @@ def visualize_before_fit(
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
         model
-            The model object, which includes model components representing the galaxies that are fitted to
+            The PyAutoFit model object, which includes model components representing the galaxies that are fitted to
             the imaging data.
         """
 
@@ -55,7 +55,7 @@ def visualize(
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization, and the pickled objects used by the aggregator output by this function.
         instance
             An instance of the model that is being fitted to the data by this analysis (whose parameters have been set
diff --git a/autogalaxy/galaxy/to_inversion.py b/autogalaxy/galaxy/to_inversion.py
index 669d62f9d..76d233ed7 100644
--- a/autogalaxy/galaxy/to_inversion.py
+++ b/autogalaxy/galaxy/to_inversion.py
@@ -15,7 +15,6 @@
 from autogalaxy.profiles.light.linear import LightProfileLinear
 from autogalaxy.galaxy.galaxy import Galaxy
 from autogalaxy.galaxy.galaxies import Galaxies
-from autogalaxy.analysis.preloads import Preloads
 
 
 class AbstractToInversion:
@@ -24,7 +23,6 @@ def __init__(
         dataset: Optional[Union[aa.Imaging, aa.Interferometer, aa.DatasetInterface]],
         adapt_images: Optional[AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
-        preloads=Preloads(),
         run_time_dict: Optional[Dict] = None,
     ):
         """
@@ -59,8 +57,6 @@ def __init__(
             the pixelization's pixels to the brightest regions of the image.
         settings_inversion
             The settings of the inversion, which controls how the linear algebra calculation is performed.
-        preloads
-            Preloads of the inversion, which are used to speed up the linear algebra calculation.
         run_time_dict
             A dictionary of run-time values used to compute the inversion, for example the noise-map normalization.
         """
@@ -82,7 +78,6 @@ def __init__(
 
         self.settings_inversion = settings_inversion
 
-        self.preloads = preloads
         self.run_time_dict = run_time_dict
 
     @property
@@ -198,7 +193,6 @@ def __init__(
         galaxies: List[Galaxy],
         adapt_images: Optional[AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
-        preloads=aa.Preloads(),
         run_time_dict: Optional[Dict] = None,
     ):
         """
@@ -234,8 +228,6 @@ def __init__(
             the pixelization's pixels to the brightest regions of the image.
         settings_inversion
             The settings of the inversion, which controls how the linear algebra calculation is performed.
-        preloads
-            Preloads of the inversion, which are used to speed up the linear algebra calculation.
         run_time_dict
             A dictionary of run-time values used to compute the inversion, for example the noise-map normalization.
         """
@@ -245,7 +237,6 @@ def __init__(
             dataset=dataset,
             adapt_images=adapt_images,
             settings_inversion=settings_inversion,
-            preloads=preloads,
             run_time_dict=run_time_dict,
         )
 
@@ -500,7 +491,6 @@ def mapper_from(
             source_plane_mesh_grid=source_plane_mesh_grid,
             image_plane_mesh_grid=image_plane_mesh_grid,
             adapt_data=adapt_galaxy_image,
-            preloads=self.preloads,
             run_time_dict=self.run_time_dict,
         )
 
@@ -592,7 +582,6 @@ def inversion(self) -> aa.AbstractInversion:
             dataset=self.dataset,
             linear_obj_list=self.linear_obj_list,
             settings=self.settings_inversion,
-            preloads=self.preloads,
             run_time_dict=self.run_time_dict,
         )
 
diff --git a/autogalaxy/gui/scribbler.py b/autogalaxy/gui/scribbler.py
index d8183f504..505562cc8 100644
--- a/autogalaxy/gui/scribbler.py
+++ b/autogalaxy/gui/scribbler.py
@@ -166,13 +166,13 @@ def show_mask(self):
         junk_mask = masks["1"]
         feature_mask = masks["2"]
         plt.ioff()
-        # plt.figure()
-        # plt.subplot(121)
-        # plt.imshow(junk_mask.astype("int"), cmap="gray")
-        # plt.title("Junk mask")
-        # plt.subplot(122)
-        # plt.imshow(feature_mask.astype("int"), cmap="gray")
-        # plt.title("Feature mask")
+        plt.figure()
+        plt.subplot(121)
+        plt.imshow(junk_mask.astype("int"), cmap="gray")
+        plt.title("Junk mask")
+        plt.subplot(122)
+        plt.imshow(feature_mask.astype("int"), cmap="gray")
+        plt.title("Feature mask")
         # plt.show()
         return junk_mask
 
diff --git a/autogalaxy/imaging/fit_imaging.py b/autogalaxy/imaging/fit_imaging.py
index ff72fe5e4..06f331691 100644
--- a/autogalaxy/imaging/fit_imaging.py
+++ b/autogalaxy/imaging/fit_imaging.py
@@ -7,7 +7,6 @@
 
 from autogalaxy.abstract_fit import AbstractFitInversion
 from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages
-from autogalaxy.analysis.preloads import Preloads
 from autogalaxy.galaxy.galaxy import Galaxy
 from autogalaxy.galaxy.galaxies import Galaxies
 from autogalaxy.galaxy.to_inversion import GalaxiesToInversion
@@ -26,7 +25,6 @@ def __init__(
         dataset_model: Optional[aa.DatasetModel] = None,
         adapt_images: Optional[AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
-        preloads: aa.Preloads = Preloads(),
         run_time_dict: Optional[Dict] = None,
     ):
         """
@@ -66,16 +64,12 @@ def __init__(
             reconstructed galaxy's morphology.
         settings_inversion
             Settings controlling how an inversion is fitted for example which linear algebra formalism is used.
-        preloads
-            Contains preloaded calculations (e.g. linear algebra matrices) which can skip certain calculations in
-            the fit.
         run_time_dict
             A dictionary which if passed to the fit records how long fucntion calls which have the `profile_func`
             decorator take to run.
         """
 
         self.galaxies = Galaxies(galaxies=galaxies, run_time_dict=run_time_dict)
-        self.preloads = preloads
 
         super().__init__(
             dataset=dataset,
@@ -135,7 +129,6 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion:
             galaxies=self.galaxies,
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
-            preloads=self.preloads,
             run_time_dict=self.run_time_dict,
         )
 
@@ -220,9 +213,9 @@ def subtracted_images_of_galaxies_dict(self) -> Dict[Galaxy, aa.Array2D]:
 
             subtracted_image = self.data - sum(other_galaxies_model_images)
 
-            subtracted_images_of_galaxies_dict[
-                self.galaxies[galaxy_index]
-            ] = subtracted_image
+            subtracted_images_of_galaxies_dict[self.galaxies[galaxy_index]] = (
+                subtracted_image
+            )
 
         return subtracted_images_of_galaxies_dict
 
@@ -285,44 +278,3 @@ def galaxies_linear_light_profiles_to_light_profiles(self) -> List[Galaxy]:
         or `GalaxyPlotter` objects.
         """
         return self.model_obj_linear_light_profiles_to_light_profiles
-
-    def refit_with_new_preloads(
-        self,
-        preloads: Preloads,
-        settings_inversion: Optional[aa.SettingsInversion] = None,
-    ) -> "FitImaging":
-        """
-        Returns a new fit which uses the dataset, galaxies and other objects of this fit, but uses a different set of
-        preloads input into this function.
-
-        This is used when setting up the preloads objects, to concisely test how using different preloads objects
-        changes the attributes of the fit.
-
-        Parameters
-        ----------
-        preloads
-            The new preloads which are used to refit the data using the
-        settings_inversion
-            Settings controlling how an inversion is fitted for example which linear algebra formalism is used.
-
-        Returns
-        -------
-        A new fit which has used new preloads input into this function but the same dataset, galaxies and other settings.
-        """
-        run_time_dict = {} if self.run_time_dict is not None else None
-
-        settings_inversion = (
-            self.settings_inversion
-            if settings_inversion is None
-            else settings_inversion
-        )
-
-        return FitImaging(
-            dataset=self.dataset,
-            galaxies=self.galaxies,
-            dataset_model=self.dataset_model,
-            adapt_images=self.adapt_images,
-            settings_inversion=settings_inversion,
-            preloads=preloads,
-            run_time_dict=run_time_dict,
-        )
diff --git a/autogalaxy/imaging/model/analysis.py b/autogalaxy/imaging/model/analysis.py
index 135bc091e..ad2bc643f 100644
--- a/autogalaxy/imaging/model/analysis.py
+++ b/autogalaxy/imaging/model/analysis.py
@@ -9,7 +9,6 @@
 
 from autogalaxy.analysis.adapt_images.adapt_image_maker import AdaptImageMaker
 from autogalaxy.analysis.analysis.dataset import AnalysisDataset
-from autogalaxy.analysis.preloads import Preloads
 from autogalaxy.cosmology.lensing import LensingCosmology
 from autogalaxy.cosmology.wrap import Planck15
 from autogalaxy.imaging.model.result import ResultImaging
@@ -85,18 +84,15 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
         model
-            The model object, which includes model components representing the galaxies that are fitted to
+            The PyAutoFit model object, which includes model components representing the galaxies that are fitted to
             the imaging data.
         """
 
         super().modify_before_fit(paths=paths, model=model)
 
-        if not paths.is_complete:
-            self.set_preloads(paths=paths, model=model)
-
         return self
 
     def log_likelihood_function(self, instance: af.ModelInstance) -> float:
@@ -153,7 +149,6 @@ def log_likelihood_function(self, instance: af.ModelInstance) -> float:
     def fit_from(
         self,
         instance: af.ModelInstance,
-        preload_overwrite: Optional[Preloads] = None,
         run_time_dict: Optional[Dict] = None,
     ) -> FitImaging:
         """
@@ -186,15 +181,12 @@ def fit_from(
 
         adapt_images = self.adapt_images_via_instance_from(instance=instance)
 
-        preloads = self.preloads if preload_overwrite is None else preload_overwrite
-
         return FitImaging(
             dataset=self.dataset,
             galaxies=galaxies,
             dataset_model=dataset_model,
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
-            preloads=preloads,
             run_time_dict=run_time_dict,
         )
 
@@ -225,7 +217,7 @@ def save_attributes(self, paths: af.DirectoryPaths):
          Parameters
          ----------
          paths
-             The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+             The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
              visualization, and the pickled objects used by the aggregator output by this function.
         """
         super().save_attributes(paths=paths)
@@ -253,7 +245,7 @@ def profile_log_likelihood_function(
         An `info_dict` is also created which stores information on aspects of the model and dataset that dictate
         run times, so the profiled times can be interpreted with this context.
 
-        The results of this profiling are then output to hard-disk in the `preloads` folder of the model-fit results,
+        The results of this profiling are then output to hard-disk in the `profiling` folder of the model-fit results,
         which they can be inspected to ensure run-times are as expected.
 
         Parameters
@@ -262,7 +254,7 @@ def profile_log_likelihood_function(
             An instance of the model that is being fitted to the data by this analysis (whose parameters have been set
             via a non-linear search).
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
 
         Returns
diff --git a/autogalaxy/imaging/model/plotter_interface.py b/autogalaxy/imaging/model/plotter_interface.py
index 2352e4e1a..564369c82 100644
--- a/autogalaxy/imaging/model/plotter_interface.py
+++ b/autogalaxy/imaging/model/plotter_interface.py
@@ -1,5 +1,4 @@
 from os import path
-from typing import List
 
 import autoarray as aa
 import autoarray.plot as aplt
@@ -14,21 +13,21 @@
 class PlotterInterfaceImaging(PlotterInterface):
     def imaging(self, dataset: aa.Imaging):
         """
-        Output visualization of an `Imaging` dataset, typically before a model-fit is performed.
+        Visualizes an `Imaging` dataset object.
 
-        Images are output to the `image` folder of the `image_path` in a subfolder called `dataset`. When used with
-        a non-linear search the `image_path` is the output folder of the non-linear search.
+        Images are output to the `image` folder of the `image_path` in a subfolder called `imaging`. When used with
+        a non-linear search the `image_path` points to the search's results folder.
         `.
         Visualization includes individual images of attributes of the dataset (e.g. the image, noise map, PSF) and a
         subplot of all these attributes on the same figure.
 
-        The images output by the `PlotterInterface` are customized using the file `config/visualize/plots.yaml` under
-        the `dataset` header.
+        The images output by the `PlotterInterface` are customized using the file `config/visualize/plots.yaml` under the
+        [dataset] header.
 
         Parameters
         ----------
         dataset
-            The imaging dataset which is visualized.
+            The imaging dataset whose attributes are visualized.
         """
 
         def should_plot(name):
@@ -158,137 +157,3 @@ def should_plot(name):
                 subtracted_image=True,
                 model_image=True,
             )
-
-    def imaging_combined(self, dataset_list: List[aa.Imaging]):
-        """
-        Output visualization of all `Imaging` datasets in a summed combined analysis, typically before a model-fit
-        is performed.
-
-        Images are output to the `image` folder of the `image_path` in a subfolder called `dataset_combined`. When
-        used with a non-linear search the `image_path` is the output folder of the non-linear search.
-        `.
-        Visualization includes individual images of attributes of each dataset (e.g. the image, noise map, PSF) on
-        a single subplot, such that the full suite of multiple datasets can be viewed on the same figure.
-
-        The images output by the `PlotterInterface` are customized using the file `config/visualize/plots.yaml` under
-        the `dataset` header.
-
-        Parameters
-        ----------
-        dataset
-            The list of imaging datasets which are visualized.
-        """
-
-        def should_plot(name):
-            return plot_setting(section=["dataset", "imaging"], name=name)
-
-        mat_plot_2d = self.mat_plot_2d_from(subfolders="combined")
-
-        dataset_plotter_list = [
-            aplt.ImagingPlotter(
-                dataset=dataset, mat_plot_2d=mat_plot_2d, include_2d=self.include_2d
-            )
-            for dataset in dataset_list
-        ]
-
-        subplot_shape = (len(dataset_list), 4)
-
-        multi_plotter = aplt.MultiFigurePlotter(
-            plotter_list=dataset_plotter_list, subplot_shape=subplot_shape
-        )
-
-        if should_plot("subplot_dataset"):
-            multi_plotter.subplot_of_figures_multi(
-                func_name_list=["figures_2d"] * 4,
-                figure_name_list=["data", "noise_map", "signal_to_noise_map", "psf"],
-                filename_suffix="dataset",
-            )
-
-            for plotter in multi_plotter.plotter_list:
-                plotter.mat_plot_2d.use_log10 = True
-
-            multi_plotter.subplot_of_figures_multi(
-                func_name_list=["figures_2d"] * 4,
-                figure_name_list=["data", "noise_map", "signal_to_noise_map", "psf"],
-                filename_suffix="dataset_log10",
-            )
-
-    def fit_imaging_combined(self, fit_list: List[FitImaging]):
-        """
-        Output visualization of all `FitImaging` objects in a summed combined analysis, typically during or after a
-        model-fit is performed.
-
-        Images are output to the `image` folder of the `image_path` in a subfolder called `combined`. When used
-        with a non-linear search the `image_path` is the output folder of the non-linear search.
-        `.
-        Visualization includes individual images of attributes of each fit (e.g. data, normalized residual-map) on
-        a single subplot, such that the full suite of multiple datasets can be viewed on the same figure.
-
-        The images output by the `PlotterInterface` are customized using the file `config/visualize/plots.yaml` under
-        the `fit` header.
-
-        Parameters
-        ----------
-        fit
-            The list of imaging fits which are visualized.
-        """
-
-        def should_plot(name):
-            return plot_setting(section=["fit", "fit_imaging"], name=name)
-
-        mat_plot_2d = self.mat_plot_2d_from(subfolders="combined")
-
-        fit_plotter_list = [
-            FitImagingPlotter(
-                fit=fit, mat_plot_2d=mat_plot_2d, include_2d=self.include_2d
-            )
-            for fit in fit_list
-        ]
-
-        subplot_shape = (len(fit_list), 5)
-
-        multi_plotter = aplt.MultiFigurePlotter(
-            plotter_list=fit_plotter_list, subplot_shape=subplot_shape
-        )
-
-        if should_plot("subplot_fit"):
-
-            def make_subplot_fit(filename_suffix):
-                multi_plotter.subplot_of_figures_multi(
-                    func_name_list=["figures_2d"] * 4,
-                    figure_name_list=[
-                        "data",
-                        "signal_to_noise_map",
-                        "model_image",
-                        "normalized_residual_map",
-                    ],
-                    filename_suffix=filename_suffix,
-                    number_subplots=len(fit_list) * 5,
-                    close_subplot=False,
-                )
-
-                for plotter in multi_plotter.plotter_list:
-                    plotter.mat_plot_2d.cmap.kwargs["vmin"] = -1.0
-                    plotter.mat_plot_2d.cmap.kwargs["vmax"] = 1.0
-
-                multi_plotter.subplot_of_figures_multi(
-                    func_name_list=["figures_2d"],
-                    figure_name_list=[
-                        "normalized_residual_map",
-                    ],
-                    filename_suffix=filename_suffix,
-                    number_subplots=len(fit_list) * 5,
-                    subplot_index_offset=4,
-                    open_subplot=False,
-                )
-
-                for plotter in multi_plotter.plotter_list:
-                    plotter.mat_plot_2d.cmap.kwargs["vmin"] = None
-                    plotter.mat_plot_2d.cmap.kwargs["vmax"] = None
-
-            make_subplot_fit(filename_suffix="fit")
-
-            for plotter in multi_plotter.plotter_list:
-                plotter.mat_plot_2d.use_log10 = True
-
-            make_subplot_fit(filename_suffix="fit_log10")
diff --git a/autogalaxy/imaging/model/visualizer.py b/autogalaxy/imaging/model/visualizer.py
index 201c72d48..fa169c0ab 100644
--- a/autogalaxy/imaging/model/visualizer.py
+++ b/autogalaxy/imaging/model/visualizer.py
@@ -20,10 +20,10 @@ def visualize_before_fit(
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
         model
-            The model object, which includes model components representing the galaxies that are fitted to
+            The PyAutoFit model object, which includes model components representing the galaxies that are fitted to
             the imaging data.
         """
 
@@ -63,7 +63,7 @@ def visualize(
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization, and the pickled objects used by the aggregator output by this function.
         instance
             An instance of the model that is being fitted to the data by this analysis (whose parameters have been set
@@ -94,82 +94,3 @@ def visualize(
         )
         if fit.inversion is not None:
             plotter.inversion(inversion=fit.inversion, during_analysis=during_analysis)
-
-    @staticmethod
-    def visualize_before_fit_combined(
-        analyses,
-        paths: af.AbstractPaths,
-        model: af.AbstractPriorModel,
-    ):
-        """
-        Performs visualization before the non-linear search begins of information which shared across all analyses
-        on a single matplotlib figure.
-
-        This function outputs visuals of all information which does not vary during the fit, for example the dataset
-        being fitted.
-
-        Parameters
-        ----------
-        analyses
-            The list of all analysis objects used for fitting via yhe non-linear search.
-        paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
-            visualization and the pickled objects used by the aggregator output by this function.
-        model
-            The model object, which includes model components representing the galaxies that are fitted to
-            the imaging data.
-        """
-
-        if analyses is None:
-            return
-
-        plotter = PlotterInterfaceImaging(
-            image_path=paths.image_path, title_prefix=analyses[0].title_prefix
-        )
-
-        dataset_list = [analysis.dataset for analysis in analyses]
-
-        plotter.imaging_combined(
-            dataset_list=dataset_list,
-        )
-
-    @staticmethod
-    def visualize_combined(
-        analyses,
-        paths: af.AbstractPaths,
-        instance: af.ModelInstance,
-        during_analysis: bool,
-    ):
-        """
-        Performs visualization during the non-linear search of information which is shared across all analyses on a
-        single matplotlib figure.
-
-        This function outputs visuals of all information which varies during the fit, for example the model-fit to
-        the dataset being fitted.
-
-        Parameters
-        ----------
-        analyses
-            The list of all analysis objects used for fitting via yhe non-linear search.
-        paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
-            visualization and the pickled objects used by the aggregator output by this function.
-        model
-            The model object, which includes model components representing the galaxies that are fitted to
-            the imaging data.
-        """
-        if analyses is None:
-            return
-
-        plotter = PlotterInterfaceImaging(
-            image_path=paths.image_path, title_prefix=analyses[0].title_prefix
-        )
-
-        fit_list = [
-            analysis.fit_from(instance=single_instance)
-            for analysis, single_instance in zip(analyses, instance)
-        ]
-
-        plotter.fit_imaging_combined(
-            fit_list=fit_list,
-        )
diff --git a/autogalaxy/interferometer/fit_interferometer.py b/autogalaxy/interferometer/fit_interferometer.py
index 7f7f6fe6b..7095ced45 100644
--- a/autogalaxy/interferometer/fit_interferometer.py
+++ b/autogalaxy/interferometer/fit_interferometer.py
@@ -7,7 +7,6 @@
 
 from autogalaxy.abstract_fit import AbstractFitInversion
 from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages
-from autogalaxy.analysis.preloads import Preloads
 from autogalaxy.galaxy.galaxy import Galaxy
 from autogalaxy.galaxy.galaxies import Galaxies
 from autogalaxy.galaxy.to_inversion import GalaxiesToInversion
@@ -21,7 +20,6 @@ def __init__(
         dataset_model: Optional[aa.DatasetModel] = None,
         adapt_images: Optional[AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
-        preloads: aa.Preloads = Preloads(),
         run_time_dict: Optional[Dict] = None,
     ):
         """
@@ -62,9 +60,6 @@ def __init__(
             reconstructed galaxy's morphology.
         settings_inversion
             Settings controlling how an inversion is fitted for example which linear algebra formalism is used.
-        preloads
-            Contains preloaded calculations (e.g. linear algebra matrices) which can skip certain calculations in
-            the fit.
         run_time_dict
             A dictionary which if passed to the fit records how long fucntion calls which have the `profile_func`
             decorator take to run.
@@ -92,8 +87,6 @@ def __init__(
         self.adapt_images = adapt_images
         self.settings_inversion = settings_inversion
 
-        self.preloads = preloads
-
     @property
     def profile_visibilities(self) -> aa.Visibilities:
         """
@@ -126,7 +119,6 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion:
             galaxies=self.galaxies,
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
-            preloads=self.preloads,
             run_time_dict=self.run_time_dict,
         )
 
@@ -223,44 +215,3 @@ def galaxies_linear_light_profiles_to_light_profiles(self) -> List[Galaxy]:
         or `GalaxyPlotter` objects.
         """
         return self.model_obj_linear_light_profiles_to_light_profiles
-
-    def refit_with_new_preloads(
-        self,
-        preloads: Preloads,
-        settings_inversion: Optional[aa.SettingsInversion] = None,
-    ) -> "FitInterferometer":
-        """
-        Returns a new fit which uses the dataset, galaxies and other objects of this fit, but uses a different set of
-        preloads input into this function.
-
-        This is used when setting up the preloads objects, to concisely test how using different preloads objects
-        changes the attributes of the fit.
-
-        Parameters
-        ----------
-        preloads
-            The new preloads which are used to refit the data using the
-        settings_inversion
-            Settings controlling how an inversion is fitted for example which linear algebra formalism is used.
-
-        Returns
-        -------
-        A new fit which has used new preloads input into this function but the same dataset, galaxies and other settings.
-        """
-        if self.run_time_dict is not None:
-            run_time_dict = {}
-        else:
-            run_time_dict = None
-
-        if settings_inversion is None:
-            settings_inversion = self.settings_inversion
-
-        return FitInterferometer(
-            dataset=self.dataset,
-            galaxies=self.galaxies,
-            dataset_model=self.dataset_model,
-            adapt_images=self.adapt_images,
-            settings_inversion=settings_inversion,
-            preloads=preloads,
-            run_time_dict=run_time_dict,
-        )
diff --git a/autogalaxy/interferometer/model/analysis.py b/autogalaxy/interferometer/model/analysis.py
index 735a73fab..7ee73d1f8 100644
--- a/autogalaxy/interferometer/model/analysis.py
+++ b/autogalaxy/interferometer/model/analysis.py
@@ -11,7 +11,6 @@
 
 from autogalaxy.analysis.adapt_images.adapt_image_maker import AdaptImageMaker
 from autogalaxy.analysis.analysis.dataset import AnalysisDataset
-from autogalaxy.analysis.preloads import Preloads
 from autogalaxy.cosmology.lensing import LensingCosmology
 from autogalaxy.cosmology.wrap import Planck15
 from autogalaxy.interferometer.model.result import ResultInterferometer
@@ -90,22 +89,15 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
         model
-            The model object, which includes model components representing the galaxies that are fitted to
+            The PyAutoFit model object, which includes model components representing the galaxies that are fitted to
             the imaging data.
         """
 
         super().modify_before_fit(paths=paths, model=model)
 
-        if not paths.is_complete:
-            logger.info(
-                "PRELOADS - Setting up preloads, may take a few minutes for fits using an inversion."
-            )
-
-            self.set_preloads(paths=paths, model=model)
-
         return self
 
     def log_likelihood_function(self, instance: af.ModelInstance) -> float:
@@ -161,7 +153,6 @@ def log_likelihood_function(self, instance: af.ModelInstance) -> float:
     def fit_from(
         self,
         instance: af.ModelInstance,
-        preload_overwrite: Optional[Preloads] = None,
         run_time_dict: Optional[Dict] = None,
     ) -> FitInterferometer:
         """
@@ -191,14 +182,11 @@ def fit_from(
 
         adapt_images = self.adapt_images_via_instance_from(instance=instance)
 
-        preloads = self.preloads if preload_overwrite is None else preload_overwrite
-
         return FitInterferometer(
             dataset=self.dataset,
             galaxies=galaxies,
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
-            preloads=preloads,
             run_time_dict=run_time_dict,
         )
 
@@ -229,7 +217,7 @@ def save_attributes(self, paths: af.DirectoryPaths):
          Parameters
          ----------
          paths
-             The paths object which manages all paths, e.g. where the non-linear search outputs are stored, visualization,
+             The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored, visualization,
              and the pickled objects used by the aggregator output by this function.
         """
         super().save_attributes(paths=paths)
@@ -259,7 +247,7 @@ def profile_log_likelihood_function(
         An `info_dict` is also created which stores information on aspects of the model and dataset that dictate
         run times, so the profiled times can be interpreted with this context.
 
-        The results of this profiling are then output to hard-disk in the `preloads` folder of the model-fit results,
+        The results of this profiling are then output to hard-disk in the `profiling` folder of the model-fit results,
         which they can be inspected to ensure run-times are as expected.
 
         Parameters
@@ -268,7 +256,7 @@ def profile_log_likelihood_function(
             An instance of the model that is being fitted to the data by this analysis (whose parameters have been set
             via a non-linear search).
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
 
         Returns
diff --git a/autogalaxy/interferometer/model/plotter_interface.py b/autogalaxy/interferometer/model/plotter_interface.py
index 335dbe25a..5686415ac 100644
--- a/autogalaxy/interferometer/model/plotter_interface.py
+++ b/autogalaxy/interferometer/model/plotter_interface.py
@@ -18,7 +18,7 @@ def interferometer(self, dataset: aa.Interferometer):
         Visualizes an `Interferometer` dataset object.
 
         Images are output to the `image` folder of the `image_path` in a subfolder called `interferometer`. When
-        used with a non-linear search the `image_path` is the output folder of the non-linear search.
+        used with a non-linear search the `image_path` points to the search's results folder.
 
         Visualization includes individual images of attributes of the dataset (e.g. the visibilities, noise map,
         uv-wavelengths) and a subplot of all these attributes on the same figure.
diff --git a/autogalaxy/interferometer/model/visualizer.py b/autogalaxy/interferometer/model/visualizer.py
index 03e5d4655..ca3b5bc1d 100644
--- a/autogalaxy/interferometer/model/visualizer.py
+++ b/autogalaxy/interferometer/model/visualizer.py
@@ -21,10 +21,10 @@ def visualize_before_fit(
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
         model
-            The model object, which includes model components representing the galaxies that are fitted to
+            The PyAutoFit model object, which includes model components representing the galaxies that are fitted to
             the imaging data.
         """
 
@@ -67,7 +67,7 @@ def visualize(
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization, and the pickled objects used by the aggregator output by this function.
         instance
             An instance of the model that is being fitted to the data by this analysis (whose parameters have been set
diff --git a/autogalaxy/operate/deflections.py b/autogalaxy/operate/deflections.py
index 28f01958c..57af22765 100644
--- a/autogalaxy/operate/deflections.py
+++ b/autogalaxy/operate/deflections.py
@@ -662,9 +662,7 @@ def einstein_mass_angular_list_from(
         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
-        ]
+        return [np.pi * einstein_radius**2 for einstein_radius in einstein_radius_list]
 
     @evaluation_grid
     def einstein_mass_angular_from(
diff --git a/autogalaxy/profiles/basis.py b/autogalaxy/profiles/basis.py
index 63747b22c..f71777a00 100644
--- a/autogalaxy/profiles/basis.py
+++ b/autogalaxy/profiles/basis.py
@@ -125,9 +125,11 @@ def image_2d_list_from(
         The image of the light profiles in the basis summed together.
         """
         return [
-            light_profile.image_2d_from(grid=grid, operated_only=operated_only)
-            if not isinstance(light_profile, lp_linear.LightProfileLinear)
-            else np.zeros((grid.shape[0],))
+            (
+                light_profile.image_2d_from(grid=grid, operated_only=operated_only)
+                if not isinstance(light_profile, lp_linear.LightProfileLinear)
+                else np.zeros((grid.shape[0],))
+            )
             for light_profile in self.light_profile_list
         ]
 
diff --git a/autogalaxy/profiles/geometry_profiles.py b/autogalaxy/profiles/geometry_profiles.py
index c75b8cdbf..64fb0c453 100644
--- a/autogalaxy/profiles/geometry_profiles.py
+++ b/autogalaxy/profiles/geometry_profiles.py
@@ -1,6 +1,11 @@
+import os
+
 from typing import Optional, Tuple, Type
 
-import numpy as np
+if os.environ.get("USE_JAX", "0") == "1":
+    import jax.numpy as np
+else:
+    import numpy as np
 
 import autoarray as aa
 
@@ -86,7 +91,7 @@ def __init__(self, centre: Tuple[float, float] = (0.0, 0.0)):
     def radial_grid_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
         """
         Convert a grid of (y, x) coordinates, to their radial distances from the profile
-        centre (e.g. :math: r = x**2 + y**2).
+        centre (e.g. :math: r = sqrt(x**2 + y**2)).
 
         Parameters
         ----------
@@ -306,7 +311,8 @@ def elliptical_radii_grid_from(
         """
         return np.sqrt(
             np.add(
-                np.square(grid[:, 1]), np.square(np.divide(grid[:, 0], self.axis_ratio))
+                np.square(grid.array[:, 1]),
+                np.square(np.divide(grid.array[:, 0], self.axis_ratio)),
             )
         )
 
@@ -329,9 +335,9 @@ 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, **kwargs).array
 
-        return np.multiply(np.sqrt(self.axis_ratio), grid_radii).view(np.ndarray)
+        return np.multiply(np.sqrt(self.axis_ratio), grid_radii)  # .view(np.ndarray)
 
     @aa.grid_dec.to_grid
     def transformed_to_reference_frame_grid_from(
diff --git a/autogalaxy/profiles/light/standard/sersic.py b/autogalaxy/profiles/light/standard/sersic.py
index 675d052b8..7ac96daed 100644
--- a/autogalaxy/profiles/light/standard/sersic.py
+++ b/autogalaxy/profiles/light/standard/sersic.py
@@ -1,4 +1,5 @@
-import numpy as np
+from autofit.jax_wrapper import numpy as np
+from numpy import seterr
 from typing import Optional, Tuple
 
 import autoarray as aa
@@ -127,7 +128,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarra
         grid_radii
             The radial distances from the centre of the profile, for each coordinate on the grid.
         """
-        np.seterr(all="ignore")
+        seterr(all="ignore")
         return np.multiply(
             self._intensity,
             np.exp(
@@ -135,7 +136,7 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarra
                     -self.sersic_constant,
                     np.add(
                         np.power(
-                            np.divide(grid_radii, self.effective_radius),
+                            np.divide(np.array(grid_radii), self.effective_radius),
                             1.0 / self.sersic_index,
                         ),
                         -1,
diff --git a/autogalaxy/profiles/mass/abstract/jax_utils.py b/autogalaxy/profiles/mass/abstract/jax_utils.py
new file mode 100644
index 000000000..dd31478cd
--- /dev/null
+++ b/autogalaxy/profiles/mass/abstract/jax_utils.py
@@ -0,0 +1,117 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from jax import custom_jvp
+from functools import partial
+
+
+def reg1(z, _, i_sqrt_pi):
+    return i_sqrt_pi / z
+
+
+def reg2(z, _, i_sqrt_pi):
+    z2 = z**2
+    return i_sqrt_pi * z / (z2 - 0.5)
+
+
+def reg3(z, _, i_sqrt_pi):
+    z2 = z**2
+    return (i_sqrt_pi / z) * (1 + 0.5 / (z2 - 1.5))
+
+
+def reg4(z, _, i_sqrt_pi):
+    z2 = z**2
+    return (i_sqrt_pi * z) * (z2 - 2.5) / (z2 * (z2 - 3.0) + 0.75)
+
+
+def reg5(z, sqrt_pi, _):
+    mz2 = -(z**2)
+    f1 = sqrt_pi
+    f2 = 1.0
+    s1 = [1.320522, 35.7668, 219.031, 1540.787, 3321.99, 36183.31]
+    s2 = [1.841439, 61.57037, 364.2191, 2186.181, 9022.228, 24322.84, 32066.6]
+
+    for s in s1:
+        f1 = s - f1 * mz2
+    for s in s2:
+        f2 = s - f2 * mz2
+
+    return jnp.exp(mz2) + 1j * z * f1 / f2
+
+
+def reg6(z, sqrt_pi, _):
+    miz = -1j * z
+    f1 = sqrt_pi
+    f2 = 1
+    s1 = [5.9126262, 30.180142, 93.15558, 181.92853, 214.38239, 122.60793]
+    s2 = [10.479857, 53.992907, 170.35400, 348.70392, 457.33448, 352.73063, 122.60793]
+
+    for s in s1:
+        f1 = s + f1 * miz
+    for s in s2:
+        f2 = s + f2 * miz
+
+    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``
+
+    # This function is a JAX conversion of
+    # "https://github.com/sibirrer/lenstronomy/tree/master/lenstronomy/LensModel/Profiles"
+    # original function written by Anowar J. Shajib (see 1906.08263)
+    # JAX conversion written by Coleman M. Krawczyk
+    """
+    sqrt_pi = 1 / jnp.sqrt(jnp.pi)
+    i_sqrt_pi = 1j * sqrt_pi
+
+    z_imag2 = z.imag**2
+    abs_z2 = z.real**2 + z_imag2
+
+    r1 = abs_z2 >= 38000.0
+    r2 = (abs_z2 >= 256.0) & (abs_z2 < 38000.0)
+    r3 = (abs_z2 >= 62.0) & (abs_z2 < 256.0)
+    r4 = (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
+    r5_1 = (abs_z2 >= 30.0) & (abs_z2 < 62.0) & (z_imag2 < 1e-13)
+    r5_2 = (abs_z2 >= 2.5) & (abs_z2 < 30.0) & (z_imag2 < 0.072)
+    r5 = r5_1 | r5_2
+    r6 = (abs_z2 < 30.0) & jnp.logical_not(r5)
+
+    args = (z, 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 = jnp.where(r4, reg4(*args), wz)
+    wz = jnp.where(r5, reg5(*args), wz)
+    wz = jnp.where(r6, reg6(*args), 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`
+    # also the derivative is defined analytically for this function so bypass
+    # auto diffing over the complex functions above.
+    (z,) = primals
+    (z_dot,) = tangents
+    primal_out = w_f_approx(z)
+    i_sqrt_pi = 1j / jnp.sqrt(jnp.pi)
+    tangent_out = z_dot * 2 * (i_sqrt_pi - z * primal_out)
+    return primal_out, tangent_out
+
+
+@partial(jax.jit, static_argnums=(0,))
+def all_comb(n):
+    i = jnp.arange(1, n, 1)
+    return jnp.cumprod((n + 1 - i) / i)
diff --git a/autogalaxy/profiles/mass/abstract/mge.py b/autogalaxy/profiles/mass/abstract/mge.py
index 3f534ff91..1958eb3ef 100644
--- a/autogalaxy/profiles/mass/abstract/mge.py
+++ b/autogalaxy/profiles/mass/abstract/mge.py
@@ -1,280 +1,8 @@
-import numpy as np
-from scipy.special import comb
-
-
-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``
-
-    # This function is copied from
-    # "https://github.com/sibirrer/lenstronomy/tree/main/lenstronomy/LensModel/Profiles"
-    # written by Anowar J. Shajib (see 1906.08263)
-    """
-
-    reg_minus_imag = z.imag < 0.0
-    z[reg_minus_imag] = np.conj(z[reg_minus_imag])
-
-    sqrt_pi = 1 / np.sqrt(np.pi)
-    i_sqrt_pi = 1j * sqrt_pi
-
-    wz = np.empty_like(z)
-
-    z_imag2 = z.imag**2
-    abs_z2 = z.real**2 + z_imag2
-
-    reg1 = abs_z2 >= 38000.0
-    if np.any(reg1):
-        wz[reg1] = i_sqrt_pi / z[reg1]
-
-    reg2 = (256.0 <= abs_z2) & (abs_z2 < 38000.0)
-    if np.any(reg2):
-        t = z[reg2]
-        wz[reg2] = i_sqrt_pi * t / (t * t - 0.5)
-
-    reg3 = (62.0 <= abs_z2) & (abs_z2 < 256.0)
-    if np.any(reg3):
-        t = z[reg3]
-        wz[reg3] = (i_sqrt_pi / t) * (1 + 0.5 / (t * t - 1.5))
-
-    reg4 = (30.0 <= abs_z2) & (abs_z2 < 62.0) & (z_imag2 >= 1e-13)
-    if np.any(reg4):
-        t = z[reg4]
-        tt = t * t
-        wz[reg4] = (i_sqrt_pi * t) * (tt - 2.5) / (tt * (tt - 3.0) + 0.75)
-
-    reg5 = (62.0 > abs_z2) & np.logical_not(reg4) & (abs_z2 > 2.5) & (z_imag2 < 0.072)
-    if np.any(reg5):
-        t = z[reg5]
-        u = -t * t
-        f1 = sqrt_pi
-        f2 = 1
-        s1 = [1.320522, 35.7668, 219.031, 1540.787, 3321.99, 36183.31]
-        s2 = [1.841439, 61.57037, 364.2191, 2186.181, 9022.228, 24322.84, 32066.6]
-
-        for s in s1:
-            f1 = s - f1 * u
-        for s in s2:
-            f2 = s - f2 * u
-
-        wz[reg5] = np.exp(u) + 1j * t * f1 / f2
-
-    reg6 = (30.0 > abs_z2) & np.logical_not(reg5)
-    if np.any(reg6):
-        t3 = -1j * z[reg6]
-
-        f1 = sqrt_pi
-        f2 = 1
-        s1 = [5.9126262, 30.180142, 93.15558, 181.92853, 214.38239, 122.60793]
-        s2 = [
-            10.479857,
-            53.992907,
-            170.35400,
-            348.70392,
-            457.33448,
-            352.73063,
-            122.60793,
-        ]
-
-        for s in s1:
-            f1 = f1 * t3 + s
-        for s in s2:
-            f2 = f2 * t3 + s
-
-        wz[reg6] = f1 / f2
-
-    # wz[reg_minus_imag] = np.conj(wz[reg_minus_imag])
-
-    return wz
-
-
-class MassProfileMGE:
-    """
-    This class speeds up deflection angle calculations of certain mass profiles by decompositing them into many
-    Gaussians.
-
-    This follows the method of Shajib 2019 - https://academic.oup.com/mnras/article/488/1/1387/5526256
-    """
-
-    def __init__(self):
-        self.count = 0
-        self.sigma_calc = 0
-        self.z = 0
-        self.zq = 0
-        self.expv = 0
-
-    @staticmethod
-    #  @aa.util.numba.jit()
-    def zeta_from(grid, amps, sigmas, axis_ratio):
-        """
-        The key part to compute the deflection angle of each Gaussian.
-        Because of my optimization, there are blocks looking weird and indirect. What I'm doing here
-        is trying to avoid big matrix operation to save time.
-        I think there are still spaces we can optimize.
-
-        It seems when using w_f_approx, it gives some errors if y < 0. So when computing for places
-        where y < 0, we first compute the value at - y, and then change its sign.
-        """
-
-        output_grid_final = np.zeros(grid.shape[0], dtype="complex128")
-
-        q2 = axis_ratio**2.0
-
-        scale_factor = axis_ratio / (sigmas[0] * np.sqrt(2.0 * (1.0 - q2)))
-
-        xs = np.array((grid[:, 1] * scale_factor).copy())
-        ys = np.array((grid[:, 0] * scale_factor).copy())
-
-        ys_minus = ys < 0.0
-        ys[ys_minus] *= -1
-        z = xs + 1j * ys
-        zq = axis_ratio * xs + 1j * ys / axis_ratio
-
-        expv = -(xs**2.0) * (1.0 - q2) - ys**2.0 * (1.0 / q2 - 1.0)
-
-        for i in range(len(sigmas)):
-            if i > 0:
-                z /= sigmas[i] / sigmas[i - 1]
-                zq /= sigmas[i] / sigmas[i - 1]
-                expv /= (sigmas[i] / sigmas[i - 1]) ** 2.0
-
-            output_grid = -1j * (w_f_approx(z) - np.exp(expv) * w_f_approx(zq))
-
-            output_grid[ys_minus] = np.conj(output_grid[ys_minus])
-
-            output_grid_final += (amps[i] * sigmas[i]) * output_grid
-
-        return output_grid_final
-
-    @staticmethod
-    def kesi(p):
-        """
-        see Eq.(6) of 1906.08263
-        """
-        n_list = np.arange(0, 2 * p + 1, 1)
-        return (2.0 * p * np.log(10) / 3.0 + 2.0 * np.pi * n_list * 1j) ** (0.5)
-
-    @staticmethod
-    def eta(p):
-        """
-        see Eq.(6) of 1906.00263
-        """
-        eta_list = np.zeros(int(2 * p + 1))
-        kesi_list = np.zeros(int(2 * p + 1))
-        kesi_list[0] = 0.5
-        kesi_list[1 : p + 1] = 1.0
-        kesi_list[int(2 * p)] = 1.0 / 2.0**p
-
-        for i in np.arange(1, p, 1):
-            kesi_list[2 * p - i] = kesi_list[2 * p - i + 1] + 2 ** (-p) * comb(p, i)
-
-        for i in np.arange(0, 2 * p + 1, 1):
-            eta_list[i] = (
-                (-1) ** i * 2.0 * np.sqrt(2.0 * np.pi) * 10 ** (p / 3.0) * kesi_list[i]
-            )
-
-        return eta_list
-
-    def decompose_convergence_via_mge(self):
-        raise NotImplementedError()
-
-    def _decompose_convergence_via_mge(
-        self, func, radii_min, radii_max, func_terms=28, func_gaussians=20
-    ):
-        """
-
-        Parameters
-        ----------
-        func : func
-            The function representing the profile that is decomposed into Gaussians.
-        normalization
-            A normalization factor tyh
-        func_terms
-            The number of terms used to approximate the input func.
-        func_gaussians
-            The number of Gaussians used to represent the input func.
-
-        Returns
-        -------
-        """
-
-        kesis = self.kesi(func_terms)  # kesi in Eq.(6) of 1906.08263
-        etas = self.eta(func_terms)  # eta in Eqr.(6) of 1906.08263
-
-        def f(sigma):
-            """Eq.(5) of 1906.08263"""
-            return np.sum(etas * np.real(target_function(sigma * kesis)))
-
-        # sigma is sampled from logspace between these radii.
-
-        log_sigmas = np.linspace(np.log(radii_min), np.log(radii_max), func_gaussians)
-        d_log_sigma = log_sigmas[1] - log_sigmas[0]
-        sigma_list = np.exp(log_sigmas)
-
-        amplitude_list = np.zeros(func_gaussians)
-
-        for i in range(func_gaussians):
-            f_sigma = np.sum(etas * np.real(func(sigma_list[i] * kesis)))
-            if (i == -1) or (i == (func_gaussians - 1)):
-                amplitude_list[i] = 0.5 * f_sigma * d_log_sigma / np.sqrt(2.0 * np.pi)
-            else:
-                amplitude_list[i] = f_sigma * d_log_sigma / np.sqrt(2.0 * np.pi)
-
-        return amplitude_list, sigma_list
-
-    def convergence_2d_via_mge_from(self, grid_radii):
-        raise NotImplementedError()
-
-    def _convergence_2d_via_mge_from(self, grid_radii):
-        """Calculate the projected convergence at a given set of arc-second gridded coordinates.
-
-        Parameters
-        ----------
-        grid
-            The grid of (y,x) arc-second coordinates the convergence is computed on.
-
-        """
-
-        self.count = 0
-        self.sigma_calc = 0
-        self.z = 0
-        self.zq = 0
-        self.expv = 0
-
-        amps, sigmas = self.decompose_convergence_via_mge()
-
-        convergence = 0.0
-
-        for i in range(len(sigmas)):
-            convergence += self.convergence_func_gaussian(
-                grid_radii=grid_radii, sigma=sigmas[i], intensity=amps[i]
-            )
-        return convergence
-
-    def convergence_func_gaussian(self, grid_radii, sigma, intensity):
-        return np.multiply(
-            intensity, np.exp(-0.5 * np.square(np.divide(grid_radii, sigma)))
-        )
-
-    def _deflections_2d_via_mge_from(self, grid, sigmas_factor=1.0):
-        axis_ratio = self.axis_ratio
-
-        if axis_ratio > 0.9999:
-            axis_ratio = 0.9999
-
-        amps, sigmas = self.decompose_convergence_via_mge()
-        sigmas *= sigmas_factor
-
-        angle = self.zeta_from(
-            grid=grid, amps=amps, sigmas=sigmas, axis_ratio=axis_ratio
-        )
-
-        angle *= np.sqrt((2.0 * np.pi) / (1.0 - axis_ratio**2.0))
-
-        return self.rotated_grid_from_reference_frame_from(
-            np.vstack((-angle.imag, angle.real)).T
-        )
+import os
+
+if os.environ.get("USE_JAX", "0") == "1":
+    USING_JAX = True
+    from .mge_jax import MassProfileMGE
+else:
+    USING_JAX = False
+    from .mge_numpy import MassProfileMGE
diff --git a/autogalaxy/profiles/mass/abstract/mge_jax.py b/autogalaxy/profiles/mass/abstract/mge_jax.py
new file mode 100644
index 000000000..2fbfd5ed7
--- /dev/null
+++ b/autogalaxy/profiles/mass/abstract/mge_jax.py
@@ -0,0 +1,162 @@
+import jax.numpy as np
+import jax
+import numpy as base_np
+
+from .jax_utils import w_f_approx, all_comb
+
+
+class MassProfileMGE:
+    """
+    This class speeds up deflection angle calculations of certain mass profiles by decompositing them into many
+    Gaussians.
+
+    This follows the method of Shajib 2019 - https://academic.oup.com/mnras/article/488/1/1387/5526256
+    """
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def zeta_from(grid, amps, sigmas, axis_ratio):
+        """
+        The key part to compute the deflection angle of each Gaussian.
+        """
+        q2 = axis_ratio**2.0
+
+        scale_factor = axis_ratio / np.sqrt(2.0 * (1.0 - q2))
+
+        xs = np.array((grid.array[:, 1] * scale_factor).copy())
+        ys = np.array((grid.array[:, 0] * scale_factor).copy())
+
+        y_sign = np.sign(ys)
+        ys = ys * y_sign
+
+        z = xs + 1j * ys
+        zq = axis_ratio * xs + 1j * ys / axis_ratio
+        expv = -(xs**2.0) * (1.0 - q2) - ys**2.0 * (1.0 / q2 - 1.0)
+        sigma_ = sigmas.reshape((-1,) + (1,) * xs.ndim)
+        inv_sigma_ = 1 / sigma_
+        amps_ = amps.reshape((-1,) + (1,) * xs.ndim)
+
+        # process as one big vectorized calculation
+        # 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 = np.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
+
+        output_grid = sigma_ * amps_ * (sigma_func_real + 1j * sigma_func_imag)
+        return output_grid.sum(axis=0)
+
+    @staticmethod
+    def kesi(p):
+        """
+        see Eq.(6) of 1906.08263
+        """
+        n_list = np.arange(0, 2 * p + 1, 1)
+        return (2.0 * p * np.log(10) / 3.0 + 2.0 * np.pi * n_list * 1j) ** (0.5)
+
+    @staticmethod
+    def eta(p):
+        """
+        see Eq.(6) of 1906.00263
+        """
+
+        i = np.arange(1, p, 1)
+        kesi_last = 1 / 2**p
+        k = kesi_last + np.cumsum(np.cumprod((p + 1 - i) / i) * kesi_last)
+
+        kesi_list = np.hstack(
+            [np.array([0.5]), np.ones(p), k[::-1], np.array([kesi_last])]
+        )
+        coef = (-1) ** np.arange(0, 2 * p + 1, 1)
+        eta_list = coef * 2.0 * np.sqrt(2.0 * np.pi) * 10 ** (p / 3.0) * kesi_list
+        return eta_list
+
+    def decompose_convergence_via_mge(self):
+        raise NotImplementedError()
+
+    def _decompose_convergence_via_mge(
+        self, func, radii_min, radii_max, func_terms=28, func_gaussians=20
+    ):
+        """
+
+        Parameters
+        ----------
+        func : func
+            The function representing the profile that is decomposed into Gaussians.
+        normalization
+            A normalization factor tyh
+        func_terms
+            The number of terms used to approximate the input func.
+        func_gaussians
+            The number of Gaussians used to represent the input func.
+
+        Returns
+        -------
+        """
+
+        kesis = self.kesi(func_terms)  # kesi in Eq.(6) of 1906.08263
+        etas = self.eta(func_terms)  # eta in Eqr.(6) of 1906.08263
+
+        # sigma is sampled from logspace between these radii.
+
+        log_sigmas = np.linspace(np.log(radii_min), np.log(radii_max), func_gaussians)
+        d_log_sigma = log_sigmas[1] - log_sigmas[0]
+        sigma_list = np.exp(log_sigmas)
+
+        amplitude_list = np.zeros(func_gaussians)
+        f_sigma = np.sum(
+            etas * np.real(func(sigma_list.reshape(-1, 1) * kesis)), axis=1
+        )
+        amplitude_list = f_sigma * d_log_sigma / np.sqrt(2.0 * np.pi)
+        amplitude_list = amplitude_list.at[0].multiply(0.5)
+        amplitude_list = amplitude_list.at[-1].multiply(0.5)
+        return amplitude_list, sigma_list
+
+    def convergence_2d_via_mge_from(self, grid_radii):
+        raise NotImplementedError()
+
+    def _convergence_2d_via_mge_from(
+        self, grid_radii, func_terms=28, func_gaussians=20
+    ):
+        """Calculate the projected convergence at a given set of arc-second gridded coordinates.
+
+        Parameters
+        ----------
+        grid
+            The grid of (y,x) arc-second coordinates the convergence is computed on.
+
+        """
+        amps, sigmas = self.decompose_convergence_via_mge(
+            func_terms=func_terms, func_gaussians=func_gaussians
+        )
+
+        convergence = 0.0
+
+        inv_sigma_ = 1 / sigmas.reshape((-1,) + (1,) * grid_radii.array.ndim)
+        amps_ = amps.reshape((-1,) + (1,) * grid_radii.array.ndim)
+        convergence = amps_ * np.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 = np.min(np.array([self.axis_ratio, 0.9999]))
+
+        amps, sigmas = self.decompose_convergence_via_mge(
+            func_terms=func_terms, func_gaussians=func_gaussians
+        )
+        sigmas *= sigmas_factor
+
+        angle = self.zeta_from(
+            grid=grid, amps=amps, sigmas=sigmas, axis_ratio=axis_ratio
+        )
+
+        angle *= np.sqrt((2.0 * np.pi) / (1.0 - axis_ratio**2.0))
+
+        return self.rotated_grid_from_reference_frame_from(
+            np.vstack((-angle.imag, angle.real)).T
+        )
diff --git a/autogalaxy/profiles/mass/abstract/mge_numpy.py b/autogalaxy/profiles/mass/abstract/mge_numpy.py
new file mode 100644
index 000000000..08fd0c50d
--- /dev/null
+++ b/autogalaxy/profiles/mass/abstract/mge_numpy.py
@@ -0,0 +1,280 @@
+import numpy as np
+from scipy.special import comb
+
+
+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``
+
+    # This function is copied from
+    # "https://github.com/sibirrer/lenstronomy/tree/main/lenstronomy/LensModel/Profiles"
+    # written by Anowar J. Shajib (see 1906.08263)
+    """
+
+    reg_minus_imag = z.imag < 0.0
+    z[reg_minus_imag] = np.conj(z[reg_minus_imag])
+
+    sqrt_pi = 1 / np.sqrt(np.pi)
+    i_sqrt_pi = 1j * sqrt_pi
+
+    wz = np.empty_like(z)
+
+    z_imag2 = z.imag**2
+    abs_z2 = z.real**2 + z_imag2
+
+    reg1 = abs_z2 >= 38000.0
+    if np.any(reg1):
+        wz[reg1] = i_sqrt_pi / z[reg1]
+
+    reg2 = (256.0 <= abs_z2) & (abs_z2 < 38000.0)
+    if np.any(reg2):
+        t = z[reg2]
+        wz[reg2] = i_sqrt_pi * t / (t * t - 0.5)
+
+    reg3 = (62.0 <= abs_z2) & (abs_z2 < 256.0)
+    if np.any(reg3):
+        t = z[reg3]
+        wz[reg3] = (i_sqrt_pi / t) * (1 + 0.5 / (t * t - 1.5))
+
+    reg4 = (30.0 <= abs_z2) & (abs_z2 < 62.0) & (z_imag2 >= 1e-13)
+    if np.any(reg4):
+        t = z[reg4]
+        tt = t * t
+        wz[reg4] = (i_sqrt_pi * t) * (tt - 2.5) / (tt * (tt - 3.0) + 0.75)
+
+    reg5 = (62.0 > abs_z2) & np.logical_not(reg4) & (abs_z2 > 2.5) & (z_imag2 < 0.072)
+    if np.any(reg5):
+        t = z[reg5]
+        u = -t * t
+        f1 = sqrt_pi
+        f2 = 1
+        s1 = [1.320522, 35.7668, 219.031, 1540.787, 3321.99, 36183.31]
+        s2 = [1.841439, 61.57037, 364.2191, 2186.181, 9022.228, 24322.84, 32066.6]
+
+        for s in s1:
+            f1 = s - f1 * u
+        for s in s2:
+            f2 = s - f2 * u
+
+        wz[reg5] = np.exp(u) + 1j * t * f1 / f2
+
+    reg6 = (30.0 > abs_z2) & np.logical_not(reg5)
+    if np.any(reg6):
+        t3 = -1j * z[reg6]
+
+        f1 = sqrt_pi
+        f2 = 1
+        s1 = [5.9126262, 30.180142, 93.15558, 181.92853, 214.38239, 122.60793]
+        s2 = [
+            10.479857,
+            53.992907,
+            170.35400,
+            348.70392,
+            457.33448,
+            352.73063,
+            122.60793,
+        ]
+
+        for s in s1:
+            f1 = f1 * t3 + s
+        for s in s2:
+            f2 = f2 * t3 + s
+
+        wz[reg6] = f1 / f2
+
+    # wz[reg_minus_imag] = np.conj(wz[reg_minus_imag])
+
+    return wz
+
+
+class MassProfileMGE:
+    """
+    This class speeds up deflection angle calculations of certain mass profiles by decompositing them into many
+    Gaussians.
+
+    This follows the method of Shajib 2019 - https://academic.oup.com/mnras/article/488/1/1387/5526256
+    """
+
+    def __init__(self):
+        self.count = 0
+        self.sigma_calc = 0
+        self.z = 0
+        self.zq = 0
+        self.expv = 0
+
+    @staticmethod
+    #  @aa.util.numba.jit()
+    def zeta_from(grid, amps, sigmas, axis_ratio):
+        """
+        The key part to compute the deflection angle of each Gaussian.
+        Because of my optimization, there are blocks looking weird and indirect. What I'm doing here
+        is trying to avoid big matrix operation to save time.
+        I think there are still spaces we can optimize.
+
+        It seems when using w_f_approx, it gives some errors if y < 0. So when computing for places
+        where y < 0, we first compute the value at - y, and then change its sign.
+        """
+
+        output_grid_final = np.zeros(grid.shape[0], dtype="complex128")
+
+        q2 = axis_ratio**2.0
+
+        scale_factor = axis_ratio / (sigmas[0] * np.sqrt(2.0 * (1.0 - q2)))
+
+        xs = np.array((grid[:, 1] * scale_factor).copy())
+        ys = np.array((grid[:, 0] * scale_factor).copy())
+
+        ys_minus = ys < 0.0
+        ys[ys_minus] *= -1
+        z = xs + 1j * ys
+        zq = axis_ratio * xs + 1j * ys / axis_ratio
+
+        expv = -(xs**2.0) * (1.0 - q2) - ys**2.0 * (1.0 / q2 - 1.0)
+
+        for i in range(len(sigmas)):
+            if i > 0:
+                z /= sigmas[i] / sigmas[i - 1]
+                zq /= sigmas[i] / sigmas[i - 1]
+                expv /= (sigmas[i] / sigmas[i - 1]) ** 2.0
+
+            output_grid = -1j * (w_f_approx(z) - np.exp(expv) * w_f_approx(zq))
+
+            output_grid[ys_minus] = np.conj(output_grid[ys_minus])
+
+            output_grid_final += (amps[i] * sigmas[i]) * output_grid
+
+        return output_grid_final
+
+    @staticmethod
+    def kesi(p):
+        """
+        see Eq.(6) of 1906.08263
+        """
+        n_list = np.arange(0, 2 * p + 1, 1)
+        return (2.0 * p * np.log(10) / 3.0 + 2.0 * np.pi * n_list * 1j) ** (0.5)
+
+    @staticmethod
+    def eta(p):
+        """
+        see Eq.(6) of 1906.00263
+        """
+        eta_list = np.zeros(int(2 * p + 1))
+        kesi_list = np.zeros(int(2 * p + 1))
+        kesi_list[0] = 0.5
+        kesi_list[1 : p + 1] = 1.0
+        kesi_list[int(2 * p)] = 1.0 / 2.0**p
+
+        for i in np.arange(1, p, 1):
+            kesi_list[2 * p - i] = kesi_list[2 * p - i + 1] + 2 ** (-p) * comb(p, i)
+
+        for i in np.arange(0, 2 * p + 1, 1):
+            eta_list[i] = (
+                (-1) ** i * 2.0 * np.sqrt(2.0 * np.pi) * 10 ** (p / 3.0) * kesi_list[i]
+            )
+
+        return eta_list
+
+    def decompose_convergence_via_mge(self):
+        raise NotImplementedError()
+
+    def _decompose_convergence_via_mge(
+        self, func, radii_min, radii_max, func_terms=28, func_gaussians=20
+    ):
+        """
+
+        Parameters
+        ----------
+        func : func
+            The function representing the profile that is decomposed into Gaussians.
+        normalization
+            A normalization factor tyh
+        func_terms
+            The number of terms used to approximate the input func.
+        func_gaussians
+            The number of Gaussians used to represent the input func.
+
+        Returns
+        -------
+        """
+
+        kesis = self.kesi(func_terms)  # kesi in Eq.(6) of 1906.08263
+        etas = self.eta(func_terms)  # eta in Eqr.(6) of 1906.08263
+
+        def f(sigma):
+            """Eq.(5) of 1906.08263"""
+            return np.sum(etas * np.real(target_function(sigma * kesis)))
+
+        # sigma is sampled from logspace between these radii.
+
+        log_sigmas = np.linspace(np.log(radii_min), np.log(radii_max), func_gaussians)
+        d_log_sigma = log_sigmas[1] - log_sigmas[0]
+        sigma_list = np.exp(log_sigmas)
+
+        amplitude_list = np.zeros(func_gaussians)
+
+        for i in range(func_gaussians):
+            f_sigma = np.sum(etas * np.real(func(sigma_list[i] * kesis)))
+            if (i == 0) or (i == (func_gaussians - 1)):
+                amplitude_list[i] = 0.5 * f_sigma * d_log_sigma / np.sqrt(2.0 * np.pi)
+            else:
+                amplitude_list[i] = f_sigma * d_log_sigma / np.sqrt(2.0 * np.pi)
+
+        return amplitude_list, sigma_list
+
+    def convergence_2d_via_mge_from(self, grid_radii):
+        raise NotImplementedError()
+
+    def _convergence_2d_via_mge_from(self, grid_radii):
+        """Calculate the projected convergence at a given set of arc-second gridded coordinates.
+
+        Parameters
+        ----------
+        grid
+            The grid of (y,x) arc-second coordinates the convergence is computed on.
+
+        """
+
+        self.count = 0
+        self.sigma_calc = 0
+        self.z = 0
+        self.zq = 0
+        self.expv = 0
+
+        amps, sigmas = self.decompose_convergence_via_mge()
+
+        convergence = 0.0
+
+        for i in range(len(sigmas)):
+            convergence += self.convergence_func_gaussian(
+                grid_radii=grid_radii, sigma=sigmas[i], intensity=amps[i]
+            )
+        return convergence
+
+    def convergence_func_gaussian(self, grid_radii, sigma, intensity):
+        return np.multiply(
+            intensity, np.exp(-0.5 * np.square(np.divide(grid_radii, sigma)))
+        )
+
+    def _deflections_2d_via_mge_from(self, grid, sigmas_factor=1.0):
+        axis_ratio = self.axis_ratio
+
+        if axis_ratio > 0.9999:
+            axis_ratio = 0.9999
+
+        amps, sigmas = self.decompose_convergence_via_mge()
+        sigmas *= sigmas_factor
+
+        angle = self.zeta_from(
+            grid=grid, amps=amps, sigmas=sigmas, axis_ratio=axis_ratio
+        )
+
+        angle *= np.sqrt((2.0 * np.pi) / (1.0 - axis_ratio**2.0))
+
+        return self.rotated_grid_from_reference_frame_from(
+            np.vstack((-angle.imag, angle.real)).T
+        )
diff --git a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py
index 9247c9c9d..838b67453 100644
--- a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py
+++ b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py
@@ -43,9 +43,7 @@ def kappa_s_and_scale_radius(
 
     ##############################
     def integrand(r):
-        return (r**2 / r**inner_slope) * (1 + r / scale_radius_kpc) ** (
-            inner_slope - 3
-        )
+        return (r**2 / r**inner_slope) * (1 + r / scale_radius_kpc) ** (inner_slope - 3)
 
     de_c = (
         (overdens / 3.0)
diff --git a/autogalaxy/profiles/mass/dark/nfw_hk24_util.py b/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
index 0c46b4cfa..71bdce62c 100644
--- a/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
+++ b/autogalaxy/profiles/mass/dark/nfw_hk24_util.py
@@ -3,6 +3,7 @@
 
 @author: felixvecchi
 """
+
 import numpy as np
 from astropy.cosmology import Planck15
 
@@ -206,16 +207,11 @@ def g1_g2_from(x1, x2, e, k_s):
 
     # Prefactors for g1
     pre_f0_g1 = (
-        ((x1 - e) ** 2 + x2**2)
-        * ((x1 + e) ** 2 + x2**2)
-        * (x1**2 - x2**2 - e**2)
+        ((x1 - e) ** 2 + x2**2) * ((x1 + e) ** 2 + x2**2) * (x1**2 - x2**2 - e**2)
     )
 
     pre_f1_g1 = (
-        2
-        * e**2
-        * (x1**2 - 1)
-        * ((x1**2 - x2**2 - e**2) ** 2 - 4 * x1**2 * x2**2)
+        2 * e**2 * (x1**2 - 1) * ((x1**2 - x2**2 - e**2) ** 2 - 4 * x1**2 * x2**2)
         - (x1**2 - x2**2 - e**2) ** 3 * (3 + e**2) / 2
         + 6 * x1**2 * x2**2 * (e**2 - 1) * (x1**2 - x2**2 - e**2)
     )
@@ -225,12 +221,7 @@ def g1_g2_from(x1, x2, e, k_s):
         - 8 * e**2 * x1**2 * x2**2
     )
 
-    pre_f3_g1 = (
-        2
-        * x1
-        * x2
-        * ((x1**2 + x2**2 + e**2) ** 2 - 4 * e**2 * (x2**2 + e**2))
-    )
+    pre_f3_g1 = 2 * x1 * x2 * ((x1**2 + x2**2 + e**2) ** 2 - 4 * e**2 * (x2**2 + e**2))
 
     # First component shear
     g1 = full_pre_factor * (
@@ -245,21 +236,12 @@ def g1_g2_from(x1, x2, e, k_s):
         * x2
         * (
             (x1**2 + x2**2 + e**2)
-            * (
-                (5 * e**2 - 3) * x1**2
-                - 3 * (1 + e**2) * x2**2
-                + (5 - 3 * e**2) * e**2
-            )
+            * ((5 * e**2 - 3) * x1**2 - 3 * (1 + e**2) * x2**2 + (5 - 3 * e**2) * e**2)
             - 4 * e**2 * x1**2 * (1 + e**2)
         )
     )
 
-    pre_f2_g2 = (
-        -2
-        * x1
-        * x2
-        * ((x1**2 + x2**2 + e**2) ** 2 - 4 * e**2 * (x2**2 + e**2))
-    )
+    pre_f2_g2 = -2 * x1 * x2 * ((x1**2 + x2**2 + e**2) ** 2 - 4 * e**2 * (x2**2 + e**2))
 
     pre_f3_g2 = -(
         (x1**2 - x2**2 - e**2) * ((x1**2 + x2**2) ** 2 - e**4)
diff --git a/autogalaxy/profiles/mass/stellar/sersic.py b/autogalaxy/profiles/mass/stellar/sersic.py
index 0d0b8ec97..a7608c903 100644
--- a/autogalaxy/profiles/mass/stellar/sersic.py
+++ b/autogalaxy/profiles/mass/stellar/sersic.py
@@ -1,5 +1,13 @@
+import os
+
+if os.environ.get("USE_JAX", "0") == "1":
+    USING_JAX = True
+    import jax.numpy as np
+else:
+    USING_JAX = False
+    import numpy as np
+
 import copy
-import numpy as np
 from scipy.integrate import quad
 from typing import List, Tuple
 
@@ -130,7 +138,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
     @aa.grid_dec.to_vector_yx
     @aa.grid_dec.transform
     @aa.grid_dec.relocate_to_radial_minimum
-    def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def deflections_2d_via_mge_from(
+        self, grid: aa.type.Grid2DLike, func_terms=28, func_gaussians=20, **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.
@@ -145,7 +155,10 @@ def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, **kwargs):
             The grid of (y,x) arc-second coordinates the convergence is computed on.
         """
         return self._deflections_2d_via_mge_from(
-            grid=grid, sigmas_factor=np.sqrt(self.axis_ratio)
+            grid=grid,
+            sigmas_factor=np.sqrt(self.axis_ratio),
+            func_terms=func_terms,
+            func_gaussians=func_gaussians,
         )
 
     @aa.grid_dec.to_vector_yx
@@ -188,7 +201,9 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
     @aa.grid_dec.to_array
     @aa.grid_dec.transform
     @aa.grid_dec.relocate_to_radial_minimum
-    def convergence_2d_via_mge_from(self, grid: aa.type.Grid2DLike, **kwargs):
+    def convergence_2d_via_mge_from(
+        self, grid: aa.type.Grid2DLike, func_terms=28, func_gaussians=20, **kwargs
+    ):
         """
         Calculate the projected convergence at a given set of arc-second gridded coordinates.
 
@@ -201,7 +216,11 @@ def convergence_2d_via_mge_from(self, grid: aa.type.Grid2DLike, **kwargs):
 
         eccentric_radii = self.eccentric_radii_grid_from(grid=grid, **kwargs)
 
-        return self._convergence_2d_via_mge_from(grid_radii=eccentric_radii)
+        return self._convergence_2d_via_mge_from(
+            grid_radii=eccentric_radii,
+            func_terms=func_terms,
+            func_gaussians=func_gaussians,
+        )
 
     @aa.over_sample
     @aa.grid_dec.to_array
@@ -247,7 +266,9 @@ def image_2d_via_radii_from(self, radius: np.ndarray):
             * (((radius / self.effective_radius) ** (1.0 / self.sersic_index)) - 1)
         )
 
-    def decompose_convergence_via_mge(self) -> Tuple[List, List]:
+    def decompose_convergence_via_mge(
+        self, func_terms=28, func_gaussians=20
+    ) -> Tuple[List, List]:
         radii_min = self.effective_radius / 100.0
         radii_max = self.effective_radius * 20.0
 
@@ -262,7 +283,11 @@ def sersic_2d(r):
             )
 
         return self._decompose_convergence_via_mge(
-            func=sersic_2d, radii_min=radii_min, radii_max=radii_max
+            func=sersic_2d,
+            radii_min=radii_min,
+            radii_max=radii_max,
+            func_terms=func_terms,
+            func_gaussians=func_gaussians,
         )
 
     def decompose_convergence_via_cse(
@@ -374,10 +399,10 @@ 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[:, index]
+            deflection_grid = self.axis_ratio * grid.array[:, index]
 
             for i in range(grid.shape[0]):
-                deflection_grid[i] *= (
+                deflection_grid = deflection_grid.at[i].multiply(
                     self.intensity
                     * self.mass_to_light_ratio
                     * quad(
@@ -385,8 +410,8 @@ def calculate_deflection_component(npow, index):
                         a=0.0,
                         b=1.0,
                         args=(
-                            grid[i, 0],
-                            grid[i, 1],
+                            grid.array[i, 0],
+                            grid.array[i, 1],
                             npow,
                             self.axis_ratio,
                             self.sersic_index,
diff --git a/autogalaxy/profiles/mass/total/isothermal.py b/autogalaxy/profiles/mass/total/isothermal.py
index e9a9ce41e..79f378930 100644
--- a/autogalaxy/profiles/mass/total/isothermal.py
+++ b/autogalaxy/profiles/mass/total/isothermal.py
@@ -1,4 +1,12 @@
-import numpy as np
+import os
+
+if os.environ.get("USE_JAX", "0") == "1":
+    USING_JAX = True
+    import jax.numpy as np
+else:
+    USING_JAX = False
+    import numpy as np
+
 from typing import Tuple
 
 import autoarray as aa
@@ -30,14 +38,21 @@ def psi_from(grid, axis_ratio, core_radius):
         The value of the Psi term.
 
     """
-    return np.sqrt(
-        np.add(
-            np.multiply(
-                axis_ratio**2.0, np.add(np.square(grid[:, 1]), core_radius**2.0)
-            ),
-            np.square(grid[:, 0]),
+    if USING_JAX:
+        return np.sqrt(
+            (axis_ratio**2.0 * (grid[:, 1] ** 2.0 + core_radius**2.0))
+            + grid[:, 0] ** 2.0
+            + 1e-16
+        )
+    else:
+        return np.sqrt(
+            np.add(
+                np.multiply(
+                    axis_ratio**2.0, np.add(np.square(grid[:, 1]), core_radius**2.0)
+                ),
+                np.square(grid[:, 0]),
+            )
         )
-    )
 
 
 class Isothermal(PowerLaw):
diff --git a/autogalaxy/profiles/mass/total/jax_utils.py b/autogalaxy/profiles/mass/total/jax_utils.py
new file mode 100644
index 000000000..5bed36c6c
--- /dev/null
+++ b/autogalaxy/profiles/mass/total/jax_utils.py
@@ -0,0 +1,47 @@
+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"))
+
+
+def body_fun(carry, n, factor, ei2phi, slope):
+    omega_nm1, partial_sum = carry
+    two_n = 2 * n
+    two_minus_slope = 2 - slope
+    ratio = (two_n - two_minus_slope) / (two_n + two_minus_slope)
+    omega_n = -factor * ratio * ei2phi * omega_nm1
+    partial_sum = partial_sum + omega_n
+    return (omega_n, partial_sum), None
+
+
+def omega(eiphi, slope, factor, n_terms=20):
+    """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
+    omega (e.g. can be used as a drop in replacement for the exp(i * phi) * special.hyp2f1
+    term in equation 13).
+
+    Parameters
+    ----------
+    eiphi:
+        `exp(i * phi)` where `phi` is the elliptical angle of the profile
+    slope:
+        The density slope of the power-law (lower value -> shallower profile, higher value
+        -> steeper profile).
+    factor:
+        The second flattening of and ellipse with axis ration q give by `f = (1 - q) / (1 + q)`
+    n_terms:
+        The number of terms to calculate for the series expansion, defaults to 20 (this should
+        be sufficient most of the time)
+    """
+    # 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),
+    )
+    return partial_sum
diff --git a/autogalaxy/profiles/mass/total/power_law.py b/autogalaxy/profiles/mass/total/power_law.py
index 9fc619e15..cdab5afe3 100644
--- a/autogalaxy/profiles/mass/total/power_law.py
+++ b/autogalaxy/profiles/mass/total/power_law.py
@@ -1,4 +1,13 @@
-import numpy as np
+import os
+
+if os.environ.get("USE_JAX", "0") == "1":
+    USING_JAX = True
+    import jax.numpy as np
+    from .jax_utils import omega
+else:
+    USING_JAX = False
+    import numpy as np
+
 from scipy import special
 from typing import Tuple
 
@@ -79,20 +88,24 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
         angle = np.arctan2(
             grid[:, 0], np.multiply(self.axis_ratio, grid[:, 1])
         )  # Note, this angle is not the position angle
-        R = np.sqrt(
-            np.add(np.multiply(self.axis_ratio**2, grid[:, 1] ** 2), grid[:, 0] ** 2)
-        )
         z = np.add(
             np.multiply(np.cos(angle), 1 + 0j), np.multiply(np.sin(angle), 0 + 1j)
         )
 
+        if USING_JAX:
+            # offset radius so calculation is finite at (0, 0)
+            R = np.sqrt((self.axis_ratio * grid[:, 1]) ** 2 + grid[:, 0] ** 2 + 1e-16)
+            zh = omega(z, slope, factor, n_terms=20)
+        else:
+            R = np.sqrt(
+                np.add(
+                    np.multiply(self.axis_ratio**2, grid[:, 1] ** 2), grid[:, 0] ** 2
+                )
+            )
+            zh = z * special.hyp2f1(1.0, 0.5 * slope, 2.0 - 0.5 * slope, -factor * z**2)
+
         complex_angle = (
-            2.0
-            * b
-            / (1.0 + self.axis_ratio)
-            * (b / R) ** (slope - 1.0)
-            * z
-            * special.hyp2f1(1.0, 0.5 * slope, 2.0 - 0.5 * slope, -factor * z**2)
+            2.0 * b / (1.0 + self.axis_ratio) * (b / R) ** (slope - 1.0) * zh
         )
 
         deflection_y = complex_angle.imag
diff --git a/autogalaxy/quantity/model/analysis.py b/autogalaxy/quantity/model/analysis.py
index 78ebb372f..aaf9e5f27 100644
--- a/autogalaxy/quantity/model/analysis.py
+++ b/autogalaxy/quantity/model/analysis.py
@@ -147,7 +147,7 @@ def save_attributes(self, paths: af.DirectoryPaths):
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization, and the pickled objects used by the aggregator output by this function.
         """
         paths.save_fits(
diff --git a/autogalaxy/quantity/model/visualizer.py b/autogalaxy/quantity/model/visualizer.py
index 2b2af686a..48b30e16a 100644
--- a/autogalaxy/quantity/model/visualizer.py
+++ b/autogalaxy/quantity/model/visualizer.py
@@ -30,7 +30,7 @@ def visualize(
         Parameters
         ----------
         paths
-            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization, and the pickled objects used by the aggregator output by this function.
         instance
             An instance of the model that is being fitted to the data by this analysis (whose parameters have been set
diff --git a/autogalaxy/util/shear_field.py b/autogalaxy/util/shear_field.py
index b0f5f02d2..179e8f09a 100644
--- a/autogalaxy/util/shear_field.py
+++ b/autogalaxy/util/shear_field.py
@@ -94,7 +94,6 @@ class ShearYX2D(aa.VectorYX2D, AbstractShearField):
 
 
 class ShearYX2DIrregular(aa.VectorYX2DIrregular, AbstractShearField):
-
     """
     An irregular shear field, which is collection of (y,x) vectors which are located on an irregular grid
     of (y,x) coordinates.
diff --git a/docs/howtogalaxy/chapter_1_introduction.rst b/docs/howtogalaxy/chapter_1_introduction.rst
index 538cd4da1..5d8cb0767 100644
--- a/docs/howtogalaxy/chapter_1_introduction.rst
+++ b/docs/howtogalaxy/chapter_1_introduction.rst
@@ -10,17 +10,14 @@ The chapter contains the following tutorials:
 `Tutorial 0: Visualization <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_1_introduction/tutorial_0_visualization.ipynb>`_
 - Setting up **PyAutoGalaxy**'s visualization library.
 
-`Tutorial 1: Grids And Galaxies <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_1_introduction/tutorial_1_grids_and_galaxies.ipynb>`_
-- Using grids of (y,x) coordinates with galaxies made up of light profiles.
+`Tutorial 1: Grids <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_1_introduction/tutorial_1_grids.ipynb>`_
+- Uniform grids of (y,x) coordinates used for ray-tracing.
 
-`Tutorial 2: Data <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_1_introduction/tutorial_2_data.ipynb>`_
-- Simulating and inspecting telescope imaging data of a galaxy.
+`Tutorial 2: Light Profiles <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_1_introduction/tutorial_2_light_profiles.ipynb>`_
+- Analytic profiles representing light distributions.
 
-`Tutorial 3: Fitting <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_1_introduction/tutorial_3_fitting.ipynb>`_
-- Fitting data with a galaxy model.
+`Tutorial 3: Galaxies <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_1_introduction/tutorial_3_galaxies.ipynb>`_
+- Making galaxy objects out of light and mass profile.
 
-`Tutorial 4: Methods <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_1_introduction/tutorial_4_methods.ipynb>`_
-- An overview of the different methods used to fit galaxies with.
-
-`Tutorial 5: Summary <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_1_introduction/tutorial_5_summary.ipynb>`_
+`Tutorial 6: Summary <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_1_introduction/tutorial_6_summary.ipynb>`_
 - A summary of the chapter.
\ No newline at end of file
diff --git a/docs/howtogalaxy/chapter_2_modeling.rst b/docs/howtogalaxy/chapter_2_modeling.rst
index 66c2432b2..d912944d6 100644
--- a/docs/howtogalaxy/chapter_2_modeling.rst
+++ b/docs/howtogalaxy/chapter_2_modeling.rst
@@ -5,27 +5,33 @@ In chapter 2, we'll take you through how to model galaxies using a non-linear se
 
 The chapter contains the following tutorials:
 
-`Tutorial 1: Non-linear Search <https://mybinder.org/v2/gh/Jammy2211/autolens_workspace/release?filepath=notebooks/howtolens/chapter_2_lens_modeling/tutorial_1_non_linear_search.ipynb>`_
-- How a non-linear search is used to fit a lens model.
+`Tutorial 1: Data <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_2_modeling/tutorial_1_data.ipynb>`_
+- Loading and inspecting telescope imaging data of a galaxy.
 
-`Tutorial 2: Parameter Space And Priors <https://mybinder.org/v2/gh/Jammy2211/autolens_workspace/release?filepath=notebooks/howtolens/chapter_2_lens_modeling/tutorial_2_parameter_space_and_priors.ipynb>`_
+`Tutorial 2: Fitting <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_2_modeling/tutorial_2_fitting.ipynb>`_
+- Fitting data with a galaxy model.
+
+`Tutorial 3: Non-linear Search <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_2_modeling/tutorial_3_non_linear_search.ipynb>`_
+- How a non-linear search is used to fit a model.
+
+`Tutorial 4: Parameter Space And Priors <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_2_modeling/tutorial_4_parameter_space_and_priors.ipynb>`_
 - The Concepts of a parameter space and priors.
 
-`Tutorial 3: Realism and Complexity <https://mybinder.org/v2/gh/Jammy2211/autolens_workspace/release?filepath=notebooks/howtolens/chapter_2_lens_modeling/tutorial_3_realism_and_complexity.ipynb>`_
-- Finding a balance between realism and complexity when composing and fitting a lens model.
+`Tutorial 5: Realism and Complexity <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_2_modeling/tutorial_5_realism_and_complexity.ipynb>`_
+- Finding a balance between realism and complexity when composing and fitting a model.
 
-`Tutorial 4: Dealing with Failure <https://mybinder.org/v2/gh/Jammy2211/autolens_workspace/release?filepath=notebooks/howtolens/chapter_2_lens_modeling/tutorial_4_dealing_with_failure.ipynb>`_
-- What to do when PyAutoLens finds an inaccurate lens model.
+`Tutorial 6: Dealing with Failure <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_2_modeling/tutorial_6_dealing_with_failure.ipynb>`_
+- What to do when PyAutoGalaxy finds an inaccurate model.
 
-`Tutorial 5: Linear Profiles <https://mybinder.org/v2/gh/Jammy2211/autolens_workspace/release?filepath=notebooks/howtolens/chapter_2_lens_modeling/tutorial_5_linear_profiles.ipynb>`_
+`Tutorial 7: Linear Profiles <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_2_modeling/tutorial_7_linear_profiles.ipynb>`_
 - Light profiles which capture complex morphologies in a reduced number of non-linear parameters.
 
-`Tutorial 6: Masking and Positions <https://mybinder.org/v2/gh/Jammy2211/autolens_workspace/release?filepath=notebooks/howtolens/chapter_2_lens_modeling/tutorial_6_masking_and_positions.ipynb>`_
-- How to mask and mark positions on your data to improve the lens model.
+`Tutorial 8: Masking <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_2_modeling/tutorial_8_masking.ipynb>`_
+- How to mask your data to improve the model.
 
-`Tutorial 7: Results <https://mybinder.org/v2/gh/Jammy2211/autolens_workspace/release?filepath=notebooks/howtolens/chapter_2_lens_modeling/tutorial_7_results.ipynb>`_
-- Overview of the results available after successfully fitting a lens model.
+`Tutorial 9: Results <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_2_modeling/tutorial_9_results.ipynb>`_
+- Overview of the results available after successfully fitting a model.
 
-`Tutorial 8: Need for Speed <https://mybinder.org/v2/gh/Jammy2211/autolens_workspace/release?filepath=notebooks/howtolens/chapter_2_lens_modeling/tutorial_8_need_for_speed.ipynb>`_
+`Tutorial 10: Need for Speed <https://mybinder.org/v2/gh/Jammy2211/autogalaxy_workspace/release?filepath=notebooks/howtogalaxy/chapter_2_modeling/tutorial_10_need_for_speed.ipynb>`_
 - How to fit complex models whilst balancing efficiency and run-time.
 
diff --git a/docs/overview/overview_1_start_here.rst b/docs/overview/overview_1_start_here.rst
index ce3046d85..5d7a399da 100644
--- a/docs/overview/overview_1_start_here.rst
+++ b/docs/overview/overview_1_start_here.rst
@@ -219,10 +219,10 @@ To further illustrate this, we create a merging galaxy system with 4 star formin
             intensity=0.1,
             effective_radius=1.6,
         ),
-        extra_galaxy_0=ag.lp.SersicSph(centre=(1.0, 1.0), intensity=0.5, effective_radius=0.2),
-        extra_galaxy_1=ag.lp.SersicSph(centre=(0.5, 0.8), intensity=0.5, effective_radius=0.2),
-        extra_galaxy_2=ag.lp.SersicSph(centre=(-1.0, -0.7), intensity=0.5, effective_radius=0.2),
-        extra_galaxy_3=ag.lp.SersicSph(centre=(-1.0, 0.4), intensity=0.5, effective_radius=0.2),
+        clump_0=ag.lp.SersicSph(centre=(1.0, 1.0), intensity=0.5, effective_radius=0.2),
+        clump_1=ag.lp.SersicSph(centre=(0.5, 0.8), intensity=0.5, effective_radius=0.2),
+        clump_2=ag.lp.SersicSph(centre=(-1.0, -0.7), intensity=0.5, effective_radius=0.2),
+        clump_3=ag.lp.SersicSph(centre=(-1.0, 0.4), intensity=0.5, effective_radius=0.2),
     )
 
     galaxy_1 = ag.Galaxy(


Test patch:
diff --git a/test_autogalaxy/analysis/analysis/test_analysis.py b/test_autogalaxy/analysis/analysis/test_analysis.py
index 387e50c8f..2e7be98d5 100644
--- a/test_autogalaxy/analysis/analysis/test_analysis.py
+++ b/test_autogalaxy/analysis/analysis/test_analysis.py
@@ -11,11 +11,10 @@
 
 def test__galaxies_via_instance(masked_imaging_7x7):
     galaxy = ag.Galaxy(redshift=0.5, light=ag.lp.Sersic(intensity=0.1))
-    extra_galaxy = ag.Galaxy(redshift=0.5, light=ag.lp.Sersic(intensity=0.2))
+    clump = ag.Galaxy(redshift=0.5, light=ag.lp.Sersic(intensity=0.2))
 
     model = af.Collection(
-        galaxies=af.Collection(galaxy=galaxy),
-        extra_galaxies=af.Collection(extra_galaxy_0=extra_galaxy),
+        galaxies=af.Collection(galaxy=galaxy), clumps=af.Collection(clump_0=clump)
     )
 
     analysis = ag.AnalysisImaging(dataset=masked_imaging_7x7)
diff --git a/test_autogalaxy/analysis/test_clump_model.py b/test_autogalaxy/analysis/test_clump_model.py
new file mode 100644
index 000000000..12d4230ae
--- /dev/null
+++ b/test_autogalaxy/analysis/test_clump_model.py
@@ -0,0 +1,99 @@
+import autofit as af
+import autogalaxy as ag
+
+
+def test__clumps():
+    centres = ag.Grid2DIrregular(values=[(1.0, 1.0)])
+
+    clump_model = ag.ClumpModel(
+        redshift=0.5, centres=centres, light_cls=ag.lp.SersicSph
+    )
+
+    clumps = clump_model.clumps
+
+    assert clumps["clump_0"].redshift == 0.5
+    assert clumps["clump_0"].light.centre == (1.0, 1.0)
+    assert isinstance(clumps["clump_0"].light.intensity, af.LogUniformPrior)
+    assert clumps["clump_0"].mass == None
+
+    clump_model = ag.ClumpModel(
+        redshift=0.5, centres=centres, mass_cls=ag.mp.IsothermalSph
+    )
+
+    clumps = clump_model.clumps
+
+    assert clumps["clump_0"].redshift == 0.5
+    assert clumps["clump_0"].mass.centre == (1.0, 1.0)
+    assert isinstance(clumps["clump_0"].mass.einstein_radius, af.UniformPrior)
+    assert clumps["clump_0"].light == None
+
+    clump_model = ag.ClumpModel(
+        redshift=0.5,
+        centres=centres,
+        light_cls=ag.lp.SersicSph,
+        mass_cls=ag.mp.IsothermalSph,
+    )
+
+    clumps = clump_model.clumps
+
+    assert clumps["clump_0"].redshift == 0.5
+    assert clumps["clump_0"].light.centre == (1.0, 1.0)
+    assert clumps["clump_0"].mass.centre == (1.0, 1.0)
+    assert isinstance(clumps["clump_0"].light.intensity, af.LogUniformPrior)
+    assert isinstance(clumps["clump_0"].mass.einstein_radius, af.UniformPrior)
+
+
+def test__clumps_light_only():
+    centres = ag.Grid2DIrregular(values=[(1.0, 1.0)])
+
+    clump_model = ag.ClumpModel(
+        redshift=0.5,
+        centres=centres,
+        light_cls=ag.lp.SersicSph,
+        mass_cls=ag.mp.IsothermalSph,
+    )
+
+    clumps_light_only = clump_model.clumps_light_only
+
+    assert clumps_light_only["clump_0"].redshift == 0.5
+    assert clumps_light_only["clump_0"].light.centre == (1.0, 1.0)
+    assert isinstance(clumps_light_only["clump_0"].light.intensity, af.LogUniformPrior)
+    assert not hasattr(clumps_light_only["clump_0"], "mass")
+
+
+def test__clumps_mass_only():
+    centres = ag.Grid2DIrregular(values=[(1.0, 1.0)])
+
+    clump_model = ag.ClumpModel(
+        redshift=0.5,
+        centres=centres,
+        light_cls=ag.lp.SersicSph,
+        mass_cls=ag.mp.IsothermalSph,
+    )
+
+    clumps_mass_only = clump_model.clumps_mass_only
+
+    assert clumps_mass_only["clump_0"].redshift == 0.5
+    assert clumps_mass_only["clump_0"].mass.centre == (1.0, 1.0)
+    assert isinstance(clumps_mass_only["clump_0"].mass.einstein_radius, af.UniformPrior)
+    assert not hasattr(clumps_mass_only["clump_0"], "light")
+
+
+def test__einstein_radius_max():
+    centres = ag.Grid2DIrregular(values=[(1.0, 1.0)])
+
+    clump_model = ag.ClumpModel(
+        redshift=0.5,
+        centres=centres,
+        mass_cls=ag.mp.IsothermalSph,
+        einstein_radius_upper_limit=1.0,
+    )
+
+    clumps = clump_model.clumps
+
+    assert isinstance(clumps["clump_0"].mass.einstein_radius, af.UniformPrior)
+    assert clumps["clump_0"].mass.einstein_radius.upper_limit == 1.0
+
+    assert clumps["clump_0"].redshift == 0.5
+    assert clumps["clump_0"].mass.centre == (1.0, 1.0)
+    assert clumps["clump_0"].light == None
diff --git a/test_autogalaxy/analysis/test_preloads.py b/test_autogalaxy/analysis/test_preloads.py
deleted file mode 100644
index 284962428..000000000
--- a/test_autogalaxy/analysis/test_preloads.py
+++ /dev/null
@@ -1,130 +0,0 @@
-import numpy as np
-from os import path
-
-import autofit as af
-
-import autogalaxy as ag
-
-
-def test__set_blurred_image():
-    # Blurred image is all zeros so preloads as zeros
-
-    fit_0 = ag.m.MockFitImaging(blurred_image=np.zeros(2))
-    fit_1 = ag.m.MockFitImaging(blurred_image=np.zeros(2))
-
-    preloads = ag.Preloads(blurred_image=1)
-    preloads.set_blurred_image(fit_0=fit_0, fit_1=fit_1)
-
-    assert preloads.blurred_image is None
-
-    # Blurred image are different, indicating the model parameters change the grid, so no preloading.
-
-    fit_0 = ag.m.MockFitImaging(blurred_image=np.array([1.0]))
-    fit_1 = ag.m.MockFitImaging(blurred_image=np.array([2.0]))
-
-    preloads = ag.Preloads(blurred_image=1)
-    preloads.set_blurred_image(fit_0=fit_0, fit_1=fit_1)
-
-    assert preloads.blurred_image is None
-
-    # Blurred images are the same meaning they are fixed in the model, so do preload.
-
-    fit_0 = ag.m.MockFitImaging(blurred_image=np.array([1.0]))
-    fit_1 = ag.m.MockFitImaging(blurred_image=np.array([1.0]))
-
-    preloads = ag.Preloads(blurred_image=1)
-    preloads.set_blurred_image(fit_0=fit_0, fit_1=fit_1)
-
-    assert (preloads.blurred_image == np.array([1.0])).all()
-
-
-def test__info():
-    file_path = path.join("{}".format(path.dirname(path.realpath(__file__))), "files")
-
-    file_preloads = path.join(file_path, "preloads.summary")
-
-    preloads = ag.Preloads(
-        blurred_image=np.zeros(3),
-        w_tilde=None,
-        use_w_tilde=False,
-        image_plane_mesh_grid_pg_list=None,
-        relocated_grid=None,
-        mapper_list=None,
-        operated_mapping_matrix=None,
-    )
-
-    af.formatter.output_list_of_strings_to_file(
-        file=file_preloads, list_of_strings=preloads.info
-    )
-
-    results = open(file_preloads)
-    lines = results.readlines()
-
-    i = 0
-
-    assert lines[i] == f"W Tilde = False\n"
-    i += 1
-    assert lines[i] == f"Use W Tilde = False\n"
-    i += 1
-    assert lines[i] == f"\n"
-    i += 1
-    assert lines[i] == f"Blurred Image = False\n"
-    i += 1
-    assert lines[i] == f"Mapper = False\n"
-    i += 1
-    assert lines[i] == f"Blurred Mapping Matrix = False\n"
-    i += 1
-    assert lines[i] == f"Inversion Linear Func (Linear Light Profile) Dicts = False\n"
-    i += 1
-    assert lines[i] == f"Curvature Matrix = False\n"
-    i += 1
-    assert lines[i] == f"Curvature Matrix Mapper Diag = False\n"
-    i += 1
-    assert lines[i] == f"Regularization Matrix = False\n"
-    i += 1
-    assert lines[i] == f"Log Det Regularization Matrix Term = False\n"
-    i += 1
-
-    preloads = ag.Preloads(
-        blurred_image=1,
-        w_tilde=1,
-        use_w_tilde=True,
-        image_plane_mesh_grid_pg_list=1,
-        mapper_list=1,
-        operated_mapping_matrix=1,
-        curvature_matrix=1,
-        regularization_matrix=1,
-        log_det_regularization_matrix_term=1,
-    )
-
-    af.formatter.output_list_of_strings_to_file(
-        file=file_preloads, list_of_strings=preloads.info
-    )
-
-    results = open(file_preloads)
-    lines = results.readlines()
-
-    i = 0
-
-    assert lines[i] == f"W Tilde = True\n"
-    i += 1
-    assert lines[i] == f"Use W Tilde = True\n"
-    i += 1
-    assert lines[i] == f"\n"
-    i += 1
-    assert lines[i] == f"Blurred Image = True\n"
-    i += 1
-    assert lines[i] == f"Mapper = True\n"
-    i += 1
-    assert lines[i] == f"Blurred Mapping Matrix = True\n"
-    i += 1
-    assert lines[i] == f"Inversion Linear Func (Linear Light Profile) Dicts = False\n"
-    i += 1
-    assert lines[i] == f"Curvature Matrix = True\n"
-    i += 1
-    assert lines[i] == f"Curvature Matrix Mapper Diag = False\n"
-    i += 1
-    assert lines[i] == f"Regularization Matrix = True\n"
-    i += 1
-    assert lines[i] == f"Log Det Regularization Matrix Term = True\n"
-    i += 1
diff --git a/test_autogalaxy/config/general.yaml b/test_autogalaxy/config/general.yaml
index 2acec5ada..26e205c6b 100644
--- a/test_autogalaxy/config/general.yaml
+++ b/test_autogalaxy/config/general.yaml
@@ -1,6 +1,5 @@
 analysis:
   n_cores: 1
-  preload_attempts: 250
 fits:
   flip_for_ds9: true
 grid:
@@ -39,6 +38,4 @@ structures:
   native_binned_only: false           # If True, data structures are only stored in their native and binned format. This is used to reduce memory usage in autocti.
 test:
   check_likelihood_function: false  # if True, when a search is resumed the likelihood of a previous sample is recalculated to ensure it is consistent with the previous run.
-  check_preloads: false
-  exception_override: true
-  preloads_check_threshold: 1.0     # If the figure of merit of a fit with and without preloads is greater than this threshold, the check preload test fails and an exception raised for a model-fit. 
+  exception_override: true
\ No newline at end of file
diff --git a/test_autogalaxy/config/visualize.yaml b/test_autogalaxy/config/visualize.yaml
index 50c40f0cb..2b170d7c2 100644
--- a/test_autogalaxy/config/visualize.yaml
+++ b/test_autogalaxy/config/visualize.yaml
@@ -8,7 +8,7 @@ include:
     einstein_radius: true
     half_light_radius: true
   include_2d:
-    border: true
+    border: false
     tangential_caustics: false
     radial_caustics: false
     tangential_critical_curves: false
diff --git a/test_autogalaxy/imaging/model/test_plotter_interface_imaging.py b/test_autogalaxy/imaging/model/test_plotter_interface_imaging.py
index 0c43849e2..a2fd39420 100644
--- a/test_autogalaxy/imaging/model/test_plotter_interface_imaging.py
+++ b/test_autogalaxy/imaging/model/test_plotter_interface_imaging.py
@@ -31,19 +31,6 @@ def test__imaging(imaging_7x7, include_2d_all, plot_path, plot_patch):
     assert path.join(plot_path, "signal_to_noise_map.png") not in plot_patch.paths
 
 
-def test__imaging_combined(imaging_7x7, plot_path, plot_patch):
-    if path.exists(plot_path):
-        shutil.rmtree(plot_path)
-
-    visualizer = PlotterInterfaceImaging(image_path=plot_path)
-
-    visualizer.imaging_combined(dataset_list=[imaging_7x7, imaging_7x7])
-
-    plot_path = path.join(plot_path, "combined")
-
-    assert path.join(plot_path, "subplot_dataset.png") in plot_patch.paths
-
-
 def test__fit_imaging(
     masked_imaging_7x7,
     fit_imaging_x2_galaxy_inversion_7x7,
@@ -74,18 +61,3 @@ def test__fit_imaging(
     )
 
     assert image.shape == (7, 7)
-
-
-def test__fit_imaging_combined(
-    fit_imaging_x2_galaxy_inversion_7x7, plot_path, plot_patch
-):
-    if path.exists(plot_path):
-        shutil.rmtree(plot_path)
-
-    visualizer = PlotterInterfaceImaging(image_path=plot_path)
-
-    visualizer.fit_imaging_combined(fit_list=2 * [fit_imaging_x2_galaxy_inversion_7x7])
-
-    plot_path = path.join(plot_path, "combined")
-
-    assert path.join(plot_path, "subplot_fit.png") in plot_patch.paths