Fix dsf_mcsolve

Author: albertomercurio

src/QuantumToolbox.jl

@@ -35,7 +35,8 @@ import SciMLBase:
     CallbackSet,
     ContinuousCallback,
     DiscreteCallback,
-    AbstractSciMLProblem
+    AbstractSciMLProblem,
+    AbstractODEIntegrator
 import StochasticDiffEq: StochasticDiffEqAlgorithm, SRA1
 import SciMLOperators:
     SciMLOperators,

src/time_evolution/callback_helpers.jl

@@ -12,6 +12,8 @@ end
 (f::SaveFuncSESolve)(integrator) = _save_func_sesolve(integrator, f.e_ops, f.is_empty_e_ops)
 (f::SaveFuncSESolve{Nothing})(integrator) = _save_func_sesolve(integrator)
 
+##
+
 # When e_ops is Nothing
 function _save_func_sesolve(integrator)
     next!(integrator.p.progr)
@@ -25,15 +27,15 @@ function _save_func_sesolve(integrator, e_ops, is_empty_e_ops)
     progr = integrator.p.progr
     if !is_empty_e_ops
         ψ = integrator.u
-        _expect = op -> dot(ψ, get_data(op), ψ)
+        _expect = op -> dot(ψ, op, ψ)
         @. expvals[:, progr.counter[]+1] = _expect(e_ops)
     end
     return _save_func_sesolve(integrator)
 end
 
 function _generate_sesolve_callback(e_ops, tlist)
     is_empty_e_ops = e_ops isa Nothing ? true : isempty(e_ops)
-    _save_affect! = SaveFuncSESolve(e_ops, is_empty_e_ops)
+    _save_affect! = SaveFuncSESolve(get_data.(e_ops), is_empty_e_ops)
     return PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
 end
 
@@ -46,12 +48,12 @@ end
 
 (f::SaveFuncMCSolve)(integrator) = _save_func_mcsolve(integrator, f.e_ops, f.is_empty_e_ops)
 
-struct LindbladJumpAffect!{T1,T2}
+struct LindbladJump{T1,T2}
     c_ops::T1
     c_ops_herm::T2
 end
 
-(f::LindbladJumpAffect!)(integrator) = _lindblad_jump_affect!(integrator, f.c_ops, f.c_ops_herm)
+(f::LindbladJump)(integrator) = _lindblad_jump_affect!(integrator, f.c_ops, f.c_ops_herm)
 
 ##
 
@@ -63,7 +65,7 @@ function _save_func_mcsolve(integrator, e_ops, is_empty_e_ops)
         copyto!(cache_mc, integrator.u)
         normalize!(cache_mc)
         ψ = cache_mc
-        _expect = op -> dot(ψ, get_data(op), ψ)
+        _expect = op -> dot(ψ, op, ψ)
         @. expvals[:, progr.counter[]+1] = _expect(e_ops)
     end
     next!(progr)
@@ -75,7 +77,7 @@ function _generate_mcsolve_kwargs(e_ops, tlist, c_ops, jump_callback, kwargs)
     c_ops_data = get_data.(c_ops)
     c_ops_herm_data = map(op -> op' * op, c_ops_data)
 
-    _affect! = LindbladJumpAffect!(c_ops_data, c_ops_herm_data)
+    _affect! = LindbladJump(c_ops_data, c_ops_herm_data)
 
     if jump_callback isa DiscreteLindbladJumpCallback
         cb1 = DiscreteCallback(_mcsolve_discrete_condition, _affect!, save_positions = (false, false))
@@ -96,8 +98,8 @@ function _generate_mcsolve_kwargs(e_ops, tlist, c_ops, jump_callback, kwargs)
         return kwargs2
     else
         is_empty_e_ops = isempty(e_ops)
-        _save_affect! = SaveFuncMCSolve(e_ops, is_empty_e_ops)
-        cb2 = PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
+        _save_affect! = SaveFuncMCSolve(get_data.(e_ops), is_empty_e_ops)
+        cb2 = _PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
         kwargs2 =
             haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(cb1, cb2, kwargs.callback),)) :
             merge(kwargs, (callback = CallbackSet(cb1, cb2),))
