Overhaul plot styling and extract fits_* output functions

- All colormap defaults changed from "jet" to None (resolves to
  the configured default colormap via autoarray)
- save_tracer_fits() and save_source_plane_images_fits() extracted
  to tracer_plots.py; plotter.py reduced to one-liner delegate calls
- Unused imports (ast, numpy, conf, hdu_list_for_output_from) removed
  from plotter.py

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

Author: 2 people

autolens/analysis/plotter.py

@@ -1,10 +1,6 @@
-import ast
 import numpy as np
 from typing import Optional
 
-from autoconf import conf
-from autoconf.fitsable import hdu_list_for_output_from
-
 import autoarray as aa
 import autogalaxy as ag
 
@@ -13,7 +9,11 @@
 from autogalaxy.analysis.plotter import Plotter as AgPlotter
 
 from autolens.lens.tracer import Tracer
-from autolens.lens.plot.tracer_plots import subplot_galaxies_images
+from autolens.lens.plot.tracer_plots import (
+    subplot_galaxies_images,
+    save_tracer_fits,
+    save_source_plane_images_fits,
+)
 from autoarray.plot.array import plot_array
 
 
@@ -63,72 +63,10 @@ def should_plot(name):
             )
 
         if should_plot("fits_tracer"):
-
-            zoom = aa.Zoom2D(mask=grid.mask)
-            mask = zoom.mask_2d_from(buffer=1)
-            grid_zoom = aa.Grid2D.from_mask(mask=mask)
-
-            image_list = [
-                tracer.convergence_2d_from(grid=grid_zoom).native,
-                tracer.potential_2d_from(grid=grid_zoom).native,
-                tracer.deflections_yx_2d_from(grid=grid_zoom).native[:, :, 0],
-                tracer.deflections_yx_2d_from(grid=grid_zoom).native[:, :, 1],
-            ]
-
-            hdu_list = hdu_list_for_output_from(
-                values_list=[image_list[0].mask.astype("float")] + image_list,
-                ext_name_list=[
-                    "mask",
-                    "convergence",
-                    "potential",
-                    "deflections_y",
-                    "deflections_x",
-                ],
-                header_dict=grid_zoom.mask.header_dict,
-            )
-
-            hdu_list.writeto(self.image_path / "tracer.fits", overwrite=True)
+            save_tracer_fits(tracer=tracer, grid=grid, output_path=self.image_path)
 
         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)
-            )
-
-            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
-            )
+            save_source_plane_images_fits(tracer=tracer, grid=grid, output_path=self.image_path)
 
     def image_with_positions(self, image: aa.Array2D, positions: aa.Grid2DIrregular):
         """

autolens/imaging/plot/fit_imaging_plots.py

@@ -6,7 +6,7 @@
 import autogalaxy as ag
 
 from autoarray.plot.array import plot_array, _zoom_array_2d
-from autoarray.plot.utils import save_figure
+from autoarray.plot.utils import save_figure, hide_unused_axes
 from autoarray.plot.utils import numpy_lines as _to_lines
 from autogalaxy.plot.plot_utils import _critical_curves_from, _caustics_from
 
@@ -39,7 +39,7 @@ def _get_source_vmax(fit):
 
 
 def _plot_source_plane(fit, ax, plane_index, zoom_to_brightest=True,
-                       colormap="jet", use_log10=False):
+                       colormap=None, use_log10=False):
     """
     Plot the source-plane image (or a blank inversion placeholder) into an axes.
 
@@ -94,7 +94,7 @@ def subplot_fit(
     fit,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
     plane_index: Optional[int] = None,
 ):
     """
@@ -214,6 +214,7 @@ def subplot_fit(
     _plot_source_plane(fit, axes_flat[11], final_plane_index, zoom_to_brightest=False,
                        colormap=colormap)
 
+    hide_unused_axes(axes_flat)
     plt.tight_layout()
     save_figure(fig, path=output_path, filename=f"subplot_fit{plane_index_tag}", format=output_format)
 
@@ -222,7 +223,7 @@ def subplot_fit_x1_plane(
     fit,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
 ):
     """
     Produce a 6-panel subplot for a single-plane tracer imaging fit.
@@ -286,7 +287,7 @@ def subplot_fit_log10(
     fit,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
     plane_index: Optional[int] = None,
 ):
     """
@@ -395,7 +396,7 @@ def subplot_fit_log10_x1_plane(
     fit,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
 ):
     """
     Produce a 6-panel log10 subplot for a single-plane tracer imaging fit.
@@ -456,7 +457,7 @@ def subplot_of_planes(
     fit,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
     plane_index: Optional[int] = None,
 ):
     """
@@ -524,7 +525,7 @@ def subplot_tracer_from_fit(
     fit,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
 ):
     """
     Produce a 9-panel tracer subplot derived from a `FitImaging` object.
@@ -600,7 +601,7 @@ def subplot_fit_combined(
     fit_list: List,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
 ):
     """
     Produce a combined multi-row subplot for a list of `FitImaging` objects.
@@ -682,7 +683,7 @@ def subplot_fit_combined_log10(
     fit_list: List,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
 ):
     """
     Produce a combined log10 multi-row subplot for a list of `FitImaging` objects.

autolens/interferometer/plot/fit_interferometer_plots.py

@@ -37,7 +37,7 @@ def _plot_yx(y, x, ax, title, xlabel="", ylabel=""):
 
 
 def _plot_source_plane(fit, ax, plane_index, zoom_to_brightest=True,
-                       colormap="jet", use_log10=False):
+                       colormap=None, use_log10=False):
     """
     Plot the source-plane image (or a blank inversion placeholder) into an axes.
 
@@ -88,7 +88,7 @@ def subplot_fit(
     fit,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
 ):
     """
     Produce a 12-panel subplot summarising an interferometer fit.
@@ -197,7 +197,7 @@ def subplot_fit_real_space(
     fit,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
 ):
     """
     Produce a real-space subplot for an interferometer fit.

autolens/lens/plot/sensitivity_plots.py

@@ -16,7 +16,7 @@ def subplot_tracer_images(
     source_image,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
     use_log10: bool = False,
 ):
     """
@@ -120,7 +120,7 @@ def subplot_sensitivity(
     data_subtracted,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
     use_log10: bool = False,
 ):
     """
@@ -248,7 +248,7 @@ def subplot_figures_of_merit_grid(
     result,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
     use_log_evidences: bool = True,
     remove_zeros: bool = True,
 ):

autolens/lens/plot/subhalo_plots.py

@@ -12,7 +12,7 @@ def subplot_detection_imaging(
     fit_imaging_with_subhalo,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
     use_log10: bool = False,
     use_log_evidences: bool = True,
     relative_to_value: float = 0.0,
@@ -103,7 +103,7 @@ def subplot_detection_fits(
     fit_imaging_with_subhalo,
     output_path: Optional[str] = None,
     output_format: str = "png",
-    colormap: str = "jet",
+    colormap: Optional[str] = None,
 ):
     """
     Produce a 6-panel subplot comparing imaging fits with and without a subhalo.