Fusion Energy Simulation: Tokamak

I think there are 3 major milestones remaining for humanity, and one of them is clean and abundant energy. Fusion energy has the potential to provide a nearly limitless source of clean energy by replicating the processes that power the sun. However, achieving controlled fusion reactions on Earth has proven to be a formidable and very challenging task.

The most well-known fusion prototype is the tokamak, which uses a magnetic field to confine the plasma within a toroidal chamber (doughnut-shaped).

Recently, I saw that DeepMind has made significant strides in this area by using RL to optimize the control of tokamak reactors [3]. I read their paper, and I wanted to summarize what I learned: physical laws, how to derive the equations, and the numerical methods used to solve them.

The equations used to model the plasma are Diffusion/Heat PDEs of this type:

The interesting thing is that by using toroidal flux coordinates, the complex 3D geometry of the tokamak can be simplified into a 1D radial coordinate system.

Source: [2]

Introduction

We have 4 PDE equations and 5 variables in normalized toroidal flux coordinates, $\hat \rho$, with $0 \leq \hat \rho \leq 1$. The important quantities being modeled are:

• ${\color{red}{T_i}}, {\color{orange}{T_e}}$: Ion and Electron Temperatures (Energy per particle). High temperature is required for fusion.
• ${\color{brown}{n_i}}, {\color{blue}{n_e}}$: Ion and Electron Densities (Particles per cubic meter). However, we do not model ion density explicitly, it is computed from electron density assuming quasi-neutrality.
• ${\color{teal}{\psi}}$: Poloidal Magnetic Flux. This describes the magnetic field structure and stability.

Here we summarize the variables and parameters used in the equations:

Boundary Conditions: The boundary conditions are as follows. All equations have a zero-derivative boundary condition at $\hat{\rho} = 0$ (Von-Neumann). The $T_i$, $T_e$, $n_e$ equations have fixed boundary conditions at $\hat{\rho} = 1$. The $\psi$ equation has a Neumann (derivative) boundary condition at $\hat{\rho} = 1$, which sets the total plasma current through the relation:

where $LCFS$ stands for Last Closed Flux Surface (i.e. $\hat{\rho} = 1$). The equations are derived from the transport equations in toroidal geometry under the assumption of static background flux surfaces [2].

1.1 Electron Particle Transport Equation

For an arbitrary plasma species $\alpha$, which may refer to electrons, main or impurity ion species, or fusion-born $\alpha$ particles, let $n_\alpha$ be the local particle density of the species and $\mathbf{u}_\alpha$ be the local velocity of the particles. The continuity equation for this species is stated as

Here $s_\alpha$ is a localized particle source. By writing the above equation inside a toroidal flux surface we obtain

A detailed derivation of the transport equations is provided in the Appendix.

We denote $\alpha=e$ (by the electron species) and denoting $\langle n_e \rangle \equiv n_e$. Using the specific form of the flux-surface averaged radial particle flux

Then

Plugging (1), we get:

Rewriting the grid-motion term using the toroidal flux $\Phi_b$ and a normalized radius $\hat\rho$. Using $\hat\rho^2 = \frac{\Phi}{\Phi_b(t)}$, so if $\Phi_b$ changes in time, by a chain-rule term:

That is exactly the operator appearing in the final equation:

Putting the pieces together gives:

1.2 Current Diffusion Equation.

In Appendix, we derived the 1D equation for the evolution of the poloidal flux $\psi$ using Ohm’s law and Faraday’s law:

Using the change of variables to $\hat\rho$ and $\Phi_b$, we get, as before, that:

The diffusion term is

Transport-code forms typically package all those geometric combinations into something like

where $g_2=\langle (\nabla V)^2/R^2\rangle$ and $g_3=\langle 1/R^2\rangle$. Therefore we obtain the final form of the current diffusion equation

1.3 Ion and Electron Heat Transport Equations

In a similar manner as the particle transport equation, an equation for energy transport can be derived [4, Eq. 7.46].

where $P_\alpha$ is the heat source for species $\alpha$ (ions or electrons), $q_\alpha$ is the heat flux, and $\Gamma_\alpha$ is the particle flux

where $E_\|$ is the parallel electric field and $G_1 := \langle (\nabla \rho)^2\rangle$. The various transport coefficients $D^\alpha_{x_\beta}, \chi^\alpha_{x_\beta}$ quantify, respectively the particle and heat flux of species $\alpha$ due to a gradient in the profile of $x_\beta$, while $D_E$ and $\chi_E$ represent (typically inward) fluxes due to parallel electric fields.

We use the same change of variables as before to express the equations in terms of $\hat\rho$ and $\Phi_b$. Using

