<!-- DO NOT EDIT. -->
<!-- THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. -->
<!-- TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: -->
<!-- "auto_examples/plug-and-play/demo_PnP_custom_optim.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-plug-and-play-demo-pnp-custom-optim-py"></a>

# PnP with custom optimization algorithm (Primal-Dual Condat-Vu)

This example shows how to define your own optimization algorithm.
For example, here, we implement the Primal-Dual Condat-Vu (CV) algorithm,
and apply it for Single Pixel Camera reconstruction.

<!-- GENERATED FROM PYTHON SOURCE LINES 9-21 -->
```Python
import deepinv as dinv
from pathlib import Path
import torch
from deepinv.models import DnCNN
from deepinv.optim.data_fidelity import L2
from deepinv.optim.prior import PnP
from deepinv.optim.optimizers import BaseOptim
from deepinv.utils import load_example
from deepinv.utils.plotting import plot, plot_curves
from deepinv.optim.optim_iterators import OptimIterator, fStep, gStep
```

<!-- GENERATED FROM PYTHON SOURCE LINES 22-36 -->

## Define a custom optimization algorithm

Creating your optimization algorithm only requires the definition of an iteration step.
The iterator should be a subclass of [`deepinv.optim.OptimIterator`](https://deepinv.org/api/stubs/deepinv.optim.OptimIterator.html.md#deepinv.optim.OptimIterator).

The Condat-Vu Primal-Dual algorithm is defined as follows:

$$
v_k &= x_k-\tau A^\top z_k \\
x_{k+1} &= \operatorname{prox}_{\tau g}(v_k) \\
u_k &= z_k + \sigma A(2x_{k+1}-x_k) \\
z_{k+1} &= \operatorname{prox}_{\sigma f^*}(u_k)

$$

where $f^*$ is the Fenchel-Legendre conjugate of $f$.

<!-- GENERATED FROM PYTHON SOURCE LINES 36-78 -->
```Python
class CVIteration(OptimIterator):
    r"""
    Single iteration of Condat-Vu Primal-Dual.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.g_step = gStepCV(**kwargs)
        self.f_step = fStepCV(**kwargs)

    def forward(self, X, cur_data_fidelity, cur_prior, cur_params, y, physics):
        r"""
        Single iteration of the Condat-Vu algorithm.

        :param dict X: Dictionary containing the current iterate and the estimated cost.
        :param deepinv.optim.DataFidelity cur_data_fidelity: Instance of the DataFidelity class defining the current data_fidelity.
        :param dict cur_prior: dictionary containing the prior-related term of interest,
            e.g. its proximal operator or gradient.
        :param dict cur_params: dictionary containing the current parameters of the model.
        :param torch.Tensor y: Input data.
        :param deepinv.physics physics: Instance of the physics modeling the data-fidelity term.
        :return: Dictionary `{"est": (x,z), "cost": F}` containing the updated current iterate
            and the estimated current cost.
        """
        x_prev, z_prev = X["est"]
        v = x_prev - cur_params["stepsize"] * physics.A_adjoint(z_prev)
        x = self.g_step(v, cur_prior, cur_params)
        u = z_prev + cur_params["stepsize"] * physics.A(2 * x - x_prev)
        z = self.f_step(u, cur_data_fidelity, cur_params, y, physics)
        F = (
            self.F_fn(x, cur_data_fidelity, cur_params, y, physics)
            if self.has_cost
            and self.F_fn is not None
            and cur_data_fidelity is not None
            and cur_prior is not None
            else None
        )
        return {"est": (x, z), "cost": F}
```

<!-- GENERATED FROM PYTHON SOURCE LINES 79-98 -->

## Define the custom fStep and gStep modules

The iterator relies on custom fStepCV (subclass of [`deepinv.optim.optim_iterators.fStep`](https://deepinv.org/api/stubs/deepinv.optim.optim_iterators.fStep.html.md#deepinv.optim.optim_iterators.fStep))
and gStepCV (subclass of [`deepinv.optim.optim_iterators.gStep`](https://deepinv.org/api/stubs/deepinv.optim.optim_iterators.gStep.html.md#deepinv.optim.optim_iterators.gStep)) modules.

In this case the fStep module is defined as follows:

$$
u_{k+1} = \operatorname{prox}_{\sigma f^*}(u_k)

$$

where $f^*$ is the Fenchel-Legendre conjugate of $f$.
The proximal operator of $f^*$ is computed using the proximal operator
of $f$ via Moreau’s identity,
and the gStep module is a simple proximal step on the prior term $\lambda g$:

$$
x_{k+1} = \operatorname{prox}_{\tau \lambda g}(v_k)
$$

<!-- GENERATED FROM PYTHON SOURCE LINES 98-147 -->
```Python
class fStepCV(fStep):
    r"""
    Condat-Vu fStep module to compute :math:`\operatorname{prox}_{\sigma f^*}(z_k)``
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def forward(self, u, cur_data_fidelity, cur_params, y, physics):
        r"""
        Single iteration on the data-fidelity term :math:`f`.

        :param torch.Tensor z: Current iterate :math:`z_k = 2Ax_{k+1}-x_k`
        :param deepinv.optim.DataFidelity cur_data_fidelity: Instance of the DataFidelity class defining the current data_fidelity.
        :param dict cur_params: Dictionary containing the current fStep parameters (keys `"stepsize"` and `"lambda"`).
        :param torch.Tensor y: Input data.
        :param deepinv.physics physics: Instance of the physics modeling the data-fidelity term.
        """
        return cur_data_fidelity.d.prox_conjugate(
            u, y, gamma=cur_params["stepsize_dual"]
        )


class gStepCV(gStep):
    r"""
    Condat-Vu gStep module to compute :math:`\operatorname{prox}_{\tau g}(v_k)`
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def forward(self, v, cur_prior, cur_params):
        r"""
        Single iteration step on the prior term :math:`\lambda g`.

        :param torch.Tensor x: Current iterate :math:`v_k = x_k-\tau A^\top u_k`.
        :param dict cur_prior: Dictionary containing the current prior.
        :param dict cur_params: Dictionary containing the current gStep parameters
            (keys `"stepsize"` and `"g_param"`).
        """
        return cur_prior.prox(
            v,
            cur_params["g_param"],
            gamma=cur_params["lambda"] * cur_params["stepsize"],
        )
```

<!-- GENERATED FROM PYTHON SOURCE LINES 148-151 -->

Define the Conva-Vu model as a subclass of [`deepinv.optim.BaseOptim`](https://deepinv.org/api/stubs/deepinv.optim.BaseOptim.html.md#deepinv.optim.BaseOptim),
in the model of other optimizer of the library, see for example [`deepinv.optim.ADMM`](https://deepinv.org/api/stubs/deepinv.optim.ADMM.html.md#deepinv.optim.ADMM).
—————————————————————————————-

<!-- GENERATED FROM PYTHON SOURCE LINES 151-186 -->
```Python
class CV(BaseOptim):
    r"""
    Primal-Dual Condat-Vu (CV) optimization algorithm.
    """

    def __init__(
        self,
        data_fidelity=None,
        prior=None,
        lambda_reg=1.0,
        stepsize=1.0,
        stepsize_dual=1.0,
        beta=1.0,
        sigma_denoiser=None,
        **kwargs,
    ):
        params_algo = {
            "lambda": lambda_reg,
            "stepsize": stepsize,
            "stepsize_dual": stepsize_dual,
            "g_param": sigma_denoiser,
            "beta": beta,
        }

        super(CV, self).__init__(
            CVIteration(),
            params_algo=params_algo,
            data_fidelity=data_fidelity,
            prior=prior,
            **kwargs,
        )
```

<!-- GENERATED FROM PYTHON SOURCE LINES 187-190 -->

## Setup paths for data loading and results.

<!-- GENERATED FROM PYTHON SOURCE LINES 190-195 -->
```Python
BASE_DIR = Path(".")
RESULTS_DIR = BASE_DIR / "results"
```

<!-- GENERATED FROM PYTHON SOURCE LINES 196-198 -->

## Load base image datasets and degradation operators.

<!-- GENERATED FROM PYTHON SOURCE LINES 198-220 -->
```Python
# Set the global random seed from pytorch to ensure reproducibility of the example.
torch.manual_seed(0)

device = dinv.utils.get_device()

# Set up the variable to fetch dataset and operators.
method = "PnP"
dataset_name = "set3c"
img_size = 64

x = load_example(
    "barbara.jpeg",
    img_size=img_size,
    grayscale=True,
    resize_mode="resize",
    device=device,
)
operation = "single_pixel"
```

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

<!-- GENERATED FROM PYTHON SOURCE LINES 221-226 -->

## Set the forward operator

We use the [`deepinv.physics.SinglePixelCamera`](https://deepinv.org/api/stubs/deepinv.physics.SinglePixelCamera.html.md#deepinv.physics.SinglePixelCamera)
class from the physics module to generate a single-pixel measurements.
The forward operator consists of the multiplication with the low frequencies of the Hadamard transform.

<!-- GENERATED FROM PYTHON SOURCE LINES 226-242 -->
```Python
noise_level_img = 0.03  # Gaussian Noise standard deviation for the degradation
n_channels = 1  # 3 for color images, 1 for gray-scale images
physics = dinv.physics.SinglePixelCamera(
    m=600,
    img_size=(1, 64, 64),
    ordering="cake_cutting",
    noise_model=dinv.physics.GaussianNoise(sigma=noise_level_img),
    device=device,
)

# Use parallel dataloader if using a GPU to speed up training,
# otherwise, as all computes are on CPU, use synchronous data loading.
num_workers = 4 if torch.cuda.is_available() else 0
```

<!-- GENERATED FROM PYTHON SOURCE LINES 243-251 -->

## Set up the PnP algorithm to solve the inverse problem.

We build the PnP model using our custom [`deepinv.optim.PDCP()`](https://deepinv.org/api/stubs/deepinv.optim.PDCP.html.md#deepinv.optim.PDCP) function.

The primal dual stepsizes $\tau$ corresponds to the `stepsize` key and $\sigma$ to the `sigma` key.

For the denoiser, we choose the 1-Lipschitz grayscale DnCNN model (see the [pretrained-weights](https://deepinv.org/user_guide/reconstruction/pretrained-models.html.md#pretrained-weights)).

<!-- GENERATED FROM PYTHON SOURCE LINES 251-283 -->
```Python
# Set up the PnP algorithm parameters :
stepsize = 0.99  # primal stepsize
stepsize_dual = 0.99  # dual stepsize
sigma_denoiser = 0.01  # denoiser parameter (noise level)
max_iter = 200
early_stop = True  # stop the algorithm when convergence is reached

# Select the data fidelity term
data_fidelity = L2()

# Specify the denoising prior
denoiser = DnCNN(
    in_channels=n_channels,
    out_channels=n_channels,
    pretrained="download_lipschitz",
    device=device,
)
prior = PnP(denoiser=denoiser)

# instantiate the algorithm class to solve the IP problem.
model = CV(
    prior=prior,
    data_fidelity=data_fidelity,
    stepsize=stepsize,
    stepsize_dual=stepsize_dual,
    sigma_denoiser=sigma_denoiser,
    early_stop=early_stop,
    max_iter=max_iter,
    verbose=True,
)
```

<!-- GENERATED FROM PYTHON SOURCE LINES 284-289 -->

## Evaluate the model on the problem and plot the results.

The model returns the output and the metrics computed along the iterations.
The ground truth image `x_gt` must be provided for computing the PSNR.

<!-- GENERATED FROM PYTHON SOURCE LINES 289-306 -->
```Python
y = physics(x)
x_lin = physics.A_adjoint(y)

# run the model on the problem. For computing the metrics along the iterations, set ``compute_metrics=True``.
x_model, metrics = model(y, physics, x_gt=x, compute_metrics=True)

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

# plot results
imgs = [x, x_lin, x_model]
plot(imgs, titles=["GT", "Linear", "Recons."], show=True)

# plot convergence curves
plot_curves(metrics, save_dir=RESULTS_DIR / "curves", show=True)
```

```none
Linear reconstruction PSNR: 23.12 dB
Model reconstruction PSNR: 25.39 dB
```

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

<a id="sphx-glr-download-auto-examples-plug-and-play-demo-pnp-custom-optim-py"></a>

[`Download Jupyter notebook: demo_PnP_custom_optim.ipynb`](https://deepinv.org/auto_examples/plug-and-play/_downloads/b33a1eb4de39658ecbc7b0bf022ff0b2/demo_PnP_custom_optim.ipynb)

[`Download Python source code: demo_PnP_custom_optim.py`](https://deepinv.org/auto_examples/plug-and-play/_downloads/824bbc71c8aa5f4974a95a73388afbe3/demo_PnP_custom_optim.py)

[`Download zipped: demo_PnP_custom_optim.zip`](https://deepinv.org/auto_examples/plug-and-play/_downloads/59cd8864b56efe5b6b42f7bb988ab466/demo_PnP_custom_optim.zip)

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