.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/sampling/demo_diffusion_sde.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_diffusion_sde.py>`
        to download the full example code.

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

.. _sphx_glr_auto_examples_sampling_demo_diffusion_sde.py:


Building your diffusion posterior sampling method using SDEs
============================================================

This demo shows you how to use
:class:`deepinv.sampling.PosteriorDiffusion` to perform posterior sampling. It also can be used to perform unconditional image generation with arbitrary denoisers, if the data fidelity term is not specified.

This method requires:

* A well-trained denoiser with varying noise levels (ideally with large noise levels) (e.g., :class:`deepinv.models.NCSNpp`).

* A (noisy) data fidelity term (e.g., :class:`deepinv.sampling.DPSDataFidelity`).

* Define a drift term :math:`f(x, t)` and a diffusion term :math:`g(t)` for the forward-time SDE. They can be defined through the :class:`deepinv.sampling.DiffusionSDE` (e.g., :class:`deepinv.sampling.VarianceExplodingDiffusion`).

The :class:`deepinv.sampling.PosteriorDiffusion` class can be used to perform posterior sampling for inverse problems.
Consider the acquisition model:

.. math::
     y = \noise{\forw{x}}

where :math:`\forw{x}` is the forward operator (e.g., a convolutional operator) and :math:`\noise{\cdot}` is the noise operator (e.g., Gaussian noise).
This class defines the reverse-time SDE for the posterior distribution :math:`p(x|y)` given the data :math:`y`:

.. math::
     d\, x_t = \left( f(x_t, t) - \frac{1 + \alpha}{2} g(t)^2 \nabla_{x_t} \log p_t(x_t | y) \right) d\,t + g(t) \sqrt{\alpha} d\, w_{t}

where :math:`f` is the drift term, :math:`g` is the diffusion coefficient and :math:`w` is the standard Brownian motion.
The drift term and the diffusion coefficient are defined by the underlying (unconditional) forward-time SDE `sde`.
In this example, we will use 2 well-known SDE in the literature: the Variance-Exploding (VE) and Variance-Preserving (VP or DDPM).

The (conditional) score function :math:`\nabla_{x_t} \log p_t(x_t | y)` can be decomposed using the Bayes' rule:

.. math::
     \nabla_{x_t} \log p_t(x_t | y) = \nabla_{x_t} \log p_t(x_t) + \nabla_{x_t} \log p_t(y | x_t).

The first term is the score function of the unconditional SDE, which is typically approximated by an MMSE denoiser (`denoiser`) using the well-known Tweedie's formula, while the
second term is approximated by the (noisy) data-fidelity term (`data_fidelity`).
We implement various data-fidelity terms in `the user guide <https://deepinv.github.io/deepinv/user_guide/reconstruction/sampling.html#id2>`_.

.. note::

    In this demo, we limit the number of diffusion steps for the sake of speed, but in practice, you should use a larger number of steps to obtain better results.

.. GENERATED FROM PYTHON SOURCE LINES 47-55

---------------------------------------------------

Let us import the necessary modules, define the denoiser and the SDE.

In this first example, we use the Variance-Exploding SDE, whose forward process is defined as:

.. math::
    d\, x_t = g(t) d\, w_t \quad \mbox{where } g(t) = \sigma_{\mathrm{min}}\left( \frac{\sigma_{\mathrm{max}}}{\sigma_{\mathrm{min}}}\right)^t

.. GENERATED FROM PYTHON SOURCE LINES 55-64

.. code-block:: Python


    import torch
    import deepinv as dinv
    from deepinv.models import NCSNpp

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float64
    figsize = 2.5
    gif_frequency = 10  # Increase this value to save the GIF saving time







.. GENERATED FROM PYTHON SOURCE LINES 65-96

.. code-block:: Python

    from deepinv.sampling import (
        PosteriorDiffusion,
        DPSDataFidelity,
        EulerSolver,
        VarianceExplodingDiffusion,
    )
    from deepinv.optim import ZeroFidelity

    # In this example, we use the pre-trained FFHQ-64 model from the
    # EDM framework: https://arxiv.org/pdf/2206.00364 .
    # The network architecture is from Song et al: https://arxiv.org/abs/2011.13456 .
    denoiser = NCSNpp(pretrained="download").to(device)

    # The solution is obtained by calling the SDE object with a desired solver (here, Euler).
    # The reproducibility of the SDE Solver class can be controlled by providing the pseudo-random number generator.
    num_steps = 150
    rng = torch.Generator(device).manual_seed(42)
    timesteps = torch.linspace(1, 0.001, num_steps)
    solver = EulerSolver(timesteps=timesteps, rng=rng)


    sigma_min = 0.005
    sigma_max = 5
    sde = VarianceExplodingDiffusion(
        sigma_max=sigma_max,
        sigma_min=sigma_min,
        alpha=0.5,
        device=device,
        dtype=dtype,
    )





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

 .. code-block:: none

    Downloading: "https://huggingface.co/deepinv/edm/resolve/main/ncsnpp-ffhq64-uncond-ve.pt?download=true" to /home/runner/.cache/torch/hub/checkpoints/ncsnpp-ffhq64-uncond-ve.pt

      0%|          | 0.00/240M [00:00<?, ?B/s]
      0%|          | 512k/240M [00:00<00:50, 4.93MB/s]
      3%|▎         | 7.62M/240M [00:00<00:05, 44.8MB/s]
     10%|█         | 24.4M/240M [00:00<00:02, 103MB/s] 
     17%|█▋        | 41.0M/240M [00:00<00:01, 131MB/s]
     24%|██▎       | 56.9M/240M [00:00<00:01, 143MB/s]
     31%|███       | 73.6M/240M [00:00<00:01, 154MB/s]
     37%|███▋      | 89.5M/240M [00:00<00:00, 158MB/s]
     44%|████▍     | 106M/240M [00:00<00:00, 162MB/s] 
     51%|█████     | 122M/240M [00:00<00:00, 164MB/s]
     57%|█████▋    | 138M/240M [00:01<00:00, 159MB/s]
     64%|██████▍   | 154M/240M [00:01<00:00, 164MB/s]
     71%|███████▏  | 171M/240M [00:01<00:00, 166MB/s]
     78%|███████▊  | 188M/240M [00:01<00:00, 169MB/s]
     85%|████████▌ | 204M/240M [00:01<00:00, 170MB/s]
     92%|█████████▏| 221M/240M [00:01<00:00, 172MB/s]
     99%|█████████▉| 238M/240M [00:01<00:00, 172MB/s]
    100%|██████████| 240M/240M [00:01<00:00, 154MB/s]




.. GENERATED FROM PYTHON SOURCE LINES 97-102

Reverse-time SDE as sampling process
--------------------------------------

When the data fidelity is not given, the posterior diffusion is equivalent to the unconditional diffusion.
Sampling is performed by solving the reverse-time SDE. To do so, we generate a reverse-time trajectory.

.. GENERATED FROM PYTHON SOURCE LINES 102-134

