VMEC++ (Variational Moments Equilibrium Code)

repository·main·Indexed 42 days ago

https://github.com/proximafusion/vmecpp

A C++ reimplementation of the Variational Moments Equilibrium Code (VMEC), providing a free-boundary ideal-MHD equilibrium solver for stellarators and tokamaks. Designed to be Python-friendly and faster than its Fortran predecessors, it supports JSON configurations, OpenMP multi-threading, and integration with the SIMSOPT optimization framework.

Tokens
20.5K
Snippets
62
Records
87
Agent score
53%

What's inside VMEC++

  1. What is VMEC++

    main

    VMEC++ is a C++ reimplementation of the Variational Moments Equilibrium Code (VMEC), designed as a free-boundary ideal-MHD equilibrium solver for stellarators and tokamaks. It is designed to be Python-friendly and is suitable for both local execution on laptops and integration into large-scale stellarator optimization pipelines.

    Key advantages over the original Fortran versions (like PARVMEC) include:

    • Zero-crash policy: Errors are reported via standard Python exceptions rather than crashing.
    • Hot-restarting: Ability to restart a run from a previously converged state.
    • Flexible Input: Supports classic INDATA format, JSON files, or programmatic construction via Python objects.
    • Performance: Typically matches or exceeds the speed of its predecessors.
    • Documentation: Includes detailed documentation of its internal numerics.
  2. Explore VMEC++ example input and output files

    main

    The examples/data/ directory contains several test cases for verifying VMEC++ functionality and experimenting with different equilibrium types. These examples include both legacy Fortran namelist files and modern JSON configurations used by VMEC++.

    Available Test Cases

    • Stellarator (CTH-like): cth_like_fixed_bdy.json
      • input.cth_like_fixed_bdy: Original Fortran namelist input.
      • cth_like_fixed_bdy.json: JSON input derived via indata2json.
      • wout_cth_like_fixed_bdy.nc: NetCDF output file (produced via PARVMEC) used to test loading wout files into VMEC++.
    • Quasi-helical (SIMSOPT): input.nfp4_QH_warm_start
      • Designed specifically for use with the SIMSOPT framework.
    • Axisymmetric Tokamak (Solov'ev): solovev
      • input.solovev: Original Fortran namelist input.
      • solovev.json: JSON input derived via indata2json.
    • Stellarator (W7-X): w7x
      • input.w7x: Original Fortran namelist input.
      • w7x.json: JSON input derived via indata2json.
  3. Key differences and limitations of VMEC++

    main

    While VMEC++ implements the iteration algorithm of Fortran VMEC 8.52, there are several functional differences and limitations to be aware of:

    Key Differences

    • Parallelization: Uses OpenMP (multi-threading) instead of MPI. It can utilize shared caches between threads but cannot parallelize across multiple nodes.
    • Error Handling: Reports issues via standard Python exceptions with a zero-crash policy.
    • Input Formats: Supports classic INDATA formats and JSON files; allows programmatic construction via Python.
    • Performance: Uses FFT kernels optimized for small mode numbers (via FFTX) for a 10-20% speedup over DFT.
    • Stability: Features reduced force spikes between multigrid stages and stable recurrence for Neumann kernel integrals.

    Current Limitations

    • Missing Logic: lgiveup/fgiveup (early termination), lbsubs (output computation), and lrfp (RFP-specific physics) are not yet implemented.
    • Profile Parameterizations: sum_cossq_s, sum_cossq_sqrts, and sum_cossq_s_free are not fully implemented.
    • Preconditioning: 2D preconditioning using block-tridiagonal solvers (precon_type and prec2d_threshold) is not implemented.
    • Convergence Behavior: VMEC++ only computes output quantities if the run converged (unless return_outputs_even_if_not_converged is set).
    • File Handling: Unlike the Fortran version, VMEC++ does not support 'extension' style input filenames; it requires a valid filename or full path to an existing file.
  4. Understanding ThreadSanitizer (TSan) and OpenMP compatibility

    main

    ThreadSanitizer (TSan) is used to detect data races in multi-threaded programs. However, standard TSan implementations often lack awareness of OpenMP constructs like #pragma omp barrier or #pragma omp critical. This lack of awareness leads to a high volume of false positives, where TSan reports data races in code that is actually correctly synchronized via OpenMP.

    Historically, users attempted to fix this by building libomp with -DLIBOMP_TSAN_SUPPORT=1 (now deprecated). The modern solution is to use Archer, which has been integrated into the LLVM project. Archer is automatically built with and loaded by the OpenMP runtime in modern LLVM distributions.

    If you are using standard Ubuntu 22.04 LTS with clang-14, be aware that it may not include working OpenMP-capable TSan support, leading to false positives even in correct code.

  5. Compare Combined vs Product Fourier Basis representations

    main

    VMEC++ maintains two distinct Fourier basis representations depending on the context:

    1. Combined Basis (External Interface)

    • Mathematical Form: $\cos(m\theta - n\zeta)$ and $\sin(m\theta - n\zeta)$
    • Usage: Used for wout files, Python API I/O, SIMSOPT compatibility, and general research/literature exchange.
    • Storage: Linear indexing by mode number mn.

    2. Product Basis (Internal Computation)

    • Mathematical Form: $\cos(m\theta)\cos(n\zeta)$, $\sin(m\theta)\sin(n\zeta)$, etc.
    • Usage: Used for internal force calculations, energy evaluations, and real-space transformations via DFTs.
    • Storage: 2D indexing by separate $m, n$ indices (layout varies between m-major and n-major depending on the class).
  6. Understand VMEC++ Fourier Basis implementation (DFT vs FFT)

    main

    VMEC++ uses Discrete Fourier Transforms (DFTs) with pre-computed basis arrays rather than Fast Fourier Transforms (FFTs). This design choice is driven by the need to support physics-specific angular discretizations (non-power-of-2 grids), incorporate integration weights directly into basis functions, and pre-apply normalization scaling factors.

    Key implementation details include:

    • Pre-computed Basis Arrays: Arrays like cosmu, sinmu, cosnv, and sinnv store pre-scaled trigonometric values.
    • Integration Weights: Surface integrals use pre-weighted arrays (cosmui, sinmui) to avoid runtime weight multiplication.
    • Performance: While DFTs are $O(N_m \cdot N_n)$, specialized FFT kernels can be compiled for specific resolutions via fftx for a 10-20% speedup.
  7. How vmecpp uses FFTX/SPIRAL codelets

    main

    vmecpp uses pre-generated FFTX/SPIRAL kernels as its toroidal FFT backend. These kernels are vendored to avoid requiring the SPIRAL toolchain (OCaml, GAP, etc.) for standard users.

    vmecpp_core links these as a static library (fftx_codelets) and uses fftx_iprdftbat_cpu_Tuple and fftx_prdftbat_cpu_Tuple to obtain {init, run, destroy} function pointers for specific shapes defined by (fftlen, nbatch, read_stride=APar=0, write_stride=APar=0).

    Supported Shapes

    The kernels cover a specific set of shapes based on ntor and mpol. If your configuration falls outside this table, vmecpp will fall back to a partial-DFT path.

    ntor (-> nZeta)mpol (-> nbatch = 12*mpol)
    6, 8, 10, 12, 14, 16, 186, 8, 10, 12, 14, 16, 18
    nZeta ∈ {16, 20, 24, 28, 32, 36, 40}nbatch ∈ {72, 96, 120, 144, 168, 192, 216}

    Monitoring Backend Selection

    You can verify which backend is being used by checking the construction-time logs:

    • FFTX active: INFO Toroidal FFT backend: FFTX (nZeta=36, 12*mpol=144).
    • Fallback active: WARN Toroidal FFT backend: partial-DFT fallback (no FFTX kernel for nZeta=50, 12*mpol=120). ...
  8. Enforce ASCII-Only Source Code

    main
    All VMEC++ source code must be ASCII-only. When documenting mathematical symbols or expressions, use LaTeX notation within comments rather than non-ASCII Unicode characters. This ensures compatibility across different development environments and toolchains.
  9. Understand the two Fourier representations in VMEC++

    main

    VMEC++ distinguishes between an Internal Product Basis used for computational efficiency and an External Combined Basis used for researcher interfaces and traditional VMEC formats.

    Internal Product Basis

    Used for internal computations to enable separable DFT operations. It uses trigonometric product pairs. The suffixes indicate the parity of the trigonometric functions:

    • _cc: cos(mθ) * cos(nζ) (even-even trig parity)
    • _ss: sin(mθ) * sin(nζ) (odd-odd trig parity)
    • _sc: sin(mθ) * cos(nζ) (odd-even trig parity)
    • _cs: cos(mθ) * sin(nζ) (even-odd trig parity)

    External Combined Basis

    Used for the researcher interface (e.g., wout files). It uses the combined basis cos(mθ - nζ) or sin(mθ - nζ).

    • rmnc: External combined basis (cos terms).
    • zmns: External combined basis (sin terms).

    Note: This trigonometric parity is distinct from poloidal mode number parity (kEvenParity/kOddParity).

    // Internal product basis (computational)
    std::vector<double> rmncc_; // cos(mθ) * cos(nζ)
    std::vector<double> rmnss_; // sin(mθ) * sin(nζ)
    
    // External combined basis (interface)
    RowMatrixXd rmnc; // cos(mθ - nζ)
    RowMatrixXd zmns; // sin(mθ - nζ)
  10. Use hot restart to accelerate convergence

    main

    VMEC++ supports 'hot restarting', which allows you to use the output of a previously converged run as the initial state for a new run. This is highly effective when running simulations on configurations that are very similar to a known converged equilibrium, as it dramatically reduces the number of iterations required.

    Important Note for Hot Restarts: When performing a hot restart, you must adjust your multigrid settings. Hot-restarted runs only allow a single multigrid step. You should truncate your ns_array, ftol_array, and niter_array to contain only the final step (the finest grid).

    import vmecpp
    
    # 1. Load input and run initial simulation
    vmec_input = vmecpp.VmecInput.from_file("w7x.json")
    vmec_output = vmecpp.run(vmec_input)
    
    # 2. Modify the input (e.g., perturb the plasma boundary)
    vmec_input.rbc[0, 0] *= 0.8
    vmec_input.rbc[1, 0] *= 1.2
    
    # 3. Adjust multigrid steps for hot restart (single step only)
    vmec_input.ns_array = vmec_input.ns_array[-1:]
    vmec_input.ftol_array = vmec_input.ftol_array[-1:]
    vmec_input.niter_array = vmec_input.niter_array[-1:]
    
    # 4. Run with hot restart using the previous output
    hot_restarted_output = vmecpp.run(vmec_input, restart_from=vmec_output)