Fix dsf_mcsolve

Author: albertomercurio

src/QuantumToolbox.jl

@@ -21,6 +21,7 @@ import SciMLBase:
     reinit!,
     remake,
     u_modified!,
+    ODEFunction,
     ODEProblem,
     SDEProblem,
     EnsembleProblem,

src/time_evolution/mcsolve.jl

@@ -117,6 +117,11 @@ function _mcsolve_generate_statistics(sol, i, states, expvals_all, jump_times, j
     return jump_which[i] = sol_i.prob.p.jump_which
 end
 
+_mcsolve_make_Heff_QobjEvo(H::QuantumObject, c_ops) = QobjEvo(H - 1im * mapreduce(op -> op' * op, +, c_ops) / 2)
+_mcsolve_make_Heff_QobjEvo(H::Tuple, c_ops) = QobjEvo((H..., -1im * mapreduce(op -> op' * op, +, c_ops) / 2))
+_mcsolve_make_Heff_QobjEvo(H::QuantumObjectEvolution, c_ops) =
+    H + QobjEvo(mapreduce(op -> op' * op, +, c_ops), -1im / 2)
+
 @doc raw"""
     mcsolveProblem(H::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject},
         ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
@@ -210,7 +215,7 @@ function mcsolveProblem(
 
     tlist = convert(Vector{_FType(ψ0)}, tlist) # Convert it to support GPUs and avoid type instabilities for OrdinaryDiffEq.jl
 
-    H_eff_evo = QobjEvo(H) + QobjEvo(mapreduce(op -> op' * op, +, c_ops) / 2, -1im)
+    H_eff_evo = _mcsolve_make_Heff_QobjEvo(H, c_ops)
 
     if e_ops isa Nothing
         expvals = Array{ComplexF64}(undef, 0, length(tlist))

src/time_evolution/sesolve.jl

@@ -124,7 +124,7 @@ function sesolveProblem(
     kwargs3 = _generate_sesolve_kwargs(e_ops, makeVal(progress_bar), tlist, kwargs2)
 
     tspan = (tlist[1], tlist[end])
-    return ODEProblem{true,FullSpecialize}(U, ψ0, tspan, p; kwargs3...)
+    return ODEProblem{true}(U, ψ0, tspan, p; kwargs3...)
 end
 
 @doc raw"""

src/time_evolution/time_evolution_dynamical.jl

@@ -572,8 +572,10 @@ function _DSF_mcsolve_Affect!(integrator)
     c_ops2 = c_ops(op_l2, dsf_params)
     @. e_ops0 = get_data(e_ops2)
     @. c_ops0 = get_data(c_ops2)
-    H_eff = H(op_l2, dsf_params).data - lmul!(convert(eltype(ψt), 0.5im), mapreduce(op -> op' * op, +, c_ops0))
-    return mul!(internal_params.U, -1im, H_eff)
+    H_nh = lmul!(convert(eltype(ψt), 0.5im), mapreduce(op -> op' * op, +, c_ops0))
+    # By doing this, we are assuming that the system is time-independent and f is a ScaledOperator
+    copyto!(integrator.f.f.L.A, H(op_l2, dsf_params).data - H_nh)
+    return u_modified!(integrator, true)
 end
 
 function _dsf_mcsolve_prob_func(prob, i, repeat)
@@ -582,7 +584,6 @@ function _dsf_mcsolve_prob_func(prob, i, repeat)
     prm = merge(
         internal_params,
         (
-            U = copy(internal_params.U),
             e_ops_mc = copy(internal_params.e_ops_mc),
             c_ops = copy(internal_params.c_ops),
             expvals = similar(internal_params.expvals),
@@ -605,7 +606,9 @@ function _dsf_mcsolve_prob_func(prob, i, repeat)
         ),
     )
 
-    return remake(prob, p = prm)
+    f = deepcopy(prob.f.f)
+
+    return remake(prob, f = f, p = prm)
 end
 
 function dsf_mcsolveEnsembleProblem(