Merge pull request #258 from Jammy2211/feature/image_mesh_jax

Feature/image mesh jax

Author: Jammy2211

autogalaxy/__init__.py

@@ -56,7 +56,7 @@
 
 from .analysis import model_util
 from .analysis.adapt_images.adapt_images import AdaptImages
-from .analysis.adapt_images.adapt_image_maker import AdaptImageMaker
+from .analysis.adapt_images.adapt_images import galaxy_name_image_dict_via_result_from
 from . import aggregator as agg
 from . import exc
 from . import plot

autogalaxy/aggregator/agg_util.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
+import numpy as np
 from typing import List, Optional
 
+from autoconf.fitsable import flip_for_ds9_from
 from autoconf.fitsable import ndarray_via_hdu_from
 
 import autofit as af
@@ -140,9 +142,24 @@ def adapt_images_from(
 
             galaxy_name_image_dict[value.header["EXTNAME"].lower()] = adapt_image
 
+        galaxy_name_image_plane_mesh_grid_dict = {}
+
+        for i, value in enumerate(fit.value(name="adapt_image_plane_mesh_grids")[1:]):
+
+            adapt_image_plane_mesh_grid = aa.Grid2DIrregular(
+                values=flip_for_ds9_from(value.data.astype("float")),
+            )
+
+            galaxy_name_image_plane_mesh_grid_dict[value.header["EXTNAME"].lower()] = (
+                adapt_image_plane_mesh_grid
+            )
+
         instance = fit.model.instance_from_prior_medians(ignore_assertions=True)
 
-        adapt_images = AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)
+        adapt_images = AdaptImages(
+            galaxy_name_image_dict=galaxy_name_image_dict,
+            galaxy_name_image_plane_mesh_grid_dict=galaxy_name_image_plane_mesh_grid_dict,
+        )
 
         adapt_images = adapt_images.updated_via_instance_from(
             instance=instance,

autogalaxy/analysis/adapt_images/adapt_image_maker.py

@@ -11,11 +11,68 @@
     from autogalaxy.galaxy.galaxy import Galaxy
 
 
+def galaxy_name_image_dict_via_result_from(
+    result, use_model_images: bool = False
+) -> "AdaptImages":
+    """
+    Returns the adapt-images from a non-linear search result.
+
+    For model-fitting, the adapt-images are typically setup using the maximum log likelihood model of the
+    previous model-fit. This means the model-fitting is used to cleanly deblend the light of the different
+    galaxies in the image (e.g. separate the lens light from the source light).
+
+    This method uses attributes of a result (e.g. dictionary mapping galaxy instances to their model-images)
+    to create the adapt-images.
+
+    This can use either:
+
+    - The model image of each galaxy in the best-fit model.
+    - The subtracted image of each galaxy in the best-fit model, where the subtracted image is the dataset
+      minus the model images of all other galaxies.
+
+    Certain models produce galaxy-images with negative flux values (e.g. a pixelization), which can cause
+    numerical issues with the adaptive schemes. To prevent this, we set a minimum flux value for each
+    galaxy-image, which is a fraction of the maximum flux value of that image defined via a config file.
+
+    Parameters
+    ----------
+    result
+        The result of a previous model-fit, which contains the model-image of each galaxy.
+    use_model_images
+        If True, the model images of the galaxies are used to create the adapt images. If False, the subtracted
+        images of the galaxies are used.
+
+    Returns
+    -------
+    The adapt-images, which are the model-image of each galaxy inferred via the previous model-fit.
+    """
+    adapt_minimum_percent = conf.instance["general"]["adapt"]["adapt_minimum_percent"]
+
+    galaxy_name_image_dict = {}
+
+    for path, galaxy in result.path_galaxy_tuples:
+        if use_model_images:
+            galaxy_image = result.model_image_galaxy_dict[path]
+        else:
+            galaxy_image = result.subtracted_signal_to_noise_map_galaxy_dict[path]
+
+        minimum_galaxy_value = adapt_minimum_percent * np.max(galaxy_image.array)
+        galaxy_image[galaxy_image < minimum_galaxy_value] = minimum_galaxy_value
+
+        galaxy_name_image_dict[path] = galaxy_image
+
+    return galaxy_name_image_dict
+
+
 class AdaptImages:
     def __init__(
         self,
         galaxy_image_dict: Optional[Dict[Galaxy, aa.Array2D]] = None,
         galaxy_name_image_dict: Optional[Dict[Tuple[str, ...], aa.Array2D]] = None,
+        galaxy_image_plane_mesh_grid_dict: Optional[Dict[Galaxy, aa.Array2D]] = None,
+        galaxy_name_image_plane_mesh_grid_dict: Optional[
+            Dict[Tuple[str, ...], aa.Grid2DIrregular]
+        ] = None,
     ):
         """
         Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the
@@ -54,6 +111,11 @@ def __init__(
         self.galaxy_image_dict = galaxy_image_dict
         self.galaxy_name_image_dict = galaxy_name_image_dict
 
+        self.galaxy_image_plane_mesh_grid_dict = galaxy_image_plane_mesh_grid_dict
+        self.galaxy_name_image_plane_mesh_grid_dict = (
+            galaxy_name_image_plane_mesh_grid_dict
+        )
+
     @property
     def mask(self) -> aa.Mask2D:
         """
@@ -85,59 +147,6 @@ def model_image(self) -> aa.Array2D:
 
         return adapt_model_image
 
-    @classmethod
-    def from_result(cls, result, use_model_images: bool = False) -> "AdaptImages":
-        """
-        Returns the adapt-images from a non-linear search result.
-
-        For model-fitting, the adapt-images are typically setup using the maximum log likelihood model of the
-        previous model-fit. This means the model-fitting is used to cleanly deblend the light of the different
-        galaxies in the image (e.g. separate the lens light from the source light).
-
-        This method uses attributes of a result (e.g. dictionary mapping galaxy instances to their model-images)
-        to create the adapt-images.
-
-        This can use either:
-
-        - The model image of each galaxy in the best-fit model.
-        - The subtracted image of each galaxy in the best-fit model, where the subtracted image is the dataset
-          minus the model images of all other galaxies.
-
-        Certain models produce galaxy-images with negative flux values (e.g. a pixelization), which can cause
-        numerical issues with the adaptive schemes. To prevent this, we set a minimum flux value for each
-        galaxy-image, which is a fraction of the maximum flux value of that image defined via a config file.
-
-        Parameters
-        ----------
-        result
-            The result of a previous model-fit, which contains the model-image of each galaxy.
-        use_model_images
-            If True, the model images of the galaxies are used to create the adapt images. If False, the subtracted
-            images of the galaxies are used.
-
-        Returns
-        -------
-        The adapt-images, which are the model-image of each galaxy inferred via the previous model-fit.
-        """
-        adapt_minimum_percent = conf.instance["general"]["adapt"][
-            "adapt_minimum_percent"
-        ]
-
-        galaxy_name_image_dict = {}
-
-        for path, galaxy in result.path_galaxy_tuples:
-            if use_model_images:
-                galaxy_image = result.model_image_galaxy_dict[path]
-            else:
-                galaxy_image = result.subtracted_signal_to_noise_map_galaxy_dict[path]
-
-            minimum_galaxy_value = adapt_minimum_percent * np.max(galaxy_image.array)
-            galaxy_image[galaxy_image < minimum_galaxy_value] = minimum_galaxy_value
-
-            galaxy_name_image_dict[path] = galaxy_image
-
-        return AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)
-
     def updated_via_instance_from(self, instance, mask=None) -> "AdaptImages":
         """
         Returns adapt-images which have been updated to map galaxy instances instead of galaxy names.
@@ -168,16 +177,37 @@ def updated_via_instance_from(self, instance, mask=None) -> "AdaptImages":
         """
         from autogalaxy.galaxy.galaxy import Galaxy
 
