Fix time evolution output when using saveat (#398)

Author: ytdHuang

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fix erroneous definition of the stochastic term in `smesolve`. ([#393])
 - Change name of `MultiSiteOperator` to `multisite_operator`. ([#394])
 - Fix `smesolve` for specifying initial state as density matrix. ([#395])
+- Fix time evolution output when using `saveat` keyword argument. ([#398])
 
 ## [v0.26.0]
 Release date: 2025-02-09
@@ -124,3 +125,4 @@ Release date: 2024-11-13
 [#393]: https://github.com/qutip/QuantumToolbox.jl/issues/393
 [#394]: https://github.com/qutip/QuantumToolbox.jl/issues/394
 [#395]: https://github.com/qutip/QuantumToolbox.jl/issues/395
+[#398]: https://github.com/qutip/QuantumToolbox.jl/issues/398

src/time_evolution/mcsolve.jl

@@ -131,12 +131,9 @@ function mcsolveProblem(
 
     T = Base.promote_eltype(H_eff_evo, ψ0)
 
-    is_empty_e_ops = e_ops isa Nothing ? true : isempty(e_ops)
-
-    saveat = is_empty_e_ops ? tlist : [tlist[end]]
     # We disable the progress bar of the sesolveProblem because we use a global progress bar for all the trajectories
-    default_values = (DEFAULT_ODE_SOLVER_OPTIONS..., saveat = saveat, progress_bar = Val(false))
-    kwargs2 = merge(default_values, kwargs)
+    default_values = (DEFAULT_ODE_SOLVER_OPTIONS..., progress_bar = Val(false))
+    kwargs2 = _merge_saveat(tlist, e_ops, default_values; kwargs...)
     kwargs3 = _generate_mcsolve_kwargs(ψ0, T, e_ops, tlist, c_ops, jump_callback, rng, kwargs2)
 
     return sesolveProblem(H_eff_evo, ψ0, tlist; params = params, kwargs3...)

src/time_evolution/mesolve.jl

@@ -79,11 +79,7 @@ function mesolveProblem(
     ρ0 = to_dense(_CType(T), mat2vec(ket2dm(ψ0).data)) # Convert it to dense vector with complex element type
     L = L_evo.data
 
-    is_empty_e_ops = (e_ops isa Nothing) ? true : isempty(e_ops)
-
-    saveat = is_empty_e_ops ? tlist : [tlist[end]]
-    default_values = (DEFAULT_ODE_SOLVER_OPTIONS..., saveat = saveat)
-    kwargs2 = merge(default_values, kwargs)
+    kwargs2 = _merge_saveat(tlist, e_ops, DEFAULT_ODE_SOLVER_OPTIONS; kwargs...)
     kwargs3 = _generate_se_me_kwargs(e_ops, makeVal(progress_bar), tlist, kwargs2, SaveFuncMESolve)
 
     tspan = (tlist[1], tlist[end])

src/time_evolution/sesolve.jl

@@ -69,11 +69,7 @@ function sesolveProblem(
     ψ0 = to_dense(_CType(T), get_data(ψ0)) # Convert it to dense vector with complex element type
     U = H_evo.data
 
-    is_empty_e_ops = (e_ops isa Nothing) ? true : isempty(e_ops)
-
-    saveat = is_empty_e_ops ? tlist : [tlist[end]]
-    default_values = (DEFAULT_ODE_SOLVER_OPTIONS..., saveat = saveat)
-    kwargs2 = merge(default_values, kwargs)
+    kwargs2 = _merge_saveat(tlist, e_ops, DEFAULT_ODE_SOLVER_OPTIONS; kwargs...)
     kwargs3 = _generate_se_me_kwargs(e_ops, makeVal(progress_bar), tlist, kwargs2, SaveFuncSESolve)
 
     tspan = (tlist[1], tlist[end])

src/time_evolution/smesolve.jl

@@ -112,11 +112,7 @@ function smesolveProblem(
 
     p = (progr = progr, times = tlist, Hdims = dims, n_sc_ops = length(sc_ops), params...)
 
-    is_empty_e_ops = (e_ops isa Nothing) ? true : isempty(e_ops)
-
-    saveat = is_empty_e_ops ? tlist : [tlist[end]]
-    default_values = (DEFAULT_SDE_SOLVER_OPTIONS..., saveat = saveat)
-    kwargs2 = merge(default_values, kwargs)
+    kwargs2 = _merge_saveat(tlist, e_ops, DEFAULT_SDE_SOLVER_OPTIONS; kwargs...)
     kwargs3 = _generate_se_me_kwargs(e_ops, makeVal(progress_bar), tlist, kwargs2, SaveFuncMESolve)
 
     tspan = (tlist[1], tlist[end])

src/time_evolution/ssesolve.jl

@@ -186,11 +186,7 @@ function ssesolveProblem(
 
     p = (progr = progr, times = tlist, Hdims = dims, n_sc_ops = length(sc_ops), params...)
 
-    is_empty_e_ops = (e_ops isa Nothing) ? true : isempty(e_ops)
-
-    saveat = is_empty_e_ops ? tlist : [tlist[end]]
-    default_values = (DEFAULT_SDE_SOLVER_OPTIONS..., saveat = saveat)
-    kwargs2 = merge(default_values, kwargs)
+    kwargs2 = _merge_saveat(tlist, e_ops, DEFAULT_SDE_SOLVER_OPTIONS; kwargs...)
     kwargs3 = _generate_se_me_kwargs(e_ops, makeVal(progress_bar), tlist, kwargs2, SaveFuncSSESolve)
     kwargs4 = _ssesolve_add_normalize_cb(kwargs3)
 

src/time_evolution/time_evolution.jl

@@ -230,6 +230,23 @@ function _check_tlist(tlist, T::Type)
     return tlist2
 end
 
+#######################################
+
+function _merge_saveat(tlist, e_ops, default_options; kwargs...)
+    is_empty_e_ops = isnothing(e_ops) ? true : isempty(e_ops)
+    saveat = is_empty_e_ops ? tlist : [tlist[end]]
+    default_values = (default_options..., saveat = saveat)
+    kwargs2 = merge(default_values, kwargs)
+
+    # DifferentialEquations.jl has this weird save_end setting
+    # So we need to do this to make sure it's consistent
+    haskey(kwargs, :save_end) && return kwargs2
+    isempty(kwargs2.saveat) && return kwargs2
+
+    save_end = tlist[end] in kwargs2.saveat
+    return merge(kwargs2, (save_end = save_end,))
+end
+
 #######################################
 #=
 Helpers for handling output of ensemble problems.

test/core-test/time_evolution.jl

@@ -26,11 +26,13 @@
 
     @testset "sesolve" begin
         tlist = range(0, 20 * 2π / g, 1000)
+        saveat_idxs = 500:900
+        saveat = tlist[saveat_idxs]
 
         prob = sesolveProblem(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
         sol = sesolve(prob)
         sol2 = sesolve(H, ψ0, tlist, progress_bar = Val(false))
-        sol3 = sesolve(H, ψ0, tlist, e_ops = e_ops, saveat = tlist, progress_bar = Val(false))
+        sol3 = sesolve(H, ψ0, tlist, e_ops = e_ops, saveat = saveat, progress_bar = Val(false))
         sol_string = sprint((t, s) -> show(t, "text/plain", s), sol)
         sol_string2 = sprint((t, s) -> show(t, "text/plain", s), sol2)
 
@@ -48,8 +50,9 @@
         @test length(sol2.states) == length(tlist)
         @test sol2.expect === nothing
         @test length(sol3.times) == length(tlist)
-        @test length(sol3.states) == length(tlist)
+        @test length(sol3.states) == length(saveat)
         @test size(sol3.expect) == (length(e_ops), length(tlist))
+        @test sol.expect[1, saveat_idxs] ≈ expect(e_ops[1], sol3.states) atol = 1e-6
         @test sol_string ==
               "Solution of time evolution\n" *
               "(return code: $(sol.retcode))\n" *
@@ -92,18 +95,20 @@
             @inferred sesolveProblem(H, ψ0_int, 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 = e_ops, saveat = saveat, 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, ssesolve and smesolve" begin
         tlist = range(0, 10 / γ, 100)
+        saveat_idxs = 50:90
+        saveat = tlist[saveat_idxs]
 
         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, ψ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))
+        sol_me3 = mesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, saveat = saveat, 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(
@@ -116,14 +121,14 @@
             progress_bar = Val(false),
             jump_callback = DiscreteLindbladJumpCallback(),
         )
-        sol_mc_states = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, saveat = tlist, progress_bar = Val(false))
+        sol_mc_states = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, saveat = saveat, progress_bar = Val(false))
         sol_mc_states2 = mcsolve(
             H,
             ψ0,
             tlist,
             c_ops,
             ntraj = 500,
-            saveat = tlist,
+            saveat = saveat,
             progress_bar = Val(false),
             jump_callback = DiscreteLindbladJumpCallback(),
         )
@@ -147,8 +152,8 @@
         @test prob_mc.prob.f.f isa MatrixOperator
         @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, vec(expect_mc_states_mean) .- vec(sol_me.expect[1, saveat_idxs])) / length(tlist) < 0.1
+        @test sum(abs, vec(expect_mc_states_mean2) .- vec(sol_me.expect[1, saveat_idxs])) / length(tlist) < 0.1
         @test sum(abs, sol_sse.expect .- sol_me.expect) / length(tlist) < 0.1
         @test sum(abs, sol_sme.expect .- sol_me.expect) / length(tlist) < 0.1
         @test length(sol_me.times) == length(tlist)
@@ -158,8 +163,9 @@
         @test length(sol_me2.states) == length(tlist)
         @test sol_me2.expect === nothing
         @test length(sol_me3.times) == length(tlist)
-        @test length(sol_me3.states) == length(tlist)
+        @test length(sol_me3.states) == length(saveat)
         @test size(sol_me3.expect) == (length(e_ops), length(tlist))
+        @test sol_me3.expect[1, saveat_idxs] ≈ expect(e_ops[1], sol_me3.states) atol = 1e-6
         @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)