# Fourier Neural Operator for Plasma Modelling

Vignesh Gopakumar
UK Atomic Energy Authority
Oxford, OX14 3EB
<EMAIL_ADDRESS>
Stanislas Pamela
UK Atomic Energy Authority
Oxford, OX14 3EB
<EMAIL_ADDRESS>
Lorenzo Zanisi
UK Atomic Energy Authority
Oxford, OX14 3EB
<EMAIL_ADDRESS>
MAST Team
UK Atomic Energy Authority
Oxford, OX14 3EB
Zongyi Li
Caltech
California, 91125
<EMAIL_ADDRESS>
Anima Anandkumar
Caltech
California, 91125
<EMAIL_ADDRESS>

###### Abstract

Predicting plasma evolution within a Tokamak is crucial to building a
sustainable fusion reactor. Whether in the simulation space or within the
experimental domain, the capability to forecast the spatio-temporal evolution
of plasma field variables rapidly and accurately could improve active control
methods on current tokamak devices and future fusion reactors. In this work,
we demonstrate the utility of using Fourier Neural Operator (FNO) to model the
plasma evolution in simulations and experiments. Our work shows that the FNO
is capable of predicting magnetohydrodynamic models governing the plasma
dynamics, 6 orders of magnitude faster than the traditional numerical solver,
while maintaining considerable accuracy (NMSE $\sim 10^{-5})$. Our work also
benchmarks the performance of the FNO against other standard surrogate models
such as Conv-LSTM and U-Net and demonstrate that the FNO takes significantly
less time to train, requires less parameters and outperforms other models. We
extend the FNO approach to model the plasma evolution observed by the cameras
positioned within the MAST spherical tokamak. We illustrate its capability in
forecasting the formation of filaments within the plasma as well as the heat
deposits. The FNO deployed to model the camera is capable of forecasting the
full length of the plasma shot within half the time of the shot duration.

##  1 Introduction

Predicting the evolution of plasma dynamics is crucial to building sustainable
fusion devices as it allows us to anticipate the plasma behaviour and take
real-time decisions on the trajectory of engineering input control parameters.
In fusion research, prediction of future plasma states is mostly undertaken
using numerical modelling codes such as BOUT++ [1], JOREK [2], JINTRAC [3] and
SOLPS [4]. However, these codes are computationally intensive, taking weeks if
not months for effective simulations to be run on high-performance computers
[1]. They are also often complex in terms of software infrastructure and
advanced solver library dependencies, making them difficult to deploy for
rapid iterative modelling.

Machine learning has offered an alternative data-driven route for obtaining
quick, inexpensive approximations to numerical simulations [5]. These machine
learning models are often physics-informed allowing us to create networks that
are interpretable and can generalise [6]. Over the past few years, the
research scope has expanded to explore machine learning based models aimed at
plasma physics[7] [8]. We build upon this existing work and extend it with
applications of neural operator learning to model plasma within fusion
devices.

Recently, a new class of data-driven models known as neural operators aim to
directly learn the solution operator of the PDE in a mesh-invariant manner [9,
10]. Unlike standard deep learning models, such as convolutional neural
networks from computer vision, neural operators are invariant to
discretization and hence, better suited for solving PDEs. They generalize the
previous finite-dimensional neural networks to learn operator mapping between
infinite-dimensional function spaces. Examples include Fourier Neural Operator
(FNO) [11], graph neural operator [12], and DeepONet [13]. We consider FNO in
this paper due to its superior cost-accuracy tradeoff [14].

Our work demonstrates the performance of Fourier Neural Operator models [11]
on solving magnetohydrodynamic (MHD) equations, and benchmark it against more
traditional surrogate model architectures such as a Convolutional-LSTM and
U-Net. In addition to numerical modelling, we demonstrate that FNOs can be
trained to predict the evolution of experimental plasmas as observed by High
Speed cameras on the MAST spherical tokamak [15]. The ability to predict
future plasma states based on camera imaging, as well as other plasma
diagnostics, could allow us to design efficient control feedback loops for
future fusion reactors. To the best of our knowledge, this the first time that
Neural Operator learning has been applied to plasma scenarios.

##  2 Problem Setting

###  2.1 MHD Simulations of Plasma Blobs

The first dataset used to explore the FNO method is a simplified Magneto-
hydrodynamics (MHD) model in toroidal geometry. In practice, advanced
simulations of plasma instabilities are run with the JOREK code, using complex
MHD models, and realistic geometries [2], however for this paper a simplified
setup was used as a proof-of-principle. The dataset is obtained using the
JOREK code [2], with a physics model similar to the Reduced-MHD model
described in [2], but with electrostatic and isothermal constraints, such that
only the density $\rho$, the electric potential $\Phi$, and the toroidal
vorticity $w$ fields are evolved. In toroidal geometry, if a toroidally
axisymmetric density blob is initialised on top of a low background density,
in the absence of any plasma current to hold the density blob in place, the
pressure gradient term in the momentum equation acts as a buoyancy effect and
causes the blob to move outwards.

