Merge branch 'main' into feature/jax_wrapper

Author: Jammy2211

autolens/__init__.py

@@ -113,6 +113,11 @@
 from . import util
 
 from autoconf import conf
+from autoconf.fitsable import ndarray_via_hdu_from
+from autoconf.fitsable import ndarray_via_fits_from
+from autoconf.fitsable import header_obj_from
+from autoconf.fitsable import output_to_fits
+from autoconf.fitsable import hdu_list_for_output_from
 
 conf.instance.register(__file__)
 

autolens/aggregator/__init__.py

@@ -21,3 +21,10 @@
 from autogalaxy.aggregator.ellipse.fit_ellipse import FitEllipseAgg
 
 from autolens.aggregator.subhalo import SubhaloAgg
+
+from autolens.aggregator.subplot import SubplotDataset as subplot_dataset
+from autolens.aggregator.subplot import SubplotTracer as subplot_tracer
+from autolens.aggregator.subplot import SubplotFit as subplot_fit
+from autolens.aggregator.subplot import SubplotFitLog10 as subplot_fit_log10
+from autolens.aggregator.subplot import FITSTracer as fits_tracer
+from autolens.aggregator.subplot import FITSFit as fits_fit

autolens/aggregator/fit_imaging.py

@@ -18,12 +18,13 @@ def _fit_imaging_from(
     settings_inversion: aa.SettingsInversion = None,
 ) -> List[FitImaging]:
     """
-    Returns a list of `FitImaging` object from a `PyAutoFit` sqlite database `Fit` object.
+    Returns a list of `FitImaging` object from a `PyAutoFit` loaded directory `Fit` or sqlite database `Fit` object.
 
-    The results of a model-fit can be stored in a sqlite database, including the following attributes of the fit:
+    The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
+    attributes of the fit:
 
     - 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/mask.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`).
 
     Each individual attribute can be loaded from the database via the `fit.value()` method.
@@ -41,7 +42,8 @@ def _fit_imaging_from(
     Parameters
     ----------
     fit
-        A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
+        A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
+        an output directory or from an sqlite database..
     instance
         A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
         randomly from the PDF).
@@ -67,7 +69,6 @@ def _fit_imaging_from(
         dataset_model_list,
         adapt_images_list,
     ):
-
         fit_dataset_list.append(
             FitImaging(
                 dataset=dataset,
@@ -88,41 +89,42 @@ def __init__(
         settings_inversion: Optional[aa.SettingsInversion] = None,
     ):
         """
-        Interfaces with an `PyAutoFit` aggregator object to create instances of `FitImaging` objects from the results
-        of a model-fit.
-
-        The results of a model-fit can be stored in a sqlite database, including the following attributes of the fit:
-
-        - 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/mask.fits`).
-        - The settings of inversions used by the fit (`dataset/settings_inversion.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.
-
-        This class's methods returns generators which create the instances of the `FitImaging` objects. This ensures
-        that large sets of results can be efficiently loaded from the hard-disk and do not require storing all
-        `FitImaging` instances in the memory at once.
-
-        For example, if the `aggregator` contains 3 model-fits, this class can be used to create a generator which
-        creates instances of the corresponding 3 `FitImaging` objects.
-
-        If multiple `FitImaging` objects were fitted simultaneously via analysis summing, the `fit.child_values()` method
-        is instead used to load lists of the data, noise-map, PSF and mask and combine them into a list of
-        `FitImaging` objects.
-
-        This can be done manually, but this object provides a more concise API.
-
-        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.
-        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.
+            Interfaces with an `PyAutoFit` aggregator object to create instances of `FitImaging` objects from the results
+            of a model-fit.
+
+            The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
+        attributes of the fit:
+
+            - 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 `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.
+
+            This class's methods returns generators which create the instances of the `FitImaging` objects. This ensures
+            that large sets of results can be efficiently loaded from the hard-disk and do not require storing all
+            `FitImaging` instances in the memory at once.
+
+            For example, if the `aggregator` contains 3 model-fits, this class can be used to create a generator which
+            creates instances of the corresponding 3 `FitImaging` objects.
+
+            If multiple `FitImaging` objects were fitted simultaneously via analysis summing, the `fit.child_values()` method
+            is instead used to load lists of the data, noise-map, PSF and mask and combine them into a list of
+            `FitImaging` objects.
+
+            This can be done manually, but this object provides a more concise API.
+
+            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.
+            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)
 
@@ -139,7 +141,8 @@ def object_via_gen_from(
         Parameters
         ----------
         fit
-            A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
+            A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
+        an output directory or from an sqlite database..
         instance
             A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
             randomly from the PDF).

autolens/aggregator/fit_interferometer.py

@@ -15,13 +15,13 @@
 def _fit_interferometer_from(
     fit: af.Fit,
     instance: Optional[af.ModelInstance] = None,
-    real_space_mask: Optional[aa.Mask2D] = None,
     settings_inversion: aa.SettingsInversion = None,
 ) -> List[FitInterferometer]:
     """
-    Returns a list of `FitInterferometer` objects from a `PyAutoFit` sqlite database `Fit` object.
+    Returns a list of `FitInterferometer` objects from a `PyAutoFit` loaded directory `Fit` or sqlite database `Fit` object.
 
-    The results of a model-fit can be stored in a sqlite database, including the following attributes of the fit:
+    The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
+    attributes of the fit:
 
     - The interferometer data, noise-map, uv-wavelengths and settings as .fits files (e.g. `dataset/data.fits`).
     - The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`).
