Toward fully automated UED operation using two-stage machine learning model

To demonstrate the feasibility of automating UED operation and diagnosing the machine performance in real time, a two-stage machine learning (ML) model based on self-consistent start-to-end simulations has been implemented. This model will not only provide the machine parameters with adequate precision, toward the full automation of the UED instrument, but also make real-time electron beam information available as single-shot nondestructive diagnostics. Furthermore, based on a deep understanding of the root connection between the electron beam properties and the features of Bragg-diffraction patterns, we have applied the hidden symmetry as model constraints, successfully improving the accuracy of energy spread prediction by a factor of five and making the beam divergence prediction two times faster. The capability enabled by the global optimization via ML provides us with better opportunities for discoveries using near-parallel, bright, and ultrafast electron beams for single-shot imaging. It also enables directly visualizing the dynamics of defects and nanostructured materials, which is impossible using present electron-beam technologies.

www.nature.com/scientificreports/ the ML model, which comprises two sequential NNs, as shown in Fig. 2b. The first NN starts from the gun to the sample, named GTS-ML, and the second one starts from the sample to the detector, named STD-ML. Among the two different types of ML configurations, the two-stage approach is preferred for the following reasons. For the first option, as one NN covers the entire beamline, the potential applications are quite limited. The BD patterns and machine parameters are the input data and output predictions of this ML model, respectively. It only can be used to automate the setup of a UED instrument as no intermediate information is provided. Alternatively, in the second option, by separating the beamline into two parts 29 , from the gun to the sample and from the sample to the detector, shown as the schematic layout of the UED beamline in Fig. 1 and the schematic layout of the NN model in Fig. 2b, we gain some advantages. This two-stage approach has much broader applications because of the availability of the electron beam properties. The two sequential NNs, GTS-ML (sample independent) and STD-ML (accelerator independent) 30 , are mutually independent. As a result, the frequency of updating the two-stage model can be lower; instead, one only needs to update the stage model with changes. Since an ML model is purposely built to solve the complex inverse problem, in our case, based on the BD patterns recorded by the detector as the input data, the two-stage ML model will ultimately make the predictions of the machine parameters, which will be used to automate the setup of a UED instrument. Furthermore, the predictions by the second NN (STD-ML), the electron beam properties at the sample, will not only provide the input to the first NN (GTS-ML) to enable the full automation of the UED instrument, but also make the real-time electron beam information available for data binning, filtering, correcting, and integrating, which enables singleshot noninvasive diagnostics for the UED instrument with high repetition rate and high fidelity. In a summary, one can apply the two-stage ML model not only to automate the UED instrument, but also to achieve real-time single-shot electron beam diagnostics while the experiments are ongoing. It is more important that the second option is built on the success of our early studies 30 . By adopting the two-stage approach, the demonstrated STD-ML model with a few new developments can be applied as the second NN.
In this report, we mainly focus on demonstrating the feasibility of the GTS-ML model and characterizing its performance. First, we perform the simulation study using the GPT code. A GTS-ML model is subsequently built based on the simulation results. We show that not only the predicted machine parameters meet the required precision, but also the prediction is extremely fast, taking less than 0.5 µs in the batch mode. In addition, to improve the STD-ML model, we have applied the hidden correlation between the electron beam properties and the features of a BD pattern as the model constraints, successfully improving the accuracy of the energy spread prediction by a factor of five and making the beam divergence prediction two times faster.

Feasibility study of the GTS-Ml model
Building a model. To demonstrate the feasibility of the GTS-ML model, GPT simulations from the gun to the sample provide the required training datasets to build the GTS-ML model. The input to the GPT simulation are several key machine parameters, including the spot size, x-jitter and y-jitter of the center position of an electron bunch emitted by the photocathode, the amplitude and phase of the RF waveform, and the current of the solenoid magnet. The charge distribution of an electron beam emitted by the photocathode is primarily determined by the characteristics of the laser pulse since the photoelectron emission from the cathode is prompt with respect to the laser light 31 . The output of the GPT simulations are the electron beam properties, including the x and y divergences, x and y spatial positions, energy, and energy spread, at the sample. The jitters of the sample can be neglected due to the following reasons: (1) drift, fluctuation, and vibration of a sample do not change the BD pattern since the electron diffraction is in Fourier space; (2) although mechanical expansion and compression of a sample could alter the BD pattern, such effects are usually much smaller compared to the peak broadening caused by the energy spread and divergences of the electron beam. The electron beam properties and the machine parameters are the input and output parameters of the GTS-ML model, respectively. For the ML models in both stages, each training data set is divided into three subsets: the training set, the validation set, and the test set.
The reasons to choose these six machine parameters as variables for each individual run of the GPT simulation while all other parameters are fixed are: (1) they are the main parameters determining the electron beam dynamics 9-12,31 ; (2) they are highly correlated with the output of the simulation (i.e., the electron beam properties at the sample); (3) their settings are absolutely necessary in determining the electron beam properties at the sample, which are unique and specific for different type of the UED experiments (i.e., the high charge single-shot and low charge accumulation modes) 3 ; (4) they are often fluctuant and, therefore, need to be constantly tuned during an experiment. Other parameters, e.g., the pulse duration of a photocathode drive laser system and the standard charge of a UED instrument, are fixed to 250 fs in full width half maximum and 0.5 pC, respectively, as they are comparably more stable in the duration of an experiment. Compared to the inherent RF-related out of sync and fluctuations 1-3 , the jitters caused by the drive laser properties and photon emission processes are often negligible (at least one order of magnitude smaller) 31,32 .
To take the 3D space charge effect into a proper consideration, we choose spacecharge3Dmesh, which is a mesh-based routine, as the 3D space-charge model. It is not only fast (computing time scales almost linearly with the number of particles) but also reliable for a 0.5 pC short electron bunch with low energy spread 9 . The time over which the particles travel from the gun to the sample is 4.18 ns and the step size is set to 5 ps. The number of particles is chosen to be 10,000, which is a good compromise between accuracy and computing speed. The GPT simulation has been bench-marked with our early study result of the influence of the focusing magnets on the beam size 1 .
To have the full coverage of the parameter space of a UED instrument, we build a table based on the Latin hypercube sampling scheme 30,33 . The Latin hypercube sampling method provides an efficient coverage of the parameter space with the maximum randomness and the highest density of the population. Each set of the www.nature.com/scientificreports/ machine parameters, which corresponds to a row in the table and determines an individual GPT simulation run, has a total number of 10,000 values, hence, 10,000 × 6 data points. The ranges of variations for these six machine parameters are determined by the experimental setup of the UED beamline 1,2 . Each GPT simulation run provides a set of the electron beam properties at the sample, including x and y divergences, x and y spatial positions, energy, and energy spread.
To build the GTS-ML model, the results of the GPT simulations form a training dataset. To be specific, for each training data set, the input data are the electron properties (p1, p2, …, p6) at the sample, including x and y spatial positions (avg(x) and avg(y)), x and y divergences (std(x′) and std(y′)), energy spread and energy (∆E/E and γ ), and the output are the six machine parameters (v1, v2, …, v6), x-jitter (x off ), y-jitter (y off ) and spot size (σ r ) of the electron bunch emitted at the photocathode, amplitude jitter (∆V RF / V RF ) and gun phase (θ gun ) of the RF waveform, and the normalized current of the solenoid magnet (I / I max ), as shown in Table 1. In GPT code. The beam energy is represented by the Lorentz factor γ.
Optimizing the model-architecture. The relationships between the input and output parameters of the training results are shown as the correlation plots in Fig. 3a. The x-axes are the data labels of those six machine parameters (v 1 , v 2 , …, v 6 ), which determine the GPT simulation conditions, and the y-axes are the electron beam properties (p 1 , p 2 , …, p 6 ), which are the output of the GPT simulation. Mean correlation ( ρ n ) vs index (n) of the machine parameters is plotted in Fig. 3b. The mean correlation is calculated via Eq. (1), where ρ n,j is the correlation between v n and p j (n, j = 1, 2, …, 6) and M is the size of the training dataset.
Since the input and output parameters are highly correlated (i.e., v 6 (0.87)), we expect the training of the GTS-ML model to be successful. The results are quite reasonable except-the predicted precision of v 4 (∆V RF / V RF ) is significantly worse than we initially expected, having a mean square error (MSE) of as large as 0.24 via a simple 1-layer ML model, shown as the black curve in Fig. 4a. The reason for the poor performance could be the space charge effect that makes the relationships between the machine parameters and the electron beam properties at the sample very nonlinear; therefore, the 1-layer ML model ( Fig. 4b) is too simple to reproduce such a complicated nonlinear dynamic process. To take the space-charge induced nonlinear effect into a proper consideration, a 4-layer ML model (Fig. 4c) has been implemented. The number of trainable parameters is 838 and 87,094, for the 1-layer model and 4-layer model, respectively. For the same training data set, the 4-layer ML model significantly improves the accuracy of the model prediction, by about 50% in terms of the MSEs of untrained data sets for all six machine parameters, shown as the red curve in Fig. 4a. On the other hand, the complexity of the 4-layer model is sufficient to fit the data, while overfitting was under control with the early stopping strategy 34 on the 10,000-point training dataset. As the result, this 4-layer NN configuration will be chosen to build the GTS-ML model.

Defining machine parameter-space.
Since v 3 (σ r ), v 4 (∆V RF /V RF ) and v 5 (θ gun ) are the machine parameters, with the most significant impact on the space charge effect, they are the most challenging ones to be predicted with high fidelity. To achieve the best possible accuracy, these three parameters are varied in GPT simulations to generate three different training datasets (named set #1, #2, and #3) for comparison, shown in Fig. 3a as red, green, and black dots, respectively. The ranges of the variations of those six training variables (v 1 , v 2 , v 3 , v 4 , v 5 , and v 6 ) and their MSE via the ML model are listed in Table 2 for data set #1, #2, and #3 with some details. Since both v 4 (∆V RF /V RF ) and v 5 (θ gun ) affect the electron beam energy, v 4 and v 5 have been deliberately varied much more in set #3 (± 2.5%, as black dots in row 4 of Fig. 3a) and set #2 (5°-60°, as green dots in row 5 of Fig. 3a), respectively. As the result, the optimal ranges of variations are determined and their corresponding MSEs are listed in the Summary columns of Table 2. With the improved 4-layer deep NN ML model, the ranges of variations of data set #2 generate the highest precisions with the MSE reduction from 30.0 µm to 22.8 µm for the spot size of the electron beam, from 0.006 to 0.003 for the RF amplitude jitter, and from 5.9° to 3.7° for the  4 , v 5 , and v 6 , with respect to the electron properties (input data), p 1 , p 2 , p 3 , p 4 , p 5 and p 6 , are listed, including the upper maximum and lower minimum limits for all parameters and their physical definitions. Compared to the RF synchronization (θ gun ) and fluctuations (∆V RF / V RF ), the laser pulse duration (σ t ) and total charge (Q tot ) are more stable and fixed in the GPT simulation. www.nature.com/scientificreports/ , v 4 (∆V RF /V RF ) and v 5 (θ gun )) that have the most significant impact on the space charge effect, are varied in GPT simulations to generate three different training datasets (named set #1, #2, and #3). For these three training data sets, the correlations between the input and output parameters are plotted. Set #1, #2 and #3 are plotted as the red, green, and black dots, respectively. x-axes are the data labels of the 6 machine parameters (v1, v2, …, v6), which correspond to row 1, 2, 3, 4, 5, and 6, respectively; y-axes are the electron beam properties (p1, p2, …, p6), which correspond to column 1, 2, 3, 4, 5, and 6, respectively. Also, the tile of each subplot provides the information of the plotted machine parameter as the x-axis and electron beam property as the y-axis. (b) Mean correlation ( ρ n ) vs index (n) of machine parameters. Error is estimated via the RMS deviation of all 3 datasets. www.nature.com/scientificreports/ gun phase. It is likely that the residual MSEs of those six machine parameters are limited by the space charge induced nonlinear effect, since the dependence of the MSE on the range of variation is not so strong. The MSEs of data set #1 and set #2 are quite similar, especially for v 4 and v 5 after unit conversion, shown as the secondary y-axis (red) in Fig. 5a and b respectively. Also, any increase of the size of a training dataset beyond 10 k doesn't help in improving the fidelity of the prediction any further.
To establish a deterministic one-to-one map between the predicted machine parameter by the ML model and the experimental setting of the UED instrument, the accuracy of the six machine parameters must match   4 , v 5 , and v 6 ), including the upper maximum and lower minimum limits, and MSEs regarding the entire dataset, including before and after unit conversion, are listed for those three data set #1, #2, and #3. The optimal ranges of variations, corresponding MSEs and targeted precisions of those six machine parameters are listed in the summary. www.nature.com/scientificreports/ well with the precision of the UED instrument. The targeted precisions of the machine parameters are listed in Summary of Table 2. As expected, we have achieved the reasonably good agreement for the 200 untrained machine parameters via the trained GTS-ML model, fulfilling the precision demanded by the automation of the UED instrument. It is evident that the output of the GTS-ML model could provide the machine parameters with the adequate precision, for the fully automated UED operation. Afterwards, a scan of each individual machine parameter around the model predicted optimal setting can further reduce the MSE in cases that higher precision of the machine parameters is required for more challenging experiments, i.e., those with the minimum energy spread and divergence of the electron beam. The ML model will predict v 3 and v 5 with the minimum σ r and θ gun , shown as the 'σ r -∆E/E' (v 3 -p 5 ) and 'θ gun -∆E/E' (v 5 -p 5 ) plots in Fig. 3a, for the minimum energy spread, and particularly set v 6 with an optimal value of 0.18 for the minimum x and y divergences, shown as '(I/I max -std(x′)' (v 6 -p 3 ) and '(I/I max -std(y′)' (v 6 -p 4 ) plots. Fine tuning via scanning σ r , θ gun , I / I max , and other parameters in-situ can further optimize the energy spread and divergence.

Improving the STD-Ml model
The STD-ML model has been successfully demonstrated in our early studies 30 . It can map every BD pattern to the corresponding underlying electron beam properties (x and y spatial positions, x and y angular divergence, energy, and energy spread). The BD patterns are labeled with the electron beam properties used to generate such patterns and are the input parameters to the ML model. To meet the continuous output variable requirement, a LeNet-style neural network architecture with a linear output activation function is used to build the ML model 35 .
We demonstrate that all electron properties except energy and energy spread have the generalization errors on the order of 0.01 with respect to the normalized range of [0, 1] for the training variables. The prediction only takes less than 0.5 µs in batch mode 30 . Based on our understanding of the connections between the electron beam properties and the features of a BD pattern 2,3 , we further improve the predictability of the most challenging electron beam parameters, energy spread and divergences, resulting in better precision and higher speed. This improvement is mainly driven by the need for real-time high-precision diagnostics. For a BD pattern, the transverse and radial peak broadenings are mainly caused by the x and y divergences and the energy spread of the electron beam, respectively, while peak positions are intrinsically correlated to the spatial positions of the beam. To improve the STD-ML model, we implement a two-step approach, which applies the brightest peak analysis (BPA) and polar coordinate transformation (PCT) in sequence. The two steps target on predicting the divergence and the energy spread, respectively. As a result, the new procedure for implementing the STD-ML model is: 1) apply BPA to predict the x and y positions and divergences; 2) center the BD pattern to the image by removing the x and y position jitters; 3) perform PCT; 4) train an ML model with the PCT BD patterns, which only predicts the energy spread and energy. Application of the new procedure results in a five times more accurate prediction of the energy spread.
PCT for a better and faster energy spread prediction. To take advantage of the fact that the electron beam energy spread only causes peak broadening in the radial direction, an early version of PCT (named PCT-v1) was applied to resample a BD pattern, extract the pixels in the ring, and cast into a rectangular image with a more than four times size reduction 30 . However, this approach suffers a significant reduction in the prediction accuracy. To overcome such a limitation, a modified PCT (named PCT-v2), one that is relative to the center of a BD pattern instead of the center of an image, has been implemented. Compared to PCT-v1, PCT-v2 reduces the MSE of the training loss by more than a factor of five.
The detailed procedure of PCT-v2 is as follows: (1) obtain x and y position jitters for each BD pattern via the method described in section III-b; (2) remove x and y position jitters from each BD pattern; (3) apply PCT-v2 to those BD patterns; (4) train an ML model with the PCT-BD patterns as the input, which only predicts the energy spread and the energy.
A set of ten thousand images with the size of 256 · 256 pixels are generated based on the initial population set by the Latin hypercube sampling 33 . The labels of those images are normalized to the range of zero to one for the purpose of standardizing the training process. One example of the input to the ML model is shown in Fig. 6a left panel. For the same dataset, the energy spread is chosen as the only training output variable (easily expandable to include both energy and energy spread). We applied PCT in four different cases, corresponding to the ones with both x and y jitters (case I), with x jitter only (case II), with y jitter only (case III), and with both x www.nature.com/scientificreports/ and y jitters removed (case IV). The difference between the top and bottom graphs in the right column of Fig. 6a is that the top graph corresponds to case I (after application of PCT-v1 and with all four peaks, labeled 1, 2, 3, and 4, misaligned in the radial direction (R) due to the x and y jitters) and the bottom graph is for case IV (after application of PCT-v2 and with these four peaks are perfectly aligned, thanks to the removal of the x and y jitters). The training and validation sets are used to fit the model and to prevent the model from overfitting, while the test set is used to check the accuracy of the trained model. For the 200 untrained images from the test set, the MSEs of case-I, II, III, and IV are plotted in Fig. 6b. The corresponding MSE values are 0.0492, 0.0183, 0.0237, and 0.0092, respectively. Training losses of the energy spread for the four cases are plotted in Fig. 6c as the light blue, light orange, dark orange, and dark blue curves, respectively. Thanks to the removal of x and y jitters, the generalization error in case IV is more than 5 times smaller than in case I. As expected, cases II and III have similar errors, which are between the errors of cases I and IV. Thus, we demonstrated that PCT-v2 has a more than a factor of five reduction in the MSE due to the removal of x and y jitters. By repositioning each BD pattern to the center of the corresponding image, we expect that the fidelity of the energy spread prediction via PCT-v2 could be immune from the impact of x and y jitters.
BPA for improving the divergence prediction. For the prediction of the divergence and the spatial position of an electron beam, we select a rectangular section in every BD pattern, which covers the brightest Bragg peak with the minimum area. The criterion for determining the size of the selected area for BPA is that the section must cover ±5 · σ x,y (Bragg peak widths in x and y directions), even when the electron beam has the maximum jitters in both x and y directions. This section is only used for the analysis of x and y divergences and spatial positions. The characteristic of the divergence is that it broadens each BD peak similarly in a BD pattern in the same direction as the divergence. In contrast, the feature of the spatial position jitter is to transversely shift the image on the detector with the same amount of the spatial position jitter. Therefore, the selected section with the brightest BD peak could be used to train an ML model, which only predicts the x and y divergences and spatial positions of the electron beam. The reduction of the image size and the selection of the peak with the highest intensity could significantly improve the speed as well as guarantee high fidelity of the prediction. Also, it is important that the availability of x and y spatial positions allows the removal of x and y position jitters from each BD pattern, which is required as the first step of PCT-v2.
We use the same set of ten thousand images. One example of the input to the ML model is shown in the left panel of Fig. 7a. The x and y divergences and spatial positions are chosen as the four training variables. We performed two types of trainings: (1) case A (standard)-for each image, the entire BD pattern ( 256 · 256 pixels) is used for the training; (2) case B (BPA) -for every image, only the area covering the brightness BD peak ( 64 · 64 www.nature.com/scientificreports/ pixels) is selected for the training, highlighted as the yellow dotted square in the left graph and also shown as the small graph in Fig. 7a top right. The average training losses of all four variables (x and y divergences and positions) are plotted in Fig. 7b as the red and green curves for case A (standard) and case B (BPA), respectively. To make a comparison between cases A and B, the MSEs regarding the test dataset for x and y jitters and divergences, including before (columns with the label "MSE") and after (columns with the label of "unit converted MSE") the unit conversion, and the time taken for the training are listed in Table 3. As expected, case B performs better than case A, with a 15% reduction in MSEs and two times faster in speed, since the selected highest intensity peak with the maximum signal-to-noise ratio contains all information required to predict the x and y divergences and spatial positions.

Method
The start-to-end simulations are performed through particle tracking via GPT from the gun to the sample and the EDP simulation of a wave-like electron being diffracted by the sample. A two-stage approach is used to build the ML models, GTS-ML from the gun to the sample and STD-ML from the sample to the detector. The BD patterns as the input and the electron beam properties as the data label are applied to train the STD-ML model. To automate the UED instrument, the machine parameters and the electron beam properties as the data label and input data, respectively, are applied to build the GTS-ML model. To improve the STD-ML model, we decouple the unique features of the beam energy spread and divergence via PCT-v2 and BPA, which preserve the radial symmetry and x-y translation invariant, respectively. The BPA method is applied to build an ML model that predicts the x and y divergences as well as spatial position jitters.
Removing the x and y position jitters makes each BD pattern to be centered in the image, hence, to be ready for PCT-v2. Then, those PCT BD patterns are used to train an ML model, which only predicts the energy spread and energy. As a result, the precisions of the energy spread and divergence predictions are improved significantly.
In practice, a training set consists of the measurements of BD patterns from a sample under different machine parameters. The previously trained STD-ML model is used to determine the beam properties at the sample, which in turn, is used to train the GTS-ML model. To use the ML model in an experiment, the users only need to provide the desired beam properties at the sample; with the real-time BD pattern recorded by the detector, the STD-ML model outputs the beam properties at the sample, and, by comparing the model predicted and user targeted beam properties, the GTS-ML model predicts the required machine parameters.

Discussion
The concept demonstrated in simulation in this study is directly applicable to automate the UED instrument in experiments. The proposed procedures are: set the machine parameters of the accelerator based on the table (10,000·6). Each set of the 6 machine parameters determines a BD pattern on the detector. We record all BD patterns. Using those BD patterns as input, we apply the pretrained STD-ML model to predict the electron beam properties at the sample, which are used as the input to the GTS-ML model as well as to provide real time diagnostic information. Subsequently, we use the machine parameters and the electron beam properties as the data label and the input data of the training dataset, respectively, to build the GTS-ML model.
A proof-of-principle simulation study has been done via the two-stage ML model. It is evident that the precision and speed of the prediction are adequate for fully automated UED operation. In the experimental case, the entire process, including setting up the UED instrument, recording BD patterns, loading the dataset to the computer, training the ML model, and updating the model in the operating system of the UED instrument, will take just a few hours (e.g., 4 h in collecting data and 10 min in training for the 10 k-data case on an RTX2080 GPU). The estimation assumes that the ML model has been established prior to the experiment and the time is only spent in updating the model. Since the GTS-ML model is sample independent, the model needs to be updated less often (e.g., once in a few weeks); it is performed only when there is a change of the machine calibration, or the commissioning of a new operational mode. Occasional retraining is needed when the beam falls out of the range on which the GTS-ML model was trained. Similarly, STD-ML needs to be updated only when there is a sample change or detector replacement.
The study demonstrates that the two-stage model enables the full automation of the UED instrument. The BPA method reduces the image size by more than a factor of 10. These images are used to train an ML model, which predicts the beam divergences and spatial position jitters with a guaranteed high fidelity. Besides, the predicted x and y position jitters are needed to center those BD patterns before the application of PCT-v2. Then, the PCT Table 3. MSEs for the electron beam properties-x and y jitters and divergences (p 1 , p 2 , p 3 , p 4 ), including before and after the unit conversion, are listed for two cases A (standard) and B (BPA). The times taken for the training process are 30 min and 14 min for case A and B, respectively. www.nature.com/scientificreports/ BD patterns are used to train an ML model, which only predicts the energy spread and energy and hence has a significantly improved precision (> 5). Furthermore, compared to the proposed method in our earlier study 1 , the global optimization enabled by the ML via Latin hypercube sampling scheme 33 provides a full coverage of the parameter space of a UED instrument, and as a result, allows more electrons to be packed (with 150 times brighter) in the most important dimensions, i.e., divergences and energy spread, as shown in Table 4. We provide a roadmap toward the full automation of the UED operation and real-time diagnosis of electron beam properties with a comprehensive set of beam parameters. The two-stage approach can be easily extended to other types of accelerator facilities. Eliminating the setup time for a UED experiment and providing the ultrahigh bright electron beam to the UED/UEM with a single-shot imaging capability can be a game changer in the ultrafast science and technology.

Data availability
The datasets generated and analyzed during the current study are not publicly available due to the reason that we want to know who has an interest in our datasets but are available from the corresponding author on reasonable request.
Received: 11 January 2022; Accepted: 25 February 2022 Table 4. The operational and ML optimized electron beam properties at the standard charge of 0.5 pC are shown as the left and right columns, respectively. We only consider the x and y divergences and normalize energy spread of the electron beam, which ultimately limit the UED/UEM instrumental resolution.