Feature/dspl fit imaging

autolens/analysis/plotter_interface.py

@@ -109,8 +109,6 @@ def should_plot(name):
                 extent=mask.geometry.extent, shape_native=tuple(shape_native)
             )
 
-            print(grid_source_plane)
-
             image_list = [grid_source_plane.mask.astype("float")]
             ext_name_list = ["mask"]
 

autolens/imaging/model/plotter_interface.py

@@ -59,12 +59,25 @@ def should_plot(name):
             fit=fit, mat_plot_2d=mat_plot_2d, include_2d=self.include_2d
         )
 
+        plane_indexes_to_plot = [i for i in fit.tracer.plane_indexes_with_images if i != 0]
+
         if should_plot("subplot_fit"):
-            fit_plotter.subplot_fit()
+
+            # This loop means that multiple subplot_fit objects are output for a double source plane lens.
+
+            if len(plane_indexes_to_plot) > 1:
+                for plane_index in plane_indexes_to_plot:
+                    fit_plotter.subplot_fit(plane_index=plane_index)
+            else:
+                fit_plotter.subplot_fit()
 
         if should_plot("subplot_fit_log10"):
             try:
-                fit_plotter.subplot_fit_log10()
+                if len(plane_indexes_to_plot) > 1:
+                    for plane_index in plane_indexes_to_plot:
+                        fit_plotter.subplot_fit_log10(plane_index=plane_index)
+                else:
+                    fit_plotter.subplot_fit_log10()
             except ValueError:
                 pass
 

autolens/imaging/plot/fit_imaging_plotters.py

@@ -422,7 +422,7 @@ def subplot(
             auto_labels=AutoLabels(filename=auto_filename),
         )
 
-    def subplot_fit(self):
+    def subplot_fit(self, plane_index: Optional[int] = None):
         """
         Standard subplot of the attributes of the plotter's `FitImaging` object.
         """
@@ -444,21 +444,22 @@ def subplot_fit(self):
         # If the lens light is not included the subplot index does not increase, so we must manually set it to 4
         self.mat_plot_2d.subplot_index = 6
 
-        final_plane_index = len(self.fit.tracer.planes) - 1
+        plane_index_tag = "" if plane_index is None else f"_{plane_index}"
+
+        plane_index = len(self.fit.tracer.planes) - 1 if plane_index is None else plane_index
 
         self.mat_plot_2d.cmap.kwargs["vmin"] = 0.0
 
         self.set_title(label="Lens Light Subtracted Image")
-        self.figures_2d_of_planes(plane_index=final_plane_index, subtracted_image=True, use_source_vmax=True)
+        self.figures_2d_of_planes(plane_index=plane_index, subtracted_image=True, use_source_vmax=True)
 
         self.set_title(label="Source Model Image (Image Plane)")
-        self.figures_2d_of_planes(plane_index=final_plane_index, model_image=True, use_source_vmax=True)
+        self.figures_2d_of_planes(plane_index=plane_index, model_image=True, use_source_vmax=True)
 
         self.mat_plot_2d.cmap.kwargs.pop("vmin")
 
         self.set_title(label="Source Plane (Zoomed)")
-        self.figures_2d_of_planes(plane_index=final_plane_index, plane_image=True, use_source_vmax=True)
-
+        self.figures_2d_of_planes(plane_index=plane_index, plane_image=True, use_source_vmax=True)
 
         self.set_title(label=None)
 