we get to the final form of the ion and electron heat transport equations.

2. Discretization and Numerical Methods

In [1], the authors use a finite-volume method to discretize the spatial derivatives in the transport equations.

2.1 Spatial discretization: Finite Volume Method (FVM)

The 1D spatial domain, $0 \leq \hat{\rho} \leq 1$, is divided into a uniform grid of $N$ cells, each with a width of $d \hat{\rho} = 1/N$. The cell centers are denoted by $\hat{\rho}_i$, where $i = 1, 2,..., N-1$, and the $N+1$ cell faces are located at $\hat{\rho}_{i\pm1/2}$. Both $\hat{\rho}=0$ and $\hat{\rho}=1$ are on the face grid.

For a generic conservation law of the form:

where $x$ is a conserved quantity, $\mathbf{\Gamma}$ is the flux, and $S$ is a source term, the FVM approach involves integrating the PDE over a control volume (named a “cell”) and applying the divergence theorem to convert volume integrals to surface integrals. For 1D systems, following dividing by the cell volume, the finite volume method reduces to a special case of finite differences:

where $x_i$ is the cell-averaged value of $x$ in cell $i$, $\Gamma_{i+1/2}$ is the flux at face $i+1/2$, and $S_i$ is the cell-averaged source term in cell $i$.

In general, the fluxes are decomposed as

where $D$ is a diffusion coefficient and $V$ now denotes a convection coefficient, leading to:

The diffusion and convection coefficients are thus calculated on the face grid. The value of $x$ on the face grid is approximated by implementing a power-law scheme for Péclet-number weighting, which smoothly transitions between central differencing and upwinding, as follows:

where the $\alpha$ weighting factor depends on the Péclet number, defined as:

where $V$ is convection and $D$ is diffusion. The power-law scheme is as follows:

The Péclet number quantifies the relative strength of convection and diffusion. If the Péclet number is small and diffusion dominates, then the weighting scheme converges to central differencing. If the absolute value of the Péclet number is large, and convection dominates, then the scheme converges to upwinding.

Boundary conditions are taken into account by introducing ghost cells $x_{N}$ and $x_{-1}$ whose values are determined by assuming linear extrapolation through the edge cells and the face boundary conditions (for Dirichlet boundary conditions), or by directly satisfying the derivative (Neumann) boundary conditions.

The above equations, when combined, define the elements of the discretization matrix and boundary condition vectors for the PDE diffusion term.

2.2 Time discretization and solver options

TORAX uses the theta method for time discretization, and employs several options for solving the discretized PDE system. These are described below.

The θ-method is a weighted average between the explicit and implicit Euler methods. For a generic ODE of the form:

where $x$ is the state vector, the theta method approximates the solution at time $t + \Delta t$ as:

where $\theta$ is a user-selected weighting parameter in the range $[0, 1]$. Different values of $\theta$ correspond to well-known solution methods: explicit Euler ($\theta = 0$), Crank-Nicolson ($\theta = 0.5$), and implicit Euler ($\theta = 1$), which is unconditionally stable.

Upon inspection of the equations, we generalize equation (1) and write the state evolution equation as:

Starting from an initial condition $\mathbf{x}_0$, equation (2) solves for $\mathbf{x}_{t+\Delta t}$ at each timestep. $\mathbf{x}_t$ is the evolving state vector at time $t$, including all variables being solved by the system, and is of length $\#N$, where $\#$ is the number of solved variables.

Source: [1]

$\mathbf{u}_t$ corresponds to all known input parameters at time $t$. This includes boundary conditions, prescribed profiles (e.g. $n_e$ in the example above), and input parameters such as heating powers or locations.

$\mathbf{\tilde{T}}$ is the transient term, where $\odot$ signifies element-wise multiplication. For example, for the $T_e$ equation, $\mathbf{\tilde{T}}=\mathbf{n_e}$, which makes the system nonlinear if $\mathbf{n_e}$ itself is an evolving variable.

$\mathbf{\bar{C}}(x_t, u_t)$ and $\mathbf{\bar{C}}(x_{t+\Delta t}, u_{t+\Delta t})$ are the discretization matrices, of size $\#N\times \#N$. In general, depending on the physics models used, $\mathbf{\bar{C}}$ depends on state variables $\mathbf{x}$, for example through state-variable dependencies of transport coefficients $\chi$, $D$, $V$, plasma conductivity, and ion-electron heat exchange, making the system nonlinear due to the $x_{t+\Delta t}$ dependence. $\mathbf{c}$ is a vector, containing source terms and boundary condition terms.

There are two solver options for solving the nonlinear evolution system of equations.

