Scalable Bayesian uncertainty quantification with data-driven priors for radio interferometric imaging

Abstract

Next-generation radio interferometers like the Square Kilometer Array have the potential to unlock scientific discoveries thanks to their unprecedented angular resolution and sensitivity. One key to unlocking their potential resides in handling the deluge and complexity of incoming data. This challenge requires building radio interferometric (RI) imaging methods that can cope with the massive data sizes and provide high-quality image reconstructions with uncertainty quantification (UQ). This work proposes a method coined quantifAI to address UQ in RI imaging with data-driven (learned) priors for high-dimensional settings. Our model, rooted in the Bayesian framework, uses a physically motivated model for the likelihood. The model exploits a data-driven convex prior potential, which can encode complex information learned implicitly from simulations and guarantee the log-concavity of the posterior. We leverage probability concentration phenomena of high-dimensional log-concave posteriors to obtain information about the posterior, avoiding MCMC sampling techniques. We rely on convex optimization methods to compute the MAP estimation, which is known to be faster and better scale with dimension than MCMC strategies. quantifAI allows us to compute local credible intervals and perform hypothesis testing of structure on the reconstructed image. We propose a novel fast method to compute pixel-wise uncertainties at different scales, which uses three and six orders of magnitude less likelihood evaluations than other UQ methods like length of the credible intervals and Monte Carlo posterior sampling, respectively. We demonstrate our method by reconstructing RI images in a simulated setting and carrying out fast and scalable UQ, which we validate with MCMC sampling. Our method shows an improved image quality and more meaningful uncertainties than the benchmark method based on a sparsity-promoting prior.

1 INTRODUCTION

Radio astronomy plays a crucial role in expanding our understanding of the Universe, offering a unique perspective on astrophysical and cosmological phenomena. Among the transformative tools in an astronomer’s toolkit, radio interferometric (RI) imaging stands out as an indispensable technique. Aperture synthesis and radio interferometry (Thompson, Moran & Swenson 2017) allow us to achieve high angular resolutions providing immense power to resolve objects. Furthermore, radio frequency signals are only weakly attenuated by our atmosphere, allowing for observations at the Earth’s surface. The unparalleled angular resolution, high sensitivity, and the different phenomena emitting in the radio wavelength regime make RI an ideal candidate to better help us understand our Universe.

The advent of the Square Kilometre Array (SKA; Dewdney et al. 2009) heralds a new era in radio astronomy (Braun et al. 2015) spanning the study from the epoch of reionization and fast radio bursts to galaxy evolution and dark energy. SKA’s vast collecting area and sensitivity promise to provide a leap forward in our observational capabilities, opening doors to discoveries. However, this transformative potential comes with the formidable computational challenge of processing and making sense of the unprecedented volume of SKA-generated data. Developing and implementing algorithms that can efficiently handle SKA’s data deluge is a challenge. In addition, achieving the high reconstruction performance required to unlock SKA’s full potential is a significant obstacle in the SKA’s data processing requirements.

The aperture synthesis techniques in RI probe the sky by acquiring specific Fourier measurements, which results in incomplete coverage of the Fourier domain of the sky’s image of interest. The incomplete Fourier coverage makes the problem of estimating the underlying sky image, which we know as RI imaging, an ill-posed inverse problem, which is further complexified by the observational noise. Having a way to quantify the uncertainty in the image reconstructions becomes essential given the uncertainties involved in the RI imaging problem. To make scientifically sound inferences and informed decisions, we need the ability to quantify these uncertainties rigorously. This motivates the development of uncertainty quantification (UQ) methods tailored to the complexities of radio interferometric data, where scalability, i.e. the computational complexity with respect to the amount of data processed, and performance play a central role. We need to ensure that our reconstructions are not only insightful but also trustworthy.

In a nutshell, we want to develop RI imaging methods that can deliver precision with uncertainty quantification and that are highly scalable. Most existing methods only tackle some of these three requirements. The widely used clean algorithm (Högbom 1974) built its success on scalability and fast inference. clean and its extensions (Cornwell 2008; Offringa et al. 2014; Offringa & Smirnov 2017) have been continuously used in many RI imaging pipelines since its inception. Despite offering limited imaging quality and reconstruction artefacts compared to other approaches, clean stands out due to its scalability. More recent approaches leverage compressed sensing theory, relying on sparse priors (often in wavelet representations) and convex optimization techniques (Wiaux et al. 2009; McEwen & Wiaux 2011; Carrillo, McEwen & Wiaux 2012, 2014; Dabbech et al. 2015, 2018; Pratley et al. 2018). These methods have been shown to improve the reconstruction quality at the expense of increased computational complexity. Considerable work has been directed to parallelization and acceleration efforts for sparsity-based methods (Onose et al. 2016; Pratley et al. 2019a; Pratley, Johnston-Hollitt & McEwen 2019b; Thouvenin et al. 2022a, b).

The deep learning revolution has introduced a powerful way to encode complex image priors in neural networks, which may be used to solve complex high-dimensional inverse problems. This data-driven or learned paradigm has gained much traction across imaging landscape problems, including RI imaging (Allam 2016; Terris et al. 2022; Aghabiglou et al. 2023; Mars, Betcke & McEwen 2023). Learned methods can improve the reconstruction quality with respect to handcrafted priors, such as sparsity-based wavelet priors, as well as provide acceleration (Terris et al. 2022; Mars et al. 2023, 2024; Aghabiglou et al. 2024) to convex optimization-based methods.

Unfortunately, none of the RI imaging methods mentioned above, learned, sparsity-based, or clean-based, provide UQ tools for a given model. Cai, Pereyra & McEwen (2018a, b) proposed methods for Bayesian UQ on RI imaging problems. Cai et al. (2018a) leverages proximal Markov chain Monte Carlo (MCMC) methods (Pereyra 2016) to provide support for sparsity-promoting priors. The proposed method allows them to reconstruct the image and provide UQ by sampling the posterior probability distribution. The drawback of the method is the high computational cost suffered by all MCMC sampling techniques. The companion paper, Cai et al. (2018b), overcomes the need for posterior sampling with maximum-a-posteriori (MAP) based UQ (Pereyra 2017) relying on convex optimization techniques. The second method (Cai et al. 2018b) provides a significant speed-up with respect to the sampling-based method (Cai et al. 2018a), but its reconstruction quality is limited to sparsity-promoting priors.

Other RI imaging methods have addressed UQ for their reconstructions. Dia et al. (2023) proposed to use score-based generative models as priors (Ho, Jain & Abbeel 2020; Song et al. 2020), which by employing the convolved likelihood approximation (Adam et al. 2022; Remy et al. 2023) are able to sample from the posterior. However, the method is computationally costly, and the sampling relies on MCMC methods. The Bayesian RI imaging method comrade (Tiede 2022) was developed for very-long-baseline interferometry (VLBI) aiming to image black holes and active galactic nuclei. The comrade method relies on MCMC sampling methods like nested sampling (Ashton et al. 2022) and Hamiltonian Monte Carlo (Xu et al. 2020) to sample from the posterior. Posterior sampling based on MCMC methods can be enough for VLBI but cannot yet cope with the dimensions of SKA-like images.

An alternative Bayesian RI imaging algorithm is the resolve method and its upgrades (Arras et al. 2018, 2021, 2022; Knollmüller, Arras & Enßlin 2023; Roth et al. 2023), introduced initially in Junklewitz et al. (2016), which can produce uncertainty maps from approximate posterior samples. The model is based on Gaussian random fields with different analytical priors used to promote certain aspects of the reconstruction, like positivity and spatial and temporal correlations. One novelty of the method is to approximate the true posterior probability distribution with variational inference and, therefore, be able to sample from the approximated distribution without resorting to MCMC methods that do not scale well in very high dimensions. The approximated posterior follows a multivariate Gaussian distribution, where their parameters are learnt by trying to maximize the information overlap between the true posterior and the model by minimizing a Kullback-Leibler (KL-) divergence. Given that the full covariance matrix scales with the squared of the number of pixels, the authors exploit an approximation in the vicinity of the mean estimate. The method is based on the metric Gaussian variational inference (MGVI) approach from Knollmüller & Enßlin (2019). resolve has been used for reconstructions with VLBI observations of M87* (Arras et al. 2022) and Sagittarius A* (Knollmüller et al. 2023) from the Event Horizon Telescope (EHT) (Event Horizon Telescope Collaboration 2019a, b), as well as observations of Cygnus A from the very large array (VLA) (Arras et al. 2018, 2021; Roth et al. 2023).

In this article, we delve into the forefront of RI imaging and propose a method coined quantifAI, based on learned convex priors, capable of delivering high-quality reconstructions with uncertainty quantification and being highly scalable. The method relies on the mathematically principled Bayesian framework to provide an understanding of the uncertainties through the posterior distribution. By restricting our model to log-concave posteriors, we can exploit recent MAP-based UQ techniques (Pereyra 2017), providing scalable optimization-based UQ. We build upon recent advances in neural-network-based convex regularizers (Goujon et al. 2023b), allowing us to improve the reconstruction quality and obtain more meaningful uncertainties with respect to Cai et al. (2018b). On top of the hypothesis tests of structure on the reconstructed image, we propose a novel fast method to estimate pixel-wise uncertainties as a function of scale.

The remainder of this article is organized as follows. In Section 2, we start by reviewing the RI imaging and techniques for the resulting inverse problem. Section 3 describes quantifAI, the proposed method, and the RI image reconstruction algorithm. In Section 4, we introduce the core of our scalable UQ and the different UQ techniques it allows. The experimental results, including the performance of quantifAI reconstruction and its UQ techniques, are presented in Section 5. Section 6 presents experimental results using more realistic observation based on simulated ungridded visibility patterns from the MeerKAT radio telescope. In Section 7 we provide a discussion on some limitations and possible extensions of the proposed methodology. We provide concluding remarks and present some future perspectives in Section 8.

2 RADIO INTERFEROMETRIC IMAGING

In this section, we start by reviewing the RI imaging inverse problem and discuss approaches to tackle it, including sparsity-based regularization, the clean method, and learned approaches. We then introduce the Bayesian framework elements needed such as MAP estimation and proximal MCMC sampling algorithms that will be later used as validation.

2.1 Radio interferometry

The interferometric measurement equation for a radio telescope (Thompson et al. 2017) in the monochromatic setting relates our observations represented by the visibility function |$\mathcal {Y}$| to the sky brightness |$\mathcal {X}$|⁠, which we want to reconstruct,

where |$\boldsymbol {u}=(u,v,w)$| are the interferometer baseline coordinates with units depending on the observation wavelength, |$\boldsymbol {l}=(l,m,n)$| are cosine sky coordinates restricted to the unit sphere, and |$\mathcal {A}$| includes direction-dependent effects (DDEs) like the primary beam of the dishes. The previous general model allows us to consider different DDEs through |$\mathcal {A}$| and non-coplanar effects through the exponential term in w. These effects become considerable when considering wide fields of view and long baselines. There exists a rich body of literature incorporating such effects, e.g. Smirnov (2011a, b, c, d), Thompson et al. (2017), and there are scalable algorithms that take them into account, e.g. Pratley et al. (2019b).

In this article, for the sake of simplicity but without loss of generality, we assume the coplanar setting, where the antennas are located in the same w plane. We also assume that we observe a small field of view such that |$1 - l^{2} - m^{2} \approx 1$|⁠. Consequently, we have |$\exp\! \left[ -2 \pi \text{i} w\! \left( \sqrt{1 - l^{2} - m^{2}} - 1 \right) \right] \approx 1$|⁠, and equation (1) reduces to

From the previous equation, we can notice the remarkable result of |$\mathcal {Y}(u,v) = \mathcal {F}(\mathcal {A} \mathcal {X})(u,v)$| where |$\mathcal {F}$| is the 2D Fourier transform.

To further simplify the problem, we will avoid using the continuous |$\mathcal {Y}$| and |$\mathcal {X}$| and work with their discrete counterparts, |$\boldsymbol {x}$| and |$\boldsymbol {y}$|⁠, respectively. The observational model we study for our RI imaging problem writes

where |$\boldsymbol {y} \in \mathbb {C}^{M}$| are the M observed complex visibilities, |$\boldsymbol {x} \in \mathbb {R}^{N}$| is the discrete sky brightness sampled on a N point grid, and |$\boldsymbol {\Phi } \in \mathbb {C}^{M \times N}$| is the linear measurement operator that models the acquisition process. Without loss of generality, the observational and instrumental noise |$\boldsymbol {n} \in \mathbb {C}^{M}$| is assumed to be independent and identically distributed (iid) white Gaussian noise with zero mean and standard deviation |$\sigma$|⁠. If the noise is not white, we can incorporate a noise whitening matrix in the |$\boldsymbol {\Phi }$| operator such that the previous white noise assumption holds.

Each pair of antennas provides us with one visibility, which is a noisy Fourier component of the intensity image. Using an array of n radio antennas allows us to sample |$\binom{n}{2}=(n^2 - n)/2$| points in the |$uv$|-plane (or Fourier plane). The distribution of these points depends on the configuration of the radio antenna array. If different time intervals are considered, the Earth’s rotation can be exploited to increase the number of |$uv$| points. The |$uv$| coverage is incomplete in all practical cases, and the measurements are noisy. Therefore, the linear operator |$\boldsymbol {\Phi }$| is ill-posed. If we also consider a large number of measurements, recovering |$\boldsymbol {x}$| from |$\boldsymbol {y}$| becomes a challenging inverse problem.

The most basic reconstruction of |$\boldsymbol {x}$| is often referred to as the naturally weighted dirty image that we will denote from now on as |$\hat{\boldsymbol {x}}_{\text{dirty}}$| and call by dirty image as calibration weights are not currently being taken into account in this work. This estimation is obtained by applying the adjoint of |$\boldsymbol {\Phi }$| to the visibilities |$\boldsymbol {y}$|⁠. To obtain a higher fidelity solution to the RI imaging inverse problem, we must regularize the problem by incorporating some prior information about the desired solutions |$\boldsymbol {x}$|⁠. A broad range of methods can be characterized by what type of prior information is used to regularize the inverse problem and which algorithm is used to compute the reconstructed image |$\hat{\boldsymbol {x}}$|⁠.

2.2 Sparsity-based regularization

The last two decades have brought us a great number of RI imaging methods based on sparse representations. The prior information exploited is that the solution |$\boldsymbol {x}$| is known to be sparsely represented in some bases or dictionaries. The bases are often built using multiscale wavelets, or a dictionary is constructed with a collection of wavelets (Mallat 2008). We can represent our image |$\boldsymbol {x}$| in a dictionary |$\boldsymbol {\Psi } \in \mathbb {C}^{N \times L}$|⁠,

where |$\boldsymbol {a} \in \mathbb {C}^{L}$| is a vector of coefficients of |$\boldsymbol {x}$| weighting the corresponding dictionary atoms of |$\boldsymbol {\Psi }$|⁠. The assumption made to regularize the inverse problem is that |$\boldsymbol {a}$| is sparse or compressible, meaning that most of the coefficients are zero-valued or near zero, respectively. An array |$\boldsymbol {a}$| is called k-sparse if it has only k non-zero elements, which can be written as |$\Vert \boldsymbol {a} \Vert _{0} = k$|⁠, where |$\Vert \cdot \Vert _{0}$| denotes the |$\ell _{0}$| pseudo-norm.

Sparsity should be ideally enforced through the |$\ell _{0}$| pseudo-norm, which is non-convex. Consequently, a convex relaxation to the |$\ell _1$| norm is used, which is a sparsity-promoting norm. The optimization problem is formulated such that its solution coincides with the inverse problem solution. Therefore, the inverse problem can be tackled with an optimization algorithm. The optimization objective comprises two competing terms: (i) a data-fidelity term |$f(\cdot )$| that promotes consistency with the observed visibilities and depends on the statistics of the noise |$\boldsymbol {n}$|⁠; and (ii) a regularization term |$r(\cdot )$| that encodes our prior knowledge of |$\boldsymbol {x}$|⁠. The optimization problem reads

where we are using a sum of regularization terms |$r_k$|⁠, each with its corresponding regularization strength parameter |$\lambda _k$|⁠. Substituting the RI data fidelity and sparsity-enforcing regularization in an overdetermined wavelet dictionary |$\Psi$| terms into equation (3) we obtain

where |$\sigma$| is the noise standard deviation.

The previous formulation is referred to as unconstrained. Other works consider the constrained formulation, which minimizes the |$\ell _1$| term with respect to a hard |$\ell _2$|-ball constraint over |$\boldsymbol {y}$| with a radius of |$\epsilon$|⁠, which is related to the noise’s |$\sigma$| (Carrillo et al. 2012; Pratley et al. 2018). In this article, we will focus on the unconstrained formulation as it has a natural Bayesian interpretation. Obtaining the solution from equation (6) involves solving a convex optimization problem, where we have the sum of a differentiable and a non-differentiable term. Proximal algorithms (Parikh & Boyd 2014) are well suited to tackle such optimization problems. Recent developments brought us a wide collection of proximal optimization algorithms, such as the forward-backward (FB) algorithm (Combettes & Pesquet 2009), the FISTA algorithm (Beck & Teboulle 2009), the alternating direction method of multipliers (ADMM; Boyd et al. 2011), and the primal-dual forward-backward algorithm (Chambolle & Pock 2011; Condat 2013), to mention a few. A rich literature exists exploiting the aforementioned concepts to tackle the RI imaging problem (Wiaux et al. 2009; McEwen & Wiaux 2011; Carrillo et al. 2012, 2014; Onose et al. 2016; Pratley et al. 2018, 2019a, b; Cai et al. 2018b; Pratley & McEwen 2019). For example, the sparsity averaging reweighted analysis (SARA) family of methods (Carrillo et al. 2012) use an overcomplete dictionary composed of a concatenation of the Dirac basis and the first eight Daubechies wavelets (Daubechies 1992) and has shown good performance in RI imaging.

2.3 CLEAN

Precursor of RI image reconstructions, the clean algorithm (Högbom 1974) is a highly successful RI imaging method and it is still being used (Event Horizon Telescope Collaboration 2019a, b) despite various negative characteristics. The clean algorithm is a non-linear iterative method that assumes a sparse sky model. clean iteratively removes the contribution of the brightest source convolved with the instrument’s point spread function or dirty beam. This method can be interpreted as a matching pursuit algorithm (Wiaux et al. 2009), or an |$\ell _0$| regularization with a basis composed of a sum of Dirac spikes. Several extensions of clean have been developed over time (Bhatnagar & Cornwell 2004; Cornwell 2008; Stewart, Fenech & Muxlow 2011; Offringa et al. 2014; Offringa & Smirnov 2017) achieving better reconstruction performance. See Rau et al. (2009) for a review of clean-based algorithms.

On top of a very early introduction, the success of clean resides in its scalability. However, clean has been shown to produce artefacts when point sources do not well describe the underlying sky model limiting clean’s image quality and justifying the need for more advanced techniques, e.g. based on Section 2.2. clean often requires manual intervention, making its use less practical. Furthermore, clean and its extensions do not provide meaningful uncertainty quantification of its reconstruction.

2.4 Learned approaches

The advent of deep learning models has affected many imaging applications, and RI imaging is no exception. Handcrafted models and priors are limited in the information they can capture or represent with respect to recent, more expressive neural networks. Learned or data-driven methods can encode complex information existing in the data, e.g. astrophysical simulations, used in their training. In general, these approaches produce reconstructions with improved quality, a computational speed-up, or both. These reasons make learned approaches very relevant to the RI imaging reconstruction problem. However, there are issues regarding the robustness of learned methods to data distribution shifts (Hendrycks et al. 2020) and scalable methods for uncertainty quantification to the reconstruction.

Allam (2016) proposed a learned method for RI imaging based on convolutional neural networks (Dong et al. 2016) originally considered for super-resolution. The approach consists of learning to post-process dirty images with variants for both, known and unknown PSF. More recently, Gheller & Vazza (2021) proposed to use a convolutional denoising autoencoder to learn to post-process radio images, e.g. the dirty image or clean’s output. Connor et al. (2022) proposed a residual deep neural network (DNN) coined POLISH that works as a learned post-processing and super-resolution network. The DNN is based on the architecture proposed in Yu et al. (2018) and takes as input dirty images at different wavelengths and resolutions. POLISH outputs a clean image at a higher resolution for each wavelength and shows a better reconstruction quality than clean. The proposed method has been applied to simulations from the upcoming Deep Synoptic Array-2000 (Hallinan et al. 2019) and real data from the very large array (VLA; Perley et al. 2011).

The Plug-and-Play (PnP) framework (Venkatakrishnan, Bouman & Wohlberg 2013) provides a way to incorporate a deep learning model into a modern optimization algorithm. The central idea is to replace a proximal regularization term with a denoising deep neural network. Ryu et al. (2019) studied conditions for the convergence of PnP algorithms. Pesquet et al. (2021) proposed a new term for the denoiser’s training loss that enforces the firm non-expansiveness of the denoiser, which is usually deep learning-based. This training procedure allows the denoiser to suit a PnP framework with theoretical convergence conditions. The PnP framework with the non-expansiveness enforced to the deep learning-based denoiser has been applied to the RI imaging problem in Terris et al. (2022), where the approach has been called AIRI for Artificial Intelligence for Regularisation in RI imaging. The approach achieved similar or better performance than competing prior-based approaches whilst providing a significant acceleration potential. The AIRI method was later validated on observations from the Australian Square Kilometre Array Pathfinder (ASKAP; Wilber et al. 2023).

Two learned approaches for an interferometric-based imager named segmented planar imaging detector for electro-optical reconnaissance (SPIDER) were proposed by Mars et al. (2023). The first approach consists of a learned post-processing step from the dirty reconstruction based on a convolutional U-Net architecture (Ronneberger, Fischer & Brox 2015). The second approach consists of a learned multiscale iterative method coined GU-Net, which incorporates the measurement operator to include measurement information at the different steps and scales of the method. GU-Net is more efficient than standard unrolling methods due to its multiscale nature. The numerical results show an improved reconstruction quality and a faster convergence than proximal optimization-based methods. In the following work (Mars et al. 2024), the GU-Net was applied to the RI imaging problem. The variations of the |$uv$|-coverage are handled by training the neural network on a broad distribution of simulated |$uv$|-coverages and subsequently fine-tuning the network for a specific sampling distribution.

