driftless-star

driftless-star is an open-source pipeline for stellarator design. It connects five physics stages (equilibrium, Boozer transform, neoclassical transport, turbulence, and profile evolution) and their relevant software into a single reproducible workflow. Given a stellarator boundary shape and initial plasma profiles, the pipeline produces transport-consistent density and temperature profiles along with fusion-power metrics.

Each stage is modular so that implementations can be swapped independently. The pipeline is designed to be closed-loop, so output profiles can be fed back as input for iterative optimization.

See Progress below.

graph LR
    IN["Boundary coefficients<br/>+ profile guesses"] --> S1["<b>Stage 1</b><br/>Equilibrium"]
    S1 -->|"wout_*.nc"| S2["<b>Stage 2</b><br/>Boozer Transform"]
    S2 -->|"boozmn_*.nc"| S3["<b>Stage 3</b><br/>Neoclassical"]
    S2 -->|"geometry"| S4["<b>Stage 4</b><br/>Turbulence"]
    S3 -->|"fluxes"| S5["<b>Stage 5</b><br/>Transport"]
    S4 -->|"heat/particle flux"| S5
    S5 --> OUT["n(r), T(r), E_r(r)<br/>P_fus, Q"]

Loading

Stages 3 and 4 run in parallel. Each stage should eventually be independently swappable (see guide).

Quick Reference

Resource	Location
Pipeline design & contributor workflow	docs/guide.md
Stage I/O specs	docs/stage{N}-{name}/spec.md
MVP I/O reference & Pixi commands	docs/mvp-pipeline.md
Physics equations & I/O contracts (TeX)	stellarator_workflow/
I/O validation methodology	docs/guide.md#io-validation
Coding standards	docs/guide.md#coding-conventions

Pipeline Stages

Stage	Physics	JAX Primary	Alternatives
1. Equilibrium	Ideal-MHD force balance	vmec_jax, DESC	VMEC++
2. Boozer Transform	Coordinate transform	booz_xform_jax	BOOZ_XFORM
3. Neoclassical	Effective ripple, drift-kinetic	NEO_JAX, sfincs_jax	NEO, SFINCS
4. Turbulence	Gyrokinetic equation	GKX	GX, GENE
5. Transport	Profile evolution, power balance	NEOPAX	Trinity3D

Where to Put Code

Phase 1 work goes into the stage spec docs (docs/stage{N}-{name}/spec.md) -- the "TO BE COMPLETED" sections.

Phase 2 adds containers and tests. Stage dependencies are managed through a Pixi workspace under stages/ (stages/pixi.toml + stages/pixi.lock), and a single templated stages/Dockerfile builds all stages via build arguments. The Snakemake orchestration environment lives in a separate root-level Pixi workspace (pixi.toml) so it can be installed on execution nodes without nesting containers. See guide for details.

Workflow

1. Fork the repository and branch from main (e.g., feat/stage1-newsoftware)
2. Work through the relevant phase in the Guide
3. Open a PR from the fork when deliverables are ready and request a review
4. After review and merge, the corresponding item below gets checked off

Progress

Phase 1: Document & Run

Install the primary code, document the API and convergence behavior, write example scripts, set up W&B tracking. Full checklist in the Guide.

• Stage 1 -- Equilibrium
  • vmec_jax
  • DESC
  • VMEC++
• Stage 2 -- Boozer Transform
  • booz_xform_jax
  • BOOZ_XFORM
• Stage 3 -- Neoclassical
  • sfincs_jax
  • NEO_JAX
  • NEO
  • SFINCS
• Stage 4 -- Turbulence
  • GKX
  • GX
  • GENE
• Stage 5 -- Transport
  • NEOPAX
  • Trinity3D

Phase 2: Containerize & Test

Containerize stages and write tests. Full checklist in the Guide.

• Stage 1 -- Equilibrium
  • vmec_jax
  • DESC
  • VMEC++
• Stage 2 -- Boozer Transform
  • booz_xform_jax
  • BOOZ_XFORM
• Stage 3 -- Neoclassical
  • NEO_JAX
  • sfincs_jax
  • NEO
  • SFINCS