.. code-block:: Python


    model = PosteriorDiffusion(
        data_fidelity=ZeroFidelity(),
        sde=sde,
        denoiser=denoiser,
        solver=solver,
        dtype=dtype,
        device=device,
        verbose=True,
    )
    x, trajectory = model(
        y=None,
        physics=None,
        x_init=(1, 3, 64, 64),
        seed=1,
        get_trajectory=True,
    )
    dinv.utils.plot(
        x,
        titles="Unconditional generation",
        save_fn="sde_sample.png",
        figsize=(figsize, figsize),
    )

    dinv.utils.save_videos(
        trajectory.cpu()[::gif_frequency],
        time_dim=0,
        titles=["VE-SDE Trajectory"],
        save_fn="sde_trajectory.gif",
        figsize=(figsize, figsize),
    )





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

 .. code-block:: none


      0%|          | 0/149 [00:00<?, ?it/s]
      1%|          | 1/149 [00:00<01:24,  1.75it/s]
      1%|▏         | 2/149 [00:01<01:21,  1.80it/s]
      2%|▏         | 3/149 [00:01<01:21,  1.79it/s]
      3%|▎         | 4/149 [00:02<01:20,  1.81it/s]
      3%|▎         | 5/149 [00:02<01:19,  1.81it/s]
      4%|▍         | 6/149 [00:03<01:18,  1.83it/s]
      5%|▍         | 7/149 [00:03<01:17,  1.82it/s]
      5%|▌         | 8/149 [00:04<01:17,  1.83it/s]
      6%|▌         | 9/149 [00:04<01:16,  1.84it/s]
      7%|▋         | 10/149 [00:05<01:15,  1.84it/s]
      7%|▋         | 11/149 [00:06<01:14,  1.84it/s]
      8%|▊         | 12/149 [00:06<01:14,  1.83it/s]
      9%|▊         | 13/149 [00:07<01:13,  1.84it/s]
      9%|▉         | 14/149 [00:07<01:13,  1.83it/s]
     10%|█         | 15/149 [00:08<01:12,  1.84it/s]
     11%|█         | 16/149 [00:08<01:12,  1.83it/s]
     11%|█▏        | 17/149 [00:09<01:11,  1.84it/s]
     12%|█▏        | 18/149 [00:09<01:11,  1.83it/s]
     13%|█▎        | 19/149 [00:10<01:10,  1.84it/s]
     13%|█▎        | 20/149 [00:10<01:10,  1.83it/s]
     14%|█▍        | 21/149 [00:11<01:09,  1.84it/s]
     15%|█▍        | 22/149 [00:12<01:09,  1.83it/s]
     15%|█▌        | 23/149 [00:12<01:08,  1.84it/s]
     16%|█▌        | 24/149 [00:13<01:08,  1.82it/s]
     17%|█▋        | 25/149 [00:13<01:07,  1.83it/s]
     17%|█▋        | 26/149 [00:14<01:07,  1.82it/s]
     18%|█▊        | 27/149 [00:14<01:06,  1.83it/s]
     19%|█▉        | 28/149 [00:15<01:06,  1.82it/s]
     19%|█▉        | 29/149 [00:15<01:05,  1.83it/s]
     20%|██        | 30/149 [00:16<01:05,  1.81it/s]
     21%|██        | 31/149 [00:16<01:04,  1.82it/s]
     21%|██▏       | 32/149 [00:17<01:04,  1.81it/s]
     22%|██▏       | 33/149 [00:18<01:03,  1.83it/s]
     23%|██▎       | 34/149 [00:18<01:03,  1.82it/s]
     23%|██▎       | 35/149 [00:19<01:02,  1.83it/s]
     24%|██▍       | 36/149 [00:19<01:02,  1.82it/s]
     25%|██▍       | 37/149 [00:20<01:01,  1.83it/s]
     26%|██▌       | 38/149 [00:20<01:00,  1.82it/s]
     26%|██▌       | 39/149 [00:21<01:00,  1.83it/s]
     27%|██▋       | 40/149 [00:21<00:59,  1.82it/s]
     28%|██▊       | 41/149 [00:22<00:58,  1.83it/s]
     28%|██▊       | 42/149 [00:22<00:58,  1.82it/s]
     29%|██▉       | 43/149 [00:23<00:57,  1.83it/s]
     30%|██▉       | 44/149 [00:24<00:57,  1.82it/s]
     30%|███       | 45/149 [00:24<00:56,  1.83it/s]
     31%|███       | 46/149 [00:25<00:56,  1.83it/s]
     32%|███▏      | 47/149 [00:25<00:55,  1.84it/s]
     32%|███▏      | 48/149 [00:26<00:55,  1.83it/s]
     33%|███▎      | 49/149 [00:26<00:54,  1.83it/s]
     34%|███▎      | 50/149 [00:27<00:54,  1.82it/s]
     34%|███▍      | 51/149 [00:27<00:53,  1.82it/s]
     35%|███▍      | 52/149 [00:28<00:53,  1.82it/s]
     36%|███▌      | 53/149 [00:29<00:52,  1.83it/s]
     36%|███▌      | 54/149 [00:29<00:52,  1.82it/s]
     37%|███▋      | 55/149 [00:30<00:51,  1.83it/s]
     38%|███▊      | 56/149 [00:30<00:50,  1.83it/s]
     38%|███▊      | 57/149 [00:31<00:50,  1.84it/s]
     39%|███▉      | 58/149 [00:31<00:49,  1.83it/s]
     40%|███▉      | 59/149 [00:32<00:49,  1.83it/s]
     40%|████      | 60/149 [00:32<00:48,  1.82it/s]
     41%|████      | 61/149 [00:33<00:47,  1.84it/s]
     42%|████▏     | 62/149 [00:33<00:47,  1.83it/s]
     42%|████▏     | 63/149 [00:34<00:46,  1.84it/s]
     43%|████▎     | 64/149 [00:35<00:46,  1.83it/s]
     44%|████▎     | 65/149 [00:35<00:46,  1.82it/s]
     44%|████▍     | 66/149 [00:36<00:45,  1.82it/s]
     45%|████▍     | 67/149 [00:36<00:44,  1.83it/s]
     46%|████▌     | 68/149 [00:37<00:44,  1.83it/s]
     46%|████▋     | 69/149 [00:37<00:43,  1.84it/s]
     47%|████▋     | 70/149 [00:38<00:43,  1.83it/s]
     48%|████▊     | 71/149 [00:38<00:42,  1.84it/s]
     48%|████▊     | 72/149 [00:39<00:41,  1.83it/s]
     49%|████▉     | 73/149 [00:39<00:41,  1.84it/s]
     50%|████▉     | 74/149 [00:40<00:40,  1.84it/s]
     50%|█████     | 75/149 [00:41<00:40,  1.84it/s]
     51%|█████     | 76/149 [00:41<00:39,  1.84it/s]
     52%|█████▏    | 77/149 [00:42<00:39,  1.84it/s]
     52%|█████▏    | 78/149 [00:42<00:38,  1.83it/s]
     53%|█████▎    | 79/149 [00:43<00:38,  1.84it/s]
     54%|█████▎    | 80/149 [00:43<00:37,  1.83it/s]
     54%|█████▍    | 81/149 [00:44<00:36,  1.84it/s]
     55%|█████▌    | 82/149 [00:44<00:36,  1.84it/s]
     56%|█████▌    | 83/149 [00:45<00:35,  1.84it/s]
     56%|█████▋    | 84/149 [00:45<00:35,  1.84it/s]
     57%|█████▋    | 85/149 [00:46<00:34,  1.83it/s]
     58%|█████▊    | 86/149 [00:47<00:34,  1.83it/s]
     58%|█████▊    | 87/149 [00:47<00:33,  1.83it/s]
     59%|█████▉    | 88/149 [00:48<00:33,  1.84it/s]
     60%|█████▉    | 89/149 [00:48<00:32,  1.83it/s]
     60%|██████    | 90/149 [00:49<00:32,  1.84it/s]
     61%|██████    | 91/149 [00:49<00:31,  1.83it/s]
     62%|██████▏   | 92/149 [00:50<00:30,  1.84it/s]
     62%|██████▏   | 93/149 [00:50<00:30,  1.83it/s]
     63%|██████▎   | 94/149 [00:51<00:29,  1.84it/s]
     64%|██████▍   | 95/149 [00:51<00:29,  1.83it/s]
     64%|██████▍   | 96/149 [00:52<00:29,  1.83it/s]
     65%|██████▌   | 97/149 [00:53<00:28,  1.82it/s]
     66%|██████▌   | 98/149 [00:53<00:27,  1.82it/s]
     66%|██████▋   | 99/149 [00:54<00:27,  1.81it/s]
     67%|██████▋   | 100/149 [00:54<00:26,  1.82it/s]
     68%|██████▊   | 101/149 [00:55<00:26,  1.81it/s]
     68%|██████▊   | 102/149 [00:55<00:25,  1.83it/s]
     69%|██████▉   | 103/149 [00:56<00:25,  1.82it/s]
     70%|██████▉   | 104/149 [00:56<00:24,  1.83it/s]
     70%|███████   | 105/149 [00:57<00:24,  1.82it/s]
     71%|███████   | 106/149 [00:57<00:23,  1.83it/s]
     72%|███████▏  | 107/149 [00:58<00:23,  1.82it/s]
     72%|███████▏  | 108/149 [00:59<00:22,  1.83it/s]
     73%|███████▎  | 109/149 [00:59<00:21,  1.82it/s]
     74%|███████▍  | 110/149 [01:00<00:21,  1.84it/s]
     74%|███████▍  | 111/149 [01:00<00:20,  1.83it/s]
     75%|███████▌  | 112/149 [01:01<00:20,  1.83it/s]
     76%|███████▌  | 113/149 [01:01<00:19,  1.81it/s]
     77%|███████▋  | 114/149 [01:02<00:19,  1.82it/s]
     77%|███████▋  | 115/149 [01:02<00:18,  1.81it/s]
     78%|███████▊  | 116/149 [01:03<00:18,  1.83it/s]
     79%|███████▊  | 117/149 [01:03<00:17,  1.82it/s]
     79%|███████▉  | 118/149 [01:04<00:17,  1.82it/s]
     80%|███████▉  | 119/149 [01:05<00:16,  1.81it/s]
     81%|████████  | 120/149 [01:05<00:15,  1.83it/s]
     81%|████████  | 121/149 [01:06<00:15,  1.82it/s]
     82%|████████▏ | 122/149 [01:06<00:14,  1.82it/s]
     83%|████████▎ | 123/149 [01:07<00:14,  1.82it/s]
     83%|████████▎ | 124/149 [01:07<00:13,  1.83it/s]
     84%|████████▍ | 125/149 [01:08<00:13,  1.82it/s]
     85%|████████▍ | 126/149 [01:08<00:12,  1.83it/s]
     85%|████████▌ | 127/149 [01:09<00:12,  1.82it/s]
     86%|████████▌ | 128/149 [01:10<00:11,  1.83it/s]
     87%|████████▋ | 129/149 [01:10<00:10,  1.82it/s]
     87%|████████▋ | 130/149 [01:11<00:10,  1.83it/s]
     88%|████████▊ | 131/149 [01:11<00:09,  1.82it/s]
     89%|████████▊ | 132/149 [01:12<00:09,  1.82it/s]
     89%|████████▉ | 133/149 [01:12<00:08,  1.82it/s]
     90%|████████▉ | 134/149 [01:13<00:08,  1.83it/s]
     91%|█████████ | 135/149 [01:13<00:07,  1.82it/s]
     91%|█████████▏| 136/149 [01:14<00:07,  1.83it/s]
     92%|█████████▏| 137/149 [01:14<00:06,  1.82it/s]
     93%|█████████▎| 138/149 [01:15<00:06,  1.83it/s]
     93%|█████████▎| 139/149 [01:16<00:05,  1.82it/s]
     94%|█████████▍| 140/149 [01:16<00:04,  1.82it/s]
     95%|█████████▍| 141/149 [01:17<00:04,  1.81it/s]
     95%|█████████▌| 142/149 [01:17<00:03,  1.83it/s]
     96%|█████████▌| 143/149 [01:18<00:03,  1.82it/s]
     97%|█████████▋| 144/149 [01:18<00:02,  1.83it/s]
     97%|█████████▋| 145/149 [01:19<00:02,  1.82it/s]
     98%|█████████▊| 146/149 [01:19<00:01,  1.83it/s]
     99%|█████████▊| 147/149 [01:20<00:01,  1.83it/s]
     99%|█████████▉| 148/149 [01:20<00:00,  1.84it/s]
    100%|██████████| 149/149 [01:21<00:00,  1.83it/s]
    100%|██████████| 149/149 [01:21<00:00,  1.83it/s]




