Differentiable Deep Learning Surrogate Models Applied to the Optimization of the IFMIF-DONES Facility

Abstract

One of the primary challenges for future nuclear fusion power plants is understanding how neutron irradiation affects reactor materials. To tackle this issue, the IFMIF-DONES project aims to build a facility capable of generating a neutron source in order to irradiate different material samples. This will be achieved by colliding a deuteron beam with a lithium jet. In this work, within the DONES-FLUX project, deep learning surrogate models are applied to the design and optimization of the IFMIF-DONES linear accelerator. Specifically, neural operators are employed to predict deuteron beam envelopes along the longitudinal axis of the accelerator and neutron irradiation effects at the end, after the beam collision. This approach has resulted in models that are able of approximating complex simulations with high accuracy (less than 17% percentage error for the worst case) and significantly reduced inference time (ranging from 2 to 6 orders of magnitude) while being differentiable. The substantial speed-up factors enable the application of online reinforcement learning algorithms, and the differentiable nature of the models allows for seamless integration with differentiable programming techniques, facilitating the solving of inverse problems to find the optimal parameters for a given objective. Overall, these results demonstrate the synergy between deep learning models and differentiable programming, offering a promising collaboration among physicists and computer scientists to further improve the design and optimization of IFMIF-DONES and other accelerator facilities. This research will lay the foundations for future projects, where optimization efforts with differentiable programming will be performed.

1. Introduction

Nuclear fusion is a promising source of sustainable energy for humanity [1]. However, the design and validation of materials that can withstand the extreme conditions within a fusion reactor remains an open issue [2]. In order to replicate this environment, the IFMIF-DONES (International Fusion Materials Irradiation Facility-DEMO Oriented Neutron Source) project is developing a linear accelerator that will produce a neutron spectrum similar to that of nuclear fusion reactions by colliding deuterons with a flowing lithium target [3]. The resulting neutrons then irradiate material samples located in a high flux test module (HFTM), allowing for the study of its effects [4].

As part of the DONES-FLUX project, this research addresses the challenge of optimizing the IFMIF-DONES facility by employing differentiable deep learning surrogate models (DLSMs). Accurate simulations of this complex environment involve vast numbers of particles and an extensive parameter space, making them computationally intensive and time-consuming, sometimes taking days to complete [5]. This computational demand renders it unfeasible to apply optimization algorithms, such as gradient descent (GD) [6] or deep reinforcement learning (DRL) [7], which may require many iterations to converge. DLSMs provide an effective solution by approximating these simulations with a high accuracy while drastically reducing the inference time, thereby enabling the application of such optimization algorithms.

Neural operators (NOs) are a class of DLSM designed to approximate mappings between function spaces, making them highly effective for solving systems of partial differential equations (PDEs) that model complex physical phenomena [8]. Unlike traditional neural networks, which typically approximate finite-dimensional mappings, NOs work in infinite-dimensional spaces, which allows them to capture spatially continuous relationships between input functions $a\left(x\right)$ and output functions $u\left(x\right)$ [9]. For this study, the Fourier neural operator (FNO) [10] and deep operator network (DeepONet) [11] architectures were selected, which have the advantages of being purely data-driven and with a proven performance across a wide range of applications [12,13,14]. Additionally, FNO has the property of discretization invariance and can predict at different resolutions with the same trained model.

In this work, NO architectures are employed as DLSMs in two critical parts of the IFMIF-DONES neutron source. First, FNO models are trained to predict beam statistical functions and beam profile distributions in the initial section of the high energy beam transport line (HEBT-S1) line of the accelerator [15]. Second, a DeepONet architecture is used to model material irradiation variables within the HFTM [16] after the collision. The efficiency and differentiability of these models enable the optimization of the quadrupoles to achieve specific operational objectives in the facility. Section 2 details the materials and methodologies used, Section 3 presents the results for both applications, and Section 4 discusses the implications of these findings.

2. Materials and Methods

2.1. General Methodology

The general workflow involves the following steps: first, identify the specific part of the facility targeted for optimization. Next, perform simulations using traditional solvers to generate a representative dataset. Then, train the differentiable DLSMs using this dataset. Finally, integrate the trained models into other applications that require speed and differentiability. This process is visually summarized in Figure 1.

Figure 1. Optmization processworkflow. (A) Select the facility part and frame the problem. (Image credit: IFMIF-DONES) (B) Generate dataset with traditional simulations. (C) Train differentiable DLSMs with the data. (D) Employ the models for different applications.

