Working mesolve

Author: albertomercurio

src/QuantumToolbox.jl

@@ -95,7 +95,7 @@ include("qobj/synonyms.jl")
 # time evolution
 include("time_evolution/time_evo_parameters.jl")
 include("time_evolution/time_evolution.jl")
-include("time_evolution/callback_helpers.jl")
+include("time_evolution/callback_helpers/callback_helpers.jl")
 include("time_evolution/mesolve.jl")
 include("time_evolution/lr_mesolve.jl")
 include("time_evolution/sesolve.jl")

src/qobj/eigsolve.jl

@@ -391,14 +391,15 @@ function eigsolve_al(
     kwargs...,
 ) where {DT1,HOpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject}}
     L_evo = _mesolve_make_L_QobjEvo(H, c_ops)
-    prob = mesolveProblem(
-        L_evo,
-        QuantumObject(ρ0, type = Operator, dims = H.dims),
-        [zero(T), T];
-        params = params,
-        progress_bar = Val(false),
-        kwargs...,
-    )
+    prob =
+        mesolveProblem(
+            L_evo,
+            QuantumObject(ρ0, type = Operator, dims = H.dims),
+            [zero(T), T];
+            params = params,
+            progress_bar = Val(false),
+            kwargs...,
+        ).prob
     integrator = init(prob, alg)
 
     # prog = ProgressUnknown(desc="Applications:", showspeed = true, enabled=progress)

src/time_evolution/callback_helpers/callback_helpers.jl

@@ -0,0 +1,71 @@
+#=
+This file contains helper functions for callbacks. The affect! function are defined taking advantage of the Julia struct, which allows to store some cache exclusively for the callback.
+=#
+
+include("sesolve_callback_helpers.jl")
+include("mesolve_callback_helpers.jl")
+include("mcsolve_callback_helpers.jl")
+
+##
+
+# Multiple dispatch depending on the progress_bar and e_ops types
+function _generate_se_me_kwargs(e_ops, progress_bar, tlist, kwargs, method)
+    cb = _generate_save_callback(e_ops, tlist, progress_bar, method)
+    return _merge_kwargs_with_callback(kwargs, cb)
+end
+_generate_se_me_kwargs(e_ops::Nothing, progress_bar::Val{false}, tlist, kwargs, method) = kwargs
+
+function _merge_kwargs_with_callback(kwargs, cb)
+    kwargs2 =
+        haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(cb, kwargs.callback),)) :
+        merge(kwargs, (callback = cb,))
+
+    return kwargs2
+end
+
+function _generate_save_callback(e_ops, tlist, progress_bar, method)
+    e_ops_data = e_ops isa Nothing ? nothing : get_data.(e_ops)
+
+    progr = getVal(progress_bar) ? ProgressBar(length(tlist), enable = getVal(progress_bar)) : nothing
+
+    expvals = e_ops isa Nothing ? nothing : Array{ComplexF64}(undef, length(e_ops), length(tlist))
+
+    _save_affect! = method(e_ops_data, progr, Ref(1), expvals)
+    return PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
+end
+
+# When e_ops is Nothing. Common for both mesolve and sesolve
+function _save_func(integrator, progr)
+    next!(progr)
+    u_modified!(integrator, false)
+    return nothing
+end
+
+# When progr is Nothing. Common for both mesolve and sesolve
+function _save_func(integrator, progr::Nothing)
+    u_modified!(integrator, false)
+    return nothing
+end
+
+##
+
+# Get the e_ops from a given AbstractODESolution. Valid for `sesolve`, `mesolve` and `ssesolve`.
+function _se_me_sse_get_expvals(sol::AbstractODESolution)
+    kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple to support Zygote.jl
+    if hasproperty(kwargs, :callback)
+        return _se_me_sse_get_expvals(kwargs.callback)
+    else
+        return nothing
+    end
+end
+function _se_me_sse_get_expvals(cb::CallbackSet)
+    _cb = cb.discrete_callbacks[1]
+    return _se_me_sse_get_expvals(_cb)
+end
+_se_me_sse_get_expvals(cb::DiscreteCallback) =
+    if (cb.affect! isa SaveFuncSESolve) || (cb.affect! isa SaveFuncMESolve)
+        return cb.affect!.expvals
+    else
+        return nothing
+    end
+_se_me_sse_get_expvals(cb::ContinuousCallback) = nothing

src/time_evolution/callback_helpers.jl renamed to src/time_evolution/callback_helpers/mcsolve_callback_helpers.jl

@@ -1,77 +1,7 @@
 #=
-This file contains helper functions for callbacks. The affect! function are defined taking advantage of the Julia struct, which allows to store some cache exclusively for the callback.
+Helper functions for the mcsolve callbacks.
 =#
 