Aghabiglou et al. (2023) recently proposed a series of DNNs that combines notions of PnP algorithms and unrolled optimization methods (Adler & Öktem 2018; Monga, Li & Eldar 2019). Each DNN is trained to transform a back-projected residual into an image residual, thus ideally improving the reconstruction of the previous iteration. The results show a significant speed-up with respect to AIRI or SARA-based methods while maintaining a similar reconstruction quality. Other recent approaches based on deep neural networks include: Wang et al. (2023) who proposed a denoising diffusion probabilistic model conditioned on the visibilities and the dirty reconstruction; and Schmidt et al. (2022) who proposed a convolutional neural network based on residual blocks that intend to inpaint the measurements, or recover the entire |$uv$| plane from an incomplete coverage.

2.5 Bayesian framework

Bayesian inference provides a principled statistical framework to solve the inverse problem in equation (3) with statistical guarantees. This framework builds upon Bayes’ famous theorem,

Bayes’ theorem relates the posterior distribution to the likelihood and prior terms that are the main constituents of a Bayesian model. The likelihood is associated with the data-fidelity term depending on the observational model and the noise statistics. The prior models expected properties of the solution |$\boldsymbol {x}$|⁠, for example, smoothness, and piecewise regularity. This prior knowledge regularizes the estimation problem.

The term in the denominator, commonly known as the Bayesian evidence, does not depend on |$\boldsymbol {x}$| as we are marginalizing over that variable, and it describes the likelihood of the observed data based on the modelling assumptions. The Bayesian evidence is crucial for making Bayesian model comparison (Robert 2007), which provides us with a consistent way to compare models. Such high dimensional integrals can be effectively estimated by, for example, nested sampling techniques (Skilling 2006; Ashton et al. 2022), or the recently introduced learned harmonic mean estimator (McEwen et al. 2021; Spurio Mancini et al. 2022; Polanska et al. 2023). Recent developments have focused on nested sampling to compute the model evidence in high-dimensional imaging problems with sparsity-based handcrafted priors (Cai, McEwen & Pereyra 2022) and deep learning-based priors (McEwen et al. 2023). Carrying out model selection is out of the scope of this work.

Under the Bayesian framework, we have the posterior distribution |$p(\boldsymbol {x}|\boldsymbol {y})$| which assigns a probability to each possible solution |$\boldsymbol {x}$| given some observations |$\boldsymbol {y}$| and a model |$\mathcal {M}$| consisting of the likelihood and prior terms. In imaging settings, explaining the information contained in the posterior distribution is not trivial due to its high-dimensional nature. The posterior distribution is generally characterized by samples computed by MCMC sampling. Efficiently sampling from high-dimensional posterior distributions is a current research topic, see e.g. Klatzer et al. (2023). Once |$p(\boldsymbol {x}|\boldsymbol {y})$| is defined, we can say that the reconstruction method will be a point estimator of the posterior that will provide us with |$\hat{\boldsymbol {x}}$|⁠. There are several choices for point estimators (Robert 2007; Arridge et al. 2019), each with advantages and drawbacks. Some examples are a sample from the posterior, |$\hat{\boldsymbol {x}} \sim p(\boldsymbol {x}|\boldsymbol {y})$|⁠, the maximum-a-posteriori estimator, |$\hat{\boldsymbol {x}} = \rm{argmax}_{\boldsymbol {x}} p(\boldsymbol {x}|\boldsymbol {y})$|⁠, or the posterior mean, |$\hat{\boldsymbol {x}} = \mathbb {E}[\boldsymbol {x}|\boldsymbol {y}]$|⁠.

The posterior also provides us with consistent ways of quantifying the uncertainty of the chosen point estimate or reconstruction (Robert 2007). For example, one way to represent uncertainty is to compute the posterior standard deviation. The pixels with a higher standard deviation are less constrained by the data and the prior allowing for more significant fluctuations.

One of the most significant drawbacks of Bayesian imaging methods is that they are known to be computationally expensive, even if there is a continuous effort targetting the scalability of these methods (Pereyra 2016; Durmus, Moulines & Pereyra 2018; Pereyra, Mieles & Zygalakis 2020; Pereyra, Vargas-Mieles & Zygalakis 2022; Klatzer et al. 2023).

2.5.1 Maximum-a-posteriori estimation

The MAP estimator is particularly interesting in high-dimensional problems like RI imaging as its formulation allows us to bypass the need for sampling from the posterior. Consequently, its computational footprint is significantly reduced. The likelihood and prior terms can be rewritten as |$p(\boldsymbol {y}|\boldsymbol {x}) = \exp [-f(\boldsymbol {x},\boldsymbol {y})]$| and |$p(\boldsymbol {x}) = \exp [-g(\boldsymbol {x})]$|⁠, respectively. The functions f and g are the likelihood and prior potentials. Using Bayes’ theorem in equation (7), we can rewrite the MAP estimation as follows

The previous optimization problem can be reformulated using the monotonicity of the logarithm as follows

One advantage of the MAP estimator is that equation (9) can be tackled efficiently with optimization algorithms. We refer the reader to Pereyra (2019) for a deeper analysis of MAP estimation.

Coming back to the RI imaging inverse problem from equation (3), we can define a (white) Gaussian likelihood,

and a sparsity-inducing Laplace-type prior defined as

Upon substitution of equations (10) and (11) into equation (9), the MAP optimization problem coincides with the one in equation (6). Therefore, the MAP reconstruction, |$\hat{\boldsymbol {x}}_{\text{MAP}}$|⁠, matches |$\hat{\boldsymbol {x}}$| from equation (6). Hence, sparsity-based approaches are MAP estimations with a prior based on the sparsity-promoting |$\ell _1$| norm in a given dictionary, e.g. wavelets.

2.5.2 Uncertainty quantification: more than a point estimate

Computing a good reconstruction for an inverse problem in the form of equation (3) can itself be challenging. Moreover, the reconstruction is often insufficient for many scientific applications that require further quantification of the result. This demand opens the door to uncertainty quantification, which provides more than a point estimate. The Bayesian framework provides us with formidable tools to do uncertainty quantification. For example, if we choose the MAP estimator as our reconstruction following the model in Section 2.5.1, we obtain the same reconstruction as in Section 2.2, which is the solution of equation (6). However, with the Bayesian framework, we can sample from the posterior and estimate the posterior standard deviation, perform a Bayesian hypothesis test of some image structure (Cai et al. 2018a; Price et al. 2021b), or compute other pixel-wise uncertainty measurements like local credible intervals (LCI; Cai et al. 2018a; Price et al. 2019).

2.5.3 Bayesian inference via MCMC sampling

Recent developments (Durmus, Moulines & Pereyra 2022) have considerably reduced the computational complexity of sampling high-dimensional posterior distributions in imaging inverse problems. Proximal MCMC sampling algorithms (Pereyra 2016; Durmus et al. 2018) extend the class of posterior distributions that can be studied by allowing the use of non-smooth terms. Sparse regularizers have been widely used in RI imaging (Carrillo et al. 2012, 2014; Pratley et al. 2018; Cai et al. 2018a), and are usually enforced through a non-smooth |$\ell _1$| term.

Let us note |$\pi$| the target probability distribution that we are interested in sampling from, which in our case will be the posterior |$p(\boldsymbol {x}|\boldsymbol {y})$|⁠. We consider a Langevin diffusion process on |$\mathbb {R}^{N}$| such that its stationary distribution is |$\pi$|⁠. Assuming that |$\pi \in \mathcal {C}^{1}$| with Lipschitz gradients, we write the Langevin diffusion as the following stochastic process

where W is a N-dimensional Brownian motion. A usual discrete-time approximation of the Langevin diffusion consists of a forward Euler approximation with a step size |$\delta$|⁠, known as the Euler-Maruyama approximation (Kloeden & Platen 2011). The resulting algorithm is known as the unadjusted Langevin algorithm (ULA),

where |$\boldsymbol {w}^{(m+1)} \sim \mathcal {N}(0,\boldsymbol {I}_{N})$| is the discrete counterpart of |$W(t)$|⁠. The ULA-based Markov chain converges to |$\pi$| with an asymptotic bias due to discretization. The bias can be accounted for with a subsequent Metropolis-Hasting (MH) accept-reject step. Adding the MH step corrects the bias but increases the algorithm’s computational complexity. The ULA algorithm with the subsequent MH step is known as the Metropolis-adjusted Langevin algorithm (MALA).

The ULA algorithm requires the target density |$\pi$| to be continuously differentiable with Lipschitz gradients. Let us now consider |$\pi (\boldsymbol {x}) \propto \exp [ -f(\boldsymbol {x}) - g(\boldsymbol {x}) ]$|⁠, where |$f \in \mathcal {C}^{1}$| with Lipschitz gradient and g is non-smooth but is a lower semicontinuous convex function that admits a proximal operator (Parikh & Boyd 2014). Proximal MCMC algorithms (Pereyra 2016) relax this assumption by approximating g, a non-smooth term in |$\pi$|⁠, by its Moreau-Yosida envelope |$g^{\gamma }$|⁠. The Moreau-Yosida approximation satisfies

where |$\gamma$| is the Moreau-Yosida approximation parameter and the proximal operator may or may not have a closed-form expression. Consequently, the non-smooth target density |$\pi$| is approximated by the smooth |$\pi _{\gamma }$|⁠, which replaces the g term with its Moreau-Yosida approximation |$g^{\gamma }$|⁠. The Markov chain targeting |$\pi _{\gamma }$| writes

and it is known as Moreau-Yosida regularized ULA (MYULA). If we add an MH step targetting the non-differentiable distribution |$\pi$|⁠, the MCMC algorithm is known as Proximal MALA (Px-MALA). The proximal MCMC algorithms previously mentioned can be further accelerated by replacing the Euler-Maruyama approximation with the more involved Runge-Kutta-Chebyshev approximation (Abdulle, Almuslimani & Vilmart 2018), giving rise to the SK-ROCK (Pereyra et al. 2020) algorithm.

Cai et al. (2018a) exploited the MYULA and Px-MALA algorithms to sample from the posterior in the RI imaging problem. The model is based on a Gaussian likelihood as in equation (10) and a sparsity promoting prior akin to equation (11). However, the framework can be used with more complex noise models (Melidonis et al. 2023), e.g. Poisson noise. In Cai et al. (2018a), the RI image reconstruction relies on the minimum mean squared error (MMSE) estimator based on the posterior mean, while in Cai et al. (2018b), the MAP is considered.

3 SCALABLE BAYESIAN DATA-DRIVEN IMAGING WITH UNCERTAINTY QUANTIFICATION

quantifAI,¹ a scalable Bayesian data-driven method with uncertainty quantification is motivated by three principles:

1. Scalability: The RI imaging inverse problem demands scalability for a method to be useful in real astronomical data scenarios such as SKA. The most time-consuming operation is evaluating the measurement operator |$\Phi$| in the likelihood function. It is, therefore, essential to minimize the number of likelihood evaluations. For these reasons, we limit ourselves to the MAP estimator for our reconstruction corresponding to the solution of a convex optimization problem which converges quickly. We need to avoid sampling-based approaches as they are prohibitively expensive in terms of computations.

2. High-quality reconstructions: To improve the quality of our reconstruction, we consider data-driven or learned priors that can better encode the expected image structures. In Section 2.4, we have already seen that data-driven approaches can better represent complex imaging priors and provide reconstructions superior to handcrafted priors, such as sparsity-promoting priors based on wavelet dictionaries.

