$ \newcommand{tr}[0]{\operatorname{tr}} \newcommand{diag}[0]{\operatorname{diag}} \newcommand{abs}[0]{\operatorname{abs}} \newcommand{pop}[0]{\operatorname{pop}} \newcommand{aux}[0]{\text{aux}} \newcommand{opt}[0]{\text{opt}} \newcommand{tgt}[0]{\text{tgt}} \newcommand{init}[0]{\text{init}} \newcommand{lab}[0]{\text{lab}} \newcommand{rwa}[0]{\text{rwa}} \newcommand{bra}[1]{\langle#1\vert} \newcommand{ket}[1]{\vert#1\rangle} \newcommand{Bra}[1]{\left\langle#1\right\vert} \newcommand{Ket}[1]{\left\vert#1\right\rangle} \newcommand{Braket}[2]{\left\langle #1\vphantom{#2}\mid{#2}\vphantom{#1}\right\rangle} \newcommand{op}[1]{\hat{#1}} \newcommand{Op}[1]{\hat{#1}} \newcommand{dd}[0]{\,\text{d}} \newcommand{Liouville}[0]{\mathcal{L}} \newcommand{DynMap}[0]{\mathcal{E}} \newcommand{identity}[0]{\mathbf{1}} \newcommand{Norm}[1]{\lVert#1\rVert} \newcommand{Abs}[1]{\left\vert#1\right\vert} \newcommand{avg}[1]{\langle#1\rangle} \newcommand{Avg}[1]{\left\langle#1\right\rangle} \newcommand{AbsSq}[1]{\left\vert#1\right\vert^2} \newcommand{Re}[0]{\operatorname{Re}} \newcommand{Im}[0]{\operatorname{Im}} $
This first example illustrates the basic use of the GRAPE.jl
by solving a simple canonical optimization problem: the transfer of population in a two level system.
using DrWatson
@quickactivate "GRAPETests"
using QuantumControl
We consider the Hamiltonian $\op{H}_{0} = - \frac{\omega}{2} \op{\sigma}_{z}$, representing a simple qubit with energy level splitting $\omega$ in the basis $\{\ket{0},\ket{1}\}$. The control field $\epsilon(t)$ is assumed to couple via the Hamiltonian $\op{H}_{1}(t) = \epsilon(t) \op{\sigma}_{x}$ to the qubit, i.e., the control field effectively drives transitions between both qubit states.
We we will use
ϵ(t) = 0.2 * QuantumControl.Shapes.flattop(t, T=5, t_rise=0.3, func=:blackman);
"""Two-level-system Hamiltonian."""
function hamiltonian(Ω=1.0, ϵ=ϵ)
σ̂_z = ComplexF64[
1 0
0 -1
]
σ̂_x = ComplexF64[
0 1
1 0
]
Ĥ₀ = -0.5 * Ω * σ̂_z
Ĥ₁ = σ̂_x
return (Ĥ₀, (Ĥ₁, ϵ))
end;
H = hamiltonian();
The control field here switches on from zero at $t=0$ to it's maximum amplitude 0.2 within the time period 0.3 (the switch-on shape is half a Blackman pulse). It switches off again in the time period 0.3 before the final time $T=5$). We use a time grid with 500 time steps between 0 and $T$:
tlist = collect(range(0, 5, length=500));
using Plots
Plots.default(
linewidth = 3,
size = (550, 300),
legend = :right,
foreground_color_legend = nothing,
background_color_legend = RGBA(1, 1, 1, 0.8)
)
function plot_control(pulse::Vector, tlist)
plot(tlist, pulse, xlabel="time", ylabel="amplitude", legend=false)
end
plot_control(ϵ::T, tlist) where {T<:Function} = plot_control([ϵ(t) for t in tlist], tlist);
fig = plot_control(H[2][2], tlist)
First, we define a convenience function for the eigenstates.
function ket(label)
result = Dict("0" => Vector{ComplexF64}([1, 0]), "1" => Vector{ComplexF64}([0, 1]))
return result[string(label)]
end;
The physical objective of our optimization is to transform the initial state $\ket{0}$ into the target state $\ket{1}$ under the time evolution induced by the Hamiltonian $\op{H}(t)$.
objectives = [Objective(initial_state=ket(0), generator=H, target_state=ket(1))];
The full control problem includes this objective, information about the time grid for the dynamics, and the functional to be used (the square modulus of the overlap $\tau$ with the target state in this case).
using QuantumControl.Functionals: J_T_sm
problem = ControlProblem(
objectives=objectives,
tlist=tlist,
pulse_options=Dict(),
iter_stop=500,
J_T=J_T_sm,
check_convergence=res -> begin
((res.J_T < 1e-3) && (res.converged = true) && (res.message = "J_T < 10⁻³"))
end,
);
Before running the optimization procedure, we first simulate the dynamics under the guess field $\epsilon_{0}(t)$. The following solves equation of motion for the defined objective, which contains the initial state $\ket{\Psi_{\init}}$ and the Hamiltonian $\op{H}(t)$ defining its evolution.
guess_dynamics = propagate_objective(
objectives[1],
problem.tlist;
storage=true,
observables=(Ψ -> abs.(Ψ) .^ 2,)
)
2×500 Matrix{Float64}: 1.0 1.0 1.0 1.0 … 0.951457 0.951459 0.951459 0.0 7.73456e-40 2.03206e-11 2.96638e-10 0.0485427 0.048541 0.048541
function plot_population(pop0::Vector, pop1::Vector, tlist)
fig = plot(tlist, pop0, label="0", xlabel="time", ylabel="population")
plot!(fig, tlist, pop1; label="1")
end;
fig = plot_population(guess_dynamics[1, :], guess_dynamics[2, :], tlist)
In the following we optimize the guess field $\epsilon_{0}(t)$ such that the intended state-to-state transfer $\ket{\Psi_{\init}} \rightarrow \ket{\Psi_{\tgt}}$ is solved.
The GRAPE package performs the optimization by calculating the gradient of $J_T$ with respect to the values of the control field at each point in time. This gradient is then fed into a backend solver that calculates an appropriate update based on that gradient.
By default, this backend is LBFGSB.jl, a wrapper around the true and tested L-BFGS-B Fortran library. L-BFGS-B is a pseudo-Hessian method: it efficiently estimates the second-order Hessian from the gradient information. The search direction determined from that Hessian dramatically improves convergence compared to using the gradient directly as a search direction. The L-BFGS-B method performs its own linesearch to determine how far to go in the search direction.
It can be quite instructive to see how the improvement in the pseudo-Hessian search direction compares to the gradient, how the linesearch finds an appropriate step width. For this purpose, we have a GRAPELinesearchAnalysis package that automatically generates plots in every iteration of the optimization showing the linesearch behavior
using GRAPELinesearchAnalysis
We feed this into the optimization as part of the info_hook
.
function store_pulses(wrk, iteration, _...)
L = length(wrk.controls)
ϵ_opt = reshape(wrk.pulsevals, L, :)
opt_pulses = [QuantumControl.Controls.discretize_on_midpoints(ϵ_opt[l, :], tlist) for l=1:L]
return Tuple(opt_pulses)
end
store_pulses (generic function with 3 methods)
opt_result_LBFGSB = @optimize_or_load(
datadir("TLS", "opt_result_LBFGSB.jld2"),
problem,
method = :grape,
force = true,
info_hook = chain_infohooks(
GRAPELinesearchAnalysis.plot_linesearch(datadir("TLS", "Linesearch", "LBFGSB")),
QuantumControl.GRAPE.print_table,
store_pulses
)
);
iter. J_T |∇J_T| ΔJ_T FG(F) secs 0 9.51e-01 5.40e-02 n/a 1(0) 0.0 1 4.45e-02 6.57e-02 -9.07e-01 4(0) 0.2 2 1.89e-02 4.38e-02 -2.56e-02 3(0) 5.6 3 5.75e-03 2.33e-02 -1.31e-02 1(0) 5.7 4 1.03e-05 1.04e-03 -5.74e-03 1(0) 11.4
When going through this tutorial locally, the generated images for the linesearch can be found in docs/TLS/Linesearch/LBFGSB
.
opt_result_LBFGSB
GRAPE Optimization Result ------------------------- - Started at 2022-10-04T11:50:17.517 - Number of objectives: 1 - Number of iterations: 4 - Number of pure func evals: 0 - Number of func/grad evals: 10 - Value of functional: 1.02765e-05 - Reason for termination: J_T < 10⁻³ - Ended at 2022-10-04T11:50:41.448 (23 seconds, 931 milliseconds)
opt_result_LBFGSB.records
5-element Vector{Tuple}: ([-2.7755575615628915e-18, 0.0004498836583447907, 0.001269014020426157, 0.0025428819003394235, 0.004322359651711154, 0.006669889670207852, 0.009655198659934522, 0.013350337952244546, 0.017824288406719662, 0.02313739119759512 … 0.023137391197595336, 0.017824288406719763, 0.0133503379522446, 0.00965519865993452, 0.006669889670207792, 0.0043223596517110874, 0.002542881900339493, 0.001269014020426182, 0.00044988365834480184, -2.7755575615628915e-18],) ([-0.6834731103016163, -0.6798273031594231, -0.6757439492371748, -0.6711378854522644, -0.6659585740144947, -0.6601439139030224, -0.6536245265916177, -0.6463287157412629, -0.6381878623405248, -0.6291419940033012 … -0.6291419940033163, -0.638187862340538, -0.6463287157412765, -0.6536245265916303, -0.660143913903035, -0.6659585740145062, -0.6711378854522161, -0.6757439492371851, -0.6798273031594334, -0.6834731103016258],) ([-0.7769734160771408, -0.7731196947454784, -0.7688190216886603, -0.7639862310337592, -0.7585707840483094, -0.752510580693903, -0.7457362454356306, -0.7381760870277739, -0.7297614937563046, -0.7204325028468093 … -0.7204325028468253, -0.7297614937563184, -0.738176087027788, -0.7457362454356435, -0.752510580693916, -0.7585707840483212, -0.7639862310337029, -0.7688190216886707, -0.7731196947454888, -0.7769734160771504],) ([-0.767700090111739, -0.7647233262624439, -0.7613001693596176, -0.7573451798356391, -0.7528075492542776, -0.7476249124968329, -0.7417276342625901, -0.7350437696099603, -0.7275044600144923, -0.7190495036382046 … -0.7190495036382144, -0.7275044600144999, -0.7350437696099681, -0.7417276342625968, -0.7476249124968395, -0.7528075492542831, -0.7573451798355769, -0.7613001693596216, -0.764723326262448, -0.7677000901117422],) ([-0.771066734951367, -0.7677770518084474, -0.7640407697576649, -0.7597725466154319, -0.754921670076591, -0.7494258696720285, -0.743215603018955, -0.7362190160880254, -0.7283673389649812, -0.7196004558062695 … -0.7196004558062824, -0.7283673389649918, -0.7362190160880363, -0.7432156030189647, -0.7494258696720383, -0.7549216700765996, -0.7597725466153727, -0.7640407697576721, -0.7677770518084547, -0.7710667349513735],)
We can plot the optimized field:
fig = plot_control(opt_result_LBFGSB.optimized_controls[1], tlist)
plot!(fig, tlist, QuantumControl.Controls.discretize(opt_result_LBFGSB.records[1][1], tlist); ls=:dash)
plot!(fig, tlist, QuantumControl.Controls.discretize(opt_result_LBFGSB.records[2][1], tlist); ls=:dash)
plot!(fig, tlist, QuantumControl.Controls.discretize(opt_result_LBFGSB.records[3][1], tlist); ls=:dash)
plot!(fig, tlist, QuantumControl.Controls.discretize(opt_result_LBFGSB.records[4][1], tlist); ls=:dash)
This notebook was generated using Literate.jl.