Source code for pyproximal.proximal.Quadratic

from typing import TYPE_CHECKING, Optional

import numpy as np
from pylops import Identity, MatrixMult
from pylops.utils.typing import NDArray
from scipy.sparse.linalg import lsqr

from pyproximal.ProxOperator import ProxOperator, _check_tau

if TYPE_CHECKING:
    from pylops.linearoperator import LinearOperator


class Quadratic(ProxOperator):
    r"""Quadratic function proximal operator.

    Proximal operator for a quadratic function: :math:`f(\mathbf{x}) =
    \frac{1}{2} \mathbf{x}^T \mathbf{Op} \mathbf{x} + \mathbf{b}^T
    \mathbf{x} + c`.

    Parameters
    ----------
    Op : :obj:`pylops.LinearOperator`, optional
        Linear operator (must be square)
    b : :obj:`np.ndarray`, optional
        Vector
    c : :obj:`float`, optional
        Scalar
    niter : :obj:`int`, optional
        Number of iterations of iterative scheme used to compute the proximal
    x0 : :obj:`np.ndarray`, optional
        Initial vector
    warm : :obj:`bool`, optional
        Warm start (``True``) or not (``False``). Uses estimate from previous
        call of ``prox`` method.

    Raises
    ------
    ValueError
        If ``Op`` is not square

    Notes
    -----
    The Quadratic proximal operator is defined as:

    .. math::

        \prox_{\tau f}(\mathbf{x}) =
        \left(\mathbf{I} + \tau  \mathbf{Op} \right)^{-1} \left(\mathbf{x} -
        \tau \mathbf{b}\right)

    when both ``Op`` and ``b`` are provided. This formula shows that the
    proximal operator requires the solution of an inverse problem. If the
    operator ``Op`` is of kind ``explicit=True``, we can solve this problem
    directly. On the other hand if ``Op`` is of kind ``explicit=False``, an
    iterative solver is employed. In this case it is possible to provide a warm
    start via the ``x0`` input parameter.

    When only ``b`` is provided, the proximal operator reduces to:

    .. math::

        \prox_{\mathbf{b}^T \mathbf{x} + c}(\mathbf{x}) =
        \mathbf{x} - \tau \mathbf{b}

    Finally if also ``b`` is not provided, the proximal operator of a constant
    function simply becomes :math:`\prox_c(\mathbf{x}) = \mathbf{x}`


    """

    def __init__(
        self,
        Op: Optional["LinearOperator"] = None,
        b: Optional[NDArray] = None,
        c: float = 0.0,
        niter: int = 10,
        x0: Optional[NDArray] = None,
        warm: bool = True,
    ) -> None:
        if Op is not None:
            if Op.shape[0] != Op.shape[1]:
                raise ValueError("Op must be square")
        super().__init__(Op, True)
        self.b = b
        if self.Op is not None and self.b is None:
            self.b = np.zeros(self.Op.shape[1], dtype=self.Op.dtype)
        self.c = c
        self.niter = niter
        self.x0 = x0
        self.warm = warm

    def __call__(self, x: NDArray) -> float:
        if self.Op is not None and self.b is not None:
            f = np.dot(x, self.Op * x) / 2.0 + np.dot(self.b, x) + self.c
        elif self.b is not None:
            f = np.dot(self.b, x) + self.c
        else:
            f = self.c
        return float(f)

    @_check_tau
    def prox(self, x: NDArray, tau: float) -> NDArray:
        if self.Op is not None and self.b is not None:
            y = x - tau * self.b
            if self.Op.explicit:
                Op1 = MatrixMult(np.eye(self.Op.shape[0]) + tau * self.Op.A)
                x = Op1.div(y)
            else:
                Op1 = Identity(self.Op.shape[0], dtype=self.Op.dtype) + tau * self.Op.A
                x = lsqr(Op1, y, iter_lim=self.niter, x0=self.x0)[0]
            if self.warm:
                self.x0 = x
        elif self.b is not None:
            x = x - tau * self.b
        return x

    def grad(self, x: NDArray) -> NDArray:
        """Compute gradient

        Parameters
        ----------
        x : :obj:`np.ndarray`
            Vector

        Returns
        -------
        g : :obj:`np.ndarray`
            Gradient vector

        """
        if self.Op is not None and self.b is not None:
            g = self.Op.matvec(x) / 2.0 + x
        elif self.b is not None:
            g = x
        else:
            g = 0.0
        return g