Quantum Optimal Control with SciML

Numerical quantum optimal control is simply optimizing pulse shapes used to drive the control lines on some controllable quantum system to drive that system through some desired state-to-state or unitary transformation. As with any numerical optimization gradient information makes the search much more efficient and not surprisingly two of the more successful approaches have used ways to calculate gradient information along with the evolution.

The Gradient Ascent Pulse Engineering (GRAPE) ¹ approach assumes piece-wise constant pulse amplitudes very similar to an ideal arbitrary waveform generator and then by saving the forwards and backward evolution up to each timestep gradients for the amplitude of the control at each time step are easy to calculate. The original paper used an approximate gradient for the timestep propagator which limited it to small timesteps and/or small Hamiltonians but of course earlier work in NMR had already shown that exact derivatives are available ² from diagonalizing the Hamiltonian which may have already been done to calculate the matrix exponential for the propagator.

The Gradient Optimization of Analytic conTrols (GOAT) ³ algorithm parametrizes the pulse shape as the coefficients of a series of basis functions. This is convenient because implementation constraints such as “the pulse should start and finish at zero amplitude” and bandwidth constraints are easy to enforce in the choice of basis functions itself whereas in GRAPE they need to be implemented as penalty functions. For GOAT the derivatives are calculated by extending the equations of motions to solve for the gradient of the unitary using the gradient of the Hamiltonian which requires some chain rule work.

All the work pouring into automatic differentiation has made implementing these two approaches rather straightforward without any extra work to implement the particular approaches of the GRAPE (saving forward and backwards evolution or GOAT (extending the equations of motion to include the derivative) approaches. The Schuster group has a paper discussing how to apply TensorFlow to GRAPE ⁴. Here, I’ll show how we can use parts of the SciML framework in Julia to implement the GOAT approach without having to actually implement GOAT.

The DRAG approach ⁵,⁶ for limiting the effects of nearby leakage levels was initially discovered from numerical quantum optimal control. The team was searching for pulse shapes for transmon qubits and kept landing on the same looking pulse which was then later justified by the full DRAG theory. The problem is a great sanity check for any quantum optimal control package: given a three level system — two qubit levels and a nearby leakage level — trying to optimize a qubit unitary should give something that looks like a DRAG pulse.

notebook nbviewer link with plots