vmec_jax.free_boundary_adjoint

Adjoint scaffolding for free-boundary vacuum solves.

Phase 1 intentionally keeps this module small and explicit. It validates the linear-solve differentiation contract that the production NESTOR replacement will need: solve the primal system in the forward pass and use transpose solves in the backward pass rather than differentiating through an iterative solver.

Functions

dense_fixed_point_solve_jax(update_fn, ...)	Solve x = update_fn(x, params) with the nonlinear implicit adjoint.
dense_mode_vacuum_solve_jax(mode_matrix, ...)	Solve a dense mode-space vacuum system and reconstruct a grid potential.
dense_nonlinear_solve_jax(residual_fn, ...)	Solve a small nonlinear residual with an implicit-root adjoint.
dense_vacuum_residual(A, x, b)	Return A @ x - b for tests and diagnostics.
dense_vacuum_solve_jax(A, b, *[, symmetric])	Solve a dense toy vacuum linear system with an implicit adjoint.
dense_vmec_nestor_mode_solve_jax(*, R, Z, ...)	Assemble and solve the dense JAX VMEC/NESTOR mode operator.
direct_coil_accepted_trace_array_controls_jax(traces)	Return stacked array-valued update controls for accepted trace replay.
direct_coil_accepted_trace_branch_metadata(...)	Return branch metadata for a fixed accepted free-boundary trace.
direct_coil_accepted_trace_controller_controls_jax(...)	Return stacked JAX-visible controls for fixed accepted trace replay.
direct_coil_accepted_trace_controller_custom_vjp_objective_jax(...)	Return a scalar stacked-controller replay objective with custom VJP.
direct_coil_accepted_trace_controller_custom_vjp_scalar_jax(...)	Return a scalar of accepted-controller replay with a custom VJP seam.
direct_coil_accepted_trace_controller_custom_vjp_scalars_jax(...)	Return several accepted-controller replay scalars with one custom VJP.
direct_coil_accepted_trace_controller_directional_check_jax(...)	Validate stacked accepted-controller replay gradients by central FD.
direct_coil_accepted_trace_controller_replay_objective_jax(...)	Replay fixed production traces through a JAX-visible accept controller.
direct_coil_accepted_trace_directional_check_jax(...)	Validate accepted-trace replay coil gradients by central FD.
direct_coil_accepted_trace_fingerprint(traces, *)	Return a branch-control fingerprint for accepted free-boundary traces.
direct_coil_accepted_trace_fingerprint_delta(...)	Compare two accepted-trace fingerprints.
direct_coil_accepted_trace_fingerprint_delta_summary(...)	Return a strict-JSON-safe accepted-trace fingerprint delta summary.
direct_coil_accepted_trace_preconditioner_controls_jax(traces)	Return stacked preconditioner/mode payloads for accepted replay.
direct_coil_accepted_trace_preconditioner_policy_segment_summary(...)	Return JSON-safe preconditioner-policy segment diagnostics.
direct_coil_accepted_trace_preconditioner_policy_segments(...)	Return consecutive static-preconditioner-policy segments.
direct_coil_accepted_trace_replay_graph_metadata(...)	Return profiling metadata for the fixed accepted-branch replay graph.
direct_coil_accepted_trace_replay_objective_jax(...)	Replay fixed accepted free-boundary traces with differentiable coils.
direct_coil_accepted_trace_scalar_controls_jax(traces)	Return stacked scalar/update controls consumed by accepted trace replay.
direct_coil_accepted_trace_step_controls_jax(...)	Return stacked state/constraint controls for direct accepted replay.
direct_coil_accepted_trace_step_policy_segment_summary(...)	Return JSON-safe diagnostics for stacked step-policy segments.
direct_coil_accepted_trace_step_policy_segments(...)	Return consecutive static step-policy segments for stacked replay.
direct_coil_adaptive_full_loop_same_branch_gate_report(...)	Report whether complete-loop FD is compatible with stacked replay AD.
direct_coil_boundary_bnormal_rms_jax(params, ...)	Replay the accepted-boundary direct-coil normal-field RMS in JAX.
direct_coil_boundary_bsqvac_from_trace_jax(...)	Replay direct-coil bsqvac on accepted geometry using trace metadata.
direct_coil_boundary_bsqvac_jax(params, *, ...)	Replay accepted-boundary direct-coil bsqvac through JAX NESTOR.
direct_coil_boundary_replay_context(static, ...)	Build static NESTOR replay data for an accepted boundary geometry.
direct_coil_boundary_replay_context_for_shape(...)	Build shape/static NESTOR replay data for accepted-boundary replay.
direct_coil_complete_solve_trace(input_path, ...)	Run a direct-coil free-boundary solve and return accepted traces.
direct_coil_fixed_trace_custom_vjp_objective_jax(...)	Return a scalar fixed-trace objective with an explicit custom VJP seam.
direct_coil_projected_mode_fixed_point_directional_check_jax(...)	Validate projected-mode fixed-point coil gradients by central FD.
direct_coil_projected_mode_fixed_point_jax(...)	Solve a small direct-coil free-boundary fixed-point validation loop.
direct_coil_projected_mode_fixed_point_objective_jax(...)	Return a scalar objective for the projected-mode fixed-point helper.
direct_coil_run_free_boundary_branch_local_scalar_value_and_grad_jax([...])	Return a production-forward branch-local scalar value and gradient.
direct_coil_run_free_boundary_branch_local_scalars_value_and_jacobian_jax([...])	Return production-forward branch-local values and a scalar Jacobian.
direct_coil_same_branch_complete_solve_fd_report(...)	Return same-branch complete-solve finite-difference diagnostics.
direct_coil_same_branch_controller_scalar_custom_vjp_report(...)	Compare a branch-local scalar custom VJP with complete-solve FD.
direct_coil_same_branch_controller_scalars_custom_vjp_report(...)	Batch same-branch custom-VJP reports for several replay scalars.
direct_coil_same_branch_physical_scalar_gate_report(...)	Return a reviewer-facing same-branch physical-scalar promotion gate.
direct_coil_same_branch_replay_gate_report(...)	Return the branch gate for promoting a fixed-trace replay derivative.
free_boundary_adjoint_trace_replay_diagnostics(...)	Return diagnostics for fixed accepted-trace free-boundary replay.
free_boundary_boundary_geometry_jax(state, ...)	Synthesize accepted free-boundary geometry through JAX.
mode_matrix_from_grpmn_jax(grpmn, *, ...[, ...])	Build the VMEC/NESTOR mode matrix from Green-function mode samples.
mode_rhs_from_gsource_jax(gsource, *, ...[, ...])	Project a VMEC/NESTOR grid source into mode-space RHS coefficients.
vacuum_boundary_fields_from_cylindrical_jax(*, ...)	JAX version of the VMEC boundary-field projection scaffold.
vacuum_boundary_fields_from_mode_coeffs_jax(...)	JAX replay of VMEC vacuum channels from NESTOR mode coefficients.
vmec_analytic_terms_from_geometry_jax(*, R, ...)	JAX VMEC/NESTOR analytic singular-source terms from analyt.f.
vmec_nonsingular_terms_from_bexni_jax(*, R, ...)	JAX VMEC/NESTOR nonsingular Green-function source/matrix assembly.
vmec_source_from_gsource_jax(gsource, *, ...)	JAX version of VMEC/NESTOR source symmetrization.