Application of Neural Ordinary Differential Equations for Tokamak Plasma Dynamics Analysis

Authors: zefang  liu
weston m stacey
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-09 20:53:13

Abstract: In the quest for controlled thermonuclear fusion, tokamaks present complex challenges in understanding burning plasma dynamics. This study introduces a multi-region multi-timescale transport model, employing Neural Ordinary Differential Equations (Neural ODEs) to simulate the intricate energy transfer processes within tokamaks. Our methodology leverages Neural ODEs for the numerical derivation of diffusivity parameters from DIII-D tokamak experimental data, enabling the precise modeling of energy interactions between electrons and ions across various regions, including the core, edge, and scrape-off layer. These regions are conceptualized as distinct nodes, capturing the critical timescales of radiation and transport processes essential for efficient tokamak operation. Validation against DIII-D plasmas under various auxiliary heating conditions demonstrates the model's effectiveness, ultimately shedding light on ways to enhance tokamak performance with deep learning.

Main paper: 
 introduction: 
 Fusion burning plasmas (Green et al., 2003) , pivotal in ITER's deuterium-tritium (D-T) fusion experiments, produce high-energy neutrons and fusion alpha particles. The latter, confined within the tokamak's magnetic field, sequentially transfer energy to electrons and ions, leading to sophisticated radiation and energy transport mechanisms like electron cyclotron radiation (ECR) and bremsstrahlung. To effectively manage the complex dynamics of energy transfer and mitigate potential thermal runaway instability, it is crucial to model the interactions between instantaneous radiation and gradual diffusion transport across the tokamak's multiple regions. This thorough understanding of energy dynamics is essential for controlling the plasma's behavior, emphasizing the need for models that accurately represent various processes across different timescales.To address these complexities, our study introduces a multi-region multi-timescale transport model (Liu & Stacey, 2021; Liu, 2022) that leverages Neural Ordinary Differential Equations (Neural ODEs) (Chen et al., 2018) for enhanced burning plasma simulation fidelity in tokamaks. Inspired by previous researches (Hill, 2019; Stacey, 2021) , our model distinguishes between the core, edge, and scrape-off layer (SOL) regions, where each region is represented as a unique node with specific energy dynamics influenced by radiation and transport processes. By employing Neural ODEs, we dynamically optimize diffusivity parameters from DIII-D tokamak experimental data, enabling precise modeling of energy interactions. This method not only facilitates a deeper understanding of the plasma dynamics but also improves model accuracy and predictive capability, particularly in capturing the intricate interactions between regions. Our application of Neural ODEs exemplifies the potential of deep learning techniques in advancing tokamak plasma analysis, providing a robust framework for future explorations in controlled thermonuclear fusion, including ITER's D-T experiments. 
tokamak plasma dynamics model: 
 This section introduces a multinodal model for tokamak plasma dynamics, starting with its geometry, then detailing particle and energy balance equations for the core, edge, and scrape-off layer (SOL) nodes, and concluding with the presentation of a parametric diffusivity model.In the multinodal model, a tokamak plasma is partitioned into three distinct regions, each modeled as a separate node: the core region (0 ≤ ρ ≤ ρ core = 0.9), the edge region (ρ core ≤ ρ ≤ ρ edge = 1.0), and the SOL region (ρ edge ≤ ρ ≤ ρ sol = 1.1), with the minor radius a and the normalized minor radius ρ = r/a, as illustrated in Figure 1 . This division simplifies the tokamak into a torus, whose cross section is considered circular, and delineates the core, edge, and SOL as toroidal shells, separated by torus surfaces at radii r core , r edge , and r sol corresponding to surfaces A core , A edge , and A sol , respectively. Distances between these regions are defined as ∆r core-edge (core to edge), ∆r edge-sol (edge to SOL), and ∆r sol-div (SOL node center to its outer surface). The multinodal plasma dynamics model for DIII-D deuterium-deuterium (D-D) plasmas employs two fundamental groups of equations to capture the complex interactions within the plasma: particle and energy balance equations. Particle balance equations quantify the conservation of particles within each node accounting for particle sources and sinks, whereas energy balance equations track the thermal energy's distribution and evolution, considering heating, radiation, and energy transport. These equations form the backbone of the model, enabling detailed simulation of plasma behavior across different regions. However, this model's diffusive approach may oversimplify plasma dynamics by omitting convective (pinch) terms, with future enhancements potentially including pinch dynamics to improve accuracy, as highlighted in Angioni et al. (2009) .Particle balance equations for deuterons in the core, edge, and SOL nodes are dn core D dt = S core D,ext + S core D,where node ∈ { core, edge, sol }, S node σ,ext is the external particle source, S node σ,IOL is the ion orbit loss (IOL) term (Stacey, 2011) , S sol σ,ion is for ionization processes, S sol σ,rec is for recombination processes, and S node σ,tran is the particle transport term. The particle transport for the core region is S core σ,tran = -n core σ -n edge σ τ core→edge P,σ, where the internodal particle transport time is τ core→edge Energy balance equations for deuterons and electrons in the core, edge, and SOL nodes are dU core D dt = P core D,aux + Q core D + P core D,where the nodal energy density is is the energy transport term. The energy transport for the core region is P core σ,tran = -U node σ = 3 2 n node σ T node σ , P node σ,U core σ -U edge σ τ core→edge E,σ, where the internodal energy transport time is τ core→edge In order to calculate internodal transport times, formulas for particle and thermal diffusivities are required. One empirical scaling for the effective thermal diffusivity (Becker, 2004) in the ELMy H-mode tokamak plasma isEQUATIONwhere the terms in this formula are the thermal diffusivity χ H98 in m 2 /s, normalized radius ρ = r/a, coefficient α H = 0.123, toroidal magnetic field B T in T, electron density n e in 10 19 m -3 , electron temperature T e in keV, electron temperature gradient ∇T e in keV/m, safety factor q = q ψ , local elongation κ, hydrogenic atomic mass number M in 1 amu, major radius R in m, and minor radius a in m. The particle and thermal diffusivities for electrons and ions (Becker & Kardaun, 2006) are assumed as χ e (ρ) = χ i (ρ) = χ H98 (ρ) and D i (ρ) = 0.6χ H98 (ρ), which can be replaced by separate formulas in future. This empirical scaling is used as the baseline in this study.Besides, a parametric diffusivity formula is proposed for particle and energy diffusivities asχ(ρ) 1 m 2 /s = α H B T 1 T α B n e (ρ) 10 19 m -3 αn T e (ρ) 1 keV α T |∇T e (ρ)| 1 keV/m α ∇T q(ρ) αq κ(ρ) ακ • M 1 amu α M R 1 m α R a 1 m αa . (4)where α H , α B , . . . , α a are diffusivity parameters for each species and node, and these parameters will be determined from experimental data. This diffusivity formula can be grouped and reformulated into a vector form: ln χ node = b node + W node ln x node , where χ node is the vector of internodal diffusivities, b node is the bias vector, W node is the weight matrix, and x node is the vector of corresponding physical values. 
