API Reference

The modules below are the supported public ownership areas. Internal modules whose names begin with _ may change without notice; use the public facades documented here.

Solver

Dense JAX block-tridiagonal monoenergetic DKE solver.

class ntx.solver.CompiledPreparedScanSolver(prepared: PreparedMonoenergeticSystem, *, batch_size: int, execution_mode: Literal['sequential', 'vectorized'], solve_batch: _CompiledBatchFunction)

Reusable fixed-shape scan solver for one prepared geometry.

The final input chunk is padded to batch_size and trimmed after the solve. Consequently, repeated calls reuse one executable instead of recompiling for each scan length.

warmup() → PreparedScanCompilationReport: Lower, compile, and execute the fixed-shape scan once explicitly.

class ntx.solver.MonoenergeticCase(nu_hat: float | Array, epsi_hat: float | Array | None = None, er_hat: float | Array | None = None): Monoenergetic DKE parameters.

class ntx.solver.PreparedDerivativeAuditResult(primal_relative_residual: Array, transpose_relative_residual: Array, direct_reverse_gradient: Array, prepared_adjoint_gradient: Array, forward_gradient: Array, finite_difference_gradient: Array, prepared_adjoint_relative_error: Array, forward_relative_error: Array, finite_difference_relative_error: Array, finite_difference_step: Array, valid: Array, coefficient: str, parameter: str, residual_tolerance: float, gradient_tolerance: float)

Residual and independent-gradient evidence for one scalar derivative.

as_dict() → dict[str, float | bool | str]: Materialize the audit as JSON-friendly scalar metadata.

class ntx.solver.PreparedMonoenergeticSystem(surface: BoozerSurface | VmecSurface, grid: GridSpec, geometry: GeometryOnGrid, d_theta: Array, d_zeta: Array): Cached geometry and derivative operators for repeated solves.

class ntx.solver.PreparedScanCompilationReport(lowering_seconds: float, compilation_seconds: float, first_execution_seconds: float, warm_execution_seconds: float, generated_code_size_bytes: int | None, argument_size_bytes: int | None, output_size_bytes: int | None, temporary_size_bytes: int | None): Ahead-of-time warmup timings and executable memory estimates.

class ntx.solver.ResidualAuditResult(tail_eliminated_l2: Array, full_system_l2: Array, retained_mode_max_abs_error: Array, n_modes: int)

Opt-in comparison of the low-memory and full Legendre solves.

property full_system_residual_l2: Array: RMS residual obtained by applying every original Legendre row.

property schur_residual_l2: Array: RMS residual of the tail-eliminated Schur system.

class ntx.solver.TransportResult(D11: Array, D31: Array, D13: Array, D33: Array, D33_spitzer: Array, f1_modes: Array, f3_modes: Array, residual_l2: Array, onsager_residual: Array)

Monoenergetic coefficients, retained modes, and solver diagnostics.

residual_l2 is retained for API compatibility. It is the RMS residual of the tail-eliminated Schur system, not the residual of every original Legendre block row. Use schur_residual_l2 for explicit new code and ntx.audit_prepared_residuals() when a full-system residual is needed.

property schur_residual_l2: Array: RMS algebraic residual of the complete tail-eliminated system.

ntx.solver.audit_prepared_coefficient_derivative(prepared: PreparedMonoenergeticSystem, case: MonoenergeticCase, *, coefficient: str = 'D11', parameter: Literal['nu_hat', 'epsi_hat', 'er_hat'] = 'er_hat', finite_difference_step: float | Array | None = None, residual_tolerance: float = 1e-10, gradient_tolerance: float = 0.0001) → PreparedDerivativeAuditResult

Audit one prepared coefficient derivative before accepting it as valid.

The audit compares direct reverse mode, forward mode, the factor-reusing prepared adjoint, and a centered finite difference. It independently evaluates the full primal and algebraic-transpose block equations.

ntx.solver.audit_prepared_residuals(prepared: PreparedMonoenergeticSystem, case: MonoenergeticCase) → ResidualAuditResult

Reconstruct all Legendre modes and independently audit every block row.

This diagnostic stores full block factors and modes. Use it for validation, not in memory-constrained production scans.

ntx.solver.compile_prepared_scan_solver(prepared: PreparedMonoenergeticSystem, *, batch_size: int | None = None, execution_mode: Literal['auto', 'sequential', 'vectorized'] = 'auto') → CompiledPreparedScanSolver

Create a reusable fixed-shape scan solver for prepared.

Automatic execution uses sequential lax.map on every backend to retain scalar-solve parity and bound memory. Explicit vectorization is available for measured crossover and accuracy studies.

ntx.solver.compile_prepared_solver(prepared: PreparedMonoenergeticSystem) → Callable[[MonoenergeticCase], TransportResult]: Return a jitted monoenergetic solver for repeated solves on one geometry.

ntx.solver.healthy_parallel_device_count() → int: Return the number of local devices that pass the NTX smoke solve.

ntx.solver.healthy_parallel_devices() → tuple: Return cached local JAX devices that pass a finite NTX smoke solve.

ntx.solver.local_parallel_device_count() → int: Return the number of local JAX devices visible to this process.

ntx.solver.prepare_monoenergetic_system(surface: BoozerSurface | VmecSurface, grid: GridSpec, *, require_resolved_geometry: bool = False) → PreparedMonoenergeticSystem: Precompute geometry and derivatives, optionally enforcing Nyquist sampling.

