Align mcsolve _get_expvals with the other solvers

Author: albertomercurio

src/time_evolution/callback_helpers/callback_helpers.jl

@@ -51,6 +51,19 @@ end
 
 ##
 
+#=
+    With this function we extract the e_ops from the SaveFuncMCSolve `affect!` function of the callback of the integrator.
+    This callback can only be a PresetTimeCallback (DiscreteCallback).
+=#
+function _get_e_ops(integrator::AbstractODEIntegrator, method::Type{SF}) where {SF<:AbstractSaveFunc}
+    cb = _get_save_callback(integrator, method)
+    if cb isa Nothing
+        return nothing
+    else
+        return cb.affect!.e_ops
+    end
+end
+
 # Get the e_ops from a given AbstractODESolution. Valid for `sesolve`, `mesolve` and `ssesolve`.
 function _get_expvals(sol::AbstractODESolution, method::Type{SF}) where {SF<:AbstractSaveFunc}
     cb = _get_save_callback(sol, method)
@@ -61,6 +74,11 @@ function _get_expvals(sol::AbstractODESolution, method::Type{SF}) where {SF<:Abs
     end
 end
 
+#=
+    _get_save_callback
+
+Return the Callback that is responsible for saving the expectation values of the system.
+=#
 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)
@@ -74,19 +92,19 @@ _get_save_callback(integrator::AbstractODEIntegrator, method::Type{SF}) where {S
 function _get_save_callback(cb::CallbackSet, method::Type{SF}) where {SF<:AbstractSaveFunc}
     cbs_discrete = cb.discrete_callbacks
     if length(cbs_discrete) > 0
-        idx = _get_save_callback_idx(method)
+        idx = _get_save_callback_idx(cb, method)
         _cb = cb.discrete_callbacks[idx]
         return _get_save_callback(_cb, method)
     else
         return nothing
     end
 end
-function _get_save_callback(cb::DiscreteCallback, method::Type{SF}) where {SF<:AbstractSaveFunc}
-    if typeof(cb.affect!) <: SF
+function _get_save_callback(cb::DiscreteCallback, ::Type{SF}) where {SF<:AbstractSaveFunc}
+    if typeof(cb.affect!) <: AbstractSaveFunc
         return cb
     end
     return nothing
 end
-_get_save_callback(cb::ContinuousCallback, method::Type{SF}) where {SF<:AbstractSaveFunc} = nothing
+_get_save_callback(cb::ContinuousCallback, ::Type{SF}) where {SF<:AbstractSaveFunc} = nothing
 
-_get_save_callback_idx(method) = 1
+_get_save_callback_idx(cb, method) = 1

src/time_evolution/callback_helpers/mcsolve_callback_helpers.jl

@@ -10,6 +10,8 @@ end
 
 (f::SaveFuncMCSolve)(integrator) = _save_func_mcsolve(integrator, f.e_ops, f.iter, f.expvals)
 
+_get_save_callback_idx(cb, ::Type{SaveFuncMCSolve}) = _mcsolve_has_continuous_jump(cb) ? 1 : 2
+
 ##
 struct LindbladJump{
     T1,
@@ -168,37 +170,6 @@ _mcsolve_discrete_condition(u, t, integrator) =
 
 ##
 
-#=
-    _mc_get_save_callback
-
-Return the Callback that is responsible for saving the expectation values of the system.
-=#
-function _mc_get_save_callback(sol::AbstractODESolution)
-    kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple to support Zygote.jl
-    return _mc_get_save_callback(kwargs.callback) # There is always the Jump callback
-end
-_mc_get_save_callback(integrator::AbstractODEIntegrator) = _mc_get_save_callback(integrator.opts.callback)
-function _mc_get_save_callback(cb::CallbackSet)
-    cbs_discrete = cb.discrete_callbacks
-
-    if length(cbs_discrete) > 0
-        idx = _mcsolve_has_continuous_jump(cb) ? 1 : 2
-        _cb = cb.discrete_callbacks[idx]
-        return _mc_get_save_callback(_cb)
-    else
-        return nothing
-    end
-end
-_mc_get_save_callback(cb::DiscreteCallback) =
-    if cb.affect! isa SaveFuncMCSolve
-        return cb
-    else
-        return nothing
-    end
-_mc_get_save_callback(cb::ContinuousCallback) = nothing
-
-##
-
 function _mc_get_jump_callback(sol::AbstractODESolution)
     kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple to support Zygote.jl
     return _mc_get_jump_callback(kwargs.callback) # There is always the Jump callback
@@ -216,8 +187,8 @@ _mc_get_jump_callback(cb::DiscreteCallback) = cb
 ##
 
 #=
-With this function we extract the c_ops and c_ops_herm from the LindbladJump `affect!` function of the callback of the integrator.
-This callback can be a DiscreteLindbladJumpCallback or a ContinuousLindbladJumpCallback.
+    With this function we extract the c_ops and c_ops_herm from the LindbladJump `affect!` function of the callback of the integrator.
+    This callback can be a DiscreteLindbladJumpCallback or a ContinuousLindbladJumpCallback.
 =#
 function _mcsolve_get_c_ops(integrator::AbstractODEIntegrator)
     cb = _mc_get_jump_callback(integrator)
@@ -228,28 +199,6 @@ function _mcsolve_get_c_ops(integrator::AbstractODEIntegrator)
     end
 end
 
-#=
-With this function we extract the e_ops from the SaveFuncMCSolve `affect!` function of the callback of the integrator.
-This callback can only be a PresetTimeCallback (DiscreteCallback).
-=#
-function _mcsolve_get_e_ops(integrator::AbstractODEIntegrator)
-    cb = _mc_get_save_callback(integrator)
-    if cb isa Nothing
-        return nothing
-    else
-        return cb.affect!.e_ops
-    end
-end
-
-function _mcsolve_get_expvals(sol::AbstractODESolution)
-    cb = _mc_get_save_callback(sol)
-    if cb isa Nothing
-        return nothing
-    else
-        return cb.affect!.expvals
-    end
-end
-
 #=
     _mcsolve_initialize_callbacks(prob, tlist)
 

src/time_evolution/callback_helpers/ssesolve_callback_helpers.jl

@@ -14,7 +14,7 @@ end
 
 _get_e_ops_data(e_ops, ::Type{SaveFuncSSESolve}) = get_data.(e_ops)
 
-_get_save_callback_idx(method::SaveFuncSSESolve) = 2 # The first one is the normalization callback
+_get_save_callback_idx(cb, ::Type{SaveFuncSSESolve}) = 1 # The first one is the normalization callback
 
 ##
 

src/time_evolution/mcsolve.jl

@@ -398,9 +398,10 @@ function mcsolve(
 
     dims = ens_prob_mc.dimensions
     _sol_1 = sol[:, 1]
-    _expvals_sol_1 = _mcsolve_get_expvals(_sol_1)
+    _expvals_sol_1 = _get_expvals(_sol_1, SaveFuncMCSolve)
 
-    _expvals_all = _expvals_sol_1 isa Nothing ? nothing : map(i -> _mcsolve_get_expvals(sol[:, i]), eachindex(sol))
+    _expvals_all =
+        _expvals_sol_1 isa Nothing ? nothing : map(i -> _get_expvals(sol[:, i], SaveFuncMCSolve), eachindex(sol))
     expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all, dims = 2) # Stack on dimension 2 to align with QuTiP
     states = map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol))
     col_times = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.col_times, eachindex(sol))

