Skip to content

Commit 3fd5a62

Browse files
albertomercurioalbertomercurio
authored
add allocations tests and clean time evolution tests (#304)
* add allocations tests and clean time evolution tests * Format Document * Remove global tlist * Fix format
1 parent 4deb2f1 commit 3fd5a62

File tree

1 file changed

+181
-112
lines changed

1 file changed

+181
-112
lines changed

‎test/core-test/time_evolution.jl‎

Lines changed: 181 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,48 @@
11
@testset "Time Evolution and Partial Trace" verbose = true begin
2+
# Global definition of the system
3+
N = 10
4+
a = kron(destroy(N), qeye(2))
5+
σm = kron(qeye(N), sigmam())
6+
σz = qeye(N) ⊗ sigmaz()
7+
8+
g = 0.01
9+
ωc = 1
10+
ωq = 0.99
11+
γ = 0.1
12+
nth = 0.001
13+
14+
# Jaynes-Cummings Hamiltonian
15+
H = ωc * a' * a + ωq / 2 * σz + g * (a' * σm + a * σm')
16+
ψ0 = kron(fock(N, 0), fock(2, 0))
17+
18+
e_ops = [a' * a, σz]
19+
c_ops = [sqrt(γ * (1 + nth)) * a, sqrt(γ * nth) * a', sqrt(γ * (1 + nth)) * σm, sqrt(γ * nth) * σm']
20+
221
@testset "sesolve" begin
3-
N = 10
4-
a_d = kron(create(N), qeye(2))
5-
a = a_d'
6-
sx = kron(qeye(N), sigmax())
7-
sy = tensor(qeye(N), sigmay())
8-
sz = qeye(N) ⊗ sigmaz()
9-
η = 0.01
10-
H = a_d * a + 0.5 * sz - 1im * η * (a - a_d) * sx
11-
psi0 = kron(fock(N, 0), fock(2, 0))
12-
t_l = LinRange(0, 1000, 1000)
13-
e_ops = [a_d * a]
14-
prob = sesolveProblem(H, psi0, t_l, e_ops = e_ops, progress_bar = Val(false))
22+
tlist = range(0, 20 * 2π / g, 1000)
23+
24+
prob = sesolveProblem(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
1525
sol = sesolve(prob)
16-
sol2 = sesolve(H, psi0, t_l, progress_bar = Val(false))
17-
sol3 = sesolve(H, psi0, t_l, e_ops = e_ops, saveat = t_l, progress_bar = Val(false))
26+
sol2 = sesolve(H, ψ0, tlist, progress_bar = Val(false))
27+
sol3 = sesolve(H, ψ0, tlist, e_ops = e_ops, saveat = tlist, progress_bar = Val(false))
1828
sol_string = sprint((t, s) -> show(t, "text/plain", s), sol)
29+
30+
## Analytical solution for the expectation value of a' * a
31+
Ω_rabi = sqrt(g^2 + ((ωc - ωq) / 2)^2)
32+
amp_rabi = g^2 / Ω_rabi^2
33+
##
34+
1935
@test prob.f.f isa MatrixOperator
20-
@test sum(abs.(sol.expect[1, :] .- sin.(η * t_l) .^ 2)) / length(t_l) < 0.1
21-
@test length(sol.times) == length(t_l)
36+
@test sum(abs.(sol.expect[1, :] .- amp_rabi .* sin.(Ω_rabi * tlist) .^ 2)) / length(tlist) < 0.1
37+
@test length(sol.times) == length(tlist)
2238
@test length(sol.states) == 1
23-
@test size(sol.expect) == (length(e_ops), length(t_l))
24-
@test length(sol2.times) == length(t_l)
25-
@test length(sol2.states) == length(t_l)
26-
@test size(sol2.expect) == (0, length(t_l))
27-
@test length(sol3.times) == length(t_l)
28-
@test length(sol3.states) == length(t_l)
29-
@test size(sol3.expect) == (length(e_ops), length(t_l))
39+
@test size(sol.expect) == (length(e_ops), length(tlist))
40+
@test length(sol2.times) == length(tlist)
41+
@test length(sol2.states) == length(tlist)
42+
@test size(sol2.expect) == (0, length(tlist))
43+
@test length(sol3.times) == length(tlist)
44+
@test length(sol3.states) == length(tlist)
45+
@test size(sol3.expect) == (length(e_ops), length(tlist))
3046
@test sol_string ==
3147
"Solution of time evolution\n" *
3248
"(return code: $(sol.retcode))\n" *
@@ -37,54 +53,58 @@
3753
"abstol = $(sol.abstol)\n" *
3854
"reltol = $(sol.reltol)\n"
3955

56+
@testset "Memory Allocations" begin
57+
allocs_tot = @allocations sesolve(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false)) # Warm-up
58+
allocs_tot = @allocations sesolve(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
59+
@test allocs_tot < 150
60+
61+
allocs_tot = @allocations sesolve(H, ψ0, tlist, saveat = [tlist[end]], progress_bar = Val(false)) # Warm-up
62+
allocs_tot = @allocations sesolve(H, ψ0, tlist, saveat = [tlist[end]], progress_bar = Val(false))
63+
@test allocs_tot < 100
64+
end
65+
4066
@testset "Type Inference sesolve" begin
41-
@inferred sesolveProblem(H, psi0, t_l, progress_bar = Val(false))
42-
@inferred sesolveProblem(H, psi0, [0, 10], progress_bar = Val(false))
43-
@inferred sesolveProblem(H, Qobj(zeros(Int64, N * 2); dims = (N, 2)), t_l, progress_bar = Val(false))
44-
@inferred sesolve(H, psi0, t_l, e_ops = e_ops, progress_bar = Val(false))
45-
@inferred sesolve(H, psi0, t_l, progress_bar = Val(false))
46-
@inferred sesolve(H, psi0, t_l, e_ops = e_ops, saveat = t_l, progress_bar = Val(false))
47-
@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
67+
@inferred sesolveProblem(H, ψ0, tlist, progress_bar = Val(false))
68+
@inferred sesolveProblem(H, ψ0, [0, 10], progress_bar = Val(false))
69+
@inferred sesolveProblem(H, Qobj(zeros(Int64, N * 2); dims = (N, 2)), tlist, progress_bar = Val(false))
70+
@inferred sesolve(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
71+
@inferred sesolve(H, ψ0, tlist, progress_bar = Val(false))
72+
@inferred sesolve(H, ψ0, tlist, e_ops = e_ops, saveat = tlist, progress_bar = Val(false))
73+
@inferred sesolve(H, ψ0, tlist, e_ops = (a' * a, a'), progress_bar = Val(false)) # We test the type inference for Tuple of different types
4874
end
4975
end
5076

5177
@testset "mesolve, mcsolve, and ssesolve" begin
52-
N = 10
53-
a = destroy(N)
54-
a_d = a'
55-
H = a_d * a
56-
c_ops = [sqrt(0.1) * a]
57-
e_ops = [a_d * a]
58-
psi0 = basis(N, 3)
59-
t_l = LinRange(0, 100, 1000)
60-
prob_me = mesolveProblem(H, psi0, t_l, c_ops, e_ops = e_ops, progress_bar = Val(false))
78+
tlist = range(0, 10 / γ, 100)
79+
80+
prob_me = mesolveProblem(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
6181
sol_me = mesolve(prob_me)
62-
sol_me2 = mesolve(H, psi0, t_l, c_ops, progress_bar = Val(false))
63-
sol_me3 = mesolve(H, psi0, t_l, c_ops, e_ops = e_ops, saveat = t_l, progress_bar = Val(false))
64-
prob_mc = mcsolveProblem(H, psi0, t_l, c_ops, e_ops = e_ops, progress_bar = Val(false))
65-
sol_mc = mcsolve(H, psi0, t_l, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false))
82+
sol_me2 = mesolve(H, ψ0, tlist, c_ops, progress_bar = Val(false))
83+
sol_me3 = mesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, saveat = tlist, progress_bar = Val(false))
84+
prob_mc = mcsolveProblem(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
85+
sol_mc = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false))
6686
sol_mc2 = mcsolve(
6787
H,
68-
psi0,
69-
t_l,
88+
ψ0,
89+
tlist,
7090
c_ops,
7191
ntraj = 500,
7292
e_ops = e_ops,
7393
progress_bar = Val(false),
7494
jump_callback = DiscreteLindbladJumpCallback(),
7595
)
76-
sol_mc_states = mcsolve(H, psi0, t_l, c_ops, ntraj = 500, saveat = t_l, progress_bar = Val(false))
96+
sol_mc_states = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, saveat = tlist, progress_bar = Val(false))
7797
sol_mc_states2 = mcsolve(
7898
H,
79-
psi0,
80-
t_l,
99+
ψ0,
100+
tlist,
81101
c_ops,
82102
ntraj = 500,
83-
saveat = t_l,
103+
saveat = tlist,
84104
progress_bar = Val(false),
85105
jump_callback = DiscreteLindbladJumpCallback(),
86106
)
87-
sol_sse = ssesolve(H, psi0, t_l, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false))
107+
sol_sse = ssesolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false))
88108