ntx.solver.solve_monoenergetic(surface: BoozerSurface | VmecSurface, grid: GridSpec, case: MonoenergeticCase, *, require_resolved_geometry: bool = False) → TransportResult: Solve one monoenergetic DKE case.

ntx.solver.solve_monoenergetic_internal(surface: BoozerSurface | VmecSurface, grid: GridSpec, case: MonoenergeticCase, *, require_resolved_geometry: bool = False) → tuple[Array, Array, Array]: Solve one monoenergetic case and return (Dij, f, s) low-order arrays.

ntx.solver.solve_monoenergetic_parallel_scan(surface: BoozerSurface | VmecSurface, grid: GridSpec, nu_hat: Array, *, epsi_hat: Array | None = None, er_hat: Array | None = None, num_devices: int | None = None, scan_batch_size: int | None = None) → dict[str, Array]

Shard a scan over healthy local JAX devices.

scan_batch_size is applied inside each device shard so device sharding and peak-memory control remain independent choices.

ntx.solver.solve_monoenergetic_scan(surface: BoozerSurface | VmecSurface, grid: GridSpec, nu_hat: Array, *, epsi_hat: Array | None = None, er_hat: Array | None = None, scan_batch_size: int | None = None) → dict[str, Array]

Scan collisionality and radial electric field on one JAX device.

scan_batch_size optionally splits the flattened scan into fixed-size batches. This preserves coefficient values while bounding peak memory on CPUs and memory-constrained accelerators.

ntx.solver.solve_prepared(prepared: PreparedMonoenergeticSystem, case: MonoenergeticCase) → TransportResult: Solve one monoenergetic case using precomputed geometry and derivatives.

ntx.solver.solve_prepared_coefficient_vector(prepared: PreparedMonoenergeticSystem, case: MonoenergeticCase) → Array: Return the coefficient vector [D11, D31, D13, D33, D33_spitzer].

ntx.solver.solve_prepared_internal(prepared: PreparedMonoenergeticSystem, case: MonoenergeticCase) → tuple[Array, Array, Array]: Solve one prepared monoenergetic case and return (Dij, f, s) low-order arrays.

ntx.solver.solve_scan(surface: BoozerSurface | VmecSurface, grid: GridSpec, cases: tuple[MonoenergeticCase, ...]) → list[TransportResult]: Solve a Python-level scan of monoenergetic cases.

Geometry

Flux-surface geometry for the monoenergetic DKE.

class ntx.geometry.BoozerSurface(m: Array, n: Array, b_cos: Array, nfp: int, iota: float | Array, psi_p: float | Array, b_theta: float | Array, b_zeta: float | Array, chi_p: float | Array | None = None, b0: float | Array | None = None, b_sin: Array | None = None, stellarator_symmetric: bool = True, source_path: Path | None = None): Single flux-surface representation in Boozer coordinates.

class ntx.geometry.GeometryOnGrid(surface_type: 'str', surface_path: 'Path | None', nfp: 'int', iota: 'float | Array', psi_p: 'float | Array | None', transport_psi_scale: 'float | Array', coefficient_psi_scale: 'float | Array', grid: 'AngularGrid', theta_2d: 'Array', zeta_2d: 'Array', b: 'Array', d_b_dtheta: 'Array', d_b_dzeta: 'Array', jacobian: 'Array', b_sub_theta: 'Array', b_sub_zeta: 'Array', b_sup_theta: 'Array', b_sup_zeta: 'Array', volume_prime: 'Array', b2_mean: 'Array', radial_drift_spatial: 'Array', b0: 'Array')

class ntx.geometry.VmecSurface(path: Path, requested_psi_n: float, psi_n: float, nfp: int, ns: int, mpol: int, ntor: int, total_mode_count: int, loaded_mode_count: int, iota: float, m: Array, n: Array, b_cos: Array, jacobian_cos: Array, b_sub_theta_cos: Array, b_sub_zeta_cos: Array, b_sup_theta_cos: Array, b_sup_zeta_cos: Array, b0: float, psi_a_hat: float, phi_edge: float, r_n: float, r_hat: float, dpsi_hat_dr_hat: float, dr_hat_dpsi_hat: float, aminor_p: float | None = None, psi_p: float | None = None, transport_psi_scale: float = 1.0, stellarator_symmetric: bool = True): Single flux-surface representation loaded from a VMEC wout file.

ntx.geometry.evaluate_boozer_modes(surface: BoozerSurface, theta: Array, zeta: Array) → tuple[Array, Array, Array]: Evaluate B, dB/dtheta, and dB/dzeta on broadcastable arrays.

ntx.geometry.evaluate_fourier_series(m: Array, n: Array, cos_coeffs: Array, theta: Array, zeta: Array, *, nfp: int, sin_coeffs: Array | None = None) → tuple[Array, Array, Array]: Evaluate a stellarator-symmetric or general Fourier series and its derivatives.

ntx.geometry.example_surface(dtype=<class 'jax.numpy.float64'>) → BoozerSurface: Return a small stellarator-symmetric test surface.

ntx.geometry.geometry_on_grid(surface: BoozerSurface | VmecSurface, spec) → GeometryOnGrid: Evaluate all geometric quantities needed by the solver.

Operators And Transport

Block operators for the Legendre-space monoenergetic DKE.

class ntx.operators.OperatorContext(surface: 'BoozerSurface | VmecSurface', geometry: 'GeometryOnGrid', nu_hat: 'Array', epsi_hat: 'Array')

