Solovev 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 Solovev configuration.
The script is located at scripts/config_scripts/solovev.py.
Problem Statement
The Solovev 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, area form) - \(\mathbf{J} = \nabla \times \mathbf{B}\) is the current density (1-form) - \(p: \Omega \to \mathbb{R}\) is the plasma pressure (0-form, scalar) - \(\mathbf{n}\) is the outward unit normal vector on the boundary - \(\Omega\) is the toroidal plasma domain
For the force-free case (\(\nabla p = 0\)), the equation becomes:
which implies \(\mathbf{J} \parallel \mathbf{B}\) (force-free condition).
Magnetic Relaxation Method
The relaxation method solves the time-dependent MHD equations:
where the electric field \(\mathbf{E}\) is given by:
and: - \(\mathbf{u}\) is the velocity field (divergence-free) - \(\eta\) is the resistivity parameter
The velocity is chosen to drive the system toward equilibrium:
where \(\mathbf{f}\) is the force imbalance. The velocity is projected onto the divergence-free space using the Leray projection \(P_{\mathrm{Leray}}\).
Time-Stepping Scheme
The time-stepping uses a midpoint rule (Crank-Nicolson):
where \(\mathbf{B}^{\mathrm{mid}} = (\mathbf{B}^n + \mathbf{B}^{n+1})/2\) is evaluated at the midpoint.
The system is solved using a Picard iteration scheme with adaptive time-stepping.
Usage:
python scripts/config_scripts/solovev.py run_name=test_run boundary_type=rotating_ellipse n_r=16 n_theta=16 n_zeta=8
The script accepts various parameters:
run_name: Name for the output filesboundary_type: Type of boundary (e.g.,rotating_ellipse,cerfon)n_r,n_theta,n_zeta: Number of splines in each directionp_r,p_theta,p_zeta: Polynomial degreesmaxit: Maximum number of iterationsAnd many more (see script for full list)
The script generates HDF5 files with simulation data and PDF plots showing:
Force trace over iterations
Pressure contours
Magnetic field structure
Discretization Parameters
This script uses a 3D DeRham sequence with mesh parameters \(n_r, n_\theta, n_\zeta\), polynomial degrees \(p_r, p_\theta, p_\zeta\), and boundary conditions (typically clamped in radial direction, periodic in poloidal and toroidal directions). For details on computing the number of DOFs \(N_0, N_1, N_2, N_3\) from these parameters, see Overview.
Finite Element Spaces
The magnetic field \(\mathbf{B}\) is represented as a 2-form (area form) in the finite element space:
The current density \(\mathbf{J} = (\nabla \times)_h \mathbf{B}\) is a 1-form, where \((\nabla \times)_h: V_2 \to V_1\) is the curl operator (weak form).
Matrix and Operator Dimensions
The 2-form mass matrix \(M_2 \in \mathbb{R}^{N_2 \times N_2}\) and 1-form mass matrix \(M_1 \in \mathbb{R}^{N_1 \times N_1}\) are used. The Leray projection \(P_{\mathrm{Leray}} \in \mathbb{R}^{N_2 \times N_2}\) ensures divergence-free velocity fields. Cross-product projections compute \(\mathbf{J} \times \mathbf{B}\) in the finite element space.
Initial Magnetic Field
The initial magnetic field is constructed from a harmonic eigenmode by solving:
The eigenvector \(\mathbf{v}\) corresponding to the smallest eigenvalue is extracted and normalized to unit L2 norm.
Code Walkthrough
Block 1: Imports and Configuration (lines 1-17)
Imports JAX, NumPy, and MRX modules for DeRham sequences, I/O, mappings, relaxation, and plotting. Enables 64-bit precision for numerical accuracy.
# %%
import os
import time
import jax
import jax.numpy as jnp
from mrx.derham_sequence import DeRhamSequence
from mrx.io import parse_args, unique_id
from mrx.mappings import cerfon_map, helical_map, rotating_ellipse_map
from mrx.plotting import generate_solovev_plots
from mrx.relaxation import MRXDiagnostics, State
from mrx.utils import run_relaxation_loop, update_config, DEFAULT_CONFIG
from mrx.utils import default_trace_dict, save_trace_dict_to_hdf5, norm_2
from mrx.utils import is_running_in_github_actions
jax.config.update("jax_enable_x64", True)
Block 2: Main Function (lines 19-41)
The main function parses command-line arguments using parse_args() and updates
the default configuration. It then calls run() with the updated configuration.
Command-line parameters override defaults using the format parameter_name=value.
def main():
"""
Runs a magnetic relaxation simulation for a Solovev configuration.
Usage: python solovev.py <parameter_name>=<parameter_value>
where <parameter_name> is one of the parameters in DEFAULT_CONFIG and <parameter_value> is the value to use.
For example:
python solovev.py run_name=test_run boundary_type=rotating_ellipse n_r=16 n_theta=16 n_zeta=8 p_r=3 p_theta=3 p_zeta=3
will run a simulation with 16 radial, 16 poloidal, and 8 toroidal splines,
with radial and poloidal splines of degree 3 and toroidal splines of degree 3, on the CPU.
"""
# Get user input
params = parse_args()
CONFIG = update_config(params, DEFAULT_CONFIG)
if is_running_in_github_actions():
CONFIG["n_r"] = 4
CONFIG["n_theta"] = 4
CONFIG["n_zeta"] = 4
CONFIG["p_r"] = 2
CONFIG["p_theta"] = 2
CONFIG["p_zeta"] = 2
run(CONFIG)
Block 3: Setup and Initialization (lines 44-91)
The run() function creates output directory structure, initializes trace dictionary,
and selects boundary mapping based on boundary_type:
tokamak: Usescerfon_map(circular cross-section)helix: Useshelical_map(helical boundary)rotating_ellipse: Usesrotating_ellipse_map(rotating elliptical cross-section)
Sets up DeRham sequence with polar coordinates and Dirichlet boundary conditions, assembles all necessary matrices (mass, derivative, cross-product projections, Leray projection), and initializes magnetic field from harmonic eigenmode:
The field is normalized to unit L2 norm: \(\hat{\mathbf{B}} = \mathbf{B}_{\mathrm{harm}} / \|\mathbf{B}_{\mathrm{harm}}\|_2\)
def run(CONFIG):
run_name = CONFIG["run_name"]
if run_name == "":
run_name = unique_id(8)
outdir = "script_outputs/solovev/" + 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
if CONFIG["boundary_type"] == "tokamak":
F = cerfon_map(CONFIG["eps"], CONFIG["kappa"], jnp.arcsin(CONFIG["delta"]))
elif CONFIG["boundary_type"] == "helix":
F = helical_map(epsilon=CONFIG["eps"], h=CONFIG["h_helix"],
nfp=CONFIG["nfp"], kappa=CONFIG["kappa"], alpha=jnp.arcsin(CONFIG["delta"]))
elif CONFIG["boundary_type"] == "rotating_ellipse":
F = rotating_ellipse_map(
eps=CONFIG["eps"], kappa=CONFIG["kappa"], nfp=CONFIG["nfp"])
else:
raise ValueError("Unknown boundary type.")
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")
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()
# Initialize initial magnetic field guess
B_harm = jnp.linalg.eigh(Seq.M2 @ Seq.dd2)[1][:, 0]
B_hat = B_harm / norm_2(B_harm, Seq)
# Initialize the state of the simulation
diagnostics = MRXDiagnostics(Seq, CONFIG["force_free"])
Block 4: Magnetic Relaxation Loop (lines 93-94)
Calls run_relaxation_loop() which performs the time-stepping iteration to relax
the magnetic field toward equilibrium. This solves the MHD force balance equation:
or for the force-free case:
where \(\mathbf{J} = \nabla \times \mathbf{B}\) is the current density. The solver uses an implicit time-stepping scheme with Picard iteration.
from mrx.relaxation import TimeStepper, MRXHessian
timestepper = TimeStepper(Seq,
Block 5: Post-processing and Output (lines 96-122)
After the relaxation loop completes:
Computes final pressure field using
diagnostics.pressure()Saves all trace data (iterations, forces, energies, etc.) to HDF5 file
Generates plots using
generate_solovev_plots()which creates: * Pressure contour plots * Force and energy convergence traces * Magnetic field visualizations (ifsave_B=True)
newton=CONFIG["precond"],
force_free=CONFIG["force_free"],
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
The Solovev configuration provides a simple analytical equilibrium that serves as a test case for the relaxation solver. The script demonstrates how to set up and run magnetic relaxation simulations for various boundary geometries.