Make time evolution solvers compatible with automatic differentiation

docs/src/resources/api.md

@@ -181,6 +181,7 @@ qeye
 ## [Time evolution](@id doc-API:Time-evolution)
 
 ```@docs
+TimeEvolutionProblem
 TimeEvolutionSol
 TimeEvolutionMCSol
 TimeEvolutionSSESol

src/QuantumToolbox.jl

@@ -23,17 +23,22 @@ import SciMLBase:
     reinit!,
     remake,
     u_modified!,
+    NullParameters,
     ODEFunction,
     ODEProblem,
     SDEProblem,
     EnsembleProblem,
     EnsembleSerial,
     EnsembleThreads,
+    EnsembleSplitThreads,
     EnsembleDistributed,
     FullSpecialize,
     CallbackSet,
     ContinuousCallback,
-    DiscreteCallback
+    DiscreteCallback,
+    AbstractSciMLProblem,
+    AbstractODEIntegrator,
+    AbstractODESolution
 import StochasticDiffEq: StochasticDiffEqAlgorithm, SRA1
 import SciMLOperators:
     SciMLOperators,
@@ -88,6 +93,7 @@ include("qobj/synonyms.jl")
 
 # time evolution
 include("time_evolution/time_evolution.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/qobj/quantum_object_evo.jl

@@ -269,7 +269,7 @@ Parse the `op_func_list` and generate the data for the `QuantumObjectEvolution`
     quote
         dims = tuple($(dims_expr...))
 
-        length(unique(dims)) == 1 || throw(ArgumentError("The dimensions of the operators must be the same."))
+        allequal(dims) || throw(ArgumentError("The dimensions of the operators must be the same."))
 
         data_expr_const = $qobj_expr_const isa Integer ? $qobj_expr_const : _make_SciMLOperator($qobj_expr_const, α)
 

src/time_evolution/callback_helpers/callback_helpers.jl

@@ -0,0 +1,93 @@
+#=
+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_e_ops_data(e_ops, method)
+
+    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
+
+_get_e_ops_data(e_ops, ::Type{SaveFuncSESolve}) = get_data.(e_ops)
+_get_e_ops_data(e_ops, ::Type{SaveFuncMESolve}) = _generate_mesolve_e_op.(e_ops)
+
+_generate_mesolve_e_op(op) = mat2vec(adjoint(get_data(op)))
+
+##
+
+# 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)
+    cb = _se_me_sse_get_save_callback(sol)
+    if cb isa Nothing
+        return nothing
+    else
+        return cb.affect!.expvals
+    end
+end
+
+function _se_me_sse_get_save_callback(sol::AbstractODESolution)
+    kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple to support Zygote.jl
+    if hasproperty(kwargs, :callback)
+        return _se_me_sse_get_save_callback(kwargs.callback)
+    else
+        return nothing
+    end
+end
+_se_me_sse_get_save_callback(integrator::AbstractODEIntegrator) = _se_me_sse_get_save_callback(integrator.opts.callback)
+function _se_me_sse_get_save_callback(cb::CallbackSet)
+    cbs_discrete = cb.discrete_callbacks
+    if length(cbs_discrete) > 0
+        _cb = cb.discrete_callbacks[1]
+        return _se_me_sse_get_save_callback(_cb)
+    else
+        return nothing
+    end
+end
+_se_me_sse_get_save_callback(cb::DiscreteCallback) =
+    if (cb.affect! isa SaveFuncSESolve) || (cb.affect! isa SaveFuncMESolve)
+        return cb
+    else
+        return nothing
+    end
+_se_me_sse_get_save_callback(cb::ContinuousCallback) = nothing