Support inner=pcg/exact in solve_nnqp_eq

src/nncg/solver.py

@@ -109,6 +109,70 @@ def _free_matvec(op: SymmetricOperator, idx: NDArray[np.int_]) -> MatVec:
     return lambda v: op.apply_free(idx, v)
 
 
+#: A single free-block solve ``(idx, rhs, x0) -> (y, iters)`` for ``A[F, F] y = rhs``.
+FreeSolve = Callable[[NDArray[np.int_], Vector, "Vector | None"], "tuple[Vector, int]"]
+
+
+def _make_free_solve(
+    op: SymmetricOperator,
+    inner: InnerSolver,
+    cg_tol: float,
+    cg_maxit: int,
+) -> FreeSolve:
+    """Build the per-free-block solver ``A[F, F] y = rhs`` for the chosen backend.
+
+    The returned callable is shared by both entry points: :func:`solve_nnqp`
+    calls it once per outer step, :func:`solve_nnqp_eq` once for the ``v0``
+    right-hand side and once per Schur-complement column. All calls on a given
+    free set share the same operator ``A_F``, so the Jacobi preconditioner is
+    read off ``op.diag`` lazily and cached across them. The ``x0`` warm start is
+    honoured only by plain CG; ``pcg`` and ``exact`` ignore it.
+
+    Args:
+        op: The SPD operator ``A``.
+        inner: Inner solver — ``"cg"``, ``"pcg"``, or ``"exact"``.
+        cg_tol: Relative residual tolerance of the CG/PCG solves.
+        cg_maxit: Iteration cap per CG/PCG solve.
+
+    Returns:
+        A callable ``(idx, rhs, x0) -> (y, iters)``; each ``"exact"`` direct
+        solve counts as one iteration.
+
+    Raises:
+        ValueError: When ``inner`` is not one of ``"cg"``, ``"pcg"``,
+            ``"exact"`` (raised eagerly), or, on the ``"exact"`` path, when a
+            free block is numerically singular (``op.rcond_free`` below 1e-12).
+        NotImplementedError: When ``inner="pcg"`` meets a backend without
+            ``diag`` (propagated from ``cvx.linalg``).
+    """
+    if inner not in ("cg", "pcg", "exact"):
+        msg = f"inner must be 'cg', 'pcg', or 'exact'; got {inner!r}"
+        raise ValueError(msg)
+    dinv: Vector | None = None  # Jacobi preconditioner, read off op.diag on first use
+    checked: NDArray[np.int_] | None = None  # last free set whose SPD-ness was verified
+
+    def free_solve(idx: NDArray[np.int_], rhs: Vector, x0: Vector | None) -> tuple[Vector, int]:
+        nonlocal dinv, checked
+        if inner == "exact":
+            # The rcond guard depends only on the free set, not the right-hand
+            # side, so verify it once per free set — solve_nnqp_eq drives p + 1
+            # solves through the same idx and must not pay for p + 1 estimates.
+            if checked is None or not np.array_equal(checked, idx):
+                rcond = op.rcond_free(idx)
+                if rcond < _RCOND_MIN:
+                    msg = f"free block of size {idx.size} is numerically singular (rcond={rcond:.2e})"
+                    raise ValueError(msg)
+                checked = idx
+            return op.solve_free(idx, rhs), 1
+        if inner == "pcg":
+            if dinv is None:
+                dinv = 1.0 / op.diag
+            return pcg(_free_matvec(op, idx), rhs, dinv[idx], tol=cg_tol, maxit=cg_maxit)
+        return cg(_free_matvec(op, idx), rhs, tol=cg_tol, maxit=cg_maxit, x0=x0)
+
+    return free_solve
+
+
 @dataclass(frozen=True)
 class Result:
     """Outcome of an active-set solve.
@@ -330,26 +394,11 @@ def solve_nnqp(
     """
     op = _require_operator(a)
     _check_dimension(op, b)
-    dinv: Vector | None = None  # Jacobi preconditioner, read off op.diag on first use
+    free_solve = _make_free_solve(op, inner, cg_tol, cg_maxit)
 
     def sub_solve(idx: NDArray[np.int_], x0: Vector | None) -> tuple[Vector, Vector | None, int]:
         """Solve the reduced system ``A_F x_F = b_F`` with the chosen inner solver."""
