Merge pull request #348 from Jammy2211/feature/source_plane_images

Feature/source plane images

Author: Jammy2211

autolens/analysis/plotter_interface.py

@@ -1,6 +1,11 @@
+import ast
+import numpy as np
+
+from autoconf import conf
 from autoconf.fitsable import hdu_list_for_output_from
 
 import autoarray as aa
+import autogalaxy as ag
 import autogalaxy.plot as aplt
 
 from autogalaxy.analysis.plotter_interface import plot_setting
@@ -86,11 +91,54 @@ def should_plot(name):
                     "deflections_y",
                     "deflections_x",
                 ],
-                header_dict=grid.mask.header_dict,
+                header_dict=grid_zoom.mask.header_dict,
             )
 
             hdu_list.writeto(self.image_path / "tracer.fits", overwrite=True)
 
+        if should_plot("fits_source_plane_images"):
+
+            shape_native = conf.instance["visualize"]["plots"]["tracer"][
+                "fits_source_plane_shape"
+            ]
+            shape_native = ast.literal_eval(shape_native)
+
+            zoom = aa.Zoom2D(mask=grid.mask)
+            mask = zoom.mask_2d_from(buffer=1)
+            grid_source_plane = aa.Grid2D.from_extent(
+                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"]
+
+            for i, plane in enumerate(tracer.planes[1:]):
+
+                if plane.has(cls=ag.LightProfile):
+
+                    image = plane.image_2d_from(
+                        grid=grid_source_plane,
+                    ).native
+
+                else:
+
+                    image = np.zeros(grid_source_plane.shape_native)
+
+                image_list.append(image)
+                ext_name_list.append(f"source_plane_image_{i+1}")
+
+            hdu_list = hdu_list_for_output_from(
+                values_list=image_list,
+                ext_name_list=ext_name_list,
+                header_dict=grid_source_plane.mask.header_dict,
+            )
+
+            hdu_list.writeto(
+                self.image_path / "source_plane_images.fits", overwrite=True
+            )
+
     def image_with_positions(self, image: aa.Array2D, positions: aa.Grid2DIrregular):
         """
         Visualizes the positions of a model-fit, where these positions are used to resample lens models where

autolens/config/visualize/plots.yaml

@@ -34,6 +34,8 @@ tracer:                                    # Settings for plots of tracers (e.g.
   subplot_tracer: true                     # Plot subplot of all quantities in each tracer (e.g. images, convergence)?
   subplot_galaxies_images: false           # Plot subplot of the image of each plane in the tracer?
   fits_tracer: false                       # Output tracer.fits file of tracer's convergence, potential, deflections_y and deflections_x?
+  fits_source_plane_images: false          # Output source_plane_images.fits file of the source-plane image (light profiles only) of each galaxy in the tracer?
+  fits_source_plane_shape: (100, 100)      # The shape of the source-plane image output in the fits_source_plane_images.fits file.
 
 inversion:                                 # Settings for plots of inversions (e.g. InversionPlotter).
   subplot_inversion: true                  # Plot subplot of all quantities in each inversion (e.g. reconstrucuted image, reconstruction)?

autolens/imaging/plot/fit_imaging_plotters.py

@@ -292,8 +292,6 @@ def figures_2d_of_planes(
 
                 elif self.tracer.planes[plane_index].has(cls=aa.Pixelization):
 
-                    pix = self.tracer.planes[plane_index].cls_list_from(cls=aa.Pixelization)[0]
-
                     inversion_plotter = self.inversion_plotter_of_plane(
                         plane_index=plane_index
                     )

test_autolens/config/visualize.yaml

@@ -99,4 +99,6 @@ plots:
   tracer:
     subplot_tracer: true
     subplot_galaxies_images: true
-    fits_tracer: true                       # Output tracer.fits file of tracer's convergence, potential, deflections_y and deflections_x?
+    fits_tracer: true                       # Output tracer.fits file of tracer's convergence, potential, deflections_y and deflections_x?
+    fits_source_plane_images: true           # Output source_plane_images.fits file of the source-plane image (light profiles only) of each galaxy in the tracer?
+    fits_source_plane_shape: (3, 3)      # The shape of the source-plane image output in the fits_source_plane_images.fits file.