Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

add allocations tests and clean time evolution tests #304

Merged
merged 4 commits into from
Nov 13, 2024
Merged

add allocations tests and clean time evolution tests #304

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
25aeba5
add allocations tests and clean time evolution tests
albertomercurio Nov 12, 2024
052f530
Format Document
albertomercurio Nov 12, 2024
c1fad94
Remove global tlist
albertomercurio Nov 13, 2024
87b32e2
Fix format
albertomercurio Nov 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
293 changes: 181 additions & 112 deletions test/core-test/time_evolution.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,32 +1,48 @@
@testset "Time Evolution and Partial Trace" verbose = true begin
# Global definition of the system
N = 10
a = kron(destroy(N), qeye(2))
σm = kron(qeye(N), sigmam())
σz = qeye(N) ⊗ sigmaz()

g = 0.01
ωc = 1
ωq = 0.99
γ = 0.1
nth = 0.001

# Jaynes-Cummings Hamiltonian
H = ωc * a' * a + ωq / 2 * σz + g * (a' * σm + a * σm')
ψ0 = kron(fock(N, 0), fock(2, 0))

e_ops = [a' * a, σz]
c_ops = [sqrt(γ * (1 + nth)) * a, sqrt(γ * nth) * a', sqrt(γ * (1 + nth)) * σm, sqrt(γ * nth) * σm']

@testset "sesolve" begin
N = 10
a_d = kron(create(N), qeye(2))
a = a_d'
sx = kron(qeye(N), sigmax())
sy = tensor(qeye(N), sigmay())
sz = qeye(N) ⊗ sigmaz()
η = 0.01
H = a_d * a + 0.5 * sz - 1im * η * (a - a_d) * sx
psi0 = kron(fock(N, 0), fock(2, 0))
t_l = LinRange(0, 1000, 1000)
e_ops = [a_d * a]
prob = sesolveProblem(H, psi0, t_l, e_ops = e_ops, progress_bar = Val(false))
tlist = range(0, 20 * 2π / g, 1000)

prob = sesolveProblem(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
sol = sesolve(prob)
sol2 = sesolve(H, psi0, t_l, progress_bar = Val(false))
sol3 = sesolve(H, psi0, t_l, e_ops = e_ops, saveat = t_l, progress_bar = Val(false))
sol2 = sesolve(H, ψ0, tlist, progress_bar = Val(false))
sol3 = sesolve(H, ψ0, tlist, e_ops = e_ops, saveat = tlist, progress_bar = Val(false))
sol_string = sprint((t, s) -> show(t, "text/plain", s), sol)

## Analytical solution for the expectation value of a' * a
Ω_rabi = sqrt(g^2 + ((ωc - ωq) / 2)^2)
amp_rabi = g^2 / Ω_rabi^2
##

@test prob.f.f isa MatrixOperator
@test sum(abs.(sol.expect[1, :] .- sin.(η * t_l) .^ 2)) / length(t_l) < 0.1
@test length(sol.times) == length(t_l)
@test sum(abs.(sol.expect[1, :] .- amp_rabi .* sin.(Ω_rabi * tlist) .^ 2)) / length(tlist) < 0.1
@test length(sol.times) == length(tlist)
@test length(sol.states) == 1
@test size(sol.expect) == (length(e_ops), length(t_l))
@test length(sol2.times) == length(t_l)
@test length(sol2.states) == length(t_l)
@test size(sol2.expect) == (0, length(t_l))
@test length(sol3.times) == length(t_l)
@test length(sol3.states) == length(t_l)
@test size(sol3.expect) == (length(e_ops), length(t_l))
@test size(sol.expect) == (length(e_ops), length(tlist))
@test length(sol2.times) == length(tlist)
@test length(sol2.states) == length(tlist)
@test size(sol2.expect) == (0, length(tlist))
@test length(sol3.times) == length(tlist)
@test length(sol3.states) == length(tlist)
@test size(sol3.expect) == (length(e_ops), length(tlist))
@test sol_string ==
"Solution of time evolution\n" *
"(return code: $(sol.retcode))\n" *
Expand All @@ -37,54 +53,58 @@
"abstol = $(sol.abstol)\n" *
"reltol = $(sol.reltol)\n"

@testset "Memory Allocations" begin
allocs_tot = @allocations sesolve(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false)) # Warm-up
allocs_tot = @allocations sesolve(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
@test allocs_tot < 150

allocs_tot = @allocations sesolve(H, ψ0, tlist, saveat = [tlist[end]], progress_bar = Val(false)) # Warm-up
allocs_tot = @allocations sesolve(H, ψ0, tlist, saveat = [tlist[end]], progress_bar = Val(false))
@test allocs_tot < 100
end

@testset "Type Inference sesolve" begin
@inferred sesolveProblem(H, psi0, t_l, progress_bar = Val(false))
@inferred sesolveProblem(H, psi0, [0, 10], progress_bar = Val(false))
@inferred sesolveProblem(H, Qobj(zeros(Int64, N * 2); dims = (N, 2)), t_l, progress_bar = Val(false))
@inferred sesolve(H, psi0, t_l, e_ops = e_ops, progress_bar = Val(false))
@inferred sesolve(H, psi0, t_l, progress_bar = Val(false))
@inferred sesolve(H, psi0, t_l, e_ops = e_ops, saveat = t_l, progress_bar = Val(false))
@inferred sesolve(H, psi0, t_l, e_ops = (a_d * a, a'), progress_bar = Val(false)) # We test the type inference for Tuple of different types
@inferred sesolveProblem(H, ψ0, tlist, progress_bar = Val(false))
@inferred sesolveProblem(H, ψ0, [0, 10], progress_bar = Val(false))
@inferred sesolveProblem(H, Qobj(zeros(Int64, N * 2); dims = (N, 2)), tlist, progress_bar = Val(false))
@inferred sesolve(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
@inferred sesolve(H, ψ0, tlist, progress_bar = Val(false))
@inferred sesolve(H, ψ0, tlist, e_ops = e_ops, saveat = tlist, progress_bar = Val(false))
@inferred sesolve(H, ψ0, tlist, e_ops = (a' * a, a'), progress_bar = Val(false)) # We test the type inference for Tuple of different types
end
end

@testset "mesolve, mcsolve, and ssesolve" begin
N = 10
a = destroy(N)
a_d = a'
H = a_d * a
c_ops = [sqrt(0.1) * a]
e_ops = [a_d * a]
psi0 = basis(N, 3)
t_l = LinRange(0, 100, 1000)
prob_me = mesolveProblem(H, psi0, t_l, c_ops, e_ops = e_ops, progress_bar = Val(false))
tlist = range(0, 10 / γ, 100)

prob_me = mesolveProblem(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
sol_me = mesolve(prob_me)
sol_me2 = mesolve(H, psi0, t_l, c_ops, progress_bar = Val(false))
sol_me3 = mesolve(H, psi0, t_l, c_ops, e_ops = e_ops, saveat = t_l, progress_bar = Val(false))
prob_mc = mcsolveProblem(H, psi0, t_l, c_ops, e_ops = e_ops, progress_bar = Val(false))
sol_mc = mcsolve(H, psi0, t_l, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false))
sol_me2 = mesolve(H, ψ0, tlist, c_ops, progress_bar = Val(false))
sol_me3 = mesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, saveat = tlist, progress_bar = Val(false))
prob_mc = mcsolveProblem(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
sol_mc = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false))
sol_mc2 = mcsolve(
H,
psi0,
t_l,
ψ0,
tlist,
c_ops,
ntraj = 500,
e_ops = e_ops,
progress_bar = Val(false),
jump_callback = DiscreteLindbladJumpCallback(),
)
sol_mc_states = mcsolve(H, psi0, t_l, c_ops, ntraj = 500, saveat = t_l, progress_bar = Val(false))
sol_mc_states = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, saveat = tlist, progress_bar = Val(false))
sol_mc_states2 = mcsolve(
H,
psi0,
t_l,
ψ0,
tlist,
c_ops,
ntraj = 500,
saveat = t_l,
saveat = tlist,
progress_bar = Val(false),
jump_callback = DiscreteLindbladJumpCallback(),
)
sol_sse = ssesolve(H, psi0, t_l, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false))
sol_sse = ssesolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false))

