quantifai
is a PyTorch-based open-source radio interferometric imaging reconstruction package with scalable Bayesian uncertainty quantification relying on data-driven (learned) priors. This package was used to produce the results of Liaudat et al. 2023. The quantifai
model relies on the data-driven convex regulariser from Goujon et al. 2022.
In this code, we bypass the need to perform Markov chain Monte Carlo (MCMC) sampling for Bayesian uncertainty quantification, and we rely on convex accelerated optimisation algorithms. The quantifai
package also includes MCMC algorithms for posterior sampling as they were used to validate our approach.
Note
This Python package is built on top of PyTorch, so a GPU can considerably accelerate all computations.
The quantifai
package relies on the convex ridge regulariser CRR from Goujon et al. 2022. The version used to generate the results from the Liaudat et al. paper is the release v0.1
from the fork github.com/tobias-liaudat/convex_ridge_regularizers. The PyTorch wavelet support relies on the release v0.1
from the fork github.com/tobias-liaudat/PyTorch-Wavelet-Toolbox.
We have not yet pushed the Python package to PyPi; therefore, the easiest way to install quantifai
is to start by cloning the repo
git clone https://github.com/astro-informatics/QuantifAI
cd QuantifAI
Continue by creating a conda environment with all the requirements already specified in environment.yml
as follows
conda env create -f environment.yml
conda activate quantifai_env
Finally, the quantifai
package can be installed by running
pip install -e .
Note
If the user does not want to create a conda environment, they can install the dependencies in the environment.yml
file. The specific version of the convex ridge regulariser, complex PyTorch support and the PyTorch wavelets used can be manually installed by running (in the following order)
pip install git+https://github.com/tobias-liaudat/[email protected]
pip install git+https://github.com/tobias-liaudat/[email protected]
pip install git+https://github.com/tobias-liaudat/[email protected]
The paper's numerical results were obtained using PyTorch version 1.13.1
.
The easiest way to get into using quantifai
is to check the different notebooks in the example/
directory, which includes:
-
Compute the MAP estimation with the
QuantifAI
model (Notebook). -
Compute the MAP estimation with the wavelet-based model (Notebook).
-
Compute the MAP-based LCIs with the
QuantifAI
model (Notebook). -
Compute the MAP-based LCIs with the wavelet-based model (Notebook).
-
Compute the MAP-based fast pixel uncertainty quantification method with
QuantifAI
(Notebook). -
Compute a hypothesis test on an inpainted surrogate image with
QuantifAI
(Notebook). -
Sample from the posterior distribution of the
QuantifAI
model using the SK-ROCK algorithm (Pereyra et al. 2020) and compare the results with sample-based LCIs (Notebook).
All the scripts and notebooks used to generate the plots of Liaudat et al. 2023 can be found in the paper/Liaudat2023/
directory and the data in the data/
directory.
The most computationally intensive results of the paper can be obtained by running the two scripts in paper/Liaudat2023/scripts/
, where UQ_SKROCK_CRR.py
corresponds to the QuantifAI
model and UQ_SKROCK_wavelets.py
to the wavelet-based model. The rest of the results and plots can be generated by running the different notebooks in paper/Liaudat2023/notebooks/
.
Should this code be used in any way, we kindly request that the following article be referenced. A BibTeX entry for this reference may look like:
@article{liaudat2023:quantifai,
author = {Tobías~I.~Liaudat and Matthijs~Mars and Matthew~A.~Price and Marcelo~Pereyra and Marta~M.~Betcke and Jason~D.~McEwen},
title = {Scalable Bayesian uncertainty quantification with data-driven priors for radio interferometric imaging},
journal = "RAS Techniques and Instruments (RASTI), submitted",
eprint = "arXiv:2312.00125",
year = "2023",
}
quantifai
is released under the GPL-3 license (see LICENSE.txt).