Dynamic compensation of stray electric fields in an ion trap using machine learning and adaptive algorithm

A gradient descent algorithm (ADAM) and a deep learning network (MLOOP) were tested for compensating stray fields in different working regimes. The source code used for the experiments is available in²⁸. The software controlled the voltages using the PXI-6713 DAQ and read the fluorescence counts from a photo-mutliplier-tube (PMT) through a time tagging counter (IDQ id800). All software was written in python and interfaced with the DAQ hardware using the library NI-DAQmx Python. A total of 44 DC electrodes and the horizontal position of the cooling laser were tuned by the program, resulting in a total of 45 input parameters.

Gradient descent optimizer

The first compensation test was performed by ADAM gradient descent algorithm. This is a first order optimizer that uses the biased first and second order moments of the gradient to update the inputs of an objective function, and was chosen for its fast convergence, versatility in multiple dimensions and tolerance to noise²³. Our goal was to maximize the fluorescence of the ion which was described by a function (f(vec alpha ,)), where (vec alpha , =(alpha _1, alpha _2, alpha _3, ldots alpha _45)) represents the array of parameters to be optimized. To find the optimal (vec alpha ,), the algorithm needs to know the values of the partial derivatives for all input parameters. Because we do not have an analytic expression for (f(vec alpha ,)), the values of its derivatives were estimated from experimental measurements by sequentially changing each input (alpha _i), and reading the associated change in fluorescence f. This data were used as inputs to ADAM for finding the optimal (vec alpha ,) which maximized f.

Before running the automated compensation, we manually adjusted the 4 weights of the voltage sets used for compensation described in the previous section. We also tried to run ADAM to optimize these 4 parameters but the increase in fluorescence was limited to 6{64d42ef84185fe650eef13e078a399812999bbd8b8ee84343ab535e62a252847}. After manual compensation, we ran ADAM on all 45 inputs with the algorithm parameters given in the source code²⁸. Each iteration took 12 s, where 9.8 s were the photon readout (0.1 s(times )2 readouts per parameter plus 2(times )0.1 s readouts at the beginning and end of the iteration), and the rest of the time was the gradient computation. If the photon count reduced by more than 40{64d42ef84185fe650eef13e078a399812999bbd8b8ee84343ab535e62a252847} of its initial value, the algorithm terminated and applied the previously found optimum. This acted as the safety net for the program, ensuring the ion was not lost while optimizing the 45 inputs. We need this safety net because if the ion is heated past the capture range for the used cooling detuning, it will be ejected from the trap. In our implementation of the algorithm we removed the reduction in the step size of the optimization algorithm as iterations progressed. This step reduction, which is present in the standard version of ADAM, is not ideal when stray fields change with time since the optimal values of the voltages also drift in time. The removal caused some fluctuations in the photon readout near the optimal settings. Adding to these fluctuations, other sources of noise, such as wavemeter laser locking³², and mechanical drift in the trap environment, resulted in daily photon count variations of around 5{64d42ef84185fe650eef13e078a399812999bbd8b8ee84343ab535e62a252847}. Fluctuation in laser power was not a concern here since the power of the cooling laser was stabilized. Despite these fluctuations, and the fact that stray fields change every day, the algorithm demonstrated an increase in fluorescence collection up to 78 ± 1{64d42ef84185fe650eef13e078a399812999bbd8b8ee84343ab535e62a252847} (Fig. 2b) when starting from a manually optimized configuration in less than 10 iterations, or 120 s.

The ADAM algorithm was fast and reliable (the ion was never lost during optimization), even in extremely volatile conditions like having time-dependent charging and stray electric field buildup. Figure 3a shows a colourmap of the voltages and laser position adjustments, where most of the improvement came from adding the same voltage to all DC electrodes indicating that the ion was not at optimal height. The volatility of the ion-trap environment causes the fluorescence rate to oscillate around the optimal point. To get the best value, instead of using the values of the final iteration, the software saved all voltage combinations and applied the setting with the highest photon count after all iterations were finished. Despite picking the best value it can be seen in Fig. 2b that the fluorescence for some iterations during the optimization are higher than the final point selected by the software. This is because when the settings are changed, the ion fluorescence rate may transiently increase and subsequently stabilize to a slightly lower value for the same voltage settings.

Deep learning network for the ion trap

The second algorithm tested was a deep learning network using the python based optimization and experimental control package MLOOP²⁰. MLOOP uses Differential Evolution³³ for exploring and sampling data. The blue points in Fig. 4a corresponds to these samples and it can be seen that even at the end of optimization, they can have non-optimum fluorescence rates. MLOOP also trains a neural network using the data collected by Differential Evolution and creates an approximate model of the experimental system. It then uses this model to predict an optimum point. The red points in Fig. 4a shows the optimum points predicted by the neural network model. It can be seen that this section starts later than Differential Evolution, as it requires some data for initial neural network training, and gradually finds the optimum and stays near it. For training of the neural network, the inbuilt ADAM optimizer is used to minimize the cost function. The sampling in MLOOP does not require a gradient calculation which greatly improves the sampling time. Even though the sampling is fast, training the network to find an optimal point requires a minimum of 100 samples and that makes MLOOP slower than ADAM. With our settings for MLOOP, each iteration took 0.7 s on average and therefore 700 s was needed to take 1000 samples shown in Fig. 4a.

