deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/optim/fixed_point.py

@@ -0,0 +1,289 @@
+import torch
+import torch.nn as nn
+import warnings
+
+
+class FixedPoint(nn.Module):
+    """
+    Fixed-point iterations module.
+
+    This module implements the fixed-point iteration algorithm given a specific fixed-point iterator (e.g.
+    proximal gradient iteration, the ADMM iteration, see :meth:`deepinv.optim.optim_iterators`), that is
+    for :math:`k=1,2,...`
+
+    .. math::
+        \qquad (x_{k+1}, u_{k+1}) = \operatorname{FixedPoint}(x_k, u_k, f, g, A, y, ...) \hspace{2cm} (1)
+
+
+    where :math:`f` is the data-fidelity term, :math:`g` is the prior, :math:`A` is the physics model, :math:`y` is the
+    data.
+
+
+    ::
+
+        # This example shows how to use the FixedPoint class to solve the problem
+        #                min_x 0.5*lambda*||Ax-y||_2^2 + ||x||_1
+        # with the PGD algorithm, where A is the identity operator, lambda = 1 and y = [2, 2].
+
+        # Create the measurement operator A
+        A = torch.tensor([[1, 0], [0, 1]], dtype=torch.float64)
+        A_forward = lambda v: A @ v
+        A_adjoint = lambda v: A.transpose(0, 1) @ v
+
+        # Define the physics model associated to this operator
+        physics = dinv.physics.LinearPhysics(A=A_forward, A_adjoint=A_adjoint)
+
+        # Define the measurement y
+        y = torch.tensor([2, 2], dtype=torch.float64)
+
+        # Define the data fidelity term
+        data_fidelity = L2()
+
+        # Define the proximity operator of the prior and store it in a dictionary
+        def prox_g(x, g_param=0.1):
+            return torch.sign(x) * torch.maximum(x.abs() - g_param, torch.tensor([0]))
+
+        prior = {"prox_g": prox_g}
+
+        # Define the parameters of the algorithm
+        params = {"g_param": 1.0, "stepsize": 1.0, "lambda": 1.0}
+
+        # Choose the iterator associated to the PGD algorithm
+        iterator = PGDIteration(data_fidelity=data_fidelity)
+
+        # Iterate the iterator
+        x_init = torch.tensor([2, 2], dtype=torch.float64)  # Define initialisation of the algorithm
+        X = {"est": (x_init ,), "cost": []}                 # Iterates are stored in a dictionary of the form {'est': (x,z), 'cost': F}
+
+        max_iter = 50
+        for it in range(max_iter):
+            X = iterator(X,  prior, params, y, physics)
+
+        # Return the solution
+        sol = X["est"][0]  # sol = [1, 1]
+
+
+    :param deepinv.optim.optim_iterators.optim_iterator iterator: function that takes as input the current iterate, as
+                                        well as parameters of the optimization problem (prior, measurements, etc.)
+    :param function update_params_fn: function that returns the parameters to be used at each iteration. Default: ``None``.
+    :param function update_prior_fn: function that returns the prior to be used at each iteration. Default: ``None``.
+    :param function init_iterate_fn: function that returns the initial iterate. Default: ``None``.
+    :param function init_metrics_fn: function that returns the initial metrics. Default: ``None``.
+    :param function check_iteration_fn: function that performs a check on the last iteration and returns a bool indicating if we can proceed to next iteration. Default: ``None``.
+    :param function check_conv_fn: function that checks the convergence after each iteration, returns a bool indicating if convergence has been reached. Default: ``None``.
+    :param int max_iter: maximum number of iterations. Default: ``50``.
+    :param bool early_stop: if True, the algorithm stops when the convergence criterion is reached. Default: ``True``.
+    :param bool anderson_acceleration: if True, the Anderson acceleration is used. Default: ``False``.
+    :param int history_size: size of the history used for the Anderson acceleration. Default: ``5``.
+    :param float beta_anderson_acc: momentum of the Anderson acceleration step. Default: ``1.0``.
+    :param float eps_anderson_acc: regularization parameter of the Anderson acceleration step. Default: ``1e-4``.
+    """
+
+    def __init__(
+        self,
+        iterator=None,
+        update_params_fn=None,
+        update_data_fidelity_fn=None,
+        update_prior_fn=None,
+        init_iterate_fn=None,
+        init_metrics_fn=None,
+        update_metrics_fn=None,
+        check_iteration_fn=None,
+        check_conv_fn=None,
+        max_iter=50,
+        early_stop=True,
+        anderson_acceleration=False,
+        history_size=5,
+        beta_anderson_acc=1.0,
+        eps_anderson_acc=1e-4,
+    ):
+        super().__init__()
+        self.iterator = iterator
+        self.max_iter = max_iter
+        self.early_stop = early_stop
+        self.update_params_fn = update_params_fn
+        self.update_data_fidelity_fn = update_data_fidelity_fn
+        self.update_prior_fn = update_prior_fn
+        self.init_iterate_fn = init_iterate_fn
+        self.init_metrics_fn = init_metrics_fn
+        self.update_metrics_fn = update_metrics_fn
+        self.check_conv_fn = check_conv_fn
+        self.check_iteration_fn = check_iteration_fn
+        self.anderson_acceleration = anderson_acceleration
+        self.history_size = history_size
+        self.beta_anderson_acc = beta_anderson_acc
+        self.eps_anderson_acc = eps_anderson_acc
+
+        if self.check_conv_fn is None and self.early_stop:
+            warnings.warn(
+                "early_stop is set to True but no check_conv_fn has been defined."
+            )
+            self.early_stop = False
+
+    def init_anderson_acceleration(self, X):
+        r"""
+        Initialize the Anderson acceleration algorithm.
+
+        :param dict X: initial iterate.
+        """
+        x = X["est"][0]
+        b, d, h, w = x.shape
+        x_hist = torch.zeros(
+            b, self.history_size, d * h * w, dtype=x.dtype, device=x.device
+        )  # history of iterates.
+        T_hist = torch.zeros(
+            b, self.history_size, d * h * w, dtype=x.dtype, device=x.device
+        )  # history of T(x_k) with T the fixed point operator.
+        H = torch.zeros(
+            b,
+            self.history_size + 1,
+            self.history_size + 1,
+            dtype=x.dtype,
+            device=x.device,
+        )  # H in the Anderson acceleration linear system Hp = q .
+        H[:, 0, 1:] = H[:, 1:, 0] = 1.0
+        q = torch.zeros(
+            b, self.history_size + 1, 1, dtype=x.dtype, device=x.device
+        )  # q in the Anderson acceleration linear system Hp = q .
+        q[:, 0] = 1
+        return x_hist, T_hist, H, q
+
+    def anderson_acceleration_step(
+        self,
+        it,
+        X_prev,
+        TX_prev,
+        x_hist,
+        T_hist,
+        H,
+        q,
+        cur_data_fidelity,
+        cur_prior,
+        cur_params,
+        *args,
+    ):
+        r"""
+        Anderson acceleration step.
+
+        :param int it: current iteration.
+        :param dict X_prev: previous iterate.
+        :param dict TX_prev: output of the fixed-point operator evaluated at X_prev
+        :param torch.Tensor x_hist: history of last ``history-size`` iterates.
+        :param torch.Tensor T_hist: history of T evlauaton at the last ``history-size``, where T is the fixed-point operator.
+        :param torch.Tensor H: H in the Anderson acceleration linear system Hp = q .
+        :param torch.Tensor q: q in the Anderson acceleration linear system Hp = q .
+        :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 args: arguments for the iterator.
+        """
+        x_prev = X_prev["est"][0]  # current iterate Tx
+        Tx_prev = TX_prev["est"][0]  # current iterate x
+        b = x_prev.shape[0]  # batchsize
+        x_hist[:, it % self.history_size] = x_prev.reshape(
+            b, -1
+        )  # prepare history of x
+        T_hist[:, it % self.history_size] = Tx_prev.reshape(
+            b, -1
+        )  # prepare history of Tx
+        m = min(it + 1, self.history_size)
+        G = T_hist[:, :m] - x_hist[:, :m]
+        H[:, 1 : m + 1, 1 : m + 1] = (
+            torch.bmm(G, G.transpose(1, 2))
+            + self.eps_anderson_acc
+            * torch.eye(m, dtype=Tx_prev.dtype, device=Tx_prev.device)[None]
+        )
+        p = torch.linalg.solve(H[:, : m + 1, : m + 1], q[:, : m + 1])[
+            :, 1 : m + 1, 0
+        ]  # solve the linear system H p = q.
+        x = (
+            self.beta_anderson_acc * (p[:, None] @ T_hist[:, :m])[:, 0]
+            + (1 - self.beta_anderson_acc) * (p[:, None] @ x_hist[:, :m])[:, 0]
+        )  # Anderson acceleration step.
+        x = x.view(x_prev.shape)
+        F = (
+            self.iterator.F_fn(x, cur_data_fidelity, cur_prior, cur_params, *args)
+            if self.iterator.has_cost
+            else None
+        )
+        est = list(TX_prev["est"])
+        est[0] = x
+        return {"est": est, "cost": F}
+
+    def forward(self, *args, compute_metrics=False, x_gt=None, **kwargs):
+        r"""
+        Loops over the fixed-point iterator as (1) and returns the fixed point.
+
+        The iterates are stored in a dictionary of the form ``X = {'est': (x_k, u_k), 'cost': F_k}`` where:
+            - ``est`` is a tuple containing the current primal and auxiliary iterates,
+            - ``cost`` is the value of the cost function at the current iterate.
+
+        Since the prior and parameters (stepsize, regularisation parameter, etc.) can change at each iteration,
+        the prior and parameters are updated before each call to the iterator.
+
+        :param bool compute_metrics: if ``True``, the metrics are computed along the iterations. Default: ``False``.
+        :param torch.Tensor x_gt: ground truth solution. Default: ``None``.
+        :param args: optional arguments for the iterator. Commonly (y,physics) where ``y`` (torch.Tensor y) is the measurement and
+                    ``physics`` (deepinv.physics) is the physics model.
+        :param kwargs: optional keyword arguments for the iterator.
+        :return tuple: ``(x,metrics)`` with ``x`` the fixed-point solution (dict) and
+                    ``metrics`` the computed along the iterations if ``compute_metrics`` is ``True`` or ``None``
+                     otherwise.
+        """
+        X = (
+            self.init_iterate_fn(*args, F_fn=self.iterator.F_fn)
+            if self.init_iterate_fn
+            else None
+        )
+        metrics = (
+            self.init_metrics_fn(X, x_gt=x_gt)
+            if self.init_metrics_fn and compute_metrics
+            else None
+        )
+        if self.anderson_acceleration:
+            x_hist, T_hist, H, q = self.init_anderson_acceleration(X)
+        it = 0
+        while it < self.max_iter:
+            cur_params = self.update_params_fn(it) if self.update_params_fn else None
+            cur_data_fidelity = (
+                self.update_data_fidelity_fn(it)
+                if self.update_data_fidelity_fn
+                else None
+            )
+            cur_prior = self.update_prior_fn(it) if self.update_prior_fn else None
+            X_prev = X
+            X = self.iterator(X_prev, cur_data_fidelity, cur_prior, cur_params, *args)
+            if self.anderson_acceleration:
+                X = self.anderson_acceleration_step(
+                    it,
+                    X_prev,
+                    X,
+                    x_hist,
+                    T_hist,
+                    H,
+                    q,
+                    cur_data_fidelity,
+                    cur_prior,
+                    cur_params,
+                    *args,
+                )
+            check_iteration = (
+                self.check_iteration_fn(X_prev, X) if self.check_iteration_fn else True
+            )
+            if check_iteration:
+                metrics = (
+                    self.update_metrics_fn(metrics, X_prev, X, x_gt=x_gt)
+                    if self.update_metrics_fn and compute_metrics
+                    else None
+                )
+                if (
+                    self.early_stop
+                    and (self.check_conv_fn is not None)
+                    and it > 1
+                    and self.check_conv_fn(it, X_prev, X)
+                ):
+                    break
+                it += 1
+            else:
+                X = X_prev
+        return X, metrics