ρt_mc = [ket2dm.(normalize.(states)) for states in sol_mc_states.states]
expect_mc_states = mapreduce(states -> expect.(Ref(e_ops[1]), states), hcat, ρt_mc)
Expand All @@ -99,26 +119,26 @@
sol_sse_string = sprint((t, s) -> show(t, "text/plain", s), sol_sse)
@test prob_me.f.f isa MatrixOperator
@test prob_mc.f.f isa MatrixOperator
@test sum(abs.(sol_mc.expect .- sol_me.expect)) / length(t_l) < 0.1
@test sum(abs.(sol_mc2.expect .- sol_me.expect)) / length(t_l) < 0.1
@test sum(abs.(vec(expect_mc_states_mean) .- vec(sol_me.expect))) / length(t_l) < 0.1
@test sum(abs.(vec(expect_mc_states_mean2) .- vec(sol_me.expect))) / length(t_l) < 0.1
@test sum(abs.(sol_sse.expect .- sol_me.expect)) / length(t_l) < 0.1
@test length(sol_me.times) == length(t_l)
@test sum(abs.(sol_mc.expect .- sol_me.expect)) / length(tlist) < 0.1
@test sum(abs.(sol_mc2.expect .- sol_me.expect)) / length(tlist) < 0.1
@test sum(abs.(vec(expect_mc_states_mean) .- vec(sol_me.expect[1, :]))) / length(tlist) < 0.1
@test sum(abs.(vec(expect_mc_states_mean2) .- vec(sol_me.expect[1, :]))) / length(tlist) < 0.1
@test sum(abs.(sol_sse.expect .- sol_me.expect)) / length(tlist) < 0.1
@test length(sol_me.times) == length(tlist)
@test length(sol_me.states) == 1
@test size(sol_me.expect) == (length(e_ops), length(t_l))
@test length(sol_me2.times) == length(t_l)
@test length(sol_me2.states) == length(t_l)
@test size(sol_me2.expect) == (0, length(t_l))
@test length(sol_me3.times) == length(t_l)
@test length(sol_me3.states) == length(t_l)
@test size(sol_me3.expect) == (length(e_ops), length(t_l))
@test length(sol_mc.times) == length(t_l)
@test size(sol_mc.expect) == (length(e_ops), length(t_l))
@test length(sol_mc_states.times) == length(t_l)
@test size(sol_mc_states.expect) == (0, length(t_l))
@test length(sol_sse.times) == length(t_l)
@test size(sol_sse.expect) == (length(e_ops), length(t_l))
@test size(sol_me.expect) == (length(e_ops), length(tlist))
@test length(sol_me2.times) == length(tlist)
@test length(sol_me2.states) == length(tlist)
@test size(sol_me2.expect) == (0, length(tlist))
@test length(sol_me3.times) == length(tlist)
@test length(sol_me3.states) == length(tlist)
@test size(sol_me3.expect) == (length(e_ops), length(tlist))
@test length(sol_mc.times) == length(tlist)
@test size(sol_mc.expect) == (length(e_ops), length(tlist))
@test length(sol_mc_states.times) == length(tlist)
@test size(sol_mc_states.expect) == (0, length(tlist))
@test length(sol_sse.times) == length(tlist)
@test size(sol_sse.expect) == (length(e_ops), length(tlist))
@test sol_me_string ==
"Solution of time evolution\n" *
"(return code: $(sol_me.retcode))\n" *
Expand Down Expand Up @@ -152,38 +172,24 @@
# Time-Dependent Hamiltonian
# ssesolve is slow to be run on CI. It is not removed from the test because it may be useful for testing in more powerful machines.