@@ -146,3 +148,74 @@ _mcsolve_continuous_condition(u, t, integrator) =
 
 _mcsolve_discrete_condition(u, t, integrator) =
     @inbounds real(dot(u, u)) < real(integrator.p.mcsolve_params.random_n[1])
+
+#=
+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_set = integrator.opts.callback # This is supposed to be a CallbackSet
+    (cb_set isa CallbackSet) || throw(ArgumentError("The callback must be a CallbackSet."))
+    cb = isempty(cb_set.continuous_callbacks) ? cb_set.discrete_callback[1] : cb_set.continuous_callbacks[1]
+    return cb.affect!.c_ops, cb.affect!.c_ops_herm
+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_set = integrator.opts.callback # This is supposed to be a CallbackSet
+    (cb_set isa CallbackSet) || throw(ArgumentError("The callback must be a CallbackSet."))
+    cb = length(cb_set.continuous_callbacks) > 0 ? cb_set.discrete_callbacks[1] : cb_set.discrete_callbacks[2]
+    return cb.affect!.e_ops
+end
+
+## Temporary function to avoid errors. Waiting for the PR In DiffEqCallbacks.jl to be merged.
+
+import SciMLBase: INITIALIZE_DEFAULT, add_tstop!
+
+function _PresetTimeCallback(
+    tstops,
+    user_affect!;
+    initialize = INITIALIZE_DEFAULT,
+    filter_tstops = true,
+    sort_inplace = false,
+    kwargs...,
+)
+    if !(tstops isa AbstractVector) && !(tstops isa Number)
+        throw(ArgumentError("tstops must either be a number or a Vector. Was $tstops"))
+    end
+
+    tstops = tstops isa Number ? [tstops] : (sort_inplace ? sort!(tstops) : sort(tstops))
+
+    condition = let
+        function (u, t, integrator)
+            if hasproperty(integrator, :dt)
+                insorted(t, tstops) && (integrator.t - integrator.dt) != integrator.t
+            else
+                insorted(t, tstops)
+            end
+        end
+    end
+
+    # Initialization: first call to `f` should be *before* any time steps have been taken:
+    initialize_preset = function (c, u, t, integrator)
+        initialize(c, u, t, integrator)
+
+        if filter_tstops
+            tdir = integrator.tdir
+            tspan = integrator.sol.prob.tspan
+            _tstops = tstops[@. tdir * tspan[1] < tdir * tstops < tdir * tspan[2]]
+        else
+            _tstops = tstops
+        end
+        for tstop in _tstops
+            add_tstop!(integrator, tstop)
+        end
+        if t ∈ tstops
+            user_affect!(integrator)
+        end
+    end
+    return DiscreteCallback(condition, user_affect!; initialize = initialize_preset, kwargs...)
+end

src/time_evolution/time_evolution_dynamical.jl

@@ -489,9 +489,9 @@ end
 # Dynamical Shifted Fock mcsolve
 
 function _DSF_mcsolve_Condition(u, t, integrator)
-    internal_params = integrator.p
-    op_list = internal_params.op_list
-    δα_list = internal_params.δα_list
+    params = integrator.p
+    op_list = params.op_list
+    δα_list = params.δα_list
 
     ψt = u
 
@@ -508,20 +508,24 @@ function _DSF_mcsolve_Condition(u, t, integrator)
 end
 
 function _DSF_mcsolve_Affect!(integrator)
-    internal_params = integrator.p
-    op_list = internal_params.op_list
-    αt_list = internal_params.αt_list
-    δα_list = internal_params.δα_list
-    H = internal_params.H_fun
-    c_ops = internal_params.c_ops_fun
-    e_ops = internal_params.e_ops_fun
-    e_ops0 = internal_params.e_ops_mc
-    c_ops0 = internal_params.c_ops
-    ψt = internal_params.dsf_cache1
-    dsf_cache = internal_params.dsf_cache2
-    expv_cache = internal_params.expv_cache
-    dsf_params = internal_params.dsf_params
-    dsf_displace_cache_full = internal_params.dsf_displace_cache_full
+    params = integrator.p
+    op_list = params.op_list
+    αt_list = params.αt_list
+    δα_list = params.δα_list
+    H = params.H_fun
+    c_ops = params.c_ops_fun
+    e_ops = params.e_ops_fun
+    ψt = params.dsf_cache1
+    dsf_cache = params.dsf_cache2
+    expv_cache = params.expv_cache
+    dsf_params = params.dsf_params
+    dsf_displace_cache_full = params.dsf_displace_cache_full
+
+    # e_ops0 = params.e_ops
+    # c_ops0 = params.c_ops
+
+    e_ops0 = _mcsolve_get_e_ops(integrator)
+    c_ops0, c_ops0_herm = _mcsolve_get_c_ops(integrator)
 
     copyto!(ψt, integrator.u)
     normalize!(ψt)
