From 74d687270699cf4d5627abd3dc4acd5e1e2ac732 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 21 Oct 2026 14:05:31 +0100 Subject: [PATCH 1/2] feat: matrix-free interferometer positivity via primal-dual interior point with active-set polish Phase 4 of #638 (#648). reconstruction_positive_via_interior_point_from solves (F + lambda H) s = D, s >= 0 with a Mehrotra predictor-corrector interior point whose Newton systems are CG-solved with the Powell preconditioner plus the barrier diagonal (sparse LU per step), then an exact free-set CG solve with primal-dual active-set corrections so held pixels are exact zeros. Module docstring lists the remaining phase-5 work. Co-Authored-By: Claude Fable 5.1 --- ...version_interferometer_matrix_free_util.py | 348 +++++++++++++++++- 1 file changed, 342 insertions(+), 6 deletions(-) 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 6315173e8..6067e188c 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_matrix_free_util.py +++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_matrix_free_util.py @@ -29,10 +29,30 @@ `W~` by its peak, which keeps it sparse; it is factorised once per `λ` by a sparse LU. Phase 1 landed the operator data term and the solver as library utilities, phase 2 the natively sparse -Matérn precision (`aa.reg.MaternSparse`), and phase 3 the evidence log-determinants (below), each with -parity tests against the dense sparse-path quantities. None of it is wired into fits and there is no -`Settings` flag. Later phases (interior-point positivity, profiling and workspace examples) follow as -their own tasks. +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 (phase 4) +-------------------- +`reconstruction_positive_via_interior_point_from` solves `(F + λH) s = D` subject to `s >= 0` by a +Mehrotra predictor-corrector primal-dual interior point method whose Newton systems are solved by CG +preconditioned by the Powell matrix plus the barrier diagonal, followed by an exact solve on the +detected support with primal-dual active-set corrections (Hintermüller, Ito & Kunisch 2002), so the +held pixels are exact zeros. On the test case (a `10 x 10` `MaternSparse` mesh on the data-dominated +dataset, 33-72% of pixels at zero) it reproduces the support of `scipy.optimize.nnls` on the whitened +dense system exactly and its values to ~1e-10 of the max in 8-10 interior-point steps and one polish +solve. + +Remaining work (phase 5) +------------------------ +- wiring as an inversion class or a `Settings` option, so fits can select the matrix-free path; +- 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; +- reconstruction uncertainty maps, the diagonal of `(F + λH)⁻¹`, by a stochastic diagonal estimator + (there is no dense inverse to read it from); +- linear light profiles and multiple mappers (every function here assumes a single mapper's `M`). The log evidence (phase 3) -------------------------- @@ -381,8 +401,8 @@ def reconstruction_via_pcg_from( preconditioner (`powell_preconditioner_from`) applied through its sparse LU solve. This is the matrix-free counterpart of the positive-negative (`use_positive_only_solver=False`) - solve of an inversion's `curvature_reg_matrix`. No positivity is imposed (an interior-point - positive solve is a later phase). + solve of an inversion's `curvature_reg_matrix`. No positivity is imposed (see + `reconstruction_positive_via_interior_point_from` for the positive solve). Parameters ---------- @@ -731,3 +751,319 @@ def log_evidence_matrix_free_from( + float(noise_normalization) ) ) + + +class PositiveReconstructionResult: + def __init__( + self, + reconstruction: np.ndarray, + support: np.ndarray, + interior_point_steps: int, + polish_iterations: int, + gap: float, + residual: float, + converged: bool, + ): + """ + The output of `reconstruction_positive_via_interior_point_from`. + + Parameters + ---------- + reconstruction + The non-negative reconstruction `s`, with exact zeros on the held (active) set. + support + Boolean mask of the free (passive) pixels, those solved for; `reconstruction` is zero + wherever it is `False`. + interior_point_steps + The number of Mehrotra predictor-corrector steps taken. + polish_iterations + The number of exact free-set solves in the active-set polish (one if the support detected + from the interior point was already optimal). + gap + The normalised duality gap `n μ / |xᵀD|` of the final interior-point iterate. + residual + The normalised dual residual `|A x - D - z| / |D|` of the final interior-point iterate. + converged + `True` if the active-set polish ended on a support satisfying the KKT conditions (no + negative free pixel and no held pixel with a negative gradient), `False` if it reached + its iteration cap and the result was clipped at zero. + """ + self.reconstruction = reconstruction + self.support = support + self.interior_point_steps = interior_point_steps + self.polish_iterations = polish_iterations + self.gap = gap + self.residual = residual + self.converged = converged + + +def _step_to_boundary(v: np.ndarray, dv: np.ndarray) -> float: + """ + The largest step `α <= 1` keeping `v + α dv >= 0` (for `v > 0`). + """ + decreasing = dv < 0.0 + + if not np.any(decreasing): + return 1.0 + + return float(min(1.0, np.min(-v[decreasing] / dv[decreasing]))) + + +def reconstruction_positive_via_interior_point_from( + operator: MatrixFreeDataOperator, + regularization_matrix_sparse: scipy.sparse.spmatrix, + coefficient: float, + max_steps: int = 60, + tol: float = 1.0e-9, + cg_rtol_min: float = 1.0e-8, + cg_rtol_max: float = 1.0e-2, + polish_rtol: float = 1.0e-10, + max_polish: int = 10, + maxiter: int = 5000, +) -> PositiveReconstructionResult: + """ + Returns the non-negative reconstruction solving + + min_s ½ sᵀ(F + λH)s - sᵀD subject to s >= 0 + + (equivalently `(F + λH) s = D` on the free pixels, zero elsewhere) without forming or factorising + the dense `(S, S)` matrix `A = F + λH`: every operation is a matrix-free product with `A` or a + sparse factorisation of the Powell preconditioner `P = κ₀ MᵀM + λH` plus a diagonal. + + This is the matrix-free counterpart of the positive-only (`use_positive_only_solver=True`) solve + of an inversion's `curvature_reg_matrix`. The method follows @HRSAstro's reference + (PyAutoArray #638, phase 4) in two stages. + + **Primal-dual interior point (Mehrotra predictor-corrector).** The KKT conditions are + `A x - D = z`, `x >= 0`, `z >= 0`, `x z = 0`. Starting from one Jacobi step, + `x₀ = max(D / diag, 0)` shifted by `10⁻² max(x₀)` into the interior (with `diag` the diagonal of + `P`, the Powell approximation of the diagonal of `A`), and a dual `z₀` from the positive part of + the gradient, each step eliminates `dz` from the linearised KKT system, leaving + + (A + Z X⁻¹) dx = -r_d - z + c, r_d = A x - D - z, + + with `c = 0` for the affine (predictor) direction and `c = (σμ - dx_aff dz_aff) / x` for the + combined direction, centring parameter `σ = (μ_aff / μ)³` and `μ = xᵀz / n`. Each of the two + solves is CG preconditioned by the barrier-diagonal-augmented Powell matrix `P + Z X⁻¹`, refactorised + by a sparse LU every step (the barrier diagonal changes); the corrector is warm-started at the + predictor. The CG tolerance is tied to the progress, `0.1 min(gap, residual + gap)`, clipped to + `[cg_rtol_min, cg_rtol_max]`, so early steps are cheap. Steps are damped to 0.99 of the distance to + the boundary. Iteration stops when the normalised gap `n μ / |xᵀD|` and dual residual + `|r_d| / |D|` are both below `tol`, or after `max_steps`. + + **Exact active-set polish.** The interior point never reaches exact zeros, so its support is read + off as the free set `{i : z_i <= diag_i x_i}` and the free pixels are solved exactly, + `A_ff s_f = D_f`, by CG preconditioned by the sparse LU of the sub-matrix `P_ff` and warm-started + at the interior-point iterate; the held pixels are exactly zero. The KKT conditions are then + checked through the multiplier `g = A s - D` and the support updated by the primal-dual active-set + rule (Hintermüller, Ito & Kunisch 2002): a pixel is held when `g_i > diag_i s_i`, which releases + held pixels with a negative gradient and holds free pixels that came out negative. This repeats + until the support is unchanged (normally after one or two solves), or `max_polish` solves, after + which the solution is clipped at zero and `converged` is `False`. + + The contributor found projected Newton and a pure active-set method stall at weak priors with + ~30% of pixels at zero; the interior point is insensitive to the size of the active set and the + polish starts from an essentially correct support, so the result matches `scipy.optimize.nnls` + on the whitened dense system. + + Parameters + ---------- + operator + The data operator providing `F` and `D`. + regularization_matrix_sparse + The regularization matrix `H` as a `scipy.sparse` matrix. + coefficient + The scalar `λ` multiplying `H`; `1.0` for an autoarray regularization matrix, which already + carries its coefficient. + max_steps + The maximum number of interior-point steps. + tol + The normalised gap and dual residual at which the interior point stops. + cg_rtol_min, cg_rtol_max + The clip range of the interior-point CG tolerance. + polish_rtol + The CG tolerance of the exact free-set solves, relative to `|D_f|`. + max_polish + The maximum number of free-set solves in the active-set polish. + maxiter + The maximum number of iterations of every CG solve. + + Returns + ------- + A `PositiveReconstructionResult` holding the reconstruction (exact zeros on the held set), the + support mask, the interior-point step and polish iteration counts and the final gap and residual. + """ + regularization_matrix_sparse = scipy.sparse.csr_matrix(regularization_matrix_sparse) + + n = operator.params + data_vector = operator.data_vector + data_vector_norm = float(np.linalg.norm(data_vector)) + + def curvature_reg_matvec(v): + return operator.curvature_matvec(v) + coefficient * ( + regularization_matrix_sparse @ v + ) + + preconditioner, _ = powell_preconditioner_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix_sparse, + coefficient=coefficient, + ) + preconditioner = preconditioner.tocsc() + diag = preconditioner.diagonal() + + def splu_from(matrix): + return scipy.sparse.linalg.splu( + matrix.tocsc(), + permc_spec="MMD_AT_PLUS_A", + diag_pivot_thresh=0.0, + options=dict(SymmetricMode=True), + ) + + # Starting point: one Jacobi step, shifted into the interior, and a positive dual. + x = np.clip(data_vector / diag, 0.0, None) + x = x + 1.0e-2 * max(float(np.max(x)), np.finfo(float).tiny) + + gradient = curvature_reg_matvec(x) - data_vector + z = np.maximum(gradient, 0.0) + 1.0e-2 * max( + float(np.max(np.abs(gradient))), np.finfo(float).tiny + ) + + interior_point_steps = 0 + gap = residual = np.inf + + for _ in range(max_steps): + dual_residual = gradient - z + mu = float(x @ z) / n + + gap = n * mu / abs(float(x @ data_vector)) + residual = float(np.linalg.norm(dual_residual)) / data_vector_norm + + if gap < tol and residual < tol: + break + + interior_point_steps += 1 + + barrier_diagonal = z / x + + factor = splu_from(preconditioner + scipy.sparse.diags(barrier_diagonal)) + + newton_operator = scipy.sparse.linalg.LinearOperator( + (n, n), + matvec=lambda v: curvature_reg_matvec(v) + barrier_diagonal * v, + dtype=np.float64, + ) + preconditioner_operator = scipy.sparse.linalg.LinearOperator( + (n, n), matvec=factor.solve, dtype=np.float64 + ) + + rtol = float(np.clip(0.1 * min(gap, residual + gap), cg_rtol_min, cg_rtol_max)) + + # Predictor (affine scaling) direction. + dx, _ = scipy.sparse.linalg.cg( + newton_operator, + -dual_residual - z, + rtol=rtol, + maxiter=maxiter, + M=preconditioner_operator, + ) + dz = -z - barrier_diagonal * dx + + mu_affine = ( + float( + (x + _step_to_boundary(x, dx) * dx) + @ (z + _step_to_boundary(z, dz) * dz) + ) + / n + ) + sigma = (mu_affine / mu) ** 3 + + # Combined centring + corrector direction, warm-started at the predictor. + correction = (sigma * mu - dx * dz) / x + + dx, _ = scipy.sparse.linalg.cg( + newton_operator, + -dual_residual - z + correction, + x0=dx, + rtol=rtol, + maxiter=maxiter, + M=preconditioner_operator, + ) + dz = -z + correction - barrier_diagonal * dx + + x = x + 0.99 * _step_to_boundary(x, dx) * dx + z = z + 0.99 * _step_to_boundary(z, dz) * dz + + gradient = curvature_reg_matvec(x) - data_vector + + # Exact active-set polish on the support detected from the interior point. + free = z <= diag * x + warm_start = x.copy() + + reconstruction = np.zeros(n) + polish_iterations = 0 + converged = False + + for _ in range(max_polish): + polish_iterations += 1 + + free_indexes = np.flatnonzero(free) + reconstruction = np.zeros(n) + + if free_indexes.size > 0: + + def free_matvec(v, free_indexes=free_indexes): + return curvature_reg_matvec( + np.bincount(free_indexes, weights=v, minlength=n) + )[free_indexes] + + sub_factor = splu_from(preconditioner[free_indexes][:, free_indexes]) + + free_operator = scipy.sparse.linalg.LinearOperator( + (free_indexes.size, free_indexes.size), + matvec=free_matvec, + dtype=np.float64, + ) + sub_preconditioner_operator = scipy.sparse.linalg.LinearOperator( + (free_indexes.size, free_indexes.size), + matvec=sub_factor.solve, + dtype=np.float64, + ) + + reconstruction_free, _ = scipy.sparse.linalg.cg( + free_operator, + data_vector[free_indexes], + x0=warm_start[free_indexes], + rtol=polish_rtol, + maxiter=maxiter, + M=sub_preconditioner_operator, + ) + + reconstruction[free_indexes] = reconstruction_free + + # Primal-dual active-set update: the multiplier of the held pixels is the gradient. + multiplier = curvature_reg_matvec(reconstruction) - data_vector + multiplier[free] = 0.0 + + free_new = ~(multiplier > diag * reconstruction) + + if np.array_equal(free_new, free): + converged = bool(np.all(reconstruction >= 0.0)) + break + + free = free_new + warm_start = np.clip(reconstruction, 0.0, None) + + if not converged: + reconstruction = np.clip(reconstruction, 0.0, None) + free = reconstruction > 0.0 + + return PositiveReconstructionResult( + reconstruction=reconstruction, + support=free, + interior_point_steps=interior_point_steps, + polish_iterations=polish_iterations, + gap=float(gap), + residual=float(residual), + converged=converged, + ) From 850ea011048109243ffa82913d0a77f3198b19a7 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 21 Oct 2026 14:05:31 +0100 Subject: [PATCH 2/2] test: interior-point positive solve parity with whitened NNLS and the sparse positive-only path Co-Authored-By: Claude Fable 5.1 --- .../interferometer/test_matrix_free_util.py | 220 +++++++++++++++++- 1 file changed, 215 insertions(+), 5 deletions(-) diff --git a/test_autoarray/inversion/inversion/interferometer/test_matrix_free_util.py b/test_autoarray/inversion/inversion/interferometer/test_matrix_free_util.py index b39a36791..1b41d7a4b 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_matrix_free_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_matrix_free_util.py @@ -1,10 +1,11 @@ """ -Parity tests for the matrix-free interferometer utilities (phases 1-3 of PyAutoArray #638). +Parity tests for the matrix-free interferometer utilities (phases 1-4 of PyAutoArray #638). The matrix-free sparse mapping matrix, curvature matrix action, data vector and PCG solve are held against the dense quantities of `InversionInterferometerSparse` on a small synthetic dataset, and (phase 3) the exact preconditioner log-determinant, the stochastic Lanczos quadrature estimate of -`ln det(F + H)` and the assembled log evidence against their dense counterparts. +`ln det(F + H)` and the assembled log evidence against their dense counterparts, and (phase 4) the +interior-point positive solve against `scipy.optimize.nnls` and the sparse path's positive-only solve. These tests deliberately carry no `importorskip("jax")`: the module is NumPy / scipy only, so the `unittest-nojax` CI leg runs them. @@ -315,7 +316,7 @@ def test__powell_preconditioner__equals_kappa0_mtm_plus_coefficient_regularizati ) -def _matern_sparse_system_from(coefficient=1.0): +def _matern_sparse_system_from(coefficient=1.0, mesh_shape=(6, 6)): """ The phase-3 case: a `MaternSparse`-regularized rectangular mesh on the data-dominated dataset (500 visibilities on long baselines), returning the dense-path inversion and the matrix-free @@ -325,7 +326,9 @@ def _matern_sparse_system_from(coefficient=1.0): regularization = aa.reg.MaternSparse(coefficient=coefficient, scale=1.0) - mapper = _rectangular_mapper_from(mask=mask, regularization=regularization) + mapper = _rectangular_mapper_from( + mask=mask, regularization=regularization, shape=mesh_shape + ) inversion = _inversion_from(dataset_sparse=dataset_sparse, mapper=mapper) operator = mf.MatrixFreeDataOperator.from_inversion( @@ -544,11 +547,211 @@ def test__log_evidence_matrix_free_from__matches_dense_sparse_path_evidence( ) +# Phase 4: the interior-point positive solve. +# +# The case is the phase-3 data-dominated dataset (500 visibilities of pure noise on long baselines) with +# a 10 x 10 `MaternSparse` rectangular mesh, so roughly half the unconstrained reconstruction is +# negative. Measured on this case (zero fraction, IP steps, polish solves; all with 0 support mismatch +# against NNLS and a value error ~1e-10 of the max, set by the polish CG tolerance): +# +# coefficient 0.01 (weak prior): 72% of pixels at zero, 8 interior-point steps, 1 polish solve +# coefficient 1e4 (strong prior): 33% of pixels at zero, 10 interior-point steps, 1 polish solve +# +# The sparse path's positive-only reconstruction is a NumPy fnnls on the dense `curvature_reg_matrix` +# (edge-zeroing off so both solve the same problem), itself within ~2e-15 of NNLS, so the matrix-free +# solve agrees with it at the same ~1e-10 level; the asserted tolerance is 1e-8 of the max. + +_POSITIVE_COEFFICIENTS = [(0.01, 0.6), (1.0e4, 0.25)] + + +def _nnls_whitened_dense_from(curvature_reg_matrix, data_vector): + """ + `scipy.optimize.nnls` on the Cholesky-whitened dense system: with `A = L Lᵀ`, + `|Lᵀs - L⁻¹D|² = sᵀAs - 2sᵀD + const`, so its minimiser over `s >= 0` is the positive solution of + `A s = D`. + """ + from scipy.optimize import nnls + + cholesky = np.linalg.cholesky(curvature_reg_matrix) + + reconstruction, _ = nnls( + cholesky.T, + np.linalg.solve(cholesky, data_vector), + maxiter=50 * curvature_reg_matrix.shape[0], + ) + + return reconstruction + + +@pytest.mark.parametrize("coefficient, min_zero_fraction", _POSITIVE_COEFFICIENTS) +def test__reconstruction_positive_via_interior_point__matches_whitened_nnls( + coefficient, min_zero_fraction +): + _, _, _, inversion, operator, regularization_matrix, _ = _matern_sparse_system_from( + coefficient=coefficient, mesh_shape=(10, 10) + ) + + reconstruction_nnls = _nnls_whitened_dense_from( + curvature_reg_matrix=np.asarray(inversion.curvature_reg_matrix), + data_vector=np.asarray(inversion.data_vector), + ) + + result = mf.reconstruction_positive_via_interior_point_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + coefficient=1.0, + ) + + reconstruction = result.reconstruction + + # A meaningful fraction of the pixels sit at zero (72% weak, 33% strong). + assert np.mean(reconstruction_nnls == 0.0) >= min_zero_fraction + + # (a) Exact zeros on exactly the NNLS support. + assert np.array_equal(reconstruction == 0.0, reconstruction_nnls == 0.0) + assert np.array_equal(result.support, reconstruction_nnls > 0.0) + assert np.all(reconstruction >= 0.0) + + # (b) Values within 1e-6 of the max (measured ~1e-10). + assert np.max(np.abs(reconstruction - reconstruction_nnls)) <= 1.0e-6 * np.max( + reconstruction_nnls + ) + + # (d) Bounded iteration counts (measured 8-10 interior-point steps, 1 polish solve). + assert result.converged + assert 0 < result.interior_point_steps <= 60 + assert result.polish_iterations <= 5 + assert result.gap < 1.0e-9 + assert result.residual < 1.0e-9 + + +@pytest.mark.parametrize("coefficient, min_zero_fraction", _POSITIVE_COEFFICIENTS) +def test__reconstruction_positive_via_interior_point__matches_sparse_path_positive_only( + coefficient, min_zero_fraction +): + """ + (c) The sparse path's own positive-only solve on the same `MaternSparse` inversion. The measured + difference is ~1e-10 of the max (the polish CG tolerance); the asserted tolerance is 1e-8. + """ + dataset_sparse, _, mapper, _, operator, regularization_matrix, _ = ( + _matern_sparse_system_from(coefficient=coefficient, mesh_shape=(10, 10)) + ) + + inversion_positive = aa.Inversion( + dataset=dataset_sparse, + linear_obj_list=[mapper], + settings=aa.Settings( + use_positive_only_solver=True, use_edge_zeroed_pixels=False + ), + ) + + assert isinstance(inversion_positive, aa.InversionInterferometerSparse) + + reconstruction_sparse_path = np.asarray(inversion_positive.reconstruction) + + result = mf.reconstruction_positive_via_interior_point_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + coefficient=1.0, + ) + + assert np.mean(reconstruction_sparse_path == 0.0) >= min_zero_fraction + assert np.array_equal( + result.reconstruction == 0.0, reconstruction_sparse_path == 0.0 + ) + assert np.max( + np.abs(result.reconstruction - reconstruction_sparse_path) + ) <= 1.0e-8 * np.max(reconstruction_sparse_path) + + +def test__reconstruction_positive_via_interior_point__truncated_interior_point__active_set_polish_recovers_nnls(): + """ + With the interior point stopped after two steps its support is wrong, so the polish must make + primal-dual active-set corrections (measured 3 solves) to reach the NNLS support and values. + """ + _, _, _, inversion, operator, regularization_matrix, _ = _matern_sparse_system_from( + coefficient=0.01, mesh_shape=(10, 10) + ) + + reconstruction_nnls = _nnls_whitened_dense_from( + curvature_reg_matrix=np.asarray(inversion.curvature_reg_matrix), + data_vector=np.asarray(inversion.data_vector), + ) + + result = mf.reconstruction_positive_via_interior_point_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + coefficient=1.0, + max_steps=2, + ) + + assert result.interior_point_steps == 2 + assert 1 < result.polish_iterations <= 10 + assert result.converged + assert np.array_equal(result.reconstruction == 0.0, reconstruction_nnls == 0.0) + assert np.max( + np.abs(result.reconstruction - reconstruction_nnls) + ) <= 1.0e-6 * np.max(reconstruction_nnls) + + +def test__reconstruction_positive_via_interior_point__all_positive_solution__equals_unconstrained_pcg(): + """ + (e) When the unconstrained solution is already positive the constraint is inactive, so the positive + solve must reduce to the unconstrained PCG solve. The dirty image is replaced by the noise-free + `d~ = W~ M s_true` of a strictly positive source, so `D = F s_true` and, with a weak prior, the + unconstrained solution `(F + H)⁻¹ F s_true` stays positive. + """ + _, _, _, _, operator, regularization_matrix, _ = _matern_sparse_system_from( + coefficient=0.01, mesh_shape=(10, 10) + ) + + source_true = 1.0 + 0.5 * np.random.default_rng(seed=5).random(operator.params) + + image_extent = np.zeros(operator.y_shape * operator.x_shape) + image_extent[operator.extent_index_for_masked_pixel] = ( + operator.mapping_matrix @ source_true + ) + + dirty_image = operator.wtilde_image_from( + image_extent.reshape(operator.y_shape, operator.x_shape) + ).ravel()[operator.extent_index_for_masked_pixel] + + operator_positive = mf.MatrixFreeDataOperator( + nufft_precision_operator=operator.nufft_precision_operator, + dirty_image=dirty_image, + mapping_matrix=operator.mapping_matrix, + mask=operator.mask, + ) + + reconstruction_pcg, info = mf.reconstruction_via_pcg_from( + operator=operator_positive, + regularization_matrix_sparse=regularization_matrix, + coefficient=1.0, + rtol=1.0e-12, + ) + + assert info == 0 + assert np.min(reconstruction_pcg) > 0.0 + + result = mf.reconstruction_positive_via_interior_point_from( + operator=operator_positive, + regularization_matrix_sparse=regularization_matrix, + coefficient=1.0, + ) + + assert np.all(result.support) + assert result.polish_iterations == 1 + assert np.max( + np.abs(result.reconstruction - reconstruction_pcg) + ) <= 1.0e-8 * np.max(np.abs(reconstruction_pcg)) + + def test__matrix_free_util__imports_and_runs_with_jax_blocked(): """ The module must work in the `unittest-nojax` leg. JAX is in `sys.modules` for the rest of a normal test session, so a subprocess hides `jax`, `jaxlib` and `nufftax` from the import system, - imports the module and runs a PCG solve and the phase-3 log-determinant / evidence functions. + imports the module and runs a PCG solve, the phase-3 log-determinant / evidence functions and the + phase-4 interior-point positive solve. """ script = textwrap.dedent( """ @@ -656,6 +859,13 @@ def _find_spec_blocking_jax(cls, name, path=None, target=None): ) assert np.isfinite(log_evidence) + result = mf.reconstruction_positive_via_interior_point_from( + operator, regularization_matrix, 1.0 + ) + assert result.converged + assert np.all(result.reconstruction >= 0.0) + assert np.all(result.reconstruction[~result.support] == 0.0) + assert not any( name.split(".")[0] in ("jax", "jaxlib", "nufftax") for name in sys.modules )