ntx.operators.apply_nullspace_condition(d_block: Array, u_block: Array | None = None) → tuple[Array, Array | None]: Replace the first row by the f^(0)(0,0)=0 constraint.

ntx.operators.build_block(coefficients: Array, d_theta: Array, d_zeta: Array) → Array: Construct a dense spatial block matrix.

ntx.operators.coefficients_for_k(ctx: OperatorContext, k: int | Array) → tuple[Array, Array, Array]: Return coefficient arrays for lower, diagonal, and upper blocks.

ntx.operators.derivative_blocks(geom: GeometryOnGrid) → tuple[Array, Array]: Return lifted theta and zeta derivative matrices for flattened fields.

ntx.operators.operator_blocks(ctx: OperatorContext, k: int | Array, d_theta: Array, d_zeta: Array) → tuple[Array, Array, Array]: Construct (L_k, D_k, U_k) dense blocks.

ntx.operators.parameter_derivative_blocks(ctx: OperatorContext, k: int | Array, d_theta: Array, d_zeta: Array) → tuple[Array, Array]: Construct dense dD_k/dnu_hat and dD_k/depsi_hat blocks.

ntx.operators.source_modes(ctx: OperatorContext, n_xi: int) → tuple[Array, Array]: Build source arrays with shape (n_xi + 1, n_fs) for s1 and s3.

Transport coefficient post-processing.

ntx.transport.coefficients_from_modes(geom: GeometryOnGrid, f1_modes: Array, f3_modes: Array, nu_hat: Array) → tuple[Array, Array, Array, Array, Array]: Compute (D11, D31, D13, D33, D33_spitzer) from modes 0, 1, and 2.

ntx.transport.onsager_error(d31: Array, d13: Array) → Array: Return |D13 + D31|, the relevant monoenergetic Onsager residual.

Resolution

Geometry sampling reports for Fourier-collocation solves.

class ntx.resolution.GeometryResolutionReport(m_min: int, m_max: int, n_min: int, n_max: int, theta_nyquist_floor: int, zeta_nyquist_floor: int, n_theta: int, n_zeta: int, theta_oversampling: float, zeta_oversampling: float, status: Literal['resolved', 'undersampled'], warnings: tuple[str, ...], errors: tuple[str, ...])

Nyquist audit for the retained geometry harmonics.

NTX evaluates phases as m * theta + n * nfp * zeta on one field period. The toroidal field-period factor therefore cancels from the sampled mode count: n_zeta resolves the retained integer n harmonics.

require_resolved() → None: Raise before geometry assembly when retained modes would alias.

ntx.resolution.geometry_resolution_report(surface: BoozerSurface | VmecSurface, grid: GridSpec, *, warning_oversampling: float = 2.25) → GeometryResolutionReport

Return retained harmonic extrema, Nyquist floors, and sampling status.

Call this host-side before passing a surface as a dynamic argument to jax.jit. Concrete surfaces are checked automatically during solver preparation; traced integer mode arrays cannot be converted into a host report without changing the transform contract.

Profiles

Profile-grade imported transport workflows built on NTX scan data.

class ntx.profiles.AmbipolarProfileFamilyResult(control: Array, er_profile: Array, ambipolar_residual: Array, bootstrap_current_response: Array, loss_history: Array)

Stacked ambipolar-profile solutions across a control-parameter family.

property bootstrap_current_proxy: Array: Compatibility alias for the reduced bootstrap-current response family.

class ntx.profiles.AmbipolarProfileResult(rho: Array, er_profile: Array, ambipolar_residual: Array, bootstrap_current_response: Array, species_particle_flux: Array, species_current_response: Array, loss_history: Array)

Solved electric-field profile and reduced monoenergetic responses.

property bootstrap_current_proxy: Array

Compatibility alias for the reduced bootstrap-current response.

New code should use bootstrap_current_response. The alias is kept so NTX 0.2.x readers do not break, but it is no longer the registered runtime field because the quantity is a reduced monoenergetic response, not a fitted bootstrap-current closure.

class ntx.profiles.MonoenergeticSpeciesProfile(charge: float | Array, nu_v: Array, A1: Array, A3: Array, particle_weight: float | Array = 1.0, current_weight: float | Array = 1.0, name: str | None = None): One-species radial inputs for reduced ambipolar/current responses.

class ntx.profiles.PrimitiveProfileTransportIterationResult(er_profile_history: Array, ambipolar_residual_history: Array, bootstrap_current_response_history: Array, transport_loss_history: Array, species_density_history: Array, species_temperature_history: Array, best_profile: AmbipolarProfileResult)

History of a primitive-profile transport relaxation workflow.

property bootstrap_current_proxy_history: Array: Compatibility alias for the reduced bootstrap-current response history.

class ntx.profiles.PrimitiveSpeciesProfile(charge: float | Array, nu_v: Array, density: Array, temperature: Array, electrostatic_prefactor: float | Array = 1.0, particle_weight: float | Array = 1.0, current_weight: float | Array = 1.0, name: str | None = None): Primitive radial inputs used to construct monoenergetic force profiles.

class ntx.profiles.ProfileBasisControlSpec(basis: Array, a1_response: Array, a3_response: Array, control_name: str = 'basis control'): Low-dimensional radial basis control applied to A1 and A3.

class ntx.profiles.ProfileBasisOptimizationResult(control_history: Array, objective_history: Array, bootstrap_objective_history: Array, residual_norm_history: Array, best_control: Array, best_profile: AmbipolarProfileResult): Optimization history for a vector profile-basis control.