@@ -561,42 +565,58 @@ function _DSF_mcsolve_Affect!(integrator)
     op_l2 = op_list .+ αt_list
     e_ops2 = e_ops(op_l2, dsf_params)
     c_ops2 = c_ops(op_l2, dsf_params)
+
+    ## By copying the data, we are assuming that the variables are Vectors and not Tuple
     @. e_ops0 = get_data(e_ops2)
     @. c_ops0 = get_data(c_ops2)
-    H_nh = lmul!(convert(eltype(ψt), 0.5im), mapreduce(op -> op' * op, +, c_ops0))
+    c_ops0_herm .= map(op -> op' * op, c_ops0)
+
+    H_nh = convert(eltype(ψt), 0.5im) * sum(c_ops0_herm)
     # By doing this, we are assuming that the system is time-independent and f is a MatrixOperator
     copyto!(integrator.f.f.A, lmul!(-1im, H(op_l2, dsf_params).data - H_nh))
     return u_modified!(integrator, true)
 end
 
 function _dsf_mcsolve_prob_func(prob, i, repeat)
-    internal_params = prob.p
+    params = prob.p
+
+    expvals = similar(params.expvals)
+    progr = ProgressBar(size(expvals, 2), enable = false)
+
+    T = eltype(expvals)
+
+    mcsolve_params = merge(
+        params.mcsolve_params,
+        (
+            random_n = T[rand()],
+            cache_mc = similar(params.mcsolve_params.cache_mc),
+            weights_mc = similar(params.mcsolve_params.weights_mc),
+            cumsum_weights_mc = similar(params.mcsolve_params.weights_mc),
+            jump_times = similar(params.mcsolve_params.jump_times),
+            jump_which = similar(params.mcsolve_params.jump_which),
+            jump_times_which_idx = T[1],
+        ),
+    )
 
     prm = merge(
-        internal_params,
+        params.params,
         (
-            e_ops_mc = deepcopy(internal_params.e_ops_mc),
-            c_ops = deepcopy(internal_params.c_ops),
-            expvals = similar(internal_params.expvals),
-            cache_mc = similar(internal_params.cache_mc),
-            weights_mc = similar(internal_params.weights_mc),
-            cumsum_weights_mc = similar(internal_params.weights_mc),
-            random_n = Ref(rand()),
-            progr_mc = ProgressBar(size(internal_params.expvals, 2), enable = false),
-            jump_times_which_idx = Ref(1),
-            jump_times = similar(internal_params.jump_times),
-            jump_which = similar(internal_params.jump_which),
-            αt_list = copy(internal_params.αt_list),
-            dsf_cache1 = similar(internal_params.dsf_cache1),
-            dsf_cache2 = similar(internal_params.dsf_cache2),
-            expv_cache = copy(internal_params.expv_cache),
-            dsf_displace_cache_full = deepcopy(internal_params.dsf_displace_cache_full), # This brutally copies also the MatrixOperators, and it is inefficient.
+            αt_list = copy(params.params.αt_list),
+            dsf_cache1 = similar(params.params.dsf_cache1),
+            dsf_cache2 = similar(params.params.dsf_cache2),
+            expv_cache = copy(params.params.expv_cache),
+            dsf_displace_cache_full = deepcopy(params.params.dsf_displace_cache_full), # This brutally copies also the MatrixOperators, and it is inefficient.
         ),
     )
 
+    p = TimeEvolutionParameters(prm, expvals, progr, mcsolve_params)
+
     f = deepcopy(prob.f.f)
 
-    return remake(prob, f = f, p = prm)
+    # We need to deepcopy the callbacks because they contain the c_ops and e_ops, which are modified in the affect function
+    cb = deepcopy(prob.kwargs[:callback])
+
+    return remake(prob, f = f, p = p, callback = cb)
 end
 
 function dsf_mcsolveEnsembleProblem(