@@ -43,7 +43,8 @@ def _fit_interferometer_from(
     Parameters
     ----------
     fit
-        A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
+        A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
+        an output directory or from an sqlite database..
     instance
         A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
         randomly from the PDF).
@@ -52,7 +53,6 @@ def _fit_interferometer_from(
     """
     dataset_list = _interferometer_from(
         fit=fit,
-        real_space_mask=real_space_mask,
     )
     tracer_list = _tracer_from(fit=fit, instance=instance)
     dataset_model_list = _dataset_model_from(fit=fit, instance=instance)
@@ -69,7 +69,6 @@ def _fit_interferometer_from(
         dataset_model_list,
         adapt_images_list,
     ):
-
         fit_dataset_list.append(
             FitInterferometer(
                 dataset=dataset,
@@ -88,45 +87,44 @@ def __init__(
         self,
         aggregator: af.Aggregator,
         settings_inversion: Optional[aa.SettingsInversion] = None,
-        real_space_mask: Optional[aa.Mask2D] = None,
     ):
         """
-        Interfaces with an `PyAutoFit` aggregator object to create instances of `FitInterferometer` objects from the
-        results of a model-fit.
+            Interfaces with an `PyAutoFit` aggregator object to create instances of `FitInterferometer` objects from the
+            results of a model-fit.
 
-        The results of a model-fit can be stored in a sqlite database, including the following attributes of the fit:
+            The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
+        attributes of the fit:
 
-        - The interferometer data, noise-map, uv-wavelengths and settings as .fits files (e.g. `dataset/data.fits`).
-        - The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`).
-        - The settings of inversions used by the fit (`dataset/settings_inversion.json`).
+            - The interferometer data, noise-map, uv-wavelengths and settings as .fits files (e.g. `dataset/data.fits`).
+            - The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`).
+            - The settings of inversions used by the fit (`dataset/settings_inversion.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 `FitInterferometer` object via the `_fit_interferometer_from` method.
+            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 `FitInterferometer` object via the `_fit_interferometer_from` method.
 
-        This class's methods returns generators which create the instances of the `FitInterferometer` objects. This ensures
-        that large sets of results can be efficiently loaded from the hard-disk and do not require storing all
-        `FitInterferometer` instances in the memory at once.
+            This class's methods returns generators which create the instances of the `FitInterferometer` objects. This ensures
+            that large sets of results can be efficiently loaded from the hard-disk and do not require storing all
+            `FitInterferometer` instances in the memory at once.
 
-        For example, if the `aggregator` contains 3 model-fits, this class can be used to create a generator which
-        creates instances of the corresponding 3 `FitInterferometer` objects.
+            For example, if the `aggregator` contains 3 model-fits, this class can be used to create a generator which
+            creates instances of the corresponding 3 `FitInterferometer` objects.
 
-        This can be done manually, but this object provides a more concise API.
+            This can be done manually, but this object provides a more concise API.
 
-        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.
-        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.
+            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.
+            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.real_space_mask = real_space_mask
 
     def object_via_gen_from(
         self, fit, instance: Optional[af.ModelInstance] = None
@@ -139,7 +137,8 @@ def object_via_gen_from(
             Parameters
             ----------
             fit
-                A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
+                A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
+        an output directory or from an sqlite database..
         instance
             A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
             randomly from the PDF).

autolens/aggregator/subplot.py

@@ -0,0 +1,100 @@
+from enum import Enum
+
+
+class FITSTracer(Enum):
+    """
+    The HDUs that can be extracted from the fit.fits file.
+    """
+
+    convergence = "CONVERGENCE"
+    potential = "POTENTIAL"
+    deflections_y = "DEFLECTIONS_Y"
+    deflections_x = "DEFLECTIONS_X"
+
+
+class FITSFit(Enum):
+    """
+    The HDUs that can be extracted from the fit.fits file.
+    """
+
+    model_image = "MODEL_IMAGE"
+    residual_map = "RESIDUAL_MAP"
+    normalized_residual_map = "NORMALIZED_RESIDUAL_MAP"
+    chi_squared_map = "CHI_SQUARED_MAP"
+
+
+class SubplotDataset(Enum):
+    """
+    The subplots that can be extracted from the subplot_fit image.
+
+    The values correspond to the position of the subplot in the 4x3 grid.
+    """
+
+    data = (0, 0)
+    data_log_10 = (1, 0)
+    noise_map = (2, 0)
+    psf = (0, 1)
+    psf_log_10 = (1, 1)
+    signal_to_noise_map = (2, 1)
+    over_sample_size_lp = (0, 2)
+    over_sample_size_pixelization = (1, 2)
+
+
+class SubplotTracer(Enum):
+    """
+    The subplots that can be extracted from the subplot_tracer image.
+
+    The values correspond to the position of the subplot in the 3x3 grid.
+    """
+
+    image = (0, 0)
+    source_image = (1, 0)
+    source_plane_image = (2, 0)
+    lens_light_image = (1, 0)
+    convergence = (1, 1)
+    potential = (1, 2)
+    magnification = (0, 2)
+    deflections_y = (1, 2)
+    deflections_x = (2, 2)
+
+
+class SubplotFit(Enum):
+    """
+    The subplots that can be extracted from the subplot_fit image.
+
+    The values correspond to the position of the subplot in the 4x3 grid.
+    """
+
+    data = (0, 0)
+    data_source_scale = (1, 0)
+    signal_to_noise_map = (2, 0)
+    model_image = (3, 0)
+    lens_light_model = (1, 0)
+    lens_light_subtracted_image = (1, 1)
+    source_model_image = (1, 2)
+    source_plane_image_zoom = (1, 3)
+    normalized_residual_map = (2, 0)
+    normalized_residual_map_one_sigma = (2, 1)
+    chi_squared_map = (2, 2)
+    source_plane_image = (2, 3)
+
+
+class SubplotFitLog10(Enum):
+    """
+    The subplots that can be extracted from the subplot_fit image.
+
+    The values correspond to the position of the subplot in the 4x3 grid.
+    """
+
+    data = (0, 0)
+    data_source_scale = (1, 0)
+    signal_to_noise_map = (2, 0)
+    model_image = (3, 0)
+    lens_light_model = (1, 0)
+    lens_light_subtracted_image = (1, 1)
+    source_model_image = (1, 2)
+    source_plane_image_zoom = (1, 3)
+    normalized_residual_map = (2, 0)
+    normalized_residual_map_one_sigma = (2, 1)
+    chi_squared_map = (2, 2)
+    source_plane_image = (2, 3)