From 539d9ffdaad1c9dc8bfc52c396527ec4b7a0a191 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 28 Aug 2026 10:48:58 -0400 Subject: [PATCH 1/2] Sparse interferometer inversion: support linear func lists and multiple mappers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `InversionInterferometerSparse.curvature_matrix` previously returned the single-mapper diagonal block only. A mixed `[AbstractLinearObjFuncList, Mapper]` list (or several mappers) with a sparse operator attached was still routed to the sparse class by the factory, so the linear-function and cross-mapper terms were silently dropped from F. This gives the interferometer sparse path parity with `InversionImagingSparse`. - `InterferometerSparseOperator`: add `curvature_matrix_off_diag_from` (A0ᵀ W~ A1), `curvature_matrix_off_diag_func_list_from` (Aᵀ W~ B), `curvature_matrix_func_list_from` (B0ᵀ W~ B1) and the `operated_matrix_slim_from` helper. Every block is formed with the same W~ = Re(Fᴴ W F) operator; because W~ already contains the inverse-variance weighting, func-list columns are passed un-weighted (unlike imaging). - `InversionInterferometerSparse`: dispatch `curvature_matrix` over single-mapper (unchanged, bit-identical) / multi-mapper / func-list+mapper paths, mirror the upper blocks, and apply `no_regularization_add_to_curvature_diag_value` as the dense path does. `data_vector = Lᵀ d~` already covered func lists and is unchanged. - Fix `__init__` overwriting the parent's `settings or Settings()` with `None`. - Tests: parity vs `InversionInterferometerMapping` for func+mapper, two mappers, func+two mappers (max rel err ~6e-16 in curvature_matrix and data_vector); operator unit tests for each new method. Closes #499 Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_012JM45sA4YGEUw6KYW8Pm96 --- .../inversion_interferometer_util.py | 248 ++++++++++++++ .../inversion/interferometer/sparse.py | 313 ++++++++++++++++-- .../interferometer/test_interferometer.py | 270 ++++++++++++++- .../test_inversion_interferometer_util.py | 179 ++++++++++ 4 files changed, 980 insertions(+), 30 deletions(-) diff --git a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py index dc693ee5f..8313d8f41 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py +++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py @@ -790,3 +790,251 @@ def body(block_i, C): C_pad = lax.fori_loop(0, n_blocks, body, C0) C = C_pad[:, :S] return 0.5 * (C + C.T) + + def curvature_matrix_off_diag_from( + self, rows0, cols0, vals0, rows1, cols1, vals1, *, S0: int, S1: int + ): + """ + Compute the off-diagonal (mapper-mapper) curvature block F01 = A0ᵀ W~ A1. + + This method mirrors `ImagingSparseOperator.curvature_matrix_off_diag_from` and is the + structural counterpart for the interferometer W~ operator. The difference between the two + is the operator itself: for imaging W = Hᵀ N⁻¹ H is a PSF correlation, whereas here + W~ = Re(Fᴴ W F) is the (translationally invariant) real-space operator of the non-uniform + Fourier transform `F`, applied via `apply_operator` on the *unmasked-extent* rectangular + grid (M = y_shape * x_shape). + + Given two sparse mapping operators: + + - A0 : (M × S0) + - A1 : (M × S1) + + this method computes F01 = A0ᵀ W~ A1 in column blocks of width `batch_size`: + + 1) Assemble Fbatch = A1[:, start:start+B] on the rectangular grid via scatter-add. + 2) Apply W~ to the block via FFT: Gbatch = W~(Fbatch). + 3) Project back with A0ᵀ via segment_sum over `cols0`. + + Parameters + ---------- + rows0, cols0, vals0 + COO triplets for A0, where `rows0` are extent-grid (flat) indices in [0, M). + rows1, cols1, vals1 + COO triplets for A1, where `rows1` are extent-grid (flat) indices in [0, M). + S0 + Number of source pixels / parameters for mapper 0. + S1 + Number of source pixels / parameters for mapper 1. + + Returns + ------- + ndarray + Off-diagonal curvature block of shape (S0, S1). + + Notes + ----- + - The result is *not* symmetrized here because it is not square in general. The symmetric + counterpart is F10 = F01ᵀ, because A0 and A1 share the same W~. + - Padding to `S1_pad = ceil(S1/B)*B` ensures `dynamic_update_slice` is always legal. + """ + import jax.numpy as jnp + from jax import lax + from jax.ops import segment_sum + + rows0 = jnp.asarray(rows0, dtype=jnp.int32) + cols0 = jnp.asarray(cols0, dtype=jnp.int32) + vals0 = jnp.asarray(vals0, dtype=jnp.float64) + + rows1 = jnp.asarray(rows1, dtype=jnp.int32) + cols1 = jnp.asarray(cols1, dtype=jnp.int32) + vals1 = jnp.asarray(vals1, dtype=jnp.float64) + + M = self.M + B = self.batch_size + + n_blocks = (S1 + B - 1) // B + S1_pad = n_blocks * B + + F01_0 = jnp.zeros((S0, S1_pad), dtype=jnp.float64) + + def body(block_i, F01): + start = block_i * B + + in_block = (cols1 >= start) & (cols1 < (start + B)) + bc = jnp.where(in_block, cols1 - start, 0).astype(jnp.int32) + v = jnp.where(in_block, vals1, 0.0) + + F = jnp.zeros((M, B), dtype=jnp.float64) + F = F.at[rows1, bc].add(v) + + G = self.apply_operator(F) # (M, B) + + contrib = vals0[:, None] * G[rows0, :] + block = segment_sum(contrib, cols0, num_segments=S0) + + width = jnp.minimum(B, jnp.maximum(0, S1 - start)) + block = block * (self.col_offsets < width)[None, :] + + return lax.dynamic_update_slice(F01, block, (0, start)) + + F01_pad = lax.fori_loop(0, n_blocks, body, F01_0) + return F01_pad[:, :S1] + + def operated_matrix_slim_from(self, matrix_slim, extent_index_for_masked_pixel): + """ + Apply the interferometer W~ operator to columns defined on the *slim masked* grid. + + The input columns are scattered from the slim masked grid onto the unmasked-extent + rectangular grid (on which W~ is defined), operated on with `apply_operator`, and gathered + back onto the slim masked grid. + + Parameters + ---------- + matrix_slim + Array of shape (M_pix, n_cols) on the slim masked grid (e.g. the real-space + `mapping_matrix` of an `AbstractLinearObjFuncList`). + extent_index_for_masked_pixel + Array of shape (M_pix,) mapping slim masked pixel indices to extent-grid flat indices. + + Returns + ------- + ndarray + Array of shape (M_pix, n_cols) equal to W~ applied to each column. + """ + import jax.numpy as jnp + + matrix_slim = jnp.asarray(matrix_slim, dtype=jnp.float64) + extent_index_for_masked_pixel = jnp.asarray( + extent_index_for_masked_pixel, dtype=jnp.int32 + ) + + grid_flat = jnp.zeros((self.M, matrix_slim.shape[1]), dtype=jnp.float64) + grid_flat = grid_flat.at[extent_index_for_masked_pixel, :].set(matrix_slim) + + return self.apply_operator(grid_flat)[extent_index_for_masked_pixel, :] + + def curvature_matrix_off_diag_func_list_from( + self, + curvature_weights, # (M_pix, n_funcs) + extent_index_for_masked_pixel, # (M_pix,) slim -> extent(flat) + rows, + cols, + vals, # triplets where rows are EXTENT indices + *, + S: int, + ): + """ + Compute the mapper–linear-function off-diagonal block Aᵀ W~ B. + + This is the interferometer counterpart of + `ImagingSparseOperator.curvature_matrix_off_diag_func_list_from`, but with one important + difference in what `curvature_weights` must contain. + + For imaging the operator is split as W = Hᵀ N⁻¹ H, so the imaging method is passed + `curvature_weights = (H B) / noise²` (the forward blur and the inverse variance are folded + into the input) and only applies Hᵀ internally. + + For an interferometer the whole operator W~ = Re(Fᴴ W F) is applied by `apply_operator`, + with the inverse-variance weighting *already inside* W~. Therefore `curvature_weights` is + the plain real-space `mapping_matrix` of the linear function list on the slim masked grid, + with **no** noise weighting and **no** forward operator applied. + + The returned matrix is: + + off_diag = Aᵀ W~ B + + which has shape (S, n_funcs). + + Parameters + ---------- + curvature_weights + Array of shape (M_pix, n_funcs) on the *slim masked* grid: the un-operated, + un-weighted real-space mapping matrix of the linear function list. + extent_index_for_masked_pixel + Array of shape (M_pix,) mapping slim masked pixel indices to extent-grid flat indices. + Used to scatter values onto the rectangular grid W~ is defined on. + rows, cols, vals + COO triplets for the mapper A, where: + - `rows` are extent-grid indices (flat), shape (nnz,) + - `cols` are source pixel indices, shape (nnz,) + - `vals` are mapping weights, shape (nnz,) + S + Number of source pixels / parameters in the mapper. + + Returns + ------- + ndarray + Off-diagonal block of shape (S, n_funcs). + + Notes + ----- + - No `batch_size` sweep is required because the operator is applied to `n_funcs` columns + (typically a handful) rather than to all S source pixels. + """ + import jax.numpy as jnp + from jax.ops import segment_sum + + curvature_weights = jnp.asarray(curvature_weights, dtype=jnp.float64) + extent_index_for_masked_pixel = jnp.asarray( + extent_index_for_masked_pixel, dtype=jnp.int32 + ) + + rows = jnp.asarray(rows, dtype=jnp.int32) + cols = jnp.asarray(cols, dtype=jnp.int32) + vals = jnp.asarray(vals, dtype=jnp.float64) + + n_funcs = curvature_weights.shape[1] + + # 1) scatter slim -> extent(flat) + grid_flat = jnp.zeros((self.M, n_funcs), dtype=jnp.float64) + grid_flat = grid_flat.at[extent_index_for_masked_pixel, :].set( + curvature_weights + ) + + # 2) apply W~ on the extent grid + operated = self.apply_operator(grid_flat) # (M, n_funcs) + + # 3) gather at the mapper's rows (extent coords) and accumulate to source pixels + contrib = vals[:, None] * operated[rows, :] + return segment_sum(contrib, cols, num_segments=S) # (S, n_funcs) + + def curvature_matrix_func_list_from( + self, + curvature_weights_0, # (M_pix, n_funcs_0) + curvature_weights_1, # (M_pix, n_funcs_1) + extent_index_for_masked_pixel, # (M_pix,) slim -> extent(flat) + ): + """ + Compute a linear-function–linear-function curvature block B0ᵀ W~ B1. + + The imaging sparse inversion forms this block as a plain dot product of noise-weighted, + PSF-convolved mapping matrices, because for imaging those matrices are already in the + data frame. For an interferometer the equivalent dense construction would require the + (expensive) visibility-space transformed mapping matrix, which the sparse formalism exists + to avoid. Because W~ = Re(Fᴴ W F) is exact and translationally invariant on the extent + grid, the block is instead formed directly through the same operator used by every other + block, which is both cheaper and keeps every block of `F` self-consistent. + + Parameters + ---------- + curvature_weights_0, curvature_weights_1 + The un-operated, un-weighted real-space `mapping_matrix` of each linear function list, + on the slim masked grid, of shape (M_pix, n_funcs). + extent_index_for_masked_pixel + Array of shape (M_pix,) mapping slim masked pixel indices to extent-grid flat indices. + + Returns + ------- + ndarray + Curvature block of shape (n_funcs_0, n_funcs_1). + """ + import jax.numpy as jnp + + curvature_weights_0 = jnp.asarray(curvature_weights_0, dtype=jnp.float64) + + operated = self.operated_matrix_slim_from( + matrix_slim=curvature_weights_1, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + ) + + return curvature_weights_0.T @ operated diff --git a/autoarray/inversion/inversion/interferometer/sparse.py b/autoarray/inversion/inversion/interferometer/sparse.py index 6853ed902..1e1f11aba 100644 --- a/autoarray/inversion/inversion/interferometer/sparse.py +++ b/autoarray/inversion/inversion/interferometer/sparse.py @@ -1,5 +1,5 @@ import numpy as np -from typing import Dict, List, Union +from typing import Dict, List, Optional, Union from autoarray import exc from autoarray.dataset.interferometer.dataset import Interferometer @@ -8,11 +8,14 @@ AbstractInversionInterferometer, ) from autoarray.inversion.linear_obj.linear_obj import LinearObj +from autoarray.inversion.linear_obj.func_list import AbstractLinearObjFuncList from autoarray.inversion.mappers import mapper_util from autoarray.settings import Settings from autoarray.inversion.mappers.abstract import Mapper from autoarray.structures.visibilities import Visibilities +from autoarray.inversion.inversion import inversion_util + class InversionInterferometerSparse(AbstractInversionInterferometer): def __init__( @@ -64,8 +67,6 @@ def __init__( preloads=preloads, ) - self.settings = settings - @property def data_vector(self) -> np.ndarray: """ @@ -78,12 +79,48 @@ def data_vector(self) -> np.ndarray: If there are multiple linear objects the `data_vectors` are concatenated ensuring their values are solved for simultaneously. - The calculation is described in more detail in `inversion_util.weighted_data_interferometer_from`. + The `data_vector` is computed as `Lᵀ d~`, where `L` is the real-space mapping matrix of every + linear object stacked horizontally and `d~ = Re(Fᴴ W d)` is the dirty image cached on the + dataset's `sparse_operator`. This single expression covers every linear object type: + + - For a `Mapper` column `a`, `aᵀ d~` is the mapper's data vector entry. + - For an `AbstractLinearObjFuncList` column `b`, `bᵀ d~ = bᵀ Re(Fᴴ W d) = Re((F b)ᴴ W d)`, + which is exactly the entry the mapping (dense) formalism computes via the linear + function's transformed mapping matrix. Linear function lists therefore require no + separate branch here. """ return self._xp.dot( self.mapping_matrix.T, self.dataset.sparse_operator.dirty_image ) + def _sparse_triplets_curvature_from(self, mapper: Mapper): + """ + Returns the sparse COO triplets `(rows, cols, vals)` of a mapper's real-space mapping + operator `A`, with `rows` indexed on the grid the interferometer `W~` operator lives on. + + The interferometer `W~` operator lives on the unmasked-extent rectangular grid + (`shape_native_masked_pixels`), not the full native grid used by the imaging path. + The triplets are therefore built with extent-flat row indices so that they match the + operator's (M = extent_y * extent_x, B) scatter buffer. + + This is the interferometer counterpart of `Mapper.sparse_triplets_curvature`, which + indexes rows on the imaging FFT grid and therefore cannot be reused here. + + Parameters + ---------- + mapper + The mapper whose mapping operator is expressed in sparse triplet form. + """ + return mapper_util.sparse_triplets_from( + pix_indexes_for_sub=mapper.pix_indexes_for_sub_slim_index, + pix_weights_for_sub=mapper.pix_weights_for_sub_slim_index, + slim_index_for_sub=mapper.slim_index_for_sub_slim_index, + fft_index_for_masked_pixel=self.mask.extent_index_for_masked_pixel, + sub_fraction_slim=mapper.over_sampler.sub_fraction.array, + return_rows_slim=False, + xp=self._xp, + ) + @property def curvature_matrix(self) -> np.ndarray: """ @@ -93,9 +130,20 @@ def curvature_matrix(self) -> np.ndarray: The linear algebra is described in the paper https://arxiv.org/pdf/astro-ph/0302587.pdf, where the curvature matrix given by equation (4) and the letter F. - If there are multiple linear objects their `operated_mapping_matrix` properties will have already been - concatenated ensuring their `curvature_matrix` values are solved for simultaneously. This includes all - diagonal and off-diagonal terms describing the covariances between linear objects. + If there are multiple linear objects their contributions are combined ensuring their `curvature_matrix` + values are solved for simultaneously. This includes all diagonal and off-diagonal terms describing the + covariances between linear objects, whether those objects are `Mapper`s, `AbstractLinearObjFuncList`s + (e.g. linear light profiles), or a mixture of both. + + Every block is formed with the same operator `W~ = Re(Fᴴ W F)`, applied by the dataset's + `sparse_operator` on the unmasked-extent rectangular grid: + + - mapper–mapper diagonal: `A_iᵀ W~ A_i` + - mapper–mapper off-diagonal: `A_iᵀ W~ A_j` + - mapper–function: `A_iᵀ W~ B_k` + - function–function: `B_kᵀ W~ B_l` + + Only the upper blocks are computed, with `curvature_matrix_mirrored_from` filling the lower ones. If a `preloads.curvature_matrix` was injected (e.g. the datacube shared-state path, where `F` is identical for every channel) it is returned directly, so the dominant `F = LᵀW̃L` build is skipped. @@ -106,7 +154,35 @@ def curvature_matrix(self) -> np.ndarray: if self._preloads is not None and self._preloads.curvature_matrix is not None: return self._preloads.curvature_matrix - return self.curvature_matrix_diag + if not self.has(cls=AbstractLinearObjFuncList) and self.total(cls=Mapper) == 1: + # The single-mapper case is the performance-critical one and its matrix is already + # square, symmetric and complete, so it bypasses the block assembly and mirroring. + curvature_matrix = self.curvature_matrix_diag + else: + if self.has(cls=AbstractLinearObjFuncList): + curvature_matrix = self._curvature_matrix_func_list_and_mapper + else: + curvature_matrix = self._curvature_matrix_multi_mapper + + curvature_matrix = inversion_util.curvature_matrix_mirrored_from( + curvature_matrix=curvature_matrix, + xp=self._xp, + ) + + if len(self.no_regularization_index_list) > 0: + if self._xp is np: + # The sparse operator returns JAX arrays, which the NumPy in-place diagonal + # update below cannot write to. + curvature_matrix = np.array(curvature_matrix) + + curvature_matrix = inversion_util.curvature_matrix_with_added_to_diag_from( + curvature_matrix=curvature_matrix, + value=self.settings.no_regularization_add_to_curvature_diag_value, + no_regularization_index_list=self.no_regularization_index_list, + xp=self._xp, + ) + + return curvature_matrix @property def curvature_matrix_diag(self) -> np.ndarray: @@ -117,23 +193,12 @@ def curvature_matrix_diag(self) -> np.ndarray: The linear algebra is described in the paper https://arxiv.org/pdf/astro-ph/0302587.pdf, where the curvature matrix given by equation (4) and the letter F. - This function computes the diagonal terms of F using the sparse linear algebra formalism. + This function computes the diagonal terms of F of the inversion's first `Mapper` using the sparse + linear algebra formalism, returning a matrix of shape [mapper.params, mapper.params]. """ mapper = self.cls_list_from(cls=Mapper)[0] - # The interferometer W~ operator lives on the unmasked-extent rectangular - # grid (shape_native_masked_pixels), not the full native grid used by - # the imaging path. Build sparse triplets with extent-flat row indices - # so they match the operator's (M = extent_y * extent_x, B) scatter buffer. - rows, cols, vals = mapper_util.sparse_triplets_from( - pix_indexes_for_sub=mapper.pix_indexes_for_sub_slim_index, - pix_weights_for_sub=mapper.pix_weights_for_sub_slim_index, - slim_index_for_sub=mapper.slim_index_for_sub_slim_index, - fft_index_for_masked_pixel=self.mask.extent_index_for_masked_pixel, - sub_fraction_slim=mapper.over_sampler.sub_fraction.array, - return_rows_slim=False, - xp=self._xp, - ) + rows, cols, vals = self._sparse_triplets_curvature_from(mapper=mapper) return self.dataset.sparse_operator.curvature_matrix_diag_from( rows=rows, @@ -142,6 +207,210 @@ def curvature_matrix_diag(self) -> np.ndarray: S=mapper.params, ) + @property + def _curvature_matrix_mapper_diag(self) -> Optional[np.ndarray]: + """ + Returns the diagonal regions of the `curvature_matrix`, a 2D matrix which uses the mappings between the data + and the linear objects to construct the simultaneous linear equations. The object is described in full in + the method `curvature_matrix`. + + This method computes the diagonal entries of all mapper objects in the `curvature_matrix`, placing each + one in the parameter range of its mapper in the full [total_params, total_params] matrix. + """ + if not self.has(cls=Mapper): + return None + + curvature_matrix = self._xp.zeros((self.total_params, self.total_params)) + + mapper_list = self.cls_list_from(cls=Mapper) + mapper_param_range_list = self.param_range_list_from(cls=Mapper) + + for mapper_index, mapper in enumerate(mapper_list): + rows, cols, vals = self._sparse_triplets_curvature_from(mapper=mapper) + + diag = self.dataset.sparse_operator.curvature_matrix_diag_from( + rows=rows, + cols=cols, + vals=vals, + S=mapper.params, + ) + + start, end = mapper_param_range_list[mapper_index] + + if self._xp is np: + curvature_matrix[start:end, start:end] = diag + else: + curvature_matrix = curvature_matrix.at[start:end, start:end].set(diag) + + return curvature_matrix + + def _curvature_matrix_off_diag_from( + self, mapper_0: Mapper, mapper_1: Mapper + ) -> np.ndarray: + """ + Returns the off-diagonal block `A_0ᵀ W~ A_1` of the `curvature_matrix` describing the covariance + between two mappers, of shape [mapper_0.params, mapper_1.params]. + """ + rows_0, cols_0, vals_0 = self._sparse_triplets_curvature_from(mapper=mapper_0) + rows_1, cols_1, vals_1 = self._sparse_triplets_curvature_from(mapper=mapper_1) + + return self.dataset.sparse_operator.curvature_matrix_off_diag_from( + rows0=rows_0, + cols0=cols_0, + vals0=vals_0, + rows1=rows_1, + cols1=cols_1, + vals1=vals_1, + S0=mapper_0.params, + S1=mapper_1.params, + ) + + @property + def _curvature_matrix_multi_mapper(self) -> np.ndarray: + """ + Returns the `curvature_matrix`, a 2D matrix which uses the mappings between the data and the linear objects to + construct the simultaneous linear equations. The object is described in full in the method `curvature_matrix`. + + This method computes the mapper entries of the `curvature_matrix` when there are multiple mapper objects in + the `Inversion`, filling in each mapper's diagonal block and the upper off-diagonal blocks describing the + covariance between every pair of mappers. The lower blocks are filled in by the mirroring performed in + `curvature_matrix`. + """ + curvature_matrix = self._curvature_matrix_mapper_diag + + if self.total(cls=Mapper) == 1: + return curvature_matrix + + mapper_list = self.cls_list_from(cls=Mapper) + mapper_param_range_list = self.param_range_list_from(cls=Mapper) + + for i in range(len(mapper_list)): + mapper_param_range_i = mapper_param_range_list[i] + + for j in range(i + 1, len(mapper_list)): + mapper_param_range_j = mapper_param_range_list[j] + + off_diag = self._curvature_matrix_off_diag_from( + mapper_0=mapper_list[i], mapper_1=mapper_list[j] + ) + + if self._xp is np: + curvature_matrix[ + mapper_param_range_i[0] : mapper_param_range_i[1], + mapper_param_range_j[0] : mapper_param_range_j[1], + ] = off_diag + else: + curvature_matrix = curvature_matrix.at[ + mapper_param_range_i[0] : mapper_param_range_i[1], + mapper_param_range_j[0] : mapper_param_range_j[1], + ].set(off_diag) + + return curvature_matrix + + @property + def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: + """ + Returns the `curvature_matrix`, a 2D matrix which uses the mappings between the data and the linear objects to + construct the simultaneous linear equations. The object is described in full in the method `curvature_matrix`. + + This method computes the `curvature_matrix` when one or more `AbstractLinearObjFuncList` objects (e.g. linear + light profiles) are fitted, optionally simultaneously with one or more `Mapper` objects. + + All mapper blocks are computed first, then the mapper–function off-diagonal blocks `A_iᵀ W~ B_k`, then the + function–function blocks `B_kᵀ W~ B_l`. Only the upper blocks are filled in, with the lower ones filled in + by the mirroring performed in `curvature_matrix`. + + Unlike the imaging sparse inversion, the linear function's `mapping_matrix` is passed to the operator + without any noise weighting or forward operation applied, because the interferometer operator + `W~ = Re(Fᴴ W F)` already contains both. + """ + if self.has(cls=Mapper): + curvature_matrix = self._curvature_matrix_multi_mapper + else: + curvature_matrix = self._xp.zeros((self.total_params, self.total_params)) + + sparse_operator = self.dataset.sparse_operator + extent_index_for_masked_pixel = self.mask.extent_index_for_masked_pixel + + mapper_list = self.cls_list_from(cls=Mapper) + mapper_param_range_list = self.param_range_list_from(cls=Mapper) + + linear_func_list = self.cls_list_from(cls=AbstractLinearObjFuncList) + linear_func_param_range_list = self.param_range_list_from( + cls=AbstractLinearObjFuncList + ) + + mapping_matrix_list = [ + linear_func.mapping_matrix for linear_func in linear_func_list + ] + + for mapper_index, mapper in enumerate(mapper_list): + mapper_param_range = mapper_param_range_list[mapper_index] + + rows, cols, vals = self._sparse_triplets_curvature_from(mapper=mapper) + + for func_index in range(len(linear_func_list)): + linear_func_param_range = linear_func_param_range_list[func_index] + + off_diag = sparse_operator.curvature_matrix_off_diag_func_list_from( + curvature_weights=mapping_matrix_list[func_index], + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + rows=rows, + cols=cols, + vals=vals, + S=mapper.params, + ) + + if self._xp is np: + curvature_matrix[ + mapper_param_range[0] : mapper_param_range[1], + linear_func_param_range[0] : linear_func_param_range[1], + ] = off_diag + else: + curvature_matrix = curvature_matrix.at[ + mapper_param_range[0] : mapper_param_range[1], + linear_func_param_range[0] : linear_func_param_range[1], + ].set(off_diag) + + # The linear func x linear func block is symmetric, so only the upper triangle of blocks is + # computed, with the mirrored block set from the transpose. + for index_0 in range(len(linear_func_list)): + linear_func_param_range_0 = linear_func_param_range_list[index_0] + + for index_1 in range(index_0, len(linear_func_list)): + linear_func_param_range_1 = linear_func_param_range_list[index_1] + + diag = sparse_operator.curvature_matrix_func_list_from( + curvature_weights_0=mapping_matrix_list[index_0], + curvature_weights_1=mapping_matrix_list[index_1], + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + ) + + if self._xp is np: + curvature_matrix[ + linear_func_param_range_0[0] : linear_func_param_range_0[1], + linear_func_param_range_1[0] : linear_func_param_range_1[1], + ] = diag + + if index_1 != index_0: + curvature_matrix[ + linear_func_param_range_1[0] : linear_func_param_range_1[1], + linear_func_param_range_0[0] : linear_func_param_range_0[1], + ] = diag.T + else: + curvature_matrix = curvature_matrix.at[ + linear_func_param_range_0[0] : linear_func_param_range_0[1], + linear_func_param_range_1[0] : linear_func_param_range_1[1], + ].set(diag) + + if index_1 != index_0: + curvature_matrix = curvature_matrix.at[ + linear_func_param_range_1[0] : linear_func_param_range_1[1], + linear_func_param_range_0[0] : linear_func_param_range_0[1], + ].set(diag.T) + + return curvature_matrix + @property def mapped_reconstructed_operated_data_dict( self, diff --git a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py index 60b28d5b5..b55cae30f 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py +++ b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py @@ -226,8 +226,7 @@ def test__operated_mapping_matrix_override__sparse_operator_raises(): linear_obj = aa.m.MockLinearObjFuncList( parameters=1, mapping_matrix=np.ones((mask.pixels_in_mask, 1)), - operated_mapping_matrix_override=(999.0 + 1.0j) - * np.ones((n_visibilities, 1)), + operated_mapping_matrix_override=(999.0 + 1.0j) * np.ones((n_visibilities, 1)), ) with pytest.raises(aa.exc.InversionException): @@ -255,9 +254,7 @@ def test__curvature_matrix__interferometer_sparse_operator__delaunay__identical_ mesh = aa.mesh.Delaunay(pixels=9) image_mesh = aa.image_mesh.Overlay(shape=(3, 3)) - image_mesh_grid = image_mesh.image_plane_mesh_grid_from( - mask=mask, adapt_data=None - ) + image_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=mask, adapt_data=None) interpolator = mesh.interpolator_from( source_plane_data_grid=grid, @@ -318,9 +315,7 @@ def test__curvature_matrix__interferometer_sparse_operator__delaunay__dft_and_nu mesh = aa.mesh.Delaunay(pixels=9) image_mesh = aa.image_mesh.Overlay(shape=(3, 3)) - image_mesh_grid = image_mesh.image_plane_mesh_grid_from( - mask=mask, adapt_data=None - ) + image_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=mask, adapt_data=None) interpolator = mesh.interpolator_from( source_plane_data_grid=grid, @@ -454,3 +449,262 @@ def test__preloads_interferometer__curvature_matrix_returned_directly_and_skips_ ) assert inversion_preloaded.curvature_matrix is curvature_matrix assert inversion_preloaded.data_vector == pytest.approx(inversion.data_vector) + + +def _sparse_parity_setup(seed=0): + """ + The 7x7 / Delaunay / TransformerDFT setup shared by the sparse-operator parity tests below, + returning the dense and sparse-operator datasets alongside the mask and a random generator. + """ + mask = aa.Mask2D( + mask=[ + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, False, True, True, True], + [True, True, False, False, False, True, True], + [True, True, True, False, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + ], + pixel_scales=2.0, + ) + + n_visibilities = 5 + rng = np.random.default_rng(seed=seed) + data = aa.Visibilities( + visibilities=rng.normal(size=(n_visibilities, 2)).astype(np.float64) + ) + noise_map = aa.VisibilitiesNoiseMap( + visibilities=np.ones((n_visibilities, 2), dtype=np.float64) + ) + uv_wavelengths = rng.normal(size=(n_visibilities, 2)).astype(np.float64) + + dataset = aa.Interferometer( + data=data, + noise_map=noise_map, + uv_wavelengths=uv_wavelengths, + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + ) + + return mask, rng, dataset, dataset.apply_sparse_operator(use_jax=False) + + +def _mapper_from(mask, pixels, shape, regularization=None): + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=1) + + mesh = aa.mesh.Delaunay(pixels=pixels) + image_mesh = aa.image_mesh.Overlay(shape=shape) + image_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=mask, adapt_data=None) + + interpolator = mesh.interpolator_from( + source_plane_data_grid=grid, + source_plane_mesh_grid=image_mesh_grid, + ) + + return aa.Mapper(interpolator=interpolator, regularization=regularization) + + +def _assert_sparse_matches_mapping(dataset, dataset_sparse, linear_obj_list): + inversion_sparse = aa.Inversion( + dataset=dataset_sparse, linear_obj_list=linear_obj_list + ) + inversion_mapping = aa.Inversion(dataset=dataset, linear_obj_list=linear_obj_list) + + assert isinstance(inversion_sparse, aa.InversionInterferometerSparse) + assert isinstance(inversion_mapping, aa.InversionInterferometerMapping) + + assert inversion_sparse.curvature_matrix == pytest.approx( + np.array(inversion_mapping.curvature_matrix), 1.0e-8 + ) + assert inversion_sparse.data_vector == pytest.approx( + np.array(inversion_mapping.data_vector), 1.0e-8 + ) + assert inversion_sparse.reconstruction == pytest.approx( + np.array(inversion_mapping.reconstruction), 1.0e-4 + ) + assert inversion_sparse.log_det_curvature_reg_matrix_term == pytest.approx( + inversion_mapping.log_det_curvature_reg_matrix_term, 1.0e-6 + ) + + +def test__interferometer_sparse_operator__func_list_and_mapper__identical_to_mapping(): + """ + A linear function list (e.g. linear light profiles) fitted simultaneously with a `Mapper` must + reproduce the dense (mapping formalism) inversion, including the mapper-function off-diagonal + blocks and the function-function block which the sparse path previously dropped entirely. + """ + mask, rng, dataset, dataset_sparse = _sparse_parity_setup() + + mapper = _mapper_from( + mask=mask, + pixels=9, + shape=(3, 3), + regularization=aa.reg.Constant(coefficient=1.0), + ) + + linear_obj = aa.m.MockLinearObjFuncList( + parameters=2, + mapping_matrix=rng.normal(size=(mask.pixels_in_mask, 2)), + ) + + _assert_sparse_matches_mapping( + dataset=dataset, + dataset_sparse=dataset_sparse, + linear_obj_list=[linear_obj, mapper], + ) + + # The linear function list is also supported when it trails the mapper in the list, where its + # parameters occupy the final rows / columns of the curvature matrix. + _assert_sparse_matches_mapping( + dataset=dataset, + dataset_sparse=dataset_sparse, + linear_obj_list=[mapper, linear_obj], + ) + + +def test__interferometer_sparse_operator__x2_mappers__identical_to_mapping(): + """ + Two `Mapper` objects fitted simultaneously require the mapper-mapper off-diagonal block + `A_0ᵀ W~ A_1`, which the sparse path previously dropped (only the first mapper was used). + """ + mask, rng, dataset, dataset_sparse = _sparse_parity_setup() + + mapper_0 = _mapper_from( + mask=mask, + pixels=9, + shape=(3, 3), + regularization=aa.reg.Constant(coefficient=1.0), + ) + mapper_1 = _mapper_from( + mask=mask, + pixels=16, + shape=(4, 4), + regularization=aa.reg.Constant(coefficient=2.0), + ) + + _assert_sparse_matches_mapping( + dataset=dataset, + dataset_sparse=dataset_sparse, + linear_obj_list=[mapper_0, mapper_1], + ) + + +def test__interferometer_sparse_operator__func_list_and_x2_mappers__identical_to_mapping(): + """ + The full mixed case: one or more linear function lists fitted simultaneously with multiple + mappers, exercising every block of the curvature matrix at once. + """ + mask, rng, dataset, dataset_sparse = _sparse_parity_setup() + + mapper_0 = _mapper_from( + mask=mask, + pixels=9, + shape=(3, 3), + regularization=aa.reg.Constant(coefficient=1.0), + ) + mapper_1 = _mapper_from( + mask=mask, + pixels=16, + shape=(4, 4), + regularization=aa.reg.Constant(coefficient=2.0), + ) + + linear_obj = aa.m.MockLinearObjFuncList( + parameters=2, + mapping_matrix=rng.normal(size=(mask.pixels_in_mask, 2)), + ) + + _assert_sparse_matches_mapping( + dataset=dataset, + dataset_sparse=dataset_sparse, + linear_obj_list=[linear_obj, mapper_0, mapper_1], + ) + + linear_obj_1 = aa.m.MockLinearObjFuncList( + parameters=1, + mapping_matrix=rng.normal(size=(mask.pixels_in_mask, 1)), + ) + + _assert_sparse_matches_mapping( + dataset=dataset, + dataset_sparse=dataset_sparse, + linear_obj_list=[linear_obj, linear_obj_1, mapper_0, mapper_1], + ) + + +def test__interferometer_sparse_operator__x1_mapper__unchanged_by_func_list_support(): + """ + The single-mapper path is the performance-critical one and must be untouched by the + func-list / multi-mapper block assembly: for a regularized mapper the `curvature_matrix` is + still exactly the `curvature_matrix_diag` build, with no mirroring or diagonal stabilisation + applied on top of it. + """ + mask, rng, dataset, dataset_sparse = _sparse_parity_setup() + + mapper = _mapper_from( + mask=mask, + pixels=9, + shape=(3, 3), + regularization=aa.reg.Constant(coefficient=1.0), + ) + + inversion = aa.Inversion(dataset=dataset_sparse, linear_obj_list=[mapper]) + + assert inversion.no_regularization_index_list == [] + assert np.array_equal( + np.array(inversion.curvature_matrix), + np.array(inversion.curvature_matrix_diag), + ) + + _assert_sparse_matches_mapping( + dataset=dataset, dataset_sparse=dataset_sparse, linear_obj_list=[mapper] + ) + + +def test__interferometer_sparse_operator__no_regularization_value_added_to_diag(): + """ + Linear function lists are typically unregularized, so their curvature diagonal receives the + `no_regularization_add_to_curvature_diag_value` stabiliser. The sparse path must apply this + exactly as the dense (mapping formalism) path does. + """ + mask, rng, dataset, dataset_sparse = _sparse_parity_setup() + + mapper = _mapper_from( + mask=mask, + pixels=9, + shape=(3, 3), + regularization=aa.reg.Constant(coefficient=1.0), + ) + + linear_obj = aa.m.MockLinearObjFuncList( + parameters=2, + mapping_matrix=rng.normal(size=(mask.pixels_in_mask, 2)), + ) + + inversion = aa.Inversion( + dataset=dataset_sparse, linear_obj_list=[linear_obj, mapper] + ) + + assert inversion.no_regularization_index_list == [0, 1] + + value = inversion.settings.no_regularization_add_to_curvature_diag_value + + curvature_matrix = np.array(inversion.curvature_matrix) + + # Rebuilding the un-stabilised blocks directly from the operator and adding the value back on + # reproduces the diagonal entries of the unregularized linear function parameters. + operator = dataset_sparse.sparse_operator + mapping_matrix = np.array(linear_obj.mapping_matrix) + + curvature_func = np.array( + operator.curvature_matrix_func_list_from( + curvature_weights_0=mapping_matrix, + curvature_weights_1=mapping_matrix, + extent_index_for_masked_pixel=mask.extent_index_for_masked_pixel, + ) + ) + + assert curvature_matrix[0, 0] == pytest.approx(curvature_func[0, 0] + value, 1.0e-8) + assert curvature_matrix[1, 1] == pytest.approx(curvature_func[1, 1] + value, 1.0e-8) + assert curvature_matrix[0, 1] == pytest.approx(curvature_func[0, 1], 1.0e-8) diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index 9a527b09b..e415b4515 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -66,3 +66,182 @@ def test__data_vector_via_transformed_mapping_matrix_from(): ) assert (data_vector_complex_via_blurred == data_vector_via_transformed).all() + + +def _sparse_operator_and_mask(): + """ + Returns a real `InterferometerSparseOperator` (and the mask it is defined on) built from a + small 7x7 `TransformerDFT` interferometer dataset. + """ + mask = aa.Mask2D( + mask=[ + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, False, True, True, True], + [True, True, False, False, False, True, True], + [True, True, True, False, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + ], + pixel_scales=2.0, + ) + + n_visibilities = 5 + rng = np.random.default_rng(seed=3) + + dataset = aa.Interferometer( + data=aa.Visibilities( + visibilities=rng.normal(size=(n_visibilities, 2)).astype(np.float64) + ), + noise_map=aa.VisibilitiesNoiseMap( + visibilities=np.ones((n_visibilities, 2), dtype=np.float64) + ), + uv_wavelengths=rng.normal(size=(n_visibilities, 2)).astype(np.float64), + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + ) + + return dataset.apply_sparse_operator(use_jax=False).sparse_operator, mask, rng + + +def _operator_dense(operator): + """ + Returns the dense [M, M] matrix of the `W~` operator, computed by applying it to the identity + on the extent grid one one-hot column at a time. + """ + return np.array(operator.apply_operator(np.eye(operator.M))) + + +def _dense_from_triplets(rows, cols, vals, M, S): + matrix = np.zeros((M, S)) + + for row, col, val in zip(rows, cols, vals): + matrix[row, col] += val + + return matrix + + +def test__interferometer_sparse_operator__curvature_matrix_off_diag_from(): + operator, mask, rng = _sparse_operator_and_mask() + + M = operator.M + + rows_0 = np.array([0, 1, 4, 4]) + cols_0 = np.array([0, 1, 0, 1]) + vals_0 = np.array([1.0, 2.0, 0.5, 0.25]) + + rows_1 = np.array([1, 3, 4, 7]) + cols_1 = np.array([0, 1, 2, 2]) + vals_1 = np.array([0.75, 1.5, 3.0, 0.125]) + + off_diag = np.array( + operator.curvature_matrix_off_diag_from( + rows0=rows_0, + cols0=cols_0, + vals0=vals_0, + rows1=rows_1, + cols1=cols_1, + vals1=vals_1, + S0=2, + S1=3, + ) + ) + + matrix_0 = _dense_from_triplets(rows_0, cols_0, vals_0, M=M, S=2) + matrix_1 = _dense_from_triplets(rows_1, cols_1, vals_1, M=M, S=3) + + off_diag_dense = matrix_0.T @ _operator_dense(operator) @ matrix_1 + + assert off_diag.shape == (2, 3) + assert off_diag == pytest.approx(off_diag_dense, 1.0e-8) + + +def test__interferometer_sparse_operator__curvature_matrix_off_diag_func_list_from(): + operator, mask, rng = _sparse_operator_and_mask() + + M = operator.M + extent_index_for_masked_pixel = np.array(mask.extent_index_for_masked_pixel) + + rows = np.array([0, 1, 4, 4, 7]) + cols = np.array([0, 1, 0, 1, 1]) + vals = np.array([1.0, 2.0, 0.5, 0.25, 3.0]) + + curvature_weights = rng.normal(size=(mask.pixels_in_mask, 3)) + + off_diag = np.array( + operator.curvature_matrix_off_diag_func_list_from( + curvature_weights=curvature_weights, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + rows=rows, + cols=cols, + vals=vals, + S=2, + ) + ) + + mapping_matrix = _dense_from_triplets(rows, cols, vals, M=M, S=2) + + # The linear function columns are scattered from the slim masked grid onto the extent grid, + # with no noise weighting applied (the inverse variance lives inside `W~`). + func_matrix = np.zeros((M, 3)) + func_matrix[extent_index_for_masked_pixel, :] = curvature_weights + + off_diag_dense = mapping_matrix.T @ _operator_dense(operator) @ func_matrix + + assert off_diag.shape == (2, 3) + assert off_diag == pytest.approx(off_diag_dense, 1.0e-8) + + +def test__interferometer_sparse_operator__curvature_matrix_func_list_from(): + operator, mask, rng = _sparse_operator_and_mask() + + M = operator.M + extent_index_for_masked_pixel = np.array(mask.extent_index_for_masked_pixel) + + curvature_weights_0 = rng.normal(size=(mask.pixels_in_mask, 2)) + curvature_weights_1 = rng.normal(size=(mask.pixels_in_mask, 3)) + + curvature_matrix = np.array( + operator.curvature_matrix_func_list_from( + curvature_weights_0=curvature_weights_0, + curvature_weights_1=curvature_weights_1, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + ) + ) + + func_matrix_0 = np.zeros((M, 2)) + func_matrix_0[extent_index_for_masked_pixel, :] = curvature_weights_0 + + func_matrix_1 = np.zeros((M, 3)) + func_matrix_1[extent_index_for_masked_pixel, :] = curvature_weights_1 + + curvature_matrix_dense = func_matrix_0.T @ _operator_dense(operator) @ func_matrix_1 + + assert curvature_matrix.shape == (2, 3) + assert curvature_matrix == pytest.approx(curvature_matrix_dense, 1.0e-8) + + +def test__interferometer_sparse_operator__operated_matrix_slim_from(): + operator, mask, rng = _sparse_operator_and_mask() + + M = operator.M + extent_index_for_masked_pixel = np.array(mask.extent_index_for_masked_pixel) + + matrix_slim = rng.normal(size=(mask.pixels_in_mask, 2)) + + operated = np.array( + operator.operated_matrix_slim_from( + matrix_slim=matrix_slim, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + ) + ) + + matrix_extent = np.zeros((M, 2)) + matrix_extent[extent_index_for_masked_pixel, :] = matrix_slim + + operated_dense = (_operator_dense(operator) @ matrix_extent)[ + extent_index_for_masked_pixel, : + ] + + assert operated.shape == (mask.pixels_in_mask, 2) + assert operated == pytest.approx(operated_dense, 1.0e-8) From ffafa86a7b739f8bd4db94ceaff2282834137580 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 28 Aug 2026 10:58:27 -0400 Subject: [PATCH 2/2] Skip JAX-only InterferometerSparseOperator unit tests when jax is absent The four new operator tests call methods that import jax internally, so the unittest-nojax leg failed with ModuleNotFoundError. Guard each with pytest.importorskip("jax"), matching test_kernel_jax_gradients.py. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_012JM45sA4YGEUw6KYW8Pm96 --- .../interferometer/test_inversion_interferometer_util.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index e415b4515..ec3fe4770 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -122,6 +122,8 @@ def _dense_from_triplets(rows, cols, vals, M, S): def test__interferometer_sparse_operator__curvature_matrix_off_diag_from(): + pytest.importorskip("jax") + operator, mask, rng = _sparse_operator_and_mask() M = operator.M @@ -157,6 +159,8 @@ def test__interferometer_sparse_operator__curvature_matrix_off_diag_from(): def test__interferometer_sparse_operator__curvature_matrix_off_diag_func_list_from(): + pytest.importorskip("jax") + operator, mask, rng = _sparse_operator_and_mask() M = operator.M @@ -193,6 +197,8 @@ def test__interferometer_sparse_operator__curvature_matrix_off_diag_func_list_fr def test__interferometer_sparse_operator__curvature_matrix_func_list_from(): + pytest.importorskip("jax") + operator, mask, rng = _sparse_operator_and_mask() M = operator.M @@ -222,6 +228,8 @@ def test__interferometer_sparse_operator__curvature_matrix_func_list_from(): def test__interferometer_sparse_operator__operated_matrix_slim_from(): + pytest.importorskip("jax") + operator, mask, rng = _sparse_operator_and_mask() M = operator.M