Linear solver

This solver addresses the nonlinearity of the PDE system with fixed-point iteration, also known as the predictor-corrector method. For $K$ iterations (user-configurable), an approximation for $\mathbf{x}_{t+\Delta t}$ is obtained by solving the following equation iteratively with $k=1,2,..,K$:

and where $\mathbf{x}_{t+\Delta t}^{0} = \mathbf{x}_t$.

By replacing $\mathbf{x}_{t+\Delta t}$ with $\mathbf{x}_{t+\Delta t}^{k-1}$ within the coefficients $\mathbf{\tilde{T}}$, $\mathbf{\bar{C}}$ and $\mathbf{c}$, these coefficients become known at every iteration step, describing a linear system of equations. $\mathbf{x}_{t+\Delta t}^k$ can then be solved using standard linear algebra methods implemented in JAX.

Newton-Raphson Solver

This solver solves the nonlinear PDE system, using a gradient-based iterative Newton-Raphson root-finding method for finding the value of $\mathbf{x}_{t+\Delta t}$ that renders the residual vector zero:

where $\mathbf{R}$ is the LHS-RHS of equation (2).

Starting from an initial guess $\mathbf{x}_{t+\Delta t}=\mathbf{x}_{t+\Delta t}^0$, the Newton-Raphson method linearizes equation (3) about iteration $\mathbf{x}_{t+\Delta t}^k$ and solves the linear system for a step $\delta\mathbf{x}$:

where $\mathbf{\bar{J}}$ is the Jacobian of $\mathbf{R}$ with respect to $\mathbf{x}_{t+\Delta t}$. Crucially, JAX automatically calculates $\mathbf{\bar{J}}$ using auto-differentiation.

With $\delta\mathbf{x} = \mathbf{x}_{t+\Delta t}^{k+1} - \mathbf{x}_{t+\Delta t}^{k}$, $\mathbf{x}_{t+\Delta t}^{k+1}$ is solved using standard linear algebra methods implemented in JAX such as LU decomposition. This process iterates until the residual falls below a user-configurable tolerance $\varepsilon$, i.e: $\| \mathbf{R}(\mathbf{x}_{t+\Delta t}^{k+1}) \|_2 < \varepsilon$, where $\|\cdot\|_2$ is the vector two-norm.

3. Code Implementation (Torax)

You can see more examples and documentation at the Torax GitHub repository.

import torax
torax_config = torax.build_torax_config_from_file('examples/iterhybrid_rampup.py') # load Iter hybrid config
torax_config.update_fields({
    "numerics.fixed_dt": 1.0,
    "numerics.t_final": 100.0, 
})
dt, state_history = torax.run_simulation(torax_config)

And after some post-processing, you can generate plots like the following:

Figure: Axial view of plasma inside a tokamak during a simulation.

The most important quantities are the ones related to fusion power density, which can be plotted as follows:

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

# Plot p_alpha_i on first subplot
dt.profiles.p_alpha_i.plot(ax=ax1)
ax1.set_xlabel('$\\hat{\\rho}$')
ax1.set_title('Ion heating power density (W/m$^3$)')

# Plot p_alpha_e on second subplot
dt.profiles.p_alpha_e.plot(ax=ax2)
ax2.set_xlabel('$\\hat{\\rho}$')
ax2.set_title('Electron heating power density (W/m$^3$)')

plt.savefig('fusion_power_density.png', dpi=300, bbox_inches='tight')
plt.show()

Figure: Ion and electron heating power density from fusion reactions during a simulation of ITER hybrid scenario ramp-up phase.

References

[0] The code used for the figures is available at: https://github.com/dani2442/dani2442_code/tree/main/torax

[1] Citrin, Jonathan, Ian Goodfellow, Akhil Raju, Jeremy Chen, Jonas Degrave, Craig Donner, Federico Felici et al. “TORAX: A fast and differentiable tokamak transport simulator in JAX.” arXiv preprint arXiv:2406.06718 (2024).

[2] Felici, Federico. “Real-time control of tokamak plasmas: from control of physics to physics-based control.” (2011).

[3] https://deepmind.google/discover/blog/bringing-ai-to-the-next-generation-of-fusion-energy/

[4] Hinton, F. L. and R. D. Hazeltine (1976). “Theory of plasma transport in toroidal confinement systems.” In: Rev. Mod. Phys. 48.2, pp. 239–308. doi: 10.1103/RevModPhys. 48.239.

Appendix: Derivation of the 1D Transport Equations

Derivation of the Poloidal Flux Diffusion Equation