deepinv/optim/optim_iterators/__init__.py

@@ -0,0 +1,9 @@
+from .optim_iterator import OptimIterator
+from .optim_iterator import fStep, gStep
+from .admm import ADMMIteration
+from .pgd import PGDIteration
+from .primal_dual_CP import CPIteration
+from .hqs import HQSIteration
+from .drs import DRSIteration
+from .gradient_descent import GDIteration
+from .optim_iterator import OptimIterator, fStep, gStep

deepinv/optim/optim_iterators/admm.py

@@ -0,0 +1,117 @@
+import torch
+from .optim_iterator import OptimIterator, fStep, gStep
+
+
+class ADMMIteration(OptimIterator):
+    r"""
+    Iterator for alternating direction method of multipliers.
+
+    Class for a single iteration of the Alternating Direction Method of Multipliers (ADMM) algorithm for
+    minimising :math:`\lambda f(x) + g(x)`.
+
+    If the attribute ``g_first`` is set to False (by default),
+    the iteration is (`see this paper <https://www.nowpublishers.com/article/Details/MAL-016>`_):
+
+    .. math::
+        \begin{equation*}
+        \begin{aligned}
+        u_{k+1} &= \operatorname{prox}_{\gamma \lambda f}(x_k - z_k) \\
+        x_{k+1} &= \operatorname{prox}_{\gamma g}(u_{k+1} + z_k) \\
+        z_{k+1} &= z_k + \beta (u_{k+1} - x_{k+1})
+        \end{aligned}
+        \end{equation*}
+
+    where :math:`\gamma>0` is a stepsize and :math:`\beta>0` is a relaxation parameter.
+
+    If the attribute ``g_first`` is set to ``True``, the functions :math:`f` and :math:`g` are
+    inverted in the previous iteration.
+
+    """
+
+    def __init__(self, **kwargs):
+        super(ADMMIteration, self).__init__(**kwargs)
+        self.g_step = gStepADMM(**kwargs)
+        self.f_step = fStepADMM(**kwargs)
+        self.requires_prox_g = True
+
+    def forward(self, X, cur_data_fidelity, cur_prior, cur_params, y, physics):
+        r"""
+        Single iteration of the ADMM 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 observation.
+        :return: Dictionary `{"est": (x, z), "cost": F}` containing the updated current iterate and the estimated current cost.
+        """
+        x, z = X["est"]
+        if z.shape != x.shape:
+            # In ADMM, the "dual" variable z is a fake dual variable as it lives in the primal, hence this line to prevent from usual initialisation
+            z = torch.zeros_like(x)
+        if self.g_first:
+            u = self.g_step(x, z, cur_prior, cur_params)
+            x = self.f_step(u, z, cur_data_fidelity, cur_params, y, physics)
+        else:
+            u = self.f_step(x, z, cur_data_fidelity, cur_params, y, physics)
+            x = self.g_step(u, z, cur_prior, cur_params)
+        z = z + cur_params["beta"] * (u - x)
+        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 fStepADMM(fStep):
+    r"""
+    ADMM fStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(fStepADMM, self).__init__(**kwargs)
+
+    def forward(self, x, z, cur_data_fidelity, cur_params, y, physics):
+        r"""
+        Single iteration step on the data-fidelity term :math:`\lambda f`.
+
+        :param torch.Tensor x: current first variable
+        :param torch.Tensor z: current second variable
+        :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.
+        """
+        if self.g_first:
+            p = x + z
+        else:
+            p = x - z
+        return cur_data_fidelity.prox(
+            p, y, physics, gamma=cur_params["lambda"] * cur_params["stepsize"]
+        )
+
+
+class gStepADMM(gStep):
+    r"""
+    ADMM gStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(gStepADMM, self).__init__(**kwargs)
+
+    def forward(self, x, z, cur_prior, cur_params):
+        r"""
+        Single iteration step on the prior term :math:`g`.
+
+        :param torch.Tensor x: current first variable
+        :param torch.Tensor z: current second variable
+        :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.
+        """
+        if self.g_first:
+            p = x - z
+        else:
+            p = x + z
+        return cur_prior.prox(p, cur_params["g_param"], gamma=cur_params["stepsize"])