-        galaxy_image_dict = {}
+        galaxy_image_dict = None
+
+        if self.galaxy_name_image_dict is not None:
+
+            galaxy_image_dict = {}
+
+            for galaxy_name, galaxy in instance.path_instance_tuples_for_class(Galaxy):
+                galaxy_name = str(galaxy_name)
 
-        for galaxy_name, galaxy in instance.path_instance_tuples_for_class(Galaxy):
-            galaxy_name = str(galaxy_name)
+                if galaxy_name in self.galaxy_name_image_dict:
+                    galaxy_image_dict[galaxy] = self.galaxy_name_image_dict[galaxy_name]
 
-            if galaxy_name in self.galaxy_name_image_dict:
-                galaxy_image_dict[galaxy] = self.galaxy_name_image_dict[galaxy_name]
+            if mask is not None:
+                for key, image in galaxy_image_dict.items():
+                    galaxy_image_dict[key] = aa.Array2D(values=image, mask=mask)
 
-        if mask is not None:
-            for key, image in galaxy_image_dict.items():
-                galaxy_image_dict[key] = aa.Array2D(values=image, mask=mask)
+        galaxy_image_plane_mesh_grid_dict = None
 
-        return AdaptImages(galaxy_image_dict=galaxy_image_dict)
+        if self.galaxy_name_image_plane_mesh_grid_dict is not None:
+
+            galaxy_image_plane_mesh_grid_dict = {}
+
+            for galaxy_name, galaxy in instance.path_instance_tuples_for_class(Galaxy):
+                galaxy_name = str(galaxy_name)
+
+                if galaxy_name in self.galaxy_name_image_plane_mesh_grid_dict:
+                    galaxy_image_plane_mesh_grid_dict[galaxy] = (
+                        self.galaxy_name_image_plane_mesh_grid_dict[galaxy_name]
+                    )
+
+        return AdaptImages(
+            galaxy_image_dict=galaxy_image_dict,
+            galaxy_image_plane_mesh_grid_dict=galaxy_image_plane_mesh_grid_dict,
+        )