This is used to create a dataset of density blobs moving radially inside a
toroidally axisymmetric slab grid, bounded by Dirichlet conditions. The blobs
evolve until they reach the wall and eventually disappear through both
convective mixing and diffusion. The density diffusion and viscosity are
fixed, while the blob’s initial conditions are varied, both in position,
width, and amplitude. A set of 120 simulations are run (100 for training, 20
for testing), all with the same time-duration achieved with 1000 simulation
timesteps, which allows for the blobs to travel all the way to the boundary
and diffuse almost entirely. The number of simulation timesteps are down-
sampled to 100 slices for the training set. Likewise, the 2D poloidal grid
resolution, which is 200$\times$200 bi-cubic finite-elements, is down-sampled
to a regular grid of 100$\times$100\. The simulations were run on the EOSC-hub
cloud, using 16 cores per run, and each simulation took approximately 160 core
hours.

MHD modelling is of importance as it helps us understand the evolution of
filaments that form across the flux surfaces of the tokamak as shown in the
right of figure 1 [16]. Accurately capturing the dynamics and evolution of the
filaments help us characterise Edge-Localised Modes (ELMs), a nonlinear
instability. ELMs can contribute damage to the reactor wall, potentially
initiate other instabilities and it is crucial to be able to gain control over
them. Through the course of this paper, the simulation modelling we undertake
is that of a relatively simple physics structures, designed to serve as a
proof-of-concept, outlining the utility of the FNO in modelling MHD scenarios.
These simulations modelled in figure 2, show the diffusion of a blob which are
a simplified scenario as that of the full ELM simulations as shown on the
right of figure 1.

###  2.2 Fast-Camera Images of the MAST Tokamak

The Mega Ampere Spherical Tokamak (MAST) was a nuclear fusion device situated
at the Culham Centre for Fusion Energy in the UK, where initial set of
experiments were demonstrated on a Tokamak design with a narrow aspect ratio,
demonstrating better efficiency over the magnetic field control [17]. MAST has
now been succeeded by MAST-Upgrade (MAST-U), which is currently operational
and experimented on for better control strategies and design of the exhaust
for potential future reactors [15].

MAST was equipped with fast camera imagining diagnostics at several points
across the Tokamak for capturing images within the visible spectrum,
displaying the plasma evolution. These Photron cameras [18] are essential
diagnostics that contributed to fundamental understandings of key plasma
phenomena in tokamaks in recent years [19]. While most of the camera data has
been used for qualitative analysis in the past, it can be exploited using more
advanced methods to provide statistical insight into plasma turbulence [20],
as well as global instabilities (disruptions) [21]. Within the scope of this
work, we focus on two camera view, one that looks at the central solenoid and
the plasma around that to the walls, providing a holistic view of the plasma,
and the so-called divertor-view, looking at the bottom of the plasma, where
energy and impurities travel from the plasma to the wall (the plasma exhaust
region). Figure-1 shows a view of the main plasma as seen by the camera(left),
together with a synthetic rendering (right) of the camera, used for the
calibration simulating the MHD physics. The cameras have a temporal resolution
of 1.2ms, with spatial resolution varying according to the desired calibration
set up for each experiment. They are tuned to visible wavelength and therefore
capture mostly the Balmer Dα light emitted mostly from the plasma edge. As
seen on Figure-1, the fast-visible cameras are able to capture complex
dynamics of filamentary eruptions at the plasma edge, the so-called Edge-
Localised Modes (ELMs) [19].

Figure 1: Left: Fast camera image of a MAST plasma during an Edge-Localised-
Mode, where plasma filaments are ejected outwards due to MHD instabilities at
the plasma edge (figure reproduced from [22]). Right: The evolution of the
filamentary structures during the multi-mode ELM simulation imaged with a
synthetic fast camera diagnostic (time given in ms). Full-scale MHD
simulations as shown in this figure demonstrates how MHD behaviours are
captured by the visible camera. [16]

##  3 Methods

###  3.1 Fourier Neural Operator

The FNO approach shown in [11] learns a mapping between two infinite
dimensional spaces from a finite collection of observed input-output pairs.
This is achieved by way of a Fourier layer, as displayed in [11]. FNO is
constructed by stacking multiple Fourier layers on top of each other. Each
Fourier layer is composed of two learnable weight matrices: $R$ and $W$. $R$
learns the mapping behaviour within the Fourier space, while $W$ learns the
mapping required within the input Euclidean space. The weight matrix $R$
within the Fourier layer enables convolution within the Fourier space,
allowing the network to be parameterised directly in it. The output of a
Fourier Layer can be expressed as:

$y=\sigma\bigg{(}\mathcal{F}^{-1}\big{(}R\mathcal{F}(x)\big{)}+Wx\bigg{)}$ (1)

where, $x$ is the input, $y$ the output and $\sigma$ is the nonlinear
activation function. $\mathcal{F}$ and $\mathcal{F}^{-1}$ represents the
Fourier and inverse Fourier transform respectively.

In this work, we will use the FNO-2d configuration, introduced in [11], where
a 2-d Fourier Neural Operator that only convolves in space is deployed in an
auto-regressive manner. The FNO takes in as its input the field values across
a 2D grid for an initial set of time instances (T_in) along with the grid
discretisation in both dimensions. It outputs a set of later time instances
(step) of the field value across the same grid. The outptut of the FNO,
coupled with the later time steps from the input are then fed back in an auto-
regressive manner to estimate further field evolution in time, until the
desired time instance (T_out) is reached. For further information regarding
the theory behind operator learning as ascribed by the FNO and its
implementation in various case settings, we invite the author to refer to the
original work in [11]. We have chosen the 2D configuration of the FNO as this
gives us flexibility in deciding the length of the temporal evolution.

The training conditions of the FNO over MHD data is given in 4.1.

###  3.2 U-Net and Convolutional LSTM

Initially demonstrated in [23], further demonstrated in [24] and further
evolved in [25], U-Net architectures have been instrumental in the design and
structure of surrogate models for a vast variety of scenarios involving
spatio-temoporal evolution. U-Net architecture is obtained by combining
convolutional auto-encoders [26] with skip connections [27] at the decoder.
The layers at each level of the encoder are skip connected and concatenated
with corresponding layers at the decoder. The architecture we deployed for
modelling reduced MHD was inspired by the work demonstrated in the PDE
benchmark dataset: PDEBench [28].

Initially demonstrated in [29], the Convolutional LSTM (ConvLSTM) replaces the
fully-connected layers in a LSTM [30] with convolutional layers. This allows
for the representation of additional spatially dependent layers within the
recurrent structure of the model. We devise a ConvLSTM consisting of 5
ConvLSTM cells and train under the same conditions as for the FNO mentioned in
section 4.1. For a detailed layout of the architecture deployed for the
ConvLSTM, refer to Appendix D.

##  4 Results

###  4.1 FNO to Predict MHD Simulations

We construct and train independent 2D Fourier Neural Operator to model the
spatio-temporal evolution of density ($\rho$), potential ($\phi$) and
vorticity ($w$) associated with the plasma blob experiment as described in
section 2.1. The architecture of the FNO was borrowed from [11] and is further
laid out in the appendix C. Each model was trained on a single A100 GPU at the
Cambridge Centre for Data Driven Discovery (CSD3).

Each FNO is trained to learn the map between 20 time instances (T_in=20) of a
given field to the next 5 time instances (step=5). The output 5 time instances
are concatenated with its previous 15 time instances to form the next 20 time
instances of the input. This auto-regressive temporal evolution of the field
is carried on until the $70^{th}$ time instance (T_out=70). Each Fourier layer
is characterised by a width of 32 and learning up to 16 modes in the Fourier
space. Prior to training, the training dataset, built as explained in in
section 2.1 is normalised using a linear range scaling to lie between -1 and 1
for each of the field variables. All models are trained using the Adam
optimizer [31] for 500 epochs with a batch size of 5. The learning rate is
initially set to 0.001 and scheduled to decrease by half after every 100
epochs. Relative LP loss function was utilised for training the model.

As shown in figure 2, individual FNOs crafted to model the evolution of each
field variable reproduce well the dynamics of the reduced MHD simulations.
Figure 2a represents the time evolution of the density of the solution from
JOREK and the FNO at the $21^{st}$, $45^{th}$ and $70^{t}h$ time instance
having been fed in the first 20 time instances. Figure 2b shows the same for
the evolution of the vorticity of the reduced MHD model (note we do not show
the electric potential $\Phi$ in Figure-2, since the vorticity $w$ is
effectively the laplacian of $\Phi$). Our model achieved a Normalised MSE of
4.2$\times$10-5 for mapping the evolution of electric potential, 5$\times$10-5
for density and 4.6$\times$10-5 for that of the vorticity. The FNO is able to
evaluate the full solution in 0.025 seconds. Plots showing the contrast of the
FNO output with that of JOREK by way of discrepancy plots can be found in
Appendix A.

(a) Density

(b) Vorticity

Figure 2: Temporal evolution of (a) the density and (b) the vorticity
variables describing the plasma evolution as obtained using the JOREK code
(top of each image) and that of the individually trained FNOs (bottom of each
figure). The spatial domain is given in toroidal geometry characterised by R
in the x-axis and Z in the y-axis.

U-Net and Conv-LSTM : Both the U-Net and the ConvLSTM with achitectures
mentioned in appendices E and D were trained under the same conditions as the
FNO (as described above), and with the same auto-regressive strategy. Each
model was trained on a single A100 GPU at the Cambridge Centre for Data Driven
Discovery (CSD3). As demonstrated in table 1, FNO takes significantly less
time to train, requires less parameters and outperforms the ConvLSTM and the
U-Net in modelling the diffusion of a blob of plasma described by reduced MHD.
For conciseness we have only showed the results from the models mapping the
evolution of potential in 1, models deployed for learning other field variable
have similar performance estimates, model composition and time consumption.
For seeing the performance across each field variable and their comparative
plots refer Appendix B.

