feat(interferometer): visualizer on array-free datasets (streaming phase 3)

autogalaxy/interferometer/fit_interferometer.py

@@ -424,6 +424,66 @@ def galaxy_image_dict(self) -> Dict[Galaxy, np.ndarray]:
 
         return {**galaxy_image_dict, **galaxy_linear_obj_image_dict}
 
+    @property
+    def model_image_natural(self) -> aa.Array2D:
+        """
+        The real-space model image `m` of the fit, on the dataset's `real_space_mask`: the image of every
+        ordinary (non-linear) light profile (`profile_image`) plus, when the fit has an inversion, the solved
+        linear objects' reconstruction (`inversion.mapped_reconstructed_data`, linear light profiles and
+        pixelizations) -- the real-space image whose visibilities are `model_data`.
+
+        It is built from these two terms rather than from `galaxy_image_dict`, whose entry for a galaxy with
+        both ordinary and linear light holds only the linear reconstruction.
+
+        It needs neither visibilities nor a transformer, so it is available on an array-free dataset (built by
+        `Interferometer.from_stream` / `from_sparse_terms`), where it is the image the natural-weighted dirty
+        model image `dirty_model_image_natural` is formed from. (There `profile_image` is all zeros, because a
+        fit with ordinary light on an array-free dataset raises before it gets here.)
+        """
+        image = np.asarray(
+            getattr(self.profile_image, "array", self.profile_image), dtype=np.float64
+        )
+
+        if self.inversion is not None:
+            reconstruction = self.inversion.mapped_reconstructed_data
+            image = image + np.asarray(
+                getattr(reconstruction, "array", reconstruction), dtype=np.float64
+            )
+
+        return aa.Array2D(
+            values=image,
+            mask=self.dataset.real_space_mask,
+        )
+
+    @property
+    def dirty_model_image_natural(self) -> aa.Array2D:
+        """
+        The naturally weighted, normalised dirty image of the model visibilities, `W~ m / sum(w)`, formed from
+        `model_image_natural` with the dataset's `sparse_operator` (see
+        `autoarray.fit.fit_interferometer.dirty_model_image_natural_from`).
+
+        It is the model counterpart of the dataset's `dirty_image_natural` and needs no visibilities, so it is
+        how a fit on an array-free dataset is visualized. It is available on any dataset carrying a
+        `sparse_operator` (array-free, or in-memory after `apply_sparse_operator()`); otherwise it raises an
+        `aa.exc.DatasetException`.
+        """
+        return aa.fit.fit_interferometer.dirty_model_image_natural_from(
+            dataset=self.dataset, image=self.model_image_natural
+        )
+
+    @property
+    def dirty_residual_map_natural(self) -> aa.Array2D:
+        """
+        The naturally weighted dirty residual map, `dirty_image_natural - dirty_model_image_natural`, which is
+        `Re(F^H W (d - F m)) / sum(w)`: the natural dirty image of the visibility residuals, computed without
+        them.
+        """
+        return aa.Array2D(
+            values=np.asarray(self.dataset.dirty_image_natural.array)
+            - np.asarray(self.dirty_model_image_natural.array),
+            mask=self.dataset.real_space_mask,
+        )
+
     @functools.cached_property
     def galaxy_model_visibilities_dict(self) -> Dict[Galaxy, np.ndarray]:
         """

autogalaxy/interferometer/model/visualizer.py

@@ -108,5 +108,5 @@ def visualize(
                     inversion=fit.inversion_with_data,
                 )
             except (IndexError, exc.InversionException):
-                logger(exc.invalid_linear_algebra_for_visualization_message())
+                logger.warning(exc.invalid_linear_algebra_for_visualization_message())
                 return

autogalaxy/interferometer/plot/fit_interferometer_plots.py

@@ -11,6 +11,52 @@
 from autogalaxy.util.plot_utils import plot_array, _save_subplot
 
 
+def _natural_panels(fit: FitInterferometer, _pf):
+    """
+    The natural-weighted dirty image, dirty model image and dirty residual map panels of a fit on an
+    array-free dataset (``fit.dataset.is_array_free``, built by ``Interferometer.from_stream`` /
+    ``from_sparse_terms``), which has no visibilities, transformer or visibility-space residuals: every
+    quantity is formed from the dataset's sparse terms and the fit's real-space model image.
+    """
+    return [
+        (fit.dataset.dirty_image_natural, _pf("Dirty Image (Natural)")),
+        (fit.dirty_model_image_natural, _pf("Dirty Model Image (Natural)")),
+        (fit.dirty_residual_map_natural, _pf("Dirty Residual Map (Natural)")),
+    ]
+
+
+def _subplot_natural(
+    fit: FitInterferometer,
+    output_path,
+    output_filename,
+    output_format,
+    colormap,
+    use_log10,
+    title_prefix,
+):
+    """
+    Write the 1x3 natural-weighted dirty-image subplot of a fit on an array-free dataset (see
+    ``_natural_panels``) to ``output_filename``.
+    """
+    _pf = (lambda t: f"{title_prefix.rstrip()} {t}") if title_prefix else (lambda t: t)
+    panels = _natural_panels(fit=fit, _pf=_pf)
+    n = len(panels)
+    fig, axes = subplots(1, n, figsize=conf_subplot_figsize(1, n))
+    axes_flat = list(axes.flatten())
+
+    for i, (array, title) in enumerate(panels):
+        plot_array(
+            array=array,
+            title=title,
+            colormap=colormap,
+            use_log10=use_log10 if i < 2 else False,
+            ax=axes_flat[i],
+        )
+
+    tight_layout()
+    _save_subplot(fig, output_path, output_filename, output_format)
+
+
 def subplot_fit(
     fit: FitInterferometer,
     output_path=None,
@@ -41,7 +87,21 @@ def subplot_fit(
         Reserved for future log-stretch support (currently unused).
     residuals_symmetric_cmap : bool
         Reserved for future symmetric-colormap support (currently unused).
+
+    On an array-free dataset (``fit.dataset.is_array_free``) there are no visibility residuals, so the
+    natural-weighted dirty image, dirty model image and dirty residual map are plotted to ``fit`` instead.
     """
