From 5c5b6c39be185f1ca916ba70adc02abfe39b4319 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 17 Jul 2026 11:56:31 +0100 Subject: [PATCH] Add potential_correction subpackage (gravitational imaging, phase 3) Adds autolens/potential_correction/ (exported as al.pc): the dpsi mesh paired to the data grid, FitDpsiImaging (dpsi-only inversion of an image residual), FitDpsiSrcImaging (joint source+dpsi inversion with full covariance), source factories, af.Analysis classes with evidence likelihoods, anchor-point rescaling and fit visualizations. Builds on the PyAutoArray operators (phase 1, #390) and PyAutoGalaxy input mass profiles (phase 2, #505). NumPy/numba only; JAX fast paths and the iterative LM engine follow as phase 4. Ported from the potential_correction package of Cao et al. 2025 (https://github.com/caoxiaoyue/lensing_potential_correction); cite via https://github.com/caoxiaoyue/potential_correction_paper. Phase 3 of PyAutoLabs/PyAutoLens#618. Co-Authored-By: Claude Fable 5 --- autolens/__init__.py | 1 + autolens/potential_correction/__init__.py | 28 + autolens/potential_correction/analysis.py | 165 +++++ autolens/potential_correction/fit.py | 659 ++++++++++++++++++ autolens/potential_correction/mesh.py | 193 +++++ autolens/potential_correction/pixelization.py | 110 +++ autolens/potential_correction/src_factory.py | 113 +++ autolens/potential_correction/util.py | 335 +++++++++ autolens/potential_correction/visualize.py | 367 ++++++++++ docs/api/potential_correction.rst | 33 + docs/index.md | 1 + .../potential_correction/__init__.py | 0 .../potential_correction/test_fit.py | 181 +++++ .../potential_correction/test_mesh.py | 69 ++ .../potential_correction/test_util.py | 179 +++++ 15 files changed, 2434 insertions(+) create mode 100644 autolens/potential_correction/__init__.py create mode 100644 autolens/potential_correction/analysis.py create mode 100644 autolens/potential_correction/fit.py create mode 100644 autolens/potential_correction/mesh.py create mode 100644 autolens/potential_correction/pixelization.py create mode 100644 autolens/potential_correction/src_factory.py create mode 100644 autolens/potential_correction/util.py create mode 100644 autolens/potential_correction/visualize.py create mode 100644 docs/api/potential_correction.rst create mode 100644 test_autolens/potential_correction/__init__.py create mode 100644 test_autolens/potential_correction/test_fit.py create mode 100644 test_autolens/potential_correction/test_mesh.py create mode 100644 test_autolens/potential_correction/test_util.py diff --git a/autolens/__init__.py b/autolens/__init__.py index 5a7e9584d..636bbb7e6 100644 --- a/autolens/__init__.py +++ b/autolens/__init__.py @@ -133,6 +133,7 @@ from . import exc from . import mock as m from . import util +from . import potential_correction as pc from autoconf import conf from autoconf.fitsable import ndarray_via_hdu_from diff --git a/autolens/potential_correction/__init__.py b/autolens/potential_correction/__init__.py new file mode 100644 index 000000000..364300406 --- /dev/null +++ b/autolens/potential_correction/__init__.py @@ -0,0 +1,28 @@ +""" +The gravitational-imaging (potential correction) technique: pixelized +corrections dpsi to the lensing potential, reconstructed on a coarse +rectangular mesh — alone from an image residual (``FitDpsiImaging``) or +jointly with the pixelized source (``FitDpsiSrcImaging``) — with the +regularization strengths of both set by maximising the Bayesian evidence. + +Ported from the ``potential_correction`` package of Cao et al. 2025 +(https://github.com/caoxiaoyue/lensing_potential_correction). If you use this +functionality in your research, please cite Cao et al. 2025; citation +materials are provided at +https://github.com/caoxiaoyue/potential_correction_paper. +""" + +from autolens.potential_correction import util +from autolens.potential_correction import visualize +from autolens.potential_correction.mesh import RegularDpsiMesh +from autolens.potential_correction.mesh import PairRegularDpsiMesh +from autolens.potential_correction.pixelization import DpsiLinearObj +from autolens.potential_correction.pixelization import DpsiPixelization +from autolens.potential_correction.pixelization import DpsiSrcPixelization +from autolens.potential_correction.src_factory import SrcFactory +from autolens.potential_correction.src_factory import AnalyticSrcFactory +from autolens.potential_correction.src_factory import PixSrcFactoryITP +from autolens.potential_correction.fit import FitDpsiImaging +from autolens.potential_correction.fit import FitDpsiSrcImaging +from autolens.potential_correction.analysis import DpsiInvAnalysis +from autolens.potential_correction.analysis import DpsiSrcInvAnalysis diff --git a/autolens/potential_correction/analysis.py b/autolens/potential_correction/analysis.py new file mode 100644 index 000000000..00ee44df9 --- /dev/null +++ b/autolens/potential_correction/analysis.py @@ -0,0 +1,165 @@ +""" +The ``af.Analysis`` classes of the gravitational-imaging (potential +correction) technique, through which a non-linear search samples the dpsi +mesh factor and regularization hyper-parameters by maximising the Bayesian +evidence of the (joint) inversion. + +Ported from the ``potential_correction`` package of Cao et al. 2025 +(https://github.com/caoxiaoyue/lensing_potential_correction). If you use this +functionality in your research, please cite Cao et al. 2025; citation +materials are provided at +https://github.com/caoxiaoyue/potential_correction_paper. +""" + +import logging +import os +from typing import Optional + +import numpy as np + +import autofit as af +import autoarray as aa +from autoarray import exc + +from autogalaxy.galaxy.galaxy import Galaxy + +from autolens.potential_correction.fit import FitDpsiImaging, FitDpsiSrcImaging +from autolens.potential_correction.pixelization import DpsiSrcPixelization +from autolens.potential_correction.src_factory import SrcFactory + +logger = logging.getLogger(__name__) + + +class DpsiInvAnalysis(af.Analysis): + def __init__( + self, + masked_imaging, + image_residual: np.ndarray, + source_gradient: np.ndarray, + anchor_points: Optional[np.ndarray] = None, + preloads: Optional[dict] = None, + ): + """ + Samples the dpsi pixelization (mesh factor + regularization + hyper-parameters) of a dpsi-only inversion of an image residual, + with the inversion's Bayesian evidence as the likelihood. + + Parameters + ---------- + masked_imaging + The masked ``al.Imaging`` dataset. + image_residual + The 1D (slim) image residual of the smooth-model fit. + source_gradient + The [n_unmasked, 2] source gradients at the ray-traced image + pixels. + anchor_points + The [3, 2] (y, x) anchor positions of the dpsi rescaling scheme. + preloads + Precomputed fit attributes shared across evaluations. + """ + self.masked_imaging = masked_imaging + self.image_residual = image_residual + self.source_gradient = source_gradient + self.anchor_points = anchor_points + self.preloads = preloads + + def log_likelihood_function(self, instance): + fit = FitDpsiImaging( + masked_imaging=self.masked_imaging, + image_residual=self.image_residual, + source_gradient=self.source_gradient, + anchor_points=self.anchor_points, + dpsi_pixelization=instance, + preloads=self.preloads, + ) + return fit.log_evidence + + +class DpsiSrcInvAnalysis(af.Analysis): + def __init__( + self, + masked_imaging, + lens_start: Galaxy, + source_start: SrcFactory, + anchor_points: Optional[np.ndarray] = None, + adapt_image=None, + src_image_mesh=None, + settings_inversion: Optional[aa.Settings] = None, + preloads: Optional[dict] = None, + ): + """ + Samples the joint source+dpsi pixelization (source pixelization and + dpsi mesh/regularization hyper-parameters) of a joint inversion, + with the inversion's Bayesian evidence as the likelihood. + + Parameters + ---------- + masked_imaging + The masked ``al.Imaging`` dataset. + lens_start + The lens galaxy of the smooth-model fit the corrections perturb. + source_start + The source factory evaluated for the source gradients. + anchor_points + The [3, 2] (y, x) anchor positions of the dpsi rescaling scheme. + adapt_image + The adapt image of the source pixelization's image mesh. + src_image_mesh + An image mesh whose image-plane mesh grid is preloaded into the + source inversion. + settings_inversion + The inversion settings; defaults to the positive-only solver + with the border relocator. + preloads + Precomputed fit attributes shared across evaluations. + """ + self.masked_imaging = masked_imaging + self.anchor_points = anchor_points + self.lens_start = lens_start + self.source_start = source_start + self.adapt_image = adapt_image + self.src_image_mesh = src_image_mesh + if settings_inversion is None: + self.settings_inversion = aa.Settings( + use_positive_only_solver=True, + use_border_relocator=True, + ) + else: + self.settings_inversion = settings_inversion + self.preloads = preloads + + def _fit_from(self, instance: DpsiSrcPixelization) -> FitDpsiSrcImaging: + return FitDpsiSrcImaging( + masked_imaging=self.masked_imaging, + anchor_points=self.anchor_points, + lens_start=self.lens_start, + source_start=self.source_start, + dpsi_pixelization=instance.dpsi_pixelization, + src_pixelization=instance.src_pixelization, + adapt_image=self.adapt_image, + src_image_mesh=self.src_image_mesh, + settings_inversion=self.settings_inversion, + preloads=self.preloads, + ) + + def log_likelihood_function(self, instance: DpsiSrcPixelization): + fit = self._fit_from(instance) + try: + return fit.log_evidence + except exc.InversionException: + # a failed inversion is a valid (very bad) sample, not a crash + logger.exception( + "InversionException during joint source+dpsi evidence evaluation; " + "returning penalty likelihood." + ) + return -1e8 + + def visualize(self, paths: af.DirectoryPaths, instance, during_analysis=True): + from autolens.potential_correction.visualize import show_fit_dpsi_src + + fit = self._fit_from(instance) + fit.log_evidence + + os.makedirs(paths.image_path, exist_ok=True) + show_fit_dpsi_src(fit=fit, output=f"{paths.image_path}/fit_dpsi_src.png") diff --git a/autolens/potential_correction/fit.py b/autolens/potential_correction/fit.py new file mode 100644 index 000000000..fb6ee2f13 --- /dev/null +++ b/autolens/potential_correction/fit.py @@ -0,0 +1,659 @@ +""" +The linear fits of the gravitational-imaging (potential correction) +technique: ``FitDpsiImaging`` inverts an image residual for pixelized +corrections dpsi to the lensing potential; ``FitDpsiSrcImaging`` jointly +inverts the image for the pixelized source and dpsi, fully accounting for +their covariance, with the Bayesian evidence used to set the regularization +strengths of both. + +Ported from the ``potential_correction`` package of Cao et al. 2025 +(https://github.com/caoxiaoyue/lensing_potential_correction). If you use this +functionality in your research, please cite Cao et al. 2025; citation +materials are provided at +https://github.com/caoxiaoyue/potential_correction_paper. +""" + +from typing import Optional + +import numpy as np +from scipy.sparse import block_diag +from scipy.spatial import Delaunay + +import autoarray as aa + +from autogalaxy.profiles.mass.input.interp import LinearNDInterpolatorExt +from autogalaxy.galaxy.galaxy import Galaxy + +from autolens.lens.tracer import Tracer +from autolens.imaging.fit_imaging import FitImaging +from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages +from autolens.potential_correction import util as pc_util +from autolens.potential_correction.pixelization import ( + DpsiLinearObj, + DpsiPixelization, +) +from autolens.potential_correction.src_factory import SrcFactory + + +class FitDpsiImaging: + def __init__( + self, + masked_imaging, + image_residual: np.ndarray, + source_gradient: np.ndarray, + dpsi_pixelization: DpsiPixelization, + anchor_points: Optional[np.ndarray] = None, + preloads: Optional[dict] = None, + ): + """ + A linear inversion of an image residual for pixelized corrections + dpsi to the lensing potential, at fixed source: the residual is + modelled as -B D_s D_psi dpsi (PSF blur matrix B, source-gradient + matrix D_s, dpsi-gradient operator D_psi), regularized on the dpsi + mesh and solved for the maximum-evidence correction. + + Parameters + ---------- + masked_imaging + The masked ``al.Imaging`` dataset. + image_residual + The 1D (slim) image residual of the smooth-model fit, of shape + [n_unmasked_data_pixels]. + source_gradient + The [n_unmasked_data_pixels, 2] (dS/dy, dS/dx) source gradients + at the ray-traced positions of the image pixels. + dpsi_pixelization + The dpsi mesh + regularization model. + anchor_points + The [3, 2] (y, x) anchor positions of the dpsi rescaling scheme. + preloads + Precomputed attributes set directly onto the fit (e.g. the + ``psf_mat`` shared across evaluations). + """ + self.masked_imaging = masked_imaging + self.input_image_residual = image_residual + self.source_gradient = source_gradient + self.anchor_points = anchor_points + self.dpsi_pixelization = dpsi_pixelization + + if preloads is not None: + for key, value in preloads.items(): + setattr(self, key, value) + + self.masked_imaging = self.masked_imaging.apply_over_sampling( + over_sample_size_lp=4, + over_sample_size_pixelization=4, + ) + + @property + def inverse_noise_covariance_matrix(self): + noise_1d = self.masked_imaging.noise_map.slim + return pc_util.inverse_covariance_matrix_from(noise_1d) + + @property + def psf_matrix(self): + return pc_util.psf_matrix_from( + np.asarray(self.masked_imaging.psf.kernel.native), + np.asarray(self.masked_imaging.mask), + ) + + @property + def source_gradient_matrix(self): + return pc_util.source_gradient_matrix_from(self.source_gradient) + + @property + def pair_dpsi_data_obj(self): + if not hasattr(self, "_pair_dpsi_data_obj"): + self._pair_dpsi_data_obj = self.dpsi_pixelization.pair_dpsi_data_mesh( + self.masked_imaging.mask, + self.masked_imaging.pixel_scales[0], + ) + return self._pair_dpsi_data_obj + + @property + def dpsi_gradient_matrix(self): + if not hasattr(self, "itp_mat"): + self.itp_mat = self.pair_dpsi_data_obj.itp_mat + return pc_util.dpsi_gradient_matrix_from( + self.itp_mat, + self.pair_dpsi_data_obj.Hx_dpsi, + self.pair_dpsi_data_obj.Hy_dpsi, + ) + + @property + def dpsi_points(self): + if not hasattr(self, "_dpsi_points"): + self._dpsi_points = np.vstack( + [ + self.pair_dpsi_data_obj.ygrid_dpsi_1d, + self.pair_dpsi_data_obj.xgrid_dpsi_1d, + ] + ).T + return self._dpsi_points + + @property + def dpsi_linear_obj(self): + return DpsiLinearObj( + mask=self.pair_dpsi_data_obj.mask_dpsi, points=self.dpsi_points + ) + + @property + def dpsi_regularization_matrix(self): + if not hasattr(self, "dpsi_reg_mat"): + self.dpsi_reg_mat = ( + self.dpsi_pixelization.regularization.regularization_matrix_from( + linear_obj=self.dpsi_linear_obj + ) + ) + return self.dpsi_reg_mat + + @property + def mapping_matrix(self): + # np.asarray guards against dense @ sparse products returning np.matrix + return np.asarray(-1.0 * self.psf_mat @ self.src_grad_mat @ self.dpsi_grad_mat) + + @property + def data_vector(self): + return np.asarray( + self.map_mat.T @ self.inv_cov_mat @ self.input_image_residual + ).ravel() + + @property + def curvature_regularization_matrix(self): + return np.asarray( + self.map_mat.T @ self.inv_cov_mat @ self.map_mat + self.reg_mat + ) + + def construct_useful_matrices(self): + if not hasattr(self, "psf_mat"): + self.psf_mat = self.psf_matrix + if not hasattr(self, "inv_cov_mat"): + self.inv_cov_mat = self.inverse_noise_covariance_matrix + if not hasattr(self, "src_grad_mat"): + self.src_grad_mat = self.source_gradient_matrix + if not hasattr(self, "dpsi_grad_mat"): + self.dpsi_grad_mat = self.dpsi_gradient_matrix + if not hasattr(self, "reg_mat"): + self.reg_mat = self.dpsi_regularization_matrix + if not hasattr(self, "map_mat"): + self.map_mat = self.mapping_matrix + if not hasattr(self, "d_vec"): + self.d_vec = self.data_vector + self.curve_reg_mat = self.curvature_regularization_matrix + + def solve_dpsi(self, return_error: bool = False): + self.construct_useful_matrices() + if return_error: + return ( + np.linalg.solve(self.curve_reg_mat, self.d_vec), + np.linalg.inv(self.curve_reg_mat), + ) + return np.linalg.solve(self.curve_reg_mat, self.d_vec) + + @property + def log_evidence(self): + """ + The Bayesian evidence of the dpsi inversion: noise normalization, + the log-determinants of the curvature+regularization and + regularization matrices, the regularization penalty of the solution + and the chi-squared of the residual fit. + """ + if not hasattr(self, "dpsi_slim") or not hasattr( + self, "model_image_residual_slim" + ): + self.dpsi_slim = self.solve_dpsi() + self.model_image_residual_slim = self.map_mat @ self.dpsi_slim + + noise_slim = np.asarray(self.masked_imaging.noise_map.slim) + + self.noise_term = float(np.sum(np.log(2 * np.pi * noise_slim**2.0))) * (-0.5) + + sign, logval = np.linalg.slogdet(self.curve_reg_mat) + if sign != 1: + raise np.linalg.LinAlgError( + "The curvature+regularization matrix is not positive definite." + ) + self.log_det_curve_reg_term = logval * (-0.5) + + try: + sign, logval = np.linalg.slogdet(self.reg_mat) + if sign != 1: + raise np.linalg.LinAlgError( + "The regularization matrix is not positive definite." + ) + self.log_det_reg_term = logval * 0.5 + except (np.linalg.LinAlgError, TypeError, ValueError): + self.log_det_reg_term = pc_util.log_det_mat(self.reg_mat, sparse=True) * 0.5 + + reg_cov_term = self.dpsi_slim.T @ self.reg_mat @ self.dpsi_slim + self.reg_cov_term = float(reg_cov_term) * (-0.5) + + residual_of_image_residual = ( + self.input_image_residual - self.model_image_residual_slim + ) + norm_residual = residual_of_image_residual / noise_slim + self.chi2_term = float(np.sum(norm_residual**2)) * (-0.5) + + return ( + self.noise_term + + self.log_det_curve_reg_term + + self.log_det_reg_term + + self.reg_cov_term + + self.chi2_term + ) + + @property + def rescaled_dpsi(self): + """ + The dpsi solution rescaled by the plane zeroing it at the three + anchor points (Suyu et al.'s scheme), plus the (a_y, a_x, c) + coefficients. Falls back to the unrescaled solution when no valid + anchor points are set. + """ + if self.anchor_points is None or np.shape(self.anchor_points) != (3, 2): + return self.dpsi_slim, 0.0, 0.0, 0.0 + if not hasattr(self, "dpsi_at_anchors"): + tri = Delaunay(np.fliplr(self.dpsi_points)) + self.dpsi_interpl = LinearNDInterpolatorExt(tri, self.dpsi_slim) + self.dpsi_at_anchors = self.dpsi_interpl( + self.anchor_points[:, 1], self.anchor_points[:, 0] + ) + ay, ax, c = pc_util.dpsi_rescale_factors_from( + self.anchor_points, self.dpsi_at_anchors + ) + dpsi_new = ( + ay * self.anchor_points[:, 0] + + ax * self.anchor_points[:, 1] + + c + + self.dpsi_slim + ) + return dpsi_new, ay, ax, c + + +class FitDpsiSrcImaging: + def __init__( + self, + masked_imaging, + lens_start: Galaxy, + source_start: SrcFactory, + dpsi_pixelization: DpsiPixelization, + src_pixelization: aa.Pixelization, + anchor_points: Optional[np.ndarray] = None, + adapt_image=None, + src_image_mesh=None, + settings_inversion: Optional[aa.Settings] = None, + preloads: Optional[dict] = None, + ): + """ + A joint linear inversion of an image for the pixelized source and + the pixelized corrections dpsi to the lensing potential, fully + accounting for their covariance: the mapping matrix is the + horizontal block [F_src | -B D_s D_psi] and the regularization the + block diagonal of the source and dpsi regularization matrices, with + the Bayesian evidence used to set both regularization strengths. + + Parameters + ---------- + masked_imaging + The masked ``al.Imaging`` dataset. + lens_start + The lens galaxy of the smooth-model fit the corrections perturb. + source_start + The source factory evaluated for the source gradients (from the + smooth-model fit's source). + dpsi_pixelization + The dpsi mesh + regularization model. + src_pixelization + The source pixelization of the joint inversion. + anchor_points + The [3, 2] (y, x) anchor positions of the dpsi rescaling scheme. + adapt_image + The adapt image of the source pixelization's image mesh. + src_image_mesh + An image mesh whose image-plane mesh grid is preloaded into the + source inversion. + settings_inversion + The inversion settings; defaults to the positive-only solver + with the border relocator. + preloads + Precomputed attributes set directly onto the fit. + """ + self.masked_imaging = masked_imaging + self.anchor_points = anchor_points + self.lens_start = lens_start + self.source_start = source_start + self.dpsi_pixelization = dpsi_pixelization + self.src_pixelization = src_pixelization + self.adapt_image = adapt_image + self.src_image_mesh = src_image_mesh + if settings_inversion is None: + self.settings_inversion = aa.Settings( + use_positive_only_solver=True, + use_border_relocator=True, + ) + else: + self.settings_inversion = settings_inversion + if preloads is not None: + for key, value in preloads.items(): + setattr(self, key, value) + + self.masked_imaging = self.masked_imaging.apply_over_sampling( + over_sample_size_lp=4, + over_sample_size_pixelization=4, + ) + + def do_source_inversion(self): + """ + Runs the standard autolens source inversion at the starting lens + model, caching its mapper, operated mapping matrix and + regularization matrix as the source blocks of the joint inversion. + """ + source_galaxy = Galaxy(redshift=1.0, pixelization=self.src_pixelization) + + adapt_kwargs = {} + if self.adapt_image is not None: + adapt_kwargs["galaxy_name_image_dict"] = {"source": self.adapt_image} + if self.src_image_mesh is not None: + self.image_plane_mesh_grid = self.src_image_mesh.image_plane_mesh_grid_from( + mask=self.masked_imaging.mask + ) + adapt_kwargs["galaxy_image_plane_mesh_grid_dict"] = { + source_galaxy: self.image_plane_mesh_grid + } + + self.adapt_images = AdaptImages(**adapt_kwargs) if adapt_kwargs else None + + tracer = Tracer(galaxies=[self.lens_start, source_galaxy]) + + self.src_fit = FitImaging( + dataset=self.masked_imaging, + tracer=tracer, + adapt_images=self.adapt_images, + settings=self.settings_inversion, + ) + + self.src_mapper = self.src_fit.inversion.linear_obj_list[0] + self.src_map_mat = self.src_fit.inversion.operated_mapping_matrix + self.src_reg_mat = self.src_fit.inversion.regularization_matrix + + @property + def inverse_noise_covariance_matrix(self): + if not hasattr(self, "inv_cov_mat"): + noise_1d = self.masked_imaging.noise_map + self.inv_cov_mat = pc_util.inverse_covariance_matrix_from(noise_1d) + return self.inv_cov_mat + + @property + def psf_matrix(self): + if not hasattr(self, "psf_mat"): + self.psf_mat = pc_util.psf_matrix_from( + np.asarray(self.masked_imaging.psf.kernel.native), + np.asarray(self.masked_imaging.mask), + ) + return self.psf_mat + + @property + def source_plane_data_grid(self): + # single-plane ray tracing is assumed here + return self.masked_imaging.grid.slim - self.lens_start.deflections_yx_2d_from( + self.masked_imaging.grid.slim + ) + + @property + def source_plane_source_gradient(self): + return self.source_start.eval_grad( + self.source_plane_data_grid[:, 1], self.source_plane_data_grid[:, 0] + ) + + @property + def source_gradient_matrix(self): + if not hasattr(self, "src_grad_mat"): + self.src_grad_mat = pc_util.source_gradient_matrix_from( + self.source_plane_source_gradient + ) + return self.src_grad_mat + + @property + def pair_dpsi_data_obj(self): + if not hasattr(self, "_pair_dpsi_data_obj"): + self._pair_dpsi_data_obj = self.dpsi_pixelization.pair_dpsi_data_mesh( + self.masked_imaging.mask, + self.masked_imaging.pixel_scales[0], + ) + return self._pair_dpsi_data_obj + + @property + def dpsi_gradient_matrix(self): + if not hasattr(self, "itp_mat"): + self.itp_mat = self.pair_dpsi_data_obj.itp_mat + if not hasattr(self, "dpsi_grad_mat"): + self.dpsi_grad_mat = pc_util.dpsi_gradient_matrix_from( + self.itp_mat, + self.pair_dpsi_data_obj.Hx_dpsi, + self.pair_dpsi_data_obj.Hy_dpsi, + ) + return self.dpsi_grad_mat + + @property + def dpsi_points(self): + if not hasattr(self, "_dpsi_points"): + self._dpsi_points = np.vstack( + [ + self.pair_dpsi_data_obj.ygrid_dpsi_1d, + self.pair_dpsi_data_obj.xgrid_dpsi_1d, + ] + ).T + return self._dpsi_points + + @property + def dpsi_linear_obj(self): + return DpsiLinearObj( + mask=self.pair_dpsi_data_obj.mask_dpsi, points=self.dpsi_points + ) + + @property + def dpsi_regularization_matrix(self): + if not hasattr(self, "dpsi_reg_mat"): + self.dpsi_reg_mat = ( + self.dpsi_pixelization.regularization.regularization_matrix_from( + linear_obj=self.dpsi_linear_obj + ) + ) + return self.dpsi_reg_mat + + @property + def dpsi_mapping_matrix(self): + if not hasattr(self, "dpsi_map_mat"): + # np.asarray guards against dense @ sparse returning np.matrix + self.dpsi_map_mat = np.asarray( + -1.0 + * self.psf_matrix + @ self.source_gradient_matrix + @ self.dpsi_gradient_matrix + ) + return self.dpsi_map_mat + + @property + def src_regularization_matrix(self): + if not hasattr(self, "src_reg_mat"): + if hasattr(self, "src_reg_base_mat"): + coeff = self.src_pixelization.regularization.coefficient + self.src_reg_mat = coeff * self.src_reg_base_mat + else: + self.do_source_inversion() + return self.src_reg_mat + + @property + def src_mapping_matrix(self): + if not hasattr(self, "src_map_mat"): + self.do_source_inversion() + return self.src_map_mat + + @property + def mapping_matrix(self): + if not hasattr(self, "map_mat"): + self.map_mat = np.hstack( + (self.src_mapping_matrix, self.dpsi_mapping_matrix) + ) + return self.map_mat + + @property + def regularization_matrix(self): + if not hasattr(self, "reg_mat"): + self.reg_mat = block_diag( + [self.src_regularization_matrix, self.dpsi_regularization_matrix] + ) + return self.reg_mat + + @property + def data_vector(self): + if not hasattr(self, "d_vec"): + self.d_vec = np.asarray( + self.mapping_matrix.T + @ self.inverse_noise_covariance_matrix + @ np.asarray(self.masked_imaging.data) + ).ravel() + return self.d_vec + + @property + def curvature_matrix(self): + if not hasattr(self, "curv_mat"): + self.curv_mat = np.asarray( + self.mapping_matrix.T + @ self.inverse_noise_covariance_matrix + @ self.mapping_matrix + ) + return self.curv_mat + + @property + def curvature_regularization_matrix(self): + return self.curvature_matrix + self.regularization_matrix + + def solve_src_dpsi(self, return_error: bool = False): + curve_reg_mat = np.asarray(self.curvature_regularization_matrix) + data_vector = np.asarray(self.data_vector) + if return_error: + return ( + np.linalg.solve(curve_reg_mat, data_vector), + np.linalg.inv(curve_reg_mat), + ) + return np.linalg.solve(curve_reg_mat, data_vector) + + @property + def log_evidence(self): + """ + The Bayesian evidence of the joint source+dpsi inversion: noise + normalization, the log-determinants of the curvature+regularization + matrix and of the source and dpsi regularization matrices, the joint + regularization penalty of the solution and the chi-squared of the + image fit. + """ + if not hasattr(self, "src_dpsi_slim") or not hasattr( + self, "model_image_slim" + ): + self.src_dpsi_slim = self.solve_src_dpsi() + self.model_image_slim = self.mapping_matrix @ self.src_dpsi_slim + + noise_slim = np.asarray(self.masked_imaging.noise_map) + + self.noise_term = float(np.sum(np.log(2 * np.pi * noise_slim**2.0))) * (-0.5) + + curve_reg_mat = np.asarray(self.curvature_regularization_matrix) + sign, logval = np.linalg.slogdet(curve_reg_mat) + if sign != 1: + raise np.linalg.LinAlgError( + "The curvature+regularization matrix is not positive definite." + ) + self.log_det_curve_reg_term = logval * (-0.5) + + self.log_det_reg_term_src = ( + pc_util.log_det_mat(self.src_regularization_matrix, sparse=True) * 0.5 + ) + try: + sign, logval = np.linalg.slogdet( + np.asarray(self.dpsi_regularization_matrix) + ) + if sign != 1: + raise np.linalg.LinAlgError( + "The dpsi regularization matrix is not positive definite." + ) + self.log_det_reg_term_dpsi = logval * 0.5 + except (np.linalg.LinAlgError, TypeError, ValueError): + self.log_det_reg_term_dpsi = ( + pc_util.log_det_mat(self.dpsi_regularization_matrix, sparse=True) * 0.5 + ) + self.log_det_reg_term = self.log_det_reg_term_src + self.log_det_reg_term_dpsi + + reg_cov_term = ( + self.src_dpsi_slim.T @ self.regularization_matrix @ self.src_dpsi_slim + ) + self.reg_cov_term = float(reg_cov_term) * (-0.5) + + image_residual = np.asarray(self.masked_imaging.data) - self.model_image_slim + norm_residual = image_residual / noise_slim + self.chi2_term = float(np.sum(norm_residual**2)) * (-0.5) + + return ( + self.noise_term + + self.log_det_curve_reg_term + + self.log_det_reg_term + + self.reg_cov_term + + self.chi2_term + ) + + def draw_random_solutions(self, n_solutions: int = 300, return_dkappa: bool = True): + """ + Draws source/dpsi solutions consistent with the data at the level + permitted by the noise and regularizations; when ``return_dkappa`` + the dpsi block is transformed to convergence corrections via the + dpsi mesh's Laplacian operator. + """ + mean, cov_mat = self.solve_src_dpsi(return_error=True) + + L = np.diag(np.ones_like(mean)) + if return_dkappa: + n_src = self.src_regularization_matrix.shape[0] + L[n_src:, n_src:] = self.pair_dpsi_data_obj.hamiltonian_dpsi.toarray() + + cov_mat = L @ cov_mat @ L.T + mean = L @ mean + return np.random.multivariate_normal( + mean, cov_mat, size=n_solutions, check_valid="warn", tol=1e-5 + ) + + @property + def best_fit_source(self): + n_s = self.src_regularization_matrix.shape[0] + return self.src_dpsi_slim[0:n_s] + + @property + def best_fit_dpsi(self): + n_s = self.src_regularization_matrix.shape[0] + return self.src_dpsi_slim[n_s:] + + @property + def rescaled_dpsi(self): + """ + The dpsi solution rescaled by the plane zeroing it at the three + anchor points (Suyu et al.'s scheme), plus the (a_y, a_x, c) + coefficients. Falls back to the unrescaled solution when no valid + anchor points are set. + """ + if self.anchor_points is None or np.shape(self.anchor_points) != (3, 2): + return self.best_fit_dpsi, 0.0, 0.0, 0.0 + if not hasattr(self, "dpsi_at_anchors"): + tri = Delaunay(np.fliplr(self.dpsi_points)) + self.dpsi_interpl = LinearNDInterpolatorExt(tri, self.best_fit_dpsi) + self.dpsi_at_anchors = self.dpsi_interpl( + self.anchor_points[:, 1], self.anchor_points[:, 0] + ) + ay, ax, c = pc_util.dpsi_rescale_factors_from( + self.anchor_points, self.dpsi_at_anchors + ) + dpsi_new = ( + ay * self.anchor_points[:, 0] + + ax * self.anchor_points[:, 1] + + c + + self.best_fit_dpsi + ) + return dpsi_new, ay, ax, c diff --git a/autolens/potential_correction/mesh.py b/autolens/potential_correction/mesh.py new file mode 100644 index 000000000..5cd14fabd --- /dev/null +++ b/autolens/potential_correction/mesh.py @@ -0,0 +1,193 @@ +""" +The dpsi mesh of the gravitational-imaging (potential correction) technique: +a rectangular mesh a factor coarser than the data grid, on whose unmasked +pixels the pixelized corrections to the lensing potential are defined, paired +to the data grid by a sparse bilinear interpolation matrix. + +Ported from the ``potential_correction`` package of Cao et al. 2025 +(https://github.com/caoxiaoyue/lensing_potential_correction). If you use this +functionality in your research, please cite Cao et al. 2025; citation +materials are provided at +https://github.com/caoxiaoyue/potential_correction_paper. +""" + +import numpy as np + +import autoarray as aa +from autoarray.operators import coarse_interp_util +from autoarray.operators import derivative_util + + +class RegularDpsiMesh: + def __init__(self, factor: int = 1): + """ + The mesh component of a ``DpsiPixelization`` model: pixelized + potential corrections are defined on a rectangular mesh ``factor`` + times coarser than the data grid. + + Parameters + ---------- + factor + How many times coarser than the data grid the dpsi mesh is. + """ + self.factor = int(factor) + + def __eq__(self, other): + return self.__dict__ == other.__dict__ and self.__class__ is other.__class__ + + def __str__(self): + return "\n".join(["{}: {}".format(k, v) for k, v in self.__dict__.items()]) + + def __repr__(self): + return "{}\n{}".format(self.__class__.__name__, str(self)) + + +class PairRegularDpsiMesh: + def __init__(self, mask, pixel_scale: float, dpsi_factor: int = 2): + """ + Pairs the dpsi mesh to the data grid: cleans the data mask, bins it + onto the coarser dpsi mesh (cleaning again), builds the bilinear + dpsi-to-data interpolation matrix and the sparse first/second + derivative operators of both grids. + + Parameters + ---------- + mask + The 2D bool data mask (``True`` = masked), typically an + annular-like region tracing the lensed arcs. + pixel_scale + The pixel size of the data grid, in arcsec. + dpsi_factor + How many times coarser than the data grid the dpsi mesh is; both + mask dimensions must be divisible by it. + """ + self.mask_data, self.diff_types_data = derivative_util.cleaned_mask_from( + np.asarray(mask) + ) + self.dpix_data = pixel_scale + self.mask_data_aa = aa.Mask2D(mask=self.mask_data, pixel_scales=self.dpix_data) + grid_data = aa.Grid2D.from_mask(mask=self.mask_data_aa) + self.xgrid_data = np.asarray(grid_data.native[:, :, 1]) + self.ygrid_data = np.asarray(grid_data.native[:, :, 0]) + limit = (self.mask_data.shape[0] * self.dpix_data) * 0.5 + self.data_bound = [-limit, limit, -limit, limit] + + self.dpsi_factor = dpsi_factor + if ( + self.mask_data.shape[0] % self.dpsi_factor != 0 + or self.mask_data.shape[1] % self.dpsi_factor != 0 + ): + raise ValueError("both mask dimensions must be divisible by dpsi_factor") + self.shape_2d_dpsi = ( + int(self.mask_data.shape[0] / self.dpsi_factor), + int(self.mask_data.shape[1] / self.dpsi_factor), + ) + self.dpix_dpsi = float(2.0 * limit / self.shape_2d_dpsi[0]) + self.mask_dpsi = coarse_interp_util.binned_mask_from( + self.mask_data, self.dpsi_factor + ) + self.mask_dpsi, self.diff_types_dpsi = derivative_util.cleaned_mask_from( + self.mask_dpsi + ) + self.mask_dpsi_aa = aa.Mask2D(mask=self.mask_dpsi, pixel_scales=self.dpix_dpsi) + grid_dpsi = aa.Grid2D.from_mask(mask=self.mask_dpsi_aa) + self.xgrid_dpsi = np.asarray(grid_dpsi.native[:, :, 1]) + self.ygrid_dpsi = np.asarray(grid_dpsi.native[:, :, 0]) + + self.grid_1d_from_mask() + + self.get_itp_box_ctr() + self.get_dpsi2data_mapping() + + self.get_gradient_operator_data() + self.get_gradient_operator_dpsi() + self.get_hamiltonian_operator_data() + self.get_hamiltonian_operator_dpsi() + + def grid_1d_from_mask(self): + self.idx_1d_data = np.where((~self.mask_data).flatten())[0] + self.xgrid_data_1d = self.xgrid_data.flatten()[self.idx_1d_data] + self.ygrid_data_1d = self.ygrid_data.flatten()[self.idx_1d_data] + + self.idx_1d_dpsi = np.where((~self.mask_dpsi).flatten())[0] + self.xgrid_dpsi_1d = self.xgrid_dpsi.flatten()[self.idx_1d_dpsi] + self.ygrid_dpsi_1d = self.ygrid_dpsi.flatten()[self.idx_1d_dpsi] + + def get_itp_box_ctr(self): + ctr_itp_box_mask = aa.Mask2D( + mask=np.full( + (self.shape_2d_dpsi[0] - 1, self.shape_2d_dpsi[1] - 1), False + ), + pixel_scales=self.dpix_dpsi, + ) + ctr_itp_box = aa.Grid2D.from_mask(mask=ctr_itp_box_mask) + self.xc_itp_box = np.asarray(ctr_itp_box.native[:, :, 1]) + self.yc_itp_box = np.asarray(ctr_itp_box.native[:, :, 0]) + + self.mask_itp_box = coarse_interp_util.interp_box_mask_from(self.mask_dpsi) + self.xc_itp_box_1d = self.xc_itp_box[~self.mask_itp_box] + self.yc_itp_box_1d = self.yc_itp_box[~self.mask_itp_box] + + if len(self.xc_itp_box_1d) == 0: + raise ValueError( + "The dpsi grid is too sparse. " + "Try decreasing the dpsi_factor to smaller values." + ) + + def get_dpsi2data_mapping(self): + """ + The sparse matrix of shape + [n_unmasked_data_pixels, n_unmasked_dpsi_pixels] mapping a vector on + the coarser dpsi mesh to the finer data grid by bilinear + interpolation. + """ + self.itp_mat = coarse_interp_util.coarse_interp_matrix_from( + self.mask_itp_box, + self.xc_itp_box, + self.yc_itp_box, + self.xgrid_data_1d, + self.ygrid_data_1d, + self.xgrid_dpsi, + self.ygrid_dpsi, + self.mask_dpsi, + ) + + def get_gradient_operator_data(self): + self.Hy_data, self.Hx_data = derivative_util.derivative_1st_operators_from( + self.mask_data, pixel_scale=self.dpix_data + ) + + def get_gradient_operator_dpsi(self): + self.Hy_dpsi, self.Hx_dpsi = derivative_util.derivative_1st_operators_from( + self.mask_dpsi, pixel_scale=self.dpix_dpsi + ) + + def get_hamiltonian_operator_data(self): + self.Hyy_data, self.Hxx_data = derivative_util.derivative_2nd_operators_from( + self.mask_data, pixel_scale=self.dpix_data + ) + self.hamiltonian_data = self.Hxx_data + self.Hyy_data + + def get_hamiltonian_operator_dpsi(self): + self.Hyy_dpsi, self.Hxx_dpsi = derivative_util.derivative_2nd_operators_from( + self.mask_dpsi, pixel_scale=self.dpix_dpsi + ) + self.hamiltonian_dpsi = self.Hxx_dpsi + self.Hyy_dpsi + + def show_grid(self, output_file="grid.png"): + from matplotlib import pyplot as plt + + plt.figure(figsize=(20, 20)) + plt.plot( + self.xgrid_data.flatten(), self.ygrid_data.flatten(), "*", color="black", label="data" + ) + plt.plot( + self.xgrid_dpsi.flatten(), self.ygrid_dpsi.flatten(), "*", color="red", label="dpsi" + ) + plt.plot(self.xgrid_data_1d, self.ygrid_data_1d, "o", color="black", label="data-unmask") + plt.plot(self.xgrid_dpsi_1d, self.ygrid_dpsi_1d, "p", color="red", label="dpsi-unmask") + plt.plot(self.xc_itp_box.flatten(), self.yc_itp_box.flatten(), "+", color="blue", label="dpsi-box") + plt.plot(self.xc_itp_box_1d, self.yc_itp_box_1d, "+", color="red", label="dpsi-box-unmask") + plt.legend() + plt.savefig(output_file, bbox_inches="tight") + plt.close() diff --git a/autolens/potential_correction/pixelization.py b/autolens/potential_correction/pixelization.py new file mode 100644 index 000000000..582c22e8d --- /dev/null +++ b/autolens/potential_correction/pixelization.py @@ -0,0 +1,110 @@ +""" +Model components of the gravitational-imaging (potential correction) +technique: the dpsi pixelization (mesh + regularization), the joint +source+dpsi pixelization, and the linear-object adapter through which any +``AbstractRegularization`` builds the dpsi regularization matrix. + +Ported from the ``potential_correction`` package of Cao et al. 2025 +(https://github.com/caoxiaoyue/lensing_potential_correction). If you use this +functionality in your research, please cite Cao et al. 2025; citation +materials are provided at +https://github.com/caoxiaoyue/potential_correction_paper. +""" + +import numpy as np + +import autoarray as aa +from autoarray.inversion.regularization.abstract import AbstractRegularization + +from autolens.potential_correction import mesh as dpsi_mesh + + +class DpsiMeshGrid: + def __init__(self, points: np.ndarray): + """ + Minimal grid wrapper exposing the ``array`` attribute the + kernel regularizations read their pixel positions from. + """ + self.array = points + + @property + def shape(self): + return self.array.shape + + def __array__(self, dtype=None): + return np.asarray(self.array, dtype=dtype) + + +class DpsiLinearObj: + def __init__(self, mask: np.ndarray, points: np.ndarray): + """ + The linear-object adapter of the dpsi mesh, exposing the attributes + the regularization schemes read: ``mask`` (the rectangular dpsi-mesh + mask, used by ``aa.reg.CurvatureMask`` / ``aa.reg.FourthOrderMask``), + ``source_plane_mesh_grid`` (the unmasked dpsi pixel positions, used + by the kernel regularizations, e.g. ``aa.reg.MaternKernel``) and + ``params`` (the parameter count, used by the regularization + weights). + + Through this adapter every dpsi regularization is built via the + standard ``regularization_matrix_from(linear_obj=...)`` interface — + no scheme-specific dispatch is needed in the fit. + + Parameters + ---------- + mask + The 2D bool mask of the dpsi mesh (``True`` = masked). + points + The [n_unmasked_dpsi, 2] array of (y, x) positions of the + unmasked dpsi pixels. + """ + self.mask = mask + self.source_plane_mesh_grid = DpsiMeshGrid(points=np.asarray(points)) + + @property + def params(self) -> int: + return int(np.count_nonzero(~np.asarray(self.mask))) + + +class DpsiPixelization: + def __init__( + self, mesh: dpsi_mesh.RegularDpsiMesh, regularization: AbstractRegularization + ): + """ + The pixelization of the potential corrections: the dpsi mesh (its + coarsening factor) plus the regularization scheme applied to the + corrections (``aa.reg.CurvatureMask``, ``aa.reg.FourthOrderMask`` or + a kernel regularization such as ``aa.reg.MaternKernel``). + + Parameters + ---------- + mesh + The ``RegularDpsiMesh`` defining the dpsi-mesh coarsening factor. + regularization + The regularization scheme applied to the potential corrections. + """ + self.mesh = mesh + self.regularization = regularization + + def pair_dpsi_data_mesh(self, mask, pixel_scale: float): + return dpsi_mesh.PairRegularDpsiMesh(mask, pixel_scale, self.mesh.factor) + + +class DpsiSrcPixelization: + def __init__( + self, dpsi_pixelization: DpsiPixelization, src_pixelization: aa.Pixelization + ): + """ + The joint pixelization of a source+dpsi inversion: the dpsi + pixelization of the potential corrections and the standard autolens + pixelization of the source. + + Parameters + ---------- + dpsi_pixelization + The pixelization of the potential corrections. + src_pixelization + The pixelization of the source reconstruction. + """ + self.dpsi_pixelization = dpsi_pixelization + self.src_pixelization = src_pixelization diff --git a/autolens/potential_correction/src_factory.py b/autolens/potential_correction/src_factory.py new file mode 100644 index 000000000..a6a705020 --- /dev/null +++ b/autolens/potential_correction/src_factory.py @@ -0,0 +1,113 @@ +""" +Source factories of the gravitational-imaging (potential correction) +technique: callables evaluating the source brightness and its gradients at +arbitrary source-plane positions, from an analytic galaxy or a pixelized +reconstruction. + +Ported from the ``potential_correction`` package of Cao et al. 2025 +(https://github.com/caoxiaoyue/lensing_potential_correction). If you use this +functionality in your research, please cite Cao et al. 2025; citation +materials are provided at +https://github.com/caoxiaoyue/potential_correction_paper. +""" + +from abc import ABC, abstractmethod + +import numpy as np +from scipy.spatial import Delaunay + +import autoarray as aa + +from autogalaxy.profiles.mass.input.interp import LinearNDInterpolatorExt +from autogalaxy.galaxy.galaxy import Galaxy + +from autolens.potential_correction import util as pc_util + +MESH2D_VORONOI_CLS = getattr(aa, "Mesh2DVoronoi", None) + + +class SrcFactory(ABC): + @abstractmethod + def eval_func(self, xgrid: np.ndarray, ygrid: np.ndarray) -> np.ndarray: + """The source brightness at the input (x, y) positions.""" + + def eval_grad(self, xgrid: np.ndarray, ygrid: np.ndarray, cross_size=0.001): + """ + The (dS/dy, dS/dx) source gradients at the input (x, y) positions, + by central differences over a cross of half-length ``cross_size``. + """ + if xgrid.shape != ygrid.shape: + raise ValueError("xgrid and ygrid must have the same shape") + origin_shape = xgrid.shape + points = np.vstack((np.ravel(ygrid), np.ravel(xgrid))).T + grad_points = pc_util.gradient_points_from(points, cross_size=cross_size) + values_at_grad_points = self.eval_func(grad_points[:, 1], grad_points[:, 0]) + src_grad_values = pc_util.source_gradient_from( + values_at_grad_points, grad_points + ) + return src_grad_values.reshape(*origin_shape, 2) + + +class AnalyticSrcFactory(SrcFactory): + def __init__(self, source_galaxy: Galaxy): + """ + A source factory evaluating an analytic source galaxy's light + profiles directly. + """ + self.source_galaxy = source_galaxy + + def eval_func(self, xgrid: np.ndarray, ygrid: np.ndarray): + grid_flat = np.vstack([np.ravel(ygrid), np.ravel(xgrid)]).T + func_values = self.source_galaxy.image_2d_from( + grid=aa.Grid2DIrregular(values=grid_flat) + ) + return np.asarray(func_values).reshape(np.shape(xgrid)) + + +class PixSrcFactoryITP(SrcFactory): + def __init__(self, points: np.ndarray, values: np.ndarray): + """ + A source factory interpolating a pixelized source reconstruction + (Delaunay linear interpolation with nearest-neighbour fallback). + Gradient evaluation uses per-point cross sizes set by the Voronoi + cell areas of the reconstruction's mesh. + + Parameters + ---------- + points + The [n_points, 2] source-plane positions of the reconstruction, + in autolens (y, x) order. + values + The reconstructed source brightness at those positions. + """ + self.points = points + self.values = values + try: + if MESH2D_VORONOI_CLS is None: + raise AttributeError("aa.Mesh2DVoronoi is unavailable") + self.vor_mesh = MESH2D_VORONOI_CLS(values=self.points) + self._split_cross = self.vor_mesh.split_cross + except Exception: + self._split_cross = pc_util.split_cross_from(self.points) + + def eval_func(self, xgrid: np.ndarray, ygrid: np.ndarray): + if not hasattr(self, "tri"): + self.tri = Delaunay(np.fliplr(self.points)) + if not hasattr(self, "interp_func"): + self.interp_func = LinearNDInterpolatorExt(self.tri, self.values) + return self.interp_func(xgrid, ygrid) + + def eval_grad(self, xgrid: np.ndarray, ygrid: np.ndarray, cross_size=None): + if cross_size is None: + grad_points = self._split_cross + else: + grad_points = pc_util.gradient_points_from( + self.points, cross_size=cross_size + ) + values_at_grad_points = self.eval_func(grad_points[:, 1], grad_points[:, 0]) + if not hasattr(self, "interp_grad_func"): + src_grad_values = pc_util.source_gradient_from( + values_at_grad_points, grad_points + ) + self.interp_grad_func = LinearNDInterpolatorExt(self.tri, src_grad_values) + return self.interp_grad_func(xgrid, ygrid) diff --git a/autolens/potential_correction/util.py b/autolens/potential_correction/util.py new file mode 100644 index 000000000..0c554d581 --- /dev/null +++ b/autolens/potential_correction/util.py @@ -0,0 +1,335 @@ +""" +Linear-algebra utilities of the gravitational-imaging (potential correction) +technique: source-gradient and dpsi-gradient operator assembly, the explicit +PSF blur matrix, noise covariance, log-determinants, anchor-point rescaling +and arc-mask generation. + +Ported from the ``potential_correction`` package of Cao et al. 2025 +(https://github.com/caoxiaoyue/lensing_potential_correction). If you use this +functionality in your research, please cite Cao et al. 2025; citation +materials are provided at +https://github.com/caoxiaoyue/potential_correction_paper. +""" + +import numpy as np +from scipy import ndimage +from scipy.sparse import csr_matrix, csc_matrix, diags, vstack +from scipy.sparse.linalg import splu + +from autoarray import numba_util +from autoarray.operators.derivative_util import cleaned_mask_from + + +def gradient_points_from(points: np.ndarray, cross_size: float = 0.001) -> np.ndarray: + """ + The cross-shaped evaluation points used to estimate function gradients at + a set of (y, x) positions by central differences. + + Parameters + ---------- + points + The [n_points, 2] array of (y, x) positions. + cross_size + The half-length of each cross arm. + + Returns + ------- + The [n_points * 4, 2] array of evaluation points, in per-point order + (+y, -y, +x, -x). + """ + points = np.asarray(points) + cross = np.array( + [ + [cross_size, 0.0], + [-cross_size, 0.0], + [0.0, cross_size], + [0.0, -cross_size], + ] + ) + return (points[:, None, :] + cross[None, :, :]).reshape((-1, 2)) + + +def source_gradient_from( + values_on_gradient_points: np.ndarray, gradient_points: np.ndarray +) -> np.ndarray: + """ + The (dS/dy, dS/dx) central-difference gradients of a source function from + its values on the cross points of ``gradient_points_from``. + """ + values_on_gradient_points = np.asarray(values_on_gradient_points).reshape((-1, 4)) + gradient_points = np.asarray(gradient_points).reshape((-1, 4, 2)) + + step_y = gradient_points[:, 0, 0] - gradient_points[:, 1, 0] + step_x = gradient_points[:, 2, 1] - gradient_points[:, 3, 1] + y_diff = ( + values_on_gradient_points[:, 0] - values_on_gradient_points[:, 1] + ) / step_y + x_diff = ( + values_on_gradient_points[:, 2] - values_on_gradient_points[:, 3] + ) / step_x + + return np.stack((y_diff, x_diff), axis=1) + + +@numba_util.jit() +def source_gradient_matrix_triplets_from(source_gradient): + """ + The (rows, cols, values) sparse triplets of the source-gradient matrix + (see ``source_gradient_matrix_from``). + """ + n_unmasked_data_points = source_gradient.shape[0] + rows_idx = np.full(n_unmasked_data_points * 2, -1, dtype=np.int64) + cols_idx = np.full(n_unmasked_data_points * 2, -1, dtype=np.int64) + values = np.full(n_unmasked_data_points * 2, 0.0, dtype=np.float64) + + count = 0 + for i in range(n_unmasked_data_points): + rows_idx[count] = i + cols_idx[count] = i * 2 + values[count] = source_gradient[i, 1] # x-derivative + count += 1 + rows_idx[count] = i + cols_idx[count] = i * 2 + 1 + values[count] = source_gradient[i, 0] # y-derivative + count += 1 + + return rows_idx, cols_idx, values + + +def source_gradient_matrix_from(source_gradient: np.ndarray) -> csr_matrix: + """ + The sparse source-gradient matrix D_s of shape + [n_unmasked_data_points, 2 * n_unmasked_data_points], which multiplies the + interleaved (x, y) dpsi-gradient vector to produce the brightness + correction of each image pixel (eq. 9 of the potential-correction + formalism). + + Parameters + ---------- + source_gradient + The [n_unmasked_data_points, 2] array of (dS/dy, dS/dx) source + gradients at the ray-traced positions of the image pixels. + """ + source_gradient = np.asarray(source_gradient) + rows_idx, cols_idx, values = source_gradient_matrix_triplets_from(source_gradient) + return csr_matrix( + (values, (rows_idx, cols_idx)), + shape=(source_gradient.shape[0], 2 * source_gradient.shape[0]), + ) + + +def dpsi_gradient_matrix_from(itp_mat, Hx, Hy) -> csr_matrix: + """ + The sparse dpsi-gradient operator D_psi of shape + [2 * n_unmasked_data_points, n_unmasked_dpsi_points]: interpolates the + dpsi mesh onto the data grid and takes its (x, y) gradients, with the + per-pixel rows interleaved as (x_0, y_0, x_1, y_1, ...) (eq. 8 of the + potential-correction formalism). + + Parameters + ---------- + itp_mat + The sparse [n_unmasked_data_points, n_unmasked_dpsi_points] + coarse-to-fine interpolation matrix. + Hx + The sparse x first-derivative operator of the dpsi mesh. + Hy + The sparse y first-derivative operator of the dpsi mesh. + """ + n_unmasked_data_points = itp_mat.shape[0] + Hx_itp = itp_mat @ Hx + Hy_itp = itp_mat @ Hy + + indices = np.empty(2 * n_unmasked_data_points, dtype=int) + indices[0::2] = np.arange(n_unmasked_data_points) + indices[1::2] = np.arange(n_unmasked_data_points, 2 * n_unmasked_data_points) + + return vstack([Hx_itp, Hy_itp]).tocsr()[indices, :] + + +@numba_util.jit() +def psf_matrix_from(psf_kernel, mask): + """ + The explicit PSF blur matrix B of shape [n_unmasked, n_unmasked]: entry + (j, i) is the fraction of pixel i's flux blurred into pixel j. The kernel + must be odd-sized and is renormalized to unit sum if required. + + Parameters + ---------- + psf_kernel + The 2D PSF kernel (odd-sized). + mask + The 2D bool mask (``True`` = masked) defining the fitted pixels. + """ + psf_hw = int(psf_kernel.shape[0] / 2) + if psf_hw * 2 + 1 != psf_kernel.shape[0]: + raise ValueError("The psf kernel size is not an odd number") + + if not np.isclose(np.sum(psf_kernel), 1.0): + psf_kernel = psf_kernel / np.sum(psf_kernel) + + mask_ext = np.ones( + (mask.shape[0] + psf_hw * 2, mask.shape[1] + psf_hw * 2), dtype="bool" + ) + mask_ext[psf_hw:-psf_hw, psf_hw:-psf_hw] = mask + image_ext_shape = mask_ext.shape + + indice_0, indice_1 = np.nonzero(~mask_ext) + n_unmasked_pix = len(indice_0) + psf_mat = np.zeros((n_unmasked_pix, n_unmasked_pix), dtype=np.float64) + + for ii in range(n_unmasked_pix): + image_unit = np.zeros(image_ext_shape, dtype=np.float64) + image_unit[ + indice_0[ii] - psf_hw : indice_0[ii] + psf_hw + 1, + indice_1[ii] - psf_hw : indice_1[ii] + psf_hw + 1, + ] = psf_kernel[:, :] + for jj in range(n_unmasked_pix): + psf_mat[jj, ii] = image_unit[indice_0[jj], indice_1[jj]] + + return psf_mat + + +def inverse_covariance_matrix_from(noise_slim) -> csr_matrix: + """ + The sparse diagonal inverse noise-covariance matrix 1/sigma_i^2 of a 1D + (slim) noise map. + """ + noise_slim = np.asarray(noise_slim) + return diags(1.0 / noise_slim**2, format="csr") + + +def log_det_mat(square_matrix, sparse: bool = False) -> float: + """ + The log-determinant of a positive-definite matrix, via sparse LU + decomposition when ``sparse`` (falling back to dense ``slogdet``). + + Raises ``np.linalg.LinAlgError`` if the determinant is not positive. + """ + if sparse: + try: + lu = splu(csc_matrix(square_matrix)) + diagL = lu.L.diagonal().astype(np.complex128) + diagU = lu.U.diagonal().astype(np.complex128) + return float(np.real(np.log(diagL).sum() + np.log(diagU).sum())) + except RuntimeError: + pass + + if not isinstance(square_matrix, np.ndarray): + square_matrix = square_matrix.toarray() + sign, logdet = np.linalg.slogdet(square_matrix) + if sign <= 0: + raise np.linalg.LinAlgError( + "The matrix is not positive definite: determinant sign is non-positive" + ) + return float(logdet) + + +def dpsi_rescale_factors_from(anchor_points, dpsi_values): + """ + The rescaling plane (a_y, a_x, c) which zeroes the dpsi solution at three + anchor points — Suyu et al.'s rescaling scheme, removing the wandering + source position and the unconstrained constant of the lensing potential. + + Parameters + ---------- + anchor_points + The [3, 2] array of (y, x) anchor positions. + dpsi_values + The dpsi values interpolated at the anchor positions. + """ + anchor_points = np.asarray(anchor_points) + dpsi_values = np.asarray(dpsi_values) + A_matrix = np.hstack( + [anchor_points, np.ones((3, 1), dtype=anchor_points.dtype)] + ) + b_vector = -1.0 * dpsi_values + a_y, a_x, c = np.linalg.solve(A_matrix, b_vector) + return a_y, a_x, c + + +def split_cross_from(points: np.ndarray) -> np.ndarray: + """ + Cross-shaped gradient-evaluation points for an irregular grid, with + per-point arm lengths set by the square root of each point's Voronoi cell + area (clipped at the 90th percentile; unbounded cells use the clip value). + + Parameters + ---------- + points + The [n_points, 2] array of (y, x) positions. + + Returns + ------- + The [n_points * 4, 2] array of evaluation points, in per-point order + (+y, -y, +x, -x). + """ + from scipy.spatial import Voronoi + + vor = Voronoi(points) + n_pixels = len(points) + region_areas = np.zeros(n_pixels) + + for i in range(n_pixels): + region_vertices_indexes = vor.regions[vor.point_region[i]] + if -1 in region_vertices_indexes or len(region_vertices_indexes) == 0: + region_areas[i] = -1 + else: + polygon = vor.vertices[region_vertices_indexes] + x = polygon[:, 0] + y = polygon[:, 1] + region_areas[i] = 0.5 * np.abs( + np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)) + ) + + max_area = np.percentile(region_areas, 90.0) + region_areas[region_areas == -1] = max_area + region_areas[region_areas > max_area] = max_area + + half_lengths = 0.5 * np.sqrt(region_areas) + + splitted_array = np.zeros((n_pixels, 4, 2)) + splitted_array[:, 0, 0] = points[:, 0] + half_lengths + splitted_array[:, 0, 1] = points[:, 1] + splitted_array[:, 1, 0] = points[:, 0] - half_lengths + splitted_array[:, 1, 1] = points[:, 1] + splitted_array[:, 2, 0] = points[:, 0] + splitted_array[:, 2, 1] = points[:, 1] + half_lengths + splitted_array[:, 3, 0] = points[:, 0] + splitted_array[:, 3, 1] = points[:, 1] - half_lengths + + return splitted_array.reshape(n_pixels * 4, 2) + + +def arc_mask_from( + snr_map: np.ndarray, + threshold: float = 3.0, + ignore_size: int = 25, + ext_size: int = 5, +) -> np.ndarray: + """ + A cleaned mask tracing the lensed arcs of a signal-to-noise map: pixels + above ``threshold`` are kept, connected islands smaller than + ``ignore_size`` pixels are dropped, the result is dilated by an + ``ext_size`` square footprint and finally cleaned so every unmasked pixel + supports a finite-difference scheme. + + Implemented with ``scipy.ndimage`` (8-connectivity labelling, matching + the original's ``skimage`` behaviour) to avoid a scikit-image dependency. + """ + bool_map = snr_map > threshold + + labels, _ = ndimage.label(bool_map, structure=np.ones((3, 3), dtype=int)) + label_sizes = np.bincount(labels.ravel()) + small_islands = np.where(label_sizes < ignore_size)[0] + small_islands = small_islands[small_islands != 0] + + mask = np.copy(bool_map) + for island_label in small_islands: + mask[labels == island_label] = 0 + mask = ndimage.binary_dilation( + mask, structure=np.ones((ext_size, ext_size), dtype=bool) + ) + + mask, _ = cleaned_mask_from(~mask, max_iter=50) + return mask diff --git a/autolens/potential_correction/visualize.py b/autolens/potential_correction/visualize.py new file mode 100644 index 000000000..895a12e29 --- /dev/null +++ b/autolens/potential_correction/visualize.py @@ -0,0 +1,367 @@ +""" +Matplotlib figures of the gravitational-imaging (potential correction) +technique: masked-data panels, irregular source-reconstruction views and the +multi-panel summaries of the dpsi-only and joint source+dpsi fits. + +Ported from the ``potential_correction`` package of Cao et al. 2025 +(https://github.com/caoxiaoyue/lensing_potential_correction). If you use this +functionality in your research, please cite Cao et al. 2025; citation +materials are provided at +https://github.com/caoxiaoyue/potential_correction_paper. +""" + +import copy +import os + +import numpy as np +import matplotlib as mpl +import matplotlib.cm as cm +import matplotlib.pyplot as plt +import matplotlib.ticker as ticker +from matplotlib.ticker import MaxNLocator +from mpl_toolkits.axes_grid1 import make_axes_locatable +from scipy.interpolate import griddata +from scipy.spatial import Voronoi, voronoi_plot_2d + + +def _plot_anchor_points(ax, anchor_points): + if anchor_points is not None and np.shape(anchor_points) == (3, 2): + anchor_points = np.asarray(anchor_points) + ax.plot(anchor_points[:, 1], anchor_points[:, 0], "rx") + + +def imshow_masked_data( + data_1d, + mask_2d, + dpix=None, + ax=None, + n_contours=None, + show_scale=None, + n_cbar_ticks=None, + centralize=False, + **kwargs, +): + """ + Shows a 1D (slim) data vector on its 2D mask, with masked pixels blanked, + a colorbar and optional contours. + """ + data_1d = np.asarray(data_1d) + mask_2d = np.asarray(mask_2d) + if centralize: + data_1d = data_1d - np.median(data_1d) + data_2d = np.zeros_like(mask_2d, dtype="float") + data_2d[~mask_2d] = data_1d + data_2d_masked = np.ma.masked_array(data_2d, mask=mask_2d) + + if "extent" in kwargs.keys(): + extent = kwargs.pop("extent") + else: + hw = mask_2d.shape[0] * dpix * 0.5 + extent = [-hw, hw, -hw, hw] + im = ax.imshow(data_2d_masked, extent=extent, **kwargs) + divider = make_axes_locatable(ax) + cax = divider.append_axes("right", size="5%", pad=0.05) + + if show_scale is None: + cbar = plt.colorbar(im, cax=cax) + else: + fmt_cbar = ticker.ScalarFormatter(useMathText=True) + fmt_cbar.set_powerlimits((show_scale, show_scale)) + cbar = plt.colorbar(im, cax=cax, format=fmt_cbar) + cbar.ax.yaxis.set_offset_position("left") + cbar.update_ticks() + if n_cbar_ticks is not None: + cbar.locator = MaxNLocator(nbins=n_cbar_ticks) + cbar.update_ticks() + + if dpix is None: + if n_contours is not None: + raise ValueError("dpix is None, cannot show contours") + else: + coord_1d = np.arange(len(mask_2d)) * dpix + coord_1d = coord_1d - np.mean(coord_1d) + xgrid, ygrid = np.meshgrid(coord_1d, coord_1d) + rgrid = np.sqrt(xgrid**2 + ygrid**2) + limit = np.max(rgrid[~mask_2d]) + ax.set_xlim(-1.0 * limit, limit) + ax.set_ylim(-1.0 * limit, limit) + if (n_contours is not None) and isinstance(n_contours, int): + # invert the grids so contours match autolens imaging orientation + xgrid = np.flipud(xgrid) + ygrid = np.flipud(ygrid) + CS = ax.contour( + xgrid, ygrid, data_2d_masked, levels=n_contours, colors="k", + corner_mask=True, + ) + if show_scale is not None: + ax.clabel( + CS, inline=True, + fmt=lambda x: f"{x * 1.0 / 10 ** show_scale:.1f}", + ) + else: + ax.clabel(CS, inline=True) + + return ax + + +def show_image_irregular_interpolate( + points, values, ax=None, enlarge_factor=1.1, npixels=100, cmap="jet", **kwargs +): + """ + Shows values on irregular (y, x) points by linear interpolation onto a + regular grid. + """ + points = np.asarray(points) + points = points[:, ::-1] # to scipy (x, y) order + + half_width = max(np.abs(points.min()), np.abs(points.max())) + half_width *= enlarge_factor + + coordinate_1d, dpix = np.linspace( + -1.0 * half_width, half_width, npixels, endpoint=True, retstep=True + ) + xgrid, ygrid = np.meshgrid(coordinate_1d, coordinate_1d) + extent = [ + -1.0 * half_width - 0.5 * dpix, + half_width + 0.5 * dpix, + -1.0 * half_width - 0.5 * dpix, + half_width + 0.5 * dpix, + ] + + source_image = griddata(points, values, (xgrid, ygrid), method="linear", fill_value=0.0) + + im = ax.imshow(source_image, origin="lower", extent=extent, cmap=cmap, **kwargs) + plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + + +def show_image_irregular( + points, values, ax=None, enlarge_factor=1.1, title="Source", cmap="jet", + minima=None, maxima=None, +): + """ + Shows values on irregular (y, x) points as coloured Voronoi cells. + """ + points = np.asarray(points) + points = points[:, ::-1] # to scipy (x, y) order + + half_width = max(np.abs(points.min()), np.abs(points.max())) + half_width *= enlarge_factor + + # far-away sentinel points close the outer Voronoi cells + points = np.append(points, [[999, 999], [-999, 999], [999, -999], [-999, -999]], axis=0) + vor = Voronoi(points) + if minima is None: + minima = min(values) + if maxima is None: + maxima = max(values) + norm = mpl.colors.Normalize(vmin=minima, vmax=maxima, clip=True) + mapper = cm.ScalarMappable(norm=norm, cmap=cmap) + mapper.set_array([]) + voronoi_plot_2d( + vor, ax=ax, show_points=False, show_vertices=False, line_width=0.05, + point_size=1, line_colors="k", line_alpha=0.2, + ) + for r in range(len(vor.point_region) - 4): + region = vor.regions[vor.point_region[r]] + if -1 not in region: + polygon = [vor.vertices[i] for i in region] + ax.fill(*zip(*polygon), color=mapper.to_rgba(values[r])) + divider = make_axes_locatable(ax) + cax = divider.append_axes("right", size="5%", pad=0.05) + plt.colorbar(mapper, cax=cax) + ax.set_xlim(-half_width, half_width) + ax.set_ylim(-half_width, half_width) + ax.set_aspect("equal", adjustable="box") + ax.set_title(title) + + +def show_fit_dpsi(fit, output="result.png"): + """ + The six-panel summary of a ``FitDpsiImaging``: residual data, model, + residual-of-residual, normalized residual, and the dpsi and dkappa maps. + """ + plt.figure(figsize=(15, 10)) + cmap = copy.copy(plt.get_cmap("jet")) + cmap.set_bad(color="white") + myargs_data = {"origin": "upper", "cmap": cmap} + myargs_data["extent"] = copy.copy(fit.pair_dpsi_data_obj.data_bound) + myargs_dpsi = copy.deepcopy(myargs_data) + xlimit = [ + fit.pair_dpsi_data_obj.xgrid_data_1d.min(), + fit.pair_dpsi_data_obj.xgrid_data_1d.max(), + ] + ylimit = [ + fit.pair_dpsi_data_obj.ygrid_data_1d.min(), + fit.pair_dpsi_data_obj.ygrid_data_1d.max(), + ] + + plt.subplot(231) + ax = plt.gca() + imshow_masked_data(fit.input_image_residual, fit.masked_imaging.mask, ax=ax, **myargs_data) + ax.set_title("Data") + _plot_anchor_points(ax, fit.anchor_points) + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + plt.subplot(232) + ax = plt.gca() + imshow_masked_data(fit.model_image_residual_slim, fit.masked_imaging.mask, ax=ax, **myargs_data) + ax.set_title("Model") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + residual_of_image_residual = fit.input_image_residual - fit.model_image_residual_slim + plt.subplot(233) + ax = plt.gca() + imshow_masked_data(residual_of_image_residual, fit.masked_imaging.mask, ax=ax, **myargs_data) + ax.set_title("Residual") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + norm_residual = residual_of_image_residual / fit.masked_imaging.noise_map.slim + plt.subplot(234) + ax = plt.gca() + imshow_masked_data(norm_residual, fit.masked_imaging.mask, ax=ax, **myargs_data) + ax.set_title("Normalized Residual") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + plt.subplot(235) + ax = plt.gca() + imshow_masked_data(fit.dpsi_slim, fit.pair_dpsi_data_obj.mask_dpsi, ax=ax, **myargs_dpsi) + _plot_anchor_points(ax, fit.anchor_points) + ax.set_title("Dpsi Map") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + dkappa_slim = fit.pair_dpsi_data_obj.hamiltonian_dpsi @ fit.dpsi_slim + plt.subplot(236) + ax = plt.gca() + imshow_masked_data(dkappa_slim, fit.pair_dpsi_data_obj.mask_dpsi, ax=ax, **myargs_dpsi) + ax.set_title("Dkappa Map") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + this_path = os.path.dirname(output) + if this_path: + os.makedirs(this_path, exist_ok=True) + plt.tight_layout() + plt.savefig(output, bbox_inches="tight") + plt.close() + + +def show_fit_dpsi_src(fit, output="result.png", show_src_grid=True, interpolate=True): + """ + The nine-panel summary of a ``FitDpsiSrcImaging``: data, noise, SNR, + model, residual, normalized residual, dpsi map, dkappa map and the + source reconstruction. + """ + fig = plt.figure(figsize=(15, 10)) + cmap = copy.copy(plt.get_cmap("jet")) + cmap.set_bad(color="white") + myargs_data = {"origin": "upper", "cmap": cmap} + myargs_data["extent"] = copy.copy(fit.pair_dpsi_data_obj.data_bound) + xlimit = [ + fit.pair_dpsi_data_obj.xgrid_data_1d.min(), + fit.pair_dpsi_data_obj.xgrid_data_1d.max(), + ] + ylimit = [ + fit.pair_dpsi_data_obj.ygrid_data_1d.min(), + fit.pair_dpsi_data_obj.ygrid_data_1d.max(), + ] + myargs_dpsi = copy.deepcopy(myargs_data) + image_plane_mesh_grid = fit.src_mapper.image_plane_mesh_grid + source_plane_mesh_grid = fit.src_mapper.source_plane_mesh_grid + + plt.subplot(331) + ax = plt.gca() + imshow_masked_data(fit.masked_imaging.data, fit.masked_imaging.mask, ax=ax, **myargs_data) + ax.set_title("Data") + _plot_anchor_points(ax, fit.anchor_points) + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + plt.subplot(332) + ax = plt.gca() + imshow_masked_data(fit.masked_imaging.noise_map, fit.masked_imaging.mask, ax=ax, **myargs_data) + ax.set_title("Noise") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + plt.subplot(333) + ax = plt.gca() + imshow_masked_data( + fit.masked_imaging.data / fit.masked_imaging.noise_map, + fit.masked_imaging.mask, ax=ax, **myargs_data, + ) + ax.set_title("SNR") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + plt.subplot(334) + ax = plt.gca() + imshow_masked_data(fit.model_image_slim, fit.masked_imaging.mask, ax=ax, **myargs_data) + if show_src_grid: + ax.scatter(image_plane_mesh_grid[:, 1], image_plane_mesh_grid[:, 0], c="black", s=0.5, alpha=0.5) + ax.set_title("Model") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + residual = fit.masked_imaging.data - fit.model_image_slim + plt.subplot(335) + ax = plt.gca() + imshow_masked_data(residual, fit.masked_imaging.mask, ax=ax, **myargs_data) + ax.set_title("Residual") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + norm_residual = residual / fit.masked_imaging.noise_map + plt.subplot(336) + ax = plt.gca() + imshow_masked_data(norm_residual, fit.masked_imaging.mask, ax=ax, **myargs_data) + ax.set_title("Norm Residual") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + n_src_pixels = fit.src_regularization_matrix.shape[0] + dpsi_slim = fit.src_dpsi_slim[n_src_pixels:] + plt.subplot(337) + ax = plt.gca() + imshow_masked_data(dpsi_slim, fit.pair_dpsi_data_obj.mask_dpsi, ax=ax, **myargs_dpsi) + _plot_anchor_points(ax, fit.anchor_points) + ax.set_title("Dpsi Map") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + dkappa_slim = fit.pair_dpsi_data_obj.hamiltonian_dpsi @ dpsi_slim + plt.subplot(338) + ax = plt.gca() + imshow_masked_data(dkappa_slim, fit.pair_dpsi_data_obj.mask_dpsi, ax=ax, **myargs_dpsi) + ax.set_title("Dkappa Map") + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + + src_slim = fit.src_dpsi_slim[0:n_src_pixels] + plt.subplot(339) + ax = plt.gca() + if interpolate: + show_image_irregular_interpolate( + source_plane_mesh_grid, src_slim, ax=ax, enlarge_factor=1.1, npixels=100, cmap="jet" + ) + else: + show_image_irregular( + source_plane_mesh_grid, src_slim, enlarge_factor=1.1, cmap="jet", ax=ax, title="Source" + ) + if show_src_grid: + ax.scatter(source_plane_mesh_grid[:, 1], source_plane_mesh_grid[:, 0], c="black", s=0.1, alpha=0.5) + ax.set_title("Source") + + plt.tight_layout() + + if output == "show": + return fig + this_path = os.path.dirname(output) + if this_path: + os.makedirs(this_path, exist_ok=True) + fig.savefig(output, bbox_inches="tight") + plt.close(fig) diff --git a/docs/api/potential_correction.rst b/docs/api/potential_correction.rst new file mode 100644 index 000000000..966eb2bc3 --- /dev/null +++ b/docs/api/potential_correction.rst @@ -0,0 +1,33 @@ +==================== +Potential Correction +==================== + +The gravitational-imaging (potential correction) technique: pixelized +corrections to the lensing potential, reconstructed alone from an image +residual or jointly with the pixelized source by maximising the Bayesian +evidence. + +Ported from the ``potential_correction`` package of Cao et al. 2025 +(https://github.com/caoxiaoyue/lensing_potential_correction). If you use this +functionality in your research, please cite Cao et al. 2025; citation +materials are provided at +https://github.com/caoxiaoyue/potential_correction_paper. + +.. currentmodule:: autolens.pc + +.. autosummary:: + :toctree: _autosummary + :template: custom-class-template.rst + :recursive: + + RegularDpsiMesh + PairRegularDpsiMesh + DpsiPixelization + DpsiSrcPixelization + SrcFactory + AnalyticSrcFactory + PixSrcFactoryITP + FitDpsiImaging + FitDpsiSrcImaging + DpsiInvAnalysis + DpsiSrcInvAnalysis diff --git a/docs/index.md b/docs/index.md index 0199a7208..1346dd71c 100644 --- a/docs/index.md +++ b/docs/index.md @@ -266,5 +266,6 @@ api/modeling api/pixelization api/point api/plot +api/potential_correction api/source ``` diff --git a/test_autolens/potential_correction/__init__.py b/test_autolens/potential_correction/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test_autolens/potential_correction/test_fit.py b/test_autolens/potential_correction/test_fit.py new file mode 100644 index 000000000..8a1b02673 --- /dev/null +++ b/test_autolens/potential_correction/test_fit.py @@ -0,0 +1,181 @@ +import numpy as np +import pytest +from scipy.sparse import csr_matrix + +import autoarray as aa +import autolens as al +from autolens.potential_correction.fit import FitDpsiImaging, FitDpsiSrcImaging +from autolens.potential_correction.pixelization import DpsiLinearObj + + +class MockNoiseMap: + def __init__(self, values): + self.slim = values + + +class MockMaskedImaging: + def __init__(self, noise_values): + self.noise_map = MockNoiseMap(noise_values) + + +def injected_dpsi_fit(): + """ + A FitDpsiImaging with every cached matrix pre-injected (the pattern of + the original package's unit tests), so the evidence terms can be + verified against hand-computed linear algebra with no imaging data. + """ + fit = object.__new__(FitDpsiImaging) + + fit.anchor_points = None + fit.input_image_residual = np.array([0.1, -0.2, 0.3, -0.1]) + fit.map_mat = np.array( + [ + [0.3, 0.1, -0.2], + [-0.1, 0.4, 0.1], + [0.2, -0.3, 0.5], + [0.0, 0.2, -0.4], + ] + ) + w = np.array([1.0, 2.0, 3.0, 4.0]) + fit.inv_cov_mat = np.diag(w) + fit.reg_mat = csr_matrix(np.diag([0.5, 0.5, 0.5])) + fit.d_vec = fit.map_mat.T @ fit.inv_cov_mat @ fit.input_image_residual + fit.curve_reg_mat = ( + fit.map_mat.T @ fit.inv_cov_mat @ fit.map_mat + fit.reg_mat.toarray() + ) + fit.psf_mat = np.eye(4) + fit.src_grad_mat = np.eye(4) + fit.dpsi_grad_mat = np.ones((4, 3)) * 0.1 + + noise = 1.0 / np.sqrt(w) + fit.masked_imaging = MockMaskedImaging(noise_values=noise) + + return fit + + +def test__fit_dpsi_imaging__solve_matches_normal_equations(): + fit = injected_dpsi_fit() + + dpsi = fit.solve_dpsi() + + assert fit.curve_reg_mat @ dpsi == pytest.approx(fit.d_vec, abs=1.0e-12) + + +def test__fit_dpsi_imaging__log_evidence_matches_hand_computed_terms(): + fit = injected_dpsi_fit() + + evidence = fit.log_evidence + + noise = fit.masked_imaging.noise_map.slim + dpsi = np.linalg.solve(fit.curve_reg_mat, fit.d_vec) + model = fit.map_mat @ dpsi + residual = fit.input_image_residual - model + + noise_term = -0.5 * np.sum(np.log(2 * np.pi * noise**2)) + log_det_curve = -0.5 * np.linalg.slogdet(fit.curve_reg_mat)[1] + log_det_reg = 0.5 * np.linalg.slogdet(fit.reg_mat.toarray())[1] + reg_cov = -0.5 * float(dpsi.T @ fit.reg_mat @ dpsi) + chi2 = -0.5 * float(np.sum((residual / noise) ** 2)) + + assert evidence == pytest.approx( + noise_term + log_det_curve + log_det_reg + reg_cov + chi2, rel=1.0e-10 + ) + + +def test__fit_dpsi_imaging__rescaled_dpsi_without_anchors_is_identity(): + fit = injected_dpsi_fit() + fit.log_evidence + + dpsi_new, ay, ax, c = fit.rescaled_dpsi + + assert dpsi_new == pytest.approx(fit.dpsi_slim) + assert (ay, ax, c) == (0.0, 0.0, 0.0) + + +def test__dpsi_linear_obj__works_for_mask_and_kernel_regularizations(): + mask = np.ones((6, 6), dtype=bool) + mask[1:5, 1:5] = False + ys, xs = np.where(~mask) + points = np.vstack([ys.astype(float), xs.astype(float)]).T + + linear_obj = DpsiLinearObj(mask=mask, points=points) + + assert linear_obj.params == 16 + + reg_curvature = aa.reg.CurvatureMask(coefficient=2.0) + matrix_curvature = reg_curvature.regularization_matrix_from(linear_obj=linear_obj) + assert matrix_curvature.shape == (16, 16) + assert matrix_curvature == pytest.approx(matrix_curvature.T, abs=1.0e-12) + + reg_kernel = aa.reg.MaternKernel(coefficient=1.0, scale=1.0, nu=0.5) + matrix_kernel = reg_kernel.regularization_matrix_from(linear_obj=linear_obj) + assert matrix_kernel.shape == (16, 16) + assert np.isfinite(matrix_kernel).all() + + +def test__fit_dpsi_src_imaging__end_to_end_evidence_is_finite(masked_imaging_7x7): + lens = al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0), + ) + source_galaxy = al.Galaxy( + redshift=1.0, + bulge=al.lp.SersicSph( + centre=(0.0, 0.0), intensity=1.0, effective_radius=0.5 + ), + ) + source_start = al.pc.AnalyticSrcFactory(source_galaxy=source_galaxy) + + dpsi_pixelization = al.pc.DpsiPixelization( + mesh=al.pc.RegularDpsiMesh(factor=1), + regularization=aa.reg.CurvatureMask(coefficient=1.0), + ) + src_pixelization = al.Pixelization( + mesh=al.mesh.RectangularUniform(shape=(3, 3)), + regularization=al.reg.Constant(coefficient=1.0), + ) + + fit = al.pc.FitDpsiSrcImaging( + masked_imaging=masked_imaging_7x7, + lens_start=lens, + source_start=source_start, + dpsi_pixelization=dpsi_pixelization, + src_pixelization=src_pixelization, + ) + + evidence = fit.log_evidence + + assert np.isfinite(evidence) + + n_src = fit.src_regularization_matrix.shape[0] + n_dpsi = np.count_nonzero(~fit.pair_dpsi_data_obj.mask_dpsi) + assert fit.best_fit_source.shape == (n_src,) + assert fit.best_fit_dpsi.shape == (n_dpsi,) + assert fit.mapping_matrix.shape[1] == n_src + n_dpsi + + +def test__fit_dpsi_imaging__end_to_end_evidence_is_finite(masked_imaging_7x7): + n_data = int(np.count_nonzero(~np.asarray(masked_imaging_7x7.mask))) + + rng = np.random.default_rng(3) + image_residual = rng.normal(scale=0.1, size=n_data) + source_gradient = np.ones((n_data, 2)) + + dpsi_pixelization = al.pc.DpsiPixelization( + mesh=al.pc.RegularDpsiMesh(factor=1), + regularization=aa.reg.CurvatureMask(coefficient=1.0), + ) + + fit = al.pc.FitDpsiImaging( + masked_imaging=masked_imaging_7x7, + image_residual=image_residual, + source_gradient=source_gradient, + dpsi_pixelization=dpsi_pixelization, + ) + + evidence = fit.log_evidence + + assert np.isfinite(evidence) + assert fit.dpsi_slim.shape == ( + np.count_nonzero(~fit.pair_dpsi_data_obj.mask_dpsi), + ) diff --git a/test_autolens/potential_correction/test_mesh.py b/test_autolens/potential_correction/test_mesh.py new file mode 100644 index 000000000..8f1d34c41 --- /dev/null +++ b/test_autolens/potential_correction/test_mesh.py @@ -0,0 +1,69 @@ +import numpy as np +import pytest + +import autolens as al + + +def circular_mask(shape=(20, 20), pixel_scale=0.5, r_min=0.0, r_max=3.6): + cy = (shape[0] - 1) / 2.0 + cx = (shape[1] - 1) / 2.0 + mask = np.ones(shape, dtype=bool) + for i in range(shape[0]): + for j in range(shape[1]): + r = np.sqrt(((i - cy) * pixel_scale) ** 2 + ((j - cx) * pixel_scale) ** 2) + if r_min <= r <= r_max: + mask[i, j] = False + return mask + + +def test__pair_regular_dpsi_mesh__shapes_and_grids(): + mask = circular_mask() + pair = al.pc.PairRegularDpsiMesh(mask, pixel_scale=0.5, dpsi_factor=2) + + n_data = np.count_nonzero(~pair.mask_data) + n_dpsi = np.count_nonzero(~pair.mask_dpsi) + + assert pair.shape_2d_dpsi == (10, 10) + assert pair.dpix_dpsi == pytest.approx(1.0) + assert pair.xgrid_data_1d.shape == (n_data,) + assert pair.ygrid_dpsi_1d.shape == (n_dpsi,) + assert pair.itp_mat.shape == (n_data, n_dpsi) + assert pair.Hx_dpsi.shape == (n_dpsi, n_dpsi) + assert pair.hamiltonian_data.shape == (n_data, n_data) + + +def test__pair_regular_dpsi_mesh__itp_matrix_partition_of_unity(): + mask = circular_mask() + pair = al.pc.PairRegularDpsiMesh(mask, pixel_scale=0.5, dpsi_factor=2) + + row_sums = np.asarray(pair.itp_mat.sum(axis=1)).ravel() + assert row_sums == pytest.approx(np.ones_like(row_sums), abs=1.0e-10) + + +def test__pair_regular_dpsi_mesh__interpolation_exact_on_linear_function(): + mask = circular_mask() + pair = al.pc.PairRegularDpsiMesh(mask, pixel_scale=0.5, dpsi_factor=2) + + dpsi_values = 2.0 * pair.ygrid_dpsi_1d + 3.0 * pair.xgrid_dpsi_1d + data_values = pair.itp_mat @ dpsi_values + + assert data_values == pytest.approx( + 2.0 * pair.ygrid_data_1d + 3.0 * pair.xgrid_data_1d, abs=1.0e-10 + ) + + +def test__pair_regular_dpsi_mesh__indivisible_shape_raises(): + mask = circular_mask(shape=(21, 21)) + with pytest.raises(ValueError): + al.pc.PairRegularDpsiMesh(mask, pixel_scale=0.5, dpsi_factor=2) + + +def test__pair_regular_dpsi_mesh__too_sparse_raises(): + mask = circular_mask(shape=(20, 20), r_min=0.7, r_max=2.0) + with pytest.raises(ValueError): + al.pc.PairRegularDpsiMesh(mask, pixel_scale=0.5, dpsi_factor=10) + + +def test__regular_dpsi_mesh__equality(): + assert al.pc.RegularDpsiMesh(factor=2) == al.pc.RegularDpsiMesh(factor=2) + assert al.pc.RegularDpsiMesh(factor=2) != al.pc.RegularDpsiMesh(factor=3) diff --git a/test_autolens/potential_correction/test_util.py b/test_autolens/potential_correction/test_util.py new file mode 100644 index 000000000..187a12cf5 --- /dev/null +++ b/test_autolens/potential_correction/test_util.py @@ -0,0 +1,179 @@ +import numpy as np +import pytest +from scipy.sparse import csr_matrix + +import autolens as al +from autolens.potential_correction import util as pc_util + + +def test__gradient_points_from__cross_layout(): + points = np.array([[1.0, 2.0], [3.0, 4.0]]) + + grad_points = pc_util.gradient_points_from(points, cross_size=0.1) + + assert grad_points.shape == (8, 2) + # per-point order is (+y, -y, +x, -x) + assert grad_points[0] == pytest.approx([1.1, 2.0]) + assert grad_points[1] == pytest.approx([0.9, 2.0]) + assert grad_points[2] == pytest.approx([1.0, 2.1]) + assert grad_points[3] == pytest.approx([1.0, 1.9]) + assert grad_points[4] == pytest.approx([3.1, 4.0]) + + +def test__source_gradient_from__exact_on_linear_function(): + points = np.array([[0.5, -0.3], [1.0, 2.0], [-1.5, 0.7]]) + grad_points = pc_util.gradient_points_from(points, cross_size=0.01) + + # S = 2 y + 3 x has gradient (2, 3) everywhere + values = 2.0 * grad_points[:, 0] + 3.0 * grad_points[:, 1] + + gradient = pc_util.source_gradient_from(values, grad_points) + + assert gradient[:, 0] == pytest.approx(2.0, abs=1.0e-10) + assert gradient[:, 1] == pytest.approx(3.0, abs=1.0e-10) + + +def test__source_gradient_matrix_from__interleaved_x_y_structure(): + source_gradient = np.array([[2.0, 3.0], [4.0, 5.0]]) # (dS/dy, dS/dx) rows + + matrix = pc_util.source_gradient_matrix_from(source_gradient).toarray() + + expected = np.array( + [ + [3.0, 2.0, 0.0, 0.0], + [0.0, 0.0, 5.0, 4.0], + ] + ) + assert matrix == pytest.approx(expected) + + +def test__dpsi_gradient_matrix_from__interleaves_x_then_y_rows(): + itp_mat = csr_matrix(np.array([[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]])) + Hx = csr_matrix(np.array([[1.0, -1.0], [0.0, 2.0]])) + Hy = csr_matrix(np.array([[3.0, 0.0], [0.0, 4.0]])) + + matrix = pc_util.dpsi_gradient_matrix_from(itp_mat, Hx, Hy).toarray() + + Hx_itp = (itp_mat @ Hx).toarray() + Hy_itp = (itp_mat @ Hy).toarray() + expected = np.empty((6, 2)) + expected[0::2] = Hx_itp + expected[1::2] = Hy_itp + assert matrix == pytest.approx(expected) + + # acting on a dpsi vector gives interleaved (x, y) gradients per pixel + dpsi = np.array([1.0, 2.0]) + result = matrix @ dpsi + assert result[0::2] == pytest.approx(Hx_itp @ dpsi) + assert result[1::2] == pytest.approx(Hy_itp @ dpsi) + + +def test__psf_matrix_from__identity_kernel_gives_identity_matrix(): + mask = np.ones((5, 5), dtype=bool) + mask[1:4, 1:4] = False + + kernel = np.zeros((3, 3)) + kernel[1, 1] = 1.0 + + psf_mat = pc_util.psf_matrix_from(kernel, mask) + + assert psf_mat == pytest.approx(np.eye(9)) + + +def test__psf_matrix_from__blur_matches_direct_convolution(): + mask = np.ones((5, 5), dtype=bool) + mask[1:4, 1:4] = False + + kernel = np.array( + [ + [0.0, 0.1, 0.0], + [0.1, 0.6, 0.1], + [0.0, 0.1, 0.0], + ] + ) + + psf_mat = pc_util.psf_matrix_from(kernel, mask) + + # blur a delta at the central unmasked pixel (index 4 of the 3x3 block) + image = np.zeros(9) + image[4] = 1.0 + blurred = psf_mat @ image + + expected = np.array([0.0, 0.1, 0.0, 0.1, 0.6, 0.1, 0.0, 0.1, 0.0]) + assert blurred == pytest.approx(expected) + + +def test__psf_matrix_from__even_kernel_raises(): + mask = np.zeros((4, 4), dtype=bool) + with pytest.raises(Exception): + pc_util.psf_matrix_from(np.ones((4, 4)) / 16.0, mask) + + +def test__inverse_covariance_matrix_from__diagonal_of_inverse_variance(): + noise = np.array([0.5, 2.0, 1.0]) + + inv_cov = pc_util.inverse_covariance_matrix_from(noise).toarray() + + assert inv_cov == pytest.approx(np.diag([4.0, 0.25, 1.0])) + + +def test__log_det_mat__matches_slogdet_dense_and_sparse(): + rng = np.random.default_rng(0) + A = rng.normal(size=(6, 6)) + spd = A @ A.T + 6.0 * np.eye(6) + + expected = np.linalg.slogdet(spd)[1] + + assert pc_util.log_det_mat(spd) == pytest.approx(expected, rel=1.0e-10) + assert pc_util.log_det_mat(csr_matrix(spd), sparse=True) == pytest.approx( + expected, rel=1.0e-10 + ) + + +def test__log_det_mat__raises_on_negative_determinant(): + with pytest.raises(np.linalg.LinAlgError): + pc_util.log_det_mat(np.diag([1.0, -1.0])) + + +def test__dpsi_rescale_factors_from__exact_plane_through_anchors(): + anchor_points = np.array([[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]]) + # dpsi at the anchors follows the plane 2 y + 3 x + 1 + dpsi_values = 2.0 * anchor_points[:, 0] + 3.0 * anchor_points[:, 1] + 1.0 + + a_y, a_x, c = pc_util.dpsi_rescale_factors_from(anchor_points, dpsi_values) + + # the rescaling plane cancels the input plane exactly + assert a_y == pytest.approx(-2.0, abs=1.0e-10) + assert a_x == pytest.approx(-3.0, abs=1.0e-10) + assert c == pytest.approx(-1.0, abs=1.0e-10) + + +def test__split_cross_from__cross_layout_and_positive_lengths(): + rng = np.random.default_rng(1) + points = rng.uniform(-1.0, 1.0, size=(12, 2)) + + split = pc_util.split_cross_from(points) + + assert split.shape == (48, 2) + split = split.reshape(12, 4, 2) + # arm 0/1 move y only; arm 2/3 move x only, symmetrically + assert split[:, 0, 1] == pytest.approx(points[:, 1]) + assert split[:, 2, 0] == pytest.approx(points[:, 0]) + assert split[:, 0, 0] - points[:, 0] == pytest.approx( + points[:, 0] - split[:, 1, 0] + ) + assert (split[:, 0, 0] > points[:, 0]).all() + + +def test__arc_mask_from__keeps_large_islands_drops_small_ones(): + snr_map = np.zeros((24, 24)) + snr_map[6:14, 6:14] = 10.0 # 64-pixel island: kept + snr_map[20, 20] = 10.0 # single pixel island: dropped + + mask = pc_util.arc_mask_from(snr_map, threshold=3.0, ignore_size=25, ext_size=3) + + # the large island (dilated) is unmasked; the small one stays masked + assert not mask[10, 10] + assert mask[20, 20] + # dilation grows the unmasked region beyond the raw island + assert not mask[5, 10]