diff --git a/autolens/__init__.py b/autolens/__init__.py index a994ba505..3de2c4ea1 100644 --- a/autolens/__init__.py +++ b/autolens/__init__.py @@ -117,6 +117,8 @@ from .point.solver.shape_solver import ShapeSolver from .quantity.fit_quantity import FitQuantity from .quantity.model.analysis import AnalysisQuantity +from .weak.dataset import WeakDataset +from .weak.simulator import SimulatorShearYX from . import exc from . import mock as m diff --git a/autolens/weak/__init__.py b/autolens/weak/__init__.py new file mode 100644 index 000000000..07552c3d9 --- /dev/null +++ b/autolens/weak/__init__.py @@ -0,0 +1,2 @@ +from autolens.weak.dataset import WeakDataset +from autolens.weak.simulator import SimulatorShearYX diff --git a/autolens/weak/dataset.py b/autolens/weak/dataset.py new file mode 100644 index 000000000..a349cadd3 --- /dev/null +++ b/autolens/weak/dataset.py @@ -0,0 +1,102 @@ +""" +Data structure for weak gravitational lensing observations. + +Weak lensing measures the small, statistical distortion of background source-galaxy shapes induced by foreground +mass. The observable is a *shear catalogue*: a set of complex shear components ``(gamma_2, gamma_1)`` measured at +the (y, x) sky positions of a population of background galaxies, together with a per-galaxy noise estimate +(typically dominated by intrinsic shape noise — each galaxy has a random unlensed ellipticity that adds to its +measured shear). + +``WeakDataset`` holds those three quantities together. It is the weak-lensing analogue of +:class:`autolens.point.dataset.PointDataset` and is the input to a :class:`autolens.weak.fit.FitWeak` (added in a +follow-up step). + +The shear catalogue is stored as a :class:`autogalaxy.util.shear_field.ShearYX2DIrregular` so the convention is +the same one pinned by ``PyAutoGalaxy`` PR #366: column 0 is :math:`\\gamma_2`, column 1 is :math:`\\gamma_1`, +and the (y, x) galaxy positions are accessible via ``shear_yx.grid``. +""" +from typing import List, Optional, Union + +import autoarray as aa + +from autogalaxy.util.shear_field import ShearYX2DIrregular + + +class WeakDataset: + def __init__( + self, + shear_yx: ShearYX2DIrregular, + noise_map: Union[float, aa.ArrayIrregular, List[float]], + name: str = "", + ): + """ + A weak-lensing shear catalogue: a ``ShearYX2DIrregular`` shear field plus a per-galaxy noise map. + + Parameters + ---------- + shear_yx + The measured (or simulated) shear at each background source-galaxy position. Shape + ``[total_galaxies, 2]`` with column 0 = :math:`\\gamma_2`, column 1 = :math:`\\gamma_1`. The (y, x) + positions of the galaxies are carried by ``shear_yx.grid``. + noise_map + The per-galaxy shear noise standard deviation (one value per galaxy). For weak lensing this is + dominated by intrinsic shape noise, typically in the range 0.2 - 0.4 per shear component. A scalar + broadcasts to a constant noise level across all galaxies. + name + Optional label, mirroring ``PointDataset.name``. Used by downstream fitting code to pair this + dataset with a corresponding model component when multiple datasets are fitted simultaneously. + """ + self.name = name + + if not isinstance(shear_yx, ShearYX2DIrregular): + raise TypeError( + "WeakDataset.shear_yx must be a ShearYX2DIrregular instance; " + f"got {type(shear_yx).__name__}." + ) + + self.shear_yx = shear_yx + + n_galaxies = len(shear_yx) + + if isinstance(noise_map, (float, int)): + noise_map = [float(noise_map)] * n_galaxies + + if not isinstance(noise_map, aa.ArrayIrregular): + noise_map = aa.ArrayIrregular(values=list(noise_map)) + + if len(noise_map) != n_galaxies: + raise ValueError( + f"WeakDataset.noise_map has length {len(noise_map)} but shear_yx has " + f"{n_galaxies} entries; the two must match." + ) + + self.noise_map = noise_map + + @property + def positions(self) -> aa.Grid2DIrregular: + """The (y, x) sky positions of the source galaxies the shear is measured at.""" + return self.shear_yx.grid + + @property + def n_galaxies(self) -> int: + """Number of source galaxies in the catalogue.""" + return len(self.shear_yx) + + @property + def info(self) -> str: + """A short human-readable summary of the dataset, mirroring ``PointDataset.info``.""" + return ( + f"name : {self.name}\n" + f"n_galaxies : {self.n_galaxies}\n" + f"shear_yx : {self.shear_yx}\n" + f"noise_map : {self.noise_map}\n" + ) + + def extent_from(self, buffer: float = 0.1) -> List[float]: + """The axis-aligned bounding box of the source-galaxy positions, padded by ``buffer`` on each side.""" + positions = self.positions + y_max = max(positions[:, 0]) + buffer + y_min = min(positions[:, 0]) - buffer + x_max = max(positions[:, 1]) + buffer + x_min = min(positions[:, 1]) - buffer + return [y_min, y_max, x_min, x_max] diff --git a/autolens/weak/simulator.py b/autolens/weak/simulator.py new file mode 100644 index 000000000..490a38c68 --- /dev/null +++ b/autolens/weak/simulator.py @@ -0,0 +1,139 @@ +""" +Simulate weak-lensing shear catalogues from a ``Tracer``. + +The shear field of a strong lens system is computed by ``Tracer.shear_yx_2d_via_hessian_from``, which differentiates +the deflection-angle field to derive the lensing Hessian and from there the (gamma_2, gamma_1) shear at any (y, x). +That gives a *noise-free* shear field; this module adds two extras needed for a realistic weak-lensing simulation: + +1. **Shape noise.** Each background source galaxy has a random unlensed ellipticity, drawn here as iid Gaussian + noise per shear component with standard deviation ``noise_sigma``. Realistic values are around 0.2 - 0.4. + +2. **Source-galaxy positions.** Either provide an explicit grid of (y, x) source positions, or let the simulator + draw a uniform-random distribution of ``n_galaxies`` positions inside a square arc-second extent — sufficient + for development work; future iterations may swap in number-density / luminosity-function-driven distributions. + +The output is a :class:`autolens.weak.dataset.WeakDataset`. +""" +from typing import Optional + +import numpy as np + +import autoarray as aa + +from autogalaxy.operate.lens_calc import LensCalc +from autogalaxy.util.shear_field import ShearYX2DIrregular + +from autolens.weak.dataset import WeakDataset + + +class SimulatorShearYX: + def __init__(self, noise_sigma: float = 0.3, seed: Optional[int] = None): + """ + Simulator for weak-lensing shear catalogues. + + Parameters + ---------- + noise_sigma + Standard deviation of the per-component Gaussian shape noise added to each measured shear vector. + Values around 0.2 - 0.4 are realistic for ground-based weak-lensing surveys; ``0.0`` produces a + noise-free shear field, useful for unit tests. + seed + Optional seed for the random number generator. When set, both the noise and the random-position + draws (for ``via_tracer_random_positions_from``) are reproducible. + """ + self.noise_sigma = float(noise_sigma) + self.seed = seed + self._rng = np.random.default_rng(seed) + + def via_tracer_from( + self, + tracer, + grid: aa.Grid2DIrregular, + name: str = "", + ) -> WeakDataset: + """ + Simulate a weak-lensing shear catalogue from a ``Tracer`` evaluated at the given (y, x) source positions. + + Computes ``tracer.shear_yx_2d_via_hessian_from(grid=grid)`` — which goes through ``LensCalc`` under the + hood — and adds iid Gaussian shape noise per component with std ``self.noise_sigma``. + + Parameters + ---------- + tracer + A PyAutoLens ``Tracer`` object exposing ``shear_yx_2d_via_hessian_from``. + grid + The (y, x) source-galaxy positions where the shear is measured. Must be an ``aa.Grid2DIrregular`` + (or convertible to one). + name + Optional label passed through to ``WeakDataset.name``. + """ + if not isinstance(grid, aa.Grid2DIrregular): + grid = aa.Grid2DIrregular(values=grid) + + true_shear = self._true_shear_yx_from(tracer=tracer, grid=grid) + + if self.noise_sigma > 0.0: + noise = self._rng.normal( + loc=0.0, scale=self.noise_sigma, size=true_shear.shape + ) + shear_array = np.asarray(true_shear) + noise + else: + shear_array = np.asarray(true_shear) + + shear_yx = ShearYX2DIrregular(values=shear_array, grid=grid) + + noise_map = aa.ArrayIrregular( + values=[self.noise_sigma] * len(grid) + ) + + return WeakDataset(shear_yx=shear_yx, noise_map=noise_map, name=name) + + def via_tracer_random_positions_from( + self, + tracer, + n_galaxies: int, + grid_extent: float = 3.0, + name: str = "", + ) -> WeakDataset: + """ + Simulate a weak-lensing shear catalogue at ``n_galaxies`` uniform-random source positions. + + Galaxy positions are drawn uniformly inside a square ``[-grid_extent, +grid_extent]`` arc-second box. + This is the simplest sensible distribution for development purposes; production weak-lensing simulations + typically use a survey-specific number density / redshift distribution. + + Parameters + ---------- + tracer + A PyAutoLens ``Tracer`` object. + n_galaxies + Number of background source galaxies to simulate. + grid_extent + Half-width of the square (in arc-seconds) inside which positions are drawn. + name + Optional label passed through to ``WeakDataset.name``. + """ + positions = self._rng.uniform( + low=-grid_extent, high=grid_extent, size=(n_galaxies, 2) + ) + grid = aa.Grid2DIrregular(values=positions) + return self.via_tracer_from(tracer=tracer, grid=grid, name=name) + + @staticmethod + def _true_shear_yx_from(tracer, grid: aa.Grid2DIrregular): + """ + Evaluate the noise-free shear field of ``tracer`` on ``grid``. + + - If the input exposes ``shear_yx_2d_via_hessian_from`` directly, use it. + - Else, if it exposes ``deflections_between_planes_from`` (it duck-types as a ``Tracer``), build a + multi-plane ``LensCalc`` via ``from_tracer``. + - Otherwise treat the input as a mass-like object (single ``MassProfile``, ``Galaxy``, ``Galaxies``) + and build a single-plane ``LensCalc`` via ``from_mass_obj``. This fallback keeps the simulator + usable in unit tests that don't want to spin up a full ``Tracer``. + """ + method = getattr(tracer, "shear_yx_2d_via_hessian_from", None) + if method is not None: + return method(grid=grid) + if hasattr(tracer, "deflections_between_planes_from"): + return LensCalc.from_tracer(tracer).shear_yx_2d_via_hessian_from(grid=grid) + return LensCalc.from_mass_obj(tracer).shear_yx_2d_via_hessian_from(grid=grid) diff --git a/test_autolens/weak/__init__.py b/test_autolens/weak/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test_autolens/weak/test_dataset.py b/test_autolens/weak/test_dataset.py new file mode 100644 index 000000000..efaf5f2eb --- /dev/null +++ b/test_autolens/weak/test_dataset.py @@ -0,0 +1,79 @@ +import numpy as np +import pytest + +import autoarray as aa +import autolens as al + +from autogalaxy.util.shear_field import ShearYX2DIrregular + + +def _make_shear_yx(positions, values): + grid = aa.Grid2DIrregular(values=positions) + return ShearYX2DIrregular(values=np.asarray(values), grid=grid) + + +def test__weak_dataset__construct_with_shear_yx_irregular_and_scalar_noise(): + shear_yx = _make_shear_yx( + positions=[(0.5, 0.5), (1.0, -0.3), (-0.7, 0.2)], + values=[(0.01, 0.02), (-0.03, 0.04), (0.05, -0.06)], + ) + + dataset = al.WeakDataset(shear_yx=shear_yx, noise_map=0.3, name="weak") + + assert dataset.name == "weak" + assert dataset.n_galaxies == 3 + assert isinstance(dataset.noise_map, aa.ArrayIrregular) + assert list(dataset.noise_map) == pytest.approx([0.3, 0.3, 0.3]) + assert dataset.shear_yx is shear_yx + + +def test__weak_dataset__construct_with_per_galaxy_noise_list(): + shear_yx = _make_shear_yx( + positions=[(0.0, 0.0), (1.0, 1.0)], values=[(0.0, 0.0), (0.1, 0.1)] + ) + + dataset = al.WeakDataset(shear_yx=shear_yx, noise_map=[0.2, 0.4]) + + assert list(dataset.noise_map) == pytest.approx([0.2, 0.4]) + + +def test__weak_dataset__rejects_non_shear_yx_input(): + grid = aa.Grid2DIrregular(values=[(0.0, 0.0)]) + bad_shear = aa.VectorYX2DIrregular(values=np.array([[0.0, 0.0]]), grid=grid) + + with pytest.raises(TypeError): + al.WeakDataset(shear_yx=bad_shear, noise_map=0.3) + + +def test__weak_dataset__rejects_mismatched_noise_length(): + shear_yx = _make_shear_yx( + positions=[(0.0, 0.0), (1.0, 1.0)], values=[(0.0, 0.0), (0.1, 0.1)] + ) + + with pytest.raises(ValueError): + al.WeakDataset(shear_yx=shear_yx, noise_map=[0.2, 0.4, 0.6]) + + +def test__weak_dataset__positions_match_shear_grid(): + shear_yx = _make_shear_yx( + positions=[(0.5, 0.5), (1.0, -0.3)], values=[(0.0, 0.0), (0.0, 0.0)] + ) + + dataset = al.WeakDataset(shear_yx=shear_yx, noise_map=0.3) + + assert dataset.positions is shear_yx.grid + + +def test__weak_dataset__extent_from_pads_position_bounding_box(): + shear_yx = _make_shear_yx( + positions=[(0.0, 0.0), (1.0, 2.0), (-0.5, -0.4)], + values=[(0.0, 0.0), (0.0, 0.0), (0.0, 0.0)], + ) + + dataset = al.WeakDataset(shear_yx=shear_yx, noise_map=0.3) + + y_min, y_max, x_min, x_max = dataset.extent_from(buffer=0.1) + assert y_min == pytest.approx(-0.6) + assert y_max == pytest.approx(1.1) + assert x_min == pytest.approx(-0.5) + assert x_max == pytest.approx(2.1) diff --git a/test_autolens/weak/test_simulator.py b/test_autolens/weak/test_simulator.py new file mode 100644 index 000000000..0007ca81f --- /dev/null +++ b/test_autolens/weak/test_simulator.py @@ -0,0 +1,117 @@ +import numpy as np +import pytest + +import autoarray as aa +import autogalaxy as ag +import autolens as al + +from autogalaxy.operate.lens_calc import LensCalc +from autogalaxy.util.shear_field import ShearYX2DIrregular + + +def _isothermal_tracer(einstein_radius=1.6, ell_comps=(0.0, 0.05)): + lens = al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.0, 0.0), + ell_comps=ell_comps, + einstein_radius=einstein_radius, + ), + ) + source = al.Galaxy(redshift=1.0) + return al.Tracer(galaxies=[lens, source]) + + +def test__simulator_shear_yx__zero_noise_matches_tracer_shear_via_hessian(): + """ + With ``noise_sigma=0`` the simulator must reproduce the tracer's Hessian-derived shear exactly — the + simulator only adds noise on top of that, never modifies the underlying shear computation. The simulator + uses ``LensCalc.from_tracer`` when the input duck-types as a ``Tracer``, so that's what we compare to. + """ + tracer = _isothermal_tracer() + grid = aa.Grid2DIrregular(values=[(0.7, 0.5), (1.0, 1.0), (-0.3, 0.6)]) + + simulator = al.SimulatorShearYX(noise_sigma=0.0, seed=0) + dataset = simulator.via_tracer_from(tracer=tracer, grid=grid, name="sim") + + expected = LensCalc.from_tracer(tracer).shear_yx_2d_via_hessian_from(grid=grid) + + assert isinstance(dataset, al.WeakDataset) + assert dataset.name == "sim" + assert dataset.n_galaxies == 3 + np.testing.assert_allclose( + np.asarray(dataset.shear_yx), np.asarray(expected), rtol=1e-6, atol=1e-9 + ) + assert list(dataset.noise_map) == pytest.approx([0.0, 0.0, 0.0]) + + +def test__simulator_shear_yx__noise_changes_values_but_preserves_shape_and_grid(): + tracer = _isothermal_tracer() + grid = aa.Grid2DIrregular(values=[(0.7, 0.5), (1.0, 1.0)]) + + truth = LensCalc.from_tracer(tracer).shear_yx_2d_via_hessian_from(grid=grid) + simulator = al.SimulatorShearYX(noise_sigma=0.3, seed=42) + dataset = simulator.via_tracer_from(tracer=tracer, grid=grid) + + assert isinstance(dataset.shear_yx, ShearYX2DIrregular) + assert np.asarray(dataset.shear_yx).shape == (2, 2) + np.testing.assert_array_equal(np.asarray(dataset.positions), np.asarray(grid)) + assert np.any(np.asarray(dataset.shear_yx) != np.asarray(truth)) + assert list(dataset.noise_map) == pytest.approx([0.3, 0.3]) + + +def test__simulator_shear_yx__seed_makes_runs_reproducible(): + tracer = _isothermal_tracer() + grid = aa.Grid2DIrregular(values=[(0.7, 0.5), (1.0, 1.0), (-0.3, 0.6)]) + + a = al.SimulatorShearYX(noise_sigma=0.3, seed=123).via_tracer_from( + tracer=tracer, grid=grid + ) + b = al.SimulatorShearYX(noise_sigma=0.3, seed=123).via_tracer_from( + tracer=tracer, grid=grid + ) + + np.testing.assert_allclose( + np.asarray(a.shear_yx), np.asarray(b.shear_yx), rtol=1e-12, atol=1e-12 + ) + + +def test__simulator_shear_yx__random_positions_within_extent_and_correct_count(): + tracer = _isothermal_tracer() + simulator = al.SimulatorShearYX(noise_sigma=0.0, seed=7) + + dataset = simulator.via_tracer_random_positions_from( + tracer=tracer, n_galaxies=20, grid_extent=2.0 + ) + + assert dataset.n_galaxies == 20 + positions = np.asarray(dataset.positions) + assert positions.shape == (20, 2) + assert positions.min() >= -2.0 + assert positions.max() <= 2.0 + + +def test__simulator_shear_yx__integration_isothermal_shear_magnitude_is_nontrivial(): + """ + Integration check: an isothermal lens at the origin produces a shear field whose magnitude near the + Einstein radius is of order ~kappa(theta_E) — i.e. clearly non-zero. This catches gross regressions + where the simulator silently returns zero or constant shear. + """ + einstein_radius = 1.6 + tracer = _isothermal_tracer(einstein_radius=einstein_radius, ell_comps=(0.0, 0.0)) + + # Galaxies sitting on a ring at the Einstein radius + angles = np.linspace(0.0, 2.0 * np.pi, 8, endpoint=False) + ring = np.stack( + [einstein_radius * np.sin(angles), einstein_radius * np.cos(angles)], axis=1 + ) + grid = aa.Grid2DIrregular(values=ring) + + simulator = al.SimulatorShearYX(noise_sigma=0.0, seed=0) + dataset = simulator.via_tracer_from(tracer=tracer, grid=grid) + + magnitudes = np.sqrt(np.sum(np.asarray(dataset.shear_yx) ** 2, axis=1)) + + assert np.all(np.isfinite(magnitudes)) + # For a singular isothermal sphere, |gamma| ~ kappa = 0.5 * theta_E / theta = 0.5 at theta_E. + assert magnitudes.mean() == pytest.approx(0.5, rel=0.2)