diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index abd4f2c..79fdf2f 100755 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -92,6 +92,7 @@ Non-Convex Log1 QuadraticEnvelopeCard QuadraticEnvelopeCardIndicator + RelaxedMumfordShah SCAD diff --git a/pyproximal/proximal/RelaxedMS.py b/pyproximal/proximal/RelaxedMS.py new file mode 100644 index 0000000..b685fb7 --- /dev/null +++ b/pyproximal/proximal/RelaxedMS.py @@ -0,0 +1,99 @@ +import numpy as np + +from pyproximal.ProxOperator import _check_tau +from pyproximal import ProxOperator +from pyproximal.proximal.L1 import _current_sigma + + +def _l2(x, alpha): + r"""Scaling operation. + + Applies the proximal of ``alpha||y - x||_2^2`` which is essentially a scaling operation. + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Vector + alpha : :obj:`float` + Scaling parameter + + Returns + ------- + y : :obj:`numpy.ndarray` + Scaled vector + + """ + y = 1 / (1 + 2 * alpha) * x + return y + + +def _current_kappa(kappa, count): + if not callable(kappa): + return kappa + else: + return kappa(count) + + +class RelaxedMumfordShah(ProxOperator): + r"""Relaxed Mumford-Shah norm proximal operator. + + Proximal operator of the relaxed Mumford-Shah norm: + :math:`\text{rMS}(x) = \min (\alpha\Vert x\Vert_2^2, \kappa)`. + + Parameters + ---------- + sigma : :obj:`float` or :obj:`list` or :obj:`np.ndarray` or :obj:`func`, optional + Multiplicative coefficient of L2 norm that controls the smoothness of the solutuon. + This can be a constant number, a list of values (for multidimensional inputs, acting + on the second dimension) 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 a + scalar (or a list of) ``sigma`` to be used. + kappa : :obj:`float` or :obj:`list` or :obj:`np.ndarray` or :obj:`func`, optional + Constant value in the rMS norm which essentially controls when the norm allows a jump. This can be a + constant number, a list of values (for multidimensional inputs, acting on the second dimension) 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 a scalar (or a list of) + ``kappa`` to be used. + + Notes + ----- + The :math:`rMS` proximal operator is defined as [1]_: + + .. math:: + \text{prox}_{\tau \text{rMS}}(x) = + \begin{cases} + \frac{1}{1+2\tau\alpha}x & \text{ if } & \vert x\vert \leq \sqrt{\frac{\kappa}{\alpha}(1 + 2\tau\alpha)} \\ + \kappa & \text{ else } + \end{cases}. + + .. [1] Strekalovskiy, E., and D. Cremers, 2014, Real-time minimization of the piecewise smooth + Mumford-Shah functional: European Conference on Computer Vision, 127–141. + + """ + def __init__(self, sigma=1., kappa=1.): + super().__init__(None, False) + self.sigma = sigma + self.kappa = kappa + self.count = 0 + + def __call__(self, x): + sigma = _current_sigma(self.sigma, self.count) + kappa = _current_sigma(self.kappa, self.count) + return np.minimum(sigma * np.linalg.norm(x) ** 2, kappa) + + 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): + sigma = _current_sigma(self.sigma, self.count) + kappa = _current_sigma(self.kappa, self.count) + + x = np.where(np.abs(x) <= np.sqrt(kappa / sigma * (1 + 2 * tau * sigma)), _l2(x, tau * sigma), x) + return x diff --git a/pyproximal/proximal/__init__.py b/pyproximal/proximal/__init__.py index c82d19c..6d64f2a 100644 --- a/pyproximal/proximal/__init__.py +++ b/pyproximal/proximal/__init__.py @@ -7,21 +7,22 @@ Box Box indicator Simplex Simplex indicator Intersection Intersection indicator - AffineSet Affines set indicator + AffineSet Affines set indicator Quadratic Quadratic function - Nonlinear Nonlinear function + Nonlinear Nonlinear function L0 L0 Norm L0Ball L0 Ball L1 L1 Norm L1Ball L1 Ball - Euclidean Euclidean Norm - EuclideanBall Euclidean Ball + Euclidean Euclidean Norm + EuclideanBall Euclidean Ball L2 L2 Norm L2Convolve L2 Norm of convolution operator L21 L2,1 Norm L21_plus_L1 L2,1 + L1 mixed-norm - Huber Huber Norm - TV Total Variation Norm + Huber Huber Norm + TV Total Variation Norm + RelaxedMumfordShah Relaxed Mumford Shah Norm Nuclear Nuclear Norm NuclearBall Nuclear Ball Orthogonal Product between orthogonal operator and vector @@ -53,6 +54,7 @@ from .L21_plus_L1 import * from .Huber import * from .TV import * +from .RelaxedMS import * from .Nuclear import * from .Orthogonal import * from .VStack import * @@ -66,8 +68,8 @@ __all__ = ['Box', 'Simplex', 'Intersection', 'AffineSet', 'Quadratic', 'Euclidean', 'EuclideanBall', 'L0', 'L0Ball', 'L1', 'L1Ball', 'L2', - 'L2Convolve', 'L21', 'L21_plus_L1', 'Huber', 'TV', 'Nuclear', - 'NuclearBall', 'Orthogonal', 'VStack', 'Nonlinear', 'SCAD', + 'L2Convolve', 'L21', 'L21_plus_L1', 'Huber', 'TV', 'RelaxedMumfordShah', + 'Nuclear', 'NuclearBall', 'Orthogonal', 'VStack', 'Nonlinear', 'SCAD', 'Log', 'Log1', 'ETP', 'Geman', 'QuadraticEnvelopeCard', 'SingularValuePenalty', 'QuadraticEnvelopeCardIndicator', 'QuadraticEnvelopeRankL2', 'Hankel'] diff --git a/pytests/test_norms.py b/pytests/test_norms.py index b4300c7..104fe1f 100644 --- a/pytests/test_norms.py +++ b/pytests/test_norms.py @@ -5,7 +5,8 @@ from pylops.basicoperators import Identity, Diagonal, MatrixMult, FirstDerivative from pyproximal.utils import moreau -from pyproximal.proximal import Box, Euclidean, L2, L1, L21, L21_plus_L1, Huber, Nuclear, TV +from pyproximal.proximal import Box, Euclidean, L2, L1, L21, L21_plus_L1, \ + Huber, Nuclear, RelaxedMumfordShah, TV par1 = {'nx': 10, 'sigma': 1., 'dtype': 'float32'} # even float32 par2 = {'nx': 11, 'sigma': 2., 'dtype': 'float64'} # odd float64 @@ -202,6 +203,22 @@ def test_TV(par): assert_array_almost_equal(tv(x), par['sigma'] * np.sum(np.abs(dx), axis=0)) +@pytest.mark.parametrize("par", [(par1), (par2)]) +def test_rMS(par): + """rMS norm and proximal/dual proximal + """ + kappa = 1. + rMS = RelaxedMumfordShah(sigma=par['sigma'], kappa=kappa) + + # norm + x = np.random.normal(0., 1., par['nx']).astype(par['dtype']) + assert rMS(x) == np.minimum(par['sigma'] * np.linalg.norm(x) ** 2, kappa) + + # prox / dualprox + tau = 2. + assert moreau(rMS, x, tau) + + def test_Nuclear_FOM(): """Nuclear norm benchmark with FOM solver """ diff --git a/testdata/marmousi_trace.npy b/testdata/marmousi_trace.npy new file mode 100644 index 0000000..90b16e9 Binary files /dev/null and b/testdata/marmousi_trace.npy differ diff --git a/tutorials/relaxed_mumford-shah.py b/tutorials/relaxed_mumford-shah.py new file mode 100644 index 0000000..781c900 --- /dev/null +++ b/tutorials/relaxed_mumford-shah.py @@ -0,0 +1,220 @@ +r""" +Relaxed Mumford-Shah regularization +=================================== +In this tutorial we will use a relaxed Mumford-Shah (rMS) functional [1]_ as regularization, +which has the following form: + +.. math:: + \text{rMS}(x) = \min (\alpha\Vert x\Vert_2^2, \kappa). + +Its corresponding proximal operator is given by + +.. math:: + \text{prox}_{\tau \text{rMS}}(x) = + \begin{cases} + \frac{1}{1+2\tau \alpha}x & \text{ if } & \vert x\vert \leq \sqrt{\frac{\kappa}{\alpha}(1 + 2\tau \alpha)} \\ + \kappa & \text{ else } + \end{cases}. + +rMS is a combination of Tikhonov and TV regularization. Once the rMS hits a certain threshold, the solution will be allowed +to jump due to the constant penalty :math:`\kappa`, and below this value rMS will be smooth due to Tikhonov regularization. +We show three denoising examples: one example that is well-suited for TV regularization and two examples where rMS +outperforms TV and Tikhonov regularization, modeled after the experiments in [2]_. + +**References** + +.. [1] Strekalovskiy, E., and D. Cremers, 2014, Real-time minimization of the piecewise smooth Mumford-Shah functional: European Conference on Computer Vision, 127–141 +.. [2] Kadu, A., and Kumar, R. and van Leeuwen, Tristan. Full-waveform inversion with Mumford-Shah regularization. SEG International Exposition and Annual Meeting, SEG-2018-2997224 + +""" +import numpy as np +import matplotlib.pyplot as plt +import pylops +import pyproximal + +np.random.seed(1) + +############################################################################### +# We start with a simple model with two jumps that is well-suited for TV +# regularization + +# Create noisy data +nx = 101 +idx_jump1 = nx // 3 +idx_jump2 = 3 * nx // 4 +x = np.zeros(nx) +x[:idx_jump1] = 2 +x[idx_jump1:idx_jump2] = 5 +n = np.random.normal(0, 0.5, nx) +y = x + n + +# Plot the model and the noisy data +fig, axs = plt.subplots(1, 1, figsize=(6, 5)) +axs.plot(x, label='True model') +axs.plot(y, label='Noisy model') +axs.legend() +plt.tight_layout() + +############################################################################### +# For both rMS and TV regularizations we use the Linearized ADMM, whilst +# for Tikhonov regularization we use LSQR + +# Define functionals +l2 = pyproximal.proximal.L2(b=y) +l1 = pyproximal.proximal.L1(sigma=5.) +Dop = pylops.FirstDerivative(nx, edge=True, kind='backward') + +# TV +L = np.real((Dop.H * Dop).eigs(neigs=1, which='LM')[0]) +tau = 1. +mu = 0.99 * tau / L +xTV, _ = pyproximal.optimization.primal.LinearizedADMM(l2, l1, Dop, tau=tau, mu=mu, + x0=np.zeros_like(x), niter=200) + +# rMS +sigma = 1e5 +kappa = 1e0 +ms_relaxed = pyproximal.proximal.RelaxedMumfordShah(sigma=sigma, kappa=kappa) +tau = 1. +mu = tau / L +xrMS, _ = pyproximal.optimization.primal.LinearizedADMM(l2, ms_relaxed, Dop, tau=tau, mu=mu, + x0=np.zeros_like(x), niter=200) + +# Tikhonov +xTikhonov = pylops.optimization.leastsquares.regularized_inversion(Op=pylops.Identity(nx), + Regs=[Dop, ], y=y, + epsRs=[6e0, ])[0] + +# Plot the results +fig, axs = plt.subplots(1, 1, figsize=(6, 5)) +axs.plot(x, label='True', linewidth=4, color='k') +axs.plot(y, '--', label='Noisy', linewidth=2, color='y') +axs.plot(xTV, label='TV') +axs.plot(xrMS, label='rMS') +axs.plot(xTikhonov, label='Tikhonov') +axs.legend() +plt.tight_layout() + +############################################################################### +# Next, we consider an example where we replace the first jump with a slope. +# As we will see, TV can not deal with this type of structure since a linear +# increase will greatly increase the TV norm, and instead TV will make a staircase. +# rMS, on the other hand, can reconstruct the model with high accuracy. + +nx = 101 +idx_jump1 = nx // 3 +idx_jump2 = 3 * nx // 4 +x = np.zeros(nx) +x[:idx_jump1] = 2 +x[idx_jump1:idx_jump2] = np.linspace(2, 4, idx_jump2 - idx_jump1) +n = np.random.normal(0, 0.25, nx) +y = x + n + +# Plot the model and the noisy data +fig, axs = plt.subplots(1, 1, figsize=(6, 5)) +axs.plot(x, label='True model') +axs.plot(y, label='Noisy model') +axs.legend() +plt.tight_layout() + +############################################################################### + +# Define functionals +l2 = pyproximal.proximal.L2(b=y) +l1 = pyproximal.proximal.L1(sigma=1.) +Dop = pylops.FirstDerivative(nx, edge=True, kind='backward') + +# TV +L = np.real((Dop.H * Dop).eigs(neigs=1, which='LM')[0]) +tau = 1. +mu = 0.99 * tau / L +xTV, _ = pyproximal.optimization.primal.LinearizedADMM(l2, l1, Dop, tau=tau, mu=mu, + x0=np.zeros_like(x), niter=200) + +# rMS +sigma = 1e1 +kappa = 1e0 +ms_relaxed = pyproximal.proximal.RelaxedMumfordShah(sigma=sigma, kappa=kappa) +tau = 1. +mu = tau / L +xrMS, _ = pyproximal.optimization.primal.LinearizedADMM(l2, ms_relaxed, Dop, tau=tau, mu=mu, + x0=np.zeros_like(x), niter=200) + +# Tikhonov +Op = pylops.Identity(nx) +Regs = [Dop, ] +epsR = [3e0, ] + +xTikhonov = pylops.optimization.leastsquares.regularized_inversion(Op=Op, Regs=Regs, y=y, epsRs=epsR)[0] + +# Plot the results +fig, axs = plt.subplots(1, 1, figsize=(6, 5)) +axs.plot(x, label='True', linewidth=4, color='k') +axs.plot(y, '--', label='Noisy', linewidth=2, color='y') +axs.plot(xTV, label='TV') +axs.plot(xrMS, label='rMS') +axs.plot(xTikhonov, label='Tikhonov') +axs.legend() +plt.tight_layout() + +############################################################################### +# Finally, we take a trace from a section of the Marmousi model. This trace shows +# rather smooth behavior with a few jumps, which makes it perfectly suited for rMS. +# TV on the other hand will artificially create a staircasing effect. + +# Get a trace from the model and add some noise +m_trace = np.load('../testdata/marmousi_trace.npy') +nz = len(m_trace) +m_trace_noisy = m_trace + np.random.normal(0, 0.1, nz) + +# Plot the model and the noisy data +fig, ax = plt.subplots(1, 1, figsize=(12, 5)) +ax.plot(m_trace, linewidth=2, label='True') +ax.plot(m_trace_noisy, label='Noisy') +ax.set_title('Trace and noisy trace') +ax.axis('tight') +ax.legend() +plt.tight_layout() + +############################################################################### + +# Define functionals +l2 = pyproximal.proximal.L2(b=m_trace_noisy) +l1 = pyproximal.proximal.L1(sigma=5e-1) +Dop = pylops.FirstDerivative(nz, edge=True, kind='backward') + +# TV +L = np.real((Dop.H * Dop).eigs(neigs=1, which='LM')[0]) +tau = 1. +mu = 0.99 * tau / L +xTV, _ = pyproximal.optimization.primal.LinearizedADMM(l2, l1, Dop, tau=tau, mu=mu, + x0=np.zeros_like(m_trace), niter=200) + +# rMS +sigma = 5e0 +kappa = 1e-1 +ms_relaxed = pyproximal.proximal.RelaxedMumfordShah(sigma=sigma, kappa=kappa) + +tau = 1. +mu = tau / L +xrMS, _ = pyproximal.optimization.primal.LinearizedADMM(l2, ms_relaxed, Dop, tau=tau, mu=mu, + x0=np.zeros_like(m_trace), niter=200) + +# Tikhonov +Op = pylops.Identity(nz) +Regs = [Dop, ] +epsR = [3e0, ] + +xTikhonov = pylops.optimization.leastsquares.regularized_inversion(Op=Op, Regs=Regs, + y=m_trace_noisy, + epsRs=epsR)[0] + +# Plot the results +fig, axs = plt.subplots(1, 1, figsize=(12, 5)) +axs.plot(m_trace, label='True', linewidth=4, color='k') +axs.plot(m_trace_noisy, '--', label='Noisy', linewidth=2, color='y') +axs.plot(xTV, label='TV') +axs.plot(xrMS, label='rMS') +axs.plot(xTikhonov, label='Tikhonov') +axs.legend() +plt.tight_layout()