Table 1: Comparative Performance across various Models for Potential. Numerical Simulation ran on 10 cpu cores, training and inference of ML models on Nvidia A100 cards. Model | Train Time (mins) | # Parameters | NMSE | Inference Time(s)
---|---|---|---|---
JOREK | - | - | - | 36000
FNO | 13 | 2106981 | 4.245e-05 | 0.025
ConvLSTM | 481 | 4874648 | 1.12e-03 | 0.14
U-Net | 26 | 7931909 | 4.48e-04 | 0.033

###  4.2 FNO to Predict Experimental Camera Data

In order to create a data-driven digital twin capable of modelling the
evolution of plasma as witnessed by the fast-cameras on MAST we invoked the 2D
FNO architecture deployed in a auto-regressive framework as we had done in
modelling the plasma described by reduced MHD. The FNO is well-suited for this
task as the fast-camera images essentially capture the larger MHD behaviours
exhibited by the experimental plasma. The utilisation of a Fourier based model
is well justified as it allows to extract the modal behaviours within the
plasma, consistent with it’s periodic properties. The camera calibrated for
specific experimental runs capture the information of the plasma onto a
uniform pixel grid, allowing us to perform Fourier transformations over them.
By using calcam, a calibration tool developed by the European Atomic Energy
Community [32] combined with 3D CAD models of MAST, we are able to determine
the domain range covered by each camera calibration. The upper and lower
limits of the domain range of these images are then fixed onto uniform grids
in both the R and Z axis, allowing us to create uniform grid discretisations
underlying each camera image.

Figure 3: Temporal evolution of the plasma around the central solenoid within
MAST as perceived by the main camera. The top layer shows the actual data
obtained from the experiment, whereas the bottom layer displays the modelling
performed by the FNO.

The FNO serving as the digital twin takes in the first 10 time instances
($T_{in}$=10) along with the calibrated grid discretisations to be able to
output the next time instance (step=1). The auto-regressive loop of feeding
the output back along with a portion of the input and the grid discretisations
is run until the $60^{th}$ time instance. The plasma takes in the first 12 ms
of the plasma within the tokamak to be able to output the next 60 ms of
evolution.

Camera viewing the central solenoid : The evolution of the plasma within the
Tokamak as seen by the main camera, along with the modelling performed by the
FNO can be seen in figure 3. The central solenoid serves the backbone of the
MAST magnet system, generating majority of the magnetic fluxes as well as the
plasma current. The main camera view as shown in figure 1 provides a panoramic
view of the reactor, showing the poloidal cross-sectional layout on both sides
of the central solenoid. The main camera has a spatial resolution of 448x640
for the experiments that we had sampled for. We had sampled a total of 59
plasma shots from the MAST database, used 50 for training and 9 for testing.
The model achieved a Normalised MSE of 0.03997.

Camera viewing the divertor : The evolution of the plasma within the Tokamak
as seen by the divertor camera, along with the modelling performed by the FNO
can be seen in figure 4. The divertor of a tokamak serves as the exhaust of
the device, where the magnetic flux surfaces converge leading to the expulsion
of the impurities from it. The divertor camera looks at the lower divertor of
MAST with a spatial resolution of 400x512. We had sampled a total of 86 plasma
shots from the MAST database, used 75 for training and 11 for testing. The
model achieved a Normalised MSE of 0.01384.

Figure 4: Temporal evolution of the plasma at the divertor within MAST as
perceived by the divertor camera. The top layer shows the actual data obtained
from the experiment, whereas the bottom layer displays the modelling performed
by the FNO.

The FNO deployed for modelling the camera data accurately preserves the global
features of the plasma evolution. It is able to clearly characterise the
locations and intensities of the heat fluxes onto the solenoid and the
divertors, allowing us to anticipate the heat flux in advance and initiate the
appropriate actuator mechanism to minimise its impact onto the various plasma
facing components. The FNO is also capable of modelling the filaments that
arise during the plasma pulse, allowing us to predict the instances of Edge
Localised Modes (ELM) that play a crucial role in the stability of the plasma.
The accurate prediction of filaments also help us forecast the instance of the
L-H transition, allowing us to potentially initiate the fusion reaction more
efficiently. While capturing the global characteristics, the FNO performs
poorly in capturing the localised phenomena with matching pixel-pixel
accuracy.

##  5 Conclusion and Future Work

Through this work we demonstrate the utility of using Fourier Neural Operator
as both a surrogate model and as a digital twin for modelling plasma
evolution. The FNO deployed as a surrogate model is capable of solving
magneto-hydrodynamics governing the transport of plasma within a fusion
device, accurately and 6 orders of magnitude faster than the numerical solver.
The FNO demonstrates utility in modelling each associated field with varying
domain ranges. Being utilised as a digital twin, the FNO is capable of
forecasting the evolution of the plasma as seen through the cameras onboard a
fusion device. The FNO camera models are capable of identifying the regions of
high heat deposits, filament structures, while predicting effectively the
plasma evolution trajectory. The FNO for the camera can predict the plasma
evolution in real-time and offers us potential in being deployed within the
planning stage for various control scenarios.

