Several proximal operators implemented in #PyTorch

proximal.py

import torch


# > https://github.com/pmelchior/proxmin/blob/master/proxmin/operators.py
def _step_gamma(step, gamma):
    """Update gamma parameter for use inside of continuous proximal operator.
    Every proximal operator for a function with a continuous parameter,
    e.g. gamma ||x||_1, needs to update that parameter to account for the
    stepsize of the algorithm.
    Returns:
        gamma * step
    """
    return gamma * step


def prox_plus(X, step):
    """Projection onto non-negative numbers
    """
    return torch.clamp(X, min=0)


def prox_soft(X, step, thresh=0):
    """Soft thresholding proximal operator
    """
    thresh_ = _step_gamma(step, thresh)
    return torch.sign(X) * prox_plus(torch.abs(X) - thresh_, step)