N = 10
a = tensor(destroy(N), qeye(2))
σm = tensor(qeye(N), sigmam())
σz = tensor(qeye(N), sigmaz())
ω = 1.0
ωd = 1.02
Δ = ω - ωd
F = 0.05
g = 0.1
γ = 0.1
nth = 0.001

# Time Evolution in the drive frame

H = Δ * a' * a + Δ * σz / 2 + g * (a' * σm + a * σm') + F * (a + a')
c_ops = [sqrt(γ * (1 + nth)) * a, sqrt(γ * nth) * a', sqrt(γ * (1 + nth)) * σm, sqrt(γ * nth) * σm']
e_ops = [a' * a, σz]

ψ0 = tensor(basis(N, 0), basis(2, 1))
tlist = range(0, 2 / γ, 1000)
H_dr_fr = H - ωd * a' * a - ωd * σz / 2 + F * (a + a')

rng = MersenneTwister(12)

sol_se = sesolve(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
sol_me = mesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
sol_mc = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
tlist = range(0, 10 / γ, 1000)

sol_se = sesolve(H_dr_fr, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
sol_me = mesolve(H_dr_fr, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
sol_mc = mcsolve(H_dr_fr, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
# sol_sse = ssesolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)

# Time Evolution in the lab frame

H = ω * a' * a + ω * σz / 2 + g * (a' * σm + a * σm')

coef1(p, t) = p.F * exp(1im * p.ωd * t)
coef2(p, t) = p.F * exp(-1im * p.ωd * t)

Expand Down Expand Up @@ -234,6 +240,87 @@
@test sol_mc.expect ≈ sol_mc_td2.expect atol = 1e-2 * length(tlist)
# @test sol_sse.expect ≈ sol_sse_td2.expect atol = 1e-2 * length(tlist)

@testset "Memory Allocations (mesolve)" begin
# We predefine the Liouvillian to avoid to count the allocations of the liouvillian function
L = liouvillian(H, c_ops)
L_td = QobjEvo((liouvillian(H, c_ops), (liouvillian(a), coef1), (liouvillian(a'), coef2)))

allocs_tot = @allocations mesolve(L, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false)) # Warm-up
allocs_tot = @allocations mesolve(L, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
@test allocs_tot < 210

allocs_tot = @allocations mesolve(L, ψ0, tlist, saveat = [tlist[end]], progress_bar = Val(false)) # Warm-up
allocs_tot = @allocations mesolve(L, ψ0, tlist, saveat = [tlist[end]], progress_bar = Val(false))
@test allocs_tot < 120

allocs_tot = @allocations mesolve(L_td, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false), params = p) # Warm-up
allocs_tot = @allocations mesolve(L_td, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false), params = p)
@test allocs_tot < 210

allocs_tot =
@allocations mesolve(L_td, ψ0, tlist, progress_bar = Val(false), saveat = [tlist[end]], params = p) # Warm-up
allocs_tot =
@allocations mesolve(L_td, ψ0, tlist, progress_bar = Val(false), saveat = [tlist[end]], params = p)
@test allocs_tot < 120
end

@testset "Memory Allocations (mcsolve)" begin
ntraj = 100
allocs_tot =
@allocations mcsolve(H, ψ0, tlist, c_ops, e_ops = e_ops, ntraj = ntraj, progress_bar = Val(false)) # Warm-up
allocs_tot =
@allocations mcsolve(H, ψ0, tlist, c_ops, e_ops = e_ops, ntraj = ntraj, progress_bar = Val(false))
@test allocs_tot < 160 * ntraj + 500 # 150 allocations per trajectory + 500 for initialization

allocs_tot = @allocations mcsolve(
H,
ψ0,
tlist,
c_ops,
ntraj = ntraj,
saveat = [tlist[end]],
progress_bar = Val(false),
) # Warm-up
allocs_tot = @allocations mcsolve(
H,
ψ0,
tlist,
c_ops,
ntraj = ntraj,
saveat = [tlist[end]],
progress_bar = Val(false),
)
@test allocs_tot < 160 * ntraj + 300 # 100 allocations per trajectory + 300 for initialization
end

@testset "Memory Allocations (ssesolve)" begin
allocs_tot =
@allocations ssesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, ntraj = 100, progress_bar = Val(false)) # Warm-up
allocs_tot =
@allocations ssesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, ntraj = 100, progress_bar = Val(false))
@test allocs_tot < 1950000 # TODO: Fix this high number of allocations

allocs_tot = @allocations ssesolve(
H,
ψ0,
tlist,
c_ops,
ntraj = 100,
saveat = [tlist[end]],
progress_bar = Val(false),
) # Warm-up
allocs_tot = @allocations ssesolve(
H,
ψ0,
tlist,
c_ops,
ntraj = 100,
saveat = [tlist[end]],
progress_bar = Val(false),
)
@test allocs_tot < 570000 # TODO: Fix this high number of allocations
end

@testset "Type Inference mesolve" begin
coef(p, t) = exp(-t)
ad_t = QobjEvo(a', coef)
Expand Down Expand Up @@ -359,34 +446,16 @@
end

@testset "mcsolve and ssesolve reproducibility" begin
N = 10
a = tensor(destroy(N), qeye(2))
σm = tensor(qeye(N), sigmam())
σp = σm'
σz = tensor(qeye(N), sigmaz())

ω = 1.0
g = 0.1
γ = 0.01
nth = 0.1

H = ω * a' * a + ω * σz / 2 + g * (a' * σm + a * σp)
c_ops = [sqrt(γ * (1 + nth)) * a, sqrt(γ * nth) * a', sqrt(γ * (1 + nth)) * σm, sqrt(γ * nth) * σp]
e_ops = [a' * a, σz]

psi0 = tensor(basis(N, 0), basis(2, 0))
tlist = range(0, 20 / γ, 1000)

rng = MersenneTwister(1234)
sol_mc1 = mcsolve(H, psi0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
sol_sse1 = ssesolve(H, psi0, tlist, c_ops, ntraj = 50, e_ops = e_ops, progress_bar = Val(false), rng = rng)
sol_mc1 = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
sol_sse1 = ssesolve(H, ψ0, tlist, c_ops, ntraj = 50, e_ops = e_ops, progress_bar = Val(false), rng = rng)

rng = MersenneTwister(1234)
sol_mc2 = mcsolve(H, psi0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
sol_sse2 = ssesolve(H, psi0, tlist, c_ops, ntraj = 50, e_ops = e_ops, progress_bar = Val(false), rng = rng)
sol_mc2 = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
sol_sse2 = ssesolve(H, ψ0, tlist, c_ops, ntraj = 50, e_ops = e_ops, progress_bar = Val(false), rng = rng)

rng = MersenneTwister(1234)
sol_mc3 = mcsolve(H, psi0, tlist, c_ops, ntraj = 510, e_ops = e_ops, progress_bar = Val(false), rng = rng)
sol_mc3 = mcsolve(H, ψ0, tlist, c_ops, ntraj = 510, e_ops = e_ops, progress_bar = Val(false), rng = rng)

@test sol_mc1.expect ≈ sol_mc2.expect atol = 1e-10
@test sol_mc1.expect_all ≈ sol_mc2.expect_all atol = 1e-10
Expand Down
Loading