diff --git a/autoarray/config/general.yaml b/autoarray/config/general.yaml index 02cbd3950..2577bd21e 100644 --- a/autoarray/config/general.yaml +++ b/autoarray/config/general.yaml @@ -10,6 +10,7 @@ inversion: nnls_target_kappa: 1.0e-11 # Central-path relaxation parameter passed to jaxnnls.solve_nnls_primal. Larger values smooth the relaxed-KKT backward pass and prevent NaN gradients on ill-conditioned Q; smaller values tighten the primal solve. Verified finite gradients across all MGE/rectangular/delaunay pipelines (imaging + interferometer) with scale invariance over 5 orders of magnitude in noise. jaxnnls's own default (1e-3) is too aggressive for the backward pass. reconstruction_vmax_factor: 0.5 # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor. log_det_method: cholesky # How the Bayesian-evidence log-determinant terms are computed. "cholesky" (default) is the historical 2*sum(log(diag(cholesky(M)))); "slogdet" uses logabsdet of slogdet(M), which is identical where M is positive-definite but finite (not NaN) where the Cholesky fails, for gradient-based searches (opt-in, non-default; does not change the default evidence). Under "slogdet" the kernel regularization schemes (Matern/Gaussian/Exponential) also compute the regularization log-det analytically from a Cholesky of their covariance instead of factorizing the formed inverse. See PyAutoArray#391. + regularization_term_method: matmul # How the Bayesian-evidence regularization term s^T H s is computed. "matmul" (default) is the historical s @ (H @ s) against the explicitly formed regularization matrix; "cho_solve" evaluates coefficient * s^T C^-1 s for the kernel schemes (Matern/Gaussian/Exponential/MaternAdapt) via one Cholesky solve of their covariance C, avoiding the explicit inverse whose round-off is amplified by cond(C) (~1e9 on clustered traced mesh vertices). Opt-in, non-default; does not change the default evidence. Schemes with no such factorization fall back to the formed matrix. numba: use_numba: true cache: true diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py index 498f23f05..eb1291ce0 100644 --- a/autoarray/inversion/inversion/abstract.py +++ b/autoarray/inversion/inversion/abstract.py @@ -691,10 +691,50 @@ def regularization_term(self) -> float: The above works include the regularization_matrix coefficient (lambda) in this calculation. In PyAutoLens, this is already in the regularization matrix and thus implicitly included in the matrix multiplication. + + Under ``regularization_term_method == "cho_solve"`` (opt-in, default off), regularization schemes + which know a factorization of their own matrix may instead supply their contribution directly via + :meth:`AbstractRegularization.regularization_term_from` — the kernel schemes (``MaternKernel`` etc.) + return ``coefficient * s^T C^-1 s`` from a single Cholesky solve of their covariance ``C``, avoiding + the round-off of contracting the explicitly formed inverse (whose error is amplified by ``cond(C)``, + ~1e9 on clustered traced mesh vertices). Because ``regularization_matrix_reduced`` is the block + diagonal of the per-object matrices when every linear object is regularized, the term is the sum of + the per-object terms; if any scheme has no shortcut (returns ``None``) the whole computation falls + back to the formed matrix. The default ``"matmul"`` path never consults the shortcut, so default + evidence values are unchanged. + + Returns + ------- + float + The regularization term of the inversion. """ if not self.has(cls=AbstractRegularization): return 0.0 + if ( + self.settings.regularization_term_method == "cho_solve" + and self.all_linear_obj_have_regularization + ): + # `reconstruction_reduced` is the full reconstruction here (the guard above is exactly the + # no-reduction case), so the per-object slices index it directly. + reconstruction = self.reconstruction_reduced + + term_list = [ + regularization.regularization_term_from( + linear_obj=linear_obj, + reconstruction=reconstruction[param_range[0] : param_range[1]], + xp=self._xp, + ) + for linear_obj, regularization, param_range in zip( + self.linear_obj_list, + self.regularization_list, + self.param_range_list_from(cls=LinearObj), + ) + ] + + if all(term is not None for term in term_list): + return sum(term_list) + return self._xp.matmul( self.reconstruction_reduced.T, self._xp.matmul( diff --git a/autoarray/inversion/regularization/abstract.py b/autoarray/inversion/regularization/abstract.py index 04ad30f84..c149f1966 100644 --- a/autoarray/inversion/regularization/abstract.py +++ b/autoarray/inversion/regularization/abstract.py @@ -213,3 +213,47 @@ def log_det_regularization_matrix_term_from( has no factorization-aware shortcut. """ return None + + def regularization_term_from( + self, linear_obj: LinearObj, reconstruction: np.ndarray, xp=np + ) -> Optional[float]: + """ + Returns this scheme's contribution to the regularization term ``s^T H s`` + computed from a factorization the scheme itself knows about, or ``None`` when + no such shortcut exists (the default). + + This is the ``s^T H s`` counterpart of + :meth:`log_det_regularization_matrix_term_from`, and exists for the same + reason. The kernel regularization schemes build ``H = coefficient * C^-1`` + from a dense covariance ``C``, so their term is + ``coefficient * s^T C^-1 s`` — obtainable from a single Cholesky *solve* + against ``s`` rather than by forming ``C^-1`` and contracting it. Forming the + explicit inverse carries round-off amplified by ``cond(C)`` (~1e9 on the + clustered traced vertices of the kNN mesh families), which then enters the + evidence through this term. + + Note this cannot remove the explicit inverse from the inversion altogether: + ``curvature_reg_matrix`` is a dense ``F + H`` feeding the dense solve for the + reconstruction, so ``H`` is still formed there. This shortcut removes the + formed inverse from the *evidence* terms only. + + The inversion consumes this ONLY when + ``Settings.regularization_term_method == "cho_solve"`` — the default + ``"matmul"`` path never calls it, so default likelihood values are unchanged. + See ``AbstractInversion.regularization_term``. + + Parameters + ---------- + linear_obj + The linear object (e.g. a ``Mapper``) whose regularization matrix the + term is of. + reconstruction + The reconstructed values ``s`` of this linear object's parameters (the + slice of the inversion's reconstruction belonging to ``linear_obj``). + + Returns + ------- + The scalar ``s^T H s`` for this linear object, or ``None`` when this scheme + has no factorization-aware shortcut. + """ + return None diff --git a/autoarray/inversion/regularization/exponential_kernel.py b/autoarray/inversion/regularization/exponential_kernel.py index 194f384c9..ec457b395 100644 --- a/autoarray/inversion/regularization/exponential_kernel.py +++ b/autoarray/inversion/regularization/exponential_kernel.py @@ -12,6 +12,7 @@ def exp_cov_matrix_from( scale: float, pixel_points: np.ndarray, # shape (N, 2) jitter: float = 1e-8, + jitter_relative: bool = False, xp=np, ) -> np.ndarray: # shape (N, N) """ @@ -54,7 +55,9 @@ def exp_cov_matrix_from( # add a small jitter on the diagonal N = pts.shape[0] - cov = cov + xp.eye(N, dtype=cov.dtype) * jitter + from autoarray.inversion.regularization.matern_kernel import apply_jitter + + cov = apply_jitter(cov, jitter=jitter, jitter_relative=jitter_relative, xp=xp) return cov @@ -65,6 +68,7 @@ def __init__( coefficient: float = 1.0, scale: float = 1.0, jitter: Optional[float] = None, + jitter_relative: bool = False, ): """ Regularization which uses an Exponential smoothing kernel to regularize the solution. @@ -99,10 +103,17 @@ def __init__( ``None`` (default) uses the historical value 1e-8 — behaviour is identical to not having this parameter (it is a fixed setting, not a free model parameter, hence the ``None`` default). + jitter_relative + If ``True`` the jitter is applied *relative* to each pixel's own variance + (``C_ii *= 1 + jitter``) rather than as a fixed absolute ``jitter * I``. + ``False`` (default) preserves the historical behaviour exactly. The absolute + convention assumes ``C_ii ~ 1``, which holds for this unweighted kernel but not + for the adaptive one; see :func:`apply_jitter` for why and when to switch. """ self.coefficient = coefficient self.scale = scale self.jitter = jitter + self.jitter_relative = jitter_relative super().__init__() @@ -150,6 +161,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray scale=self.scale, pixel_points=linear_obj.source_plane_mesh_grid.array, jitter=self.jitter_value, + jitter_relative=self.jitter_relative, xp=xp, ) @@ -173,6 +185,7 @@ def log_det_regularization_matrix_term_from( scale=self.scale, pixel_points=linear_obj.source_plane_mesh_grid.array, jitter=self.jitter_value, + jitter_relative=self.jitter_relative, xp=xp, ) @@ -181,3 +194,33 @@ def log_det_regularization_matrix_term_from( ) return linear_obj.params * np.log(self.coefficient) - log_det_covariance + + def regularization_term_from( + self, linear_obj: LinearObj, reconstruction: np.ndarray, xp=np + ) -> float: + """ + The regularization term ``s^T H s`` from a single Cholesky solve of the kernel + covariance: ``H = coefficient * C^-1``, so + ``s^T H s = coefficient * s^T C^-1 s``, with the quadratic form evaluated by + solving ``C x = s`` rather than by forming ``C^-1``. + + Consumed by the inversion only when + ``Settings.regularization_term_method == "cho_solve"`` (see + :meth:`AbstractRegularization.regularization_term_from`); the default + ``"matmul"`` path contracts the formed ``H`` and is unchanged. + """ + from autoarray.inversion.regularization.matern_kernel import ( + quadratic_form_via_cholesky, + ) + + covariance_matrix = exp_cov_matrix_from( + scale=self.scale, + pixel_points=linear_obj.source_plane_mesh_grid.array, + jitter=self.jitter_value, + jitter_relative=self.jitter_relative, + xp=xp, + ) + + return self.coefficient * quadratic_form_via_cholesky( + covariance_matrix, reconstruction, xp=xp + ) diff --git a/autoarray/inversion/regularization/gaussian_kernel.py b/autoarray/inversion/regularization/gaussian_kernel.py index 710915e56..ce9a94231 100644 --- a/autoarray/inversion/regularization/gaussian_kernel.py +++ b/autoarray/inversion/regularization/gaussian_kernel.py @@ -12,6 +12,7 @@ def gauss_cov_matrix_from( scale: float, pixel_points: np.ndarray, # shape (N, 2) jitter: float = 1e-8, + jitter_relative: bool = False, xp=np, ) -> np.ndarray: """ @@ -46,7 +47,9 @@ def gauss_cov_matrix_from( # Add tiny jitter on the diagonal N = pts.shape[0] - cov = cov + xp.eye(N, dtype=cov.dtype) * jitter + from autoarray.inversion.regularization.matern_kernel import apply_jitter + + cov = apply_jitter(cov, jitter=jitter, jitter_relative=jitter_relative, xp=xp) return cov @@ -57,6 +60,7 @@ def __init__( coefficient: float = 1.0, scale: float = 1.0, jitter: Optional[float] = None, + jitter_relative: bool = False, ): """ Regularization which uses a Gaussian smoothing kernel to regularize the solution. @@ -89,10 +93,17 @@ def __init__( ``None`` (default) uses the historical value 1e-8 — behaviour is identical to not having this parameter (it is a fixed setting, not a free model parameter, hence the ``None`` default). + jitter_relative + If ``True`` the jitter is applied *relative* to each pixel's own variance + (``C_ii *= 1 + jitter``) rather than as a fixed absolute ``jitter * I``. + ``False`` (default) preserves the historical behaviour exactly. The absolute + convention assumes ``C_ii ~ 1``, which holds for this unweighted kernel but not + for the adaptive one; see :func:`apply_jitter` for why and when to switch. """ self.coefficient = coefficient self.scale = scale self.jitter = jitter + self.jitter_relative = jitter_relative super().__init__() @property @@ -139,6 +150,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray scale=self.scale, pixel_points=linear_obj.source_plane_mesh_grid.array, jitter=self.jitter_value, + jitter_relative=self.jitter_relative, xp=xp, ) @@ -157,9 +169,10 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray N = regularization_matrix.shape[0] diag_mean = xp.mean(xp.diag(regularization_matrix)) h_jitter = 1e-8 * xp.abs(diag_mean) - regularization_matrix = regularization_matrix + xp.eye( - N, dtype=regularization_matrix.dtype - ) * h_jitter + regularization_matrix = ( + regularization_matrix + + xp.eye(N, dtype=regularization_matrix.dtype) * h_jitter + ) return regularization_matrix @@ -185,6 +198,7 @@ def log_det_regularization_matrix_term_from( scale=self.scale, pixel_points=linear_obj.source_plane_mesh_grid.array, jitter=self.jitter_value, + jitter_relative=self.jitter_relative, xp=xp, ) @@ -193,3 +207,39 @@ def log_det_regularization_matrix_term_from( ) return linear_obj.params * np.log(self.coefficient) - log_det_covariance + + def regularization_term_from( + self, linear_obj: LinearObj, reconstruction: np.ndarray, xp=np + ) -> float: + """ + The regularization term ``s^T H s`` from a single Cholesky solve of the kernel + covariance: ``H = coefficient * C^-1``, so + ``s^T H s = coefficient * s^T C^-1 s``, with the quadratic form evaluated by + solving ``C x = s`` rather than by forming ``C^-1``. + + As with :meth:`log_det_regularization_matrix_term_from`, this is the term of + the analytic ``coefficient * C^-1``: it excludes both the symmetrisation and + the trace-scaled stabilisation jitter that :meth:`regularization_matrix_from` + applies to the formed matrix, since both exist only to guard the + factorization of the explicit inverse that this shortcut avoids entirely. + + Consumed by the inversion only when + ``Settings.regularization_term_method == "cho_solve"`` (see + :meth:`AbstractRegularization.regularization_term_from`); the default + ``"matmul"`` path contracts the formed ``H`` and is unchanged. + """ + from autoarray.inversion.regularization.matern_kernel import ( + quadratic_form_via_cholesky, + ) + + covariance_matrix = gauss_cov_matrix_from( + scale=self.scale, + pixel_points=linear_obj.source_plane_mesh_grid.array, + jitter=self.jitter_value, + jitter_relative=self.jitter_relative, + xp=xp, + ) + + return self.coefficient * quadratic_form_via_cholesky( + covariance_matrix, reconstruction, xp=xp + ) diff --git a/autoarray/inversion/regularization/matern_adapt_kernel.py b/autoarray/inversion/regularization/matern_adapt_kernel.py index efb495cc8..3a1f7acfc 100644 --- a/autoarray/inversion/regularization/matern_adapt_kernel.py +++ b/autoarray/inversion/regularization/matern_adapt_kernel.py @@ -10,6 +10,9 @@ from autoarray.inversion.regularization.matern_kernel import matern_kernel from autoarray.inversion.regularization.matern_kernel import matern_cov_matrix_from from autoarray.inversion.regularization.matern_kernel import inv_via_cholesky +from autoarray.inversion.regularization.matern_kernel import ( + quadratic_form_via_cholesky, +) from autoarray.inversion.regularization.adapt import adapt_regularization_weights_from @@ -22,6 +25,7 @@ def __init__( outer_coefficient: float = 1.0, signal_scale: float = 1.0, jitter: Optional[float] = None, + jitter_relative: bool = False, ): """ Regularization which uses a Matern smoothing kernel to regularize the solution with regularization weights @@ -63,8 +67,25 @@ def __init__( Controls how strongly the kernel weights adapt to pixel brightness. Larger values make bright pixels receive significantly higher weights (and faint pixels lower weights), while smaller values produce a more uniform weighting. Typical values are of order unity (e.g. 0.5–2.0). + jitter + The small value added to the covariance diagonal for numerical stability. + ``None`` (default) uses the historical value 1e-8. + jitter_relative + If ``True`` the jitter is applied *relative* to each pixel's own variance + (``C_ii *= 1 + jitter``) rather than as a fixed absolute ``jitter * I``. + ``False`` (default) preserves the historical behaviour exactly. **This scheme is + the one the absolute convention breaks down on**: its ``C_ii = w_i^2`` spans the + adaptive-weight dynamic range, so at wide inner/outer coefficients a fixed + ``1e-8`` can reach 100% of a faint pixel's variance, destroying its kernel + structure. See :func:`apply_jitter`. """ - super().__init__(coefficient=0.0, scale=scale, nu=nu, jitter=jitter) + super().__init__( + coefficient=0.0, + scale=scale, + nu=nu, + jitter=jitter, + jitter_relative=jitter_relative, + ) self.inner_coefficient = inner_coefficient self.outer_coefficient = outer_coefficient self.signal_scale = signal_scale @@ -111,6 +132,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray nu=self.nu, weights=kernel_weights, jitter=self.jitter_value, + jitter_relative=self.jitter_relative, xp=xp, ) @@ -138,7 +160,40 @@ def log_det_regularization_matrix_term_from( nu=self.nu, weights=kernel_weights, jitter=self.jitter_value, + jitter_relative=self.jitter_relative, xp=xp, ) return -2.0 * xp.sum(xp.log(xp.diag(xp.linalg.cholesky(covariance_matrix)))) + + def regularization_term_from( + self, linear_obj: LinearObj, reconstruction: np.ndarray, xp=np + ) -> float: + """ + The regularization term ``s^T H s`` from a single Cholesky solve of the + weighted kernel covariance. This scheme's ``H = C_w^-1`` carries no + coefficient scaling (the adaptive weights are inside ``C_w``), so the term is + ``s^T C_w^-1 s`` — which is why this override is needed rather than the + inherited :class:`MaternKernel` one, whose ``coefficient`` is fixed at 0.0 + here. + + Consumed by the inversion only when + ``Settings.regularization_term_method == "cho_solve"`` (see + :meth:`AbstractRegularization.regularization_term_from`); the default + ``"matmul"`` path contracts the formed ``H`` and is unchanged. + """ + kernel_weights = 1.0 / self.regularization_weights_from( + linear_obj=linear_obj, xp=xp + ) + + covariance_matrix = matern_cov_matrix_from( + scale=self.scale, + pixel_points=linear_obj.source_plane_mesh_grid.array, + nu=self.nu, + weights=kernel_weights, + jitter=self.jitter_value, + jitter_relative=self.jitter_relative, + xp=xp, + ) + + return quadratic_form_via_cholesky(covariance_matrix, reconstruction, xp=xp) diff --git a/autoarray/inversion/regularization/matern_kernel.py b/autoarray/inversion/regularization/matern_kernel.py index ca6d6e0f2..ffa61ab3d 100644 --- a/autoarray/inversion/regularization/matern_kernel.py +++ b/autoarray/inversion/regularization/matern_kernel.py @@ -90,6 +90,7 @@ def matern_cov_matrix_from( pixel_points, weights=None, jitter: float = 1e-8, + jitter_relative: bool = False, xp=np, ): """ @@ -112,6 +113,9 @@ def matern_cov_matrix_from( Optional array-like of shape [N]. If None, treated as all ones. jitter The small value added to the covariance diagonal for numerical stability. + jitter_relative + If ``True`` the jitter is scaled by each pixel's own variance (``C_ii *= 1 + jitter``) + rather than added as a fixed absolute value; see :func:`apply_jitter`. xp Backend (numpy or jax.numpy). @@ -148,12 +152,58 @@ def matern_cov_matrix_from( # -------------------------------- # Add diagonal jitter (JAX-safe) # -------------------------------- - pixels = pts.shape[0] - covariance_matrix = covariance_matrix + jitter * xp.eye( - pixels, dtype=covariance_matrix.dtype + return apply_jitter( + covariance_matrix, jitter=jitter, jitter_relative=jitter_relative, xp=xp ) - return covariance_matrix + +def apply_jitter( + covariance_matrix, jitter: float, jitter_relative: bool = False, xp=np +): + """ + Add the stabilisation jitter to a kernel covariance's diagonal. + + Two conventions, selected by ``jitter_relative``: + + - ``False`` (default, historical) — an **absolute** ``jitter * I``. This is only + meaningful when the covariance's diagonal is itself ~1, which holds for the + unweighted kernels (``K(0) == 1``) but *not* for the weighted ones, where + ``C_ii = w_i^2`` spans the dynamic range of the adaptive weights. + - ``True`` — a **relative** ``jitter * diag(diag(C))``, i.e. ``C_ii *= (1 + jitter)``. + Writing ``C = D^(1/2) R D^(1/2)`` for the correlation matrix ``R`` (unit diagonal), + this is exactly ``D^(1/2) (R + jitter I) D^(1/2)`` — the jitter is applied to the + *correlation* matrix, so every pixel receives the same relative protection whatever + its own scale, and the conditioning of the correlation part is bounded exactly as in + the absolute convention. + + The absolute convention silently breaks down on ``MaternAdaptKernel`` at wide + adaptive-weight ranges: with ``inner_coefficient=0.1``/``outer_coefficient=100`` the + faintest pixel's ``C_ii`` reaches ``1e-8``, at which point a fixed ``1e-8`` jitter is + 100% of that pixel's variance and its kernel structure is destroyed. Under the + relative convention that pixel is perturbed by ``1e-8`` relative, like every other. + + Parameters + ---------- + covariance_matrix + The (N, N) kernel covariance, before jitter. + jitter + The jitter magnitude — absolute or relative per the flag below. + jitter_relative + If ``True`` scale the jitter by each pixel's own variance; if ``False`` (default) + add it as a fixed absolute value. + xp + Backend (numpy or jax.numpy). + + Returns + ------- + The covariance matrix with the jitter added to its diagonal. + """ + if jitter_relative: + return covariance_matrix + jitter * xp.diag(xp.diag(covariance_matrix)) + + pixels = covariance_matrix.shape[0] + + return covariance_matrix + jitter * xp.eye(pixels, dtype=covariance_matrix.dtype) def inv_via_cholesky(C, xp=np): @@ -173,6 +223,44 @@ def inv_via_cholesky(C, xp=np): return jla.cho_solve((L, True), I) +def quadratic_form_via_cholesky(C, s, xp=np): + """ + The quadratic form ``s^T C^-1 s`` evaluated through a single Cholesky solve. + + This is the implicit counterpart of ``s @ inv_via_cholesky(C) @ s``: it never + forms ``C^-1``, solving ``C x = s`` for the single right-hand side ``s`` instead + of the ``N`` right-hand sides of the identity. Both accuracy and cost improve — + the explicit inverse's round-off is amplified by ``cond(C)``, which reaches ~1e9 + on the clustered traced vertices of the kNN mesh families, and it performs ``N`` + triangular solves where one suffices. + + Parameters + ---------- + C + The (N, N) symmetric positive-definite kernel covariance matrix. + s + The (N,) vector the quadratic form is taken over (the reconstruction). + xp + Backend (numpy or jax.numpy). + + Returns + ------- + The scalar ``s^T C^-1 s``. + """ + # NumPy + if xp is np: + import scipy.linalg as la + + cho = la.cho_factor(C, lower=True, check_finite=False) + return s @ la.cho_solve(cho, s, check_finite=False) + + # JAX + import jax.scipy.linalg as jla + + L = xp.linalg.cholesky(C) + return s @ jla.cho_solve((L, True), s) + + class MaternKernel(AbstractRegularization): def __init__( self, @@ -180,6 +268,7 @@ def __init__( scale: float = 1.0, nu: float = 0.5, jitter: Optional[float] = None, + jitter_relative: bool = False, ): """ Regularization which uses a Matern smoothing kernel to regularize the solution. @@ -221,12 +310,19 @@ def __init__( ``None`` (default) uses the historical value 1e-8 — behaviour is identical to not having this parameter (it is a fixed setting, not a free model parameter, hence the ``None`` default). + jitter_relative + If ``True`` the jitter is applied *relative* to each pixel's own variance + (``C_ii *= 1 + jitter``) rather than as a fixed absolute ``jitter * I``. + ``False`` (default) preserves the historical behaviour exactly. The absolute + convention assumes ``C_ii ~ 1``, which holds for this unweighted kernel but not + for the adaptive one; see :func:`apply_jitter` for why and when to switch. """ self.coefficient = coefficient self.scale = scale self.nu = nu self.jitter = jitter + self.jitter_relative = jitter_relative super().__init__() @property @@ -272,6 +368,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray pixel_points=linear_obj.source_plane_mesh_grid.array, nu=self.nu, jitter=self.jitter_value, + jitter_relative=self.jitter_relative, xp=xp, ) @@ -294,6 +391,7 @@ def log_det_regularization_matrix_term_from( pixel_points=linear_obj.source_plane_mesh_grid.array, nu=self.nu, jitter=self.jitter_value, + jitter_relative=self.jitter_relative, xp=xp, ) @@ -302,3 +400,30 @@ def log_det_regularization_matrix_term_from( ) return linear_obj.params * np.log(self.coefficient) - log_det_covariance + + def regularization_term_from( + self, linear_obj: LinearObj, reconstruction: np.ndarray, xp=np + ) -> float: + """ + The regularization term ``s^T H s`` from a single Cholesky solve of the kernel + covariance: ``H = coefficient * C^-1``, so + ``s^T H s = coefficient * s^T C^-1 s``, with the quadratic form evaluated by + solving ``C x = s`` rather than by forming ``C^-1``. + + Consumed by the inversion only when + ``Settings.regularization_term_method == "cho_solve"`` (see + :meth:`AbstractRegularization.regularization_term_from`); the default + ``"matmul"`` path contracts the formed ``H`` and is unchanged. + """ + covariance_matrix = matern_cov_matrix_from( + scale=self.scale, + pixel_points=linear_obj.source_plane_mesh_grid.array, + nu=self.nu, + jitter=self.jitter_value, + jitter_relative=self.jitter_relative, + xp=xp, + ) + + return self.coefficient * quadratic_form_via_cholesky( + covariance_matrix, reconstruction, xp=xp + ) diff --git a/autoarray/settings.py b/autoarray/settings.py index 392c4e558..29c1116e2 100644 --- a/autoarray/settings.py +++ b/autoarray/settings.py @@ -18,6 +18,7 @@ def __init__( nnls_solver_tol: Optional[float] = None, nnls_max_iter: Optional[int] = None, log_det_method: Optional[str] = None, + regularization_term_method: Optional[str] = None, ): """ The settings of an Inversion, customizing how a linear set of equations are solved for. @@ -118,6 +119,32 @@ def __init__( ``cond(C)``, ~1e-6 absolute in the evidence on clustered traced mesh vertices) and finite at any regularization coefficient (``C``'s conditioning does not depend on it). See :meth:`AbstractRegularization.log_det_regularization_matrix_term_from`. + regularization_term_method + How the Bayesian-evidence regularization term ``s^T H s`` + (``AbstractInversion.regularization_term``) is computed. `None` (default) reads the + packaged value ``"matmul"``. + + - ``"matmul"`` (default) — ``s @ (H @ s)`` against the explicitly formed regularization + matrix, the historical computation. + - ``"cho_solve"`` — for the kernel regularization schemes (``MaternKernel``, + ``GaussianKernel``, ``ExponentialKernel``, ``MaternAdaptKernel``), whose + ``H = coefficient * C^-1``, evaluate ``coefficient * s^T C^-1 s`` by solving + ``C x = s`` through one Cholesky of the kernel covariance ``C`` instead of forming + ``C^-1`` and contracting it. More accurate — the explicit inverse's round-off is + amplified by ``cond(C)``, which reaches ~1e9 on the clustered traced vertices of the + kNN mesh families — and cheaper, being one triangular solve rather than ``N``. This + is an **opt-in, non-default** alternative; schemes with no such factorization + (``Constant``, ``Adapt``, the split families) have no shortcut and fall back to the + formed matrix, so mixed inversions stay correct. See + :meth:`AbstractRegularization.regularization_term_from`. + + This is deliberately separate from ``log_det_method``: the two terms can be moved + onto their exact factorizations independently, which is what makes it possible to + attribute an evidence shift to one term rather than both. + + Note neither option removes the explicit inverse from the inversion as a whole — + ``curvature_reg_matrix`` is a dense ``F + H`` feeding the dense solve for the + reconstruction, so ``H`` is still formed there regardless. """ self.use_mixed_precision = use_mixed_precision self.nnls_solver_tol = nnls_solver_tol @@ -129,6 +156,7 @@ def __init__( no_regularization_add_to_curvature_diag_value ) self._log_det_method = log_det_method + self._regularization_term_method = regularization_term_method @property def use_positive_only_solver(self): @@ -166,3 +194,10 @@ def log_det_method(self): return conf.instance["general"]["inversion"]["log_det_method"] return self._log_det_method + + @property + def regularization_term_method(self): + if self._regularization_term_method is None: + return conf.instance["general"]["inversion"]["regularization_term_method"] + + return self._regularization_term_method diff --git a/test_autoarray/inversion/regularizations/test_kernel_jax_gradients.py b/test_autoarray/inversion/regularizations/test_kernel_jax_gradients.py new file mode 100644 index 000000000..8d9813599 --- /dev/null +++ b/test_autoarray/inversion/regularizations/test_kernel_jax_gradients.py @@ -0,0 +1,163 @@ +""" +JAX leg of the gate for the kernel-scheme linear-algebra work: the +`quadratic_form_via_cholesky` term shortcut and the `apply_jitter` conventions must be +`jit`-safe, differentiable, and agree with the NumPy path. + +These mirror, at library level, what +`autolens_workspace_test/scripts/imaging/jax_grad/regularization.py` certifies at +workspace level. Skipped when JAX is absent (it is an optional dependency). + +The Matern kernel is deliberately not covered here: its JAX path needs the modified +Bessel `K_nu` from `tfp-nightly`, which is an optional-of-an-optional. The Gaussian and +Exponential kernels exercise the same shared code (`apply_jitter`, +`quadratic_form_via_cholesky`) without it. +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +jnp = pytest.importorskip("jax.numpy") + +jax.config.update("jax_enable_x64", True) + +from autoarray.inversion.regularization.matern_kernel import ( # noqa: E402 + apply_jitter, + inv_via_cholesky, + quadratic_form_via_cholesky, +) +from autoarray.inversion.regularization.gaussian_kernel import ( # noqa: E402 + gauss_cov_matrix_from, +) +from autoarray.inversion.regularization.exponential_kernel import ( # noqa: E402 + exp_cov_matrix_from, +) + +POINTS = np.random.default_rng(0).normal(size=(10, 2)) +VECTOR = np.random.default_rng(1).normal(size=10) + + +@pytest.mark.parametrize("jitter_relative", [False, True]) +@pytest.mark.parametrize("cov_from", [gauss_cov_matrix_from, exp_cov_matrix_from]) +def test__apply_jitter__is_jit_safe_in_both_conventions(jitter_relative, cov_from): + """ + `jitter_relative` is a static Python bool, so the branch inside `apply_jitter` must + resolve at trace time rather than on a tracer. + """ + + @jax.jit + def build(points): + covariance = cov_from(scale=1.0, pixel_points=points, jitter=0.0, xp=jnp) + return apply_jitter( + covariance, jitter=1e-8, jitter_relative=jitter_relative, xp=jnp + ) + + covariance = build(jnp.asarray(POINTS)) + + assert covariance.shape == (POINTS.shape[0], POINTS.shape[0]) + assert bool(jnp.all(jnp.isfinite(covariance))) + + +@pytest.mark.parametrize("cov_from", [gauss_cov_matrix_from, exp_cov_matrix_from]) +def test__quadratic_form_via_cholesky__gradient_is_finite_difference_certified( + cov_from, +): + """ + The certification the workspace `jax_grad/regularization.py` script exists to give: + autodiff through the Cholesky solve must match a central finite difference. + """ + + def term(scale): + covariance = cov_from( + scale=scale, pixel_points=jnp.asarray(POINTS), jitter=1e-8, xp=jnp + ) + return quadratic_form_via_cholesky(covariance, jnp.asarray(VECTOR), xp=jnp) + + autodiff = float(jax.grad(term)(1.3)) + + step = 1.0e-6 + finite_difference = float((term(1.3 + step) - term(1.3 - step)) / (2.0 * step)) + + assert autodiff == pytest.approx(finite_difference, rel=1.0e-6) + + +@pytest.mark.parametrize("cov_from", [gauss_cov_matrix_from, exp_cov_matrix_from]) +def test__quadratic_form_via_cholesky__eager_and_jit_agree(cov_from): + def term(scale): + covariance = cov_from( + scale=scale, pixel_points=jnp.asarray(POINTS), jitter=1e-8, xp=jnp + ) + return quadratic_form_via_cholesky(covariance, jnp.asarray(VECTOR), xp=jnp) + + assert float(jax.jit(term)(1.3)) == pytest.approx(float(term(1.3)), rel=1.0e-12) + + +@pytest.mark.parametrize("cov_from", [gauss_cov_matrix_from, exp_cov_matrix_from]) +def test__quadratic_form_via_cholesky__matches_explicit_inverse_and_numpy(cov_from): + """ + The shortcut must be the same quantity as the explicit-inverse contraction it + replaces, on both backends. + """ + covariance_jax = cov_from( + scale=1.3, pixel_points=jnp.asarray(POINTS), jitter=1e-8, xp=jnp + ) + covariance_numpy = cov_from(scale=1.3, pixel_points=POINTS, jitter=1e-8) + + implicit = float( + quadratic_form_via_cholesky(covariance_jax, jnp.asarray(VECTOR), xp=jnp) + ) + explicit = float( + jnp.asarray(VECTOR) + @ (inv_via_cholesky(covariance_jax, xp=jnp) @ jnp.asarray(VECTOR)) + ) + numpy_implicit = float(quadratic_form_via_cholesky(covariance_numpy, VECTOR)) + + assert implicit == pytest.approx(explicit, rel=1.0e-9) + assert implicit == pytest.approx(numpy_implicit, rel=1.0e-9) + + +def test__regularization_term_from__is_differentiable_end_to_end_under_jax(): + """ + The scheme-level hook, not just the helper: gradients must reach the mesh grid the + term is built from. + """ + import autoarray as aa + + regularization = aa.reg.GaussianKernel(coefficient=3.0, scale=1.0) + + class _Obj: + def __init__(self, array): + self.source_plane_mesh_grid = type("_G", (), {"array": array})() + + def term(points): + return regularization.regularization_term_from( + linear_obj=_Obj(points), reconstruction=jnp.asarray(VECTOR), xp=jnp + ) + + value = float(term(jnp.asarray(POINTS))) + gradient = jax.grad(term)(jnp.asarray(POINTS)) + + assert np.isfinite(value) + assert bool(jnp.all(jnp.isfinite(gradient))) + assert float(jnp.linalg.norm(gradient)) > 0.0 + + +@pytest.mark.parametrize("jitter_relative", [False, True]) +def test__relative_jitter__does_not_break_gradients(jitter_relative): + def term(scale): + covariance = gauss_cov_matrix_from( + scale=scale, + pixel_points=jnp.asarray(POINTS), + jitter=1e-8, + jitter_relative=jitter_relative, + xp=jnp, + ) + return quadratic_form_via_cholesky(covariance, jnp.asarray(VECTOR), xp=jnp) + + autodiff = float(jax.grad(term)(1.3)) + + step = 1.0e-6 + finite_difference = float((term(1.3 + step) - term(1.3 - step)) / (2.0 * step)) + + assert np.isfinite(autodiff) + assert autodiff == pytest.approx(finite_difference, rel=1.0e-6) diff --git a/test_autoarray/inversion/regularizations/test_kernel_jitter_relative.py b/test_autoarray/inversion/regularizations/test_kernel_jitter_relative.py new file mode 100644 index 000000000..5fabe6670 --- /dev/null +++ b/test_autoarray/inversion/regularizations/test_kernel_jitter_relative.py @@ -0,0 +1,223 @@ +import pytest + +import autoarray as aa +import numpy as np + +from autoarray.inversion.regularization.adapt import adapt_regularization_weights_from +from autoarray.inversion.regularization.matern_kernel import ( + apply_jitter, + matern_cov_matrix_from, +) + + +def _mapper(): + source_plane_mesh_grid = aa.Grid2D.no_mask( + values=[[0.1, 0.1], [1.1, 0.6], [2.1, 0.1], [0.4, 1.1], [1.1, 7.1], [2.1, 1.1]], + shape_native=(3, 2), + pixel_scales=1.0, + ) + return aa.m.MockMapper(source_plane_mesh_grid=source_plane_mesh_grid) + + +def _adapt_mapper(): + source_plane_mesh_grid = aa.Grid2D.no_mask( + values=[[0.1, 0.1], [1.1, 0.6], [2.1, 0.1], [0.4, 1.1], [1.1, 7.1], [2.1, 1.1]], + shape_native=(3, 2), + pixel_scales=1.0, + ) + return aa.m.MockMapper( + source_plane_mesh_grid=source_plane_mesh_grid, + pixel_signals=np.array([0.0, 0.2, 0.4, 0.6, 0.8, 1.0]), + ) + + +@pytest.mark.parametrize( + "reg_absolute, reg_relative", + [ + ( + aa.reg.MaternKernel(coefficient=3.0, scale=2.0, nu=2.0), + aa.reg.MaternKernel( + coefficient=3.0, scale=2.0, nu=2.0, jitter_relative=True + ), + ), + ( + aa.reg.ExponentialKernel(coefficient=3.0, scale=2.0), + aa.reg.ExponentialKernel(coefficient=3.0, scale=2.0, jitter_relative=True), + ), + ( + aa.reg.GaussianKernel(coefficient=3.0, scale=0.5), + aa.reg.GaussianKernel(coefficient=3.0, scale=0.5, jitter_relative=True), + ), + ], +) +def test__unweighted_kernels__relative_and_absolute_jitter_agree( + reg_absolute, reg_relative +): + """ + The unweighted kernels have ``K(0) == 1``, so their covariance diagonal is ~1 and the + absolute and relative conventions coincide. This pins that switching the flag on is + not a silent science change for these three schemes. + """ + mapper = _mapper() + + absolute = reg_absolute.regularization_matrix_from(linear_obj=mapper) + relative = reg_relative.regularization_matrix_from(linear_obj=mapper) + + assert relative == pytest.approx(absolute, rel=1.0e-7) + + +def test__default_is_byte_identical_to_absolute_jitter(): + """ + `jitter_relative` defaults to False on every scheme, so default-constructed + regularization matrices must be bit-for-bit what they were before the flag existed. + """ + mapper = _mapper() + adapt_mapper = _adapt_mapper() + + for reg, linear_obj in [ + (aa.reg.MaternKernel(coefficient=3.0, scale=2.0, nu=2.0), mapper), + (aa.reg.ExponentialKernel(coefficient=3.0, scale=2.0), mapper), + (aa.reg.GaussianKernel(coefficient=3.0, scale=0.5), mapper), + ( + aa.reg.MaternAdaptKernel( + scale=2.0, nu=2.0, inner_coefficient=0.5, outer_coefficient=4.0 + ), + adapt_mapper, + ), + ]: + assert reg.jitter_relative is False + + covariance = matern_cov_matrix_from( + scale=2.0, + nu=2.0, + pixel_points=linear_obj.source_plane_mesh_grid.array, + jitter=0.0, + ) + pixels = covariance.shape[0] + + assert np.array_equal( + apply_jitter(covariance, jitter=1e-8, jitter_relative=False), + covariance + 1e-8 * np.eye(pixels), + ) + + +def test__apply_jitter_relative__is_a_pure_rescaling_of_the_diagonal(): + covariance = matern_cov_matrix_from( + scale=1.0, + nu=2.0, + pixel_points=np.random.default_rng(0).normal(size=(12, 2)), + jitter=0.0, + weights=np.linspace(0.1, 5.0, 12), + ) + + jittered = apply_jitter(covariance, jitter=1e-3, jitter_relative=True) + + # Diagonal scaled by exactly (1 + jitter); off-diagonal untouched. + assert np.diag(jittered) == pytest.approx(np.diag(covariance) * (1.0 + 1e-3)) + + off_diagonal = ~np.eye(covariance.shape[0], dtype=bool) + assert np.array_equal(jittered[off_diagonal], covariance[off_diagonal]) + + +def test__adaptive_weights__absolute_jitter_swamps_faint_pixels__relative_does_not(): + """ + The bug 2b fixes. `MaternAdaptKernel`'s covariance diagonal is `C_ii = w_i^2`, which + spans the adaptive-weight dynamic range. At wide inner/outer coefficients the fixed + absolute 1e-8 reaches 100% of the faintest pixel's variance and destroys its kernel + structure; the relative convention perturbs every pixel by 1e-8 whatever its scale. + """ + points = np.random.default_rng(0).normal(size=(40, 2)) + pixel_signals = np.linspace(0.0, 1.0, 40) + + kernel_weights = 1.0 / adapt_regularization_weights_from( + inner_coefficient=0.1, outer_coefficient=100.0, pixel_signals=pixel_signals + ) + + covariance = matern_cov_matrix_from( + scale=1.0, nu=2.0, pixel_points=points, weights=kernel_weights, jitter=0.0 + ) + + faintest_variance = np.diag(covariance).min() + + # The regime the absolute convention breaks in: the faintest pixel's own variance has + # fallen to the jitter's magnitude. + assert faintest_variance == pytest.approx(1.0e-8, rel=0.5) + + absolute = apply_jitter(covariance, jitter=1e-8, jitter_relative=False) + relative = apply_jitter(covariance, jitter=1e-8, jitter_relative=True) + + absolute_distortion = ( + np.diag(absolute) - np.diag(covariance) + ).max() / faintest_variance + relative_distortion = np.max( + (np.diag(relative) - np.diag(covariance)) / np.diag(covariance) + ) + + # Absolute: order-unity corruption of the faint pixel. Relative: 1e-8, as intended. + assert absolute_distortion > 0.5 + assert relative_distortion == pytest.approx(1.0e-8, rel=1.0e-6) + + +def test__relative_jitter__preserves_conditioning_on_clustered_vertices(): + """ + The jitter is load-bearing for the Cholesky on clustered (traced) mesh vertices. The + relative convention must not weaken that protection — on an unweighted kernel, where + the diagonal is ~1, the two conventions must give the same conditioning. + """ + rng = np.random.default_rng(42) + points = np.concatenate( + [ + rng.normal(scale=3.0e-3, size=(24, 2)), + rng.normal(loc=1.0, scale=3.0e-3, size=(24, 2)), + ] + ) + + covariance = matern_cov_matrix_from( + scale=1.0, nu=2.5, pixel_points=points, jitter=0.0 + ) + + absolute = apply_jitter(covariance, jitter=1e-8, jitter_relative=False) + relative = apply_jitter(covariance, jitter=1e-8, jitter_relative=True) + + # Both factorize, and the conditioning is materially the same. + np.linalg.cholesky(absolute) + np.linalg.cholesky(relative) + + assert np.linalg.cond(relative) == pytest.approx( + np.linalg.cond(absolute), rel=1.0e-3 + ) + + +def test__adapt_kernel_end_to_end__flag_reaches_the_regularization_matrix(): + """ + The flag must be threaded to every covariance call site, not just the constructor. + """ + mapper = _adapt_mapper() + + absolute = aa.reg.MaternAdaptKernel( + scale=2.0, nu=2.0, inner_coefficient=0.1, outer_coefficient=100.0 + ) + relative = aa.reg.MaternAdaptKernel( + scale=2.0, + nu=2.0, + inner_coefficient=0.1, + outer_coefficient=100.0, + jitter_relative=True, + ) + + assert not np.array_equal( + absolute.regularization_matrix_from(linear_obj=mapper), + relative.regularization_matrix_from(linear_obj=mapper), + ) + + # The log-det and term shortcuts must see the same covariance the matrix did. + assert absolute.log_det_regularization_matrix_term_from( + linear_obj=mapper + ) != relative.log_det_regularization_matrix_term_from(linear_obj=mapper) + + reconstruction = np.array([0.3, -1.2, 0.7, 2.1, -0.4, 1.5]) + assert absolute.regularization_term_from( + linear_obj=mapper, reconstruction=reconstruction + ) != relative.regularization_term_from( + linear_obj=mapper, reconstruction=reconstruction + ) diff --git a/test_autoarray/inversion/regularizations/test_kernel_regularization_term.py b/test_autoarray/inversion/regularizations/test_kernel_regularization_term.py new file mode 100644 index 000000000..6c73ae97b --- /dev/null +++ b/test_autoarray/inversion/regularizations/test_kernel_regularization_term.py @@ -0,0 +1,226 @@ +import pytest + +import autoarray as aa +import numpy as np + + +def _mapper_and_grid(): + source_plane_mesh_grid = aa.Grid2D.no_mask( + values=[[0.1, 0.1], [1.1, 0.6], [2.1, 0.1], [0.4, 1.1], [1.1, 7.1], [2.1, 1.1]], + shape_native=(3, 2), + pixel_scales=1.0, + ) + mapper = aa.m.MockMapper(source_plane_mesh_grid=source_plane_mesh_grid) + return mapper + + +RECONSTRUCTION = np.array([0.3, -1.2, 0.7, 2.1, -0.4, 1.5]) + + +@pytest.mark.parametrize( + "reg, rtol", + [ + (aa.reg.MaternKernel(coefficient=3.0, scale=2.0, nu=2.0), 1.0e-8), + (aa.reg.ExponentialKernel(coefficient=3.0, scale=2.0), 1.0e-8), + # GaussianKernel's formed matrix carries an extra trace-scaled diagonal + # stabilisation jitter (and a symmetrisation) which the analytic shortcut + # deliberately excludes, exactly as its log-det shortcut does, so the + # agreement is at the jitter scale rather than machine precision. + (aa.reg.GaussianKernel(coefficient=3.0, scale=0.5), 1.0e-6), + ], +) +def test__kernel_regularization_term_shortcut__matches_formed_matrix(reg, rtol): + mapper = _mapper_and_grid() + + formed = reg.regularization_matrix_from(linear_obj=mapper) + expected = RECONSTRUCTION @ (formed @ RECONSTRUCTION) + + shortcut = reg.regularization_term_from( + linear_obj=mapper, reconstruction=RECONSTRUCTION + ) + + assert shortcut == pytest.approx(expected, rel=rtol) + + +def test__matern_adapt_kernel_regularization_term_shortcut__matches_formed_matrix(): + source_plane_mesh_grid = aa.Grid2D.no_mask( + values=[[0.1, 0.1], [1.1, 0.6], [2.1, 0.1], [0.4, 1.1], [1.1, 7.1], [2.1, 1.1]], + shape_native=(3, 2), + pixel_scales=1.0, + ) + mapper = aa.m.MockMapper( + source_plane_mesh_grid=source_plane_mesh_grid, + pixel_signals=np.array([0.1, 0.4, 0.2, 0.9, 0.5, 0.3]), + ) + + reg = aa.reg.MaternAdaptKernel( + scale=2.0, nu=2.0, inner_coefficient=0.5, outer_coefficient=4.0 + ) + + formed = reg.regularization_matrix_from(linear_obj=mapper) + expected = RECONSTRUCTION @ (formed @ RECONSTRUCTION) + + shortcut = reg.regularization_term_from( + linear_obj=mapper, reconstruction=RECONSTRUCTION + ) + + assert shortcut == pytest.approx(expected, rel=1.0e-8) + + +def test__matern_adapt_kernel__does_not_inherit_the_zero_coefficient_matern_term(): + """ + `MaternAdaptKernel` passes `coefficient=0.0` to `MaternKernel.__init__` (its adaptive + weights live inside the covariance instead). Were it to inherit `MaternKernel`'s + `regularization_term_from`, the `self.coefficient` factor would zero the term — this + pins the override that prevents that. + """ + source_plane_mesh_grid = aa.Grid2D.no_mask( + values=[[0.1, 0.1], [1.1, 0.6], [2.1, 0.1], [0.4, 1.1], [1.1, 7.1], [2.1, 1.1]], + shape_native=(3, 2), + pixel_scales=1.0, + ) + mapper = aa.m.MockMapper( + source_plane_mesh_grid=source_plane_mesh_grid, + pixel_signals=np.array([0.1, 0.4, 0.2, 0.9, 0.5, 0.3]), + ) + + reg = aa.reg.MaternAdaptKernel( + scale=2.0, nu=2.0, inner_coefficient=0.5, outer_coefficient=4.0 + ) + + assert reg.coefficient == 0.0 + + term = reg.regularization_term_from( + linear_obj=mapper, reconstruction=RECONSTRUCTION + ) + + assert term != 0.0 + assert np.isfinite(term) + + +def test__non_kernel_regularization__has_no_regularization_term_shortcut(): + mapper = _mapper_and_grid() + + assert ( + aa.reg.Constant(coefficient=1.0).regularization_term_from( + linear_obj=mapper, reconstruction=RECONSTRUCTION + ) + is None + ) + + +def test__inversion_regularization_term__kernel_shortcut_only_under_cho_solve(): + mapper = _mapper_and_grid() + reg = aa.reg.MaternKernel(coefficient=3.0, scale=2.0, nu=2.0) + mapper = aa.m.MockMapper( + source_plane_mesh_grid=mapper.source_plane_mesh_grid, + regularization=reg, + ) + + formed = reg.regularization_matrix_from(linear_obj=mapper) + + # Default ("matmul") path: byte-identical to contracting the formed matrix — the + # shortcut is never consulted. + default_inversion = aa.m.MockInversion( + linear_obj_list=[mapper], + regularization_matrix=formed, + reconstruction=RECONSTRUCTION, + ) + expected_default = RECONSTRUCTION @ (formed @ RECONSTRUCTION) + assert default_inversion.regularization_term == pytest.approx( + expected_default, rel=1.0e-12 + ) + + # Opt-in "cho_solve" path: the Cholesky-solve shortcut. + cho_solve_inversion = aa.m.MockInversion( + linear_obj_list=[mapper], + regularization_matrix=formed, + reconstruction=RECONSTRUCTION, + settings=aa.Settings(regularization_term_method="cho_solve"), + ) + expected_shortcut = reg.regularization_term_from( + linear_obj=mapper, reconstruction=RECONSTRUCTION + ) + assert cho_solve_inversion.regularization_term == pytest.approx( + expected_shortcut, rel=1.0e-12 + ) + + # On a well-conditioned fixture the two agree to high precision. + assert cho_solve_inversion.regularization_term == pytest.approx( + default_inversion.regularization_term, rel=1.0e-8 + ) + + +def test__inversion_regularization_term__falls_back_when_a_scheme_has_no_shortcut(): + """ + A `Constant` scheme returns `None`, so even under "cho_solve" the whole computation + must fall back to the formed matrix rather than dropping that object's contribution. + """ + mapper = _mapper_and_grid() + reg = aa.reg.Constant(coefficient=1.0) + mapper = aa.m.MockMapper( + source_plane_mesh_grid=mapper.source_plane_mesh_grid, + regularization=reg, + ) + + # Any SPD matrix serves as the formed `H` here — the point of the test is that the + # `Constant` scheme has no shortcut, so the formed matrix is what must be used. It is + # built from a kernel scheme only because `Constant`'s own matrix needs mesh + # neighbours this mock fixture does not carry. + formed = aa.reg.MaternKernel( + coefficient=3.0, scale=2.0, nu=2.0 + ).regularization_matrix_from(linear_obj=mapper) + + inversion = aa.m.MockInversion( + linear_obj_list=[mapper], + regularization_matrix=formed, + reconstruction=RECONSTRUCTION, + settings=aa.Settings(regularization_term_method="cho_solve"), + ) + + assert inversion.regularization_term == pytest.approx( + RECONSTRUCTION @ (formed @ RECONSTRUCTION), rel=1.0e-12 + ) + + +def test__cho_solve_beats_explicit_inverse_on_an_ill_conditioned_covariance(): + """ + The reason the shortcut exists. On clustered vertices `cond(C)` is large and the + explicitly formed `C^-1` carries round-off amplified by it; the Cholesky solve does + not. + + Graded against an exactly-known reference: choosing the right-hand side as + `s = C v` makes `C^-1 s = v`, so the true quadratic form is `s^T v = v^T C v` — a + well-conditioned expression needing no inverse at all. + """ + from autoarray.inversion.regularization.matern_kernel import ( + matern_cov_matrix_from, + inv_via_cholesky, + quadratic_form_via_cholesky, + ) + + # Deliberately clustered points -> small minimum pairwise separation -> large cond(C), + # mirroring the traced vertices of the kNN mesh families. + rng = np.random.default_rng(42) + points = np.concatenate( + [ + rng.normal(scale=3.0e-3, size=(24, 2)), + rng.normal(loc=1.0, scale=3.0e-3, size=(24, 2)), + ] + ) + + covariance = matern_cov_matrix_from(scale=1.0, nu=2.5, pixel_points=points) + + assert np.linalg.cond(covariance) > 1.0e8 + + v = rng.normal(size=points.shape[0]) + s = covariance @ v + reference = s @ v + + explicit = s @ (inv_via_cholesky(covariance) @ s) + implicit = quadratic_form_via_cholesky(covariance, s) + + explicit_error = abs(explicit - reference) / abs(reference) + implicit_error = abs(implicit - reference) / abs(reference) + + assert implicit_error < explicit_error