1import numpy as np
  2
  3
  4def cg(matvec, rhs, tol=1e-10, max_iter=None):
  5    """Conjugate gradients, returning x, iterations, and final relative residual."""
  6    rhs = np.asarray(rhs, float)
  7    n = len(rhs)
  8    if max_iter is None:
  9        max_iter = max(4 * n, 1000)
 10    x = np.zeros(n)
 11    r = rhs - matvec(x)
 12    p = r.copy()
 13    rr = float(r @ r)
 14    r0 = max(np.sqrt(rr), 1e-30)
 15    if r0 == 0:
 16        return x, 0, 0.0
 17    for k in range(1, max_iter + 1):
 18        Ap = matvec(p)
 19        den = float(p @ Ap)
 20        if den <= 0 or not np.isfinite(den):
 21            raise RuntimeError("CG encountered a non-SPD operator")
 22        alpha = rr / den
 23        x += alpha * p
 24        r -= alpha * Ap
 25        nr = np.linalg.norm(r)
 26        if nr <= tol * r0:
 27            return x, k, nr / r0
 28        rr_new = float(r @ r)
 29        p = r + (rr_new / rr) * p
 30        rr = rr_new
 31    return x, max_iter, np.linalg.norm(r) / r0
 32
 33
 34def active_set_simplex(Z, b, lam=1e-3, tol=1e-10, pivot_tol=1e-9,
 35                       max_pivots=1000, cg_tol=1e-11):
 36    """Solve min .5*x'(Z'Z+lam I)x-b'x, x>=0, sum(x)=1.
 37
 38    Products use Z and Z.T only. The equality constrained free solve uses two
 39    CG solves (A_F y=b_F and A_F q=1), followed by a scalar Schur correction.
 40    """
 41    Z = np.asarray(Z, float); b = np.asarray(b, float)
 42    n, E = Z.shape
 43    F = list(range(E)); pivots = 0; cg_iters = 0
 44    x = np.ones(E) / E
 45    last_nu = 0.0
 46    def Avec(v, inds):
 47        return Z[:, inds].T @ (Z[:, inds] @ v) + lam * v
 48    while pivots <= max_pivots:
 49        inds = np.array(F, dtype=int)
 50        Af = lambda v: Avec(v, inds)
 51        y, it1, _ = cg(Af, b[inds], tol=cg_tol, max_iter=max(20, 8*len(F)))
 52        q, it2, _ = cg(Af, np.ones(len(F)), tol=cg_tol, max_iter=max(20, 8*len(F)))
 53        cg_iters += it1 + it2
 54        denom = float(np.ones(len(F)) @ q)
 55        nu = (float(np.ones(len(F)) @ y) - 1.0) / denom
 56        xf = y - nu * q
 57        x[:] = 0.0; x[inds] = xf; last_nu = nu
 58        bad = [(v, i) for v, i in zip(xf, inds) if v < -pivot_tol]
 59        if bad:
 60            # Bland tie-break after most-negative pivot.
 61            m = min(v for v, _ in bad)
 62            i = min(i for v, i in bad if v <= m + 1e-14)
 63            F.remove(i); pivots += 1
 64            if not F: raise RuntimeError("empty free set")
 65            continue
 66        g = Z.T @ (Z @ x) + lam*x - b + last_nu
 67        blocked = [i for i in range(E) if i not in F and g[i] < -pivot_tol]
 68        if blocked:
 69            m = min(g[i] for i in blocked)
 70            i = min(i for i in blocked if g[i] <= m + 1e-14)
 71            F.append(i); F.sort(); pivots += 1
 72            continue
 73        # Numerical cleanup is only after all sign decisions.
 74        x[np.abs(x) < tol] = 0.0
 75        return x, {"pivots": pivots, "cg_matvecs": cg_iters, "free": len(F),
 76                   "nu": last_nu, "gradient": g}
 77    raise RuntimeError("active-set pivot limit reached")
 78
 79
 80def projected_gradient(Z, b, lam=1e-3, steps=5000, tol=1e-12):
 81    """Reference projected-gradient solve, with simplex projection."""
 82    E = Z.shape[1]
 83    L = np.linalg.eigvalsh(Z.T @ Z + lam*np.eye(E))[-1]
 84    x = np.ones(E)/E
 85    for k in range(steps):
 86        grad = Z.T @ (Z @ x) + lam*x - b
 87        xn = simplex_project(x - grad/L)
 88        if np.linalg.norm(xn-x) < tol:
 89            return xn, k+1
 90        x = xn
 91    return x, steps
 92
 93
 94def simplex_project(v):
 95    u = np.sort(v)[::-1]
 96    cssv = np.cumsum(u) - 1
 97    ind = np.arange(1, len(v)+1)
 98    rho = np.nonzero(u - cssv/ind > 0)[0]
 99    rho = rho[-1] if len(rho) else 0
100    theta = cssv[rho]/(rho+1)
101    return np.maximum(v-theta, 0.0)
102
103
104def objective(Z, b, x, lam):
105    return 0.5*(np.linalg.norm(Z@x)**2 + lam*np.dot(x,x)) - np.dot(b,x)