Support inner=pcg/exact in solve_nnqp_eq - #17
Merged
Merged
Conversation
Extend the equality-augmented solver with the same inner-solver choice solve_nnqp already offers. Both entry points now share a single free-block solver factory (_make_free_solve), which dispatches cg/pcg/exact, caches the Jacobi preconditioner across a free set, and guards the exact path's rcond check — removing the duplicated dispatch that had drifted between the two. solve_nnqp_eq gains `inner` and `cg_maxit`; all p+1 Schur right-hand sides go through the chosen backend (exact via solve_free replaces all p+1 CG calls; pcg preconditions each once). Warm start stays cg-only, as before. Tests: exact matches the cg trajectory, pcg recovers the optimum, pcg beats cg on a new diagonally ill-scaled eq generator (make_scaled_eq_problem), and an unknown inner is rejected. Coverage 100%. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends the equality-constrained solver solve_nnqp_eq to support the same inner-solver choices as solve_nnqp ("cg", "pcg", "exact"), factoring shared free-block solve logic into a common factory to avoid duplicated dispatch.
Changes:
- Add
innerandcg_maxitparameters tosolve_nnqp_eq, enabling"cg"/"pcg"/"exact"backends. - Introduce
_make_free_solve(...)to centralize inner-solver dispatch and (for PCG) reuse a Jacobi preconditioner across solves. - Add new equality-path tests and a new scaled equality problem generator to validate
"exact"/"pcg"behavior and error handling.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
src/nncg/solver.py |
Adds _make_free_solve and wires solve_nnqp_eq to accept/dispatch inner backends. |
tests/test_eq.py |
Adds tests covering inner="exact", inner="pcg", iteration comparisons, and unknown-inner rejection. |
tests/problems.py |
Adds make_scaled_eq_problem to generate diagonally ill-scaled equality problems for preconditioning tests. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+151
to
+165
| dinv: Vector | None = None # Jacobi preconditioner, read off op.diag on first use | ||
|
|
||
| def free_solve(idx: NDArray[np.int_], rhs: Vector, x0: Vector | None) -> tuple[Vector, int]: | ||
| 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, 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) |
Comment on lines
421
to
425
| p_max: int = 3, | ||
| inner: InnerSolver = "cg", | ||
| cg_maxit: int = 100_000, | ||
| warm: tuple[NDArray[np.bool_], Vector] | None = None, | ||
| ) -> Result: |
- #18: memoise the exact-path rcond_free guard per free set in _make_free_solve, so solve_nnqp_eq's p+1 same-idx solves pay for one conditioning estimate, not p+1. - #19: replace the bare `r_pcg.inner < r_cg.inner` assertions with a 0.7x margin on large diagonal spreads, robust to BLAS-dependent iteration drift. - #20: forward track and max_outer from solve_nnqp_eq to _active_set_loop, matching solve_nnqp; add trajectory-recording and outer-cap tests. Closes #18 Closes #19 Closes #20 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-up to #16 (now merged):
solve_nnqp_eqgains the same inner-solver choicesolve_nnqpoffers.What
solve_nnqp_eqnow takesinner("cg"/"pcg"/"exact") andcg_maxit. Previously it hardcoded plain CG.Both entry points now share one free-block solver factory
_make_free_solve(op, inner, cg_tol, cg_maxit) -> (idx, rhs, x0) -> (y, iters), which dispatchescg/pcg/exact, lazily caches the Jacobi preconditioner across a free set, guards theexactpath'srcond_freecheck, and raisesValueErroreagerly on an unknowninner. This removes the inner-solver dispatch that had been duplicated (and had drifted) between the two solvers.Behaviour on the equality path
All
p + 1Schur right-hand sides (v0+pcolumns) go through the chosen backend:exact—op.solve_freereplaces allp+1CG calls (natural, since they shareA_F).pcg— one Jacobi preconditioner offop.diag, reused across allp+1RHS.cg-only (as insolve_nnqp).Tests (coverage 100%)
test_eq_exact_inner_matches_cg— exact recovers the optimum and visits the same free set as CG (equality analogue of the inexactness lemma).test_eq_pcg_inner_recovers_optimum— PCG recovers the planted optimum, feasible.test_eq_pcg_beats_cg_under_diagonal_scaling— on a newmake_scaled_eq_problemgenerator (planted eq optimum under a bad diagonal scaling), PCG needs fewer inner iterations than CG.test_eq_rejects_unknown_inner— badinnerraises.Gates
make typecheck(ty + mypy strict),make test(61 passed, 100% coverage),make fmt,make deptry— all green.🤖 Generated with Claude Code