Applications of Fourier Neural Operators in the Ifmif-Dones Accelerator

Authors: guillermo rodríguez llorente
galo  gallardo
roberto  gómez-espinosa martín
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-09 10:29:35

Abstract: In this work, Fourier Neural Operators are employed to improve control and optimization of an experimental module of the IFMIF-DONES linear accelerator, otherwise hindered by its simulations high complexity. The models are trained to predict beam envelopes along the lattice’s longitudinal axis, considering variations in quadrupole strengths and particle injections. They serve three purposes: enabling fast inference of beam envelopes, creating an environment for training a Deep Reinforcement Learning agent responsible for shaping the beam, and developing an optimizer for identifying optimal accelerator parameters. The resulting models offer significantly faster predictions (up to 3 orders of magnitude) com- pared to traditional simulators, with maximum percentage errors below 2 %. This accelerated simulation capability makes it feasible to train control agents, since the time per step taken is reduced from 3s to 4 × 10e−3s. Additionally, Stochastic Gradient Descent was applied to optimize one of the models itself, determining the best parameters for a given target and thus solving the inverse problem within seconds. These results demonstrate the synergy and potential of these Deep Learning models, offering promising pathways to advance control and optimization strategies in the IFMIF-DONES accelerator and in other complex scientific facilities.

Main paper: 
 introduction: 
 One of the foremost challenges in future nuclear fusion power plants is studying the effects of neutron irradiation on reactor materials. In addressing this issue, the IFMIF-DONES facility aims to build a linear accelerator in order to generate a neutron spectrum similar to that of a fusion reaction by the collision of deuterons with a lithium circuit (Bernardi et al., 2022) . This study, conducted within the DONES-FLUX project, explores the joint application of Deep Learning Surrogate Models (DLSMs) and Deep Reinforcement Learning (DRL) in order to improve the optimization and control of the facility.DLSMs are models of physical, chemical, or biological processes based on Artificial Intelligence (AI), that enable rapid and accurate high-fidelity simulations. This technique provides a method to accelerate complex simulations, governed by partial differential equations, with data (Hao et al., 2023) . On the other hand, DRL is an AI paradigm that is applied to solve sequential decisionmaking problems (Graesser & Keng, 2019) . In this approach, an agent interacts with an environment (usually simulated) that provides feedback in the form of rewards or penalties based on its actions. DRL-based controllers can improve the development and efficiency of traditional systems, allowing agents to achieve goals set by experts instead of constantly making manual adjustments (Degrave et al., 2022) . Training DRL agents using traditional simulators becomes impractical due to the substantial time required to process each step. In this context, the use of DLSMs helps expedite execution times, enabling the online learning of such agents.In this work, we train two Fourier Neural Operator (FNO) (Li et al., 2020) models with data from the OPAL simulator (Adelmann et al., 2019) . The selection of this architecture is motivated by the fact that it is purely data driven, since there is a lack of comprehensive knowledge regarding the differential equations that govern the transfer of beam bunches across the accelerator. Moreover, the FNO is discretization invariant which enables multi-resolution approaches, an advantage compared with a neural network Azizzadenesheli et al. (2024) . The simulator makes use of the Monte Carlo method for the interactions of particles along the lattice, which is stochastic. This is a computationally expensive process which incentives the development of model capable of predicting the macroscopic distributions and envelopes of the particles. The distribution of particles along the accelerator is a two-dimensional Gaussian Wiedemann (2015) , for this reason for our purposes of training DRL agents, predicting only the beam envelope (standard deviation along the two transversal axes) can be seen as a form of feature engineering, in which we are reducing an observation space of 20000 particles to a two-dimensional space, otherwise untraceable.The trained models are able to predict the deuteron envelopes along the longitude of an experimental module of the linear accelerator, displayed in the lower part of Fig. 1 . These models find three different applications: fast prediction of beam envelopes, training of a DRL agent that controls the geometry of the beam at the end of the lattice, and building an optimizer to find the best parameters for a given configuration. The problem can be set as design problem if the geometry of the beam is static or as a real time control problem if the geometry is dynamic (for instance, the currents of the quadrupoles can be controlled). For the second option, the DRL agent can be used to optimize the design of the lattice and to correct the beam geometry in real time, as the footprint at the target needs to have a specific shape for the reaction. For details on FNOs and DRL, refer to Section 2. Section 3 presents the methodology and obtained results, while the discussion is encapsulated in Section 4. 