-        nonlocal dinv
-        if inner == "exact":
-            rcond = op.rcond_free(idx)
-            if rcond < _RCOND_MIN:
-                msg = f"free block of size {idx.size} is numerically singular (rcond={rcond:.2e})"
-                raise ValueError(msg)
-            return op.solve_free(idx, b[idx]), None, 1
-        if inner == "pcg":
-            if dinv is None:
-                dinv = 1.0 / op.diag
-            xf, k_step = pcg(_free_matvec(op, idx), b[idx], dinv[idx], tol=cg_tol, maxit=cg_maxit)
-            return xf, None, k_step
-        if inner != "cg":
-            msg = f"inner must be 'cg', 'pcg', or 'exact'; got {inner!r}"
-            raise ValueError(msg)
-        xf, k_step = cg(_free_matvec(op, idx), b[idx], tol=cg_tol, maxit=cg_maxit, x0=x0)
+        xf, k_step = free_solve(idx, b[idx], x0)
         return xf, None, k_step
 
     def reduced_gradient(x: Vector, lam: Vector | None) -> Vector:  # noqa: ARG001
@@ -376,14 +425,18 @@ def solve_nnqp_eq(
     tol: float = 1e-8,
     cg_tol: float = 1e-10,
     p_max: int = 3,
+    inner: InnerSolver = "cg",
+    track: bool = False,
+    cg_maxit: int = 100_000,
+    max_outer: int | None = None,
     warm: tuple[NDArray[np.bool_], Vector] | None = None,
 ) -> Result:
     """Solve ``min 1/2 x^T A x - b^T x`` subject to ``x >= 0`` and ``B x = c``.
 
     On each free set the saddle system is solved by eliminating the multiplier
     ``lambda`` in R^p through the p-by-p Schur complement
     ``S = B_F A_F^{-1} B_F^T``: the ``p + 1`` right-hand sides share the
-    operator ``A_F`` and are each one matrix-free CG solve, then
+    operator ``A_F`` and are each one inner solve (see ``inner``), then
     ``S lambda = c - B_F v0`` fixes the multipliers in closed form. The single
     normalisation ``1^T x = beta`` is the ``p = 1`` case. ``B`` must have full
     row rank on the visited free sets (automatic for ``p = 1``).
@@ -394,36 +447,49 @@ def solve_nnqp_eq(
         b_eq: Equality matrix ``B`` of shape ``(p, n)``, full row rank.
         c_eq: Equality right-hand side ``c`` of shape ``(p,)``.
         tol: Threshold of the primal and dual violator tests.
-        cg_tol: Relative residual tolerance of the inner CG solves.
+        cg_tol: Relative residual tolerance of the inner CG/PCG solves.
         p_max: Patience budget of the batch fast path.
+        inner: Inner solver for each free block, applied to all ``p + 1``
+            right-hand sides — ``"cg"`` (matrix-free), ``"pcg"`` (Jacobi from
+            ``op.diag``), or ``"exact"`` (direct ``op.solve_free``); the same
+            choice :func:`solve_nnqp` offers.
+        track: Record the free-set trajectory in ``Result.traj``.
+        cg_maxit: Iteration cap per inner CG/PCG solve.
+        max_outer: Optional cap on outer steps; when hit, the current iterate
+            is returned with ``converged=False``.
         warm: Optional ``(free_mask, x_prev)`` pair from a previous solve —
             the same tuple :func:`solve_nnqp` accepts. Starts the loop from
             that free set and seeds the ``v0`` solve of every saddle step
             from the newest iterate; the ``v1`` columns are re-solved cold.
-            Across a support-stable parameter step the loop then terminates
-            in a single outer step.
+            Only ``inner="cg"`` consumes the warm start; ``pcg``/``exact``
+            ignore it. Across a support-stable parameter step the loop then
+            terminates in a single outer step.
 
     Returns:
         A :class:`Result` with the multipliers in ``lam``. The reduced
         gradient underlying the dual test is ``s = A x - b - B^T lam``.
 
     Raises:
         TypeError: When ``a`` is not a :class:`cvx.linalg.SymmetricOperator`.
-        ValueError: When the operator dimension does not match ``len(b)``.
+        ValueError: When the operator dimension does not match ``len(b)``;
+            when ``inner`` is not one of ``"cg"``, ``"pcg"``, ``"exact"``; or
+            when ``inner="exact"`` meets a numerically singular free block.
+        NotImplementedError: When ``inner="pcg"`` meets a backend that does
+            not expose ``diag`` (propagated from ``cvx.linalg``).
     """
     op = _require_operator(a)
     _check_dimension(op, b)
     p = b_eq.shape[0]
+    free_solve = _make_free_solve(op, inner, cg_tol, cg_maxit)
 
     def sub_solve(idx: NDArray[np.int_], x0: Vector | None) -> tuple[Vector, Vector | None, int]:
         """Solve the saddle system on the free set via the p-by-p Schur complement."""
-        matvec_f = _free_matvec(op, idx)
         b_f = b_eq[:, idx]