deepinv/optim/optim_iterators/drs.py

@@ -0,0 +1,115 @@
+import torch
+
+from .optim_iterator import OptimIterator, fStep, gStep
+
+
+class DRSIteration(OptimIterator):
+    r"""
+    Iterator for Douglas-Rachford Splitting.
+
+    Class for a single iteration of the Douglas-Rachford Splitting (DRS) algorithm for minimising
+    :math:`\lambda f(x) + g(x)`.
+
+    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}_{\gamma \lambda f}(z_k) \\
+        x_{k+1} &= \operatorname{prox}_{\gamma g}(2*u_{k+1}-z_k) \\
+        z_{k+1} &= z_k + \beta (x_{k+1} - u_{k+1})
+        \end{aligned}
+        \end{equation*}
+
+    where :math:`\gamma>0` is a stepsize and :math:`\beta>0` is a relaxation parameter.
+
+    If the attribute ``g_first`` is set to True, the functions :math:`f` and :math:`g` are inverted in the previous iteration.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.g_step = gStepDRS(**kwargs)
+        self.f_step = fStepDRS(**kwargs)
+        self.requires_prox_g = True
+
+    def forward(self, X, cur_data_fidelity, cur_prior, cur_params, y, physics):
+        r"""
+        Single iteration of the DRS 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 observation.
+        :return: Dictionary `{"est": (x, z), "cost": F}` containing the updated current iterate and the estimated current cost.
+        """
+        x, z = X["est"]
+        if z.shape != x.shape:
+            # In DRS, the "dual" variable z is a fake dual variable as it lives in the primal, hence this line to prevent from usual initialisation
+            z = torch.zeros_like(x)
+        if self.g_first:
+            u = self.g_step(x, z, cur_prior, cur_params)
+            x = self.f_step(u, z, cur_data_fidelity, cur_params, y, physics)
+        else:
+            u = self.f_step(x, z, cur_data_fidelity, cur_params, y, physics)
+            x = self.g_step(u, z, cur_prior, cur_params)
+        z = z + cur_params["beta"] * (x - u)
+        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 fStepDRS(fStep):
+    r"""
+    DRS fStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(fStepDRS, self).__init__(**kwargs)
+
+    def forward(self, x, z, cur_data_fidelity, cur_params, y, physics):
+        r"""
+        Single iteration step on the data-fidelity term :math:`f`.
+
+        :param torch.Tensor x: Current first variable.
+        :param torch.Tensor z: Current second variable.
+        :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 self.g_first:
+            p = 2 * x - z
+        else:
+            p = z
+        return cur_data_fidelity.prox(
+            p, y, physics, gamma=cur_params["lambda"] * cur_params["stepsize"]
+        )
+
+
+class gStepDRS(gStep):
+    r"""
+    DRS gStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(gStepDRS, self).__init__(**kwargs)
+
+    def forward(self, x, z, cur_prior, cur_params):
+        r"""
+        Single iteration step on the prior term :math:`g`.
+
+        :param torch.Tensor x:  Current first variable.
+        :param torch.Tensor z: Current second variable.
+        :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.
+        """
+        if self.g_first:
+            p = z
+        else:
+            p = 2 * x - z
+        return cur_prior.prox(p, cur_params["g_param"], gamma=cur_params["stepsize"])