models: 
 Neural operators (NOs) represent an emerging branch of DLSMs, extending the architecture of neural networks for the approximation of operators capable of mapping between infinite-dimensional function spaces (Kovachki et al., 2023) . This generalization becomes particularly relevant for solving systems of partial differential equations, which can ultimately represent physical phenomena. The main architecture of these models takes data from input functions a(x) to predict output functions u(x), as can be seen in the upper part of Fig. 1 . The layer P maps the input function to its first hidden representation, which is then passed through layers of kernels with non-linear activation functions, and finally taken to the output representation with the layer Q. The FNO model employs the Fourier transform for computing the kernels (Li et al., 2020) . On the other hand, DRL techniques consist of two fundamental elements: the environment and the agent. The environment represents the model of the problem to be solved, with states (S t ) containing information about the system at time t; whereas the agent is an algorithm that uses information obtained from the environment to make decisions (Graesser & Keng, 2019) . Following the execution of an action (A t ), the environment transitions to the subsequent state (S t+1 ), and a reward function assesses whether the agent is progressing toward or deviating from the objective. These challenges are framed as Markov Decision Processes, where learning occurs through the Bellman equations and the agent's experiences through multiple iterations. The primary objective for the agent is to learn a policy function (π(S t )) determining optimal actions based on observed states. This function, essentially the agent's brain, is approximated through neural networks. 
fno predictions: 
 The data required to train the FNO models was generated with simulations of deuteron bunches1 travelling through the experimental lattice with varying conditions. The input functions are taken to be the constant parameters controlling the shape of the beam. First, the injected Gaussian distribution of particles (s x (z = 0), s y (z = 0)) lies within the interval [0.25, 0.75] mm. Second, the magnetic strengths of the quadrupoles (k 1 , k 2 , k 3 ) range between [-2.5, 2.5] T m -1 . The output functions that the model predicts are the beam envelopes (in this case, the standard deviations), s x (z) and s y (z), along the lattice longitudinal axis z. The end of the line is located at a distance z = L. The FNO model with variable injection is used to train a DRL agent in order to sequentially reach different beam shapes at the end of the line. This is achieved through the manipulation of quadrupoles acting as actuators. In each step, which corresponds to a deuteron bunch crossing, the agent is able to incrementally modify the 3 quadrupoles (∆k 1 , ∆k 2 , ∆k 3 ) with a set limit 3 . The default state of the quadrupoles is (k 1 = 0, k 2 = 0, k 3 = 0). The observation space for this agent is defined with the standard deviations (s x (L), s y (L)) at the target , the absolute state of the quadrupoles and the target configuration to reach. The reason to limit the observation space of standard deviations to z = L is that in real applications only the distribution at the target would be available using a camera or other sensor. Finally, the reward function was defined as the normalized negative euclidean distance between the target and the current configuration, with a 10 % distance threshold for the winning condition.The wrapper for the DRL environment was built with the Gymnasium tool (Towers et al., 2023) and then a Proximal Policy Optimization (PPO) (Schulman et al., 2017) agent was trained with the Stable-Baselines 3 framework (Raffin et al., 2021 ). Results for one example can be seen in Fig. 3 , where the agent reaches the desired configuration (s x (L) = 3.7, s y (L) = 2.0) within 9 actions. The 2-dimensional deuteron distributions are shown for the initial and final states at the target. Diagrams with the training loop, the environment and the policy network can be seen in Appendix E.Tables with the hyperparameters and the observations and actions spaces are shown in Appendix C. Additionally, Table 6 displays the agent evaluation and performance on other targets.Another application of the FNO model comes from the fact that it is a differentiable PyTorch model, thus the gradients can be utilized in inverse problems to find the optimal input parameters for a given output (Zhou et al., 2024) . In this case, we built a loss function with the target parameters of the beam at the end of the lattice and used Stochastic Gradient Descent (SGD) to obtain the optimal values of the quadrupoles and particles injections.We applied the SDG optimization to both constant and variable injection models, with euclidean distance to the target parameters as loss function, and a early stopping tolerance of 0.005, constraining the inputs to the domain in which the FNO has been trained. This method lead with both FNO models to the optimal design in less than 32 seconds, a fast process compared with other techniques, such as evolutionary algorithms or DRL methods. Optimization loop diagram can be found in Appendix D and tables with results for different runs and the parameters used by the optimizer are shown in Appendix F. 
conclusions and outlook: 
 The trained FNO models exhibit a remarkable capability to predict beam envelopes at a speed 3 orders of magnitude faster than traditional simulations with OPAL, achieving a maximum percentage error below 2 %. This accelerated prediction not only facilitates rapid inference of beam properties but also enables the training of DRL agents within feasible time frames. The DRL agent achieves the optimal policy in less than 20 minutes, a substantial improvement compared to the estimated 225 hours required with OPAL for 270k steps. Additionally, the differentiable nature of the FNO model is leveraged to develop an optimizer for identifying optimal accelerator parameters, leading to significantly reduced optimization times-merely seconds compared to alternative methods such as DRL optimization or evolutionary algorithms.These results underscore the synergistic potential of combining DLSMs with other Deep Learning techniques, presenting promising avenues for advancing control and optimization strategies not only in the IFMIF-DONES accelerator but also in other complex scientific facilities. In the context of nuclear fusion power plant development, this study highlights the promising role of AI techniques in advancing fusion research, as can also be seen in the work of Degrave et al. (2022) . Finally, the next step of this research will be to apply these techniques to each of the modules of the IFMIF-DONES facility.A APPENDIX: FNO MODEL SUMMARY  
 
 Appendix: Each bunch consists ofparticles. 2 https://github.com/NVIDIA/modulus-sym 3 It is a hyperparameter set to 0.25 T m -1 . 


