DESC Stellarator Code Suite

DESC (Stellarator Optimization Package) solves for and optimizes 3D MHD equilibria using pseudo-spectral numerical methods and automatic differentiation.

Key capabilities include:

• Equilibrium Computation: Fixed boundary, advanced (continuation/perturbation), and free boundary (vacuum or finite beta with external fields).
• Optimization: Basic and advanced optimization including constraints, precise quasi-symmetry, and constrained optimization.
• Specialized Constraints: Near-axis constraints (loading solutions from QSC/QIC) and coil optimization (second-stage optimization of magnetic coils).
• Data Handling: Analysis, plotting, and saving outputs to VMEC format.

DESC uses JAX for acceleration. Choosing between jax.numpy (jnp) and numpy (np) depends on the intended use case:

• Use jnp if:
  • The function is used for optimization (part of an objective function).
  • The function needs to be JIT compiled (use jnp with jit). Note: It is best to jit only the outermost function.
  • The function needs to be differentiated (use jnp with jacfwd, jacrev, or grad).
  • You are unsure (in doubt, jnp is the safe bet).
• Use np if:
  • You are certain the function is only used during initialization or post-processing (e.g., plotting).

Warning: jax.numpy is almost a drop-in replacement for numpy, but there are subtle and important differences.

The external/ directory contains objectives and utilities that depend on software outside of the DESC repository.

Important Considerations:

• Maintenance: These utilities are not routinely tested and may not be compatible with all versions of DESC. The DESC team does not guarantee regular maintenance or documentation for these external components.
• Availability: External codes are not included with DESC. You must obtain access to them independently, and they may require specific licenses.
• Risk: Use these utilities at your own risk.

Last Tested Compatibility:

• _terpsichore.py: Tested on Sept 8, 2025, with DESC v0.15.0 and a TERPSICHORE version compiled on Perlmutter (Jan 2025).
• neo.py: Tested on Oct 23, 2025, with the version of NEO on Stellar.
• paraview.py: Tested on Nov 24, 2025, with DESC v0.16.0, Paraview v5.13.3, and pyvista v0.46.4.

The jac_chunk_size controls the trade-off between memory usage and computation time:

• Memory usage $\approx m_0 + m_1 \times \text{jac_chunk_size}$
• Computation time $\approx t_0 + t_1 / \text{jac_chunk_size}$

Strategy:

1. Check the optimization output for the Number of parameters and Number of objectives to understand the total Jacobian shape.
2. Set jac_chunk_size to a value smaller than the Number of parameters that allows the Jacobian to fit in memory.
3. If using jac_chunk_size="auto", DESC will use a heuristic to estimate a size that fits in memory.
4. Adjust the value based on the performance of your optimization.

An optimizer wrapper must return a scipy.optimize.OptimizeResult object. Key attributes include:

• x: The solution array.
• success: A Boolean flag indicating if the optimizer exited successfully.
• message: A string describing the cause of termination.
• fun: The function value at the solution.
• nit: Number of iterations.
• nfev: Number of function evaluations.
• njev: Number of gradient evaluations.

You can control how JAX/XLA allocates GPU memory using environment variables. This is useful for resolving Out of Memory (OOM) errors or for memory profiling.

• Pre-allocate a specific fraction of memory: Use XLA_PYTHON_CLIENT_MEM_FRACTION to allow a single process to claim a larger portion of the GPU (e.g., 0.90 for 90%). This can help resolve OOM errors.
• Allocate memory as needed: Use XLA_PYTHON_CLIENT_ALLOCATOR="platform" to prevent pre-allocation. This is recommended for debugging or memory profiling, though it may incur a slight performance penalty due to frequent allocation/de-allocation.

import os

# Option 1: Pre-allocate 90% of GPU memory
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.90"

# Option 2: Allocate memory on-demand (for profiling/debugging)
os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"

To regenerate the figures (Figures 1-6) shown in the paper, use the plotter.py script.

Prerequisites: You must checkout version v0.7.2 of the DESC code.

Outputs: All generated plots are saved in the data sub-folder:

• f_B.png (Figure 1)
• f_C.png (Figure 2)
• f_T.png (Figure 3)
• boundaries.png (Figure 4)
• Booz.png (Figure 5)
• errors.png (Figure 6)

python plotter.py

To generate the figures used in the paper "Magnetic Fields with General Omnigenity", run the plotter.py script. This script is compatible with the latest commit of the repository.

python plotter.py

If you are making your first code contribution, look for issues labeled with:

• good first issue: Tasks requiring only a few lines of code and a test or two.
• help wanted: More involved tasks for contributors beyond the absolute beginner level.

To enable saving 3D plots within DESC, you must install the kaleido package via pip after installing DESC.

pip install kaleido

Bugs are tracked via GitHub issues. To submit a high-quality report, include:

• A clear, descriptive title.
• Exact, detailed steps to reproduce the problem (explain how you did it).
• Specific examples using Markdown code blocks.
• The observed behavior vs. the expected behavior.
• Plots of incorrect results if applicable.
• For crashes: the Python stack trace and full error message (use a code block, file attachment, or a Gist link).

Provide context regarding:

• Whether the problem is new (e.g., after an update) or persistent.
• If it is reliably reproducible.
• Your environment: DESC version, JAX version, OS name/version, hardware (CPU/GPU/RAM), and whether you are running locally, on a cluster, or in a VM.