class ntx.profiles.ProfileControlOptimizationResult(control_history: Array, objective_history: Array, bootstrap_objective_history: Array, residual_norm_history: Array, best_control: Array, best_profile: AmbipolarProfileResult): Optimization history for a scalar profile control.

class ntx.profiles.ProfileControlSpec(a1_response: Array, a3_response: Array, control_name: str = 'control'): Linear control map applied to A1 and A3 for each species.

class ntx.profiles.ProfileTransportClosureSpec(particle_relaxation: Array, current_relaxation: Array, particle_target: float | Array = 0.0, current_target: float | Array = 0.0, particle_source: float | Array = 0.0, current_source: float | Array = 0.0, normalization_floor: float | Array = 1.0, max_normalized_update: float | Array = 0.35, density_relaxation: float | Array = 0.0, temperature_relaxation: float | Array = 0.0, density_target: float | Array = 0.0, temperature_target: float | Array = 0.0, density_source: float | Array = 0.0, temperature_source: float | Array = 0.0, primitive_normalization_floor: float | Array = 1.0, max_primitive_normalized_update: float | Array = 0.2, radial_smoothing_strength: float | Array = 0.0, closure_name: str = 'transport loop'): Relaxation closure for iterating reduced responses toward transport targets.

class ntx.profiles.ProfileTransportIterationResult(er_profile_history: Array, ambipolar_residual_history: Array, bootstrap_current_response_history: Array, transport_loss_history: Array, species_a1_history: Array, species_a3_history: Array, best_profile: AmbipolarProfileResult)

History of a simple self-consistent profile transport relaxation loop.

property bootstrap_current_proxy_history: Array: Compatibility alias for the reduced bootstrap-current response history.

ntx.profiles.advance_primitive_profile_transport(primitive_profiles: tuple[PrimitiveSpeciesProfile, ...], profile: AmbipolarProfileResult, closure_spec: ProfileTransportClosureSpec) → tuple[PrimitiveSpeciesProfile, ...]: Apply one explicit transport-relaxation update to primitive profiles.

ntx.profiles.advance_profile_transport(species_profiles: tuple[MonoenergeticSpeciesProfile, ...], profile: AmbipolarProfileResult, closure_spec: ProfileTransportClosureSpec) → tuple[MonoenergeticSpeciesProfile, ...]: Apply one explicit transport-relaxation update to A1 and A3.

ntx.profiles.ambipolar_residual_profile(scan: NeopaxScan, species_profiles: tuple[MonoenergeticSpeciesProfile, ...], *, er_profile: Array) → Array: Return the charge-weighted monoenergetic ambipolar residual profile.

ntx.profiles.apply_profile_basis_control(species_profiles: tuple[MonoenergeticSpeciesProfile, ...], control: Array, control_spec: ProfileBasisControlSpec) → tuple[MonoenergeticSpeciesProfile, ...]: Apply a low-dimensional radial basis control to A1 and A3.

ntx.profiles.apply_profile_control(species_profiles: tuple[MonoenergeticSpeciesProfile, ...], control: Array | float, control_spec: ProfileControlSpec) → tuple[MonoenergeticSpeciesProfile, ...]: Apply a scalar control to the A1 and A3 profiles for each species.

ntx.profiles.bootstrap_current_objective(rho: Array, current_response: Array, *, weight: Array | None = None) → Array: Return a weighted quadratic objective for a reduced current response.

ntx.profiles.build_species_profile_from_primitives(rho: Array, primitive: PrimitiveSpeciesProfile, *, er_profile: Array) → MonoenergeticSpeciesProfile: Construct A1(r) and A3(r) from primitive density/temperature profiles.

ntx.profiles.build_species_profiles_from_primitives(rho: Array, primitives: tuple[PrimitiveSpeciesProfile, ...], *, er_profile: Array) → tuple[MonoenergeticSpeciesProfile, ...]: Vectorized helper for primitive-to-monoenergetic profile construction.

ntx.profiles.current_response_objective(rho: Array, current_response: Array, *, weight: Array | None = None) → Array: Return a weighted quadratic objective for a reduced current response.

ntx.profiles.evaluate_scan_channel(scan: NeopaxScan, channel: str, rho: Array, nu_v: Array, er_profile: Array) → Array: Interpolate one NTX scan channel over (rho, nu_v, E_r).

ntx.profiles.evaluate_species_current_response(scan: NeopaxScan, species: MonoenergeticSpeciesProfile, *, rho: Array | None = None, er_profile: Array) → Array: Return the reduced monoenergetic parallel-current response for one species.

ntx.profiles.evaluate_species_particle_flux(scan: NeopaxScan, species: MonoenergeticSpeciesProfile, *, rho: Array | None = None, er_profile: Array) → Array: Return the reduced monoenergetic particle-flux response for one species.

ntx.profiles.optimize_profile_basis_control(scan: NeopaxScan, species_profiles: tuple[MonoenergeticSpeciesProfile, ...], control_spec: ProfileBasisControlSpec, *, control_initial: Array, learning_rate: float = 0.12, optimization_steps: int = 12, solve_steps: int = 16, damping: float = 0.8, smoothing_strength: float = 0.0, weight: Array | None = None, residual_penalty: float = 1.0, control_penalty: float = 0.01, control_bound: float | Array | None = 0.4, backtracking_steps: int = 5) → ProfileBasisOptimizationResult: Optimize a vector radial-basis control against the profile objective.

