Plot improvements: rename fits functions, new output filenames, fix imports

autolens/analysis/plotter.py

@@ -12,8 +12,8 @@
 from autolens.lens.plot.tracer_plots import (
     subplot_tracer,
     subplot_galaxies_images,
-    save_tracer_fits,
-    save_source_plane_images_fits,
+    fits_tracer,
+    fits_source_plane_images,
 )
 from autoarray.plot.array import plot_array
 
@@ -72,10 +72,10 @@ def should_plot(name):
             )
 
         if should_plot("fits_tracer"):
-            save_tracer_fits(tracer=tracer, grid=grid, output_path=self.image_path)
+            fits_tracer(tracer=tracer, grid=grid, output_path=self.image_path)
 
         if should_plot("fits_source_plane_images"):
-            save_source_plane_images_fits(tracer=tracer, grid=grid, output_path=self.image_path)
+            fits_source_plane_images(tracer=tracer, grid=grid, output_path=self.image_path)
 
     def image_with_positions(self, image: aa.Array2D, positions: aa.Grid2DIrregular):
         """

autolens/imaging/model/plotter.py

@@ -3,7 +3,11 @@
 from typing import List
 
 from autogalaxy.imaging.model.plotter import PlotterImaging as AgPlotterImaging
-from autogalaxy.imaging.model.plotter import fits_to_fits
+from autogalaxy.imaging.plot.fit_imaging_plots import (
+    fits_fit,
+    fits_galaxy_images,
+    fits_model_galaxy_images,
+)
 
 from autolens.analysis.plotter import Plotter
 from autolens.imaging.fit_imaging import FitImaging
@@ -75,7 +79,14 @@ def should_plot(name):
         if should_plot("subplot_of_planes"):
             subplot_of_planes(fit, output_path=output_path, output_format=fmt)
 
-        fits_to_fits(should_plot=should_plot, image_path=self.image_path, fit=fit)
+        if should_plot("fits_fit"):
+            fits_fit(fit=fit, output_path=self.image_path)
+
+        if should_plot("fits_galaxy_images"):
+            fits_galaxy_images(fit=fit, output_path=self.image_path)
+
+        if should_plot("fits_model_galaxy_images"):
+            fits_model_galaxy_images(fit=fit, output_path=self.image_path)
 
     def fit_imaging_combined(
             self,

autolens/imaging/plot/fit_imaging_plots.py

@@ -518,7 +518,7 @@ def subplot_of_planes(
         _plot_source_plane(fit, axes_flat[3], pidx, colormap=colormap)
 
         plt.tight_layout()
-        save_figure(fig, path=output_path, filename=f"of_plane_{pidx}", format=output_format)
+        save_figure(fig, path=output_path, filename=f"fit_of_plane_{pidx}", format=output_format)
 
 
 def subplot_tracer_from_fit(

autolens/interferometer/model/plotter.py

@@ -2,8 +2,11 @@
     PlotterInterferometer as AgPlotterInterferometer,
 )
 
-from autogalaxy.interferometer.model.plotter import fits_to_fits
 from autogalaxy.interferometer.plot import fit_interferometer_plots as ag_fit_interferometer_plots
+from autogalaxy.interferometer.plot.fit_interferometer_plots import (
+    fits_galaxy_images,
+    fits_dirty_images,
+)
 
 from autolens.interferometer.fit_interferometer import FitInterferometer
 from autolens.interferometer.plot.fit_interferometer_plots import (
@@ -54,8 +57,8 @@ def should_plot(name):
         if should_plot("subplot_fit_real_space"):
             subplot_fit_real_space(fit, output_path=output_path, output_format=fmt)
 
-        fits_to_fits(
-            should_plot=should_plot,
-            image_path=self.image_path,
-            fit=fit,
-        )
+        if should_plot("fits_galaxy_images"):
+            fits_galaxy_images(fit=fit, output_path=self.image_path)
+
+        if should_plot("fits_dirty_images"):
+            fits_dirty_images(fit=fit, output_path=self.image_path)

autolens/lens/plot/tracer_plots.py

@@ -231,7 +231,7 @@ def subplot_galaxies_images(
     save_figure(fig, path=output_path, filename="galaxies_images", format=output_format)
 
 
-def save_tracer_fits(
+def fits_tracer(
     tracer,
     grid: aa.type.Grid2DLike,
     output_path,
@@ -274,7 +274,7 @@ def save_tracer_fits(
     hdu_list.writeto(output_path / "tracer.fits", overwrite=True)
 
 
-def save_source_plane_images_fits(
+def fits_source_plane_images(
     tracer,
     grid: aa.type.Grid2DLike,
     output_path,

test_autolens/analysis/test_plotter.py

@@ -25,7 +25,7 @@ def test__tracer(masked_imaging_7x7, tracer_x2_plane_7x7, plot_path, plot_patch)
         grid=masked_imaging_7x7.grids.lp,
     )
 
-    assert path.join(plot_path, "subplot_galaxies_images.png") in plot_patch.paths
+    assert path.join(plot_path, "galaxies_images.png") in plot_patch.paths
 
     image = al.ndarray_via_fits_from(
         file_path=path.join(plot_path, "tracer.fits"), hdu=0

test_autolens/imaging/model/test_plotter_imaging.py

@@ -26,9 +26,9 @@ def test__fit_imaging(
         fit=fit_imaging_x2_plane_inversion_7x7,
     )
 
-    assert path.join(plot_path, "subplot_tracer.png") in plot_patch.paths
-    assert path.join(plot_path, "subplot_fit.png") in plot_patch.paths
-    assert path.join(plot_path, "subplot_fit_log10.png") in plot_patch.paths
+    assert path.join(plot_path, "tracer.png") in plot_patch.paths
+    assert path.join(plot_path, "fit.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_log10.png") in plot_patch.paths
 
     image = al.ndarray_via_fits_from(
         file_path=path.join(plot_path, "fit.fits"), hdu=0
@@ -52,4 +52,4 @@ def test__fit_imaging_combined(
 
     visualizer.fit_imaging_combined(fit_list=2 * [fit_imaging_x2_plane_inversion_7x7])
 
-    assert path.join(plot_path, "subplot_fit_combined.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_combined.png") in plot_patch.paths

test_autolens/imaging/plot/test_fit_imaging_plots.py

@@ -31,7 +31,7 @@ def test__subplot_fit__two_plane_tracer__output_file_created(
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_fit.png") in plot_patch.paths
+    assert path.join(plot_path, "fit.png") in plot_patch.paths
 
 
 def test__subplot_fit__single_plane_tracer__delegates_to_x1_plane_and_creates_file(
@@ -42,7 +42,7 @@ def test__subplot_fit__single_plane_tracer__delegates_to_x1_plane_and_creates_fi
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_fit_x1_plane.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_x1_plane.png") in plot_patch.paths
 
 
 def test__subplot_fit_x1_plane__single_plane_tracer__output_file_created(
@@ -53,7 +53,7 @@ def test__subplot_fit_x1_plane__single_plane_tracer__output_file_created(
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_fit_x1_plane.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_x1_plane.png") in plot_patch.paths
 
 
 def test__subplot_fit_log10__two_plane_tracer__output_file_created(
@@ -64,7 +64,7 @@ def test__subplot_fit_log10__two_plane_tracer__output_file_created(
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_fit_log10.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_log10.png") in plot_patch.paths
 
 
 def test__subplot_fit_log10__single_plane_tracer__delegates_to_x1_plane_and_creates_file(
@@ -75,7 +75,7 @@ def test__subplot_fit_log10__single_plane_tracer__delegates_to_x1_plane_and_crea
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_fit_log10.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_log10.png") in plot_patch.paths
 
 
 def test__subplot_fit_log10_x1_plane__single_plane_tracer__output_file_created(
@@ -86,7 +86,7 @@ def test__subplot_fit_log10_x1_plane__single_plane_tracer__output_file_created(
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_fit_log10.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_log10.png") in plot_patch.paths
 
 
 def test__subplot_of_planes__no_plane_index_specified__all_plane_files_created(
@@ -98,8 +98,8 @@ def test__subplot_of_planes__no_plane_index_specified__all_plane_files_created(
         output_format="png",
     )
 
-    assert path.join(plot_path, "subplot_of_plane_0.png") in plot_patch.paths
-    assert path.join(plot_path, "subplot_of_plane_1.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_of_plane_0.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_of_plane_1.png") in plot_patch.paths
 
 
 def test__subplot_of_planes__plane_index_0_specified__only_plane_0_file_created(
@@ -112,8 +112,8 @@ def test__subplot_of_planes__plane_index_0_specified__only_plane_0_file_created(
         plane_index=0,
     )
 
-    assert path.join(plot_path, "subplot_of_plane_0.png") in plot_patch.paths
-    assert path.join(plot_path, "subplot_of_plane_1.png") not in plot_patch.paths
+    assert path.join(plot_path, "fit_of_plane_0.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_of_plane_1.png") not in plot_patch.paths
 
 
 def test__subplot_tracer_from_fit__two_plane_tracer__output_file_created(
@@ -124,7 +124,7 @@ def test__subplot_tracer_from_fit__two_plane_tracer__output_file_created(
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_tracer.png") in plot_patch.paths
+    assert path.join(plot_path, "tracer.png") in plot_patch.paths
 
 
 def test__subplot_fit_combined__list_of_two_fits__output_file_created(
@@ -135,7 +135,7 @@ def test__subplot_fit_combined__list_of_two_fits__output_file_created(
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_fit_combined.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_combined.png") in plot_patch.paths
 
 
 def test__subplot_fit_combined_log10__list_of_two_fits__output_file_created(

test_autolens/interferometer/model/test_plotter_interferometer.py

@@ -27,9 +27,9 @@ def test__fit_interferometer(
         fit=fit_interferometer_x2_plane_7x7,
     )
 
-    assert path.join(plot_path, "subplot_fit.png") in plot_patch.paths
-    assert path.join(plot_path, "subplot_fit_real_space.png") in plot_patch.paths
-    assert path.join(plot_path, "subplot_fit_dirty_images.png") in plot_patch.paths
+    assert path.join(plot_path, "fit.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_real_space.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_dirty_images.png") in plot_patch.paths
 
     image = al.ndarray_via_fits_from(
         file_path=path.join(plot_path, "galaxy_images.fits"), hdu=0

test_autolens/interferometer/plot/test_fit_interferometer_plots.py

@@ -22,7 +22,7 @@ def test__subplot_fit(
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_fit.png") in plot_patch.paths
+    assert path.join(plot_path, "fit.png") in plot_patch.paths
 
 
 def test__subplot_fit_real_space(
@@ -36,4 +36,4 @@ def test__subplot_fit_real_space(
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_fit_real_space.png") in plot_patch.paths
+    assert path.join(plot_path, "fit_real_space.png") in plot_patch.paths

test_autolens/lens/plot/test_tracer_plots.py

@@ -7,8 +7,8 @@
     subplot_tracer,
     subplot_lensed_images,
     subplot_galaxies_images,
-    save_tracer_fits,
-    save_source_plane_images_fits,
+    fits_tracer,
+    fits_source_plane_images,
 )
 
 directory = path.dirname(path.realpath(__file__))
@@ -33,7 +33,7 @@ def test__subplot_tracer__two_plane_tracer__output_file_created(
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_tracer.png") in plot_patch.paths
+    assert path.join(plot_path, "tracer.png") in plot_patch.paths
 
 
 def test__subplot_galaxies_images__two_plane_tracer__output_file_created(
@@ -45,7 +45,7 @@ def test__subplot_galaxies_images__two_plane_tracer__output_file_created(
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_galaxies_images.png") in plot_patch.paths
+    assert path.join(plot_path, "galaxies_images.png") in plot_patch.paths
 
 
 def test__subplot_lensed_images__two_plane_tracer__output_file_created(
@@ -57,24 +57,24 @@ def test__subplot_lensed_images__two_plane_tracer__output_file_created(
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_lensed_images.png") in plot_patch.paths
+    assert path.join(plot_path, "lensed_images.png") in plot_patch.paths
 
 
-def test__save_tracer_fits__two_plane_tracer__tracer_fits_file_written(
+def test__fits_tracer__two_plane_tracer__tracer_fits_file_written(
     tracer_x2_plane_7x7, grid_2d_7x7, tmp_path
 ):
-    save_tracer_fits(
+    fits_tracer(
         tracer=tracer_x2_plane_7x7,
         grid=grid_2d_7x7,
         output_path=tmp_path,
     )
     assert (tmp_path / "tracer.fits").exists()
 
 
-def test__save_source_plane_images_fits__two_plane_tracer__source_plane_fits_file_written(
+def test__fits_source_plane_images__two_plane_tracer__source_plane_fits_file_written(
     tracer_x2_plane_7x7, grid_2d_7x7, tmp_path
 ):
-    save_source_plane_images_fits(
+    fits_source_plane_images(
         tracer=tracer_x2_plane_7x7,
         grid=grid_2d_7x7,
         output_path=tmp_path,

test_autolens/point/model/test_plotter_point.py

@@ -21,4 +21,4 @@ def test__fit_point(fit_point_dataset_x2_plane, plot_path, plot_patch):
 
     plotter.fit_point(fit=fit_point_dataset_x2_plane)
 
-    assert path.join(plot_path, "subplot_fit.png") in plot_patch.paths
+    assert path.join(plot_path, "fit.png") in plot_patch.paths

test_autolens/point/plot/test_fit_point_plots.py

@@ -23,4 +23,4 @@ def test__subplot_fit(fit_point_dataset_x2_plane, plot_path, plot_patch):
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_fit.png") in plot_patch.paths
+    assert path.join(plot_path, "fit.png") in plot_patch.paths

test_autolens/point/plot/test_point_dataset_plots.py

@@ -23,4 +23,4 @@ def test__subplot_dataset(point_dataset, plot_path, plot_patch):
         output_path=plot_path,
         output_format="png",
     )
-    assert path.join(plot_path, "subplot_dataset_point.png") in plot_patch.paths
+    assert path.join(plot_path, "dataset_point.png") in plot_patch.paths