deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/optim/optim_iterators/hqs.py

@@ -0,0 +1,74 @@
+from .optim_iterator import OptimIterator, fStep, gStep
+
+
+class HQSIteration(OptimIterator):
+    r"""
+    Single iteration of half-quadratic splitting.
+
+    Class for a single iteration of the Half-Quadratic Splitting (HQS) algorithm for minimising :math:`\lambda f(x) + g(x)`.
+    The iteration is given by
+
+
+    .. math::
+        \begin{equation*}
+        \begin{aligned}
+        u_{k} &= \operatorname{prox}_{\gamma \lambda f}(x_k) \\
+        x_{k+1} &= \operatorname{prox}_{\sigma g}(u_k).
+        \end{aligned}
+        \end{equation*}
+
+
+    where :math:`\gamma` and :math:`\sigma` are step-sizes. Note that this algorithm does not converge to
+    a minimizer of :math:`\lambda f(x) + g(x)`, but instead to a minimizer of
+    :math:`\lambda \gamma\, ^1f+\sigma g`, where :math:`^1f` denotes
+    the Moreau envelope of :math:`f`
+
+    """
+
+    def __init__(self, **kwargs):
+        super(HQSIteration, self).__init__(**kwargs)
+        self.g_step = gStepHQS(**kwargs)
+        self.f_step = fStepHQS(**kwargs)
+        self.requires_prox_g = True
+
+
+class fStepHQS(fStep):
+    r"""
+    HQS fStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(fStepHQS, self).__init__(**kwargs)
+
+    def forward(self, x, cur_data_fidelity, cur_params, y, physics):
+        r"""
+        Single proximal step on the data-fidelity term :math:`f`.
+
+        :param torch.Tensor x: Current iterate :math:`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 parameters of the algorithm.
+        :param torch.Tensor y: Input data.
+        :param deepinv.physics physics: Instance of the physics modeling the data-fidelity term.
+        """
+        return cur_data_fidelity.prox(
+            x, y, physics, gamma=cur_params["lambda"] * cur_params["stepsize"]
+        )
+
+
+class gStepHQS(gStep):
+    r"""
+    HQS gStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(gStepHQS, self).__init__(**kwargs)
+
+    def forward(self, x, cur_prior, cur_params):
+        r"""
+        Single proximal step on the prior term :math:`g`.
+
+        :param torch.Tensor x: Current iterate :math:`x_k`.
+        :param dict cur_prior: Class containing the current prior.
+        :param dict cur_params: Dictionary containing the current parameters of the algorithm.
+        """
+        return cur_prior.prox(x, cur_params["g_param"], gamma=cur_params["stepsize"])

deepinv/optim/optim_iterators/optim_iterator.py

@@ -0,0 +1,141 @@
+import torch
+import torch.nn as nn
+from deepinv.optim.data_fidelity import L2
+
+
+class OptimIterator(nn.Module):
+    r"""
+    Base class for all :meth:`Optim` iterators.
+
+    An optim iterator is an object that implements a fixed point iteration for minimizing the sum of two functions
+    :math:`F = \lambda*f + g` where :math:`f` is a data-fidelity term  that will be modeled by an instance of physics
+    and g is a regularizer. The fixed point iteration takes the form
+
+    .. math::
+        \qquad (x_{k+1}, z_{k+1}) = \operatorname{FixedPoint}(x_k, z_k, f, g, A, y, ...)
+
+    where :math:`x` is a "primal" variable converging to the solution of the minimization problem, and
+    :math:`z` is a "dual" variable.
+
+
+    .. note::
+        By an abuse of terminology, we call "primal" and "dual" variables the variables that are updated
+        at each step and which may correspond to the actual primal and dual variables from 
+        (for instance in the case of the PD algorithm), but not necessarily (for instance in the case of the
+        PGD algorithm).
+
+
+    The implementation of the fixed point algorithm in :meth:`deepinv.optim`  is split in two steps, alternating between
+    a step on f and a step on g, that is for :math:`k=1,2,...`
+
+    .. math::
+        z_{k+1} = \operatorname{step}_f(x_k, z_k, y, A, ...)\\
+        x_{k+1} = \operatorname{step}_g(x_k, z_k, y, A, ...)
+
+    where :math:`\operatorname{step}_f` and :math:`\operatorname{step}_g` are the steps on f and g respectively.
+
+    :param bool g_first: If True, the algorithm starts with a step on g and finishes with a step on f.
+    :param F_fn: function that returns the function F to be minimized at each iteration. Default: None.
+    :param bool has_cost: If True, the function F is computed at each iteration. Default: False.
+     """
+
+    def __init__(self, g_first=False, F_fn=None, has_cost=False):
+        super(OptimIterator, self).__init__()
+        self.g_first = g_first
+        self.F_fn = F_fn
+        self.has_cost = has_cost
+        if self.F_fn is None:
+            self.has_cost = False
+        self.f_step = fStep(g_first=self.g_first)
+        self.g_step = gStep(g_first=self.g_first)
+        self.requires_grad_g = False
+        self.requires_prox_g = False
+
+    def relaxation_step(self, u, v, beta):
+        r"""
+        Performs a relaxation step of the form :math:`\beta u + (1-\beta) v`.
+
+        :param torch.Tensor u: First tensor.
+        :param torch.Tensor v: Second tensor.
+        :param float beta: Relaxation parameter.
+        :return: Relaxed tensor.
+        """
+        return beta * u + (1 - beta) * v
+
+    def forward(self, X, cur_data_fidelity, cur_prior, cur_params, y, physics):
+        r"""
+        General form of a single iteration of splitting algorithms for minimizing :math:`F = \lambda f + g`, alternating
+        between a step on :math:`f` and a step on :math:`g`.
+        The primal and dual variables as well as the estimated cost at the current iterate are stored in a dictionary
+        $X$ of the form `{'est': (x,z), 'cost': F}`.
+
+        :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 deepinv.optim.prior cur_prior: Instance of the Prior class defining the current prior.
+        :param dict cur_params: Dictionary containing the current parameters of the algorithm.
+        :param torch.Tensor y: Input data.
+        :param deepinv.physics physics: Instance of the physics modeling the observation.
+        :return: Dictionary `{"est": (x, z), "cost": F}` containing the updated current iterate and the estimated current cost.
+        """
+        x_prev = X["est"][0]
+        if not self.g_first:
+            z = self.f_step(x_prev, cur_data_fidelity, cur_params, y, physics)
+            x = self.g_step(z, cur_prior, cur_params)
+        else:
+            z = self.g_step(x_prev, cur_prior, cur_params)
+            x = self.f_step(z, cur_data_fidelity, cur_params, y, physics)
+        x = self.relaxation_step(x, x_prev, cur_params["beta"])
+        F = (
+            self.F_fn(x, cur_data_fidelity, cur_prior, cur_params, y, physics)
+            if self.has_cost
+            else None
+        )
+        return {"est": (x, z), "cost": F}
+
+
+class fStep(nn.Module):
+    r"""
+    Module for the single iteration steps on the data-fidelity term :math:`f`.
+
+    :param bool g_first: If True, the algorithm starts with a step on g and finishes with a step on f. Default: False.
+    :param kwargs: Additional keyword arguments.
+    """
+
+    def __init__(self, g_first=False, **kwargs):
+        super(fStep, self).__init__()
+        self.g_first = g_first
+
+        def forward(self, x, cur_data_fidelity, cur_params, y, physics):
+            r"""
+            Single iteration step on the data-fidelity term :math:`f`.
+
+            :param torch.Tensor x: Current iterate.
+            :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 parameters of the algorithm.
+            :param torch.Tensor y: Input data.
+            :param deepinv.physics physics: Instance of the physics modeling the observation.
+            """
+            pass
+
+
+class gStep(nn.Module):
+    r"""
+    Module for the single iteration steps on the prior term :math:`g`.
+
+    :param bool g_first: If True, the algorithm starts with a step on g and finishes with a step on f. Default: False.
+    :param kwargs: Additional keyword arguments.
+    """
+
+    def __init__(self, g_first=False, **kwargs):
+        super(gStep, self).__init__()
+        self.g_first = g_first
+
+        def forward(self, x, cur_prior, cur_params):
+            r"""
+            Single iteration step on the prior term :math:`g`.
+
+            :param torch.Tensor x: Current iterate.
+            :param deepinv.optim.prior cur_prior: Instance of the Prior class defining the current prior.
+            :param dict cur_params: Dictionary containing the current parameters of the algorithm.
+            """
+            pass

deepinv/optim/optim_iterators/pgd.py

@@ -0,0 +1,91 @@
+from .optim_iterator import OptimIterator, fStep, gStep
+from .utils import gradient_descent_step
+
+
+class PGDIteration(OptimIterator):
+    r"""
+    Iterator for proximal gradient descent.
+
+    Class for a single iteration of the Proximal Gradient Descent (PGD) algorithm for minimising :math:`\lambda f(x) + g(x)`.
+
+    The iteration is given by
+
+    .. math::
+        \begin{equation*}
+        \begin{aligned}
+        u_{k} &= x_k - \lambda \gamma \nabla f(x_k) \\
+        x_{k+1} &= \operatorname{prox}_{\gamma g}(u_k),
+        \end{aligned}
+        \end{equation*}
+
+
+    where :math:`\gamma` is a stepsize that should satisfy :math:`\lambda \gamma \leq 2/\operatorname{Lip}(\|\nabla f\|)`.
+
+    """
+
+    def __init__(self, **kwargs):
+        super(PGDIteration, self).__init__(**kwargs)
+        self.g_step = gStepPGD(**kwargs)
+        self.f_step = fStepPGD(**kwargs)
+        if self.g_first:
+            self.requires_grad_g = True
+        else:
+            self.requires_prox_g = True
+
+
+class fStepPGD(fStep):
+    r"""
+    PGD fStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(fStepPGD, self).__init__(**kwargs)
+
+    def forward(self, x, cur_data_fidelity, cur_params, y, physics):
+        r"""
+         Single PGD iteration step on the data-fidelity term :math:`f`.
+
+         :param torch.Tensor x: Current iterate :math:`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 parameters of the algorithm.
+         :param torch.Tensor y: Input data.
+         :param deepinv.physics physics: Instance of the physics modeling the data-fidelity term.
+        """
+        if not self.g_first:
+            # if cur_params["lambda"] >= 2:
+            #     raise ValueError("lambda must be smaller than 2")
+            grad = (
+                cur_params["lambda"]
+                * cur_params["stepsize"]
+                * cur_data_fidelity.grad(x, y, physics)
+            )
+            return gradient_descent_step(x, grad)
+        else:
+            return cur_data_fidelity.prox(
+                x, y, physics, gamma=cur_params["lambda"] * cur_params["stepsize"]
+            )
+
+
+class gStepPGD(gStep):
+    r"""
+    PGD gStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(gStepPGD, self).__init__(**kwargs)
+
+    def forward(self, x, cur_prior, cur_params):
+        r"""
+        Single iteration step on the prior term :math:`g`.
+
+        :param torch.Tensor x: Current iterate :math:`x_k`.
+        :param dict cur_prior: Dictionary containing the current prior.
+        :param dict cur_params: Dictionary containing the current parameters of the algorithm.
+        """
+        if not self.g_first:
+            return cur_prior.prox(
+                x, cur_params["g_param"], gamma=cur_params["stepsize"]
+            )
+        else:
+            grad = cur_params["stepsize"] * cur_prior.grad(x, cur_params["g_param"])
+            return gradient_descent_step(x, grad)