ntx.profiles.optimize_profile_control(scan: NeopaxScan, species_profiles: tuple[MonoenergeticSpeciesProfile, ...], control_spec: ProfileControlSpec, *, control_initial: float | Array = 0.0, learning_rate: float = 0.15, optimization_steps: int = 12, solve_steps: int = 16, damping: float = 0.8, smoothing_strength: float = 0.0, weight: Array | None = None, residual_penalty: float = 1.0, control_bound: float | Array | None = 0.6, backtracking_steps: int = 5) → ProfileControlOptimizationResult: Optimize a scalar profile control against the current-response objective.

ntx.profiles.primitive_profile_transport_loss(profile: AmbipolarProfileResult, primitive_profiles: tuple[PrimitiveSpeciesProfile, ...], closure_spec: ProfileTransportClosureSpec) → Array: Combined profile-transport loss including primitive source/target closure.

ntx.profiles.profile_transport_loss(profile: AmbipolarProfileResult, closure_spec: ProfileTransportClosureSpec) → Array: Quadratic transport mismatch loss for a solved ambipolar profile.

ntx.profiles.solve_ambipolar_er_profile(scan: NeopaxScan, species_profiles: tuple[MonoenergeticSpeciesProfile, ...], *, er_initial: Array | None = None, steps: int = 16, damping: float = 0.8, smoothing_strength: float = 0.0) → AmbipolarProfileResult: Solve a smooth ambipolar E_r(r) profile from an NTX scan.

ntx.profiles.solve_ambipolar_profile_family(scan: NeopaxScan, species_profiles_family: tuple[tuple[MonoenergeticSpeciesProfile, ...], ...], *, control: Array | None = None, er_initial: Array | None = None, steps: int = 16, damping: float = 0.8, smoothing_strength: float = 0.0) → AmbipolarProfileFamilyResult: Solve a family of ambipolar profiles across explicit profile controls.

ntx.profiles.solve_primitive_profile_transport_loop(scan: NeopaxScan, primitive_profiles: tuple[PrimitiveSpeciesProfile, ...], closure_spec: ProfileTransportClosureSpec, *, iterations: int = 8, er_initial: Array | None = None, solve_steps: int = 16, damping: float = 0.8, smoothing_strength: float = 0.0) → PrimitiveProfileTransportIterationResult: Iterate a primitive density/temperature transport closure.

ntx.profiles.solve_profile_transport_loop(scan: NeopaxScan, species_profiles: tuple[MonoenergeticSpeciesProfile, ...], closure_spec: ProfileTransportClosureSpec, *, iterations: int = 8, er_initial: Array | None = None, solve_steps: int = 16, damping: float = 0.8, smoothing_strength: float = 0.0) → ProfileTransportIterationResult: Iterate a simple self-consistent profile transport closure.

NEOPAX Bridge

Explicit NTX-to-NEOPAX mapping helpers.

class ntx.neopax.DifferentiableNeopaxField(n_r: int, a_b: Array, Psia_value: Array, rho_grid: Array, rho_grid_half: Array, r_grid: Array, r_grid_half: Array, dr: Array, Vprime: Array, Vprime_half: Array, overVprime: Array, epsilon_t: Array, B0: Array, B_10: Array, enlogation: Array, iota: Array, R0: Array, B0prime: Array, curvature: Array, G_PS: Array, sqrtg00_value: Array, Bsqav: Array, I_value: Array, G_value: Array): JAX-safe field payload compatible with the NEOPAX flux routines.

class ntx.neopax.NeopaxMonoenergeticArrays(a_b: Array, rho: Array, nu_log: Array, Er_list: Array, D11_log: Array, D13: Array, D33: Array): Pure-array NEOPAX mapping payload for differentiable imported workflows.

class ntx.neopax.NeopaxScan(rho: Array, nu_v: Array, Er: Array, Es: Array, drds: Array, D11: Array, D13: Array, D33: Array, D33_spitzer: Array | None = None, D31: Array | None = None, Er_tilde: Array | None = None, Er_to_Ertilde: Array | None = None, dr_tildedr: Array | None = None, dr_tildeds: Array | None = None, a_b: float | None = None, psia: float | None = None, b00: Array | None = None, r00: Array | None = None, boozer_i: Array | None = None, boozer_g: Array | None = None, iota: Array | None = None, fac_reference_to_sfincs_11: Array | None = None, fac_reference_to_sfincs_31: Array | None = None, fac_reference_to_sfincs_33: Array | None = None, fac_monkes_to_sfincs_11: Array | None = None, fac_monkes_to_sfincs_31: Array | None = None, fac_monkes_to_sfincs_33: Array | None = None, fac_sfincs_to_dkes_11: Array | None = None, fac_sfincs_to_dkes_31: Array | None = None, fac_sfincs_to_dkes_33: Array | None = None, fac_dkes_to_d11star: Array | None = None, fac_dkes_to_d31star: Array | None = None, fac_dkes_to_d33star: Array | None = None, source_name: str | None = None): Monoenergetic scan data shaped for NEOPAX.

ntx.neopax.build_differentiable_neopax_field(*, n_r: int, rho_half, rho_full, volume_p, vp, iotaf, Psia, bmnc_b, rmnc_b, gmnc_b, xm_b, xn_b, bvco, buco, r0_override=None) → DifferentiableNeopaxField

Mirror NEOPAX.Field(…) with pure JAX array operations.

The implementation intentionally follows the external constructor closely so raw-array parity can be checked directly, while avoiding NumPy scalar conversions that break JAX tracing.

ntx.neopax.build_differentiable_neopax_field_from_vmec_booz_files(n_r: int, vmec_path, booz_path) → DifferentiableNeopaxField