89109
ρt_mc = [ket2dm.(normalize.(states)) for states in sol_mc_states.states]
90110
expect_mc_states = mapreduce(states -> expect.(Ref(e_ops[1]), states), hcat, ρt_mc)
@@ -99,26 +119,26 @@
99119
sol_sse_string = sprint((t, s) -> show(t, "text/plain", s), sol_sse)
100120
@test prob_me.f.f isa MatrixOperator
101121
@test prob_mc.f.f isa MatrixOperator
102-
@test sum(abs.(sol_mc.expect .- sol_me.expect)) / length(t_l) < 0.1
103-
@test sum(abs.(sol_mc2.expect .- sol_me.expect)) / length(t_l) < 0.1
104-
@test sum(abs.(vec(expect_mc_states_mean) .- vec(sol_me.expect))) / length(t_l) < 0.1
105-
@test sum(abs.(vec(expect_mc_states_mean2) .- vec(sol_me.expect))) / length(t_l) < 0.1
106-
@test sum(abs.(sol_sse.expect .- sol_me.expect)) / length(t_l) < 0.1
107-
@test length(sol_me.times) == length(t_l)
122+
@test sum(abs.(sol_mc.expect .- sol_me.expect)) / length(tlist) < 0.1
123+
@test sum(abs.(sol_mc2.expect .- sol_me.expect)) / length(tlist) < 0.1
124+
@test sum(abs.(vec(expect_mc_states_mean) .- vec(sol_me.expect[1, :]))) / length(tlist) < 0.1
125+
@test sum(abs.(vec(expect_mc_states_mean2) .- vec(sol_me.expect[1, :]))) / length(tlist) < 0.1
126+
@test sum(abs.(sol_sse.expect .- sol_me.expect)) / length(tlist) < 0.1
127+
@test length(sol_me.times) == length(tlist)
108128
@test length(sol_me.states) == 1
109-
@test size(sol_me.expect) == (length(e_ops), length(t_l))
110-
@test length(sol_me2.times) == length(t_l)
111-
@test length(sol_me2.states) == length(t_l)
112-
@test size(sol_me2.expect) == (0, length(t_l))
113-
@test length(sol_me3.times) == length(t_l)
114-
@test length(sol_me3.states) == length(t_l)
115-
@test size(sol_me3.expect) == (length(e_ops), length(t_l))
116-
@test length(sol_mc.times) == length(t_l)
117-
@test size(sol_mc.expect) == (length(e_ops), length(t_l))
118-
@test length(sol_mc_states.times) == length(t_l)
119-
@test size(sol_mc_states.expect) == (0, length(t_l))
120-
@test length(sol_sse.times) == length(t_l)
121-
@test size(sol_sse.expect) == (length(e_ops), length(t_l))
129+
@test size(sol_me.expect) == (length(e_ops), length(tlist))
130+
@test length(sol_me2.times) == length(tlist)
131+
@test length(sol_me2.states) == length(tlist)
132+
@test size(sol_me2.expect) == (0, length(tlist))
133+
@test length(sol_me3.times) == length(tlist)
134+
@test length(sol_me3.states) == length(tlist)
135+
@test size(sol_me3.expect) == (length(e_ops), length(tlist))
136+
@test length(sol_mc.times) == length(tlist)
137+
@test size(sol_mc.expect) == (length(e_ops), length(tlist))
138+
@test length(sol_mc_states.times) == length(tlist)
139+
@test size(sol_mc_states.expect) == (0, length(tlist))
140+
@test length(sol_sse.times) == length(tlist)
141+
@test size(sol_sse.expect) == (length(e_ops), length(tlist))
122142
@test sol_me_string ==
123143
"Solution of time evolution\n" *
124144
"(return code: $(sol_me.retcode))\n" *
@@ -152,38 +172,24 @@
152172
# Time-Dependent Hamiltonian
153173
# 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.
154174

