add sesolve_map

Author: ytdHuang

src/time_evolution/sesolve.jl

@@ -1,4 +1,4 @@
-export sesolveProblem, sesolve
+export sesolveProblem, sesolve, sesolve_map
 
 _sesolve_make_U_QobjEvo(H) = -1im * QuantumObjectEvolution(H, type = Operator())
 
@@ -157,12 +157,16 @@ end
 function sesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit5(); kwargs...)
     sol = solve(prob.prob, alg; kwargs...)
 
-    ψt = map(ϕ -> QuantumObject(ϕ, type = Ket(), dims = prob.dimensions), sol.u)
+    return _gen_sesolve_solution(sol, prob.times, prob.dimensions)
+end
+
+function _gen_sesolve_solution(sol, times, dimensions)
+    ψt = map(ϕ -> QuantumObject(ϕ, type = Ket(), dims = dimensions), sol.u)
 
     kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
 
     return TimeEvolutionSol(
-        prob.times,
+        times,
         sol.t,
         ψt,
         _get_expvals(sol, SaveFuncSESolve),
@@ -172,3 +176,55 @@ function sesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit
         kwargs.reltol,
     )
 end
+
+@doc raw"""
+    sesolve_map
+
+(TBA)
+"""
+function sesolve_map(
+    H::Union{AbstractQuantumObject{Operator},Tuple},
+    ψ0::Vector{<:QuantumObject{Ket}},
+    tlist::AbstractVector;
+    alg::OrdinaryDiffEqAlgorithm = Tsit5(),
+    ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
+    e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
+    params = [NullParameters()],
+    progress_bar::Union{Val,Bool} = Val(true),
+    kwargs...,
+)
+    # mapping initial states and parameters
+    ψ0_iter = map(get_data, ψ0)
+    iter = collect(Iterators.product(ψ0_iter, params...))
+    ntraj = length(iter)
+
+    # we disable the progress bar of the sesolveProblem because we use a global progress bar for all the trajectories
+    prob = sesolveProblem(
+        H,
+        first(ψ0),
+        tlist;
+        e_ops = e_ops,
+        params = first(iter)[2:end],
+        progress_bar = Val(false),
+        kwargs...,
+    )
+
+    # generate and solve ensemble problem
+    _output_func = _ensemble_dispatch_output_func(ensemblealg, progress_bar, ntraj) # setup global progress bar
+    ens_prob = TimeEvolutionProblem(
+        EnsembleProblem(
+            prob.prob,
+            prob_func = (prob, i, repeat) -> remake(prob, u0 = iter[i][1], p = iter[i][2:end]),
+            safetycopy = false,
+        ),
+        prob.times,
+        prob.dimensions,
+        (progr = _output_func[2], channel = _output_func[3]),
+    )
+    sol = _ensemble_dispatch_solve(ens_prob, alg, ensemblealg, ntraj)
+
+    # handle solution and make it become an Array of TimeEvolutionSol
+    return reshape(map(i -> _gen_sesolve_solution(sol[:, i], prob.times, prob.dimensions), eachindex(sol)), size(iter))
+end
+sesolve_map(H::Union{AbstractQuantumObject{Operator},Tuple}, ψ0::QuantumObject{Ket}, tlist::AbstractVector; kwargs...) =
+    sesolve_map(H, [ψ0], tlist; kwargs...)

src/time_evolution/time_evolution.jl

@@ -358,12 +358,14 @@ This is very useful especially for dispatching which method to use to update the
 =#
 
 # Output function with progress bar update
+_ensemble_output_func_progress(sol, i, progr, ::Nothing) = next!(progr)
 function _ensemble_output_func_progress(sol, i, progr, output_func)
     next!(progr)
     return output_func(sol, i)
 end
 
 # Output function with distributed channel update for progress bar
+_ensemble_output_func_distributed(sol, i, channel, ::Nothing) = put!(channel, true)
 function _ensemble_output_func_distributed(sol, i, channel, output_func)
     put!(channel, true)
     return output_func(sol, i)
@@ -373,7 +375,7 @@ function _ensemble_dispatch_output_func(
     ::ET,
     progress_bar,
     ntraj,
-    output_func,
+    output_func = nothing,
 ) where {ET<:Union{EnsembleSerial,EnsembleThreads}}
     if getVal(progress_bar)
         progr = ProgressBar(ntraj, enable = getVal(progress_bar))
@@ -387,7 +389,7 @@ function _ensemble_dispatch_output_func(
     ::ET,
     progress_bar,
     ntraj,
-    output_func,
+    output_func = nothing,
 ) where {ET<:Union{EnsembleSplitThreads,EnsembleDistributed}}
     if getVal(progress_bar)
         progr = ProgressBar(ntraj, enable = getVal(progress_bar))

test/core-test/time_evolution.jl

@@ -147,6 +147,48 @@ end
     end
 end
 
+@testitem "sesolve_map" setup=[TESetup] begin
+
+    # Get parameters from TESetup to simplify the code
+    a = TESetup.a
+    σz = TESetup.σz
+    σm = TESetup.σm
+    ψ0 = TESetup.ψ0
+    e_ops = TESetup.e_ops
+
+    g = 0.01
+    ω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)
+
+    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))
+
+    @test size(sols1) == (1, 3, 4)
+    @test typeof(sols1) isa Array{<:TimeEvolutionSol}
+    @test size(sols2) == (2, 3, 4)
+    @test typeof(sols2) isa Array{<:TimeEvolutionSol}
+    for (i, ωc) in enumerate(ωc_list)
+        for (j, ωq) in enumerate(ωq_list)
+            sol = sols1[1, 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
+        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))
+    end
+end
+
 @testitem "mesolve" setup=[TESetup] begin
     using SciMLOperators