computational framework: 
 In this section, we develop a computational framework for simulating plasma dynamics in tokamaks. This framework consists of a series of interconnected modules designed to process experimental data, simulate plasma behavior, and optimize model parameters. Initially, a data module reads experimental inputs, such as two-dimensional plasma profiles and one-dimensional global parameters, from the OMFIT (Meneghini et al., 2015) for DIII-D data. A preprocessing module then standardizes these inputs into uniform time sequences and performs volume averaging on two-dimensional signals to obtain nodal particle densities and temperatures.The core of the framework includes a diffusivity model that calculates particle and thermal diffusivities based on experimental conditions, and a transport time model that determines the timescales for particle and energy transport between nodes. These models feed into a reactor simulation module that integrates sources, sinks, and transport terms into a dynamical system, which is then solved using the Neural Ordinary Differential Equation (Neural ODE) solver (Chen et al., 2018) . This dynamical system solver outputs estimated particle densities and temperatures, which are refined through an optimization module. This module computes the mean square error (MSE) between model predictions and experimental data, utilizing back-propagation for gradient computation and parameter updating via gradient descent, ensuring the model's parameters are optimized for accurate plasma behavior representation.A workflow diagram in Figure 2 visually represents the framework's structure, where cylinders are datasets, rectangles are modules, solid lines are forward flows to solve the problem, and dashed lines are back propagation processes to optimize the parameters in the diffusivity model. This architecture enables a comprehensive analysis of tokamak plasma dynamics, leveraging modern computational techniques for enhanced modeling accuracy and predictive capability in fusion research.  
simulation methods: 
 In the development of the multinodal model for DIII-D plasma simulations, 25 shots from experimental data listed in Appendix A Table 2 are utilized, adhering to the criteria set by Hill & Stacey (2017) for ELMing non-RMP H-mode. Particle density and temperature profiles, derived from ZIP-FIT (Peng et al., 2002) integrated with EFIT and Thomson scattering data, are used to establish nodal densities and temperatures. To ensure the model's generalizability across various experimental conditions and prevent overfitting to specific shots, the data are split into two sets: a training dataset comprising 20 shots for optimizing the parametric diffusivity model, and a testing dataset of 5 shots for model performance evaluation. The training process involves solving the multinodal model's dynamical system using the training dataset, where the diffusivity parameters are initialized with the χ H98 model (Becker, 2004) , and the nodal source and sink terms are combined into a dynamical system and solved for simulated nodal densities and temperatures. The mean squared error (MSE) between the model's predictions and experimental data is calculated to guide the optimization of diffusivity parameters through gradient descent, iterating over multiple epochs until a satisfactory MSE is achieved. Once optimized, the model's effectiveness is assessed on the testing dataset by comparing its predictions against experimental measurements, thereby evaluating the computational method's overall accuracy and reliability.Following the training of the parametric diffusivity model with the designated training shots, the multinodal model's performance is evaluated on testing shots by comparing the mean squared errors (MSEs) between the model's solutions and experimental measurements, where the particle densities and temperatures are normalized by 10 19 m -3 and 1 keV respectively. The multinodal model with the empirical diffusivity scaling (Becker, 2004 ) is used as the baseline. This comparison, detailed in Table 1 , highlights the improvements achieved through optimization, with MSE reductions exceeding 96% across individual shots and an overall average MSE decrease of more than 98%. Such significant reductions in MSEs, achieved without the model being directly trained on the testing dataset, underscore the computational method's efficacy and the optimized multinodal model's capability to generalize and accurately predict the behavior of new, unseen shots. Detailed results of the model's performance for two specific testing shots are provided in Appendix B, including experiment signals and model solutions for the core, edge, and SOL nodes. The model with optimized parameters outperforms its original, unoptimized version in performance, especially in the core and edge nodes, where the original model exhibits significant deviations between simulation results and experimental measurements.  
