Iterative Island Calculation
Note
For general information about finite element discretization, basis functions, mesh parameters, polynomial degrees, boundary conditions, and matrix/operator dimensions, see Overview.
This script performs iterative island calculations for magnetic configurations.
The script is located at scripts/config_scripts/iter_islands.py.
Problem Statement
This script is similar to solovev.py but configured for studying magnetic islands
in tokamak configurations. It 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 tokamak plasma domain
Tokamak Geometry
The tokamak uses a cerfon mapping (circular cross-section) with parameters: - \(\epsilon = 0.33\): Inverse aspect ratio - \(\kappa = 1.7\): Elongation parameter - \(\delta = 0.33\): Triangularity parameter
Initial Magnetic Field
The initial magnetic field is defined in physical cylindrical coordinates:
where: - \(\tau = q_* \kappa (1 + \kappa^2) / (\kappa + 1)\) is the safety factor parameter - \(R_0 = 1.0\) is the major radius - \(\kappa\) is the elongation parameter
Perturbation for Island Seeding
A perturbation is applied to seed 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
The perturbation strength is typically small (\(\sim 10^{-5}\)) to avoid disrupting the equilibrium while still seeding islands.
Usage:
python scripts/config_scripts/iter_islands.py run_name=test_run
The script generates output files with island calculation results.
Code Walkthrough
Block 1: Imports and Configuration (lines 1-16)
Imports modules including cerfon_map for tokamak boundaries.
# %%
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
from mrx.relaxation import MRXDiagnostics, State
from mrx.utils import update_config, DEFAULT_CONFIG, run_relaxation_loop
from mrx.utils import default_trace_dict, save_trace_dict_to_hdf5, norm_2
from mrx.plotting import trace_plot
jax.config.update("jax_enable_x64", True)
Block 2: Main Function (lines 19-33)
Sets default configuration for tokamak geometry:
boundary_type="tokamak": Fixed for tokamakeps=0.33: Inverse aspect ratiokappa=1.7: Elongationdelta=0.33: Triangularityq_star=1.54: Safety factorpert_strength=2e-5: Small perturbation amplitude for island seedingsave_every=10: Save magnetic field snapshots every 10 iterations
def main():
# Get user input
params = parse_args()
CONFIG = DEFAULT_CONFIG.copy()
CONFIG["boundary_type"] = "tokamak"
CONFIG["eps"] = 0.33
CONFIG["kappa"] = 1.7
CONFIG["delta"] = 0.33
CONFIG["delta_B"] = 0.2
CONFIG["q_star"] = 1.54
CONFIG["dt"] = 1e-4
CONFIG["pert_strength"] = 2e-5
CONFIG["save_every"] = 10
CONFIG = update_config(params, CONFIG)
run(CONFIG)
Block 3: Setup and Initial Condition (lines 36-111)
The run() function:
Creates output directory:
script_outputs/iter/{run_name}/Sets up cerfon mapping (circular tokamak cross-section)
Defines initial magnetic field \(\mathbf{B}_0(p)\) with tokamak-specific scaling
Defines perturbation
dB_xyz(p)with Gaussian radial profile and poloidal/toroidal mode structureProjects field into FEM space and applies Leray projection
Applies perturbation to seed magnetic islands if
apply_pert_after=0
def run(CONFIG):
run_name = CONFIG["run_name"]
if run_name == "":
run_name = unique_id(8)
outdir = "script_outputs/iter/" + run_name + "/"
os.makedirs(outdir, exist_ok=True)
print("Running simulation " + run_name + "...")
kappa = CONFIG["kappa"]
eps = CONFIG["eps"]
alpha = jnp.arcsin(CONFIG["delta"])
start_time = time.time()
F = cerfon_map(eps, kappa, alpha)
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"] * kappa * (1 + kappa**2) / (kappa + 1)
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!"
Seq.evaluate_1d()
Seq.assemble_all()
Seq.build_crossproduct_projections()
Seq.assemble_leray_projection()
trace_dict = default_trace_dict.copy()
trace_dict["start_time"] = start_time
def B_0(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 = - (kappa**2 / 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])
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
# Initialize the initial magnetic field in the FEM space
B_dof = jnp.linalg.solve(Seq.M2, Seq.P2(B_0))
B_dof = Seq.P_Leray @ B_dof
B_dof /= norm_2(B_dof, 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_dof = jnp.linalg.solve(Seq.M2, Seq.P2(dB_xyz))
dB_dof = Seq.P_Leray @ dB_dof
dB_dof /= norm_2(dB_dof, Seq)
B_dof += CONFIG["pert_strength"] * dB_dof
diagnostics = MRXDiagnostics(Seq, CONFIG["force_free"])
# Create TimeStepper to compute initial force (matching relaxation_loop)
Block 4: Magnetic Relaxation Loop (lines 114-116)
Runs relaxation with perturbation function available for dynamic application.
gamma=CONFIG["gamma"],
newton=CONFIG["precond"],
force_free=CONFIG["force_free"],
Block 5: Post-processing (lines 119-142)
Saves final state and generates plots. The perturbation creates magnetic islands that can be tracked through the relaxation process.
The script is designed to study how magnetic islands evolve during relaxation, which is important for understanding plasma stability and transport in tokamaks.
# Compute initial force and hessian
F, _, _ = timestepper.compute_force(B_dof)
if CONFIG["precond"]:
hessian_assembler = MRXHessian(Seq)
H = hessian_assembler.assemble(B_dof)
else:
H = None
# Create state
state = State(B_n=B_dof,
B_nplus1=B_dof,
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
state = timestepper.update_field(state, "v", v)