Merge pull request #364 from Jammy2211/feature/jax_simplify_visualization

feature/jax_simplify_visualization

Author: Jammy2211

README.rst

@@ -26,6 +26,18 @@ PyAutoLens: Open-Source Strong Lensing
 .. |arXiv| image:: https://img.shields.io/badge/arXiv-1708.07377-blue
     :target: https://arxiv.org/abs/1708.07377
 
+.. image:: https://www.repostatus.org/badges/latest/active.svg
+    :target: https://www.repostatus.org/#active
+    :alt: Project Status: Active
+
+.. image:: https://img.shields.io/pypi/pyversions/autolens
+    :target: https://pypi.org/project/autolens/
+    :alt: Python Versions
+
+.. image:: https://img.shields.io/pypi/v/autolens.svg
+    :target: https://pypi.org/project/autolens/
+    :alt: PyPI Version
+
 |binder| |RTD| |Tests| |Build| |code-style| |JOSS| |arXiv|
 
 `Installation Guide <https://pyautolens.readthedocs.io/en/latest/installation/overview.html>`_ |

autolens/analysis/plotter_interface.py

@@ -1,5 +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
@@ -31,7 +32,12 @@ class PlotterInterface(AgPlotterInterface):
         The path on the hard-disk to the `image` folder of the non-linear searches results.
     """
 
-    def tracer(self, tracer: Tracer, grid: aa.type.Grid2DLike):
+    def tracer(
+        self,
+        tracer: Tracer,
+        grid: aa.type.Grid2DLike,
+        visuals_2d_of_planes_list: Optional[aplt.Visuals2D] = None,
+    ):
         """
         Visualizes a `Tracer` object.
 
@@ -63,7 +69,7 @@ def should_plot(name):
             tracer=tracer,
             grid=grid,
             mat_plot_2d=mat_plot_2d,
-            include_2d=self.include_2d,
+            visuals_2d_of_planes_list=visuals_2d_of_planes_list,
         )
 
         if should_plot("subplot_galaxies_images"):
@@ -169,7 +175,6 @@ def should_plot(name):
             image_plotter = aplt.Array2DPlotter(
                 array=image,
                 mat_plot_2d=mat_plot_2d,
-                include_2d=self.include_2d,
                 visuals_2d=visuals_2d,
             )
 

autolens/imaging/model/plotter_interface.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 import autoarray.plot as aplt
 
@@ -19,7 +19,7 @@ class PlotterInterfaceImaging(PlotterInterface):
     imaging_combined = AgPlotterInterfaceImaging.imaging_combined
 
     def fit_imaging(
-        self, fit: FitImaging,
+        self, fit: FitImaging, visuals_2d_of_planes_list : Optional[aplt.Visuals2D] = None
     ):
         """
         Visualizes a `FitImaging` object, which fits an imaging dataset.
@@ -45,7 +45,7 @@ def fit_imaging(
             mat_plot_2d = self.mat_plot_2d_from()
 
             fit_plotter = FitImagingPlotter(
-                fit=fit, mat_plot_2d=mat_plot_2d, include_2d=self.include_2d
+                fit=fit, mat_plot_2d=mat_plot_2d, visuals_2d_of_planes_list=visuals_2d_of_planes_list,
             )
 
             fit_plotter.subplot_tracer()
@@ -56,7 +56,7 @@ def should_plot(name):
         mat_plot_2d = self.mat_plot_2d_from()
 
         fit_plotter = FitImagingPlotter(
-            fit=fit, mat_plot_2d=mat_plot_2d, include_2d=self.include_2d
+            fit=fit, mat_plot_2d=mat_plot_2d, visuals_2d_of_planes_list=visuals_2d_of_planes_list,
         )
 
         plane_indexes_to_plot = [i for i in fit.tracer.plane_indexes_with_images if i != 0]
@@ -72,6 +72,7 @@ def should_plot(name):
                 fit_plotter.subplot_fit()
 
         if should_plot("subplot_fit_log10"):
+
             try:
                 if len(fit.tracer.planes) > 2:
                     for plane_index in plane_indexes_to_plot:
@@ -81,6 +82,7 @@ def should_plot(name):
             except ValueError:
                 pass
 
+
         if should_plot("subplot_of_planes"):
             fit_plotter.subplot_of_planes()
 
@@ -92,7 +94,7 @@ def should_plot(name):
 
         fits_to_fits(should_plot=should_plot, image_path=self.image_path, fit=fit)
 
-    def fit_imaging_combined(self, fit_list: List[FitImaging]):
+    def fit_imaging_combined(self, fit_list: List[FitImaging], visuals_2d_of_planes_list : Optional[aplt.Visuals2D] = None):
         """
         Output visualization of all `FitImaging` objects in a summed combined analysis, typically during or after a
         model-fit is performed.
@@ -119,7 +121,7 @@ def should_plot(name):
 
         fit_plotter_list = [
             FitImagingPlotter(
-                fit=fit, mat_plot_2d=mat_plot_2d, include_2d=self.include_2d
+                fit=fit, mat_plot_2d=mat_plot_2d, visuals_2d_of_planes_list=visuals_2d_of_planes_list,
             )
             for fit in fit_list
         ]

autolens/imaging/model/visualizer.py

@@ -4,6 +4,8 @@
 import autogalaxy as ag
 
 from autolens.imaging.model.plotter_interface import PlotterInterfaceImaging
+
+from autolens.lens import tracer_util
 from autolens import exc
 
 
@@ -109,15 +111,6 @@ def visualize(
             except exc.InversionException:
                 return
 
-        plotter_interface = PlotterInterfaceImaging(
-            image_path=paths.image_path, title_prefix=analysis.title_prefix
-        )
-
-        try:
-            plotter_interface.fit_imaging(fit=fit)
-        except exc.InversionException:
-            pass
-
         tracer = fit.tracer_linear_light_profiles_to_light_profiles
 
         zoom = ag.Zoom2D(mask=fit.mask)
@@ -127,8 +120,28 @@ def visualize(
 
         grid = ag.Grid2D.from_extent(extent=extent, shape_native=shape_native)
 
+        visuals_2d_of_planes_list = tracer_util.visuals_2d_of_planes_list_from(
+            tracer=fit.tracer,
+            grid=fit.grids.lp.mask.derive_grid.all_false
+        )
+
+        plotter_interface = PlotterInterfaceImaging(
+            image_path=paths.image_path,
+            title_prefix=analysis.title_prefix,
+        )
+
+        try:
+            plotter_interface.fit_imaging(
+                fit=fit,
+                visuals_2d_of_planes_list=visuals_2d_of_planes_list
+            )
+        except exc.InversionException:
+            pass
+
         plotter_interface.tracer(
-            tracer=tracer, grid=grid,
+            tracer=tracer,
+            grid=grid,
+            visuals_2d_of_planes_list=visuals_2d_of_planes_list
         )
         plotter_interface.galaxies(
             galaxies=tracer.galaxies,