Build the NTX Boozer field object directly from VMEC and Boozer files.

The Boozer coefficients are tabulated in normalized radius. This reader uses build_differentiable_neopax_field, whose B0 and Bsqav normalization is evaluated on that same normalized-radius support.

ntx.neopax.build_differentiable_neopax_field_from_vmex_boundary_params(context: VmecJaxBoundaryContext, params, *, n_r: int, vmec_project: bool = True, max_iter: int = 50, step_size: float = 1.0, ftol: float | None = None, implicit=None, mboz: int = 12, nboz: int = 12, apply_boozer_sign_convention: bool = True) → DifferentiableNeopaxField: Solve a fixed boundary and build a tracer-safe NEOPAX field.

ntx.neopax.build_differentiable_neopax_field_from_vmex_state(*, state, static, indata, signgs: int, n_r: int, mboz: int = 12, nboz: int = 12, apply_boozer_sign_convention: bool = True) → DifferentiableNeopaxField: Build a tracer-safe NEOPAX field from an in-memory vmex state.

ntx.neopax.build_ntx_neopax_scan(surface_loader: Callable[[float], BoozerSurface | VmecSurface], *, rho: Array, nu_v: Array, Es: Array | None = None, Er: Array | None = None, drds: Array, grid: GridSpec, source_name: str | None = None) → NeopaxScan

Build a NEOPAX-style scan from NTX surfaces.

Parameters:

surface_loader – Callable receiving one rho value and returning the corresponding NTX surface object.
rho – Arrays following the same conventions as NEOPAX’s reference HDF5 files.
nu_v – Arrays following the same conventions as NEOPAX’s reference HDF5 files.
Es – Arrays following the same conventions as NEOPAX’s reference HDF5 files.
Er – Arrays following the same conventions as NEOPAX’s reference HDF5 files.
drds – Arrays following the same conventions as NEOPAX’s reference HDF5 files.
grid – NTX angular and Legendre resolution for the solve.

ntx.neopax.build_ntx_neopax_scan_from_surfaces(surfaces: tuple[BoozerSurface | VmecSurface, ...], *, rho: Array, nu_v: Array, Es: Array | None = None, Er: Array | None = None, drds: Array, grid: GridSpec, source_name: str | None = None) → NeopaxScan

Build a NEOPAX-style scan from an explicit tuple of NTX surfaces.

This is the intended imported path when the caller already has surface objects in memory and wants to avoid a Python callback boundary.

ntx.neopax.build_ntx_neopax_scan_from_vmex_boundary_params(context: VmecJaxBoundaryContext, params, *, rho: Array, nu_v: Array, Es: Array | None = None, Er: Array | None = None, drds: Array, grid: GridSpec, source_name: str | None = None, vmec_project: bool = True, max_iter: int = 50, step_size: float = 1.0, ftol: float | None = None, implicit=None, mboz: int = 12, nboz: int = 12, psi_p: float = 1.0, min_bmn_to_load: float = 0.0) → NeopaxScan: Solve a fixed boundary and build a NEOPAX-style scan from the result.

ntx.neopax.build_ntx_neopax_scan_from_vmex_state(*, state, static, indata, signgs: int, rho: Array, nu_v: Array, Es: Array | None = None, Er: Array | None = None, drds: Array, grid: GridSpec, source_name: str | None = None, mboz: int = 12, nboz: int = 12, psi_p: float = 1.0, min_bmn_to_load: float = 0.0) → NeopaxScan: Build a NEOPAX-style scan directly from an in-memory vmex state.

ntx.neopax.get_differentiable_neopax_fluxes(species, grid, field, database)

Evaluate NEOPAX no-momentum fluxes with an axis-safe radial block.

The reference monoenergetic databases do not include the magnetic axis exactly. For integrated objectives such as total bootstrap current, the axis contribution is weighted by Vprime[0] = 0, so copying the first interior radial block into the axis block removes an AD-only singularity without changing the physical integral.

ntx.neopax.load_neopax_reference_scan(path: str | Path) → NeopaxScan: Load a NEOPAX-style HDF5 monoenergetic table.

ntx.neopax.neopax_scan_requires_rebuild(path: str | Path) → bool: Return whether a cached NEOPAX-style scan is missing required fields.

ntx.neopax.scan_to_neopax_arrays(scan: NeopaxScan, *, a_b: float | Array, d33_mode: str = 'raw') → NeopaxMonoenergeticArrays

Map NTX scan data into the pure arrays consumed by NEOPAX.Monoenergetic.

The default raw branch preserves the historical database convention used by the integrated workflow. The spitzer and conductivity_difference branches are explicit audit/stress-test choices and should not be promoted as global defaults without a transfer gate.

ntx.neopax.to_neopax_monoenergetic(scan: NeopaxScan, *, a_b: float | Array, d33_mode: str = 'raw'): Construct NEOPAX.Monoenergetic from NTX scan data.

ntx.neopax.write_neopax_scan_hdf5(scan: NeopaxScan, path: str | Path) → Path: Write a NEOPAX-style HDF5 file from a scan payload.

Autodiff

Autodiff-oriented NTX demonstrations and analysis helpers.

