A general Bayesian algorithm for the autonomous alignment of beamlines

Abstract

Autonomous methods to align beamlines can decrease the amount of time spent on diagnostics, and also uncover better global optima leading to better beam quality. The alignment of these beamlines is a high-dimensional expensive-to-sample optimization problem involving the simultaneous treatment of many optical elements with correlated and nonlinear dynamics. Bayesian optimization is a strategy of efficient global optimization that has proved successful in similar regimes in a wide variety of beamline alignment applications, though it has typically been implemented for particular beamlines and optimization tasks. In this paper, we present a basic formulation of Bayesian inference and Gaussian process models as they relate to multi-objective Bayesian optimization, as well as the practical challenges presented by beamline alignment. We show that the same general implementation of Bayesian optimization with special consideration for beamline alignment can quickly learn the dynamics of particular beamlines in an online fashion through hyperparameter fitting with no prior information. We present the implementation of a concise software framework for beamline alignment and test it on four different optimization problems for experiments on X-ray beamlines at the National Synchrotron Light Source II and the Advanced Light Source, and an electron beam at the Accelerator Test Facility, along with benchmarking on a simulated digital twin. We discuss new applications of the framework, and the potential for a unified approach to beamline alignment at synchrotron facilities.

Keywords: Bayesian optimization; automated alignment; digital twins; machine learning; synchrotron radiation.

Figure 1

The prior distribution, noiseless posterior distribution and noisy posterior distributions for a GP with covariance 〈f(x_i)f(x_j)〉 = M_5/2(|x_i − x_j|/4), where M_ν(r) is the Matérn function as defined in equation (19). For each distribution, we draw four random functions (colored lines). The black line represents the mean of each distribution, while the dark- and light-shaded regions represent the 1σ and 2σ intervals, respectively.

Figure 2

An example of an iteration of a Bayesian optimization algorithm trying to maximize the negated Himmelblau function f(x₁, x₂) = , whose true global optima are marked as white circles. Using existing data points (far left) and the assumption that the function is distributed as a GP, we can use Bayesian inference to compute a posterior consisting of a mean (center left) and error (center right), upon which we can compute an acquisition function (far right) which informs us of the best points to sample. The black-edged diamonds superimposed on the acquisition function show the best eight points to sample, optimized in parallel and with the optimal routing represented by the red line.

Figure 3

(Upper left) The result of changing the positions of two coupled dimensions of the TES beamline. (Upper right) A quasi-random sample of 16 points from the ground truth. (Lower left) A non-latent GP fitted to the parameter space fitted to the sampled points. (Lower right) A latent GP fitted to the same points, which correctly infers the latent dimensions.

Figure 4

A schematic of the TES (8-BM) beamline at NSLS-II. This representation shows the many optical components that make up modern beamlines, with each optical component having many degrees of freedom that must be optimized in concert in order to carry out experiments effectively.

Figure 5

Four different beam configurations on the NSLS-II TES (8-BM) beamline, where the upper left-hand panel shows the initial beam and the lower right represents the optimal alignment. In this alignment test, we adjust the translation and rotation of each of the horizontal and vertical Kirkpatrick–Baez mirrors and the pitch and vertical translation of a toroidal mirror, for a total of six degrees of freedom to maximize the flux density of the beam.

Figure 6

Four different beam configurations on the NSLS-II ISS (8-ID) beamline during automated alignment, where the upper left-hand panel shows the initial beam and the lower right represents the optimal alignment. In this alignment test, we adjust the translation of a central crystal and the translation and pitch of two ancillary crystals for a total of five degrees of freedom to maximize the flux density of the total beam on the area detector.

Figure 7

A schematic of beamline 5.3.1 at the Advanced Light Source. The beamline has four degrees of freedom (toroidal mirror pitch and bend, and monochromator angle and height) and four constraints (four-jaw slits).

Figure 8

Four different beam configurations on the Advanced Light Source Beamline 5.3.1 during automated alignment. In total, the photon transport has four active degrees of freedom: the focusing mirror pitch and tangential bend, and the channel-cut crystal angle and height. The upper left-hand panel shows the initial manually aligned beam and the lower right the final beam after automated alignment. The upper right and lower left panels show intermediate points collected in the automated alignment process.

Figure 9

Four different electron beam configurations at the Brookhaven National Laboratory ATF at different stages of automated alignment, where the upper left-hand panel shows the starting beam and the lower right the optimal beam. In this alignment test, we tune the current of four quadrupole electromagnets to maximize the objective in equation (23).

Figure 10

The eight-dimensional optimization of the simulated TES beamline, where the degrees of freedom comprise the toroidal and Kirkpatrick–Baez mirrors. The colors show different varieties of Bayesian optimization algorithms both with and without latent inputs and composite outputs, with both the cumulative maximum of all individual runs (thin lines) and the median cumulative maximum (thick line). Each variety starts out with a quasi-random sampling of 32 points (shaded light blue) and then performs a Bayesian optimization loop with the expected improvement acquisition function. The benefit of using both latent inputs and composite outputs is shown, as we can achieve a better optimum more robustly and more quickly.

Figure 11

The four-dimensional optimization of just the K-B mirrors, whose motors are each misaligned by up to 0.05 mm. After an initial quasi-random sample of 16 points (shaded light blue), the agent is able almost instantly to return to the optimal alignment.