Skip to content

Commit 5d6bbb9

Browse files
Jammy2211claude
andcommitted
test: matrix-free log-det and SLQ evidence parity with the dense sparse path
ln det P vs slogdet, SLQ within 0.5% of dense slogdet(F + H) and deterministic, error shrinking from 2 to 10 probes, factor/matrix mismatch guard, evidence parity with fit_util.log_evidence_from of the sparse inversion, and the nojax subprocess leg extended to the phase-3 functions. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1 parent 9feb4f3 commit 5d6bbb9

1 file changed

Lines changed: 270 additions & 5 deletions

File tree

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

Lines changed: 270 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
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
44
The 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
79
These tests deliberately carry no `importorskip("jax")`: the module is NumPy / scipy only, so the
810
`unittest-nojax` CI leg runs them.
@@ -26,10 +28,16 @@
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+
310547
def 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

Comments
 (0)