3. Uncertainty quantification: There are many ways to quantify uncertainty based on sampling the posterior distribution. However, using sampling-based methods is prohibitively expensive, and one of our key criteria is computational scalability. Therefore, we need to restrict ourselves to log-concave posteriors, which is equivalent to saying that the addition of our potentials |$f + g$| has to be convex, and to explicit potentials. As we will later describe in more detail in Section 4, the first restriction enables the use of efficient methods relying on the concentration of probability for high-dimensional log-concave distribution (Pereyra 2017). Consequently, we can use approximate posterior information bypassing sampling methods. These methods are orders of magnitude faster resulting in a scalable Bayesian UQ method. In a nutshell, we require the posterior potential to be convex and explicit for scalable UQ. The likelihood is typically convex for RI imaging problems so we will enforce the prior potential g to be convex and explicit. The requirement of explicit potentials will be explained in Section 4.

We continue by introducing the data-driven convex regularizers and the optimization algorithm used to compute the MAP estimation for the proposed method.

3.1 Learned convex regularizers

As stated before, we need an expressive regularizer that is convex and has an explicit potential. Modern regularizers based on deep neural networks, like convolutional neural networks, used in RI imaging reconstruction methods satisfy neither of the two constraints. This last constraint, i.e. with an explicit potential required by the UQ approach, excludes a range of denoisers whose potentials are defined implicitly. PnP approaches (Terris et al. 2022) only require the denoising of the image without explicitly computing the regularization potential. For example, a typical iteration from a PnP algorithm writes

where D is the denoiser, f is the data-fidelity term, |$\gamma$| is the step size, and k is the iteration number. The algorithm’s convergence can be assured if D and the stepsize satisfies some conditions (Ryu et al. 2019; Pesquet et al. 2021). Even if the denoiser D is convex, we cannot use it for our approach as we must evaluate the potential.

Mukherjee et al. (2020) proposed a learned convex regularizer parametrized by the architecture of a deep input-convex neural network (ICNN; Amos, Xu & Zico Kolter 2016), which is convex by construction. The training of the regularizer is done with an adversarial framework introduced by Lunz, Öktem & Schönlieb (2018).

Very recently, a learnable convex-ridge regularizer neural network (CRR-NN;² Goujon et al. 2023b) has been proposed, which comes with the required properties of being convex and having an explicit potential. In addition, the model focuses on being reliable and interpretable while still being expressive enough to provide excellent reconstruction quality. The CRR-NN regularizer, |$R_{\boldsymbol {\theta }}$|⁠, has the form

where |$\boldsymbol {h}_{n}$| are learnable 2D convolution kernels, |$( \boldsymbol {h}_{i} \ast \boldsymbol {x} )[k]$| denotes the k-th pixel of the resulting convolution, |$N_{\text{C}}$| is the number of channels or kernels, |$\psi _i: \mathbb {R} \mapsto \mathbb {R}$| are learnable non-linear convex profile functions with a Lipschitz continuous derivative, i.e. |$\psi _i \in C^{1,1}(\mathbb {R})$|⁠, and |$\boldsymbol {\theta }$| in |$R_{\boldsymbol {\theta }}$| represents all learnable parameters. The convexity constraint on the learnable activation functions, |$\psi _i$|⁠, is enforced by making the pointwise |$\sigma _i: \mathbb {R} \rightarrow \mathbb {R}$| monotonically increasing, with |$\psi _{i}^{\prime } = \sigma _{i}$|⁠, where |$\sigma _{i} \in C^{0,1}_{\uparrow }(\mathbb {R})$|⁠, and |$C^{0,1}_{\uparrow }(\mathbb {R})$| is the set of scalar Lipschitz continuous and increasing functions on |$\mathbb {R}$|⁠. The |$\sigma _i$| functions are chosen as learnable linear splines, which have been shown to be more expressive than ReLU functions in Propositions 3.3 and 3.5 from Goujon et al. (2023b, section III.B). The main difference between a prior based on the CRR-NN and a wavelet dictionary is that the kernels (or filters) and the activation (or thresholding) functions are learnt in the first one. In the second one, they are fixed or handcrafted. We refer the reader to Goujon et al. (2023b, figs 5 and 6) for examples of learned kernels |$\boldsymbol {h}_{n}$| and activation functions for two trainings with two different noise levels. See Bohra et al. (2020) and Goujon et al. (2023b), for more information on learnable splines.

In the spirit of PnP approaches, the CRR-NN training is based on the denoising problem that reads

where |$\boldsymbol {y}$| is a noisy version of |$\boldsymbol {x}$|⁠, and |$\lambda$| is a parameter controlling the regularization strength. The denoising problem is addressed through the fixed point of the problem, which given the convexity assumptions, is unique. A gradient step of equation (19) reads

where |$\alpha$| is the stepsize. Convergence can be guaranteed if the stepsize satisfies |$\alpha \in (0, 2/(1 + \lambda \, \text{Lip}(\nabla R_{\boldsymbol {\theta }})))$|⁠, where |$\text{Lip}(\cdot )$| denotes the Lipschitz constant. By composing t gradient descent updates of equation (20), i.e. a t-fold composition, we obtain a multigradient step denoiser that we denote |$T_{R_{\boldsymbol {\theta }}, \lambda , \alpha }^{t}$| following the notation of Goujon et al. (2023b).

The denoising problem in equation (19) can be formulated as a fix point problem for the t-step denoiser |$T_{R_{\boldsymbol {\theta }}, \lambda , \alpha }^{t}$| as follows,

We build the CRR-NN training by penalizing the residual of the fix point problem in equation (21) with a loss function |$\mathcal {L}$|⁠, for a training set of pairs of noiseless and noisy images |$\lbrace \boldsymbol {x}^{(m)}, \boldsymbol {y}^{(m)}\rbrace _{m=1}^{M}$|⁠, and reads

After having trained the denoiser, we define our prior potential as

where we have dropped the |$\boldsymbol {\theta }_{t}^{\ast }, \lambda _{t}^{\ast }$| notation for |$\boldsymbol {\theta }, \lambda$| and added a scaling parameter, μ, to boost performance following Goujon et al. (2023b). For the optimization algorithm, we need the Lipschitz constant of the gradient of the potential in equation (23), which can be expressed as

which is calculated in Goujon et al. (2023b, Prop. IV.1), and |$\Sigma _{\infty } = \text{diag}(\Vert {\sigma }_{1}^{\prime } \Vert _{\infty }, \ldots , \Vert \sigma _{N_C}^{\prime } \Vert _{\infty } )$|⁠, and |$\boldsymbol {W} = [\boldsymbol {w}_{1} \cdots \boldsymbol {w}_{N_C}]^{T}$| where |$\boldsymbol {w}_{i}$| corresponds to the filter |$\boldsymbol {h}_{i}$|⁠: |$\boldsymbol {h}_{i} \ast \boldsymbol {x} = \boldsymbol {w}_{i}^{T} \boldsymbol {x}$|⁠.

3.2 Computing our reconstruction: the MAP

In our case, computing the MAP reduces to solving a convex optimization problem. Following equation (9), the optimization problem we address is the following one,

where in addition we include |$\iota _{\mathbb {R}^{N}}$|⁠, an indicator function enforcing the reconstructed image to be real. The proximal operator of the indicator function to a convex set is known and it amounts to a projection onto that convex set. In the last term of equation (25) the proximal operator of the reality constraint is the projection of the vector to the real number, which is written as |$\text{Re}(\cdot )$|⁠. We have assumed a (white) Gaussian likelihood and the prior term is based on a previously trained CRR-NN. The CRR-NN is smooth with Lipschitz continuous gradients. However, the non-smoothness of the reality enforcing constraint forces us to rely on proximal algorithms (Parikh & Boyd 2014) instead of an accelerated gradient descent method (Nesterov 2018). In this case, we use the FISTA algorithm (Beck & Teboulle 2009).

For the optimization, we need the gradient of the likelihood and prior terms

where, in our case, |$(\cdot )^{\dagger }$| is the complex conjugate transpose.

To ensure the algorithm’s convergence we use the stepsize |$\tau = 1/L$|⁠, where |$L = \text{Lip}(\nabla _{\boldsymbol {x}} f(\boldsymbol {x},\boldsymbol {y}) + \nabla g(\boldsymbol {x}))$|⁠. We can estimate a simple bound for the Lipschitz constant as follows

where we have exploited the result from equation (24), and |$\Vert \Phi ^{\dagger } \Phi \Vert$| denotes the spectral norm, which in the case of a linear operator coincides with its maximum singular value. In the simplified problem we are considering in Section 2.1 with gridded visibilities, we have that |$\Vert \Phi ^{\dagger } \Phi \Vert = 1$|⁠. If a more realistic linear operator should be considered, the maximum singular value could be computed iteratively via the power method (Golub & van Loan 2013).

We initialize the optimization with the naturally weighted dirty image, |$\boldsymbol {x}_{(0)} = \text{Re}(\Phi ^{\dagger }\boldsymbol {y})$|⁠. The optimization procedure is summarized in Algorithm 1. We optimize for a fixed number of iterations |$N_{\text{max}}$|⁠, or until a tolerance criterion of |$\xi$| is reached. The stepsize is computed using the bound from equation (28).

4 SCALABLE UNCERTAINTY QUANTIFICATION

Enforcing the posterior’s convexity and explicit potential opens the door to scalable UQ methodology that was unreachable otherwise. The restriction to log-concave posteriors is the price we pay to gain great scalability. Our approach is based on the work from Pereyra (2017), which exploits concentration phenomena occurring in high-dimensional log-concave posteriors. The Bayesian high-posterior-density region can be approximated in log-concave models as the posterior probability mass tends to concentrate in particular regions on the parameter space. The approximation requires the MAP estimation, |$\hat{\boldsymbol {x}}_{\text{MAP}}$|⁠, which we have already computed as it is the chosen point estimate for our reconstruction. This result allows us to estimate information from the posterior probability density function without MCMC sampling. In this section, we introduce the main result we exploit for UQ. We then describe the proposed scalable UQ methods and how to validate our results with Langevin-based MCMC sampling algorithms.

4.1 Highest posterior density regions

Let us define a posterior credible region with a credible level of |$100(1-\alpha ){{\ \rm per\ cent}}$| as a set |$C_{\alpha } \in \mathbb {R}^{N}$| satisfying

with |$\mathbb {1}_{C_{\alpha }}$| being being unity if |$\boldsymbol {x} \in C_{\alpha }$| and zero otherwise. There are many regions satisfying the previous equation. We will focus on the highest posterior density region (HPD), which is defined as

where f and g are the potentials of our likelihood and prior terms, and |$\gamma _{\alpha }$| is a constant that defines a level-set of the log-posterior such that equation (29) holds. The HPD region has the property of having minimum volume (Section 5.5; Robert 2007).

Our posterior |$p(\boldsymbol {x} | \boldsymbol {y}) = \exp [-f(\boldsymbol {x})-g(\boldsymbol {x})]/Z$| is log-concave on |$\mathbb {R}^{N}$|⁠, where Z is the Bayesian evidence. Then, following Pereyra (2017, Theorem 3.1), for any |$\alpha \in (4 \exp [(-N/3)], 1)$|⁠, the HPD region |$C_{\alpha }$| from equation (30) is contained in

where

with a positive constant |$\tau _{\alpha } = \sqrt{16 \log (3/\alpha )}$| independent of |$p(\boldsymbol {x} | \boldsymbol {y})$|⁠.

