Merge pull request #508 from PyAutoLabs/feature/potential-correction-interferometer

feat: zero-fill extrapolation for input pixelized mass profiles

Author: Jammy2211

autogalaxy/profiles/mass/input/input_deflections.py

@@ -18,6 +18,7 @@ def __init__(
         deflections_x: np.ndarray,
         image_plane_grid: aa.type.Grid2DLike,
         mask: aa.type.Mask2D,
+        extrapolate: str = "nearest",
         Hy: Optional[spmatrix] = None,
         Hx: Optional[spmatrix] = None,
     ):
@@ -56,6 +57,14 @@ def __init__(
             The cleaned 2D mask defining the unmasked pixels (see
             ``aa.util.derivative.cleaned_mask_from``); its ``pixel_scale``
             sets the finite-difference step of the derived convergence.
+        extrapolate
+            The extrapolation behaviour outside the unmasked pixels' convex
+            hull: ``"nearest"`` (default; the field continues beyond the
+            grid) or ``"zero"`` (the field vanishes outside it — required
+            when the profile represents a localized correction on a
+            sub-region of a larger grid, e.g. an arc-restricted dpsi mesh,
+            where nearest extrapolation would produce spurious constant
+            deflections everywhere else).
         Hy
             The sparse first-derivative operator along y of the mask; built
             from the mask if not input.
@@ -69,6 +78,7 @@ def __init__(
         self.deflections_x = np.asarray(deflections_x)
         self.image_plane_grid = np.asarray(image_plane_grid)
         self.mask = mask
+        self.extrapolate = extrapolate
         self.Hy = Hy
         self.Hx = Hx
 
@@ -85,9 +95,9 @@ def _build_interpolators(self):
         ) * 0.5
 
         self.tri = Delaunay(np.fliplr(self.image_plane_grid))
-        self.interp_defl_y = LinearNDInterpolatorExt(self.tri, self.deflections_y)
-        self.interp_defl_x = LinearNDInterpolatorExt(self.tri, self.deflections_x)
-        self.interp_kappa = LinearNDInterpolatorExt(self.tri, self.convergence_slim)
+        self.interp_defl_y = LinearNDInterpolatorExt(self.tri, fill=self.extrapolate, values=self.deflections_y)
+        self.interp_defl_x = LinearNDInterpolatorExt(self.tri, fill=self.extrapolate, values=self.deflections_x)
+        self.interp_kappa = LinearNDInterpolatorExt(self.tri, fill=self.extrapolate, values=self.convergence_slim)
 
     @aa.decorators.to_array
     def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):

autogalaxy/profiles/mass/input/input_potential.py

@@ -17,6 +17,7 @@ def __init__(
         lensing_potential: np.ndarray,
         image_plane_grid: aa.type.Grid2DLike,
         mask: aa.type.Mask2D,
+        extrapolate: str = "nearest",
         Hy: Optional[spmatrix] = None,
         Hx: Optional[spmatrix] = None,
         Hyy: Optional[spmatrix] = None,
@@ -54,6 +55,14 @@ def __init__(
             ``aa.util.derivative.cleaned_mask_from``); its ``pixel_scale``
             sets the finite-difference step of the derived deflections and
             convergence.
+        extrapolate
+            The extrapolation behaviour outside the unmasked pixels' convex
+            hull: ``"nearest"`` (default; the field continues beyond the
+            grid) or ``"zero"`` (the field vanishes outside it — required
+            when the profile represents a localized correction on a
+            sub-region of a larger grid, e.g. an arc-restricted dpsi mesh,
+            where nearest extrapolation would produce spurious constant
+            deflections everywhere else).
         Hy
             The sparse first-derivative operator along y of the mask; built
             from the mask if not input.
@@ -72,6 +81,7 @@ def __init__(
         self.lensing_potential = np.asarray(lensing_potential)
         self.image_plane_grid = np.asarray(image_plane_grid)
         self.mask = mask
+        self.extrapolate = extrapolate
         self.Hy = Hy
         self.Hx = Hx
         self.Hyy = Hyy
@@ -96,10 +106,10 @@ def _build_interpolators(self):
         ) * 0.5
 
         self.tri = Delaunay(np.fliplr(self.image_plane_grid))
-        self.interp_psi = LinearNDInterpolatorExt(self.tri, self.lensing_potential)
-        self.interp_defl_y = LinearNDInterpolatorExt(self.tri, self.deflections_y)
-        self.interp_defl_x = LinearNDInterpolatorExt(self.tri, self.deflections_x)
-        self.interp_kappa = LinearNDInterpolatorExt(self.tri, self.convergence_slim)
+        self.interp_psi = LinearNDInterpolatorExt(self.tri, fill=self.extrapolate, values=self.lensing_potential)
+        self.interp_defl_y = LinearNDInterpolatorExt(self.tri, fill=self.extrapolate, values=self.deflections_y)
+        self.interp_defl_x = LinearNDInterpolatorExt(self.tri, fill=self.extrapolate, values=self.deflections_x)
+        self.interp_kappa = LinearNDInterpolatorExt(self.tri, fill=self.extrapolate, values=self.convergence_slim)
 
     @aa.decorators.to_array
     def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):

