feat: preconditioning tools (precond.py) with Jacobi and randomized Nyström

[tschm]

What

Adds src/nncg/precond.py next to krylov.py and solver.py, collecting the package's preconditioning tools, and generalises the CG inner solver to a callable preconditioner M⁻¹: v → v.

Why

Preconditioning was previously a single inline dinv = 1.0 / op.diag in solver.py. This factors it out and opens the door to non-diagonal preconditioners — specifically a randomized Nyström / Woodbury low-rank preconditioner, which cannot be expressed as a dinv vector and needs the callable form.

Changes

precond.py (new) — preconditioners as callables r → M⁻¹r:

• inverse_diagonal(op) — 1/diag(A) off op.diag, with an SPD-positivity guard (rejects non-positive/non-finite diagonals instead of leaking an inf).
• identity(), diagonal(dinv), jacobi(op).
• nystrom(matvec, n, rank, oversample=10, shift=None, seed=None) — randomized Nyström sketch A ≈ U diag(λ) Uᵀ applied as a scalar shift + low-rank Woodbury correction in O(nr), built once and captured in the closure. Default shift is the largest uncaptured eigenvalue (the deflation choice — measured 1.8–2.4× vs the smallest-captured alternative), overridable.

krylov.py

• pcg takes a callable preconditioner (None → identity), applying z = precond(r); added the Preconditioner type alias.
• cg collapses to a thin wrapper over the single pcg loop — plain CG is identity-PCG, removing the duplicated loop.

solver.py

• inner="pcg" passes diagonal(dinv[F]); behavior unchanged.
• New inner="nystrom" mode with nystrom_rank/nystrom_oversample/nystrom_shift/nystrom_seed knobs (seed fixed by default for reproducibility). Builds a fresh sketch per free set, cached across the p+1 right-hand sides of an equality solve.

Cost, stated plainly

Unlike Jacobi (a cached diag(A)[F] slice), the Nyström sketch rebuilds per outer step — each free block A_F differs — costing rank + oversample matrix-free products per free set on top of CG iterations. Worth it only when each free block's spectrum decays steeply. Documented in _make_free_solve and both public docstrings.

Tests

New tests/test_precond.py (guards, identity == cg, Jacobi, Nyström correctness / win / full-rank / error cases), plus inner="nystrom" tests for solve_nnqp and solve_nnqp_eq (optimum recovery, KKT certificate, same free-set trajectory as CG per the inexactness lemma, eager rank guard).

All gates green: 85 tests pass, ty + mypy strict clean, deptry clean, fmt clean. No new runtime dependencies.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR introduces a dedicated preconditioning module and updates the Krylov inner solvers to accept callable preconditioners (r -> M⁻¹r), enabling non-diagonal options such as a randomized Nyström/Woodbury preconditioner. It then wires these options into the NNQP active-set solver and adds targeted tests for correctness and trajectory invariance.

Changes:

• Added src/nncg/precond.py with Jacobi/diagonal utilities plus a randomized Nyström preconditioner builder.
• Generalized pcg to accept a callable preconditioner and refactored cg to delegate to pcg (identity preconditioner).
• Extended solve_nnqp / solve_nnqp_eq with inner="nystrom" and added tests validating recovery, guards, and free-set trajectory matching.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
src/nncg/precond.py	New preconditioner builders (identity/diagonal/Jacobi/Nyström) as callables for PCG.
src/nncg/krylov.py	PCG now takes a callable preconditioner; CG is a wrapper over PCG with identity.
src/nncg/solver.py	Adds inner="nystrom" mode and threads Nyström parameters into free-set solves with per-free-set caching for eq solves.
tests/test_precond.py	New unit tests for preconditioner builders, guards, and Nyström behavior.
tests/test_solver.py	Adds NNQP tests for inner="nystrom" correctness and trajectory invariance.
tests/test_eq.py	Adds equality-constrained solve test covering Nyström path and cache behavior.
tests/test_krylov.py	Updates PCG tests to use diagonal(...) callable preconditioner rather than a raw vector.

Comments suppressed due to low confidence (1)

src/nncg/krylov.py:116

• pcg() applies the preconditioner before checking for a zero RHS or an already-converged warm start. This does unnecessary work (and can call an expensive preconditioner even when the solve should return immediately). Consider moving the bnorm/initial residual checks before the first preconditioner application.

    z = precond(r)
    p = z.copy()
    rz = float(r @ z)
    bnorm = float(np.linalg.norm(rhs))
    if bnorm == 0.0:

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.