.. GENERATED FROM PYTHON SOURCE LINES 152-167

We obtain the following unconditional sample

.. container:: image-row

   .. image-sg-ignore:: /auto_examples/images/sde_sample.png
      :alt: example of unconditional sample
      :srcset: /auto_examples/images/sde_sample.png
      :class: custom-img
      :ignore_missing: true

   .. image-sg-ignore:: /auto_examples/images/sde_trajectory.gif
      :alt: example of unconditional trajectory
      :srcset: /auto_examples/images/sde_trajectory.gif
      :class: custom-gif
      :ignore_missing: true

.. GENERATED FROM PYTHON SOURCE LINES 169-172

When the data fidelity is given, together with the measurements and the physics, this class can be used to perform posterior sampling for inverse problems.
For example, consider the inpainting problem, where we have a noisy image and we want to recover the original image.
We can use the :class:`deepinv.sampling.DPSDataFidelity` as the data fidelity term.

.. GENERATED FROM PYTHON SOURCE LINES 173-219

.. code-block:: Python


    del trajectory  # clean memory
    mask = torch.ones_like(x)
    mask[..., 24:40, 24:40] = 0.0
    physics = dinv.physics.Inpainting(img_size=x.shape[1:], mask=mask, device=device)
    y = physics(x)

    weight = 1.0  # guidance strength
    dps_fidelity = DPSDataFidelity(denoiser=denoiser, weight=weight)

    model = PosteriorDiffusion(
        data_fidelity=dps_fidelity,
        denoiser=denoiser,
        sde=sde,
        solver=solver,
        dtype=dtype,
        device=device,
        verbose=True,
    )

    # To perform posterior sampling, we need to provide the measurements, the physics and the solver.
    # Moreover, when the physics is given, the initial point can be inferred from the physics if not given explicitly.
    seed_1 = 11
    x_hat, trajectory = model(
        y,
        physics,
        seed=seed_1,
        get_trajectory=True,
    )
    # Here, we plot the original image, the measurement and the posterior sample
    dinv.utils.plot(
        [x, y, x_hat],
        show=True,
        titles=["Original", "Measurement", "Posterior sample"],
        save_fn="posterior_sample.png",
        figsize=(figsize * 3, figsize),
    )
    # We can also save the trajectory of the posterior sample
    dinv.utils.save_videos(
        trajectory[::gif_frequency],
        time_dim=0,
        titles=["Posterior sample with VE"],
        save_fn="posterior_trajectory.gif",
        figsize=(figsize, figsize),
    )





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

 .. code-block:: none


      0%|          | 0/149 [00:00<?, ?it/s]
      1%|          | 1/149 [00:01<04:09,  1.68s/it]
      1%|▏         | 2/149 [00:03<04:03,  1.65s/it]
      2%|▏         | 3/149 [00:04<04:00,  1.65s/it]
      3%|▎         | 4/149 [00:06<03:58,  1.64s/it]
      3%|▎         | 5/149 [00:08<03:56,  1.64s/it]
      4%|▍         | 6/149 [00:09<03:54,  1.64s/it]
      5%|▍         | 7/149 [00:11<03:53,  1.64s/it]
      5%|▌         | 8/149 [00:13<03:50,  1.64s/it]
      6%|▌         | 9/149 [00:14<03:49,  1.64s/it]
      7%|▋         | 10/149 [00:16<03:47,  1.63s/it]
      7%|▋         | 11/149 [00:18<03:45,  1.64s/it]
      8%|▊         | 12/149 [00:19<03:43,  1.63s/it]
      9%|▊         | 13/149 [00:21<03:42,  1.63s/it]
      9%|▉         | 14/149 [00:22<03:40,  1.63s/it]
     10%|█         | 15/149 [00:24<03:38,  1.63s/it]
     11%|█         | 16/149 [00:26<03:36,  1.63s/it]
     11%|█▏        | 17/149 [00:27<03:35,  1.63s/it]
     12%|█▏        | 18/149 [00:29<03:33,  1.63s/it]
     13%|█▎        | 19/149 [00:31<03:32,  1.64s/it]
     13%|█▎        | 20/149 [00:32<03:31,  1.64s/it]
     14%|█▍        | 21/149 [00:34<03:29,  1.63s/it]
     15%|█▍        | 22/149 [00:36<03:30,  1.65s/it]
     15%|█▌        | 23/149 [00:37<03:27,  1.65s/it]
     16%|█▌        | 24/149 [00:39<03:25,  1.64s/it]
     17%|█▋        | 25/149 [00:40<03:23,  1.64s/it]
     17%|█▋        | 26/149 [00:42<03:21,  1.64s/it]
     18%|█▊        | 27/149 [00:44<03:19,  1.64s/it]
     19%|█▉        | 28/149 [00:45<03:18,  1.64s/it]
     19%|█▉        | 29/149 [00:47<03:16,  1.64s/it]
     20%|██        | 30/149 [00:49<03:14,  1.64s/it]
     21%|██        | 31/149 [00:50<03:12,  1.63s/it]
     21%|██▏       | 32/149 [00:52<03:11,  1.64s/it]
     22%|██▏       | 33/149 [00:54<03:09,  1.63s/it]
     23%|██▎       | 34/149 [00:55<03:09,  1.65s/it]
     23%|██▎       | 35/149 [00:57<03:07,  1.64s/it]
     24%|██▍       | 36/149 [00:59<03:05,  1.64s/it]
     25%|██▍       | 37/149 [01:00<03:03,  1.64s/it]
     26%|██▌       | 38/149 [01:02<03:01,  1.64s/it]
     26%|██▌       | 39/149 [01:03<02:59,  1.64s/it]
     27%|██▋       | 40/149 [01:05<02:58,  1.64s/it]
     28%|██▊       | 41/149 [01:07<02:56,  1.63s/it]
     28%|██▊       | 42/149 [01:08<02:54,  1.63s/it]
     29%|██▉       | 43/149 [01:10<02:52,  1.63s/it]
     30%|██▉       | 44/149 [01:12<02:51,  1.63s/it]
     30%|███       | 45/149 [01:13<02:49,  1.63s/it]
     31%|███       | 46/149 [01:15<02:48,  1.63s/it]
     32%|███▏      | 47/149 [01:16<02:46,  1.64s/it]
     32%|███▏      | 48/149 [01:18<02:44,  1.63s/it]
     33%|███▎      | 49/149 [01:20<02:44,  1.65s/it]
     34%|███▎      | 50/149 [01:21<02:42,  1.65s/it]
     34%|███▍      | 51/149 [01:23<02:41,  1.64s/it]
     35%|███▍      | 52/149 [01:25<02:39,  1.64s/it]
     36%|███▌      | 53/149 [01:26<02:37,  1.64s/it]
     36%|███▌      | 54/149 [01:28<02:35,  1.64s/it]
     37%|███▋      | 55/149 [01:30<02:34,  1.64s/it]
     38%|███▊      | 56/149 [01:31<02:32,  1.64s/it]
     38%|███▊      | 57/149 [01:33<02:30,  1.64s/it]
     39%|███▉      | 58/149 [01:35<02:29,  1.64s/it]
     40%|███▉      | 59/149 [01:36<02:27,  1.64s/it]
     40%|████      | 60/149 [01:38<02:25,  1.64s/it]
     41%|████      | 61/149 [01:39<02:24,  1.64s/it]
     42%|████▏     | 62/149 [01:41<02:22,  1.63s/it]
     42%|████▏     | 63/149 [01:43<02:20,  1.63s/it]
     43%|████▎     | 64/149 [01:44<02:18,  1.63s/it]
     44%|████▎     | 65/149 [01:46<02:17,  1.63s/it]
     44%|████▍     | 66/149 [01:48<02:15,  1.63s/it]
     45%|████▍     | 67/149 [01:49<02:13,  1.63s/it]
     46%|████▌     | 68/149 [01:51<02:11,  1.63s/it]
     46%|████▋     | 69/149 [01:52<02:10,  1.63s/it]
     47%|████▋     | 70/149 [01:54<02:08,  1.63s/it]
     48%|████▊     | 71/149 [01:56<02:07,  1.63s/it]
     48%|████▊     | 72/149 [01:57<02:05,  1.63s/it]
     49%|████▉     | 73/149 [01:59<02:03,  1.63s/it]
     50%|████▉     | 74/149 [02:01<02:02,  1.63s/it]
     50%|█████     | 75/149 [02:02<02:00,  1.63s/it]
     51%|█████     | 76/149 [02:04<01:59,  1.64s/it]
     52%|█████▏    | 77/149 [02:06<01:57,  1.63s/it]
     52%|█████▏    | 78/149 [02:07<01:56,  1.64s/it]
     53%|█████▎    | 79/149 [02:09<01:54,  1.63s/it]
     54%|█████▎    | 80/149 [02:10<01:52,  1.64s/it]
     54%|█████▍    | 81/149 [02:12<01:51,  1.63s/it]
     55%|█████▌    | 82/149 [02:14<01:49,  1.63s/it]
     56%|█████▌    | 83/149 [02:15<01:47,  1.63s/it]
     56%|█████▋    | 84/149 [02:17<01:46,  1.63s/it]
     57%|█████▋    | 85/149 [02:19<01:44,  1.63s/it]
     58%|█████▊    | 86/149 [02:20<01:42,  1.63s/it]
     58%|█████▊    | 87/149 [02:22<01:41,  1.63s/it]
     59%|█████▉    | 88/149 [02:23<01:39,  1.63s/it]
     60%|█████▉    | 89/149 [02:25<01:37,  1.63s/it]
     60%|██████    | 90/149 [02:27<01:36,  1.63s/it]
     61%|██████    | 91/149 [02:28<01:34,  1.63s/it]
     62%|██████▏   | 92/149 [02:30<01:33,  1.64s/it]
     62%|██████▏   | 93/149 [02:32<01:31,  1.64s/it]
     63%|██████▎   | 94/149 [02:33<01:29,  1.64s/it]
     64%|██████▍   | 95/149 [02:35<01:28,  1.64s/it]
     64%|██████▍   | 96/149 [02:37<01:26,  1.63s/it]
     65%|██████▌   | 97/149 [02:38<01:25,  1.64s/it]
     66%|██████▌   | 98/149 [02:40<01:23,  1.63s/it]
     66%|██████▋   | 99/149 [02:41<01:21,  1.64s/it]
     67%|██████▋   | 100/149 [02:43<01:20,  1.63s/it]
     68%|██████▊   | 101/149 [02:45<01:18,  1.64s/it]
     68%|██████▊   | 102/149 [02:46<01:17,  1.64s/it]
     69%|██████▉   | 103/149 [02:48<01:15,  1.64s/it]
     70%|██████▉   | 104/149 [02:50<01:13,  1.64s/it]
     70%|███████   | 105/149 [02:51<01:12,  1.64s/it]
     71%|███████   | 106/149 [02:53<01:10,  1.63s/it]
     72%|███████▏  | 107/149 [02:55<01:08,  1.64s/it]
     72%|███████▏  | 108/149 [02:56<01:06,  1.63s/it]
     73%|███████▎  | 109/149 [02:58<01:05,  1.63s/it]
     74%|███████▍  | 110/149 [02:59<01:03,  1.63s/it]
     74%|███████▍  | 111/149 [03:01<01:01,  1.63s/it]
     75%|███████▌  | 112/149 [03:03<01:00,  1.63s/it]
     76%|███████▌  | 113/149 [03:04<00:58,  1.63s/it]
     77%|███████▋  | 114/149 [03:06<00:57,  1.63s/it]
     77%|███████▋  | 115/149 [03:08<00:55,  1.63s/it]
     78%|███████▊  | 116/149 [03:09<00:53,  1.63s/it]
     79%|███████▊  | 117/149 [03:11<00:52,  1.63s/it]
     79%|███████▉  | 118/149 [03:12<00:50,  1.63s/it]
     80%|███████▉  | 119/149 [03:14<00:48,  1.63s/it]
     81%|████████  | 120/149 [03:16<00:47,  1.63s/it]
     81%|████████  | 121/149 [03:17<00:45,  1.63s/it]
     82%|████████▏ | 122/149 [03:19<00:44,  1.63s/it]
     83%|████████▎ | 123/149 [03:21<00:42,  1.63s/it]
     83%|████████▎ | 124/149 [03:22<00:40,  1.64s/it]
     84%|████████▍ | 125/149 [03:24<00:39,  1.63s/it]
     85%|████████▍ | 126/149 [03:26<00:37,  1.64s/it]
     85%|████████▌ | 127/149 [03:27<00:35,  1.63s/it]
     86%|████████▌ | 128/149 [03:29<00:34,  1.63s/it]
     87%|████████▋ | 129/149 [03:30<00:32,  1.64s/it]
     87%|████████▋ | 130/149 [03:32<00:31,  1.64s/it]
     88%|████████▊ | 131/149 [03:34<00:29,  1.63s/it]
     89%|████████▊ | 132/149 [03:35<00:27,  1.63s/it]
     89%|████████▉ | 133/149 [03:37<00:26,  1.63s/it]
     90%|████████▉ | 134/149 [03:39<00:24,  1.63s/it]
     91%|█████████ | 135/149 [03:40<00:22,  1.63s/it]
     91%|█████████▏| 136/149 [03:42<00:21,  1.63s/it]
     92%|█████████▏| 137/149 [03:44<00:19,  1.63s/it]
     93%|█████████▎| 138/149 [03:45<00:17,  1.63s/it]
     93%|█████████▎| 139/149 [03:47<00:16,  1.64s/it]
     94%|█████████▍| 140/149 [03:48<00:14,  1.63s/it]
     95%|█████████▍| 141/149 [03:50<00:13,  1.64s/it]
     95%|█████████▌| 142/149 [03:52<00:11,  1.64s/it]
     96%|█████████▌| 143/149 [03:53<00:09,  1.64s/it]
     97%|█████████▋| 144/149 [03:55<00:08,  1.64s/it]
     97%|█████████▋| 145/149 [03:57<00:06,  1.64s/it]
     98%|█████████▊| 146/149 [03:58<00:04,  1.64s/it]
     99%|█████████▊| 147/149 [04:00<00:03,  1.64s/it]
     99%|█████████▉| 148/149 [04:02<00:01,  1.63s/it]
    100%|██████████| 149/149 [04:03<00:00,  1.64s/it]
    100%|██████████| 149/149 [04:03<00:00,  1.64s/it]




