.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/06-learning/01-plug-and-play.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_06-learning_01-plug-and-play.py>`
        to download the full example code.

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

.. _sphx_glr_auto_examples_06-learning_01-plug-and-play.py:


================================
10. Plug-and-play reconstruction
================================

In lesson 9 the prior knowledge was a signal model; in lesson 5 it was an
explicit penalty, a wavelet or total-variation term with a weight. In this
lesson the prior is a pretrained denoising network, used in place of the
penalty inside the same ADMM and FISTA iterations.

ADMM and FISTA use the penalty :math:`g` only through its proximal operator,

.. math::

   \operatorname{prox}_{\gamma g}(v) = \arg\min_x \; \tfrac12 \|x - v\|_2^2
   + \gamma\, g(x).

This is the maximum a posteriori estimate of an image observed as
:math:`v` in white Gaussian noise, under the prior :math:`\exp(-g)`: the
proximal operator is a denoiser. Plug-and-play (PnP) regularization
[#venkatakrishnan]_ [#ahmad]_ replaces it by an image denoiser
:math:`D_\sigma`, without an explicit :math:`g`. Each iteration alternates a
data-consistency step on the measured k-space and a denoising step, and the
denoiser's noise level :math:`\sigma` takes the role of the regularization
weight.

The denoiser here is DRUNet [#zhang]_, with the weights ``deepinv``
distributes. It was trained to remove Gaussian noise from grayscale
photographs, not from MR images.

**Learning objectives**

- Wrap a pretrained denoiser as :class:`bartorch.priors.ImplicitPrior` and
  hand it to :func:`bartorch.optim.admm` and :func:`bartorch.optim.fista` in
  place of a :mod:`bartorch.priors` term.
- Compare the result with total variation on the same data, in the images,
  the error maps and an enlarged region.
- Change the denoiser's noise level and recognise under- and
  over-regularization.

Previous: :doc:`../05-model-based/01-quantitative-models`. Next:
:doc:`02-modl-with-admm`, where you train the denoiser through the iteration
instead of using a pretrained one.

The pretrained weights, about 125 MB, are downloaded the first time you run
this. The network runs once per iteration, and on a CPU that is most of the
run time.

.. GENERATED FROM PYTHON SOURCE LINES 51-191








.. GENERATED FROM PYTHON SOURCE LINES 192-207

Simulated acquisition
---------------------

The object is the BrainWeb slice [#brainweb]_ of lessons 2 and 5, a
:math:`T_1`-weighted spin-echo image, here at 128 by 128; the code that
builds it is hidden. You acquire a quarter of the phase encodes
(:math:`R = 4`), drawn at random around a 16-line ACS block, with eight
channels of BART's analytical head coil and complex Gaussian noise of
variance :math:`10^{-3}`. At that acceleration and noise level, both noise
amplification and incoherent aliasing limit a reconstruction without a
prior.

The SENSE operator uses the sensitivities the data were simulated with, so
that the reconstructions below differ only in the prior. Lesson 3 is about
estimating them.

.. GENERATED FROM PYTHON SOURCE LINES 207-235

.. code-block:: Python

    import torch

    import bartorch
    import bartorch.tools as bt

    from bartorch import linop, optim, priors

    COILS = 8
    ACCELERATION = 4
    CALIBRATION = 16

    sensitivities = bt.coils(t=bt.grid(D=(SIZE, SIZE, 1)), n=COILS)[:, 0]
    sensitivities /= bartorch.rss(sensitivities, axes=(0,), keepdim=True)
    pattern = simulate(ACCELERATION, CALIBRATION)  # (y, 1)

    A = linop.CartesianSense(sensitivities, (SIZE, SIZE), pattern=pattern)
    data = bt.noise(A(image), n=1e-3, s=42) * pattern
    zero_filled = A.H(data)





.. image-sg:: /auto_examples/06-learning/images/sphx_glr_01-plug-and-play_001.png
   :alt: the object, sampling pattern, zero-filled
   :srcset: /auto_examples/06-learning/images/sphx_glr_01-plug-and-play_001.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 236-247

The zero-filled image, :math:`A^H y`, carries the incoherent aliasing of
the random sampling, spread over the whole field of view, on top of the
thermal noise. The orange box marks the region enlarged below.

Explicit penalty
----------------

The reference is total variation under ADMM, as in lesson 5. Of the
weights 0.002, 0.005, 0.01 and 0.02, 0.005 gives the lowest NRMSE against
the phantom and 0.02 the highest SSIM; 0.01 is within 0.001 of that NRMSE
with an SSIM close to the highest, and is used here.

.. GENERATED FROM PYTHON SOURCE LINES 247-263

.. code-block:: Python

    total_variation = optim.admm(
        data,
        A,
        priors.TotalVariation((-1, -2), 0.01),
        maxiter=60,
        rho=0.1,
    )





.. image-sg:: /auto_examples/06-learning/images/sphx_glr_01-plug-and-play_002.png
   :alt: total variation, its error
   :srcset: /auto_examples/06-learning/images/sphx_glr_01-plug-and-play_002.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 264-281

Total variation removes most of the aliasing and the noise. At this weight
it still leaves a blotchy texture across the brain and flattens gradual
intensity changes into patches, the piecewise-constant images that total
variation favours.

Denoiser as prior
-----------------

The network takes real planes of shape ``(n, channels, y, x)``; your image
is complex. :class:`~bartorch.priors.ImplicitPrior` bridges the two:
``spatial=2`` says the network works on 2D planes, and around each call the
prior scales the image to unit peak modulus, denoises its real and imaginary
parts as two grayscale planes, and scales the result back. So ``sigma`` is
in units of the image's peak.

Nothing is trained here, so you evaluate the network without gradients.
You run the same prior under ADMM and under FISTA.

.. GENERATED FROM PYTHON SOURCE LINES 281-314

.. code-block:: Python

    from deepinv.models import DRUNet

    ITERATIONS = 30

    denoiser = DRUNet(in_channels=1, out_channels=1, pretrained="download")
    denoiser = denoiser.eval()
    prior = priors.ImplicitPrior(denoiser, sigma=0.05, spatial=2)

    with torch.no_grad():
        admm = optim.admm(data, A, prior, maxiter=ITERATIONS, rho=0.2)
        fista = optim.fista(data, A, prior, maxiter=ITERATIONS)





.. rst-class:: sphx-glr-horizontal


    *

      .. image-sg:: /auto_examples/06-learning/images/sphx_glr_01-plug-and-play_003.png
         :alt: total variation error, DRUNet, ADMM error, DRUNet, FISTA error
         :srcset: /auto_examples/06-learning/images/sphx_glr_01-plug-and-play_003.png
         :class: sphx-glr-multi-img

    *

      .. image-sg:: /auto_examples/06-learning/images/sphx_glr_01-plug-and-play_004.png
         :alt: truth, total variation, DRUNet, ADMM, DRUNet, FISTA
         :srcset: /auto_examples/06-learning/images/sphx_glr_01-plug-and-play_004.png
         :class: sphx-glr-multi-img


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

 .. code-block:: none

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

      0%|          | 0.00/125M [00:00<?, ?B/s]
      0%|          | 128k/125M [00:00<05:03, 430kB/s]
      0%|          | 512k/125M [00:00<01:25, 1.52MB/s]
      1%|▏         | 1.62M/125M [00:00<00:27, 4.61MB/s]
      4%|▍         | 4.75M/125M [00:00<00:09, 12.9MB/s]
     11%|█         | 13.6M/125M [00:00<00:03, 36.1MB/s]
     31%|███       | 38.2M/125M [00:00<00:00, 99.6MB/s]
     58%|█████▊    | 72.0M/125M [00:00<00:00, 167MB/s] 
     99%|█████████▉| 123M/125M [00:01<00:00, 273MB/s] 
    100%|██████████| 125M/125M [00:01<00:00, 124MB/s]




.. GENERATED FROM PYTHON SOURCE LINES 315-333

The error maps are on one scale. Under ADMM the denoiser removes the noise
and the aliasing and keeps the tissue boundaries, and its error map is
darker inside the brain than the total-variation one. In the enlarged region
the ventricles and the larger cortical folds are there, though the finest
sulci are lost. Under FISTA, at the same :math:`\sigma`, the same denoiser
gives a much smoother image: the cortical folds are gone, and its error runs
along every tissue boundary.

The two iterations differ in the step before the denoiser. FISTA takes a
gradient step of fixed length on the data term; ADMM solves a least-squares
problem that holds the image to the data with weight :math:`\rho`. The
denoiser is therefore applied to different images, and the same
:math:`\sigma` has a different effective strength. With a
fixed :math:`\sigma` neither iteration minimises a known objective, so the
iteration count and :math:`\rho` are part of the result, and you choose
them as you would choose the weight of a penalty.

The numbers:

.. GENERATED FROM PYTHON SOURCE LINES 333-335

.. code-block:: Python

    report(reconstructions)





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

 .. code-block:: none

         zero-filled  NRMSE 0.199  SSIM 0.482
     total variation  NRMSE 0.161  SSIM 0.754
        DRUNet, ADMM  NRMSE 0.103  SSIM 0.893
       DRUNet, FISTA  NRMSE 0.135  SSIM 0.898




.. GENERATED FROM PYTHON SOURCE LINES 336-346

SSIM ranks the over-smoothed FISTA image above the ADMM one. One figure of
merit does not replace looking at the images.

Exercise: denoiser noise level
------------------------------

:math:`\sigma` is the strength of the prior. A :math:`\sigma` below the
noise level of the iterate leaves noise and aliasing in place; one above it
removes image detail as well. Run ADMM at a smaller and a larger
:math:`\sigma` and compare with 0.05:

.. GENERATED FROM PYTHON SOURCE LINES 346-368

.. code-block:: Python

    levels = (0.02, 0.05, 0.12)
    sweep = {0.05: admm}
    with torch.no_grad():
        for sigma in (0.02, 0.12):
            sweep[sigma] = optim.admm(
                data,
                A,
                priors.ImplicitPrior(denoiser, sigma=sigma, spatial=2),
                maxiter=ITERATIONS,
                rho=0.2,
            )

    report({f"sigma {sigma:.2f}": sweep[sigma] for sigma in levels})





.. image-sg:: /auto_examples/06-learning/images/sphx_glr_01-plug-and-play_005.png
   :alt: $\sigma$ = 0.02 (too small), $\sigma$ = 0.05, $\sigma$ = 0.12 (too large)
   :srcset: /auto_examples/06-learning/images/sphx_glr_01-plug-and-play_005.png
   :class: sphx-glr-single-img


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

 .. code-block:: none

          sigma 0.02  NRMSE 0.139  SSIM 0.734
          sigma 0.05  NRMSE 0.103  SSIM 0.893
          sigma 0.12  NRMSE 0.127  SSIM 0.886




.. GENERATED FROM PYTHON SOURCE LINES 369-396

At the smallest :math:`\sigma` the noise and the aliasing remain as a
mottled texture; at the largest the cortex is smoothed into uniform white
matter and small structures disappear. The printed errors have their
minimum at 0.05. You can repeat the sweep under FISTA and compare where
its minimum lies.

DRUNet was trained neither on MR images nor on the residual aliasing an
undersampled scan leaves, which is not white Gaussian noise. In lesson 11
you train a network for exactly that, inside the iteration.

As a spec
---------

What this lesson built, stated the way you would ask an agent for it:

.. code-block:: text

   Reconstruct a 2D Cartesian SENSE acquisition (128x128, 8 coils, R=4,
   16-line ACS, known sensitivities) with bartorch, using a pretrained
   denoiser as the prior. Build A with linop.CartesianSense. Load deepinv's
   DRUNet (1 channel, pretrained weights) in eval mode and wrap it as
   priors.ImplicitPrior(denoiser, sigma=0.05, spatial=2). Under
   torch.no_grad, run optim.admm (30 iterations, rho=0.2) and optim.fista
   (30 iterations) with it. Compare against optim.admm with
   priors.TotalVariation over the last two axes (weight 0.01, 60
   iterations, rho=0.1). Report NRMSE and SSIM of the magnitudes after a
   global least-squares scale, and repeat ADMM at sigma 0.02 and 0.12.

.. GENERATED FROM PYTHON SOURCE LINES 398-419

References
----------

.. [#brainweb] Collins DL, Zijdenbos AP, Kollokian V, Sled JG, Kabani NJ, Holmes CJ,
   Evans AC. Design and construction of a realistic digital brain phantom.
   *IEEE Trans Med Imaging* 17(3):463-468 (1998).
   https://doi.org/10.1109/42.712135

.. [#venkatakrishnan] Venkatakrishnan SV, Bouman CA, Wohlberg B. Plug-and-play
   priors for model based reconstruction. *IEEE Global Conference on Signal and
   Information Processing*, 945-948 (2013).
   https://doi.org/10.1109/GlobalSIP.2013.6737048

.. [#ahmad] Ahmad R, Bouman CA, Buzzard GT, Chan S, Liu S, Reehorst ET,
   Schniter P. Plug-and-play methods for magnetic resonance imaging: using
   denoisers for image recovery. *IEEE Signal Process Mag* 37(1):105-116
   (2020). https://doi.org/10.1109/MSP.2019.2949470

.. [#zhang] Zhang K, Li Y, Zuo W, Zhang L, Van Gool L, Timofte R. Plug-and-play
   image restoration with deep denoiser prior. *IEEE Trans Pattern Anal Mach
   Intell* 44(10):6360-6376 (2022). https://doi.org/10.1109/TPAMI.2021.3088914


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

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


.. _sphx_glr_download_auto_examples_06-learning_01-plug-and-play.py:

.. only:: html

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

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

      :download:`Download Jupyter notebook: 01-plug-and-play.ipynb <01-plug-and-play.ipynb>`

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

      :download:`Download Python source code: 01-plug-and-play.py <01-plug-and-play.py>`

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

      :download:`Download zipped: 01-plug-and-play.zip <01-plug-and-play.zip>`


.. only:: html

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

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