feat(plot): interferometer plots on array-free datasets — natural-weighted dirty panels, inversion subplot fix (streaming phase 3)

autoarray/dataset/plot/interferometer_plots.py

@@ -9,6 +9,44 @@
 from autoarray.structures.grids.irregular_2d import Grid2DIrregular
 
 
+def _subplot_natural_dataset(
+    dataset,
+    output_path,
+    output_filename,
+    output_format,
+    colormap,
+    use_log10,
+    title_prefix=None,
+):
+    """
+    1x2 subplot of the natural-weighted dirty image and dirty beam of an array-free
+    ``Interferometer`` (built by ``from_stream`` / ``from_sparse_terms``), which carries no
+    visibilities, uv-wavelengths or transformer, so only these real-space terms can be drawn.
+    """
+    _pf = (lambda t: f"{title_prefix.rstrip()} {t}") if title_prefix else (lambda t: t)
+
+    fig, axes = subplots(1, 2, figsize=conf_subplot_figsize(1, 2))
+
+    plot_array(
+        dataset.dirty_image_natural,
+        ax=axes[0],
+        title=_pf("Dirty Image (Natural)"),
+        colormap=colormap,
+        use_log10=use_log10,
+    )
+    plot_array(
+        dataset.dirty_beam,
+        ax=axes[1],
+        title=_pf("Dirty Beam (Natural)"),
+        colormap=colormap,
+        use_log10=use_log10,
+    )
+
+    hide_unused_axes(axes)
+    tight_layout()
+    subplot_save(fig, output_path, output_filename, output_format)
+
+
 def subplot_interferometer_dataset(
     dataset,
     output_path: Optional[str] = None,
@@ -24,6 +62,9 @@ def subplot_interferometer_dataset(
     Panels: Visibilities | UV-Wavelengths | Amplitudes vs UV-distances |
             Phases vs UV-distances | Dirty Image | Dirty S/N Map
 
+    An array-free dataset (``dataset.is_array_free``) has no visibilities, so a 1x2 subplot
+    of its ``dirty_image_natural`` and ``dirty_beam`` is written to the same filename instead.
+
     Parameters
     ----------
     dataset
@@ -39,6 +80,17 @@ def subplot_interferometer_dataset(
     use_log10
         Apply log10 normalisation to image panels.
     """
+    if dataset.is_array_free:
+        return _subplot_natural_dataset(
+            dataset,
+            output_path=output_path,
+            output_filename=output_filename,
+            output_format=output_format,
+            colormap=colormap,
+            use_log10=use_log10,
+            title_prefix=title_prefix,
+        )
+
     _pf = (lambda t: f"{title_prefix.rstrip()} {t}") if title_prefix else (lambda t: t)
 
     fig, axes = subplots(2, 3, figsize=conf_subplot_figsize(2, 3))
@@ -104,6 +156,9 @@ def subplot_interferometer_dirty_images(
     """
     1x3 subplot of dirty image, dirty noise map, and dirty S/N map.
 
+    An array-free dataset (``dataset.is_array_free``) has no visibilities, so a 1x2 subplot
+    of its ``dirty_image_natural`` and ``dirty_beam`` is written to the same filename instead.
+
     Parameters
     ----------
     dataset
@@ -119,6 +174,16 @@ def subplot_interferometer_dirty_images(
     use_log10
         Apply log10 normalisation.
     """
+    if dataset.is_array_free:
+        return _subplot_natural_dataset(
+            dataset,
+            output_path=output_path,
+            output_filename=output_filename,
+            output_format=output_format,
+            colormap=colormap,
+            use_log10=use_log10,
+        )
+
     fig, axes = subplots(1, 3, figsize=conf_subplot_figsize(1, 3))
 
     plot_array(
@@ -164,7 +229,9 @@ def fits_interferometer(
       ``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``).
+      ``uv_wavelengths``). An array-free dataset (``from_stream`` /
+      ``from_sparse_terms``) has none of these, so its natural-weighted
+      ``dirty_image_natural`` and ``dirty_beam`` are written instead.
 
     Parameters
     ----------
@@ -183,8 +250,9 @@ def fits_interferometer(
         values_list = []
         ext_name_list = []
 
-        values_list.append(np.asarray(dataset.data.in_array))
-        ext_name_list.append("data")
+        if dataset.data is not None:
+            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))
@@ -194,13 +262,19 @@ def fits_interferometer(
             values_list.append(np.asarray(dataset.uv_wavelengths))
             ext_name_list.append("uv_wavelengths")
 
+        if dataset.is_array_free:
+            values_list.append(np.asarray(dataset.dirty_image_natural.native))
+            ext_name_list.append("dirty_image_natural")
+            values_list.append(np.asarray(dataset.dirty_beam.native))
+            ext_name_list.append("dirty_beam")
+
         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:
+        if dataset.data is not None and data_path is not None:
             output_to_fits(
                 values=np.asarray(dataset.data.in_array),
                 file_path=data_path, overwrite=overwrite,

autoarray/fit/fit_interferometer.py

@@ -13,6 +13,54 @@
 from autoarray import type as ty
 
 
+def dirty_model_image_natural_from(dataset, image) -> Array2D:
+    """
+    Returns the naturally weighted, normalised dirty image of the model visibilities of a real-space
+    `image`, `Re(F^H W F m) / sum(w)`, computed without a transformer or any visibility-sized array.
+
+    `W~ = Re(F^H W F)` is the operator cached on the dataset's `sparse_operator` (the one the sparse
+    inversion uses for its curvature matrix), so `W~ m / sum(w)` is one FFT convolution on the real-space
+    grid. It is the model counterpart of `Interferometer.dirty_image_natural` (the natural dirty image of
+    the data, `Re(F^H W d) / sum(w)`): their difference is the natural dirty residual map. It is available
+    on both dataset types -- an array-free one built by `from_stream` / `from_sparse_terms` and an in-memory
+    one after `apply_sparse_operator()` -- and is how the visualizers draw a model on an array-free dataset.
+
+    Parameters
+    ----------
+    dataset
+        The `Interferometer` dataset, which must carry a `sparse_operator`.
+    image
+        The model image `m` on the slim masked real-space grid of the dataset's `real_space_mask`.
+    """
+    sparse_operator = getattr(dataset, "sparse_operator", None)
+
+    if sparse_operator is None:
+        raise exc.DatasetException(
+            "The natural dirty model image `W~ m / sum(w)` needs the dataset's `sparse_operator`; call "
+            "`apply_sparse_operator()` on the dataset (an array-free dataset built by from_stream / "
+            "from_sparse_terms always carries one)."
+        )
+
+    sparse_terms = getattr(dataset, "sparse_terms", None)
+
+    if sparse_terms is not None:
+        sum_weights = float(sparse_terms.sum_weights)
+    else:
+        sum_weights = float(np.sum(dataset.noise_map.array.real**-2.0))
+
+    image = np.asarray(getattr(image, "array", image), dtype=np.float64)
+
+    operated_image = sparse_operator.operated_matrix_slim_from(
+        matrix_slim=image[:, None],
+        extent_index_for_masked_pixel=dataset.real_space_mask.extent_index_for_masked_pixel,
+        xp=np,
+    )[:, 0]
+
+    return Array2D(
+        values=np.asarray(operated_image) / sum_weights, mask=dataset.real_space_mask
+    )
+
+
 class FitInterferometer(FitDataset):
     def __init__(
         self,

autoarray/fit/plot/fit_interferometer_plots.py

@@ -7,6 +7,78 @@
 from autoarray.plot.utils import subplots, subplot_save, symmetric_vmin_vmax, hide_unused_axes, conf_subplot_figsize, tight_layout
 
 
+def _subplot_fit_natural(
+    fit,
+    model_image,
+    output_path,
+    output_filename,
+    output_format,
+    colormap,
+    use_log10,
+    residuals_symmetric_cmap,
+):
+    """
+    Subplot of the natural-weighted dirty images of a ``FitInterferometer`` whose dataset is
+    array-free (built by ``from_stream`` / ``from_sparse_terms``), which has no visibilities,
+    transformer or visibility-space residuals.
+
+    Panels: Dirty Image (Natural) | Dirty Model Image (Natural) | Dirty Residual Map (Natural),
+    the last two only when the real-space ``model_image`` is given. The dirty model image is
+    ``W~ m / sum(w)`` (see ``dirty_model_image_natural_from``) and the dirty residual map is the
+    dirty image minus it.
+    """
+    from autoarray.fit.fit_interferometer import dirty_model_image_natural_from
+
+    dataset = fit.dataset
+    dirty_image = dataset.dirty_image_natural
+
+    if model_image is None:
+        fig, axes = subplots(1, 1, figsize=conf_subplot_figsize(1, 1))
+        axes = [axes]
+    else:
+        fig, axes = subplots(1, 3, figsize=conf_subplot_figsize(1, 3))
+
+    plot_array(
+        dirty_image,
+        ax=axes[0],
+        title="Dirty Image (Natural)",
+        colormap=colormap,
+        use_log10=use_log10,
+    )
+
+    if model_image is not None:
+        dirty_model_image = dirty_model_image_natural_from(
+            dataset=dataset, image=model_image
+        )
+        dirty_residual_map = dirty_image - dirty_model_image
+
+        if residuals_symmetric_cmap:
+            vmin_r, vmax_r = symmetric_vmin_vmax(dirty_residual_map)
+        else:
+            vmin_r = vmax_r = None
+
+        plot_array(
+            dirty_model_image,
+            ax=axes[1],
+            title="Dirty Model Image (Natural)",
+            colormap=colormap,
+            use_log10=use_log10,
+        )
+        plot_array(
+            dirty_residual_map,
+            ax=axes[2],
+            title="Dirty Residual Map (Natural)",
+            colormap=colormap,
+            use_log10=False,
+            vmin=vmin_r,
+            vmax=vmax_r,
+        )
+
+    hide_unused_axes(axes)
+    tight_layout()
+    subplot_save(fig, output_path, output_filename, output_format)
+
+
 def subplot_fit_interferometer(
     fit,
     output_path: Optional[str] = None,
@@ -15,6 +87,7 @@ def subplot_fit_interferometer(
     colormap=None,
     use_log10: bool = False,
     residuals_symmetric_cmap: bool = True,
+    model_image=None,
 ):
     """
     2×3 subplot of ``FitInterferometer`` residuals in UV-plane.
@@ -38,7 +111,24 @@ def subplot_fit_interferometer(
     residuals_symmetric_cmap
         Not used here (UV-plane residuals are scatter plots); kept for API
         consistency.
+    model_image
+        The real-space model image, used only when the fit's dataset is array-free
+        (``fit.dataset.is_array_free``): the visibility-space panels cannot be drawn, so
+        the natural-weighted dirty image, dirty model image and dirty residual map are
+        plotted to the same filename instead (the last two only if this is given).
     """
+    if fit.dataset.is_array_free:
+        return _subplot_fit_natural(
+            fit,
+            model_image=model_image,
+            output_path=output_path,
+            output_filename=output_filename,
+            output_format=output_format,
+            colormap=colormap,
+            use_log10=use_log10,
+            residuals_symmetric_cmap=residuals_symmetric_cmap,
+        )
+
     fig, axes = subplots(2, 3, figsize=conf_subplot_figsize(2, 3))
     axes = axes.flatten()
 
@@ -110,6 +200,7 @@ def subplot_fit_interferometer_dirty_images(
     colormap=None,
     use_log10: bool = False,
     residuals_symmetric_cmap: bool = True,
+    model_image=None,
 ):
     """
     2×3 subplot of ``FitInterferometer`` dirty-image components.
@@ -133,7 +224,24 @@ def subplot_fit_interferometer_dirty_images(
         Apply log10 normalisation to non-residual panels.
     residuals_symmetric_cmap
         Centre residual colour scale symmetrically around zero.
+    model_image
+        The real-space model image, used only when the fit's dataset is array-free
+        (``fit.dataset.is_array_free``): the natural-weighted dirty image, dirty model
+        image and dirty residual map are plotted to the same filename instead (the last
+        two only if this is given).
     """
+    if fit.dataset.is_array_free:
+        return _subplot_fit_natural(
+            fit,
+            model_image=model_image,
+            output_path=output_path,
+            output_filename=output_filename,
+            output_format=output_format,
+            colormap=colormap,
+            use_log10=use_log10,
+            residuals_symmetric_cmap=residuals_symmetric_cmap,
+        )
+
     fig, axes = subplots(2, 3, figsize=conf_subplot_figsize(2, 3))
     axes = axes.flatten()