In continuing with this line of work, we will be exploring the utility of the
FNO in more complex MHD cases while allowing for modelling multiple vector
fields simultaneously. In this work we had chosen a MHD case with a simplified
geometry, however we will be looking to expand to unstructured grids as
demonstrated in [33]. We shall also be expanding our work on the camera data
with a larger dataset covering broader experimental scenarios with varying
time lengths. Once proofed on a larger dataset we shall be integrating the
prediction module with a plasma confinement mode classifier, allowing us to
forecast the time instance of an L-H transition [34] within a tokamak.

## References

  * [1] Fabio Riva, Fulvio Militello, Sarah Elmore, John T Omotani, Ben Dudson, Nick R Walkden, and the MAST team.  Three-dimensional plasma edge turbulence simulations of the mega ampere spherical tokamak and comparison with experimental measurements.  Plasma Physics and Controlled Fusion, 61(9):095013, aug 2019.
  * [2] M. Hoelzl, G.T.A. Huijsmans, S.J.P. Pamela, M. Bécoulet, E. Nardon, F.J. Artola, B. Nkonga, C.V. Atanasiu, V. Bandaru, A. Bhole, D. Bonfiglio, A. Cathey, O. Czarny, A. Dvornova, T. Fehér, A. Fil, E. Franck, S. Futatani, M. Gruca, H. Guillard, J.W. Haverkort, I. Holod, D. Hu, S.K. Kim, S.Q. Korving, L. Kos, I. Krebs, L. Kripner, G. Latu, F. Liu, P. Merkel, D. Meshcheriakov, V. Mitterauer, S. Mochalskyy, J.A. Morales, R. Nies, N. Nikulsin, F. Orain, J. Pratt, R. Ramasamy, P. Ramet, C. Reux, K. Särkimäki, N. Schwarz, P. Singh Verma, S.F. Smith, C. Sommariva, E. Strumberger, D.C. van Vugt, M. Verbeek, E. Westerhof, F. Wieschollek, and J. Zielinski.  The jorek non-linear extended mhd code and applications to large-scale instabilities and their control in magnetically confined fusion plasmas.  Nuclear Fusion, 61(6):065001, may 2021.
  * [3] Michele Romanelli, Gerard Corrigan, Vassili Parail, Sven Wiesen, Roberto Ambrosino, Paula Da Silva Aresta Belo, Luca Garzotti, Derek Harting, Florian Köchl, Tuomas Koskela, Laura Lauro-Taroni, Chiara Marchetto, Massimiliano Mattei, Elina Militello-Asp, Maria Filomena Ferreira Nave, Stanislas Pamela, Antti Salmi, Pär Strand, Gabor Szepesi, and EFDA-JET Contributors.  Jintrac: A system of codes for integrated simulation of tokamak scenarios.  Plasma and Fusion Research, 9:3403023–3403023, 2014.
  * [4] S. Wiesen, D. Reiter, V. Kotov, M. Baelmans, W. Dekeyser, A.S. Kukushkin, S.W. Lisgo, R.A. Pitts, V. Rozhansky, G. Saibene, I. Veselova, and S. Voskoboynikov.  The new solps-iter code package.  Journal of Nuclear Materials, 463:480–484, 2015.  PLASMA-SURFACE INTERACTIONS 21.
  * [5] Alexander Lavin, David Krakauer, Hector Zenil, Justin Gottschlich, Tim Mattson, Johann Brehmer, Anima Anandkumar, Sanjay Choudry, Kamil Rocki, Atılım Güneş Baydin, Carina Prunkl, Brooks Paige, Olexandr Isayev, Erik Peterson, Peter L. McMahon, Jakob Macke, Kyle Cranmer, Jiaxin Zhang, Haruko Wainwright, Adi Hanuka, Manuela Veloso, Samuel Assefa, Stephan Zheng, and Avi Pfeffer.  Simulation intelligence: Towards a new generation of scientific methods, 2021.
  * [6] George Em Karniadakis, Ioannis G. Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang.  Physics-informed machine learning.  Nature Reviews Physics, 3(6):422–440, Jun 2021.
  * [7] K. L. van de Plassche, J. Citrin, C. Bourdelle, Y. Camenen, F. J. Casson, V. I. Dagnelie, F. Felici, A. Ho, and S. Van Mulders and.  Fast modeling of turbulent transport in fusion plasmas using neural networks.  Physics of Plasmas, 27(2):022310, feb 2020.
  * [8] Vignesh Gopakumar and D Samaddar.  Image mapping the temporal evolution of edge characteristics in tokamaks using neural networks.  Machine Learning: Science and Technology, 1(1):015006, feb 2020\.
  * [9] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Neural operator: Learning maps between function spaces.  arXiv preprint arXiv:2108.08481, 2021.
  * [10] Zongyi Li, Hongkai Zheng, Nikola Borislavov Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Andrew Stuart, Kamyar Azizzadenesheli, and Anima Anandkumar.  Physics-informed neural operator for learning partial differential equations, 2022.
  * [11] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Fourier neural operator for parametric partial differential equations, 2020.
  * [12] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Multipole graph neural operator for parametric partial differential equations, 2020.
  * [13] Lu Lu, Pengzhan Jin, and George Em Karniadakis.  Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators.  arXiv preprint arXiv:1910.03193, 2019.
  * [14] Maarten De Hoop, Daniel Zhengyu Huang, Elizabeth Qian, and Andrew M Stuart.  The cost-accuracy trade-off in operator learning with neural networks.  arXiv preprint arXiv:2203.13181, 2022.
  * [15] MAST tokamak.
  * [16] S.F. Smith, S.J.P. Pamela, A. Fil, M. Hölzl, G.T.A. Huijsmans, A. Kirk, D. Moulton, O. Myatra, A.J. Thornton, and H.R. Wilson and.  Simulations of edge localised mode instabilities in MAST-u super-x tokamak plasmas.  Nuclear Fusion, 60(6):066021, may 2020.
  * [17] I.T. Chapman, J. Adamek, R.J. Akers, S. Allan, L. Appel, O. Asunta, M. Barnes, N. Ben Ayed, T. Bigelow, W. Boeglin, J. Bradley, J. Brünner, P. Cahyna, M. Carr, J. Caughman, M. Cecconello, C. Challis, S. Chapman, J. Chorley, G. Colyer, N. Conway, W.A. Cooper, M. Cox, N. Crocker, B. Crowley, G. Cunningham, A. Danilov, D. Darrow, R. Dendy, A. Diallo, D. Dickinson, S. Diem, W. Dorland, B. Dudson, D. Dunai, L. Easy, S. Elmore, A. Field, G. Fishpool, M. Fox, E. Fredrickson, S. Freethy, L. Garzotti, Y.C. Ghim, K. Gibson, J. Graves, C. Gurl, W. Guttenfelder, C. Ham, J. Harrison, D. Harting, E. Havlickova, J. Hawke, N. Hawkes, T. Hender, S. Henderson, E. Highcock, J. Hillesheim, B. Hnat, J. Holgate, J. Horacek, J. Howard, B. Huang, K. Imada, O. Jones, S. Kaye, D. Keeling, A. Kirk, I. Klimek, M. Kocan, H. Leggate, M. Lilley, B. Lipschultz, S. Lisgo, Y.Q. Liu, B. Lloyd, B. Lomanowski, I. Lupelli, G. Maddison, J. Mailloux, R. Martin, G. McArdle, K. McClements, B. McMillan, A. Meakins, H. Meyer, C. Michael, F. Militello, J. Milnes, A.W. Morris, G. Motojima, D. Muir, E. Nardon, V. Naulin, G. Naylor, A. Nielsen, M. O’Brien, T. O’Gorman, Y. Ono, H. Oliver, S. Pamela, L. Pangione, F. Parra, A. Patel, W. Peebles, M. Peng, R. Perez, S. Pinches, L. Piron, M. Podesta, M. Price, M. Reinke, Y. Ren, C. Roach, J. Robinson, M. Romanelli, V. Rozhansky, S. Saarelma, S. Sangaroon, A. Saveliev, R. Scannell, A. Schekochihin, S. Sharapov, R. Sharples, V. Shevchenko, S. Silburn, J. Simpson, J. Storrs, Y. Takase, H. Tanabe, H. Tanaka, D. Taylor, G. Taylor, D. Thomas, N. Thomas-Davies, A. Thornton, M. Turnyanskiy, M. Valovic, R. Vann, N. Walkden, H. Wilson, L.V. Wyk, T. Yamada, S. Zoletnik, MAST, and MAST Upgrade Teams.  Overview of mast results.  Nuclear Fusion, 55(10):104008, jun 2015.
  * [18] Photron Limited.  Photron high speed cameras for slow motion analysis, 2022.
  * [19] A Kirk, N Ben Ayed, G Counsell, B Dudson, T Eich, A Herrmann, B Koch, R Martin, A Meakins, S Saarelma, R Scannell, S Tallents, M Walsh, H R Wilson, and the MAST team.  Filament structures at the plasma edge on mast.  Plasma Physics and Controlled Fusion, 48(12B):B433, nov 2006.
  * [20] Nicholas Walkden, Fabio Riva, James Harrison, Fulvio Militello, Thomas Farley, John Omotani, and Bruce Lipschultz.  The physics of turbulence localised to the tokamak divertor volume.  Communications Physics, 5(1):139, Jun 2022.
  * [21] C.Ham et al.  Insights on disruption physics in mast using high speed visible camera data.  In IAEA Second Technical Meeting on Plasma Disruptions and their Mitigation, 2022.
  * [22] G. Hommen, M. de Baar, P. Nuij, G. McArdle, R. Akers, and M. Steinbuch.  Optical boundary reconstruction of tokamak plasmas for feedback control of plasma position and shape.  Review of Scientific Instruments, 81:113504, may 2010.
  * [23] Mateusz Buda, Ashirbani Saha, and Maciej A. Mazurowski.  Association of genomic subtypes of lower-grade gliomas with shape features automatically extracted by a deep learning algorithm.  Computers in Biology and Medicine, 109:218–225, 2019.
  * [24] Zhihao Jiang, Pejman Tahmasebi, and Zhiqiang Mao.  Deep residual u-net convolution neural networks with autoregressive strategy for fluid flow predictions in large-scale geosystems.  Advances in Water Resources, 150:103878, 2021.
  * [25] Gege Wen, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M. Benson.  U-fno – an enhanced fourier neural operator-based deep-learning model for multiphase flow, 2021.
  * [26] G. E. Hinton and R. R. Salakhutdinov.  Reducing the dimensionality of data with neural networks.  Science, 313(5786):504–507, 2006.
  * [27] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770–778, 2016.
  * [28] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert.  PDEBench Datasets, 2022.
  * [29] Xingjian Shi, Zhourong Chen, Hao Wang, Dit-Yan Yeung, Wai-kin Wong, and Wang-chun Woo.  Convolutional lstm network: A machine learning approach for precipitation nowcasting, 2015.
  * [30] Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  Neural Comput., 9(8):1735–1780, nov 1997.
  * [31] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization, 2014.
  * [32] Scott Silburn, James Harrison, Tom Farley, Jordan Cavalier, Sam Van Stroud, Jamie McGowan, Augustin Marignier, Emily Nurse, Christian Gutschow, Mark Smithies, Alasdair Wynn, and Rhys Doyle.  Calcam.  July 2022.
  * [33] Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar.  Fourier neural operator with learned deformations for pdes on general geometries, 2022.
  * [34] K H Burrell, T N Carlstrom, E J Doyle, D Finkenthal, P Gohil, R J Groebner, D L Hillis, J Kim, H Matsumoto, R A Moyer, T H Osborne, C L Rettig, W A Peebles, T L Rhodes, H StJohn, R D Stambaugh, M R Wade, and J G Watkins.  Physics of the l-mode to h-mode transition in tokamaks.  Plasma Physics and Controlled Fusion, 34(13):1859, dec 1992.

