Increasing memory and runtime performance of GRAPE for control in large quantum systems

Gradient Ascent Pulse Engineering (GRAPE) is a popular technique in quantum optimal control. Recent implementations of the GRAPE algorithm are based on automatic differentiation (AD). AD helps avoid the need for coding of analytical gradients but incurs a large memory cost. Specifically, AD stores intermediate states and propagators at all time steps, thus posing a severe bottleneck for quantum systems with large Hilbert space dimension. To address this issue, we implement hard-coded analytical gradients in a scheme that avoids propagator storage and significantly reduces the storage of states. We further succeed in enhancing runtime performance with an improved algorithm for state propagation and propagator derivatives. We benchmark the performance and memory cost of our code against AD-based implementations for large-dimensional state transfer and gate optimization problems. Results confirm the expected improvements that will allow us to tackle optimal control of much larger quantum systems in the future.