Merge pull request #412 from Jammy2211/feature/visualization_final

Visualization final: interferometer tracer, plot consolidation

Author: Jammy2211

CLAUDE.md

@@ -9,33 +9,44 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 pip install -e ".[dev]"
 ```
 
-### Run Tests
-```bash
-# All tests
-python -m pytest test_autolens/
+### Run Tests
+```bash
+# All tests
+python -m pytest test_autolens/
 
 # Single test file
 python -m pytest test_autolens/lens/test_tracer.py
 
-# With output
-python -m pytest test_autolens/imaging/test_fit_imaging.py -s
-```
-
-### Codex / sandboxed runs
-
-When running Python from Codex or any restricted environment, set writable cache directories so `numba` and `matplotlib` do not fail on unwritable home or source-tree paths:
-
-```bash
-NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib python -m pytest test_autolens/
-```
-
-This workspace is often imported from `/mnt/c/...` and Codex may not be able to write to module `__pycache__` directories or `/home/jammy/.cache`, which can cause import-time `numba` caching failures without this override.
-
-### Formatting
-```bash
-black autolens/
+# With output
+python -m pytest test_autolens/imaging/test_fit_imaging.py -s
+```
+
+### Codex / sandboxed runs
+
+When running Python from Codex or any restricted environment, set writable cache directories so `numba` and `matplotlib` do not fail on unwritable home or source-tree paths:
+
+```bash
+NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib python -m pytest test_autolens/
+```
+
+This workspace is often imported from `/mnt/c/...` and Codex may not be able to write to module `__pycache__` directories or `/home/jammy/.cache`, which can cause import-time `numba` caching failures without this override.
+
+### Formatting
+```bash
+black autolens/
+```
+
+### Plot Output Mode
+
+Set `PYAUTOARRAY_OUTPUT_MODE=1` to capture every figure produced by a script into numbered PNG files in `./output_mode/<script_name>/`. This is useful for visually inspecting all plots from an integration test without needing a display.
+
+```bash
+PYAUTOARRAY_OUTPUT_MODE=1 python scripts/my_script.py
+# -> ./output_mode/my_script/0_fit.png, 1_tracer.png, ...
 ```
 
+When this env var is set, all `save_figure`, `subplot_save`, and `_save_subplot` calls are intercepted — the normal output path is bypassed and figures are written sequentially to the output_mode directory instead.
+
 ## Architecture
 
 **PyAutoLens** is the gravitational lensing layer built on top of PyAutoGalaxy. It adds multi-plane ray-tracing, the `Tracer` object, and lensing-specific fit classes. It depends on:

autolens/analysis/plotter.py

@@ -10,12 +10,11 @@
 
 from autolens.lens.tracer import Tracer
 from autolens.lens.plot.tracer_plots import (
-    subplot_tracer,
     subplot_galaxies_images,
     fits_tracer,
     fits_source_plane_images,
 )
-from autoarray.plot.array import plot_array
+from autogalaxy.plot.plot_utils import plot_array
 
 
 class Plotter(AgPlotter):
@@ -67,18 +66,6 @@ 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

@@ -6,7 +6,8 @@
 import autoarray as aa
 import autogalaxy as ag
 
-from autoarray.plot.array import plot_array, _zoom_array_2d
+from autogalaxy.plot.plot_utils import plot_array
+from autoarray.plot.array import _zoom_array_2d
 from autoarray.plot.utils import save_figure, hide_unused_axes, conf_subplot_figsize
 from autoarray.plot.utils import numpy_lines as _to_lines
 from autoarray.inversion.mappers.abstract import Mapper
@@ -42,11 +43,11 @@ def _compute_critical_curve_lines(tracer, grid):
         _rad_ca_lines = _to_lines(list(rad_ca) if rad_ca is not None else []) or []
         image_plane_lines = (_tan_cc_lines + _rad_cc_lines) or None
         image_plane_line_colors = (
-            ["black"] * len(_tan_cc_lines) + ["white"] * len(_rad_cc_lines)
+            ["white"] * len(_tan_cc_lines) + ["yellow"] * len(_rad_cc_lines)
         )
         source_plane_lines = (_tan_ca_lines + _rad_ca_lines) or None
         source_plane_line_colors = (
-            ["black"] * len(_tan_ca_lines) + ["white"] * len(_rad_ca_lines)
+            ["white"] * len(_tan_ca_lines) + ["yellow"] * len(_rad_ca_lines)
         )
         return image_plane_lines, image_plane_line_colors, source_plane_lines, source_plane_line_colors
     except Exception:
@@ -134,9 +135,17 @@ def _plot_source_plane(fit, ax, plane_index, zoom_to_brightest=True,
     """
     tracer = fit.tracer_linear_light_profiles_to_light_profiles
     if not tracer.planes[plane_index].has(cls=aa.Pixelization):
+        if zoom_to_brightest:
+            grid = fit.mask.derive_grid.all_false
+        else:
+            zoom = aa.Zoom2D(mask=fit.mask)
+            grid = aa.Grid2D.from_extent(
+                extent=zoom.extent_from(buffer=0),
+                shape_native=zoom.shape_native,
+            )
         image = plane_image_from(
             galaxies=tracer.planes[plane_index],
-            grid=fit.mask.derive_grid.all_false,
+            grid=grid,
             zoom_to_brightest=zoom_to_brightest,
         )
         plot_array(
@@ -713,6 +722,7 @@ def subplot_tracer_from_fit(
     else:
         axes_flat[1].axis("off")
 
+
     # Panel 2: Source Plane (No Zoom) (same as subplot_fit panel 12)
     _plot_source_plane(fit, axes_flat[2], final_plane_index, zoom_to_brightest=False,
                        colormap=colormap, title="Source Plane (No Zoom)",
@@ -742,17 +752,14 @@ def subplot_tracer_from_fit(
 
     # Panel 6: Deflections Y
     plot_array(array=deflections_y, ax=axes_flat[6], title="Deflections Y",
-               lines=image_plane_lines, line_colors=image_plane_line_colors,
                colormap=colormap)
 
     # Panel 7: Deflections X
     plot_array(array=deflections_x, ax=axes_flat[7], title="Deflections X",
-               lines=image_plane_lines, line_colors=image_plane_line_colors,
                colormap=colormap)
 
     # Panel 8: Magnification
     plot_array(array=magnification, ax=axes_flat[8], title="Magnification",
-               lines=image_plane_lines, line_colors=image_plane_line_colors,
                colormap=colormap)
 
     plt.tight_layout()

autolens/interferometer/model/plotter.py

@@ -14,6 +14,7 @@
     subplot_fit,
     subplot_fit_dirty_images,
     subplot_fit_real_space,
+    subplot_tracer_from_fit,
     _compute_critical_curve_lines,
 )
 from autolens.analysis.plotter import Plotter
@@ -59,10 +60,7 @@ def should_plot(name):
         # Use pre-computed critical curves if provided, otherwise compute once here.
         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.dataset.real_space_mask)
-            _cc_grid = aa.Grid2D.from_extent(
-                extent=_zoom.extent_from(buffer=0), shape_native=_zoom.shape_native
-            )
+            _cc_grid = fit.dataset.real_space_mask.derive_grid.all_false
             ip_lines, ip_colors, sp_lines, sp_colors = _compute_critical_curve_lines(tracer, _cc_grid)
         else:
             ip_lines, ip_colors, sp_lines, sp_colors = (
@@ -77,6 +75,13 @@ def should_plot(name):
                 source_plane_lines=sp_lines, source_plane_line_colors=sp_colors,
             )
 
+        if plot_setting(section="tracer", name="subplot_tracer"):
+            subplot_tracer_from_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:
             subplot_fit_dirty_images(
                 fit, output_path=output_path, output_format=fmt,

autolens/interferometer/model/visualizer.py

@@ -101,10 +101,7 @@ def visualize(
         )
 
         # Compute grid and critical curves once for all plot functions.
-        zoom = ag.Zoom2D(mask=fit.dataset.real_space_mask)
-        grid = ag.Grid2D.from_extent(
-            extent=zoom.extent_from(buffer=0), shape_native=zoom.shape_native
-        )
+        grid = fit.dataset.real_space_mask.derive_grid.all_false
         ip_lines, ip_colors, sp_lines, sp_colors = _compute_critical_curve_lines(
             tracer, grid
         )