class ntx.autodiff.BootstrapOptimizationResult(scale_history: 'Array', gradient_history: 'Array', objective_history: 'Array', scale_grid: 'Array', objective_landscape: 'Array', rho: 'Array', baseline_scale: 'Array', optimized_scale: 'Array', baseline_current_profile: 'Array', optimized_current_profile: 'Array', baseline_d13_profile: 'Array', optimized_d13_profile: 'Array', baseline_d33_profile: 'Array', optimized_d33_profile: 'Array', current_sensitivity: 'Array', harmonic_m: 'Array', harmonic_n: 'Array', harmonic_reference_value: 'Array', nu_value: 'Array', serial_seconds: 'Array', parallel_seconds: 'Array')

class ntx.autodiff.DerivativeAuditResult(nu_hat: 'Array', er_hat_scan: 'Array', amplitude_value: 'Array', er_reference: 'Array', nu_reference: 'Array', autodiff_d11_da: 'Array', finite_difference_d11_da: 'Array', autodiff_d33_da: 'Array', finite_difference_d33_da: 'Array', autodiff_d11_der: 'Array', finite_difference_d11_der: 'Array', autodiff_d33_der: 'Array', finite_difference_d33_der: 'Array')

class ntx.autodiff.InverseProblemResult(amplitude_history: 'Array', gradient_history: 'Array', loss_history: 'Array', nu_hat: 'Array', target_response: 'Array', initial_response: 'Array', fitted_response: 'Array', sensitivity: 'Array', inferred_amplitude: 'Array', target_amplitude: 'Array')

class ntx.autodiff.NeopaxProfileAutodiffResult(parameter_history: 'Array', loss_history: 'Array', rho: 'Array', target_er_profile: 'Array', fitted_er_profile: 'Array', target_d11_profile: 'Array', fitted_d11_profile: 'Array', target_d33_profile: 'Array', fitted_d33_profile: 'Array', sensitivity_matrix: 'Array', nu_value: 'Array')

class ntx.autodiff.NeopaxProfileUncertaintyResult(rho: 'Array', fitted_er_profile: 'Array', fitted_d33_profile: 'Array', target_d33_profile: 'Array', sensitivity_matrix: 'Array', parameter_covariance: 'Array', parameter_std: 'Array', parameter_correlation: 'Array', linearized_d33_std: 'Array', monte_carlo_d33_mean: 'Array', monte_carlo_d33_std: 'Array', monte_carlo_d33_quantile_low: 'Array', monte_carlo_d33_quantile_high: 'Array', fisher_matrix: 'Array', fisher_eigenvalues: 'Array', hessian_vector_probe: 'Array', gauss_newton_vector_probe: 'Array', hessian_probe_relative_error: 'Array', sample_count: 'Array')

class ntx.autodiff.RobustBootstrapOptimizationResult(scale_history: 'Array', gradient_history: 'Array', objective_history: 'Array', scale_grid: 'Array', deterministic_objective_landscape: 'Array', robust_objective_landscape: 'Array', rho: 'Array', baseline_scale: 'Array', optimized_scale: 'Array', baseline_current_profile: 'Array', optimized_current_profile: 'Array', optimized_current_mean: 'Array', optimized_current_std: 'Array', optimized_current_quantile_low: 'Array', optimized_current_quantile_high: 'Array', harmonic_m: 'Array', harmonic_n: 'Array', harmonic_reference_value: 'Array', nu_value: 'Array', uncertainty_sigma: 'Array', risk_aversion: 'Array')

ntx.autodiff.example_bootstrap_current_optimization(surfaces: tuple, *, rho: Array, nu_v: Array, Es: Array, Er: Array, drds: Array, grid, a_b: float = 1.0, nu_index: int = 1, learning_rate: float = 0.2, steps: int = 48, regularization: float = 0.005, serial_seconds: float = 0.0, parallel_seconds: float = 0.0) → BootstrapOptimizationResult: Optimize a reduced bootstrap-current response through one geometry control.

ntx.autodiff.example_bootstrap_current_robust_optimization(surfaces: tuple, *, rho: Array, nu_v: Array, Es: Array, Er: Array, drds: Array, grid, a_b: float = 1.0, nu_index: int = 1, learning_rate: float = 0.2, steps: int = 40, regularization: float = 0.005, uncertainty_sigma: float = 0.06, risk_aversion: float = 0.35, scale_grid_size: int = 29, quadrature_order: int = 5) → RobustBootstrapOptimizationResult: Optimize a reduced bootstrap-current response under control uncertainty.

ntx.autodiff.example_derivative_audit(*, grid: GridSpec | None = None, coefficient_index: int = 1, amplitude_value: float = 0.085, nu_hat: Array | None = None, er_hat_scan: Array | None = None, er_reference: float = 0.001, nu_reference: float = 0.0003, fd_step_amplitude: float = 0.0001, fd_step_er: float = 1e-05) → DerivativeAuditResult: Compare direct JAX gradients against finite differences.

ntx.autodiff.example_inverse_problem(*, grid: GridSpec | None = None, nu_hat: Array | None = None, er_hat: float = 0.001, target_amplitude: float = 0.085, initial_amplitude: float = 0.03, learning_rate: float = 0.5, steps: int = 24, coefficient_index: int = 1) → InverseProblemResult: Recover one Boozer-harmonic amplitude from synthetic transport data.

ntx.autodiff.example_neopax_profile_autodiff(surfaces: tuple, *, rho: Array, nu_v: Array, Es: Array, Er: Array, drds: Array, grid, a_b: float = 1.0, nu_index: int = 1, learning_rate: float = 0.25, steps: int = 32, use_neopax_package: bool = False, jacobian_chunk_size: int | str | None = None) → NeopaxProfileAutodiffResult: Infer an electric-field profile with optional chunked sensitivities.

