feat: LEGGOS area magnification on LensCalc (arc phase 6a)

autogalaxy/operate/lens_calc.py

@@ -808,6 +808,161 @@ def magnification_2d_via_hessian_from(self, grid, xp=np) -> aa.ArrayIrregular:
             return aa.ArrayIrregular(values=1.0 / det_A)
         return 1.0 / det_A
 
+    def _sub_pixel_magnifications_from(self, mask: aa.Mask2D, sub_size: int):
+        """
+        The signed Hessian magnification at the centre of every sub-pixel of the mask's unmasked pixels, and the area
+        of one sub-pixel, for the area magnification calculations below.
+        """
+        grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=sub_size)
+
+        sub_pixel_centres = aa.Grid2DIrregular(
+            values=np.asarray(grid.over_sampled.array)
+        )
+
+        magnifications = np.asarray(
+            self.magnification_2d_via_hessian_from(grid=sub_pixel_centres)
+        )
+
+        sub_pixel_area = mask.pixel_scales[0] * mask.pixel_scales[1] / sub_size**2
+
+        return magnifications, sub_pixel_area
+
+    def source_plane_area_from(self, mask: aa.Mask2D, sub_size: int = 4) -> float:
+        """
+        Returns the area of the source-plane region that the unmasked pixels of an image-plane mask map from.
+
+        Each small image-plane area `dA` maps from a source-plane area `dA / |mu|`, where `mu` is the local
+        magnification. Splitting every unmasked pixel into `sub_size x sub_size` sub-pixels and summing their areas
+        divided by the Hessian magnification at their centres gives the source-plane area (Eq. 7 of the LEGGOS II
+        paper, arXiv:2606.20804):
+
+        `A_src = sum_i A_pix,i / |mu_i|`
+
+        For a mask that covers one lensed image of a source (e.g. an arc), this is the area of the part of the source
+        that forms that image, which links the area magnification to the physical size of the source.
+
+        A mask covering several images of the same source counts that source once per image; use
+        `area_magnification_list_from` to split it into its images first.
+
+        Parameters
+        ----------
+        mask
+            The image-plane mask whose unmasked pixels define the region (e.g. one image of an arc).
+        sub_size
+            Each unmasked pixel is split into `sub_size x sub_size` sub-pixels, at whose centres the magnification is
+            evaluated. Increase it until the result converges, which matters most for pixels close to a critical
+            curve, where the magnification changes rapidly.
+        """
+        magnifications, sub_pixel_area = self._sub_pixel_magnifications_from(
+            mask=mask, sub_size=sub_size
+        )
+
+        return float(np.sum(sub_pixel_area / np.abs(magnifications)))
+
+    def area_magnification_from(self, mask: aa.Mask2D, sub_size: int = 4) -> float:
+        """
+        Returns the area magnification of the image-plane region defined by a mask: the area of the region divided by
+        the area of the source-plane region it maps from (Eq. 6 of the LEGGOS II paper, arXiv:2606.20804):
+
+        `mu_area = A_img / A_src`,  `A_src = sum_i A_pix,i / |mu_i|`
+
+        This is the standard magnification of an arc in cluster lensing. It needs only the local magnification map
+        (no source model and no forward ray-shooting of the source shape), and it is unsigned.
+
+        The mask must cover ONE lensed image: for a mask spanning several images of the same source, `A_src` counts
+        the source once per image and the ratio is not the magnification of any image. Use
+        `area_magnification_list_from` to split such a mask into its images.
+
+        Parameters
+        ----------
+        mask
+            The image-plane mask whose unmasked pixels define the region (e.g. one image of an arc).
+        sub_size
+            Each unmasked pixel is split into `sub_size x sub_size` sub-pixels, at whose centres the magnification is
+            evaluated (see `source_plane_area_from`).
+        """
+        magnifications, sub_pixel_area = self._sub_pixel_magnifications_from(
+            mask=mask, sub_size=sub_size
+        )
+
+        image_plane_area = sub_pixel_area * magnifications.shape[0]
+        source_plane_area = np.sum(sub_pixel_area / np.abs(magnifications))
+
+        return float(image_plane_area / source_plane_area)
+
+    def image_masks_from(self, mask: aa.Mask2D) -> List[aa.Mask2D]:
+        """
+        Splits the unmasked pixels of an image-plane mask into the separate lensed images they contain.
+
+        Images of a source are separated by critical curves, across which the magnification changes sign (the image
+        parity flips). Pixels are grouped by the sign of the Hessian magnification at their centres, and each group
+        is split into connected regions (pixels sharing an edge). Each region is returned as its own mask, largest
+        first.
+
+        Two images of the same parity that touch, or one image cut by the mask edge into separate pieces, are not
+        distinguished by this rule; check the returned masks for such cases.
+
+        Parameters
+        ----------
+        mask
+            The image-plane mask whose unmasked pixels contain one or more lensed images.
+        """
+        from scipy import ndimage
+
+        grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=1)
+
+        magnifications = np.asarray(
+            self.magnification_2d_via_hessian_from(
+                grid=aa.Grid2DIrregular(values=np.asarray(grid.array))
+            )
+        )
+
+        unmasked = ~np.asarray(mask)
+        parity = np.zeros(mask.shape_native)
+        parity[unmasked] = np.sign(magnifications)
+
+        image_masks = []
+
+        for sign in (1.0, -1.0):
+            labels, total = ndimage.label(parity == sign)
+
+            for label in range(1, total + 1):
+                image_masks.append(
+                    aa.Mask2D(
+                        mask=labels != label,
+                        pixel_scales=mask.pixel_scales,
+                        origin=mask.origin,
+                    )
+                )
+
+        return sorted(
+            image_masks,
+            key=lambda image_mask: np.sum(~np.asarray(image_mask)),
+            reverse=True,
+        )
+
+    def area_magnification_list_from(
+        self, mask: aa.Mask2D, sub_size: int = 4
+    ) -> List[float]:
+        """
+        Returns the area magnification of every lensed image in an image-plane mask, in the order of
+        `image_masks_from` (largest image first).
+
+        The mask is split into its images with `image_masks_from`, and `area_magnification_from` is computed for each.
+
+        Parameters
+        ----------
+        mask
+            The image-plane mask whose unmasked pixels contain one or more lensed images.
+        sub_size
+            Each unmasked pixel is split into `sub_size x sub_size` sub-pixels, at whose centres the magnification is
+            evaluated (see `source_plane_area_from`).
+        """
+        return [
+            self.area_magnification_from(mask=image_mask, sub_size=sub_size)
+            for image_mask in self.image_masks_from(mask=mask)
+        ]
+
     def contour_list_from(self, grid, contour_array):
         grid_contour = aa.Grid2DContour(
             grid=grid,

test_autogalaxy/operate/test_area_magnification.py

@@ -0,0 +1,139 @@
+"""
+Area magnification (LEGGOS II, arXiv:2606.20804, Eq. 6-7) on `LensCalc`.
+
+The singular isothermal sphere is the analytic witness: its magnification is
+`mu = theta / (theta - theta_E)`, so an image-plane annulus `r1 <= theta <= r2`
+outside the Einstein radius maps from the source-plane annulus
+`r1 - theta_E <= beta <= r2 - theta_E` and
+
+    mu_area = (r2^2 - r1^2) / ((r2 - theta_E)^2 - (r1 - theta_E)^2)
+
+exactly; inside the Einstein radius the image is parity-flipped (mu < 0).
+"""
+
+import numpy as np
+import pytest
+
+import autoarray as aa
+import autogalaxy as ag
+
+EINSTEIN_RADIUS = 1.0
+
+
+def _lens_calc():
+    mass = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=EINSTEIN_RADIUS)
+    return ag.LensCalc.from_mass_obj(mass)
+
+
+def _annulus_mask(inner, outer, pixel_scale=0.02):
+    shape = int(2 * (outer + 0.2) / pixel_scale)
+    return aa.Mask2D.circular_annular(
+        shape_native=(shape, shape),
+        pixel_scales=pixel_scale,
+        inner_radius=inner,
+        outer_radius=outer,
+    )
+
+
+def _analytic(r1, r2):
+    return (r2**2 - r1**2) / ((r2 - EINSTEIN_RADIUS) ** 2 - (r1 - EINSTEIN_RADIUS) ** 2)
+
+
+def test__sis_annulus__matches_the_analytic_area_magnification():
+    r1, r2 = 1.5, 2.0
+
+    mu_area = _lens_calc().area_magnification_from(
+        mask=_annulus_mask(inner=r1, outer=r2), sub_size=4
+    )
+
+    # measured -4.7e-4; the pixelized annulus edge, not the method, sets it
+    assert mu_area == pytest.approx(_analytic(r1, r2), rel=2.0e-3)
+
+
+def test__source_plane_area__matches_the_analytic_source_annulus():
+    r1, r2 = 1.5, 2.0
+
+    source_plane_area = _lens_calc().source_plane_area_from(
+        mask=_annulus_mask(inner=r1, outer=r2), sub_size=4
+    )
+
+    analytic = np.pi * ((r2 - EINSTEIN_RADIUS) ** 2 - (r1 - EINSTEIN_RADIUS) ** 2)
+
+    # measured -1.5e-4
+    assert source_plane_area == pytest.approx(analytic, rel=1.0e-3)
+
+
+def test__sub_size__converges():
+    """
+    On a coarse mask the per-pixel area integral is evaluated at sub-pixel
+    centres; refining the sub-pixels must converge on the same answer.
+    """
+    mask = _annulus_mask(inner=1.2, outer=1.8, pixel_scale=0.1)
+    lens_calc = _lens_calc()
+
+    reference = lens_calc.source_plane_area_from(mask=mask, sub_size=32)
+
+    errors = [
+        abs(lens_calc.source_plane_area_from(mask=mask, sub_size=s) / reference - 1)
+        for s in (1, 2, 4, 8)
+    ]
+
+    # measured 3.9e-4, 9.7e-5, 2.4e-5, 5.7e-6: second order in the sub-pixel size
+    assert all(a > b for a, b in zip(errors, errors[1:]))
+    assert errors[2] < 1.0e-4
+
+
+def test__image_masks__split_a_mask_across_the_critical_curve_by_parity():
+    """
+    An annulus straddling the Einstein ring contains the positive-parity outer
+    image ring and the negative-parity inner ring; they are returned separately,
+    and the outer one's area magnification matches the analytic value.
+    """
+    lens_calc = _lens_calc()
+
+    mask = _annulus_mask(inner=0.6, outer=1.6)
+
+    image_masks = lens_calc.image_masks_from(mask=mask)
+
+    assert len(image_masks) == 2
+
+    total = sum(np.sum(~np.asarray(image_mask)) for image_mask in image_masks)
+    assert total == np.sum(~np.asarray(mask))
+
+    signs = []
+    for image_mask in image_masks:
+        grid = aa.Grid2D.from_mask(mask=image_mask, over_sample_size=1)
+        mu = np.asarray(
+            lens_calc.magnification_2d_via_hessian_from(
+                grid=aa.Grid2DIrregular(values=np.asarray(grid.array))
+            )
+        )
+        assert np.all(np.sign(mu) == np.sign(mu[0]))
+        signs.append(np.sign(mu[0]))
+
+    assert sorted(signs) == [-1.0, 1.0]
+
+    mu_list = lens_calc.area_magnification_list_from(mask=mask, sub_size=4)
+
+    assert len(mu_list) == 2
+    assert all(mu > 0.0 for mu in mu_list)
+
+    # the positive-parity image is the outer ring, theta_E <= theta <= 1.6
+    # (measured -6e-4; the Einstein-ring boundary is pixelized)
+    outer = mu_list[signs.index(1.0)]
+    assert outer == pytest.approx(_analytic(EINSTEIN_RADIUS, 1.6), rel=3.0e-3)
+
+
+def test__a_mask_spanning_two_images_differs_from_each_image():
+    """
+    The documented trap: summing A_src over two images counts the source twice,
+    so the whole-mask ratio is not either image's magnification.
+    """
+    lens_calc = _lens_calc()
+
+    mask = _annulus_mask(inner=0.6, outer=1.6)
+
+    whole = lens_calc.area_magnification_from(mask=mask, sub_size=4)
+    per_image = lens_calc.area_magnification_list_from(mask=mask, sub_size=4)
+
+    assert all(abs(whole / mu - 1.0) > 0.01 for mu in per_image)