Tutorial: Optimization with CRAB

Author: Michael Goerz

CRAB [7] is a gradient-free optimization that assumes that the control pulse can be parametrized with a few (a “chopped” selection of) randomized spectral components. This tutorial demonstrates an optimization using CRAB for a reduced model for two transmon qubits.

The model is taken from Ref [8], section II.C. In that paper, the same system was optimized using Krotov’s method.

• [7] Doria et al., Phys. Rev. Lett. 106, 190501 (2011)

• [8] Goerz et al., Phys. Rev. A. 91, 062307 (2015)

In addition to the QDYN, we will heavily utilize the scientific Python ecosystem in modeling the system and performing the optimization, including the QDYN Python support package (qdynpylib) in revision c3ca3eb

%matplotlib inline
### If you would like to show the version information, you can do this by installing
### the version_information ipython magic from goerz's github page by the following command:
##!pip install git+git://github.com/goerz/version_information.git

### After that (or if you already installed it), you can use it by commenting in the folling lines:
#%load_ext version_information
#%version_information scipy, numpy, matplotlib, qutip

Software	Version
Python	3.7.6 64bit [GCC 7.3.0]
IPython	7.15.0
OS	Linux 4.15.0 88 generic x86_64 with debian buster sid
scipy	1.4.1
numpy	1.18.5
matplotlib	3.2.2
qutip	4.5.0
Fri Jun 26 11:43:24 2020 CEST

The Transmon Hamiltonian and logical basis

We will use the qutip framework to help us construct quantum objects:

import qutip
from qutip import tensor, basis

It is straighforward to write out the drift Hamiltonian H0 and the control Hamiltonian H1 depending on the qubit frequencies (w1, w2), the frequency of the rotating frame wd, the anharmonicities (alpha1, alpha2), the effective qubit-qubit couplint Jeff, and the relative coupling strength lamda (sic!, to not conflict with the lambda keyword in Python). We assume the transmon is cut off after N levels

def H0_H1(w1, w2, wd, alpha1, alpha2, Jeff, lamda, N):
    b1 = qutip.destroy(N); b1_dag = b1.dag()
    b2 = qutip.destroy(N); b2_dag = b2.dag()
    Id = qutip.identity(N)
    w1 = w1 - wd; w2 = w2 - wd  # RWA
    H0 = (tensor(((w1 + 0.5 * alpha1) * b1_dag * b1 +
                  (0.5 * alpha1) * b1_dag * b1 * b1_dag * b1), Id) +
          tensor(Id, ((w2 + 0.5 * alpha2) * b2_dag * b2 +
                      (0.5 * alpha2) * b2_dag * b2 * b2_dag * b2)) +
          Jeff * (tensor(b1_dag, b2) + tensor(b1, b2_dag)))
    H1 = tensor((b1  + b1_dag), Id) + tensor(Id, lamda*(b2 + b2_dag))
    return H0, H1

In correspondence to Ref [1] Table 1 (respectively, the original data used in the calculations for that paper), we use the following parameters:

w1     = 4.3796         # GHz
w2     = 4.61368        # GHz
wd     = 4.49848        # GHz
alpha1 = -0.210         # GHz
alpha2 = -0.215         # GHz
Jeff   = -0.00299394    # GHz
lamda  = 1.0326725681   # GHz

N = 5 # number of levels

Note that a factor of \(2 \pi\) is implicit (QDYN knows how to use GHz as an energy unit)

H0, H1 = H0_H1(w1, w2, wd, alpha1, alpha2, Jeff, lamda, N)

For simplificity, we will use the canonical bare states to encode the qubits:

ket00 = tensor(basis(N, 0), basis(N, 0))
ket01 = tensor(basis(N, 0), basis(N, 1))
ket10 = tensor(basis(N, 1), basis(N, 0))
ket11 = tensor(basis(N, 1), basis(N, 1))

The CRAB control pulse

We choose a fixed gate duration \(T\), and a number of time grid point nt that will allow us to accurately propagate our pulse in the rotating wave approximation.

T = 300 # ns
nt = 3000

The QDYN Python library already knows about the formula

\begin{equation} S(t) = \sum_{n} (a_n \cos(\omega_n t) + b_n \cos(\omega_n t)) \end{equation}

of the basic CRAB envelope

from qdyn.pulse import CRAB_carrier

We choose a fixed number \(n\) of frequency components (this is the “chopped” part of CRAB)

n_crab = 10

The frequencies \(\omega_n\) themselves are also fixed. They are chosen about the first \(n\) principal harmonics of the total pulse duration \(T\):

import numpy as np

harmonics = (1.0/T) * np.arange(1,n_crab+1)

harmonics # GHz (2pi implicit)

array([0.00333333, 0.00666667, 0.01      , 0.01333333, 0.01666667,
       0.02      , 0.02333333, 0.02666667, 0.03      , 0.03333333])

