Data-driven feedback optimization for particle accelerator application

1.1 Contributions

The contributions of this paper are twofold:

1. Section 2.3 proposes an approximation scheme for the Hessian matrix of the optimization problem that can be evaluated from first-order derivative information of the underlying plant. Using this approximation, it becomes possible to take advantage of the Hessian for optimization in model-free settings, significantly improving the convergence speed of the algorithm.

2. We introduce the beam position control problem in particle accelerators as an application of feedback optimization in Section 4 that is challenging to solve with other methods such as Bayesian optimization.

1.2 Notation

We use I _n to denote the n × n identity matrix. In addition, S + n is the set of symmetric n × n matrices that are positive definite. For vectors z ∈ R n , z M : = z ⊤ M z denotes the weighted 2-norm with weight M ∈ S + n . If the Euclidean norm with M = I _n is used, the subscript is dropped for conciseness. Finally, M₁ ⊗ M₂ is the Kronecker product of the matrices M₁ and M₂, and vec denotes the column-wise vectorization operator, i.e., vec(M) is the vector consisting of the stacked columns of the matrix M.

2.1 Problem statement

Consider a stable physical plant with input u ∈ R p and output y ∈ R n , and a twice continuously differentiable cost function Φ : R p × R n → R mapping u and y into a scalar cost. Our goal is to drive the plant to a steady-state pair (u*, y*) that minimizes Φ.

More precisely, we consider the physical plant to be expressed by an unknown or at least uncertain input-output steady-state map h : R p → R n such that y = h(u) holds for all steady-state pairs (u, y). Furthermore, suppose there are polyhedral constraints Au ≤ b and Cy ≤ d on the admissible inputs and outputs, described by appropriately sized matrices A, C and vectors b, d, respectively, where we assume that the constraints are decoupled for simplicity. The desired steady-state pair (u*, y*) is then obtained from the solution of the optimization problem

Note that the cost Φ is a design parameter meaning that we assume perfect knowledge of the function.

There exist a variety of approaches for obtaining local solutions to optimization problems in the form of (1), e.g., projected gradient descent or sequential quadratic programming [1], while finding global solutions is generally intractable due to non-convexity. These methods can be roughly categorized to belong to one of two groups: On the one hand, there are optimization methods that perform a feedforward offline optimization based on an a priori known model of the plant. However, since no information from the physical plant is taken into account, the obtained input u will typically be suboptimal when actually applied. On the other hand, online methods like the Nelder-Mead algorithm or Bayesian optimization [1] use output measurements to find the optimal input but can only partially include available knowledge of the plant.

2.2 Feedback optimization

In its current formulation, the equality constraint (1b) implies that the output y can be eliminated from the optimization problem by substitution, leading to the new cost function Φ ̃ ( u ) : = Φ ( u , h ( u ) ) in terms of only the input u. As such, measurements from the physical plant are not taken into account.

Instead, feedback can be included in the optimization process by utilizing measurements of the output y in place of knowledge of the steady-state map h [9]. The idea is to take advantage of the chain rule to calculate the Jacobian of the composite function Φ ̃ as

followed by a substitution of h(u) by y. Note, that for computing the derivative of the composite cost Φ ̃ it is therefore only necessary to know the Jacobian ∂ h ∂ u , not the full steady-state map h.

Based on this reformulation [16], implements the control law u_k+1 = u _k + ρκ(u _k , y _k ) with step size ρ > 0 and control update κ : R p × R n → R p , where u _k is the plant input after k steps and y _k is its corresponding measured output. For initialization, a suitable (with respect to the constraints) input u₀ can be chosen arbitrarily. At every step, the control update κ is determined by solving the parametric quadratic program

where ν ∈ R p is the optimization variable that can be thought of as the next step direction, G : R p × R n → S + p is a parameter-dependent weight, and

is the Jacobian (2) of the composite cost evaluated at the measured output y instead of the steady-state map h(u). As such, the only model information required to evaluate (3) is the Jacobian ∂ h ∂ u . An approach without any knowledge of the plant is discussed in Section 3.

The structure of problem (3) is similar to other algorithms for solving nonlinear optimization problems but has the special property that it is partially evaluated at the measured output y. In particular, the weight G takes the role of the approximated Hessian that is common in quasi-Newton schemes [1]. For example, with G(u, y) ≡ I and strictly inside the feasible set, i.e., both constraints inactive, the optimization variable ν corresponds to the negative gradient.

Under suitable assumptions, i.e., feasibility of the problem, regularity of the constraint set, and a sufficiently small step size ρ, [16], Theorem 3] shows that the incremental control law drives the plant into a local minimum of the composite cost function Φ ̃ . However, the result does not discuss the convergence speed of the method, especially with respect to the weight G. For iterative optimization algorithms, it is well known that including second-order information in form of the Hessian (or an approximation thereof) can drastically increase the speed of convergence close to the optimum, e.g., quasi-Newton methods or the Levenberg-Marquardt algorithm [1]. Hence, while convergence can be guaranteed also for G(u, y) = I _p , it is highly desirable to include second-order information when solving (3) if available.

