femto-PIXAR: a self-supervised neural network method for reconstructing femtosecond X-ray free electron laser pulses

Fluctuations in microscopic systems are directly related to their fundamental excitations – a concept dating back to Einstein’s description of the Brownian motion of pollen particles. This relationship has motivated the need for measuring fluctuations within microscopic systems to directly compare with first-principles theoretical models. However, the spatial and temporal scales of quantum systems far exceed those accessible by classical microscopy, and despite the advances in modern light sources, the energy scales pertinent to these fluctuations remain challenging to probe. One notable exception to this has been in inelastic neutron scattering, where small energy changes of scattered neutrons can be observed, and has led to understanding of dispersion in a myriad of areas, from frustrated magnets and high temperature superconductors, to biological tissue and exotic topological structures, such as skyrmions. However, the reliance of this on large crystals and its limitations with strongly absorbing elements restrict its applicability. The advent of the free electron laser has changed all of this. Its capability to generate femtosecond (fs) laser pulses with extremely high pulse energies at X-ray wavelengths has ushered in a new era of science, with new developments rapidly evolving in many fields [1, 2] including structural biology [3, 4], enzyme catalysis [5], and astrophysics [6, 7]. Furthermore, the development of coherent, two-pulse configurations brings new potential for fluctuation studies within quantum materials at their relevant timescales, i.e. at the meV energy scale and below [8, 9, 10, 11, 12, 13].

Importantly, with X-FELs it is now possible to deliver coherent X-ray pulses at sub-picosecond spacing for experiments, resolving snapshots in structures with sub-Angstrom resolution, and comparing these scattering events to measure fluctuations, but critically, no current method exists to temporally resolve these individual incident pulses at such fine timescales. This highlights the urgent need for reliable characterization of X-FEL laser pulse intensities with femtosecond resolution. Such methods must be robust across various lasing configurations and experimental setups. Notably, X-FELs are highly sensitive to operational parameters which govern the durations, energies, and sequences of X-ray pulses [14, 15]. Compounding this complexity, X-FELs exhibit intrinsic intensity fluctuations, causing shot-to-shot variations in the pulse energy spectrum and photon density [16, 17]. These beam intensity fluctuations from the X-FEL present significant challenges for experiments where accurate characterization of pulse properties, including pulse duration and energy density, is essential to extract meaningful results. These challenges have so far hindered use of this new capability to deliver new science in the observation of fluctuations for understanding properties in materials.

We address this need by developing a novel self-supervised machine learning (ML) framework called femto-PIXAR: femtosecond Power Inference of X-ray pulses using AI References, which provides sufficient sensitivity to resolve microjoule-scale pulses with femtosecond pulse separation. We show our approach accurately predicts pulse energies for an experimental dataset, verified against total pulse energy monitors in real X-FEL data, and confirm the predictions match the expected pulse delays. This demonstration will make fluctuation studies possible in the sub-picosecond regime, and will also provide a tool for a wide array of X-FEL experiments in many other fields as well.

The experimental setup was designed to produce two coherent X-ray pulses starting from a single, long electron bunch, itself generated from two overlapping injector lasers. The two X-ray pulses were created using a double-slotted V-shaped foil (for time domain control) [18] [19] and hard X-ray self-seeding [20] (for coherence and energy selection). To our knowledge this configuration, described in Fig. 1 a), has never been used before. The experiment was conducted at a photon energy of 8.35 keV (Ni $K$-edge), with slotted-foil pulse separations of $0\,\mathrm{fs}$ (single pulse), $\approx 20\,\mathrm{fs}$, and $\approx 30\,\mathrm{fs}$. Different stages of the electron bunch shaping sequence are labeled $e^{-}_{\text{I}}$, $e^{-}_{\text{II}}$, $e^{-}_{\text{III}}$, $e^{-}_{\text{IV}}$, representing the same electron distribution at different instances as it proceeds in time. A 3 $\mu$m thick aluminum foil with two slots of variable separation from 1 mm to 5 mm was placed in the middle of the dispersive section of the bunch compressor BC2 [18]. At the locations where the electron bunch $e^{-}_{\text{I}}$ passes through the foil the electron emittance was ‘spoiled’ (i.e. increased) due to coulomb scattering. The emittance growth suppresses FEL amplification in the undulator, so that only the electrons $e^{-}_{\text{II}}$ that pass through the slits of the V-shaped foil produce X-ray pulses.

Downstream of the bunch compressor the relative distance between the unspoiled electron distributions $e^{-}_{\text{III}}$ is in the direction of travel, resulting in an arrival time difference $\Delta t$, that can be changed by shifting the slotted foil. $e^{-}_{\text{IV}}$ passes through the undulator magnets resulting in X-ray lasing. Between the first and second undulator segment a hard X-ray self-seeding scheme is applied (not shown here). Downstream of the undulators, the X-band Transverse deflecting mode CAVity (XTCAV) streaks the electron beam on a target to image the final longitudinal phase space [21]. The X-ray laser pulses are unaffected by the XTCAV deflector and continue downstream to the X-ray gas energy monitor (GEM), where the total X-ray pulse energy is transmissively monitored on a shot-to-shot basis using a photomultiplier tube (PMT) to measure the fluorescence induced in a small volume of N₂ gas [22]. After the GEM, the X-rays continue down the beamline to the experimental endstations. This entire process repeats at a 120 Hz repetition rate.

The XTCAV images the longitudinal (time-energy) distribution of the post-lasing electron bunch. By comparison with a lasing-off reference, it is possible to infer the temporal profile of the photon pulse. However, because the lasing-on and lasing-off references cannot be measured simultaneously, this poses a fundamental challenge in determining the change in the electron phase space due to the undulators. The “classical” approach has been to record a set of images while suppressing lasing, cluster these images by shape, and then select the most similar cluster for each lasing-on image by comparing the projected electron current [21, 23, 24]. While the classical approach has been demonstrated for various two-pulse modes [19, 25, 14], it struggles for small pulse energies, when shot-to-shot differences between the true lasing-off electron bunch and the selected reference cluster exceed the lasing-induced changes. As a result, when combining seeding with two-pulse modes, the classical method breaks down.

A deep learning method proposed by Ren et al. [26] sidesteps the need for matched lasing-off references by training a neural network (NN) directly to predict power from a lasing-on image; however, the NN in  [26] relies on simulated labels for training, and it is particularly challenging to simulate the complex electron beam distribution produced by the slotted foil. On the other hand, for the slotted foil setup, only a narrow segment of the beam lases. We can exploit this feature to approximate a reference in the short lasing region by using information in the nearby non-lasing regions. For example, Zeng et al. [27] suggested to use a polynomial regression across the lasing region as a reference. However, the polynomials cannot model the complex phase space of this setup and fail to reproduce the weak lasing pulses.

Here, we instead propose an approach using a U-net [28] to generate the corresponding lasing-off reference image. We trained the NN using self-supervision, eliminating the need for simulations. Specifically, we masked regions of lasing-off pulses, and trained the NN to regenerate those regions. Once trained, the network receives a lasing-on image with masked lasing regions and reconstructs the best lasing-off reference. We then calculated the center of mass (COM) to retrieve the X-ray pulse distribution [24]. The network can also detect the lasing region by identifying areas with the largest reconstruction discrepancy. Our method provides high-fidelity lasing-off references without relying on the accuracy of computationally-expensive simulations while retaining interpretability through direct observation of the generated reference.

Fig. 1 b) and c) shows the core steps. The inputs to the network were cropped, background subtracted and normalized lasing off images, where all values within two random mask-regions were set to zero (black bars in Fig. 1 b) and c)). The width of an individual mask was chosen to be 16 pixels and the distance range for the two masks was set from 0 to 40 pixels, corresponding to approx. 0 to 38 fs. We then trained a U-net to reconstruct the electron phase space of the lasing-off data (Fig. 1 b) ). As the center of mass is the critical value for evaluation, we used a loss that combines the Mean Squared Error of pixel values with the absolute error of the COM. After training, we masked the lasing regions in a lasing-on data set. As the network is trained only on the lasing-off data, it reconstructs the best matching lasing-off electron phase space exactly for a given image (Fig. 1 c). This lasing-off image will also exactly match the original $x$ and $y$ position, so no shifting in the time axis or normalization with the GEM (as in the classical approach) is necessary.

A crucial step of this process was to determine the correct mask positions in the lasing data. The separation of the two pulses was approximately fixed by the slit separation, but due to shot-to-shot differences in electron bunch energy and compression of the electron bunch, the position of the features relative to the center of mass of the electron bunch was not constant. Our method scans the image to find two locations where input and reconstruction differ most, centering the masks on the maximum signal. We then use the methods from Behrens et. al [24], marked as $f(a,b)$ in Fig. 1 to obtain the power profile using the spectral COM from both images.

Fig. 2 shows example reconstructions for the three different configurations and one lasing-off example, all randomly selected from the top 1$\%$ of data with the highest GEM values (ensuring a clearly visible signal). None of these images were part of the network training process. To demonstrate the efficacy of the proposed method, we benchmark the predicted power and delay against existing diagnostics. First, we compare the integral of the predicted profile with the total energy of both pulses measured by the GEM, which cannot resolve pulses on the fs scale. In Fig. 3 one can see the correlation of the integral of the power profile with the GEM for the three different slotted foil configurations. We note that while there is a strong linear correlation between the predicted power and GEM monitor, there is a disagreement in proportionality. Further experiments are necessary to determine if it is due to either the GEM or XTCAV calibration, or bias from the method. However, as the typical user experiment usually relies primarily on the relative power between the pulse pairs, this discrepancy does not affect the experimental analysis.

We compare our predictions to two existing state-of-the-art XTCAV analysis methods. In order to evaluate the correlation with the GEM, in the classical approach (see the previous section), the energy jitter (shift in the y-axis) cannot be compensated for by normalization with the GEM but must instead be corrected using a scheme that matches the head and tail of the COM profiles for the lasing-on and lasing-off images. We also compare our results with the polynomial regression method of Zeng et al. [27]. Some example predictions of this method can be found in the supplemental material. Fig. 3 compares the GEM correlation for all three delays, neither of which results in the expected correlation.

Fig. 4 shows the temporal and energy characteristics of the X-ray laser pulses as retrieved by the U-net reconstruction scheme. The time assignments of each shot were obtained by fitting the X-ray power trace with two Gaussian functions. For these plots we took only the data where both Gaussian curve fits indicated a signal larger than 15 ${\mu}$J. We find that the predicted time difference between both pulses closely matches the delay given by the slotted foil condition. The temporal separation jitter aligns well with the 9$\%$ jitter observed by Ding et al. [19]. As an additional benchmark we checked the predicted pulse energies for lasing-off data with randomly placed masks and find that 80$\%$ of the data has a deviation from zero of less than 15 $\mu$J, and 99.8$\%$ less than 25 $\mu$J.

In conclusion, femto-PIXAR is a novel X-FEL diagnostic which we demonstrate by reconstructing the X-ray power profiles from weak X-FEL pulses in a non-standard two-pulse, self-seeding configuration. For this configuration, the classical method of XTCAV analysis, which uses a matching algorithm to subtract similar lasing-off shots for the energy loss calculation, fails to reproduce meaningful correlations with the X-ray gas detector. By contrast, our method is able to distinguish the relative amplitude of each X-ray pulse on a shot-to-shot basis which is necessary for X-ray photon fluctuation spectroscopy experiments and other future experiments at X-FEL facilities. Our approach has various benefits, avoiding the need for energy calibration with an external reference. A neural network based scheme also scales well for the large datasets expected from high-throughput experiments enabled by next generation X-FEL facilities such as the LCLS-II [29]. As opposed to previous deep-learning work, our approach uses self-supervision with a physics-motivated loss function, avoiding the need for high-fidelity simulations. Our method also reconstructs the full 2D phase space, providing interpretability of the network and giving crucial information about when the network might not be applicable for a specific dataset. Future work will consider extensions to more general X-FEL configurations beyond the slotted foil. For example, using diagnostics beyond the XTCAV could enable analysis of configurations with lasing across the full electron bunch [30], e.g. to provide the femtosecond resolution needed for temporal ghost imaging [31]. The analysis chain after reconstructing the lasing-off reference can also make use of machine learning, e.g. incorporating spectral measurements to refine the X-ray pulse profile [32].

Use of the Linac Coherent Light Source (LCLS), SLAC National Accelerator Laboratory, is supported by the US Department of Energy, Office of Science, Office of Basic Energy Sciences under Contract No. DE-AC02-76SF00515. The authors acknowledge support from DESY (Hamburg, Germany). This work has in part been funded by the IVF project InternLabs-0011 (HIR3X). R. Plumley acknowledges support from the US Department of Energy, Office of Science, Basic Energy Sciences, for the Materials Sciences and Engineering Division under contract DE-AC02-76SF00515. J.J.Turner acknowledges support from the U.S. Department of Energy (DOE), Office of Science, Basic Energy Sciences under Award No. DE-SC0022216.

The centered XTCAV images have a shape of 240 x 240 pixel. A median dark image is subtracted and they are normalized to have intensities roughly between 0 and 1 before they are used in the network. Before obtaining the power profiles a median filter is used to identify the signal region and set the other regions to zero. This approach is just used to identify the signal region, no median operation is applied to the image.

We used an U-net architecture as described in [28]. The optimizer is RMSprop of pytorch with a weight decay of 1e-8. Out of some manual tries, we picked a learning rate of 5e-7 and a batch size of 2. We found that we get good reconstructions with this parameter set. A more extensive hyper parameter search might provide slightly better reconstructions or faster conversion of the network. However we found that the data preparation like scaling and centering is way more important than the exact hyper parameters. For the loss we used an MSE loss on the reconstruction, combined with an absolute error on the COM. The COM has a relative low weight of 1e-4 as we found that higher weights significantly worsen the reconstruction.
In all our lasing-on datasets we exclude images with weak XGMD signal ($<5\,\mu$J). With this threshold we discarded about 15% of our data, resulting in about $4\times 10^{5}$ samples with approx. $20$ fs difference, about $1.5\times 10^{6}$ samples with approx. $30$ fs difference and $9\times 10^{4}$ with only one pulse. Weak shots (where essentially no lasing occurs) with $E_{\text{GMD}}<5\,\mu$J are included in the lasing-off dataset to get more training data. This results in about 30K lasing-off samples.

For the classical analysis the signal region is identified as described in Appendix A, and then divided into 120 slices so that all lasing-on and lasing-off data has the same length in slices and lasing-on and lasing-off data can be compared directly. HierarchicalClustering (sklearn.cluster.AgglomerativeClustering) with euclidean distance is used to cluster the current profiles into 500 groups. Normally, a few hundred lasing-offs are used to build a reference set, but to have a fairer comparison all lasing-off images that are used in the training of the network are provided to the algorithm. The COM profile of these groups are averaged and used as reference.

For the polynomial regression we used an algorithm described in [33]. Zeng et al. [27] uses a 13 dimensional polynomial, we find that our algorithm does not converge in this case. In Fig. A1 we show one example of the regression method with increasing polynomial degrees. In our benchmark we use the highest degree where the algorithm converges, with a maximum of 13. In Fig. A2 we show resulting regressions for ten different samples. We found that the predictions have a better correlation with the XGMD if we ignore the (obvious wrong) parts of the predicted profile that are negative.

An additional method to benchmark our method is to evaluate the generated power profiles for data where no lasing process took place. Ideally, for this setup, the integral of the power profiles would be zero. Fig. A3 shows a histogram of the resulting energies of 4600 non-lasing data points. We used 460 lasing-off samples that were not part of the network training and used each of them 10 times with randomly set masks. 80$\%$ of the data has a deviation from zero of less than 15 $\mu$J, 99.8 $\%$ less than 25 $\mu$J. As we are interested in the pulse ratio between the two peaks it is important that not one of the pulses is systematically under and the other over estimated. We check this on the right side of Fig. A3 and see that the pulse energy difference between the two pulses is centered well around zero. This indicates, that, for this dataset, not one pulse is overestimated over the other.

Using a double-Gaussian fit over the reconstructed transient X-ray power profiles, we are able to determine the arrival times of the two pulses at the XTCAV target. The time axis in A4 is centered with respect to the COM of the spectrograph in the temporal domain. The reconstructed pulses obtained by the Gaussian fits produce bimodal time distributions that match the X-FEL setups, with average time jitter of 3.3 fs.