Skip to content

Solve interior with exact-JVP Newton-Krylov#66

Open
krystophny wants to merge 38 commits into
developfrom
ad-29-newton-interior-solve
Open

Solve interior with exact-JVP Newton-Krylov#66
krystophny wants to merge 38 commits into
developfrom
ad-29-newton-interior-solve

Conversation

@krystophny

Copy link
Copy Markdown
Member

Replaces the finite-difference-Jacobian hybr interior re-solve in the SIMSOPT wrappers with a damped Newton-Krylov solver.

  • Matvecs come from State.raw_residual_jvp (exact on Enzyme builds, centered FD on default builds), preconditioned by State.apply_preconditioner; no dense O(n^2) FD Jacobian remains anywhere on the solve path.
  • The equilibrium map now drives the interior raw residual to force balance instead of preserving the fixture's residual target, matching VMEC++'s solve_interior semantics: the energy becomes stationary with respect to interior solve error, which makes finite-difference cross-checks of the boundary gradient well conditioned.
  • Backtracking line search with stall-accept (VMEC++ pattern) plus an acceptance ceiling: a stalled solve above accept_tol * |F0| raises and triggers the penalty contract.
  • Counters residual_evaluations and jvp_evaluations make solver cost measurable.

Default builds keep working: the same Newton-Krylov runs with FD matvecs and the public API is unchanged.

Verification

Before

GvecEnergyProblem.solve_equilibrium called scipy.optimize.root(method="hybr"), whose MINPACK Jacobian costs O(n^2) FD residual evaluations per solve (~20 s per solve on the converted solovev case) and fails on stiff cases.

After

Enzyme-on isolated build (fixed RZ, LLVM/Flang 22.1.6):

$ pytest test-CI/unit-pygvec/test_ad_simsopt.py test-CI/unit-pygvec/test_ad_objectives.py -q --require-pygvec
16 passed in 2.66s

Enzyme-off isolated build:

$ pytest test-CI/unit-pygvec/test_ad_simsopt.py test-CI/unit-pygvec/test_ad_hvp.py test-CI/unit-pygvec/test_ad_boundary.py test-CI/unit-pygvec/test_ad_external_optimizers.py -q --require-pygvec
27 passed, 2 skipped in 25.92s

test_newton_interior_solve_reaches_force_balance drives the fixture (initial interior residual norm 3.78) below 1e-3 relative with fewer than 20 residual evaluations per Newton iteration. test_simsopt_gradient_matches_fd_resolve validates the adjoint boundary gradient against centered FD with full re-solves at rtol=1e-5.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant