Pseudocode for Krotov’s Method

Michael Goerz
(Dated: December 11, 2019)

For reference, Algorithm 1 shows the complete pseu- U in lines 7, 24 37). This is under the assumption that
docode of an optimization with Krotov’s method, as evaluating U dominates the application of the operator
implemented in the krotov package (https://github. (i−1) (i)
µlkn to the state φk (tn−1 ) and the evaluation of the
com/qucontrol/krotov). inner product of two states, lines 31, 34. This condition
Variables are color coded. Scalars are set in blue, e.g. is fulfilled for any non-trivial Hilbert space dimension.
(0)
ln . States (Hilbert space states or vectorized density
matrices) are set in purple, e.g. φinit Loops over the index k are parallelizable, in particu-
k . They may be
annotated with light gray superscripts to indicate the lar in a shared-memory (multi-threaded) parallelization
iteration-index i of the control under which state was environment like OpenMP. In a (multi-process) method-
propagated, and with light gray time arguments. These passing environment like MPI, some care must be taken
annotations serve only to connect the variables to the to minimize communication overhead from passing large
(0) (0) state vectors. For some (but not all) functionals, inter-
equations of motion: φk (tn ) and φk (tn−1 ) are the same process communication can be reduced to only the scalar
variable φk . Operators acting on states are set in green, values constituting the sum over k in lines 31, 34.
e.g. µlkn . These may be implemented as a sparse matrix
or implicitly as a function that returns the result of ap- The memory requirements of the algorithm are dom-
plying the operator to a state. Lastly, storage arrays are inated by the storage arrays Φ0 , Φ1 , and X. Each of
set in red, e.g. Φ0 . Each element of a storage array is a these must store N (NT + 1) full state vectors (a full time
state. propagation for each of the N objectives). Each state
The Python implementation groups several of the al- vector is typically an array of double-precision complex
gorithm’s input parameters by introducing a list of N numbers. For a Hilbert space dimension d, a state vec-
“objectives”. The objectives are indexed by k, and each tor thus requires 16d bytes of memory, or 16d2 bytes
objective contains the initial state φinit
k , the Hamiltonian for a density matrix. Under certain conditions, the use
or Liouvillian Hk to be used by the propagator U and for of Φ0 and Φ1 can be avoided: both are required only
the operators µlkn , and possibly a “target” to be taken when the second order update is used (σ(t) 6= 0). When
into account by the function χ. In many applications, the first order update is sufficient, Φ1 may overwrite Φ0
Hk ≡ H is the same in all objectives, and µlkn ≡ µl if H so that the two collapse into a single forward-storage
is linear in the controls in addition. The subscript n and Φ. The states stored in Φ are only used for the inho-
(i−1) mogeneity ∂gb /∂ hφk | in Eq. (3), and no storage Φ of
the superscript (i−1) for µlkn in lines 31, 34 comes into
play only if H is not linear in the control. Mathemat- forward-propagated states at all is required if gb ≡ 0.
ically, µlkn would then have to be evaluated using the Thus, in most examples, only the storage X of the
updated control. Since the update is not yet known, the backward-propagated states remains. In principle, if the
guess control may be used as an approximation (valid for time propagation U is unitary (i.e., invertible), the states
sufficiently large λa,l ). stored in X could be recovered by forward-propagation
(i−1)
The CPU resources required for the optimization are of {χk (t = 0)}, eliminating X at the (considerable)
dominated by the time propagation (calls to the function runtime cost of an additional time propagation.

Optimization Functional and Equations of Motion

XZ T Z T
(i) (i) (i) (i) (i)
J[{|φk (t)i}, {l (t)}] = JT ({|φk (T )i}) + ga (l (t)) dt + gb ({φk (t)}) dt (1)
l 0 0

∂  (i) E i 
 (i)
E
φk (t) = − Ĥ (i) (t) φk (t) (2)
∂t h̄


∂  (i−1) E i † (i−1)  (i−1) E ∂gb 
(t) = − Ĥ (t) (t) + (3)
∂t k
χk
h̄ ∂ hφk | (i−1)
χ


 (i−1)
E ∂JT 
with χk (T ) = − (4)
∂ hφk (T )| (i−1)
 2

Algorithm 1 Krotov’s Method for Quantum Optimal Control

Input:
(0) (0)
1. list of guess control values {ln } where ln is the value of the l’th control field on the n’th interval of the propagation
(0) (0)
time grid (t0 = 0, . . . , tNT = T ), i.e., ln ≡ l (t̃n−1 ) with n ∈ [1, NT ] and t̃n ≡ (tn + tn+1 )/2
2. list of update-shape values {Sln } with each Sln ∈ [0, 1]
3. list of update step size values {λa,l }
4. list of N initial states {φinit
k } at t = t0 = 0

5. propagator function U that in “forward mode” receives a state φk (tn ) and a list of control values {ln } and returns
φk (tn+1 ) by solving the differential equation (2), respectively in “backward mode” (indicated as U † ) receives a state
χk (tn ) and returns χk (tn−1 ) by solving the differential equation (3)
6. list of operators µlkn = ∂Hk  , where Hk is the right-hand-side of the equation of motion of φk (t), up to a factor

∂l (t) ln
of (−i/h̄), cf. Eq. (2)
7. function χ that receives a list of states {φk (T )} and returns a list of states {χk (T )} according to Eq. (4)
8. optionally, if a second order construction of the pulse update is necessary: function σ(t)
(opt) (opt) (0)
Output: optimized control values {ln }, such that J[{ln }] ≤ J[{ln }], with J defined in Eq. (1).
(0)
1: procedure KrotovOptimization({ln }, {Sln }, {λa,l }, {φinit k }, U , {µlkn }, χ, σ )
2: i←0 . iteration number
3: allocate forward storage array Φ0 [1 . . . N, 0 . . . NT ]
4: for k ← 1, . . . , N do . initial forward-propagation
(0)
5: Φ0 [k, 0] ← φk (t0 ) ← φinit
k
6: for n ← 1, . . . , NT do
(0) (0) (0)
7: Φ0 [k, n] ← φk (tn ) ← U (φk (tn−1 ), {ln }) . propagate and store
8: end for
9: end for
10: while not converged do . optimization loop
11: i←i+1
(i) (i−1)
12: Φ1 , {ln } ← KrotovIteration(Φ0 , {ln }, . . . )
13: Φ0 ← Φ1
14: end while
(opt) (i)
15: ∀l, ∀n : ln ← ln . final optimized controls
16: end procedure

(i−1)
17: procedure KrotovIteration(Φ0 , {ln }, {Sln }, {λa,l }, {φinit
k }, U , {µlkn }, χ, σ )
(i−1)
18: ∀k : φk (T ) ← Φ0 [k, NT ]
(i−1) (i−1)
19: {χk (T )} ← χ({φk (T )}) . backward boundary condition
20: allocate backward storage array X[1 . . . N, 0 . . . NT ].
21: for k ← 1, . . . , N do
(i−1)
22: X[k, NT ] ← χk (T )
23: for n ← NT , . . . , 1 do . backward-propagate and store
(i−1) (i−1) (i−1)
24: X[k, n − 1] ← χk (tn−1 ) ← U † (χk (tn ), {ln }, Φ0 )
25: end for
26: end for
27: allocate forward storage array Φ1 [1 . . . N, 0 . . . NT ]
(i)
28: ∀k : Φ1 [k, 0] ← φk (t0 ) ← φinit k
29: for n ← 1, . . . , NT do . sequential update loop
(i−1)
30: ∀k : χk (tn−1 ) ← X[k, n − 1]
P
(i−1)  (i−1)  (i)
∀l : ∆ln ← λSa,l ln

31: Im k χk (tn−1 )µlkn φk (tn−1 ) . first order
32: if σ(t) 6= 0 then . second order
(i) (i)
33: ∀k : ∆φk (tn−1 ) ← φk (tn−1 ) − Φ0 [k, n − 1]
Sln

(i)  (i−1)  (i)
∀l : ∆ln ← ∆ln + λa,l Im k 21 σ(t̃n ) ∆φk (tn−1 )µlkn φk (tn−1 )
P 
34:
35: end if
(i) (i−1)
36: ∀l : ln ← ln + ∆ln . apply update
(i) (i) (i)
37: ∀k : Φ1 [k, n] ← φk (tn ) ← U (φk (tn−1 ), {ln }) . propagate and store
38: end for
39: if σ(t) 6= 0 then
40: Update internal parameters of σ(t) if necessary
41: end if
42: end procedure
 3

Notes:
• The index k numbers the independent states to be propagated, respectively the independent “objectives” (see
text for details), l numbers the independent control fields, and n numbers the intervals on the time grid. All of
these indices start at 1.
• The optimization loop may be stopped if the optimization functional or the change of functional falls below a
pre-defined threshold, a maximum number of iterations is reached, or any other criterion.
(i−1)
• The braket notation in line 31 indicates the (Hilbert-Schmidt) inner product of the state χk (tn − 1) and the
(i−1) (i)
state resulting from applying µ lkn to φ k (tn−1 ). In Hilbert space, this is the standard braket. In Liouville
space, it is tr χk † µlkn [φk ] with density matrices χk , φk and a super-operator µlkn .

(i−1)
• For numerical stability, the states χk (T ) in line 19 may be normalized. This norm then has to taken into
account in the pulse update, line 31.
• In line 24, the storage array Φ0 is passed to U † only to account for the inhomogeneity due to a possible state-
dependent constraint, ∂gb /∂ hφk | in Eq. (3). If gb ≡ 0, the parameter can be omitted.