2.2. Optimizing the HEBT-S1 Lattice

The first problem is to determine the optimal values for the magnets in the HEBT-S1 section, which is illustrated in the left side of Figure 2. The lattice comprises six quadrupoles (from now on represented by the integer j) and their corresponding drift spaces. The lengths, apertures, and field strengths of each component are defined as in [17] and are detailed in Table A4 of Appendix A. A set of coordinates $(x,y,z)$ is employed, where $(x,y)$ represent the transversal axes and z describes the accelerator longitude. This lattice configuration was implemented in the OPAL simulator [18].

Figure 2. (Left):The IFMIF-DONES HEBT, with the S1 components, quadrupoles, and drift zones. Taken from [15]. (Right): Two-dimensional slices of the geometry of the HFTM. Taken from [16].

A total number of 8192 deuteron beam simulations were conducted, each with a different set of magnetic field values $k_j$, spanning the 4–8 ${\mathrm{Tm}}^{-1}$ range. The simulation outputs included the one-dimensional statistical beam functions $f_x\left(z\right)$ and $f_y\left(z\right)$(Root mean square (RMS) and maximum size of the envelopes (MAX)) and the normalized number of particles $n\left(z\right)$. The two-dimensional beam profile distributions, ${\rho }_x(x,z)$ and ${\rho }_y(y,z)$, were also recorded. Each simulation corresponded to a single deuteron bunch crossing the lattice, with 100,000 injected particles generated from Gaussian distributions (see Table A1 of Appendix A for the injection parameters).

From the entire dataset of simulations, 20% was reserved for testing, while the remaining 80% was used for training the models. Two different models were developed: a 1D-FNO to predict the statistical beam functions and particle counts, and a 2D-FNO to predict the beam profile distributions. A summary of the models is presented in Equation (1) and the full architectures with their hyperparameters are detailed in Table A6 of Appendix C. NVIDIA Modulus, version 1.5.0 [19] and PyTorch 2.3.0 [20] were selected as frameworks because of the FNO implementation and differentiable computational graph capability.

$$\begin{array}{cc}\hfill k_j\left(z\right)& ⟶{\mathrm{FNO}}_{1D}⟶f_x\left(z\right),f_y\left(z\right)\hfill \\ \hfill k_j(x,z),k_j(y,z)& ⟶{\mathrm{FNO}}_{2D}⟶{\rho }_x(x,z),{\rho }_y(y,z)\hfill \end{array}$$

(1)

Finally, the 1D model was combined with a GD optimization algorithm, Adam [21], to explore the parameter space and solve the inverse problem of determining the optimal quadrupole settings needed to achieve the desired beam configurations at the end of the lattice (related to the beam footprint, the RMS, MAX, and n functions evaluated at $z=10$ m). This method is effective because the trained models are differentiable with respect to the input parameters, $k_j$. Details of the optimizer and other algorithm hyperparameters are provided in Table A7 of Appendix C, whereas the loss function is described in Equation (A1) of Appendix E. The different terms in the loss function minimize the discrepancy between the objective and current beam configurations and impose penalties on out-of-bound solutions. A summary of the software and hardware resources used in this work is presented in Table A5 of Appendix B.

2.3. Surrogates for Materials Irradiation Variables in HFTM

As previously mentioned, the accelerated deuteron beam interacts with liquid lithium, initiating the Li(d, xn) stripping reaction [22], which generates high-energy neutrons to replicate the radiation conditions of a nuclear fusion reactor. In order to evaluate the mechanical properties of the materials intended for these reactors, the samples were placed inside the HFTM, shown on the right side of Figure 2, where they were exposed to varying degrees of neutron flux to cover different irradiation quantities. Preliminary studies on the variables influenced by radiation must be conducted through simulations. In IFMIF-DONES, Monte Carlo (MC) codes deliver the most accurate results, albeit at a significant computational cost. This is due to the statistical nature of the method, which requires a large number of source particles to achieve statistically reliable solutions, resulting in prolonged computational times [23].

