diff --git a/autoarray/__init__.py b/autoarray/__init__.py index d1b7f08c6..ca5fd05b7 100644 --- a/autoarray/__init__.py +++ b/autoarray/__init__.py @@ -51,6 +51,9 @@ from .inversion.inversion.interferometer.sparse import ( InversionInterferometerSparse, ) +from .inversion.inversion.interferometer.matrix_free import ( + InversionInterferometerMatrixFree, +) from .inversion.inversion.interferometer.mapping import InversionInterferometerMapping from .inversion.inversion.interferometer.inversion_interferometer_util import ( InterferometerSparseOperator, diff --git a/autoarray/config/general.yaml b/autoarray/config/general.yaml index 35af5233a..4f6b67df7 100644 --- a/autoarray/config/general.yaml +++ b/autoarray/config/general.yaml @@ -19,7 +19,7 @@ inversion: log_det_method: cholesky # How the Bayesian-evidence log-determinant terms are computed. "cholesky" (default) is the historical 2*sum(log(diag(cholesky(M)))); "slogdet" uses logabsdet of slogdet(M), which is identical where M is positive-definite but finite (not NaN) where the Cholesky fails, for gradient-based searches (opt-in, non-default; does not change the default evidence). Under "slogdet" the kernel regularization schemes (Matern/Gaussian/Exponential) also compute the regularization log-det analytically from a Cholesky of their covariance instead of factorizing the formed inverse. See PyAutoArray#391. regularization_term_method: matmul # How the Bayesian-evidence regularization term s^T H s is computed. "matmul" (default) is the historical s @ (H @ s) against the explicitly formed regularization matrix; "cho_solve" evaluates coefficient * s^T C^-1 s for the kernel schemes (Matern/Gaussian/Exponential/MaternAdapt) via one Cholesky solve of their covariance C, avoiding the explicit inverse whose round-off is amplified by cond(C) (~1e9 on clustered traced mesh vertices). Opt-in, non-default; does not change the default evidence. Schemes with no such factorization fall back to the formed matrix. interferometer_numba_nnz_per_source_max: null # How a NumPy (xp=np) single-mapper interferometer inversion chooses between the numba `direct_conv` curvature path (InversionInterferometerSparseNumba) and the FFT one (InversionInterferometerSparse). null (default) = cost model: predicts both routes' single-thread time from source pixels S, total mapping non-zeros and the extent rectangle (FFT ~ S x rfft2+irfft2 cost of the padded (2ny,2nx) extent, direct_conv ~ nnz x ny x nx) and picks the faster; calibrated on autolens_profiling#226 and the Aris 512x512 mask (PyAutoArray#628). A positive value is a hard gate in mean non-zeros per source column (mapper.pix_sizes_for_sub_slim_index.sum() / mapper.params): numba at or below it, FFT above (the pre-#628 behaviour, default was 60.0). Set 0 to disable the numba path entirely. - interferometer_route: null # Force the interferometer curvature route for benchmarking: null (default) routes by interferometer_numba_nnz_per_source_max; "fft" always uses the FFT route; "numba" uses the numba direct_conv route whenever it can represent the inversion (NumPy, one mapper, no linear functions, no over-sampling), overriding the gate and the 0 kill switch. + interferometer_route: null # Force the interferometer curvature route for benchmarking: null (default) routes by interferometer_numba_nnz_per_source_max; "fft" always uses the FFT route; "numba" uses the numba direct_conv route whenever it can represent the inversion (NumPy, one mapper, no linear functions, no over-sampling), overriding the gate and the 0 kill switch. "matrix_free" uses InversionInterferometerMatrixFree, which never forms the dense curvature or regularization matrix (PCG / interior-point solve, stochastic Lanczos log-det); it needs NumPy, one RectangularUniform mapper with MaternSparse regularization, no linear functions and (positive-only) use_edge_zeroed_pixels false, raises on anything else, and has no dense attributes such as reconstruction_noise_map (PyAutoArray#657). interferometer_direct_conv_symmetric_max_work_fraction: 0.7 # When the serial (numba parallel: false) direct_conv curvature route runs, use the symmetric row-restricted kernel (curvature_direct_conv_symmetric: source columns ranked by their top extent row, each convolving only the extent rows later-ranked columns span and only each row's mask x-range, then a mirror pass) when its geometry work fraction (its convolution work / the full kernel's) is at or below this value. Same F to round-off (computed entries bit-identical). Synthetic Aris-512 geometry (compact columns): f ~0.41, 1.81-1.90x single-thread on RAL; Aris's real 512 Delaunay mesh (multiply imaged, tall columns): f ~0.51, ~1.46x measured (PyAutoArray#635). The route cost model scales t_dc by min(1, 1.32 x work fraction), so the kernel wins below f ~0.76; 0.7 keeps a predicted gain >= 1.08x. Set 0 to always use the full kernel. The parallel kernel is unaffected. numba: use_numba: true # If True, use numba-accelerated mapper and Sibson interpolation paths when eligible. diff --git a/autoarray/inversion/inversion/factory.py b/autoarray/inversion/inversion/factory.py index fd4af6932..9b310b56b 100644 --- a/autoarray/inversion/inversion/factory.py +++ b/autoarray/inversion/inversion/factory.py @@ -1,6 +1,7 @@ import numpy as np from typing import List, Optional, Tuple, Union +from autoarray import exc from autoarray.dataset.imaging.dataset import Imaging from autoarray.dataset.interferometer.dataset import Interferometer from autoarray.inversion.inversion.imaging.mapping import InversionImagingMapping @@ -17,6 +18,9 @@ from autoarray.inversion.inversion.interferometer_numba.sparse import ( InversionInterferometerSparseNumba, ) +from autoarray.inversion.inversion.interferometer.matrix_free import ( + InversionInterferometerMatrixFree, +) from autoarray.inversion.inversion.interferometer_numba import ( sparse as interferometer_numba_sparse, ) @@ -224,8 +228,25 @@ def inversion_interferometer_from( # `transform_mapping_matrix` of shape (N_vis, S), which dominates run time and memory at large N_vis, and # `InversionInterferometerSparse` handles zero mappers (the func-func curvature blocks and the data vector # need no mapper). `_use_interferometer_numba` returns False whenever a func-list is present. + route_is_matrix_free = ( + settings is not None and settings.interferometer_route == "matrix_free" + ) + if dataset.sparse_operator is not None: + # The matrix-free route is opt-in only. `InversionInterferometerMatrixFree` raises on + # every configuration it cannot represent (JAX, linear functions, several mappers, a + # non-rectangular mesh, a regularization without a sparse form, ...); there is no silent + # fallback to the dense routes. + if route_is_matrix_free: + return InversionInterferometerMatrixFree( + dataset=dataset, + linear_obj_list=linear_obj_list, + settings=settings, + xp=xp, + preloads=preloads, + ) + # The numba `direct_conv` kernel uses the mapper weights as-is and cannot apply a primary # beam, so a dataset with a non-trivial one (any value other than 1) always takes the FFT # sparse route. `P = ones` keeps the numba route, so it is bit-for-bit `P = None`. @@ -260,6 +281,14 @@ def inversion_interferometer_from( preloads=preloads, ) + if route_is_matrix_free: + raise exc.InversionException( + '`Settings(interferometer_route="matrix_free")` was set but the interferometer ' + "dataset has no `sparse_operator`. The matrix-free route reads the w-tilde kernel " + "and dirty image from it; call `dataset.apply_sparse_operator(...)` first, or use " + '`interferometer_route="fft"` (or leave it unset).' + ) + return InversionInterferometerMapping( dataset=dataset, linear_obj_list=linear_obj_list, diff --git a/autoarray/inversion/inversion/interferometer/inversion_interferometer_matrix_free_util.py b/autoarray/inversion/inversion/interferometer/inversion_interferometer_matrix_free_util.py index 6067e188c..1376980c7 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_matrix_free_util.py +++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_matrix_free_util.py @@ -30,8 +30,8 @@ Phase 1 landed the operator data term and the solver as library utilities, phase 2 the natively sparse Matérn precision (`aa.reg.MaternSparse`), phase 3 the evidence log-determinants (below) and phase 4 the -positivity constraint (below), each with parity tests against the dense sparse-path quantities. None of -it is wired into fits and there is no `Settings` flag. +positivity constraint (below), each with parity tests against the dense sparse-path quantities. Phase 5a +wires them into fits (`InversionInterferometerMatrixFree`, `Settings(interferometer_route="matrix_free")`). Positivity (phase 4) -------------------- @@ -46,7 +46,8 @@ Remaining work (phase 5) ------------------------ -- wiring as an inversion class or a `Settings` option, so fits can select the matrix-free path; +- wiring into fits: landed in phase 5a (#657) as `InversionInterferometerMatrixFree` + (`interferometer/matrix_free.py`), selected by `Settings(interferometer_route="matrix_free")`; - a profiling note, dense vs matrix-free at `50^2`, `64^2`, `128^2` and `192^2` mesh pixels, ingested by PyAutoPulse; - a workspace example; diff --git a/autoarray/inversion/inversion/interferometer/matrix_free.py b/autoarray/inversion/inversion/interferometer/matrix_free.py new file mode 100644 index 000000000..9ba53d20f --- /dev/null +++ b/autoarray/inversion/inversion/interferometer/matrix_free.py @@ -0,0 +1,413 @@ +import numpy as np +from typing import Dict, List, Union + +from autonerves import cached_property + +from autoarray import exc +from autoarray.dataset.interferometer.dataset import Interferometer +from autoarray.inversion.inversion.dataset_interface import DatasetInterface +from autoarray.inversion.inversion.interferometer.sparse import ( + InversionInterferometerSparse, +) +from autoarray.inversion.inversion.interferometer import ( + inversion_interferometer_matrix_free_util as matrix_free_util, +) +from autoarray.inversion.linear_obj.func_list import AbstractLinearObjFuncList +from autoarray.inversion.linear_obj.linear_obj import LinearObj +from autoarray.inversion.mappers.abstract import Mapper +from autoarray.settings import Settings + +_ROUTE_HINT = ( + 'Use `Settings(interferometer_route="fft")` (or leave `interferometer_route` unset) ' + "for this model." +) + +# The regularization methods the matrix-free route needs; only `aa.reg.MaternSparse` supplies +# all three. +_SPARSE_REGULARIZATION_METHODS = ( + "regularization_matrix_sparse_from", + "regularization_factor_sparse_from", + "log_det_regularization_matrix_term_from", +) + + +class InversionInterferometerMatrixFree(InversionInterferometerSparse): + + # The PCG relative residual tolerance and iteration cap of the unconstrained solve. + pcg_rtol = 1.0e-10 + pcg_maxiter = 5000 + + # The stochastic Lanczos quadrature settings of `ln det(F + H)` (the util module defaults). + slq_n_probes = 10 + slq_n_steps = 40 + slq_seed = 20221 + + def __init__( + self, + dataset: Union[Interferometer, DatasetInterface], + linear_obj_list: List[LinearObj], + settings: Settings = None, + xp=np, + preloads=None, + ): + """ + The single-mapper interferometer inversion solved matrix-free: neither the dense `(S, S)` + curvature matrix `F = Aᵀ W~ A` nor the dense regularization matrix `H` is ever formed or + factorised. + + `InversionInterferometerSparse` forms `F` and `H` densely and factorises `F + H`, which caps + a pixelized interferometer source at roughly `100^2` mesh pixels. This class instead uses the + operator form of `inversion_interferometer_matrix_free_util` (PyAutoArray #638, read that + module's docstring for the method): + + - `F s = Mᵀ W~ (M s)` is applied by `MatrixFreeDataOperator.curvature_matvec`, with `M` the + sparse mapping matrix and `W~` an FFT convolution, and the data vector is `D = Mᵀ d~`; + - `H` is the natively sparse Matérn precision of `aa.reg.MaternSparse`, which already + carries its coefficient (so the utilities' `λ` is `1.0`); + - the reconstruction solves `(F + H) s = D` by preconditioned conjugate gradients (Powell + preconditioner), or, when `settings.use_positive_only_solver` is `True`, by the + interior-point positive solve with an exact active-set polish; + - `ln det(F + H)` is the exact `ln det P` of the preconditioner plus a stochastic Lanczos + quadrature estimate of `ln det(P⁻¹(F + H))` (10 probes, 40 steps, fixed seed 20221, so it + is deterministic and smooth in the model parameters). It is the only approximate evidence + term: within 0.5% of the dense value on the util module's test case, with an accuracy set + by how well `P` approximates `F + H`; + - `ln det H` is exact (`MaternSparse.log_det_regularization_matrix_term_from`). + + The mapping matrix carries any over-sampling (`sub_fraction`) and primary beam exactly as the + parent's sparse triplets do, so both are supported and match the FFT route. + + The class raises on every configuration it cannot represent rather than silently working + around it (see `_check_preconditions`), and the inversion factory routes here only when + `Settings(interferometer_route="matrix_free")` is set, with no fallback. Unsupported: + + - a JAX (`xp=jnp`) array module -- the utilities are NumPy / scipy only; + - linear-function lists (linear light profiles, MGE bases) and more than one mapper; + - a mesh other than `RectangularUniform`; + - a regularization other than `MaternSparse` (one without a sparse matrix, factor and exact + log determinant); + - edge-zeroed pixels (`settings.use_edge_zeroed_pixels`) with the positive-only solver: the + base class's edge-zeroing subsets the dense system, which this class never forms. Pass + `use_edge_zeroed_pixels=False`. + + Quantities that need the dense matrices -- `curvature_matrix`, `curvature_matrix_diag`, + `curvature_reg_matrix`, `curvature_reg_matrix_reduced`, `reconstruction_covariance_matrix`, + `reconstruction_noise_map` and `reconstruction_noise_map_with_covariance` -- raise an + `InversionException`. Reconstruction uncertainty maps need a stochastic diagonal estimator + of `(F + H)⁻¹`, which is future work. + + Parameters + ---------- + dataset + The interferometer dataset (or `DatasetInterface`) being reconstructed. It must carry a + `sparse_operator`, whose `nufft_precision_operator` is the `W~` kernel. + linear_obj_list + The linear objects reconstructing the data. Exactly one `Mapper`, and nothing else. + settings + The inversion settings (`autoarray.settings.Settings`). + xp + The array module. Must be `numpy`. + preloads + Optional `AbstractPreloads`, forwarded to the parent unchanged. + """ + if xp is not np: + raise exc.InversionException( + '`InversionInterferometerMatrixFree` (the `interferometer_route="matrix_free"` ' + f"route) was passed a non-NumPy array module ({xp!r}). The matrix-free " + "utilities are NumPy / scipy only (scipy CG, sparse LU), with no JAX path.\n\n" + + _ROUTE_HINT + ) + + super().__init__( + dataset=dataset, + linear_obj_list=linear_obj_list, + settings=settings, + xp=xp, + preloads=preloads, + ) + + self._check_preconditions() + + def _check_preconditions(self) -> None: + """ + Raise on every configuration the matrix-free route cannot represent. + + These are raised, not worked around: each would otherwise give a different (silently wrong) + reconstruction or evidence, or none at all. + """ + from autoarray.inversion.mesh.mesh.rectangular_uniform import ( + RectangularUniform, + ) + + if self.has(cls=AbstractLinearObjFuncList): + raise exc.InversionException( + "A linear-function list (e.g. a linear light profile or MGE basis) was passed to " + '`InversionInterferometerMatrixFree` (the `interferometer_route="matrix_free"` ' + "route). The matrix-free operator represents a single mapper's sparse mapping " + "matrix only and has no mapper x function or function x function block.\n\n" + + _ROUTE_HINT + ) + + total_mappers = self.total(cls=Mapper) + + if total_mappers != 1: + raise exc.InversionException( + f'`InversionInterferometerMatrixFree` (the `interferometer_route="matrix_free"` ' + f"route) was passed {total_mappers} mappers. The matrix-free operator represents " + "exactly one mapper and has no off-diagonal mapper x mapper block.\n\n" + + _ROUTE_HINT + ) + + mapper = self.cls_list_from(cls=Mapper)[0] + + if not isinstance(mapper.mesh, RectangularUniform): + raise exc.InversionException( + '`InversionInterferometerMatrixFree` (the `interferometer_route="matrix_free"` ' + f"route) was passed a mapper with a {type(mapper.mesh).__name__} mesh. The route " + "requires a `RectangularUniform` mesh, the mesh the natively sparse Matérn " + "regularization (`aa.reg.MaternSparse`) is built and tested on.\n\n" + + _ROUTE_HINT + ) + + regularization = mapper.regularization + + missing = [ + method + for method in _SPARSE_REGULARIZATION_METHODS + if not hasattr(regularization, method) + ] + + if missing: + raise exc.InversionException( + '`InversionInterferometerMatrixFree` (the `interferometer_route="matrix_free"` ' + f"route) was passed a {type(regularization).__name__} regularization, which lacks " + f"{', '.join(f'`{method}`' for method in missing)}. The route needs the " + "regularization matrix as a sparse matrix, its sparse factor (for the evidence " + "probes) and its exact log determinant; only `aa.reg.MaternSparse` supplies " + "them.\n\n" + "Use `aa.reg.MaternSparse`, or " + _ROUTE_HINT + ) + + if ( + self.settings.use_positive_only_solver + and self.settings.use_edge_zeroed_pixels + ): + raise exc.InversionException( + '`InversionInterferometerMatrixFree` (the `interferometer_route="matrix_free"` ' + "route) was passed `use_positive_only_solver=True` with " + "`use_edge_zeroed_pixels=True`. Edge-zeroing subsets the dense " + "`curvature_reg_matrix`, which this route never forms, so it is not applied " + "here.\n\n" + "Pass `Settings(..., use_edge_zeroed_pixels=False)` to solve the full positive " + "system matrix-free, or " + _ROUTE_HINT + ) + + @property + def _mapper(self) -> Mapper: + return self.cls_list_from(cls=Mapper)[0] + + @cached_property + def operator(self) -> matrix_free_util.MatrixFreeDataOperator: + """ + The matrix-free data operator of the single mapper: `F s` by `curvature_matvec` and + `D = Mᵀ d~` by `data_vector`. + """ + return matrix_free_util.MatrixFreeDataOperator.from_inversion( + inversion=self, mapper=self._mapper + ) + + @cached_property + def regularization_matrix_sparse(self): + """ + The regularization matrix `H` as a `scipy.sparse` matrix. It already carries the + regularization coefficient, so every utility is called with `coefficient=1.0`. + """ + mapper = self._mapper + + return mapper.regularization.regularization_matrix_sparse_from(mapper) + + @cached_property + def regularization_factor_sparse(self): + """ + The sparse factor `R` of exactly `H = RᵀR`: `MaternSparse.regularization_factor_sparse_from` + returns the factor of the unscaled precision, so the coefficient is folded in as + `sqrt(coefficient)`. + """ + mapper = self._mapper + regularization = mapper.regularization + + return np.sqrt( + regularization.coefficient + ) * regularization.regularization_factor_sparse_from(mapper) + + @cached_property + def data_vector(self) -> np.ndarray: + """ + The data vector `D = Mᵀ d~`, from the matrix-free operator. + """ + return self.operator.data_vector + + @cached_property + def reconstruction(self) -> np.ndarray: + """ + The reconstruction solving `(F + H) s = D` without forming `F` or `H` densely. + + With `settings.use_positive_only_solver` it is the non-negative solution + (`reconstruction_positive_via_interior_point_from`, exact zeros on the held pixels), stored + with its diagnostics on `positive_result`; otherwise the unconstrained PCG solution + (`reconstruction_via_pcg_from`), which raises an `InversionException` if CG does not converge + within `pcg_maxiter` iterations. + """ + if self.settings.use_positive_only_solver: + result = matrix_free_util.reconstruction_positive_via_interior_point_from( + operator=self.operator, + regularization_matrix_sparse=self.regularization_matrix_sparse, + coefficient=1.0, + ) + + self.positive_result = result + + return np.asarray(result.reconstruction, dtype=np.float64) + + reconstruction, info = matrix_free_util.reconstruction_via_pcg_from( + operator=self.operator, + regularization_matrix_sparse=self.regularization_matrix_sparse, + coefficient=1.0, + rtol=self.pcg_rtol, + maxiter=self.pcg_maxiter, + ) + + if info != 0: + raise exc.InversionException( + "The matrix-free PCG solve of `InversionInterferometerMatrixFree` did not " + f"converge to rtol={self.pcg_rtol} within {info} iterations." + ) + + return np.asarray(reconstruction, dtype=np.float64) + + @property + def regularization_term(self) -> float: + """ + The regularization term `sᵀ H s`, with the sparse `H` (which carries the coefficient, as the + dense `regularization_matrix` does). + """ + s = self.reconstruction + + return float(s @ (self.regularization_matrix_sparse @ s)) + + @cached_property + def log_det_curvature_reg_matrix_term(self) -> float: + """ + `ln det(F + H)`: the exact `ln det P` of the Powell preconditioner plus the stochastic Lanczos + quadrature estimate of `ln det(P⁻¹(F + H))`, both returned together by + `log_det_curvature_reg_via_slq_from` (the estimate's per-probe values already include + `ln det P`). + """ + log_det, _ = matrix_free_util.log_det_curvature_reg_via_slq_from( + operator=self.operator, + regularization_matrix_sparse=self.regularization_matrix_sparse, + regularization_factor_sparse=self.regularization_factor_sparse, + coefficient=1.0, + n_probes=self.slq_n_probes, + n_steps=self.slq_n_steps, + seed=self.slq_seed, + ) + + return log_det + + @cached_property + def log_det_regularization_matrix_term(self) -> float: + """ + The exact `ln det H` from the regularization's sparse factorisation + (`MaternSparse.log_det_regularization_matrix_term_from`, which includes the coefficient). + """ + mapper = self._mapper + + return float( + mapper.regularization.log_det_regularization_matrix_term_from(mapper) + ) + + @property + def fast_chi_squared(self): + """ + The chi-squared `sᵀ F s - 2 sᵀ D + dᵀ N⁻¹ d`, as `AbstractInversionInterferometer.fast_chi_squared` + but with `F s` applied matrix-free (`operator.curvature_matvec`). + """ + s = self.reconstruction + + chi_squared_term_1 = float(s @ self.operator.curvature_matvec(s)) + + chi_squared_term_2 = -2.0 * float(s @ self.data_vector) + + if self.dataset.data is None: + # The interface carries no visibilities, so term 3 is a precomputed scalar: the + # interface's own `data_term` (that of profile-subtracted visibilities) when given, + # else the `d^T N^-1 d` the sparse operator cached when it was built from the raw data + # (e.g. a pixelization-only fit, where nothing was subtracted). + chi_squared_term_3 = getattr(self.dataset, "data_term", None) + + if chi_squared_term_3 is None: + chi_squared_term_3 = getattr( + self.dataset.sparse_operator, "data_term", None + ) + + if chi_squared_term_3 is None: + raise exc.InversionException( + "The dataset input to this interferometer inversion has `data=None`, which " + "is only valid when its `sparse_operator` carries a precomputed `data_term` " + "(sum(d_r^2/sigma_r^2) + sum(d_i^2/sigma_i^2)) -- e.g. one built by " + "`Interferometer.apply_sparse_operator` or " + "`Interferometer.apply_sparse_operator_from_chunks`. This dataset's " + "sparse operator has none, so `fast_chi_squared` cannot be computed. " + "Pass the visibilities as `data` instead." + ) + else: + chi_squared_term_3 = np.sum( + self.dataset.data.array.real**2.0 + / self.dataset.noise_map.array.real**2.0 + ) + np.sum( + self.dataset.data.array.imag**2.0 + / self.dataset.noise_map.array.imag**2.0 + ) + + return chi_squared_term_1 + chi_squared_term_2 + chi_squared_term_3 + + def _dense_unavailable(self, name: str): + return exc.InversionException( + f"`{name}` is not available on `InversionInterferometerMatrixFree` (the " + '`interferometer_route="matrix_free"` route), which never forms the dense curvature ' + "or regularization matrices. Reconstruction uncertainties need a stochastic diagonal " + "estimator of (F + H)^-1, which is not implemented yet.\n\n" + _ROUTE_HINT + ) + + @property + def curvature_matrix(self): + raise self._dense_unavailable("curvature_matrix") + + @property + def curvature_matrix_diag(self): + raise self._dense_unavailable("curvature_matrix_diag") + + @property + def curvature_reg_matrix(self): + raise self._dense_unavailable("curvature_reg_matrix") + + @property + def curvature_reg_matrix_reduced(self): + raise self._dense_unavailable("curvature_reg_matrix_reduced") + + @property + def reconstruction_covariance_matrix(self): + raise self._dense_unavailable("reconstruction_covariance_matrix") + + @property + def reconstruction_noise_map_with_covariance(self): + raise self._dense_unavailable("reconstruction_noise_map_with_covariance") + + @property + def reconstruction_noise_map(self): + raise self._dense_unavailable("reconstruction_noise_map") + + @property + def reconstruction_noise_map_dict(self) -> Dict[LinearObj, np.ndarray]: + raise self._dense_unavailable("reconstruction_noise_map_dict") diff --git a/autoarray/inversion/plot/inversion_plots.py b/autoarray/inversion/plot/inversion_plots.py index 22caa36ef..cde4de88b 100644 --- a/autoarray/inversion/plot/inversion_plots.py +++ b/autoarray/inversion/plot/inversion_plots.py @@ -192,7 +192,8 @@ def _recon_array(): mesh_grid=mesh_grid, lines=lines, ) - except (KeyError, TypeError, np.linalg.LinAlgError): + except (KeyError, TypeError, np.linalg.LinAlgError, exc.InversionException): + # `InversionException`: the matrix-free interferometer route has no noise map. pass # panel 7: regularization weights @@ -512,11 +513,12 @@ def save_reconstruction_csv( try: noise_map = inversion.reconstruction_noise_map_dict[mapper] - except np.linalg.LinAlgError: + except (np.linalg.LinAlgError, exc.InversionException) as error: logger.warning( f"save_reconstruction_csv: could not compute the reconstruction noise map for " - f"mapper {i} (singular curvature_reg_matrix); writing the noise_map column of " - f"source_plane_reconstruction_{i}.csv as nan." + f"mapper {i} ({type(error).__name__}: e.g. a singular curvature_reg_matrix, or " + f"an inversion route with no noise map such as the matrix-free one); writing " + f"the noise_map column of source_plane_reconstruction_{i}.csv as nan." ) noise_map = None diff --git a/autoarray/settings.py b/autoarray/settings.py index 58b0124a7..7771e898b 100644 --- a/autoarray/settings.py +++ b/autoarray/settings.py @@ -209,6 +209,15 @@ def __init__( (`InversionInterferometerSparseNumba`) whenever it can represent the inversion (NumPy backend, one mapper, no linear functions, no over-sampling, numba installed), overriding the geometry gate and the `0` kill switch. + ``"matrix_free"`` uses `InversionInterferometerMatrixFree`, which never forms the + dense curvature matrix `F` or regularization matrix `H`: `F s` is applied as an FFT + convolution, the reconstruction is a preconditioned CG (or interior-point positive) + solve and `ln det(F + H)` a stochastic Lanczos quadrature estimate. It requires the + NumPy backend, one `RectangularUniform` mapper with `aa.reg.MaternSparse` + regularization, no linear functions and, with the positive-only solver, + `use_edge_zeroed_pixels=False`; it raises on anything else rather than falling + back, and dense-only attributes (e.g. `curvature_matrix`, + `reconstruction_noise_map`) are unavailable. interferometer_direct_conv_symmetric_max_work_fraction When the serial numba `direct_conv` route runs, the largest geometry work fraction at which it uses the symmetric, row-restricted kernel @@ -449,18 +458,21 @@ def interferometer_numba_nnz_per_source_max(self) -> Optional[float]: def interferometer_route(self) -> Optional[str]: """ A forced interferometer curvature route: `None` (route by - `interferometer_numba_nnz_per_source_max`), ``"numba"`` or ``"fft"``. + `interferometer_numba_nnz_per_source_max`), ``"numba"``, ``"fft"`` or + ``"matrix_free"``. See the constructor docstring; ``"numba"`` still requires an inversion the - `direct_conv` kernel can represent. + `direct_conv` kernel can represent, and ``"matrix_free"`` raises on an inversion + `InversionInterferometerMatrixFree` cannot represent. """ value = self._interferometer_route if value is None: value = self._inversion_config_value("interferometer_route", None) - if value not in (None, "numba", "fft"): + if value not in (None, "numba", "fft", "matrix_free"): raise ValueError( - f"interferometer_route={value!r} is invalid; expected None, 'numba' or 'fft'." + f"interferometer_route={value!r} is invalid; expected None, 'numba', 'fft' " + "or 'matrix_free'." ) return value diff --git a/test_autoarray/inversion/inversion/interferometer/test_matrix_free.py b/test_autoarray/inversion/inversion/interferometer/test_matrix_free.py new file mode 100644 index 000000000..5b843764d --- /dev/null +++ b/test_autoarray/inversion/inversion/interferometer/test_matrix_free.py @@ -0,0 +1,411 @@ +""" +Tests of `InversionInterferometerMatrixFree`, the fit-selectable matrix-free interferometer inversion +(`Settings(interferometer_route="matrix_free")`, PyAutoArray #657, phase 5a of #638). + +Every fit quantity is held against the FFT sparse route (`InversionInterferometerSparse`) on the +`MaternSparse` rectangular-mesh cases of `test_matrix_free_util.py`, whose fixture helpers are reused. +Like that module, these tests carry no module-level JAX dependency, so the `unittest-nojax` leg runs +them. +""" + +import numpy as np +import pytest + +import autoarray as aa +from autoarray import exc +from autoarray.fit import fit_util +from autoarray.inversion.inversion.factory import inversion_interferometer_from + +from test_autoarray.inversion.inversion.interferometer.test_matrix_free_util import ( + _dataset_sparse_from, + _delaunay_mapper_from, + _rectangular_mapper_from, +) + + +def _settings_from(route, use_positive_only_solver=False): + return aa.Settings( + interferometer_route=route, + use_positive_only_solver=use_positive_only_solver, + use_edge_zeroed_pixels=False, + ) + + +def _matern_case_from(coefficient=1.0, mesh_shape=(6, 6)): + """ + The util tests' data-dominated dataset (500 visibilities on long baselines) and a `MaternSparse` + rectangular mapper (over-sampled, `over_sample_size=2`). + """ + mask, dataset_sparse = _dataset_sparse_from(n_visibilities=500, uv_scale=1.0e6) + + mapper = _rectangular_mapper_from( + mask=mask, + regularization=aa.reg.MaternSparse(coefficient=coefficient, scale=1.0), + shape=mesh_shape, + ) + + return mask, dataset_sparse, mapper + + +def _inversion_pair_from( + dataset_sparse, mapper, use_positive_only_solver=False, linear_obj_list=None +): + linear_obj_list = linear_obj_list or [mapper] + + inversion_matrix_free = aa.Inversion( + dataset=dataset_sparse, + linear_obj_list=linear_obj_list, + settings=_settings_from("matrix_free", use_positive_only_solver), + ) + inversion_fft = aa.Inversion( + dataset=dataset_sparse, + linear_obj_list=linear_obj_list, + settings=_settings_from("fft", use_positive_only_solver), + ) + + return inversion_matrix_free, inversion_fft + + +def _assert_relative_close(value, reference, rtol): + value = np.asarray(value) + reference = np.asarray(reference) + + assert np.max(np.abs(value - reference)) <= rtol * np.max(np.abs(reference)) + + +def test__factory__routes_by_interferometer_route(): + _, dataset_sparse, mapper = _matern_case_from() + + inversion = inversion_interferometer_from( + dataset=dataset_sparse, + linear_obj_list=[mapper], + settings=_settings_from("matrix_free"), + ) + + assert isinstance(inversion, aa.InversionInterferometerMatrixFree) + + inversion = inversion_interferometer_from( + dataset=dataset_sparse, + linear_obj_list=[mapper], + settings=_settings_from("fft"), + ) + + assert type(inversion) is aa.InversionInterferometerSparse + + inversion_matrix_free, _ = _inversion_pair_from( + dataset_sparse=dataset_sparse, mapper=mapper + ) + + assert isinstance(inversion_matrix_free, aa.InversionInterferometerMatrixFree) + + +def test__factory__matrix_free_without_sparse_operator__raises(): + mask = aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=0.5, radius=2.5) + + rng = np.random.default_rng(seed=1) + + dataset = aa.Interferometer( + data=aa.Visibilities(visibilities=rng.normal(size=(20, 2))), + noise_map=aa.VisibilitiesNoiseMap(visibilities=np.ones((20, 2))), + uv_wavelengths=rng.normal(size=(20, 2)) * 1.0e4, + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + ) + + mapper = _rectangular_mapper_from( + mask=mask, regularization=aa.reg.MaternSparse(coefficient=1.0, scale=1.0) + ) + + with pytest.raises(exc.InversionException, match="sparse_operator"): + aa.Inversion( + dataset=dataset, + linear_obj_list=[mapper], + settings=_settings_from("matrix_free"), + ) + + +@pytest.mark.parametrize("coefficient", [1.0, 0.01]) +def test__fit_quantities__match_fft_route(coefficient): + """ + The reconstruction (PCG, rtol 1e-10) matches the dense solve to ~1e-11 of the max; the + chi-squared, regularization term and exact `ln det H` to round-off; `ln det(F + H)` is the SLQ + estimate, within its documented 0.5% tolerance, and so is the assembled evidence (the tolerance + logic of `test__log_evidence_matrix_free_from__matches_dense_sparse_path_evidence`). + """ + _, dataset_sparse, mapper = _matern_case_from(coefficient=coefficient) + + inversion_matrix_free, inversion_fft = _inversion_pair_from( + dataset_sparse=dataset_sparse, mapper=mapper + ) + + reconstruction = inversion_matrix_free.reconstruction + + assert isinstance(reconstruction, np.ndarray) + assert reconstruction.shape == (inversion_matrix_free.operator.params,) + + _assert_relative_close(reconstruction, inversion_fft.reconstruction, rtol=1.0e-6) + + _assert_relative_close( + inversion_matrix_free.data_vector, inversion_fft.data_vector, rtol=1.0e-10 + ) + + for name in [ + "regularization_term", + "fast_chi_squared", + "log_det_regularization_matrix_term", + ]: + assert float(getattr(inversion_matrix_free, name)) == pytest.approx( + float(getattr(inversion_fft, name)), rel=1.0e-8 + ), name + + log_det_curvature_reg_dense = float(inversion_fft.log_det_curvature_reg_matrix_term) + + assert inversion_matrix_free.log_det_curvature_reg_matrix_term == pytest.approx( + log_det_curvature_reg_dense, rel=0.005 + ) + + noise_normalization = dataset_sparse.sparse_operator.noise_normalization + + def log_evidence_from(inversion): + return fit_util.log_evidence_from( + chi_squared=inversion.fast_chi_squared, + regularization_term=inversion.regularization_term, + log_curvature_regularization_term=inversion.log_det_curvature_reg_matrix_term, + log_regularization_term=inversion.log_det_regularization_matrix_term, + noise_normalization=noise_normalization, + ) + + log_evidence_dense = log_evidence_from(inversion_fft) + log_evidence_matrix_free = log_evidence_from(inversion_matrix_free) + + assert abs(log_evidence_matrix_free - log_evidence_dense) <= 0.5 * 0.005 * abs( + log_det_curvature_reg_dense + ) + + +def test__log_det_curvature_reg_matrix_term__is_deterministic(): + _, dataset_sparse, mapper = _matern_case_from() + + inversion_0, _ = _inversion_pair_from(dataset_sparse=dataset_sparse, mapper=mapper) + inversion_1, _ = _inversion_pair_from(dataset_sparse=dataset_sparse, mapper=mapper) + + assert ( + inversion_0.log_det_curvature_reg_matrix_term + == inversion_1.log_det_curvature_reg_matrix_term + ) + + +def test__mapped_reconstructed_data_and_image__match_fft_route(): + _, dataset_sparse, mapper = _matern_case_from() + + inversion_matrix_free, inversion_fft = _inversion_pair_from( + dataset_sparse=dataset_sparse, mapper=mapper + ) + + _assert_relative_close( + inversion_matrix_free.mapped_reconstructed_data.array, + inversion_fft.mapped_reconstructed_data.array, + rtol=1.0e-6, + ) + _assert_relative_close( + inversion_matrix_free.mapped_reconstructed_operated_data.array, + inversion_fft.mapped_reconstructed_operated_data.array, + rtol=1.0e-6, + ) + _assert_relative_close( + inversion_matrix_free.reconstruction_dict[mapper], + inversion_fft.reconstruction_dict[mapper], + rtol=1.0e-6, + ) + + +@pytest.mark.parametrize("coefficient", [0.01, 1.0e4]) +def test__positive_only__matches_fft_route_positive_only(coefficient): + """ + The setup of `test__reconstruction_positive_via_interior_point__matches_sparse_path_positive_only` + (10 x 10 `MaternSparse` mesh, edge-zeroing off), through the factory. + """ + _, dataset_sparse, mapper = _matern_case_from( + coefficient=coefficient, mesh_shape=(10, 10) + ) + + inversion_matrix_free, inversion_fft = _inversion_pair_from( + dataset_sparse=dataset_sparse, mapper=mapper, use_positive_only_solver=True + ) + + reconstruction = inversion_matrix_free.reconstruction + reconstruction_fft = np.asarray(inversion_fft.reconstruction) + + assert np.mean(reconstruction_fft == 0.0) >= 0.25 + assert np.all(reconstruction >= 0.0) + assert np.array_equal(reconstruction == 0.0, reconstruction_fft == 0.0) + assert inversion_matrix_free.positive_result.converged + + _assert_relative_close(reconstruction, reconstruction_fft, rtol=1.0e-6) + + +def test__primary_beam__matches_fft_route(): + """ + The sparse mapping matrix is built from the parent's beamed triplets, so a non-trivial primary beam + is supported and matches the FFT route. + """ + mask = aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=0.5, radius=2.5) + + rng = np.random.default_rng(seed=1) + + n_visibilities = 500 + + dataset = aa.Interferometer( + data=aa.Visibilities(visibilities=rng.normal(size=(n_visibilities, 2))), + noise_map=aa.VisibilitiesNoiseMap(visibilities=np.ones((n_visibilities, 2))), + uv_wavelengths=rng.normal(size=(n_visibilities, 2)) * 1.0e6, + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + primary_beam=aa.Array2D( + values=np.random.default_rng(seed=2).uniform( + 0.2, 1.0, size=mask.pixels_in_mask + ), + mask=mask, + ), + ) + + dataset_sparse = dataset.apply_sparse_operator(use_jax=False) + + mapper = _rectangular_mapper_from( + mask=mask, regularization=aa.reg.MaternSparse(coefficient=1.0, scale=1.0) + ) + + inversion_matrix_free, inversion_fft = _inversion_pair_from( + dataset_sparse=dataset_sparse, mapper=mapper + ) + + _assert_relative_close( + inversion_matrix_free.reconstruction, inversion_fft.reconstruction, rtol=1.0e-6 + ) + assert inversion_matrix_free.fast_chi_squared == pytest.approx( + float(inversion_fft.fast_chi_squared), rel=1.0e-8 + ) + _assert_relative_close( + inversion_matrix_free.mapped_reconstructed_data.array, + inversion_fft.mapped_reconstructed_data.array, + rtol=1.0e-6, + ) + + +def test__unsupported_configurations__raise(): + mask, dataset_sparse, mapper = _matern_case_from() + + def inversion_from(linear_obj_list, settings=None): + return aa.Inversion( + dataset=dataset_sparse, + linear_obj_list=linear_obj_list, + settings=settings or _settings_from("matrix_free"), + ) + + # Two mappers. + mapper_1 = _rectangular_mapper_from( + mask=mask, regularization=aa.reg.MaternSparse(coefficient=1.0, scale=1.0) + ) + + with pytest.raises(exc.InversionException, match="2 mappers"): + inversion_from([mapper, mapper_1]) + + # A linear function list. + func_list = aa.m.MockLinearObjFuncList( + parameters=1, + mapping_matrix=np.ones((mask.pixels_in_mask, 1)), + ) + + with pytest.raises(exc.InversionException, match="linear-function list"): + inversion_from([func_list, mapper]) + + # A Delaunay mesh. + with pytest.raises(exc.InversionException, match="Delaunay mesh"): + inversion_from([_delaunay_mapper_from(mask=mask)]) + + # A regularization without a sparse form. + mapper_constant = _rectangular_mapper_from(mask=mask, coefficient=1.0) + + with pytest.raises(exc.InversionException, match="Constant regularization"): + inversion_from([mapper_constant]) + + # Edge-zeroing with the positive-only solver. + with pytest.raises(exc.InversionException, match="use_edge_zeroed_pixels"): + inversion_from( + [mapper], + settings=aa.Settings( + interferometer_route="matrix_free", + use_positive_only_solver=True, + use_edge_zeroed_pixels=True, + ), + ) + + # Every message names the route and points at the FFT route. + with pytest.raises(exc.InversionException) as error: + inversion_from([mapper_constant]) + + assert "matrix_free" in str(error.value) + assert 'interferometer_route="fft"' in str(error.value) + + +def test__jax_array_module__raises(): + jnp = pytest.importorskip("jax.numpy") + + _, dataset_sparse, mapper = _matern_case_from() + + with pytest.raises(exc.InversionException, match="non-NumPy array module"): + aa.InversionInterferometerMatrixFree( + dataset=dataset_sparse, + linear_obj_list=[mapper], + settings=_settings_from("matrix_free"), + xp=jnp, + ) + + +@pytest.mark.parametrize( + "name", + [ + "curvature_matrix", + "curvature_matrix_diag", + "curvature_reg_matrix", + "curvature_reg_matrix_reduced", + "reconstruction_covariance_matrix", + "reconstruction_noise_map", + "reconstruction_noise_map_with_covariance", + "reconstruction_noise_map_dict", + ], +) +def test__dense_attributes__raise(name): + _, dataset_sparse, mapper = _matern_case_from() + + inversion, _ = _inversion_pair_from(dataset_sparse=dataset_sparse, mapper=mapper) + + with pytest.raises(exc.InversionException, match=name): + getattr(inversion, name) + + +def test__save_reconstruction_csv__writes_nan_noise_map(tmp_path): + import autoarray.plot # noqa: F401 (initialises the plot package before its submodule) + from autoarray.inversion.plot.inversion_plots import save_reconstruction_csv + + _, dataset_sparse, mapper = _matern_case_from() + + inversion, _ = _inversion_pair_from(dataset_sparse=dataset_sparse, mapper=mapper) + + save_reconstruction_csv(inversion=inversion, output_path=tmp_path) + + lines = (tmp_path / "source_plane_reconstruction_0.csv").read_text().splitlines() + + assert lines[0] == "y,x,reconstruction,noise_map" + assert len(lines) == 1 + mapper.params + assert all(line.endswith(",nan") for line in lines[1:]) + + +def test__settings__interferometer_route_validation(): + assert ( + aa.Settings(interferometer_route="matrix_free").interferometer_route + == "matrix_free" + ) + + with pytest.raises(ValueError): + aa.Settings(interferometer_route="typo")