autogalaxy/analysis/adapt_images/adapt_images.py

@@ -7,7 +7,6 @@
 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.cosmology.lensing import LensingCosmology
 from autogalaxy.analysis.analysis.analysis import Analysis
@@ -22,7 +21,7 @@ class AnalysisDataset(Analysis):
     def __init__(
         self,
         dataset: Union[aa.Imaging, aa.Interferometer],
-        adapt_image_maker: Optional[AdaptImageMaker] = None,
+        adapt_images: Optional[AdaptImages] = None,
         cosmology: LensingCosmology = None,
         settings_inversion: aa.SettingsInversion = None,
         preloads: aa.Preloads = None,
@@ -41,8 +40,8 @@ def __init__(
         ----------
         dataset
             The dataset that is the model is fitted too.
-        adapt_image_maker
-            Makes the adapt-model image and galaxies images of a previous result in a model-fitting pipeline, which are
+        adapt_images
+            The adapt-model image and galaxies images of a previous result in a model-fitting pipeline, which are
             used by certain classes for adapting the analysis to the properties of the dataset.
         cosmology
             The Cosmology assumed for this analysis.
@@ -61,26 +60,12 @@ def __init__(
         )
 
         self.dataset = dataset
-        self.adapt_image_maker = adapt_image_maker
-        self._adapt_images = None
+        self.adapt_images = adapt_images
 
         self.settings_inversion = settings_inversion or aa.SettingsInversion()
 
         self.title_prefix = title_prefix
 
-    @property
-    def adapt_images(self):
-
-        if self._adapt_images is not None:
-            return self._adapt_images
-
-        if self.adapt_image_maker is None:
-            return None
-
-        self._adapt_images = self.adapt_image_maker.adapt_images
-
-        return self._adapt_images
-
     def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
         """
         This function is called immediately before the non-linear search begins and performs final tasks and checks
@@ -119,10 +104,6 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
 
         self.dataset.grids.border_relocator
 
-        if self.adapt_image_maker is not None:
-            if not paths.is_complete:
-                self._adapt_images = self.adapt_image_maker.adapt_images
-
         super().modify_before_fit(paths=paths, model=model)
 
         return self

autogalaxy/analysis/analysis/dataset.py

@@ -109,7 +109,6 @@ def source_custom_model_from(result: Result, source_is_model: bool = False) -> a
     if source_is_model:
         pixelization = af.Model(
             aa.Pixelization,
-            image_mesh=result.instance.galaxies.source.pixelization.image_mesh,
             mesh=result.instance.galaxies.source.pixelization.mesh,
             regularization=result.model.galaxies.source.pixelization.regularization,
         )
@@ -122,7 +121,6 @@ def source_custom_model_from(result: Result, source_is_model: bool = False) -> a
 
     pixelization = af.Model(
         aa.Pixelization,
-        image_mesh=result.instance.galaxies.source.pixelization.image_mesh,
         mesh=result.instance.galaxies.source.pixelization.mesh,
         regularization=result.instance.galaxies.source.pixelization.regularization,
     )