##  Appendix A Error Plots

(a) Electric Potential

(b) Density

(c) Vorticity

Figure 5: Temporal evolution of the field variables describing the plasma
evolution as obtained using the code suite JOREK (top of each image) and that
of the individually trained FNOs (middle of each figure) and the absolute
error across both (bottom of each figures). The spatial domain is given in
toroidal geometry characterised by R in the x-axis and Z in the y-axis. Figure
(a) represents the temporal evolution of electric potential, figure (b) shows
density and figure (c) that of the vorticity.

##  Appendix B Performance plots across various models

Figure 6: Evolution of the Electric Potential modelled by various surrogate models. Each surrogate model takes in the first 20 time instances to output in total the next 50 time instances. The Initial distribution denotes the 21st time instance, Middle the 45th and the Final 70th time instance. The top layer denotes the actual solution, the second layer that modelled by the FNO, the third by the U-Net and the bottom layer the output from the Conv-LSTM.  Figure 7: Evolution of the Density modelled by various surrogate models. Each surrogate model takes in the first 20 time instances to output in total the next 50 time instances. The Initial distribution denotes the 21st time instance, Middle the 45th and the Final 70th time instance. The top layer denotes the actual solution, the second layer that modelled by the FNO, the third by the U-Net and the bottom layer the output from the Conv-LSTM.  Figure 8: Evolution of the  vorticity modelled by various surrogate models. Each surrogate model takes in the first 20 time instances to output in total the next 50 time instances. The Initial distribution denotes the 21st time instance, Middle the 45th and the Final 70th time instance. The top layer denotes the actual solution, the second layer that modelled by the FNO, the third by the U-Net and the bottom layer the output from the Conv-LSTM.  Table 2: Comparative Performance across various Models Model | Field | Training Time (mins) | # Parameters | NMSE
---|---|---|---|---
FNO |  | 13 | 2106981 | 4.245e-05
ConvLSTM | Potential | 481 | 4874648 | 1.12e-03
U-Net |  | 26 | 7931909 | 4.48e-04
FNO |  | 13 | 2106981 | 5.03e-05
ConvLSTM | Density | 481 | 4874648 | 7.93e-03
U-Net |  | 26 | 7931909 | 1.34e-04
FNO |  | 13 | 2106981 | 4.65e-05
ConvLSTM | Vorticity | 479 | 4874648 | 7.98e-05
U-Net |  | 26 | 7931909 | 5.45e-05

