Fast, Efficient and Flexible Particle Accelerator Optimisation Using Densely Connected and Invertible Neural Networks

Abstract

Particle accelerators are enabling tools for scientific exploration and discovery in various disciplines. However, finding optimised operation points for these complex machines is a challenging task due to the large number of parameters involved and the underlying non-linear dynamics. Here, we introduce two families of data-driven surrogate models, based on deep and invertible neural networks, that can replace the expensive physics computer models. These models are employed in multi-objective optimisations to find Pareto optimal operation points for two fundamentally different types of particle accelerators. Our approach reduces the time-to-solution for a multi-objective accelerator optimisation up to a factor of 640 and the computational cost up to 98%. The framework established here should pave the way for future online and real-time multi-objective optimisation of particle accelerators.

1. Introduction

Advances in accelerator science and technology have enabled discoveries in particle physics and other fields—from chemistry and biology to medical applications—for more than a century [1], with no end in sight [2]. Particle accelerators consist of a multitude of building blocks, and the relationship between changes in machine settings, known as design variables, and the corresponding particle-beam response is often non-linear. For this reason, the development and operation of a particle accelerator rely heavily on computational models. These models use first principles of physics to state the equations of motion, while numerical algorithms are then employed to solve them. These models provide valuable insight, but their high computational cost is prohibitive for many applications. For example, a single simulation of the Argonne Wakefield Accelerator (AWA) (Figure 1) takes approximately ten minutes with the high-fidelity physics model Object-Oriented Parallel Accelerator Library (OPAL) [3], which renders such models unsuitable for real-time usage. In order to achieve optimal operation points of these machines, genetic algorithms (GAs) are typically the method of choice for solving multi-objective optimisation problems [4,5,6,7,8]. However, an optimisation sometimes requires thousands of model evaluations. For this reason, the time-to-solution can easily be in the order of days. Recent work [9] addressed this issue by training a data-driven surrogate model to approximate the computer model at a single position of interest. Evaluating such a surrogate model takes less than a second and results in predictions that are very close to the high-fidelity physics model.

Related to these works, Scheinker et al. developed and demonstrated an online multi-timescale multi-objective optimisation algorithm that performs real-time feedback on particle accelerators [10]. In contrast to our contribution, they demonstrate a general control theoretical motivated feedback algorithm, capable of performing online multi-objective optimisation. Bayesian optimisation in combination with particle-in-cell simulations [11] is used to tune a plasma accelerator autonomously to the beam energy spread to the sub-percent at a given energy and intensity.

In this work, we introduce two new surrogate models, a forward model and an invertible model. Both are designed to model the machine at any longitudinal position, overcoming a main limitation of existing approaches. Moreover, the proposed methods for building these models are general enough to be applicable to any kind of particle accelerators. We demonstrate the generality of the ansatz by considering a linear accelerator and a quasi-circular machine, i.e., a cyclotron. Furthermore, the fast surrogate models enable optimisations on time scales suitable for online and real-time multi-objective optimisations of particle accelerators.

The forward model resembles existing models described, e.g., in [9]. While previous work focused on approximating OPAL only at one position (usually at the end) of the accelerator (as in [9]), our new model provides an approximation to OPAL at a multitude of positions along the accelerator, allowing us to estimate beam properties along the entire accelerator depicted in Figure 1d. To achieve this goal, our model takes both the design variables and the position along the accelerator as input and predicts the beam properties at that position. Consequently, our model is a complete replacement for OPAL and does not limit our options of objectives and constraints for beam optimisations. Existing models are not capable of optimising the beam properties at multiple positions at the same time, but our forward model enables such sophisticated optimisations, which we demonstrate empirically.

The invertible model is capable of solving both the forward and the inverse problem by providing two kinds of predictions, called the forward and the inverse prediction. The forward prediction approximates OPAL, i.e., it calculates the beam properties corresponding to a given machine configuration. The inverse prediction takes a user-imposed target beam and returns a machine configuration that realises an approximation to such a beam. In order to incorporate expert intuition and knowledge in the optimisation, the result of the inverse prediction initialises a GA-based multi-objective optimisation, which then requires fewer generations to converge than when the usual random initialisation is used.

We demonstrate the capabilities of the two surrogate models with the two real-world examples depicted in Figure 1. The first example is the AWA, which is a linear accelerator. The second one is the cyclotron used for the proposed Isotope Decay-at-Rest experiment, a proposed very-high-intensity electron-antineutrino source (from here on called the IsoDAR machine). We developed a forward model and an invertible model for both of these fundamentally different accelerators. Furthermore, we used them for a multi-objective accelerator optimisation to find machine configurations for desired beam properties.

1.1. Physics Models and Datasets

For both accelerators we used the same underlying physics model, based on the OPAL accelerator simulation framework. The dynamics of an accelerator are given by the following function:

$$\begin{array}{cc}\hfill \mathbf{f}:{\mathbb{R}}^m\times \mathbb{R}& \to {\mathbb{R}}^n\hfill \\ \hfill \mathbf{f}(\mathbf{x},s)& =\mathbf{y}\left(s\right),\hfill \end{array}$$

where the quantity $\mathbf{x}$ represents a vector of machine settings (design variables), the scalar $s\in \mathbb{R}$ denotes the position along the accelerator and the vector $\mathbf{y}\left(s\right)\in {\mathbb{R}}^n$ denotes the beam properties (quantities of interest). The function $\mathbf{f}$ represents the OPAL high-fidelity physics model.

Our surrogate models learn the relationship between design variables and quantities of interest based on examples, which means that we need to incorporate the physics of an accelerator into a dataset. We achieve this by sampling the space of design variables randomly and evaluating the sampled configurations of interest at multiple positions along the machine. One sample in the dataset consists of the design variables, a position along the machine and the corresponding quantities of interest.

1.2. Specifics of the Argonne Wakefield Accelerator Model

The AWA accelerator [12] consists mainly of an electron gun, several radio-frequency cavities for accelerating the electrons and magnets for keeping the beam focused. The electron gun generates bunches that form a train of Gaussians in the longitudinal direction, separated by a peak-to-peak distance $\lambda$. Other electron-gun variables are the gun phase $\varphi$, the charge Q, the laser spot size (SIGXY) and the current in the buck focusing and matching solenoids (BFS and MS). For details, we refer to Figure 1d. In total, we have nine design variables, denoted $x\in {\mathbb{R}}^9$, that define the various operation points of the AWA. The ranges of $x$ are determined by the physical limitations of the accelerator and listed in

Table A3. The design variables and their ranges for the AWA models.

Bound	IBF (A)	IM (A)	$\mathit{\varphi }$ (°)	ILS1 (A)	ILS2 (A)	ILS3 (A)	Q (nC)	$\mathit{\lambda }$ (ps)	SIGXY (mm)
Lower	450	100	−50	0	0	0	0.3	0.3	1.5
Upper	550	260	10	250	200	200	5	2	12.5

. The quantities of interest (QoIs), $y\in {\mathbb{R}}^8$ are as follows: the transversal root-mean-square beam sizes ${\sigma }_x,{\sigma }_y$, the transversal normalised emittances ${ϵ}_x,{ϵ}_y$, the mean bunch energy E, the energy spread of the beam $\Delta E$ and the correlations between transversal position and momentum $\mathrm{Corr}(x,p_x),\mathrm{Corr}(y,p_y)$. A summary and more details about the used QoIs are given in

Table A2. Quantities predicted by the various models. Checkmarks ✓ mean that the model predicts the quantity, while crosses x mean that the quantity is not predicted by the model.

	AWA		IsoDAR
	Forward	Invertible	Forward	Invertible
