Feature/xp namespace

autolens/__init__.py

@@ -29,6 +29,7 @@
 from autoarray.operators.transformer import TransformerDFT
 from autoarray.operators.transformer import TransformerNUFFT
 from autoarray.preloads import Preloads
+from autoarray.preloads import mapper_indices_from
 from autoarray.structures.arrays.uniform_1d import Array1D
 from autoarray.structures.arrays.uniform_2d import Array2D
 from autoarray.structures.arrays.rgb import Array2DRGB
@@ -115,8 +116,6 @@
 from .quantity.fit_quantity import FitQuantity
 from .quantity.model.analysis import AnalysisQuantity
 
-from .analysis.preloads import mapper_indices_from
-
 from . import exc
 from . import mock as m
 from . import util

autolens/aggregator/subplot.py

@@ -69,8 +69,8 @@ class SubplotFitX1Plane(Enum):
     signal_to_noise_map = (1, 0)
     model_data = (2, 0)
     lens_light_subtracted_image = (0, 1)
-    lens_light_subtracted_image_zero = (0, 1)
-    normalized_residual_map = (0, 2)
+    lens_light_subtracted_image_zero = (1, 1)
+    normalized_residual_map = (2, 1)
 
 
 class SubplotFit(Enum):

autolens/analysis/analysis/lens.py

@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import logging
 import numpy as np
 from typing import List, Optional
@@ -96,7 +95,7 @@ def tracer_via_instance_from(
         )
 
     def log_likelihood_penalty_from(
-        self, instance: af.ModelInstance
+        self, instance: af.ModelInstance, xp=np
     ) -> Optional[float]:
         """
         Call the positions overwrite log likelihood function, which add a penalty term to the likelihood if the
@@ -117,15 +116,15 @@ def log_likelihood_penalty_from(
         The penalty value of the positions log likelihood, if the positions do not trace close in the source plane,
         else a None is returned to indicate there is no penalty.
         """
-        log_likelihood_penalty = jnp.array(0.0)
+        log_likelihood_penalty = xp.array(0.0)
 
         if self.positions_likelihood_list is not None:
 
             try:
                 for positions_likelihood in self.positions_likelihood_list:
                     log_likelihood_penalty = (
                         positions_likelihood.log_likelihood_penalty_from(
-                            instance=instance, analysis=self
+                            instance=instance, analysis=self, xp=xp
                         )
                     )
 

autolens/analysis/model_util.py

@@ -1,130 +1,7 @@
-import numpy as np
-from typing import Optional, Tuple
-
 import autofit as af
 import autolens as al
 