##  Appendix C FNO 2D Architeture

The architecture of the FNO 2D deployed in an auto-regressive manner for
modelling the Reduced MHD equations. The model initially lifts the input data
into a higher dimensional space using a linear layer. This is followed by 4
Fourier layers that learn the necessary modal mappings within the spatial
domain, before being projected onto the lower dimensional space by a series of
linear layers. Fourier 2d represents the 2D Fourier operator, Conv2d
represents the 2D convolutional layer, Add operation adds the outputs from the
Fourier layer and the convolutional layer together, GELU denotes the
activation function: Gaussian Error Linear Unit.

Table 3: Architecture of the FNO 2D deployed for modelling Reduced MHD Part | Layer | Output Shape
---|---|---
Input | - | (5, 100, 100, 22)
Lifting | Linear | (5, 100, 100, 32)
Fourier 1 | Fourier2d/Conv2d/Add/GELU | (5, 32, 100, 100)
Fourier 2 | Fourier2d/Conv2d/Add/GELU | (5, 32, 100, 100)
Fourier 3 | Fourier2d/Conv2d/Add/GELU | (5, 32, 100, 100)
Fourier 4 | Fourier2d/Conv2d/Add/GELU | (5, 32, 100, 100)
Projection 1 | Linear | (5, 100, 100, 128)
Projection 2 | Linear | (5, 100, 100, 5)

