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
44The matrix-free sparse mapping matrix, curvature matrix action, data vector and PCG solve are held
55against 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
910These 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+
547749def 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