Visualization cleanup: critical curves once, interferometer/point passdown

- Lift critical curve computation to visualizer level in interferometer and point
  so it is computed once and passed through the plotter call chain
- Add image_plane_lines/source_plane_lines parameters to interferometer subplot_fit,
  subplot_fit_dirty_images, subplot_fit_real_space and point subplot_fit
- Thread critical curve parameters through PlotterInterferometer.fit_interferometer,
  Plotter.tracer, PlotterPoint.fit_point and VisualizerPoint.visualize
- Sync fit_imaging_plots with interferometer equivalent for consistency

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: 2 people

autolens/analysis/plotter.py

@@ -37,6 +37,10 @@ def tracer(
         self,
         tracer: Tracer,
         grid: aa.type.Grid2DLike,
+        image_plane_lines=None,
+        image_plane_line_colors=None,
+        source_plane_lines=None,
+        source_plane_line_colors=None,
     ):
         """
         Visualizes a `Tracer` object.
@@ -47,6 +51,14 @@ def tracer(
             The maximum log likelihood `Tracer` of the non-linear search.
         grid
             A 2D grid of (y,x) arc-second coordinates used to perform ray-tracing.
+        image_plane_lines
+            Pre-computed critical-curve lines to overlay on image-plane panels.
+        image_plane_line_colors
+            Colours for each image-plane line.
+        source_plane_lines
+            Pre-computed caustic lines to overlay on source-plane panels.
+        source_plane_line_colors
+            Colours for each source-plane line.
         """
 
         def should_plot(name):
@@ -55,6 +67,18 @@ def should_plot(name):
         output_path = str(self.image_path)
         fmt = self.fmt
 
