Merge pull request #247 from Jammy2211/feature/visualization_final

Visualization final: config origin, fits API, output mode

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_autoarray/
+### Run Tests
+```bash
+# All tests
+python -m pytest test_autoarray/
 
 # Single test file
 python -m pytest test_autoarray/structures/test_arrays.py
 
-# With output
-python -m pytest test_autoarray/structures/test_arrays.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_autoarray/
-```
-
-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 autoarray/
+# With output
+python -m pytest test_autoarray/structures/test_arrays.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_autoarray/
+```
+
+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 autoarray/
+```
+
+### 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
 
 **PyAutoArray** is the low-level data structures and numerical utilities package for the PyAuto ecosystem. It provides:

autoarray/abstract_ndarray.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 
-from autoconf.fitsable import output_to_fits
+
 
 from typing import TYPE_CHECKING
 
@@ -135,11 +135,6 @@ def with_new_array(self, array: np.ndarray) -> "AbstractNDArray":
         new_array._array = array
         return new_array
 
-    def flip_hdu_for_ds9(self, values):
-        if conf.instance["general"]["fits"]["flip_for_ds9"]:
-            return self._xp.flipud(values)
-        return values
-
     def copy(self):
         new = copy(self)
         return new
@@ -272,24 +267,6 @@ def native(self) -> Structure:
         Returns the data structure in its `native` format which contains all unmaksed values to the native dimensions.
         """
 
-    def output_to_fits(self, file_path: str, overwrite: bool = False):
-        """
-        Output the grid to a .fits file.
-
-        Parameters
-        ----------
-        file_path
-            The path the file is output to, including the filename and the .fits extension, e.g. '/path/to/filename.fits'
-        overwrite
-            If a file already exists at the path, if overwrite=True it is overwritten else an error is raised.
-        """
-        output_to_fits(
-            values=self.native.array.astype("float"),
-            file_path=file_path,
-            overwrite=overwrite,
-            header_dict=self.mask.header_dict,
-        )
-
     @property
     def shape(self):
         try:

autoarray/config/general.yaml

@@ -1,5 +1,3 @@
-fits:
-  flip_for_ds9: false                 # If True, the image is flipped before output to a .fits file, which is useful for viewing in DS9.
 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.
 inversion:

autoarray/dataset/imaging/dataset.py

@@ -542,38 +542,3 @@ def apply_sparse_operator_cpu(
             sparse_operator=sparse_operator,
         )
 
-    def output_to_fits(
-        self,
-        data_path: Union[Path, str],
-        psf_path: Optional[Union[Path, str]] = None,
-        noise_map_path: Optional[Union[Path, str]] = None,
-        overwrite: bool = False,
-    ):
-        """
-        Output an imaging 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').
-        psf_path
-            The path to the psf .fits file where the psf is output (e.g. '/path/to/psf.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.psf is not None and psf_path is not None:
-            self.psf.kernel.output_to_fits(file_path=psf_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)

autoarray/dataset/interferometer/dataset.py

@@ -2,7 +2,7 @@
 import numpy as np
 from typing import Optional
 
-from autoconf.fitsable import ndarray_via_fits_from, output_to_fits
+from autoconf.fitsable import ndarray_via_fits_from
 from autoconf import cached_property
 
 from autoarray.dataset.abstract.dataset import AbstractDataset
@@ -424,47 +424,6 @@ def signal_to_noise_map(self):
             signal_to_noise_map_real + 1j * signal_to_noise_map_imag
         )
 