In our test the neural network in MLOOP had 5 layers with 45 nodes each, all with Gaussian error correction. The neural network structure (number of layers and cells) was manually optimized and tested on a 45-dimensional positive definite quadratic function before being used for the experiment. Once the ion was trapped, positioned above the integrated mirror²², and photon counts were read, the program started sampling 100 different voltage combinations around its initial point. Then, the network started training on the initial data and making predictions for the voltages that maximise fluorescence. Since the ion trap setup is very sensitive to changes in the electric field, the voltages were set to move a maximum of 1{64d42ef84185fe650eef13e078a399812999bbd8b8ee84343ab535e62a252847} of their previous value in each iteration to reduce the chance of losing the ion. As a step size value could not be explicitly defined, this percentage was chosen to make the changes similar to the step size used for ADAM.

A small percentage of our initial trials with the maximum change of a few percent (instead of 1{64d42ef84185fe650eef13e078a399812999bbd8b8ee84343ab535e62a252847}) led to an unstable ion during the parameter search sequence. This is because MLOOP is a global optimizer and can set the voltages to values far from the stable starting point. Since the ion trap is a complicated system that can only be modelled for a specific range of configurations, moving away from these settings can lead to unpredictable and usually unstable behavior. MLOOP also has an in-built mechanism that handles function noise using a predefined expected uncertainty. We set this uncertainty to the peak-to-peak noise of the photon readout when no optimization was running.

Since MLOOP is a global optimizer it was able to find optimum points different from the points found by ADAM. For trials where low numbers of initial training data points were used, these configurations proved to be unstable and in most cases resulted in the loss of the ion. Unstable states were also observed occasionally if the optimizer was run for too long. With moderate-size training sets, MLOOP was able to find voltage settings with fluorescence rates similar or higher than optimum points found by ADAM as shown in Fig. 4a. Considering the long duration of the MLOOP iteration sequence and the possibility of finding unstable settings in volatile conditions, the test of optimization with induced changing stray fields (“Testing under poor trap conditions”) was only performed with the ADAM optimizer as the gradient based search method proved to be more robust against fluctuations in the ion environment.

Table 1 Saturation power ((P_sat)), detection efficiency ((eta )) and average background count rate (Bkg) before and after optimization with different methods.

To test the effectiveness of the protocols, the saturation power, (P_sat), was measured before and after the optimization process. The (P_sat) is the laser power at which the fluorescence rate of a two-level system is half the fluorescence at infinite laser power. We also measured the overall detection efficiency (eta ), the fraction of emitted photons which resulted in detection events. Table 1 shows (P_sat) decreased (ion photon absorption was improved) using both ADAM and MLOOP. The detection efficiency was approximately the same for all runs as expected.

Another test was done by measuring fluorescence versus laser detuning before and after optimization. Figure 4b shows that the measured values follows the expected Lorentzian profile^29,30,31 and associated linewidth before and after optimization. This indicates that the initial micromotion magnitude (beta ) was sufficiently small for fluorescence to be a good optimization proxy. Clear increase in florescence can be seen after optimizing 44 electrodes individually both with ADAM and MLOOP. The fit residual curve (difference between the experimental values and the theoretical fit) shows that optimizing individual electrodes, resulted in slight increase in heating instability near the resonance.

Testing under poor trap conditions

To test the live performance of the optimization protocol in a non-ideal situation, we deliberately charged the trap by shining 369.5 nm UV laser light onto the chip for 70 min. The power of the laser was (200pm 15mu W) and the Gaussian diameter of the focus was (120pm 10mu m). This process ejects electrons due to the photo-electric effect³⁴ and produces irregular and potentially unpredictable slow time varying electric fields within the trap. The process charged the trap significantly and made a noticeable reduction to the photon count. The ADAM algorithm was then tested both during charging and after charging was stopped. In both cases an improvement of fluorescence rate was observed.

The first experiment was performed to test the optimizing process after charging. In this test, starting with the optimal manual setting, ADAM individual electrode optimizer was able to obtain 27{64d42ef84185fe650eef13e078a399812999bbd8b8ee84343ab535e62a252847} improvement in the fluorescence rate (blue points on the left side of Fig. 5a). Then charging was induced onto the trap for 70 min and a clear decrease in photon count was seen that went even lower than the initial value (red points in Fig. 5a). At this point charging was stopped and ADAM was run again and fluorescence rate returned back to the previous optimum, within the error, in approximately 12 min. During the second optimization, the fluorescence goes higher than the stable final value for some iterations before the final. This is because of the same effect explained in “Gradient descent optimizer” section that the fluorescence might spike right after a change but go down slightly after stabilizing. Looking at the changes of individual electrodes, shown in Fig. 3b, we see that the main electrodes adjusted were those around the ion and some throughout the trap. The change in the laser horizontal position was negligible.

Another experiment was done by running ADAM during continuous charging for real-time compensation. Since we induce charging via laser scattering from the trap, the collected photons are both from the ion and the scattered laser and fluctuations in the intensity of scattered light confuses the optimizer. Despite that the optimizer did not lose the ion nor needed to abort the process. Figure 5b shows that the fluorescence rate, even after a 70-min charging session, remained near the optimum value. After stopping the charging, the ion remained trapped for more than 8 h and was intentionally removed from the trap after this time.