+        if should_plot("subplot_tracer"):
+            subplot_tracer(
+                tracer=tracer,
+                grid=grid,
+                output_path=output_path,
+                output_format=fmt,
+                image_plane_lines=image_plane_lines,
+                image_plane_line_colors=image_plane_line_colors,
+                source_plane_lines=source_plane_lines,
+                source_plane_line_colors=source_plane_line_colors,
+            )
+
         if should_plot("subplot_galaxies_images"):
             subplot_galaxies_images(
                 tracer=tracer,

autolens/imaging/plot/fit_imaging_plots.py

@@ -53,6 +53,19 @@ def _compute_critical_curve_lines(tracer, grid):
         return None, None, None, None
 
 
+def _compute_critical_curves_from_fit(fit):
+    """Compute critical-curve and caustic lines from a FitImaging object.
+
+    Convenience wrapper around :func:`_compute_critical_curve_lines` that
+    derives the tracer and grid from *fit* directly, using the fully unmasked
+    image-plane grid so the curves cover the whole image extent.
+
+    Returns the same 4-tuple as :func:`_compute_critical_curve_lines`.
+    """
+    tracer = fit.tracer_linear_light_profiles_to_light_profiles
+    return _compute_critical_curve_lines(tracer, fit.mask.derive_grid.all_false)
+
+
 def _get_source_vmax(fit):
     """
     Return the colour-scale maximum for source-plane panels.
@@ -217,14 +230,8 @@ def subplot_fit(
     source_vmax = _get_source_vmax(fit)
 
     if image_plane_lines is None and source_plane_lines is None:
-        tracer = fit.tracer_linear_light_profiles_to_light_profiles
-        _zoom = aa.Zoom2D(mask=fit.mask)
-        _cc_grid = aa.Grid2D.from_extent(
-            extent=_zoom.extent_from(buffer=0),
-            shape_native=_zoom.shape_native,
-        )
         image_plane_lines, image_plane_line_colors, source_plane_lines, source_plane_line_colors = (
-            _compute_critical_curve_lines(tracer, _cc_grid)
+            _compute_critical_curves_from_fit(fit)
         )
 
     fig, axes = plt.subplots(3, 4, figsize=conf_subplot_figsize(3, 4))
@@ -422,14 +429,8 @@ def subplot_fit_log10(
     source_vmax = _get_source_vmax(fit)
 
     if image_plane_lines is None and source_plane_lines is None:
-        tracer = fit.tracer_linear_light_profiles_to_light_profiles
-        _zoom = aa.Zoom2D(mask=fit.mask)
-        _cc_grid = aa.Grid2D.from_extent(
-            extent=_zoom.extent_from(buffer=0),
-            shape_native=_zoom.shape_native,
-        )
         image_plane_lines, image_plane_line_colors, source_plane_lines, source_plane_line_colors = (
-            _compute_critical_curve_lines(tracer, _cc_grid)
+            _compute_critical_curves_from_fit(fit)
         )
 
     fig, axes = plt.subplots(3, 4, figsize=conf_subplot_figsize(3, 4))
@@ -673,14 +674,11 @@ def subplot_tracer_from_fit(
     tracer = fit.tracer_linear_light_profiles_to_light_profiles
 
     # --- grid ---
-    zoom = aa.Zoom2D(mask=fit.mask)
-    grid = aa.Grid2D.from_extent(
-        extent=zoom.extent_from(buffer=0), shape_native=zoom.shape_native
-    )
+    grid = fit.mask.derive_grid.all_false
 
     if image_plane_lines is None and source_plane_lines is None:
         image_plane_lines, image_plane_line_colors, source_plane_lines, source_plane_line_colors = (
-            _compute_critical_curve_lines(tracer, grid)
+            _compute_critical_curves_from_fit(fit)
         )
 
     source_vmax = _get_source_vmax(fit)
@@ -689,7 +687,7 @@ def subplot_tracer_from_fit(
     lens_galaxies = ag.Galaxies(galaxies=tracer.planes[0])
     lens_image = lens_galaxies.image_2d_from(grid=traced_grids[0])
 
-    deflections = lens_galaxies.deflections_yx_2d_from(grid=grid)
+    deflections = tracer.deflections_yx_2d_from(grid=grid)
     deflections_y = aa.Array2D(values=deflections.slim[:, 0], mask=grid.mask)
     deflections_x = aa.Array2D(values=deflections.slim[:, 1], mask=grid.mask)
 

autolens/interferometer/model/plotter.py

@@ -1,8 +1,9 @@
+import autoarray as aa
+
 from autogalaxy.interferometer.model.plotter import (
     PlotterInterferometer as AgPlotterInterferometer,
 )
 
-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,
@@ -11,7 +12,9 @@
 from autolens.interferometer.fit_interferometer import FitInterferometer
 from autolens.interferometer.plot.fit_interferometer_plots import (
     subplot_fit,
+    subplot_fit_dirty_images,
     subplot_fit_real_space,
+    _compute_critical_curve_lines,
 )
 from autolens.analysis.plotter import Plotter
 
@@ -41,21 +44,35 @@ def should_plot(name):
         output_path = str(self.image_path)
         fmt = self.fmt
 
+        # Compute critical curves and caustics once for all subplot functions.
+        tracer = fit.tracer_linear_light_profiles_to_light_profiles
+        _zoom = aa.Zoom2D(mask=fit.dataset.real_space_mask)
+        _cc_grid = aa.Grid2D.from_extent(
+            extent=_zoom.extent_from(buffer=0), shape_native=_zoom.shape_native
+        )
+        ip_lines, ip_colors, sp_lines, sp_colors = _compute_critical_curve_lines(tracer, _cc_grid)
+
         if should_plot("subplot_fit"):
-            subplot_fit(fit, output_path=output_path, output_format=fmt)
+            subplot_fit(
+                fit, output_path=output_path, output_format=fmt,
+                image_plane_lines=ip_lines, image_plane_line_colors=ip_colors,
+                source_plane_lines=sp_lines, source_plane_line_colors=sp_colors,
+            )
 
         if should_plot("subplot_fit_dirty_images") or quick_update:
-            ag_fit_interferometer_plots.subplot_fit_dirty_images(
-                fit=fit,
-                output_path=self.image_path,
-                output_format=self.fmt,
+            subplot_fit_dirty_images(
+                fit, output_path=output_path, output_format=fmt,
+                image_plane_lines=ip_lines, image_plane_line_colors=ip_colors,
             )
 
         if quick_update:
             return
 
         if should_plot("subplot_fit_real_space"):
-            subplot_fit_real_space(fit, output_path=output_path, output_format=fmt)
+            subplot_fit_real_space(
+                fit, output_path=output_path, output_format=fmt,
+                source_plane_lines=sp_lines, source_plane_line_colors=sp_colors,
+            )
 
         if should_plot("fits_galaxy_images"):
             fits_galaxy_images(fit=fit, output_path=self.image_path)