2.3 Evaluating and approximating the Hessian

In the previous section, it was discussed how the chain rule can be employed to partially evaluate the Jacobian of the composite cost Φ ̃ at a measured output y. Importantly, the terms in (2) can be evaluated without knowledge of the full steady-state map h but only its Jacobian. However, no analogous steps have been employed on the Hessian, and the original formulation of the optimization problem with the weight G only depending on the input u does not allow for its evaluation in terms of the output y.

As an augmentation of optimization scheme discussed in Section 2.2, we therefore propose to perform a similar partial evaluation at the measured output y for the Hessian. Taking the Jacobian of ∂ Φ ̃ ∂ u ⊤ , we obtain

Note that ∂ 2 h ∂ u 2 in the last term is a rank 3 tensor with dimension n × p × p since h is a R n -valued function. Since we can again assume perfect knowledge of the cost function Φ, it is sufficient to know the Jacobian ∂ h ∂ u in order to evaluate the first four terms of (4), the knowledge of which is already assumed to solve problem (3). Furthermore, the reformulation allows for substituting the measured output y for the evaluation of the steady-state map h(u). Hence, only the final term of (4) requires information beyond the Jacobian of h to compute, which is in particular its Hessian ∂ 2 h ∂ u 2 .

4.1 Linear beamline tracking

A particle in the accelerator can be described relative to a reference particle by a 6 + 1-dimensional state

where x and y are the transverse position in horizontal and vertical direction, p _x and p _y are the respective normalized momenta relative to the reference momentum p₀, and

are the longitudinal position and energy offset, respectively, with c being the speed of light. Finally, a constant 1 is added as the last component to cover affine effects in the dynamics.

In the context of this paper, three types of accelerator lattice elements are of special interest. We will therefore introduce a linear beamline model in the following, neglecting collective effects, e.g., space charge effects, and fringe fields at the element boundaries. Note, however, that the linear effects dominate at ultra-relativistic particle energies, such that the linear model is reasonable to apply to the particular EuXFEL section of interest, where the beam has already reached GeV-level energies [28]. A fast implementation of particle tracking routines including non-linear and inter-particle effects is provided by Cheetah [29], [30], which will be employed for the simulations below.

4.2 Dump beamline feedback with known sensitivity

Having introduced the relevant beamline elements and their purpose, we arrive at the specifics of the control problem at hand. Along the beamline section of the main electron dump, there are twelve steering magnets, six in the horizontal plane and six in the vertical plane, in alternating fashion. In between the magnets, ten BPMs are placed to determine the transverse position of the beam, each measuring the horizontal and vertical component. Furthermore, quadrupole magnets are placed along the beamline. Our goal is to find a tuple of steering angles (α₀, …, α₁₂) that minimize the deviation of the particles from the center position of the BPMs.

The prevalent disturbance in the problem is the state of the particle beam entering the dump beamline section. Since the main purpose of the beam is generating photons, the particle positions at the entrance to the dump section are not controlled, causing the mean position to potentially drift over time. Furthermore, the alignment of the quadrupoles and BPMs with respect to the reference beam is difficult, presenting an uncertainty in the beamline model. The task of the optimizer is therefore to reject the disturbance and steer the beam to the center position of the BPMs. As such, we select the quadratic cost function

with tuning parameter η > 0, and set η = 10⁻³ in the following. Note that, depending on the incoming beam, it may be impossible to center the beam at all BPMs due to the arrangement of the position monitors and magnets. In addition, the optimization problem is constrained to keep the beam within ±2 cm of the reference at the BPMs to prevent hitting the surrounding beam pipe.

For the first demonstration, we study the effect of including the Hessian in the optimization problem (3) as discussed in Section 2.3. However, since the curvature term ∂ 2 h ∂ u 2 can be difficult to obtain, we propose to neglect the final component of (4) while calculating the Hessian and therefore set

for the parameter-varying weight in the optimization problem. Note that ∂ 2 h ∂ u 2 = 0 in case of the dump beamline position feedback problem since the dipole response matrices T_sh and T_sv are linear in the steering angle α, such that no error is incurred by this approximation for the current scenario. Furthermore, G(u, y) is guaranteed to be positive definite because the chosen Φ does not contain cross terms between u and y and satisfies η > 0.