.. GENERATED FROM PYTHON SOURCE LINES 238-252

We obtain the following posterior sample and trajectory

.. container:: image-col

   .. image-sg-ignore:: /auto_examples/images/posterior_sample.png
      :alt: example of posterior sample
      :srcset: /auto_examples/images/posterior_sample.png
      :ignore_missing: true

   .. image-sg-ignore:: /auto_examples/images/posterior_trajectory.gif
      :alt: example of posterior trajectory
      :srcset: /auto_examples/images/posterior_trajectory.gif
      :ignore_missing: true
      :class: custom-gif

.. GENERATED FROM PYTHON SOURCE LINES 255-263

.. note::

    **Reproducibility**: To ensure the reproducibility, if the parameter `rng` is given, the same sample will
    be generated when the same seed is used.
    One can obtain varying samples by using a different seed.

    **Parallel sampling**: one can draw multiple samples in parallel by giving the initial shape, e.g., `x_init = (B, C, H, W)`


.. GENERATED FROM PYTHON SOURCE LINES 265-273

Varying the SDE
---------------

One can also change the underlying SDE for sampling.
For example, we can also use the Variance-Preserving (VP or DDPM) in :class:`deepinv.sampling.VariancePreservingDiffusion`, whose forward drift and diffusion term are defined as:

.. math::
    f(x_t, t) = -\frac{1}{2} \beta(t)x_t \qquad \mbox{ and } \qquad g(t) = \beta(t)  \qquad \mbox{ with } \beta(t) = \beta_{\mathrm{min}}  + t \left( \beta_{\mathrm{max}} - \beta_{\mathrm{min}} \right).

.. GENERATED FROM PYTHON SOURCE LINES 273-316