Theorem 3.2 in Pereyra (2017) gives the error analysis of the approximation, and we see that |$0 \le \hat{\gamma }_{\alpha } - \gamma _{\alpha } \le \tau _{\alpha }\sqrt{N}+N$|⁠. Therefore, the error upper bound grows linearly with the dimension N, making |$\hat{C}_{\alpha }$| a stable approximation of |$C_{\alpha }$|⁠. The error lower bound along with the convexity of |$f+g$| guarantees the inclusion |$C \subseteq \hat{C}$| and consequently the resulting approximation is a conservative one where the errors are overestimated.

After showing the main result allowing us to do UQ bypassing posterior sampling methods, it is clear from where the constraints of the prior come. The convexity is required to guarantee a log-concave posterior, as the likelihood potential is convex. The prior potential g needs to be explicit to compute the approximate HPD region using equation (32).

4.2 MAP-based UQ methods

We now explore different scalable UQ schemes based on the fast approximate implicit representation of the HPD region. For all the methods presented, we assume that we have already computed the |$\hat{\boldsymbol {x}}_{\text{MAP}}$| estimation and the approximated HPD region threshold, |$\hat{\gamma }_{\alpha }$|⁠.

4.2.1 Bayesian hypothesis testing of structure

A useful UQ tool is to perform a knock-out hypothesis test to asses if a surrogate image still belongs to the HPD region (Cai et al. 2018a, b; Price et al. 2021b). First, the surrogate image |$\boldsymbol {x}_{\text{sgt}}$| is constructed by modifying the reconstruction, |$\hat{\boldsymbol {x}}_{\text{MAP}}$|⁠. Then, it suffices to check if

If the inequality is satisfied, we cannot draw conclusions on the test we made, as |$\boldsymbol {x}_{\text{sgt}}$| still belongs to the HPD region. However, if the inequality does not hold, we can conclude that |$\boldsymbol {x}_{\text{sgt}}$| is out from the HDP region with a |$100(1-\alpha ){{\ \rm per\ cent}}$| confidence level.

This test can answer a variety of questions about our reconstructed image. One example is to interrogate some structure in the image to see if it is a reconstruction artefact or is physically motivated. For this question, the surrogate image would be composed of an image with the region of interest artificially inpainted with surrounding information. We need to take the inpainted image as our surrogate and evaluate equation (33) to see if the test is conclusive.

The image inpainting algorithm is built similarly as in Cai et al. (2018b) but adapted to the CRR-NN-based prior. We start by selecting a region of interest |$\Omega _{\text{D}}$|⁠, which is a subset of (typically contiguous) pixels from the image, where |$\Omega _{\text{D}} \subseteq \Omega$|⁠, where |$\Omega$| denotes the set of all the image pixels. The region |$\Omega _{\text{D}}$| will be inpainted with background information. We then inpaint this region iteratively minimizing |$R_{\boldsymbol {\theta }}$| based on the following scheme

where |$\mathbb {1}$| are indicator functions, and |$\mathbb {1}_{\Omega - \Omega _{\text{D}}}$| is a shorthand for |$\mathbb {1}_{\Omega } - \mathbb {1}_{\Omega _{\text{D}}}$|⁠. We carry out a gradient step with the CRR-NN on the surrogate image and only update the region of interest. The hyperparameters, |$\alpha$|⁠, |$\lambda$|⁠, and μ are set as in Algorithm 1.

Alternatively, Repetti, Pereyra & Wiaux (2019) presented a more sophisticated method to perform hypothesis testing of structure, which also exploits the approximations in equations (31)–(32). The method is dubbed Bayesian uncertainty quantification by optimization (BUQO), and to answer the hypothesis test, it aims to study the intersection of two sets. The first one is defined in equation (31) that corresponds to the MAP estimate. The second one describes the set of feasible inpainted images given a region of interest and a set of constraints of desired properties. If the set intersection is empty, the structure of interest is considered present in the image with confidence |$\alpha$| from equation (31). Tang & Repetti (2023) proposed an extension of the BUQO method to inpaint with data-driven models. However, these methods involve solving an expensive optimization problem that does not scale with the high-dimensional settings we are considering in this work.

Another example is to interrogate the reconstruction to see if the fine structure of the image is physical or likely an artefact. To construct the surrogate image we convolve the region of interest, |$\Omega _{\text{D}}$|⁠, with a Gaussian smoothing kernel |$G(0,\Sigma )$|⁠,

where |$\ast$| denotes the 2D convolution operation and test equation (33).

4.2.2 Local credible intervals

Local credible intervals (LCIs) provide a tool to quantify spatial uncertainty per pixel at different resolutions. The LCIs are interpreted as Bayesian error bars for each pixel or superpixel, where with superpixel, we refer to a group of contiguous pixels. Cai et al. (2018a) computed LCIs using MCMC methods and then extended it in Cai et al. (2018b) to compute them based on the approximated HPD region based on the MAP. Price et al. (2019) also exploited MAP-based LCIs in another imaging inverse problem, mass-mapping, for weak gravitational lensing convergence reconstruction.

Let us write |$\Omega = \lbrace \Omega _i \rbrace _{i=1}^{M}$| the set of superpixels that partition the domain of |$\boldsymbol {x}$|⁠. This partition is such that |$\Omega _i \cap \Omega _j = \emptyset , \forall i \ne j$| and |$\Omega = \cup _i \Omega _i$|⁠. We denote the indicator of the superpixel |$\Omega _i$| as |$\boldsymbol {\zeta }_{\Omega _i}$|⁠, that is one if the pixel belongs to the superpixel |$\Omega _i$| and zero otherwise. The use of smaller or bigger superpixel sizes, i.e. |$\Vert \boldsymbol {\zeta }_{\Omega _i} \Vert _{0}$|⁠, allows us to visualize the LCIs at different scales. The calculation of the LCIs is based on computing an upper and lower bound for each superpixel. Each bound is defined by the constant value we need to add or subtract to the superpixel region so that the modified image exits the approximate HPD credible region |$\hat{C}_{\alpha }$|⁠. In other words, we compute the values that saturate the HPD region for each superpixel.

We can isolate the superpixel region by defining the following surrogate image

where |$\bar{\boldsymbol {x}}_{\text{MAP}, \Omega _i}$| corresponds to the mean value of the pixels in the superpixel |$\Omega _i$|⁠, and |$\xi \in \mathbb {R}$|⁠. We vary the superpixel value from its mean by a uniform value |$\xi$|⁠. The bounds for a superpixel |$\Omega _i$| are computed as

where we use the threshold |$\hat{\gamma }_{\alpha }$| defined in equation (32). The calculation of each bound can be recast as a problem of finding the zero of a function. The class of root-finding algorithms is well adapted for this root-finding problem, and, in practice, we use the bisection method (Burden & Faires 1989). Price, Pratley & McEwen (2021a) proposed a faster way to compute the superpixel bounds by exploiting linearity that we could use to further accelerate the computation of |$\xi _{+, \Omega _i}$| and |$\xi _{-, \Omega _i}$|⁠.

Once the bounds have been computed, we can collate the results for all superpixels and use the length of the LCIs to visualize the reconstruction uncertainty. The length of the LCI for superpixel |$\Omega _i$| is defined as |$l_{i} = \xi _{+, \Omega _i} - \xi _{-, \Omega _i}$|⁠, which we can visualize as an uncertainty image

The choice of using the mean on |$\bar{\boldsymbol {x}}_{\text{MAP}, \Omega _i}$| for the region of the superpixel that will be studied constitutes a deviation from the original MAP reconstruction. We find it more physical to move the averaged superpixel rather than moving the original pixels belonging to the superpixel by a constant value. This choice constitutes another approximation to the proposed scheme that has already approximated the HPD region in equation (31). Using superpixels allows us to gain sensibility and computing time at the expense of lowering the resolution of the LCI map, which can be a sensible trade-off for very large images.

We will later validate the computed LCIs using the posterior samples obtained from computing the posterior standard deviation at different superpixel sizes. The method requires turning each posterior sample into an image with M superpixels. We set the value of the superpixel to the mean of the values of belonging pixels.

4.2.3 Fast pixel uncertainty quantification at different scales

The MAP-based LCIs described in the previous section are orders of magnitude faster than their MCMC-based counterparts (Cai et al. 2018a, b). Nevertheless, to compute the LCIs, we still need to evaluate the likelihood potential, f, several times for each superpixel. As mentioned, evaluating the likelihood potential is by far the most time-consuming operation. The fact that we need to evaluate f several times for each subpixel might make the LCIs impractical for SKA-scale problems.

To overcome this issue, we propose a new approach relying on wavelet decomposition of the MAP reconstruction that reads

with a wavelet dictionary |$\boldsymbol {\Psi }$|⁠. We define the hard thresholding operator for |$\boldsymbol {a} \in \mathbb {C}^{L}$| with a threshold of |$\xi _{\text{th}}$|⁠,

as the point-wise application of the following hard-thresholding function

Let |$\hat{\xi }_{\text{th}}$| be the thresholded value for which the thresholded image saturate the HPD region as follows

We can compute the threshold bound with a root-finding method, as was the case for the LCIs. The advantage of this formulation is that we are solving only one root-finding problem as opposed to one per superpixel in the LCIs calculation. This change considerably reduces the number of likelihood evaluations and, therefore, the computational complexity of the UQ method.

By computing the difference between the MAP, |$\hat{\boldsymbol {x}}_{\text{MAP}}$|⁠, and the thresholded surrogate, |$\hat{\boldsymbol {x}}_{\text{MAP}, \, \hat{\xi }_{\text{th}}}$|⁠, we obtain an estimation of the solution’s uncertainty and this can give us information about possible errors in the MAP. Furthermore, we can compare the MAP and the thresholded surrogate image to estimate errors as a function of scale, thus exposing the different structures of our reconstruction.

Let us consider our wavelet transformation as a multiscale transform of level |$J+1$| (Mallat 2008; Starck, Murtagh & Fadili 2010). We can rewrite equation (40) showcasing the multiscale coefficients as follows

where |$\hat{\boldsymbol {a}}_{\text{MAP}, l}$| are the coefficients corresponding to the l-th level of decomposition, and the zeroth level corresponds to the coarse scale. We can now build thresholded surrogate images at different scales by replacing the MAP wavelet coefficients at level l from equation (44) with the coefficients of the thresholded surrogate image |$\hat{\boldsymbol {x}}_{\text{MAP}, \, \hat{\xi }_{\text{th}}}$|⁠. Let us write the thresholded surrogate image at level j as follows

where |$\hat{\boldsymbol {a}}_{\text{MAP}, \, \hat{\xi }_{\text{th}}, \, j}$| corresponds to the wavelet coefficients of the thresholded surrogate image |$\hat{\boldsymbol {x}}_{\text{MAP}, \, \hat{\xi }_{\text{th}}}$| at level j. The errors at level j can be approximated by the difference between |$\hat{\boldsymbol {x}}_{\text{MAP}}$| and |$\hat{\boldsymbol {x}}_{\text{MAP}, \, \hat{\xi }_{\text{th}}, \, j}$|⁠.

There are two main advantages of this approach to pixel-based UQ with respect to the LCIs described in Section 4.2.2. The first one is the reduced computational complexity, as we only need to solve a single root-finding problem, significantly reducing the number of likelihood evaluations. The second is that when we saturate the HPD region, we consider the entire image simultaneously. In the LCI computation, we only change a small group of pixels until it saturates the HDP region that corresponds to the entire image. This behaviour can be problematic as, for example, the LCI error bounds will be larger if the size of the image grows and the superpixel size is kept the same, which is an undesirable property. Consequently, the predicted errors from the new pixel UQ approach are faster to compute and considerably tighter than the LCIs.

