feat(weak-lensing): add WeakDataset + SimulatorShearYX (step 1)

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

autolens/weak/__init__.py

@@ -0,0 +1,2 @@
+from autolens.weak.dataset import WeakDataset
+from autolens.weak.simulator import SimulatorShearYX

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]

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)

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)