The actual frequencies are shifted from these principal harmonics in a randomized fashion (the “randomized” part of CRAB)

w_n = harmonics + (1.0/T) * (np.random.random(n_crab) - 0.5)

w_n # GHz

array([0.00188079, 0.00652263, 0.01081657, 0.01490174, 0.01573941,
       0.01942468, 0.02315663, 0.02649059, 0.03030823, 0.03172236])

We can now define the total CRAB pulse amplitude in terms of the free coefficient arrays \(a_n\) and \(b_n\). We will multiply the entire pulse with a “flattop” shape that switches the pulse on and off in a 7.5 ns window, to ensure that the result remains physical.

from qdyn.pulse import flattop
A = 1.0  # amplitude scaling parameter
def crab_amplitude(tgrid, a_n, b_n):
    return (
        A *
        CRAB_carrier(tgrid, time_unit='ns', freq=w_n, freq_unit='GHz',
                     a=a_n, b=b_n) *
        flattop(tgrid, t_start=0, t_stop=T, t_rise=0.025*T))

QDYN provides data structures for working with such analytical pulses

from qdyn.analytical_pulse import AnalyticalPulse
AnalyticalPulse.register_formula('crab', crab_amplitude)

Our guess pulse will start from a random choice of coefficients

parameters = {
    'a_n': np.random.random(n_crab),
    'b_n': np.random.random(n_crab)
}

The pulse should be normalized (by redefining A) to a reasonable peak amplitude:

E0 = 0.5 # GHz

max_ampl = np.max(np.abs(
    AnalyticalPulse('crab', parameters=parameters, time_unit='ns', ampl_unit='GHz', T=T, nt=nt)
    .to_num_pulse().amplitude))
A = E0 / max_ampl

Note that as a variation, one may also make the w_n free parameters, instead of keeping them fixed. This gives more flexibility, but also increases the number of control parameters, which may slow down a Nelder-Mead simplex search significantly.

crab_pulse = AnalyticalPulse('crab', parameters=parameters, time_unit='ns', ampl_unit='GHz', T=T, nt=nt)

crab_pulse.to_num_pulse().show(show_spectrum=False)

For our optimization, it crucial that we have easy access all the free parameters. The AnalyticalPulse class can give the concatenation of the coefficients \(a_n\) and \(b_n\), i.e. the total array of control parameters:

x0 = crab_pulse.parameters_to_array()
x0

array([0.21946456, 0.39159222, 0.43977924, 0.35734139, 0.85263094,
       0.16836106, 0.12676104, 0.54339242, 0.94889672, 0.10252433,
       0.33185406, 0.31316668, 0.40038705, 0.21063435, 0.57617372,
       0.90368758, 0.02605315, 0.66170273, 0.64102805, 0.02600543])

We will be able to use the inverse array_to_parameters method when modifying these paremeters.

The QDYN numerical model

The QDYN Python library provides routines for setting up a full numerical model

import qdyn.model

model = qdyn.model.LevelModel()

model.add_ham(H0, op_unit='GHz', op_type='pot')
model.add_ham(H1, pulse=crab_pulse.to_num_pulse(),
              op_unit='dimensionless', op_type='dip')
model.add_state(ket00, '00')
model.add_state(ket01, '01')
model.add_state(ket10, '10')
model.add_state(ket11, '11')
model.set_propagation(T, nt, time_unit='ns', prop_method='cheby')

Since we are doing gate optimization, we will need independent runfolders for propagating all of the two-qubit basis states

model.write_to_runfolder('./rf00')
model.write_to_runfolder('./rf10')
model.write_to_runfolder('./rf01')
model.write_to_runfolder('./rf11')

We can now propagate the guess pulse and analyze its figure of merit

Propagation of the guess pulse

For the propagation, we can make use of the qdyn_prop_traj utility that comes as part of QDYN

! qdyn_prop_traj -v

QDYN 2.0dev revision cc7301ddca2f8019bfe82fc7553c1f51ccf9679c (44-check-example-qdynpylib)
  features: no-check-cheby, no-check-newton, no-parallel-ham, use-mkl=sequential, parallel-oct, backtraces, no-debug, no-no-ipo
  compiled with ifort on Fri Jun 26 10:58:17 2020 on host talpiyot
1

Since we are interested in gates, we must propagate all four two-qubit basis states. We can do this in parallel, using the subprocess module of the Python standard library

import subprocess as sp
import os

The following wrapper propagates the runfolders ./rf00, ./rf01, ./rf10, and ./rf11 in their current state, leaving a file psi_final.dat.1 in each of them that contains the result of the propagation

