Skip to content

Commit 850ea01

Browse files
Jammy2211claude
andcommitted
test: interior-point positive solve parity with whitened NNLS and the sparse positive-only path
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1 parent 74d6872 commit 850ea01

1 file changed

Lines changed: 215 additions & 5 deletions

File tree

‎test_autoarray/inversion/inversion/interferometer/test_matrix_free_util.py‎

Lines changed: 215 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
"""
2-
Parity tests for the matrix-free interferometer utilities (phases 1-3 of PyAutoArray #638).
2+
Parity tests for the matrix-free interferometer utilities (phases 1-4 of PyAutoArray #638).
33
44
The matrix-free sparse mapping matrix, curvature matrix action, data vector and PCG solve are held
55
against the dense quantities of `InversionInterferometerSparse` on a small synthetic dataset, and
66
(phase 3) the exact preconditioner log-determinant, the stochastic Lanczos quadrature estimate of
7-
`ln det(F + H)` and the assembled log evidence against their dense counterparts.
7+
`ln det(F + H)` and the assembled log evidence against their dense counterparts, and (phase 4) the
8+
interior-point positive solve against `scipy.optimize.nnls` and the sparse path's positive-only solve.
89
910
These tests deliberately carry no `importorskip("jax")`: the module is NumPy / scipy only, so the
1011
`unittest-nojax` CI leg runs them.
@@ -315,7 +316,7 @@ def test__powell_preconditioner__equals_kappa0_mtm_plus_coefficient_regularizati
315316
)
316317

317318

318-
def _matern_sparse_system_from(coefficient=1.0):
319+
def _matern_sparse_system_from(coefficient=1.0, mesh_shape=(6, 6)):
319320
"""
320321
The phase-3 case: a `MaternSparse`-regularized rectangular mesh on the data-dominated dataset
321322
(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):
325326

326327
regularization = aa.reg.MaternSparse(coefficient=coefficient, scale=1.0)
327328

328-
mapper = _rectangular_mapper_from(mask=mask, regularization=regularization)
329+
mapper = _rectangular_mapper_from(
330+
mask=mask, regularization=regularization, shape=mesh_shape
331+
)
329332
inversion = _inversion_from(dataset_sparse=dataset_sparse, mapper=mapper)
330333