• Stage 4 -- Turbulence
  • GKX
  • GX
  • GENE
• Stage 5 -- Transport
  • NEOPAX
  • Trinity3D

Phase 3: Integrate

Snakemake DAG, end-to-end tests, and publishing. Details in the Guide.

• config.yaml + Snakemake DAG
• Swappability patterns (single-stage, multi-stage, end-to-end)
• End-to-end integration tests
• Pipeline-level W&B aggregation
• GHCR image publishing

Usage

A run is a folder under inputs/ holding its run config (config.yaml) and every stage input. A fresh clone ships one ready-to-run example, inputs/quick_run/. common_input.toml in the run folder is the shared transport config read by Stages 3, 4, and 5.

driftless-star iterates toward transport-consistent profiles by chaining forward passes. Each pass's Stage 5 transport solution feeds the next one three ways: as a boundary refit from the evolved pressure, as kinetic profiles prescribed to Stages 3, 4, and 5, and as the advanced transport clock.

pixi run driftless-star --config inputs/quick_run/config.yaml --max-iters 3 --cores 4

Each iteration is a full pipeline run under its own outputs/<run>/loop/iter_N/ tree (outputs/ is gitignored), and the driver stops early once the pressure profile settles within the config's convergence.pressure_rel_tol. Stages listed under loop.rerun as false are frozen, so iterations after the first reuse their iteration 1 artifacts. See docs/mvp-pipeline.md.

Run a single forward pass

One traversal from Stage 1 to Stage 5, without the transport feedback the loop adds:

pixi run driftless-star-fwd --configfile inputs/quick_run/config.yaml --cores 4

Its artifacts land under outputs/<run>/stageN_<name>/, one directory per stage, which is the flat layout the stage specs reference. A loop iteration writes that same tree one level down, under outputs/<run>/loop/iter_N/output/.

Define your own run

Copy the example, repoint input_dir/output_dir, and edit the inputs. Either entry point takes the new config the same way:

cp -r inputs/quick_run inputs/my_run
# in inputs/my_run/config.yaml, set input_dir: inputs/my_run and output_dir: outputs/my_run,
# then edit the boundary, profiles, and resolution as needed
pixi run driftless-star-fwd --configfile inputs/my_run/config.yaml --cores 4

Run on GPUs

Two top-level run-config keys pick the -gpu image variant for every stage and the device pool its containers may use:

gpu_ids: "4,5,6,7"   # null runs CPU images; "all" uses every GPU the execution host reports
jobs_per_gpu: 2      # concurrent jobs allowed per GPU

Every concurrent job is pinned to one free slot of that pool, so the pool offers pool size times jobs_per_gpu slots and saturating it takes at least that many cores. Either key can be overridden per invocation, on a forward pass or on the loop:

pixi run driftless-star-fwd --configfile inputs/quick_run/config.yaml --cores 8 --config gpu_ids=4,5,6,7 jobs_per_gpu=2
pixi run driftless-star --config inputs/quick_run/config.yaml --cores 8 --gpu-ids 4,5,6,7 --jobs-per-gpu 2

GPU mode needs an NVIDIA host with nvidia-container-toolkit configured on the docker daemon. See docs/mvp-pipeline.md for how the pinning works, how to share a host with other users, and the current limitations.

Visualize the pipeline graph

Render the file-flow graph (files as nodes, rules as edges) including the closed-loop post-processing step by targeting the convergence signal file:

pixi run -e pipeline bash -c 'snakemake --filegraph outputs/quick_run/stage5_post_processing/converge_status.json --configfile inputs/quick_run/config.yaml | dot -Tpdf > docs/figs/stellaforge_filegraph.pdf'

Omit the target to graph the plain forward pass (stops at Stage 5). Needs a one-time pixi run -e pipeline dot -c; see docs/mvp-pipeline.md for PNG/SVG and --rulegraph/--dag variants, including drawing the complete per-surface job DAG (the Stage 3/4 fan-out appears in --dag output only after the prepare manifests exist).

License

MIT