Projection onto the probability simplex using isotonic regression.

projection_simplex_isotonic.py

# Author: Mathieu Blondel
# License: BSD 3 clause

import numpy as np
from sklearn.isotonic import isotonic_regression


def projection_simplex(x, z=1):
    """
    Compute argmin_{p : p >= 0 and \sum_i p_i = z} ||p - x||
    """
    perm = np.argsort(x)[::-1]
    x = x[perm]
    inv_perm = np.zeros(len(x), dtype=np.int)
    inv_perm[perm] = np.arange(len(x))
    x[0] -= z
    dual_sol = isotonic_regression(x, increasing=False)
    x[0] += z
    primal_sol = x - dual_sol
    return primal_sol[inv_perm]


if __name__ == '__main__':
    rng = np.random.RandomState(None)
    x = rng.rand(10)
    p = projection_simplex(x)
    print(p)
    print(np.sum(p))