From cd26722b4f192d211dfb294c2ac4fdc01c8604e9 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 29 Oct 2026 12:45:14 +0000 Subject: [PATCH] feat: add introductory four-panel imaging plot with shared stretch --- autoarray/dataset/plot/imaging_plots.py | 91 +++++++++++++++++++++++ autoarray/plot/array.py | 7 +- autoarray/plot/utils.py | 37 +++++++-- test_autoarray/plot/test_imaging_intro.py | 89 ++++++++++++++++++++++ 4 files changed, 218 insertions(+), 6 deletions(-) create mode 100644 test_autoarray/plot/test_imaging_intro.py diff --git a/autoarray/dataset/plot/imaging_plots.py b/autoarray/dataset/plot/imaging_plots.py index 2f8beb899..6d5805c46 100644 --- a/autoarray/dataset/plot/imaging_plots.py +++ b/autoarray/dataset/plot/imaging_plots.py @@ -146,6 +146,97 @@ def subplot_imaging_dataset( +def _intro_panel(array, radius=None, excluded_pixels=None): + """Prepare a display-only crop with excluded values missing, not zeroed. + + Preserve the native pixel/coordinate registration; fitting arrays are never + mutated. Coordinates, ticks and colour mapping remain the renderer's job. + """ + import numpy as np + from autoarray.plot.utils import zoom_array + + if radius is not None and (not np.isfinite(radius) or radius <= 0): + raise ValueError("radius must be a finite positive number") + if radius is None and excluded_pixels is None: + array = zoom_array(array) + values = np.array(array.native.array, dtype=float, copy=True) + values[np.asarray(array.mask.array, dtype=bool)] = np.nan + if excluded_pixels is not None: + excluded = np.asarray(excluded_pixels, dtype=bool) + if excluded.shape != values.shape: + raise ValueError("excluded_pixels must have the dataset image shape") + values[excluded] = np.nan + xmin, xmax, ymin, ymax = array.geometry.extent + if radius is not None: + ny, nx = values.shape + dx, dy = (xmax - xmin) / nx, (ymax - ymin) / ny + x = xmin + (np.arange(nx) + 0.5) * dx + y = ymax - (np.arange(ny) + 0.5) * dy + rows, cols = np.flatnonzero(abs(y) <= radius), np.flatnonzero(abs(x) <= radius) + if not len(rows) or not len(cols): + raise ValueError("radius does not include any pixel centres") + r0, r1, c0, c1 = rows[0], rows[-1] + 1, cols[0], cols[-1] + 1 + values = values[r0:r1, c0:c1] + xmin, xmax, ymin, ymax = xmin + c0 * dx, xmin + c1 * dx, ymax - r1 * dy, ymax - r0 * dy + return values, (xmin, xmax, ymin, ymax) + + +def subplot_imaging_dataset_intro( + dataset, + output_path: Optional[str] = None, + output_filename: str = "dataset_intro", + output_format: str = None, + colormap=None, + stretch: Optional[float] = 0.01, + radius: Optional[float] = None, + excluded_pixels=None, +): + """Four-panel introduction: Image, Noise map, Signal-to-noise map and PSF. + + Uses the ordinary plot renderer and configured colormap, fonts, titles, + arcsecond ticks, colourbars and saving. Unlike the diagnostic 3x3 subplot, + it omits duplicate log panels and over-sampling maps. Numerical sampling + and the dataset itself are unchanged. + + ``stretch`` is the image's asinh linear-width fraction: smaller values + reveal faint structure next to bright light; None selects linear display. + The S/N and PSF panels use asinh fractions 0.1 and 0.003, respectively; + noise remains linear. All colourbars retain original measurement values. + ``colormap=None`` inherits configuration on every panel. + + ``radius`` optionally crops the first three panels about (0, 0), in + arcseconds, for display only. None follows the usual mask-zoom convention. + ``excluded_pixels`` is a native-shaped boolean array (True = exclude) for + deliberately down-weighted neighbours. These pixels, dataset-masked values + and non-finite measurements use the colormap's standard bad-value colour; + they do not determine the display scale. No clipping or data mutation. + """ + from autoarray.plot.array import plot_array + from autoarray.plot.utils import apply_labels + + panels = [ + (dataset.data, "Image", stretch, None), + (dataset.noise_map, "Noise map", None, None), + (dataset.signal_to_noise_map, "Signal-to-noise map", 0.1, ""), + ] + prepared = [(_intro_panel(array, radius, excluded_pixels), title, scale, unit) + for array, title, scale, unit in panels] + fig, axes = subplots(2, 2, figsize=conf_subplot_figsize(2, 2)) + for ax, ((values, extent), title, scale, unit) in zip(axes.flat, prepared): + plot_array(values, ax=ax, extent=extent, title=title, + colormap=colormap, stretch=scale, cb_unit=unit) + if dataset.psf is not None: + plot_array(dataset.psf.kernel, ax=axes[1, 1], title="PSF", + colormap=colormap, stretch=0.003, cb_unit="") + else: + apply_labels(axes[1, 1], title="PSF", is_subplot=True) + axes[1, 1].text(0.5, 0.5, "No PSF supplied", ha="center", va="center", + transform=axes[1, 1].transAxes) + tight_layout() + subplot_save(fig, output_path, output_filename, output_format) + return fig + + def subplot_imaging_dataset_list( dataset_list, output_path=None, diff --git a/autoarray/plot/array.py b/autoarray/plot/array.py index 1852dac2c..e1cb407e6 100644 --- a/autoarray/plot/array.py +++ b/autoarray/plot/array.py @@ -69,6 +69,7 @@ def plot_array( output_path: Optional[str] = None, output_filename: str = "array", output_format: str = None, + stretch: Optional[float] = None, ) -> None: """ Plot a 2D array (image) using ``plt.imshow``. @@ -127,6 +128,10 @@ def plot_array( Explicit color scale limits. use_log10 When ``True`` a ``LogNorm`` is applied. + stretch + Optional positive asinh linear-width fraction of the displayed range, + e.g. 0.01 to reveal faint structure beside a bright peak. Only changes + the colour mapping; incompatible with ``use_log10=True``. origin_imshow Passed directly to ``imshow`` (``"upper"`` or ``"lower"``). arcsec_after_decimal @@ -194,7 +199,7 @@ def plot_array( fig = ax.get_figure() # --- colour normalisation -------------------------------------------------- - norm = norm_from(array=array, use_log10=use_log10, vmin=vmin, vmax=vmax) + norm = norm_from(array=array, use_log10=use_log10, vmin=vmin, vmax=vmax, stretch=stretch) # Compute the axes-box aspect ratio from the data extent so that the # physical cell is correctly shaped and tight_layout has no whitespace diff --git a/autoarray/plot/utils.py b/autoarray/plot/utils.py index 9dde85441..32be106af 100644 --- a/autoarray/plot/utils.py +++ b/autoarray/plot/utils.py @@ -255,6 +255,17 @@ def symmetric_cmap_from(array, symmetric_value=None): If given, fix the half-range to this value (``vmin=-symmetric_value``, ``vmax=+symmetric_value``). + stretch + Optional positive asinh linear-width fraction of the displayed range. + Smaller values bring up faint structure; negative values remain defined. + Mutually exclusive with ``use_log10``. ``None`` preserves the existing + linear/log behavior. + + stretch : float or None + Positive asinh linear-width fraction of the displayed value range. + Mutually exclusive with ``use_log10``. Non-finite values are ignored + when deriving limits. ``None`` retains the existing linear/log scale. + Returns ------- matplotlib.colors.Normalize or None @@ -272,7 +283,7 @@ def symmetric_cmap_from(array, symmetric_value=None): return colors.Normalize(vmin=-abs_max, vmax=abs_max) -def norm_from(array=None, use_log10=False, vmin=None, vmax=None): +def norm_from(array=None, use_log10=False, vmin=None, vmax=None, stretch=None): """Build the matplotlib colour norm for *array* from flat arguments. The single implementation behind every colour scale the plot functions @@ -307,8 +318,8 @@ def norm_from(array=None, use_log10=False, vmin=None, vmax=None): Parameters ---------- array - The values being coloured. Only read when *use_log10* is ``True`` and - no finite *vmax* is given; ``None`` is accepted and takes the fallback. + The values being coloured. Read for log or asinh scaling when limits + are omitted; ``None`` is accepted and takes a finite fallback. use_log10 When ``True`` a ``LogNorm`` is returned, floored at the configured ``log10_min_value``. @@ -320,6 +331,22 @@ def norm_from(array=None, use_log10=False, vmin=None, vmax=None): ------- matplotlib.colors.Normalize or None """ + if stretch is not None: + if use_log10: + raise ValueError("stretch and use_log10 are mutually exclusive") + if not np.isfinite(stretch) or stretch <= 0: + raise ValueError("stretch must be a finite positive fraction") + from matplotlib.colors import AsinhNorm + + values = np.ma.masked_invalid(np.ma.asarray(array, dtype=float)).compressed() if array is not None else np.array([]) + lo = float(vmin) if vmin is not None else (float(values.min()) if values.size else 0.0) + hi = float(vmax) if vmax is not None else (float(values.max()) if values.size else 1.0) + if not np.isfinite(lo) or not np.isfinite(hi): + raise ValueError("asinh colour limits must be finite") + if hi <= lo: + hi = lo + max(abs(lo) * 0.01, 1e-6) + return AsinhNorm(linear_width=stretch * (hi - lo), vmin=lo, vmax=hi) + if use_log10: log10_min = _conf_log10_min_value() @@ -693,8 +720,8 @@ def _colorbar_tick_values(norm) -> Optional[List[float]]: import matplotlib.colors as mcolors lo, hi = float(norm.vmin), float(norm.vmax) - if isinstance(norm, mcolors.LogNorm): - mid = 10 ** ((np.log10(lo) + np.log10(hi)) / 2.0) + if isinstance(norm, (mcolors.LogNorm, mcolors.AsinhNorm)): + mid = float(norm.inverse(0.5)) else: mid = (lo + hi) / 2.0 return [lo, mid, hi] diff --git a/test_autoarray/plot/test_imaging_intro.py b/test_autoarray/plot/test_imaging_intro.py new file mode 100644 index 000000000..3b86e3ed9 --- /dev/null +++ b/test_autoarray/plot/test_imaging_intro.py @@ -0,0 +1,89 @@ +"""The intro compositor uses the ordinary renderer, style and configuration.""" +import numpy as np +import matplotlib.pyplot as plt +import pytest +from matplotlib.colors import AsinhNorm +from autonerves import conf +import autoarray as aa +from autoarray.plot.array import plot_array +from autoarray.plot.utils import norm_from, _colorbar_tick_values +from autoarray.dataset.plot.imaging_plots import subplot_imaging_dataset_intro, _intro_panel + + +def dataset(): + values = np.arange(81, dtype=float).reshape(9, 9) - 3 + return aa.Imaging( + data=aa.Array2D.no_mask(values, pixel_scales=0.1), + noise_map=aa.Array2D.no_mask(np.ones((9, 9)), pixel_scales=0.1), + psf=aa.Convolver(kernel=aa.Array2D.no_mask(np.ones((3, 3)) / 9, pixel_scales=0.1)), + ) + + +def test_intro_inherits_config_and_standard_style(tmp_path, monkeypatch): + general = conf.instance["visualize"]["general"] + monkeypatch.setitem(general, "colormap", "cividis") + d = dataset() + original = d.data.native.array.copy() + fig = subplot_imaging_dataset_intro(d, output_path=tmp_path, output_format="png") + panels = [ax for ax in fig.axes if ax.images] + assert [ax.get_title() for ax in panels] == ["Image", "Noise map", "Signal-to-noise map", "PSF"] + assert all(ax.images[0].get_cmap().name == "cividis" for ax in panels) + direct, ax = plt.subplots() + values, extent = _intro_panel(d.data) + plot_array(values, extent=extent, ax=ax, title="Image", stretch=0.01) + assert ax.title.get_fontsize() == panels[0].title.get_fontsize() + assert [t.get_text() for t in ax.get_xticklabels()] == [t.get_text() for t in panels[0].get_xticklabels()] + assert ax.get_yticklabels()[0].get_rotation() == panels[0].get_yticklabels()[0].get_rotation() + assert isinstance(panels[0].images[0].norm, AsinhNorm) + np.testing.assert_array_equal(d.data.native.array, original) + assert (tmp_path / 'dataset_intro.png').is_file() + plt.close(direct) + + +def test_explicit_colormap_and_excluded_noise(tmp_path): + d = dataset() + excluded = np.zeros((9, 9), dtype=bool) + excluded[2, 3] = True + noise = np.ones((9, 9)) + noise[2, 3] = 1e8 + d = aa.Imaging(data=d.data, noise_map=aa.Array2D.no_mask(noise, pixel_scales=.1), psf=d.psf) + fig = subplot_imaging_dataset_intro(d, colormap="magma", radius=0.25, + excluded_pixels=excluded, output_path=tmp_path, output_format="png") + panels = [ax for ax in fig.axes if ax.images] + assert all(ax.images[0].get_cmap().name == "magma" for ax in panels) + assert panels[1].images[0].norm.vmax < 2 + assert np.ma.getmaskarray(panels[0].images[0].get_array()).any() + assert d.noise_map.native.array[2, 3] == 1e8 + + +def test_crop_preserves_pixel_coordinates(): + d = dataset() + values, extent = _intro_panel(d.data, radius=0.25) + assert values.shape == (5, 5) + np.testing.assert_allclose(extent, [-0.25, 0.25, -0.25, 0.25]) + assert values[2, 2] == d.data.native.array[4, 4] + + +@pytest.mark.parametrize('values', [np.zeros((2, 2)), np.full((2, 2), np.nan), np.array([-2., 0., 100.])]) +def test_asinh_range_is_finite_and_invertible(values): + norm = norm_from(values, stretch=0.01) + assert np.isfinite([norm.vmin, norm.vmax]).all() + assert norm.vmax > norm.vmin + ticks = _colorbar_tick_values(norm) + np.testing.assert_allclose(norm(ticks), [0, 0.5, 1], atol=1e-10) + + +@pytest.mark.parametrize('stretch', [0, -1, np.nan, np.inf]) +def test_invalid_stretch_rejected(stretch): + with pytest.raises(ValueError, match='stretch'): + norm_from(np.ones(3), stretch=stretch) + + +def test_asinh_cannot_be_combined_with_log(): + with pytest.raises(ValueError, match='mutually exclusive'): + norm_from([1, 2], stretch=.01, use_log10=True) + + +def test_ordinary_norm_remains_unchanged(): + assert norm_from([1, 2]) is None + assert not isinstance(norm_from([1, 2], use_log10=True), AsinhNorm)