autogalaxy/profiles/mass/input/interp.py

@@ -4,11 +4,20 @@
 
 
 class LinearNDInterpolatorExt:
-    def __init__(self, points, values):
+    def __init__(self, points, values, fill: str = "nearest"):
         """
         Linear interpolation over a Delaunay triangulation of scattered 2D
-        points, falling back to nearest-neighbour interpolation outside the
-        convex hull so extrapolated values are never NaN.
+        points, with a choice of extrapolation behaviour outside the convex
+        hull so extrapolated values are never NaN:
+
+        - ``fill="nearest"`` (default): nearest-neighbour extrapolation —
+          appropriate when the sampled field genuinely continues beyond the
+          hull (e.g. a source brightness evaluated slightly off-mesh).
+        - ``fill="zero"``: zero extrapolation — appropriate when the field is
+          only defined on the sampled region and must vanish outside it
+          (e.g. localized potential corrections: nearest extrapolation would
+          smear constant non-zero values — and for their deflections,
+          spurious constant deflections — across the whole grid).
 
         Ported from the ``potential_correction`` package of Cao et al. 2025
         (https://github.com/caoxiaoyue/lensing_potential_correction; cite via
@@ -22,13 +31,23 @@ def __init__(self, points, values):
             triangulation of them.
         values
             The values interpolated.
+        fill
+            The extrapolation behaviour outside the convex hull:
+            ``"nearest"`` or ``"zero"``.
         """
+        if fill not in ("nearest", "zero"):
+            raise ValueError(f"fill must be 'nearest' or 'zero', got {fill!r}")
+        self.fill = fill
         self.funcinterp = LinearNDInterpolator(points, values)
-        self.funcnearest = NearestNDInterpolator(points, values)
+        self.funcnearest = (
+            NearestNDInterpolator(points, values) if fill == "nearest" else None
+        )
 
     def __call__(self, *args):
         z = self.funcinterp(*args)
         chk = np.isnan(z)
         if chk.any():
-            return np.where(chk, self.funcnearest(*args), z)
+            if self.fill == "nearest":
+                return np.where(chk, self.funcnearest(*args), z)
+            return np.where(chk, 0.0, z)
         return z

test_autogalaxy/profiles/mass/input/test_input_potential.py

@@ -80,3 +80,49 @@ def test__operators_can_be_preloaded():
 
     deflections = np.asarray(profile.deflections_yx_2d_from(grid=grid))
     assert deflections[:, 0] == pytest.approx(1.0, abs=1.0e-8)
+
+
+def test__zero_extrapolation__deflections_vanish_outside_mesh():
+    mask, grid = masked_setup()
+    grid_arr = np.asarray(grid)
+    potential = 2.0 * grid_arr[:, 0] + 3.0 * grid_arr[:, 1]
+
+    profile_nearest = ag.mp.InputPotential(
+        lensing_potential=potential, image_plane_grid=grid_arr, mask=mask
+    )
+    profile_zero = ag.mp.InputPotential(
+        lensing_potential=potential,
+        image_plane_grid=grid_arr,
+        mask=mask,
+        extrapolate="zero",
+    )
+
+    far_outside = aa.Grid2DIrregular(values=[(6.0, 6.0), (-7.0, 5.0)])
+
+    deflections_nearest = np.asarray(
+        profile_nearest.deflections_yx_2d_from(grid=far_outside)
+    )
+    deflections_zero = np.asarray(profile_zero.deflections_yx_2d_from(grid=far_outside))
+
+    # nearest extrapolation smears constant non-zero deflections outward;
+    # zero extrapolation vanishes
+    assert not np.allclose(deflections_nearest, 0.0)
+    assert deflections_zero == pytest.approx(0.0, abs=1.0e-12)
+
+    # inside the mesh the two modes agree exactly
+    inside = aa.Grid2DIrregular(values=[(0.1, -0.2), (0.7, 0.4)])
+    assert np.asarray(profile_zero.deflections_yx_2d_from(grid=inside)) == pytest.approx(
+        np.asarray(profile_nearest.deflections_yx_2d_from(grid=inside)), abs=1.0e-12
+    )
+
+
+def test__invalid_extrapolate_raises():
+    mask, grid = masked_setup()
+    grid_arr = np.asarray(grid)
+    with pytest.raises(ValueError):
+        ag.mp.InputPotential(
+            lensing_potential=np.ones(grid_arr.shape[0]),
+            image_plane_grid=grid_arr,
+            mask=mask,
+            extrapolate="taper",
+        )