Merge branch 'main' into feature/ellipse_utils

Author: samlange04

autogalaxy/__init__.py

@@ -10,22 +10,17 @@
 from autoarray.dataset.imaging.dataset import Imaging  # noqa
 from autoarray.dataset.interferometer.dataset import Interferometer  # noqa
 from autoarray.dataset.dataset_model import DatasetModel
-from autoarray.inversion.pixelization import mesh  # noqa
+from autoarray.inversion.mesh import mesh  # noqa
 from autoarray.inversion import regularization as reg  # noqa
-from autoarray.inversion.pixelization import image_mesh
-from autoarray.inversion.pixelization.mappers.abstract import AbstractMapper  # noqa
-from autoarray.inversion.inversion.settings import SettingsInversion  # noqa
+from autoarray.inversion.mesh import image_mesh
+from autoarray.inversion.mappers.abstract import Mapper  # noqa
+from autoarray.settings import Settings  # noqa
 from autoarray.inversion.inversion.factory import inversion_from as Inversion  # noqa
-from autoarray.inversion.pixelization.image_mesh.abstract import AbstractImageMesh
-from autoarray.inversion.pixelization.mesh.abstract import AbstractMesh
+from autoarray.inversion.mesh.image_mesh.abstract import AbstractImageMesh
+from autoarray.inversion.mesh.mesh.abstract import AbstractMesh
 from autoarray.inversion.regularization.abstract import AbstractRegularization
-from autoarray.inversion.pixelization.pixelization import Pixelization  # noqa
-from autoarray.inversion.pixelization.mappers.abstract import AbstractMapper
-from autoarray.inversion.pixelization.mappers.mapper_grids import MapperGrids  # noqa
-from autoarray.inversion.pixelization.mappers.factory import (
-    mapper_from as Mapper,
-)  # noqa
-from autoarray.inversion.pixelization.border_relocator import BorderRelocator
+from autoarray.inversion.pixelization import Pixelization  # noqa
+from autoarray.inversion.mesh.border_relocator import BorderRelocator
 from autoarray.preloads import Preloads
 from autoarray.preloads import mapper_indices_from
 from autoarray.mask.mask_1d import Mask1D  # noqa
@@ -43,18 +38,22 @@
 from autoarray.structures.grids.uniform_2d import Grid2D  # noqa
 from autoarray.structures.grids.irregular_2d import Grid2DIrregular  # noqa
 from autoarray.operators.over_sampling.over_sampler import OverSampler  # noqa
