From 8aa9513dd09d0525f5863781b681232fa7016070 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 17 Jul 2026 09:43:20 +0100 Subject: [PATCH] Add masked-grid derivative operators, coarse-mesh interpolation and mask regularizations (potential correction phase 1) Ports the generic linear-algebra layer of the gravitational-imaging (potential correction) technique of Cao et al. 2025 (https://github.com/caoxiaoyue/lensing_potential_correction) into autoarray: sparse finite-difference derivative operators on masked 2D grids, coarse-mesh bilinear interpolation matrices, and the CurvatureMask / FourthOrderMask regularizations. Parity-certified against the original; cite via https://github.com/caoxiaoyue/potential_correction_paper. Phase 1 of PyAutoLabs/PyAutoLens#618. Co-Authored-By: Claude Fable 5 --- .../inversion/regularization/__init__.py | 2 + .../regularization/curvature_mask.py | 108 +++ .../regularization/fourth_order_mask.py | 111 ++++ autoarray/operators/coarse_interp_util.py | 265 ++++++++ autoarray/operators/derivative_util.py | 616 ++++++++++++++++++ autoarray/util/__init__.py | 2 + .../regularizations/test_curvature_mask.py | 74 +++ .../regularizations/test_fourth_order_mask.py | 70 ++ .../operators/test_coarse_interp_util.py | 181 +++++ .../operators/test_derivative_util.py | 159 +++++ 10 files changed, 1588 insertions(+) create mode 100644 autoarray/inversion/regularization/curvature_mask.py create mode 100644 autoarray/inversion/regularization/fourth_order_mask.py create mode 100644 autoarray/operators/coarse_interp_util.py create mode 100644 autoarray/operators/derivative_util.py create mode 100644 test_autoarray/inversion/regularizations/test_curvature_mask.py create mode 100644 test_autoarray/inversion/regularizations/test_fourth_order_mask.py create mode 100644 test_autoarray/operators/test_coarse_interp_util.py create mode 100644 test_autoarray/operators/test_derivative_util.py diff --git a/autoarray/inversion/regularization/__init__.py b/autoarray/inversion/regularization/__init__.py index aebffb8b0..71804aeca 100644 --- a/autoarray/inversion/regularization/__init__.py +++ b/autoarray/inversion/regularization/__init__.py @@ -7,6 +7,8 @@ from .adapt_split import AdaptSplit from .brightness_zeroth import BrightnessZeroth from .adapt_split_zeroth import AdaptSplitZeroth +from .curvature_mask import CurvatureMask +from .fourth_order_mask import FourthOrderMask from .gaussian_kernel import GaussianKernel from .exponential_kernel import ExponentialKernel from .matern_kernel import MaternKernel diff --git a/autoarray/inversion/regularization/curvature_mask.py b/autoarray/inversion/regularization/curvature_mask.py new file mode 100644 index 000000000..ed97222fd --- /dev/null +++ b/autoarray/inversion/regularization/curvature_mask.py @@ -0,0 +1,108 @@ +from __future__ import annotations +import numpy as np +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from autoarray.inversion.linear_obj.linear_obj import LinearObj + +from autoarray.inversion.regularization.abstract import AbstractRegularization +from autoarray.operators import derivative_util + + +def curvature_reg_matrix_via_mask_from(mask, pixel_scale: float = 1.0) -> np.ndarray: + """ + The curvature regularization matrix of the unmasked pixels of a + rectangular masked 2D grid, H = Hxx^T Hxx + Hyy^T Hyy, where Hxx / Hyy are + forward second-difference operators which degrade gracefully to first / + zeroth order at the mask edge (see + ``derivative_util.forward_difference_operators_from``). + + This is the curvature regularization scheme used by the + gravitational-imaging (potential correction) technique of Cao et al. 2025 + (https://github.com/caoxiaoyue/lensing_potential_correction; cite via + https://github.com/caoxiaoyue/potential_correction_paper), applied to the + pixelized corrections of the lensing potential defined on a coarse + rectangular mesh. + + Parameters + ---------- + mask + The 2D bool mask (``True`` = masked) of the rectangular grid whose + unmasked pixels are regularized. + pixel_scale + The finite-difference step; regularization matrices are + conventionally built with 1.0, the coefficient absorbing the scale. + + Returns + ------- + The [n_unmasked, n_unmasked] regularization matrix. + """ + return derivative_util.forward_difference_reg_matrix_from( + mask=mask, pixel_scale=pixel_scale, max_order=2 + ).toarray() + + +class CurvatureMask(AbstractRegularization): + def __init__(self, coefficient: float = 1.0): + """ + Curvature regularization on the unmasked pixels of a rectangular + masked 2D grid. + + Each unmasked pixel is regularized with a forward second-difference + stencil along both grid directions, degrading to first / zeroth order + where the mask edge truncates the stencil, penalising curvature in + the reconstructed solution. This contrasts mapper-based schemes + (e.g. ``Constant``), which regularize via mesh-neighbour differences: + here the linear object is defined on a masked rectangular grid and + must expose its ``mask``. + + This is the curvature regularization scheme of the + gravitational-imaging (potential correction) technique, applied to + pixelized corrections of the lensing potential; it is ported from the + ``potential_correction`` package of Cao et al. 2025 + (https://github.com/caoxiaoyue/lensing_potential_correction). If you + use it in your research, please cite Cao et al. 2025; citation + materials are provided at + https://github.com/caoxiaoyue/potential_correction_paper. + + Parameters + ---------- + coefficient + The regularization coefficient which multiplies the matrix, + setting the strength of the smoothing. + """ + self.coefficient = coefficient + + super().__init__() + + def regularization_weights_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray: + """ + Returns the regularization weights of this regularization scheme, + which are equal for every parameter. + + Parameters + ---------- + linear_obj + The linear object which uses these weights when performing + regularization. + """ + return self.coefficient * xp.ones(linear_obj.params) + + def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray: + """ + Returns the regularization matrix with shape [pixels, pixels]. + + Parameters + ---------- + linear_obj + The linear object which uses this matrix to perform + regularization. It must expose the ``mask`` of the rectangular + masked 2D grid its parameters are defined on. + + Returns + ------- + The regularization matrix. + """ + return self.coefficient * curvature_reg_matrix_via_mask_from( + mask=linear_obj.mask + ) diff --git a/autoarray/inversion/regularization/fourth_order_mask.py b/autoarray/inversion/regularization/fourth_order_mask.py new file mode 100644 index 000000000..9e960b3f0 --- /dev/null +++ b/autoarray/inversion/regularization/fourth_order_mask.py @@ -0,0 +1,111 @@ +from __future__ import annotations +import numpy as np +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from autoarray.inversion.linear_obj.linear_obj import LinearObj + +from autoarray.inversion.regularization.abstract import AbstractRegularization +from autoarray.operators import derivative_util + + +def fourth_order_reg_matrix_via_mask_from( + mask, pixel_scale: float = 1.0 +) -> np.ndarray: + """ + The fourth-order regularization matrix of the unmasked pixels of a + rectangular masked 2D grid, H = H4x^T H4x + H4y^T H4y, where H4x / H4y + are forward fourth-difference operators which degrade gracefully to + third / second / first / zeroth order at the mask edge (see + ``derivative_util.forward_difference_operators_from``). + + This is the fourth-order regularization scheme used by the + gravitational-imaging (potential correction) technique of Cao et al. 2025 + (https://github.com/caoxiaoyue/lensing_potential_correction; cite via + https://github.com/caoxiaoyue/potential_correction_paper), applied to the + pixelized corrections of the lensing potential defined on a coarse + rectangular mesh. + + Parameters + ---------- + mask + The 2D bool mask (``True`` = masked) of the rectangular grid whose + unmasked pixels are regularized. + pixel_scale + The finite-difference step; regularization matrices are + conventionally built with 1.0, the coefficient absorbing the scale. + + Returns + ------- + The [n_unmasked, n_unmasked] regularization matrix. + """ + return derivative_util.forward_difference_reg_matrix_from( + mask=mask, pixel_scale=pixel_scale, max_order=4 + ).toarray() + + +class FourthOrderMask(AbstractRegularization): + def __init__(self, coefficient: float = 1.0): + """ + Fourth-order regularization on the unmasked pixels of a rectangular + masked 2D grid. + + Each unmasked pixel is regularized with a forward fourth-difference + stencil along both grid directions, degrading to third / second / + first / zeroth order where the mask edge truncates the stencil. + Penalising the fourth derivative permits solutions with curvature + (e.g. localised perturbations) while still suppressing high-frequency + noise, making it a weaker prior than ``CurvatureMask``. The linear + object regularized must be defined on a masked rectangular grid and + expose its ``mask``. + + This is the fourth-order regularization scheme of the + gravitational-imaging (potential correction) technique, applied to + pixelized corrections of the lensing potential; it is ported from the + ``potential_correction`` package of Cao et al. 2025 + (https://github.com/caoxiaoyue/lensing_potential_correction). If you + use it in your research, please cite Cao et al. 2025; citation + materials are provided at + https://github.com/caoxiaoyue/potential_correction_paper. + + Parameters + ---------- + coefficient + The regularization coefficient which multiplies the matrix, + setting the strength of the smoothing. + """ + self.coefficient = coefficient + + super().__init__() + + def regularization_weights_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray: + """ + Returns the regularization weights of this regularization scheme, + which are equal for every parameter. + + Parameters + ---------- + linear_obj + The linear object which uses these weights when performing + regularization. + """ + return self.coefficient * xp.ones(linear_obj.params) + + def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray: + """ + Returns the regularization matrix with shape [pixels, pixels]. + + Parameters + ---------- + linear_obj + The linear object which uses this matrix to perform + regularization. It must expose the ``mask`` of the rectangular + masked 2D grid its parameters are defined on. + + Returns + ------- + The regularization matrix. + """ + return self.coefficient * fourth_order_reg_matrix_via_mask_from( + mask=linear_obj.mask + ) diff --git a/autoarray/operators/coarse_interp_util.py b/autoarray/operators/coarse_interp_util.py new file mode 100644 index 000000000..e9477babc --- /dev/null +++ b/autoarray/operators/coarse_interp_util.py @@ -0,0 +1,265 @@ +""" +Sparse bilinear interpolation from a coarse regular mesh to the unmasked +pixels of a finer 2D grid. + +Used by the gravitational-imaging (potential correction) technique, where +pixelized corrections to the lensing potential are defined on a mesh a factor +coarser than the data grid and interpolated onto it. The interpolation is a +sparse matrix of shape [n_unmasked_fine_pixels, n_unmasked_coarse_pixels], +each row holding the bilinear weights of the four coarse-mesh corners of the +box enclosing (or nearest to) the fine pixel. + +The implementations in this module are ported from the ``potential_correction`` +package of Cao et al. 2025 (https://github.com/caoxiaoyue/lensing_potential_correction). +If you use this functionality in your research, please cite Cao et al. 2025; +citation materials are provided at +https://github.com/caoxiaoyue/potential_correction_paper. +""" + +import numpy as np +from scipy.sparse import csr_matrix + +from autoarray import exc +from autoarray import numba_util + + +@numba_util.jit() +def binned_image_from(arr, bin_factor=1): + """ + Bins a 2D array by averaging over ``bin_factor`` x ``bin_factor`` blocks. + """ + n0 = arr.shape[0] // bin_factor + n1 = arr.shape[1] // bin_factor + binned_arr = np.zeros((n0, n1), dtype=np.float64) + n_per_bin = bin_factor**2 + for i in range(n0): + for j in range(n1): + for m in range(i * bin_factor, (i + 1) * bin_factor): + for n in range(j * bin_factor, (j + 1) * bin_factor): + binned_arr[i, j] += arr[m, n] + binned_arr[i, j] = binned_arr[i, j] / n_per_bin + return binned_arr + + +@numba_util.jit() +def binned_mask_from(mask, bin_factor): + """ + The mask of a grid binned coarser by ``bin_factor``: a coarse pixel is + unmasked only if *every* fine pixel inside it is unmasked. + + The fine mask's shape must be divisible by ``bin_factor``. + """ + unmask = (~mask).astype(np.float64) + binned_unmask = binned_image_from(unmask, bin_factor) + return ~np.isclose(binned_unmask, 1.0) + + +@numba_util.jit() +def interp_box_mask_from(mask): + """ + The mask of the interpolation boxes of a coarse mesh: box (i, j) — whose + corners are mesh pixels (i, j), (i, j+1), (i+1, j), (i+1, j+1) — is + unmasked only if all four corners are unmasked. + + Returns a bool array of shape [n0 - 1, n1 - 1] for a mask of shape + [n0, n1]. + """ + itp_box_mask = np.ones((mask.shape[0] - 1, mask.shape[1] - 1), dtype="bool") + for i in range(mask.shape[0] - 1): + for j in range(mask.shape[1] - 1): + if ( + (~mask[i, j]) + and (~mask[i + 1, j]) + and (~mask[i, j + 1]) + and (~mask[i + 1, j + 1]) + ): + itp_box_mask[i, j] = False + return itp_box_mask + + +@numba_util.jit() +def bilinear_weights_from_box(box_x, box_y, position=(0.0, 0.0)): + """ + The bilinear interpolation (or extrapolation) weights of ``position`` + inside the square box with corner coordinates ``box_x`` / ``box_y``. + + Parameters + ---------- + box_x + The 4 corner x coordinates, in [top-left, top-right, bottom-left, + bottom-right] order. + box_y + The 4 corner y coordinates, in the same order. + position + The (y, x) coordinates at which the weights are evaluated. + + Returns + ------- + An array of shape [4,] holding the weights in the same corner order. + """ + y, x = position + box_size = box_x[1] - box_x[0] + wx = (x - box_x[0]) / box_size + wy = (y - box_y[2]) / box_size + + weight_top_left = (1 - wx) * wy + weight_top_right = wx * wy + weight_bottom_left = (1 - wx) * (1 - wy) + weight_bottom_right = wx * (1 - wy) + + return np.array( + [weight_top_left, weight_top_right, weight_bottom_left, weight_bottom_right] + ) + + +@numba_util.jit() +def coarse_interp_triplets_from( + mask_itp_box, + xc_itp_box, + yc_itp_box, + xgrid_fine_1d, + ygrid_fine_1d, + xgrid_coarse, + ygrid_coarse, + mask_coarse, +): + """ + The (rows, cols, values) sparse triplets of the coarse-to-fine bilinear + interpolation matrix (see ``coarse_interp_matrix_from``). + + Each unmasked fine pixel is assigned the unmasked interpolation box whose + centre is nearest to it (searching outward if the naively nearest box is + masked), and receives the bilinear weights of that box's four corners. + + Parameters + ---------- + mask_itp_box + The interpolation-box mask (see ``interp_box_mask_from``). + xc_itp_box, yc_itp_box + The 2D x / y coordinates of the interpolation-box centres. + xgrid_fine_1d, ygrid_fine_1d + The 1D x / y coordinates of the unmasked fine-grid pixels. + xgrid_coarse, ygrid_coarse + The 2D x / y coordinates of the coarse mesh. + mask_coarse + The coarse-mesh mask. + """ + unmask_coarse = ~mask_coarse + i_indices_unmasked, j_indices_unmasked = np.where(unmask_coarse) + n_unmasked_coarse_pixels = len(i_indices_unmasked) + index_dict_coarse = {} + for count in range(n_unmasked_coarse_pixels): + i, j = i_indices_unmasked[count], j_indices_unmasked[count] + index_dict_coarse[(i, j)] = count + + n_unmasked_fine_pixels = len(xgrid_fine_1d) + rows_itp_mat = np.full(n_unmasked_fine_pixels * 4, -1, dtype=np.int64) + cols_itp_mat = np.full(n_unmasked_fine_pixels * 4, -1, dtype=np.int64) + data_itp_mat = np.full(n_unmasked_fine_pixels * 4, 0.0, dtype=np.float64) + count_itp_mat = 0 + for count in range(n_unmasked_fine_pixels): + this_x_fine = xgrid_fine_1d[count] + this_y_fine = ygrid_fine_1d[count] + + j_min = np.argmin(np.abs(xc_itp_box[0, :] - this_x_fine)) + i_min = np.argmin(np.abs(yc_itp_box[:, 0] - this_y_fine)) + + if ~mask_itp_box[i_min, j_min]: + i = i_min + j = j_min + else: + search_width = 2 + i = -1 + j = -1 + dist_tmp = 1e8 + while i == -1: + m_lo = max(0, i_min - search_width) + m_hi = min(mask_itp_box.shape[0], i_min + search_width + 1) + n_lo = max(0, j_min - search_width) + n_hi = min(mask_itp_box.shape[1], j_min + search_width + 1) + for m in range(m_lo, m_hi): + for n in range(n_lo, n_hi): + if ~mask_itp_box[m, n]: + this_dist = np.sqrt( + (xc_itp_box[m, n] - this_x_fine) ** 2 + + (yc_itp_box[m, n] - this_y_fine) ** 2 + ) + if this_dist < dist_tmp: + dist_tmp = this_dist + i = m + j = n + search_width += 1 + + itp_box_corners_x = ( + xgrid_coarse[i, j], + xgrid_coarse[i, j + 1], + xgrid_coarse[i + 1, j], + xgrid_coarse[i + 1, j + 1], + ) + itp_box_corners_y = ( + ygrid_coarse[i, j], + ygrid_coarse[i, j + 1], + ygrid_coarse[i + 1, j], + ygrid_coarse[i + 1, j + 1], + ) + itp_weights = bilinear_weights_from_box( + itp_box_corners_x, itp_box_corners_y, position=(this_y_fine, this_x_fine) + ) + itp_idx = ( + index_dict_coarse[(i, j)], + index_dict_coarse[(i, j + 1)], + index_dict_coarse[(i + 1, j)], + index_dict_coarse[(i + 1, j + 1)], + ) + + for k in range(4): + rows_itp_mat[count_itp_mat] = count + cols_itp_mat[count_itp_mat] = itp_idx[k] + data_itp_mat[count_itp_mat] = itp_weights[k] + count_itp_mat += 1 + + return rows_itp_mat, cols_itp_mat, data_itp_mat + + +def coarse_interp_matrix_from( + mask_itp_box, + xc_itp_box, + yc_itp_box, + xgrid_fine_1d, + ygrid_fine_1d, + xgrid_coarse, + ygrid_coarse, + mask_coarse, +): + """ + The sparse bilinear interpolation matrix mapping a vector on the unmasked + pixels of a coarse regular mesh to the unmasked pixels of a finer grid. + + See ``coarse_interp_triplets_from`` for the parameters. + + Returns + ------- + A ``scipy.sparse.csr_matrix`` of shape + [n_unmasked_fine_pixels, n_unmasked_coarse_pixels]. + """ + mask_itp_box = np.asarray(mask_itp_box) + if np.count_nonzero(~mask_itp_box) == 0: + raise exc.MeshException( + "The coarse mesh has no unmasked interpolation box (no 2x2 block of " + "unmasked coarse pixels) — it is too sparse for the fine grid. " + "Decrease the coarsening factor." + ) + rows, cols, data = coarse_interp_triplets_from( + mask_itp_box, + np.asarray(xc_itp_box), + np.asarray(yc_itp_box), + np.asarray(xgrid_fine_1d), + np.asarray(ygrid_fine_1d), + np.asarray(xgrid_coarse), + np.asarray(ygrid_coarse), + mask_coarse, + ) + n_unmasked_coarse = np.count_nonzero(~np.asarray(mask_coarse)) + return csr_matrix( + (data, (rows, cols)), shape=(len(xgrid_fine_1d), n_unmasked_coarse) + ) diff --git a/autoarray/operators/derivative_util.py b/autoarray/operators/derivative_util.py new file mode 100644 index 000000000..62bf96471 --- /dev/null +++ b/autoarray/operators/derivative_util.py @@ -0,0 +1,616 @@ +""" +Sparse finite-difference derivative operators on masked 2D grids. + +Each operator is a sparse matrix of shape [n_unmasked_pixels, n_unmasked_pixels] +which, acting on a 1D vector of values on the unmasked pixels of a 2D grid, +returns the finite-difference derivative of those values on the same pixels. +They are the building blocks of the gravitational-imaging (potential +correction) technique, where derivatives of the lensing potential and its +pixelized corrections are taken on masked image grids. + +The implementations in this module are 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. + +Conventions (matching the rest of autoarray): + +- ``mask`` is a bool 2D array where ``True`` means *masked* (excluded). +- Grid coordinates are (y, x), with y *decreasing* as the array index along + axis 0 increases, and x increasing with the index along axis 1. The y-step + of ``-pixel_scale`` used below encodes this. +- ``diff_types`` is an int array of shape [total_pixels, total_pixels, 2] + whose last dimension is (y, x) and whose values select the finite-difference + scheme available at each unmasked pixel: 0 = backward, 1 = central, + 2 = forward, -1 = none (pixel cannot support a derivative in the direction). +""" + +import numpy as np +from scipy.sparse import csr_matrix + +from autoarray import exc +from autoarray import numba_util + + +@numba_util.jit() +def clean_mask_iteration_from(mask): + """ + One iteration of mask cleaning: masks every unmasked pixel which cannot + support a first-derivative finite-difference scheme (backward, central or + forward) in *both* the y and x directions, and records the scheme + available at each surviving pixel. + + Parameters + ---------- + mask + The 2D bool mask (``True`` = masked). Must have shape of at least + (3, 3) so a difference scheme can exist. + + Returns + ------- + The cleaned mask and the ``diff_types`` array (see module docstring). + """ + in_mask = np.copy(mask) + in_mask = in_mask.astype("bool") + out_mask = np.ones_like(in_mask, dtype="bool") + n1, n2 = in_mask.shape + diff_types = np.full((n1, n2, 2), -1, dtype="int") + + for i in range(n1): + for j in range(n2): + if not in_mask[i, j]: + if_remove_y = True + if_remove_x = True + + if i == 0: + if ~in_mask[i + 1, j] and ~in_mask[i + 2, j]: + if_remove_y = False + diff_types[i, j, 0] = 2 + elif i == 1: + if ~in_mask[i - 1, j] and ~in_mask[i + 1, j]: + if_remove_y = False + diff_types[i, j, 0] = 1 + elif ~in_mask[i + 1, j] and ~in_mask[i + 2, j]: + if_remove_y = False + diff_types[i, j, 0] = 2 + elif i == n1 - 1: + if ~in_mask[i - 1, j] and ~in_mask[i - 2, j]: + if_remove_y = False + diff_types[i, j, 0] = 0 + elif i == n1 - 2: + if ~in_mask[i - 1, j] and ~in_mask[i + 1, j]: + if_remove_y = False + diff_types[i, j, 0] = 1 + elif ~in_mask[i - 1, j] and ~in_mask[i - 2, j]: + if_remove_y = False + diff_types[i, j, 0] = 0 + else: + if ~in_mask[i - 1, j] and ~in_mask[i + 1, j]: + if_remove_y = False + diff_types[i, j, 0] = 1 + elif ~in_mask[i - 1, j] and ~in_mask[i - 2, j]: + if_remove_y = False + diff_types[i, j, 0] = 0 + elif ~in_mask[i + 1, j] and ~in_mask[i + 2, j]: + if_remove_y = False + diff_types[i, j, 0] = 2 + + if j == 0: + if ~in_mask[i, j + 1] and ~in_mask[i, j + 2]: + if_remove_x = False + diff_types[i, j, 1] = 2 + elif j == 1: + if ~in_mask[i, j - 1] and ~in_mask[i, j + 1]: + if_remove_x = False + diff_types[i, j, 1] = 1 + elif ~in_mask[i, j + 1] and ~in_mask[i, j + 2]: + if_remove_x = False + diff_types[i, j, 1] = 2 + elif j == n2 - 1: + if ~in_mask[i, j - 1] and ~in_mask[i, j - 2]: + if_remove_x = False + diff_types[i, j, 1] = 0 + elif j == n2 - 2: + if ~in_mask[i, j - 1] and ~in_mask[i, j + 1]: + if_remove_x = False + diff_types[i, j, 1] = 1 + elif ~in_mask[i, j - 1] and ~in_mask[i, j - 2]: + if_remove_x = False + diff_types[i, j, 1] = 0 + else: + if ~in_mask[i, j - 1] and ~in_mask[i, j + 1]: + if_remove_x = False + diff_types[i, j, 1] = 1 + elif ~in_mask[i, j - 1] and ~in_mask[i, j - 2]: + if_remove_x = False + diff_types[i, j, 1] = 0 + elif ~in_mask[i, j + 1] and ~in_mask[i, j + 2]: + if_remove_x = False + diff_types[i, j, 1] = 2 + + if not (if_remove_y or if_remove_x): + out_mask[i, j] = False + + return out_mask, diff_types + + +def cleaned_mask_from(mask, max_iter: int = 50): + """ + Iteratively cleans a mask until every unmasked pixel supports a + first-derivative finite-difference scheme in both directions (see + ``clean_mask_iteration_from``), returning the cleaned mask and the + per-pixel scheme types. + + Parameters + ---------- + mask + The 2D bool mask (``True`` = masked). + max_iter + The maximum number of cleaning iterations before raising. + + Returns + ------- + The cleaned mask and the ``diff_types`` array (see module docstring). + """ + old_mask = np.copy(np.asarray(mask)) + clean_success = False + + for _ in range(max_iter): + new_mask, diff_types = clean_mask_iteration_from(old_mask) + if (new_mask == old_mask).all(): + clean_success = True + break + old_mask = new_mask + + if not clean_success: + raise exc.MaskException( + f"The mask is not fully cleaned after {max_iter} iterations" + ) + + return new_mask, diff_types + + +def _diff_types_of_cleaned_mask_from(mask): + """ + Returns the ``diff_types`` of a mask, raising if the mask is not already + cleaned (derivative operators are only defined on cleaned masks). + """ + mask = np.asarray(mask).astype(bool) + new_mask, diff_types = cleaned_mask_from(mask) + if not (new_mask == mask).all(): + raise exc.MaskException( + "The mask has not been fully cleaned: one or more unmasked pixels do " + "not support a finite-difference scheme. Clean it first via " + "cleaned_mask_from(mask)." + ) + return mask, diff_types + + +@numba_util.jit() +def derivative_1st_triplets_from(mask, diff_types, dpix=1.0): + """ + The (rows, cols, values) sparse triplets of the first-derivative operators + Hx and Hy of a cleaned mask (see ``derivative_1st_operators_from``). + """ + unmask = ~mask + i_indices_unmasked, j_indices_unmasked = np.where(unmask) + n_unmasked_pixels = len(i_indices_unmasked) + + rows_hx = np.full(n_unmasked_pixels * 2, -1, dtype=np.int64) + cols_hx = np.full(n_unmasked_pixels * 2, -1, dtype=np.int64) + data_hx = np.full(n_unmasked_pixels * 2, 0.0, dtype=np.float64) + rows_hy = np.full(n_unmasked_pixels * 2, -1, dtype=np.int64) + cols_hy = np.full(n_unmasked_pixels * 2, -1, dtype=np.int64) + data_hy = np.full(n_unmasked_pixels * 2, 0.0, dtype=np.float64) + + # y decreases as the index along axis 0 increases; x increases with axis 1. + step_y = -1.0 * dpix + step_x = 1.0 * dpix + + index_dict = {} + for count in range(n_unmasked_pixels): + i, j = i_indices_unmasked[count], j_indices_unmasked[count] + index_dict[(i, j)] = count + + count_sparse_hy = 0 + count_sparse_hx = 0 + for count in range(n_unmasked_pixels): + i, j = i_indices_unmasked[count], j_indices_unmasked[count] + + if diff_types[i, j, 0] == 0: + rows_hy[count_sparse_hy] = count + cols_hy[count_sparse_hy] = index_dict[(i - 1, j)] + data_hy[count_sparse_hy] = -1.0 / step_y + count_sparse_hy += 1 + rows_hy[count_sparse_hy] = count + cols_hy[count_sparse_hy] = index_dict[(i, j)] + data_hy[count_sparse_hy] = 1.0 / step_y + count_sparse_hy += 1 + elif diff_types[i, j, 0] == 1: + rows_hy[count_sparse_hy] = count + cols_hy[count_sparse_hy] = index_dict[(i - 1, j)] + data_hy[count_sparse_hy] = -1.0 / (2 * step_y) + count_sparse_hy += 1 + rows_hy[count_sparse_hy] = count + cols_hy[count_sparse_hy] = index_dict[(i + 1, j)] + data_hy[count_sparse_hy] = 1.0 / (2 * step_y) + count_sparse_hy += 1 + elif diff_types[i, j, 0] == 2: + rows_hy[count_sparse_hy] = count + cols_hy[count_sparse_hy] = index_dict[(i, j)] + data_hy[count_sparse_hy] = -1.0 / step_y + count_sparse_hy += 1 + rows_hy[count_sparse_hy] = count + cols_hy[count_sparse_hy] = index_dict[(i + 1, j)] + data_hy[count_sparse_hy] = 1.0 / step_y + count_sparse_hy += 1 + + if diff_types[i, j, 1] == 0: + rows_hx[count_sparse_hx] = count + cols_hx[count_sparse_hx] = index_dict[(i, j - 1)] + data_hx[count_sparse_hx] = -1.0 / step_x + count_sparse_hx += 1 + rows_hx[count_sparse_hx] = count + cols_hx[count_sparse_hx] = index_dict[(i, j)] + data_hx[count_sparse_hx] = 1.0 / step_x + count_sparse_hx += 1 + elif diff_types[i, j, 1] == 1: + rows_hx[count_sparse_hx] = count + cols_hx[count_sparse_hx] = index_dict[(i, j - 1)] + data_hx[count_sparse_hx] = -1.0 / (2 * step_x) + count_sparse_hx += 1 + rows_hx[count_sparse_hx] = count + cols_hx[count_sparse_hx] = index_dict[(i, j + 1)] + data_hx[count_sparse_hx] = 1.0 / (2 * step_x) + count_sparse_hx += 1 + elif diff_types[i, j, 1] == 2: + rows_hx[count_sparse_hx] = count + cols_hx[count_sparse_hx] = index_dict[(i, j)] + data_hx[count_sparse_hx] = -1.0 / step_x + count_sparse_hx += 1 + rows_hx[count_sparse_hx] = count + cols_hx[count_sparse_hx] = index_dict[(i, j + 1)] + data_hx[count_sparse_hx] = 1.0 / step_x + count_sparse_hx += 1 + + return rows_hx, cols_hx, data_hx, rows_hy, cols_hy, data_hy + + +def derivative_1st_operators_from(mask, pixel_scale: float = 1.0): + """ + The sparse first-derivative operators (Hy, Hx) of a cleaned mask. + + Hy (Hx) has shape [n_unmasked_pixels, n_unmasked_pixels]; acting on a 1D + vector of values on the unmasked pixels it returns their first + y (x) derivative, using the central / backward / forward scheme each + pixel's neighbours permit. + + Parameters + ---------- + mask + The cleaned 2D bool mask (``True`` = masked); raises if not cleaned. + pixel_scale + The pixel size (e.g. in arcsec) setting the finite-difference step. + + Returns + ------- + The (Hy, Hx) operators as ``scipy.sparse.csr_matrix``. + """ + mask, diff_types = _diff_types_of_cleaned_mask_from(mask) + rows_hx, cols_hx, data_hx, rows_hy, cols_hy, data_hy = ( + derivative_1st_triplets_from(mask, diff_types, dpix=pixel_scale) + ) + + n_unmasked = np.count_nonzero(~mask) + Hx = csr_matrix((data_hx, (rows_hx, cols_hx)), shape=(n_unmasked, n_unmasked)) + Hy = csr_matrix((data_hy, (rows_hy, cols_hy)), shape=(n_unmasked, n_unmasked)) + return Hy, Hx + + +@numba_util.jit() +def derivative_2nd_triplets_from(mask, diff_types, dpix=1.0): + """ + The (rows, cols, values) sparse triplets of the second-derivative + operators Hxx and Hyy of a cleaned mask (see + ``derivative_2nd_operators_from``). + """ + unmask = ~mask + i_indices_unmasked, j_indices_unmasked = np.where(unmask) + n_unmasked_pixels = len(i_indices_unmasked) + + rows_hxx = np.full(n_unmasked_pixels * 3, -1, dtype=np.int64) + cols_hxx = np.full(n_unmasked_pixels * 3, -1, dtype=np.int64) + data_hxx = np.full(n_unmasked_pixels * 3, 0.0, dtype=np.float64) + rows_hyy = np.full(n_unmasked_pixels * 3, -1, dtype=np.int64) + cols_hyy = np.full(n_unmasked_pixels * 3, -1, dtype=np.int64) + data_hyy = np.full(n_unmasked_pixels * 3, 0.0, dtype=np.float64) + + step_y = -1.0 * dpix + step_x = 1.0 * dpix + + index_dict = {} + for count in range(n_unmasked_pixels): + i, j = i_indices_unmasked[count], j_indices_unmasked[count] + index_dict[(i, j)] = count + + count_sparse_hyy = 0 + count_sparse_hxx = 0 + for count in range(n_unmasked_pixels): + i, j = i_indices_unmasked[count], j_indices_unmasked[count] + + if diff_types[i, j, 0] == 0: + rows_hyy[count_sparse_hyy] = count + cols_hyy[count_sparse_hyy] = index_dict[(i - 2, j)] + data_hyy[count_sparse_hyy] = 1.0 / step_y**2 + count_sparse_hyy += 1 + rows_hyy[count_sparse_hyy] = count + cols_hyy[count_sparse_hyy] = index_dict[(i - 1, j)] + data_hyy[count_sparse_hyy] = -2.0 / step_y**2 + count_sparse_hyy += 1 + rows_hyy[count_sparse_hyy] = count + cols_hyy[count_sparse_hyy] = index_dict[(i, j)] + data_hyy[count_sparse_hyy] = 1.0 / step_y**2 + count_sparse_hyy += 1 + elif diff_types[i, j, 0] == 1: + rows_hyy[count_sparse_hyy] = count + cols_hyy[count_sparse_hyy] = index_dict[(i - 1, j)] + data_hyy[count_sparse_hyy] = 1.0 / step_y**2 + count_sparse_hyy += 1 + rows_hyy[count_sparse_hyy] = count + cols_hyy[count_sparse_hyy] = index_dict[(i, j)] + data_hyy[count_sparse_hyy] = -2.0 / step_y**2 + count_sparse_hyy += 1 + rows_hyy[count_sparse_hyy] = count + cols_hyy[count_sparse_hyy] = index_dict[(i + 1, j)] + data_hyy[count_sparse_hyy] = 1.0 / step_y**2 + count_sparse_hyy += 1 + elif diff_types[i, j, 0] == 2: + rows_hyy[count_sparse_hyy] = count + cols_hyy[count_sparse_hyy] = index_dict[(i, j)] + data_hyy[count_sparse_hyy] = 1.0 / step_y**2 + count_sparse_hyy += 1 + rows_hyy[count_sparse_hyy] = count + cols_hyy[count_sparse_hyy] = index_dict[(i + 1, j)] + data_hyy[count_sparse_hyy] = -2.0 / step_y**2 + count_sparse_hyy += 1 + rows_hyy[count_sparse_hyy] = count + cols_hyy[count_sparse_hyy] = index_dict[(i + 2, j)] + data_hyy[count_sparse_hyy] = 1.0 / step_y**2 + count_sparse_hyy += 1 + + if diff_types[i, j, 1] == 0: + rows_hxx[count_sparse_hxx] = count + cols_hxx[count_sparse_hxx] = index_dict[(i, j - 2)] + data_hxx[count_sparse_hxx] = 1.0 / step_x**2 + count_sparse_hxx += 1 + rows_hxx[count_sparse_hxx] = count + cols_hxx[count_sparse_hxx] = index_dict[(i, j - 1)] + data_hxx[count_sparse_hxx] = -2.0 / step_x**2 + count_sparse_hxx += 1 + rows_hxx[count_sparse_hxx] = count + cols_hxx[count_sparse_hxx] = index_dict[(i, j)] + data_hxx[count_sparse_hxx] = 1.0 / step_x**2 + count_sparse_hxx += 1 + elif diff_types[i, j, 1] == 1: + rows_hxx[count_sparse_hxx] = count + cols_hxx[count_sparse_hxx] = index_dict[(i, j - 1)] + data_hxx[count_sparse_hxx] = 1.0 / step_x**2 + count_sparse_hxx += 1 + rows_hxx[count_sparse_hxx] = count + cols_hxx[count_sparse_hxx] = index_dict[(i, j)] + data_hxx[count_sparse_hxx] = -2.0 / step_x**2 + count_sparse_hxx += 1 + rows_hxx[count_sparse_hxx] = count + cols_hxx[count_sparse_hxx] = index_dict[(i, j + 1)] + data_hxx[count_sparse_hxx] = 1.0 / step_x**2 + count_sparse_hxx += 1 + elif diff_types[i, j, 1] == 2: + rows_hxx[count_sparse_hxx] = count + cols_hxx[count_sparse_hxx] = index_dict[(i, j)] + data_hxx[count_sparse_hxx] = 1.0 / step_x**2 + count_sparse_hxx += 1 + rows_hxx[count_sparse_hxx] = count + cols_hxx[count_sparse_hxx] = index_dict[(i, j + 1)] + data_hxx[count_sparse_hxx] = -2.0 / step_x**2 + count_sparse_hxx += 1 + rows_hxx[count_sparse_hxx] = count + cols_hxx[count_sparse_hxx] = index_dict[(i, j + 2)] + data_hxx[count_sparse_hxx] = 1.0 / step_x**2 + count_sparse_hxx += 1 + + return rows_hxx, cols_hxx, data_hxx, rows_hyy, cols_hyy, data_hyy + + +def derivative_2nd_operators_from(mask, pixel_scale: float = 1.0): + """ + The sparse second-derivative operators (Hyy, Hxx) of a cleaned mask. + + Hyy (Hxx) has shape [n_unmasked_pixels, n_unmasked_pixels]; acting on a + 1D vector of values on the unmasked pixels it returns their second + y (x) derivative. The sum ``Hyy + Hxx`` is the Laplacian (Hamiltonian) + operator of the masked grid, which converts a lensing potential to a + convergence (times two). + + Parameters + ---------- + mask + The cleaned 2D bool mask (``True`` = masked); raises if not cleaned. + pixel_scale + The pixel size (e.g. in arcsec) setting the finite-difference step. + + Returns + ------- + The (Hyy, Hxx) operators as ``scipy.sparse.csr_matrix``. + """ + mask, diff_types = _diff_types_of_cleaned_mask_from(mask) + rows_hxx, cols_hxx, data_hxx, rows_hyy, cols_hyy, data_hyy = ( + derivative_2nd_triplets_from(mask, diff_types, dpix=pixel_scale) + ) + + n_unmasked = np.count_nonzero(~mask) + Hxx = csr_matrix((data_hxx, (rows_hxx, cols_hxx)), shape=(n_unmasked, n_unmasked)) + Hyy = csr_matrix((data_hyy, (rows_hyy, cols_hyy)), shape=(n_unmasked, n_unmasked)) + return Hyy, Hxx + + +@numba_util.jit() +def forward_difference_triplets_from(mask, dpix=1.0, max_order=2): + """ + The (rows, cols, values) sparse triplets of the order-capped forward + finite-difference operators used to build regularization matrices (see + ``forward_difference_operators_from``). + + For every unmasked pixel a forward-difference stencil is emitted in each + direction, of the highest order the run of consecutive unmasked pixels + ahead of it permits, capped at ``max_order``: order L uses the L+1 point + stencil with coefficients (-1)^(L-k) C(L, k) / step^L (k = 0..L), and + order 0 degrades to the identity (a zeroth-order regularization of the + pixel itself). + """ + unmask = ~mask + n1, n2 = unmask.shape + i_indices_unmasked, j_indices_unmasked = np.where(unmask) + n_unmasked_pixels = len(i_indices_unmasked) + + # Rows L = 0..4 hold the forward-difference stencil coefficients of order L. + coeffs = np.array( + [ + [1.0, 0.0, 0.0, 0.0, 0.0], + [-1.0, 1.0, 0.0, 0.0, 0.0], + [1.0, -2.0, 1.0, 0.0, 0.0], + [-1.0, 3.0, -3.0, 1.0, 0.0], + [1.0, -4.0, 6.0, -4.0, 1.0], + ] + ) + + max_entries = n_unmasked_pixels * (max_order + 1) + rows_hx = np.full(max_entries, -1, dtype=np.int64) + cols_hx = np.full(max_entries, -1, dtype=np.int64) + data_hx = np.full(max_entries, 0.0, dtype=np.float64) + rows_hy = np.full(max_entries, -1, dtype=np.int64) + cols_hy = np.full(max_entries, -1, dtype=np.int64) + data_hy = np.full(max_entries, 0.0, dtype=np.float64) + + step_y = -1.0 * dpix + step_x = 1.0 * dpix + + index_dict = {} + for count in range(n_unmasked_pixels): + i, j = i_indices_unmasked[count], j_indices_unmasked[count] + index_dict[(i, j)] = count + + count_sparse_hy = 0 + count_sparse_hx = 0 + for count in range(n_unmasked_pixels): + i, j = i_indices_unmasked[count], j_indices_unmasked[count] + + order_x = 0 + for k in range(1, max_order + 1): + if j + k > n2 - 1: + break + if not unmask[i, j + k]: + break + order_x = k + + norm_x = step_x**order_x + for k in range(order_x + 1): + rows_hx[count_sparse_hx] = count + cols_hx[count_sparse_hx] = index_dict[(i, j + k)] + data_hx[count_sparse_hx] = coeffs[order_x, k] / norm_x + count_sparse_hx += 1 + + order_y = 0 + for k in range(1, max_order + 1): + if i + k > n1 - 1: + break + if not unmask[i + k, j]: + break + order_y = k + + norm_y = step_y**order_y + for k in range(order_y + 1): + rows_hy[count_sparse_hy] = count + cols_hy[count_sparse_hy] = index_dict[(i + k, j)] + data_hy[count_sparse_hy] = coeffs[order_y, k] / norm_y + count_sparse_hy += 1 + + return ( + rows_hx[:count_sparse_hx], + cols_hx[:count_sparse_hx], + data_hx[:count_sparse_hx], + rows_hy[:count_sparse_hy], + cols_hy[:count_sparse_hy], + data_hy[:count_sparse_hy], + ) + + +def forward_difference_operators_from( + mask, pixel_scale: float = 1.0, max_order: int = 2 +): + """ + The sparse order-capped forward finite-difference operators (Hy, Hx) of a + mask, used to build derivative-penalising regularization matrices. + + Unlike the pure derivative operators above, every unmasked pixel receives + a row: the stencil order degrades gracefully (``max_order`` → ... → 1 → 0) + with the run of consecutive unmasked pixels ahead of it, so pixels at the + mask edge are regularized at lower order rather than dropped. With + ``max_order=2`` this reproduces the curvature regularization scheme of + potential-correction methods; ``max_order=4`` the fourth-order scheme. + + Parameters + ---------- + mask + The 2D bool mask (``True`` = masked). + pixel_scale + The pixel size setting the finite-difference step (regularization + matrices are conventionally built with the default of 1.0, the + regularization coefficient absorbing the scale). + max_order + The highest stencil order to emit (2 or 4). + + Returns + ------- + The (Hy, Hx) operators as ``scipy.sparse.csr_matrix``. + """ + if max_order not in (1, 2, 3, 4): + raise ValueError(f"max_order must be in 1..4, got {max_order}") + + mask = np.asarray(mask).astype(bool) + rows_hx, cols_hx, data_hx, rows_hy, cols_hy, data_hy = ( + forward_difference_triplets_from(mask, dpix=pixel_scale, max_order=max_order) + ) + + n_unmasked = np.count_nonzero(~mask) + Hx = csr_matrix((data_hx, (rows_hx, cols_hx)), shape=(n_unmasked, n_unmasked)) + Hy = csr_matrix((data_hy, (rows_hy, cols_hy)), shape=(n_unmasked, n_unmasked)) + return Hy, Hx + + +def forward_difference_reg_matrix_from( + mask, pixel_scale: float = 1.0, max_order: int = 2 +): + """ + The regularization matrix H = Hx^T Hx + Hy^T Hy built from the + order-capped forward finite-difference operators of a mask (see + ``forward_difference_operators_from``). + + Parameters + ---------- + mask + The 2D bool mask (``True`` = masked). + pixel_scale + The finite-difference step (default 1.0 for regularization use). + max_order + The highest stencil order (2 = curvature, 4 = fourth-order). + + Returns + ------- + The [n_unmasked, n_unmasked] regularization matrix as a + ``scipy.sparse.csr_matrix``. + """ + Hy, Hx = forward_difference_operators_from( + mask, pixel_scale=pixel_scale, max_order=max_order + ) + return (Hx.T @ Hx + Hy.T @ Hy).tocsr() diff --git a/autoarray/util/__init__.py b/autoarray/util/__init__.py index 540440a52..ee417913e 100644 --- a/autoarray/util/__init__.py +++ b/autoarray/util/__init__.py @@ -28,5 +28,7 @@ ) from autoarray.operators import transformer_util as transformer +from autoarray.operators import derivative_util as derivative +from autoarray.operators import coarse_interp_util as coarse_interp from autoarray.util import misc_util as misc from autoarray.util import dataset_util as dataset diff --git a/test_autoarray/inversion/regularizations/test_curvature_mask.py b/test_autoarray/inversion/regularizations/test_curvature_mask.py new file mode 100644 index 000000000..0500df061 --- /dev/null +++ b/test_autoarray/inversion/regularizations/test_curvature_mask.py @@ -0,0 +1,74 @@ +import numpy as np +import pytest + +import autoarray as aa + + +class MockMaskedLinearObj: + def __init__(self, mask): + self.mask = mask + + @property + def params(self): + return int(np.count_nonzero(~self.mask)) + + +def test__regularization_matrix_from__matches_hand_computed_operators(): + mask = np.array([[False, False, False]]) + + reg = aa.reg.CurvatureMask(coefficient=1.0) + + regularization_matrix = reg.regularization_matrix_from( + linear_obj=MockMaskedLinearObj(mask=mask) + ) + + # Along x: pixel 0 carries the 2nd-order stencil [1, -2, 1], pixel 1 the + # 1st-order [-1, 1] and pixel 2 the 0th-order [1]. Along y no pixel has a + # neighbour below, so every row is 0th-order (the identity). + hx = np.array([[1.0, -2.0, 1.0], [0.0, -1.0, 1.0], [0.0, 0.0, 1.0]]) + hy = np.eye(3) + expected = hx.T @ hx + hy.T @ hy + + assert regularization_matrix == pytest.approx(expected, abs=1.0e-10) + + +def test__regularization_matrix_from__symmetric_positive_semi_definite(): + mask = np.ones((8, 8), dtype=bool) + mask[2:6, 2:6] = False + + reg = aa.reg.CurvatureMask(coefficient=1.0) + + regularization_matrix = reg.regularization_matrix_from( + linear_obj=MockMaskedLinearObj(mask=mask) + ) + + assert regularization_matrix == pytest.approx(regularization_matrix.T, abs=1.0e-12) + assert np.linalg.eigvalsh(regularization_matrix).min() > -1.0e-10 + + +def test__regularization_matrix_from__scales_linearly_with_coefficient(): + mask = np.ones((8, 8), dtype=bool) + mask[2:6, 2:6] = False + linear_obj = MockMaskedLinearObj(mask=mask) + + matrix_coeff_1 = aa.reg.CurvatureMask(coefficient=1.0).regularization_matrix_from( + linear_obj=linear_obj + ) + matrix_coeff_3 = aa.reg.CurvatureMask(coefficient=3.0).regularization_matrix_from( + linear_obj=linear_obj + ) + + assert matrix_coeff_3 == pytest.approx(3.0 * matrix_coeff_1, abs=1.0e-10) + + +def test__regularization_weights_from(): + mask = np.ones((8, 8), dtype=bool) + mask[2:6, 2:6] = False + + reg = aa.reg.CurvatureMask(coefficient=2.0) + + weights = reg.regularization_weights_from( + linear_obj=MockMaskedLinearObj(mask=mask) + ) + + assert weights == pytest.approx(2.0 * np.ones(16), abs=1.0e-10) diff --git a/test_autoarray/inversion/regularizations/test_fourth_order_mask.py b/test_autoarray/inversion/regularizations/test_fourth_order_mask.py new file mode 100644 index 000000000..9a38692ce --- /dev/null +++ b/test_autoarray/inversion/regularizations/test_fourth_order_mask.py @@ -0,0 +1,70 @@ +import numpy as np +import pytest + +import autoarray as aa + + +class MockMaskedLinearObj: + def __init__(self, mask): + self.mask = mask + + @property + def params(self): + return int(np.count_nonzero(~self.mask)) + + +def test__regularization_matrix_from__matches_hand_computed_operators(): + mask = np.array([[False, False, False, False, False]]) + + reg = aa.reg.FourthOrderMask(coefficient=1.0) + + regularization_matrix = reg.regularization_matrix_from( + linear_obj=MockMaskedLinearObj(mask=mask) + ) + + # Along x the stencil order degrades with the run of pixels ahead: + # pixel 0 carries the 4th-order stencil, pixel 1 the 3rd, pixel 2 the + # 2nd, pixel 3 the 1st and pixel 4 the 0th. Along y every row is + # 0th-order (the identity). + hx = np.array( + [ + [1.0, -4.0, 6.0, -4.0, 1.0], + [0.0, -1.0, 3.0, -3.0, 1.0], + [0.0, 0.0, 1.0, -2.0, 1.0], + [0.0, 0.0, 0.0, -1.0, 1.0], + [0.0, 0.0, 0.0, 0.0, 1.0], + ] + ) + hy = np.eye(5) + expected = hx.T @ hx + hy.T @ hy + + assert regularization_matrix == pytest.approx(expected, abs=1.0e-10) + + +def test__regularization_matrix_from__symmetric_positive_semi_definite(): + mask = np.ones((10, 10), dtype=bool) + mask[2:8, 2:8] = False + + reg = aa.reg.FourthOrderMask(coefficient=1.0) + + regularization_matrix = reg.regularization_matrix_from( + linear_obj=MockMaskedLinearObj(mask=mask) + ) + + assert regularization_matrix == pytest.approx(regularization_matrix.T, abs=1.0e-12) + assert np.linalg.eigvalsh(regularization_matrix).min() > -1.0e-10 + + +def test__regularization_matrix_from__scales_linearly_with_coefficient(): + mask = np.ones((10, 10), dtype=bool) + mask[2:8, 2:8] = False + linear_obj = MockMaskedLinearObj(mask=mask) + + matrix_coeff_1 = aa.reg.FourthOrderMask( + coefficient=1.0 + ).regularization_matrix_from(linear_obj=linear_obj) + matrix_coeff_3 = aa.reg.FourthOrderMask( + coefficient=3.0 + ).regularization_matrix_from(linear_obj=linear_obj) + + assert matrix_coeff_3 == pytest.approx(3.0 * matrix_coeff_1, abs=1.0e-10) diff --git a/test_autoarray/operators/test_coarse_interp_util.py b/test_autoarray/operators/test_coarse_interp_util.py new file mode 100644 index 000000000..cab407804 --- /dev/null +++ b/test_autoarray/operators/test_coarse_interp_util.py @@ -0,0 +1,181 @@ +import numpy as np +import pytest + +import autoarray as aa +from autoarray import exc +from autoarray.operators import coarse_interp_util +from autoarray.operators import derivative_util + + +def test__binned_mask_from__coarse_pixel_unmasked_only_if_all_fine_unmasked(): + mask = np.zeros((4, 4), dtype=bool) + mask[0, 1] = True + + binned_mask = coarse_interp_util.binned_mask_from(mask, 2) + + assert binned_mask.shape == (2, 2) + assert binned_mask[0, 0] + assert not binned_mask[0, 1] + assert not binned_mask[1, 0] + assert not binned_mask[1, 1] + + +def test__interp_box_mask_from__box_unmasked_only_if_all_four_corners_unmasked(): + mask = np.zeros((3, 3), dtype=bool) + mask[0, 0] = True + + box_mask = coarse_interp_util.interp_box_mask_from(mask) + + assert box_mask.shape == (2, 2) + assert box_mask[0, 0] + assert not box_mask[0, 1] + assert not box_mask[1, 0] + assert not box_mask[1, 1] + + +def test__bilinear_weights_from_box__corners_and_centre(): + box_x = np.array([0.0, 1.0, 0.0, 1.0]) + box_y = np.array([1.0, 1.0, 0.0, 0.0]) + + weights = coarse_interp_util.bilinear_weights_from_box( + box_x, box_y, position=(1.0, 0.0) + ) + assert weights == pytest.approx([1.0, 0.0, 0.0, 0.0], abs=1.0e-10) + + weights = coarse_interp_util.bilinear_weights_from_box( + box_x, box_y, position=(0.5, 0.5) + ) + assert weights == pytest.approx([0.25, 0.25, 0.25, 0.25], abs=1.0e-10) + + +def coarse_fine_setup(fine_shape=(20, 20), fine_pixel_scale=0.5, factor=2): + """ + A circular fine mask with its coarse binned counterpart and all the grids + the interpolation matrix needs, mirroring how the potential-correction + mesh pairs a coarse dpsi mesh to the data grid. + """ + cy = (fine_shape[0] - 1) / 2.0 + cx = (fine_shape[1] - 1) / 2.0 + mask_fine = np.ones(fine_shape, dtype=bool) + for i in range(fine_shape[0]): + for j in range(fine_shape[1]): + r = np.sqrt( + ((i - cy) * fine_pixel_scale) ** 2 + ((j - cx) * fine_pixel_scale) ** 2 + ) + if r <= 3.6: + mask_fine[i, j] = False + mask_fine, _ = derivative_util.cleaned_mask_from(mask_fine) + + grid_fine = aa.Grid2D.uniform(shape_native=fine_shape, pixel_scales=fine_pixel_scale) + xgrid_fine_1d = np.array(grid_fine.native[:, :, 1])[~mask_fine] + ygrid_fine_1d = np.array(grid_fine.native[:, :, 0])[~mask_fine] + + coarse_shape = (fine_shape[0] // factor, fine_shape[1] // factor) + coarse_pixel_scale = fine_pixel_scale * factor + mask_coarse = coarse_interp_util.binned_mask_from(mask_fine, factor) + + grid_coarse = aa.Grid2D.uniform( + shape_native=coarse_shape, pixel_scales=coarse_pixel_scale + ) + xgrid_coarse = np.array(grid_coarse.native[:, :, 1]) + ygrid_coarse = np.array(grid_coarse.native[:, :, 0]) + + box_centres = aa.Grid2D.uniform( + shape_native=(coarse_shape[0] - 1, coarse_shape[1] - 1), + pixel_scales=coarse_pixel_scale, + ) + xc_itp_box = np.array(box_centres.native[:, :, 1]) + yc_itp_box = np.array(box_centres.native[:, :, 0]) + mask_itp_box = coarse_interp_util.interp_box_mask_from(mask_coarse) + + return ( + mask_itp_box, + xc_itp_box, + yc_itp_box, + xgrid_fine_1d, + ygrid_fine_1d, + xgrid_coarse, + ygrid_coarse, + mask_coarse, + ) + + +def test__coarse_interp_matrix_from__rows_are_partition_of_unity(): + setup = coarse_fine_setup() + + itp_mat = coarse_interp_util.coarse_interp_matrix_from(*setup) + + row_sums = np.asarray(itp_mat.sum(axis=1)).ravel() + assert row_sums == pytest.approx(np.ones_like(row_sums), abs=1.0e-10) + + +def test__coarse_interp_matrix_from__raises_when_no_unmasked_interp_box(): + setup = list(coarse_fine_setup()) + setup[0] = np.ones_like(setup[0]) # mask every interpolation box + + with pytest.raises(exc.MeshException): + coarse_interp_util.coarse_interp_matrix_from(*setup) + + +def test__coarse_interp_matrix_from__box_search_stays_in_bounds_at_mesh_edge(): + # A coarse mesh whose only unmasked box sits in one corner: fine pixels + # far from it force the nearest-box search to widen well beyond the mesh + # bounds, which must clamp (the original implementation wrapped around + # via negative indexing and crashed). + mask_coarse = np.ones((5, 5), dtype=bool) + mask_coarse[0:2, 0:2] = False + + coarse_pixel_scale = 1.0 + grid_coarse = aa.Grid2D.uniform(shape_native=(5, 5), pixel_scales=coarse_pixel_scale) + xgrid_coarse = np.array(grid_coarse.native[:, :, 1]) + ygrid_coarse = np.array(grid_coarse.native[:, :, 0]) + + box_centres = aa.Grid2D.uniform(shape_native=(4, 4), pixel_scales=coarse_pixel_scale) + xc_itp_box = np.array(box_centres.native[:, :, 1]) + yc_itp_box = np.array(box_centres.native[:, :, 0]) + mask_itp_box = coarse_interp_util.interp_box_mask_from(mask_coarse) + + # fine positions in the far opposite corner of the mesh + xgrid_fine_1d = np.array([2.0, 1.5]) + ygrid_fine_1d = np.array([-2.0, -1.5]) + + itp_mat = coarse_interp_util.coarse_interp_matrix_from( + mask_itp_box, + xc_itp_box, + yc_itp_box, + xgrid_fine_1d, + ygrid_fine_1d, + xgrid_coarse, + ygrid_coarse, + mask_coarse, + ) + + row_sums = np.asarray(itp_mat.sum(axis=1)).ravel() + assert row_sums == pytest.approx(np.ones_like(row_sums), abs=1.0e-10) + + +def test__coarse_interp_matrix_from__exact_on_bilinear_function(): + setup = coarse_fine_setup() + ( + mask_itp_box, + xc_itp_box, + yc_itp_box, + xgrid_fine_1d, + ygrid_fine_1d, + xgrid_coarse, + ygrid_coarse, + mask_coarse, + ) = setup + + def f(y, x): + return 1.5 + 2.0 * x - 3.0 * y + 0.5 * x * y + + values_coarse = f(ygrid_coarse[~mask_coarse], xgrid_coarse[~mask_coarse]) + + itp_mat = coarse_interp_util.coarse_interp_matrix_from(*setup) + + values_fine = itp_mat @ values_coarse + + assert values_fine == pytest.approx( + f(ygrid_fine_1d, xgrid_fine_1d), abs=1.0e-10 + ) diff --git a/test_autoarray/operators/test_derivative_util.py b/test_autoarray/operators/test_derivative_util.py new file mode 100644 index 000000000..6d957ecb1 --- /dev/null +++ b/test_autoarray/operators/test_derivative_util.py @@ -0,0 +1,159 @@ +import numpy as np +import pytest + +import autoarray as aa +from autoarray import exc +from autoarray.operators import derivative_util + + +def circular_mask_from(shape, radius, pixel_scale=1.0): + cy = (shape[0] - 1) / 2.0 + cx = (shape[1] - 1) / 2.0 + mask = np.ones(shape, dtype=bool) + for i in range(shape[0]): + for j in range(shape[1]): + r = np.sqrt(((i - cy) * pixel_scale) ** 2 + ((j - cx) * pixel_scale) ** 2) + if r <= radius: + mask[i, j] = False + return mask + + +def grids_of_unmasked_pixels_from(mask, pixel_scale=1.0): + grid = aa.Grid2D.uniform(shape_native=mask.shape, pixel_scales=pixel_scale) + ygrid = np.array(grid.native[:, :, 0])[~mask] + xgrid = np.array(grid.native[:, :, 1])[~mask] + return ygrid, xgrid + + +def test__cleaned_mask_from__removes_pixels_without_difference_scheme(): + mask = np.ones((7, 7), dtype=bool) + mask[2:5, 2:5] = False + mask[0, 0] = False # isolated unmasked pixel, no neighbours + + cleaned, diff_types = derivative_util.cleaned_mask_from(mask) + + assert cleaned[0, 0] + assert not cleaned[2:5, 2:5].any() + assert (diff_types[~cleaned] >= 0).all() + + +def test__cleaned_mask_from__already_clean_mask_is_unchanged(): + mask = circular_mask_from(shape=(11, 11), radius=4.0) + cleaned, _ = derivative_util.cleaned_mask_from(mask) + + cleaned_again, _ = derivative_util.cleaned_mask_from(cleaned) + + assert (cleaned_again == cleaned).all() + + +def test__derivative_1st_operators_from__exact_on_linear_function(): + pixel_scale = 0.5 + mask = circular_mask_from(shape=(12, 12), radius=2.4, pixel_scale=pixel_scale) + mask, _ = derivative_util.cleaned_mask_from(mask) + + ygrid, xgrid = grids_of_unmasked_pixels_from(mask, pixel_scale=pixel_scale) + values = 2.0 * ygrid + 3.0 * xgrid + + Hy, Hx = derivative_util.derivative_1st_operators_from( + mask, pixel_scale=pixel_scale + ) + + assert Hy @ values == pytest.approx(2.0 * np.ones_like(values), abs=1.0e-10) + assert Hx @ values == pytest.approx(3.0 * np.ones_like(values), abs=1.0e-10) + + +def test__derivative_1st_operators_from__raises_on_unclean_mask(): + mask = np.ones((7, 7), dtype=bool) + mask[2:5, 2:5] = False + mask[0, 0] = False + + with pytest.raises(exc.MaskException): + derivative_util.derivative_1st_operators_from(mask) + + +def test__derivative_2nd_operators_from__exact_on_quadratic_function(): + pixel_scale = 0.5 + mask = circular_mask_from(shape=(12, 12), radius=2.4, pixel_scale=pixel_scale) + mask, _ = derivative_util.cleaned_mask_from(mask) + + ygrid, xgrid = grids_of_unmasked_pixels_from(mask, pixel_scale=pixel_scale) + values = ygrid**2 + 0.5 * xgrid**2 + + Hyy, Hxx = derivative_util.derivative_2nd_operators_from( + mask, pixel_scale=pixel_scale + ) + + assert Hyy @ values == pytest.approx(2.0 * np.ones_like(values), abs=1.0e-9) + assert Hxx @ values == pytest.approx(1.0 * np.ones_like(values), abs=1.0e-9) + + +def test__forward_difference_operators_from__order_degrades_at_grid_edge(): + mask = np.zeros((4, 4), dtype=bool) + + Hy, Hx = derivative_util.forward_difference_operators_from(mask, max_order=2) + + values = np.ones(16) + + # Differences of any order >= 1 annihilate a constant; only the + # zeroth-order rows (pixels with no unmasked pixel ahead) return it. + result_x = Hx @ values + result_y = Hy @ values + + for count, (i, j) in enumerate(np.argwhere(np.ones((4, 4), dtype=bool))): + assert result_x[count] == pytest.approx(1.0 if j == 3 else 0.0, abs=1.0e-10) + assert result_y[count] == pytest.approx(1.0 if i == 3 else 0.0, abs=1.0e-10) + + +def test__forward_difference_operators_from__fourth_order_annihilates_cubic(): + mask = np.zeros((10, 10), dtype=bool) + pixel_scale = 1.0 + + ygrid, xgrid = grids_of_unmasked_pixels_from(mask, pixel_scale=pixel_scale) + values = xgrid**3 + + Hy, Hx = derivative_util.forward_difference_operators_from( + mask, pixel_scale=pixel_scale, max_order=4 + ) + + result_x = Hx @ values + + # Rows with the full 4-point run ahead of them carry the fourth-order + # stencil, which annihilates a cubic exactly. + for count, (i, j) in enumerate(np.argwhere(np.ones((10, 10), dtype=bool))): + if j <= 5: + assert result_x[count] == pytest.approx(0.0, abs=1.0e-9) + + +def test__forward_difference_operators_from__masked_run_truncates_stencil_order(): + mask = np.zeros((1, 5), dtype=bool) + mask[0, 3] = True + + Hy, Hx = derivative_util.forward_difference_operators_from(mask, max_order=2) + + # Unmasked pixels are columns j = 0, 1, 2, 4. Runs ahead: j=0 has j=1, 2 + # (order 2); j=1 has j=2 only (j=3 masked, order 1); j=2 has none + # (order 0); j=4 is at the edge (order 0). + expected_hx = np.array( + [ + [1.0, -2.0, 1.0, 0.0], + [0.0, -1.0, 1.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + ) + + assert Hx.toarray() == pytest.approx(expected_hx, abs=1.0e-10) + + +def test__forward_difference_reg_matrix_from__symmetric_positive_semi_definite(): + mask = circular_mask_from(shape=(10, 10), radius=3.5) + mask, _ = derivative_util.cleaned_mask_from(mask) + + for max_order in (2, 4): + reg_matrix = derivative_util.forward_difference_reg_matrix_from( + mask, max_order=max_order + ).toarray() + + assert reg_matrix == pytest.approx(reg_matrix.T, abs=1.0e-12) + eigenvalues = np.linalg.eigvalsh(reg_matrix) + assert eigenvalues.min() > -1.0e-10