Decision: Dear Authors,

Thank you for submitting the draft.

Both reviewers agree that the work contains an interesting contribution. It is expected that authors will be addressing comments by the reviewers in the final draft. 

regards

AC

Reviews: review: title: title: Review of FNO paper
main_review: main_review: This is a review of the manuscript, "Applications of Fourier neural operators in the IFMIF-DONES accelerator," submitted to the ICLR 2024 Workshop on AI4DifferentialEquations In Science. The paper describes a Fourier neural operator (FNO) trained to predict the envelope parameters of neutron beams in the IFMIF-DONES facility. Further, a deep reinforcement learning (DRL) model is trained for optimization and/or control purposes. My review is summarized below in point form.

1. The FNO provides a nice speed-up, but it is also (presumably) trained on $\{s_x, s_y\}$ data as a function of $z$. Are $\{s_x, s_y\}$ data actually available as a (near-)continuous function of $z$? Or are individual particle trajectories somehow measured? What is the expected fidelity of such data, and how does the FNO's performance deteriorate in the context of realistic uncertainties? For instance, does parameterizing the beam envelope in terms of the standard deviation of the particle distribution lead to potential model errors? The authors should provide a more comprehensive phenomenological description of the physical system. They should answer such questions, and better motivate their chosen surrogate model form. Doing so will help situate readers from other fields.

2. Did the authors consider alternative operator network architectures? In many cases, a simple multilayer perceptron (MLP) will suffice. Especially in the present context, which features low-dimensional input and output spaces and a limited number of control parameters. A simpler network might perform equally well! The authors should test an MLP of similar depth, having the same number of trainable parameters as their FNOs, to assess whether the FNO configuration is beneficial, per se.

3. The authors use stochastic gradient descent to optimize $\{s_x(0), s_y(0), k_1, k_2, k_3\}$ to achieve a target output envelope, $\{s_x(L), s_y(L)\}$. As presented, this seems like a very simple optimization problem that could be solved using vanilla gradient descent. Why is it necessary to use the stochastic version?

4. Related to the first and third points, the need for a DRL controller is unclear. Is the controller required for real-time applications? And if $\{s_x, s_y\}$ data are available as a function of $z$ to train the model (i.e., in the absence of ``comprehensive knowledge regarding the [governing equations]''), why does the controller only have access to the data at $z = L$? Wouldn't it be feasible to reduce the number of requisite actions by feeding the controller a richer observation?

5. The authors should comment on the performance of their FNO outside the training envelope. (It can be difficult to operationalize the notion of extrapolation in many deep learning settings, of course, but the present demonstration is relatively low-dimensional.)
title: title: Review of “Applications of Fourier Neural Operators in the IFMIF-DONES Accelerator”.
main_review: main_review: This uses data driven methods to help with shaping the beam envelopes in the IFMIF-DONES linear accelerator. The authors do this in two steps. They first train a neural operator to predict the beam envelope given the tuning parameters. They then show that this low-cost map between tuning choice and outcome can be used to cost effectively train a control agent to choose tuning parameters given the desired envelope. The authors also show that the neural operator map can also be directly optimized over to produce the desired tuning parameters. Outside of the contribution to the IMFIF-DONES workflow, the paper introduces a novel workflow for reinforcement learning by first training a low-cost model to replace a numerical integrator and then using it to perform the large number of iterations necessary for reinforcement learning. 

My knowledge on the problem domain of nuclear fusion is minimal, but the paper is well written and demonstrates clear success with the implementation proposed. Additionally, the pipeline used to train the model may be useful in other applications where low-cost replacements for numerical integrators may allow reinforcement learning techniques to be more easily applied. My only hesitation is my lack of familiarity with the challenge and scope of the physical system being studied or existing best practices in the field.

Minor Comments:
1)	In section 2 you reference figure 6 where I think you intend to reference figure 1.
2)	In the second paragraph of section 3.3 “both constant an variable” should be “both constant and variable”
3)	As someone outside of the field, understanding the structure of the data being explored is a bit of a challenge. It may be helpful to provide additional visualizations of the range of envelopes in the appendix. 
4)	There is no baseline method provided, it would be helpful to discuss how the tuning parameters are selected in the current state-of-the-art.