-from autoarray.structures.mesh.rectangular_2d import Mesh2DRectangular  # noqa
-from autoarray.structures.mesh.rectangular_2d_uniform import (
-    Mesh2DRectangularUniform,
+from autoarray.inversion.mesh.interpolator.rectangular import (
+    InterpolatorRectangular,
+)  # noqa
+from autoarray.inversion.mesh.interpolator.delaunay import (
+    InterpolatorDelaunay,
 )  # noqa
-from autoarray.structures.mesh.delaunay_2d import Mesh2DDelaunay  # noqa
 from autoarray.structures.vectors.uniform import VectorYX2D  # noqa
 from autoarray.structures.vectors.irregular import VectorYX2DIrregular  # noqa
 from autoarray.layout.region import Region1D  # noqa
 from autoarray.layout.region import Region2D  # noqa
 from autoarray.structures.arrays.kernel_2d import Kernel2D  # noqa
 from autoarray.structures.visibilities import Visibilities  # noqa
 from autoarray.structures.visibilities import VisibilitiesNoiseMap  # noqa
+from autoarray.inversion.mesh.mesh_geometry.rectangular import (
+    rectangular_edge_pixel_list_from,
+)
 
 from .analysis import model_util
 from .analysis.adapt_images.adapt_images import AdaptImages

autogalaxy/abstract_fit.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 from typing import TYPE_CHECKING, Dict, Optional
 
+import numpy as np
+
 from autofit import ModelInstance
 
 if TYPE_CHECKING:
@@ -13,11 +15,7 @@
 
 
 class AbstractFitInversion:
-    def __init__(
-        self,
-        model_obj,
-        settings_inversion: aa.SettingsInversion,
-    ):
+    def __init__(self, model_obj, settings: aa.Settings, xp=np):
         """
         An abstract fit object which fits to datasets (e.g. imaging, interferometer) inherit from.
 
@@ -30,11 +28,20 @@ def __init__(
             The object which contains the model components (e.g. light profiles, galaxies, etc) which are used to
             create the model-data that fits the data. In PyAutoGalaxy this is a list of galaxies and PyAutoLens
             it is a `Tracer`.
-        settings_inversion
+        settings
             Settings controlling how an inversion is fitted for example which linear algebra formalism is used.
         """
         self.model_obj = model_obj
-        self.settings_inversion = settings_inversion
+        self.settings = settings or aa.Settings()
+        self.use_jax = xp is not np
+
+    @property
+    def _xp(self):
+        if self.use_jax:
+            import jax.numpy as jnp
+
+            return jnp
+        return np
 
     @property
     def total_mappers(self) -> int:
@@ -68,7 +75,7 @@ def perform_inversion(self) -> bool:
     def sparse_operator(self) -> Optional[aa.ImagingSparseOperator]:
         """
         Only call the `sparse_operator` property of a dataset used to perform efficient linear algebra calculations if
-        the SettingsInversion()` object has `use_sparse_operator=True`, to avoid unnecessary computation.
+        the Settings()` object has `use_sparse_operator=True`, to avoid unnecessary computation.
 
         Returns
         -------
@@ -95,6 +102,14 @@ def linear_light_profile_intensity_dict(
 
         This function returns a dictionary which maps every linear light profile instance to its solved for
         `intensity` value in the inversion, so that the intensity value of every light profile can be accessed.
+
+        Type casting is complicated by JAX. When this function is used in a JAX.jit (e.g. computed latent varialbes)
+        it requires the reconstruction values to be JAX arrays, but when it is used outside of JAX certain taks
+        requires the reconstruction values to be floats.
+
+        An example of the latter is using a tracer inferred in one search to pass the solved for intensity values of
+        linear light profiles to a subsequent search, for example setting up the intensities of the mass components
+        of a light dark model.
         """
 
         if self.inversion is None:
@@ -110,12 +125,20 @@ def linear_light_profile_intensity_dict(
             reconstruction = self.inversion.reconstruction_dict[linear_obj_func]
 
             for i, light_profile in enumerate(linear_obj_func.light_profile_list):
-                linear_light_profile_intensity_dict[light_profile] = reconstruction[i]
+                if self.use_jax:
+                    linear_light_profile_intensity_dict[light_profile] = reconstruction[
+                        i
+                    ]
+                else:
+                    linear_light_profile_intensity_dict[light_profile] = float(
+                        reconstruction[i]
+                    )
 
         return linear_light_profile_intensity_dict
 
     def galaxy_linear_obj_data_dict_from(
-        self, use_image: bool = False
+        self,
+        use_operated: bool = True,
     ) -> Dict[Galaxy, aa.Array2D]:
         """
         Returns a dictionary mapping every galaxy containing a linear
@@ -127,16 +150,17 @@ def galaxy_linear_obj_data_dict_from(
         This is used to create the overall `galaxy_model_image_dict`, which maps every galaxy to its
         overall `model_data` (e.g. including the `model_data` of orindary light profiles too).
 
-        If `use_image=False`, the `reconstructed_data` of the inversion (e.g. an image for dataset data,
+        If `use_operated=False`, the `reconstructed_data` of the inversion (e.g. an image for dataset data,
         visibilities for  interferometer data) is input in the dictionary.
 
-        if `use_image=True`, the `reconstructed_image` of the inversion (e.g. the image for dataset data, the
+        if `use_operated=True`, the `reconstructed_operated_data` of the inversion (e.g. the image for dataset data, the
         real-space image for interferometer data) is input in the dictionary.
 
         Parameters
         ----------
-        use_image
-            Whether to put the reconstructed data or images in the dictionary.
+        use_operated
+            Whether to use the operated (e.g PSF convolved) images of the linear objects in the dictionary, or
+            the unoperated images.
 
         Returns
         -------
@@ -154,13 +178,12 @@ def galaxy_linear_obj_data_dict_from(
             except KeyError:
                 continue
 
-            if not use_image:
-                mapped_reconstructed = self.inversion.mapped_reconstructed_data_dict[
-                    linear_obj
-                ]
-
+            if use_operated:
+                mapped_reconstructed = (
+                    self.inversion.mapped_reconstructed_operated_data_dict[linear_obj]
+                )
             else:
-                mapped_reconstructed = self.inversion.mapped_reconstructed_image_dict[
+                mapped_reconstructed = self.inversion.mapped_reconstructed_data_dict[
                     linear_obj
                 ]
 

autogalaxy/aggregator/ellipse/fit_ellipse.py

@@ -33,7 +33,7 @@ def _fit_ellipse_from(
     is instead used to load lists of the data, noise-map, PSF and mask and combine them into a list of
     `FitEllipse` objects.
 
-    The settings of an inversion can be overwritten by inputting a `settings_inversion` object, for example
+    The settings of an inversion can be overwritten by inputting a `settings` object, for example
     if you want to use a grid with a different inversion solver.
 
     Parameters

autogalaxy/aggregator/imaging/fit_imaging.py

@@ -16,7 +16,7 @@
 def _fit_imaging_from(
     fit: af.Fit,
     instance: Optional[af.ModelInstance] = None,
-    settings_inversion: aa.SettingsInversion = None,
+    settings: aa.Settings = None,
 ) -> List[FitImaging]:
     """
     Returns a list of `FitImaging` objects from a `PyAutoFit` loaded directory `Fit` or sqlite database `Fit` object.
@@ -26,7 +26,7 @@ def _fit_imaging_from(
 
     - The imaging data, noise-map, PSF and settings as .fits files (e.g. `dataset/data.fits`).
     - The mask used to mask the `Imaging` data structure in the fit (`dataset.fits[hdu=0]`).
-    - The settings of inversions used by the fit (`dataset/settings_inversion.json`).
+    - The settings of inversions used by the fit (`dataset/settings.json`).
 
     Each individual attribute can be loaded from the database via the `fit.value()` method.
 
@@ -37,7 +37,7 @@ def _fit_imaging_from(
     is instead used to load lists of the data, noise-map, PSF and mask and combine them into a list of
     `FitImaging` objects.
 
-    The settings of an inversion can be overwritten by inputting a `settings_inversion` object, for example
+    The settings of an inversion can be overwritten by inputting a `settings` object, for example
     if you want to use a grid with a different inversion solver.
 
     Parameters
@@ -48,8 +48,8 @@ def _fit_imaging_from(
     instance
         A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
         randomly from the PDF).
-    settings_inversion
-        Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
+    settings
+        Optionally overwrite the `Settings` of the `Inversion` object that is created from the fit.
     """
 
     from autogalaxy.imaging.fit_imaging import FitImaging
@@ -62,7 +62,7 @@ def _fit_imaging_from(
 
     adapt_images_list = agg_util.adapt_images_from(fit=fit)
 
-    settings_inversion = settings_inversion or fit.value(name="settings_inversion")
+    settings = settings or fit.value(name="settings")
 
     fit_dataset_list = []
 
@@ -79,7 +79,7 @@ def _fit_imaging_from(
                 galaxies=galaxies,
                 dataset_model=dataset_model,
                 adapt_images=adapt_images,
-                settings_inversion=settings_inversion,
+                settings=settings,
             )
         )
 
@@ -90,7 +90,7 @@ class FitImagingAgg(af.AggBase):
     def __init__(
         self,
         aggregator: af.Aggregator,
-        settings_inversion: Optional[aa.SettingsInversion] = None,
+        settings: Optional[aa.Settings] = None,
     ):
         """
             Interfaces with an `PyAutoFit` aggregator object to create instances of `FitImaging` objects from the results
@@ -101,7 +101,7 @@ def __init__(
 
             - The imaging data, noise-map, PSF and settings as .fits files (e.g. `dataset/data.fits`).
             - The mask used to mask the `Imaging` data structure in the fit (`dataset.fits[hdu=0]`).
-            - The settings of inversions used by the fit (`dataset/settings_inversion.json`).
+            - The settings of inversions used by the fit (`dataset/settings.json`).
 
             The `aggregator` contains the path to each of these files, and they can be loaded individually. This class
             can load them all at once and create an `FitImaging` object via the `_fit_imaging_from` method.
@@ -123,19 +123,19 @@ def __init__(
             ----------
             aggregator
                 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.
+            settings
+                Optionally overwrite the `Settings` of the `Inversion` object that is created from the fit.
 
         Parameters
         ----------
         aggregator
             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.
+        settings
+            Optionally overwrite the `Settings` of the `Inversion` object that is created from the fit.
         """
         super().__init__(aggregator=aggregator)
 
-        self.settings_inversion = settings_inversion
+        self.settings = settings
 
     def object_via_gen_from(
         self, fit, instance: Optional[af.ModelInstance] = None
@@ -157,5 +157,5 @@ def object_via_gen_from(
         return _fit_imaging_from(
             fit=fit,
             instance=instance,
-            settings_inversion=self.settings_inversion,
+            settings=self.settings,
         )