booz_xform_jax

Install from PyPI:

pip install booz_xform_jax

booz_xform_jax is a pure-Python, JAX-based implementation of the Boozer coordinate transformation used to post-process VMEC equilibria. The package keeps the familiar xbooz_xform workflow for command-line use while also providing a differentiable Python and JAX API for optimization and analysis pipelines.

What This Package Covers

• VMEC wout ingestion through booz_xform_jax.Booz_xform.read_wout() and booz_xform_jax.Booz_xform.read_wout_data().

• STELLOPT-style in_booz.* command-line execution through booz_xform_jax, xbooz_xform, and xbooz_xform_jax.

• Boozer-spectrum computation, boozmn NetCDF writing, and round-trip file reading.

• JAX-native functional transforms for JIT compilation, batching, and differentiation.

• Plotting helpers for surfplots, symmetry plots, mode plots, and wireframe visualizations.

The documentation below combines project-specific API details with the established BOOZ_XFORM conventions used in STELLOPT and the modern HiddenSymmetries implementation.

User Guide

• Quickstart
  • Command-Line Workflow
  • Object-Oriented Python API
  • Functional JAX API
  • Examples
• Theory And Derivations
  • Overview
  • Coordinate Conventions
  • Derivation As Implemented Here
  • Magnetic-Field Component Conventions
  • Resolution Guidance
  • Related References
• Inputs And Outputs
  • Accepted Inputs
  • Array Layout And Surface Handling
  • CLI Input File Format
  • Generated Outputs
  • boozmn Variable Mapping
  • Compatibility Notes
• Numerics And Performance
  • Numerical Strategy
  • Precision
  • Grid Construction
  • Vectorized Versus Streamed Fourier Accumulation
  • JIT Compilation And Differentiation
  • Surface Resolution Guidance
  • Plotting And Diagnostics
  • Platform Considerations
• JAX Geometry And Jacobian Plan
  • Purpose
  • Physics Contract
  • Milestones
  • Downstream FAX Contract
  • Implementation Notes
  • References
• Examples
  • Bundled Example Scripts
  • CLI Example
  • Python Example
  • Using The Plot Helpers
• STELLOPT Compatibility
  • Reference Material
  • Theory Conventions Used In STELLOPT
  • Legacy Input File Format
  • Execution
  • Output Format
  • Visualization Notes
  • Compilation And Packaging
  • Why This Matters
• API And Source Reference
  • High-Level Class
  • Functional JAX API
  • CLI Driver
  • VMEC Input Utilities
  • boozmn File I/O
  • Plotting Utilities
• Citations
  • Primary Literature
  • Software And Documentation References
  • Suggested Citation Practice