-    def output_to_fits(
-        self,
-        data_path=None,
-        noise_map_path=None,
-        uv_wavelengths_path=None,
-        overwrite=False,
-    ):
-        """
-        Output the interferometer dataset to multiple .fits files.
-
-        Each component (visibilities, noise map, uv_wavelengths) is saved to its own .fits file.
-        Any path set to `None` means that component is not saved.
-
-        Parameters
-        ----------
-        data_path
-            The path where the visibility data is saved (e.g. '/path/to/visibilities.fits').
-            If `None`, the visibilities are not saved.
-        noise_map_path
-            The path where the noise map is saved (e.g. '/path/to/noise_map.fits').
-            If `None`, the noise map is not saved.
-        uv_wavelengths_path
-            The path where the uv_wavelengths array is saved (e.g. '/path/to/uv_wavelengths.fits').
-            If `None`, the uv_wavelengths are not saved.
-        overwrite
-            If `True`, existing .fits files are overwritten. If `False`, an exception is raised
-            if a file already exists at the given path.
-        """
-        if data_path is not None:
-            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)
-
-        if self.uv_wavelengths is not None and uv_wavelengths_path is not None:
-            output_to_fits(
-                values=self.uv_wavelengths,
-                file_path=uv_wavelengths_path,
-                overwrite=overwrite,
-            )
-
     @property
     def psf(self):
         """

autoarray/dataset/plot/imaging_plots.py

@@ -181,3 +181,75 @@ def subplot_imaging_dataset_list(
         plot_array(dataset.signal_to_noise_map, ax=axes[i][2], title="Signal-To-Noise Map")
     plt.tight_layout()
     subplot_save(fig, output_path, output_filename, output_format)
+
+
+def fits_imaging(
+    dataset,
+    file_path=None,
+    data_path=None,
+    psf_path=None,
+    noise_map_path=None,
+    overwrite=False,
+):
+    """Write an ``Imaging`` dataset to FITS.
+
+    Supports two modes:
+
+    * **Separate files** — pass ``data_path``, ``psf_path``, ``noise_map_path``
+      to write each component to its own single-HDU FITS file.
+    * **Single multi-HDU file** — pass ``file_path`` to write all components
+      into one FITS file with named extensions (``mask``, ``data``, ``psf``,
+      ``noise_map``).
+
+    Parameters
+    ----------
+    dataset
+        The ``Imaging`` dataset to write.
+    file_path : str or Path, optional
+        Path for a single multi-HDU FITS file.
+    data_path, psf_path, noise_map_path : str or Path, optional
+        Paths for individual component files.
+    overwrite : bool
+        If ``True`` existing files are replaced.
+    """
+    from autoconf.fitsable import output_to_fits, hdu_list_for_output_from, write_hdu_list
+
+    header_dict = dataset.data.mask.header_dict if hasattr(dataset.data.mask, "header_dict") else None
+
+    if file_path is not None:
+        values_list = [dataset.data.mask.astype("float")]
+        ext_name_list = ["mask"]
+
+        values_list.append(dataset.data.native.array.astype("float"))
+        ext_name_list.append("data")
+
+        if dataset.psf is not None:
+            values_list.append(dataset.psf.kernel.native.array.astype("float"))
+            ext_name_list.append("psf")
+
+        if dataset.noise_map is not None:
+            values_list.append(dataset.noise_map.native.array.astype("float"))
+            ext_name_list.append("noise_map")
+
+        hdu_list = hdu_list_for_output_from(
+            values_list=values_list,
+            ext_name_list=ext_name_list,
+            header_dict=header_dict,
+        )
+        write_hdu_list(hdu_list, file_path=file_path, overwrite=overwrite)
+    else:
+        if data_path is not None:
+            output_to_fits(
+                values=dataset.data.native.array.astype("float"),
+                file_path=data_path, overwrite=overwrite, header_dict=header_dict,
+            )
+        if dataset.psf is not None and psf_path is not None:
+            output_to_fits(
+                values=dataset.psf.kernel.native.array.astype("float"),
+                file_path=psf_path, overwrite=overwrite,
+            )
+        if dataset.noise_map is not None and noise_map_path is not None:
+            output_to_fits(
+                values=dataset.noise_map.native.array.astype("float"),
+                file_path=noise_map_path, overwrite=overwrite, header_dict=header_dict,
+            )

autoarray/dataset/plot/interferometer_plots.py

@@ -19,7 +19,7 @@ def subplot_interferometer_dataset(
     use_log10: bool = False,
 ):
     """
-    2×3 subplot of interferometer dataset components.
+    2x3 subplot of interferometer dataset components.
 
     Panels: Visibilities | UV-Wavelengths | Amplitudes vs UV-distances |
             Phases vs UV-distances | Dirty Image | Dirty S/N Map
@@ -100,7 +100,7 @@ def subplot_interferometer_dirty_images(
     use_log10: bool = False,
 ):
     """
-    1×3 subplot of dirty image, dirty noise map, and dirty S/N map.
+    1x3 subplot of dirty image, dirty noise map, and dirty S/N map.
 
     Parameters
     ----------
@@ -144,3 +144,72 @@ def subplot_interferometer_dirty_images(
     hide_unused_axes(axes)
     plt.tight_layout()
     subplot_save(fig, output_path, output_filename, output_format)
+
+
+def fits_interferometer(
+    dataset,
+    file_path=None,
+    data_path=None,
+    noise_map_path=None,
+    uv_wavelengths_path=None,
+    overwrite=False,
+):
+    """Write an ``Interferometer`` dataset to FITS.
+
+    Supports two modes:
+
+    * **Separate files** -- pass ``data_path``, ``noise_map_path``,
+      ``uv_wavelengths_path`` to write each component to its own FITS file.
+    * **Single multi-HDU file** -- pass ``file_path`` to write all components
+      into one FITS file with named extensions (``data``, ``noise_map``,
+      ``uv_wavelengths``).
+
+    Parameters
+    ----------
+    dataset
+        The ``Interferometer`` dataset to write.
+    file_path : str or Path, optional
+        Path for a single multi-HDU FITS file.
+    data_path, noise_map_path, uv_wavelengths_path : str or Path, optional
+        Paths for individual component files.
+    overwrite : bool
+        If ``True`` existing files are replaced.
+    """
+    from autoconf.fitsable import output_to_fits, hdu_list_for_output_from, write_hdu_list
+
+    if file_path is not None:
+        values_list = []
+        ext_name_list = []
+
+        values_list.append(np.asarray(dataset.data.in_array))
+        ext_name_list.append("data")
+
+        if dataset.noise_map is not None:
+            values_list.append(np.asarray(dataset.noise_map.in_array))
+            ext_name_list.append("noise_map")
+
+        if dataset.uv_wavelengths is not None:
+            values_list.append(np.asarray(dataset.uv_wavelengths))
+            ext_name_list.append("uv_wavelengths")
+
+        hdu_list = hdu_list_for_output_from(
+            values_list=values_list,
+            ext_name_list=ext_name_list,
+        )
+        write_hdu_list(hdu_list, file_path=file_path, overwrite=overwrite)
+    else:
+        if data_path is not None:
+            output_to_fits(
+                values=np.asarray(dataset.data.in_array),
+                file_path=data_path, overwrite=overwrite,
+            )
+        if dataset.noise_map is not None and noise_map_path is not None:
+            output_to_fits(
+                values=np.asarray(dataset.noise_map.in_array),
+                file_path=noise_map_path, overwrite=overwrite,
+            )
+        if dataset.uv_wavelengths is not None and uv_wavelengths_path is not None:
+            output_to_fits(
+                values=dataset.uv_wavelengths,
+                file_path=uv_wavelengths_path, overwrite=overwrite,
+            )