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.

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. Magnetic element dimensions, length, and edge where it starts relative to the beginning of the lattice. The aperture radius of all the elements is r = 50 mm. This setup is displayed in Figure 2.

Element	Length (m)	Edge (m)	Element Type
DR1	1.00	0.00	Drift
Q1	0.25	1.00	Quadrupole
DR2	0.60	1.25	Drift
Q2	0.25	1.85	Quadrupole
DR3	0.60	2.10	Drift
Q3	0.25	2.70	Quadrupole
DR4	1.65	2.95	Drift
Q4F	0.25	4.60	Quadrupole
DR5	1.50	4.85	Drift
Q5	0.25	6.35	Quadrupole
DR6	1.50	6.60	Drift
Q6	0.25	8.10	Quadrupole
DR7	1.65	8.35	Drift

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].

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. Summary of the key parameters for the OPAL-T configuration used in the simulation. Data employed in this work should be fully reproducible using the set-up displayed in this Section.

Variable	Value	Description
Particle	Deuteron	Type of particle used for simulation.
Number of particles	100,000	Number of particles in the simulation.
Energy	40 MeV	Energy of the particles in the simulation.
Injection type	Gauss	Type of distribution used for the injection
${\sigma }_{x,y,z}$	28 mm	Standard deviations for the spatial distribution in x, y, and  z.
${\sigma }_{p_x,p_y}$	0.023 $\times {10}^{-3}$	Standard deviations for the momentum distribution in x and y.
${\sigma }_{p_z}$	0.039 $\times {10}^{-3}$	Standard deviation for the momentum distribution in z.
MAXSTEPS	10,000	Maximum number of simulation steps.
DT	1.65 $\times {10}^{-10}$	Time step for the simulation.
ZSTOP	10	Stop condition for the z-coordinate during the simulation.
METHOD	PARALLEL-T	The OPAL mode.
FIELDSOLVER	FFT	Field solver method used in the simulation.
$M_{X,Y,T}$	16	Number of grid points in the x, y, and t directions.
PARFFTX,Y,T	true	Whether the FFT solver is parallelized for x, y, and t.
BCFFTX,Y,T	open	Boundary conditions used for the FFT solver in x, y, and t.
BBOXINCR	1	Increment for the bounding box size.
GREENSF	STANDARD	Green’s function type used in the simulation.

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. Hyperparameters for the FNO architectures in NVIDIA Modulus are shared between models when a single value is specified. For parameters with two values, the first corresponds to the 1D model, and the second to the 2D model.

Hyperparameter	Value
scheduler	tf_exponential_lr
optimizer	adam
loss	sum
decoder.nr_layers	1
decoder.layer_size	256–512
fno.dimension	1–2
fno.nr_fno_layers	4
fno.fno_modes	12
scheduler.decay_rate	0.95
scheduler.decay_steps	1000
training.max_steps	10,000
batch_size.grid	32–8
batch_size.validation	32–8

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. Gradient descent optimization hyperparameters.

Hyperparameter	Value
scheduler	exponential_lr
optimizer	adam
loss	custom distance
scheduler.gamma	0.999
max_steps	100,000
early_stop_loss	0.005
learning_rate	1

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. System and software specifications.

Category	Specification
System Specifications
GPU Model	NVIDIA GeForce RTX 3060
CUDA Version	12.2
CPU Model	Intel(R) Core(TM) i9-10900K @ 3.70 GHz
CPUs	12
Software Specifications
OS	Ubuntu 22.04.4 LTS (Docker)
nvidia-modulus.sym	1.5.0
OPAL	2022.1.0
openmpi	4.1.4
Torch	2.3.0
Geant4	11.2.1
OpenMC	0.14.0
DeepXDE	1.12.0

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. 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$

). 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.

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. Summary of the tally setup used in the OpenMC simulation.