To simulate that process and gather the necessary data for training the model, a coupling between Geant4 [24] and OpenMC [25] was established, as illustrated in Figure 3. First, Geant4 models the Li(d, xn) stripping reaction, where deuterium and lithium interact to generate high-energy neutrons (>14 MeV). This is achieved by simulating a deuteron beam interacting with a 25 mm lithium curtain, positioned in front of a detector. The detector captures key data, including the position, angle, and energy of particles from the reaction, particularly the neutrons. Next, the collected data serve as the input for the second simulator, OpenMC, which simulates neutron transport in HFTM and calculates magnitudes such as flux, damage energy, heating, and hydrogen and helium production. Figure 4 illustrates the Geant4 and OpenMC simulations. On the left, a simple representation of the lithium curtain and the deuteron beam is shown, while the right side of the figure depicts a vertical cross-section of the HFTM model. The simulation process is repeated with varying parameters (see ranges in Table 1). This simulation setup is a simplification of the IFMIF-DONES setup, as the main purpose of this study is to validate the data-driven model. For a high resolution physical model, refer to [26]. Future works will integrate more precise physical representations of the system.

Figure 3. Neutronic DLSM process. Using the spatial and energy distribution of the deuteron beam, a Geant4 simulation is conducted to determine the neutron flux. This flux data are then input into the OpenMC to model the variables influencing the materials within the HFTM.

Figure 4. (Left): Lithium curtain and deuteron beam in Geant4. The red tracks correspond to neutrons, while the green tracks to photons. (Right): Vertical section of the HFTM modelled in OpenMC.

Table 1. Input parameter ranges for the deuteron beam. See Table A2 for details on the simulation arguments.

Parameter	Ranges	Step
Energy ($ϵ$)	20–40 MeV	5 MeV
Beam width (L)	10–20 cm	2 cm
Parameter	Values
Number of deuterons (D)	$1\times {10}^7$, $5\times {10}^7$, $1\times {10}^8$, $5\times {10}^8$, $1\times {10}^9$

The OpenMC simulator has been validated for the IFMIF-DONES facility in other works [27] where the McDeLicious extension [22] for the Li(d, xn) stripping reaction for MCNP has been adapted [28]. However, this extension is not Open Source, and for reproducibility purposes, the alternative use of Geant4 has been choosen, acknowledging the decriment in the generation of reliable physics data. See Table A3 for a summary of the OpenMC tallies used in this work.

The NO model utilized for this task is based on the DeepONet architecture [11], implemented using the DeepXDE library [29]. This architecture comprises two distinct sub-networks: the branch network and the trunk network. The input data are divided, such that neutron source parameters $(ϵ,L,D)$ are processed by the branch, while spatial coordinates $(x,y)$ are managed by the trunk. This operation and the result of merging both nets to produce an output as an operator G can be seen in Equation (2).

$$\begin{array}{cc}\hfill (ϵ,L,D)& ⟶f_{\mathrm{branch}}(ϵ,L,D)\hfill \\ \hfill (x,y)& ⟶f_{\mathrm{trunk}}(x,y)\hfill \\ \hfill f_{\mathrm{branch}}(ϵ,L,D)\otimes f_{\mathrm{trunk}}(x,y)& ⟶G(ϵ,L,D)(x,y)\hfill \end{array}$$

(2)

Five separate models are trained, each corresponding to a specific variable (flux, damage energy, heating, hydrogen production, and helium production). To evaluate improvements in inference time, each model is trained with 100 training samples, 25 validation samples, and 25 test samples (only 150 simulations were utilized due to the significant computational resources needed to calculate the neutronic tallies. In contrast, the deuteron models required a substantially larger number of simulations to accurately capture the beam dynamics behavior).

3. Results

3.1. Quadrupoles Optimization in the HEBT-S1 Lattice

The 1D-FNO model is able to predict all of the beam statistical functions and the number of particles ${10}^3$ times faster than the OPAL code, with a trade-off mean average percentage error (MAPE) of 3.60% across all variables (see Table A10 of Appendix F for the individual errors). Figure 5 displays a test example of the results obtained by the model, compared with the simulations. Some instances for the predicted number of particles are also shown in Figure A3 of Appendix G.

Figure 5. Prediction results for the 1D beam statistics functions along the accelerator’s longitudinal axis, shown for the x and y axes. Blue lines represent OPAL simulations, while red lines indicate FNO predictions.