.. code-block:: Python


    from deepinv.sampling import VariancePreservingDiffusion

    del trajectory
    sde = VariancePreservingDiffusion(device=device, dtype=dtype)
    model = PosteriorDiffusion(
        data_fidelity=dps_fidelity,
        denoiser=denoiser,
        sde=sde,
        solver=solver,
        device=device,
        dtype=dtype,
        verbose=True,
    )

    x_hat_vp, trajectory = model(
        y,
        physics,
        seed=111,
        timesteps=torch.linspace(1, 0.001, 150),
        get_trajectory=True,
    )
    dinv.utils.plot(
        [x_hat, x_hat_vp],
        titles=[
            "posterior sample with VE",
            "posterior sample with VP",
        ],
        save_fn="posterior_sample_ve_vp.png",
        figsize=(figsize * 3, figsize),
    )


    # We can also save the trajectory of the posterior sample
    dinv.utils.save_videos(
        trajectory[::gif_frequency],
        time_dim=0,
        titles=["Posterior sample with VP"],
        save_fn="posterior_trajectory_vp.gif",
        figsize=(figsize, figsize),
    )






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

 .. code-block:: none


      0%|          | 0/149 [00:00<?, ?it/s]
      1%|          | 1/149 [00:01<04:03,  1.65s/it]
      1%|▏         | 2/149 [00:03<04:00,  1.63s/it]
      2%|▏         | 3/149 [00:04<03:59,  1.64s/it]
      3%|▎         | 4/149 [00:06<03:56,  1.63s/it]
      3%|▎         | 5/149 [00:08<03:55,  1.63s/it]
      4%|▍         | 6/149 [00:09<03:53,  1.63s/it]
      5%|▍         | 7/149 [00:11<03:51,  1.63s/it]
      5%|▌         | 8/149 [00:13<03:49,  1.63s/it]
      6%|▌         | 9/149 [00:14<03:48,  1.63s/it]
      7%|▋         | 10/149 [00:16<03:47,  1.64s/it]
      7%|▋         | 11/149 [00:17<03:45,  1.63s/it]
      8%|▊         | 12/149 [00:19<03:43,  1.63s/it]
      9%|▊         | 13/149 [00:21<03:41,  1.63s/it]
      9%|▉         | 14/149 [00:22<03:40,  1.63s/it]
     10%|█         | 15/149 [00:24<03:38,  1.63s/it]
     11%|█         | 16/149 [00:26<03:37,  1.64s/it]
     11%|█▏        | 17/149 [00:27<03:35,  1.63s/it]
     12%|█▏        | 18/149 [00:29<03:34,  1.64s/it]
     13%|█▎        | 19/149 [00:31<03:32,  1.64s/it]
     13%|█▎        | 20/149 [00:32<03:31,  1.64s/it]
     14%|█▍        | 21/149 [00:34<03:29,  1.64s/it]
     15%|█▍        | 22/149 [00:35<03:27,  1.64s/it]
     15%|█▌        | 23/149 [00:37<03:26,  1.64s/it]
     16%|█▌        | 24/149 [00:39<03:24,  1.64s/it]
     17%|█▋        | 25/149 [00:40<03:22,  1.64s/it]
     17%|█▋        | 26/149 [00:42<03:21,  1.64s/it]
     18%|█▊        | 27/149 [00:44<03:19,  1.63s/it]
     19%|█▉        | 28/149 [00:45<03:18,  1.64s/it]
     19%|█▉        | 29/149 [00:47<03:16,  1.63s/it]
     20%|██        | 30/149 [00:49<03:15,  1.64s/it]
     21%|██        | 31/149 [00:50<03:12,  1.63s/it]
     21%|██▏       | 32/149 [00:52<03:11,  1.64s/it]
     22%|██▏       | 33/149 [00:53<03:09,  1.63s/it]
     23%|██▎       | 34/149 [00:55<03:08,  1.63s/it]
     23%|██▎       | 35/149 [00:57<03:06,  1.63s/it]
     24%|██▍       | 36/149 [00:58<03:05,  1.64s/it]
     25%|██▍       | 37/149 [01:00<03:03,  1.63s/it]
     26%|██▌       | 38/149 [01:02<03:01,  1.63s/it]
     26%|██▌       | 39/149 [01:03<02:59,  1.63s/it]
     27%|██▋       | 40/149 [01:05<02:57,  1.63s/it]
     28%|██▊       | 41/149 [01:07<02:56,  1.63s/it]
     28%|██▊       | 42/149 [01:08<02:54,  1.63s/it]
     29%|██▉       | 43/149 [01:10<02:53,  1.64s/it]
     30%|██▉       | 44/149 [01:11<02:51,  1.64s/it]
     30%|███       | 45/149 [01:13<02:50,  1.64s/it]
     31%|███       | 46/149 [01:15<02:48,  1.64s/it]
     32%|███▏      | 47/149 [01:16<02:47,  1.64s/it]
     32%|███▏      | 48/149 [01:18<02:45,  1.64s/it]
     33%|███▎      | 49/149 [01:20<02:43,  1.64s/it]
     34%|███▎      | 50/149 [01:21<02:42,  1.64s/it]
     34%|███▍      | 51/149 [01:23<02:40,  1.64s/it]
     35%|███▍      | 52/149 [01:25<02:38,  1.64s/it]
     36%|███▌      | 53/149 [01:26<02:37,  1.64s/it]
     36%|███▌      | 54/149 [01:28<02:35,  1.64s/it]
     37%|███▋      | 55/149 [01:29<02:34,  1.64s/it]
     38%|███▊      | 56/149 [01:31<02:32,  1.64s/it]
     38%|███▊      | 57/149 [01:33<02:31,  1.64s/it]
     39%|███▉      | 58/149 [01:34<02:29,  1.64s/it]
     40%|███▉      | 59/149 [01:36<02:27,  1.64s/it]
     40%|████      | 60/149 [01:38<02:25,  1.63s/it]
     41%|████      | 61/149 [01:39<02:24,  1.64s/it]
     42%|████▏     | 62/149 [01:41<02:22,  1.63s/it]
     42%|████▏     | 63/149 [01:43<02:20,  1.64s/it]
     43%|████▎     | 64/149 [01:44<02:18,  1.63s/it]
     44%|████▎     | 65/149 [01:46<02:17,  1.64s/it]
     44%|████▍     | 66/149 [01:47<02:15,  1.64s/it]
     45%|████▍     | 67/149 [01:49<02:14,  1.64s/it]
     46%|████▌     | 68/149 [01:51<02:12,  1.64s/it]
     46%|████▋     | 69/149 [01:52<02:10,  1.64s/it]
     47%|████▋     | 70/149 [01:54<02:08,  1.63s/it]
     48%|████▊     | 71/149 [01:56<02:07,  1.64s/it]
     48%|████▊     | 72/149 [01:57<02:05,  1.63s/it]
     49%|████▉     | 73/149 [01:59<02:04,  1.64s/it]
     50%|████▉     | 74/149 [02:01<02:02,  1.63s/it]
     50%|█████     | 75/149 [02:02<02:00,  1.63s/it]
     51%|█████     | 76/149 [02:04<01:59,  1.63s/it]
     52%|█████▏    | 77/149 [02:05<01:57,  1.63s/it]
     52%|█████▏    | 78/149 [02:07<01:55,  1.63s/it]
     53%|█████▎    | 79/149 [02:09<01:54,  1.63s/it]
     54%|█████▎    | 80/149 [02:10<01:52,  1.63s/it]
     54%|█████▍    | 81/149 [02:12<01:50,  1.63s/it]
     55%|█████▌    | 82/149 [02:14<01:49,  1.63s/it]
     56%|█████▌    | 83/149 [02:15<01:47,  1.63s/it]
     56%|█████▋    | 84/149 [02:17<01:46,  1.64s/it]
     57%|█████▋    | 85/149 [02:19<01:44,  1.63s/it]
     58%|█████▊    | 86/149 [02:20<01:43,  1.64s/it]
     58%|█████▊    | 87/149 [02:22<01:41,  1.64s/it]
     59%|█████▉    | 88/149 [02:23<01:39,  1.64s/it]
     60%|█████▉    | 89/149 [02:25<01:38,  1.64s/it]
     60%|██████    | 90/149 [02:27<01:36,  1.64s/it]
     61%|██████    | 91/149 [02:28<01:35,  1.64s/it]
     62%|██████▏   | 92/149 [02:30<01:33,  1.64s/it]
     62%|██████▏   | 93/149 [02:32<01:31,  1.64s/it]
     63%|██████▎   | 94/149 [02:33<01:30,  1.64s/it]
     64%|██████▍   | 95/149 [02:35<01:28,  1.64s/it]
     64%|██████▍   | 96/149 [02:37<01:26,  1.64s/it]
     65%|██████▌   | 97/149 [02:38<01:25,  1.64s/it]
     66%|██████▌   | 98/149 [02:40<01:23,  1.64s/it]
     66%|██████▋   | 99/149 [02:41<01:21,  1.64s/it]
     67%|██████▋   | 100/149 [02:43<01:20,  1.64s/it]
     68%|██████▊   | 101/149 [02:45<01:18,  1.63s/it]
     68%|██████▊   | 102/149 [02:46<01:16,  1.63s/it]
     69%|██████▉   | 103/149 [02:48<01:15,  1.63s/it]
     70%|██████▉   | 104/149 [02:50<01:13,  1.64s/it]
     70%|███████   | 105/149 [02:51<01:11,  1.63s/it]
     71%|███████   | 106/149 [02:53<01:10,  1.64s/it]
     72%|███████▏  | 107/149 [02:55<01:08,  1.64s/it]
     72%|███████▏  | 108/149 [02:56<01:07,  1.64s/it]
     73%|███████▎  | 109/149 [02:58<01:05,  1.63s/it]
     74%|███████▍  | 110/149 [02:59<01:03,  1.64s/it]
     74%|███████▍  | 111/149 [03:01<01:02,  1.65s/it]
     75%|███████▌  | 112/149 [03:03<01:00,  1.65s/it]
     76%|███████▌  | 113/149 [03:04<00:59,  1.64s/it]
     77%|███████▋  | 114/149 [03:06<00:57,  1.65s/it]
     77%|███████▋  | 115/149 [03:08<00:55,  1.64s/it]
     78%|███████▊  | 116/149 [03:09<00:54,  1.64s/it]
     79%|███████▊  | 117/149 [03:11<00:52,  1.64s/it]
     79%|███████▉  | 118/149 [03:13<00:50,  1.64s/it]
     80%|███████▉  | 119/149 [03:14<00:49,  1.64s/it]
     81%|████████  | 120/149 [03:16<00:47,  1.64s/it]
     81%|████████  | 121/149 [03:17<00:45,  1.64s/it]
     82%|████████▏ | 122/149 [03:19<00:44,  1.63s/it]
     83%|████████▎ | 123/149 [03:21<00:42,  1.64s/it]
     83%|████████▎ | 124/149 [03:22<00:40,  1.63s/it]
     84%|████████▍ | 125/149 [03:24<00:39,  1.63s/it]
     85%|████████▍ | 126/149 [03:26<00:37,  1.63s/it]
     85%|████████▌ | 127/149 [03:27<00:35,  1.63s/it]
     86%|████████▌ | 128/149 [03:29<00:34,  1.64s/it]
     87%|████████▋ | 129/149 [03:31<00:32,  1.64s/it]
     87%|████████▋ | 130/149 [03:32<00:31,  1.64s/it]
     88%|████████▊ | 131/149 [03:34<00:29,  1.64s/it]
     89%|████████▊ | 132/149 [03:35<00:27,  1.64s/it]
     89%|████████▉ | 133/149 [03:37<00:26,  1.64s/it]
     90%|████████▉ | 134/149 [03:39<00:24,  1.64s/it]
     91%|█████████ | 135/149 [03:40<00:22,  1.64s/it]
     91%|█████████▏| 136/149 [03:42<00:21,  1.64s/it]
     92%|█████████▏| 137/149 [03:44<00:19,  1.64s/it]
     93%|█████████▎| 138/149 [03:45<00:18,  1.64s/it]
     93%|█████████▎| 139/149 [03:47<00:16,  1.64s/it]
     94%|█████████▍| 140/149 [03:49<00:14,  1.64s/it]
     95%|█████████▍| 141/149 [03:50<00:13,  1.64s/it]
     95%|█████████▌| 142/149 [03:52<00:11,  1.64s/it]
     96%|█████████▌| 143/149 [03:54<00:09,  1.64s/it]
     97%|█████████▋| 144/149 [03:55<00:08,  1.64s/it]
     97%|█████████▋| 145/149 [03:57<00:06,  1.64s/it]
     98%|█████████▊| 146/149 [03:58<00:04,  1.64s/it]
     99%|█████████▊| 147/149 [04:00<00:03,  1.64s/it]
     99%|█████████▉| 148/149 [04:02<00:01,  1.63s/it]
    100%|██████████| 149/149 [04:03<00:00,  1.64s/it]
    100%|██████████| 149/149 [04:03<00:00,  1.64s/it]