-
-def mge_model_from(
-    mask_radius: float,
-    total_gaussians: int = 30,
-    gaussian_per_basis: int = 1,
-    centre_prior_is_uniform: bool = True,
-    centre: Tuple[float, float] = (0.0, 0.0),
-    centre_fixed: Optional[Tuple[float, float]] = None,
-    use_spherical: bool = False,
-) -> af.Collection:
-    """
-    Construct a Multi-Gaussian Expansion (MGE) for the lens or source galaxy light
-
-    This model is designed as a "start here" configuration for lens modeling:
-
-    - The lens and source light are represented by a Basis object composed of many
-      Gaussian light profiles with fixed logarithmically spaced widths (`sigma`).
-    - All Gaussians within each basis share common centres and ellipticity
-      components, reducing degeneracy while retaining flexibility.
-
-    - Users can combine with a lens mass model of their choiuce.
-
-    The resulting model provides a good balance of speed, flexibility, and accuracy
-    for fitting most galaxy-scale strong lenses.
-
-    This code is mostly to make the API simple for new users, hiding the technical
-    details of setting up an MGE. More advanced users may wish to customize the
-    model further.
-
-    Parameters
-    ----------
-    mask_radius
-        The outer radius (in arcseconds) of the circular mask applied to the data.
-        This determines the maximum Gaussian width (`sigma`) used in the lens MGE.
-    lens_total_gaussians
-        Total number of Gaussian light profiles used in the lens MGE basis.
-    source_total_gaussians
-        Total number of Gaussian light profiles used in the source MGE basis.
-    lens_gaussian_per_basis
-        Number of separate Gaussian bases to include for the lens light profile.
-        Each basis has `lens_total_gaussians` components.
-    source_gaussian_per_basis
-        Number of separate Gaussian bases to include for the source light profile.
-        Each basis has `source_total_gaussians` components.
-
-    Returns
-    -------
-    model : af.Collection
-        An `autofit.Collection` containing:
-        - A lens galaxy at redshift 0.5, with:
-          * bulge light profile: MGE basis of Gaussians
-          * mass profile: Isothermal ellipsoid
-          * external shear
-        - A source galaxy at redshift 1.0, with:
-          * bulge light profile: MGE basis of Gaussians
-
-    Notes
-    -----
-    - Lens light Gaussians have widths (sigma) logarithmically spaced between 0.01"
-      and the mask radius.
-    - Source light Gaussians have widths logarithmically spaced between 0.01" and 1.0".
-    - Gaussian centres are free parameters but tied across all components in each
-      basis to reduce dimensionality.
-    - This function is a convenience utility: it hides the technical setup of MGE
-      composition and provides a ready-to-use lens model for quick experimentation.
-    """
-
-    # The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0".
-    log10_sigma_list = np.linspace(-4, np.log10(mask_radius), total_gaussians)
-
-    # By defining the centre here, it creates two free parameters that are assigned below to all Gaussians.
-
-    if centre_fixed is not None:
-        centre_0 = centre[0]
-        centre_1 = centre[1]
-    elif centre_prior_is_uniform:
-        centre_0 = af.UniformPrior(
-            lower_limit=centre[0] - 0.1, upper_limit=centre[0] + 0.1
-        )
-        centre_1 = af.UniformPrior(
-            lower_limit=centre[1] - 0.1, upper_limit=centre[1] + 0.1
-        )
-    else:
-        centre_0 = af.GaussianPrior(mean=centre[0], sigma=0.3)
-        centre_1 = af.GaussianPrior(mean=centre[1], sigma=0.3)
-
-    if use_spherical:
-        model_cls = al.lp_linear.GaussianSph
-    else:
-        model_cls = al.lp_linear.Gaussian
-
-    bulge_gaussian_list = []
-
-    for j in range(gaussian_per_basis):
-        # A list of Gaussian model components whose parameters are customized belows.
-
-        gaussian_list = af.Collection(
-            af.Model(model_cls) for _ in range(total_gaussians)
-        )
-
-        # Iterate over every Gaussian and customize its parameters.
-
-        for i, gaussian in enumerate(gaussian_list):
-            gaussian.centre.centre_0 = centre_0  # All Gaussians have same y centre.
-            gaussian.centre.centre_1 = centre_1  # All Gaussians have same x centre.
-            if not use_spherical:
-                gaussian.ell_comps = gaussian_list[
-                    0
-                ].ell_comps  # All Gaussians have same elliptical components.
-            gaussian.sigma = (
-                10 ** log10_sigma_list[i]
-            )  # All Gaussian sigmas are fixed to values above.
-
-        bulge_gaussian_list += gaussian_list
-
-    # The Basis object groups many light profiles together into a single model component.
-
-    return af.Model(
-        al.lp_basis.Basis,
-        profile_list=bulge_gaussian_list,
-    )
+from autogalaxy.analysis.model_util import mge_model_from
 
 
 def simulator_start_here_model_from(

autolens/analysis/positions.py

@@ -1,10 +1,7 @@
 import jax
-import jax.numpy as jnp
 import numpy as np
-from typing import Optional, Union
+from typing import Optional
 from os import path
-import os
-from typing import Dict
 
 import autoarray as aa
 import autofit as af
@@ -137,8 +134,8 @@ def output_positions_info(
             f.write("")
 
     def log_likelihood_penalty_from(
-        self, instance: af.ModelInstance, analysis: AnalysisDataset
-    ) -> jnp.array:
+        self, instance: af.ModelInstance, analysis: AnalysisDataset, xp=np
+    ) -> np.array:
         """
         Returns a log-likelihood penalty used to constrain lens models where multiple image-plane
         positions do not trace to within a threshold distance of one another in the source-plane.
@@ -176,16 +173,17 @@ def log_likelihood_penalty_from(
         tracer = analysis.tracer_via_instance_from(instance=instance)
 
         if not tracer.has(cls=ag.mp.MassProfile) or len(tracer.planes) == 1:
-            return (jnp.array(0.0),)
+            return xp.array(0.0)
 
         positions_fit = SourceMaxSeparation(
             data=self.positions,
             noise_map=None,
             tracer=tracer,
             plane_redshift=self.plane_redshift,
+            xp=xp
         )
 
-        max_separation = jnp.max(
+        max_separation = xp.max(
             positions_fit.furthest_separations_of_plane_positions.array
         )
 
@@ -194,5 +192,5 @@ def log_likelihood_penalty_from(
         return jax.lax.cond(
             max_separation > self.threshold,
             lambda: penalty,
-            lambda: jnp.array(0.0),
+            lambda: xp.array(0.0),
         )

autolens/config/visualize/plots.yaml

@@ -66,5 +66,3 @@ fit_quantity:                              # Settings for plots of fit quantitie
 galaxies:                                  # Settings for plots of galaxies (e.g. GalaxiesPlotter).
   subplot_galaxies: true                   # Plot subplot of all quantities in each galaxies group (e.g. images, convergence)?
   subplot_galaxy_images: false             # Plot subplot of the image of each galaxy in the model?
-  subplot_galaxies_1d: false               # Plot subplot of all quantities in 1D of each galaxies group (e.g. images, convergence)?
-  subplot_galaxies_1d_decomposed: false    # Plot subplot of all quantities in 1D decomposed of each galaxies group (e.g. images, convergence)?

autolens/exc.py

@@ -32,7 +32,7 @@ class PixelizationException(af.exc.FitException):
     """
     Raises exceptions associated with the `inversion/pixelization` modules and `Pixelization` classes.
 
-    For example if a `Rectangular` mesh has dimensions below 3x3.
+    For example if a `RectangularMagnification` mesh has dimensions below 3x3.
 
     This exception overwrites `autoarray.exc.PixelizationException` in order to add a `FitException`. This means that
     if this exception is raised during a model-fit in the analysis class's `log_likelihood_function` that model

autolens/imaging/fit_imaging.py

@@ -24,6 +24,7 @@ def __init__(
         adapt_images: Optional[ag.AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
         preloads: aa.Preloads = None,
+        xp=np
     ):
         """
         Fits an imaging dataset using a `Tracer` object.
@@ -64,7 +65,7 @@ def __init__(
             Settings controlling how an inversion is fitted for example which linear algebra formalism is used.
         """
 
-        super().__init__(dataset=dataset, dataset_model=dataset_model)
+        super().__init__(dataset=dataset, dataset_model=dataset_model, xp=xp)
         AbstractFitInversion.__init__(
             self=self, model_obj=tracer, settings_inversion=settings_inversion
         )
@@ -86,12 +87,14 @@ def blurred_image(self) -> aa.Array2D:
         ):
             return self.tracer.image_2d_from(
                 grid=self.grids.lp,
+                xp=self._xp,
             )
 
         return self.tracer.blurred_image_2d_from(
             grid=self.grids.lp,
             psf=self.dataset.psf,
             blurring_grid=self.grids.blurring,
+            xp=self._xp,
         )
 
     @property
@@ -118,6 +121,7 @@ def tracer_to_inversion(self) -> TracerToInversion:
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
+            xp=self._xp,
         )
 
     @cached_property

