def prox_max_entropy(X, step, gamma=1):
    """Proximal operator for maximum entropy regularization.

    g(x) = gamma \sum_i x_i ln(x_i)

    has the analytical solution of gamma W(1/gamma exp((X-gamma)/gamma)), where
    W is the Lambert W function.
    """
    from scipy.special import lambertw
    gamma_ = _step_gamma(step, gamma)
    # minimize entropy: return gamma_ * np.real(lambertw(np.exp((X - gamma_) / gamma_) / gamma_))
    above = X > 0
    X[above] = gamma_ * np.real(lambertw(np.exp(X[above]/gamma_ - 1) / gamma_))
    return X
Proximal operator for maximum entropy regularization.

    g(x) = gamma \sum_i x_i ln(x_i)

    has the analytical solution of gamma W(1/gamma exp((X-gamma)/gamma)), where
    W is the Lambert W function.