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 990371c19..6315173e8 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_matrix_free_util.py +++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_matrix_free_util.py @@ -28,10 +28,30 @@ `P = κ₀ MᵀM + λ H`, where `κ₀` is the kernel's zero-lag value (the diagonal of `W~`). `P` replaces `W~` by its peak, which keeps it sparse; it is factorised once per `λ` by a sparse LU. -This phase lands the operator data term and the solver as library utilities, with parity tests -against the dense sparse-path quantities. It does **not** wire them into fits and adds no -`Settings` flag. Later phases (sparse Matérn precision, stochastic Lanczos log-determinant for the -evidence, interior-point positivity, profiling and workspace examples) follow as their own tasks. +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. + +The log evidence (phase 3) +-------------------------- +The Bayesian evidence also needs `ln det(F + λ H)` and `ln det(λ H)`. `ln det H` is exact from a +sparse factorisation (for `MaternSparse`, `log_det_regularization_matrix_term_from`). For the curvature +term, `ln det(F + λ H) = ln det P + ln det(P⁻¹(F + λ H))`: `ln det P` is exact from the Powell +preconditioner's sparse LU (`log_det_preconditioner_from`) and the second term is estimated by +stochastic Lanczos quadrature in the `P` inner product (`log_det_curvature_reg_via_slq_from`), with +probes `v ~ N(0, P⁻¹)` drawn from a sparse factor `R` of `H = RᵀR` and a fixed seed so the estimate is +deterministic and smooth in the model parameters. `log_evidence_matrix_free_from` assembles the five +evidence terms with the conventions of `FitInterferometer.log_evidence`. + +Tolerance: the estimate is only as good as the preconditioner. On the test case (a `6 x 6` +`MaternSparse` rectangular mesh, 500 visibilities on long baselines, so the dirty beam is compact and +`P⁻¹(F + H)` has eigenvalues in `~[0.6, 1.6]`) it lies within 0.5% of the dense `slogdet(F + H)` at +the defaults (10 probes, 40 Lanczos steps; measured ~0.1-0.15% for the default seed), and the error +falls as `1 / sqrt(n_probes)`. When `W~` is far from `κ₀ I` (few visibilities, fringes broader than +the mask) the eigenvalues spread over orders of magnitude and the per-probe variance grows by an +order of magnitude or more; the per-probe values returned give the Monte-Carlo error of the mean. The regularization matrix ------------------------- @@ -414,3 +434,300 @@ def reconstruction_via_pcg_from( maxiter=maxiter, M=preconditioner_operator, ) + + +def log_det_preconditioner_from(factor) -> float: + """ + Returns the exact `ln det P` of the Powell preconditioner from its sparse LU factorisation (the + second return value of `powell_preconditioner_from`). + + SuperLU factorises `Pr P Pc = L U` with row and column permutations `Pr`, `Pc` and a unit-diagonal + `L`, so `|det P| = prod_i |U_ii|` (a permutation has determinant `+-1`). `P` is symmetric positive + definite, so `det P > 0` and + + ln det P = sum_i ln |U_ii|. + + Parameters + ---------- + factor + The `scipy.sparse.linalg.SuperLU` factorisation of `P`. + + Returns + ------- + The log determinant `ln det P`. + """ + return float(np.sum(np.log(np.abs(factor.U.diagonal())))) + + +def log_det_curvature_reg_via_slq_from( + operator: MatrixFreeDataOperator, + regularization_matrix_sparse: scipy.sparse.spmatrix, + regularization_factor_sparse: scipy.sparse.spmatrix, + coefficient: float, + n_probes: int = 10, + n_steps: int = 40, + seed: int = 20221, +): + """ + Returns a stochastic Lanczos quadrature (SLQ, Ubaru, Chen & Saad 2017) estimate of + `ln det(F + λ H)`, the curvature-regularization log determinant of the Bayesian evidence, without + forming the dense `(S, S)` matrix `A = F + λ H`. + + With the Powell preconditioner `P = κ₀ MᵀM + λ H` (`powell_preconditioner_from`), + + ln det A = ln det P + ln det(P⁻¹ A), + + where `ln det P` is exact from its sparse LU (`log_det_preconditioner_from`) and only the second + term, whose eigenvalues cluster near one because `P` captures the dominant part of `A`, is + estimated stochastically: + + - **Probes** `v ~ N(0, P⁻¹)`. Writing `H = RᵀR` (`R` the `regularization_factor_sparse`), + `P = BᵀB` with `B = [√κ₀ M; √λ R]`, so `v = P⁻¹ Bᵀ g = P⁻¹(√κ₀ Mᵀ g₁ + √λ Rᵀ g₂)` with `g` + standard normal has covariance `P⁻¹ BᵀB P⁻¹ = P⁻¹`. Then `E[vᵀ P f(P⁻¹A) v] = tr f(P⁻¹A)` + for `f = ln`, i.e. `E = ln det(P⁻¹A)`. + - **Lanczos** on `P⁻¹A`, which is self-adjoint in the `P` inner product `_P = xᵀ P y`. + Each probe starts from `q₁ = v / ‖v‖_P`; the recurrence keeps every basis vector + `P`-orthonormal with full reorthogonalisation (the Krylov space is small, `n_steps` vectors) and + stops early once the off-diagonal falls below `1e-10` of the diagonal (an invariant subspace). + - **Gauss quadrature** on the tridiagonal `T` with eigenpairs `(θ_k, u_k)`: the probe's value is + `‖v‖²_P Σ_k u_k[0]² ln θ_k`. + + The random generator is `np.random.default_rng(seed)` with a fixed default seed, so the estimate + is deterministic and, because the same probes are reused, a smooth function of `λ` (and of any + other parameter) -- which a non-linear search over the evidence needs. + + Accuracy: on the small rectangular-mesh test cases the estimate lies within 0.5% of the dense + `slogdet(F + H)` at the defaults (10 probes, 40 steps); the per-probe spread returned gives the + Monte-Carlo error of the mean (`std / sqrt(n_probes)`). + + Parameters + ---------- + operator + The data operator providing `F`, `M` and `κ₀`. + regularization_matrix_sparse + The regularization matrix `H` as a `scipy.sparse` matrix. + regularization_factor_sparse + A sparse factor `R` of exactly that matrix, `H = RᵀR`. For `MaternSparse`, whose + `regularization_matrix_sparse_from` returns `coefficient * R₀ᵀR₀` with + `R₀ = regularization_factor_sparse_from(...)`, this is `sqrt(regularization.coefficient) * R₀`. + It is checked against `H` on one random vector. + coefficient + The scalar `λ` multiplying `H`, as in `powell_preconditioner_from`; `1.0` for an autoarray + regularization matrix, which already carries its coefficient. + n_probes + The number of Gaussian probe vectors. + n_steps + The maximum number of Lanczos steps per probe (capped at `S`). + seed + The seed of the probe generator. + + Returns + ------- + The estimate of `ln det(F + λ H)` and the `(n_probes,)` array of per-probe estimates (each + `ln det P` plus that probe's quadrature value), whose mean is the estimate. + """ + regularization_matrix_sparse = scipy.sparse.csr_matrix(regularization_matrix_sparse) + regularization_factor_sparse = scipy.sparse.csr_matrix(regularization_factor_sparse) + + n = operator.params + + rng = np.random.default_rng(seed) + + check = rng.standard_normal(n) + h_check = regularization_matrix_sparse @ check + rtr_check = regularization_factor_sparse.T @ (regularization_factor_sparse @ check) + + if np.linalg.norm(rtr_check - h_check) > 1.0e-8 * np.linalg.norm(h_check): + raise ValueError( + "regularization_factor_sparse R does not factor regularization_matrix_sparse H " + "(H != RᵀR); include the regularization coefficient in R (e.g. " + "sqrt(regularization.coefficient) * regularization.regularization_factor_sparse_from(mapper))." + ) + + preconditioner, factor = powell_preconditioner_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix_sparse, + coefficient=coefficient, + ) + + log_det_preconditioner = log_det_preconditioner_from(factor) + + def curvature_reg_matvec(v): + return operator.curvature_matvec(v) + coefficient * ( + regularization_matrix_sparse @ v + ) + + n_steps = min(int(n_steps), n) + + sqrt_kappa0 = np.sqrt(operator.kappa0) + sqrt_coefficient = np.sqrt(coefficient) + + probe_values = np.zeros(n_probes) + + for probe in range(n_probes): + g_image = rng.standard_normal(operator.mapping_matrix.shape[0]) + g_regularization = rng.standard_normal(regularization_factor_sparse.shape[0]) + + v = factor.solve( + sqrt_kappa0 * (operator.mapping_matrix_transpose @ g_image) + + sqrt_coefficient * (regularization_factor_sparse.T @ g_regularization) + ) + + pv = preconditioner @ v + norm2 = float(v @ pv) + norm = np.sqrt(norm2) + + q = v / norm + pq = pv / norm + q_prev = np.zeros(n) + beta_prev = 0.0 + + basis = [] + basis_p = [] + alpha = [] + beta = [] + + for step in range(n_steps): + basis.append(q) + basis_p.append(pq) + + aq = curvature_reg_matvec(q) + a = float(q @ aq) + + r = factor.solve(aq) - a * q - beta_prev * q_prev + + for q_j, pq_j in zip(basis, basis_p): + r -= float(pq_j @ r) * q_j + + pr = preconditioner @ r + b = float(np.sqrt(max(float(r @ pr), 0.0))) + + alpha.append(a) + + if b < 1.0e-10 * abs(a) or step == n_steps - 1: + break + + beta.append(b) + + q_prev, q, pq, beta_prev = q, r / b, pr / b, b + + k = len(alpha) + + tridiagonal = ( + np.diag(alpha) + np.diag(beta[: k - 1], 1) + np.diag(beta[: k - 1], -1) + ) + + theta, eigenvectors = np.linalg.eigh(tridiagonal) + + probe_values[probe] = norm2 * float( + np.sum(eigenvectors[0] ** 2 * np.log(np.clip(theta, 1.0e-300, None))) + ) + + log_det_probes = log_det_preconditioner + probe_values + + return float(np.mean(log_det_probes)), log_det_probes + + +def log_evidence_matrix_free_from( + operator: MatrixFreeDataOperator, + regularization_matrix_sparse: scipy.sparse.spmatrix, + regularization_factor_sparse: scipy.sparse.spmatrix, + coefficient: float, + reconstruction: np.ndarray, + log_det_regularization: float, + data_term: float, + noise_normalization: float, + n_probes: int = 10, + n_steps: int = 40, + seed: int = 20221, + log_det_curvature_reg: float = None, +) -> float: + """ + Returns the log Bayesian evidence of a matrix-free interferometer inversion, assembling the same + five terms with the same conventions as `FitInterferometer.log_evidence` / + `fit_util.log_evidence_from`: + + ln Z = -0.5 * [χ² + sᵀ(λH)s + ln det(F + λH) - ln det(λH) + Σ ln(2π σ²)] + + with + + - `χ² = dᵀN⁻¹d - 2 sᵀD + sᵀFs`, the dense path's `fast_chi_squared` (`data_term = dᵀN⁻¹d`, + `D = operator.data_vector`, `F s = operator.curvature_matvec(s)`); + - `sᵀ(λH)s` the regularization term; + - `ln det(F + λH)` by stochastic Lanczos quadrature (`log_det_curvature_reg_via_slq_from`), the + only approximate term (within 0.5% at the defaults on the test cases), unless + `log_det_curvature_reg` is given; + - `ln det(λH) = S ln λ + log_det_regularization`, exact (for `MaternSparse`, + `log_det_regularization_matrix_term_from`); + - the noise normalisation `Σ ln(2π σ²)`. + + With `λ = coefficient = 1.0` and an autoarray regularization matrix (which already carries its + coefficient) every term is the dense sparse-path quantity. + + Parameters + ---------- + operator + The data operator providing `F` and `D`. + regularization_matrix_sparse + The regularization matrix `H` as a `scipy.sparse` matrix. + regularization_factor_sparse + A sparse factor `R` with `H = RᵀR` (see `log_det_curvature_reg_via_slq_from`). + coefficient + The scalar `λ` multiplying `H`; `1.0` for an autoarray regularization matrix. + reconstruction + The reconstruction `s` (e.g. from `reconstruction_via_pcg_from`). + log_det_regularization + The exact `ln det H` of `regularization_matrix_sparse` (without `λ`). + data_term + `dᵀN⁻¹d = Σ d_r²/σ_r² + Σ d_i²/σ_i²`, as cached on the dataset's `sparse_operator.data_term` + (or the dataset interface's own `data_term` for profile-subtracted visibilities). + noise_normalization + `Σ ln(2π σ²)` over the real and imaginary visibility noise, as cached on the dataset's + `sparse_operator.noise_normalization`. + n_probes, n_steps, seed + The SLQ settings (see `log_det_curvature_reg_via_slq_from`). + log_det_curvature_reg + An optional precomputed `ln det(F + λH)` used in place of the SLQ estimate (e.g. the dense + value, for parity checks of the other terms). + + Returns + ------- + The log evidence. + """ + regularization_matrix_sparse = scipy.sparse.csr_matrix(regularization_matrix_sparse) + + s = np.asarray(reconstruction, dtype=np.float64) + + chi_squared = ( + float(data_term) + - 2.0 * float(s @ operator.data_vector) + + float(s @ operator.curvature_matvec(s)) + ) + + regularization_term = coefficient * float(s @ (regularization_matrix_sparse @ s)) + + if log_det_curvature_reg is None: + log_det_curvature_reg, _ = log_det_curvature_reg_via_slq_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix_sparse, + regularization_factor_sparse=regularization_factor_sparse, + coefficient=coefficient, + n_probes=n_probes, + n_steps=n_steps, + seed=seed, + ) + + log_det_regularization_term = s.size * np.log(coefficient) + float( + log_det_regularization + ) + + return float( + -0.5 + * ( + chi_squared + + regularization_term + + log_det_curvature_reg + - log_det_regularization_term + + float(noise_normalization) + ) + ) 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 913259b1e..b39a36791 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_matrix_free_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_matrix_free_util.py @@ -1,8 +1,10 @@ """ -Parity tests for the matrix-free interferometer utilities (phase 1 of PyAutoArray #638). +Parity tests for the matrix-free interferometer utilities (phases 1-3 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. +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. These tests deliberately carry no `importorskip("jax")`: the module is NumPy / scipy only, so the `unittest-nojax` CI leg runs them. @@ -26,10 +28,16 @@ ) -def _dataset_sparse_from(seed=1, n_visibilities=40): +def _dataset_sparse_from(seed=1, n_visibilities=40, uv_scale=1.0e4): """ A circular-masked dataset, so the slim grid is a strict subset of its unmasked extent and the scatter / gather between them is exercised, with random visibilities of unit noise. + + `uv_scale` sets the baseline lengths in wavelengths. At the default `1e4` the fringes are far + broader than the mask, so `W~` is close to a constant and nowhere near the Powell preconditioner's + `κ₀ I`; with many visibilities on long baselines (`500`, `1e6`) the dirty beam is compact, as for + real data, and `P⁻¹(F + H)` has eigenvalues near one, the regime the stochastic log-determinant + is built for. """ mask = aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=0.5, radius=2.5) @@ -38,7 +46,7 @@ def _dataset_sparse_from(seed=1, n_visibilities=40): dataset = aa.Interferometer( data=aa.Visibilities(visibilities=rng.normal(size=(n_visibilities, 2))), noise_map=aa.VisibilitiesNoiseMap(visibilities=np.ones((n_visibilities, 2))), - uv_wavelengths=rng.normal(size=(n_visibilities, 2)) * 1.0e4, + uv_wavelengths=rng.normal(size=(n_visibilities, 2)) * uv_scale, real_space_mask=mask, transformer_class=aa.TransformerDFT, ) @@ -307,11 +315,240 @@ def test__powell_preconditioner__equals_kappa0_mtm_plus_coefficient_regularizati ) +def _matern_sparse_system_from(coefficient=1.0): + """ + 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 + operator, `H` and its factor `R` with `H = RᵀR` (the coefficient folded into `R`). + """ + mask, dataset_sparse = _dataset_sparse_from(n_visibilities=500, uv_scale=1.0e6) + + regularization = aa.reg.MaternSparse(coefficient=coefficient, scale=1.0) + + mapper = _rectangular_mapper_from(mask=mask, regularization=regularization) + inversion = _inversion_from(dataset_sparse=dataset_sparse, mapper=mapper) + + operator = mf.MatrixFreeDataOperator.from_inversion( + inversion=inversion, mapper=mapper + ) + + regularization_matrix = regularization.regularization_matrix_sparse_from(mapper) + regularization_factor = np.sqrt( + regularization.coefficient + ) * regularization.regularization_factor_sparse_from(mapper) + + return ( + dataset_sparse, + regularization, + mapper, + inversion, + operator, + regularization_matrix, + regularization_factor, + ) + + +@pytest.mark.parametrize("coefficient", [1.0, 0.01]) +def test__log_det_preconditioner_from__equals_dense_slogdet(coefficient): + _, _, _, _, operator, regularization_matrix, _ = _matern_sparse_system_from( + coefficient=coefficient + ) + + preconditioner, factor = mf.powell_preconditioner_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + coefficient=1.0, + ) + + sign, log_det = np.linalg.slogdet(preconditioner.toarray()) + + assert sign == 1.0 + assert mf.log_det_preconditioner_from(factor) == pytest.approx(log_det, abs=1.0e-10) + + +@pytest.mark.parametrize("coefficient", [1.0, 0.01]) +def test__log_det_curvature_reg_via_slq__within_half_percent_of_dense_and_deterministic( + coefficient, +): + ( + _, + _, + _, + inversion, + operator, + regularization_matrix, + regularization_factor, + ) = _matern_sparse_system_from(coefficient=coefficient) + + sign, log_det_dense = np.linalg.slogdet(np.asarray(inversion.curvature_reg_matrix)) + + assert sign == 1.0 + + estimate, log_det_probes = mf.log_det_curvature_reg_via_slq_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + regularization_factor_sparse=regularization_factor, + coefficient=1.0, + n_probes=10, + n_steps=40, + ) + + assert log_det_probes.shape == (10,) + assert estimate == pytest.approx(np.mean(log_det_probes), 1.0e-14) + assert abs(estimate - log_det_dense) <= 0.005 * abs(log_det_dense) + + estimate_repeat, log_det_probes_repeat = mf.log_det_curvature_reg_via_slq_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + regularization_factor_sparse=regularization_factor, + coefficient=1.0, + n_probes=10, + n_steps=40, + ) + + assert estimate_repeat == estimate + assert np.array_equal(log_det_probes_repeat, log_det_probes) + + +def test__log_det_curvature_reg_via_slq__error_shrinks_with_probes(): + """ + Loose convergence check: over a handful of seeds the mean error at 10 probes is below that at 2, + and the Monte-Carlo standard error of the mean shrinks. + """ + ( + _, + _, + _, + inversion, + operator, + regularization_matrix, + regularization_factor, + ) = _matern_sparse_system_from() + + _, log_det_dense = np.linalg.slogdet(np.asarray(inversion.curvature_reg_matrix)) + + def errors_and_standard_errors(n_probes): + errors, standard_errors = [], [] + + for seed in range(5): + estimate, log_det_probes = mf.log_det_curvature_reg_via_slq_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + regularization_factor_sparse=regularization_factor, + coefficient=1.0, + n_probes=n_probes, + seed=seed, + ) + errors.append(abs(estimate - log_det_dense)) + standard_errors.append(np.std(log_det_probes) / np.sqrt(n_probes)) + + return np.mean(errors), np.mean(standard_errors) + + error_2, standard_error_2 = errors_and_standard_errors(n_probes=2) + error_10, standard_error_10 = errors_and_standard_errors(n_probes=10) + + assert error_10 < error_2 + assert standard_error_10 < standard_error_2 + + +def test__log_det_curvature_reg_via_slq__factor_not_matching_regularization_matrix__raises(): + _, regularization, mapper, _, operator, regularization_matrix, _ = ( + _matern_sparse_system_from(coefficient=0.01) + ) + + with pytest.raises(ValueError): + mf.log_det_curvature_reg_via_slq_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + regularization_factor_sparse=regularization.regularization_factor_sparse_from( + mapper + ), + coefficient=1.0, + ) + + +@pytest.mark.parametrize("coefficient", [1.0, 0.01]) +def test__log_evidence_matrix_free_from__matches_dense_sparse_path_evidence( + coefficient, +): + """ + The dense reference is `FitInterferometer.log_evidence`'s assembly, `fit_util.log_evidence_from`, + of the unconstrained sparse-path inversion's terms. With the dense `ln det(F + H)` substituted the + other four terms agree to round-off; with the SLQ estimate the total agrees within the SLQ + tolerance (0.5% of `ln det(F + H)`). + """ + from autoarray.fit import fit_util + + ( + dataset_sparse, + regularization, + mapper, + inversion, + operator, + regularization_matrix, + regularization_factor, + ) = _matern_sparse_system_from(coefficient=coefficient) + + noise_normalization = dataset_sparse.sparse_operator.noise_normalization + + log_evidence_dense = fit_util.log_evidence_from( + chi_squared=inversion.fast_chi_squared, + regularization_term=inversion.regularization_term, + log_curvature_regularization_term=inversion.log_det_curvature_reg_matrix_term, + log_regularization_term=inversion.log_det_regularization_matrix_term, + noise_normalization=noise_normalization, + ) + + reconstruction, info = mf.reconstruction_via_pcg_from( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + coefficient=1.0, + rtol=1.0e-12, + ) + + assert info == 0 + + log_det_regularization = regularization.log_det_regularization_matrix_term_from( + mapper + ) + + assert log_det_regularization == pytest.approx( + inversion.log_det_regularization_matrix_term, abs=1.0e-10 + ) + + kwargs = dict( + operator=operator, + regularization_matrix_sparse=regularization_matrix, + regularization_factor_sparse=regularization_factor, + coefficient=1.0, + reconstruction=reconstruction, + log_det_regularization=log_det_regularization, + data_term=dataset_sparse.sparse_operator.data_term, + noise_normalization=noise_normalization, + ) + + log_evidence_exact_log_det = mf.log_evidence_matrix_free_from( + **kwargs, log_det_curvature_reg=inversion.log_det_curvature_reg_matrix_term + ) + + assert log_evidence_exact_log_det == pytest.approx( + log_evidence_dense, abs=1.0e-10 * abs(log_evidence_dense) + ) + + log_evidence_slq = mf.log_evidence_matrix_free_from(**kwargs) + + _, log_det_dense = np.linalg.slogdet(np.asarray(inversion.curvature_reg_matrix)) + + assert abs(log_evidence_slq - log_evidence_dense) <= 0.5 * 0.005 * abs( + log_det_dense + ) + + 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. + imports the module and runs a PCG solve and the phase-3 log-determinant / evidence functions. """ script = textwrap.dedent( """ @@ -391,6 +628,34 @@ def _find_spec_blocking_jax(cls, name, path=None, target=None): assert info == 0 assert np.allclose(reconstruction, np.asarray(inversion.reconstruction)) + regularization_matrix = scipy.sparse.csr_matrix( + np.asarray(inversion.regularization_matrix) + ) + regularization_factor = scipy.sparse.csr_matrix( + np.linalg.cholesky(regularization_matrix.toarray()).T + ) + + _, factor = mf.powell_preconditioner_from(operator, regularization_matrix, 1.0) + assert np.isfinite(mf.log_det_preconditioner_from(factor)) + + log_det, _ = mf.log_det_curvature_reg_via_slq_from( + operator, regularization_matrix, regularization_factor, 1.0, n_probes=2 + ) + assert np.isfinite(log_det) + + log_evidence = mf.log_evidence_matrix_free_from( + operator, + regularization_matrix, + regularization_factor, + 1.0, + reconstruction, + np.linalg.slogdet(regularization_matrix.toarray())[1], + 1.0, + 1.0, + n_probes=2, + ) + assert np.isfinite(log_evidence) + assert not any( name.split(".")[0] in ("jax", "jaxlib", "nufftax") for name in sys.modules )