Merge pull request #523 from PyAutoLabs/feature/weak-visualization

feat(weak): aplt plotters for WeakDataset shear catalogues (#496)

Author: Jammy2211

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,

autolens/weak/plot/__init__.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,
+    )

autolens/weak/plot/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