deepinv/optim/optim_iterators/gradient_descent.py

@@ -0,0 +1,90 @@
+from .optim_iterator import OptimIterator, fStep, gStep
+from .utils import gradient_descent_step
+
+
+class GDIteration(OptimIterator):
+    r"""
+    Iterator for Gradient Descent.
+
+    Class for a single iteration of the gradient descent (GD) algorithm for minimising :math:`\lambda f(x) + g(x)`.
+
+    The iteration is given by
+
+
+    .. math::
+        \begin{equation*}
+        \begin{aligned}
+        v_{k} &= \nabla f(x_k) + \nabla g(x_k) \\
+        x_{k+1} &= x_k-\gamma v_{k}
+        \end{aligned}
+        \end{equation*}
+
+
+   where :math:`\gamma` is a stepsize.
+    """
+
+    def __init__(self, **kwargs):
+        super(GDIteration, self).__init__(**kwargs)
+        self.g_step = gStepGD(**kwargs)
+        self.f_step = fStepGD(**kwargs)
+        self.requires_grad_g = True
+
+    def forward(self, X, cur_data_fidelity, cur_prior, cur_params, y, physics):
+        r"""
+        Single gradient descent iteration on the objective :math:`\lambda f(x) + g(x)`.
+
+        :param dict X: Dictionary containing the current iterate :math:`x_k`.
+        :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.
+        :return: Dictionary `{"est": (x, ), "cost": F}` containing the updated current iterate and the estimated current cost.
+        """
+        x_prev = X["est"][0]
+        grad = cur_params["stepsize"] * (
+            self.g_step(x_prev, cur_prior, cur_params)
+            + self.f_step(x_prev, cur_data_fidelity, cur_params, y, physics)
+        )
+        x = gradient_descent_step(x_prev, grad)
+        F = self.F_fn(x, cur_prior, cur_params, y, physics) if self.has_cost else None
+        return {"est": (x,), "cost": F}
+
+
+class fStepGD(fStep):
+    r"""
+    GD fStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(fStepGD, self).__init__(**kwargs)
+
+    def forward(self, x, cur_data_fidelity, cur_params, y, physics):
+        r"""
+        Single gradient descent iteration on the data fit 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_params["lambda"] * cur_data_fidelity.grad(x, y, physics)
+
+
+class gStepGD(gStep):
+    r"""
+    GD gStep module.
+    """
+
+    def __init__(self, **kwargs):
+        super(gStepGD, 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 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.
+        """
+        return cur_prior.grad(x, cur_params["g_param"])