Visualization final: fits API, plot consolidation

CLAUDE.md

@@ -9,36 +9,47 @@ 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_autogalaxy/
+### Run Tests
+```bash
+# All tests
+python -m pytest test_autogalaxy/
 
 # Single test file
 python -m pytest test_autogalaxy/galaxy/test_galaxy.py
 
 # Single test
 python -m pytest test_autogalaxy/galaxy/test_galaxy.py::TestGalaxy::test_name
 
-# With output
-python -m pytest test_autogalaxy/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_autogalaxy/
-```
-
-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 autogalaxy/
+# With output
+python -m pytest test_autogalaxy/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_autogalaxy/
+```
+
+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 autogalaxy/
+```
+
+### 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
 
 **PyAutoGalaxy** is a Bayesian galaxy morphology fitting library. It depends on two sibling packages:
@@ -223,4 +234,4 @@ find . -type f -name "*.py" | xargs dos2unix
 ```
 
 Prefer simple shell commands.
-Avoid chaining with && or pipes.
+Avoid chaining with && or pipes.

autogalaxy/aggregator/agg_util.py

@@ -23,7 +23,6 @@
 import numpy as np
 from typing import List, Optional
 
-from autoconf.fitsable import flip_for_ds9_from
 from autoconf.fitsable import ndarray_via_hdu_from
 
 import autofit as af
@@ -168,7 +167,7 @@ def adapt_images_from(
         for i, value in enumerate(fit.value(name="adapt_image_plane_mesh_grids")[1:]):
 
             adapt_image_plane_mesh_grid = aa.Grid2DIrregular(
-                values=flip_for_ds9_from(value.data.astype("float")),
+                values=value.data.astype("float"),
             )
 
             galaxy_name_image_plane_mesh_grid_dict[value.header["EXTNAME"].lower()] = (

autogalaxy/config/general.yaml

@@ -1,5 +1,3 @@
-fits:
-  flip_for_ds9: false
 psf:
   use_fft_default: true              # If True, PSFs are convolved using FFTs by default, which is faster and uses less memory in all cases except for very small PSFs, False uses direct convolution.
 updates:

autogalaxy/gui/clicker.py

@@ -3,6 +3,7 @@
 
 import autoarray as aa
 import autoarray.plot as aplt
+from autoarray.plot.utils import _conf_imshow_origin
 
 from autogalaxy import exc
 
@@ -29,7 +30,7 @@ def start(self, data, pixel_scales):
         fig = plt.figure(figsize=(14, 14))
         cmap = aplt.Cmap(cmap="jet", norm="log", vmin=1.0e-3, vmax=np.max(data) / 3.0)
         norm = cmap.norm_from(array=data, use_log10=True)
-        plt.imshow(data.native, cmap="jet", norm=norm, extent=ext)
+        plt.imshow(data.native, cmap="jet", norm=norm, extent=ext, origin=_conf_imshow_origin())
         if not data.mask.is_all_false:
             grid = data.mask.derive_grid.edge
             plt.scatter(y=grid[:, 0], x=grid[:, 1], c="k", marker="x", s=10)

autogalaxy/gui/scribbler.py

@@ -4,6 +4,8 @@
 import matplotlib.pyplot as plt
 from typing import Tuple
 
+from autoarray.plot.utils import _conf_imshow_origin
+
 
 class Scribbler:
     def __init__(
@@ -69,15 +71,15 @@ def __init__(
             self.ax = self.figure.add_subplot(121)
             plt.axis(extent)
             plt.axis("off")
-            plt.imshow(rgb_image)
+            plt.imshow(rgb_image, origin=_conf_imshow_origin())
         self.ax = self.figure.add_subplot(111)
 
         if cmap is None:
-            plt.imshow(image, interpolation="none")
+            plt.imshow(image, interpolation="none", origin=_conf_imshow_origin())
         else:
             norm = cmap.norm_from(array=image)
             cmap_name = getattr(cmap, "cmap_name", None) or cmap.config_dict.get("cmap", "viridis")
-            plt.imshow(image, cmap=cmap_name, norm=norm)
+            plt.imshow(image, cmap=cmap_name, norm=norm, origin=_conf_imshow_origin())
 
         if mask_overlay is not None:
             grid = mask_overlay.derive_grid.edge

autogalaxy/plot/__init__.py

@@ -1,10 +1,14 @@
-from autogalaxy.plot.plot_utils import plot_array, plot_grid
+from autogalaxy.plot.plot_utils import plot_array, plot_grid, fits_array
 
 from autoarray.dataset.plot.imaging_plots import (
     subplot_imaging_dataset,
     subplot_imaging_dataset_list,
+    fits_imaging,
+)
+from autoarray.dataset.plot.interferometer_plots import (
+    subplot_interferometer_dirty_images,
+    fits_interferometer,
 )
-from autoarray.dataset.plot.interferometer_plots import subplot_interferometer_dirty_images
 
 from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image
 

autogalaxy/plot/plot_utils.py

@@ -69,27 +69,21 @@ def _to_positions(*items):
 
 
 def _save_subplot(fig, output_path, output_filename, output_format="png",
-                  dpi=300, structure=None):
+                  dpi=300):
     """Save a subplot figure to disk (or show it if output_path is falsy).
 
-    Mirrors the interface of ``autoarray.plot.plots.utils.save_figure``.
-    When ``output_format`` is ``"fits"`` the *structure* argument is used to
-    write a FITS file via its ``output_to_fits`` method.
+    For FITS output use the dedicated ``fits_*`` functions instead.
     """
+    from autoarray.plot.utils import _output_mode_save
+
+    if _output_mode_save(fig, output_filename):
+        return
+
     fmt = output_format[0] if isinstance(output_format, (list, tuple)) else (output_format or "png")
     if output_path:
         os.makedirs(str(output_path), exist_ok=True)
         fpath = os.path.join(str(output_path), f"{output_filename}.{fmt}")
-        if fmt == "fits":
-            if structure is not None and hasattr(structure, "output_to_fits"):
-                structure.output_to_fits(file_path=fpath, overwrite=True)
-            else:
-                logger.warning(
-                    f"_save_subplot: fits format requested for {output_filename} "
-                    "but no compatible structure was provided; skipping."
-                )
-        else:
-            fig.savefig(fpath, dpi=dpi, bbox_inches="tight", pad_inches=0.1)
+        fig.savefig(fpath, dpi=dpi, bbox_inches="tight", pad_inches=0.1)
     else:
         plt.show()
     plt.close(fig)
@@ -122,7 +116,7 @@ def _numpy_grid(grid):
 
 def plot_array(
     array,
-    title,
+    title="",
     output_path=None,
     output_filename="array",
     output_format="png",
@@ -133,6 +127,7 @@ def plot_array(
     symmetric=False,
     positions=None,
     lines=None,
+    line_colors=None,
     grid=None,
     cb_unit=None,
     ax=None,
@@ -171,46 +166,43 @@ def plot_array(
         Point positions to scatter-plot over the image.
     lines : list or array-like or None
         Line coordinates to overlay on the image.
+    line_colors : list or None
+        Colours for each entry in *lines*.
     grid : array-like or None
         An additional grid of points to overlay.
     ax : matplotlib.axes.Axes or None
         Existing ``Axes`` to draw into.  When provided the figure is *not*
         saved — the caller is responsible for saving.
     """
     from autoarray.plot import plot_array as _aa_plot_array
-    from autoarray.plot import zoom_array, auto_mask_edge
 
     colormap = _resolve_colormap(colormap)
     output_format = _resolve_format(output_format)
-    array = zoom_array(array)
-
-    try:
-        arr = array.native.array
-        extent = array.geometry.extent
-    except AttributeError:
-        arr = np.asarray(array)
-        extent = None
-
-    mask = auto_mask_edge(array) if hasattr(array, "mask") else None
 
     if symmetric:
+        try:
+            arr = array.native.array
+        except AttributeError:
+            arr = np.asarray(array)
         finite = arr[np.isfinite(arr)]
         abs_max = float(np.max(np.abs(finite))) if len(finite) > 0 else 1.0
         vmin, vmax = -abs_max, abs_max
 
     _positions_list = positions if isinstance(positions, list) else _to_positions(positions)
     _lines_list = lines if isinstance(lines, list) else _to_lines(lines)
 
-    _output_path = None if ax is not None else output_path
+    if ax is not None:
+        _output_path = None
+    else:
+        _output_path = output_path if output_path is not None else "."
 
     _aa_plot_array(
-        array=arr,
+        array=array,
         ax=ax,
-        extent=extent,
-        mask=mask,
         grid=_numpy_grid(grid),
         positions=_positions_list,
         lines=_lines_list,
+        line_colors=line_colors,
         title=title or "",
         colormap=colormap,
         use_log10=use_log10,
@@ -223,9 +215,65 @@ def plot_array(
     )
 
 
+def _fits_values_and_header(array):
+    """Extract raw numpy values and header dict from an autoarray object.
+
+    Returns ``(values, header_dict, ext_name)`` where *header_dict* and
+    *ext_name* may be ``None`` for plain arrays.
+    """
+    from autoarray.structures.visibilities import AbstractVisibilities
+    from autoarray.mask.abstract_mask import Mask
+
+    if isinstance(array, AbstractVisibilities):
+        return np.asarray(array.in_array), None, None
+    if isinstance(array, Mask):
+        header = array.header_dict if hasattr(array, "header_dict") else None
+        return np.asarray(array.astype("float")), header, "mask"
+    if hasattr(array, "native"):
+        try:
+            header = array.mask.header_dict
+        except (AttributeError, TypeError):
+            header = None
+        return np.asarray(array.native.array).astype("float"), header, None
+
+    return np.asarray(array), None, None
+
+
+def fits_array(array, file_path, overwrite=False, ext_name=None):
+    """Write an autoarray ``Array2D``, ``Mask2D``, or array-like to a ``.fits`` file.
+
+    Handles header metadata (pixel scales, origin) automatically for
+    autoarray objects.
+
+    Parameters
+    ----------
+    array
+        The data to write.
+    file_path : str or Path
+        Full path including filename and ``.fits`` extension.
+    overwrite : bool
+        If ``True`` an existing file at *file_path* is replaced.
+    ext_name : str or None
+        FITS extension name.  Auto-detected for masks (``"mask"``).
+    """
+    from autoconf.fitsable import output_to_fits
+
+    values, header_dict, auto_ext_name = _fits_values_and_header(array)
+    if ext_name is None:
+        ext_name = auto_ext_name
+
+    output_to_fits(
+        values=values,
+        file_path=file_path,
+        overwrite=overwrite,
+        header_dict=header_dict,
+        ext_name=ext_name,
+    )
+
+
 def plot_grid(
     grid,
-    title,
+    title="",
     output_path=None,
     output_filename="grid",
     output_format="png",
@@ -258,7 +306,11 @@ def plot_grid(
     from autoarray.plot import plot_grid as _aa_plot_grid
 
     output_format = _resolve_format(output_format)
-    _output_path = None if ax is not None else output_path
+
+    if ax is not None:
+        _output_path = None
+    else:
+        _output_path = output_path if output_path is not None else "."
 
     _aa_plot_grid(
         grid=np.array(grid.array),

autogalaxy/quantity/dataset_quantity.py

@@ -234,32 +234,3 @@ def shape_native(self):
     def pixel_scales(self):
         return self.data.pixel_scales
 
-    def output_to_fits(
-        self,
-        data_path: Union[Path, str],
-        noise_map_path: Optional[Union[Path, str]] = None,
-        overwrite: bool = False,
-    ):
-        """
-        Output a quantity dataset to multiple .fits file.
-
-        For each attribute of the imaging data (e.g. `data`, `noise_map`) the path to
-        the .fits can be specified, with `hdu=0` assumed automatically.
-
-        If the `data` has been masked, the masked data is output to .fits files. A mask can be separately output to
-        a file `mask.fits` via the `Mask` objects `output_to_fits` method.
-
-        Parameters
-        ----------
-        data_path
-            The path to the data .fits file where the image data is output (e.g. '/path/to/data.fits').
-        noise_map_path
-            The path to the noise_map .fits where the noise_map is output (e.g. '/path/to/noise_map.fits').
-        overwrite
-            If `True`, the .fits files are overwritten if they already exist, if `False` they are not and an
-            exception is raised.
-        """
-        self.data.output_to_fits(file_path=data_path, overwrite=overwrite)
-
-        if self.noise_map is not None and noise_map_path is not None:
-            self.noise_map.output_to_fits(file_path=noise_map_path, overwrite=overwrite)

test_autogalaxy/config/general.yaml

@@ -1,7 +1,5 @@
 analysis:
   n_cores: 1
-fits:
-  flip_for_ds9: false
 psf:
   use_fft_default: false              # If True, PSFs are convolved using FFTs by default, which is faster and uses less memory in all cases except for very small PSFs, False uses direct convolution. Real space used for unit tests.
 inversion: