<!-- DO NOT EDIT. -->
<!-- THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. -->
<!-- TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: -->
<!-- "auto_examples/sampling/demo_sampling.py" -->
<!-- LINE NUMBERS ARE GIVEN BELOW. -->

#### NOTE
New to DeepInverse? Get started with the basics with the
[5 minute quickstart tutorial](https://deepinv.org/auto_examples/basics/demo_quickstart.html.md#sphx-glr-auto-examples-basics-demo-quickstart-py)..

<a id="sphx-glr-auto-examples-sampling-demo-sampling-py"></a>

# Uncertainty quantification with PnP-ULA.

This code shows you how to use sampling algorithms to quantify uncertainty of a reconstruction
from incomplete and noisy measurements.

ULA obtains samples by running the following iteration:

$$
x_{k+1} = x_k +  \alpha \eta \nabla \log p_{\sigma}(x_k) + \eta \nabla \log p(y|x_k)  + \sqrt{2 \eta} z_k
$$

where $z_k \sim \mathcal{N}(0, I)$ is a Gaussian random variable, $\eta$ is the step size and
$\alpha$ is a parameter controlling the regularization.

The PnP-ULA method is described in the paper Laumont *et al.*<sup>[1](#footcite-laumont2022bayesian)</sup>.

<!-- GENERATED FROM PYTHON SOURCE LINES 21-26 -->
```Python
import deepinv as dinv
from deepinv.utils.plotting import plot
import torch
from deepinv.utils import load_example
```

<!-- GENERATED FROM PYTHON SOURCE LINES 27-31 -->

## Load image from the internet

This example uses an image of Messi.

<!-- GENERATED FROM PYTHON SOURCE LINES 31-36 -->
```Python
device = dinv.utils.get_device()

x = load_example("messi.jpg", img_size=32).to(device)
```

```none
Selected GPU 0 with 24085.25 MiB free memory
```

<!-- GENERATED FROM PYTHON SOURCE LINES 37-41 -->

## Define forward operator and noise model

This example uses inpainting as the forward operator and Gaussian noise as the noise model.

<!-- GENERATED FROM PYTHON SOURCE LINES 41-49 -->
```Python
sigma = 0.1  # noise level
physics = dinv.physics.Inpainting(mask=0.5, img_size=x.shape[1:], device=device)
physics.noise_model = dinv.physics.GaussianNoise(sigma=sigma)

# Set the global random seed from pytorch to ensure reproducibility of the example.
torch.manual_seed(0)
```

```none
<torch._C.Generator object at 0x7f2ceabc8d10>
```

<!-- GENERATED FROM PYTHON SOURCE LINES 50-57 -->

## Define the likelihood

Since the noise model is Gaussian, the negative log-likelihood is the L2 loss.

$$
-\log p(y|x) \propto \frac{1}{2\sigma^2} \|y-Ax\|^2

$$

<!-- GENERATED FROM PYTHON SOURCE LINES 57-61 -->
```Python
# load Gaussian Likelihood
likelihood = dinv.optim.data_fidelity.L2(sigma=sigma)
```

<!-- GENERATED FROM PYTHON SOURCE LINES 62-80 -->

## Define the prior

The score a distribution can be approximated using Tweedie’s formula via the
[`deepinv.optim.ScorePrior`](https://deepinv.org/api/stubs/deepinv.optim.ScorePrior.html.md#deepinv.optim.ScorePrior) class.

$$
\nabla \log p_{\sigma}(x) \approx \frac{1}{\sigma^2} \left(D(x,\sigma)-x\right)
$$

This example uses a pretrained DnCNN model.
From a Bayesian point of view, the score plays the role of the gradient of the
negative log prior
The hyperparameter `sigma_denoiser` ($sigma$) controls the strength of the prior.

In this example, we use a pretrained DnCNN model using the [`deepinv.loss.FNEJacobianSpectralNorm`](https://deepinv.org/api/stubs/deepinv.loss.FNEJacobianSpectralNorm.html.md#deepinv.loss.FNEJacobianSpectralNorm) loss,
which makes sure that the denoiser is firmly non-expansive (see Terris *et al.*<sup>[2](#footcite-terris2020building)</sup>), and helps to
stabilize the sampling algorithm.

<!-- GENERATED FROM PYTHON SOURCE LINES 80-86 -->
```Python
sigma_denoiser = 2 / 255
prior = dinv.optim.ScorePrior(
    denoiser=dinv.models.DnCNN(pretrained="download_lipschitz")
).to(device)
```

<!-- GENERATED FROM PYTHON SOURCE LINES 87-95 -->

## Create the MCMC sampler

Here we use the Unadjusted Langevin Algorithm (ULA) to sample from the posterior defined in
[`deepinv.sampling.ULAIterator`](https://deepinv.org/api/stubs/deepinv.sampling.ULAIterator.html.md#deepinv.sampling.ULAIterator).
The hyperparameter `step_size` controls the step size of the MCMC sampler,
`regularization` controls the strength of the prior and
`iterations` controls the number of iterations of the sampler.

<!-- GENERATED FROM PYTHON SOURCE LINES 95-114 -->
```Python
regularization = 0.9
step_size = 0.01 * (sigma**2)
iterations = int(5e3) if torch.cuda.is_available() else 10
params = {
    "step_size": step_size,
    "alpha": regularization,
    "sigma": sigma_denoiser,
}
f = dinv.sampling.sampling_builder(
    "ULA",
    prior=prior,
    data_fidelity=likelihood,
    max_iter=iterations,
    params_algo=params,
    thinning=1,
    verbose=True,
)
```

<!-- GENERATED FROM PYTHON SOURCE LINES 115-118 -->

## Generate the measurement

We apply the forward model to generate the noisy measurement.

<!-- GENERATED FROM PYTHON SOURCE LINES 118-122 -->
```Python
y = physics(x)
```

<!-- GENERATED FROM PYTHON SOURCE LINES 123-127 -->

## Run sampling algorithm and plot results

The sampling algorithm returns the posterior mean and variance.
We compare the posterior mean with a simple linear reconstruction.

<!-- GENERATED FROM PYTHON SOURCE LINES 127-146 -->
```Python
mean, var = f.sample(y, physics)

# compute linear inverse
x_lin = physics.A_adjoint(y)

# compute PSNR
print(f"Linear reconstruction PSNR: {dinv.metric.PSNR()(x, x_lin).item():.2f} dB")
print(f"Posterior mean PSNR: {dinv.metric.PSNR()(x, mean).item():.2f} dB")

# plot results
error = (mean - x).abs().sum(dim=1).unsqueeze(1)  # per pixel average abs. error
std = var.sum(dim=1).unsqueeze(1).sqrt()  # per pixel average standard dev.
imgs = [x_lin, x, mean, std / std.flatten().max(), error / error.flatten().max()]
plot(
    imgs,
    titles=["measurement", "ground truth", "post. mean", "post. std", "abs. error"],
)
```

```none
  0%|          | 0/5000 [00:00<?, ?it/s]
  3%|▎         | 134/5000 [00:00<00:03, 1333.41it/s]
  6%|▌         | 289/5000 [00:00<00:03, 1454.83it/s]
  9%|▉         | 439/5000 [00:00<00:03, 1472.25it/s]
 12%|█▏        | 589/5000 [00:00<00:02, 1480.42it/s]
 15%|█▍        | 739/5000 [00:00<00:02, 1484.93it/s]
 18%|█▊        | 889/5000 [00:00<00:02, 1487.72it/s]
 21%|██        | 1038/5000 [00:00<00:02, 1485.85it/s]
 24%|██▎       | 1187/5000 [00:00<00:02, 1478.10it/s]
 27%|██▋       | 1335/5000 [00:00<00:02, 1472.99it/s]
 30%|██▉       | 1483/5000 [00:01<00:02, 1469.54it/s]
 33%|███▎      | 1630/5000 [00:01<00:02, 1467.16it/s]
 36%|███▌      | 1777/5000 [00:01<00:02, 1465.52it/s]
 38%|███▊      | 1924/5000 [00:01<00:02, 1464.40it/s]
 41%|████▏     | 2071/5000 [00:01<00:02, 1463.60it/s]
 44%|████▍     | 2218/5000 [00:01<00:01, 1463.05it/s]
 47%|████▋     | 2365/5000 [00:01<00:01, 1462.67it/s]
 50%|█████     | 2512/5000 [00:01<00:01, 1462.39it/s]
 53%|█████▎    | 2659/5000 [00:01<00:01, 1462.21it/s]
 56%|█████▌    | 2806/5000 [00:01<00:01, 1462.07it/s]
 59%|█████▉    | 2953/5000 [00:02<00:01, 1461.99it/s]
 62%|██████▏   | 3100/5000 [00:02<00:01, 1461.96it/s]
 65%|██████▍   | 3247/5000 [00:02<00:01, 1461.90it/s]
 68%|██████▊   | 3394/5000 [00:02<00:01, 1461.88it/s]
 71%|███████   | 3541/5000 [00:02<00:00, 1461.84it/s]
 74%|███████▍  | 3688/5000 [00:02<00:00, 1461.83it/s]
 77%|███████▋  | 3835/5000 [00:02<00:00, 1461.83it/s]
 80%|███████▉  | 3982/5000 [00:02<00:00, 1461.79it/s]
 83%|████████▎ | 4129/5000 [00:02<00:00, 1461.76it/s]
 86%|████████▌ | 4276/5000 [00:02<00:00, 1461.73it/s]
 88%|████████▊ | 4423/5000 [00:03<00:00, 1461.79it/s]
 91%|█████████▏| 4570/5000 [00:03<00:00, 1461.78it/s]
 94%|█████████▍| 4717/5000 [00:03<00:00, 1461.77it/s]
 97%|█████████▋| 4864/5000 [00:03<00:00, 1461.78it/s]
100%|██████████| 5000/5000 [00:03<00:00, 1464.54it/s]
Iteration 4999, current converge crit. = 1.42E-05, objective = 1.00E-03
Iteration 4999, current converge crit. = 3.42E-04, objective = 1.00E-03
Linear reconstruction PSNR: 8.55 dB
Posterior mean PSNR: 22.32 dB
```

<!-- GENERATED FROM PYTHON SOURCE LINES 147-150 -->
* **References:**

* <a id='footcite-laumont2022bayesian'>**[1]**</a> Rémi Laumont, Valentin De Bortoli, Andrés Almansa, Julie Delon, Alain Durmus, and Marcelo Pereyra. Bayesian imaging using plug & play priors: when langevin meets tweedie. *SIAM Journal on Imaging Sciences*, 15(2):701–737, 2022.
* <a id='footcite-terris2020building'>**[2]**</a> Matthieu Terris, Audrey Repetti, Jean-Christophe Pesquet, and Yves Wiaux. Building firmly nonexpansive convolutional neural networks. In *ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)*, 8658–8662. IEEE, 2020.

**Total running time of the script:** (0 minutes 3.571 seconds)

<a id="sphx-glr-download-auto-examples-sampling-demo-sampling-py"></a>

[`Download Jupyter notebook: demo_sampling.ipynb`](https://deepinv.org/auto_examples/sampling/_downloads/5a89e1000770b922d57eed794f82857e/demo_sampling.ipynb)

[`Download Python source code: demo_sampling.py`](https://deepinv.org/auto_examples/sampling/_downloads/5596c4a992481456ea481739b5152618/demo_sampling.py)

[`Download zipped: demo_sampling.zip`](https://deepinv.org/auto_examples/sampling/_downloads/199de052057d8e7a8b39ef1dbe23d4c9/demo_sampling.zip)

[Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)
