r"""
Concave penalties
=================
This example considers proximal operators of concave separable penalties.

"""
import matplotlib.pyplot as plt
import numpy as np

import pyproximal

plt.close('all')
x = np.linspace(-5, 5, 101)


def compare_penalty_and_proximal_operator(penalty):
    fig, ax = plt.subplots(1, 2)
    ax[0].plot(x, penalty.elementwise(x), label=penalty.__class__.__name__)
    ax[0].plot(x, np.abs(x), 'k--', label='l1')
    ax[0].set_aspect(2)
    ax[0].set_xlim(-5, 5)
    ax[0].set_ylim(0, 5)
    ax[0].set_title('Penalty')
    ax[0].legend(loc='upper center')
    for tau in [0.25, 0.50, 0.75, 1.00]:
        ax[1].plot(x, penalty.prox(x, tau), label=f'tau={tau:.2f}')
    ax[1].plot(x, x, 'k--')
    ax[1].set_aspect('equal', 'box')
    ax[1].set_xlim(-5, 5)
    ax[1].set_ylim(-5, 5)
    ax[1].set_title('Proximal operator')
    ax[1].legend(loc='upper left')


###############################################################################
# The SCAD penalty combines both soft-thresholding and hard-thresholding in a
# continuous manner
scad = pyproximal.SCAD(1, 3.7)
compare_penalty_and_proximal_operator(scad)

###############################################################################
# The Log penalty encourages sparsity more than the l1-penalty and
# parametrizes a family of functions which lie between l0 and l1-penalties.
log = pyproximal.Log(1, 0.5)
compare_penalty_and_proximal_operator(log)

###############################################################################
# The ETP penalty is similar to the Log penalty in that it tends to the
# l1-penalty and the l0-penalty at its extremes.
etp = pyproximal.ETP(1, 0.25)
compare_penalty_and_proximal_operator(etp)


###############################################################################
# The Geman penalty
geman = pyproximal.Geman(3, 1.2)
compare_penalty_and_proximal_operator(geman)


###############################################################################
# The quadratic envelope of the l0-penalty
f_mu = pyproximal.QuadraticEnvelopeCard(1.5)
compare_penalty_and_proximal_operator(f_mu)