Figure 6 presents the predictions of the beam profile distributions model for a test example, alongside the corresponding OPAL simulations and residual errors. The model achieves a mean absolute error (MAE) of 5.15 $\times {10}^{-3}$ on the test set. The inference speed boost factor matches that of the previous model, as both rely on the same set of simulations. A comprehensive summary of error metrics for each function across both models is provided in Table A10 of Appendix F, with additional test examples of the model predictions also available in Figure A5 of Appendix G.

Figure 6. Test prediction results for the 2D beam profile distributions along the accelerator’s longitudinal axis, displayed for the x and y axes. The (left) panel shows OPAL simulations, the (middle) panel presents FNO predictions, and the (right) panel illustrates the residual errors.

Table A9 of Appendix D summarizes the optimization outcomes for different beam objectives, presenting the target and achieved configurations for the statistical beam functions evaluated at $z=10$ m. The table also details the optimized quadrupole values, the computational time required to converge to a solution, and the particle count as reported by OPAL during the evaluation. A consistent threshold of $n\left(10\right)=98\%$ was applied as a loss optimization criterion, ensuring the particle losses remained below 2%. The optimization process reliably achieved the desired objectives within approximately 10 min.

3.2. Neutronic Irradiation Models for HFTM

The prediction results are shown in Figure 7 and its corresponding errors in Figure 8. Using 100 training samples, DeepONet successfully learns the geometry of the HFTM and delivers accurate predictions. The 150 instances of the dataset took a total of $7.78\times {10}^5$ s to generate, resulting in an average time per simulation of $5.187\times {10}^3$ s. The inference time of DeepONet per simulation was $5.5\times {10}^{-4}$ s, resulting in a speedup factor of $9.43\times {10}^6$. The performance metrics achieved by DeepONet are summarized in Table A11, in Appendix F. Finally, the numerical simulation results used to train these models were not validated and should not be referenced in future works.

Figure 7. Neutronic variables predictions. (Left): Geant4 and OpenMC coupled output. (Right): The DeepONet prediction. For each predicted variable, a different model produces the output.

Figure 8. Errors in the neutronic predictions. (Left): Residuals for each magnitude. (Right): Symmetric Absolute Percentage Error (SAPE) for every magnitude. This metric has the visible limitation where it sets to 100% when a 0 true value is encountered.

4. Discussion

In this study, NO operator models were trained as differentiable DLSMs to significantly enhance the inference speed of traditional simulations. These properties enable their use in optimization loops and applications that would otherwise require an impractical number of time steps. The developed models significantly reduced the inference times, with factors ranging from ${10}^3$ to ${10}^5$, while maintaining minimal error metrics. However, it is important to note that these models create predictions at unseen points within the trained input space and, thus, do not extrapolate. This would require architectures like Physics-Informed Neural Operators (PINOs) [30], which integrate physical constraints for improved generalization. Additionally, the deuteron beam optimization procedure manages to reach all of the established objectives mantaining particle counts and ensuring the desired footprint geometry, which is a key factor for the collision.

These results not only demonstrate the potential of DLSMs for optimization tasks, but their applications extend beyond solving inverse problems. As shown in other studies, including [31], the significant reduction in simulation time enables the online training of deep reinforcement learning agents for control and tuning. These algorithms often require millions of time steps to converge and the efficiency of surrogate models makes this feasible.

While the results are promising, there are certain limitations that are being currently addressed. The optimization was restricted to the HEBT-S1 section due to the exponential growth in computational resources required as the number of magnetic elements increased. Further works will include new sampling strategies more suited to high-dimensional spaces of parameters and other optimization strategies. Furthermore, incorporating additional channels into the models, such as particle momentum data, could enhance optimization by providing phase-space information, a critical factor in particle accelerator operations [32]. This refinement has the potential to improve the model performance and to broaden their applicability. Future works will address this with the inclusion of quality beam outputs, like emmitance, in model predictions, and including these features in the objective optimization function to find not only the parameters that provide an objective shape, but also the ones that have the best beam quality.

The limitations of the current DeepOnet model for neutrionics predictions are that if the geometry changes, a new re-training has to be created. To address this issue, future works will include the use of novel geometry-aware models, like Geom-DeepOnet [33], with a generalization capability for different geometries. Addionally, the HFTM model will be updated with the final IFMIF-DONES version for a more precise representation.

In conclusion, the trained NO models represent a significant step forward in beam dynamics simulation and neutron irradiation. By leveraging their speed, accuracy, and differentiability, they offer compelling solutions for the challenges faced in the IFMIF-DONES neutron source and other scientific facilities.