From 3700bc2464e94bacb52c367166f55fc27743e5a5 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 18 May 2026 20:18:08 +0100 Subject: [PATCH] feat(weak): aplt plotters for WeakDataset shear catalogues MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a `weak/plot/` package with five module-level helpers for visualising a WeakDataset, re-exported into the `aplt` namespace: - plot_shear_yx_2d — headless-quiver of the shear field (spin-2 convention) - plot_ellipticities, plot_phis — scalar-coloured scatter via plot_grid - plot_noise_map — per-galaxy noise scatter - subplot_weak_dataset — 2x2 mosaic combining the four panels Plotters access the shear field through .ellipticities / .phis only, so the [gamma_2, gamma_1] storage convention pinned by PyAutoGalaxy #366 stays encapsulated. Five matching tests under test_autolens/weak/plot/ exercise each helper via the existing plot_patch fixture. Step 2 of the weak-lensing series (issue #496). Workspace TODO replacement in autolens_workspace/scripts/weak/simulator.py follows in the workspace PR. Co-Authored-By: Claude Opus 4.7 (1M context) --- autolens/plot/__init__.py | 8 + autolens/weak/plot/__init__.py | 0 autolens/weak/plot/weak_dataset_plots.py | 233 ++++++++++++++++++ test_autolens/weak/plot/__init__.py | 0 .../weak/plot/test_weak_dataset_plots.py | 92 +++++++ 5 files changed, 333 insertions(+) create mode 100644 autolens/weak/plot/__init__.py create mode 100644 autolens/weak/plot/weak_dataset_plots.py create mode 100644 test_autolens/weak/plot/__init__.py create mode 100644 test_autolens/weak/plot/test_weak_dataset_plots.py diff --git a/autolens/plot/__init__.py b/autolens/plot/__init__.py index 8b0c05c95..3e6e1a991 100644 --- a/autolens/plot/__init__.py +++ b/autolens/plot/__init__.py @@ -65,6 +65,14 @@ from autolens.point.plot.fit_point_plots import subplot_fit as subplot_fit_point from autolens.point.plot.point_dataset_plots import subplot_dataset as subplot_point_dataset +from autolens.weak.plot.weak_dataset_plots import ( + plot_shear_yx_2d, + plot_ellipticities, + plot_phis, + plot_noise_map, + subplot_weak_dataset, +) + from autolens.lens.plot.subhalo_plots import ( subplot_detection_imaging, subplot_detection_fits, diff --git a/autolens/weak/plot/__init__.py b/autolens/weak/plot/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/autolens/weak/plot/weak_dataset_plots.py b/autolens/weak/plot/weak_dataset_plots.py new file mode 100644 index 000000000..47e66c9b2 --- /dev/null +++ b/autolens/weak/plot/weak_dataset_plots.py @@ -0,0 +1,233 @@ +""" +Module-level matplotlib helpers for visualising a ``WeakDataset``. + +A shear catalogue is a set of complex shear measurements ``(gamma_2, gamma_1)`` +at the ``(y, x)`` positions of background source galaxies. The natural way to +draw it is matplotlib's ``quiver`` with **headless line segments**, because +shear is a spin-2 quantity — a 180-degree rotation maps the shear back to +itself, so an arrowhead would suggest a directionality the data does not +have. This is the same convention used in weak-lensing science papers +(e.g. KiDS, DES). + +The plotters access the shear field exclusively through the derived properties +``.ellipticities`` (``|gamma|``) and ``.phis`` (position angle, in **degrees**) +defined on ``AbstractShearField``. Indexing the underlying ``[:, 0]`` / +``[:, 1]`` storage directly is deliberately avoided so the plotters keep +working if the ``[gamma_2, gamma_1]`` convention pinned by PyAutoGalaxy PR +#366 ever changes. +""" +from typing import Optional + +import numpy as np + +from autoarray.plot.grid import plot_grid +from autoarray.plot.utils import ( + subplots, + save_figure, + conf_subplot_figsize, + tight_layout, +) + + +def _positions_yx(shear_yx) -> np.ndarray: + """Return the ``(N, 2)`` ``[y, x]`` position array for a shear field.""" + grid = shear_yx.grid + return np.array(grid.array if hasattr(grid, "array") else grid) + + +def plot_shear_yx_2d( + shear_yx, + ax=None, + title: str = "Shear Field", + output_path: Optional[str] = None, + output_filename: str = "shear_yx", + output_format: Optional[str] = None, +): + """ + Plot a shear field as a quiver of headless line segments at galaxy positions. + + Each segment is centred on the galaxy position (``pivot="middle"``), has a + length proportional to the shear magnitude ``|gamma|`` and is oriented at + the shear position angle ``phi``. Segments are colour-coded by ``|gamma|``. + + Parameters + ---------- + shear_yx + A ``ShearYX2D`` / ``ShearYX2DIrregular`` carrying the shear vectors and + the ``(y, x)`` galaxy grid. + ax + Existing ``Axes`` to draw onto; ``None`` creates a new figure. + title + Figure title. + output_path, output_filename, output_format + Standard workspace output controls. When ``ax`` is supplied the saving + is the caller's responsibility (typically ``subplot_weak_dataset``). + """ + positions = _positions_yx(shear_yx) + y, x = positions[:, 0], positions[:, 1] + + mag = np.asarray(shear_yx.ellipticities) + phi_rad = np.deg2rad(np.asarray(shear_yx.phis)) + + u = mag * np.cos(phi_rad) + v = mag * np.sin(phi_rad) + + standalone = ax is None + if standalone: + fig, ax = subplots(1, 1) + + ax.quiver( + x, + y, + u, + v, + mag, + pivot="middle", + headwidth=0, + headlength=0, + headaxislength=0, + cmap="viridis", + ) + ax.set_xlabel('x (")') + ax.set_ylabel('y (")') + ax.set_title(title) + ax.set_aspect("equal") + + if standalone: + tight_layout() + save_figure( + fig, + path=output_path, + filename=output_filename, + format=output_format, + ) + + +def plot_ellipticities( + shear_yx, + ax=None, + title: str = r"Shear Magnitude $|\gamma|$", + output_path: Optional[str] = None, + output_filename: str = "shear_ellipticities", + output_format: Optional[str] = None, +): + """ + Plot a colour-coded scatter of the shear magnitude ``|gamma|`` at each galaxy. + + Delegates to ``autoarray.plot.grid.plot_grid`` with ``color_array`` set to + the per-galaxy ellipticities. + """ + plot_grid( + grid=_positions_yx(shear_yx), + ax=ax, + color_array=np.asarray(shear_yx.ellipticities), + colormap="viridis", + title=title, + output_path=output_path if ax is None else None, + output_filename=output_filename, + output_format=output_format, + ) + + +def plot_phis( + shear_yx, + ax=None, + title: str = r"Shear Position Angle $\phi$", + output_path: Optional[str] = None, + output_filename: str = "shear_phis", + output_format: Optional[str] = None, +): + """ + Plot a colour-coded scatter of the shear position angle ``phi`` at each galaxy. + + Position angles are cyclic, so a cyclic colormap (``twilight``) is used. + """ + plot_grid( + grid=_positions_yx(shear_yx), + ax=ax, + color_array=np.asarray(shear_yx.phis), + colormap="twilight", + title=title, + output_path=output_path if ax is None else None, + output_filename=output_filename, + output_format=output_format, + ) + + +def plot_noise_map( + dataset, + ax=None, + title: str = "Noise Map", + output_path: Optional[str] = None, + output_filename: str = "noise_map", + output_format: Optional[str] = None, +): + """ + Plot a colour-coded scatter of the per-galaxy shear noise at each position. + + Takes the full ``WeakDataset`` (not just the shear field) because the noise + map lives on the dataset. + """ + plot_grid( + grid=_positions_yx(dataset.shear_yx), + ax=ax, + color_array=np.asarray(dataset.noise_map), + colormap="magma", + title=title, + output_path=output_path if ax is None else None, + output_filename=output_filename, + output_format=output_format, + ) + + +def subplot_weak_dataset( + dataset, + output_path: Optional[str] = None, + output_filename: str = "subplot_weak_dataset", + output_format: Optional[str] = None, + title_prefix: Optional[str] = None, +): + """ + Produce a 2x2 subplot mosaic visualising a ``WeakDataset``. + + Panels: shear field, noise map, shear magnitude, shear position angle. + """ + fig, axes = subplots(2, 2, figsize=conf_subplot_figsize(2, 2)) + ax_quiver, ax_noise, ax_mag, ax_phi = ( + axes[0, 0], + axes[0, 1], + axes[1, 0], + axes[1, 1], + ) + + _prefix = f"{title_prefix.rstrip()} " if title_prefix else "" + name_part = f" — {dataset.name}" if dataset.name else "" + + plot_shear_yx_2d( + shear_yx=dataset.shear_yx, + ax=ax_quiver, + title=f"{_prefix}Shear Field{name_part}", + ) + plot_noise_map( + dataset=dataset, + ax=ax_noise, + title=f"{_prefix}Noise Map{name_part}", + ) + plot_ellipticities( + shear_yx=dataset.shear_yx, + ax=ax_mag, + title=f"{_prefix}Shear Magnitude{name_part}", + ) + plot_phis( + shear_yx=dataset.shear_yx, + ax=ax_phi, + title=f"{_prefix}Position Angle{name_part}", + ) + + tight_layout() + save_figure( + fig, + path=output_path, + filename=output_filename, + format=output_format, + ) diff --git a/test_autolens/weak/plot/__init__.py b/test_autolens/weak/plot/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test_autolens/weak/plot/test_weak_dataset_plots.py b/test_autolens/weak/plot/test_weak_dataset_plots.py new file mode 100644 index 000000000..64ed388b0 --- /dev/null +++ b/test_autolens/weak/plot/test_weak_dataset_plots.py @@ -0,0 +1,92 @@ +from pathlib import Path + +import autoarray as aa +import autolens as al + +import pytest + +from autolens.weak.plot.weak_dataset_plots import ( + plot_shear_yx_2d, + plot_ellipticities, + plot_phis, + plot_noise_map, + subplot_weak_dataset, +) + +directory = Path(__file__).resolve().parent + + +def _isothermal_tracer(): + lens = al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.6), + ) + source = al.Galaxy(redshift=1.0) + return al.Tracer(galaxies=[lens, source]) + + +@pytest.fixture(name="weak_dataset") +def make_weak_dataset(): + """Deterministic 4-galaxy WeakDataset built from an Isothermal lens.""" + grid = aa.Grid2DIrregular( + values=[(0.7, 0.5), (1.0, 1.0), (-0.3, 0.6), (-1.1, -0.8)] + ) + simulator = al.SimulatorShearYX(noise_sigma=0.0, seed=0) + return simulator.via_tracer_from( + tracer=_isothermal_tracer(), grid=grid, name="test" + ) + + +@pytest.fixture(name="plot_path") +def make_plot_path(): + return directory / "files" / "plots" / "weak_dataset" + + +def test__plot_shear_yx_2d(weak_dataset, plot_path, plot_patch): + plot_shear_yx_2d( + shear_yx=weak_dataset.shear_yx, + output_path=plot_path, + output_filename="shear_yx", + output_format="png", + ) + assert str(plot_path / "shear_yx.png") in plot_patch.paths + + +def test__plot_ellipticities(weak_dataset, plot_path, plot_patch): + plot_ellipticities( + shear_yx=weak_dataset.shear_yx, + output_path=plot_path, + output_filename="shear_ellipticities", + output_format="png", + ) + assert str(plot_path / "shear_ellipticities.png") in plot_patch.paths + + +def test__plot_phis(weak_dataset, plot_path, plot_patch): + plot_phis( + shear_yx=weak_dataset.shear_yx, + output_path=plot_path, + output_filename="shear_phis", + output_format="png", + ) + assert str(plot_path / "shear_phis.png") in plot_patch.paths + + +def test__plot_noise_map(weak_dataset, plot_path, plot_patch): + plot_noise_map( + dataset=weak_dataset, + output_path=plot_path, + output_filename="noise_map", + output_format="png", + ) + assert str(plot_path / "noise_map.png") in plot_patch.paths + + +def test__subplot_weak_dataset(weak_dataset, plot_path, plot_patch): + subplot_weak_dataset( + dataset=weak_dataset, + output_path=plot_path, + output_filename="subplot_weak_dataset", + output_format="png", + ) + assert str(plot_path / "subplot_weak_dataset.png") in plot_patch.paths