# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.1.0
#     language: julia
#     name: julia-1.1
# ---

# +
using Plots
# pyplot()
plotly()

using QSimulator
# -

# # Transmon Ramsey Oscillations
#
# Starting from a coherence between two energy eigenstates we expect oscilltions at the detuning between the energy levels.  We'll work with natural units of GHz and ns for numerical stability reasons. We'll start with a single 3 level transmon in the lab frame with a 5 GHz qubit frequency and -200 MHz anharmonicity.

# we create a specific QSystem
q0 = DuffingTransmon("q0", 3, DuffingSpec(5, -0.2))
# we can ask for the Hamiltonian of an QSystem as a Matrix
hamiltonian(q0)

# a CompositeQSystems is a tensor product structure of QSystem's and is what all the solvers are built around
cqs = CompositeQSystem([q0]);
add_hamiltonian!(cqs, q0)
hamiltonian(cqs)

# evolve an initial superposition state for 1 ns
times = range(0, stop=1, length=201) |> collect
ψ0 = (1/sqrt(2)) * ComplexF64[1; 1; 0]
ψs = unitary_state(cqs, times, ψ0);

# plot the projection on to the initial state and we expect to see 5 GHz oscillations
signal = Float64[real(ψ0'*ψ) for ψ in ψs]
expected = 0.5 .+ 0.5*cos.(2π*5 * times)
p = plot(times, signal, linewidth=2, label="simulated")
plot!(p, times, expected, label="ideal")
xlabel!(p, "Time (ns)")

# # Rabi Oscillations
#
# Driving Rabi oscillations in the lab frame is a good example of a parametric time dependent Hamiltonian. The drive electric field couples to the transmon dipole or $X$ operator.
#
# ## Constant Drive

qubit_freq = 5.0
q0 = DuffingTransmon("q0", 3, DuffingSpec(5, -0.2))
cqs = CompositeQSystem([q0]);
add_hamiltonian!(cqs, q0)
add_hamiltonian!(cqs, dipole_drive(q0, t -> 0.02*cos(2π*qubit_freq * t)), q0);
ψ_init = ComplexF64[1; 0; 0]
times = range(0, stop=100, length=101) |> collect
ψs = unitary_state(cqs, times, ψ_init);

p = plot(times, [abs2(s[1]) for s in ψs], label="Ground State Pop.")
plot!(p, times, [abs2(s[2]) for s in ψs], label="Excited State Pop.")
xlabel!(p, "Time (ns)")

# ## Variable Amplitude Gaussian Pulse

# +
# write a helper function that returns the drive Hamiltonian at a particular point in time
function gaussian(pulse_length, pulse_freq, t; cutoff=2.5)
    σ = pulse_length/2/cutoff
    pulse = exp(-0.5*((t-pulse_length/2)/σ)^2)
    pulse * cos(2π*pulse_freq * t)
end

function flat(pulse_freq,  t)
   cos(2π*pulse_freq * t)
end

# +
ψ0 = ComplexF64[1; 0; 0]
states_flat = []
states_gaussian = []
amps = 0.1 * range(0, stop=1, length=51)
pulse_length = 25.0
qubit_freq = 5.0

for amp = amps
    # first do flat pulse
    # three level transmon in the lab frame
    q0 = DuffingTransmon("q0", 3, DuffingSpec(5, -0.2))
    cqs = CompositeQSystem([q0]);
    add_hamiltonian!(cqs, q0)
    add_hamiltonian!(cqs, dipole_drive(q0, t -> amp*flat(qubit_freq, t)), q0);
    ψs = unitary_state(cqs, [0, pulse_length], ψ0)
    push!(states_flat, ψs[end])

    # now gaussian
    cqs = CompositeQSystem([q0]);
    add_hamiltonian!(cqs, q0)
    add_hamiltonian!(cqs, dipole_drive(q0, t -> amp*gaussian(pulse_length, qubit_freq, t)), q0);
    ψs = unitary_state(cqs, [0, pulse_length], ψ0)
    push!(states_gaussian, ψs[end])
end
# -

p1 = plot()
for ct = 0:2
    plot!(p1, amps*1e3, [abs2(s[ct+1]) for s in states_flat], label="$ct State Pop.")
end
xlabel!(p1, "Nutation Strength (MHz)")
title!(p1, "Flat Pulse")
p2 = plot()
for ct = 0:2
    plot!(p2, amps*1e3, [abs2(s[ct+1]) for s in states_gaussian], label="$ct State Pop.")
end
xlabel!(p2, "Peak Nutation Strength (MHz)")
title!(p2, "Gaussian Pulse")
plot(p1,p2, layout=(1,2), size=(800,400))

# # Two Qubit  Gates - Parametric Gates in the Lab Frame
#
# $$ \mathcal{H}(t) = \omega_0 \hat{n}_0 + \Delta_0 \Pi_{2_0} + \omega_1(t)\hat{n}_1 + \Delta_1 \Pi_{2_1} + gX_0X_1$$
#
# ## iSWAP

# +
# parameters from Blue Launch paper
q0 = DuffingTransmon("q0", 3, DuffingSpec(3.94015, -0.1807))
q1 = PerturbativeTransmon("q1", 3, TransmonSpec(.172, 12.71, 3.69))

freqs = 114:0.5:124
times = 0.0:5:600

dims = (3,3)
basis = TensorProductBasis(dims)
ψ₀₁ = TensorProductBasisState(basis, (0,1))
ψ₁₀ = TensorProductBasisState(basis, (1,0))

ψ₀ = vec(ψ₁₀) # start in 10 state
pop₀₁ = []
pop₁₀ = []

amp = 0.323

for freq = freqs
    cqs = CompositeQSystem([q0, q1])
    add_hamiltonian!(cqs, q0)
    add_hamiltonian!(cqs, 0.006*X([q0, q1]), [q0,q1])
    add_hamiltonian!(cqs, parametric_drive(q1, t -> amp*sin(2π*freq/1e3*t)), q1)
    ψs = unitary_state(cqs, times, ψ₀);
    push!(pop₀₁, [abs2(ψ[QSimulator.index(ψ₀₁)]) for ψ in ψs])
    push!(pop₁₀, [abs2(ψ[QSimulator.index(ψ₁₀)]) for ψ in ψs])
end
# -

p1 = contour(freqs, times, cat(pop₀₁..., dims=2), fill=true)
xlabel!(p1, "Frequency (MHz)")
ylabel!(p1, "Time (ns)")
p2 = contour(freqs, times, cat(pop₁₀..., dims=2), fill=true)
xlabel!(p2, "Frequency (MHz)")
plot(p1,p2, layout=(1,2), size=(600,300))

# +
# Look more finely at a slice along time to show we are getting full contrast and look at lab frame jaggedness
q0 = DuffingTransmon("q0", 3, DuffingSpec(3.94015, -0.1807))
q1 = PerturbativeTransmon("q1", 3, TransmonSpec(.172, 12.71, 3.69))

# should get an iSWAP interaction at ≈ 122 MHz
freq = 118.5/1e3
amp = 0.323
times = 0:200

dims = (3,3)
basis = TensorProductBasis(dims)
ψ₀₁ = TensorProductBasisState(basis, (0,1))
ψ₁₀ = TensorProductBasisState(basis, (1,0))

ψ₀ = vec(ψ₁₀) # start in 10 state
pop₀₁ = []
pop₁₀ = []

cqs = CompositeQSystem([q0, q1])
add_hamiltonian!(cqs, q0)
add_hamiltonian!(cqs, 0.006*X([q0, q1]), [q0,q1])
add_hamiltonian!(cqs, parametric_drive(q1, t -> amp*sin(2π*freq*t)), q1)
ψs = unitary_state(cqs, times, ψ₀);
pop₀₁ = [abs2(ψ[QSimulator.index(ψ₀₁)]) for ψ in ψs]
pop₁₀ = [abs2(ψ[QSimulator.index(ψ₁₀)]) for ψ in ψs];
# -

p = plot(times, pop₀₁, label="01")
plot!(p, times, pop₁₀, label="10")

# # CZ02

# +
# parameters from Blue Launch paper
q0 = DuffingTransmon("q0", 3, DuffingSpec(3.94015, -0.1807))
q1 = PerturbativeTransmon("q1", 3, TransmonSpec(.172, 12.71, 3.69))

freqs = 90:1:115
times = 0.0:10:300

dims = (3,3)
basis = TensorProductBasis(dims)
ψ₁₁ = TensorProductBasisState(basis, (1,1))
ψ₀₂ = TensorProductBasisState(basis, (0,2))
ψ₀ = vec(ψ₁₁) # start in 11 state
pop₁₁ = []
pop₀₂ = []

amp = 0.245

for freq = freqs
    cqs = CompositeQSystem([q0, q1])
    add_hamiltonian!(cqs, q0)
    add_hamiltonian!(cqs, 0.006*X([q0, q1]), [q0,q1])
    add_hamiltonian!(cqs, parametric_drive(q1, t -> amp*sin(2π*freq/1e3*t)), q1)
    ψs = unitary_state(cqs, times, ψ₀);
    push!(pop₁₁, [abs2(ψ[QSimulator.index(ψ₁₁)]) for ψ in ψs])
    push!(pop₀₂, [abs2(ψ[QSimulator.index(ψ₀₂)]) for ψ in ψs])
end
pop₁₁ = cat(pop₁₁..., dims=2)
pop₀₂ = cat(pop₀₂..., dims=2);
# -

p1 = contour(freqs, times, pop₁₁, fill=true)
xlabel!(p1, "Frequency (MHz)")
ylabel!(p1, "Time (ns)")
p2 = contour(freqs, times, pop₀₂, fill=true)
xlabel!(p2, "Frequency (MHz)")
plot(p1,p2, layout=(1,2), size=(600,300))

# # Two Qubit  Gates - Parametric Gates in the Rotating Frame
#
# We can move into the doubly rotating frame. The dipole coupling becomes time dependent and we discard the flip-flip (flop-flop) terms in the Hamiltonian.
#
# $$ \mathcal{H}(t) =  \Delta_0 \Pi_{2_0} + \omega_1(t)\hat{n}_1 + \Delta_1 \Pi_{2_1} - \omega_1(0)\hat{n}_1 + e^{i\delta t}\sigma_+\sigma_- + e^{-i\delta t}\sigma_-\sigma_+$$

# +
q0 = DuffingTransmon("q0", 3, DuffingSpec(3.94015, -0.1807))
q1 = PerturbativeTransmon("q1", 3, TransmonSpec(.172, 12.71, 3.69))

# should get an iSWAP interaction at ≈ 122 MHz
freq = 118.5/1e3
amp = 0.323
times = 0:200

dims = (3,3)
basis = TensorProductBasis(dims)
ψ₀₁ = TensorProductBasisState(basis, (0,1))
ψ₁₀ = TensorProductBasisState(basis, (1,0))
ψ₀ = vec(ψ₁₀) # start in 10 state

cqs = CompositeQSystem([q0, q1])
add_hamiltonian!(cqs, q0)
# add rotating frame Hamiltonian shifts
add_hamiltonian!(cqs, -spec(q0).frequency*number(q0), q0)
q1_freq = hamiltonian(q1, 0.0)[2,2]
add_hamiltonian!(cqs, -q1_freq*number(q1), q1)

diff_freq = spec(q0).frequency - q1_freq
add_hamiltonian!(cqs, t -> 0.006*.5*X_Y([q0, q1], [diff_freq*t, 0.0]), [q0,q1])
add_hamiltonian!(cqs, parametric_drive(q1, t -> amp*sin(2π*freq*t)), q1)

ψs = unitary_state(cqs, times, ψ₀);
pop₀₁ = [abs2(ψ[QSimulator.index(ψ₀₁)]) for ψ in ψs]
pop₁₀ = [abs2(ψ[QSimulator.index(ψ₁₀)]) for ψ in ψs];
# -

p = plot(times, pop₀₁, label="01")
plot!(p, times, pop₁₀, label="10")

# +
q0 = DuffingTransmon("q0", 3, DuffingSpec(3.94015, -0.1807))
q1 = PerturbativeTransmon("q1", 3, TransmonSpec(.172, 12.71, 3.69))

# should get an CZ02 interaction at ≈ 115 MHz
freq = 102.5/1e3
amp = 0.245
times = 0:0.5:200

dims = (3,3)
basis = TensorProductBasis(dims)
ψ₁₁ = TensorProductBasisState(basis, (1,1))
ψ₀₂ = TensorProductBasisState(basis, (0,2))
ψ₀ = vec(ψ₁₁) # start in 11 state

cqs = CompositeQSystem([q0, q1])
add_hamiltonian!(cqs, q0)
# add rotating frame Hamiltonian shifts
add_hamiltonian!(cqs, -spec(q0).frequency*number(q0), q0)
q1_freq = hamiltonian(q1, 0.0)[2,2]
add_hamiltonian!(cqs, -q1_freq*number(q1), q1)

diff_freq = spec(q0).frequency - q1_freq
add_hamiltonian!(cqs, t -> 0.006*.5*X_Y([q0, q1], [diff_freq*t, 0.0]), [q0,q1])
add_hamiltonian!(cqs, parametric_drive(q1, t -> amp*sin(2π*freq*t)), q1)

ψs = unitary_state(cqs, times, ψ₀);
pop₁₁ = [abs2(ψ[QSimulator.index(ψ₁₁)]) for ψ in ψs]
pop₀₂ = [abs2(ψ[QSimulator.index(ψ₀₂)]) for ψ in ψs];
# -

p = plot(times, pop₁₁, label="11")
plot!(p, times, pop₀₂, label="02")

# # Two Qubit  Gates - Soft Shoulders
#
# Look at how soft shoulders distort the pulse shape.
#

# +
using SpecialFunctions: erf

q0 = DuffingTransmon("q0", 3, DuffingSpec(3.94015, -0.1807))
q1 = PerturbativeTransmon("q1", 3, TransmonSpec(.172, 12.71, 3.69))

# should get an CZ02 interaction at ≈ 115 MHz
freqs = 1e-3*(90:1:115)
amp = 0.245
risetime = 50
times = (2*risetime):10:500

dims = (3,3)
basis = TensorProductBasis(dims)
ψ₁₁ = TensorProductBasisState(basis, (1,1))
ψ₀₂ = TensorProductBasisState(basis, (0,2))
ψ₀ = vec(ψ₁₁) # start in 11 state
pop₁₁ = Float64[]
pop₀₂ = Float64[]

for tmax = times

    # erfsquared pulse parameters
    fwhm = 0.5 * risetime
    t₁ = fwhm
    t₂ = tmax - fwhm
    σ = 0.5 * fwhm / sqrt(2*log(2))
    erf_squared(t) = 0.5 * (erf((t - t₁)/σ) - erf((t - t₂)/σ) )

    for freq = freqs
        cqs = CompositeQSystem([q0, q1])
        add_hamiltonian!(cqs, q0)
        # add rotating frame Hamiltonian shifts
        add_hamiltonian!(cqs, -spec(q0).frequency*number(q0), q0)
        q1_freq = hamiltonian(q1, 0.0)[2,2]
        add_hamiltonian!(cqs, -q1_freq*number(q1), q1)

        diff_freq = spec(q0).frequency - q1_freq
        add_hamiltonian!(cqs, t -> 0.006*.5*X_Y([q0, q1], [diff_freq*t, 0.0]), [q0,q1])

        add_hamiltonian!(cqs, parametric_drive(q1, t -> amp*erf_squared(t)*sin(2π*freq*t)), q1)

        ψ = unitary_state(cqs, tmax, ψ₀);
        push!(pop₁₁, abs2(ψ[QSimulator.index(ψ₁₁)]))
        push!(pop₀₂, abs2(ψ[QSimulator.index(ψ₀₂)]))
    end
end
pop₁₁ = reshape(pop₁₁, length(freqs), length(times))
pop₀₂ = reshape(pop₀₂, length(freqs), length(times));
# -

p1 = contour(times, 1e3*freqs, pop₀₂, fill=true)
ylabel!(p1, "Frequency (MHz)")
xlabel!(p1, "Time (ns)")
title!(p1, "Population 11")
p2 = contour(times, 1e3*freqs, pop₀₂, fill=true)
xlabel!(p2, "Time (ns)")
title!(p2,  "Population 02")
plot(p1,p2, layout=(1,2), size=(800,400))

# # Pulse Amplitude Noise
#
# We can estimate the effect of fluctuations on the pulse amplitude with a Krauss map sum of unitaries weighted by a normal distribution.
#
# $$ \mathcal{S}(\rho) = \sum_k \lambda_k U_k\rho U_k^\dagger $$

# +
# pick an operating point of the plots above

freq = 103/1e3
amp = 0.245
tmax = 190

q0 = DuffingTransmon("q0", 3, DuffingSpec(3.94015, -0.1807))
q1 = PerturbativeTransmon("q1", 3, TransmonSpec(.172, 12.71, 3.69))

risetime = 50
fwhm = 0.5 * risetime
t₁ = fwhm
t₂ = tmax - fwhm
σ = 0.5 * fwhm / sqrt(2*log(2))
erf_squared(t) = 0.5 * (erf((t - t₁)/σ) - erf((t - t₂)/σ) )

cqs = CompositeQSystem([q0, q1])
add_hamiltonian!(cqs, q0)
# add rotating frame Hamiltonian shifts
add_hamiltonian!(cqs, -spec(q0).frequency*number(q0), q0)
q1_freq = hamiltonian(q1, 0.0)[2,2]
add_hamiltonian!(cqs, -q1_freq*number(q1), q1)

diff_freq = spec(q0).frequency - q1_freq
add_hamiltonian!(cqs, t -> 0.006*.5*X_Y([q0, q1], [diff_freq*t, 0.0]), [q0,q1])

add_hamiltonian!(cqs, parametric_drive(q1, t -> amp*erf_squared(t)*sin(2π*freq*t)), q1)

U = unitary_propagator(cqs, float(tmax));
# -

# project out into the qubit space
projector = [1,0]*[1, 0, 0]' + [0, 1]*[0, 1, 0]' 
projector = projector ⊗ projector
U_proj = projector * U * projector'

import QuantumInfo: liou, avgfidelity, kraus2liou
import Cliffords: Z
using Optim
using LinearAlgebra: diagm

Zrot = θ -> exp(-1im * θ * π * Z)

CZ = diagm(0 => [1.0,1.0,1.0,-1.0])
res = optimize(zs -> 1 - avgfidelity(liou((Zrot(zs[1]) ⊗ Zrot(zs[2])) * U_proj), liou(CZ)), [0.0, 0.0])

# +
# pick an operating point of the plots above
freq = 103/1e3
amp = 0.245
tmax = 190

q0 = DuffingTransmon("q0", 3, DuffingSpec(3.94015, -0.1807))
q1 = PerturbativeTransmon("q1", 3, TransmonSpec(.172, 12.71, 3.69))

risetime = 50
fwhm = 0.5 * risetime
t₁ = fwhm
t₂ = tmax - fwhm
σ = 0.5 * fwhm / sqrt(2*log(2))
erf_squared(t) = 0.5 * (erf((t - t₁)/σ) - erf((t - t₂)/σ) )

Us = []

amp_noises = range(-0.005, stop=0.005, length=101)
for amp_noise = amp_noises
    cqs = CompositeQSystem([q0, q1])
    add_hamiltonian!(cqs, q0)
    # add rotating frame Hamiltonian shifts
    add_hamiltonian!(cqs, -spec(q0).frequency*number(q0), q0)
    q1_freq = hamiltonian(q1, 0.0)[2,2]
    add_hamiltonian!(cqs, -q1_freq*number(q1), q1)

    diff_freq = spec(q0).frequency - q1_freq
    add_hamiltonian!(cqs, t -> 0.006*.5*X_Y([q0, q1], [diff_freq*t, 0.0]), [q0,q1])

    add_hamiltonian!(cqs, parametric_drive(q1, t -> amp*(1+amp_noise)*erf_squared(t)*sin(2π*freq*t)), q1)

    push!(Us, unitary_propagator(cqs, float(tmax)));
end
z_corrs = res.minimizer
Us = [projector * U * projector' for U = Us]
Us = [Zrot(z_corrs[1]) ⊗ Zrot(z_corrs[2]) * U for U = Us];
# -

p = plot(1 .+ amp_noises, [avgfidelity(liou(U), liou(CZ)) for U = Us])
xlabel!(p, "Pulse Amplitude")
ylabel!(p, "CZ Fidelity")

σ = 0.0005
distribution = exp.(-0.5*(amp_noises/σ).^2)
distribution /= sum(distribution)
plot(amp_noises, distribution)

fids = Float64[]
sigmas = range(1e-4, stop=3e-3, length=101)
for σ = sigmas
    distribution = exp.(-0.5*(amp_noises/σ).^2)
    distribution /= sum(distribution)
    fid = avgfidelity(kraus2liou([√λ*U for (λ,U) = zip(distribution, Us)]), liou(CZ))
    push!(fids, fid)
end
p = plot(sigmas, fids)
xlabel!(p, "Relative Pulse Amplitude σ")
ylabel!(p, "CZ Average Fidelity")



StarcoderdataJupyterScriptsDedupFiltered