vmec_jax.optimization

Optimization-oriented helpers for vmec_jax workflows.

Functions

apply_boundary_params(boundary, specs, params)	Apply parameter updates to a boundary coefficient set.
boundary_param_names(specs)	Return the parameter names for a list of specs.
boundary_param_specs(boundary, modes, *[, ...])	Build parameter specifications for boundary optimization.
create_x_scale(specs, *[, alpha])	Compute per-parameter exponential spectral scaling weights.
extend_boundary_for_max_mode(indata, static, ...)	Extend indata, static, and boundary to support max_mode DOFs.
gauss_newton_least_squares(residual_fun, ...)	Solve a nonlinear least-squares problem with a concrete Gauss-Newton loop.
lift_boundary_params(source_specs, ...)	Lift a parameter vector defined on one boundary basis to another.
make_qh_residuals_fn(static, indata, *[, ...])	Build a residuals_from_state callable for quasi-helical symmetry.
make_qs_residuals_fn(static, indata, *[, ...])	General quasisymmetry residuals factory supporting QH and QA objectives.
parse_surface_list(text)	Parse a comma-separated surface list into floats/ints.
prepare_fixed_boundary_context(*, static, ...)	Precompute common fixed-boundary inputs for optimization loops.
rebuild_indata_with_resolution(indata, *, ...)	Return a copy of indata with updated VMEC spectral resolution.
smooth_min_abs_iota_residual(iota, minimum, *)	Smooth residual for the differentiable constraint abs(iota) >= minimum.
surface_indices_from_s(s_half, surfaces)	Map surface requests to half-mesh indices.
surface_indices_from_static(static, surfaces)	Map surface requests to indices using a VMEC static object.
truncate_indata_boundary_modes(indata, *, ...)	Return a copy of indata with boundary modes above max_mode zeroed.

Classes

BoundaryParamSpec(name, kind, index, m, n)	Descriptor for a boundary Fourier coefficient parameter.
FixedBoundaryContext(st_guess, signgs, flux, ...)	Bundled inputs for repeated fixed-boundary solves.
FixedBoundaryExactOptimizer(static, indata, ...)	End-to-end optimizer for fixed-boundary VMEC equilibria.