Skip to content

Commit 8aa9513

Browse files
Jammy2211claude
authored and
Jammy2211
committed
Add masked-grid derivative operators, coarse-mesh interpolation and mask regularizations (potential correction phase 1)
Ports the generic linear-algebra layer of the gravitational-imaging (potential correction) technique of Cao et al. 2025 (https://github.com/caoxiaoyue/lensing_potential_correction) into autoarray: sparse finite-difference derivative operators on masked 2D grids, coarse-mesh bilinear interpolation matrices, and the CurvatureMask / FourthOrderMask regularizations. Parity-certified against the original; cite via https://github.com/caoxiaoyue/potential_correction_paper. Phase 1 of PyAutoLabs/PyAutoLens#618. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 059eb22 commit 8aa9513

10 files changed

Lines changed: 1588 additions & 0 deletions

File tree

‎autoarray/inversion/regularization/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
from .adapt_split import AdaptSplit
88
from .brightness_zeroth import BrightnessZeroth
99
from .adapt_split_zeroth import AdaptSplitZeroth
10+
from .curvature_mask import CurvatureMask
11+
from .fourth_order_mask import FourthOrderMask
1012
from .gaussian_kernel import GaussianKernel
1113
from .exponential_kernel import ExponentialKernel
1214
from .matern_kernel import MaternKernel

‎autoarray/inversion/regularization/curvature_mask.py‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
from __future__ import annotations
2+
import numpy as np
3+
from typing import TYPE_CHECKING
4+
5+
if TYPE_CHECKING:
6+
from autoarray.inversion.linear_obj.linear_obj import LinearObj
7+
8+
from autoarray.inversion.regularization.abstract import AbstractRegularization
9+
from autoarray.operators import derivative_util
10+
11+
12+
def curvature_reg_matrix_via_mask_from(mask, pixel_scale: float = 1.0) -> np.ndarray:
13+
"""
14+
The curvature regularization matrix of the unmasked pixels of a
15+
rectangular masked 2D grid, H = Hxx^T Hxx + Hyy^T Hyy, where Hxx / Hyy are
16+
forward second-difference operators which degrade gracefully to first /
17+
zeroth order at the mask edge (see
18+
``derivative_util.forward_difference_operators_from``).
19+
20+
This is the curvature regularization scheme used by the
21+
gravitational-imaging (potential correction) technique of Cao et al. 2025
22+
(https://github.com/caoxiaoyue/lensing_potential_correction; cite via
23+
https://github.com/caoxiaoyue/potential_correction_paper), applied to the
24+
pixelized corrections of the lensing potential defined on a coarse
25+
rectangular mesh.
26+
27+
Parameters
28+
----------
29+
mask
30+
The 2D bool mask (``True`` = masked) of the rectangular grid whose
31+
unmasked pixels are regularized.
32+
pixel_scale
33+
The finite-difference step; regularization matrices are
34+
conventionally built with 1.0, the coefficient absorbing the scale.
35+
36+
Returns
37+
-------
38+
The [n_unmasked, n_unmasked] regularization matrix.
39+
"""
40+
return derivative_util.forward_difference_reg_matrix_from(
41+
mask=mask, pixel_scale=pixel_scale, max_order=2
42+
).toarray()
43+
44+
45+
class CurvatureMask(AbstractRegularization):
46+
def __init__(self, coefficient: float = 1.0):
47+
"""
48+
Curvature regularization on the unmasked pixels of a rectangular
49+
masked 2D grid.
50+
51+
Each unmasked pixel is regularized with a forward second-difference
52+
stencil along both grid directions, degrading to first / zeroth order
53+
where the mask edge truncates the stencil, penalising curvature in
54+
the reconstructed solution. This contrasts mapper-based schemes
55+
(e.g. ``Constant``), which regularize via mesh-neighbour differences:
56+
here the linear object is defined on a masked rectangular grid and
57+
must expose its ``mask``.
58+
59+
This is the curvature regularization scheme of the
60+
gravitational-imaging (potential correction) technique, applied to
61+
pixelized corrections of the lensing potential; it is ported from the
62+
``potential_correction`` package of Cao et al. 2025
63+
(https://github.com/caoxiaoyue/lensing_potential_correction). If you
64+
use it in your research, please cite Cao et al. 2025; citation
65+
materials are provided at
66+
https://github.com/caoxiaoyue/potential_correction_paper.
67+
68+
Parameters
69+
----------
70+
coefficient
71+
The regularization coefficient which multiplies the matrix,
72+
setting the strength of the smoothing.
73+
"""
74+
self.coefficient = coefficient
75+
76+
super().__init__()
77+
78+
def regularization_weights_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray:
79+
"""
80+
Returns the regularization weights of this regularization scheme,
81+
which are equal for every parameter.
82+
83+
Parameters
84+
----------
85+
linear_obj
86+
The linear object which uses these weights when performing
87+
regularization.
88+
"""
89+
return self.coefficient * xp.ones(linear_obj.params)
90+
91+
def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray:
92+
"""
93+
Returns the regularization matrix with shape [pixels, pixels].
94+
95+
Parameters
96+
----------
97+
linear_obj
98+
The linear object which uses this matrix to perform
99+
regularization. It must expose the ``mask`` of the rectangular
100+
masked 2D grid its parameters are defined on.
101+
102+
Returns
103+
-------
104+
The regularization matrix.
105+
"""
106+
return self.coefficient * curvature_reg_matrix_via_mask_from(
107+
mask=linear_obj.mask
108+
)

‎autoarray/inversion/regularization/fourth_order_mask.py‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
from __future__ import annotations
2+
import numpy as np
3+
from typing import TYPE_CHECKING
4+
5+
if TYPE_CHECKING:
6+
from autoarray.inversion.linear_obj.linear_obj import LinearObj
7+
8+
from autoarray.inversion.regularization.abstract import AbstractRegularization
9+
from autoarray.operators import derivative_util
10+
11+
12+
def fourth_order_reg_matrix_via_mask_from(
13+
mask, pixel_scale: float = 1.0
14+
) -> np.ndarray:
15+
"""
16+
The fourth-order regularization matrix of the unmasked pixels of a
17+
rectangular masked 2D grid, H = H4x^T H4x + H4y^T H4y, where H4x / H4y
18+
are forward fourth-difference operators which degrade gracefully to
19+
third / second / first / zeroth order at the mask edge (see
20+
``derivative_util.forward_difference_operators_from``).
21+
22+
This is the fourth-order regularization scheme used by the
23+
gravitational-imaging (potential correction) technique of Cao et al. 2025
24+
(https://github.com/caoxiaoyue/lensing_potential_correction; cite via
25+
https://github.com/caoxiaoyue/potential_correction_paper), applied to the
26+
pixelized corrections of the lensing potential defined on a coarse
27+
rectangular mesh.
28+
29+
Parameters
30+
----------
31+
mask
32+
The 2D bool mask (``True`` = masked) of the rectangular grid whose
33+
unmasked pixels are regularized.
34+
pixel_scale
35+
The finite-difference step; regularization matrices are
36+
conventionally built with 1.0, the coefficient absorbing the scale.
37+
38+
Returns
39+
-------
40+
The [n_unmasked, n_unmasked] regularization matrix.
41+
"""
42+
return derivative_util.forward_difference_reg_matrix_from(
43+
mask=mask, pixel_scale=pixel_scale, max_order=4
44+
).toarray()
45+
46+
47+
class FourthOrderMask(AbstractRegularization):
48+
def __init__(self, coefficient: float = 1.0):
49+
"""
50+
Fourth-order regularization on the unmasked pixels of a rectangular
51+
masked 2D grid.
52+
53+
Each unmasked pixel is regularized with a forward fourth-difference
54+
stencil along both grid directions, degrading to third / second /
55+
first / zeroth order where the mask edge truncates the stencil.
56+
Penalising the fourth derivative permits solutions with curvature
57+
(e.g. localised perturbations) while still suppressing high-frequency
58+
noise, making it a weaker prior than ``CurvatureMask``. The linear
59+
object regularized must be defined on a masked rectangular grid and
60+
expose its ``mask``.
61+
62+
This is the fourth-order regularization scheme of the
63+
gravitational-imaging (potential correction) technique, applied to
64+
pixelized corrections of the lensing potential; it is ported from the
65+
``potential_correction`` package of Cao et al. 2025
66+
(https://github.com/caoxiaoyue/lensing_potential_correction). If you
67+
use it in your research, please cite Cao et al. 2025; citation
68+
materials are provided at
69+
https://github.com/caoxiaoyue/potential_correction_paper.
70+
71+
Parameters
72+
----------
73+
coefficient
74+
The regularization coefficient which multiplies the matrix,
75+
setting the strength of the smoothing.
76+
"""
77+
self.coefficient = coefficient
78+
79+
super().__init__()
80+
81+
def regularization_weights_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray:
82+
"""
83+
Returns the regularization weights of this regularization scheme,
84+
which are equal for every parameter.
85+
86+
Parameters
87+
----------
88+
linear_obj
89+
The linear object which uses these weights when performing
90+
regularization.
91+
"""
92+
return self.coefficient * xp.ones(linear_obj.params)
93+
94+
def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray:
95+
"""
96+
Returns the regularization matrix with shape [pixels, pixels].
97+
98+
Parameters
99+
----------
100+
linear_obj
101+
The linear object which uses this matrix to perform
102+
regularization. It must expose the ``mask`` of the rectangular
103+
masked 2D grid its parameters are defined on.
104+
105+
Returns
106+
-------
107+
The regularization matrix.
108+
"""
109+
return self.coefficient * fourth_order_reg_matrix_via_mask_from(
110+
mask=linear_obj.mask
111+
)

0 commit comments

Comments
 (0)