From 67110ff0a81b457e92b2f56bfece25810e63a633 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 30 Sep 2026 20:09:53 +0100 Subject: [PATCH] feat(interferometer): visualizer on array-free datasets (streaming phase 3, PyAutoArray#596) FitInterferometer.model_image_natural / dirty_model_image_natural / dirty_residual_map_natural; subplot_fit, subplot_fit_dirty_images, subplot_fit_real_space and fits_dirty_images draw the natural-weighted set on array-free datasets under the same filenames; the visualizer's InversionException fallback called logger(...) (TypeError) and now logs a warning. In-memory output unchanged. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01JZZksyZ8LTA4LLxoZjQMNF --- .../interferometer/fit_interferometer.py | 60 ++++++++ autogalaxy/interferometer/model/visualizer.py | 2 +- .../plot/fit_interferometer_plots.py | 116 ++++++++++++++- .../model/test_plotter_interferometer.py | 133 ++++++++++++++++++ .../plot/test_fit_interferometer_plotters.py | 54 +++++++ .../interferometer/test_fit_interferometer.py | 128 +++++++++++++++++ 6 files changed, 486 insertions(+), 7 deletions(-) diff --git a/autogalaxy/interferometer/fit_interferometer.py b/autogalaxy/interferometer/fit_interferometer.py index 9fd3da1ea..c6ecb6412 100644 --- a/autogalaxy/interferometer/fit_interferometer.py +++ b/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]: """ diff --git a/autogalaxy/interferometer/model/visualizer.py b/autogalaxy/interferometer/model/visualizer.py index 21f1f89bd..ccfaef274 100644 --- a/autogalaxy/interferometer/model/visualizer.py +++ b/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 diff --git a/autogalaxy/interferometer/plot/fit_interferometer_plots.py b/autogalaxy/interferometer/plot/fit_interferometer_plots.py index cadfce8f0..044b837dd 100644 --- a/autogalaxy/interferometer/plot/fit_interferometer_plots.py +++ b/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,6 +291,10 @@ 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 @@ -219,6 +302,27 @@ def fits_dirty_images(fit: FitInterferometer, output_path) -> None: 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, diff --git a/test_autogalaxy/interferometer/model/test_plotter_interferometer.py b/test_autogalaxy/interferometer/model/test_plotter_interferometer.py index 2a28a48b8..9010de943 100644 --- a/test_autogalaxy/interferometer/model/test_plotter_interferometer.py +++ b/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), + ) diff --git a/test_autogalaxy/interferometer/plot/test_fit_interferometer_plotters.py b/test_autogalaxy/interferometer/plot/test_fit_interferometer_plotters.py index b532f7baa..a56418130 100644 --- a/test_autogalaxy/interferometer/plot/test_fit_interferometer_plotters.py +++ b/test_autogalaxy/interferometer/plot/test_fit_interferometer_plotters.py @@ -32,3 +32,57 @@ def test__fit_sub_plot_real_space( ) assert str(plot_path / "fit_real_space.png") in plot_patch.paths + + +def test__fit_sub_plots__array_free_dataset( + interferometer_7, plot_path, plot_patch, monkeypatch +): + pytest.importorskip("nufftax") + + import autoarray as aa + import autogalaxy as ag + from autogalaxy.interferometer.plot import fit_interferometer_plots + + dataset = aa.Interferometer.from_stream( + [(interferometer_7.uv_wavelengths, interferometer_7.data, interferometer_7.noise_map)], + real_space_mask=interferometer_7.real_space_mask, + transformer_class=type(interferometer_7.transformer), + ) + + pixelization = ag.Pixelization( + mesh=ag.mesh.RectangularUniform(shape=(3, 3)), + regularization=ag.reg.Constant(coefficient=1.0), + ) + + fit = ag.FitInterferometer( + dataset=dataset, + galaxies=[ag.Galaxy(redshift=0.5, pixelization=pixelization)], + ) + + titles = [] + plot_array = fit_interferometer_plots.plot_array + + def _plot_array(*args, title=None, **kwargs): + titles.append(title) + return plot_array(*args, title=title, **kwargs) + + monkeypatch.setattr(fit_interferometer_plots, "plot_array", _plot_array) + + natural_titles = [ + "Dirty Image (Natural)", + "Dirty Model Image (Natural)", + "Dirty Residual Map (Natural)", + ] + + for function, filename in ( + (aplt.subplot_fit_interferometer, "fit"), + (aplt.subplot_fit_dirty_images, "fit_dirty_images"), + (aplt.subplot_fit_real_space, "fit_real_space"), + ): + titles.clear() + plot_patch.paths = [] + + function(fit=fit, output_path=plot_path, output_format="png") + + assert str(plot_path / f"{filename}.png") in plot_patch.paths + assert titles == natural_titles diff --git a/test_autogalaxy/interferometer/test_fit_interferometer.py b/test_autogalaxy/interferometer/test_fit_interferometer.py index 26d2cccb5..f222d237c 100644 --- a/test_autogalaxy/interferometer/test_fit_interferometer.py +++ b/test_autogalaxy/interferometer/test_fit_interferometer.py @@ -996,3 +996,131 @@ def test__fit_figure_of_merit__array_free_dataset__linear_light_only__matches_in ag.FitInterferometer(dataset=dataset_sparse, galaxies=galaxies).figure_of_merit, rel=1.0e-8, ) + + +def test__natural_dirty_images__array_free_pixelization_only_fit(interferometer_7): + pytest.importorskip("nufftax") + + dataset = _array_free_dataset_from(interferometer_7) + + fit = ag.FitInterferometer(dataset=dataset, galaxies=_pixelization_only_galaxies()) + + assert fit.inversion.transformer is None + + model_image = fit.model_image_natural + + assert isinstance(model_image, aa.Array2D) + np.testing.assert_allclose( + model_image.array, + fit.inversion.mapped_reconstructed_data.array, + rtol=1.0e-12, + atol=1.0e-12 * np.abs(model_image.array).max(), + ) + + dirty_model_image = fit.dirty_model_image_natural + dirty_residual_map = fit.dirty_residual_map_natural + + np.testing.assert_allclose( + dirty_residual_map.array, + dataset.dirty_image_natural.array - dirty_model_image.array, + rtol=1.0e-12, + ) + + # The in-memory sparse fit of the same visibilities gives the same natural images, and + # its dirty model image equals the transformer's natural dirty image of the model + # visibilities `F m`. + dataset_memory = interferometer_7.apply_sparse_operator() + + fit_memory = ag.FitInterferometer( + dataset=dataset_memory, galaxies=_pixelization_only_galaxies() + ) + + expected = _natural_dirty_image_of_model_data( + dataset=dataset_memory, fit=fit_memory + ) + + np.testing.assert_allclose( + fit_memory.dirty_model_image_natural.array, + expected.array, + rtol=1.0e-8, + atol=1.0e-8 * np.abs(expected.array).max(), + ) + np.testing.assert_allclose( + dirty_model_image.array, + fit_memory.dirty_model_image_natural.array, + rtol=1.0e-6, + atol=1.0e-6 * np.abs(expected.array).max(), + ) + + +def _natural_dirty_image_of_model_data(dataset, fit): + """ + The independent reference for `fit.dirty_model_image_natural`: the transformer's natural-weighted dirty + image of the fit's actual model visibilities `fit.model_data`, `Re(F^H (w m_vis)) / sum(w)` with + `w = 1 / sigma^2` per component (the weighting of `Interferometer.dirty_image_natural`). + """ + noise_map = dataset.noise_map.array + visibilities = np.asarray(fit.model_data.array) + weighted = aa.Visibilities( + visibilities=visibilities.real * noise_map.real**-2.0 + + 1j * visibilities.imag * noise_map.imag**-2.0 + ) + return dataset.transformer.image_from(visibilities=weighted) / float( + np.sum(noise_map.real**-2.0) + ) + + +@pytest.mark.parametrize("linear_component", ["linear_light_profile", "pixelization"]) +def test__natural_dirty_images__galaxy_with_ordinary_and_linear_light__matches_model_data( + interferometer_7, linear_component +): + """ + A galaxy with both an ordinary light profile and a linear component (a linear light profile, or a + pixelization) has only its linear reconstruction in `galaxy_image_dict`; `model_image_natural` must + still include the ordinary light, so `dirty_model_image_natural` is the natural dirty image of the + fit's model visibilities `model_data`. + """ + dataset = interferometer_7.apply_sparse_operator(use_jax=False) + + if linear_component == "linear_light_profile": + galaxy = ag.Galaxy( + redshift=0.5, + bulge=ag.lp.Sersic(intensity=0.001, centre=(0.05, 0.05)), + disk=ag.lp_linear.Sersic(centre=(0.0, 0.1)), + ) + else: + galaxy = ag.Galaxy( + redshift=0.5, + bulge=ag.lp.Sersic(intensity=1.0e-4), + pixelization=ag.Pixelization( + mesh=ag.mesh.RectangularUniform(shape=(3, 3)), + regularization=ag.reg.Constant(coefficient=1.0e-3), + ), + ) + + fit = ag.FitInterferometer(dataset=dataset, galaxies=[galaxy]) + + # Both the ordinary and the linear light contribute, so omitting either is detectable. + assert np.abs(fit.profile_image.array).max() > 0.0 + assert np.abs(fit.inversion.mapped_reconstructed_data.array).max() > 0.0 + + np.testing.assert_allclose( + fit.model_image_natural.array, + fit.profile_image.array + fit.inversion.mapped_reconstructed_data.array, + rtol=1.0e-12, + ) + + expected = _natural_dirty_image_of_model_data(dataset=dataset, fit=fit) + + np.testing.assert_allclose( + fit.dirty_model_image_natural.array, + expected.array, + rtol=1.0e-8, + atol=1.0e-8 * np.abs(expected.array).max(), + ) + np.testing.assert_allclose( + fit.dirty_residual_map_natural.array, + dataset.dirty_image_natural.array - expected.array, + rtol=1.0e-8, + atol=1.0e-8 * np.abs(expected.array).max(), + )