Stellarator Configuration
Note
For general information about finite element discretization, basis functions, mesh parameters, polynomial degrees, boundary conditions, and matrix/operator dimensions, see Overview.
This script runs a magnetic relaxation simulation for a stellarator configuration.
The script is located at scripts/config_scripts/stell.py.
Problem Statement
The stellarator configuration solves the magnetohydrostatic (MHS) equilibrium equation:
with boundary conditions:
and the constraint:
where: - \(\mathbf{B}: \Omega \to \mathbb{R}^3\) is the magnetic field (2-form) - \(\mathbf{J} = \nabla \times \mathbf{B}\) is the current density (1-form) - \(p: \Omega \to \mathbb{R}\) is the plasma pressure (0-form) - \(\mathbf{n}\) is the outward unit normal vector on the boundary - \(\Omega\) is the stellarator plasma domain
The stellarator geometry uses a rotating elliptical boundary with 3D structure:
where: - \(R(r, \chi, \zeta) = R_0 + \epsilon r \cos(2\pi\chi + n_{\mathrm{fp}} \zeta)\) is the major radius - \(Z(r, \chi, \zeta) = \epsilon \kappa r \sin(2\pi\chi + n_{\mathrm{fp}} \zeta)\) is the vertical coordinate - \(\phi = 2\pi\zeta\) is the toroidal angle - \(R_0 = 1.0\) is the major radius - \(\epsilon = 0.33\) is the inverse aspect ratio - \(\kappa = 1.1\) is the elongation parameter - \(n_{\mathrm{fp}} = 3\) is the number of field periods (rotational symmetry)
Initial Magnetic Field
The initial magnetic field is defined in physical cylindrical coordinates:
where \(\tau = q_* = 1.54\) is the safety factor.
Perturbation Function
A perturbation can be added to introduce magnetic islands:
where: - \(a(r) = \exp(-(r - r_0)^2/(2\sigma^2))\) is a Gaussian radial profile - \(m_{\mathrm{pol}}\) is the poloidal mode number - \(n_{\mathrm{tor}}\) is the toroidal mode number - \(r_0\) is the radial location of the perturbation - \(\sigma\) is the radial width
Usage:
python scripts/config_scripts/stell.py run_name=test_run n_r=16 n_theta=16 n_zeta=8
The script accepts various parameters similar to the Solovev script. It generates HDF5 files with simulation data and PDF plots.
Code Walkthrough
Block 1: Imports and Configuration (lines 1-15)
Imports necessary modules. Note that this script uses rotating_ellipse_map
specifically for stellarator boundaries.
# %%
import os
import time
import jax
import jax.numpy as jnp
from mrx.mappings import rotating_ellipse_map
from mrx.derham_sequence import DeRhamSequence
from mrx.io import parse_args, unique_id
from mrx.plotting import trace_plot
from mrx.relaxation import MRXDiagnostics, State
from mrx.utils import update_config, default_trace_dict, save_trace_dict_to_hdf5, norm_2, run_relaxation_loop, DEFAULT_CONFIG
jax.config.update("jax_enable_x64", True)
Block 2: Main Function (lines 17-33)
Parses command-line arguments and sets default configuration parameters:
boundary_type="rotating_ellipse": Fixed for stellarator geometryeps=0.33: Ellipticity parameterkappa=1.1: Elongation parameterq_star=1.54: Safety factorn_fp=3: Number of field periods (rotational symmetry)Default discretization:
n_r=8, n_theta=8, n_zeta=4dt=1e-4: Time step size
def main():
# Get user input
params = parse_args()
CONFIG = DEFAULT_CONFIG.copy()
# Specific configuration parameters for the simulation that are not the default configuration
# NOTE: these can be overridden by the user-supplied parameters from the command line.
CONFIG["boundary_type"] = "rotating_ellipse"
CONFIG["eps"] = 0.33
CONFIG["kappa"] = 1.1
CONFIG["q_star"] = 1.54
CONFIG["n_fp"] = 3
CONFIG["n_r"] = 8
CONFIG["n_theta"] = 8
CONFIG["n_zeta"] = 4
CONFIG["dt"] = 1e-4
CONFIG = update_config(params, CONFIG)
run(CONFIG)
Block 3: Setup and Initial Condition (lines 36-108)
The run() function:
Creates output directory:
script_outputs/stell/{run_name}/Sets up rotating ellipse mapping with specified parameters
Configures DeRham sequence with polar coordinates
Defines initial magnetic field \(\mathbf{B}_{\mathrm{xyz}}(p)\) in physical space
Defines perturbation function \(\delta\mathbf{B}_{\mathrm{xyz}}(p)\) for adding magnetic islands
Projects initial field into FEM space and applies Leray projection to ensure divergence-free condition
Optionally applies perturbation if
apply_pert_after=0andpert_strength>0
def run(CONFIG):
run_name = CONFIG["run_name"]
if run_name == "":
run_name = unique_id(8)
outdir = "script_outputs/stell/" + run_name + "/"
os.makedirs(outdir, exist_ok=True)
print("Running simulation " + run_name + "...")
start_time = time.time()
# Initialize the trace dictionary
trace_dict = default_trace_dict.copy()
trace_dict["start_time"] = start_time
F = rotating_ellipse_map(CONFIG["eps"], CONFIG["kappa"], CONFIG["n_fp"])
ns = (CONFIG["n_r"], CONFIG["n_theta"], CONFIG["n_zeta"])
ps = (CONFIG["p_r"], CONFIG["p_theta"], 0
if CONFIG["n_zeta"] == 1 else CONFIG["p_zeta"])
q = max(ps)
types = ("clamped", "periodic",
"constant" if CONFIG["n_zeta"] == 1 else "periodic")
tau = CONFIG["q_star"]
print("Setting up FEM spaces...")
Seq = DeRhamSequence(ns, ps, q, types, F, polar=True, dirichlet=True)
assert jnp.min(Seq.J_j) > 0, "Mapping is singular!"
# Assemble the FEM spaces and build the projections
Seq.evaluate_1d()
Seq.assemble_all()
Seq.build_crossproduct_projections()
Seq.assemble_leray_projection()
# Define the initial magnetic field in physical space
def B_xyz(p):
x, y, z = F(p)
R = (x**2 + y**2)**0.5
phi = jnp.arctan2(y, x)
BR = z * R
Bphi = tau / R
Bz = - (1 / 2 * (R**2 - 1**2) + z**2)
Bx = BR * jnp.cos(phi) - Bphi * jnp.sin(phi)
By = BR * jnp.sin(phi) + Bphi * jnp.cos(phi)
return jnp.array([Bx, By, Bz])
# Define the perturbation to the initial magnetic field in physical space
def dB_xyz(p):
r, θ, ζ = p
DFx = jax.jacfwd(F)(p)
def a(r):
return jnp.exp(- (r - CONFIG["pert_radial_loc"])**2 / (2 * CONFIG["pert_radial_width"]**2))
B_rad = a(r) * jnp.sin(2 * jnp.pi * θ * CONFIG["pert_pol_mode"]) * \
jnp.sin(2 * jnp.pi * ζ * CONFIG["pert_tor_mode"]) * DFx[:, 0]
return B_rad
# Solve for the magnetic field in the FEM space
B_hat = jnp.linalg.solve(Seq.M2, Seq.P2(B_xyz))
B_hat = Seq.P_Leray @ B_hat
B_hat /= norm_2(B_hat, Seq)
# Apply a perturbation to the initial magnetic field if specified
if CONFIG["apply_pert_after"] == 0 and CONFIG["pert_strength"] > 0:
print("Applying perturbation to initial condition...")
dB_hat = jnp.linalg.solve(Seq.M2, Seq.P2(dB_xyz))
dB_hat = Seq.P_Leray @ dB_hat
dB_hat /= norm_2(dB_hat, Seq)
B_hat += CONFIG["pert_strength"] * dB_hat
# Initialize the state and diagnostics of the simulation
diagnostics = MRXDiagnostics(Seq, CONFIG["force_free"])
Block 4: Magnetic Relaxation Loop (lines 111-113)
Runs the relaxation loop with perturbation function stored in CONFIG for later use.
timestepper = TimeStepper(Seq,
gamma=CONFIG["gamma"],
newton=CONFIG["precond"],
Block 5: Post-processing (lines 115-139)
Computes final pressure, saves data to HDF5, and generates trace plots.
Note: This script uses trace_plot() directly instead of generate_solovev_plots(),
so it only generates convergence traces, not pressure contour plots.
The stellarator configuration uses a rotating elliptical boundary that creates 3D magnetic field structure with inherent rotational transform, making it distinct from axisymmetric tokamak configurations.
picard_tol=CONFIG["solver_tol"],
picard_k_restart=CONFIG["solver_maxit"])
# Compute initial force and hessian
F, _, _ = timestepper.compute_force(B_hat)
if CONFIG["precond"]:
hessian_assembler = MRXHessian(Seq)
H = hessian_assembler.assemble(B_hat)
else:
H = None
# Create state
state = State(B_n=B_hat,
B_nplus1=B_hat,
dt=CONFIG["dt"],
eta=CONFIG["eta"],
hessian=H,
v=F,
F_norm=norm_2(F, Seq))
# Update v and v_norm
if CONFIG["precond"]:
v = timestepper.apply_inverse_hessian(state, F)
else:
v = F