Working sesolve case

Author: albertomercurio

src/time_evolution/callback_helpers/callback_helpers.jl

@@ -139,7 +139,7 @@ Return the Callback that is responsible for saving the expectation values of the
 =#
 function _get_save_callback(sol::AbstractODESolution, method::Type{SF}) where {SF<:AbstractSaveFunc}
     kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple to support Zygote.jl
-    if hasproperty(kwargs, :callback)
+    if hasproperty(kwargs, :callback) && !isnothing(kwargs.callback)
         return _get_save_callback(kwargs.callback, method)
     else
         return nothing

src/time_evolution/sesolve.jl

@@ -183,13 +183,13 @@ end
 (TBA)
 """
 function sesolve_map(
-    H::Union{AbstractQuantumObject{Operator},Tuple},
-    ψ0::Vector{<:QuantumObject{Ket}},
+    H::Union{QuantumObjectEvolution{Operator},Tuple},
+    ψ0::AbstractVector{<:QuantumObject{Ket}},
     tlist::AbstractVector;
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
-    params = [NullParameters()],
+    params::Tuple = (NullParameters(),),
     progress_bar::Union{Val,Bool} = Val(true),
     kwargs...,
 )
@@ -214,7 +214,13 @@ function sesolve_map(
     ens_prob = TimeEvolutionProblem(
         EnsembleProblem(
             prob.prob,
-            prob_func = (prob, i, repeat) -> remake(prob, u0 = iter[i][1], p = iter[i][2:end]),
+            prob_func = (prob, i, repeat) -> remake(
+                prob,
+                f = deepcopy(prob.f.f),
+                u0 = iter[i][1],
+                p = iter[i][2:end],
+                callback = haskey(prob.kwargs, :callback) ? deepcopy(prob.kwargs[:callback]) : nothing,
+            ),
             safetycopy = false,
         ),
         prob.times,

test/core-test/time_evolution.jl

@@ -157,35 +157,43 @@ end
     e_ops = TESetup.e_ops
 
     g = 0.01
+
+    ψ_0_e = tensor(fock(N, 0), basis(2, 0))
+    ψ_1_g = tensor(fock(N, 1), basis(2, 1))
+
+    ψ0_list = [ψ_0_e, ψ_1_g]
     ωc_list = [1, 1.01, 1.02]
     ωq_list = [0.96, 0.97, 0.98, 0.99]
+
     tlist = range(0, 20 * 2π / g, 1000)
 
-    ωc(p, t) = p[1]
-    ωq(p, t) = p[2]
-    H = g * (a' * σm + a * σm') + QobjEvo(a' * a, ωc) + QobjEvo(σz / 2, ωq)
+    ωc_fun(p, t) = p[1]
+    ωq_fun(p, t) = p[2]
+    H = QobjEvo(a' * a, ωc_fun) + QobjEvo(σz / 2, ωq_fun) + g * (a' * σm + a * σm')
 
-    sols1 = sesolve_map(H, ψ0, tlist; e_ops = e_ops, params = [ωc_list, ωq_list])
-    sols2 = sesolve_map(H, [ψ0, ψ0], tlist; e_ops = e_ops, params = [ωc_list, ωq_list], progress_bar = Val(false))
+    sols1 = sesolve_map(H, ψ_0_e, tlist; e_ops = e_ops, params = (ωc_list, ωq_list))
+    sols2 = sesolve_map(H, ψ0_list, tlist; e_ops = e_ops, params = (ωc_list, ωq_list), progress_bar = Val(false))
 
     @test size(sols1) == (1, 3, 4)
-    @test typeof(sols1) isa Array{<:TimeEvolutionSol}
+    @test sols1 isa Array{<:TimeEvolutionSol}
     @test size(sols2) == (2, 3, 4)
-    @test typeof(sols2) isa Array{<:TimeEvolutionSol}
+    @test sols2 isa Array{<:TimeEvolutionSol}
     for (i, ωc) in enumerate(ωc_list)
         for (j, ωq) in enumerate(ωq_list)
-            sol = sols1[1, i, j]
+            sol_0_e = sols2[1, i, j]
+            sol_1_g = sols2[2, i, j]
 
             ## Analytical solution for the expectation value of a' * a
             Ω_rabi = sqrt(g^2 + ((ωc - ωq) / 2)^2)
             amp_rabi = g^2 / Ω_rabi^2
 
-            @test sum(abs.(sol.expect[1, :] .- amp_rabi .* sin.(Ω_rabi * tlist) .^ 2)) / length(tlist) < 0.1
+            @test sol_0_e.expect[1, :] ≈ amp_rabi .* sin.(Ω_rabi * tlist) .^ 2 atol = 1e-2
+            @test sol_1_g.expect[1, :] ≈ 1 .- amp_rabi .* sin.(Ω_rabi * tlist) .^ 2 atol = 1e-2
         end
     end
 
     @testset "Type Inference sesolve_map" begin
-        @inferred sesolve_map(H, [ψ0, ψ0], tlist; e_ops = e_ops, params = [ωc_list, ωq_list], progress_bar = Val(false))
+        @inferred sesolve_map(H, ψ0_list, tlist; e_ops = e_ops, params = (ωc_list, ωq_list), progress_bar = Val(false))
     end
 end