Implementing Proximal Operators

The proximal operator of a function $g : \mathbb{R}^n \to \mathbb{R} \cup \{+\infty\}$ with step size $\lambda > 0$ is:

$$\mathrm{prox}_{\lambda g}(\mathbf{v}) = \arg\min_{\mathbf{x}} \left\{ g(\mathbf{x}) + \frac{1}{2\lambda}\|\mathbf{x} - \mathbf{v}\|_2^2 \right\}$$

Interpretation: find the point that balances minimizing $g$ against staying close to $\mathbf{v}$. When $g = 0$, the proximal operator is the identity. When $g$ is the indicator function of a set $C$ ($g(\mathbf{x}) = 0$ if $\mathbf{x} \in C$, $+\infty$ otherwise), the proximal operator is the projection onto $C$.

The proximal operator of $g(\mathbf{x}) = \|\mathbf{x}\|_1$ is the soft-thresholding (shrinkage) operator, applied element-wise:

$$\mathcal{S}_\lambda(v_i) = \mathrm{sign}(v_i) \max(|v_i| - \lambda, 0)$$

def soft_threshold(v, lam):
    """Proximal operator of lam * ||x||_1."""
    return np.sign(v) * np.maximum(np.abs(v) - lam, 0)

This shrinks each component toward zero by $\lambda$. Components with $|v_i| \le \lambda$ are set exactly to zero, producing sparsity — the mechanism behind LASSO and compressed sensing.

The proximal operator of the indicator function $\iota_C$ of a closed convex set $C$ is the Euclidean projection:

$$\mathrm{prox}_{\iota_C}(\mathbf{v}) = \Pi_C(\mathbf{v}) = \arg\min_{\mathbf{x} \in C} \|\mathbf{x} - \mathbf{v}\|_2$$

Common projections:

# Projection onto non-negative orthant
def proj_nonneg(v):
    return np.maximum(v, 0)

# Projection onto L2 ball of radius r
def proj_l2_ball(v, r):
    norm_v = np.linalg.norm(v)
    return v * min(1, r / norm_v) if norm_v > 0 else v

# Projection onto simplex {x >= 0, sum(x) = 1}
def proj_simplex(v):
    n = len(v)
    u = np.sort(v)[::-1]
    cssv = np.cumsum(u) - 1
    rho = np.max(np.where(u > cssv / np.arange(1, n+1)))
    theta = cssv[rho] / (rho + 1)
    return np.maximum(v - theta, 0)

For the group LASSO penalty $\sum_g \|\mathbf{x}_g\|_2$ (mixed $\ell_{2,1}$ norm), the proximal operator applies block-wise shrinkage:

$$[\mathrm{prox}_{\lambda g}(\mathbf{v})]_g = \mathbf{v}_g \cdot \max\!\left(1 - \frac{\lambda}{\|\mathbf{v}_g\|_2}, 0\right)$$

This sets entire groups to zero (group sparsity) rather than individual elements:

def group_soft_threshold(v, lam, groups):
    """Proximal of lam * sum_g ||x_g||_2."""
    result = np.zeros_like(v)
    for g in groups:
        vg = v[g]
        norm_vg = np.linalg.norm(vg)
        if norm_vg > lam:
            result[g] = vg * (1 - lam / norm_vg)
    return result

The total variation of a 1D signal is $\mathrm{TV}(\mathbf{x}) = \sum_{i=1}^{n-1} |x_{i+1} - x_i|$, which promotes piecewise-constant solutions.

The proximal operator of TV has no simple closed form but can be computed in $O(n)$ time using the taut-string algorithm or via iterative methods:

def prox_tv_1d(v, lam, n_iter=100):
    """Proximal of lam * TV(x) via dual (Chambolle)."""
    n = len(v)
    p = np.zeros(n - 1)  # dual variable
    for _ in range(n_iter):
        # Forward difference of primal
        x = v - lam * _adj_diff(p, n)
        # Gradient step on dual
        dx = np.diff(x)
        p = p + dx / (2 * lam)
        # Project onto [-1, 1]
        p = np.clip(p, -1, 1)
    return v - lam * _adj_diff(p, n)

def _adj_diff(p, n):
    """Adjoint of forward difference: -div."""
    d = np.zeros(n)
    d[0] = -p[0]
    d[1:-1] = p[:-1] - p[1:]
    d[-1] = p[-1]
    return d