ntx.autodiff.example_neopax_profile_uncertainty(surfaces: tuple, *, rho: Array, nu_v: Array, Es: Array, Er: Array, drds: Array, grid: GridSpec, a_b: float = 1.0, nu_index: int = 1, learning_rate: float = 0.25, steps: int = 32, parameter_std: Array | None = None, parameter_correlation: float = -0.35, target_params: Array | None = None, initial_params: Array | None = None, hessian_probe: Array | None = None, monte_carlo_samples: int = 64, random_seed: int = 0, jacobian_chunk_size: int | str | None = None) → NeopaxProfileUncertaintyResult

Compare linearized covariance propagation against Monte Carlo on a profile fit.

jacobian_chunk_size controls both the profile sensitivity and local residual Jacobians. Use it for large profile bases after measuring the runtime-memory tradeoff on the target workload.

Input And Output

Input/output helpers for Boozer-surface files.

ntx.io.load_boozer_modes_csv(path: str | Path, *, nfp: int, iota: float, psi_p: float, b_theta: float, b_zeta: float) → BoozerSurface: Load columns m,n,b_cos[,b_sin] from a CSV or whitespace-delimited file.

ntx.io.load_dkes_surface(path: str | Path) → BoozerSurface

Load a DKES-format ddkes2.data flux surface.

The scalar geometry is stored in the datain namelist and the Fourier entries are interpreted as borbi(n, m) = B_mn.

ntx.io.load_magnetic_configuration_surface(path: str | Path) → BoozerSurface: Load a text magnetic-configuration file with m n B_mn Fourier rows.

ntx.io.load_vmec_surface(path: str | Path, *, psi_n: float, vmec_radial_option: int = 0, vmec_nyquist_option: int = 1, vmec_mode_convention: str = 'reduced', min_bmn_to_load: float = 0.0) → VmecSurface

Load one VMEC flux surface from a wout_*.nc file.

NTX now sources the VMEC data through vmex rather than through a separate local netCDF parser.

ntx.io.write_result_jsonable(result) → dict[str, float]: Convert a result to a small JSON-serializable mapping.

Validation

Validation registries and benchmark metadata for NTX.

class ntx.validation.AngularOversamplingAudit(theta_nyquist_floor: int, zeta_nyquist_floor: int, n_xi: int, coefficient_atol: float, recommended_oversampling: float, points: tuple[AngularOversamplingPoint, ...])

Oversampling ladder relative to its finest angular reference grid.

property recommended_point: AngularOversamplingPoint: Return the first measured point at or above the recommendation.

class ntx.validation.AngularOversamplingPoint(requested_ratio: float, n_theta: int, n_zeta: int, theta_oversampling: float, zeta_oversampling: float, coefficients: tuple[float, float, float], relative_errors: tuple[float, float, float], max_relative_error: float, schur_residual_l2: float, prepare_seconds: float, lowering_seconds: float, compilation_seconds: float, first_execution_seconds: float, warm_execution_seconds: float, temporary_size_bytes: int | None): One angular grid, coefficient error, and compiled-cost measurement.

class ntx.validation.BenchmarkEntry(id: str, lane: Literal['monoenergetic', 'bootstrap-current', 'integrated-workflow', 'autodiff', 'profile-workflow', 'performance', 'geometry-breadth'], maturity: Literal['positive-gate', 'stress-gate', 'software-gate', 'planned-lane'], title: str, claim_scope: str, literature_anchors: tuple[str, ...], scripts: tuple[str, ...], tests: tuple[str, ...], artifacts: tuple[str, ...], manuscript_figures: tuple[str, ...], docs: tuple[str, ...], open_work: tuple[str, ...] = ()): A maintained map from a research claim to code, tests, and artifacts.

class ntx.validation.BenchmarkEvaluation(entry: 'BenchmarkEntry', path_status: 'tuple[BenchmarkPathStatus, ...]')

class ntx.validation.BenchmarkPathStatus(kind: "Literal['script', 'test', 'artifact', 'doc']", path: 'str', exists: 'bool')

class ntx.validation.PhysicsGate(name: 'str', category: 'GateCategory', metric: 'str', relation: 'GateRelation', threshold: 'float | None', source: 'str', rationale: 'str')

class ntx.validation.PhysicsGateResult(gate: 'PhysicsGate', value: 'float | None', status: 'GateStatus', details: 'str' = '')

ntx.validation.audit_angular_oversampling(surface: BoozerSurface | VmecSurface, case: MonoenergeticCase, *, ratios: tuple[float, ...] = (1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 2.5), n_xi: int = 16, coefficient_atol: float = 1e-12, recommended_oversampling: float = 2.25, repeats: int = 2) → AngularOversamplingAudit

Measure angular-grid error and compiled cost against the finest grid.

The ratio multiplies the odd Nyquist floor independently in each angle. The final ratio is the numerical reference, not an analytical solution. Research promotion still requires successive-grid convergence and an independent-code comparison where one is available.

ntx.validation.benchmark_matrix() → tuple[BenchmarkEntry, ...]: Return the maintained NTX benchmark matrix.

ntx.validation.evaluate_artifact_gates(root: Path) → list[PhysicsGateResult]

Evaluate committed artifact-backed gates below a repository root.

Missing artifacts produce missing results rather than exceptions so validation reports can distinguish absent evidence from failed thresholds.

ntx.validation.physics_gate_registry() → tuple[PhysicsGate, ...]: Return analytical and artifact-backed physics gates in stable order.