# GRAPE code for Khaneja et al. "Optimal Control of Coupled Spin Dynamics: Design of NMR Pulse Sequences by Gradient Ascent Algorithms"

I want to write gradient ascent pulse engineering (GRAPE) and I try the simple Hamiltonian H=\Omega*|0\rangle\langle 1| + \Omega*|1\rangle\langle 0|. The initial state is |0\rangle\langle 0| and the final state is |1\rangle\langle 1|. However, I cannot get the right result. I have two questions: (1) I cannot show my \Omega pulse in figure; (2) The GRAPE dose not achieve the desired outcome. I cannot find what my question is. Here is my code:

# "gradient ascent pulse engineering (GRAPE) coed demo"
using LinearAlgebra, PyPlot

# Define variables
dim = 2
Natom = 1

# Define basis
s0 = [1; 0]
s1 = [0; 1]

tr01 = s0 * s1'
tr10 = s1 * s0'

# Start state
global Ļ0 = s0 * s0'

# Target state
global Ļf = s1 * s1'

# Time list
t_list = range(0, 1, length=101)
dt = t_list[2] - t_list[1]

# Initial Pulses
global Ī© = rand(length(t_list)) .* (2Ļ) .- Ļ

# Initialize
fidelity = []
Ļ0_list = Vector{Any}(undef, length(t_list))
Ļf_list = Vector{Any}(undef, length(t_list))

function PWC_pulse(pwc_pulse)
pwc_pulse = [0; pwc_pulse; 0]
time_steps = 0:(length(pwc_pulse)-1)
time_steps_stair = repeat(time_steps, 1, 2)[1:end-1]
pwc_pulse_stair = repeat(pwc_pulse, 1, 2)[1:end-1]
return pwc_pulse_stair
end

for iter in 1:1501
# Bidirectional Evolution
for i in 1:length(t_list)
H = Ī©[i] * (s0 * s1' + s1 * s0')
U_indt = exp(-im * H * dt)
global Ļ0 = U_indt * Ļ0 * U_indt'
Ļ0_list[i] = Ļ0
end
for i in length(t_list):-1:1
H = Ī©[i] * (s0 * s1' + s1 * s0')
U_indt = exp(im * H * dt)
global Ļf = U_indt' * Ļf * U_indt
Ļf_list[i] = Ļf
end

for i in 1:length(t_list)
grad_Ī©[i] = -real(tr(Ļf_list[i]' * 1im * dt * ((tr01 + tr10) * Ļ0_list[i] - Ļ0_list[i] * (tr01 + tr10))))
end

# Update Pulses
lr = 0.01 - iter / 100

# Fidelity
global fidelity = abs(tr(Ļf * (s1 * s1')))
global infid = 1 - fidelity

# Plot Pulses
if mod(iter, 100) == 0
figure()
plot(pwc_pulse_stair)
xlabel("time")
ylabel("pulse strength")
title("$(iter)-th fidelity:$(round(fidelity, digits=4))")
if iter == 300
savefig("pulses_300.png")
end

sleep(0.01)
end
end

figure()
plot(1:length(infid), log.(infid))
xlabel("iteration")
ylabel("infidelity")
savefig("infidelity.png")
gcf()


An optimization like GRAPE is not guaranteed to achieve the desired target state. It is a gradient-based optimization, thus it might end up in a local minimum or it might be the case that your Hamiltonian is simply physically incapable of achieving the final state you desire.

A typical debugging tool in such situations would be to plot your fidelity as a function of the iteration step (i.e. each time you make a gradient-following step, record the current fidelity). If the fidelity plateaus, then probably you have reached a local minimum or your Hamiltonian does not permit the evolution you want. If the fidelity goes down, then probably something is wrong with your gradient calculation. If the fidelity goes up but has not reached a saturation point, then you probably need to run the optimization for longer.

Also, consider using pre-packaged quantum control tools like:

These tools are tested by many people and less likely to contain bugs.

Thanks for your suggestion. In fact, I have also seen these packages (but I havenāt figured out their usage yet). However, I would rather figure it out the logic of GRAPE algorithm compared to using ready-made packaged. If you have any knowledge in this area, I would like to seek your help. I believe you have clearly written out the gradient, but why cannot it be optimized correctly

    # Gradient Calculation
for i in 1:length(t_list)
grad_Ī©[i] = -real(tr(Ļf_list[i]' * 1im * dt * ((tr01 + tr10) * Ļ0_list[i] - Ļ0_list[i] * (tr01 + tr10))))
end


Try the debugging plot I suggested. That would be a good way to check your gradient for correctness.

Also: just use āfinite differencesā and check numerically whether your gradient is correct.

I followed your advice, but it seems like Iām at a loss because the programs I wrote are all from that paper and I canāt detect any issues with it. I have revised a part of my program again, and it can now display all images normally, but still cannot see the results I want. Here is the code

# "gradient ascent pulse engineering (GRAPE) coed demo"
using LinearAlgebra, PyPlot
pygui(true)

# Define variables
dim = 2
Natom = 1

# Define basis
s0 = [1; 0]
s1 = [0; 1]

tr01 = s0 * s1'
tr10 = s1 * s0'

# Start state
global Ļ0 = s0 * s0'

# Target state
global Ļf = s1 * s1'

# Time list
t_list = range(0, pi, length=101)
dt = t_list[2] - t_list[1]

# Initial Pulses
global Ī© = rand(length(t_list)) # .* (2Ļ) .- Ļ

# Initialize
fidelity = []
Ļ0_list = Vector{Any}(undef, length(t_list))
Ļf_list = Vector{Any}(undef, length(t_list))
infid = zeros(301)

for iter in 1:301
# Bidirectional Evolution
for i in 1:length(t_list)
H = Ī©[i] * (s0 * s1' + s1 * s0')
U_indt = exp(-im * H * dt)
global Ļ0 = U_indt * Ļ0 * U_indt'
Ļ0_list[i] = Ļ0
end
for i in length(t_list):-1:1
H = Ī©[i] * (s0 * s1' + s1 * s0')
U_indt = exp(-im * H * dt)
global Ļf = U_indt' * Ļf * U_indt
Ļf_list[i] = Ļf
end

for i in 1:length(t_list)
grad_Ī©[i] = -real(tr(Ļf_list[i]' * 1im * dt * ((tr01 + tr10) * Ļ0_list[i] - Ļ0_list[i] * (tr01 + tr10))))
end

# Update Pulses
lr = 0.5 - iter / 100

# Fidelity
global fidelity = abs(tr(Ļf * (s1 * s1')))
infid[iter] = 1 - fidelity

# Plot Pulses
if mod(iter, 100) == 0
figure(iter)

# subplot(1,10,mm)
xlabel("time")
ylabel("pulse strength")
title("$(iter)-th fidelity:$(round(fidelity, digits=4))")

show()
# if iter == 100
#     savefig("pulses_100.png")
# end
sleep(0.1)
# close()  # Close the figure
end
end
figure()
# subplot(1,2,2)
plot(infid)
xlabel("iteration")
ylabel("infidelity")
savefig("infidelity.png")
show()
# gcf()


and here is the figure

You seem to be showing the optimized pulse now, so Iām assuming that part is solvedā¦

(2) The GRAPE dose not achieve the desired outcome

Even though Iām not surprised, based on the implementation, youāll have to be a bit more specific. In what way does it not achieve the desired outcome?

Here is my code

I would strongly recommend not writing scripts with global variables like that. You will want to put your GRAPE implementation (specifically, the loop for iter ā¦) in a function. Global variables in Julia kill performance. Donāt get into that habit, not even for prototyping.

# Start state
global Ļ0 = s0 * s0'


Since this is a closed system, why are you working with density matrices? Not that it really matters for a two-level system, but itās something to keep in mind when scaling to real systems.

global Ī© = rand(length(t_list)) .* (2Ļ) .- Ļ

Thatās okay for playing around, but for most systems, the piecewise-constant pulses are an approximation of physically time-continuous pulses. So setting them to random values isnāt going to be physical. Also, note that the choice of guess pulse has a major impact on convergence. So, you generally want to start with something thatās physically motivated.

    # Gradient Calculation
for i in 1:length(t_list)
grad_Ī©[i] = -real(tr(Ļf_list[i]' * 1im * dt * ((tr01 + tr10) * Ļ0_list[i] - Ļ0_list[i] * (tr01 + tr10))))
end


GRAPE is known to require high-precision gradients to converge. This first-order approximation will generally not do. Have a look at J. Chem. Phys. 131, 234108 (2009) and J. Magnet. Res. 212, 412 (2011). You can, of course, play around with the first-order Taylor series you have here, but donāt be surprised if you donāt see convergence. At the very least, use a higher-order Taylor series. The state-of-the art would be the approach described in J. Chem. Phys. 143, 084113 (2015) and Quantum 6 , 871 (2022), also implemented in QuantumGradientGenerators.jl.

    # Update Pulses
lr = 0.01 - iter / 100


Even though thatās along the lines of the original GRAPE paper, there is absolutely no guarantee here that going in the direction of the gradient by exactly lr will get you to a lower error (or higher fidelity: it looks like youāre maximizing fidelity instead of the usual minimization; in any case, make sure you use the appropriate sign for the direction). You really want to determine the step width with a proper linesearch algorithm. Actually, what you want to do is get the gradient, and then feed it into any existing gradient-optimization toolbox. Preferably LBFGS: Using the quasi-Hessian that LBFGS gives you āfor freeā makes a huge difference. Nobody uses GRAPE with a plain gradient, in real applications. Again, nothing wrong with playing around with that, but donāt expect it to give very useful results.

Also, truncating the pulses in the last line can break your convergence. Usually it works, but be careful with it. You can also use LBFGS-B (where the -B is for āboundsā, which incorporates bounds in a mathematically more sound way).

Technically true, as a general statement for gradient-based optimization, but for the quantum control problem it has been shown (by Hersh Rabitzā gazillion papers) that there are no local minima if the controls are not constrained. The PWC approximation is a constraint, so this holds only if the pulses are reasonably good approximations of time-continuous fields. Adding bounds on the amplitude is definitely a constraint, and Iāve seen local traps in the case.

or it might be the case that your Hamiltonian is simply physically incapable of achieving the final state you desire.

Thatās very true. You have to check ācontrollabilityā first. But for typical quantum control problems like this one, theyāre controllable, so thatās probably not something to worry about for this example.

A typical debugging tool in such situations would be to plot your fidelity as a function of the iteration step

Thatās always useful, and a proper GRAPE implementation (with a high-precision gradient and a line search) should show monotonic convergence

If the fidelity plateaus, then probably you have reached a local minimum or your Hamiltonian does not permit the evolution you want.

Even though there are no local minima, there can be very flat areas in the optimization landscape. Itās not uncommon for the optimization to slow down and then pick up speed again. Especially if you start with a bad guess pulse. So, always give it a few 10000 iterations.

Another graphical debugging tool I would recommend is plotting out the linesearch like I do with the GRAPELinesearchAnalysis.jl package. That will tell you if your gradient is correct.

Also, consider using pre-packaged quantum control tools like:
JuliaQuantumControl/QuantumControl.jl: Julia Framework for Quantum Dynamics and Control

Not just QuantumControl.jl, but specifically GRAPE.jl. Thereās nothing wrong with implementing your own GRAPE as a learning experience, of course. In fact, I would encourage it. But, you probably donāt want to use that implementation āin productionā.

These tools are tested by many people

I wish that was true

To be fair, the Julia implementation evolved from a Fortran implementation, which has been tested by many people.

less likely to contain bugs.

That I will stand by: The GRAPE implementation should be pretty solid. And there are a few subtleties in GRAPE, some of which Iāve hinted at above. So bugs would definitely be expected in a ānaiveā implementation.

I would point @fqguo to the graphical scheme for GRAPE from Quantum 6 , 871 (2022):

Thereās also pseudo-code in that paper that might be helpful, and it describes how to calculate exact gradients.

You could do that for a toy problem with just a few time steps, for debugging. But even in this example, there are 100 time steps, and Iām not sure finite difference will be able to handle that in any reasonable amount of time. I do recommend plotting fidelity against step size, in the direction of the gradient. If it doesnāt to up for a small enough step size, the gradient is wrong.

2 Likes

Thanks for your answer. After reading it, I gained a lot from it. If I have any further questions after reading your recommendation, I hope to seek your help. Thank you very much for your help! @goerz and @Krastanov

2 Likes