def propagate():
    """Propagate all runfolders in their current state"""
    def proc(label, out_fh):
        """Background process that propagates for the given label"""
        return sp.Popen(
            ['qdyn_prop_traj', "--internal-units=GHz_units.txt",
             "--write-final-state=psi_final.dat",
             "--state-label=%s" % label, "."],
            cwd="./rf%s" % label, stdout=out_fh
        )
    labels = ['00', '01', '10', '11']
    # cleanup of previous propagation
    for rf in ["./rf%s" for label in labels]:
        try:
            os.unlink(os.path.join(rf, 'psi_final.dat.1'))
        except OSError:
            pass
    # run
    procs = []
    for label in labels:
        with open("%s.log" % label, 'w') as log_fh:
            procs.append(proc(label, out_fh=log_fh))
    for p in procs:
        p.wait()

propagate()

The four propagated states can be read using the read_psi_amplitudes routine from QDYN, and we can easily construct the two-qubit gate (as a QDYN.gate2q.Gate2Q object).

from qdyn.io import read_psi_amplitudes

def get_U():
    """Construct Gate2Q instance from current state of runfolders"""
    basis = ket00, ket01, ket10, ket11
    np_basis = [np.array(ket.data.todense()).flatten()
                for ket in basis]
    ket_T = [read_psi_amplitudes(
                 "./rf%s/psi_final.dat.1" % label, N*N)
             for label in ['00', '01', '10', '11']]
    return qutip.Qobj(
        [[np.inner(np_basis[i], ket_T[j])
          for j in range(4)]
         for i in range(4)])

U_guess = get_U()

Our figure of merit, to be maximized, is the concurrence of the gate:

import weylchamber

def concurrence(U):
    return weylchamber.concurrence(*weylchamber.c1c2c3(U))
concurrence(U_guess)

0.8084839972538467

However, we must also take into account the loss of population from the logical subspace

def pop_loss(U):
    return 1-np.trace((U.dag()*U).full().real)/U.shape[0]
pop_loss(U_guess)

0.09468761461282993

Optimization

The core idea of the CRAB method is fully expressed in the way that the pulse is parametrized. The actual optimization is then a (Nelder-Mead) simplex-search over the free parameters. We can use the “black box” algorithm that the scipy package provides.

from scipy.optimize import minimize

In order to conform to the interface of minimize, we must define a function that takes an array of control parameters and returns a value of the optimization functional, to be minimized

def fun(x):
    crab_pulse.array_to_parameters(x)
    for label in ['00', '01', '10', '11']:
        rf = "./rf%s" % label
        pulse_file = os.path.join(rf, 'pulse1.dat')
        crab_pulse.to_num_pulse().write(pulse_file)
    propagate()
    U = get_U()
    f = 1.0 - concurrence(U) + pop_loss(U)
    with open("simplex.log", 'a') as log_fh:
        log_fh.write("%s -> %s\n" % (x, f))
    return f

The starting point of the optimization is the earlier defined x0.

fun(x0)

0.28620361735898325

The optimization itself is now a one-liner:

res = minimize(fun, x0, method='Nelder-Mead')

res.message

'Maximum number of function evaluations has been exceeded.'

After the optimization ends, we can compare the initial and final control parameters, and the control field that they imply

x0

array([0.21946456, 0.39159222, 0.43977924, 0.35734139, 0.85263094,
       0.16836106, 0.12676104, 0.54339242, 0.94889672, 0.10252433,
       0.33185406, 0.31316668, 0.40038705, 0.21063435, 0.57617372,
       0.90368758, 0.02605315, 0.66170273, 0.64102805, 0.02600543])

res.x

array([0.15429222, 0.59462772, 0.48740891, 0.19549337, 0.73368345,
       0.19649435, 0.11174224, 0.50744211, 0.6519327 , 0.10224495,
       0.44648893, 0.41122166, 0.36985051, 0.23245543, 0.77741703,
       1.17952445, 0.02941732, 0.43869892, 0.63524623, 0.02429511])

crab_pulse.array_to_parameters(x0)
crab_pulse.to_num_pulse().show(show_spectrum=False)

crab_pulse.array_to_parameters(res.x)
crab_pulse.to_num_pulse().show(show_spectrum=False)

The number of evaluations (i.e., the number of propagations) is available in res.nfev:

res.nfev

4000

The optimized solution induces the gate

U_opt = get_U()

As a reminder, the starting point for our figure of merit was

concurrence(U_guess)

0.8084839972538467

pop_loss(U_guess)

0.09468761461282993

This has now improved to

concurrence(U_opt)

0.9914968914386503

pop_loss(U_opt)

0.004283928347898924

While this is far from perfect, it is a signficant improvement. We may get better results by varying the relative weigth between the concurrence and the population loss in the optimization functional. We may also project the pulse onto a different CRAB parametrization (“DCRAB” [9]), and do another optimization. This approach has been shown to elimintate local traps.

Typically, it is also a good idea to run several CRAB optimizations in parallel (with different w_n), and to pick the best outcome at the end.

• [9] Rach et al., Phys. Rev. A 92, 062343 (2015)