##  Appendix D ConvLSTM Architecture

The architecture of the 2D Convolutional-LSTM deployed for modelling the
Reduced MHD equations. The model consists of 5 Convolutional-LSTM Cells with
features scaled and lifted to 256, which is progressively scaled down to the
desired output in the following manner ([256, 128, 64, 32, 5]). Each Conv-LSTM
Cell consists of an input gate, forget gate and an output gate, where the
weight matrices associated with each gate are given by 2D convolutional layers
rather than being linear layers. For an understanding of the equations
governing the various gates in an LSTM we refer the author to [30].

Table 4: Architecture of the ConvLSTM deployed for modelling Reduced MHD Part | Layer | Output Shape
---|---|---
Input | - | (5, 20, 100, 100)
Conv-LSTM Cell 1 | Conv2D/Sigmoid/Hadamard Product/Tanh | (5, 256, 100, 100)
Conv-LSTM Cell 2 | Conv2D/Sigmoid/Hadamard Product/Tanh | (5, 128, 100, 100)
Conv-LSTM Cell 3 | Conv2D/Sigmoid/Hadamard Product/Tanh | (5, 64, 100, 100)
Conv-LSTM Cell 4 | Conv2D/Sigmoid/Hadamard Product/Tanh | (5, 32, 100, 100)
Conv-LSTM Cell 5 | Conv2D/Sigmoid/Hadamard Product/Tanh | (5, 5, 100, 100)

##  Appendix E U-Net Architecture

The architecture of the 2D U-Net deployed in an auto-regressive manner for
modelling the Reduced MHD equations. The model consists of 4 encoding and 4
decoding layers. Each encoding layer is given by a 2D convolution (Conv2d), 2D
Batch Normalisation (BatchNorm2d), non-linear activation (Tanh), followed by
max-pooling layer (MaxPool2d). Successive to the encoding layers, we establish
a bottle neck layer which is identical to an encoding layer but without the
pooling element. In each Decoder, the information is up-scaled via a 2D
transposed convolution (ConvTranspose2D), the output of which is concatenated
with that of the output of the corresponding encoder and then passed through a
series of 2D convolutional layers. The final output is rescaled to the desired
dimensions using 2D convolutions.

Table 5: Architecture of the 2D U-Net deployed for modelling Reduced MHD Part | Layer | Output Shape
---|---|---
Input | - | (5, 20, 100, 100)
Encoder 1 | Conv2d/BatchNorm2d/Tanh | (5, 32, 100, 100)
Pool 1 | MaxPool2d | (5, 32, 50, 50)
Encoder 2 | Conv2d/BatchNorm2d/Tanh | (5, 64, 50, 50)
Pool 2 | MaxPool2d | (5, 64, 25, 25)
Encoder 3 | Conv2d/BatchNorm2d/Tanh | (5, 128, 25, 25)
Pool 3 | MaxPool2d | (5, 128, 12, 12)
Encoder 4 | Conv2d/BatchNorm2d/Tanh | (5, 256, 12, 12)
Pool 4 | MaxPool2d | (5, 256, 6, 6)
Bottleneck | Conv2d/BatchNorm2d/Tanh | (5, 512, 6, 6)
Decoder 4 | ConvTranspose2D/Encoder 4 | (5, 256, 12, 12)
Decoder 3 | ConvTranspose2D/Encoder 3 | (5, 128, 25, 25)
Decoder 2 | ConvTranspose2D/Encoder 2 | (5, 64, 50, 50)
Decoder 1 | ConvTranspose2D/Encoder 1 | (5, 32, 100, 100)
Rescale | Conv2D | (5, 20, 100, 100)