331334
operator = mf.MatrixFreeDataOperator.from_inversion(
@@ -544,11 +547,211 @@ def test__log_evidence_matrix_free_from__matches_dense_sparse_path_evidence(
544547
)
545548

546549

550+
# Phase 4: the interior-point positive solve.
551+
#
552+
# The case is the phase-3 data-dominated dataset (500 visibilities of pure noise on long baselines) with
553+
# a 10 x 10 `MaternSparse` rectangular mesh, so roughly half the unconstrained reconstruction is
554+
# negative. Measured on this case (zero fraction, IP steps, polish solves; all with 0 support mismatch
555+
# against NNLS and a value error ~1e-10 of the max, set by the polish CG tolerance):
556+
#
557+
# coefficient 0.01 (weak prior): 72% of pixels at zero, 8 interior-point steps, 1 polish solve
558+
# coefficient 1e4 (strong prior): 33% of pixels at zero, 10 interior-point steps, 1 polish solve
559+
#
560+
# The sparse path's positive-only reconstruction is a NumPy fnnls on the dense `curvature_reg_matrix`
561+
# (edge-zeroing off so both solve the same problem), itself within ~2e-15 of NNLS, so the matrix-free
562+
# solve agrees with it at the same ~1e-10 level; the asserted tolerance is 1e-8 of the max.
563+
564+
_POSITIVE_COEFFICIENTS = [(0.01, 0.6), (1.0e4, 0.25)]
565+
566+
567+
def _nnls_whitened_dense_from(curvature_reg_matrix, data_vector):
568+
"""
569+
`scipy.optimize.nnls` on the Cholesky-whitened dense system: with `A = L Lᵀ`,
570+
`|Lᵀs - L⁻¹D|² = sᵀAs - 2sᵀD + const`, so its minimiser over `s >= 0` is the positive solution of
571+
`A s = D`.
572+
"""
573+
from scipy.optimize import nnls
574+
575+
cholesky = np.linalg.cholesky(curvature_reg_matrix)
576+
577+
reconstruction, _ = nnls(
578+
cholesky.T,
579+
np.linalg.solve(cholesky, data_vector),
580+
maxiter=50 * curvature_reg_matrix.shape[0],
581+
)
582+
583+
return reconstruction
584+
585+
586+
@pytest.mark.parametrize("coefficient, min_zero_fraction", _POSITIVE_COEFFICIENTS)
587+
def test__reconstruction_positive_via_interior_point__matches_whitened_nnls(
588+
coefficient, min_zero_fraction
589+
):
590+
_, _, _, inversion, operator, regularization_matrix, _ = _matern_sparse_system_from(
591+
coefficient=coefficient, mesh_shape=(10, 10)
592+
)
593+
594+
reconstruction_nnls = _nnls_whitened_dense_from(
595+
curvature_reg_matrix=np.asarray(inversion.curvature_reg_matrix),
596+
data_vector=np.asarray(inversion.data_vector),
597+
)
598+
599+
result = mf.reconstruction_positive_via_interior_point_from(
600+
operator=operator,
601+
regularization_matrix_sparse=regularization_matrix,
602+
coefficient=1.0,
603+
)
604+
605+
reconstruction = result.reconstruction
606+
607+
# A meaningful fraction of the pixels sit at zero (72% weak, 33% strong).
608+
assert np.mean(reconstruction_nnls == 0.0) >= min_zero_fraction
609+
610+
# (a) Exact zeros on exactly the NNLS support.
611+
assert np.array_equal(reconstruction == 0.0, reconstruction_nnls == 0.0)
612+
assert np.array_equal(result.support, reconstruction_nnls > 0.0)
613+
assert np.all(reconstruction >= 0.0)
614+
615+
# (b) Values within 1e-6 of the max (measured ~1e-10).
616+
assert np.max(np.abs(reconstruction - reconstruction_nnls)) <= 1.0e-6 * np.max(
617+
reconstruction_nnls
618+
)
619+
620+
# (d) Bounded iteration counts (measured 8-10 interior-point steps, 1 polish solve).
621+
assert result.converged
622+
assert 0 < result.interior_point_steps <= 60
623+
assert result.polish_iterations <= 5
624+
assert result.gap < 1.0e-9
625+
assert result.residual < 1.0e-9
626+
627+
628+
@pytest.mark.parametrize("coefficient, min_zero_fraction", _POSITIVE_COEFFICIENTS)
629+
def test__reconstruction_positive_via_interior_point__matches_sparse_path_positive_only(
630+
coefficient, min_zero_fraction
631+
):
632+
"""
633+
(c) The sparse path's own positive-only solve on the same `MaternSparse` inversion. The measured
634+
difference is ~1e-10 of the max (the polish CG tolerance); the asserted tolerance is 1e-8.
635+
"""
636+
dataset_sparse, _, mapper, _, operator, regularization_matrix, _ = (
637+
_matern_sparse_system_from(coefficient=coefficient, mesh_shape=(10, 10))
638+
)
639+
640+
inversion_positive = aa.Inversion(
641+
dataset=dataset_sparse,
642+
linear_obj_list=[mapper],
643+
settings=aa.Settings(
644+
use_positive_only_solver=True, use_edge_zeroed_pixels=False
645+
),
646+
)
647+
648+
assert isinstance(inversion_positive, aa.InversionInterferometerSparse)
649+
650+
reconstruction_sparse_path = np.asarray(inversion_positive.reconstruction)
651+
652+
result = mf.reconstruction_positive_via_interior_point_from(
653+
operator=operator,
654+
regularization_matrix_sparse=regularization_matrix,
655+
coefficient=1.0,
656+
)
657+
658+
assert np.mean(reconstruction_sparse_path == 0.0) >= min_zero_fraction
659+
assert np.array_equal(
660+
result.reconstruction == 0.0, reconstruction_sparse_path == 0.0
661+
)
662+
assert np.max(
663+
np.abs(result.reconstruction - reconstruction_sparse_path)
664+
) <= 1.0e-8 * np.max(reconstruction_sparse_path)
665+
666+
667+
def test__reconstruction_positive_via_interior_point__truncated_interior_point__active_set_polish_recovers_nnls():
668+
"""
669+
With the interior point stopped after two steps its support is wrong, so the polish must make
670+
primal-dual active-set corrections (measured 3 solves) to reach the NNLS support and values.
671+
"""
672+
_, _, _, inversion, operator, regularization_matrix, _ = _matern_sparse_system_from(
673+
coefficient=0.01, mesh_shape=(10, 10)
674+
)
675+
676+
reconstruction_nnls = _nnls_whitened_dense_from(
677+
curvature_reg_matrix=np.asarray(inversion.curvature_reg_matrix),
678+
data_vector=np.asarray(inversion.data_vector),
679+
)
680+
681+
result = mf.reconstruction_positive_via_interior_point_from(
682+
operator=operator,
683+
regularization_matrix_sparse=regularization_matrix,
684+
coefficient=1.0,
685+
max_steps=2,
686+
)
687+
688+
assert result.interior_point_steps == 2
689+
assert 1 < result.polish_iterations <= 10
690+
assert result.converged
691+
assert np.array_equal(result.reconstruction == 0.0, reconstruction_nnls == 0.0)
692+
assert np.max(
693+
np.abs(result.reconstruction - reconstruction_nnls)
694+
) <= 1.0e-6 * np.max(reconstruction_nnls)
695+
696+
697+
def test__reconstruction_positive_via_interior_point__all_positive_solution__equals_unconstrained_pcg():
698+
"""
699+
(e) When the unconstrained solution is already positive the constraint is inactive, so the positive
700+
solve must reduce to the unconstrained PCG solve. The dirty image is replaced by the noise-free
701+
`d~ = W~ M s_true` of a strictly positive source, so `D = F s_true` and, with a weak prior, the
702+
unconstrained solution `(F + H)⁻¹ F s_true` stays positive.
703+
"""
704+
_, _, _, _, operator, regularization_matrix, _ = _matern_sparse_system_from(
705+
coefficient=0.01, mesh_shape=(10, 10)
706+
)
707+
708+
source_true = 1.0 + 0.5 * np.random.default_rng(seed=5).random(operator.params)
709+
710+
image_extent = np.zeros(operator.y_shape * operator.x_shape)
711+
image_extent[operator.extent_index_for_masked_pixel] = (
712+
operator.mapping_matrix @ source_true
713+
)
714+
715+
dirty_image = operator.wtilde_image_from(
716+
image_extent.reshape(operator.y_shape, operator.x_shape)
717+
).ravel()[operator.extent_index_for_masked_pixel]
718+
719+
operator_positive = mf.MatrixFreeDataOperator(
720+
nufft_precision_operator=operator.nufft_precision_operator,
721+
dirty_image=dirty_image,
722+
mapping_matrix=operator.mapping_matrix,
723+
mask=operator.mask,
724+
)
725+
726+
reconstruction_pcg, info = mf.reconstruction_via_pcg_from(
727+
operator=operator_positive,
728+
regularization_matrix_sparse=regularization_matrix,
729+
coefficient=1.0,
730+
rtol=1.0e-12,
731+
)
732+
733+
assert info == 0
734+
assert np.min(reconstruction_pcg) > 0.0
735+
736+
result = mf.reconstruction_positive_via_interior_point_from(
737+
operator=operator_positive,
738+
regularization_matrix_sparse=regularization_matrix,
739+
coefficient=1.0,
740+
)
741+
742+
assert np.all(result.support)
743+
assert result.polish_iterations == 1
744+
assert np.max(
745+
np.abs(result.reconstruction - reconstruction_pcg)
746+
) <= 1.0e-8 * np.max(np.abs(reconstruction_pcg))
747+
748+
547749
def test__matrix_free_util__imports_and_runs_with_jax_blocked():
548750
"""
549751
The module must work in the `unittest-nojax` leg. JAX is in `sys.modules` for the rest of a
550752
normal test session, so a subprocess hides `jax`, `jaxlib` and `nufftax` from the import system,
551-
imports the module and runs a PCG solve and the phase-3 log-determinant / evidence functions.
753+
imports the module and runs a PCG solve, the phase-3 log-determinant / evidence functions and the
754+
phase-4 interior-point positive solve.
552755
"""
553756
script = textwrap.dedent(
554757
"""
@@ -656,6 +859,13 @@ def _find_spec_blocking_jax(cls, name, path=None, target=None):
656859
)
657860
assert np.isfinite(log_evidence)
658861
862+
result = mf.reconstruction_positive_via_interior_point_from(
863+
operator, regularization_matrix, 1.0
864+
)
865+
assert result.converged
866+
assert np.all(result.reconstruction >= 0.0)
867+
assert np.all(result.reconstruction[~result.support] == 0.0)
868+
659869
assert not any(
660870
name.split(".")[0] in ("jax", "jaxlib", "nufftax") for name in sys.modules
661871
)

0 commit comments

Comments
 (0)