.. GENERATED FROM PYTHON SOURCE LINES 336-357

We can comparing the sampling trajectory depending on the underlying SDE

.. container:: image-col

   .. image-sg-ignore:: /auto_examples/images/posterior_sample_ve_vp.png
      :alt: posterior sample with VP
      :srcset: /auto_examples/images/posterior_sample_ve_vp.png
      :ignore_missing: true
   .. container:: image-row

      .. image-sg-ignore:: /auto_examples/images/posterior_trajectory.gif
          :alt: posterior trajectory with VE
          :srcset: /auto_examples/images/posterior_trajectory.gif
          :ignore_missing: true
          :class: custom-gif

      .. image-sg-ignore:: /auto_examples/images/posterior_trajectory_vp.gif
          :alt: posterior trajectory with VP
          :srcset: /auto_examples/images/posterior_trajectory_vp.gif
          :ignore_missing: true
          :class: custom-gif

.. GENERATED FROM PYTHON SOURCE LINES 359-365

Plug-and-play Posterior Sampling with arbitrary denoisers
---------------------------------------------------------

The :class:`deepinv.sampling.PosteriorDiffusion` class can be used together with any (well-trained) denoisers for posterior sampling.
For example, we can use the :class:`deepinv.models.DRUNet` for posterior sampling.
We can also change the underlying SDE, for example change the `sigma_max` value.

.. GENERATED FROM PYTHON SOURCE LINES 365-429

.. code-block:: Python


    del trajectory  # clean memory
    sigma_min = 0.001
    sigma_max = 10.0
    rng = torch.Generator(device)
    dtype = torch.float32
    timesteps = torch.linspace(1, 0.001, 250)
    solver = EulerSolver(timesteps=timesteps, rng=rng)
    denoiser = dinv.models.DRUNet(pretrained="download").to(device)

    sde = VarianceExplodingDiffusion(
        sigma_max=sigma_max, sigma_min=sigma_min, alpha=0.75, device=device, dtype=dtype
    )
    x = dinv.utils.load_example(
        "butterfly.png",
        img_size=256,
        resize_mode="resize",
    ).to(device)

    mask = torch.ones_like(x)
    mask[..., 70:150, 120:180] = 0
    physics = dinv.physics.Inpainting(
        mask=mask,
        img_size=x.shape[1:],
        device=device,
    )

    y = physics(x)
    model = PosteriorDiffusion(
        data_fidelity=DPSDataFidelity(denoiser=denoiser, weight=0.3),
        denoiser=denoiser,
        sde=sde,
        solver=solver,
        dtype=dtype,
        device=device,
        verbose=True,
    )

    # To perform posterior sampling, we need to provide the measurements, the physics and the solver.
    x_hat, trajectory = model(
        y=y,
        physics=physics,
        seed=12,
        get_trajectory=True,
    )

    # Here, we plot the original image, the measurement and the posterior sample
    dinv.utils.plot(
        [x, y, x_hat.clip(0, 1)],
        titles=["Original", "Measurement", "Posterior sample DRUNet"],
        figsize=(figsize * 3, figsize),
        save_fn="posterior_sample_DRUNet.png",
    )

    # We can also save the trajectory of the posterior sample
    dinv.utils.save_videos(
        trajectory[::gif_frequency].clip(0, 1),
        time_dim=0,
        titles=["Posterior trajectory DRUNet"],
        save_fn="posterior_sample_DRUNet.gif",
        figsize=(figsize, figsize),
    )






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

 .. code-block:: none


      0%|          | 0/249 [00:00<?, ?it/s]
      0%|          | 1/249 [00:04<20:21,  4.93s/it]
      1%|          | 2/249 [00:09<20:11,  4.91s/it]
      1%|          | 3/249 [00:14<20:04,  4.90s/it]
      2%|▏         | 4/249 [00:19<19:59,  4.90s/it]
      2%|▏         | 5/249 [00:24<19:55,  4.90s/it]
      2%|▏         | 6/249 [00:29<19:51,  4.90s/it]
      3%|▎         | 7/249 [00:34<19:46,  4.90s/it]
      3%|▎         | 8/249 [00:39<19:41,  4.90s/it]
      4%|▎         | 9/249 [00:44<19:36,  4.90s/it]
      4%|▍         | 10/249 [00:49<19:30,  4.90s/it]
      4%|▍         | 11/249 [00:53<19:25,  4.90s/it]
      5%|▍         | 12/249 [00:58<19:19,  4.89s/it]
      5%|▌         | 13/249 [01:03<19:17,  4.90s/it]
      6%|▌         | 14/249 [01:08<19:12,  4.90s/it]
      6%|▌         | 15/249 [01:13<19:06,  4.90s/it]
      6%|▋         | 16/249 [01:18<19:02,  4.90s/it]
      7%|▋         | 17/249 [01:23<18:56,  4.90s/it]
      7%|▋         | 18/249 [01:28<18:51,  4.90s/it]
      8%|▊         | 19/249 [01:33<18:46,  4.90s/it]
      8%|▊         | 20/249 [01:38<18:42,  4.90s/it]
      8%|▊         | 21/249 [01:42<18:37,  4.90s/it]
      9%|▉         | 22/249 [01:47<18:32,  4.90s/it]
      9%|▉         | 23/249 [01:52<18:28,  4.90s/it]
     10%|▉         | 24/249 [01:57<18:21,  4.90s/it]
     10%|█         | 25/249 [02:02<18:18,  4.90s/it]
     10%|█         | 26/249 [02:07<18:13,  4.90s/it]
     11%|█         | 27/249 [02:12<18:07,  4.90s/it]
     11%|█         | 28/249 [02:17<18:03,  4.90s/it]
     12%|█▏        | 29/249 [02:22<17:58,  4.90s/it]
     12%|█▏        | 30/249 [02:27<17:52,  4.90s/it]
     12%|█▏        | 31/249 [02:31<17:48,  4.90s/it]
     13%|█▎        | 32/249 [02:36<17:43,  4.90s/it]
     13%|█▎        | 33/249 [02:41<17:37,  4.89s/it]
     14%|█▎        | 34/249 [02:46<17:32,  4.90s/it]
     14%|█▍        | 35/249 [02:51<17:26,  4.89s/it]
     14%|█▍        | 36/249 [02:56<17:21,  4.89s/it]
     15%|█▍        | 37/249 [03:01<17:17,  4.89s/it]
     15%|█▌        | 38/249 [03:06<17:11,  4.89s/it]
     16%|█▌        | 39/249 [03:11<17:06,  4.89s/it]
     16%|█▌        | 40/249 [03:15<17:02,  4.89s/it]
     16%|█▋        | 41/249 [03:20<16:58,  4.90s/it]
     17%|█▋        | 42/249 [03:25<16:51,  4.89s/it]
     17%|█▋        | 43/249 [03:30<16:47,  4.89s/it]
     18%|█▊        | 44/249 [03:35<16:43,  4.89s/it]
     18%|█▊        | 45/249 [03:40<16:38,  4.89s/it]
     18%|█▊        | 46/249 [03:45<16:35,  4.90s/it]
     19%|█▉        | 47/249 [03:50<16:30,  4.90s/it]
     19%|█▉        | 48/249 [03:55<16:27,  4.91s/it]
     20%|█▉        | 49/249 [04:00<16:24,  4.92s/it]
     20%|██        | 50/249 [04:05<16:18,  4.92s/it]
     20%|██        | 51/249 [04:09<16:14,  4.92s/it]
     21%|██        | 52/249 [04:14<16:09,  4.92s/it]
     21%|██▏       | 53/249 [04:19<16:05,  4.93s/it]
     22%|██▏       | 54/249 [04:24<16:00,  4.92s/it]
     22%|██▏       | 55/249 [04:29<15:56,  4.93s/it]
     22%|██▏       | 56/249 [04:34<15:51,  4.93s/it]
     23%|██▎       | 57/249 [04:39<15:46,  4.93s/it]
     23%|██▎       | 58/249 [04:44<15:41,  4.93s/it]
     24%|██▎       | 59/249 [04:49<15:35,  4.92s/it]
     24%|██▍       | 60/249 [04:54<15:29,  4.92s/it]
     24%|██▍       | 61/249 [04:59<15:27,  4.94s/it]
     25%|██▍       | 62/249 [05:04<15:21,  4.93s/it]
     25%|██▌       | 63/249 [05:09<15:13,  4.91s/it]
     26%|██▌       | 64/249 [05:13<15:08,  4.91s/it]
     26%|██▌       | 65/249 [05:18<15:03,  4.91s/it]
     27%|██▋       | 66/249 [05:23<14:57,  4.91s/it]
     27%|██▋       | 67/249 [05:28<14:52,  4.90s/it]
     27%|██▋       | 68/249 [05:33<14:47,  4.91s/it]
     28%|██▊       | 69/249 [05:38<14:42,  4.90s/it]
     28%|██▊       | 70/249 [05:43<14:39,  4.91s/it]
     29%|██▊       | 71/249 [05:48<14:34,  4.91s/it]
     29%|██▉       | 72/249 [05:53<14:29,  4.91s/it]
     29%|██▉       | 73/249 [05:58<14:24,  4.91s/it]
     30%|██▉       | 74/249 [06:03<14:19,  4.91s/it]
     30%|███       | 75/249 [06:07<14:15,  4.92s/it]
     31%|███       | 76/249 [06:12<14:10,  4.92s/it]
     31%|███       | 77/249 [06:17<14:05,  4.91s/it]
     31%|███▏      | 78/249 [06:22<14:00,  4.91s/it]
     32%|███▏      | 79/249 [06:27<13:55,  4.91s/it]
     32%|███▏      | 80/249 [06:32<13:50,  4.91s/it]
     33%|███▎      | 81/249 [06:37<13:45,  4.91s/it]
     33%|███▎      | 82/249 [06:42<13:39,  4.91s/it]
     33%|███▎      | 83/249 [06:47<13:34,  4.91s/it]
     34%|███▎      | 84/249 [06:52<13:31,  4.92s/it]
     34%|███▍      | 85/249 [06:57<13:24,  4.91s/it]
     35%|███▍      | 86/249 [07:01<13:21,  4.91s/it]
     35%|███▍      | 87/249 [07:06<13:16,  4.91s/it]
     35%|███▌      | 88/249 [07:11<13:11,  4.91s/it]
     36%|███▌      | 89/249 [07:16<13:06,  4.92s/it]
     36%|███▌      | 90/249 [07:21<13:01,  4.91s/it]
     37%|███▋      | 91/249 [07:26<12:56,  4.91s/it]
     37%|███▋      | 92/249 [07:31<12:53,  4.93s/it]
     37%|███▋      | 93/249 [07:36<12:49,  4.93s/it]
     38%|███▊      | 94/249 [07:41<12:45,  4.94s/it]
     38%|███▊      | 95/249 [07:46<12:41,  4.94s/it]
     39%|███▊      | 96/249 [07:51<12:36,  4.94s/it]
     39%|███▉      | 97/249 [07:56<12:31,  4.94s/it]
     39%|███▉      | 98/249 [08:01<12:27,  4.95s/it]
     40%|███▉      | 99/249 [08:06<12:22,  4.95s/it]
     40%|████      | 100/249 [08:11<12:18,  4.95s/it]
     41%|████      | 101/249 [08:16<12:13,  4.95s/it]
     41%|████      | 102/249 [08:21<12:08,  4.95s/it]
     41%|████▏     | 103/249 [08:25<12:03,  4.95s/it]
     42%|████▏     | 104/249 [08:30<11:58,  4.96s/it]
     42%|████▏     | 105/249 [08:35<11:52,  4.95s/it]
     43%|████▎     | 106/249 [08:40<11:48,  4.95s/it]
     43%|████▎     | 107/249 [08:45<11:43,  4.96s/it]
     43%|████▎     | 108/249 [08:50<11:39,  4.96s/it]
     44%|████▍     | 109/249 [08:55<11:33,  4.95s/it]
     44%|████▍     | 110/249 [09:00<11:28,  4.95s/it]
     45%|████▍     | 111/249 [09:05<11:22,  4.95s/it]
     45%|████▍     | 112/249 [09:10<11:17,  4.95s/it]
     45%|████▌     | 113/249 [09:15<11:13,  4.95s/it]
     46%|████▌     | 114/249 [09:20<11:08,  4.95s/it]
     46%|████▌     | 115/249 [09:25<11:04,  4.96s/it]
     47%|████▋     | 116/249 [09:30<10:59,  4.96s/it]
     47%|████▋     | 117/249 [09:35<10:53,  4.95s/it]
     47%|████▋     | 118/249 [09:40<10:48,  4.95s/it]
     48%|████▊     | 119/249 [09:45<10:42,  4.94s/it]
     48%|████▊     | 120/249 [09:50<10:37,  4.94s/it]
     49%|████▊     | 121/249 [09:55<10:33,  4.95s/it]
     49%|████▉     | 122/249 [10:00<10:28,  4.95s/it]
     49%|████▉     | 123/249 [10:04<10:22,  4.94s/it]
     50%|████▉     | 124/249 [10:09<10:17,  4.94s/it]
     50%|█████     | 125/249 [10:14<10:12,  4.94s/it]
     51%|█████     | 126/249 [10:19<10:08,  4.94s/it]
     51%|█████     | 127/249 [10:24<10:02,  4.94s/it]
     51%|█████▏    | 128/249 [10:29<09:58,  4.95s/it]
     52%|█████▏    | 129/249 [10:34<09:54,  4.95s/it]
     52%|█████▏    | 130/249 [10:39<09:48,  4.94s/it]
     53%|█████▎    | 131/249 [10:44<09:44,  4.95s/it]
     53%|█████▎    | 132/249 [10:49<09:38,  4.95s/it]
     53%|█████▎    | 133/249 [10:54<09:34,  4.96s/it]
     54%|█████▍    | 134/249 [10:59<09:28,  4.95s/it]
     54%|█████▍    | 135/249 [11:04<09:24,  4.95s/it]
     55%|█████▍    | 136/249 [11:09<09:19,  4.95s/it]
     55%|█████▌    | 137/249 [11:14<09:13,  4.94s/it]
     55%|█████▌    | 138/249 [11:19<09:08,  4.94s/it]
     56%|█████▌    | 139/249 [11:24<09:03,  4.94s/it]
     56%|█████▌    | 140/249 [11:29<08:58,  4.94s/it]
     57%|█████▋    | 141/249 [11:33<08:53,  4.94s/it]
     57%|█████▋    | 142/249 [11:38<08:48,  4.94s/it]
     57%|█████▋    | 143/249 [11:43<08:41,  4.92s/it]
     58%|█████▊    | 144/249 [11:48<08:36,  4.92s/it]
     58%|█████▊    | 145/249 [11:53<08:31,  4.92s/it]
     59%|█████▊    | 146/249 [11:58<08:25,  4.91s/it]
     59%|█████▉    | 147/249 [12:03<08:21,  4.92s/it]
     59%|█████▉    | 148/249 [12:08<08:18,  4.93s/it]
     60%|█████▉    | 149/249 [12:13<08:13,  4.94s/it]
     60%|██████    | 150/249 [12:18<08:08,  4.94s/it]
     61%|██████    | 151/249 [12:23<08:04,  4.94s/it]
     61%|██████    | 152/249 [12:28<07:59,  4.94s/it]
     61%|██████▏   | 153/249 [12:33<07:53,  4.94s/it]
     62%|██████▏   | 154/249 [12:38<07:49,  4.94s/it]
     62%|██████▏   | 155/249 [12:43<07:45,  4.95s/it]
     63%|██████▎   | 156/249 [12:47<07:39,  4.94s/it]
     63%|██████▎   | 157/249 [12:52<07:34,  4.94s/it]
     63%|██████▎   | 158/249 [12:57<07:29,  4.94s/it]
     64%|██████▍   | 159/249 [13:02<07:25,  4.95s/it]
     64%|██████▍   | 160/249 [13:07<07:19,  4.94s/it]
     65%|██████▍   | 161/249 [13:12<07:15,  4.94s/it]
     65%|██████▌   | 162/249 [13:17<07:10,  4.95s/it]
     65%|██████▌   | 163/249 [13:22<07:05,  4.95s/it]
     66%|██████▌   | 164/249 [13:27<07:00,  4.95s/it]
     66%|██████▋   | 165/249 [13:32<06:55,  4.95s/it]
     67%|██████▋   | 166/249 [13:37<06:50,  4.94s/it]
     67%|██████▋   | 167/249 [13:42<06:45,  4.94s/it]
     67%|██████▋   | 168/249 [13:47<06:39,  4.94s/it]
     68%|██████▊   | 169/249 [13:52<06:35,  4.94s/it]
     68%|██████▊   | 170/249 [13:57<06:30,  4.94s/it]
     69%|██████▊   | 171/249 [14:02<06:26,  4.95s/it]
     69%|██████▉   | 172/249 [14:07<06:20,  4.94s/it]
     69%|██████▉   | 173/249 [14:12<06:15,  4.94s/it]
     70%|██████▉   | 174/249 [14:16<06:10,  4.95s/it]
     70%|███████   | 175/249 [14:21<06:05,  4.94s/it]
     71%|███████   | 176/249 [14:26<06:00,  4.94s/it]
     71%|███████   | 177/249 [14:31<05:56,  4.95s/it]
     71%|███████▏  | 178/249 [14:36<05:51,  4.95s/it]
     72%|███████▏  | 179/249 [14:41<05:46,  4.95s/it]
     72%|███████▏  | 180/249 [14:46<05:41,  4.95s/it]
     73%|███████▎  | 181/249 [14:51<05:36,  4.95s/it]
     73%|███████▎  | 182/249 [14:56<05:32,  4.96s/it]
     73%|███████▎  | 183/249 [15:01<05:26,  4.95s/it]
     74%|███████▍  | 184/249 [15:06<05:21,  4.95s/it]
     74%|███████▍  | 185/249 [15:11<05:16,  4.95s/it]
     75%|███████▍  | 186/249 [15:16<05:12,  4.95s/it]
     75%|███████▌  | 187/249 [15:21<05:07,  4.95s/it]
     76%|███████▌  | 188/249 [15:26<05:02,  4.96s/it]
     76%|███████▌  | 189/249 [15:31<04:57,  4.95s/it]
     76%|███████▋  | 190/249 [15:36<04:52,  4.96s/it]
     77%|███████▋  | 191/249 [15:41<04:47,  4.95s/it]
     77%|███████▋  | 192/249 [15:46<04:42,  4.95s/it]
     78%|███████▊  | 193/249 [15:51<04:37,  4.95s/it]
     78%|███████▊  | 194/249 [15:55<04:32,  4.95s/it]
     78%|███████▊  | 195/249 [16:00<04:27,  4.96s/it]
     79%|███████▊  | 196/249 [16:05<04:22,  4.95s/it]
     79%|███████▉  | 197/249 [16:10<04:17,  4.96s/it]
     80%|███████▉  | 198/249 [16:15<04:12,  4.96s/it]
     80%|███████▉  | 199/249 [16:20<04:07,  4.96s/it]
     80%|████████  | 200/249 [16:25<04:02,  4.95s/it]
     81%|████████  | 201/249 [16:30<03:57,  4.96s/it]
     81%|████████  | 202/249 [16:35<03:52,  4.95s/it]
     82%|████████▏ | 203/249 [16:40<03:47,  4.96s/it]
     82%|████████▏ | 204/249 [16:45<03:43,  4.96s/it]
     82%|████████▏ | 205/249 [16:50<03:37,  4.95s/it]
     83%|████████▎ | 206/249 [16:55<03:32,  4.95s/it]
     83%|████████▎ | 207/249 [17:00<03:28,  4.95s/it]
     84%|████████▎ | 208/249 [17:05<03:22,  4.95s/it]
     84%|████████▍ | 209/249 [17:10<03:17,  4.95s/it]
     84%|████████▍ | 210/249 [17:15<03:12,  4.94s/it]
     85%|████████▍ | 211/249 [17:20<03:08,  4.95s/it]
     85%|████████▌ | 212/249 [17:25<03:02,  4.94s/it]
     86%|████████▌ | 213/249 [17:30<02:58,  4.95s/it]
     86%|████████▌ | 214/249 [17:35<02:53,  4.96s/it]
     86%|████████▋ | 215/249 [17:39<02:48,  4.96s/it]
     87%|████████▋ | 216/249 [17:44<02:43,  4.96s/it]
     87%|████████▋ | 217/249 [17:49<02:38,  4.95s/it]
     88%|████████▊ | 218/249 [17:54<02:33,  4.95s/it]
     88%|████████▊ | 219/249 [17:59<02:28,  4.94s/it]
     88%|████████▊ | 220/249 [18:04<02:23,  4.94s/it]
     89%|████████▉ | 221/249 [18:09<02:18,  4.94s/it]
     89%|████████▉ | 222/249 [18:14<02:13,  4.94s/it]
     90%|████████▉ | 223/249 [18:19<02:08,  4.94s/it]
     90%|████████▉ | 224/249 [18:24<02:03,  4.93s/it]
     90%|█████████ | 225/249 [18:29<01:58,  4.93s/it]
     91%|█████████ | 226/249 [18:34<01:53,  4.93s/it]
     91%|█████████ | 227/249 [18:39<01:48,  4.93s/it]
     92%|█████████▏| 228/249 [18:44<01:43,  4.93s/it]
     92%|█████████▏| 229/249 [18:49<01:38,  4.94s/it]
     92%|█████████▏| 230/249 [18:54<01:33,  4.94s/it]
     93%|█████████▎| 231/249 [18:58<01:28,  4.94s/it]
     93%|█████████▎| 232/249 [19:03<01:24,  4.94s/it]
     94%|█████████▎| 233/249 [19:08<01:18,  4.94s/it]
     94%|█████████▍| 234/249 [19:13<01:14,  4.94s/it]
     94%|█████████▍| 235/249 [19:18<01:09,  4.94s/it]
     95%|█████████▍| 236/249 [19:23<01:04,  4.94s/it]
     95%|█████████▌| 237/249 [19:28<00:59,  4.94s/it]
     96%|█████████▌| 238/249 [19:33<00:54,  4.93s/it]
     96%|█████████▌| 239/249 [19:38<00:49,  4.93s/it]
     96%|█████████▋| 240/249 [19:43<00:44,  4.93s/it]
     97%|█████████▋| 241/249 [19:48<00:39,  4.93s/it]
     97%|█████████▋| 242/249 [19:53<00:34,  4.93s/it]
     98%|█████████▊| 243/249 [19:58<00:29,  4.93s/it]
     98%|█████████▊| 244/249 [20:03<00:24,  4.93s/it]
     98%|█████████▊| 245/249 [20:08<00:19,  4.93s/it]
     99%|█████████▉| 246/249 [20:12<00:14,  4.93s/it]
     99%|█████████▉| 247/249 [20:17<00:09,  4.93s/it]
    100%|█████████▉| 248/249 [20:22<00:04,  4.93s/it]
    100%|██████████| 249/249 [20:27<00:00,  4.93s/it]
    100%|██████████| 249/249 [20:27<00:00,  4.93s/it]




.. GENERATED FROM PYTHON SOURCE LINES 451-465

We obtain the following posterior trajectory

.. container:: image-col

   .. image-sg-ignore:: /auto_examples/images/posterior_sample_DRUNet.png
      :alt: posterior sample DRUNet
      :srcset: /auto_examples/images/posterior_sample_DRUNet.png
      :ignore_missing: true

   .. image-sg-ignore:: /auto_examples/images/posterior_sample_DRUNet.gif
      :alt: posterior trajectory DRUNet
      :srcset: /auto_examples/images/posterior_sample_DRUNet.gif
      :ignore_missing: true
      :class: custom-gif


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

   **Total running time of the script:** (30 minutes 49.614 seconds)


.. _sphx_glr_download_auto_examples_sampling_demo_diffusion_sde.py:

.. only:: html

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

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

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

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

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

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

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


.. only:: html

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

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