Matrix-free solution to make very large meshes more efficient for Interferometer

[HRSAstro]

Summary

For large source meshes the inversion is limited by the dense n × n curvature and regularization matrices. Memory scales as n² and each Cholesky or NNLS solve as ~n³. For interferometer data with the w-tilde formalism, this puts a practical ceiling at around 100 × 100 source pixels. I have implemented a matrix-free alternative in PyUVimage (built on autoarray/autogalaxy), following Powell et al. (2021), and would like to suggest it as an option in autoarray. I'm happy to share the code.

The problem

At 192 × 192 (36,864 pixels), F + H takes 11 GB in double precision and each factorisation is ~2 × 10¹³ flops. The non-negative solver refactorises many times per solve. On our test case (NGC 4274, 2.6 M visibilities), the first positive solve extrapolated to days on 2 cores.

The approach

No n × n matrix is formed. Each piece is replaced by an operator or a sparse matrix:

1. Data term. F s = Mᵀ W̃ (M s). W̃ is a convolution on the image grid, so it can be applied with one FFT pair on a 2× zero-padded grid, using the w-tilde kernel autoarray already computes. M is the mapper's interpolation matrix, built as a sparse matrix from pix_indexes and weights. D = Mᵀ (dirty image). Both agree with the dense path to 1e-16.
2. Prior. A Matérn covariance has a sparse precision matrix when α = ν + 1 is an integer (Lindgren, Rue & Lindström 2011). For α = 2 (ν = 1): C⁻¹ = c L², with L = κ² − ∇² (5-point Laplacian), κ = √(2ν)/ℓ, and c set so the variance is 1. Per-pixel weights (adaptive or brightness-based priors) enter as W⁻¹ c L² W⁻¹. As far as I can tell, the current MaternKernel regularization builds the dense covariance and inverts it. The sparse form only exists for integer α, so ν = 1.5 becomes ν = 1.
3. Solve. Preconditioned conjugate gradients on (F + λC⁻¹) s = D. The preconditioner is Powell et al.'s P = κ₀ MᵀM + λC⁻¹, where κ₀ is the peak of W̃. P is sparse and is factorised once per λ. For unweighted priors, an FFT preconditioner built from the gridded uv weight density works about as well.
4. Evidence. ln det H is exact from the sparse factor. ln det(F + H) = ln det P + ln det(P⁻¹(F + H)), with the second term from stochastic Lanczos quadrature (Ubaru, Chen & Saad 2017; 10 fixed probe vectors, ≤ 40 Lanczos steps). This agrees with the dense evidence to ~0.2% of ln Z on our tests. Powell et al. stop at ln det P, which they report can be 1–20% off.
5. Positivity. A primal-dual interior point method (Mehrotra predictor-corrector). Each step is a CG solve preconditioned by P plus the barrier's diagonal. Once the set of pixels held at zero has settled, the free pixels are solved exactly, with a few primal-dual active-set corrections (Hintermüller, Ito & Kunisch 2002). This gives exact zeros and matches scipy's NNLS on small problems. I also tried a projected Newton method and a pure active-set method, but both stalled at weak priors where ~30% of pixels sit at zero.

Results

Real ALMA data, 2 cores, visibilities streamed, adaptive prior, regularization strength set by a residual criterion. "w-tilde" is the current dense path.

dataset	mesh	w-tilde, positive	matrix-free, positive	matrix-free, unconstrained
lensed SMG	50²	34 s, 1.07 GB	5.1 s, 0.37 GB	4.0 s
lensed SMG	58²	77 s, 1.46 GB	15 s, 0.39 GB	4.4 s
[CII] galaxy	54²	69 s, 1.21 GB	7.0 s, 0.37 GB	40 s
NGC 4274	64²	429 s, 2.00 GB	49 s, 0.40 GB	12 s
NGC 4274	192²	not feasible	266 s, 0.80 GB	60 s

On the compact sources, the positive matrix-free models agree with the positive w-tilde models to 0.2–0.7% of the peak after smoothing to the beam. Most of the remaining difference on extended emission comes from ν = 1 vs 1.5.

Limitations

• It needs a regular rectangular mesh, because the prior is built on the grid. Delaunay or Voronoi meshes could use the finite-element form of the same prior on the triangulation (Lindgren et al.'s original setting), which I haven't tried.
• Only integer α, so ν = 1, 2, …
• Uncertainty maps need diag((F + H)⁻¹), which would need a stochastic diagonal estimator. Not implemented yet.
• Linear light profiles / point components are not handled on this path yet.
• The number of CG iterations depends on how strong the prior is relative to the data, so weak priors cost more.

Possible shape

Perhaps a separate inversion class (or a Settings option) for interferometer data on a rectangular mesh, where the curvature-plus-regularization matrix is a LinearOperator. The reconstruction would come from CG or the interior point method, and the log det terms from the sparse factor plus Lanczos quadrature. Much of the inversion interface reads the dense matrices directly, so a separate path may be simpler to start with.

References

• Powell D. M. et al., 2021, MNRAS, 501, 515
• Lindgren F., Rue H., Lindström J., 2011, J. R. Stat. Soc. B, 73, 423
• Ubaru S., Chen J., Saad Y., 2017, SIAM J. Matrix Anal. Appl., 38, 1075
• Mehrotra S., 1992, SIAM J. Optim., 2, 575
• Hintermüller M., Ito K., Kunisch K., 2002, SIAM J. Optim., 13, 865

[Jammy2211]

@HRSAstro thank you, this is a very clear proposal and the benchmark table makes the case. We would like to incorporate it and have opened PyAutoLabs/PyAutoArray#638 to track it, with your write-up quoted; we will plan the phases there shortly and would value your involvement.

To guide that, do you have working Python code we could look at, even just snippets showing the operator data term, the sparse Matérn precision, the SLQ evidence and the interior-point solve running on your build? Pointers into pyuvimage are ideal. We will post an update here when work starts.

[HRSAstro]

upstream_snippets.py

I attached some code snippets that includes the matrix-free solution implementation, with evidence assembly and postive-only solve. I ran it against pyuvimage's implementation and a dense solve on a small mock:

check	result
F against the dense matrix	agrees to 6e-16
sparse Matérn prior against PyUVimage’s	identical
CG solution against a dense solve	agrees to 3e-5
ln det(F + H), Lanczos estimate against exact	within 0.5%
positive solve against scipy NNLS	same pixels held at zero

[HRSAstro]

Note these snippets also include the primary beam implementation (#35)

[Jammy2211]

@HRSAstro thanks, the snippets and the check table made planning straightforward. Phase 1 (operator data term and preconditioned CG with parity tests against the current sparse path) is now up as PyAutoLabs/PyAutoArray#645 under PyAutoLabs/PyAutoArray#638; the sparse Matérn precision, Lanczos evidence and interior-point positivity follow as phases 2 to 4. We will keep that issue updated.