-########## SESOLVE ##########
-
-struct SaveFuncSESolve{TE,PT<:Union{Nothing,ProgressBar},IT,TEXPV<:Union{Nothing,AbstractMatrix}}
-    e_ops::TE
-    progr::PT
-    iter::IT
-    expvals::TEXPV
-end
-
-(f::SaveFuncSESolve)(integrator) = _save_func_sesolve(integrator, f.e_ops, f.progr, f.iter, f.expvals)
-(f::SaveFuncSESolve{Nothing})(integrator) = _save_func_sesolve(integrator, f.progr)
-
-##
-
-# When e_ops is Nothing
-function _save_func_sesolve(integrator, progr)
-    next!(progr)
-    u_modified!(integrator, false)
-    return nothing
-end
-
-# When progr is Nothing
-function _save_func_sesolve(integrator, progr::Nothing)
-    u_modified!(integrator, false)
-    return nothing
-end
-
-# When e_ops is a list of operators
-function _save_func_sesolve(integrator, e_ops, progr, iter, expvals)
-    ψ = integrator.u
-    _expect = op -> dot(ψ, op, ψ)
-    @. expvals[:, iter[]] = _expect(e_ops)
-    iter[] += 1
-
-    return _save_func_sesolve(integrator, progr)
-end
-
-function _generate_sesolve_callback(e_ops, tlist, progress_bar)
-    e_ops_data = e_ops isa Nothing ? nothing : get_data.(e_ops)
-
-    progr = getVal(progress_bar) ? ProgressBar(length(tlist), enable = getVal(progress_bar)) : nothing
-
-    expvals = e_ops isa Nothing ? nothing : Array{ComplexF64}(undef, length(e_ops), length(tlist))
-
-    _save_affect! = SaveFuncSESolve(e_ops_data, progr, Ref(1), expvals)
-    return PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
-end
-
-function _sesolve_get_expvals(sol::AbstractODESolution)
-    kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple to support Zygote.jl
-    if hasproperty(kwargs, :callback)
-        return _sesolve_get_expvals(kwargs.callback)
-    else
-        return nothing
-    end
-end
-function _sesolve_get_expvals(cb::CallbackSet)
-    _cb = cb.discrete_callbacks[1]
-    return _sesolve_get_expvals(_cb)
-end
-_sesolve_get_expvals(cb::DiscreteCallback) =
-    if cb.affect! isa SaveFuncSESolve
-        return cb.affect!.expvals
-    else
-        return nothing
-    end
-_sesolve_get_expvals(cb::ContinuousCallback) = nothing
-
-########## MCSOLVE ##########
-
 struct SaveFuncMCSolve{TE,IT,TEXPV}
     e_ops::TE
     iter::IT

src/time_evolution/callback_helpers/mesolve_callback_helpers.jl

@@ -0,0 +1,29 @@
+#=
+Helper functions for the mesolve callbacks.
+=#
+
+struct SaveFuncMESolve{TE,PT<:Union{Nothing,ProgressBar},IT,TEXPV<:Union{Nothing,AbstractMatrix}}
+    e_ops::TE
+    progr::PT
+    iter::IT
+    expvals::TEXPV
+end
+
+(f::SaveFuncMESolve)(integrator) = _save_func_mesolve(integrator, f.e_ops, f.progr, f.iter, f.expvals)
+(f::SaveFuncMESolve{Nothing})(integrator) = _save_func(integrator, f.progr)
+
+##
+
+# When e_ops is a list of operators
+function _save_func_mesolve(integrator, e_ops, progr, iter, expvals)
+    # This is equivalent to tr(op * ρ), when both are matrices.
+    # The advantage of using this convention is that We don't need
+    # to reshape u to make it a matrix, but we reshape the e_ops once.
+
+    ρ = integrator.u
+    _expect = op -> dot(op, ρ)
+    @. expvals[:, iter[]] = _expect(e_ops)
+    iter[] += 1
+
+    return _save_func(integrator, progr)
+end

src/time_evolution/callback_helpers/sesolve_callback_helpers.jl

@@ -0,0 +1,25 @@
+#=
+Helper functions for the sesolve callbacks.
+=#
+
+struct SaveFuncSESolve{TE,PT<:Union{Nothing,ProgressBar},IT,TEXPV<:Union{Nothing,AbstractMatrix}}
+    e_ops::TE
+    progr::PT
+    iter::IT
+    expvals::TEXPV
+end
+
+(f::SaveFuncSESolve)(integrator) = _save_func_sesolve(integrator, f.e_ops, f.progr, f.iter, f.expvals)
+(f::SaveFuncSESolve{Nothing})(integrator) = _save_func(integrator, f.progr) # Common for both mesolve and sesolve
+
+##
+
+# When e_ops is a list of operators
+function _save_func_sesolve(integrator, e_ops, progr, iter, expvals)
+    ψ = integrator.u
+    _expect = op -> dot(ψ, op, ψ)
+    @. expvals[:, iter[]] = _expect(e_ops)
+    iter[] += 1
+
+    return _save_func(integrator, progr)
+end