+    if fit.dataset.is_array_free:
+        return _subplot_natural(
+            fit=fit,
+            output_path=output_path,
+            output_filename="fit",
+            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)
     panels = [
         (fit.residual_map, _pf("Residual Map")),
@@ -89,7 +149,21 @@ def subplot_fit_dirty_images(
         Apply a log₁₀ stretch to the plotted values.
     residuals_symmetric_cmap : bool
         Reserved for future symmetric-colormap support (currently unused).
+
+    On an array-free dataset (``fit.dataset.is_array_free``) the unweighted dirty images cannot be formed,
+    so the natural-weighted dirty image, dirty model image and dirty residual map are plotted instead.
     """
+    if fit.dataset.is_array_free:
+        return _subplot_natural(
+            fit=fit,
+            output_path=output_path,
+            output_filename="fit_dirty_images",
+            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)
     panels = [
         (fit.dirty_image, _pf("Dirty Image"), None),
@@ -135,7 +209,9 @@ def subplot_fit_real_space(
       grid.
     - **With pixelization**: shows three dirty-image panels (dirty image, dirty
       model image, dirty residual map), which are the best real-space
-      representation available when a pixelized source is used.
+      representation available when a pixelized source is used. On an array-free
+      dataset these are the natural-weighted dirty image, dirty model image and
+      dirty residual map.
 
     Parameters
     ----------
@@ -165,11 +241,14 @@ def subplot_fit_real_space(
         )
     else:
         _pf = (lambda t: f"{title_prefix.rstrip()} {t}") if title_prefix else (lambda t: t)
-        panels = [
-            (fit.dirty_image, _pf("Dirty Image")),
-            (fit.dirty_model_image, _pf("Dirty Model Image")),
-            (fit.dirty_residual_map, _pf("Dirty Residual Map")),
-        ]
+        if fit.dataset.is_array_free:
+            panels = _natural_panels(fit=fit, _pf=_pf)
+        else:
+            panels = [
+                (fit.dirty_image, _pf("Dirty Image")),
+                (fit.dirty_model_image, _pf("Dirty Model Image")),
+                (fit.dirty_residual_map, _pf("Dirty Residual Map")),
+            ]
         n = len(panels)
         fig, axes = subplots(1, n, figsize=conf_subplot_figsize(1, n))
         axes_flat = list(axes.flatten())
@@ -212,13 +291,38 @@ def fits_dirty_images(fit: FitInterferometer, output_path) -> None:
     Extensions: ``mask``, ``dirty_image``, ``dirty_noise_map``, ``dirty_model_image``,
     ``dirty_residual_map``, ``dirty_normalized_residual_map``, ``dirty_chi_squared_map``.
 
+    On an array-free dataset (``fit.dataset.is_array_free``) the extensions are instead ``mask``,
+    ``dirty_image_natural``, ``dirty_beam``, ``dirty_model_image_natural`` and
+    ``dirty_residual_map_natural``.
+
     Parameters
     ----------
     fit
         The interferometer fit whose dirty-image arrays are saved.
     output_path
         Directory in which to write ``fit_dirty_images.fits``.
     """
+    if fit.dataset.is_array_free:
+        image_list = [
+            fit.dataset.dirty_image_natural.native_for_fits,
+            fit.dataset.dirty_beam.native_for_fits,
+            fit.dirty_model_image_natural.native_for_fits,
+            fit.dirty_residual_map_natural.native_for_fits,
+        ]
+        hdu_list = hdu_list_for_output_from(
+            values_list=[image_list[0].mask.astype("float")] + image_list,
+            ext_name_list=[
+                "mask",
+                "dirty_image_natural",
+                "dirty_beam",
+                "dirty_model_image_natural",
+                "dirty_residual_map_natural",
+            ],
+            header_dict=fit.dataset.real_space_mask.header_dict,
+        )
+        hdu_list.writeto(Path(output_path) / "fit_dirty_images.fits", overwrite=True)
+        return
+
     image_list = [
         fit.dirty_image.native_for_fits,
         fit.dirty_noise_map.native_for_fits,

test_autogalaxy/interferometer/model/test_plotter_interferometer.py

@@ -48,3 +48,136 @@ def test__fit_interferometer(
     )
 
     assert image.shape == (5, 5)
+
+
+def _array_free_dataset_from(dataset):
+    import autoarray as aa
+
+    return aa.Interferometer.from_stream(
+        [(dataset.uv_wavelengths, dataset.data, dataset.noise_map)],
+        real_space_mask=dataset.real_space_mask,
+        transformer_class=type(dataset.transformer),
+    )
+
+
+def _pixelization_only_model():
+    import autofit as af
+
+    pixelization = ag.Pixelization(
+        mesh=ag.mesh.RectangularUniform(shape=(3, 3)),
+        regularization=ag.reg.Constant(coefficient=1.0),
+    )
+
+    return af.Collection(
+        galaxies=af.Collection(
+            galaxy=af.Model(ag.Galaxy, redshift=0.5, pixelization=pixelization)
+        )
+    )
+
+
+def _visualize(dataset, image_path):
+    """
+    Run the interferometer visualizer's `visualize_before_fit` and `visualize` for a
+    pixelization-only model on `dataset`, as a non-linear search would.
+    """
+    from types import SimpleNamespace
+
+    from autogalaxy.interferometer.model.visualizer import VisualizerInterferometer
+
+    model = _pixelization_only_model()
+    instance = model.instance_from_prior_medians()
+
+    analysis = ag.AnalysisInterferometer(dataset=dataset, use_jax=False)
+    paths = SimpleNamespace(image_path=image_path)
+
+    VisualizerInterferometer.visualize_before_fit(
+        analysis=analysis, paths=paths, model=model
+    )
+    VisualizerInterferometer.visualize(
+        analysis=analysis, paths=paths, instance=instance, during_analysis=False
+    )
+
+    return analysis.fit_from(instance=instance)
+
+
+def _ext_names_from(file_path):
+    from astropy.io import fits
+
+    with fits.open(file_path) as hdu_list:
+        return [hdu.name for hdu in hdu_list]
+
+
+def test__visualizer__array_free_dataset(interferometer_7, tmp_path, plot_patch):
+    pytest.importorskip("nufftax")
+
+    import numpy as np
+
+    dataset = _array_free_dataset_from(interferometer_7)
+
+    fit = _visualize(dataset=dataset, image_path=tmp_path)
+
+    for filename in ("dataset", "fit", "galaxies", "inversion_0_0"):
+        assert str(tmp_path / f"{filename}.png") in plot_patch.paths, filename
+
+    assert (tmp_path / "galaxy_images.fits").exists()
+    assert _ext_names_from(tmp_path / "fit_dirty_images.fits") == [
+        "MASK",
+        "DIRTY_IMAGE_NATURAL",
+        "DIRTY_BEAM",
+        "DIRTY_MODEL_IMAGE_NATURAL",
+        "DIRTY_RESIDUAL_MAP_NATURAL",
+    ]
+
+    for hdu, array in (
+        (1, dataset.dirty_image_natural),
+        (2, dataset.dirty_beam),
+        (3, fit.dirty_model_image_natural),
+        (4, fit.dirty_residual_map_natural),
+    ):
+        np.testing.assert_allclose(
+            ag.ndarray_via_fits_from(
+                file_path=tmp_path / "fit_dirty_images.fits", hdu=hdu
+            ),
+            array.native_for_fits,
+            rtol=1.0e-6,
+            atol=1.0e-12,
+        )
+
+
+def test__visualizer__in_memory_dataset__fit_dirty_images_unchanged(
+    interferometer_7, tmp_path, plot_patch
+):
+    import numpy as np
+
+    fit = _visualize(dataset=interferometer_7, image_path=tmp_path)
+
+    for filename in ("dataset", "fit", "galaxies", "inversion_0_0"):
+        assert str(tmp_path / f"{filename}.png") in plot_patch.paths, filename
+
+    assert _ext_names_from(tmp_path / "fit_dirty_images.fits") == [
+        "MASK",
+        "DIRTY_IMAGE",
+        "DIRTY_NOISE_MAP",
+        "DIRTY_MODEL_IMAGE",
+        "DIRTY_RESIDUAL_MAP",
+        "DIRTY_NORMALIZED_RESIDUAL_MAP",
+        "DIRTY_CHI_SQUARED_MAP",
+    ]
+
+    for hdu, array in enumerate(
+        (
+            fit.dirty_image,
+            fit.dirty_noise_map,
+            fit.dirty_model_image,
+            fit.dirty_residual_map,
+            fit.dirty_normalized_residual_map,
+            fit.dirty_chi_squared_map,
+        ),
+        start=1,
+    ):
+        np.testing.assert_array_equal(
+            ag.ndarray_via_fits_from(
+                file_path=tmp_path / "fit_dirty_images.fits", hdu=hdu
+            ),
+            np.asarray(array.native_for_fits),
+        )