E (MeV)	✓	✓	✓	✓
$\Delta E$ (MeV)	✓	✓	✓	✓
${\sigma }_x,{\sigma }_y$ (m)	✓	✓	✓	✓
${ϵ}_x,{ϵ}_y$ (mm mrad)	✓	✓	✓	✓
$\mathrm{Corr}(x,p_x),\mathrm{Corr}(y,p_y)$	✓	x	x	x
${\sigma }_z$ (m	x	x	✓	✓
${ϵ}_z$ (mm mrad)	x	x	✓	✓
$N_l$	x	x	✓	✓
$h_x,h_y,h_z$	x	x	✓	✓

.

We built a labeled dataset of the AWA by using OPAL and the latin hypercube procedure [13] in order to sample the search space of $x$, which results in a training/validation set of 18,065 points and a test set of 913 points.

1.3. Specifics of the IsoDAR Cyclotron Model

The simulations are based on the latest IsoDAR [14,15,16] cyclotron design, with the nominal beam intensity of 5 mA ${\mathrm{H}}_2^{+}$ beam (equivalent to 10 mA of protons). Ongoing modeling efforts of IsoDAR consider the radial momenta $p_{r0}$, the injection radius $r_0$, the phase ${\varphi }_{rf}$ of the radio frequency of the acceleration cavities and the root mean square beam sizes ${\sigma }_x,{\sigma }_y,{\sigma }_z$ at injection. The physical ranges of the machine settings are given in

Table A4. The design variables and their ranges for the IsoDAR accelerator models.

Bound	${\mathit{p}}_{\mathit{r}0}$ ($\mathit{\beta }\mathit{\gamma }$)	${\mathit{r}}_0$ (mm)	${\mathit{\varphi }}_{\mathit{rf}}$ (°)	${\mathit{\sigma }}_{\mathit{x}}$ (mm)	${\mathit{\sigma }}_{\mathit{y}}$ (mm)	${\mathit{\sigma }}_{\mathit{z}}$ (mm)
Lower	0.002254	115.9	283.0	0.95	2.85	4.75
Upper	0.002346	119.9	287.0	1.05	3.15	5.25

. The surrogate models predict the transversal and longitudinal root mean square beam sizes ${\sigma }_x,{\sigma }_y,{\sigma }_z$, the projected emittances ${ϵ}_x,{ϵ}_y,{ϵ}_z$, the beam halo parameters $h_x,h_y,h_z$, the energy E, the energy spread of the beam $\Delta E$ and the particle loss $N_l$ (see Table A2). For this example, we have six machine settings $\mathbf{x}\in {\mathbb{R}}^6$ and twelve quantities of interest $\mathbf{y}\in {\mathbb{R}}^{12}$. All the samples are taken at turn 95, in the vicinity of the extraction channel.

The labeled dataset, obtained with OPAL, consists of 5000 random machine configurations. Among these, 1000 samples are randomly selected and used as the test set, 800 for the validation set and the remaining 3200 samples form the training set.

2. Results

2.1. Forward and Invertible Models

We assess the quality of the surrogate models by evaluating the adjusted coefficient of determination (${\overline{R}}^2$) and the relative prediction error at 95% confidence on the test set; i.e., data that have not been used during the development of the models. For the AWA model, the ${\overline{R}}^2$ values for the forward model are very close to one at nearly all positions along the machine, except for the solenoids and cavities. As a consequence, the variance of the dataset is explained well, as can be observed in Figure 2a. Due to the presence of angular momenta inside the solenoid region, the variance of the emittance is not captured well. In OPAL this effect is not suppressed, however, Bush’s theorem assures that, at the exit, the effect is reversed and does not contribute to our quantities of interest. In consequence, only the variance of the emittance is not captured well inside the solenoids and cavities, respectively. However, the relative prediction error for this quantity is still less than 25% for 95% confidence, as depicted in Figure 2c. The relative errors for the other quantities are even lower than 15% and 10% for ${\sigma }_x$ (Figure 2b) and $\Delta E$ (Figure 2d), respectively.

For the forward surrogate model of IsoDAR, the variance of the dataset is captured very well with a minimum ${\overline{R}}^2$ value of 0.95 among all predicted values, except for the quantities in z-direction. Here, it ranges between 0.82 and 0.86. The relative prediction error of the forward model is at most 5.6%.

Since we are mainly interested in the beam parameters, we evaluated the invertible model by performing first the inversion, computing machine settings from beam parameters, with the invertible model and then a forward prediction with the forward model. For the AWA model the error at 95% confidence is almost 40% for the beam size, emittance and energy spread. Nevertheless, the sampling error is only 20% for three out of four target beams in the test set. The energy, on the other hand, is obtained almost with perfect accuracy.

For the IsoDAR model, the relative errors for 95% confidence range from 0.08% to 11%. The corresponding ${\overline{R}}^2$ values lie between 0.79 and 0.96 except for the longitudinal quantities, for which we obtained values between 0.65 and 0.77. The detailed values for the IsoDAR model can be found in

Table A1. ${\overline{R}}^2$ values of the IsoDAR models. The values of the row “inv FP” refer to the forward prediction of the invertible model, whereas the values in “inv IP” refer to its inverse prediction. The latter ones are calculated with the help of the forward model.

	E	${\mathit{\sigma }}_{\mathit{x}}$	${\mathit{\sigma }}_{\mathit{y}}$	${\mathit{\sigma }}_{\mathit{z}}$	${\mathit{h}}_{\mathit{x}}$	${\mathit{h}}_{\mathit{y}}$	${\mathit{h}}_{\mathit{z}}$	${\mathit{ϵ}}_{\mathit{x}}$	${\mathit{ϵ}}_{\mathit{y}}$	${\mathit{ϵ}}_{\mathit{z}}$	$\mathbf{\Delta }\mathit{E}$	${\mathit{N}}_{\mathit{l}}$
forward	1.0	0.97	0.96	0.85	0.9	0.95	0.86	0.99	0.99	0.82	0.98	1.0
inv FP	0.97	0.92	0.91	0.75	0.84	0.88	0.81	0.93	0.95	0.69	0.94	0.9
inv IP	0.96	0.94	0.88	0.77	0.83	0.82	0.77	0.91	0.91	0.65	0.91	0.79

, Figure A1 and Figure A2.

Further details about the datasets and the model parameters can be found in Appendix B and Appendix C.

2.2. Multi-Objective Optimisation

To find good operation points for the accelerators, we solved multi-objective optimisation problems [5,6,9,17]. The standard approach for that is visualized in Figure 3a: A physics based model, such as OPAL, together with a GA and random initialization is used. We investigate two other approaches depicted in Figure 3b. First, we use a GA together with the forward model, instead of OPAL, and a random initialisation. In the second approach, the GA and the forward model are still employed, but the invertible model additionally provides us with a good initial guess for the optimisation.

2.3. Performance Metrics for the Optimisations

By characterising an approximation to the Pareto front, the set of optimal solutions encompasses two main aspects. First, we strive to converge quickly, using as few iterations as possible, towards non-dominated quantities of interest. The range of the objectives in the non-dominated set is chosen to analyse the convergence behaviour. Thus, for each generation the ranges of all objectives individually are compared for the two optimisation approaches. Second, we demand that the non-dominated quantities of interest are diverse, i.e., well spread out over the entire approximated Pareto front. The convex hull volume $V_{ch}$ (see Appendix A.1 for more details), the number of solutions of the non-dominated set, as well as the generational distance [18] and the hypervolume difference (i.e., the difference between the hypervolume of the optimal solution and the hypervolumes of the Pareto front of each generation, with a reference point chosen close to the Nadir point [19]) assess this quality aspect. Finally, plotting projections of the non-dominating set provides a holistic view to compare the effect of the initialisation strategies (random vs. invertible model initialisation) on both the convergence and the diversity of the non-dominated set. These plots show the evolution of the solutions of the optimisation problem across the generations directly.

2.4. Multi-Objective Optimisation for AWA

The constrained multi-objective optimisation problem of the following:

$$\begin{array}{ccccc}& \min \hfill & & {\sigma }_x\left(26\mathrm{m}\right)\hfill & \\ & \min \hfill & & {ϵ}_x\left(26\mathrm{m}\right)\hfill & \\ & \min \hfill & & \Delta E\left(26\mathrm{m}\right)\hfill & \\ & \mathrm{s}.\mathrm{t}.\hfill & \hfill {\sigma }_x\left(s\right)& \ge 2\mathrm{mm},\hfill & \hfill s\in 0,0.25,0.5,\cdots ,10\mathrm{m}\\ & & \hfill {\sigma }_x\left(s\right)& \le 5\mathrm{mm},\hfill & \hfill s\in 0,0.25,0.5,\cdots ,10\mathrm{m}\\ & & \hfill \left|\mathrm{Corr}(x,p_x)\right|& \le 0.1,\hfill & \hfill \mathrm{at}26\mathrm{m},\end{array}$$

represents a generic multi-objective optimisation task, found in many of the past and future AWA experiments [12]. The optimisation of beam parameters was required at 26 m for a pilot experiment performed in 2020, regarding a new scheme for electron cooling. In addition, there were multiple constraints at various positions along the accelerator. We solved the optimisation problem twice, once initialising the GA population randomly and once initialising the GA population with the invertible model outcome, where we utilized target values for the beam parameters as listed in

Table A6. Target vector for the biased AWA optimisation.

E	${\mathit{\sigma }}_{\mathit{x}}$	${\mathit{\sigma }}_{\mathit{y}}$	${\mathit{ϵ}}_{\mathit{x}}$	${\mathit{ϵ}}_{\mathit{y}}$	$\mathbf{\Delta }\mathit{E}$	s
47.5 MeV	2 mm	2 mm	3 mm mrad	3 mm mrad	60 keV	13.7 m

.

The objective values for different generations are depicted in Figure 4a. The optimisation initialised with the invertible model converged almost entirely after 10 generations. If random initialisation is used, more than 100 generations are needed to arrive at the same optimal configurations. The ranges of the beam properties in the non-dominated set confirm this observation (Figure 4b–d). Using the invertible model results in a bigger hypervolume even when prediction uncertainty is accounted for (Figure 4e). Furthermore, the optimisation initialised with the invertible model finds more solutions during the first 500 generations (Figure 4f). These two statements imply that the initialisation with the invertible model finds both more and more diverse non-dominated design variables.

The non-dominated feasible machine settings are evaluated with OPAL (dots in Figure 5a–c) in order to validate the predicted optimal objective values. All the OPAL evaluated points lie within the region of uncertainty for the optimal objective values at 90% confidence.

In Figure 5d, deeper insights into the constraint on the correlation parameter ($Corr(x,P_x)$) is given. Although the surrogate models found many configurations lying outside the allowed region, both initialisation strategies result in at least some feasible configurations. Importantly, no machine configuration in the training/validation set is feasible according to the optimisation problem, that means that no sample in these datasets fulfills the optimisation constraints. Nevertheless, some optimal points fulfill the correlation constraint.

2.5. Multi-Objective Optimisation for IsoDAR

For the IsoDAR case, we solved a two-objective optimisation problem, similar to that introduced by Edelen et al. [9]. We aim to minimise simultaneously the projected emittance ${ϵ}_x$ and the energy spread $\Delta E$ of the beam without any constraints.

$$\begin{array}{ccc}\hfill & \min \hfill & \hfill \Delta E\\ \hfill & \min \hfill & \hfill {ϵ}_x.\end{array}$$

As before, the optimisation problem is solved with a GA by using the forward surrogate model. First, a random initialisation is used and then an initialisation using the invertible model with the target beam parameters in

Table A5. Target vector for the biased IsoDAR optimisation.

E	${\mathit{\sigma }}_{\mathit{x}}$	${\mathit{\sigma }}_{\mathit{y}}$	${\mathit{\sigma }}_{\mathit{z}}$	${\mathit{h}}_{\mathit{x}}$	${\mathit{h}}_{\mathit{y}}$
112.0 MeV	2.4 mm	2.1 mm	1.5 mm	4.5	3.6
$h_z$	${ϵ}_x$	${ϵ}_y$	${ϵ}_z$	$\Delta E$	$N_l$
3.4	4.2 mm mrad	4.5 mm mrad	2.1 mm mrad	146.2 keV	7090.0

is performed.

The convergence of the objective values for the two optimisation approaches is depicted in Figure 6a. Both approaches result in the same non-dominated set. The optimisation with the invertible model initialisation converges faster and finds more optimal points already in the first few generations. The random initialisation approach needs approximately 40 generations until it is at the same level. Figure 6b shows that both initialisation schemes result in the end to nearly the same number of solutions. The difference of the objective values for the two approaches in the first generations of the optimisation can also be observed in Figure 6d, which depicts the ranges from minimum to maximum values of the non-dominated set for the two objectives.

The performance-metrics plots Figure 6c,g show the behaviour of the two approaches quantitatively. The generational distance for the invertible model initialisation starts at a much lower value. Therefore, the non-dominated sets found in the first generation are already closer to the final non-dominated set than the ones found with the random initialisation. The lower hypervolume difference and the higher convex hull volume for the invertible model initialisation indicate that the initial non-dominated set is more widely spread, but these advantages vanish after approximately 40 generations.

In order to validate the optimisation results, we randomly chose 100 points of the non-dominated sets of both approaches and simulated them with OPAL. We did not evaluate all non-dominated machine settings because of the high computational cost of OPAL for the IsoDAR example. As can be seen in Figure 6e,f, nearly all of the OPAL points lie within the 90-percent confidence region of the prediction with the test set, although only few points of the original OPAL dataset are in this area.

2.6. Computational Advantages

We compare solving an optimisation problem with OPAL to our approach of using a forward surrogate model instead of OPAL. There are two quantities to be considered for the speedup analysis. First, the time-to-solution t measured in hours and, second, the computational cost c measured in CPU hours. The detailed calculations for the improvement of both quantities can be found in Appendix D in the Appendix. Previous work [9] also calculated speedups, but these calculations do not include the computational cost and the entire cost of developing a surrogate model because they neglect the hyperparameter scan. Moreover, the existing speedups refer to a model that is only capable of predicting the quantities of interest at one position. Our calculations refer to our novel forward surrogate model, which is capable of predicting quantities of interest at a plethora of positions and do not only include the generation of the dataset but also the training, hyperparameter tuning and the cost of running the optimisations themselves. The invertible model reduces the number of generations. However, we calculated the speedup for a fixed number of generations; therefore, the effect of the invertible model is not included in the speedup calculations.

For the AWA (see

Table A8. Parameters related to the speedup calculation. For the AWA, the prediction time refers to predicting the beam properties with a spacing of 5 cm, i.e., at 520 positions. For the IsoDAR accelerator, we modelled only 1 position. The times are measured by measuring the prediction time for predicting $N_m$ machines and then dividing by $N_m$. For the AWA case, $N_m=64$; for IsoDAR $N_m=2^{20}$.

Quantity	Symbol	AWA	IsoDAR
Time for one OPAL evaluation	$t_o$	10 min	1.8 h
Time to train one forward model	$t_t$	49 h	1 h
Time to predict one machine with the surrogate	$t_p$	21 ms	26 μs
Number of unique machine settings in the dataset	n	21.000	5.000
Number of hyperparameters to try	$n_h$	100	120
Number of CPU cores per OPAL evaluation	$r_o$	4	1
Number of CPU cores to train and evaluate the surrogate	$r_t$	12	1
Number of generations	$n_g$	1.000	1.000
Number of individuals per generation	$n_i$	200	300

), we calculated the relative improvement $t_{\mathrm{OPAL}}/t_{\mathrm{surr}}\approx 3.39$ in terms of time-to-solution and $c_{\mathrm{OPAL}}/c_{\mathrm{surr}}\approx 1.83$ in terms of computational cost. This approach is more than three times faster and requires approximately 45% less computational resources (see Equation (A1) in Appendix D.3). For IsoDAR  the benefits are even more pronounced: $t_{\mathrm{OPAL}}/t_{\mathrm{surr}}\approx 640.00$ and $c_{\mathrm{OPAL}}/c_{\mathrm{surr}}\approx 59.20$. Hence, the usage of a surrogate model saves more than 98% of the computational cost and time resources compared to the approach using OPAL.

3. Methods

3.1. Forward Model

The forward surrogate model $\tilde{f}$ (Figure 3c) is a fast-to-evaluate approximation to the expensive OPAL-based physics model f.

$$\begin{array}{cc}\hfill \tilde{\mathbf{f}}:{\mathbb{R}}^m\times \mathbb{R}& \to {\mathbb{R}}^n\hfill \\ \hfill \tilde{\mathbf{f}}(\mathbf{x},s)& =\tilde{\mathbf{y}}\left(s\right)\hfill \\ \hfill \mathrm{s}.\mathrm{t}.\tilde{\mathbf{f}}(\mathbf{x},s)& \approx \mathbf{f}(\mathbf{x},s)\forall \mathbf{x},s.\hfill \end{array}$$

While existing work [9] is restricted to approximating the accelerator at a single position s, usually at the end, we built a model for the AWA that approximates the dynamics of the machine at all positions. Without loss of generality, the IsoDAR model predicts the beam properties only at the end of the machine, as ongoing design work requires. We follow the work by Edelen et al. [9] and use densely connected feedforward neural networks [20] as candidates for the forward model. We focus only on architectures where all hidden layers are of the same width in order to simplify the hyperparameter scan. For the AWA model, the hidden layers use the Rectified Linear Unit (ReLU) activation and the output layer uses the tanh activation function, whereas for the IsoDAR model tanh is used as an activation for the hidden layers, and the output layer activation is linear. We use the mean absolute error (MAE) as the loss function for the AWA model and the mean squared error (MSE) for the IsoDAR model.

We optimise the trainable parameters by using the Adam algorithm with default parameters. The non-trainable parameters are selected using grid search and can be found in

Table A7. Parameters and properties of our surrogate models. EOM means end of the machine.

	AWA		IsoDAR
	Forward	Invertible	Forward	Invertible
Modelled region	$s\in [0,26]\mathrm{m}$ (entire machine)	$s\in [0,13.7]\mathrm{m}$	EOM	EOM
Preprocessing $\mathbf{x}$	Scale to $[-1,1]$	QuantileScaler Scale to $[-1,1]$	Scale to $[0,1]$	Scale to $[0,1]$
Preprocessing $\mathbf{y}$	Shift to be positive Apply $log\left(\right)$ Scale to $[-1,1]$	Clip ${ϵ}_x\in [0,200]\mathrm{mm}\mathrm{mrad}$ Scale to $[-1,1]$	Scale to $[0,1]$	Scale to $[0,1]$
Dimension of the latent space	-	1	-	1
Nominal Dimension	-	12	-	14
Distribution of the latent space	-	Unif(-1, 1)	-	Unif(−1, 1)
Loss function	MAE	${\mathcal{L}}_{\mathrm{inv}}$ (Equation (1)) with weights: $w_x=400$ $w_y=400$ $w_z=400$ $w_r=3$ $w_{\mathrm{artificial}}=1$	MSE	${\mathcal{L}}_{\mathrm{inv}}$ (Equation (1)) with weights: $w_x=400$ $w_y=400$ $w_z=400$ $w_r=10$ $w_{\mathrm{artificial}}$ = 1
Training algorithm	Adam	Adam	Adam	Adam
Learning rate	${10}^{-4}$	${10}^{-4}$	${10}^{-4}$	${10}^{-3}$
Batch size	256	256	256	8
Number of epochs	56	15	5000	30
Architecture	$7\times 500$	$8\times 3\times 100$	$6\times 80$	$5\times 2\times 60$
Activation of hidden neurons	ReLU	ReLU	tanh	ReLU
Number of trainable parameters	1,512,618	688,192	33,932	91,340
CPU cores for training	12	12	1	1
Time for training	49 h	31 h	1 h	1 h

, along with further information about the models. All the models (forward and invertible) are implemented by using the TensorFlow framework [21] in version 2.0. For the hyperparameter scan of the IsoDAR model the RAY TUNE library [22] is used in addition. The models are trained and evaluated on the Merlin6 cluster at the Paul Scherrer Institut, using 12 CPU cores for the AWA dataset and one CPU core for the IsoDAR case.

3.2. Invertible Model

The invertible model (see Figure 3d) is capable of performing two kinds of predictions: A forward prediction approximating OPAL and an inverse prediction aiming to solve the following inverse problem. Given a target beam, the invertible model predicts a vector of design variables such that the corresponding beam is a good approximation of the target beam.

We follow the ansatz by Ardizzone et al. [23] and refer to their work for the details of the architecture of the inverse model. Ardizzone et al. build a neural network consisting of invertible layers, so-called affine coupling blocks (see Figure 3d). Each of them contains two neural networks, which are sharing parameters. We decided that all hidden layers of the internal networks have the same number of neurons, and the internal networks of all affine coupling blocks have the same architecture. Each internal network consists of the same number of hidden layers of the same size. All hidden neurons of the internal networks use the ReLU activation function.

The solution of the inverse problem is not unique, hence, a mechanism to select one solution is needed. This is realised by mapping the machine settings not only to the design variables but also to a latent space $\mathcal{Z}\subset {\mathbb{R}}^{d_z}$ of dimension $d_z$. This space follows a known probability distribution (e.g., uniform distribution as used in this work), which allows sampling random points in it. Sampling a point in the latent space corresponds to selecting one solution of the inverse problem. For this reason, the inverse prediction is also called sampling.

For technical reasons, an invertible neural network requires that the input and output vectors have the same length. This is generally not fulfilled. If the input and/or output vectors are not big enough, vectors containing noise of small magnitude are added so that the total dimension of input and output vectors is d. We follow the advice of Ardizzone et al. [23] and allow padding not only on the smaller vector but also on both in order to increase the network width and, therefore, the model capacity. The total dimension d is a hyperparameter. We denote the padding vectors ${\mathbf{x}}_{\mathrm{pad}}\in {\mathbb{R}}^p,{\mathbf{y}}_{\mathrm{pad}}\in {\mathbb{R}}^q$ and obtain $d=m+1+p=n+1+\dim \left(\mathcal{Z}\right)+q$.

The mathematical formulation of the forward prediction is as follows.

$$\begin{array}{cc}\hfill \widehat{\mathbf{f}}:{\mathbb{R}}^m\times \mathbb{R}\times {\mathbb{R}}^p& \to {\mathbb{R}}^n\times \mathbb{R}\times \mathcal{Z}\times {\mathbb{R}}^q\hfill \\ \hfill \widehat{\mathbf{f}}(\mathbf{x},s,{\mathbf{x}}_{\mathrm{pad}})& =(\widehat{\mathbf{y}},s,\mathbf{z},{\mathbf{y}}_{\mathrm{pad}})\hfill \\ \hfill \mathrm{s}.\mathrm{t}.\widehat{\mathbf{y}}& \approx \mathbf{f}(\mathbf{x},s).\hfill \end{array}$$

The inverse prediction is now written as follows.

$$\begin{array}{cc}\hfill {\widehat{\mathbf{f}}}^{-\mathbf{1}}:{\mathbb{R}}^n\times \mathbb{R}\times \mathcal{Z}\times {\mathbb{R}}^q& \to {\mathbb{R}}^m\times \mathbb{R}\times {\mathbb{R}}^p\hfill \\ \hfill {\widehat{\mathbf{f}}}^{-\mathbf{1}}(\mathbf{y},s,\mathbf{z},{\mathbf{y}}_{\mathrm{pad}})& =(\widehat{\mathbf{x}},s,{\mathbf{x}}_{\mathrm{pad}})\hfill \\ \hfill \mathrm{s}.\mathrm{t}.\mathbf{f}(\widehat{\mathbf{x}},s)& \approx \mathbf{y}.\hfill \end{array}$$

In order to improve the performance of the invertible model  we use a best-of-n strategy. For each inverse prediction (at inference time only), $n_{\mathrm{tries}}$ design variable configurations are sampled, evaluated with the forward pass and the best configuration is chosen as the final prediction. Note that we do not need to rely on a forward model but only on the invertible model itself. We chose $n_{\mathrm{tries}}=32$ for both the IsoDAR model and the AWA, as we observed no significant improvement for bigger values. All prediction errors are calculated with this particular choice. The loss function consists of multiple parts:

$${\mathcal{L}}_{\mathrm{inv}}=w_x{\mathcal{L}}_x+w_y{\mathcal{L}}_y+w_z{\mathcal{L}}_z+w_r{\mathcal{L}}_r+w_{\mathrm{artificial}}{\mathcal{L}}_{\mathrm{artificial}},$$

(1)

where the ${\mathcal{L}}_x$ loss ensures that the sampled machine setting vectors $\widehat{\mathbf{x}}$ follow the same distribution as the machine settings in the dataset, the ${\mathcal{L}}_z$ loss ensures that the latent space vectors follow the desired distribution. Both loss functions are realised as mean field discrepancy [24]. The ${\mathcal{L}}_y$ loss is the mean squared error between $\widehat{\mathbf{y}}$ and ${\mathbf{y}}_{\mathrm{true}}$, and the ${\mathcal{L}}_{\mathrm{artificial}}$ loss is the sum of the MSEs of both ${\mathbf{x}}_{\mathrm{pad}}$ and ${\mathbf{y}}_{\mathrm{pad}}$. The reconstruction loss ${\mathcal{L}}_r$ ensures that the inverse prediction is robust with respect to perturbations of small amplitude $ϵ$.

$${\mathcal{L}}_r={\displaystyle \sum _{i=1}^N}{∥{\widehat{\mathbf{f}}}^{-\mathbf{1}}\left(\widehat{\mathbf{f}}({\mathbf{x}}_i,s_i)+ϵ\right)-{\mathbf{x}}_i∥}^2.$$

The inverse prediction aims to generate machine settings $\widehat{\mathbf{x}}={\widehat{\mathbf{f}}}^{-1}{(\mathbf{y},s)}_{1:m}$ such that the resulting beam $\widehat{\mathbf{y}}=\mathbf{f}{(\widehat{\mathbf{x}},s)}_{1:n}$ comes close to desired target-beam properties. In order to estimate the prediction accuracy of the models, we reproduce vectors y from the test set. It is not reasonable to compare the sampled machine settings with the corresponding setting $\mathbf{x}$ in the test set because multiple machine settings might realise similar beams. Instead, we performed a forward prediction with the forward surrogate model. This allows us to estimate which beam properties the sampled configuration results in i.e. $\tilde{\mathbf{y}}=\tilde{\mathbf{f}}(\widehat{\mathbf{x}},s)$. That quantity is then used to describe the error of the inverse prediction.

4. Discussion and Outlook

We have introduced a novel flexible forward surrogate model that is capable of simulating high-fidelity physics models at a plethora of positions along the accelerator. We have shown that forward surrogate models reduce both the time (by a speedup factor of 3.39 for the AWA and 640 for IsoDAR) and the computational cost (by 45% and 98%) of beam optimisations. The benefits increase if several optimisations need to be performed because the time-consuming and computationally expensive part is the development of the model. Once the models are developed, the optimisations are almost free and can be performed in approximately 10 min on 12 cores. This allows experimenters to adapt to new situations by changing constraints and objectives and to rerun the optimisation with very little computational resources. Furthermore, the surrogate models themselves are useful in control-room settings, as they are fast enough to be used in a graphical user interface. This provides operators the possibility to quickly try out different machine settings. Despite these benefits of the surrogate model approach, one needs to be careful when applying them to problems with narrow constraints. Neural networks can only accurately model the machine-setting regimes represented in the dataset on which they are trained. If the feasible region of an optimisation problem is narrow, the models will hardly be able to find configurations that fulfill all constraints. One solution to this issue is to sample the feasible region more densely instead of sampling the design space uniformly. This is difficult because the feasible region in design space is usually a priori unknown. We have also presented an invertible model, which makes it possible to simulate the forward direction as well as the inversion, thus predicting machine settings that result in desired beam parameters. The invertible model alone has the potential to support the design of accelerators, and it could also be used to implement a failure-prediction system or an adaptive control system. Another use-case of invertible surrogate models was demonstrated in this paper. Invertible surrogate models were shown to be able to bias the initialisation of a GA towards the optimal region, thereby incorporating prior knowledge and experience of experimenters. This reduces the number of generations needed for convergence of the GA when a forward surrogate model is used.

As the optimisation is almost free when a forward surrogate is employed, there is not much incentive to go through the labour of developing an invertible model. However, if OPAL is used to evaluate all individuals, the reduced number of generations might reduce the time and cost of the optimisation significantly. Using OPAL solves the problem of the under-represented feasible space because OPAL is capable of modelling the relevant physics of a particle accelerator. The approach of applying the biased initialisation to an OPAL optimisation might combine both the advantage of needing fewer generations with the ability to accurately represent the feasible space. The accompanying savings might justify the cost of developing the invertible model. Further research is needed to investigate this prediction. Finally, the choice of the target vector is important, and additional research is needed to investigate its impact.

This research is a proof of concept and shows that the forward and inverse models of particle accelerators can be approximated with neural networks. So far, no real machine/experiment data are included in this work. In particular, the influence of measurement errors on the behavior of the neural network models is of great interest and needs further investigation. A first step to assess that would be to train the neural network models with simulated data and to test them with real measurement data. This would be possible only for the AWA but not for IsoDAR since the accelerator has not been built yet.

The societal impact of accelerator science and technology will continue, with no end in sight [1,2,16,25]. With this research study, we add new computational tools and, hence, contribute to the quest of finding optimal accelerator designs and machine configurations, which in turn are likely to greatly reduce construction and operational costs and to improve physics performance.