autolens/imaging/model/analysis.py

@@ -1,5 +1,5 @@
 import logging
-
+import numpy as np
 
 import autofit as af
 import autogalaxy as ag
@@ -46,7 +46,7 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
 
         return self
 
-    def log_likelihood_function(self, instance: af.ModelInstance) -> float:
+    def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
         instance to the imaging dataset.
@@ -86,14 +86,15 @@ def log_likelihood_function(self, instance: af.ModelInstance) -> float:
         """
 
         log_likelihood_penalty = self.log_likelihood_penalty_from(
-            instance=instance
+            instance=instance,
+            xp=xp
         )
 
-        return self.fit_from(instance=instance).figure_of_merit - log_likelihood_penalty
+        return self.fit_from(instance=instance, xp=xp).figure_of_merit - log_likelihood_penalty
 
     def fit_from(
         self,
-        instance: af.ModelInstance,
+        instance: af.ModelInstance, xp=np
     ) -> FitImaging:
         """
         Given a model instance create a `FitImaging` object.
@@ -130,7 +131,8 @@ def fit_from(
             dataset_model=dataset_model,
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
-            preloads=self.preloads
+            preloads=self.preloads,
+            xp=xp
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):

autolens/interferometer/fit_interferometer.py

@@ -21,6 +21,7 @@ def __init__(
         adapt_images: Optional[ag.AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
         preloads: aa.Preloads = None,
+        xp=np,
     ):
         """
         Fits an interferometer dataset using a `Tracer` object.
@@ -76,12 +77,14 @@ def __init__(
         super().__init__(
             dataset=dataset,
             dataset_model=dataset_model,
+            xp=xp,
         )
         AbstractFitInversion.__init__(
             self=self, model_obj=tracer, settings_inversion=settings_inversion
         )
 
         self.preloads = preloads
+        self._xp = xp
 
     @property
     def profile_visibilities(self) -> aa.Visibilities:
@@ -90,7 +93,7 @@ def profile_visibilities(self) -> aa.Visibilities:
         transform to the sum of light profile images.
         """
         return self.tracer.visibilities_from(
-            grid=self.grids.lp, transformer=self.dataset.transformer
+            grid=self.grids.lp, transformer=self.dataset.transformer, xp=self._xp
         )
 
     @property
@@ -117,6 +120,7 @@ def tracer_to_inversion(self) -> TracerToInversion:
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
+            xp=self._xp,
         )
 
     @cached_property