4.3 MCMC sampling and UQ validation

As stated before, MCMC sampling is not yet scalable to the high dimensions of the RI imaging problems we target. However, sampling is still helpful in validating the UQ approaches of this paper. If we sample from the posterior distribution, we can compute the posterior standard deviation, providing a visual representation of the posterior, including the learned convex regularizer. Sampling from the posterior also allows us to compare the MAP estimator with another widely known estimator, the posterior mean (Arridge et al. 2019), which coincides with the minimum mean squared error (MMSE) estimator under some assumptions.

The log-posterior distribution of the quantifAI model with the CRR-NN reads

with the first two terms belonging to |$\mathcal {C}^{1}$| with Lipschitz gradients, we do not need to use any approximation, e.g. the MY envelope, to sample from it. The reality constraint is enforced directly when evaluating the gradient of the log-likelihood. The Langevin diffusion sampling algorithms reviewed in Section 2.5.3 require the gradient of the log-posterior distribution, which have been computed in equations (26) and (27). In practice, we will use the SK-ROCK algorithm (Pereyra et al. 2020) as it provides a faster convergence than the ULA algorithm. We omit the subsequent MH step to accelerate the calculations motivated by the consistent results from Cai et al. (2018a).

The log-posterior distribution of the analysis formulation of the model from Cai et al. (2018a) with a wavelet-based sparsity enforcing prior reads

which includes the non-smooth prior term with the |$\ell _1$| norm, and the reality constraint which we again apply to the gradient of the log-likelihood. We resort to the MY envelope |$\gamma$|-approximation of the sparsity-inducing prior term as shown in equation (14). The proximal operator of the prior term has a closed-form solution that reads

where |$\boldsymbol {a} = \Psi ^{\dagger } \boldsymbol {x}$| and we have applied element-wise the soft-thresholding function

The threshold |$\beta _{\text{th}}$| used in practice is |$\lambda \gamma$|⁠, the product of the regularization strength and the parameter of the MY approximation. See Cai et al. (2018a) for more details on sampling the model with a wavelet-based regularization. In practice, we again rely on the SK-ROCK algorithm for sampling and avoid using an MH step for the reasons mentioned above.

5 EXPERIMENTAL RESULTS

In this section, we demonstrate the quantifAI model against the wavelet-based model presented in Cai et al. (2018a, b) as it is one of the few RI imaging methods providing UQ capabilities. We use a simulated setup with real reconstructed RI images considered as the ground truth. We focus on the UQ capabilities of the methods, while also considering reconstruction performance.

5.1 Data set

The base images used in our experiment are the ones from Cai et al. (2018a): (i) the H i region of the M31 galaxy in Fig. 1(a), (ii) the W28 supernova remnant in Fig. 2(a), (iii) the 3C288 radio galaxy in Fig. 3(a), and (iv) the Cygnus A radio galaxy in Fig. 4(a). All the images are |$256 \times 256$| pixels, except for the Cygnus A galaxy, which is |$256 \times 512$|⁠. As the ground truth images are reconstructed from real observations, some minor original residuals and backgrounds are more noticeable in the log scale images; for example, see Fig. 3(b). The ground truth images’ values have been normalized to a unitless range between 0 and 1, and therefore, the colour bars in the reconstruction figures follow this range.

The previous images correspond to |$\boldsymbol {x}$| in our observational model from equation (3). The complex Gaussian noise |$\boldsymbol {n} \in \mathbb {C}^{M}$| is composed of independent and identically distributed (i.i.d.) samples. Each sample is simulated following a complex Gaussian distribution, |$n_i \sim \mathcal {N}_{\mathcal {C}}(0, \sigma ^2)$|⁠, which implies that |$\text{Re}(\boldsymbol {n}), \text{Im}(\boldsymbol {n}) \sim \mathcal {N}(0, \sigma /\sqrt{2})$| (Tse & Viswanath 2005). The noise is set such that we get a specific input signal-to-noise ratio (ISNR) on each image. For all the experiments, we set the ISNR to 30 dB, and the noise standard deviation is computed as follows

To mimic the |$uv$| plane coverage, we reuse the Fourier mask from Cai et al. (2018a, fig. 2) and use it to generate the visibilities from |$\boldsymbol {y}$|⁠. The variable sampling density profile was taken from Puy, Vandergheynst & Wiaux (2011) and represents a |$10{{\ \rm per\ cent}}$| coverage of the Fourier plane. In the experiments in the current section, we work with gridded visibilities where we have around |$1.3 \times 10^{4}$| visibilities for Cygnus A and |$6.5 \times 10^{3}$| visibilities for the rest of the images. The validation of the UQ techniques through MCMC sampling requires a large amount of iterations. The use of gridded visibilities allows us to base the forward operator |$\boldsymbol {\Phi }$| from equation (3) on the FFT (Cooley & Tukey 1965), helping us to alleviate the computational burden of the validation. Section 6 presents results with ungridded visibilities.

5.2 Models and experiment settings

5.2.1 RI imaging models

The CRR-NN in the quantifAI model is a pre-trained model with |$t=5$|⁠, Gaussian white noise with standard deviation |$\sigma =5$|⁠, and parameters |$\mu =20$|⁠, |$\lambda =5 \times 10^{4}$|⁠. The model was trained on a set of natural images from the BSD500 data set (Arbeláez et al. 2011) containing images of landscapes, people, animals, and objects among others. The images are scaled to the [0,255] range, using |$\ell _1$| norm as the loss function in equation (22) with the Adam optimizer (Kingma & Ba 2017). The training parameters followed Goujon et al. (2023b, section VI.A).

The wavelet dictionary |$\Psi$| used in the wavelet-based model is composed of Daubechies 8 wavelets (Daubechies 1992) with a multiresolution level |$J=4$| following the setup from Cai et al. (2018a, b). The regularization parameter |$\lambda _{\text{WAV}}$| was set to |$1 \times 10^{2}$|⁠.

The regularization strengths of both models, |$\lambda$| and |$\lambda _{\text{WAV}}$|⁠, were set to maximize the MAP reconstruction SNR using a grid search. We observed that quantifAI’s reconstruction SNR is significantly more robust with respect to the choice of regularization strength than the wavelet-based models.

5.2.2 Optimization settings

For quantifAI, we used the optimization algorithm shown in Algorithm 1. The wavelet-based model also requires a proximal algorithm due to its non-smooth component and to provide a fair comparison we used the FISTA algorithm (Beck & Teboulle 2009) presented with more detail in Appendix  A. In these experiments, we assumed that the noise level |$\sigma$| is known. If the noise level is unknown, it may be estimated by standard techniques (Price et al. 2021b). Both algorithms’ tolerance criterion |$\xi$| was set to |$10^{-5}$|⁠, and the total number of iterations to |$1.5 \times 10^{4}$|⁠. Nevertheless, both optimization algorithms converged before the total number of iterations was reached.

5.2.3 MCMC sampling settings

We generate |$5 \times 10^{4}$| samples from each posterior distribution, with |$5 \times 10^{4}$| burn-in iterations and a thinning factor of 10. The burn-in iterations consist of generating several samples that will be discarded so that the chain passes its transient period. The thinning factor corresponds to the number of samples that need to be generated between two samples so that they can be considered independently drawn from the target distribution. The sampling algorithm produced a total of |$5.5 \times 10^{5}$| samples for each model. We have set to 10 the number of stages for the SK-ROCK algorithm (Pereyra et al. 2020), which is one of its main hyperparameters. The sampling of the posterior probability distributions is used as a validation, and therefore we set the sampling parameters focusing on good reconstructions and posterior samples rather than speed.

The wavelet-based model requires the MY envelope approximation to guarantee the chain’s convergence, as described in Section 2.5.3 and Section 4.3. The MY approximation parameter |$\gamma$| was set to the inverse of the likelihood gradient’s Lipschitz constant, c.f. the first term of equation (28).

The choice of the step sizes is critical to ensure the chains’ convergence to the target distribution in a reasonable amount of time. The step size is chosen as a function of each posterior gradient’s Lipschitz constant. The step sizes |$\delta _{\rm {\small Q}}$| and |$\delta _{\text{W}}$|⁠, corresponding to the quantifAI and wavelet-based models, respectively, are computed as follows

where the Lipschitz constant bounds are shown in equation (28), and |$\kappa _{\rm {\small Q}}$| and |$\kappa _{\text{W}}$|⁠, are two positive constants smaller than one, here set to 0.98. We have followed the advise from Durmus et al. (2018) and Cai et al. (2018a) to set the sampling parameters.

5.2.4 UQ settings

We set |$\alpha = 0.01$| in all the UQ methods, so the confidence level is |$99{{\ \rm per\ cent}}$|⁠. We used the bisection algorithm to compute the LCIs and the fast pixel UQ at different scales, with tolerance |$10^{-4}$| and maximum number of iterations 200, for both models. We used the same wavelet dictionary as in the wavelet-based model for the fast pixel UQ at different scales.

The inpainting algorithm uses the same stopping criterion as Algorithm 1. In this case, the tolerance is set to |$5 \times 10^{-6}$|⁠, and the total number of iterations to |$1.5 \times 10^{4}$|⁠. The CRR-NN used for the inpainting is the same one used in the quantifAI model.

The Gaussian blurring kernel |$G(0,\Sigma )$| from equation (35) is set using |$\Sigma =\sigma _{G}^{2} I_{2 \times 2}$|⁠, with |$\sigma _{G}$| being 3.5 pixels and a truncation radius of 7 pixels, giving a kernel |$G \in \mathbb {R}^{15 \times 15}$|⁠.

5.3 Image reconstruction

We present the RI image reconstructions of the four ground truth test images in Figs 1, 2, 3, and 4. In each figure, we compare the wavelet-based and quantifAI models, and we include the dirty reconstruction as a reference. The metric used to compare the RI image reconstruction is the SNR expressed in dB defined as follows

where |$\boldsymbol {x}_{\text{gt}}$| corresponds to the reference or ground truth, and |$\boldsymbol {x}$| to the estimation, and |$\Vert \cdot \Vert _{2}$| is the usual |$\ell _2$| norm.

The quantitative reconstruction performance results are presented in Table 1. The MAP reconstruction from quantifAI performs significantly better than the wavelet-based counterpart in every image from our data set. The performance gain lies between 1.9 and 12.7 dB, with an average gain of 7 dB. It is difficult to see the quantifAI improvements by eye when inspecting reconstructed images. However, when observing the errors in the fourth column, the improved quality of quantifAI’s reconstructions becomes evident. Shifting towards the sampling results, we observe a similar behaviour of the MMSE reconstruction in favour of quantifAI’s images. The MAP is considerably faster than the MMSE, relying on optimization rather than posterior sampling. Recall that the MMSE is built as averaging posterior samples. In addition, the MAP consistently provides improved reconstruction performance with respect to the MMSE.

The posterior standard deviation provides a qualitative way to validate the posterior model and its uncertainties. The comparison of the posterior standard deviation with the MAP reconstruction error shows a higher correlation for the quantifAI model than the wavelet-based model. In addition, the posterior standard deviation of quantifAI shows lower variance than its wavelet-based counterpart, which is in agreement with quantifAI’s smaller reconstruction error. For example, in image W28 in Fig. 2, we observe in sub-Fig. 2(j) that the posterior standard deviation value is large near the edges of the ground truth image. It is reassuring that quantifAI’s reconstruction error also shows the same behaviour.

