vmec_jax.implicit

Implicit differentiation utilities.

This module provides custom VJP wrappers for equilibrium sub-solves so that outer objectives can differentiate through equilibrium states without backpropagating through many optimization iterations.

Initial scope

Implicit differentiation for the lambda-only fixed-geometry solve.

Future work

Extend to the full fixed-boundary solve over (R,Z,lambda) using the same implicit-function machinery.

Functions

solve_fixed_boundary_state_implicit(state0, ...)	Fixed-boundary solve with a custom VJP using implicit differentiation.
solve_fixed_boundary_state_implicit_vmec_residual(...)	Implicitly differentiate a VMEC residual fixed-point solve.
solve_lambda_state_implicit(state0, static, ...)	Solve lambda with a custom VJP that uses implicit differentiation.

Classes

ImplicitFixedBoundaryOptions([cg_max_iter, ...])	Controls for the implicit backward pass (fixed-boundary solve).
ImplicitLambdaOptions([cg_max_iter, cg_tol, ...])	Controls for the implicit backward pass.