11"""
2- Parity tests for the matrix-free interferometer utilities (phase 1 of PyAutoArray #638).
2+ Parity tests for the matrix-free interferometer utilities (phases 1-3 of PyAutoArray #638).
33
44The matrix-free sparse mapping matrix, curvature matrix action, data vector and PCG solve are held
5- against the dense quantities of `InversionInterferometerSparse` on a small synthetic dataset.
5+ against the dense quantities of `InversionInterferometerSparse` on a small synthetic dataset, and
6+ (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.
68
79These tests deliberately carry no `importorskip("jax")`: the module is NumPy / scipy only, so the
810`unittest-nojax` CI leg runs them.
2628)
2729
2830
29- def _dataset_sparse_from (seed = 1 , n_visibilities = 40 ):
31+ def _dataset_sparse_from (seed = 1 , n_visibilities = 40 , uv_scale = 1.0e4 ):
3032 """
3133 A circular-masked dataset, so the slim grid is a strict subset of its unmasked extent and the
3234 scatter / gather between them is exercised, with random visibilities of unit noise.
35+
36+ `uv_scale` sets the baseline lengths in wavelengths. At the default `1e4` the fringes are far
37+ broader than the mask, so `W~` is close to a constant and nowhere near the Powell preconditioner's
38+ `κ₀ I`; with many visibilities on long baselines (`500`, `1e6`) the dirty beam is compact, as for
39+ real data, and `P⁻¹(F + H)` has eigenvalues near one, the regime the stochastic log-determinant
40+ is built for.
3341 """
3442 mask = aa .Mask2D .circular (shape_native = (12 , 12 ), pixel_scales = 0.5 , radius = 2.5 )
3543
@@ -38,7 +46,7 @@ def _dataset_sparse_from(seed=1, n_visibilities=40):
3846 dataset = aa .Interferometer (
3947 data = aa .Visibilities (visibilities = rng .normal (size = (n_visibilities , 2 ))),
4048 noise_map = aa .VisibilitiesNoiseMap (visibilities = np .ones ((n_visibilities , 2 ))),
41- uv_wavelengths = rng .normal (size = (n_visibilities , 2 )) * 1.0e4 ,
49+ uv_wavelengths = rng .normal (size = (n_visibilities , 2 )) * uv_scale ,
4250 real_space_mask = mask ,
4351 transformer_class = aa .TransformerDFT ,
4452 )
@@ -307,11 +315,240 @@ def test__powell_preconditioner__equals_kappa0_mtm_plus_coefficient_regularizati
307315 )
308316
309317
318+ def _matern_sparse_system_from (coefficient = 1.0 ):
319+ """
320+ The phase-3 case: a `MaternSparse`-regularized rectangular mesh on the data-dominated dataset
321+ (500 visibilities on long baselines), returning the dense-path inversion and the matrix-free
322+ operator, `H` and its factor `R` with `H = RᵀR` (the coefficient folded into `R`).
323+ """
324+ mask , dataset_sparse = _dataset_sparse_from (n_visibilities = 500 , uv_scale = 1.0e6 )
325+
326+ regularization = aa .reg .MaternSparse (coefficient = coefficient , scale = 1.0 )
327+
328+ mapper = _rectangular_mapper_from (mask = mask , regularization = regularization )
329+ inversion = _inversion_from (dataset_sparse = dataset_sparse , mapper = mapper )
330+
331+ operator = mf .MatrixFreeDataOperator .from_inversion (
332+ inversion = inversion , mapper = mapper
333+ )
334+
335+ regularization_matrix = regularization .regularization_matrix_sparse_from (mapper )
336+ regularization_factor = np .sqrt (
337+ regularization .coefficient
338+ ) * regularization .regularization_factor_sparse_from (mapper )
339+
340+ return (
341+ dataset_sparse ,
342+ regularization ,
343+ mapper ,
344+ inversion ,
345+ operator ,
346+ regularization_matrix ,
347+ regularization_factor ,
348+ )
349+
350+
351+ @pytest .mark .parametrize ("coefficient" , [1.0 , 0.01 ])
352+ def test__log_det_preconditioner_from__equals_dense_slogdet (coefficient ):
353+ _ , _ , _ , _ , operator , regularization_matrix , _ = _matern_sparse_system_from (
354+ coefficient = coefficient
355+ )
356+
357+ preconditioner , factor = mf .powell_preconditioner_from (
358+ operator = operator ,
359+ regularization_matrix_sparse = regularization_matrix ,
360+ coefficient = 1.0 ,
361+ )
362+
363+ sign , log_det = np .linalg .slogdet (preconditioner .toarray ())
364+
365+ assert sign == 1.0
366+ assert mf .log_det_preconditioner_from (factor ) == pytest .approx (log_det , abs = 1.0e-10 )
367+
368+
369+ @pytest .mark .parametrize ("coefficient" , [1.0 , 0.01 ])
370+ def test__log_det_curvature_reg_via_slq__within_half_percent_of_dense_and_deterministic (
371+ coefficient ,
372+ ):
373+ (
374+ _ ,
375+ _ ,
376+ _ ,
377+ inversion ,
378+ operator ,
379+ regularization_matrix ,
380+ regularization_factor ,
381+ ) = _matern_sparse_system_from (coefficient = coefficient )
382+
383+ sign , log_det_dense = np .linalg .slogdet (np .asarray (inversion .curvature_reg_matrix ))
384+
385+ assert sign == 1.0
386+
387+ estimate , log_det_probes = mf .log_det_curvature_reg_via_slq_from (
388+ operator = operator ,
389+ regularization_matrix_sparse = regularization_matrix ,
390+ regularization_factor_sparse = regularization_factor ,
391+ coefficient = 1.0 ,
392+ n_probes = 10 ,
393+ n_steps = 40 ,
394+ )
395+
396+ assert log_det_probes .shape == (10 ,)
397+ assert estimate == pytest .approx (np .mean (log_det_probes ), 1.0e-14 )
398+ assert abs (estimate - log_det_dense ) <= 0.005 * abs (log_det_dense )
399+
400+ estimate_repeat , log_det_probes_repeat = mf .log_det_curvature_reg_via_slq_from (
401+ operator = operator ,
402+ regularization_matrix_sparse = regularization_matrix ,
403+ regularization_factor_sparse = regularization_factor ,
404+ coefficient = 1.0 ,
405+ n_probes = 10 ,
406+ n_steps = 40 ,
407+ )
408+
409+ assert estimate_repeat == estimate
410+ assert np .array_equal (log_det_probes_repeat , log_det_probes )
411+
412+
413+ def test__log_det_curvature_reg_via_slq__error_shrinks_with_probes ():
414+ """
415+ Loose convergence check: over a handful of seeds the mean error at 10 probes is below that at 2,
416+ and the Monte-Carlo standard error of the mean shrinks.
417+ """
418+ (
419+ _ ,
420+ _ ,
421+ _ ,
422+ inversion ,
423+ operator ,
424+ regularization_matrix ,
425+ regularization_factor ,
426+ ) = _matern_sparse_system_from ()
427+
428+ _ , log_det_dense = np .linalg .slogdet (np .asarray (inversion .curvature_reg_matrix ))
429+
430+ def errors_and_standard_errors (n_probes ):
431+ errors , standard_errors = [], []
432+
433+ for seed in range (5 ):
434+ estimate , log_det_probes = mf .log_det_curvature_reg_via_slq_from (
435+ operator = operator ,
436+ regularization_matrix_sparse = regularization_matrix ,
437+ regularization_factor_sparse = regularization_factor ,
438+ coefficient = 1.0 ,
439+ n_probes = n_probes ,
440+ seed = seed ,
441+ )
442+ errors .append (abs (estimate - log_det_dense ))
443+ standard_errors .append (np .std (log_det_probes ) / np .sqrt (n_probes ))
444+
445+ return np .mean (errors ), np .mean (standard_errors )
446+
447+ error_2 , standard_error_2 = errors_and_standard_errors (n_probes = 2 )
448+ error_10 , standard_error_10 = errors_and_standard_errors (n_probes = 10 )
449+
450+ assert error_10 < error_2
451+ assert standard_error_10 < standard_error_2
452+
453+
454+ def test__log_det_curvature_reg_via_slq__factor_not_matching_regularization_matrix__raises ():
455+ _ , regularization , mapper , _ , operator , regularization_matrix , _ = (
456+ _matern_sparse_system_from (coefficient = 0.01 )
457+ )
458+
459+ with pytest .raises (ValueError ):
460+ mf .log_det_curvature_reg_via_slq_from (
461+ operator = operator ,
462+ regularization_matrix_sparse = regularization_matrix ,
463+ regularization_factor_sparse = regularization .regularization_factor_sparse_from (
464+ mapper
465+ ),
466+ coefficient = 1.0 ,
467+ )
468+
469+
470+ @pytest .mark .parametrize ("coefficient" , [1.0 , 0.01 ])
471+ def test__log_evidence_matrix_free_from__matches_dense_sparse_path_evidence (
472+ coefficient ,
473+ ):
474+ """
475+ The dense reference is `FitInterferometer.log_evidence`'s assembly, `fit_util.log_evidence_from`,
476+ of the unconstrained sparse-path inversion's terms. With the dense `ln det(F + H)` substituted the
477+ other four terms agree to round-off; with the SLQ estimate the total agrees within the SLQ
478+ tolerance (0.5% of `ln det(F + H)`).
479+ """
480+ from autoarray .fit import fit_util
481+
482+ (
483+ dataset_sparse ,
484+ regularization ,
485+ mapper ,
486+ inversion ,
487+ operator ,
488+ regularization_matrix ,
489+ regularization_factor ,
490+ ) = _matern_sparse_system_from (coefficient = coefficient )
491+
492+ noise_normalization = dataset_sparse .sparse_operator .noise_normalization
493+
494+ log_evidence_dense = fit_util .log_evidence_from (
495+ chi_squared = inversion .fast_chi_squared ,
496+ regularization_term = inversion .regularization_term ,
497+ log_curvature_regularization_term = inversion .log_det_curvature_reg_matrix_term ,
498+ log_regularization_term = inversion .log_det_regularization_matrix_term ,
499+ noise_normalization = noise_normalization ,
500+ )
501+
502+ reconstruction , info = mf .reconstruction_via_pcg_from (
503+ operator = operator ,
504+ regularization_matrix_sparse = regularization_matrix ,
505+ coefficient = 1.0 ,
506+ rtol = 1.0e-12 ,
507+ )
508+
509+ assert info == 0
510+
511+ log_det_regularization = regularization .log_det_regularization_matrix_term_from (
512+ mapper
513+ )
514+
515+ assert log_det_regularization == pytest .approx (
516+ inversion .log_det_regularization_matrix_term , abs = 1.0e-10
517+ )
518+
519+ kwargs = dict (
520+ operator = operator ,
521+ regularization_matrix_sparse = regularization_matrix ,
522+ regularization_factor_sparse = regularization_factor ,
523+ coefficient = 1.0 ,
524+ reconstruction = reconstruction ,
525+ log_det_regularization = log_det_regularization ,
526+ data_term = dataset_sparse .sparse_operator .data_term ,
527+ noise_normalization = noise_normalization ,
528+ )
529+
530+ log_evidence_exact_log_det = mf .log_evidence_matrix_free_from (
531+ ** kwargs , log_det_curvature_reg = inversion .log_det_curvature_reg_matrix_term
532+ )
533+
534+ assert log_evidence_exact_log_det == pytest .approx (
535+ log_evidence_dense , abs = 1.0e-10 * abs (log_evidence_dense )
536+ )
537+
538+ log_evidence_slq = mf .log_evidence_matrix_free_from (** kwargs )
539+
540+ _ , log_det_dense = np .linalg .slogdet (np .asarray (inversion .curvature_reg_matrix ))
541+
542+ assert abs (log_evidence_slq - log_evidence_dense ) <= 0.5 * 0.005 * abs (
543+ log_det_dense
544+ )
545+
546+
310547def test__matrix_free_util__imports_and_runs_with_jax_blocked ():
311548 """
312549 The module must work in the `unittest-nojax` leg. JAX is in `sys.modules` for the rest of a
313550 normal test session, so a subprocess hides `jax`, `jaxlib` and `nufftax` from the import system,
314- imports the module and runs a PCG solve.
551+ imports the module and runs a PCG solve and the phase-3 log-determinant / evidence functions .
315552 """
316553 script = textwrap .dedent (
317554 """
@@ -391,6 +628,34 @@ def _find_spec_blocking_jax(cls, name, path=None, target=None):
391628 assert info == 0
392629 assert np.allclose(reconstruction, np.asarray(inversion.reconstruction))
393630
631+ regularization_matrix = scipy.sparse.csr_matrix(
632+ np.asarray(inversion.regularization_matrix)
633+ )
634+ regularization_factor = scipy.sparse.csr_matrix(
635+ np.linalg.cholesky(regularization_matrix.toarray()).T
636+ )
637+
638+ _, factor = mf.powell_preconditioner_from(operator, regularization_matrix, 1.0)
639+ assert np.isfinite(mf.log_det_preconditioner_from(factor))
640+
641+ log_det, _ = mf.log_det_curvature_reg_via_slq_from(
642+ operator, regularization_matrix, regularization_factor, 1.0, n_probes=2
643+ )
644+ assert np.isfinite(log_det)
645+
646+ log_evidence = mf.log_evidence_matrix_free_from(
647+ operator,
648+ regularization_matrix,
649+ regularization_factor,
650+ 1.0,
651+ reconstruction,
652+ np.linalg.slogdet(regularization_matrix.toarray())[1],
653+ 1.0,
654+ 1.0,
655+ n_probes=2,
656+ )
657+ assert np.isfinite(log_evidence)
658+
394659 assert not any(
395660 name.split(".")[0] in ("jax", "jaxlib", "nufftax") for name in sys.modules
396661 )
0 commit comments