research2020–2022first authorbuilt the CNNs

UED beam diagnostics with CNNs

A small CNN that reads each electron shot's pointing and energy jitter straight off its diffraction pattern.

No two shots at the MeV electron diffraction setup at Brookhaven's ATF are quite the same. It shoots short electron bunches through a thin crystal and records where they scatter. Each bunch points a little differently and carries a little more or less energy, and over many shots that blurs the pattern you came for. Xi Yang at Brookhaven could read the beam off the pattern by fitting Gaussians to the spots, at about 0.4 s a frame, too slow once the gun fires more than a few times a second. Xi's idea was to train a network on frames the fit had already labeled and let it read every new shot. I built and trained the CNNs, and I'm first author on the two Scientific Reports papers that came out of it (2021 and 2022).

One shot

Electrons that scatter off the crystal land as bright spots on a phosphor screen. Where the spots sit depends on where the beam points and on its energy. How wide they are depends on its divergence and energy spread. So in principle one frame holds all six numbers.

Two false-color detector images. Left, a full frame about 512 pixels square: scattered bright diffraction spots on a dark blue background, a dotted box around the central region and a small dotted circle in the middle. Right, the 256 by 256 crop from inside the box: a ring of about a dozen bright spots around a large dark disk.
One recorded shot, the whole frame on the left (the dotted circle is the hole in the mirror) and the 256×256 crop the network gets on the right, with the middle masked out.

A small CNN

We started in simulation. Xi tracked the beam from the gun to the sample with GPT, and Xi and Lijun Wu turned each simulated beam into a diffraction pattern with Lijun's code. That gave 10,000 patterns of a strontium titanate crystal, each labeled with the six numbers that made it.

I went with a LeNet-style network, about as plain as a CNN gets. Five rounds of convolution and pooling shrink the frame to a stack of 6×6 maps, then one dense layer feeds a linear output per number.

Schematic of the convolutional network: a 1 by 256 by 256 input, five pairs of convolution and 2 by 2 max-pooling layers with 2, 4, 8, 16 and 32 feature maps shrinking to 32 maps of 6 by 6, then a flatten to 1152 units, a fully connected layer of 120 units and 6 outputs.
The whole network, from one frame down to six numbers.

Energy spread

The first network learned all six at once. Pointing and divergence came out within 2% of their range on patterns it had never seen. Energy was only so-so. Energy spread was close to a guess.

Two panels over 200 simulated test patterns. Top, vertical pointing: the orange predicted crosses sit on the blue label line, RMS error 0.0177, with a narrow error histogram. Bottom, energy spread: the crosses scatter widely around the label line, RMS error 0.2093, with a broad error histogram.
Two of the six from that first run, vertical pointing on top following its labels and energy spread below all over the place.

The network wasn't really at fault. A spot's width, measured against how far out the spot sits, comes from both the divergence and the energy spread.

relative spot width=(Δθθ)2+(ΔEE)2

Δθ is the beam's divergence, θ the Bragg angle of that spot, and ΔE the spread in its energy E.

In this set the divergence was far bigger. In the example we worked through in the paper, the two terms split 95% to 5% before they go under the root. The spread is in the picture, just buried.

So we made a second set where only the energy spread changed and the divergence was squeezed to a few microradians. Now each spot stretches along the radius, and the same network read the spread to about 0.2% of its range. The real beamline can't squeeze the divergence that far. We proposed condenser lenses and an aperture for it and never tried them.

Three simulated diffraction patterns side by side, labeled energy spread 9.4e-05, 2.5e-02 and 4.4e-02. Each has four bright spots in the corners, circled by dotted ellipses, around a central dark disk; from left to right the spots grow longer, each along the line pointing out from the center.
Three simulated patterns with the divergence held small and more energy spread from left to right, where the spots only get longer along the radius.

Shots from 2020

Then the real thing, a few thousand shots the BNL group recorded in 2020 on a MoTe2 sample. The detector only catches the first ring of spots, and a Gaussian fit places a spot far better than it measures its width. So only the pointing in x and y and the energy jitter had labels (Xi's fits, shot by shot). Divergence and energy spread had nothing to learn from, so I cut the output layer to three and trained again.

On shots it hadn't seen, the readings sit right on the fitted labels. Vertical pointing came out about as well as horizontal, and the energy jitter to about 3 × 10−5 RMS.

Three panels over 200 real test shots, for x pointing, y pointing and energy jitter. In each, orange predicted crosses sit on the blue label line, with RMS errors 0.0033, 0.0034 and 0.0043 in the titles, and a narrow error histogram on the right.
Two hundred real shots the network never saw, fitted labels as the blue line and its readings as orange crosses.

How much data does it need? I retrained on smaller and bigger sets, and the error flattened out after three or four thousand shots.

The ring, unrolled

Training still took long enough that I wanted it shorter. Nearly everything useful sits in the ring of spots, about 38% of the pixels. So I unrolled the ring into a strip, radius one way and angle the other, and trained a smaller network on that. The input shrank by 80%, and training took less than half as long.

Left, a real frame with the ring of spots marked by dashed circles and an arrow showing the polar axes, and below it the ring unrolled into a long thin strip with the spots in a row. Right, two loss plots labeled origin and polar, each with training and validation MSE falling over about 80 thousand batches; the polar validation curve levels off noticeably higher.
The ring from a real shot unrolled into a strip, with training and validation loss for the full frame and for the strip (the strip's validation loss settles higher).

It cost accuracy, and the errors roughly doubled. I had unrolled around the center of the image, but pointing jitter shifts the whole pattern, so the strip scrambles where the spots sit relative to each other. Pointing and energy are read from exactly that layout, and a network as shallow as mine couldn't learn it back.

Back to the knobs

Xi's next question was whether a second network could go one step further back, from the beam at the sample to the machine settings that made it. The 2022 paper chains the two. The first, reading the detector, is my CNN. The second maps the six beam numbers to six settings, from where the laser lands on the cathode to the gun phase and the solenoid current. It trained on 10,000 of Xi's GPT runs.

Two network diagrams with the beam going right to left. Left, one network runs from the gun input through hidden layers to the diffraction-pattern measurement. Right, the same span split into two chained networks, GTS-ML from the gun to the electron properties at the sample and STD-ML from there to the diffraction pattern.
One network from the detector all the way back to the gun (left) against two in a chain with the beam at the sample in between (right), with the beam traveling right to left.

A one-layer network did worst on the RF amplitude, probably because space charge makes that part of the map strongly nonlinear. I built a four-layer network instead, and it was better on all six. Two settings still missed the precision the machine would need. The gun-phase error was 3.7° against a target of 3°, and the RF amplitude fell short too.

Offline

Nothing here ran on the machine. The 2021 network was trained and tested on shots recorded in 2020, so "real time" in that paper is what the timing says it could do. The 2022 paper is simulation from start to end, and its plan for the real machine (four hours of data, ten minutes of training) is an estimate. Nor was a virtual diagnostic a new idea. Sanchez-Gonzalez and others did it for X-ray pulses at an FEL in 2017, and our first title even echoes theirs.

Numbers and figures are from Zhang et al., Sci. Rep. 11, 13890 (2021) and Sci. Rep. 12, 4240 (2022), both open access under CC BY 4.0; the figures are cropped.