.. 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.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_examples_sampling_demo_sampling.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_sampling_demo_sampling.py:


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:

.. math::

    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 :math:`z_k \sim \mathcal{N}(0, I)` is a Gaussian random variable, :math:`\eta` is the step size and
:math:`\alpha` is a parameter controlling the regularization.

The PnP-ULA method is described in the paper `"Bayesian imaging using Plug & Play priors: when Langevin meets Tweedie
" <https://arxiv.org/abs/2103.04715>`_.

.. GENERATED FROM PYTHON SOURCE LINES 21-27

.. code-block:: Python


    import deepinv as dinv
    from deepinv.utils.plotting import plot
    import torch
    from deepinv.utils.demo import load_url_image








.. GENERATED FROM PYTHON SOURCE LINES 28-32

Load image from the internet
--------------------------------------------

This example uses an image of Lionel Messi from Wikipedia.

.. GENERATED FROM PYTHON SOURCE LINES 32-41

.. code-block:: Python


    device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"

    url = (
        "https://upload.wikimedia.org/wikipedia/commons/b/b4/"
        "Lionel-Messi-Argentina-2022-FIFA-World-Cup_%28cropped%29.jpg"
    )
    x = load_url_image(url=url, img_size=32).to(device)








.. GENERATED FROM PYTHON SOURCE LINES 42-46

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 46-54

.. code-block:: Python


    sigma = 0.1  # noise level
    physics = dinv.physics.Inpainting(mask=0.5, tensor_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)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    <torch._C.Generator object at 0x7f1764340770>



.. GENERATED FROM PYTHON SOURCE LINES 55-62

Define the likelihood
--------------------------------------------------------------

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

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

.. GENERATED FROM PYTHON SOURCE LINES 62-66

.. code-block:: Python


    # load Gaussian Likelihood
    likelihood = dinv.optim.data_fidelity.L2(sigma=sigma)








.. GENERATED FROM PYTHON SOURCE LINES 67-86

Define the prior
-------------------------------------------

The score a distribution can be approximated using Tweedie's formula via the
:class:`deepinv.optim.ScorePrior` class.

.. math::

           \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`` (:math:`sigma`) controls the strength of the prior.

In this example, we use a pretrained DnCNN model using the :class:`deepinv.loss.FNEJacobianSpectralNorm` loss,
which makes sure that the denoiser is firmly non-expansive (see
`"Building firmly nonexpansive convolutional neural networks" <https://hal.science/hal-03139360>`_), and helps to
stabilize the sampling algorithm.

.. GENERATED FROM PYTHON SOURCE LINES 86-92

.. code-block:: Python


    sigma_denoiser = 2 / 255
    prior = dinv.optim.ScorePrior(
        denoiser=dinv.models.DnCNN(pretrained="download_lipschitz")
    ).to(device)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Downloading: "https://huggingface.co/deepinv/dncnn/resolve/main/dncnn_sigma2_lipschitz_color.pth?download=true" to /home/runner/.cache/torch/hub/checkpoints/dncnn_sigma2_lipschitz_color.pth

      0%|          | 0.00/2.56M [00:00<?, ?B/s]
    100%|██████████| 2.56M/2.56M [00:00<00:00, 43.4MB/s]




.. GENERATED FROM PYTHON SOURCE LINES 93-101

Create the MCMC sampler
--------------------------------------------------------------

Here we use the Unadjusted Langevin Algorithm (ULA) to sample from the posterior defined in
:class:`deepinv.sampling.ULA`.
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 101-115

.. code-block:: Python


    regularization = 0.9
    step_size = 0.01 * (sigma**2)
    iterations = int(5e3) if torch.cuda.is_available() else 10
    f = dinv.sampling.ULA(
        prior=prior,
        data_fidelity=likelihood,
        max_iter=iterations,
        alpha=regularization,
        step_size=step_size,
        verbose=True,
        sigma=sigma_denoiser,
    )








.. GENERATED FROM PYTHON SOURCE LINES 116-119

Generate the measurement
--------------------------------------------------------------
We apply the forward model to generate the noisy measurement.

.. GENERATED FROM PYTHON SOURCE LINES 119-123

.. code-block:: Python


    y = physics(x)









.. GENERATED FROM PYTHON SOURCE LINES 124-128

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 128-146

.. code-block:: Python


    mean, var = f(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"],
    )



.. image-sg:: /auto_examples/sampling/images/sphx_glr_demo_sampling_001.png
   :alt: measurement, ground truth, post. mean, post. std, abs. error
   :srcset: /auto_examples/sampling/images/sphx_glr_demo_sampling_001.png
   :class: sphx-glr-single-img


.. rst-class:: sphx-glr-script-out

 .. code-block:: none


      0%|          | 0/10 [00:00<?, ?it/s]
     90%|█████████ | 9/10 [00:00<00:00, 83.97it/s]
    100%|██████████| 10/10 [00:00<00:00, 84.65it/s]
    Monte Carlo sampling finished! elapsed time=0.12 seconds
    Linear reconstruction PSNR: 8.79 dB
    Posterior mean PSNR: 8.86 dB





.. rst-class:: sphx-glr-timing

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


.. _sphx_glr_download_auto_examples_sampling_demo_sampling.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: demo_sampling.ipynb <demo_sampling.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: demo_sampling.py <demo_sampling.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: demo_sampling.zip <demo_sampling.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_