Stellamorph

Generative modeling for stellarator design with diffusion models and flow matching.

This repository covers the standard workflow:

training -> sampling -> physics evaluation

The codebase includes DDPM and flow-matching generators, sampling scripts, and VMEC-based evaluation utilities for quasi-symmetry metrics.

Installation

Create the conda environment:

conda env create -f environment.yml
conda activate distel

If the environment already exists:

conda env update -f environment.yml --prune
conda activate distel

Install local package hooks:

./install.sh

VMEC-related utilities:

./tools/setup/install_vmec2000.sh
./tools/dev/check_vmec.sh
./tools/dev/verify_vmec_env.py --instantiate

Standard Usage

Train a model:

./scripts/train.sh
./scripts/train.sh --schedule ddpm
./scripts/train.sh --schedule flow
./scripts/train.sh --conditions mean_iota aspect_ratio nfp helicity --use_pca --pca_size 50

Sample new configurations:

./scripts/sample.sh --model_dir <run_dir>
./scripts/sample.sh --model_dir <run_dir> --n_samples 64 --target_nfp 4
./scripts/sample.sh --model_dir <run_dir> --schedule flow --flow_formulation ode --flow_ode_order 3
./scripts/sample.sh --model_dir <run_dir> --schedule flow --flow_formulation sde

Evaluate sampled configurations:

./scripts/eval.sh --sample_dir <sample_dir>

Run the full pipeline from a config file:

./scripts/pipeline.sh configs/pipeline.example.conf.sh

Search scripts are in /scripts/search folder, please refer to /scripts/search/readme.md for execution.

Example Results

Main comparison across the evaluated nfp range:

Weight ablation for search:

Runtime sweep:

References

1. Misha Padidar, Teresa Huang, Andrew Giuliani, and Marina Spivak. Diffusion for Fusion: Designing Stellarators with Generative AI. arXiv:2511.20445, 2025. https://arxiv.org/abs/2511.20445
2. Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow Matching for Generative Modeling. arXiv:2210.02747, 2022. https://arxiv.org/abs/2210.02747