@@ -480,7 +481,7 @@ def subplot_fit(self):
 
         self.set_title(label="Source Plane (No Zoom)")
         self.figures_2d_of_planes(
-            plane_index=final_plane_index,
+            plane_index=plane_index,
             plane_image=True,
             zoom_to_brightest=False,
             use_source_vmax=True
@@ -489,11 +490,11 @@ def subplot_fit(self):
         self.set_title(label=None)
 
         self.mat_plot_2d.output.subplot_to_figure(
-            auto_filename="subplot_fit"
+            auto_filename=f"subplot_fit{plane_index_tag}"
         )
         self.close_subplot_figure()
 
-    def subplot_fit_log10(self):
+    def subplot_fit_log10(self, plane_index: Optional[int] = None):
         """
         Standard subplot of the attributes of the plotter's `FitImaging` object.
         """
@@ -530,21 +531,22 @@ def subplot_fit_log10(self):
         # If the lens light is not included the subplot index does not increase, so we must manually set it to 4
         self.mat_plot_2d.subplot_index = 6
 
-        final_plane_index = len(self.fit.tracer.planes) - 1
+        plane_index_tag = "" if plane_index is None else f"_{plane_index}"
+
+        plane_index = len(self.fit.tracer.planes) - 1 if plane_index is None else plane_index
 
         self.mat_plot_2d.cmap.kwargs["vmin"] = 0.0
 
         self.set_title(label="Lens Light Subtracted Image")
-        self.figures_2d_of_planes(plane_index=final_plane_index, subtracted_image=True, use_source_vmax=True)
+        self.figures_2d_of_planes(plane_index=plane_index, subtracted_image=True, use_source_vmax=True)
 
         self.set_title(label="Source Model Image (Image Plane)")
-        self.figures_2d_of_planes(plane_index=final_plane_index, model_image=True, use_source_vmax=True)
+        self.figures_2d_of_planes(plane_index=plane_index, model_image=True, use_source_vmax=True)
 
         self.mat_plot_2d.cmap.kwargs.pop("vmin")
 
         self.set_title(label="Source Plane (Zoomed)")
-        self.figures_2d_of_planes(plane_index=final_plane_index, plane_image=True, use_source_vmax=True)
-
+        self.figures_2d_of_planes(plane_index=plane_index, plane_image=True, use_source_vmax=True)
 
         self.set_title(label=None)
 
@@ -570,7 +572,7 @@ def subplot_fit_log10(self):
 
         self.set_title(label="Source Plane (No Zoom)")
         self.figures_2d_of_planes(
-            plane_index=final_plane_index,
+            plane_index=plane_index,
             plane_image=True,
             zoom_to_brightest=False,
             use_source_vmax=True
@@ -579,7 +581,7 @@ def subplot_fit_log10(self):
         self.set_title(label=None)
 
         self.mat_plot_2d.output.subplot_to_figure(
-            auto_filename="subplot_fit_log10"
+            auto_filename=f"subplot_fit_log10{plane_index_tag}"
         )
         self.close_subplot_figure()
 

autolens/lens/tracer.py

@@ -802,7 +802,18 @@ def cls_list_from(self, cls: Type) -> List:
         return cls_list
 
     @property
-    def plane_indexes_with_pixelizations(self):
+    def plane_indexes_with_pixelizations(self) -> List[int]:
+        """
+        Returns a list of integer indexes of the indexes of planes which use a `Pixelization` to reconstruct the
+        source galaxy.
+
+        This list is used to set up an inversion, whereby each pixelization is extracted from the tracer with its
+        corresponding ray-traced grid and passed to the PyAutoArray `inversion` module.
+
+        Returns
+        -------
+        The list of integer indexes of the planes which use a `Pixelization` to reconstruct the source galaxy.
+        """
         plane_indexes_with_inversions = [
             plane_index if plane.has(cls=aa.Pixelization) else None
             for (plane_index, plane) in enumerate(self.planes)
@@ -813,6 +824,32 @@ def plane_indexes_with_pixelizations(self):
             if plane_index is not None
         ]
 
+    @property
+    def plane_indexes_with_images(self):
+        """
+        Returns a list of integer indexes of the indexes of planes which create an image, meaning they either
+        have a `LightProfile` or `Pixelization`.
+
+        This list is used to visualize double source plane lenses, whereby a fit for every plane with a
+        `LightProfile` or `Pixelization` is created.
+
+        Returns
+        -------
+        The list of integer indexes of the planes which create an image.
+        """
+        plane_indexes_with_images = [
+            plane_index if plane.has(cls=ag.LightProfile) else None
+            for (plane_index, plane) in enumerate(self.planes)
+        ] + self.plane_indexes_with_pixelizations
+
+        plane_indexes_with_images = list(dict.fromkeys(plane_indexes_with_images))
+
+        return [
+            plane_index
+            for plane_index in plane_indexes_with_images
+            if plane_index is not None
+        ]
+
     @property
     def perform_inversion(self) -> bool:
         """