src/time_evolution/time_evolution_dynamical.jl

@@ -522,7 +522,7 @@ function _DSF_mcsolve_Affect!(integrator)
     # e_ops0 = params.e_ops
     # c_ops0 = params.c_ops
 
-    e_ops0 = _mcsolve_get_e_ops(integrator)
+    e_ops0 = _get_e_ops(integrator, SaveFuncMCSEsolve)
     c_ops0, c_ops0_herm = _mcsolve_get_c_ops(integrator)
 
     copyto!(ψt, integrator.u)

test/runtests.jl

@@ -7,23 +7,23 @@ const testdir = dirname(@__FILE__)
 
 # Put core tests in alphabetical order
 core_tests = [
-    # "block_diagonal_form.jl",
-    # "correlations_and_spectrum.jl",
-    # "dynamical_fock_dimension_mesolve.jl",
-    # "dynamical-shifted-fock.jl",
-    # "eigenvalues_and_operators.jl",
-    # "entanglement.jl",
-    # "generalized_master_equation.jl",
-    # "low_rank_dynamics.jl",
-    # "negativity_and_partial_transpose.jl",
-    # "progress_bar.jl",
-    # "quantum_objects.jl",
-    # "quantum_objects_evo.jl",
-    # "states_and_operators.jl",
-    # "steady_state.jl",
+    "block_diagonal_form.jl",
+    "correlations_and_spectrum.jl",
+    "dynamical_fock_dimension_mesolve.jl",
+    "dynamical-shifted-fock.jl",
+    "eigenvalues_and_operators.jl",
+    "entanglement.jl",
+    "generalized_master_equation.jl",
+    "low_rank_dynamics.jl",
+    "negativity_and_partial_transpose.jl",
+    "progress_bar.jl",
+    "quantum_objects.jl",
+    "quantum_objects_evo.jl",
+    "states_and_operators.jl",
+    "steady_state.jl",
     "time_evolution.jl",
-    # "utilities.jl",
-    # "wigner.jl",
+    "utilities.jl",
+    "wigner.jl",
 ]
 
 if (GROUP == "All") || (GROUP == "Core")