conclusion: 
 This study introduces a multinodal plasma dynamics model for tokamak analysis, segmenting the plasma into core, edge, and SOL regions as distinct nodes with tailored timescales for various phenomena. We detail computational techniques for refining diffusivity parameters within this model, demonstrating its validation and parameter optimization using DIII-D plasma data. The enhanced performance of the optimized model signifies its advancement over previous models. This work illustrates the effectiveness of employing Neural Ordinary Differential Equations (Neural ODEs) in tokamak plasma dynamics research, highlighting an advancement in simulation methodologies.Looking forward, applying these techniques to ITER fusion plasmas is promising but challenging due to the absence of direct training sets, potentially requiring extrapolation with derived parameters to address ITER's distinct conditions.0.6 0.8I P [MA]Shot 1311901.81.9 Figure 3 : Signals of shot 131190, including (from top to bottom, left to right) plasma current I P , toroidal magnetic field B 0 , safety factor q 95 , gas puffing rate GAS, ohmic heating power P Ω , neutral beam injection (NBI) power P NBI , electron cyclotron heating (ECH) power P ECH , and ion cyclotron heating (ICH) power P ICH . Figure 7 : Signals of shot 140418, including (from top to bottom, left to right) plasma current I P , toroidal magnetic field B 0 , safety factor q 95 , gas puffing rate GAS, ohmic heating power P Ω , neutral beam injection (NBI) power P NBI , electron cyclotron heating (ECH) power P ECH , and ion cyclotron heating (ICH) power P ICH . |B 0 | [T] 6 8 10 q 95 1 2 3 4 5 t [s] 
 
 Appendix:  


Decision: Thanks to reviewers for their honest feedback. Reviewer LpgF  made great suggestions on re-organizing the paper contents and address more useful information in the main text while moving detailed info to appendix. The authors are required to address their questions/comments in the final version. The study seems to be interesting, in terms of using Neural ODEs to simulate Tokamak plasma dynamics. Under the condition that authors will re0structure their paper for better readability, we decide to accept this as poster.

Reviews: review: title: title: Better baselines are needed to accurately judge the proposed dynamics model's performance
main_review: main_review: Paper proposes a plasma dynamics model with a parametric module to estimate particle and energy diffusivities. The dynamical system is solved with NeuralODE and the parameters are optimized through standard backpropagation through the solver. 

Since it is not my expertise, I am unable to judge the novelty and correctness of the proposed dynamics model in Section 2.
Feedback on the simulation results in Section 4:
1. Predictions for particle density $n$ seem quite far from the ground truth in the tested examples (e.g., in Figures 4, 9). But there is a mismatch with the reported MSE values in Table 1: they seem too low when compared with these predictions and the fact that $n$ is in the range of $10^{19}$. 
2. The baseline (original model) seems to be the one with randomly initialized parameters. If that is indeed the case, it is a weak baseline. A better comparison would be to use other dynamical models proposed in the literature or purely neural network based methods. Without an appropriate baseline, it is hard to judge whether the reported predictions/errors are reasonable for this task.
title: title: Review of "Application of Neural Ordinary Differential Equations for Tokamak Plasma Dynamics Analysis"
main_review: main_review: This paper presents a study in which Neural ODEs are used to simulate energy transfer processes in tokamaks.  The proposed models are trained using experimental data.  Validation of the proposed data-driven models is performed under various auxiliary heating conditions.

This paper is interesting and generally well-written.  I really like that the authors are using experimental data to train/validate their models, something that is not generally done in the literature.  I do feel that there is a bit of misbalance in the paper.  The very complex equations in sections 2.2-2.3 could go in an appendix.  I would recommend more detail about the architecture and how it was optimized / trained in the body of the paper, since this is for an AI workshop. and since you are showing a tremendous improvement by optimizing your model in Table 1.  There is too much right now in the Appendix - I would suggest moving some of the figures / results here into the body of the paper, since you should really be focusing on the AI/ML.

I cannot comment on the novelty of the proposed methods for the tokamak plasma application, as this is not my area of expertise.