In order to evaluate the difference in convergence speed between the two setups, we run 30 scenarios with randomly positioned incoming beams and quadrupole misalignment using the beam dynamics simulation in Cheetah and the quadratic program solver qpth [32]. Furthermore, we take advantage of the differentiability of Cheetah to get access to the sensitivity ∂ h ∂ u with the reference beam and perfect quadrupole alignment, not the actual incoming beam and misalignment. Both the beam center and the misalignment are sampled from zero mean normal distributions with a standard deviation of 1 mm for the beam and 250 μm for the magnets. All 30 scenarios are run twice, once with only first-order information, i.e., G(u, y) = I, and step size ρ_jac = 10⁻⁵ and once with the (approximated) Hessian as given in (14) and ρ_hes = 3 × 10⁻¹. In both cases, the step size has been manually tuned for fast convergence. Increasing the step size in the first-order case to above 10⁻⁴ renders the closed-loop unstable, leading to divergence of the algorithm. The deviation is evaluated as the root mean square error (RMSE) taken over all scenarios and summed over the BPMs and is plotted on a logarithmic scale in Figure 2.

Figure 2:

Comparison of the convergence speed of the closed-loop feedback optimization with purely first-order information (Jacobian only) and with the approximated Hessian in (14). The RMSE is taken over all BPMs and with 30 random perturbations of the incoming electron beam and quadrupole misalignments, while the shading indicates the best- and worst-case scenarios.

As seen in Figure 2, the RMSE converges to its final value within 20 steps when employing the second-order information of the Hessian. On the other hand, the convergence is much slower when only using the Jacobian, obtaining an RMSE that is over a magnitude worse than the optimal value of the first case after 500 iterations., demonstrating the importance of including second-order information when available. Furthermore, Table 1 shows the time required for computing the derivatives, solving the quadratic program, and in total, with and without employing the Hessian matrix. Even though the variant including the Hessian matrix takes approximately twice as long to evaluate, it is well within the 100 ms sampling interval of the EuXFEL. As such, the improved convergence outweighs the increased computational demand.

4.3 Online identification and control

For the second scenario, we assume that the sensitivity ∂ h ∂ u is unknown and has to be estimated as described in Section 3. The excitation signal z _k is taken as zero mean white Gaussian noise with standard deviation 10⁻⁷. Furthermore, we choose G as given in (14) with the estimated sensitivity ∇ h ̃ k substituted for the Jacobian ∂ h ∂ u and use the step size ρ = 3 × 10⁻¹. For the Kalman filter, the noise covariance matrices are set as Q_e = 10I and R_e = 10⁻⁷I, respectively, where the magnitude of R_e has to be seen relative to the excitation z _k [9].

As before, we simulate 30 runs with randomly sampled incoming beam and quadrupole misalignment, taking the same samples as in the previous section for better comparability. In addition, the initial estimate x ̂ 0 of the sensitivity is chosen at random for each run. The behavior of the RMSE for each position measurement is shown in Figure 3 on a logarithmic scale. Contrasting the previous scenario with known sensitivity ∂ h ∂ u , the optimizer is initially steering the beam away from the BPM center positions due to the randomly initialized sensitivities. However, the RMSE starts improving after less than 20 iterations, quickly recovering the lost performance and converging to an improved solution within 50 steps.

Figure 3:

RMSE over 30 random incoming electron beams and quadrupole misalignment in the x and y component of all ten BPMs.

This behavior demonstrates how using the Hessian (14) can improve the convergence speed even with estimated sensitivity, resulting in a better performance than the first-order case shown in the previous section and a similar final RMSE to the scenario with Hessian. This is visualized in Figure 4. For all 30 scenarios, the difference between the two schemes is not more than 20 μm in RMSE, and in every second scenario the difference is close to 0, indicating that the estimated sensitivity achieves an accuracy comparable to the case with known model.

Figure 4:

Histogram of the difference in final RMSE between optimization with known model from Section 4.2 and with the sensitivity estimation from Section 4.3. Negative values indicate that the estimated sensitivity achieved worse accuracy.

For comparison, the same 30 scenarios have also been optimized using the Nelder-Mead algorithm which is routinely deployed at the EuXFEL through OCELOT [6]. The implementation is taken from the Python scipy library [33], [34] and the resulting RMSE is shown in Figure 5. After the initial increase, both algorithms start optimizing at approximately the same RMSE. However, while the Nelder-Mead simplex algorithm is monotonically improving its solution, the feedback optimizer is converging at a much faster rate.

Figure 5:

Comparison of the total RMSE between feedback optimization and the Nelder-Mead algorithm. The shaded areas indicate the best and worst case of the 30 scenarios.

Finally, consider Figure 6, which shows the Frobenius norm of the difference between the actual Jacobian ∂ h ∂ u of the accelerator and the estimated sensitivity ∇ h ̃ k . Similar to the RMSE of the position, the sensitivity error shows a fast decrease at the beginning. This is caused by the strong excitation of the plant due to the feedback optimizer searching for the optimal input. The filter continues to improve its estimate based on the excitation z _k once the initial plateau is reached by the optimizer, albeit at a slower pace, before converging after 28,000 iterations. Note, however, that the final error and the convergence speed depend on the excitation z _k and the chosen covariance matrices Q_e and R_e.

Figure 6:

Frobenius norm of the difference between the estimated and actual sensitivity plotted over time. The solid line indicates the RMSE, while the shading visualizes the best- and worst-case.