Replace _save_subplot with autoarray's save_figure

- Remove _save_subplot from plot_utils.py (now just 3 autolens-specific
  helpers: _to_lines, _critical_curves_from, _caustics_from)
- All callers now import save_figure from autoarray.plot.plots.utils
- autoarray's save_figure updated to accept format as list or single string

https://claude.ai/code/session_01CzJBy8KvFXiNchoNdk5i9k

Author: claude

autolens/imaging/plot/fit_imaging_plots.py

@@ -6,9 +6,9 @@
 import autogalaxy as ag
 
 from autoarray.plot.plots.array import plot_array, _zoom_array_2d
+from autoarray.plot.plots.utils import save_figure
 from autolens.plot.plot_utils import (
     _to_lines,
-    _save_subplot,
     _critical_curves_from,
     _caustics_from,
 )
@@ -139,7 +139,7 @@ def subplot_fit(
                        colormap=colormap)
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, f"subplot_fit{plane_index_tag}", output_format)
+    save_figure(fig, path=output_path, filename=f"subplot_fit{plane_index_tag}", format=output_format)
 
 
 def subplot_fit_x1_plane(
@@ -177,7 +177,7 @@ def subplot_fit_x1_plane(
                colormap=colormap, vmin=-_abs_max, vmax=_abs_max)
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_fit_x1_plane", output_format)
+    save_figure(fig, path=output_path, filename="subplot_fit_x1_plane", format=output_format)
 
 
 def subplot_fit_log10(
@@ -260,7 +260,7 @@ def subplot_fit_log10(
                        colormap=colormap, use_log10=True)
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, f"subplot_fit_log10{plane_index_tag}", output_format)
+    save_figure(fig, path=output_path, filename=f"subplot_fit_log10{plane_index_tag}", format=output_format)
 
 
 def subplot_fit_log10_x1_plane(
@@ -300,7 +300,7 @@ def subplot_fit_log10_x1_plane(
                colormap=colormap, use_log10=True)
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_fit_log10", output_format)
+    save_figure(fig, path=output_path, filename="subplot_fit_log10", format=output_format)
 
 
 def subplot_of_planes(
@@ -339,7 +339,7 @@ def subplot_of_planes(
         _plot_source_plane(fit, axes_flat[3], pidx, colormap=colormap)
 
         plt.tight_layout()
-        _save_subplot(fig, output_path, f"subplot_of_plane_{pidx}", output_format)
+        save_figure(fig, path=output_path, filename=f"subplot_of_plane_{pidx}", format=output_format)
 
 
 def subplot_tracer_from_fit(
@@ -389,7 +389,7 @@ def subplot_tracer_from_fit(
         axes_flat[i].axis("off")
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_tracer", output_format)
+    save_figure(fig, path=output_path, filename="subplot_tracer", format=output_format)
 
 
 def subplot_fit_combined(
@@ -444,7 +444,7 @@ def subplot_fit_combined(
                    title="Normalized Residual Map", colormap=colormap)
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_fit_combined", output_format)
+    save_figure(fig, path=output_path, filename="subplot_fit_combined", format=output_format)
 
 
 def subplot_fit_combined_log10(
@@ -501,7 +501,7 @@ def subplot_fit_combined_log10(
                    title="Normalized Residual Map", colormap=colormap)
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "fit_combined_log10", output_format)
+    save_figure(fig, path=output_path, filename="fit_combined_log10", format=output_format)
 
 
 def _symmetric_vmax(array) -> float:

autolens/interferometer/plot/fit_interferometer_plots.py

@@ -6,9 +6,9 @@
 import autogalaxy as ag
 
 from autoarray.plot.plots.array import plot_array
+from autoarray.plot.plots.utils import save_figure
 from autolens.plot.plot_utils import (
     _to_lines,
-    _save_subplot,
     _critical_curves_from,
 )
 
@@ -115,7 +115,7 @@ def subplot_fit(
                        zoom_to_brightest=False, colormap=colormap)
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_fit", output_format)
+    save_figure(fig, path=output_path, filename="subplot_fit", format=output_format)
 
 
 def subplot_fit_real_space(
@@ -158,4 +158,4 @@ def subplot_fit_real_space(
         axes_flat[1].set_title("Source Reconstruction")
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_fit_real_space", output_format)
+    save_figure(fig, path=output_path, filename="subplot_fit_real_space", format=output_format)

autolens/lens/plot/sensitivity_plots.py

@@ -6,7 +6,7 @@
 import autoarray as aa
 
 from autoarray.plot.plots.array import plot_array
-from autolens.plot.plot_utils import _save_subplot
+from autoarray.plot.plots.utils import save_figure
 
 
 def subplot_tracer_images(
@@ -73,7 +73,7 @@ def subplot_tracer_images(
                colormap=colormap, use_log10=use_log10, lines=no_perturb_cc_lines)
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_lensed_images", output_format)
+    save_figure(fig, path=output_path, filename="subplot_lensed_images", format=output_format)
 
 
 def subplot_sensitivity(
@@ -162,7 +162,7 @@ def subplot_sensitivity(
         pass
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_sensitivity", output_format)
+    save_figure(fig, path=output_path, filename="subplot_sensitivity", format=output_format)
 
 
 def subplot_figures_of_merit_grid(
@@ -183,4 +183,4 @@ def subplot_figures_of_merit_grid(
     plot_array(array=figures_of_merit, ax=ax, title="Increase in Log Evidence",
                colormap=colormap)
     plt.tight_layout()
-    _save_subplot(fig, output_path, "sensitivity", output_format)
+    save_figure(fig, path=output_path, filename="sensitivity", format=output_format)

autolens/lens/plot/subhalo_plots.py

@@ -3,7 +3,7 @@
 from typing import Optional
 
 from autoarray.plot.plots.array import plot_array
-from autolens.plot.plot_utils import _save_subplot
+from autoarray.plot.plots.utils import save_figure
 from autolens.imaging.plot.fit_imaging_plots import _plot_source_plane
 
 
@@ -57,7 +57,7 @@ def subplot_detection_imaging(
     )
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_detection_imaging", output_format)
+    save_figure(fig, path=output_path, filename="subplot_detection_imaging", format=output_format)
 
 
 def subplot_detection_fits(
@@ -101,4 +101,4 @@ def subplot_detection_fits(
                        colormap=colormap)
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_detection_fits", output_format)
+    save_figure(fig, path=output_path, filename="subplot_detection_fits", format=output_format)

autolens/lens/plot/tracer_plots.py

@@ -6,10 +6,10 @@
 import autogalaxy as ag
 
 from autoarray.plot.plots.array import plot_array
+from autoarray.plot.plots.utils import save_figure
 from autolens.plot.plot_utils import (
     _to_lines,
     _to_positions,
-    _save_subplot,
     _critical_curves_from,
     _caustics_from,
 )
@@ -90,7 +90,7 @@ def subplot_tracer(
                lines=image_plane_lines, colormap=colormap)
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_tracer", output_format)
+    save_figure(fig, path=output_path, filename="subplot_tracer", format=output_format)
 
 
 def subplot_lensed_images(
@@ -120,7 +120,7 @@ def subplot_lensed_images(
         )
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_lensed_images", output_format)
+    save_figure(fig, path=output_path, filename="subplot_lensed_images", format=output_format)
 
 
 def subplot_galaxies_images(
@@ -177,4 +177,4 @@ def subplot_galaxies_images(
             idx += 1
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_galaxies_images", output_format)
+    save_figure(fig, path=output_path, filename="subplot_galaxies_images", format=output_format)

autolens/plot/plot_utils.py

@@ -1,6 +1,4 @@
-import os
 import numpy as np
-import matplotlib.pyplot as plt
 
 
 def _to_lines(*items):
@@ -32,26 +30,6 @@ def _to_positions(*items):
     return _to_lines(*items)
 
 
-def _save_subplot(fig, output_path, filename, output_format="png"):
-    """Save a subplot figure to disk (or show it when output_path is None)."""
-    if isinstance(output_format, (list, tuple)):
-        fmts = output_format
-    else:
-        fmts = [output_format]
-
-    if output_path:
-        os.makedirs(output_path, exist_ok=True)
-        for fmt in fmts:
-            fig.savefig(
-                os.path.join(output_path, f"{filename}.{fmt}"),
-                bbox_inches="tight",
-                pad_inches=0.1,
-            )
-    else:
-        plt.show()
-    plt.close(fig)
-
-
 def _critical_curves_from(tracer, grid):
     """Return (tangential_critical_curves, radial_critical_curves) as lists of arrays."""
     from autolens.lens import tracer_util

autolens/point/plot/fit_point_plots.py

@@ -2,7 +2,7 @@
 import numpy as np
 from typing import Optional
 
-from autolens.plot.plot_utils import _save_subplot
+from autoarray.plot.plots.utils import save_figure
 
 
 def subplot_fit(
@@ -57,4 +57,4 @@ def subplot_fit(
         )
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_fit", output_format)
+    save_figure(fig, path=output_path, filename="subplot_fit", format=output_format)

autolens/point/plot/point_dataset_plots.py

@@ -2,7 +2,7 @@
 import numpy as np
 from typing import Optional
 
-from autolens.plot.plot_utils import _save_subplot
+from autoarray.plot.plots.utils import save_figure
 
 
 def subplot_dataset(
@@ -49,4 +49,4 @@ def subplot_dataset(
         )
 
     plt.tight_layout()
-    _save_subplot(fig, output_path, "subplot_dataset_point", output_format)
+    save_figure(fig, path=output_path, filename="subplot_dataset_point", format=output_format)