Symplectic neural surrogate models for beam dynamics

Development of robust machine-learning (ML) based surrogates for particle accelerators can significantly benefit the modeling, design, optimization, monitoring and control of such accelerators. It is desirable that the surrogate models embed fundamental physical constraints to the interaction and dynamics of the beams, for which an accelerator must be designed to operate upon. We implement and train a class of phase space structure-preserving neural networks — Henon Neural Networks (HenonNets) [1], for nonlinear beam dynamics problems. It is demonstrated that the trained HenonNet model predicts the beam transfer matrix to a reasonably good accuracy while strongly maintaining the symplecticity. To explore such model’s applicability and flexibility for high brightness or intensity beams, we further test it with beam dynamics in the presence of electrostatic and radiative collective effects. Our results indicate that HenonNet may be used as a base ML model for the surrogate of complex beam dynamics, thus opening up a wide range of applications.


ML models for beam dynamics
Machine learning models are increasingly being employed in industrial and scientific applications, where there is often a need to make fast prediction based on observation or simulation data.For particle accelerators, many ML models have recently been proposed and developed [2].Over the past decade, some general common practice also emerge which help guide the development of successful ML models.These include, • A good underlying parameterized model of the data; • Sufficient data to cover the parameter region; • Effective optimization method to fit data to the model.
Many ML models rely on the generic universal approximation theorem of neural network [3] which can be used to approximate any nonlinear function mapping the input and output.In that sense, a generic ML model can be used for our purpose, but such a model usually lacks the efficiency in training and accuracy in prediction, particularly in regard to the quality of prediction from the interpolation among trained data and for the unseen data, as well as the model interpretability.In physics-informed ML models, one possible way is to enforce the physical constraints to the underlying ML model, e.g., in the form of conservation laws.The simplest way is to enforce such constraints in the loss function, thus one hopes to use constrained optimization to incorporate the knowledge of the physical system in the prediction.Another type of ML models build directly on the mathematical structure of the physics problem, e.g., for a Hamiltonian system, ML models can be used to approximate either the Hamiltonian, or the generating functions of its canonical transform.The beam dynamics falls into this category, where the evolution of the beam in 6-dimensional phase space (x, y, z, p x , p y , p z ) follows Hamiltonian dynamics in the absence of the collective effects.Furthermore, a ML model that respects such Hamiltonian structure may be used as the base model for the beam dynamics with collective effects.In this work, we present our beam dynamics surrogate models utilizing a particular ML model with built-in Hamiltonian structure -the HenonNet.

Hamiltonian Beam Dynamics
Here we provide a brief summary of the Hamiltonian formalism for a particle beam.For the dynamics of a single particle of charge q e , coordinate ξ(t) = (q(t), p(t)) in an external field with potential ϕ and A , the Hamiltonian H(ξ(t), t) is where the Poisson matrix J = 0 I −I 0 .Here ξ is time-dependent but it can be also s-dependent as in accelerator lattice with corresponding Hamiltonian.

HenonNet
The architecture of HenonNet is proposed in [1], where a fast surrogate model to predict Poincaré plot of a magnetic field is considered.The network is inspired by a universal approximation theorem [4] that for given ϵ > 0 and a C r symplectic diffeomorphism F, there exists N and a Henon-like map H [5] such that ∥F − H 4N ∥ C r < ϵ.Ref. [4] however only provides an existence theorem, while we rely on ML optimization to construct such a map [1].Since developed, HenonNet has been extended to other applications such as discovering equations of state in the symplectic manifold and surrogate models with a provable adiabatic invariant for multi-scale Hamiltonian systems [6].Such a methodology has recently been applied to beam dynamics [7].However, one additional challenge is that there exist a large number of physics parameters in the system, while all the previous HenonNet studies only considered a Hamiltonian under a set of fixed parameters.A naive application of HenonNet inevitably requires to retrain the network when parameters change, which is not practical.To overcome such a difficulty, we generalize the previous HenonNet to incorporate physical parameters into the network.For a parameteric Hamiltonian with the canonical phase space of R 2N and a parameter vector µ = (µ 1 , µ 2 , . . .µ m ) in domain P ∈ R m , we first construct a scalar function V W : R N × P → R by a feed-forward neural network, where W stands for its trainable weights.The input of this network is the canonical variable p and the parameter vector µ.The Henon layer with potential V W and bias η ∈ R N is the neural network layer HL (W ,η) : R 2N × P → R 2N given by where the individual map H (W ,η) : (q, p; µ) → (q, p) is given by It is easy to show the Henon layer is a symplectic map with respect to the phase space R 2N for any given weights W , bias η, and parameter µ.The HenonNet is a sequential composition of several Henon layers.The concept of HenonNets is therefore similar with the conventional feedforward neural network except a custom layer is used.The advantage of using the custom Henon layer is its built-in property of symplecticity, a fundamental property of Hamiltonian flows, which hence guarantees symplecticity of HenonNet.Note that in practice different potential functions (feed-forward networks) are used in each Henon layer.The modified network has a parametric symplectic universal approximation theorem, which suggests a uniform error bound for the closed parameter set.This theorem will be discussed more carefully in a forthcoming ML theory paper.Here, we focus on demonstrating its power for fast prediction of the beam dynamics in a range of physical parameters.Note that predictions are always significantly better when the network essentially performs interpolations instead of extrapolations in the parameter space.

Surrogate model for linac
We use HenonNet to build a surrogate model of the Coupled-Cavity-Linac (CCL) for the LANSCE accelerator [8].We first test the surrogate model for a single CCL RF tank (tank 5.1) with mismatched and shifted beams, and analyze its accuracy.Beam dynamics for this CCL tank without the space-charge effect are simulated using the BEAMPATH [9] particlein-cell code for the following cases (Fig. 1): ( 1 Our HenonNet model consists of 5 layers with 10 weights per dimension per layer and a total of 265 trainable weights including the bias.50000 simulation particles are used for training and validation.The training batch size is 1250.A custom learning rate schedule is also used.Figure 2 shows the comparison of the HenonNet model prediction with ground truth when case 2 is used for training.For all models, training/validation data are random selected with 50%:50% ratio to avoid over-fitting.Good agreements between prediction and ground truth are observed except for the beam tails in the p z − z phase space of case 3, which is likely due to the lack of training data there.Nonetheless, for case 3, the model has already demonstrated some level of extrapolation beyond the training data (e.g., for particles ∼ z = −1.5 cm).
As the HenonNet model is built using neural network, it is possible to use automatic differentiation provided by the ML tools to obtain the Jacobian of the map between its input and x,y,z .The 6D learned transfer matrix also has very small off-block-diagonal components.This result confirms that the model indeed captures the essential structure of the Hamiltonian beam dynamics and the level of accuracy achieved is ensured by the symplecticity of the model.Another metric of the prediction accuracy is the relative error of the emittance δϵ ≡ |(ϵ Henon − ϵ CCL )/ϵ CCL |.Such errors of the predicted beam distributions for all 4 cases, when using case 3 for training, are summarized in Table 1.These errors are also consistent with the error obtained from the linear transfer matrix.

Parameterized models for space charge effect
To build a surrogate model for the entire LANSCE CCL, where the beam is accelerated from 100 MeV to 800 MeV in ∼ 700 m, we break it into 4 sections of equal length and train their corresponding parameterized HenonNet models, each with 5 layers and 25 weights per layer.The parameter is the beam current I and the training data are obtained from 3 simulations (I = 0, 10, 20 mA).Currently only models for the longitudinal phase space with space charge

Model for coherent synchrotron radiation
We also demonstrate the use of a basic HenonNet model as surrogate for the Coherent Synchrotron Radiation (CSR) effect in the dynamics of high brightness electron beams.Figure 4 shows the effect on the phase spaces (where δ = ∆E/E and τ = ct − z), which manifests as the kinks, from the 1D CSR wakefield for a 130 MeV, 10 nC electron beam in a magnetic bend.A HenonNet of 50 layers and 10 weights/layer are used for the surrogate model.Good validation agreement with the ground truth is observed in Fig. 4.

Conclusion
We have built a preliminary surrogate model of the beam dynamics based on the HenonNet, using a major part (the 700 meter Coupled-Cavity Linac) of the LANSCE accelerator as test bed and demonstrated that it achieves comparable accuracy to detailed beam dynamics simulations.Such model can make predictions in 15 millisecond -a 10 6 times improvement that opens up possibilities for applications such as real-time feedback control and fast design optimization.Preliminary validation result for beam collective effects are also explored, which will be extended to the study of more realistic and robust surrogate models in the future.

Figure 1 .
Figure 1.Longitudinal phase spaces for the four beams at the (a) entrance and (b) exit of the CCL tank 5.1.The synchronous particle is shown for reference.

14thFigure 2 .= 1
Figure 2. Validation of mismatched beam (case 2) and prediction of strongly mismatched beam (case 3) in the CCL tank shown in Fig. 1.All results are relative to the synchronous particle, which is shown in Fig. 1.

Figure 3 shows
the model validation result for the I = 10 mA beam.We have also built a 2-parameter (i.e., RF field strength E and beam current I) surrogate model for the same CCL tank in Fig.1.Figure3also shows the validation result for this 2-parameter model.I=10 mAE=1.1 MV/m, I=10 mA

Figure 3 .
Figure 3. Validation of parametric HenonNet models for CCL with 1 and 2 parameters.

Table 1 .
Relative error for the emittance The full CCL model is then composited from the models for each section.