feat: input pixelized mass profiles (potential correction phase 2)

autogalaxy/config/notation.yaml

@@ -69,6 +69,9 @@ label:
     zeroth_signal_scale: V
   superscript:
     ExternalShear: ext
+    GaussianRandomField: grf
+    InputDeflections: input
+    InputPotential: input
     Mesh: mesh
     Point: point
     SMBH: smbh

autogalaxy/profiles/mass/__init__.py

@@ -66,3 +66,9 @@
     ChameleonSph,
 )
 from .sheets import ExternalPotential, ExternalShear, MassSheet
+from .input import (
+    InputDeflections,
+    InputPotential,
+    GaussianRandomField,
+    LinearNDInterpolatorExt,
+)

autogalaxy/profiles/mass/input/__init__.py

@@ -0,0 +1,4 @@
+from .interp import LinearNDInterpolatorExt
+from .input_deflections import InputDeflections
+from .input_potential import InputPotential
+from .gaussian_random_field import GaussianRandomField

autogalaxy/profiles/mass/input/gaussian_random_field.py

@@ -0,0 +1,127 @@
+from functools import cached_property
+
+import numpy as np
+
+import autoarray as aa
+
+from autogalaxy.profiles.mass.abstract.abstract import MassProfile
+from autogalaxy.profiles.mass.input.input_potential import InputPotential
+
+
+def gaussian_random_field_from(
+    shape_native, pixel_scale: float, power_amplitude: float, power_slope: float, seed: int
+) -> np.ndarray:
+    """
+    A real Gaussian random field realization with isotropic power-law power
+    spectrum P(k) = power_amplitude * k^(-power_slope).
+
+    The field is generated by filtering white Gaussian noise in Fourier space
+    (multiplying its transform by sqrt(P(k)) and inverse transforming), which
+    guarantees a real field whose power spectrum follows P(k); the k = 0 mode
+    is zeroed so the field has zero mean. Note this normalization convention
+    differs from the ``powerbox`` package used by the original
+    implementation — the spectrum's shape and seed-reproducibility are what
+    potential-correction validation relies on.
+
+    Parameters
+    ----------
+    shape_native
+        The 2D shape of the realization.
+    pixel_scale
+        The pixel size, setting the physical frequencies k.
+    power_amplitude
+        The amplitude of the power spectrum.
+    power_slope
+        The (positive) slope of the power-law spectrum, P(k) ~ k^-slope.
+    seed
+        The random seed, making the realization reproducible.
+    """
+    rng = np.random.default_rng(seed)
+    white_noise = rng.normal(size=shape_native)
+
+    ky = np.fft.fftfreq(shape_native[0], d=pixel_scale) * 2.0 * np.pi
+    kx = np.fft.fftfreq(shape_native[1], d=pixel_scale) * 2.0 * np.pi
+    k_grid = np.sqrt(ky[:, None] ** 2 + kx[None, :] ** 2)
+
+    power = np.zeros_like(k_grid)
+    nonzero = k_grid > 0
+    power[nonzero] = power_amplitude * k_grid[nonzero] ** (-power_slope)
+
+    field_ft = np.fft.fft2(white_noise) * np.sqrt(power)
+    return np.fft.ifft2(field_ft).real
+
+
+class GaussianRandomField(MassProfile):
+    def __init__(
+        self,
+        mask: aa.Mask2D,
+        power_amplitude: float = 1.0,
+        power_slope: float = 1.0,
+        seed: int = 1,
+    ):
+        """
+        A mass profile whose lensing potential is a Gaussian random field
+        realization with power spectrum P(k) = power_amplitude * k^(-power_slope),
+        used to simulate extended perturbations of a smooth lens-mass model
+        (e.g. for validating potential-correction reconstructions).
+
+        The realization is evaluated on the unmasked pixels of the input mask
+        and wrapped in an ``InputPotential``, from which the deflection angles
+        and convergence are derived via the mask's sparse derivative
+        operators.
+
+        Ported from the ``potential_correction`` package of Cao et al. 2025
+        (https://github.com/caoxiaoyue/lensing_potential_correction). If you
+        use this profile in your research, please cite Cao et al. 2025;
+        citation materials are provided at
+        https://github.com/caoxiaoyue/potential_correction_paper. The
+        realization here uses a plain-numpy Fourier filter rather than the
+        original's ``powerbox`` dependency (see
+        ``gaussian_random_field_from``).
+
+        Parameters
+        ----------
+        mask
+            The cleaned 2D mask (an ``aa.Mask2D`` carrying the pixel scale)
+            on whose unmasked pixels the potential is defined (see
+            ``aa.util.derivative.cleaned_mask_from``).
+        power_amplitude
+            The amplitude of the potential's power spectrum.
+        power_slope
+            The (positive) slope of the power-law spectrum, P(k) ~ k^-slope.
+        seed
+            The random seed, making the realization reproducible.
+        """
+        self.mask = mask
+        self.power_amplitude = power_amplitude
+        self.power_slope = power_slope
+        self.seed = seed
+        super().__init__()
+
+    @cached_property
+    def lensing_potential_native(self) -> np.ndarray:
+        return gaussian_random_field_from(
+            shape_native=self.mask.shape_native,
+            pixel_scale=self.mask.pixel_scale,
+            power_amplitude=self.power_amplitude,
+            power_slope=self.power_slope,
+            seed=self.seed,
+        )
+
+    @cached_property
+    def input_potential(self) -> InputPotential:
+        grid = aa.Grid2D.from_mask(mask=self.mask)
+        return InputPotential(
+            lensing_potential=self.lensing_potential_native[~np.asarray(self.mask)],
+            image_plane_grid=np.asarray(grid),
+            mask=self.mask,
+        )
+
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return self.input_potential.convergence_2d_from(grid=grid, xp=xp, **kwargs)
+
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return self.input_potential.potential_2d_from(grid=grid, xp=xp, **kwargs)
+
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return self.input_potential.deflections_yx_2d_from(grid=grid, xp=xp, **kwargs)