-        v0, k0 = cg(matvec_f, b[idx], tol=cg_tol, x0=x0)
+        v0, k0 = free_solve(idx, b[idx], x0)
         v1 = np.zeros((idx.size, p))
         k_cols = 0
         for j in range(p):
-            v1[:, j], kj = cg(matvec_f, b_f[j], tol=cg_tol)
+            v1[:, j], kj = free_solve(idx, b_f[j], None)
             k_cols += kj
         schur = b_f @ v1  # p-by-p Schur complement, SPD
         lam = cholesky_solve(schur, c_eq - b_f @ v0)
@@ -435,4 +501,13 @@ def reduced_gradient(x: Vector, lam: Vector | None) -> Vector:
         correction = b_eq.T @ lam if lam is not None else np.zeros_like(b)
         return op.matvec(x) - b - correction
 
-    return _active_set_loop(len(b), sub_solve, reduced_gradient, tol=tol, p_max=p_max, warm=warm)
+    return _active_set_loop(
+        len(b),
+        sub_solve,
+        reduced_gradient,
+        tol=tol,
+        p_max=p_max,
+        track=track,
+        max_outer=max_outer,
+        warm=warm,
+    )

tests/problems.py

@@ -132,6 +132,55 @@ def make_adversarial(
     return a, m.T @ rng.standard_normal(n)
 
 
+def make_scaled_eq_problem(
+    n: int,
+    kappa_core: float,
+    spread: float,
+    p: int,
+    support_frac: float = 0.5,
+    seed: int = 0,
+) -> tuple[Matrix, Vector, Matrix, Vector, Vector, Vector, Vector]:
+    """Equality-augmented planted problem under a bad diagonal scaling.
+
+    Combines :func:`make_eq_problem`'s planted KKT triple with
+    :func:`make_scaled_problem`'s ill-conditioned-by-scaling ``A``: a
+    well-conditioned core (condition number ``kappa_core``) wrapped in a
+    diagonal scaling ``D`` spanning ``[1, spread]``. Jacobi preconditioning
+    removes the scaling, so PCG inner solves run at the core's conditioning
+    while plain CG pays for the full spread — the equality-constrained analogue
+    of the ``make_scaled_problem`` preconditioning win.
+
+    Args:
+        n: Problem dimension.
+        kappa_core: Condition number of the unscaled core.
+        spread: Ratio between the largest and smallest diagonal scale.
+        p: Number of equality constraints (rows of ``B``).
+        support_frac: Fraction of indices in the optimal support.
+        seed: Seed of the random generator.
+
+    Returns:
+        The tuple ``(A, b, B, c, x_star, lam_star, s_star)``.
+    """
+    rng = np.random.default_rng(seed)
+    eig = np.geomspace(1.0, kappa_core, n)
+    q, _ = np.linalg.qr(rng.standard_normal((n, n)))
+    core = 0.5 * ((q * eig) @ q.T + ((q * eig) @ q.T).T)
+    d = rng.permutation(np.geomspace(1.0, spread, n))
+    a = core * np.sqrt(np.outer(d, d))
+    k = max(p + 1, round(support_frac * n))
+    perm = rng.permutation(n)
+    supp, off = perm[:k], perm[k:]
+    x_star = np.zeros(n)
+    x_star[supp] = rng.uniform(0.5, 1.5, size=k)
+    s_star = np.zeros(n)
+    s_star[off] = rng.uniform(0.5, 1.5, size=n - k)
+    b_eq = rng.standard_normal((p, n))
+    lam_star = rng.standard_normal(p)
+    b = a @ x_star - b_eq.T @ lam_star - s_star
+    c_eq = b_eq @ x_star
+    return a, b, b_eq, c_eq, x_star, lam_star, s_star
+
+
 def make_scaled_problem(
     n: int,
     kappa_core: float,

tests/test_eq.py

@@ -5,7 +5,7 @@
 from cvx.linalg import DenseOperator
 
 from nncg import solve_nnqp_eq
-from tests.problems import make_eq_problem
+from tests.problems import make_eq_problem, make_scaled_eq_problem
 
 
 @pytest.mark.parametrize("p", [1, 3, 8])
@@ -70,6 +70,71 @@ def test_eq_warm_start_survives_support_drift() -> None:
     assert np.linalg.norm(b_eq @ warm.x - c_eq) < 1e-9
 
 
+def test_eq_exact_inner_matches_cg() -> None:
+    """The direct (``inner="exact"``) eq solve matches the CG solve.
+
+    Both recover the planted optimum and settle on the same free set — the
+    equality analogue of the inexactness lemma: the inner accuracy does not
+    change the sign decisions of the outer loop.
+    """
+    a, b, b_eq, c_eq, x_star, _, _ = make_eq_problem(80, 1e4, 3, seed=5)
+    op = DenseOperator(a)
+    r_cg = solve_nnqp_eq(op, b, b_eq, c_eq)
+    r_ex = solve_nnqp_eq(op, b, b_eq, c_eq, inner="exact")
+    assert r_ex.converged
+    assert np.array_equal(r_ex.free, r_cg.free)
+    assert np.max(np.abs(r_ex.x - x_star)) < 1e-6
+    assert np.linalg.norm(b_eq @ r_ex.x - c_eq) < 1e-9
+
+
+def test_eq_pcg_inner_recovers_optimum() -> None:
+    """The Jacobi-preconditioned eq solve recovers the same planted optimum."""
+    a, b, b_eq, c_eq, x_star, _, _ = make_eq_problem(80, 1e4, 3, seed=6)
+    r_pcg = solve_nnqp_eq(DenseOperator(a), b, b_eq, c_eq, inner="pcg")
+    assert r_pcg.converged
+    assert np.max(np.abs(r_pcg.x - x_star)) < 1e-6
+    assert np.linalg.norm(b_eq @ r_pcg.x - c_eq) < 1e-9
+
+
+def test_eq_pcg_beats_cg_under_diagonal_scaling() -> None:
+    """On a diagonally ill-scaled eq problem PCG needs fewer inner iterations."""
+    a, b, b_eq, c_eq, x_star, _, _ = make_scaled_eq_problem(80, 1e2, 1e6, 3, seed=7)
+    op = DenseOperator(a)
+    r_cg = solve_nnqp_eq(op, b, b_eq, c_eq, inner="cg")
+    r_pcg = solve_nnqp_eq(op, b, b_eq, c_eq, inner="pcg")
+    assert r_pcg.converged
+    assert np.max(np.abs(r_pcg.x - x_star)) < 1e-6
+    assert np.max(np.abs(r_cg.x - x_star)) < 1e-6
+    # A comfortable margin, not a bare `<`: the 1e6 diagonal spread makes the
+    # Jacobi win large enough to survive BLAS-dependent iteration-count drift.
+    assert r_pcg.inner <= 0.7 * r_cg.inner
+
+
+def test_eq_track_records_trajectory() -> None:
+    """``track=True`` records the visited free sets, ending on the converged one."""
+    a, b, b_eq, c_eq, _, _, _ = make_eq_problem(60, 1e3, 3, seed=8)
+    res = solve_nnqp_eq(DenseOperator(a), b, b_eq, c_eq, track=True)
+    assert res.converged
+    assert res.traj is not None
+    assert len(res.traj) == res.outer  # one recorded free set per outer step
+    assert np.array_equal(res.traj[-1], np.flatnonzero(res.free))  # last = converged support
+
+
+def test_eq_max_outer_caps_the_loop() -> None:
+    """``max_outer`` stops the loop early and reports non-convergence."""
+    a, b, b_eq, c_eq, _, _, _ = make_eq_problem(80, 1e5, 3, seed=9)
+    res = solve_nnqp_eq(DenseOperator(a), b, b_eq, c_eq, max_outer=1)
+    assert not res.converged
+    assert res.outer == 1
+
+
+def test_eq_rejects_unknown_inner() -> None:
+    """An unrecognised inner solver is rejected up front, not run as CG."""
+    a, b, b_eq, c_eq, _, _, _ = make_eq_problem(20, 1e2, 1, seed=0)
+    with pytest.raises(ValueError, match="inner must be"):
+        solve_nnqp_eq(DenseOperator(a), b, b_eq, c_eq, inner="nope")  # type: ignore[arg-type]
+
+
 def test_eq_reduced_gradient_certifies() -> None:
     """At the exit, x >= 0 and s = Ax - b - B^T lam >= 0 hold to tolerance.
 

tests/test_solver.py

@@ -53,7 +53,9 @@ def test_pcg_inner_solve() -> None:
     r_pcg = solve_nnqp(DenseOperator(a), b, inner="pcg")
     assert np.max(np.abs(r_cg.x - x_star)) < 1e-6
     assert np.max(np.abs(r_pcg.x - x_star)) < 1e-6
-    assert r_pcg.inner < r_cg.inner
+    # A comfortable margin, not a bare `<`: the exact counts vary with the BLAS
+    # backend, but Jacobi removes a 1e4 diagonal spread, so the win is large.
+    assert r_pcg.inner <= 0.7 * r_cg.inner
 
 
 def test_max_outer_cap_reports_nonconvergence() -> None: