Metadata-Version: 2.4
Name: vmex
Version: 0.9.0
Summary: JAX implementation of VMEC2000 with differentiable fixed-boundary and branch-local free-boundary research paths.
Author-email: Rogerio Jorge <rogerio.jorge@wisc.edu>
License-Expression: MIT
Project-URL: Homepage, https://github.com/uwplasma/vmex
Project-URL: Documentation, https://vmex.readthedocs.io/en/latest/
Project-URL: Repository, https://github.com/uwplasma/vmex
Project-URL: Issues, https://github.com/uwplasma/vmex/issues
Project-URL: Changelog, https://github.com/uwplasma/vmex/releases
Keywords: vmec,stellarator,mhd,jax,optimization
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: jax>=0.9.2
Requires-Dist: jaxlib>=0.9.2
Requires-Dist: scipy>=1.16
Requires-Dist: netCDF4
Requires-Dist: h5py
Requires-Dist: matplotlib
Requires-Dist: packaging
Requires-Dist: filelock
Requires-Dist: booz_xform_jax>=0.1.1
Requires-Dist: solvax>=0.20.0
Provides-Extra: optimizers
Requires-Dist: jaxopt; extra == "optimizers"
Requires-Dist: optax; extra == "optimizers"
Provides-Extra: neoclassical
Requires-Dist: neo-jax>=1.0.2; extra == "neoclassical"
Provides-Extra: freeb
Requires-Dist: virtual-casing-jax>=0.0.5; extra == "freeb"
Provides-Extra: turbulence
Requires-Dist: gkx>=1.7.1; extra == "turbulence"
Requires-Dist: jax>=0.10.1; extra == "turbulence"
Requires-Dist: jaxlib>=0.10.1; extra == "turbulence"
Provides-Extra: coils
Requires-Dist: essos>=0.16; extra == "coils"
Provides-Extra: docs
Requires-Dist: sphinx; extra == "docs"
Requires-Dist: sphinx-copybutton; extra == "docs"
Requires-Dist: furo; extra == "docs"
Requires-Dist: myst-parser; extra == "docs"
Requires-Dist: sphinx-design; extra == "docs"
Requires-Dist: sphinxext-rediraffe; extra == "docs"
Provides-Extra: dev
Requires-Dist: build; extra == "dev"
Requires-Dist: furo; extra == "dev"
Requires-Dist: mypy; extra == "dev"
Requires-Dist: myst-parser; extra == "dev"
Requires-Dist: pytest; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: pytest-xdist; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Requires-Dist: sphinx; extra == "dev"
Requires-Dist: sphinx-copybutton; extra == "dev"
Requires-Dist: sphinx-design; extra == "dev"
Requires-Dist: sphinxext-rediraffe; extra == "dev"
Requires-Dist: types-setuptools; extra == "dev"
Dynamic: license-file

# VMEX

[![PyPI version](https://img.shields.io/pypi/v/vmex.svg)](https://pypi.org/project/vmex/)
[![Python](https://img.shields.io/badge/python-3.11%20%7C%203.12-blue.svg)](pyproject.toml)
[![License](https://img.shields.io/github/license/uwplasma/vmex)](LICENSE)
[![CI](https://img.shields.io/github/actions/workflow/status/uwplasma/vmex/ci.yml?branch=main&label=ci)](https://github.com/uwplasma/vmex/actions/workflows/ci.yml)
[![Coverage](https://codecov.io/gh/uwplasma/vmex/branch/main/graph/badge.svg)](https://codecov.io/gh/uwplasma/vmex)
[![Docs](https://img.shields.io/readthedocs/vmex/latest?label=docs)](https://vmex.readthedocs.io/en/latest/)

VMEX computes stellarator and tokamak ideal-MHD equilibria in JAX. It reads
VMEC input files, solves fixed- and free-boundary problems, and writes standard
`wout_*.nc` files. Implicit derivatives connect equilibria to boundary, profile,
and coil optimization.

- **Use existing workflows:** VMEC input/output, multigrid continuation and hot restart.
- **Design with gradients:** SciPy or JAX optimizers, scalar adjoints and residual Jacobians.
- **Inspect the physics:** Boozer transforms, magnetic fields and spatial derivatives,
  quasisymmetry, quasi-isodynamicity and stability diagnostics.
- **Choose the hardware:** CPU or GPU equilibrium solves (optimization gradients default to CPU),
  reusable compilation and independent-case ensembles.
- **Connect coils:** ESSOS fields, NESTOR free boundary and finite-beta exterior fields.

![VMEX equilibria and diagnostics](docs/_static/figures/readme_equilibrium_showcase.webp)

The [capability reference](https://vmex.readthedocs.io/en/latest/reference/capabilities.html)
defines supported models and validation limits. VMEX's toroidal equilibrium
model assumes nested flux surfaces; mirror and high-order polishing features
have narrower validation scopes.

## Install

```console
pip install vmex
vmex --doctor
vmex --test
```

Python 3.11–3.12 is tested. CPU JAX is included; for GPUs follow the
[JAX installation guide](https://docs.jax.dev/en/latest/installation.html) and
[VMEX GPU guide](https://vmex.readthedocs.io/en/latest/howto/run-on-gpu.html).
Optional extras include `vmex[coils]` (ESSOS), `vmex[freeb]` (virtual casing),
`vmex[neoclassical]` (effective ripple), and `vmex[optimizers]` (JAXopt/Optax).
See [installation](https://vmex.readthedocs.io/en/latest/installation.html) for dependencies.

## Solve, plot and restart

With your own VMEC input file:

```console
vmex input.my_case --plot
vmex --plot wout_my_case.nc
vmex --booz wout_my_case.nc
vmex input.nearby --restart wout_my_case.nc
```

`--plot` writes six PNGs beside the input or in `--outdir`: the summary below, flux-surface cross-sections,
`|B|` in VMEC angles, radial profiles, Mercier stability and the 3-D LCFS. The summary adds Boozer `|B|`, a `J` map,
`D_R` and the DESC-normalized force balance (effective ripple needs `vmex[neoclassical]`). The QA and QI panels are
`vmex examples/data/input.nfp2_QA_finite_beta --plot` and `vmex examples/data/input.nfp4_QI_finite_beta --plot`.

![vmex --plot summary of the bundled finite-beta NFP=2 QA equilibrium](docs/_static/figures/readme_diagnostics_qa.webp)
![vmex --plot summary of the bundled finite-beta NFP=4 QI equilibrium](docs/_static/figures/readme_diagnostics_summary.webp)

VMEX follows the deck's `NS_ARRAY`, `FTOL_ARRAY` and `NITER_ARRAY`. In Python:

```python
import vmex as vj

inp = vj.VmecInput.from_file("input.my_case")
result = vj.solve_multigrid(inp, verbose=True)
print(result.converged, result.fsqr, result.fsqz, result.fsql)
wout = vj.wout_from_state(
    inp=inp, state=result.state, fsqr=result.fsqr, fsqz=result.fsqz,
    fsql=result.fsql, niter=result.iterations, converged=result.converged)
vj.write_wout("wout_my_case.nc", wout)
```

Pass `initial_state=result.state` for a nearby Python solve, or `restart_from=`
for a saved WOUT. [Restart](https://vmex.readthedocs.io/en/latest/howto/restart-from-previous-run.html)
and [CLI](https://vmex.readthedocs.io/en/latest/reference/cli.html) guides cover
resolution changes, devices, profiles and output controls.

`vmex equilibrium.h5` reads DESC text inputs and HDF5/pickle outputs without installing DESC, using the final stage or equilibrium; it writes `input.equilibrium` and solves it to write `wout_equilibrium.nc`.
`--desc-tol 0` retains all boundary modes. The default 1% boundary tolerance does not guarantee magnetic-field accuracy. WOUT iota has the opposite sign to DESC.

## Differentiate and optimize

Compute a boundary/profile derivative without differentiating through every forward iteration:

```python
import jax
from vmex.core import implicit

params = implicit.params_from_input(inp)
gradient = jax.grad(lambda p: implicit.run(inp, p).aspect)(params)
```

For a simple boundary optimization, supply objectives and let SciPy choose the
steps. Each tuple is `(function, target, cost_weight)`:

```python
from scipy.optimize import least_squares
from vmex import optimize as opt

problem = opt.VmecProblem.from_tuples(
    inp, [(opt.aspect_ratio, 4.0, 1.0)], max_mode=1, use_ess=True)
fit = least_squares(
    problem.residual, problem.x0, jac=problem.residual_jac,
    x_scale=problem.scales, max_nfev=20)
problem.input_from_x(fit.x).to_indata("input.optimized")
equilibrium = problem.equilibrium_from_x(fit.x)
vj.write_wout("wout_optimized.nc", equilibrium.wout)
```

`problem.value_and_grad` supplies scalar objectives to BFGS/L-BFGS-B;
`problem.jax_value_and_grad` supports JAX optimizers. `problem.evaluate(x)`
reports solve effort and derivative status. Implicit derivatives describe the
chosen discrete equilibrium equations: small linear residuals alone do not
establish continuum accuracy. See the
[gradient tutorial](https://vmex.readthedocs.io/en/latest/tutorials/first-gradient.html)
and [optimization guide](https://vmex.readthedocs.io/en/latest/howto/optimize-a-boundary.html)
for convergence checks, constraints, scaling and finite-difference verification.

## What you can build

Every figure below is regenerated by the script named beside it, from decks in this repository.

### Quasisymmetric and quasi-isodynamic boundaries

![QA at nfp 2, QH at nfp 4 and QP at nfp 2](docs/_static/figures/readme_optimization.webp)

Boundary optimization for quasi-axisymmetry, quasi-helical symmetry and quasi-poloidal symmetry, each at its own
field-period count, from `examples/optimization/QA_optimization.py` and its QH and QP siblings.
`python examples/plot_optimized_families.py` draws the panel.

![QI at nfp 1, 2, 3 and 4](docs/_static/figures/readme_qi.webp)

Quasi-isodynamic designs across four field-period counts, showing how the
`|B|` contours close poloidally as `nfp` rises. The QI scripts in
`examples/optimization/` cover vacuum, finite beta, bootstrap-consistent and
maximum-`J` continuation variants, and the same families have scalar-adjoint,
SciPy, JAXopt and Optax drivers.

### Free boundary with coils

![Free-boundary beta ramp and Shafranov shift](docs/_static/figures/readme_essos_beta_scan.webp)

NESTOR free boundary driven by ESSOS coils, ramping beta and tracking the
Shafranov shift: `python examples/free_boundary_essos_coils.py`. Free-boundary
runs also accept an MGRID table, and `VmecExtender` evaluates the exterior
field with the plasma's own virtual-casing contribution, within the distance
limits [below](#fields-coils-and-free-boundary).

### Single-stage plasma and coil design

`examples/optimization/single_stage_optimization.py` adjusts the plasma
boundary and the coils against one weighted objective, solving the equilibrium
implicitly at every step; `single_stage_free_boundary_optimization.py` couples
them through a true free-boundary solve. Both print final plasma and coil
metrics; `single_stage_optimization.py` also states whether it met its
rotational-transform and normal-field targets. A lower weighted penalty with
unmet targets is not a design.

### Open mirrors and stellarator-mirror hybrids

![Fixed-boundary non-axisymmetric mirror](docs/_static/figures/mirror_fixed_boundary_3d.webp)

Fixed-boundary open mirrors, from `examples/mirror/mirror_fixed_boundary_nonaxisymmetric.py`.

![Free-boundary mirror beta scan](docs/_static/figures/mirror_free_boundary_beta_scan.webp)

The free-boundary mirror beta scan, `examples/mirror/mirror_free_boundary_beta_scan.py`,
over the validated 0 to 10 percent beta range.

![Stellarator-mirror hybrid](docs/_static/figures/stellarator_mirror_hybrid.webp)

Periodic stellarator-mirror hybrids,
`examples/mirror/stellarator_mirror_hybrid.py`, with a quasi-isodynamic variant
in `qi_mirror_hybrid_fourier_vs_bspline.py`. Hybrids and anisotropy are
research scopes: see the
[mirror guide](https://vmex.readthedocs.io/en/latest/howto/mirror-machines.html)
for what is validated.

### Agreement with VMEC2000 and VMEC++

![Force-residual traces for VMEX, VMEC2000 and VMEC++](docs/_static/figures/readme_convergence.webp)

The same deck through all three codes, `python benchmarks/make_readme_figures.py --only convergence`.
Cross-code agreement, its norms and its limits are in the [validation record](docs/explanation/validation.md).

### Running the examples

```console
git clone https://github.com/uwplasma/vmex
cd vmex
pip install -e .
vmex examples/data/input.circular_tokamak --plot
python examples/take_gradients.py
```

| Application | Runnable starting point |
|---|---|
| Tokamak or stellarator equilibrium | `vmex examples/data/input.circular_tokamak --plot`; other decks in [examples/data](examples/data/) |
| QA, QH, QP or QI boundary design | [examples/optimization](examples/optimization/), including bootstrap, ballooning, Mercier and maximum-`J` variants |
| Asymmetric boundary design | [stellarator_asymmetry](examples/optimization/stellarator_asymmetry/) vacuum and finite-beta scripts |
| Single-stage plasma and coils | `single_stage_optimization.py`, `single_stage_free_boundary_optimization.py` |
| Fields and spatial derivatives | `python examples/vmex_get_B_gradB.py` |
| ESSOS coils and a free-boundary beta scan | `python examples/free_boundary_essos_coils.py` |
| Finite-beta exterior field lines | `python examples/vmex_fieldline_tracing_finite_beta.py` |
| Effective ripple | `python examples/epsilon_effective.py` |
| Open mirrors | `mirror/mirror_fixed_boundary_nonaxisymmetric.py`, `mirror/mirror_free_boundary_beta_scan.py` |
| Research force-balance polishing | `python examples/force_balance_polishing.py` |

The optimization scripts expose resolutions, objective weights and iteration budgets near the top.
Inspect those settings before a research run; advanced coil examples need the optional dependencies
and versions in the [ESSOS guide](https://vmex.readthedocs.io/en/latest/howto/use-essos-fields-and-coils.html).

## Fields, coils and free boundary

The live equilibrium exposes Cartesian `B()`, `gradB()`, `gradgradB()` and
`gradgradgradB()`, with corresponding VJPs in the originating problem's degrees
of freedom. Use `set_points_xyz(...)` or `set_points_flux(...)` to select
interior evaluation points.

For an exterior field, `vj.VmecExtender.from_file("wout_my_case.nc",
external_field=coils.B)` combines the plasma's virtual-casing contribution with
the supplied coil field. The plasma part is a quadrature over a source grid on
the plasma surface (32 points per field period in each angle by default) whose
error grows rapidly near that surface: evaluate at distances of at least about
twice the toroidal source-grid spacing from the plasma surface. Closer in,
`with_near_surface_continuation` uses a first-order continuation of the
on-surface field. Targets must also stay away from coil filaments, and an MGRID
field has a finite tabulated domain. See the [exterior-field explanation](https://vmex.readthedocs.io/en/latest/explanation/nestor-vacuum.html)
and [field and coil usage](https://vmex.readthedocs.io/en/latest/howto/use-essos-fields-and-coils.html).

![Exterior field lines of a finite-beta QA with coils only and with coils plus plasma](docs/_static/figures/readme_extender_exterior_islands.webp)

Field lines seeded within 5 mm outside a finite-beta QA boundary, in the coil
field alone and with the plasma's field added, stopped 55 mm out where the
continuation ends: `python examples/vmex_fieldline_tracing_finite_beta.py` (needs ESSOS).

Joint boundary/coil optimization and the boundary-Schur adjoint remain advanced
workflows with substantial solve costs. Open mirrors support defined isotropic
fixed/free-boundary cases; the shipped free-boundary 0–10% beta range is the
supported range, while higher beta, anisotropy and periodic hybrids need further
validation. See the [mirror guide](https://vmex.readthedocs.io/en/latest/howto/mirror-machines.html).

## Accuracy and optional polishing

A small VMEC `FSQR/FSQZ/FSQL` means the discrete solve converged. It does not
by itself bound the continuous force error `J × B − ∇p`. Optional spline-based
polishing is disabled by default and remains a research feature: recorded
certified cases are axisymmetric; a generally accurate, affordable 3-D polished
solve is still an open goal.

```console
vmex examples/data/input.shaped_tokamak_pressure_polished --polish auto --plot
```

`AUTO` estimates the Gauss–Newton work against `--polish-budget`; this is an
admission estimate, not an enforced end-to-end timeout. Inspect
`result.polish_report` when using Python. The legacy pointwise `eps_F` metric
is bounded above by 2 by construction and can saturate in vacuum; read the
dimensional and volume-normalized metrics with it.

The [validation record](docs/explanation/validation.md) explains the measured
near-axis improvement, failed 3-D attempts, native DESC comparison and export
errors. The [polishing reference](https://vmex.readthedocs.io/en/latest/explanation/high-order-force-balance.html)
defines the method and certificate. Exported-and-refitted WOUT comparisons
measure reconstruction error as well as solver error.

## Performance and parallel execution

JAX compilation is reused for matching array structures. Measure first-call,
cache-reload and warm costs separately, including refinement and gradients for
optimization. CPU/GPU performance depends on resolution and workload; see
[benchmark evidence](https://vmex.readthedocs.io/en/latest/explanation/validation.html)
and [profiling tools](benchmarks/).

`vj.parallel.solve_ensemble(inputs, workers=None)` distributes independent
cases; `workers=1` provides a serial baseline. Set worker and device budgets
using the [ensemble guide](https://vmex.readthedocs.io/en/latest/howto/parallel-ensembles.html).
Multi-device kernel/AD tests do not yet establish a scalable distributed
nonlinear equilibrium solve.

## Documentation, development and citation

Start with [your first equilibrium](https://vmex.readthedocs.io/en/latest/tutorials/first-equilibrium.html),
then [your first optimization](https://vmex.readthedocs.io/en/latest/tutorials/first-optimization.html).
The [API](https://vmex.readthedocs.io/en/latest/reference/api/basic.html),
[VMEC compatibility](https://vmex.readthedocs.io/en/latest/reference/vmec2000-compatibility.html)
and [troubleshooting](https://vmex.readthedocs.io/en/latest/howto/troubleshoot.html)
pages cover research use.

For development, install `pip install -e ".[dev]"` and run
`python tools/preflight.py --static`; the [test manifest](tests/manifest.json)
defines numerical suites. Report issues with the input deck and `vmex --doctor`
output. See [contributing](CONTRIBUTING.md), [citation metadata](CITATION.cff),
[license](LICENSE) and the [plan/logbook](plan.md). A new release will follow
integration and validation of the agreed priorities.