autogalaxy/profiles/mass/input/input_deflections.py

@@ -0,0 +1,116 @@
+from typing import Optional
+
+import numpy as np
+from scipy.sparse import spmatrix
+from scipy.spatial import Delaunay
+
+import autoarray as aa
+from autoarray.operators import derivative_util
+
+from autogalaxy.profiles.mass.abstract.abstract import MassProfile
+from autogalaxy.profiles.mass.input.interp import LinearNDInterpolatorExt
+
+
+class InputDeflections(MassProfile):
+    def __init__(
+        self,
+        deflections_y: np.ndarray,
+        deflections_x: np.ndarray,
+        image_plane_grid: aa.type.Grid2DLike,
+        mask: aa.type.Mask2D,
+        Hy: Optional[spmatrix] = None,
+        Hx: Optional[spmatrix] = None,
+    ):
+        """
+        A pixelized mass model defined by known deflection angles on the
+        unmasked pixels of an image-plane grid (e.g. from a previous lens
+        model, a particle simulation, or a potential-correction
+        reconstruction).
+
+        Deflections at arbitrary positions are evaluated by linear
+        interpolation over a Delaunay triangulation of the unmasked pixels
+        (nearest-neighbour fallback outside the convex hull). The convergence
+        is derived from the input deflections via the sparse first-derivative
+        operators of the mask, kappa = 0.5 * (dalpha_y/dy + dalpha_x/dx). The
+        lensing potential is not derivable from deflections alone and returns
+        zeros.
+
+        Ported from the ``potential_correction`` package of Cao et al. 2025
+        (https://github.com/caoxiaoyue/lensing_potential_correction). If you
+        use this profile in your research, please cite Cao et al. 2025;
+        citation materials are provided at
+        https://github.com/caoxiaoyue/potential_correction_paper.
+
+        Parameters
+        ----------
+        deflections_y
+            The 1D (slim) array of the y components of the deflection angles
+            on the unmasked pixels.
+        deflections_x
+            The 1D (slim) array of the x components of the deflection angles
+            on the unmasked pixels.
+        image_plane_grid
+            The [n_unmasked, 2] (y, x) grid of the unmasked pixels the
+            deflection angles are defined on.
+        mask
+            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.
+        Hy
+            The sparse first-derivative operator along y of the mask; built
+            from the mask if not input.
+        Hx
+            The sparse first-derivative operator along x of the mask; built
+            from the mask if not input.
+        """
+        super().__init__()
+
+        self.deflections_y = np.asarray(deflections_y)
+        self.deflections_x = np.asarray(deflections_x)
+        self.image_plane_grid = np.asarray(image_plane_grid)
+        self.mask = mask
+        self.Hy = Hy
+        self.Hx = Hx
+
+        self._build_interpolators()
+
+    def _build_interpolators(self):
+        if self.Hy is None or self.Hx is None:
+            self.Hy, self.Hx = derivative_util.derivative_1st_operators_from(
+                np.asarray(self.mask), pixel_scale=self.mask.pixel_scale
+            )
+
+        self.convergence_slim = (
+            self.Hy @ self.deflections_y + self.Hx @ self.deflections_x
+        ) * 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)
+
+    @aa.decorators.to_array
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        grid = np.asarray(grid)
+        return self.interp_kappa(grid[:, 1], grid[:, 0])
+
+    @aa.decorators.to_array
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        return np.zeros(shape=np.asarray(grid).shape[0])
+
+    @aa.decorators.to_vector_yx
+    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
+        """
+        Calculate the deflection angles at a given set of arc-second gridded
+        coordinates, by interpolating the input deflection angles.
+
+        Parameters
+        ----------
+        grid
+            The grid of (y,x) arc-second coordinates the deflection angles
+            are computed on.
+        """
+        grid = np.asarray(grid)
+        deflections_y = self.interp_defl_y(grid[:, 1], grid[:, 0])
+        deflections_x = self.interp_defl_x(grid[:, 1], grid[:, 0])
+        return np.stack((deflections_y, deflections_x), axis=-1)