The performance results showcase the expressive power of the CRR-NN-based prior even if the regularizer is constrained to be convex. The results also confirm the generalization power of the CRR-NN-based prior. Even if trained on natural images, the CRR-NN can provide remarkable reconstruction performances for astronomical images and meaningful posterior standard deviations.

The reconstructions using the wavelet-based prior model do exhibit some low-intensity artefacts in the Cygnus A and 3C288 images, as shown in Figs 4(h) and 3(h). These artefacts are due to the patterns in the ground truth images, which originate from real observations, and the thresholding of the orthogonal wavelet basis. Such patterns are absent in the M31 and W28 images because the noise was removed in a pre-processing step, as seen in Figs 1(b) and 2(b) in comparison to Figs 4(b) and 3(b). The regularization strength for the wavelet-prior model was selected to maximize the reconstruction SNR. This chosen value is lower than in Cai et al. (2018b), explaining the observed patterns and the finer details in our reconstructions. Employing a wavelet dictionary instead of an orthogonal wavelet basis and adding a positivity-enforcing constraint could mitigate the appearance of these artefacts.

5.4 Hypothesis testing of image structure

We start by carrying out hypothesis tests of image structure, which are the most scalable UQ techniques we will study. First, a surrogate image is created by modifying one region of interest. It only takes one further evaluation of the likelihood and prior potential to carry out the hypothesis test. The test helps us to quantitatively answer a scientific question with a |$100(1 - \alpha ){{\ \rm per\ cent}}$| confidence level. The scientific question targeted depends on the constructed surrogate image, and in this work, we consider two scenarios.

In the first scenario, we consider a particular structure in the reconstructed intensity image. We can query whether the structure’s origin is physical or not. For example, the structure could be a reconstruction artefact or a physical process. Fig. 5 shows this option, where we have analysed different regions of the four images. The first four inpainted regions correspond to physical structures, and the fifth region, i.e. region number 2 of image 3C288, does not correspond to a physical structure. The surrogate images are produced with an inpainting algorithm using quantifAI’s prior so that the inpainted region agrees with the prior.

The second scenario is to blur the finer structure in the reconstructed image and perform a hypothesis test to elucidate the question of whether the blurred structure is physical or not. The test is illustrated in Fig. 6. In this case, all four blurred images represent physical structures.

In both cases, we compare the hypothesis test using a MAP-based approach described in this work and a sampling-based approach for validation. In the MAP-based approach, we build the HPD region in equation (31) with the approximation in equation (32) and use the MAP estimation as our reconstruction. In the sampling-based approach, we use the MMSE as the reconstruction, i.e. the mean of the posterior samples, and compute the threshold defining the HPD region using the quantile function on the potentials of the posterior samples following Cai et al. (2018a, Section 5.2).

Table 2 presents the results for the inpainting hypothesis test, where the inpainted surrogates are shown in Fig. 5. The MAP- and sampling-based results are consistent in all the images studied, where the threshold computed with the posterior samples is slightly tighter than the MAP-based approximation. The hypothesis tests correctly classify the structure in images M31, W28, and 3C288, including the two cases of the latter image. The UQ methods cannot make a strong statistical statement about the structures in the Cygnus A image. In this image, where the inpainted region has a tiny physical structure, the potentials of the inpainted surrogate image rest close to the MAP and MMSE estimators. We include the hypothesis test results of the same inpainting experiment for the wavelet-based model in Appendix B1 to provide a comparison between the models. We used the wavelet prior to inpaint the region of interest to allow for a fair comparison. All results from the wavelet-based model are in agreement with quantifAI.

The results from the blurred surrogates of Fig. 6 are presented in Table 3. In all the images, the hypothesis test concludes that the blurred fine structure is physical as the potential falls out of the HPD region. The MAP- and sampling-based results are consistent with each other.

The different hypothesis tests have shown consistent results between the sampling-based and highly scalable MAP-based results. In addition, the results from the hypothesis tests are coherent between the quantifAI and wavelet-based model. We remark that the approach based on the MAP requires one further measurement operator evaluation to carry out the hypothesis test. The test provides a highly scalable way to answer scientific questions about the uncertainty of the RI imaging reconstructions.

5.5 Local credible intervals

We have exploited the approximation of the HPD region from Section 4.1 based on the MAP estimations and a credible level of |$99{{\ \rm per\ cent}}$|⁠. The approximate HPD regions were then used to compute the LCIs, whose lengths per pixel are visualized as an image, c.f. Fig. 7. The LCI lengths are displayed after subtracting the mean LCI length overall superpixels in the image, which is shown in the top left corner of the image. The UQ results for quantifAI are presented for two superpixel sizes, |$4 \times 4$| and |$8 \times 8$|⁠. We have omitted LCIs from the wavelet-based prior for conciseness. The posterior standard deviations at the two superpixel sizes are included for comparison with the significantly faster MAP-based UQ technique of the LCIs. We find a reasonable agreement between the structure in the LCI plots and the posterior standard deviation. For example, the 3C288 image with superpixel size |$8 \times 8$| yields tighter LCIs in the two elliptical regions and in the small connecting structure in the centre of the image. The corresponding posterior standard deviation is smaller in the aforementioned regions, which is expected as most of the observed signal concentrates there. The LCIs and the posterior standard deviation represent different quantiles, so we would not expect an exact agreement even without any approximation in the computation of the LCIs.

We observe, as expected, that the larger superpixels have tighter LCIs, as seen in the mean LCIs shown on the top left corner of the subfigures in Fig. 7. The reconstructions are naturally less uncertain on the larger scales due to the properties of our measurement operator, as the visibilities are generally concentrated towards the low frequencies. In addition, varying the value of a larger superpixel saturates the HPD region faster than for a small superpixel. We have also computed the LCIs for the superpixels of size |$16 \times 16$|⁠, which we have not included for conciseness. The corresponding mean LCI values are 0.20, 0.08, 0.24, and 0.07 for the images in the same order as in Fig. 7.

When comparing the mean value of the LCIs from the four reconstructions from Fig. 7 we notice that two of them, M31 and 3C288, have higher uncertainty than the rest. The higher the uncertainty, the larger the mean value of the LCI gets, as the superpixel values need larger changes before they saturate the HPD region. Image 3C288, with a superpixel size of |$4 \times 4$|⁠, is an example where the LCIs have saturated as the mean is close to unity;³ therefore, the LCI image’s detailed structure is lost due to the saturation. This saturation highlights the need for superpixel sizes to be selected appropriately, depending on the case at hand.

5.6 Fast pixel uncertainty quantification at different scales

The fast pixel UQ method results for the images M31 and W28 are reported in Fig. 8. We use the error between the MAP estimation and the ground truth image, i.e. true error, to validate the predicted uncertainty of the fast UQ method. The true error at different scales can be computed following equation (45),

where |$\boldsymbol {a}_{\text{GT}, l}$| are the wavelet decomposition coefficients of the ground truth image at multi-resolution level l. We have replaced the ground truth image’s wavelet coefficient at a single level with the coefficients from the MAP reconstruction.

We observe a good agreement between the predicted and ground truth errors at the different multi-resolution levels. There is an overestimation of the errors, which can come from two sources. First, the approximation of the HPD region is conservative, as it has been discussed in Pereyra (2017). Secondly, the MAP estimation is already missing some of the fine or high-frequency structures in the ground truth images. This fact can be seen in the MAP reconstruction errors in subFigs 1(h) and 2(h). The missing high-frequency structure is expected due to the properties of the measurement operator discussed in Section 2.1.

The structure from the chosen wavelet representation, |$\boldsymbol {\Psi }$|⁠, underpinning the UQ method can be observed in the predicted errors. This structure is visible mainly in the higher frequencies of the W28, where point sources are in the image. The wavelet structure should be taken into account when analysing the reconstruction errors.

This fast pixel UQ method allows us to approximate the reconstruction errors made at different scales for a fraction of the computational cost of the LCI pixel UQ method. The evaluations of the measurement operators are reduced by three orders of magnitude, resulting in an ultra-fast and truly scalable pixel UQ method. Furthermore, a single non-linear equation solve, e.g. root finding problem, of the new pixel UQ method suffices to predict the errors at all scales, while with LCIs, we are required to repeat the process for each superpixel size.

5.7 Computation time

The computation wall-clock time for both models, quantifAI and the wavelet-based, are summarized in Table 4. We include the results for only the W28 image in Tables 4 and 5 as they are representative of the other images. All the computations for both models were done using an Nvidia-A100 40GB GPU using pytorch (Paszke et al. 2019). We observe a lower computation time for the quantifAI model with respect to its wavelet-based counterpart. One reason is the lightweight CRR-NN model that quickly evaluates its gradient and potential. Note that the regularization strength has an impact on the number of iterations and it could be changed to favour a faster convergence. The regularization strength was chosen to optimize MAP reconstruction quality.

The results shown in Table 4 highlight the importance of relying on optimization-based rather than sampling-based reconstructions when focusing on the scalability of the method. There is a difference of approximately four orders of magnitude in the computation time of the MAP and the MMSE which relies on MCMC sampling techniques. Focusing on UQ, the posterior sampling is 60 times slower than the computation of the LCIs with |$8 \times 8$| superpixels and more than 37 500 times slower than the fast pixel UQ proposed in this work. The new fast pixel UQ provides an extremely rapid approach to providing pixel-based UQ, over 630 times faster than the |$8 \times 8$| LCIs.

The evaluation of the measurement operator is the most time-consuming operation in a real large-scale RI imaging scenario. If we target scalability, we need to monitor the number of measurement operator evaluations. Table 5 summarizes the number of measurement operator evaluations required for the UQ techniques. The results are only shown for the quantifAI model as they are representative of the wavelet-based model. As mentioned before, we note the reduction of evaluations between optimization and sampling-based reconstructions. We remark on the reduction in the number of evaluations for the UQ tasks, approximately three orders of magnitude between the sampling and the LCIs, and three subsequent orders of magnitude between the LCIs and the fast pixel UQ. These results make the fast pixel UQ six orders of magnitude faster than MCMC sampling. The MAP estimation for the CRR required 1082 measurement operator evaluations. However, the algorithm’s settings were chosen to maximize the reconstruction SNR. By modifying the regularization parameter of the CRR-based prior, we can reduce the number of evaluations by an order of magnitude. Recent developments in code parallelization for RI imaging reconstruction algorithms⁴ (Pratley et al. 2019a; Pratley & McEwen 2019) could be integrated to push the scalability of the method further.

6 EXPERIMENTAL RESULTS WITH UNGRIDDED VISIBILITIES

In the previous section, we have validated the proposed UQ methods from quantifAI in a simple setting with gridded visibilities. This choice lets us use the FFT algorithm for the forward model, which allows us to run MCMC algorithms for UQ validation in a sensible amount of time. In this section, we showcase quantifAI in an experiment using simulated visibility patterns from the MeerKAT radio telescope (Jonas & MeerKAT Team 2016). The main difference is that the visibility patterns are ungridded following a realistic distribution, which obliges us to rely on the Non-Uniform FFT (NUFFT) for the forward model. We remark that in this experiment, we are not addressing many of the challenges of dealing with real data, which are beyond the scope of this article.

6.1 Data set and experiment settings

