优化将n维向量投影到有效概率空间的函数

Question

我有一个python函数，它将概率列表投射到有效的概率空间中，每个概率的最小阈值th。例如，假设th=0.05，输入[0.3, 0.3, 0.4]返回相同的列表，而输入[1, 0, 1]返回[0.475, 0.05, 0.475]。我需要在TensorFlow中使用此代码，我正在尝试消除所有循环并尽可能地对其进行矢量化。

def projection(pi_in, th=0):
    pi = numpy.array(pi_in, dtype=np.float64)
    len_pi = len(pi_in)

    t = (1.0 - sum(pi)) / len_pi
    pi += t                   # sum(pi)=1, but values might be < th

    while True:
        n = 0
        excess = 0
        for i in range(len_pi):
            if pi[i] < th:    # fix values below th
                excess += th - pi[i]
                pi[i] = th
            elif pi[i] > th:  # count values above th
                n += 1

        if excess == 0:       # break when all values are valid
            break
        else:                 # otherwise decrement large values
            for i in range(len_pi):
                if pi[i] > th:
                    pi[i] -= excess / n

    return pi

这是代码的半优化版本。但是，我在处理代码和循环本身的条件时遇到了困难。如何优化此代码以使其高效（尽可能快），避免for循环和条件（工作向量），删除while循环（这可能是不可能的），并且可以轻松地移植到{ {1}}框架？