deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/optim/prior.py

@@ -0,0 +1,288 @@
+import torch
+import torch.nn as nn
+
+from deepinv.optim.utils import gradient_descent
+
+
+class Prior(nn.Module):
+    r"""
+    Prior term :math:`g(x)`.
+
+    This is the base class for the prior term :math:`g(x)`. Similarly to the :meth:`deepinv.optim.DataFidelity` class,
+    this class comes with methods for computing
+    :math:`\operatorname{prox}_{g}` and :math:`\nabla g`.
+    To implement a custom prior, for an explicit prior, overwrite :math:`g` (do not forget to specify
+    `self.explicit_prior = True`)
+
+    This base class is also used to implement implicit priors. For instance, in PnP methods, the method computing the
+    proximity operator is overwritten by a method performing denoising. For an implicit prior, overwrite `grad`
+    or `prox`.
+
+
+    .. note::
+
+        The methods for computing the proximity operator and the gradient of the prior rely on automatic
+        differentiation. These methods should not be used when the prior is not differentiable, although they will
+        not raise an error.
+
+
+    :param callable g: Prior function :math:`g(x)`.
+    """
+
+    def __init__(self, g=None):
+        super().__init__()
+        self._g = g
+        self.explicit_prior = False if self._g is None else True
+
+    def g(self, x, *args, **kwargs):
+        r"""
+        Computes the prior :math:`g(x)`.
+
+        :param torch.tensor x: Variable :math:`x` at which the prior is computed.
+        :return: (torch.tensor) prior :math:`g(x)`.
+        """
+        return self._g(x, *args, **kwargs)
+
+    def forward(self, x, *args, **kwargs):
+        r"""
+        Computes the prior :math:`g(x)`.
+
+        :param torch.tensor x: Variable :math:`x` at which the prior is computed.
+        :return: (torch.tensor) prior :math:`g(x)`.
+        """
+        return self.g(x, *args, **kwargs)
+
+    def grad(self, x, *args, **kwargs):
+        r"""
+        Calculates the gradient of the prior term :math:`g` at :math:`x`.
+        By default, the gradient is computed using automatic differentiation.
+
+        :param torch.tensor x: Variable :math:`x` at which the gradient is computed.
+        :return: (torch.tensor) gradient :math:`\nabla_x g`, computed in :math:`x`.
+        """
+        with torch.enable_grad():
+            x = x.requires_grad_()
+            grad = torch.autograd.grad(
+                self.g(x, *args, **kwargs), x, create_graph=True, only_inputs=True
+            )[0]
+        return grad
+
+    def prox(
+        self,
+        x,
+        *args,
+        gamma=1.0,
+        stepsize_inter=1.0,
+        max_iter_inter=50,
+        tol_inter=1e-3,
+        **kwargs,
+    ):
+        r"""
+        Calculates the proximity operator of :math:`g` at :math:`x`. By default, the proximity operator is computed using internal gradient descent.
+
+        :param torch.tensor x: Variable :math:`x` at which the proximity operator is computed.
+        :param float gamma: stepsize of the proximity operator.
+        :param float stepsize_inter: stepsize used for internal gradient descent
+        :param int max_iter_inter: maximal number of iterations for internal gradient descent.
+        :param float tol_inter: internal gradient descent has converged when the L2 distance between two consecutive iterates is smaller than tol_inter.
+        :return: (torch.tensor) proximity operator :math:`\operatorname{prox}_{\gamma g}(x)`, computed in :math:`x`.
+        """
+        grad = lambda z: gamma * self.grad(z, *args, **kwargs) + (z - x)
+        return gradient_descent(
+            grad, x, step_size=stepsize_inter, max_iter=max_iter_inter, tol=tol_inter
+        )
+
+    def prox_conjugate(self, x, *args, gamma=1.0, lamb=1.0, **kwargs):
+        r"""
+        Calculates the proximity operator of the convex conjugate :math:`(\lambda g)^*` at :math:`x`, using the Moreau formula.
+
+        ::Warning:: Only valid for convex :math:`g`
+
+        :param torch.tensor x: Variable :math:`x` at which the proximity operator is computed.
+        :param float gamma: stepsize of the proximity operator.
+        :param float lamb: math:`\lambda` parameter in front of :math:`f`
+        :return: (torch.tensor) proximity operator :math:`\operatorname{prox}_{\gamma \lambda g)^*}(x)`, computed in :math:`x`.
+        """
+        return x - gamma * self.prox(x / gamma, lamb / gamma, *args, **kwargs)
+
+
+class PnP(Prior):
+    r"""
+    Plug-and-play prior :math:`\operatorname{prox}_{\gamma g}(x) = \operatorname{D}_{\sigma}(x)`.
+
+
+    :param callable denoiser: Denoiser :math:`\operatorname{D}_{\sigma}`.
+    """
+
+    def __init__(self, denoiser, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.denoiser = denoiser
+        self.explicit_prior = False
+
+    def prox(self, x, sigma_denoiser, *args, **kwargs):
+        r"""
+        Uses denoising as the proximity operator of the PnP prior :math:`g` at :math:`x`.
+
+        :param torch.tensor x: Variable :math:`x` at which the proximity operator is computed.
+        :param float sigma_denoiser: noise level parameter of the denoiser.
+        :return: (torch.tensor) proximity operator at :math:`x`.
+        """
+        return self.denoiser(x, sigma_denoiser)
+
+
+class RED(Prior):
+    r"""
+    Regularization-by-Denoising (RED) prior :math:`\nabla g(x) = \operatorname{Id} - \operatorname{D}_{\sigma}(x)`.
+
+
+    :param callable denoiser: Denoiser :math:`\operatorname{D}_{\sigma}`.
+    """
+
+    def __init__(self, denoiser, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.denoiser = denoiser
+        self.explicit_prior = False
+
+    def grad(self, x, sigma_denoiser, *args, **kwargs):
+        r"""
+        Calculates the gradient of the prior term :math:`g` at :math:`x`.
+        By default, the gradient is computed using automatic differentiation.
+
+        :param torch.Tensor x: Variable :math:`x` at which the gradient is computed.
+        :return: (:class:`torch.Tensor`) gradient :math:`\nabla_x g`, computed in :math:`x`.
+        """
+        return x - self.denoiser(x, sigma_denoiser)
+
+
+class ScorePrior(Prior):
+    r"""
+    Score via MMSE denoiser :math:`\nabla g(x)=\left(x-\operatorname{D}_{\sigma}(x)\right)/\sigma^2`.
+
+    This approximates the score of a distribution using Tweedie's formula, i.e.,
+
+    .. math::
+
+        - \nabla \log p_{\sigma}(x) \propto \left(x-D(x,\sigma)\right)/\sigma^2
+
+    where :math:`p_{\sigma} = p*\mathcal{N}(0,I\sigma^2)` is the prior convolved with a Gaussian kernel,
+    :math:`D(\cdot,\sigma)` is a (trained or model-based) denoiser with noise level :math:`\sigma`,
+    which is typically set to a low value.
+
+    .. note::
+
+        If math:`\sigma=1`, this prior is equal to :class:`deepinv.optim.RED`, which is defined in
+        `Regularization by Denoising (RED) <https://arxiv.org/abs/1611.02862>`_ and doesn't require the normalization.
+
+
+    .. note::
+
+        This class can also be used with maximum-a-posteriori (MAP) denoisers,
+        but :math:`p_{\sigma}(x)` is not given by the convolution with a Gaussian kernel, but rather
+        given by the Moreau-Yosida envelope of :math:`p(x)`, i.e.,
+
+        .. math::
+
+            p_{\sigma}(x)=e^{- \inf_z \left(-\log p(z) + \frac{1}{2\sigma}\|x-z\|^2 \right)}.
+
+
+    """
+
+    def __init__(self, denoiser, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.denoiser = denoiser
+        self.explicit_prior = False
+
+    def forward(self, x, sigma):
+        r"""
+        Applies the denoiser to the input signal.
+
+        :param torch.Tensor x: the input tensor.
+        :param float sigma: the noise level.
+        """
+        return (1 / sigma**2) * (x - self.denoiser(x, sigma))
+
+
+class Tikhonov(Prior):
+    r"""
+    Tikhonov regularizer :math:`g(x) = \frac{1}{2}\| x \|_2^2`.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.explicit_prior = True
+
+    def g(self, x, ths=1.0):
+        r"""
+        Computes the Tikhonov regularizer :math:`g(x) = \frac{\tau}{2}\| x \|_2^2`.
+
+        :param torch.Tensor x: Variable :math:`x` at which the prior is computed.
+        :param float ths: regularization parameter :math:`\tau`.
+        :return: (torch.Tensor) prior :math:`g(x)`.
+        """
+        return (
+            0.5
+            * ths
+            * torch.norm(x.contiguous().view(x.shape[0], -1), p=2, dim=-1) ** 2
+        )
+
+    def grad(self, x):
+        r"""
+        Calculates the gradient of the Tikhonov regularization term :math:`g` at :math:`x`.
+
+        :param torch.Tensor x: Variable :math:`x` at which the gradient is computed.
+        :return: (torch.Tensor) gradient at :math:`x`.
+        """
+        return x
+
+    def prox(self, x, ths=1.0, gamma=1.0):
+        r"""
+        Calculates the proximity operator of the Tikhonov regularization term :math:`\gamma \tau g` at :math:`x`.
+
+        :param torch.Tensor x: Variable :math:`x` at which the proximity operator is computed.
+        :param float ths: regularization parameter :math:`\tau`.
+        :param float gamma: stepsize of the proximity operator.
+        :return: (torch.Tensor) proximity operator at :math:`x`.
+        """
+        return (1 / (ths * gamma + 1)) * x
+
+
+class L1Prior(Prior):
+    r"""
+    :math:`\ell_1` prior :math:`g(x) = \| x \|_1`.
+
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.explicit_prior = True
+
+    def g(self, x, ths=1.0):
+        r"""
+        Computes the regularizer :math:`g(x) = \tau\| x \|_1`.
+
+        :param torch.Tensor x: Variable :math:`x` at which the prior is computed.
+        :param float ths: threshold parameter :math:`\tau`.
+        :return: (torch.Tensor) prior :math:`g(x)`.
+        """
+        return ths * torch.norm(x.contiguous().view(x.shape[0], -1), p=1, dim=-1)
+
+    def prox(self, x, ths=1.0, gamma=1.0):
+        r"""
+        Calculates the proximity operator of the l1 regularization term :math:`g` at :math:`x`.
+
+        More precisely, it computes
+
+        .. math::
+            \operatorname{prox}_{\gamma \tau g}(x) = \operatorname{sign}(x) \max(|x| - \gamma \tau, 0)
+
+
+        where :math:`\tau` is the threshold parameter and :math:`\gamma` is a stepsize.
+
+        :param torch.Tensor x: Variable :math:`x` at which the proximity operator is computed.
+        :param float ths: threshold parameter :math:`\tau`.
+        :param float gamma: stepsize of the proximity operator.
+        :return: (torch.Tensor) proximity operator at :math:`x`.
+        """
+        return torch.sign(x) * torch.max(
+            torch.abs(x) - ths * gamma, torch.zeros_like(x)
+        )

deepinv/optim/utils.py

@@ -0,0 +1,80 @@
+from deepinv.utils import zeros_like
+
+
+def check_conv(X_prev, X, it, crit_conv="residual", thres_conv=1e-3, verbose=False):
+    if crit_conv == "residual":
+        if isinstance(X_prev, dict):
+            X_prev = X_prev["est"][0]
+        if isinstance(X, dict):
+            X = X["est"][0]
+        crit_cur = (X_prev - X).norm() / (X.norm() + 1e-06)
+    elif crit_conv == "cost":
+        F_prev = X_prev["cost"]
+        F = X["cost"]
+        crit_cur = (F_prev - F).norm() / (F.norm() + 1e-06)
+    else:
+        raise ValueError("convergence criteria not implemented")
+    if crit_cur < thres_conv:
+        if verbose:
+            print(
+                f"Iteration {it}, current converge crit. = {crit_cur:.2E}, objective = {thres_conv:.2E} \r"
+            )
+        return True
+    else:
+        return False
+
+
+def conjugate_gradient(A, b, max_iter=1e2, tol=1e-5):
+    """
+    Standard conjugate gradient algorithm to solve Ax=b
+        see: http://en.wikipedia.org/wiki/Conjugate_gradient_method
+    :param A: Linear operator as a callable function, has to be square!
+    :param b: input tensor
+    :param max_iter: maximum number of CG iterations
+    :param tol: absolute tolerance for stopping the CG algorithm.
+    :return: torch tensor x verifying Ax=b
+
+    """
+
+    def dot(s1, s2):
+        dot = (s1 * s2).flatten().sum()
+        return dot
+
+    x = zeros_like(b)
+
+    r = b
+    p = r
+    rsold = dot(r, r)
+
+    for i in range(int(max_iter)):
+        Ap = A(p)
+        alpha = rsold / dot(p, Ap)
+        x = x + p * alpha
+        r = r + Ap * (-alpha)
+        rsnew = dot(r, r)
+        # print(rsnew.sqrt())
+        if rsnew.sqrt() < tol:
+            break
+        p = r + p * (rsnew / rsold)
+        rsold = rsnew
+
+    return x
+
+
+def gradient_descent(grad_f, x, step_size=1.0, max_iter=1e2, tol=1e-5):
+    """
+    Standard gradient descent algorithm to solve min_x f(x)
+    :param grad_f: gradient of function to bz minimized as a callable function.
+    :param x: input tensor
+    :param step_size: (constant) step size of the gradient descent algorithm.
+    :param max_iter: maximum number of iterations
+    :param tol: absolute tolerance for stopping the algorithm.
+    :return: torch tensor x verifying min_x f(x)
+
+    """
+    for i in range(int(max_iter)):
+        x_prev = x
+        x = x - grad_f(x) * step_size
+        if check_conv(x_prev, x, i, thres_conv=tol):
+            break
+    return x

deepinv/physics/__init__.py

@@ -0,0 +1,18 @@
+from .inpainting import Inpainting
+from .compressed_sensing import CompressedSensing
+from .blur import Blur, BlindBlur, Downsampling, BlurFFT
+from .range import Decolorize
+from .haze import Haze
+from .forward import Denoising, Physics, LinearPhysics, DecomposablePhysics
+from .noise import (
+    GaussianNoise,
+    PoissonNoise,
+    PoissonGaussianNoise,
+    UniformNoise,
+    UniformGaussianNoise,
+)
+from .mri import MRI
+from .tomography import Tomography
+from .lidar import SinglePhotonLidar
+from .singlepixel import SinglePixelCamera
+from .remote_sensing import Pansharpen