155-
N = 10
156-
a = tensor(destroy(N), qeye(2))
157-
σm = tensor(qeye(N), sigmam())
158-
σz = tensor(qeye(N), sigmaz())
159-
ω = 1.0
160175
ωd = 1.02
161-
Δ = ω - ωd
162176
F = 0.05
163-
g = 0.1
164-
γ = 0.1
165-
nth = 0.001
166177

167178
# Time Evolution in the drive frame
168179

169-
H = Δ * a' * a + Δ * σz / 2 + g * (a' * σm + a * σm') + F * (a + a')
170-
c_ops = [sqrt(γ * (1 + nth)) * a, sqrt(γ * nth) * a', sqrt(γ * (1 + nth)) * σm, sqrt(γ * nth) * σm']
171-
e_ops = [a' * a, σz]
172-
173-
ψ0 = tensor(basis(N, 0), basis(2, 1))
174-
tlist = range(0, 2 / γ, 1000)
180+
H_dr_fr = H - ωd * a' * a - ωd * σz / 2 + F * (a + a')
175181

176182
rng = MersenneTwister(12)
177183

178-
sol_se = sesolve(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
179-
sol_me = mesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
180-
sol_mc = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
184+
tlist = range(0, 10 / γ, 1000)
185+
186+
sol_se = sesolve(H_dr_fr, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
187+
sol_me = mesolve(H_dr_fr, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
188+
sol_mc = mcsolve(H_dr_fr, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
181189
# sol_sse = ssesolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
182190

183191
# Time Evolution in the lab frame
184192

185-
H = ω * a' * a + ω * σz / 2 + g * (a' * σm + a * σm')
186-
187193
coef1(p, t) = p.F * exp(1im * p.ωd * t)
188194
coef2(p, t) = p.F * exp(-1im * p.ωd * t)
189195

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

243+
@testset "Memory Allocations (mesolve)" begin
244+
# We predefine the Liouvillian to avoid to count the allocations of the liouvillian function
245+
L = liouvillian(H, c_ops)
246+
L_td = QobjEvo((liouvillian(H, c_ops), (liouvillian(a), coef1), (liouvillian(a'), coef2)))
247+
248+
allocs_tot = @allocations mesolve(L, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false)) # Warm-up
249+
allocs_tot = @allocations mesolve(L, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
250+
@test allocs_tot < 210
251+
252+
allocs_tot = @allocations mesolve(L, ψ0, tlist, saveat = [tlist[end]], progress_bar = Val(false)) # Warm-up
253+
allocs_tot = @allocations mesolve(L, ψ0, tlist, saveat = [tlist[end]], progress_bar = Val(false))
254+
@test allocs_tot < 120
255+
256+
allocs_tot = @allocations mesolve(L_td, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false), params = p) # Warm-up
257+
allocs_tot = @allocations mesolve(L_td, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false), params = p)
258+
@test allocs_tot < 210
259+
260+
allocs_tot =
261+
@allocations mesolve(L_td, ψ0, tlist, progress_bar = Val(false), saveat = [tlist[end]], params = p) # Warm-up
262+
allocs_tot =
263+
@allocations mesolve(L_td, ψ0, tlist, progress_bar = Val(false), saveat = [tlist[end]], params = p)
264+
@test allocs_tot < 120
265+
end
266+
267+
@testset "Memory Allocations (mcsolve)" begin
268+
ntraj = 100
269+
allocs_tot =
270+
@allocations mcsolve(H, ψ0, tlist, c_ops, e_ops = e_ops, ntraj = ntraj, progress_bar = Val(false)) # Warm-up
271+
allocs_tot =
272+
@allocations mcsolve(H, ψ0, tlist, c_ops, e_ops = e_ops, ntraj = ntraj, progress_bar = Val(false))
273+
@test allocs_tot < 160 * ntraj + 500 # 150 allocations per trajectory + 500 for initialization
274+
275+
allocs_tot = @allocations mcsolve(
276+
H,
277+
ψ0,
278+
tlist,
279+
c_ops,
280+
ntraj = ntraj,
281+
saveat = [tlist[end]],
282+
progress_bar = Val(false),
283+
) # Warm-up
284+
allocs_tot = @allocations mcsolve(
285+
H,
286+
ψ0,
287+
tlist,
288+
c_ops,
289+
ntraj = ntraj,
290+
saveat = [tlist[end]],
291+
progress_bar = Val(false),
292+
)
293+
@test allocs_tot < 160 * ntraj + 300 # 100 allocations per trajectory + 300 for initialization
294+
end
295+
296+
@testset "Memory Allocations (ssesolve)" begin
297+
allocs_tot =
298+
@allocations ssesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, ntraj = 100, progress_bar = Val(false)) # Warm-up
299+
allocs_tot =
300+
@allocations ssesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, ntraj = 100, progress_bar = Val(false))
301+
@test allocs_tot < 1950000 # TODO: Fix this high number of allocations
302+
303+
allocs_tot = @allocations ssesolve(
304+
H,
305+
ψ0,
306+
tlist,
307+
c_ops,
308+
ntraj = 100,
309+
saveat = [tlist[end]],
310+
progress_bar = Val(false),
311+
) # Warm-up
312+
allocs_tot = @allocations ssesolve(
313+
H,
314+
ψ0,
315+
tlist,
316+
c_ops,
317+
ntraj = 100,
318+
saveat = [tlist[end]],
319+
progress_bar = Val(false),
320+
)
321+
@test allocs_tot < 570000 # TODO: Fix this high number of allocations
322+
end
323+
237324
@testset "Type Inference mesolve" begin
238325
coef(p, t) = exp(-t)
239326
ad_t = QobjEvo(a', coef)
@@ -359,34 +446,16 @@
359446
end
360447

361448
@testset "mcsolve and ssesolve reproducibility" begin
362-
N = 10
363-
a = tensor(destroy(N), qeye(2))
364-
σm = tensor(qeye(N), sigmam())
365-
σp = σm'
366-
σz = tensor(qeye(N), sigmaz())
367-
368-
ω = 1.0
369-
g = 0.1
370-
γ = 0.01
371-
nth = 0.1
372-
373-
H = ω * a' * a + ω * σz / 2 + g * (a' * σm + a * σp)
374-
c_ops = [sqrt(γ * (1 + nth)) * a, sqrt(γ * nth) * a', sqrt(γ * (1 + nth)) * σm, sqrt(γ * nth) * σp]
375-
e_ops = [a' * a, σz]
376-
377-
psi0 = tensor(basis(N, 0), basis(2, 0))
378-
tlist = range(0, 20 / γ, 1000)
379-
380449
rng = MersenneTwister(1234)
381-
sol_mc1 = mcsolve(H, psi0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
382-
sol_sse1 = ssesolve(H, psi0, tlist, c_ops, ntraj = 50, e_ops = e_ops, progress_bar = Val(false), rng = rng)
450+
sol_mc1 = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
451+
sol_sse1 = ssesolve(H, ψ0, tlist, c_ops, ntraj = 50, e_ops = e_ops, progress_bar = Val(false), rng = rng)
383452

384453
rng = MersenneTwister(1234)
385-
sol_mc2 = mcsolve(H, psi0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
386-
sol_sse2 = ssesolve(H, psi0, tlist, c_ops, ntraj = 50, e_ops = e_ops, progress_bar = Val(false), rng = rng)
454+
sol_mc2 = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng)
455+
sol_sse2 = ssesolve(H, ψ0, tlist, c_ops, ntraj = 50, e_ops = e_ops, progress_bar = Val(false), rng = rng)
387456

388457
rng = MersenneTwister(1234)
389-
sol_mc3 = mcsolve(H, psi0, tlist, c_ops, ntraj = 510, e_ops = e_ops, progress_bar = Val(false), rng = rng)
458+
sol_mc3 = mcsolve(H, ψ0, tlist, c_ops, ntraj = 510, e_ops = e_ops, progress_bar = Val(false), rng = rng)
390459

391460
@test sol_mc1.expect ≈ sol_mc2.expect atol = 1e-10
392461
@test sol_mc1.expect_all ≈ sol_mc2.expect_all atol = 1e-10

0 commit comments

Comments
 (0)