Profile Newton Propagation

profile_propagate.jl

# Include this file in a devrepl of 
module TransmonModel

using LinearAlgebra
using SparseArrays
using QuantumControl.Shapes: flattop

const GHz = 2π
const MHz = 0.001GHz
const ns = 1.0
const μs = 1000ns

⊗ = kron
const 𝕚 = 1im
const N = 6  # levels per transmon

function hamiltonian(;
    Ωre,
    Ωim,
    N=N,  # levels per transmon
    ω₁=4.380GHz,
    ω₂=4.614GHz,
    ωd=4.498GHz,
    α₁=-210MHz,
    α₂=-215MHz,
    J=-3MHz,
    λ=1.03,
    use_sparse=:auto
)
    𝟙 = SparseMatrixCSC{ComplexF64,Int64}(sparse(I, N, N))
    b̂₁ = spdiagm(1 => complex.(sqrt.(collect(1:N-1)))) ⊗ 𝟙
    b̂₂ = 𝟙 ⊗ spdiagm(1 => complex.(sqrt.(collect(1:N-1))))
    b̂₁⁺ = sparse(b̂₁')
    b̂₂⁺ = sparse(b̂₂')
    n̂₁ = sparse(b̂₁' * b̂₁)
    n̂₂ = sparse(b̂₂' * b̂₂)
    n̂₁² = sparse(n̂₁ * n̂₁)
    n̂₂² = sparse(n̂₂ * n̂₂)
    b̂₁⁺_b̂₂ = sparse(b̂₁' * b̂₂)
    b̂₁_b̂₂⁺ = sparse(b̂₁ * b̂₂')

    ω̃₁ = ω₁ - ωd
    ω̃₂ = ω₂ - ωd

    Ĥ₀ = sparse(
        (ω̃₁ - α₁ / 2) * n̂₁ +
        (α₁ / 2) * n̂₁² +
        (ω̃₂ - α₂ / 2) * n̂₂ +
        (α₂ / 2) * n̂₂² +
        J * (b̂₁⁺_b̂₂ + b̂₁_b̂₂⁺)
    )

    Ĥ₁re = (1 / 2) * (b̂₁ + b̂₁⁺ + λ * b̂₂ + λ * b̂₂⁺)
    Ĥ₁im = (𝕚 / 2) * (b̂₁⁺ - b̂₁ + λ * b̂₂⁺ - λ * b̂₂)

    if ((N < 5) && (use_sparse ≢ true)) || use_sparse ≡ false
        H = (Array(Ĥ₀), (Array(Ĥ₁re), Ωre), (Array(Ĥ₁im), Ωim))
    else
        H = (Ĥ₀, (Ĥ₁re, Ωre), (Ĥ₁im, Ωim))
    end
    return H

end

function guess_pulses(; T=400ns, E₀=35MHz, dt=0.1ns, t_rise=15ns)

    tlist = collect(range(0, T, step=dt))
    Ωre = t -> E₀ * flattop(t, T=T, t_rise=t_rise)
    Ωim = t -> 0.0

    return tlist, Ωre, Ωim

end

function ket(i::Int64; N=N)
    Ψ = zeros(ComplexF64, N)
    Ψ[i+1] = 1
    return Ψ
end

function ket(indices::Int64...; N=N)
    Ψ = ket(indices[1]; N=N)
    for i in indices[2:end]
        Ψ = Ψ ⊗ ket(i; N=N)
    end
    return Ψ
end

function ket(label::AbstractString; N=N)
    indices = [parse(Int64, digit) for digit in label]
    return ket(indices...; N=N)
end


end  # TransmonModel


module TransmonBenchmarks

using InteractiveUtils
using Profile
using ..TransmonModel: hamiltonian, ket, guess_pulses
using QuantumPropagators: _propagate
using QuantumControlBase:
    objective_genfunc, initobjpropwrk, MinimalObjective, _propagate_objective
using BenchmarkTools
using PProf


suite = BenchmarkGroup(["method"])


N = 6
tlist, Ωre, Ωim = guess_pulses()
H = hamiltonian(; Ωre, Ωim, N=6, use_sparse=false)
Ψ₀ = ket("00"; N=6)
obj = MinimalObjective(initial_state=Ψ₀, generator=H)
genfunc = objective_genfunc(obj, tlist)
wrk_cheby = initobjpropwrk(obj, tlist, Val(:cheby))
wrk_newton = initobjpropwrk(obj, tlist, Val(:newton))

function run_benchmark()
    _propagate(Ψ₀, genfunc, tlist, wrk_newton)
    Profile.clear()
    @pprof _propagate(Ψ₀, genfunc, tlist, wrk_newton)
end

end # TransmonBenchmarks

TransmonBenchmarks.run_benchmark()