torchmfbd: a flexible multi-object multi-frame blind deconvolution code

TL;DR

Torchmfbd delivers a modern, open-source MOMFBD implementation built on PyTorch, enabling GPU-accelerated, flexible restoration of ground-based solar images with spatially variant PSFs. It combines MAP-based inference with wavefront-based and data-driven (NMF) PSF parameterizations, and incorporates both implicit and explicit regularization, destretching, and patch-based processing. The framework supports large fields of view and provides efficient optimization using downstream GPU acceleration, demonstrated across IMaX, CRISP, HiFI, and CHROMIS data, with competitive image quality and reduced computation time. The code is designed to be extensible, allowing new reconstruction and regularization methods to be added, with future work including linear equality constraints and proximal projection approaches to handle more complex acquisition schemes.

Abstract

Post-facto image restoration techniques are essential for improving the quality of ground-based astronomical observations, which are affected by atmospheric turbulence. Multi-object multi-frame blind deconvolution (MOMFBD) methods are widely used in solar physics to achieve diffraction-limited imaging. We present torchmfbd, a new open-source code for MOMFBD that leverages the PyTorch library to provide a flexible, GPU-accelerated framework for image restoration. The code is designed to handle spatially variant point spread functions (PSFs) and includes advanced regularization techniques. The code implements the MOMFBD method using a maximum a-posteriori estimation framework. It supports both wavefront-based and data-driven PSF parameterizations, including a novel experimental approach using non-negative matrix factorization. Regularization techniques, such as smoothness and sparsity constraints, can be incorporated to stabilize the solution. The code also supports dividing large fields of view into patches and includes tools for apodization and destretching. The code architecture is designed to become a flexible platform over which new reconstruction and regularization methods can also be implemented straightforwardly. We demonstrate the capabilities of torchmfbd on real solar observations, showing its ability to produce high-quality reconstructions efficiently. The GPU acceleration significantly reduces computation time, making the code suitable for large datasets. The code is publicly available at https://github.com/aasensio/torchmfbd.

Figures (8)

• Figure 1: Basis functions of the NMF dictionary for the CRISP instrument at 8542 Å. The first 100 basis functions, out of a total of 150, are shown.
• Figure 2: Reconstruction of PSFs with different values of $r_0$ using the NMF dictionary for the CRISP instrument at 8542 Å. The PSF is shown in the rightmost column, while the NMF reconstructions with different values of $M$ are shown in the other columns.
• Figure 3: Computing time (solid lines) and memory consumption (dashed lines) for reconstructing a pair of WB and NB images with 15 frames in an NVIDIA RTX 4090 GPU as a function of the number of batches. Patches of different size are considered with different colors.
• Figure 4: Azimuthally averaged power spectra for relevant observations shown in Sec. \ref{['sec:results']}. The average is only performed for angles in the Fourier plane with $-45\pm 15^\circ$ and $45\pm 15^\circ$. The vertical line marks the diffraction limit. Note that the power spectrum is computed in the full FoV, so that some high frequency artifacts can appear when the FoV is built from mosaicked patches. To avoid showing noise-dominated artifacts, the power below an estimated noise floor (noise level per frame divided by the number of frames) is not displayed.
• Figure 5: Results of the reconstruction of IMaX data using single-frame phase diversity. The upper panels show the focused and defocused frames. The lower panels display the results from torchmfbd alongside the reconstructed image of the original data release.