GRAPE Demos
using Piccolo
using Optim
using LinearAlgebra
using SparseArrays
using CairoMakie
# useful
const ⊗ = kronkron (generic function with 61 methods)Goals
- Learn the quantum isomorphisms that map variables to real-valued state vectors
- Study how gradient descent and Newton's method can be used to optimize quantum controls.
I. Isomorphisms
Piccolo isomorphisms
- The standard quantum states are kets, $|\psi\rangle$, and Unitaries, $U$.
- Open quantum system require density matrices, $\rho$, and quantum channels, $\Phi$.
- Standard quantum states have an open system counterpart,
\[\begin{align} \text{closed} &\longrightarrow \text{open} \\ \hline |\psi\rangle &\longrightarrow |\psi\rangle \langle \psi | \\ U &\longrightarrow U \cdot U^\dagger \end{align}\]
🚧 ⚠️ If you are seeing a lot of boxes like Ũ⃗, it is very useful to have the JuliaMono fonts for Piccolo. Install and change the default font family.
# Ok, so it's not technically a wavefunction
ψ = [1; 2] + im * [3; 4]
ψ̃ = ket_to_iso(ψ)4-element Vector{Int64}:
1
2
3
4iso_to_ket(ψ̃)2-element Vector{Complex{Int64}}:
1 + 3im
2 + 4im# We often need to convert a complex matrix U to a real vector, Ũ⃗.
U = [1 5; 2 6] + im * [3 7; 4 8]2×2 Matrix{Complex{Int64}}:
1+3im 5+7im
2+4im 6+8imRemember what you learned about Julia arrays! Why would I write the matrix this way?
Ũ⃗ = operator_to_iso_vec(U)8-element Vector{Int64}:
1
2
3
4
5
6
7
8iso_vec_to_operator(Ũ⃗)2×2 Matrix{Complex{Int64}}:
1+3im 5+7im
2+4im 6+8imPhysics check: What's an efficiency that we might be able to leverage when storing $\rho$ that you don't see here?
# Warning: The isomorphism `density_to_iso_vec` is not the same as `operator_to_iso_vec`.
ρ = [1 2; 3 4] + im * [5 6; 7 8]
ρ̃⃗ = density_to_iso_vec(ρ)8-element Vector{Int64}:
1
3
2
4
5
7
6
8Exercise
- Just how big are these vectors for a single qubit state? A two qubit state?
- What about quantum channels?
II. Quantum dynamics
Quantum systems
First up, we are going to look at some dynamics convenience functions in Piccolo.
- Let's flip a qubit from the ground state to the excited state.
- Introduce the isomorphisms that make quantum dynamics real-valued.
- Use PiccoloQuantumObjects to make a quantum system.
- Use a rollout to integrate the quantum system forward in time.
\[H(u(t)) = \underbrace{u_1(t) XI + u_2(t) YI}_\text{qubit 1} + \underbrace{u_3(t) IX + u_4(t) IY}_\text{qubit 2} + \underbrace{u_5(t) XX}_\text{coupling}\]
H_drives = [
PAULIS.X ⊗ PAULIS.I,
PAULIS.Y ⊗ PAULIS.I,
PAULIS.I ⊗ PAULIS.X,
PAULIS.I ⊗ PAULIS.Y,
PAULIS.X ⊗ PAULIS.X
]
system = QuantumSystem(H_drives)QuantumSystem: levels = 4, n_drives = 5- Quantum systems contain the operators we need, including the real valued versions.
get_drift(system)4×4 SparseArrays.SparseMatrixCSC{ComplexF64, Int64} with 0 stored entries:
⋅ ⋅ ⋅ ⋅
⋅ ⋅ ⋅ ⋅
⋅ ⋅ ⋅ ⋅
⋅ ⋅ ⋅ ⋅ - Quick check: What do we expect to see?
get_drives(system)[1]4×4 SparseArrays.SparseMatrixCSC{ComplexF64, Int64} with 4 stored entries:
⋅ ⋅ 1.0+0.0im ⋅
⋅ ⋅ ⋅ 1.0+0.0im
1.0+0.0im ⋅ ⋅ ⋅
⋅ 1.0+0.0im ⋅ ⋅ system.H(randn(system.n_drives))4×4 SparseArrays.SparseMatrixCSC{ComplexF64, Int64} with 12 stored entries:
⋅ 1.03208-0.0300109im … -1.35906+0.0im
1.03208+0.0300109im ⋅ 0.610696-1.0478im
0.610696+1.0478im -1.35906+0.0im 1.03208-0.0300109im
-1.35906+0.0im 0.610696+1.0478im ⋅ - Quick check: How big will this operator be?
system.G(randn(system.n_drives))8×8 SparseArrays.SparseMatrixCSC{Float64, Int64} with 40 stored entries:
⋅ -0.748511 -0.45111 … 1.12976 -1.67731 -0.0282423
0.748511 ⋅ ⋅ ⋅ -0.0282423 -1.67731
0.45111 ⋅ ⋅ -0.0282423 ⋅ 1.12976
⋅ 0.45111 0.748511 -1.67731 1.12976 ⋅
⋅ -1.12976 1.67731 -0.748511 -0.45111 ⋅
-1.12976 ⋅ 0.0282423 … ⋅ ⋅ -0.45111
1.67731 0.0282423 ⋅ ⋅ ⋅ -0.748511
0.0282423 1.67731 -1.12976 0.45111 0.748511 ⋅ - We can use a system to perform a rollout.
# Timing information (e.g. 20 ns superconducting qubit gate)
T = 40
Δt = 0.5
timesteps = fill(Δt, T)40-element Vector{Float64}:
0.5
0.5
0.5
0.5
0.5
0.5
0.5
0.5
0.5
0.5
⋮
0.5
0.5
0.5
0.5
0.5
0.5
0.5
0.5
0.5# Controls
controls = randn(system.n_drives, T + 1);5×41 Matrix{Float64}:
-0.192976 1.51629 -1.27529 … -0.809342 0.791013 0.771139
1.68979 -0.279634 -0.767858 -0.0849856 0.404975 -0.0782606
0.315975 -0.23213 0.375845 -0.301382 0.582139 0.46862
0.166577 -0.0600409 -1.25386 0.542118 -1.86174 0.836567
-0.260813 -1.46123 1.65408 0.0265883 0.00184781 0.95562unitary_rollout(controls, timesteps, system)32×41 Matrix{Float64}:
1.0 0.641823 0.457662 … -0.011918 0.127227 0.120502
0.0 0.0435325 0.329316 0.372037 0.220352 0.399733
0.0 0.750614 0.366133 0.544586 0.228119 0.485619
0.0 0.0722386 0.32301 -0.0348248 0.142998 -0.374511
0.0 0.00457179 -0.347449 0.564748 0.747483 0.348779
0.0 -0.102446 0.343685 … -0.0946765 0.125117 -0.498224
0.0 0.0836819 -0.222096 -0.470652 -0.511475 -0.189459
0.0 0.00305514 0.394387 0.119336 0.163644 0.209377
0.0 -0.0654048 0.290355 0.22644 0.103573 0.368296
1.0 0.641824 0.405682 0.170323 0.367169 0.00992179
⋮ ⋱ ⋮
0.0 -0.102446 -0.373989 -0.379692 -0.0704556 -0.702964
0.0 0.0776745 -0.317856 0.587627 0.41328 0.698375
0.0 -0.714618 -0.274229 … 0.281902 0.215013 -0.0745987
0.0 -0.0654048 0.314768 0.0877565 0.197407 -0.128976
1.0 0.641823 0.376557 -0.523956 -0.597694 -0.440975
0.0 0.22295 0.448272 -0.14678 -0.0760651 -0.0885984
0.0 0.0826485 -0.256464 0.245497 0.0765019 0.147365
0.0 -0.102543 -0.410124 … 0.224435 0.208211 0.481541
0.0 0.00457179 0.384667 0.401011 0.575996 0.184841# Entangling gate
U_goal = GATES.CX
# How'd we do?
println("ℱ = ", unitary_rollout_fidelity(U_goal, controls, timesteps, system))ℱ = 0.04016281101271073We have all the pieces we need to solve!
Let's put Piccolo to work.
# Piccolo (we'll learn more about this later)
prob = UnitarySmoothPulseProblem(system, U_goal, T, Δt);DirectTrajOptProblem
timesteps = 40
duration = 19.5
variable names = (:Ũ⃗, :a, :da, :dda, :Δt)
knot point dimension = 48
# save these initial controls for later
a_init = prob.trajectory.a
plot(prob.trajectory, :a)
solve!(
prob,
max_iter=20, print_level=1, verbose=false, options=IpoptOptions(eval_hessian=false)
)
ℱ = unitary_rollout_fidelity(prob.trajectory, system)
println("The fidelity is ", ℱ)
******************************************************************************
This program contains Ipopt, a library for large-scale nonlinear optimization.
Ipopt is released as open source code under the Eclipse Public License (EPL).
For more information visit https://github.com/coin-or/Ipopt
******************************************************************************
The fidelity is 0.9791586907388493a_final = prob.trajectory.a
plot(prob.trajectory, :a)
III. GRAPE
The GRAPE algorithm comes from NMR in 2004, and there is a Julia version. We'll reproduce GRAPE in this example.
# We work with timesteps between knot points
timesteps = fill(Δt, T)
# Let's use our previous function to compute the fidelity
GRAPE(controls) = abs(1 - unitary_rollout_fidelity(U_goal, controls, timesteps, system))GRAPE (generic function with 1 method)Automatic differentiation
- It's quick to test! Compare different algorithms, e.g.,
Newton(),GradientDescent(),LBFGS() - If you switch from gradient descent to a quasi-Newton method, you get to write another paper.
result_GRAPE = optimize(GRAPE, collect(a_init), LBFGS()) * Status: success
* Candidate solution
Final objective value: 1.332268e-15
* Found with
Algorithm: L-BFGS
* Convergence measures
|x - x'| = 9.21e-09 ≰ 0.0e+00
|x - x'|/|x'| = 8.31e-09 ≰ 0.0e+00
|f(x) - f(x')| = 6.66e-16 ≰ 0.0e+00
|f(x) - f(x')|/|f(x')| = 5.00e-01 ≰ 0.0e+00
|g(x)| = 8.46e-09 ≤ 1.0e-08
* Work counters
Seconds run: 68 (vs limit Inf)
Iterations: 59
f(x) calls: 176
∇f(x) calls: 176
a_GRAPE = Optim.minimizer(result_GRAPE)
println("The fidelity is ", unitary_rollout_fidelity(U_goal, a_GRAPE, timesteps, system))The fidelity is 0.9999999999999987- What do we think we'll see here?
series(cumsum(timesteps), a_GRAPE)
Analytic gradients
Calculus practice
- We can combine forward and backward rollouts to compute the gradients,
\[\begin{align} \frac{\partial U(T)}{\partial u_k(t)} &= U(T, t) (-i H_k \Delta t) U(t) \\ \Rightarrow \langle\psi_\text{goal} | \frac{\partial U(T)}{\partial u_k(t)} |\psi_\text{init.}\rangle &= -i \Delta t \langle\psi_\text{goal}^\text{bwd.}(t) | H_k |\psi_\text{init.}^\text{fwd.}(t) \rangle. \end{align}\]
Exercise
- Implement gradient descent using the analytic gradients.
- Sometimes, there are insights you can only get by opening up the black box, e.g. d-GRAPE.
IV. Function Spaces
- Pick a function basis for the controls and optimize the coefficients. Some choices are trig functions or Slepians.
- Our optimization parameters are now coefficients of the basis,
\[u(t) = u_0 + \sum_{j=1}^{n} c_j a_j(t)\]
- The modes $a_j(t)$ stay fixed, and the coefficients $c_j$ are optimized.
# First n = 5 entries in a Fourier series, including the constant term
n = 5
fourier_series = [cos.(π * j * (0:T-1) / T .- π/2) for j in 0:n-1]
function get_controls(coefficients)
a(c) = sum(cⱼ * aⱼ for (cⱼ, aⱼ) in zip(c, fourier_series))
return stack([a(c) for c in eachrow(coefficients)], dims=1)
end
function GRAFS(coefficients)
controls = get_controls(coefficients)
return abs(1 - unitary_rollout_fidelity(U_goal, controls, timesteps, system))
endGRAFS (generic function with 1 method)c_init = rand(system.n_drives, n)
result_GRAFS = optimize(GRAFS, c_init, LBFGS()) * Status: success
* Candidate solution
Final objective value: 0.000000e+00
* Found with
Algorithm: L-BFGS
* Convergence measures
|x - x'| = 7.44e-09 ≰ 0.0e+00
|x - x'|/|x'| = 5.32e-09 ≰ 0.0e+00
|f(x) - f(x')| = 0.00e+00 ≤ 0.0e+00
|f(x) - f(x')|/|f(x')| = NaN ≰ 0.0e+00
|g(x)| = 3.54e-08 ≰ 1.0e-08
* Work counters
Seconds run: 29 (vs limit Inf)
Iterations: 91
f(x) calls: 379
∇f(x) calls: 379
a_GRAFS = Optim.minimizer(result_GRAFS)
println("The fidelity is ", 1 - unitary_rollout_fidelity(U_goal, a_GRAFS, timesteps, system))The fidelity is 0.993771770586151series(cumsum(timesteps), get_controls(a_GRAFS))
- These shapes are a lot nicer! But performance depends a lot on the expressivity and initial condition.
c_init = randn(system.n_drives, n)
result_GRAFS_2 = optimize(GRAFS, c_init, LBFGS()) * Status: success
* Candidate solution
Final objective value: 0.000000e+00
* Found with
Algorithm: L-BFGS
* Convergence measures
|x - x'| = 4.81e-13 ≰ 0.0e+00
|x - x'|/|x'| = 2.21e-13 ≰ 0.0e+00
|f(x) - f(x')| = 0.00e+00 ≤ 0.0e+00
|f(x) - f(x')|/|f(x')| = NaN ≰ 0.0e+00
|g(x)| = 9.65e-08 ≰ 1.0e-08
* Work counters
Seconds run: 29 (vs limit Inf)
Iterations: 82
f(x) calls: 365
∇f(x) calls: 365
a_GRAFS_2 = Optim.minimizer(result_GRAFS_2)
println("The fidelity is ", 1 - unitary_rollout_fidelity(U_goal, a_GRAFS_2, timesteps, system))The fidelity is 0.6996978325918379f = Figure()
ax = Axis(f[1,1])
series!(ax, cumsum(timesteps), get_controls(a_GRAFS))
ax = Axis(f[2,1])
series!(ax, cumsum(timesteps), get_controls(a_GRAFS_2))
f
Exercise: A filtering approach
- Pass the controls through a spectral filter: Look up Slepians and consider how to bound the bandwidth by choice of basis.
- How might we shape the bandwidth of the controls? (Remember, we can just rely on automatic differentiation!)
V. States in costs
Exercise:
- Let's switch to a transmon, which has more than two levels and can be leaky.
\[H(u(t)) = \tfrac{1}{2} \eta a^\dagger a^\dagger a a + u_1(t) (a + a^\dagger) - i u_2(t) (a - a^\dagger)\]
The optimizer can exploit the higher levels!
Add a leakage penalty to a guard state. Notice that working with states can be awkward.
T = 40
Δt = 0.25
timesteps = fill(Δt, T)
function Transmon(n)
a = annihilate(n)
x = a + a'
p = -im * (a - a')
η = 0.1
return QuantumSystem(1/2 * a'a'a*a, [x, p])
end
transmon_2 = Transmon(2)
transmon_4 = Transmon(4)QuantumSystem: levels = 4, n_drives = 2function TransmonGRAFS(
goal::AbstractPiccoloOperator, coefficients, timesteps, sys::AbstractQuantumSystem
)
controls = get_controls(coefficients)
return abs(1 - unitary_rollout_fidelity(goal, controls, timesteps, sys))
endTransmonGRAFS (generic function with 1 method)- Quick aside: Embedded operators
U_emb(n) = EmbeddedOperator(GATES.X, 1:2, n)U_emb (generic function with 1 method)U_emb(4).operator4×4 Matrix{ComplexF64}:
0.0+0.0im 1.0+0.0im 0.0+0.0im 0.0+0.0im
1.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im
0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0im
0.0+0.0im 0.0+0.0im 0.0+0.0im 0.0+0.0imunembed(U_emb(4))2×2 Matrix{ComplexF64}:
0.0+0.0im 1.0+0.0im
1.0+0.0im 0.0+0.0imsys2, U2 = Transmon(2), U_emb(2)
c_init = randn(sys2.n_drives, n)
result_GRAFS_3 = optimize(a -> TransmonGRAFS(U2, a, timesteps, sys2), c_init, LBFGS()) * Status: success
* Candidate solution
Final objective value: 8.881784e-16
* Found with
Algorithm: L-BFGS
* Convergence measures
|x - x'| = 2.16e-08 ≰ 0.0e+00
|x - x'|/|x'| = 1.02e-08 ≰ 0.0e+00
|f(x) - f(x')| = 2.44e-15 ≰ 0.0e+00
|f(x) - f(x')|/|f(x')| = 2.75e+00 ≰ 0.0e+00
|g(x)| = 2.93e-10 ≤ 1.0e-08
* Work counters
Seconds run: 0 (vs limit Inf)
Iterations: 11
f(x) calls: 34
∇f(x) calls: 34
a_GRAFS_3 = get_controls(Optim.minimizer(result_GRAFS_3))
println("The fidelity is ", unitary_rollout_fidelity(U2, a_GRAFS_3, timesteps, sys2))The fidelity is 1.0000000000000009- Quick check: What might happen now?
println(
"The fidelity is ", unitary_rollout_fidelity(U_emb(4), a_GRAFS_3, timesteps, Transmon(4))
)The fidelity is 0.08858572794785598TODO:
- Add an L2 penalty to states that are not in the computational basis.
- Use a modified GRAPE cost to penalize leakage while maintaining fidelity.
- Study how leakage and fidelity change with the penalty.
- Study how the anharmonicity η affects leakage.