diff --git a/autoarray/inversion/regularization/__init__.py b/autoarray/inversion/regularization/__init__.py index 5bcc2d5b0..e4e26b310 100644 --- a/autoarray/inversion/regularization/__init__.py +++ b/autoarray/inversion/regularization/__init__.py @@ -15,5 +15,6 @@ from .gaussian_kernel import GaussianKernel from .exponential_kernel import ExponentialKernel from .matern_kernel import MaternKernel +from .matern_sparse import MaternSparse from .matern_adapt_kernel import MaternAdaptKernel from .matern_adapt_power_kernel import MaternAdaptPowerKernel diff --git a/autoarray/inversion/regularization/matern_sparse.py b/autoarray/inversion/regularization/matern_sparse.py new file mode 100644 index 000000000..2d7582483 --- /dev/null +++ b/autoarray/inversion/regularization/matern_sparse.py @@ -0,0 +1,424 @@ +""" +A sparse Matérn regularization for rectangular meshes, built from the SPDE construction of the Matérn +field (Lindgren, Rue & Lindström 2011). + +This is phase 2 of the matrix-free interferometer inversion epic (PyAutoArray #638, community proposal by +@HRSAstro, https://github.com/PyAutoLabs/.github/discussions/39). The dense `MaternKernel` builds the +`(S, S)` Matérn covariance `C` and inverts it, so its regularization matrix `H = coefficient * C^-1` is +fully dense. For integer `alpha = nu + 1` (in 2D) the Matérn *precision* `C^-1` is instead sparse: it is +a polynomial in the discrete Laplacian. This module builds it for `nu = 1` (`alpha = 2`), giving an `H` +with the 13-point stencil of `L^2`, which the matrix-free utilities of +`inversion_interferometer_matrix_free_util` take directly as a `scipy.sparse` matrix. + +The module is NumPy / scipy only and imports no JAX. +""" + +from __future__ import annotations +import numpy as np +import scipy.sparse +import scipy.sparse.linalg +from typing import Optional, Tuple, TYPE_CHECKING + +if TYPE_CHECKING: + from autoarray.inversion.linear_obj.linear_obj import LinearObj + +from autoarray.inversion.regularization.abstract import AbstractRegularization +from autoarray.inversion.regularization.abstract import validate_coefficient + + +def rectangular_laplacian_sparse_from( + shape_native: Tuple[int, int], pixel_scales: Tuple[float, float] +) -> scipy.sparse.csr_matrix: + """ + Returns the 5-point discrete Laplacian of a rectangular mesh with reflecting (zero-flux Neumann) + edges, as a `scipy.sparse` CSR matrix of shape `(pixels, pixels)`. + + Mesh pixels are indexed in row-major order, which is the order of a rectangular mesh's + `source_plane_mesh_grid` (built by `grid_2d_slim_via_shape_native_not_mask_from`): pixel + `k = i * total_x_pixels + j` is in row `i` (row 0 is the top row, the largest y) and column `j` + (column 0 is the left column, the smallest x). Pixel `k` therefore neighbours `k +- 1` along x (within + its row) and `k +- total_x_pixels` along y. + + The Laplacian is the adjacency matrix minus its row sums, each axis divided by the square of its + own pixel scale, so for an interior pixel + + (lap s)_k = (s_{k-1} - 2 s_k + s_{k+1}) / dx^2 + (s_{k-nx} - 2 s_k + s_{k+nx}) / dy^2 + + and an edge pixel simply drops its missing neighbours (and the matching diagonal term), which is the + reflecting boundary condition. The rows of the matrix sum to zero. + + Parameters + ---------- + shape_native + The 2D dimensions of the rectangular mesh, `(total_y_pixels, total_x_pixels)`. + pixel_scales + The `(y, x)` spacing between neighbouring mesh pixel centres. + + Returns + ------- + The sparse Laplacian. + """ + total_y_pixels, total_x_pixels = int(shape_native[0]), int(shape_native[1]) + pixel_scale_y, pixel_scale_x = float(pixel_scales[0]), float(pixel_scales[1]) + + pixels = total_y_pixels * total_x_pixels + + index = np.arange(pixels).reshape(total_y_pixels, total_x_pixels) + + rows = [] + cols = [] + data = [] + + for left, right, pixel_scale in ( + (index[:, :-1], index[:, 1:], pixel_scale_x), + (index[:-1, :], index[1:, :], pixel_scale_y), + ): + value = 1.0 / pixel_scale**2 + + rows += [left.ravel(), right.ravel()] + cols += [right.ravel(), left.ravel()] + data += [np.full(2 * left.size, value)] + + adjacency = scipy.sparse.csr_matrix( + (np.concatenate(data), (np.concatenate(rows), np.concatenate(cols))), + shape=(pixels, pixels), + ) + + return ( + adjacency - scipy.sparse.diags(np.asarray(adjacency.sum(axis=1)).ravel()) + ).tocsr() + + +def _splu_symmetric_from(matrix: scipy.sparse.spmatrix): + """ + Sparse LU of a symmetric positive definite matrix with the symmetric-mode options (`MMD_AT_PLUS_A` + ordering, no pivoting threshold), matching `powell_preconditioner_from`. + """ + return scipy.sparse.linalg.splu( + scipy.sparse.csc_matrix(matrix), + permc_spec="MMD_AT_PLUS_A", + diag_pivot_thresh=0.0, + options=dict(SymmetricMode=True), + ) + + +class MaternSparse(AbstractRegularization): + def __init__( + self, + coefficient: float = 1.0, + scale: float = 1.0, + nu: float = 1.0, + weights: Optional[np.ndarray] = None, + ): + """ + Regularization which uses a Matérn prior on the source, expressed through its sparse precision + matrix (the SPDE construction of Lindgren, Rue & Lindström 2011) rather than a dense covariance. + + The dense `MaternKernel` regularizes every pixel with every other pixel by building the Matérn + covariance `C` and inverting it, so its regularization matrix is dense and costs `O(S^3)` to form. + For integer `alpha = nu + 1` a Matérn field on a 2D grid is the solution of the stochastic PDE + `L^(alpha / 2) f = noise` with + + L = kappa^2 I - lap, kappa = sqrt(2 nu) / scale, + + where `lap` is the discrete Laplacian. Its precision is then `C^-1 = c L^alpha`, a sparse matrix. + This class implements `nu = 1` (`alpha = 2`): + + H = coefficient * W^-1 c L^2 W^-1 = coefficient * R^T R, R = sqrt(c) L W^-1, + + with: + + - `lap` the 5-point Laplacian with reflecting edges (`rectangular_laplacian_sparse_from`), each + axis scaled by its own mesh pixel scale; + - `c = e^T L^-2 e` with `e` the indicator of the mesh's centre pixel, so the prior variance at the + mesh centre, `e^T (c L^2)^-1 e = (e^T L^-2 e) / c`, is 1, matching the unit zero-lag value `K(0) = 1` of the dense kernel (edge + pixels, with fewer neighbours, have a larger variance under reflecting edges); + - `W = diag(w)` the per-pixel kernel weights (`kernel_weights_from`), which scale the prior + standard deviation of each pixel exactly as `C_ij = K(d_ij) w_i w_j` does for the adaptive dense + kernels (`MaternAdaptKernel` uses `w = 1 / regularization_weights`). They are unit by default. + + `H` has the 13-point stencil of `L^2`. It is exactly symmetric positive definite (`L` is, since + `kappa^2 > 0`), so unlike the dense kernels no diagonal jitter is added. + + **Meshes**: only rectangular meshes (`RectangularUniform`, `RectangularRTUAdaptDensity` and their + subclasses) are supported, because the 5-point Laplacian needs the mesh's grid topology. The + Laplacian is built on the mapper's `source_plane_mesh_grid` (the uniform overlay grid, also the + grid the dense kernel measures distances on), whose pixels are in row-major order from the + top-left (see `rectangular_laplacian_sparse_from`). Other meshes (e.g. `Delaunay`) raise + `NotImplementedError`. + + **Smoothness**: only `nu = 1` is implemented. Higher integer `nu` is `c L^(nu + 1)`, a follow-up; + non-integer `alpha` has no exact sparse precision. + + **Matrix-free use**: `regularization_matrix_sparse_from` returns `H` as a `scipy.sparse` matrix, + to be passed to `powell_preconditioner_from` / `reconstruction_via_pcg_from` with `coefficient=1.0` + (it already carries this class's coefficient), and `regularization_factor_sparse_from` returns + `R`, which draws the evidence probes `R^T z`. `regularization_matrix_from` returns the dense array + of the same `H`, so the class is a drop-in on today's dense inversion path. + + The class is NumPy / scipy only: `xp` is honoured only for the dense array returned by + `regularization_matrix_from`, and none of its methods can be traced by JAX. + + A full description of regularization and this matrix can be found in the parent + `AbstractRegularization` class. + + Parameters + ---------- + coefficient + The regularization coefficient which controls the degree of smoothing of the inversion + reconstruction. + scale + The correlation length of the Matérn prior, in the units of the mesh coordinates. + nu + The smoothness of the Matérn prior. Only `nu = 1` is implemented. + weights + Optional per-pixel kernel weights `w` of shape `[pixels]` (the prior standard deviation of each + pixel, relative to the unit centre variance). `None` gives unit weights. Subclasses coupling the + weights to the data (e.g. adaptive brightness) override `kernel_weights_from` instead. + """ + if nu != 1.0: + raise NotImplementedError( + f"MaternSparse implements only nu = 1 (alpha = nu + 1 = 2, precision c L^2); " + f"nu = {nu} was input. Higher integer nu (precision c L^(nu + 1)) is a follow-up, and " + f"non-integer alpha has no exact sparse precision -- use MaternKernel for those." + ) + + validate_coefficient(coefficient=coefficient) + + self.coefficient = coefficient + self.scale = scale + self.nu = nu + self.weights = weights + + super().__init__() + + def regularization_weights_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray: + """ + Returns the regularization weights of this regularization scheme. + + As for the dense `MaternKernel`, these are the coefficient for every pixel; the per-pixel prior + scaling of this scheme is held by `kernel_weights_from`. + + Parameters + ---------- + linear_obj + The linear object (e.g. a ``Mapper``) which uses these weights when performing regularization. + + Returns + ------- + The regularization weights. + """ + return self.coefficient * xp.ones(linear_obj.params) + + def kernel_weights_from(self, linear_obj: LinearObj) -> np.ndarray: + """ + Returns the per-pixel kernel weights `w` entering `H = coefficient * W^-1 c L^2 W^-1`: the input + `weights` if given, otherwise ones. + + Parameters + ---------- + linear_obj + The linear object (e.g. a ``Mapper``) which uses these weights when performing regularization. + + Returns + ------- + The kernel weights, shape `[pixels]`. + """ + if self.weights is None: + return np.ones(linear_obj.params) + + weights = np.asarray(self.weights, dtype=float) + + if weights.shape != (linear_obj.params,): + raise ValueError( + f"MaternSparse weights have shape {weights.shape} but the mesh has " + f"{linear_obj.params} pixels." + ) + + return weights + + @staticmethod + def mesh_shape_and_pixel_scales_from( + linear_obj: LinearObj, + ) -> Tuple[Tuple[int, int], Tuple[float, float]]: + """ + Returns the `(total_y_pixels, total_x_pixels)` shape and the `(y, x)` spacing of pixel centres of + a linear object's rectangular mesh, read from its `source_plane_mesh_grid`. + + Raises `NotImplementedError` if the linear object is not a mapper with a rectangular mesh. + """ + from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_density import ( + RectangularRTUAdaptDensity, + ) + + mesh = getattr(linear_obj, "mesh", None) + + if not isinstance(mesh, RectangularRTUAdaptDensity): + raise NotImplementedError( + f"MaternSparse supports only rectangular meshes (RectangularUniform, " + f"RectangularRTUAdaptDensity and their subclasses), because its sparse precision uses the " + f"5-point Laplacian of the mesh grid; got a {type(mesh).__name__} mesh. Use MaternKernel " + f"for other meshes (e.g. Delaunay)." + ) + + shape_native = (int(mesh.shape[0]), int(mesh.shape[1])) + + grid = np.asarray(linear_obj.source_plane_mesh_grid.array).reshape( + shape_native[0], shape_native[1], 2 + ) + + # Row-major from the top-left: y decreases down the rows, x increases along the columns. + pixel_scales = ( + float(grid[0, 0, 0] - grid[1, 0, 0]), + float(grid[0, 1, 1] - grid[0, 0, 1]), + ) + + return shape_native, pixel_scales + + def _laplacian_operator_and_c_from(self, linear_obj: LinearObj): + """ + Returns `L = kappa^2 I - lap` (CSC), its sparse LU and the centre-variance normalization `c`. + """ + shape_native, pixel_scales = self.mesh_shape_and_pixel_scales_from( + linear_obj=linear_obj + ) + + pixels = shape_native[0] * shape_native[1] + + kappa = np.sqrt(2.0 * self.nu) / self.scale + + laplacian_operator = ( + kappa**2 * scipy.sparse.identity(pixels, format="csr") + - rectangular_laplacian_sparse_from( + shape_native=shape_native, pixel_scales=pixel_scales + ) + ).tocsc() + + factor = _splu_symmetric_from(laplacian_operator) + + centre = np.zeros(pixels) + centre[(shape_native[0] // 2) * shape_native[1] + shape_native[1] // 2] = 1.0 + + l_inv_centre = factor.solve(centre) + + # The covariance (c L^2)^-1 has centre variance (e^T L^-2 e) / c, so c = e^T L^-2 e makes it one; + # e^T L^-2 e = |L^-1 e|^2 since L is symmetric. + c = float(l_inv_centre @ l_inv_centre) + + return laplacian_operator, factor, c + + def regularization_factor_sparse_from( + self, linear_obj: LinearObj + ) -> scipy.sparse.csr_matrix: + """ + Returns the factor `R = sqrt(c) L W^-1` of the unscaled precision, `c W^-1 L^2 W^-1 = R^T R`, as a + `scipy.sparse` CSR matrix (the regularization matrix is `coefficient * R^T R`). + + `R^T z` with `z ~ N(0, I)` is a draw with covariance `R^T R`, which is how the stochastic evidence + log-determinant of the matrix-free path draws its probes. + + Parameters + ---------- + linear_obj + The linear object (a ``Mapper`` with a rectangular mesh). + + Returns + ------- + The sparse factor `R`. + """ + laplacian_operator, _, c = self._laplacian_operator_and_c_from( + linear_obj=linear_obj + ) + + weights = self.kernel_weights_from(linear_obj=linear_obj) + + return ( + np.sqrt(c) * laplacian_operator @ scipy.sparse.diags(1.0 / weights) + ).tocsr() + + def regularization_matrix_sparse_from( + self, linear_obj: LinearObj + ) -> scipy.sparse.csr_matrix: + """ + Returns the regularization matrix `H = coefficient * R^T R = coefficient * W^-1 c L^2 W^-1` as a + `scipy.sparse` CSR matrix of shape `[pixels, pixels]`. + + This is the matrix the matrix-free utilities take; pass it with `coefficient=1.0` there, since it + already carries this class's coefficient. + + Parameters + ---------- + linear_obj + The linear object (a ``Mapper`` with a rectangular mesh). + + Returns + ------- + The sparse regularization matrix. + """ + factor = self.regularization_factor_sparse_from(linear_obj=linear_obj) + + return (self.coefficient * (factor.T @ factor)).tocsr() + + def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray: + """ + Returns the regularization matrix with shape [pixels, pixels], as the dense array of + `regularization_matrix_sparse_from`, so this scheme is a drop-in on the dense inversion path. + + Parameters + ---------- + linear_obj + The linear object (a ``Mapper`` with a rectangular mesh). + + Returns + ------- + The regularization matrix. + """ + return xp.asarray( + self.regularization_matrix_sparse_from(linear_obj=linear_obj).toarray() + ) + + def log_det_regularization_matrix_term_from( + self, linear_obj: LinearObj, xp=np + ) -> float: + """ + The analytically exact ``log det H`` from the sparse LU of `L`: + + log det H = pixels * log(coefficient) + pixels * log(c) + 2 * sum(log|U_ii|) - 2 * sum(log w) + + since `H = coefficient * c W^-1 L^2 W^-1`, `det(L^2) = det(L)^2` and the LU's `L` factor has a + unit diagonal (`L` is SPD, so `det L > 0`). This is the same quantity the dense `MaternKernel` + returns (``log det H``, there ``pixels * log(coefficient) - log det C``). + + Consumed by the inversion only when ``Settings.log_det_method == "slogdet"`` (see + :meth:`AbstractRegularization.log_det_regularization_matrix_term_from`); the default evidence path + factorizes the formed ``H``. + """ + _, factor, c = self._laplacian_operator_and_c_from(linear_obj=linear_obj) + + weights = self.kernel_weights_from(linear_obj=linear_obj) + + pixels = linear_obj.params + + return float( + pixels * np.log(self.coefficient) + + pixels * np.log(c) + + 2.0 * np.sum(np.log(np.abs(factor.U.diagonal()))) + - 2.0 * np.sum(np.log(weights)) + ) + + def regularization_term_from( + self, linear_obj: LinearObj, reconstruction: np.ndarray, xp=np + ) -> float: + """ + The regularization term ``s^T H s``, evaluated as ``coefficient * |R s|^2`` with the sparse factor + `R`, so no dense matrix is formed. + + 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``. + """ + factor = self.regularization_factor_sparse_from(linear_obj=linear_obj) + + r_s = factor @ np.asarray(reconstruction, dtype=float) + + return float(self.coefficient * (r_s @ r_s)) diff --git a/test_autoarray/inversion/inversion/interferometer/test_matrix_free_util.py b/test_autoarray/inversion/inversion/interferometer/test_matrix_free_util.py index 955c66194..913259b1e 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_matrix_free_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_matrix_free_util.py @@ -46,7 +46,7 @@ def _dataset_sparse_from(seed=1, n_visibilities=40): return mask, dataset.apply_sparse_operator(use_jax=False) -def _rectangular_mapper_from(mask, coefficient=1.0, shape=(6, 6)): +def _rectangular_mapper_from(mask, coefficient=1.0, shape=(6, 6), regularization=None): grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=2) mesh = aa.mesh.RectangularUniform(shape=shape) @@ -61,7 +61,7 @@ def _rectangular_mapper_from(mask, coefficient=1.0, shape=(6, 6)): return aa.Mapper( interpolator=interpolator, - regularization=aa.reg.Constant(coefficient=coefficient), + regularization=regularization or aa.reg.Constant(coefficient=coefficient), ) @@ -194,6 +194,56 @@ def test__reconstruction_via_pcg__matches_unconstrained_sparse_path(coefficient) ) <= 1.0e-8 * np.max(np.abs(reconstruction_sparse_path)) +@pytest.mark.parametrize("coefficient", [1.0, 0.01]) +def test__reconstruction_via_pcg__matern_sparse_matrix__matches_unconstrained_sparse_path( + coefficient, +): + """ + Phase 2: the natively sparse Matérn precision of `MaternSparse` is the `H` of the matrix-free solve. + The dense-path inversion uses the same `MaternSparse` instance as the pixelization's regularization + (through its dense `regularization_matrix_from`), and since `H` already carries the coefficient the + matrix-free `λ` is 1. + """ + mask, dataset_sparse = _dataset_sparse_from() + + regularization = aa.reg.MaternSparse(coefficient=coefficient, scale=1.0) + + mapper = _rectangular_mapper_from(mask=mask, regularization=regularization) + inversion = _inversion_from(dataset_sparse=dataset_sparse, mapper=mapper) + + operator = mf.MatrixFreeDataOperator.from_inversion( + inversion=inversion, mapper=mapper + ) + + regularization_matrix = regularization.regularization_matrix_sparse_from(mapper) + + assert np.array_equal( + regularization_matrix.toarray(), np.asarray(inversion.regularization_matrix) + ) + + preconditioner, _ = mf.powell_preconditioner_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + coefficient=1.0, + ) + + assert scipy.sparse.issparse(preconditioner) + + reconstruction, info = mf.reconstruction_via_pcg_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + coefficient=1.0, + rtol=1.0e-12, + ) + + reconstruction_sparse_path = np.asarray(inversion.reconstruction) + + assert info == 0 + assert np.max( + np.abs(reconstruction - reconstruction_sparse_path) + ) <= 1.0e-8 * np.max(np.abs(reconstruction_sparse_path)) + + def test__reconstruction_via_pcg__coefficient_scales_regularization_matrix(): mask, dataset_sparse = _dataset_sparse_from() mapper = _rectangular_mapper_from(mask=mask) diff --git a/test_autoarray/inversion/regularizations/test_matern_sparse.py b/test_autoarray/inversion/regularizations/test_matern_sparse.py new file mode 100644 index 000000000..8dd8368d8 --- /dev/null +++ b/test_autoarray/inversion/regularizations/test_matern_sparse.py @@ -0,0 +1,353 @@ +""" +Tests of the sparse Matérn (SPDE) precision regularization `MaternSparse` (phase 2 of PyAutoArray #638). + +These carry no `importorskip("jax")`: the module is NumPy / scipy only, so the `unittest-nojax` CI leg +runs them. +""" + +import os +import subprocess +import sys +import textwrap + +import numpy as np +import pytest +import scipy.sparse + +import autoarray as aa + + +def _rectangular_mapper_from(shape=(4, 3), mesh_cls=None): + """ + A mapper with a rectangular mesh laid over a circular-masked grid. The mesh grid's y and x spacings + differ (1.1875 and 1.5833 for the 4x3 mesh), so the per-axis Laplacian scaling is exercised. + """ + mask = aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=0.5, radius=2.5) + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=2) + + mesh = (mesh_cls or aa.mesh.RectangularUniform)(shape=shape) + + return aa.Mapper( + interpolator=mesh.interpolator_from( + source_plane_data_grid=grid, source_plane_mesh_grid=None + ), + regularization=aa.reg.Constant(), + ) + + +def _laplacian_operator_by_hand_from(mapper, scale, nu=1.0): + """ + `L = kappa^2 I - lap` built pixel by pixel from the mesh's (y, x) coordinates: neighbours are found + geometrically (same row and one x-spacing apart, or same column and one y-spacing apart) rather than + from the row-major index arithmetic the class uses, so this also checks the mesh ordering. + """ + grid = np.asarray(mapper.source_plane_mesh_grid.array) + pixels = grid.shape[0] + + ys = np.unique(grid[:, 0]) + xs = np.unique(grid[:, 1]) + dy = ys[1] - ys[0] + dx = xs[1] - xs[0] + + kappa = np.sqrt(2.0 * nu) / scale + + L = kappa**2 * np.eye(pixels) + + for k in range(pixels): + for m in range(pixels): + if k == m: + continue + same_row = np.isclose(grid[k, 0], grid[m, 0]) + same_col = np.isclose(grid[k, 1], grid[m, 1]) + if same_row and np.isclose(abs(grid[k, 1] - grid[m, 1]), dx): + L[k, m] -= 1.0 / dx**2 + L[k, k] += 1.0 / dx**2 + if same_col and np.isclose(abs(grid[k, 0] - grid[m, 0]), dy): + L[k, m] -= 1.0 / dy**2 + L[k, k] += 1.0 / dy**2 + + return L + + +def _c_by_hand_from(L, shape): + centre = (shape[0] // 2) * shape[1] + shape[1] // 2 + # The covariance (c L^2)^-1 has centre variance [L^-2]_cc / c, so c = [L^-2]_cc makes it one. + return np.linalg.inv(L @ L)[centre, centre] + + +def test__mesh_ordering__source_plane_mesh_grid_is_row_major_from_top_left(): + mapper = _rectangular_mapper_from(shape=(4, 3)) + + grid = np.asarray(mapper.source_plane_mesh_grid.array).reshape(4, 3, 2) + + # y is constant along a row and decreases down the rows; x increases along a row. + assert np.allclose(grid[:, :, 0], grid[:, :1, 0]) + assert np.all(np.diff(grid[:, 0, 0]) < 0.0) + assert np.allclose(grid[:, :, 1], grid[:1, :, 1]) + assert np.all(np.diff(grid[0, :, 1]) > 0.0) + + shape_native, pixel_scales = aa.reg.MaternSparse.mesh_shape_and_pixel_scales_from( + mapper + ) + + assert shape_native == (4, 3) + assert pixel_scales == pytest.approx(mapper.mesh_geometry.pixel_scales, 1.0e-12) + + +def test__regularization_matrix__4x3_mesh__equals_hand_built_coefficient_c_l_squared(): + mapper = _rectangular_mapper_from(shape=(4, 3)) + + reg = aa.reg.MaternSparse(coefficient=1.7, scale=2.0) + + L = _laplacian_operator_by_hand_from(mapper=mapper, scale=2.0) + c = _c_by_hand_from(L=L, shape=(4, 3)) + + regularization_matrix = reg.regularization_matrix_sparse_from(mapper) + + assert isinstance(regularization_matrix, scipy.sparse.csr_matrix) + assert regularization_matrix.toarray() == pytest.approx( + 1.7 * c * L @ L, rel=1.0e-12, abs=1.0e-14 + ) + + # The dense-path matrix is the same array. + assert np.array_equal( + reg.regularization_matrix_from(mapper), regularization_matrix.toarray() + ) + + # R^T R is the unscaled precision. + factor = reg.regularization_factor_sparse_from(mapper) + assert (factor.T @ factor).toarray() == pytest.approx( + c * L @ L, rel=1.0e-12, abs=1.0e-14 + ) + + +def test__regularization_matrix__symmetric_positive_definite(): + mapper = _rectangular_mapper_from(shape=(5, 6)) + + regularization_matrix = aa.reg.MaternSparse( + coefficient=0.3, scale=1.5 + ).regularization_matrix_from(mapper) + + assert np.allclose( + regularization_matrix, regularization_matrix.T, rtol=0.0, atol=0.0 + ) + assert np.min(np.linalg.eigvalsh(regularization_matrix)) > 0.0 + np.linalg.cholesky(regularization_matrix) + + +@pytest.mark.parametrize("shape", [(4, 3), (7, 8)]) +def test__log_det__matches_slogdet(shape): + mapper = _rectangular_mapper_from(shape=shape) + + reg = aa.reg.MaternSparse(coefficient=2.5, scale=1.2) + + sign, log_det = np.linalg.slogdet(reg.regularization_matrix_from(mapper)) + + assert sign == 1.0 + assert reg.log_det_regularization_matrix_term_from(mapper) == pytest.approx( + log_det, abs=1.0e-10 + ) + + +def test__coefficient_scaling__matrix_doubles_and_log_det_shifts_by_n_log_2(): + mapper = _rectangular_mapper_from(shape=(4, 3)) + + reg = aa.reg.MaternSparse(coefficient=1.3, scale=2.0) + reg_double = aa.reg.MaternSparse(coefficient=2.6, scale=2.0) + + assert reg_double.regularization_matrix_from(mapper) == pytest.approx( + 2.0 * reg.regularization_matrix_from(mapper), rel=1.0e-14 + ) + assert reg_double.log_det_regularization_matrix_term_from( + mapper + ) - reg.log_det_regularization_matrix_term_from(mapper) == pytest.approx( + mapper.params * np.log(2.0), abs=1.0e-10 + ) + + +def test__weights__equal_hand_built_w_inv_c_l_squared_w_inv(): + mapper = _rectangular_mapper_from(shape=(4, 3)) + + weights = np.random.default_rng(seed=1).uniform(0.5, 2.0, size=mapper.params) + + reg = aa.reg.MaternSparse(coefficient=0.8, scale=2.0, weights=weights) + + L = _laplacian_operator_by_hand_from(mapper=mapper, scale=2.0) + c = _c_by_hand_from(L=L, shape=(4, 3)) + W_inv = np.diag(1.0 / weights) + + regularization_matrix = reg.regularization_matrix_from(mapper) + + assert regularization_matrix == pytest.approx( + 0.8 * W_inv @ (c * L @ L) @ W_inv, rel=1.0e-12, abs=1.0e-14 + ) + + sign, log_det = np.linalg.slogdet(regularization_matrix) + assert reg.log_det_regularization_matrix_term_from(mapper) == pytest.approx( + log_det, abs=1.0e-10 + ) + + with pytest.raises(ValueError): + aa.reg.MaternSparse(weights=np.ones(3)).regularization_matrix_from(mapper) + + +def test__regularization_term__equals_s_h_s(): + mapper = _rectangular_mapper_from(shape=(4, 3)) + + weights = np.random.default_rng(seed=2).uniform(0.5, 2.0, size=mapper.params) + reg = aa.reg.MaternSparse(coefficient=0.8, scale=2.0, weights=weights) + + s = np.random.default_rng(seed=3).normal(size=mapper.params) + + assert reg.regularization_term_from(mapper, s) == pytest.approx( + s @ reg.regularization_matrix_from(mapper) @ s, rel=1.0e-12 + ) + + +def test__regularization_weights__coefficient_times_ones(): + mapper = _rectangular_mapper_from(shape=(4, 3)) + + assert np.array_equal( + aa.reg.MaternSparse(coefficient=3.0).regularization_weights_from(mapper), + 3.0 * np.ones(12), + ) + + +def test__adaptive_rectangular_mesh__supported(): + mapper = _rectangular_mapper_from( + shape=(4, 3), mesh_cls=aa.mesh.RectangularRTUAdaptDensity + ) + + reg = aa.reg.MaternSparse(coefficient=1.0, scale=2.0) + + L = _laplacian_operator_by_hand_from(mapper=mapper, scale=2.0) + c = _c_by_hand_from(L=L, shape=(4, 3)) + + assert reg.regularization_matrix_from(mapper) == pytest.approx( + c * L @ L, rel=1.0e-12, abs=1.0e-14 + ) + + +def test__non_rectangular_mesh__raises_not_implemented(): + mask = aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=0.5, radius=2.5) + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=1) + + image_mesh_grid = aa.image_mesh.Overlay(shape=(5, 5)).image_plane_mesh_grid_from( + mask=mask, adapt_data=None + ) + + mapper = aa.Mapper( + interpolator=aa.mesh.Delaunay( + pixels=image_mesh_grid.shape[0] + ).interpolator_from( + source_plane_data_grid=grid, source_plane_mesh_grid=image_mesh_grid + ), + regularization=aa.reg.Constant(), + ) + + reg = aa.reg.MaternSparse(coefficient=1.0, scale=1.0) + + with pytest.raises(NotImplementedError, match="rectangular"): + reg.regularization_matrix_from(mapper) + + with pytest.raises(NotImplementedError, match="rectangular"): + reg.log_det_regularization_matrix_term_from(mapper) + + +def test__nu_not_one__raises_not_implemented(): + with pytest.raises(NotImplementedError, match="nu = 1"): + aa.reg.MaternSparse(nu=2.0) + + with pytest.raises(NotImplementedError): + aa.reg.MaternSparse(nu=0.5) + + +def test__covariance__unit_variance_at_centre_and_decays_along_row(): + """ + A loose sanity check of the prior the precision encodes, not an exactness test: `C = (c L^2)^-1` + has unit variance at the mesh centre by construction (`c`), and the correlation with the centre + falls monotonically with distance along the centre row, as a Matérn kernel does. Edge effects of + the reflecting boundary mean `C` is not exactly the stationary Matérn kernel, so no values beyond + these are asserted. + """ + shape = (11, 11) + mapper = _rectangular_mapper_from(shape=shape) + + reg = aa.reg.MaternSparse(coefficient=1.0, scale=1.5) + + covariance = np.linalg.inv(reg.regularization_matrix_from(mapper)) + + centre_row, centre_col = shape[0] // 2, shape[1] // 2 + centre = centre_row * shape[1] + centre_col + + assert covariance[centre, centre] == pytest.approx(1.0, rel=1.0e-10) + + along_row = covariance[ + centre, centre_row * shape[1] + np.arange(centre_col, shape[1]) + ] + + assert np.all(np.diff(along_row) < 0.0) + assert np.all(along_row > 0.0) + + +def test__matern_sparse__imports_and_runs_with_jax_blocked(): + """ + The module must work in the `unittest-nojax` leg; a subprocess hides `jax`, `jaxlib` and `nufftax` + from the import system (the same pattern as `test_matrix_free_util.py`), imports the class and + builds the matrix and its log-det. + """ + script = textwrap.dedent( + """ + import importlib.machinery + import sys + + _find_spec = importlib.machinery.PathFinder.find_spec.__func__ + + + def _find_spec_blocking_jax(cls, name, path=None, target=None): + if name.split(".")[0] in ("jax", "jaxlib", "nufftax"): + return None + return _find_spec(cls, name, path, target) + + + importlib.machinery.PathFinder.find_spec = classmethod(_find_spec_blocking_jax) + + try: + import jax + except ModuleNotFoundError: + pass + else: + raise AssertionError("jax was not blocked") + + import numpy as np + + import autoarray as aa + from autoarray.inversion.regularization.matern_sparse import MaternSparse + + mask = aa.Mask2D.circular(shape_native=(8, 8), pixel_scales=1.0, radius=3.0) + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=1) + mapper = aa.Mapper( + interpolator=aa.mesh.RectangularUniform(shape=(4, 4)).interpolator_from( + source_plane_data_grid=grid, source_plane_mesh_grid=None + ), + regularization=aa.reg.Constant(), + ) + + reg = MaternSparse(coefficient=1.0, scale=1.0) + H = reg.regularization_matrix_from(mapper) + assert np.isclose( + reg.log_det_regularization_matrix_term_from(mapper), np.linalg.slogdet(H)[1] + ) + + assert not any( + name.split(".")[0] in ("jax", "jaxlib", "nufftax") for name in sys.modules + ) + """ + ) + + environment = dict(os.environ) + environment["PYAUTO_DISABLE_JAX"] = "1" + + subprocess.run( + [sys.executable, "-c", script], check=True, env=environment, timeout=600 + )