import numpy as np
from scipy.linalg import cho_factor, cho_solve
from scipy.sparse.linalg import lsqr as sp_lsqr
from pylops import MatrixMult, Identity
from pylops.optimization.basic import lsqr
from pylops.utils.backend import get_array_module, get_module_name
from pyproximal.ProxOperator import _check_tau
from pyproximal import ProxOperator
[docs]class L2(ProxOperator):
r"""L2 Norm proximal operator.
The Proximal operator of the :math:`\ell_2` norm is defined as: :math:`f(\mathbf{x}) =
\frac{\sigma}{2} ||\mathbf{Op}\mathbf{x} - \mathbf{b}||_2^2`
and :math:`f_\alpha(\mathbf{x}) = f(\mathbf{x}) +
\alpha \mathbf{q}^T\mathbf{x}`.
Parameters
----------
Op : :obj:`pylops.LinearOperator`, optional
Linear operator
b : :obj:`numpy.ndarray`, optional
Data vector
q : :obj:`numpy.ndarray`, optional
Dot vector
sigma : :obj:`int`, optional
Multiplicative coefficient of L2 norm
alpha : :obj:`float`, optional
Multiplicative coefficient of dot product
qgrad : :obj:`bool`, optional
Add q term to gradient (``True``) or not (``False``)
niter : :obj:`int` or :obj:`func`, optional
Number of iterations of iterative scheme used to compute the proximal.
This can be a constant number or a function that is called passing a
counter which keeps track of how many times the ``prox`` method has
been invoked before and returns the ``niter`` to be used.
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.
densesolver : :obj:`str`, optional
Use ``numpy``, ``scipy``, or ``factorize`` when dealing with explicit
operators. The former two rely on dense solvers from either library,
whilst the last computes a factorization of the matrix to invert and
avoids to do so unless the :math:`\tau` or :math:`\sigma` paramets
have changed. Choose ``densesolver=None`` when using PyLops versions
earlier than v1.18.1 or v2.0.0
**kwargs_solver : :obj:`dict`, optional
Dictionary containing extra arguments for
:py:func:`scipy.sparse.linalg.lsqr` solver when using
numpy data (or :py:func:`pylops.optimization.solver.lsqr` and
when using cupy data)
Notes
-----
The L2 proximal operator is defined as:
.. math::
prox_{\tau f_\alpha}(\mathbf{x}) =
\left(\mathbf{I} + \tau \sigma \mathbf{Op}^T \mathbf{Op} \right)^{-1}
\left( \mathbf{x} + \tau \sigma \mathbf{Op}^T \mathbf{b} -
\tau \alpha \mathbf{q}\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, ``Op`` is assumed to be an Identity operator
and the proximal operator reduces to:
.. math::
\prox_{\tau f_\alpha}(\mathbf{x}) =
\frac{\mathbf{x} + \tau \sigma \mathbf{b} - \tau \alpha \mathbf{q}}
{1 + \tau \sigma}
If ``b`` is not provided, the proximal operator reduces to:
.. math::
\prox_{\tau f_\alpha}(\mathbf{x}) =
\frac{\mathbf{x} - \tau \alpha \mathbf{q}}{1 + \tau \sigma}
Finally, note that the second term in :math:`f_\alpha(\mathbf{x})` is added
because this combined expression appears in several problems where Bregman
iterations are used alongside a proximal solver.
"""
def __init__(self, Op=None, b=None, q=None, sigma=1., alpha=1.,
qgrad=True, niter=10, x0=None, warm=True,
densesolver=None, kwargs_solver=None):
super().__init__(Op, True)
self.b = b
self.q = q
self.sigma = sigma
self.alpha = alpha
self.qgrad = qgrad
self.niter = niter
self.x0 = x0
self.warm = warm
self.densesolver = densesolver
self.count = 0
self.kwargs_solver = {} if kwargs_solver is None else kwargs_solver
# when using factorize, store the first tau*sigma=0 so that the
# first time it will be recomputed (as tau cannot be 0)
if self.densesolver == 'factorize':
self.tausigma = 0
# create data term
if self.Op is not None and self.b is not None:
self.OpTb = self.sigma * self.Op.H @ self.b
# create A.T A upfront for explicit operators
if self.Op.explicit:
self.ATA = np.conj(self.Op.A.T) @ self.Op.A
def __call__(self, x):
if self.Op is not None and self.b is not None:
f = (self.sigma / 2.) * (np.linalg.norm(self.Op * x - self.b) ** 2)
elif self.b is not None:
f = (self.sigma / 2.) * (np.linalg.norm(x - self.b) ** 2)
else:
f = (self.sigma / 2.) * (np.linalg.norm(x) ** 2)
if self.q is not None:
f += self.alpha * np.dot(self.q, x)
return f
def _increment_count(func):
"""Increment counter
"""
def wrapped(self, *args, **kwargs):
self.count += 1
return func(self, *args, **kwargs)
return wrapped
@_increment_count
@_check_tau
def prox(self, x, tau):
# define current number of iterations
if isinstance(self.niter, int):
niter = self.niter
else:
niter = self.niter(self.count)
# solve proximal optimization
if self.Op is not None and self.b is not None:
y = x + tau * self.OpTb
if self.q is not None:
y -= tau * self.alpha * self.q
if self.Op.explicit:
if self.densesolver != 'factorize':
Op1 = MatrixMult(np.eye(self.Op.shape[1]) +
tau * self.sigma * self.ATA)
if self.densesolver is None:
# to allow backward compatibility with PyLops versions earlier
# than v1.18.1 and v2.0.0
x = Op1.div(y)
else:
x = Op1.div(y, densesolver=self.densesolver)
else:
if self.tausigma != tau * self.sigma:
# recompute factorization
self.tausigma = tau * self.sigma
ATA = np.eye(self.Op.shape[1]) + \
self.tausigma * self.ATA
self.cl = cho_factor(ATA)
x = cho_solve(self.cl, y)
else:
Op1 = Identity(self.Op.shape[1], dtype=self.Op.dtype) + \
float(tau * self.sigma) * (self.Op.H * self.Op)
if get_module_name(get_array_module(x)) == 'numpy':
x = sp_lsqr(Op1, y, iter_lim=niter, x0=self.x0,
**self.kwargs_solver)[0]
else:
x = lsqr(Op1, y, niter=niter, x0=self.x0,
**self.kwargs_solver)[0].ravel()
if self.warm:
self.x0 = x
elif self.b is not None:
num = x + tau * self.sigma * self.b
if self.q is not None:
num -= tau * self.alpha * self.q
x = num / (1. + tau * self.sigma)
else:
num = x
if self.q is not None:
num -= tau * self.alpha * self.q
x = num / (1. + tau * self.sigma)
return x
def grad(self, x):
if self.Op is not None and self.b is not None:
g = self.sigma * self.Op.H @ (self.Op @ x - self.b)
elif self.b is not None:
g = self.sigma * (x - self.b)
else:
g = self.sigma * x
if self.q is not None and self.qgrad:
g += self.alpha * self.q
return g
[docs]class L2Convolve(ProxOperator):
r"""L2 Norm proximal operator with convolution operator
Proximal operator for the L2 norm defined as: :math:`f(\mathbf{x}) =
\frac{\sigma}{2} ||\mathbf{h} * \mathbf{x} - \mathbf{b}||_2^2` where
:math:`\mathbf{h}` is the kernel of a convolution operator and
:math:`*` represents convolution
Parameters
----------
h : :obj:`np.ndarray`, optional
Kernel of convolution operator
b : :obj:`numpy.ndarray`, optional
Data vector
b : :obj:`int`, optional
Fourier transform number of samples
sigma : :obj:`int`, optional
Multiplicative coefficient of L2 norm
dims : :obj:`tuple`, optional
Number of samples for each dimension
(``None`` if only one dimension is available)
dir : :obj:`int`, optional
Direction along which smoothing is applied.
Notes
-----
The L2Convolve proximal operator is defined as:
.. math::
prox_{\tau f}(\mathbf{x}) =
F^{-1}\left(\frac{\tau\sigma F(\mathbf{h})^* F(\mathbf{b}) + F(\mathbf{x})}
{1 + \tau\sigma F(\mathbf{h})^* F(\mathbf{h})} \right)
"""
def __init__(self, h, b=None, nfft=2**10, sigma=1., dims=None, dir=None):
super().__init__(None, True)
self.nfft = nfft
self.sigma = sigma
self.dims = dims
self.dir = -1 if dir is None else dir
# convert data and filter to Fourier domain
self.bf = np.fft.fft(b, self.nfft, axis=self.dir)
self.hf = np.fft.fft(h, self.nfft, axis=self.dir)
# expand dimensions of filters
if self.dims is not None:
self.bf = self.bf.reshape(self.dims)
self.dimsf = list(dims).copy()
self.dimsf[dir] = nfft
ndims = len(dims)
for _ in range(dir - 1):
self.hf = np.expand_dims(self.hf, axis=0)
for _ in range(ndims - dir - 1):
self.hf = np.expand_dims(self.hf, axis=-1)
# precompute terms for prox
self.hbf = np.conj(self.hf) * self.bf
self.h2f = np.abs(self.hf) ** 2
def __call__(self, x):
if self.dims is not None:
x = x.reshape(self.dims)
xf = np.fft.fft(x, self.nfft, axis=self.dir)
f = (self.sigma / 2.) * np.linalg.norm(np.fft.ifft(self.bf - self.hf * xf,
axis=self.dir)) ** 2
return f
@_check_tau
def prox(self, x, tau):
if self.dims is not None:
x = x.reshape(self.dims)
xf = np.fft.fft(x, self.nfft, axis=self.dir)
yf = (xf + self.sigma * tau * self.hbf) / \
(1. + self.sigma * tau * self.h2f)
y = np.fft.ifft(yf, axis=self.dir)
if self.dims is None:
y = y[:len(x)]
else:
y = np.take(y, range(self.dims[self.dir]), axis=self.dir).ravel()
return y.ravel()
def grad(self, x):
if self.dims is not None:
x = x.reshape(self.dims)
xf = np.fft.fft(x, self.nfft, axis=self.dir)
f = self.sigma * np.fft.ifft(np.conj(self.hf) * (self.hf * xf - self.bf),
axis=self.dir)
return f.ravel()