We have simulated four single-frequency ungridded visibility patterns of differing synthesis times for MeerKAT. The start frequency is set to 1400 MHz with a channel width of 10 MHz. The pointing position has been randomly selected and set to (13h18m54.86s, |$-15$|d36m04.25s) in the J2000 reference, and it was maintained for the four generated data sets. We have used a publicly available code⁵ that is based on the casa simulation software (The CASA Team 2022). The synthesis times used are 1, 2, 4, and 8 h, with a constant integration time of 240 s. Each data set has a field of view of approximately 1 deg|$^2$|⁠. The number of visibilities of each data set is |$3\times 10^{4}$|⁠, |$6\times 10^{4}$|⁠, |$1.2\times 10^{5}$|⁠, and |$2.4\times 10^{5}$|⁠, correspondingly. Fig. 9 presents the simulated visibility patterns. We reuse the images described in Section 5.1 as the ground truth, keeping their original dimensions, which are unrealistically small to be representative of a MeerKAT observation grid size.

To cope with the ungridded visibilities, we have to change the forward operator |$\boldsymbol {\Phi }$| from equation (3) that before was based on the FFT. We rely on the pytorch-based NUFFT implementation from Muckley et al. (2020)⁶ based on Kaiser-Bessel gridding (Fessler & Sutton 2003). We have used the same images, Gaussian noise model presented in Section 5.1, and trained CRR-NN in quantifAI. We have reused the previously introduced hyperparameter values of quantifAI except for |$\lambda$|⁠, which we have tuned to maximize the reconstruction’s SNR for the four data sets to |$10^{4}$|⁠, |$1.4 \times 10^{4}$|⁠, |$1.9 \times 10^{4}$|⁠, and |$2.25 \times 10^{4}$|⁠, respectively.

6.2 Results

The reconstructions and the fast pixel UQ maps for the image M31 are shown in Fig. 10. Each column corresponds to each of the four data sets. The results for the other images are postponed to Appendix  C. Table 6 presents quantitative results regarding reconstruction SNR, measurement operator evaluations, and wall-clock computing time. The reconstructions present good quality in terms of SNR, which increases with longer synthesis times, which was expected as the Fourier coverage increases, as seen in Fig. 9. The wall-clock reconstruction time increases with the longer synthesis times as the time for each evaluation of the measurement operator becomes longer due to the larger number of visibilities.

Even in this more challenging setting, we find a good correlation between the predicted and the oracle error. The pixel UQ maps still show some patterns characteristic of the multiscale wavelet representation used and should, therefore, be considered in the analysis of the pixel UQ maps. The fast pixel UQ represents a fraction of the time and number of measurement operator evaluations required for the MAP reconstruction.

7 DISCUSSION

In this work, we have worked with synthetic data sets, starting with gridded visibilities and ending with realistic ungridded visibility patterns. The handling of real or realistic data is beyond the scope of this work, which focuses on the methodology presented and the validation of the UQ techniques. Some problems faced while handling large-scale observation include the calibration of direction-independent and direction-dependent effects. Further studies of the performance of quantifAI to more realistic images with differing large dynamic ranges and bright point sources are left for the future. The extension of quantifAI to incorporate a frequency axis in the reconstruction to cope with multi-frequency observations is also left for further study. Even if quantifAI could reconstruct images with UQ maps with |$\sim 10^{5}$| visibilities in less than two minutes, there is ongoing work to exploit existing C++ parallelization capabilities to scale the method further. A detailed performance and computing time benchmark is expected for the future implementation.

The current approach to set the regularization strength of the CRR-NN, |$\lambda$|⁠, with a grid search is not compatible with real data as we require access to the ground truth. There are several ways to circumvent this problem in the large-scale setting. One way forward is to consider a subset of the observations to alleviate the computational burden and rely on the empirical Bayesian approach from Vidal et al. (2020) and De Bortoli et al. (2020) to estimate the regularization parameter directly from the observed data. Another way forward is to follow a heuristic approach similar to Terris et al. (2022). The study of the best strategy for quantifAI is left for further study.

We have not included a positivity constraint in the quantifAI model, as is the case for Cai et al. (2018b), which we use for comparison. However, there is no impediment to adding a positivity constraint to quantifAI; it amounts to changing the indicator function from equation (25) to an indicator function of the real positive orthant. The proximal operator associated with this indicator function has a closed form: the projection to the positive orthant. Consequently, we can compute the corresponding MAP solution and evaluate the potential using the modified indicator function.

The fast pixel UQ proposed in this article improved the tightness of the UQ bounds with respect to the LCIs by considering the entire image when saturating the HPD region. Nevertheless, these pixel-level UQ maps are intended to be shared as visual aids accompanying the reconstructions. The pixel UQ maps can drive the astronomer’s attention to a specific region in the image where a consequent hypothesis test can be carried out using the techniques described in this work. The choice of wavelet basis impacts the fast pixel UQ maps produced, so it should be considered when analysing the maps. A way to improve the fast pixel UQ maps would be to replace the orthogonal wavelet basis used with a more performant dictionary as in SARA (Carrillo et al. 2012), which has shown to be well adapted for RI images.

8 CONCLUSIONS

In this work, we propose a new method coined quantifAI that addresses uncertainty quantification in radio-interferometric (RI) imaging with data-driven (learned) priors in very high-dimensional settings. We have focused on three fundamental points in the RI imaging pipeline: scalability, estimation performance, and uncertainty quantification (UQ).

Our model builds upon a principled Bayesian framework for the UQ analysis, which is known to be computationally expensive when exploiting MCMC sampling methods. However, in this work, we leverage convex optimization techniques to estimate the MAP, the point estimate of the posterior distribution we use as reconstruction. We restrict our model to a log-concave posterior distribution to remain highly scalable and have Bayesian UQ techniques. This restriction is equivalent to having convex potentials for our likelihood and prior. In this scenario, we can exploit an approximation of the HPD region, which only requires the MAP estimation (Pereyra 2017) and bypasses expensive sampling techniques.

We want to include data-driven priors that can encode complex information learned implicitly from training data making them more expressive. Consequently, the learned priors allow us to improve performance with respect to previous models based on handcrafted priors (Cai et al. 2018b), e.g. wavelet-based sparsity-promoting priors. To support fast UQ techniques, our models must be convex, hence we adopt the recently introduced learnable convex-ridge regularizer neural network (CRR-NN; Goujon et al. 2023b). The CRR-NN-based prior is performant, reliable, and has shown to be robust to data distribution shifts. The quantifAI model uses an analytic physically motivated model for the likelihood and the learned CRR-NN-based prior. In this work, we are focusing on the methodology, which is why we have only considered small problems, i.e. images of |$256 \times 256$|⁠. Nevertheless, quantifAI can be integrated into the distributed frameworks (Pratley et al. 2019a; Pratley & McEwen 2019), which is the focus of ongoing work.

Numerical experiments are conducted with four images representative of RI imaging. We compare the quantifAI model with the model containing a wavelet-based prior of Cai et al. (2018b). Our results show a considerable improvement in the reconstruction performance for quantifAI. We validate our results against posterior samples from MCMC sampling algorithms and compute the posterior standard deviation. We found that quantifAI produced more meaningful posterior standard deviations in comparison to the wavelet-based model. We also included numerical experiments with simulated MeerKAT ungridded visibilities, where we present quantifAI’s performance and computing times going up to |$\sim 10^{5}$| visibilities.

We explore several MAP-based UQ techniques that rely on the approximate HPD region. We carry out hypothesis tests of image structure to asses if some structures observed in the reconstructions are physical. We then computed LCIs (c.f. Bayesian error bars) to measure the pixel-wise uncertainty. These two approaches were proposed by Cai et al. (2018b), and in this work, we validated them with MCMC posterior sampling results. Even if LCIs represent an already scalable alternative to sampling-based methods to provide pixel-wise UQ, they remain expensive for SKA-size data. Therefore, we proposed a novel pixel-wise UQ technique to approximate pixel errors at different scales that is three orders of magnitude faster than the LCIs. The new approach is based on thresholding the coefficients of a wavelet representation of the reconstruction until the HPD region saturates and is six orders of magnitude faster than sampling-based techniques.

quantifAI proposed an approach with the potential to be highly scalable and performant to address UQ in RI imaging. In this work, we have compared quantifAI to a wavelet-based model using numerical experiments and a variety of metrics. However, as both models rely on the Bayesian framework, we could make a Bayesian model comparison, a principled approach to model selection and determine which model the data favours. Recent developments in McEwen et al. (2023) extend the model comparison to the learnt setting, with data-driven priors. The focus of ongoing work is to implement the proposed methodology in existing RI imaging frameworks purify  ⁷ (Carrillo et al. 2014; Pratley et al. 2018, 2019b) and sopt  ⁸ (Carrillo et al. 2012; Onose et al. 2016) to exploit massively parallelized computing environment (Pratley et al. 2019a; Pratley & McEwen 2019) and to realize the potential of scalability. In the near future, we plan to benchmark the speed and scalability of quantifAI in a highly realistic setting.

A new perspective is to relax the convexity constraint of the prior by exploiting the fact that the posterior potential needs to be convex (rather than the prior) and that the RI imaging likelihood is already strongly convex. The relaxation of the CRR regularizer has been studied in a very recent work (Goujon, Neumayer & Unser 2023a), where a weakly-convex-ridge-regularizer neural network (WCRR-NN) has been proposed. If the WCRR-NN is adopted, it could further enhance the expressiveness of the regularizer and the reconstruction performance of quantifAI in the RI imaging problem.

Acknowledgement

This work is supported by the UK Research and Innovation (UKRI) and Engineering and Physical Sciences Research Council (EPSRC) grant numbers EP/W007673/1 (LEXCI) and EP/T007346/1 (BOLT). MM is supported by the UCL Centre for Doctoral Training in Data Intensive Science (STFC Training grant ST/P006736/1).

We acknowledge the Python packages used in this work: ipython (Pérez & Granger 2007), jupyter (Kluyver et al. 2016), matplotlib (Hunter 2007), pytorch (Paszke et al. 2019), numpy (Harris et al. 2020), astropy (Astropy Collaboration 2022), scikit image (van der Walt et al. 2014).

DATA AVAILABILITY

We provide the quantifaIpytorch-based Python package in a publicly available repository.⁹ We are in favour of reproducible research, which is why the images, visibilities, scripts, trained CRR-NN model and code to reproduce this article’s experiments can be found in the aforementioned repository.

Footnotes

References

APPENDIX A: MAP CALCULATION FOR THE WAVELET-BASED PRIOR MODEL

The MAP estimation for the wavelet-based model can recasted as the following optimization problem

where |$\lambda _{\text{wav}}$| corresponds to the regularization strength of the wavelet-based prior. The FISTA algorithm to estimate the MAP is presented in Algorithm 2 where we use the soft thresholding operator, |$\text{soft}(\cdot )$|⁠, defined in equation (49). The Lipschitz constant used to define the step size can be set as |$L_{\text{wav}} = \Vert \Phi ^{\dagger } \Phi \Vert / \sigma ^{2}$|⁠.

APPENDIX B: WAVELET-BASED BAYESIAN UNCERTAINTY QUANTIFICATION

B1 Hypothesis testing of image structure

Fig. B1 and Table B1 present the results of the hypothesis testing of structure for the model with the sparsity-promoting prior.

APPENDIX C: MORE REALISTIC EXPERIMENT RESULTS

Figs C1, C2, and C3 present the results using the realistic MeerKAT ungridded visibility pattern of the images W28, Cygnus A, and 3C288, correspondingly. Tables C1, C2, and C3 show the quantitative results of the experiment.