xp refactor complete

Author: Jammy2211

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/analysis/analysis/lens.py

@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import logging
 import numpy as np
 from typing import List, Optional

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

autolens/analysis/preloads.py

@@ -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/config/visualize/plots.yaml

@@ -87,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
@@ -119,7 +121,7 @@ def tracer_to_inversion(self) -> TracerToInversion:
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
-            xp=self.xp,
+            xp=self._xp,
         )
 
     @cached_property

autolens/imaging/fit_imaging.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

autolens/interferometer/fit_interferometer.py

@@ -114,7 +114,7 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
 
         return self
 
-    def log_likelihood_function(self, instance):
+    def log_likelihood_function(self, instance, xp=np):
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
         instance to the interferometer dataset.
@@ -153,14 +153,16 @@ def log_likelihood_function(self, instance):
             The log likelihood indicating how well this model instance fitted the interferometer data.
         """
 
-        log_likelihood_penalty = self.log_likelihood_penalty_from(instance=instance)
+        log_likelihood_penalty = self.log_likelihood_penalty_from(
+            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,
-    ) -> FitInterferometer:
+    def fit_from(self, instance: af.ModelInstance, xp=np) -> FitInterferometer:
         """
         Given a model instance create a `FitInterferometer` object.
 
@@ -196,6 +198,7 @@ def fit_from(
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
+            xp=xp,
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):

autolens/interferometer/model/analysis.py

@@ -1,3 +1,6 @@
+import numpy as np
+
+
 class MockTracer:
     def __init__(
         self, traced_grid_2d_list_from=None, image_plane_mesh_grid_pg_list=None
@@ -61,5 +64,5 @@ def einstein_radius_from(self, grid):
     def einstein_mass_angular_from(self, grid):
         return self.einstein_mass
 
-    def time_delays_from(self, grid):
+    def time_delays_from(self, grid, xp=np):
         return self.time_delays