We derive the poloidal flux diffusion equation, describing the temporal evolution of the poloidal flux under the assumption of static background flux surfaces. We follow [2] Preliminaries. We use the useful relation

Here, we used that the average of $\nabla \cdot \mathbf{F}$ over a flux surface can be represented by the derivative with respect to the enclosed volume $V$ of the volume integral of $\nabla \cdot \mathbf{F}$. Then we change to flux coordinates

where $\phi$ is the toroidal angle and $\ell_p$ the poloidal length along the flux surface. Then we use Gauss divergence theorem to convert volume to surface integral,

Because $\phi$ is symmetric (axis) then the integral over $\phi$ gives a factor $2\pi$. Finally, the differential surface area of the flux surface is

Next, by the chain rule, we have

Thus,

And this finishes the proof of the relation.

Consider a surface of constant poloidal flux whose boundary moves with velocity $\mathbf{u}_\psi$. For this surface:

For a scalar field $F(t, \mathbf{x})$, define $H(t) = \int_V F\,dV$, where $V$ is the volume enclosed by $\psi = \text{const}$ moving with $\mathbf{u}_\psi$. By the Reynolds Transport Theorem, we have

Using this equality, the time rate of change of toroidal flux $\Phi$ enclosed by $\psi = \text{const}$, i.e., $F=\mathbf{B}\cdot\nabla\phi$ and $\Phi(t)=\int_V \mathbf{B}\cdot\nabla\phi\, dV$, is:

1. Poloidal electric field.

Using Faraday’s law $\partial_t \mathbf{B} = -\nabla\times\mathbf{E}$, we can rewrite the first volume integral (A.2) as:

2. Toroidal electric field.

We can rewrite the second integral of (A.2) using Ampère’s law:

Using that $\mathbf{B} \cdot \nabla \psi = 0$, we have

Then,

Thus, using previous equation and (A.1):

Combining previous results into (A.2):

3. Rate of change of poloidal flux.

Define the equivalent cylindrical fields:

Then, for $\dot{B}_0=0$:

4. Ohm’s law.

We can write the flux-surface-averaged Ohm’s law as:

where $\mathbf{j}_{ni} = \mathbf{j}_{bs} + \mathbf{j}_{cd}$ is the non-inductive current density, which includes the bootstrap current density $\mathbf{j}_{bs}$ and the external current drive density $\mathbf{j}_{cd}$. Equivalently

5. Parallel current.

Finally, we rewrite the parallel current $j_\parallel = \frac{\langle \mathbf{j}\cdot\mathbf{B}\rangle}{B_0}$ in terms of the poloidal flux $\psi$:

The key step is converting flux-surface averages of divergences into 1D (radial) derivatives:

Taking $\mathbf{F}=\nabla\psi/R^2$ and using $\psi=\psi(\rho)$ gives $\nabla\psi=(\partial\psi/\partial\rho)\nabla\rho$ and hence $\mathbf{F}\cdot\nabla\rho=((\nabla\rho)^2/R^2)(\partial\psi/\partial\rho)$. Also, since $T=T(\rho)$, $\langle \nabla T\cdot\nabla\psi/R^2\rangle = \langle (\nabla\rho)^2/R^2\rangle(\partial T/\partial\rho)(\partial\psi/\partial\rho)$. Defining $G^2\equiv \frac{V'}{4\pi^2}\left\langle \frac{(\nabla\rho)^2}{R^2} \right\rangle$ packages the terms and the product rule yields the final compact form. with

The poloidal flux diffusion equation. Combining (A.3), (A.4) and (A.5):

$$ \sigma_\parallel \left( \frac{\partial\psi}{\partial t} + \rho\frac{\dot{B}_0}{2B_0}\frac{\partial\psi}{\partial\rho} \right) = \frac{R_0 J^2}{\mu_0 \rho}\frac{\partial}{\partial\rho} \left(\frac{G^2}{J}\frac{\partial\psi}{\partial\rho}\right)- \frac{V'}{2\pi\rho}(j_{bs} + j_{cd}) $$

Derivation of the Particle Transport Equation

For an arbitrary plasma species $\alpha$, let $n_\alpha$ be the density, $\mathbf{u}_\alpha$ the velocity, $s_\alpha$ the source term. The continuity equation is:

Integrating over the volume enclosed by a flux surface (and using Gauss theorem for the second term):

Using Eq. (A.1b) with $F=n_\alpha$ and $\Phi=\text{const}$ and the fact that $\nabla \rho \| \nabla \Phi$:

Using flux surface averaging:

Then taking the derivative with respect to $\rho$:

where $\Gamma_\alpha$ is the particle flux and $S_\alpha$ the source term. Finally, at constant $\rho$: