API reference¶

@dataclass(frozen=True)
class ActiveSetConfig:
    """Configuration of the active-set outer loop (:class:`ActiveSetSolver`).

    Bundles the outer-loop knobs into one argument; the inner solver and its
    tolerances live in :class:`nncg.inner.InnerSolver`, and the warm start stays
    a separate argument.

    Attributes:
        tol: Threshold of the primal and dual KKT violator tests.
        p_max: Patience budget — non-improving batch steps tolerated before a
            least-index Bland fallback pivot. Any value gives finite termination.
        track: Record the visited free-set trajectory in ``Result.traj``.
        max_outer: Optional cap on outer steps; when hit, the current iterate is
            returned with ``converged=False``.
    """

    tol: float = 1e-8
    p_max: int = 3
    track: bool = False
    max_outer: int | None = None

@dataclass(frozen=True)
class ActiveSetSolver:
    """The primal-dual active-set outer loop for the non-negative quadratic program.

    Holds the outer-loop :class:`ActiveSetConfig` and an
    :class:`nncg.inner.InnerSolver`, and drives the guarded block-pivot loop
    around the per-free-block solve the inner solver provides. It never touches a
    preconditioner — everything about CG/PCG/Nyström lives in ``inner``.

    Attributes:
        inner: The inner solver for each free block — e.g. :class:`nncg.inner.CG`
            (plain CG), :class:`nncg.inner.Jacobi`, :class:`nncg.inner.Nystrom`
            or :class:`nncg.inner.Exact`.
        config: Outer-loop configuration (violator tolerance, patience,
            trajectory tracking, outer-step cap).

    Examples:
        ``A`` enters as an operator, never as a bare array:

        >>> import numpy as np
        >>> from cvx.linalg import DenseOperator
        >>> from nncg import ActiveSetSolver, CG, kkt_violation
        >>> a = DenseOperator(np.array([[2.0, 0.0], [0.0, 2.0]]))
        >>> b = np.array([2.0, -2.0])

        The unconstrained minimiser would be ``(1, -1)``, so the bound binds on
        the second coordinate and the loop returns ``(1, 0)`` with that
        coordinate active:

        >>> res = ActiveSetSolver(inner=CG()).solve(a, b)
        >>> res.converged
        True
        >>> bool(np.allclose(res.x, [1.0, 0.0]))
        True
        >>> res.free.tolist()
        [True, False]

        ``converged`` is the KKT exit, which :func:`nncg.kkt_violation` scores
        independently — zero certifies the unique global minimiser:

        >>> round(kkt_violation(a, b, res.x), 12)
        0.0

        Generic data never needs the Bland fallback; that it stayed dormant is
        reported rather than assumed:

        >>> res.fallback
        0

        Swap the inner solver freely — the outer loop is unchanged, and on this
        problem so is the answer:

        >>> from nncg import Exact
        >>> direct = ActiveSetSolver(inner=Exact()).solve(a, b)
        >>> bool(np.allclose(direct.x, res.x))
        True
    """

    inner: InnerSolver
    config: ActiveSetConfig = field(default_factory=ActiveSetConfig)

    def solve(
        self,
        a: SymmetricOperator,
        b: Vector,
        warm: tuple[NDArray[np.bool_], Vector] | None = None,
    ) -> Result:
        """Minimise ``1/2 x^T A x - b^T x`` over ``x >= 0`` by the active-set loop.

        Each free-block solve is delegated to :attr:`inner`; the reduced matrix
        is never materialised and ``A`` is never refactorised. The batch
        block-pivot fast path is guarded by a least-index Bland fallback, so
        termination at the unique global minimiser is unconditional.

        Args:
            a: The SPD operator ``A`` (a :class:`cvx.linalg.SymmetricOperator`) —
                ``DenseOperator`` for an explicit array, ``GramOperator(M, ridge)``
                for ``A = M^T M + ridge I`` whose Gram matrix is never formed.
            b: The linear term ``b``.
            warm: Optional ``(free_mask, x_prev)`` pair from a previous solve.
                Starts the loop from that free set and warm-starts every inner
                solve from the newest iterate (the :class:`nncg.inner.Exact`
                inner solver is direct, so it has nothing to seed but still
                starts from the warm free set) — across a support-stable
                parameter step the loop then terminates in a single outer step.

        Returns:
            A :class:`Result`; ``converged`` is True iff the KKT system was
            satisfied to ``config.tol``, which certifies the unique global
            minimiser.

        Raises:
            TypeError: When ``a`` is not a :class:`cvx.linalg.SymmetricOperator`.
            ValueError: When the operator dimension does not match ``len(b)``, or
                on the inner solver's own conditions in
                :meth:`InnerSolver.solve`.
            NotImplementedError: When a diagonal-preconditioned inner solver
                (:class:`nncg.inner.Jacobi`) meets a backend without ``diag``
                (propagated from ``cvx.linalg``).
        """
        _require_operator(a, b)

        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."""
            xf, k_step = self.inner.solve(a, idx, b[idx], x0)
            return xf, None, k_step

        def reduced_gradient(x: Vector, lam: Vector | None) -> Vector:  # noqa: ARG001
            """Return the reduced gradient ``s = A x - b``."""
            return a.matvec(x) - b

        return self._run(len(b), sub_solve, reduced_gradient, warm)

    def solve_eq(
        self,
        a: SymmetricOperator,
        b: Vector,
        b_eq: Matrix,
        c_eq: Vector,
        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 inner solve, 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``).

        Args:
            a: The SPD operator ``A`` (a :class:`cvx.linalg.SymmetricOperator`).
            b: The linear term ``b``.
            b_eq: Equality matrix ``B`` of shape ``(p, n)``, full row rank.
            c_eq: Equality right-hand side ``c`` of shape ``(p,)``.
            warm: Optional ``(free_mask, x_prev)`` pair from a previous solve.
                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 (their right-hand sides are the rows of ``B_F``,
                unrelated to ``x_prev``). 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)``, or
                on the inner solver's own conditions in
                :meth:`InnerSolver.solve`.
            NotImplementedError: When a diagonal-preconditioned inner solver
                (:class:`nncg.inner.Jacobi`) meets a backend without ``diag``
                (propagated from ``cvx.linalg``).
        """
        _require_operator(a, b)

        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."""
            return _saddle_solve(self.inner, a, b, b_eq, c_eq, idx, x0)

        def reduced_gradient(x: Vector, lam: Vector | None) -> Vector:
            """Return the constrained reduced gradient ``s = A x - b - B^T lam``."""
            correction = b_eq.T @ lam if lam is not None else np.zeros_like(b)
            return a.matvec(x) - b - correction

        return self._run(len(b), sub_solve, reduced_gradient, warm)

    def _run(
        self,
        n: int,
        sub_solve: SubSolve,
        reduced_gradient: ReducedGradient,
        warm: tuple[NDArray[np.bool_], Vector] | None,
    ) -> Result:
        """Run the guarded primal-dual active-set loop.

        The driver owns everything the termination proof depends on: the primal
        and dual violator tests, the batch exchange with its patience counter,
        and the least-index Bland fallback. What is solved on each free set — a
        single reduced system (:meth:`solve`), or the equality-augmented saddle
        system (:meth:`solve_eq`) — enters through the ``sub_solve`` callback,
        with ``reduced_gradient`` supplying the matching dual test quantity. The
        thresholds (``tol``, ``p_max``, ``track``, ``max_outer``) are read from
        :attr:`config`.

        Args:
            n: Problem dimension.
            sub_solve: Callback ``(idx, x0) -> (x_F, lam, inner_iters)`` solving
                the subproblem on the free set ``idx``. ``x0`` is a warm inner
                guess restricted to ``idx`` (None on a cold start); ``lam`` are
                the equality multipliers (None for the bound-only problem).
            reduced_gradient: Callback ``(x, lam) -> s`` computing the reduced
                gradient that drives the dual violator test.
            warm: Optional ``(free_mask, x_prev)`` pair from a previous solve.
                Starts the loop from that free set and seeds every subproblem
                solve from the newest iterate.

        Returns:
            A :class:`Result`; ``lam`` is whatever the last subproblem returned.
        """
        cfg = self.config
        x, outer, inner_total, fallback, converged, free, lam, traj = _drive(
            cfg.tol, cfg.p_max, cfg.track, cfg.max_outer, n, sub_solve, reduced_gradient, warm
        )
        return Result(
            x=x,
            outer=outer,
            inner=inner_total,
            fallback=fallback,
            converged=converged,
            free=free,
            lam=lam,
            traj=traj,
        )

class InnerSolver(Protocol):
    """The inner-solver interface the active-set loop depends on (dependency inversion).

    Structural (a :class:`typing.Protocol`): anything with a matching
    :meth:`solve` is an inner solver, so implementations need neither import
    nor subclass this — this module (the high-level loop) owns the interface, and
    the implementations depend on it, not the other way round. The built-ins live
    in :mod:`nncg.inner` (:class:`~nncg.inner.CG`, :class:`~nncg.inner.Jacobi`,
    :class:`~nncg.inner.Nystrom`, :class:`~nncg.inner.Exact`); further ones —
    Clarabel- or KKT-equation-based — live in Jebel-Quant/mean_variance_solvers.
    """

    def solve(self, op: SymmetricOperator, idx: NDArray[np.int_], rhs: Vector, x0: Vector | None) -> tuple[Vector, int]:
        """Solve the free-block system ``A[F, F] y = rhs``, warm-started at ``x0``.

        Returns the free-block solution and the inner iteration count (each
        direct solve counts as one). Called once per outer step by the
        bound-constrained loop, and once per ``p + 1`` right-hand side per outer
        step by the equality-augmented loop.
        """
        ...

@dataclass(frozen=True)
class Result:
    """Outcome of an active-set solve.

    Attributes:
        x: The minimiser (or the final iterate if ``converged`` is False).
        outer: Number of outer active-set steps taken.
        inner: Total inner (CG/PCG) iterations across all outer steps; each
            direct inner solve counts as one.
        fallback: Number of least-index Bland fallback pivots taken.
        converged: True when the KKT exit was reached; False when an
            ``max_outer`` cap stopped the loop first.
        free: Boolean mask of the final free set.
        lam: Multipliers of the equality constraints (equality-augmented
            solves only; None otherwise).
        traj: The sequence of visited free sets as index tuples when
            trajectory tracking was requested; None otherwise.
    """

    x: Vector
    outer: int
    inner: int
    fallback: int
    converged: bool
    free: NDArray[np.bool_]
    lam: Vector | None = None
    traj: list[tuple[int, ...]] | None = None

@dataclass(frozen=True)
class MPRGP:
    """The MPRGP solver for the non-negative quadratic program.

    A matrix-free, factorisation-free alternative to
    :class:`nncg.solver.ActiveSetSolver` on the bound-constrained problem
    ``min_{x>=0} 1/2 x^T A x - b^T x``: it interleaves conjugate-gradient,
    expansion and proportioning steps under the proportioning test, never forming
    or factorising ``A``. Holds only its :class:`MPRGPConfig`; the operator and
    right-hand side are passed to :meth:`solve`.

    Attributes:
        config: Solver configuration (tolerance, proportioning constant,
            projected-gradient step, iteration cap, seed).

    Examples:
        The same operator interface as the active-set loop:

        >>> import numpy as np
        >>> from cvx.linalg import DenseOperator
        >>> from nncg import MPRGP, MPRGPConfig, kkt_violation
        >>> a = DenseOperator(np.array([[2.0, 0.0], [0.0, 2.0]]))
        >>> b = np.array([2.0, -2.0])
        >>> res = MPRGP().solve(a, b)
        >>> res.converged
        True
        >>> bool(np.allclose(res.x, [1.0, 0.0]))
        True
        >>> round(kkt_violation(a, b, res.x), 12)
        0.0

        The counts break the run down by move, and always sum to
        ``iterations``; ``hessian_products`` is the honest matrix-free cost —
        at least one product per step, plus the initial gradient:

        >>> res.iterations == res.cg_steps + res.expansion_steps + res.proportioning_steps
        True
        >>> res.hessian_products > res.iterations
        True

        Nothing is factorised, so the whole configuration is a handful of
        scalars — here a tighter tolerance and an explicit projected-gradient
        step, which skips the power-iteration estimate of ``1/||A||``:

        >>> tight = MPRGP(config=MPRGPConfig(tol=1e-12, alpha_bar=0.5)).solve(a, b)
        >>> tight.converged
        True
        >>> bool(np.allclose(tight.x, [1.0, 0.0]))
        True
    """

    config: MPRGPConfig = _DEFAULT_MPRGP

    def solve(self, a: SymmetricOperator, b: Vector, x0: Vector | None = None) -> MPRGPResult:
        """Minimise ``1/2 x^T A x - b^T x`` over ``x >= 0`` by MPRGP.

        Args:
            a: The SPD operator ``A`` (a :class:`cvx.linalg.SymmetricOperator`) —
                ``DenseOperator`` for an explicit array, ``GramOperator(M, ridge)``
                for ``A = M^T M + ridge I`` whose Gram matrix is never formed.
            b: The linear term ``b``.
            x0: Optional feasible warm start; it is projected onto ``x >= 0`` and
                the iteration begins there. ``None`` starts from the origin.

        Returns:
            An :class:`MPRGPResult`; ``converged`` is True iff the projected
            gradient fell below ``config.tol * ||b||``, which certifies the unique
            global minimiser.

        Raises:
            TypeError: When ``a`` is not a :class:`cvx.linalg.SymmetricOperator`.
            ValueError: When the operator dimension does not match ``len(b)``, or
                when ``config.alpha_bar`` is set but not strictly positive.
        """
        _require_operator(a, b)
        cfg = self.config
        alpha_bar = _resolve_alpha_bar(a, cfg.alpha_bar, cfg.seed)
        return _mprgp(a.matvec, b, x0, alpha_bar, cfg.gamma, cfg.tol, cfg.max_iter)

@dataclass(frozen=True)
class MPRGPConfig:
    """Configuration of the MPRGP solver (:class:`MPRGP`).

    Attributes:
        tol: Relative stopping tolerance on the projected gradient — the loop
            exits when ``||phi(x) + beta(x)|| <= tol * ||b||`` (``||b||`` replaced
            by ``1`` when ``b = 0``), which certifies the KKT conditions.
        gamma: Proportioning constant ``Gamma > 0`` balancing expansion against
            proportioning. ``1.0`` is the standard, near-optimal choice; larger
            values expand more eagerly, smaller ones release more eagerly.
        alpha_bar: Fixed projected-gradient step, which must satisfy
            ``0 < alpha_bar <= 2/||A||`` for convergence. ``None`` estimates the
            safe default ``1/||A||`` by matrix-free power iteration.
        max_iter: Iteration cap; the current iterate is returned with
            ``converged=False`` when it is hit.
        seed: Seed of the power-iteration ``||A||`` estimate (only used when
            ``alpha_bar is None``), fixed so a solve is reproducible.

    Raises:
        ValueError: When ``gamma`` is not strictly positive.
    """

    tol: float = 1e-8
    gamma: float = 1.0
    alpha_bar: float | None = None
    max_iter: int = 100_000
    seed: int = 0

    def __post_init__(self) -> None:
        """Validate that the proportioning constant is strictly positive."""
        if self.gamma <= 0.0:
            msg = f"gamma must be strictly positive; got {self.gamma:.2e}"
            raise ValueError(msg)

@dataclass(frozen=True)
class MPRGPResult:
    """Outcome of an MPRGP solve.

    The iteration counts are broken out by move because they carry the algorithm's
    signature: expansion and proportioning steps are the ones that change the
    active set, while a run of conjugate-gradient steps is plain CG on a fixed
    face. ``hessian_products`` is the honest cost of a matrix-free method — one
    product per CG or proportioning step, two per expansion step, plus one for the
    initial gradient.

    Attributes:
        x: The minimiser (or the final iterate if ``converged`` is False).
        iterations: Total MPRGP steps taken (the sum of the three move counts).
        hessian_products: Number of operator matrix-vector products consumed.
        cg_steps: Conjugate-gradient (minimisation-within-the-face) steps.
        expansion_steps: Expansion (bound-hitting projected-gradient) steps.
        proportioning_steps: Proportioning (constraint-releasing) steps.
        converged: True when the projected-gradient stopping test was met; False
            when ``max_iter`` stopped the loop first.
        free: Boolean mask of the final free set (``x > 0``).
    """

    x: Vector
    iterations: int
    hessian_products: int
    cg_steps: int
    expansion_steps: int
    proportioning_steps: int
    converged: bool
    free: NDArray[np.bool_]