deepinv/optim/optim_iterators/primal_dual_CP.py

@@ -0,0 +1,145 @@
+import torch
+
+from .optim_iterator import OptimIterator, fStep, gStep
+
+
+class CPIteration(OptimIterator):
+    r"""
+    Iterator for Chambolle-Pock.
+
+    Class for a single iteration of the `Chambolle-Pock <https://hal.science/hal-00490826/document>`_ Primal-Dual (PD)
+    algorithm for minimising :math:`\lambda F(Kx) + G(x)` or :math:`\lambda F(x) + G(Kx)` for generic functions :math:`F` and :math:`G`.
+    Our implementation corresponds to Algorithm 1 of `<https://hal.science/hal-00490826/document>`_.
+
+    If the attribute ``g_first`` is set to ``False`` (by default), the iteration is given by
+
+    .. math::
+        \begin{equation*}
+        \begin{aligned}
+        u_{k+1} &= \operatorname{prox}_{\sigma (\lambda F)^*}(u_k + \sigma K z_k) \\
+        x_{k+1} &= \operatorname{prox}_{\tau G}(x_k-\tau K^\top u_{k+1}) \\
+        z_{k+1} &= x_{k+1} + \beta(x_{k+1}-x_k) \\
+        \end{aligned}
+        \end{equation*}
+
+    where :math:`(\lambda F)^*` is the Fenchel-Legendre conjugate of :math:`\lambda F`, :math:`\beta>0` is a relaxation parameter, and :math:`\sigma` and :math:`\tau` are step-sizes that should
+    satisfy :math:`\sigma \tau \|K\|^2 \leq 1`.
+
+    If the attribute ``g_first`` is set to ``True``, the functions :math:`F` and :math:`G` are inverted in the previous iteration.
+
+    In particular, setting :math:`F = \distancename`, :math:`K = A` and :math:`G = \regname`, the above algorithms solves
+
+    .. math::
+
+        \begin{equation*}
+        \underset{x}{\operatorname{min}} \,\, \lambda \distancename(Ax, y) + \regname(x)
+        \end{equation*}
+
+
+    with a splitting on :math:`\distancename`, with not differentiability assumption needed on :math:`\distancename`
+    or :math:`\regname`, not any invertibility assumption on :math:`A`.
+
+    Note that the algorithm requires an intiliazation of the three variables :math:`x_0`, :math:`z_0` and :math:`u_0`.
+    """
+
+    def __init__(self, **kwargs):
+        super(CPIteration, self).__init__(**kwargs)
+        self.g_step = gStepCP(**kwargs)
+        self.f_step = fStepCP(**kwargs)
+
+    def forward(self, X, cur_data_fidelity, cur_prior, cur_params, y, physics):
+        r"""
+        Single iteration of the Chambolle-Pock 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 deepinv.optim.Prior cur_prior: Instance of the Prior class defining the current prior.
+        :param dict cur_params: dictionary containing the current parameters of the algorithm.
+        :param torch.Tensor y: Input data.
+        :param deepinv.physics physics: Instance of the physics modeling the data-fidelity term.
+        :return: Dictionary `{"est": (x, ), "cost": F}` containing the updated current iterate and the estimated current cost.
+        """
+        x_prev, z_prev, u_prev = X["est"]  # x : primal, z : relaxed primal, u : dual
+        K = lambda x: cur_params["K"](x) if "K" in cur_params.keys() else x
+        K_adjoint = (
+            lambda x: cur_params["K_adjoint"](x)
+            if "K_adjoint" in cur_params.keys()
+            else x
+        )
+        if self.g_first:
+            u = self.g_step(u_prev, K(z_prev), cur_prior, cur_params)
+            x = self.f_step(
+                x_prev, K_adjoint(u), cur_data_fidelity, y, physics, cur_params
+            )
+        else:
+            u = self.f_step(
+                u_prev, K(z_prev), cur_data_fidelity, y, physics, cur_params
+            )
+            x = self.g_step(x_prev, K_adjoint(u), cur_prior, cur_params)
+        z = x + cur_params["beta"] * (x - x_prev)
+        F = (
+            self.F_fn(x, cur_data_fidelity, cur_prior, cur_params, y, physics)
+            if self.has_cost
+            else None
+        )
+        return {"est": (x, z, u), "cost": F}
+
+
+class fStepCP(fStep):
+    r"""
+    Chambolle-Pock fStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(fStepCP, self).__init__(**kwargs)
+
+    def forward(self, x, w, cur_data_fidelity, y, physics, cur_params):
+        r"""
+        Single Chambolle-Pock iteration step on the data-fidelity term :math:`\lambda f`.
+
+        :param torch.Tensor x: Current first variable :math:`x` if `"g_first"` and :math:`u` otherwise.
+        :param torch.Tensor w: Current second variable :math:`A^\top u` if `"g_first"` and :math:`A z` otherwise.
+        :param deepinv.optim.DataFidelity cur_data_fidelity: Instance of the DataFidelity class defining the current data_fidelity.
+        :param torch.Tensor y: Input data.
+        :param deepinv.physics physics: Instance of the physics modeling the data-fidelity term.
+        :param dict cur_params: Dictionary containing the current fStep parameters (keys `"stepsize_dual"` (or `"stepsize"`) and `"lambda"`).
+        """
+        if self.g_first:
+            p = x - cur_params["stepsize"] * w
+            return cur_data_fidelity.prox(
+                p, y, physics, gamma=cur_params["stepsize"] * cur_params["lambda"]
+            )
+        else:
+            p = x + cur_params["stepsize_dual"] * w
+            return cur_data_fidelity.prox_d_conjugate(
+                p, y, gamma=cur_params["stepsize_dual"], lamb=cur_params["lambda"]
+            )
+
+
+class gStepCP(gStep):
+    r"""
+    Chambolle-Pock gStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(gStepCP, self).__init__(**kwargs)
+
+    def forward(self, x, w, cur_prior, cur_params):
+        r"""
+        Single Chambolle-Pock iteration step on the prior term :math:`g`.
+
+        :param torch.Tensor x: Current first variable :math:`u` if `"g_first"` and :math:`x` otherwise.
+        :param torch.Tensor w: Current second variable :math:`A z` if `"g_first"` and :math:`A^\top u` otherwise.
+        :param deepinv.optim.prior cur_prior: Instance of the Prior class defining the current prior.
+        :param dict cur_params: Dictionary containing the current gStep parameters (keys `"prox_g"`, `"stepsize"` (or `"stepsize_dual"`) and `"g_param"`).
+        """
+        if self.g_first:
+            p = x + cur_params["stepsize_dual"] * w
+            return cur_prior.prox_conjugate(
+                p, cur_params["g_param"], gamma=cur_params["stepsize_dual"]
+            )
+        else:
+            p = x - cur_params["stepsize"] * w
+            return cur_prior.prox(
+                p, cur_params["g_param"], gamma=cur_params["stepsize"]
+            )

deepinv/optim/optim_iterators/utils.py

@@ -0,0 +1,17 @@
+def gradient_descent_step(x, grad, bregman_potential="L2"):
+    r"""
+    Performs a single step of gradient descent on the Bregman divergence.
+
+    :param torch.Tensor x: Current iterate.
+    :param torch.Tensor grad: Gradient of the Bregman divergence.
+    :param str bregman_potential: Bregman potential used in the Bregman divergence.
+    """
+    if bregman_potential == "L2":
+        grad_step = x - grad
+    elif bregman_potential == "Burg_entropy":
+        grad_step = x / (1 + x * grad)
+    else:
+        raise ValueError(
+            f"Gradient Descent with bregman potential {bregman_potential} not implemented"
+        )
+    return grad_step