Efficient Dynamic and Momentum Aperture Optimization for Lattice Design Using Multipoint Bayesian Algorithm Execution

TL;DR

The paper tackles the computational bottleneck in maximizing dynamic and momentum aperture in storage-ring design. It introduces multipointBAX, which uses neural-field surrogates to model DA/MA maps at the single-particle level and guides acquisition with a MeanBAX strategy that focuses on boundary points near the Pareto front. The approach yields two-order-of-magnitude speedups over state-of-the-art methods and achieves equivalent Pareto-front results with far fewer tracking simulations, demonstrated on the SSRL-X lattice with batch-accelerated, early-stopped optimization. This has practical impact for designing future light sources, colliders, and large-scale facilities by enabling robust, scalable, and interpretable multipoint optimization across complex accelerator lattices.

Abstract

We demonstrate that multipoint Bayesian algorithm execution can overcome fundamental computational challenges in storage ring design optimization. Dynamic (DA) and momentum (MA) optimization is a multipoint, multiobjective design task for storage rings, ultimately informing the flux of x-ray sources and luminosity of colliders. Current state-of-art black-box optimization methods require extensive particle-tracking simulations for each trial configuration; the high computational cost restricts the extent of the search to $\sim 10^3$ configurations, and therefore limits the quality of the final design. We remove this bottleneck using multipointBAX, which selects, simulates, and models each trial configuration at the single particle level. We demonstrate our approach on a novel design for a fourth-generation light source, with neural-network powered multipointBAX achieving equivalent Pareto front results using more than two orders of magnitude fewer tracking computations compared to genetic algorithms. The significant reduction in cost positions multipointBAX as a promising alternative to black-box optimization, and we anticipate multipointBAX will be instrumental in the design of future light sources, colliders, and large-scale scientific facilities.

Figures (9)

• Figure 1: Comparison of Bayesian optimization and BAX for a single objective optimization. There are two key differences: first, in multipointBAX we simulate a single particle at each acquisition, whereas BO simulates a full map of more than one thousand particles. Second, the BO surrogate models the objective function (configuration $\rightarrow$ objective), whereas the BAX surrogate models the underlying physics function (i.e. configuration $\rightarrow$ map). Because of the increased complexity and large number of data points to model, BAX also requires use of a NN model whereas BO can use a simpler Gaussian process model.
• Figure 2: Schematic of the multipointBAX implementation for DAMA optimization. For more details, see description in Appendix \ref{['appendix:algorithm']} and Algo \ref{['alg:mpbax']}.
• Figure 3: Model learning process: Left scatter plot shows the configuration input distribution projected on two dimensions (u0 and u1) using a UMAP umap transform fit on initial samples of BAX. Gray points show all sampled configurations during the BAX run, blue points show final PF from GA, red points show final PF from BAX, and green points show initial PF from BAX. Right heatmaps: Each picture shows DA (left three columns) and MA (right three columns) maps predicted by the NN. 'spos' is the position along the accelerator. Top row in the orange box shows a 'good' configuration predicted to lie on the Pareto front (orange on the UMAP), while bottom row in the purple box shows a random 'bad' configuration (purple on the UMAP) far from the predicted Pareto front. For both DA and MA, from left to right: Prediction after initial training (Initial DA/MA), prediction after final training (Final DA/MA), and groundtruth (GT DA/MA).
• Figure 4: BAX vs. NSGA-II performance as a function of iteration number using hypervolume. We run both BAX and NSGA-II twice, and show the mean performance (dashed and solid curves) with the corresponding $2\sigma$ error bands (light and dim shaded areas). Dashed lines indicate optimization cost, and solid lines include the cost of the final evaluation. Red stars denote BAX’s early-stopping points, and green star indicates where NSGA-II reaches comparable performance.
• Figure 5: Neural net architectures for DA and MA. The fc layers denoted by shaded blue area are fully connected layers. The lightgray blocks show the change of the input data shape when it goes through the layers. The lightblue block is the particle position part of the input. We concatenate the input again with the fc4 layer output to weight the inputs more, although the observed impact on training was not significant.