diff --git a/autolens/potential_correction/__init__.py b/autolens/potential_correction/__init__.py index 364300406..54924b240 100644 --- a/autolens/potential_correction/__init__.py +++ b/autolens/potential_correction/__init__.py @@ -26,3 +26,6 @@ from autolens.potential_correction.fit import FitDpsiSrcImaging from autolens.potential_correction.analysis import DpsiInvAnalysis from autolens.potential_correction.analysis import DpsiSrcInvAnalysis +from autolens.potential_correction import dense_util +from autolens.potential_correction.iterative import IterFitDpsiSrcImaging +from autolens.potential_correction.iterative import IterDpsiSrcInvAnalysis diff --git a/autolens/potential_correction/dense_util.py b/autolens/potential_correction/dense_util.py new file mode 100644 index 000000000..6a26273ff --- /dev/null +++ b/autolens/potential_correction/dense_util.py @@ -0,0 +1,447 @@ +""" +Dense linear-algebra kernels of the gravitational-imaging (potential +correction) technique, written once against the PyAuto ``xp`` API: every +function takes ``xp=np`` and runs identically under numpy (the default) or +``jax.numpy`` (pass ``xp=jnp`` for JIT-able, accelerator-ready execution). +Sparse-matrix inputs are densified on entry; no scipy.sparse arithmetic +happens inside the kernels. + +These kernels replace the scipy.sparse + numpy paths of ``fit.py`` in the +performance-critical settings of the iterative solver (``iterative.py``) and +evidence-based hyper-parameter sampling. + +Ported from the ``potential_correction`` package of Cao et al. 2025 +(https://github.com/caoxiaoyue/lensing_potential_correction), refactored from +its ``jax_ops`` module onto the ecosystem ``xp`` convention. 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 + + +def as_dense(matrix, xp=np): + """ + Converts a numpy / scipy-sparse / jax array to a dense 2D ``xp`` array. + """ + toarray = getattr(matrix, "toarray", None) + if toarray is not None: + matrix = toarray() + return xp.asarray(matrix) + + +def dense_block_diag_from(*blocks, xp=np): + """ + A dense block-diagonal matrix from two or more dense square blocks. + """ + blocks = tuple(as_dense(b, xp=xp) for b in blocks) + sizes = [b.shape[0] for b in blocks] + total = sum(sizes) + rows = [] + offset = 0 + for b, size in zip(blocks, sizes): + left = xp.zeros((size, offset), dtype=b.dtype) + right = xp.zeros((size, total - offset - size), dtype=b.dtype) + rows.append(xp.concatenate([left, b, right], axis=1)) + offset += size + return xp.concatenate(rows, axis=0) + + +def inverse_noise_variance_from(noise_slim, xp=np): + """ + The diagonal of the inverse noise covariance as a 1D vector, 1/sigma^2. + """ + noise = xp.asarray(noise_slim) + return 1.0 / (noise * noise) + + +def source_gradient_matrix_dense_from(source_gradient, xp=np): + """ + The dense source-gradient matrix D_s of shape [n_data, 2 * n_data]: + column 2i holds the x-derivative and column 2i+1 the y-derivative of the + source at data pixel i (the dense form of + ``util.source_gradient_matrix_from``). + """ + sg = xp.asarray(source_gradient) + n_data = sg.shape[0] + rows = xp.arange(n_data) + matrix = xp.zeros((n_data, 2 * n_data), dtype=sg.dtype) + if xp is np: + matrix[rows, 2 * rows] = sg[:, 1] + matrix[rows, 2 * rows + 1] = sg[:, 0] + else: + matrix = matrix.at[rows, 2 * rows].set(sg[:, 1]) + matrix = matrix.at[rows, 2 * rows + 1].set(sg[:, 0]) + return matrix + + +def dpsi_gradient_matrix_dense_from(itp_mat, Hx, Hy, xp=np): + """ + The dense dpsi-gradient operator D_psi of shape [2 * n_data, n_dpsi], + interleaving the x and y gradient rows per data pixel (the dense form of + ``util.dpsi_gradient_matrix_from``). + """ + itp = as_dense(itp_mat, xp=xp) + Hx_d = as_dense(Hx, xp=xp) + Hy_d = as_dense(Hy, xp=xp) + Hx_itp = itp @ Hx_d + Hy_itp = itp @ Hy_d + return xp.stack((Hx_itp, Hy_itp), axis=1).reshape( + 2 * itp.shape[0], Hx_d.shape[1] + ) + + +def dpsi_mapping_matrix_from(psf_mat, src_grad_mat, dpsi_grad_mat, xp=np): + """ + The dpsi mapping matrix -B D_s D_psi of shape [n_data, n_dpsi]. + """ + psf = as_dense(psf_mat, xp=xp) + sg = as_dense(src_grad_mat, xp=xp) + dg = as_dense(dpsi_grad_mat, xp=xp) + return -1.0 * psf @ sg @ dg + + +def log_evidence_joint_dense_from( + data_slim, noise_slim, mapping_matrix, src_reg_matrix, dpsi_reg_matrix, xp=np +): + """ + The Bayesian evidence of the joint source+dpsi inversion, decomposed into + its terms, from dense inputs. + + Returns + ------- + A dict with keys ``evidence``, ``valid`` (positive determinant signs), + ``solution``, ``model_image``, ``curvature_reg_matrix``, ``noise_term``, + ``logdet_curve``, ``logdet_src``, ``logdet_dpsi``, ``chi2`` and + ``reg_cov``. Under ``xp=jax.numpy`` invalid decompositions surface as + ``valid=False`` rather than raising (JIT-compatible); callers must check + it. + """ + data = xp.asarray(data_slim) + noise = xp.asarray(noise_slim) + mapping = as_dense(mapping_matrix, xp=xp) + src_reg = as_dense(src_reg_matrix, xp=xp) + dpsi_reg = as_dense(dpsi_reg_matrix, xp=xp) + + reg_matrix = dense_block_diag_from(src_reg, dpsi_reg, xp=xp) + inv_var = inverse_noise_variance_from(noise, xp=xp) + + weighted_mapping = mapping * inv_var[:, None] + curvature_matrix = mapping.T @ weighted_mapping + data_vector = mapping.T @ (inv_var * data) + + curvature_reg_matrix = curvature_matrix + reg_matrix + solution = xp.linalg.solve(curvature_reg_matrix, data_vector) + model_image = mapping @ solution + + sign_curve, logdet_curve = xp.linalg.slogdet(curvature_reg_matrix) + sign_src, logdet_src = xp.linalg.slogdet(src_reg) + sign_dpsi, logdet_dpsi = xp.linalg.slogdet(dpsi_reg) + + residual = data - model_image + reg_cov = solution @ (reg_matrix @ solution) + chi2 = xp.sum((residual / noise) ** 2.0) + noise_term = -0.5 * xp.sum(xp.log(2.0 * xp.pi * noise * noise)) + + evidence = ( + noise_term + - 0.5 * logdet_curve + + 0.5 * (logdet_src + logdet_dpsi) + - 0.5 * reg_cov + - 0.5 * chi2 + ) + + valid = (sign_curve > 0) & (sign_src > 0) & (sign_dpsi > 0) + + return { + "evidence": evidence, + "valid": valid, + "solution": solution, + "model_image": model_image, + "curvature_reg_matrix": curvature_reg_matrix, + "noise_term": noise_term, + "logdet_curve": logdet_curve, + "logdet_src": logdet_src, + "logdet_dpsi": logdet_dpsi, + "chi2": chi2, + "reg_cov": reg_cov, + } + + +def log_evidence_from_fixed_curvature( + curvature_matrix, + data_vector, + data_slim, + mapping_matrix, + inv_var, + noise_term, + src_reg_matrix, + dpsi_reg_matrix, + xp=np, +): + """ + The joint evidence with the curvature matrix M^T C^-1 M and data vector + precomputed — the fast path when only the regularization matrices change + between hyper-parameter samples. Log-determinants come from Cholesky + factors (one O(n^3) each); a failed factorization surfaces as + ``valid=False``. + + Returns the same dict as ``log_evidence_joint_dense_from`` (without + ``curvature_reg_matrix``). + """ + if xp is np: + from scipy.linalg import cho_solve + else: + from jax.scipy.linalg import cho_solve + + F = as_dense(curvature_matrix, xp=xp) + d_vec = xp.asarray(data_vector) + data = xp.asarray(data_slim) + mapping = as_dense(mapping_matrix, xp=xp) + iv = xp.asarray(inv_var) + src_reg = as_dense(src_reg_matrix, xp=xp) + dpsi_reg = as_dense(dpsi_reg_matrix, xp=xp) + + reg_matrix = dense_block_diag_from(src_reg, dpsi_reg, xp=xp) + curve_reg = F + reg_matrix + + if xp is np: + try: + L_cr = np.linalg.cholesky(curve_reg) + L_src = np.linalg.cholesky(src_reg) + L_dpsi = np.linalg.cholesky(dpsi_reg) + except np.linalg.LinAlgError: + return {"evidence": -np.inf, "valid": False} + else: + L_cr = xp.linalg.cholesky(curve_reg) + L_src = xp.linalg.cholesky(src_reg) + L_dpsi = xp.linalg.cholesky(dpsi_reg) + + solution = cho_solve((L_cr, True), d_vec) + logdet_curve = 2.0 * xp.sum(xp.log(xp.diag(L_cr))) + logdet_src = 2.0 * xp.sum(xp.log(xp.diag(L_src))) + logdet_dpsi = 2.0 * xp.sum(xp.log(xp.diag(L_dpsi))) + + model_image = mapping @ solution + residual = data - model_image + chi2 = xp.sum(iv * residual * residual) + reg_cov = solution @ (reg_matrix @ solution) + + evidence = ( + noise_term + - 0.5 * logdet_curve + + 0.5 * (logdet_src + logdet_dpsi) + - 0.5 * reg_cov + - 0.5 * chi2 + ) + + if xp is np: + valid = bool( + np.all(np.diag(L_cr) > 0) + & np.all(np.diag(L_src) > 0) + & np.all(np.diag(L_dpsi) > 0) + ) + else: + valid = ( + xp.all(xp.diag(L_cr) > 0) + & xp.all(xp.diag(L_src) > 0) + & xp.all(xp.diag(L_dpsi) > 0) + ) + + return { + "evidence": evidence, + "valid": valid, + "solution": solution, + "model_image": model_image, + "noise_term": noise_term, + "logdet_curve": logdet_curve, + "logdet_src": logdet_src, + "logdet_dpsi": logdet_dpsi, + "chi2": chi2, + "reg_cov": reg_cov, + } + + +def log_evidence_dpsi_dense_from( + data_slim, noise_slim, mapping_matrix, reg_matrix, xp=np +): + """ + The evidence of the dpsi-only inversion of an image residual, from dense + inputs. + + Returns a dict with ``evidence``, ``valid``, ``solution``, + ``model_image`` and ``curvature_reg_matrix``. + """ + data = xp.asarray(data_slim) + noise = xp.asarray(noise_slim) + mapping = as_dense(mapping_matrix, xp=xp) + reg = as_dense(reg_matrix, xp=xp) + + inv_var = inverse_noise_variance_from(noise, xp=xp) + weighted_mapping = mapping * inv_var[:, None] + curvature_matrix = mapping.T @ weighted_mapping + data_vector = mapping.T @ (inv_var * data) + + curvature_reg_matrix = curvature_matrix + reg + solution = xp.linalg.solve(curvature_reg_matrix, data_vector) + model_image = mapping @ solution + + sign_curve, logdet_curve = xp.linalg.slogdet(curvature_reg_matrix) + sign_reg, logdet_reg = xp.linalg.slogdet(reg) + + residual = data - model_image + reg_cov = solution @ (reg @ solution) + chi2 = xp.sum((residual / noise) ** 2.0) + noise_term = -0.5 * xp.sum(xp.log(2.0 * xp.pi * noise * noise)) + + evidence = ( + noise_term + - 0.5 * logdet_curve + + 0.5 * logdet_reg + - 0.5 * reg_cov + - 0.5 * chi2 + ) + + return { + "evidence": evidence, + "valid": (sign_curve > 0) & (sign_reg > 0), + "solution": solution, + "model_image": model_image, + "curvature_reg_matrix": curvature_reg_matrix, + } + + +def lm_cost_from(data_slim, inv_var, s, dpsi, L, src_reg_matrix, dpsi_reg_matrix, xp=np): + """ + The Levenberg-Marquardt cost of a candidate (s, dpsi): + 0.5 chi^2 + 0.5 s^T R_s s + 0.5 dpsi^T R_dpsi dpsi. + + Returns (cost, chi2, reg_s, reg_dpsi). + """ + data = xp.asarray(data_slim) + iv = xp.asarray(inv_var) + s_arr = xp.asarray(s) + dpsi_arr = xp.asarray(dpsi) + L_d = as_dense(L, xp=xp) + src_reg = as_dense(src_reg_matrix, xp=xp) + dpsi_reg = as_dense(dpsi_reg_matrix, xp=xp) + + residual = data - L_d @ s_arr + chi2 = 0.5 * xp.sum(iv * residual * residual) + reg_s = 0.5 * (s_arr @ (src_reg @ s_arr)) + reg_dpsi = 0.5 * (dpsi_arr @ (dpsi_reg @ dpsi_arr)) + cost = chi2 + reg_s + reg_dpsi + + return cost, chi2, reg_s, reg_dpsi + + +def lm_hessian_and_gradient_from( + data_slim, inv_var, x, L, J_dpsi, src_reg_matrix, dpsi_reg_matrix, xp=np +): + """ + The LM Hessian H = J^T C^-1 J + R, the negated gradient + J^T C^-1 r - R x, and the current residual/cost terms, for the combined + state x = [s | dpsi]. + + Returns (H, minus_gradient, residual, chi2, reg_s, reg_dpsi, cost). + """ + data = xp.asarray(data_slim) + iv = xp.asarray(inv_var) + x_arr = xp.asarray(x) + L_d = as_dense(L, xp=xp) + J_d = as_dense(J_dpsi, xp=xp) + src_reg = as_dense(src_reg_matrix, xp=xp) + dpsi_reg = as_dense(dpsi_reg_matrix, xp=xp) + + n_s = src_reg.shape[0] + s = x_arr[:n_s] + dpsi = x_arr[n_s:] + + J_combined = xp.concatenate([L_d, J_d], axis=1) + R = dense_block_diag_from(src_reg, dpsi_reg, xp=xp) + + residual = data - L_d @ s + weighted_J = J_combined * iv[:, None] + H = J_combined.T @ weighted_J + R + minus_gradient = J_combined.T @ (iv * residual) - R @ x_arr + + chi2 = 0.5 * xp.sum(iv * residual * residual) + reg_s = 0.5 * (s @ (src_reg @ s)) + reg_dpsi = 0.5 * (dpsi @ (dpsi_reg @ dpsi)) + cost = chi2 + reg_s + reg_dpsi + + return H, minus_gradient, residual, chi2, reg_s, reg_dpsi, cost + + +def solve_lm_step_from(H, minus_gradient, mu, constraint_matrix=None, x=None, xp=np): + """ + The damped LM step delta_x solving (H + mu I) dx = -g, or — when a + ``constraint_matrix`` C is given — the equality-constrained step via the + KKT system enforcing C (x + dx) = 0. + """ + H_d = as_dense(H, xp=xp) + g = xp.asarray(minus_gradient) + n_x = H_d.shape[0] + H_lm = H_d + mu * xp.eye(n_x, dtype=H_d.dtype) + + if constraint_matrix is None: + return xp.linalg.solve(H_lm, g) + + C = xp.asarray(constraint_matrix) + x_arr = xp.asarray(x) + n_c = C.shape[0] + top = xp.concatenate([H_lm, C.T], axis=1) + bottom = xp.concatenate( + [C, xp.zeros((n_c, n_c), dtype=H_d.dtype)], axis=1 + ) + H_kkt = xp.concatenate([top, bottom], axis=0) + rhs = xp.concatenate([g, -(C @ x_arr)]) + solution = xp.linalg.solve(H_kkt, rhs) + return solution[:n_x] + + +def log_evidence_lm_from( + data_slim, + noise_slim, + s, + dpsi, + L, + J_s, + J_dpsi, + src_reg_matrix, + dpsi_reg_matrix, + xp=np, +): + """ + The Laplace-approximation log evidence at an LM solution (s, dpsi), + without the noise-normalization term (add it separately): + 0.5 [ logdet R_s + logdet R_dpsi - logdet H - chi^2 - s^T R_s s + - dpsi^T R_dpsi dpsi ]. + """ + data = xp.asarray(data_slim) + noise = xp.asarray(noise_slim) + s_arr = xp.asarray(s) + dpsi_arr = xp.asarray(dpsi) + L_d = as_dense(L, xp=xp) + J_s_d = as_dense(J_s, xp=xp) + J_dpsi_d = as_dense(J_dpsi, xp=xp) + src_reg = as_dense(src_reg_matrix, xp=xp) + dpsi_reg = as_dense(dpsi_reg_matrix, xp=xp) + + inv_var = inverse_noise_variance_from(noise, xp=xp) + + residual = data - L_d @ s_arr + chi2 = xp.sum(inv_var * residual * residual) + + reg_val = s_arr @ (src_reg @ s_arr) + dpsi_arr @ (dpsi_reg @ dpsi_arr) + + J_combined = xp.concatenate([J_s_d, J_dpsi_d], axis=1) + R = dense_block_diag_from(src_reg, dpsi_reg, xp=xp) + H = J_combined.T @ (J_combined * inv_var[:, None]) + R + + _, logdet_H = xp.linalg.slogdet(H) + _, logdet_Rs = xp.linalg.slogdet(src_reg) + _, logdet_Rd = xp.linalg.slogdet(dpsi_reg) + + return 0.5 * ((logdet_Rs + logdet_Rd) - logdet_H - chi2 - reg_val) diff --git a/autolens/potential_correction/iterative.py b/autolens/potential_correction/iterative.py new file mode 100644 index 000000000..0129a7540 --- /dev/null +++ b/autolens/potential_correction/iterative.py @@ -0,0 +1,739 @@ +""" +The iterative solver of the gravitational-imaging (potential correction) +technique: ``IterFitDpsiSrcImaging`` jointly optimizes the pixelized source +and the pixelized potential corrections dpsi with a Levenberg-Marquardt +loop, re-ray-tracing the image grid through the corrected lens (the macro +model plus an ``InputPotential`` built from the current dpsi) at every step. + +The linear algebra runs through the ``xp``-parameterized kernels of +``dense_util``: pass ``xp=jax.numpy`` to ``solve_joint_optimization`` / +``log_evidence`` for accelerator-ready execution, or leave the numpy default. + +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 +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 autogalaxy.profiles.mass.input.input_potential import InputPotential +from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages + +from autolens.lens.tracer import Tracer +from autolens.imaging.fit_imaging import FitImaging +from autolens.potential_correction import dense_util +from autolens.potential_correction import util as pc_util +from autolens.potential_correction.pixelization import ( + DpsiLinearObj, + DpsiPixelization, + DpsiSrcPixelization, +) +from autolens.potential_correction.src_factory import PixSrcFactoryITP, SrcFactory + +logger = logging.getLogger(__name__) + + +class IterFitDpsiSrcImaging: + def __init__( + self, + masked_imaging, + lens_start: Galaxy, + dpsi_pixelization: DpsiPixelization, + src_pixelization: aa.Pixelization, + gauge_constraints: bool = False, + adapt_image=None, + src_image_mesh=None, + settings_inversion: Optional[aa.Settings] = None, + preloads: Optional[dict] = None, + n_iter: int = 20, + tol: float = 1e-6, + verbose: bool = False, + visualize_output_dir: Optional[str] = None, + visualize_every_n: int = 1000000, + ): + """ + Iteratively solves for the pixelized source s and potential + corrections dpsi minimizing the penalized objective + P(s, dpsi) = 0.5 [ (d - L(psi0+dpsi) s)^T C^-1 (d - L(psi0+dpsi) s) + + s^T R_s s + dpsi^T R_dpsi dpsi ], with a Levenberg-Marquardt loop + on the combined state x = [s | dpsi]. Unlike the one-shot + ``FitDpsiSrcImaging`` linearization, each accepted step re-ray-traces + the image grid through the corrected lens, so the corrections feed + back into the source mapping. + + Parameters + ---------- + masked_imaging + The masked ``al.Imaging`` dataset. + lens_start + The macro lens galaxy the corrections perturb. + dpsi_pixelization + The dpsi mesh + regularization model. + src_pixelization + The source pixelization. + gauge_constraints + Whether to impose the gauge constraints = = + = 0 (removing the constant / deflection-drift + degeneracies) via an equality-constrained KKT step. + 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 border relocator without + the positive-only solver (the LM state is signed). + preloads + Precomputed attributes set directly onto the fit. + n_iter + The maximum number of outer LM iterations. + tol + The step-norm convergence tolerance. + verbose + Whether to log per-iteration costs. + visualize_output_dir + Directory the intermediate ``show`` figures are written to. + visualize_every_n + Visualize every n-th iteration (default effectively never). + """ + self.masked_imaging = masked_imaging + self.lens_start = lens_start + self.dpsi_pixelization = dpsi_pixelization + self.src_pixelization = src_pixelization + self.gauge_constraints = gauge_constraints + self.adapt_image = adapt_image + self.src_image_mesh = src_image_mesh + self.n_iter = int(n_iter) + self.tol = float(tol) + self.verbose = bool(verbose) + self.visualize_output_dir = visualize_output_dir + self.visualize_every_n = int(visualize_every_n) + + if settings_inversion is None: + self.settings_inversion = aa.Settings( + use_positive_only_solver=False, + 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, + ) + + # ----------------------------- + # Cached low-level operators + # ----------------------------- + + @property + def inverse_noise_variance(self): + if not hasattr(self, "inv_noise_var"): + self.inv_noise_var = np.asarray( + 1.0 / self.masked_imaging.noise_map.slim**2 + ) + return self.inv_noise_var + + @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 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, "dpsi_grad_mat"): + self.dpsi_grad_mat = pc_util.dpsi_gradient_matrix_from( + self.pair_dpsi_data_obj.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_regularization_matrix(self): + if not hasattr(self, "dpsi_reg_mat"): + linear_obj = DpsiLinearObj( + mask=self.pair_dpsi_data_obj.mask_dpsi, points=self.dpsi_points + ) + self.dpsi_reg_mat = ( + self.dpsi_pixelization.regularization.regularization_matrix_from( + linear_obj=linear_obj + ) + ) + return self.dpsi_reg_mat + + # ----------------------------- + # Per-iteration constructions + # ----------------------------- + + def _updated_lens_galaxies_from_dpsi(self, dpsi_vec: np.ndarray): + """ + The macro lens plus a pixelized ``InputPotential`` correction built + from the current dpsi vector, as two galaxies on the lens plane. + """ + mask_dpsi_aa = self.pair_dpsi_data_obj.mask_dpsi_aa + grid_dpsi = aa.Grid2D.from_mask(mask=mask_dpsi_aa) + + pix_mass_profile = InputPotential( + lensing_potential=dpsi_vec, + image_plane_grid=np.asarray(grid_dpsi), + mask=mask_dpsi_aa, + ) + + lens_macro = self.lens_start + lens_pix = Galaxy( + redshift=getattr(lens_macro, "redshift", 0.2), mass=pix_mass_profile + ) + + return [lens_macro, lens_pix] + + def _build_pix_src_tracer(self, lens_galaxies): + 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 + } + + adapt_images = AdaptImages(**adapt_kwargs) if adapt_kwargs else None + + tracer = Tracer(galaxies=[*lens_galaxies, source_galaxy]) + + return tracer, adapt_images + + def _init_source_inversion(self, lens_galaxies): + """ + Runs one full source inversion at the given lens galaxies, caching + the source mesh, mapping and regularization matrices the iteration + keeps fixed (the source mesh does not move between LM steps). + """ + tracer, adapt_images = self._build_pix_src_tracer(lens_galaxies) + + src_fit = FitImaging( + dataset=self.masked_imaging, + tracer=tracer, + adapt_images=adapt_images, + settings=self.settings_inversion, + ) + + self.tracer = tracer + mapper = src_fit.inversion.linear_obj_list[0] + self.source_plane_mesh_grid = mapper.source_plane_mesh_grid + self.image_plane_mesh_grid = mapper.image_plane_mesh_grid + self._src_mesh = mapper.mesh + self._border_relocator = ( + self.masked_imaging.grids.border_relocator + if self.settings_inversion.use_border_relocator + else None + ) + self._src_regularization = self.src_pixelization.regularization + self.src_reg_mat = src_fit.inversion.regularization_matrix + self.src_map_mat = src_fit.inversion.operated_mapping_matrix + + def _update_source_inversion(self, lens_galaxies): + """ + Re-ray-traces the data grid through the updated lens, keeping the + source mesh fixed, and rebuilds the PSF-convolved source mapping + matrix. + """ + tracer, _ = self._build_pix_src_tracer(lens_galaxies) + self.tracer = tracer + + source_plane_data_grid = tracer.traced_grid_2d_list_from( + grid=self.masked_imaging.grid + )[-1] + + interpolator = self._src_mesh.interpolator_from( + source_plane_data_grid=source_plane_data_grid, + source_plane_mesh_grid=self.source_plane_mesh_grid, + border_relocator=self._border_relocator, + adapt_data=None, + ) + + mapper = aa.Mapper( + interpolator=interpolator, + regularization=self._src_regularization, + settings=self.settings_inversion, + image_plane_mesh_grid=self.image_plane_mesh_grid, + ) + + self.src_map_mat = self.masked_imaging.psf.convolved_mapping_matrix_from( + mapping_matrix=mapper.mapping_matrix, + mask=self.masked_imaging.mask, + ) + + def _source_plane_data_grid(self, tracer): + defl = tracer.deflections_yx_2d_from(self.masked_imaging.grid.slim) + return self.masked_imaging.grid.slim - defl + + def _source_gradient_matrix(self, tracer, source_factory: SrcFactory): + src_grid = self._source_plane_data_grid(tracer) + src_grad_vals = source_factory.eval_grad(src_grid[:, 1], src_grid[:, 0]) + return pc_util.source_gradient_matrix_from(src_grad_vals) + + def _source_factory_from(self, s: np.ndarray): + return PixSrcFactoryITP(points=self.source_plane_mesh_grid, values=s) + + def get_L_Js_Jdpsi(self, s, dpsi, xp=np): + """ + The source mapping matrix L and the Jacobians (J_s = L, J_dpsi) of + the model image w.r.t. the source and dpsi at the state (s, dpsi), + after updating the lens for the current corrections. + """ + lens_galaxies = self._updated_lens_galaxies_from_dpsi(np.asarray(dpsi)) + source_factory = self._source_factory_from(np.asarray(s)) + self._update_source_inversion(lens_galaxies) + src_grad_mat = self._source_gradient_matrix(self.tracer, source_factory) + dpsi_map_mat = dense_util.dpsi_mapping_matrix_from( + self.psf_matrix, src_grad_mat, self.dpsi_gradient_matrix, xp=xp + ) + L = dense_util.as_dense(self.src_map_mat, xp=xp) + return L, L, dpsi_map_mat + + def _regularization_matrix(self, xp=np): + return dense_util.dense_block_diag_from( + self.src_reg_mat, self.dpsi_regularization_matrix, xp=xp + ) + + def _init_joint_optimization(self): + n_dpsi = self.dpsi_regularization_matrix.shape[0] + lens_galaxies = self._updated_lens_galaxies_from_dpsi(np.zeros(n_dpsi)) + self._init_source_inversion(lens_galaxies) + n_s = self.src_reg_mat.shape[0] + return n_s, n_dpsi + + # ----------------------------- + # The Levenberg-Marquardt loop + # ----------------------------- + + def solve_joint_optimization(self, xp=np): + """ + Runs the Levenberg-Marquardt loop on the combined state + x = [s | dpsi], accepting steps that decrease the penalized cost and + adapting the damping mu. Returns the optimized (s, dpsi); they are + also stored as ``s_opt`` / ``dpsi_opt``. + + Parameters + ---------- + xp + The array backend of the dense kernels: ``numpy`` (default) or + ``jax.numpy`` for accelerator-ready execution. The autolens + ray-tracing / source-mesh updates between steps always run on + the host. + """ + n_s, n_dpsi = self._init_joint_optimization() + + x = xp.zeros(n_s + n_dpsi) + + data_slim = xp.asarray(np.asarray(self.masked_imaging.data.slim)) + inv_var = xp.asarray(self.inverse_noise_variance) + src_reg_mat = dense_util.as_dense(self.src_reg_mat, xp=xp) + dpsi_reg_mat = dense_util.as_dense(self.dpsi_regularization_matrix, xp=xp) + + constraint_matrix = None + if self.gauge_constraints: + # gauge: = = = 0, removing the + # constant and deflection-drift degeneracies of the potential + G = np.zeros((3, n_dpsi)) + G[0, :] = 1.0 / n_dpsi + G[1, :] = self.dpsi_points[:, 1] / n_dpsi + G[2, :] = self.dpsi_points[:, 0] / n_dpsi + constraint_matrix = xp.asarray( + np.hstack([np.zeros((3, n_s)), G]) + ) + + s = np.asarray(x[:n_s]) + dpsi = np.asarray(x[n_s:]) + L, J_s, J_dpsi = self.get_L_Js_Jdpsi(s, dpsi, xp=xp) + + current_cost, chi2, reg_s, reg_dpsi = ( + float(v) + for v in dense_util.lm_cost_from( + data_slim, inv_var, x[:n_s], x[n_s:], L, src_reg_mat, dpsi_reg_mat, + xp=xp, + ) + ) + + mu = 1.0 + + if self.verbose: + logger.info( + "Starting joint LM optimization: %d iterations, initial cost %.4e", + self.n_iter, + current_cost, + ) + + for i in range(self.n_iter): + H, minus_gradient, residual, chi2, reg_s, reg_dpsi, current_cost = ( + dense_util.lm_hessian_and_gradient_from( + data_slim, inv_var, x, L, J_dpsi, src_reg_mat, dpsi_reg_mat, + xp=xp, + ) + ) + current_cost = float(current_cost) + + if self.visualize_output_dir is not None and i % self.visualize_every_n == 0: + self.show( + np.asarray(L), np.asarray(x[:n_s]), np.asarray(x[n_s:]), + output=f"iter_{i}.png", + ) + + if self.verbose: + logger.info( + "Iter %d: cost=%.4e chi2=%.4e reg_s=%.4e mu=%.1e", + i, current_cost, float(chi2), float(reg_s), mu, + ) + + step_accepted = False + while not step_accepted: + delta_x = None + try: + delta_x = dense_util.solve_lm_step_from( + H, minus_gradient, mu, + constraint_matrix=constraint_matrix, x=x, xp=xp, + ) + if np.any(np.isnan(np.asarray(delta_x))): + delta_x = None + except np.linalg.LinAlgError: + delta_x = None + + if delta_x is not None: + x_new = x + delta_x + s_new = np.asarray(x_new[:n_s]) + dpsi_new = np.asarray(x_new[n_s:]) + + L_new, J_s_new, J_dpsi_new = self.get_L_Js_Jdpsi( + s_new, dpsi_new, xp=xp + ) + cost_new, chi2_new, reg_s_new, reg_dpsi_new = ( + dense_util.lm_cost_from( + data_slim, inv_var, x_new[:n_s], x_new[n_s:], + L_new, src_reg_mat, dpsi_reg_mat, xp=xp, + ) + ) + cost_new = float(cost_new) + + if cost_new < current_cost: + x = x_new + L, J_s, J_dpsi = L_new, J_s_new, J_dpsi_new + current_cost = cost_new + mu = max(1e-15, mu / 3.0) + step_accepted = True + + if float(xp.linalg.norm(delta_x)) < self.tol: + if self.verbose: + logger.info( + "Converged at iteration %d (step tolerance).", i + ) + self.s_opt = np.asarray(x[:n_s]) + self.dpsi_opt = np.asarray(x[n_s:]) + return self.s_opt, self.dpsi_opt + else: + mu *= 5.0 + if mu > 1e15: + logger.warning( + "LM damping parameter exceeded 1e15; stopping." + ) + self.s_opt = np.asarray(x[:n_s]) + self.dpsi_opt = np.asarray(x[n_s:]) + return self.s_opt, self.dpsi_opt + else: + mu *= 5.0 + if mu > 1e15: + logger.warning("LM solver failed repeatedly; stopping.") + self.s_opt = np.asarray(x[:n_s]) + self.dpsi_opt = np.asarray(x[n_s:]) + return self.s_opt, self.dpsi_opt + + self.s_opt = np.asarray(x[:n_s]) + self.dpsi_opt = np.asarray(x[n_s:]) + return self.s_opt, self.dpsi_opt + + def log_evidence( + self, + s: Optional[np.ndarray] = None, + dpsi: Optional[np.ndarray] = None, + include_noise_normalization: bool = True, + xp=np, + ) -> float: + """ + The Laplace-approximation log evidence at (s, dpsi) (defaulting to + the optimized state of ``solve_joint_optimization``). + """ + if s is None or dpsi is None: + if hasattr(self, "s_opt") and hasattr(self, "dpsi_opt"): + s = self.s_opt + dpsi = self.dpsi_opt + else: + raise ValueError( + "Provide s and dpsi, or run solve_joint_optimization() first." + ) + + L, J_s, J_dpsi = self.get_L_Js_Jdpsi(s, dpsi, xp=xp) + + log_evidence = float( + dense_util.log_evidence_lm_from( + data_slim=np.asarray(self.masked_imaging.data.slim), + noise_slim=np.asarray(self.masked_imaging.noise_map.slim), + s=s, + dpsi=dpsi, + L=L, + J_s=J_s, + J_dpsi=J_dpsi, + src_reg_matrix=self.src_reg_mat, + dpsi_reg_matrix=self.dpsi_regularization_matrix, + xp=xp, + ) + ) + + if include_noise_normalization: + noise_1d = np.asarray(self.masked_imaging.noise_map.slim) + log_evidence += float(np.sum(np.log(2.0 * np.pi * noise_1d**2.0))) * ( + -0.5 + ) + + return log_evidence + + def show(self, L, s, dpsi, output="show", interpolate=False, show_src_grid=False): + """ + The nine-panel summary of the current LM state (data, noise, SNR, + model, residuals, dpsi/dkappa maps and the source). + """ + import copy + import os + + from matplotlib import pyplot as plt + + from autolens.potential_correction.visualize import ( + imshow_masked_data, + show_image_irregular, + show_image_irregular_interpolate, + ) + + fig = plt.figure(figsize=(15, 10)) + cmap = copy.copy(plt.get_cmap("jet")) + cmap.set_bad(color="white") + myargs_data = { + "origin": "upper", + "cmap": cmap, + "extent": self.pair_dpsi_data_obj.data_bound, + } + xlimit = [ + self.pair_dpsi_data_obj.xgrid_data_1d.min(), + self.pair_dpsi_data_obj.xgrid_data_1d.max(), + ] + ylimit = [ + self.pair_dpsi_data_obj.ygrid_data_1d.min(), + self.pair_dpsi_data_obj.ygrid_data_1d.max(), + ] + + def plot_subplot(pos, data, mask, title): + plt.subplot(pos) + ax = plt.gca() + imshow_masked_data(data, mask, ax=ax, **myargs_data) + ax.set_title(title) + ax.set_xlim(*xlimit) + ax.set_ylim(*ylimit) + return ax + + plot_subplot(331, self.masked_imaging.data, self.masked_imaging.mask, "Data") + plot_subplot(332, self.masked_imaging.noise_map, self.masked_imaging.mask, "Noise") + plot_subplot( + 333, + self.masked_imaging.data / self.masked_imaging.noise_map, + self.masked_imaging.mask, + "SNR", + ) + + model_image_slim = np.asarray(L) @ np.asarray(s) + ax = plot_subplot(334, model_image_slim, self.masked_imaging.mask, "Model") + if show_src_grid: + ax.scatter( + self.image_plane_mesh_grid[:, 1], self.image_plane_mesh_grid[:, 0], + c="black", s=0.5, alpha=0.5, + ) + + residual = self.masked_imaging.data - model_image_slim + plot_subplot(335, residual, self.masked_imaging.mask, "Residual") + plot_subplot( + 336, residual / self.masked_imaging.noise_map, self.masked_imaging.mask, + "Norm Residual", + ) + + plot_subplot(337, dpsi, self.pair_dpsi_data_obj.mask_dpsi, "Dpsi Map") + plot_subplot( + 338, + self.pair_dpsi_data_obj.hamiltonian_dpsi @ np.asarray(dpsi), + self.pair_dpsi_data_obj.mask_dpsi, + "Dkappa Map", + ) + + plt.subplot(339) + ax = plt.gca() + if interpolate: + show_image_irregular_interpolate( + self.source_plane_mesh_grid, s, ax=ax, enlarge_factor=1.1, + npixels=100, cmap="jet", + ) + else: + show_image_irregular( + self.source_plane_mesh_grid, s, enlarge_factor=1.1, cmap="jet", + ax=ax, title="Source", + ) + if show_src_grid: + ax.scatter( + self.source_plane_mesh_grid[:, 1], self.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 + if self.visualize_output_dir is None: + raise ValueError("visualize_output_dir must be set when saving figures") + os.makedirs(self.visualize_output_dir, exist_ok=True) + fig.savefig( + os.path.join(self.visualize_output_dir, output), bbox_inches="tight" + ) + plt.close(fig) + + +class IterDpsiSrcInvAnalysis(af.Analysis): + def __init__( + self, + masked_imaging, + lens_start: Galaxy, + n_iter: int = 20, + tol: float = 1e-6, + adapt_image=None, + settings_inversion: Optional[aa.Settings] = None, + preloads: Optional[dict] = None, + gauge_constraints: bool = True, + verbose: bool = False, + xp=np, + ): + """ + Samples the joint source+dpsi pixelization hyper-parameters with the + iterative LM solver's Laplace evidence as the likelihood. + + Parameters + ---------- + masked_imaging + The masked ``al.Imaging`` dataset. + lens_start + The macro lens galaxy the corrections perturb. + n_iter, tol, gauge_constraints, verbose + Forwarded to ``IterFitDpsiSrcImaging``. + adapt_image + The adapt image of the source pixelization's image mesh. + settings_inversion + The inversion settings; defaults to the border relocator without + the positive-only solver, with the adaptive image-mesh + thresholds of the original implementation. + preloads + Precomputed fit attributes shared across evaluations. + xp + The array backend of the dense LM kernels (``numpy`` or + ``jax.numpy``). + """ + self.masked_imaging = masked_imaging + self.lens_start = lens_start + self.n_iter = int(n_iter) + self.tol = float(tol) + self.adapt_image = adapt_image + self.gauge_constraints = bool(gauge_constraints) + self.verbose = bool(verbose) + self.xp = xp + if settings_inversion is None: + self.settings_inversion = aa.Settings( + use_positive_only_solver=False, + use_border_relocator=True, + image_mesh_min_mesh_pixels_per_pixel=3, + image_mesh_min_mesh_number=5, + image_mesh_adapt_background_percent_threshold=0.1, + image_mesh_adapt_background_percent_check=0.8, + ) + else: + self.settings_inversion = settings_inversion + self.preloads = preloads + + def _fit_from(self, instance: DpsiSrcPixelization, visualize_output_dir=None): + return IterFitDpsiSrcImaging( + masked_imaging=self.masked_imaging, + lens_start=self.lens_start, + dpsi_pixelization=instance.dpsi_pixelization, + src_pixelization=instance.src_pixelization, + adapt_image=self.adapt_image, + gauge_constraints=self.gauge_constraints, + n_iter=self.n_iter, + tol=self.tol, + verbose=self.verbose, + settings_inversion=self.settings_inversion, + preloads=self.preloads, + visualize_output_dir=visualize_output_dir, + ) + + def log_likelihood_function(self, instance: DpsiSrcPixelization): + fit = self._fit_from(instance) + try: + s_opt, dpsi_opt = fit.solve_joint_optimization(xp=self.xp) + return fit.log_evidence(s=s_opt, dpsi=dpsi_opt, xp=self.xp) + except exc.InversionException: + # a failed inversion is a valid (very bad) sample, not a crash + logger.exception( + "InversionException during iterative source+dpsi optimization; " + "returning penalty likelihood." + ) + return -1e30 + + def visualize(self, paths: af.DirectoryPaths, instance, during_analysis=True): + fit = self._fit_from(instance, visualize_output_dir=str(paths.image_path)) + s_opt, dpsi_opt = fit.solve_joint_optimization(xp=self.xp) + L, _, _ = fit.get_L_Js_Jdpsi(s_opt, dpsi_opt, xp=self.xp) + fit.show(L, s_opt, dpsi_opt, output="final.png") diff --git a/autolens/potential_correction/util.py b/autolens/potential_correction/util.py index 0c554d581..51fd39690 100644 --- a/autolens/potential_correction/util.py +++ b/autolens/potential_correction/util.py @@ -282,8 +282,18 @@ def split_cross_from(points: np.ndarray) -> np.ndarray: 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 + # the percentile must be taken over bounded cells only: with many + # unbounded cells (small meshes) a percentile over the -1 sentinels goes + # negative and the arm lengths become NaN + positive_areas = region_areas[region_areas > 0] + if positive_areas.size == 0: + from scipy.spatial import cKDTree + + distances, _ = cKDTree(points).query(points, k=2) + max_area = float(np.median(distances[:, 1])) ** 2 + else: + max_area = np.percentile(positive_areas, 90.0) + region_areas[region_areas <= 0] = max_area region_areas[region_areas > max_area] = max_area half_lengths = 0.5 * np.sqrt(region_areas) diff --git a/test_autolens/potential_correction/test_dense_util.py b/test_autolens/potential_correction/test_dense_util.py new file mode 100644 index 000000000..135164dc8 --- /dev/null +++ b/test_autolens/potential_correction/test_dense_util.py @@ -0,0 +1,264 @@ +import numpy as np +import pytest +from scipy.linalg import block_diag as scipy_block_diag +from scipy.sparse import csr_matrix + +from autolens.potential_correction import dense_util +from autolens.potential_correction import util as pc_util + + +def test__as_dense__handles_sparse_and_dense(): + dense = np.array([[1.0, 2.0], [3.0, 4.0]]) + + assert dense_util.as_dense(dense) == pytest.approx(dense) + assert dense_util.as_dense(csr_matrix(dense)) == pytest.approx(dense) + + +def test__dense_block_diag_from__matches_scipy(): + a = np.array([[1.0, 2.0], [3.0, 4.0]]) + b = np.array([[5.0]]) + c = np.eye(3) * 7.0 + + result = dense_util.dense_block_diag_from(a, b, c) + + assert result == pytest.approx(scipy_block_diag(a, b, c)) + + +def test__source_gradient_matrix_dense_from__matches_sparse_version(): + rng = np.random.default_rng(0) + source_gradient = rng.normal(size=(5, 2)) + + dense = dense_util.source_gradient_matrix_dense_from(source_gradient) + sparse = pc_util.source_gradient_matrix_from(source_gradient).toarray() + + assert dense == pytest.approx(sparse) + + +def test__dpsi_gradient_matrix_dense_from__matches_sparse_version(): + rng = np.random.default_rng(1) + itp = csr_matrix(rng.normal(size=(6, 4))) + Hx = csr_matrix(rng.normal(size=(4, 4))) + Hy = csr_matrix(rng.normal(size=(4, 4))) + + dense = dense_util.dpsi_gradient_matrix_dense_from(itp, Hx, Hy) + sparse = pc_util.dpsi_gradient_matrix_from(itp, Hx, Hy).toarray() + + assert dense == pytest.approx(sparse) + + +def test__dpsi_mapping_matrix_from__is_minus_psf_srcgrad_dpsigrad(): + rng = np.random.default_rng(2) + psf = rng.normal(size=(4, 4)) + src_grad = rng.normal(size=(4, 8)) + dpsi_grad = rng.normal(size=(8, 3)) + + result = dense_util.dpsi_mapping_matrix_from(psf, src_grad, dpsi_grad) + + assert result == pytest.approx(-1.0 * psf @ src_grad @ dpsi_grad) + + +def joint_problem(n_data=6, n_src=3, n_dpsi=2, seed=3): + rng = np.random.default_rng(seed) + data = rng.normal(size=n_data) + noise = rng.uniform(0.5, 1.5, size=n_data) + mapping = rng.normal(size=(n_data, n_src + n_dpsi)) + A = rng.normal(size=(n_src, n_src)) + src_reg = A @ A.T + n_src * np.eye(n_src) + B = rng.normal(size=(n_dpsi, n_dpsi)) + dpsi_reg = B @ B.T + n_dpsi * np.eye(n_dpsi) + return data, noise, mapping, src_reg, dpsi_reg + + +def reference_joint_evidence(data, noise, mapping, src_reg, dpsi_reg): + """The phase-3 numpy evidence formulation, computed directly.""" + reg = scipy_block_diag(src_reg, dpsi_reg) + inv_cov = np.diag(1.0 / noise**2) + curve_reg = mapping.T @ inv_cov @ mapping + reg + d_vec = mapping.T @ inv_cov @ data + solution = np.linalg.solve(curve_reg, d_vec) + model = mapping @ solution + + noise_term = -0.5 * np.sum(np.log(2 * np.pi * noise**2)) + logdet_curve = -0.5 * np.linalg.slogdet(curve_reg)[1] + logdet_reg = 0.5 * ( + np.linalg.slogdet(src_reg)[1] + np.linalg.slogdet(dpsi_reg)[1] + ) + reg_cov = -0.5 * float(solution @ reg @ solution) + chi2 = -0.5 * float(np.sum(((data - model) / noise) ** 2)) + return noise_term + logdet_curve + logdet_reg + reg_cov + chi2, solution + + +def test__log_evidence_joint_dense_from__matches_reference_formulation(): + data, noise, mapping, src_reg, dpsi_reg = joint_problem() + + result = dense_util.log_evidence_joint_dense_from( + data, noise, mapping, src_reg, dpsi_reg + ) + + expected, solution = reference_joint_evidence( + data, noise, mapping, src_reg, dpsi_reg + ) + assert bool(result["valid"]) + assert float(result["evidence"]) == pytest.approx(expected, rel=1.0e-10) + assert result["solution"] == pytest.approx(solution, rel=1.0e-10) + + +def test__log_evidence_from_fixed_curvature__matches_joint_dense(): + data, noise, mapping, src_reg, dpsi_reg = joint_problem() + + full = dense_util.log_evidence_joint_dense_from( + data, noise, mapping, src_reg, dpsi_reg + ) + + inv_var = 1.0 / noise**2 + curvature = mapping.T @ (mapping * inv_var[:, None]) + data_vector = mapping.T @ (inv_var * data) + noise_term = -0.5 * np.sum(np.log(2 * np.pi * noise**2)) + + fast = dense_util.log_evidence_from_fixed_curvature( + curvature_matrix=curvature, + data_vector=data_vector, + data_slim=data, + mapping_matrix=mapping, + inv_var=inv_var, + noise_term=noise_term, + src_reg_matrix=src_reg, + dpsi_reg_matrix=dpsi_reg, + ) + + assert bool(fast["valid"]) + assert float(fast["evidence"]) == pytest.approx( + float(full["evidence"]), rel=1.0e-10 + ) + assert fast["solution"] == pytest.approx(np.asarray(full["solution"]), rel=1.0e-8) + + +def test__log_evidence_from_fixed_curvature__invalid_on_non_pd_regularization(): + data, noise, mapping, src_reg, dpsi_reg = joint_problem() + + result = dense_util.log_evidence_from_fixed_curvature( + curvature_matrix=mapping.T @ mapping, + data_vector=mapping.T @ data, + data_slim=data, + mapping_matrix=mapping, + inv_var=1.0 / noise**2, + noise_term=0.0, + src_reg_matrix=-1.0 * src_reg, # not positive definite + dpsi_reg_matrix=dpsi_reg, + ) + + assert not bool(result["valid"]) + + +def test__log_evidence_dpsi_dense_from__matches_reference_formulation(): + rng = np.random.default_rng(4) + n_data, n_dpsi = 5, 3 + data = rng.normal(size=n_data) + noise = rng.uniform(0.5, 1.5, size=n_data) + mapping = rng.normal(size=(n_data, n_dpsi)) + reg = np.diag(rng.uniform(0.5, 2.0, size=n_dpsi)) + + result = dense_util.log_evidence_dpsi_dense_from(data, noise, mapping, reg) + + inv_cov = np.diag(1.0 / noise**2) + curve_reg = mapping.T @ inv_cov @ mapping + reg + solution = np.linalg.solve(curve_reg, mapping.T @ inv_cov @ data) + model = mapping @ solution + expected = ( + -0.5 * np.sum(np.log(2 * np.pi * noise**2)) + - 0.5 * np.linalg.slogdet(curve_reg)[1] + + 0.5 * np.linalg.slogdet(reg)[1] + - 0.5 * float(solution @ reg @ solution) + - 0.5 * float(np.sum(((data - model) / noise) ** 2)) + ) + + assert bool(result["valid"]) + assert float(result["evidence"]) == pytest.approx(expected, rel=1.0e-10) + + +def test__lm_cost_and_hessian__hand_computed(): + data, noise, mapping, src_reg, dpsi_reg = joint_problem() + n_src = src_reg.shape[0] + inv_var = 1.0 / noise**2 + L = mapping[:, :n_src] + J_dpsi = mapping[:, n_src:] + + rng = np.random.default_rng(5) + x = rng.normal(size=mapping.shape[1]) + s, dpsi = x[:n_src], x[n_src:] + + cost, chi2, reg_s, reg_dpsi = dense_util.lm_cost_from( + data, inv_var, s, dpsi, L, src_reg, dpsi_reg + ) + residual = data - L @ s + assert float(chi2) == pytest.approx(0.5 * np.sum(inv_var * residual**2)) + assert float(reg_s) == pytest.approx(0.5 * s @ src_reg @ s) + assert float(cost) == pytest.approx(float(chi2) + float(reg_s) + float(reg_dpsi)) + + H, minus_gradient, res, chi2_h, _, _, cost_h = ( + dense_util.lm_hessian_and_gradient_from( + data, inv_var, x, L, J_dpsi, src_reg, dpsi_reg + ) + ) + J = np.hstack([L, J_dpsi]) + R = scipy_block_diag(src_reg, dpsi_reg) + assert H == pytest.approx(J.T @ (J * inv_var[:, None]) + R) + assert minus_gradient == pytest.approx(J.T @ (inv_var * residual) - R @ x) + assert float(cost_h) == pytest.approx(float(cost)) + + +def test__solve_lm_step_from__unconstrained_and_constrained(): + data, noise, mapping, src_reg, dpsi_reg = joint_problem() + n_src = src_reg.shape[0] + inv_var = 1.0 / noise**2 + x = np.zeros(mapping.shape[1]) + + H, minus_gradient, *_ = dense_util.lm_hessian_and_gradient_from( + data, inv_var, x, mapping[:, :n_src], mapping[:, n_src:], src_reg, dpsi_reg + ) + + mu = 0.7 + step = dense_util.solve_lm_step_from(H, minus_gradient, mu) + assert (H + mu * np.eye(H.shape[0])) @ step == pytest.approx( + minus_gradient, rel=1.0e-8 + ) + + # a single constraint: the dpsi block sums to zero after the step + C = np.zeros((1, mapping.shape[1])) + C[0, n_src:] = 1.0 + step_c = dense_util.solve_lm_step_from( + H, minus_gradient, mu, constraint_matrix=C, x=x + ) + assert float((C @ (x + step_c))[0]) == pytest.approx(0.0, abs=1.0e-8) + + +def test__log_evidence_lm_from__matches_hand_computed(): + data, noise, mapping, src_reg, dpsi_reg = joint_problem() + n_src = src_reg.shape[0] + L = mapping[:, :n_src] + J_dpsi = mapping[:, n_src:] + + rng = np.random.default_rng(6) + s = rng.normal(size=n_src) + dpsi = rng.normal(size=dpsi_reg.shape[0]) + + result = dense_util.log_evidence_lm_from( + data, noise, s, dpsi, L, L, J_dpsi, src_reg, dpsi_reg + ) + + inv_var = 1.0 / noise**2 + residual = data - L @ s + chi2 = np.sum(inv_var * residual**2) + reg_val = s @ src_reg @ s + dpsi @ dpsi_reg @ dpsi + J = np.hstack([L, J_dpsi]) + R = scipy_block_diag(src_reg, dpsi_reg) + H = J.T @ (J * inv_var[:, None]) + R + expected = 0.5 * ( + np.linalg.slogdet(src_reg)[1] + + np.linalg.slogdet(dpsi_reg)[1] + - np.linalg.slogdet(H)[1] + - chi2 + - reg_val + ) + + assert float(result) == pytest.approx(expected, rel=1.0e-10) diff --git a/test_autolens/potential_correction/test_iterative.py b/test_autolens/potential_correction/test_iterative.py new file mode 100644 index 000000000..ffaaecf8b --- /dev/null +++ b/test_autolens/potential_correction/test_iterative.py @@ -0,0 +1,87 @@ +import numpy as np +import pytest + +import autoarray as aa +import autolens as al + + +def iter_fit_from(masked_imaging, gauge_constraints=False, n_iter=2): + lens = al.Galaxy( + redshift=0.5, + mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0), + ) + 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), + ) + return al.pc.IterFitDpsiSrcImaging( + masked_imaging=masked_imaging, + lens_start=lens, + dpsi_pixelization=dpsi_pixelization, + src_pixelization=src_pixelization, + gauge_constraints=gauge_constraints, + n_iter=n_iter, + ) + + +def test__solve_joint_optimization__returns_finite_state_and_decreases_cost( + masked_imaging_7x7, +): + fit = iter_fit_from(masked_imaging_7x7) + + s_opt, dpsi_opt = fit.solve_joint_optimization() + + n_dpsi = np.count_nonzero(~fit.pair_dpsi_data_obj.mask_dpsi) + assert s_opt.shape == (fit.src_reg_mat.shape[0],) + assert dpsi_opt.shape == (n_dpsi,) + assert np.isfinite(s_opt).all() + assert np.isfinite(dpsi_opt).all() + + # the optimized state must beat the zero starting state + data = np.asarray(fit.masked_imaging.data.slim) + inv_var = fit.inverse_noise_variance + L, _, _ = fit.get_L_Js_Jdpsi(s_opt, dpsi_opt) + from autolens.potential_correction import dense_util + + cost_opt, *_ = dense_util.lm_cost_from( + data, inv_var, s_opt, dpsi_opt, L, fit.src_reg_mat, + fit.dpsi_regularization_matrix, + ) + cost_zero = 0.5 * float(np.sum(inv_var * data**2)) + assert float(cost_opt) < cost_zero + + +def test__solve_joint_optimization__gauge_constraints_are_satisfied( + masked_imaging_7x7, +): + fit = iter_fit_from(masked_imaging_7x7, gauge_constraints=True) + + s_opt, dpsi_opt = fit.solve_joint_optimization() + + n_dpsi = dpsi_opt.shape[0] + G = np.zeros((3, n_dpsi)) + G[0, :] = 1.0 / n_dpsi + G[1, :] = fit.dpsi_points[:, 1] / n_dpsi + G[2, :] = fit.dpsi_points[:, 0] / n_dpsi + + assert G @ dpsi_opt == pytest.approx(np.zeros(3), abs=1.0e-6) + + +def test__log_evidence__finite_at_optimum(masked_imaging_7x7): + fit = iter_fit_from(masked_imaging_7x7) + + fit.solve_joint_optimization() + evidence = fit.log_evidence() + + assert np.isfinite(evidence) + + +def test__log_evidence__requires_state_or_solve(masked_imaging_7x7): + fit = iter_fit_from(masked_imaging_7x7) + + with pytest.raises(ValueError): + fit.log_evidence() diff --git a/test_autolens/potential_correction/test_util.py b/test_autolens/potential_correction/test_util.py index 187a12cf5..6668fc812 100644 --- a/test_autolens/potential_correction/test_util.py +++ b/test_autolens/potential_correction/test_util.py @@ -165,6 +165,19 @@ def test__split_cross_from__cross_layout_and_positive_lengths(): assert (split[:, 0, 0] > points[:, 0]).all() +def test__split_cross_from__all_unbounded_cells_yield_finite_arms(): + # with 4 points every Voronoi cell is unbounded; the arm length must fall + # back to the nearest-neighbour spacing rather than go NaN (the original + # implementation took a percentile over the -1 sentinels) + points = np.array([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]]) + + split = pc_util.split_cross_from(points) + + assert np.isfinite(split).all() + split = split.reshape(4, 4, 2) + 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