Variable	Description
Tally Scores
Flux	Tracks the total particle flux.
Heating (MT = 301)	Nuclear heating associated with the reaction rates.
Damage-energy (MT = 444)	Calculates the damage energy produced by the reactions.
H Production (n, p)	Tally for hydrogen (H) production via (n, p) reactions.
He Production (n, a)	Tally for helium (He) production through (n, $\alpha$) reactions.
Filters
MeshFilter	Covers the entire volume of the HFTM geometry.

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. FNO model errors evaluated in all the test set. The first part shows the 1D model errors, and the second part shows the 2D model errors. No MAPE is displayed for the 2D model due to the very small values of the pixels.

FNO 1D Model Errors
Variable	MAPE (%)	MAE
RMSx$\left(z\right)$	4.38	0.11
RMSy$\left(z\right)$	5.63	0.16
MAXx$\left(z\right)$	3.92	0.30
MAXy$\left(z\right)$	3.76	0.39
$n\left(z\right)$	0.92	0.01
FNO 2D Model Errors
Variable	MAPE (%)	MAE
${\rho }_x(x,z)$	-	4.7 $\times {10}^{-4}$
${\rho }_y(y,z)$	-	5.6 $\times {10}^{-4}$

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 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.

Table A9. Optimization examples of beam configurations chosen to showcase different geometric shapes. The algorithm is capable of finding solutions as long as they are within the feasible solution space. All of the functions are evaluated at the footprint location, $z=10$ m.

Target (mm)				Achieved (mm)				${\mathit{k}}_{\mathit{j}}$ (Tm−1)						t (min)	n
RMSx	RMSy	MAXx	MAXy	RMSx	RMSy	MAXx	MAXy	k1	k2	k3	k4	k5	k6
5.4	2.4	18.0	8.0	5.4	2.4	17.9	8.0	5.9	−7.1	6.9	−5.2	5.0	−6.3	10.42	1.00
4.7	2.6	16.0	7.0	4.7	2.6	15.6	6.9	5.1	−5.7	7.6	−6.4	5.1	−6.0	3.42	0.99
2.5	3.4	9.0	12.0	2.6	3.1	8.9	12.0	5.0	−8.0	7.4	−7.7	5.0	−6.9	10.44	1.00
3.0	3.0	11.0	10.0	3.0	3.0	10.5	9.8	5.3	−6.7	7.9	−8.0	5.0	−6.0	1.52	0.99
6.4	8.5	22.0	30.0	6.3	8.5	21.5	27.0	5.4	−6.2	7.9	−6.3	5.4	−7.2	3.34	0.99
11.8	10.4	38.9	33.0	11.7	10.3	39.0	32.5	5.0	−6.3	7.9	−6.6	5.9	−7.5	10.45	0.99
16.5	1.0	50.0	3.4	16.5	1.0	49.8	3.3	6.5	−7.5	7.1	−5.6	6.5	−5.8	2.30	0.99
9.8	10.2	33.4	37.3	9.8	10.2	32.5	31.3	5.0	−6.0	7.6	−6.3	5.6	−7.4	3.7	0.99
7.0	6.6	23.5	19.5	7.0	6.0	23.2	19.5	6.4	−7.8	6.2	−5.0	5.0	−7.8	10.43	0.99
13.3	4.3	48.6	14.2	13.3	4.3	45.7	13.6	5.5	−6.5	7.7	−6.5	6.2	−6.4	1.48	0.99
7.7	11.6	25.9	29.9	7.7	11.6	25.9	29.4	6.0	−5.0	5.1	−6.0	5.0	−7.3	4.31	0.98
4.2	13.1	14.5	37.5	4.1	13.1	14.5	36.7	5.0	−5.5	5.0	−6.9	5.0	−7.6	10.47	0.99

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. Mean sMAPE, NMAE, and R² of the DeepONet model.

Variable	sMAPE (%)	NMAE	R2
Flux	2.3	0.031	0.996
Damage E.	9.2	0.120	0.968
Heating	8.9	0.111	0.972
H prod.	16.8	0.139	0.967
He prod.	8.3	0.095	0.977

, in Appendix F. Finally, the numerical simulation results used to train these models were not validated and should not be referenced in future works.

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.