vmecpp.autodiff_wout module

JAX port of output_quantities.cc: wout quantities from a converged geometry.

The solver returns the equilibrium in the internal product basis. From this we can compute all the physics quantities in a wout.

Only the geometry leaves are traced. Radial and angular grids, mode tables and the flux and mass profiles are concrete NumPy arrays taken from the input, so wout_quantities() works under jax.jit and jax.grad.

vmecpp.autodiff_wout.PROFILE_PARAMETERS = {'ac': 'pcurr_type', 'ai': 'piota_type', 'am': 'pmass_type', 'curtor': 'pcurr_type', 'pres_scale': 'pmass_type'}

The differentiable profile fields and the profile type each belongs to.

vmecpp.autodiff_wout.UNKNOWN_DIAGNOSTICS: dict[str, Any] = {'delbsq': array([], dtype=float64), 'force_residual_lambda': array([], dtype=float64), 'force_residual_r': array([], dtype=float64), 'force_residual_z': array([], dtype=float64), 'fsql': nan, 'fsqr': nan, 'fsqt': array([], dtype=float64), 'fsqz': nan, 'ier_flag': -1, 'itfsq': 0, 'lcfs_force': array([], dtype=float64), 'niter': 0, 'restart_reason_timetrace': array([], dtype=int64), 'wdot': array([], dtype=float64)}

Solver diagnostics of a VmecWOut whose solve is not observable, e.g. under jax.jit: ier_flag = -1, NaN residuals and empty traces.

vmecpp.autodiff_wout.WOUT_QUANTITIES = ('rmnc', 'zmns', 'lmns', 'lmns_full', 'gmnc', 'bmnc', 'bsubumnc', 'bsubvmnc', 'bsubsmns', 'bsupumnc', 'bsupvmnc', 'currumnc', 'currvmnc', 'rmns', 'zmnc', 'lmnc', 'lmnc_full', 'gmns', 'bmns', 'bsubumns', 'bsubvmns', 'bsubsmnc', 'bsupumns', 'bsupvmns', 'currumns', 'currvmns', 'raxis_cc', 'zaxis_cs', 'raxis_cs', 'zaxis_cc', 'iotas', 'iotaf', 'q_factor', 'phi', 'phipf', 'phips', 'chi', 'chipf', 'presf', 'pres', 'mass', 'vp', 'buco', 'bvco', 'jcuru', 'jcurv', 'jdotb', 'bdotb', 'bdotgradv', 'specw', 'over_r', 'beta_vol', 'equif', 'DMerc', 'DShear', 'DWell', 'DCurr', 'DGeod', 'wb', 'wp', 'rmax_surf', 'rmin_surf', 'zmax_surf', 'aspect', 'betatotal', 'betapol', 'betator', 'betaxis', 'b0', 'rbtor0', 'rbtor', 'IonLarmor', 'volavgB', 'ctor', 'Aminor_p', 'Rmajor_p', 'volume')

The VmecWOut fields computed by wout_quantities().

The non-stellarator-symmetric fields are None for lasym=False.

vmecpp.autodiff_wout.half_grid_profiles(vmec_input, ns, parameters)

The C++ half-grid (mu_0 p, iota, currH) as a JAX function of power-series parameters, stacked to shape (3, ns - 1); iota is before the theta flip.

Parameters:
Return type:

Array

vmecpp.autodiff_wout.mass_profile(vmec_input, s_half)

The half-grid mass profile mu_0 p of the C++ radial profiles.

Parameters:

  • vmec_input (Any)

  • s_half (ndarray)

Return type:

ndarray

vmecpp.autodiff_wout.static_fields(vmec_input)

The VmecWOut fields fixed by the input: echoes, sizes and mode tables.

Free-boundary fields are those of a fixed-boundary run.

Parameters:

vmec_input (Any)

Return type:

dict[str, Any]

vmecpp.autodiff_wout.toroidal_flux_derivative(vmec_input, s)

d phi / d s of the input flux profile, before the 2 pi and sign factors.

Parameters:

  • vmec_input (Any)

  • s (ndarray)

Return type:

ndarray

vmecpp.autodiff_wout.wout_quantities(geometry, vmec_input, *, iota_half=None, current_half=None, mass_half=None)

Map a converged geometry to the physics quantities of its wout file.

Parameters:

  • geometry (Geometry) – The solved equilibrium in the internal product basis.

  • vmec_input (Any) – The vmecpp.VmecInput the geometry was solved for; supplies the grid sizes, mode tables and the flux and mass profiles.

  • iota_half – Rotational transform on the half grid. Defaults to the ratio of the poloidal and toroidal flux increments of geometry.

  • current_half – Enclosed toroidal current <B_u> on the half grid, for ncurr = 1: chi’ is then solved from it as in the C++ computeBContra, and iota follows. Excludes iota_half.

  • mass_half – The half-grid mass profile mu_0 p dV/ds^gamma. Defaults to mass_profile() of the input, which is not differentiated.

Return type:

dict[str, Array | None]

Returns:

The fields listed in WOUT_QUANTITIES, in the layout of vmecpp.VmecWOut.