[no ci] Make dWdt in-place

Author: albertomercurio

src/time_evolution/callback_helpers/callback_helpers.jl

@@ -74,7 +74,9 @@ function _generate_stochastic_save_callback(e_ops, sc_ops, tlist, store_measurem
     expvals = e_ops isa Nothing ? nothing : Array{ComplexF64}(undef, length(e_ops), length(tlist))
     m_expvals = getVal(store_measurement) ? Array{Float64}(undef, length(sc_ops), length(tlist) - 1) : nothing
 
-    _save_func = method(store_measurement, e_ops_data, m_ops_data, progr, Ref(1), expvals, m_expvals, tlist)
+    _save_func_cache = Array{Float64}(undef, length(sc_ops))
+    _save_func =
+        method(store_measurement, e_ops_data, m_ops_data, progr, Ref(1), expvals, m_expvals, tlist, _save_func_cache)
     return FunctionCallingCallback(_save_func, funcat = tlist)
 end
 
@@ -169,11 +171,11 @@ _get_save_callback_idx(cb, method) = 1
 
 # TODO: Add some cache mechanism to avoid memory allocations
 # TODO: To improve. See https://github.com/SciML/DiffEqNoiseProcess.jl/issues/214
-function _homodyne_dWdt(integrator, tlist, iter)
+function _homodyne_dWdt!(dWdt_cache, integrator, tlist, iter)
     idx = findfirst(>=(tlist[iter[]-1]), integrator.W.t)
 
     # We are assuming that the last element is tlist[iter[]]
-    @inbounds _dWdt = (integrator.W.u[end] .- integrator.W.u[idx]) ./ (integrator.W.t[end] - integrator.W.t[idx])
+    @inbounds dWdt_cache .= (integrator.W.u[end] .- integrator.W.u[idx]) ./ (integrator.W.t[end] - integrator.W.t[idx])
 
-    return _dWdt
+    return nothing
 end

src/time_evolution/callback_helpers/smesolve_callback_helpers.jl

@@ -11,6 +11,7 @@ struct SaveFuncSMESolve{
     TEXPV<:Union{Nothing,AbstractMatrix},
     TMEXPV<:Union{Nothing,AbstractMatrix},
     TLT<:AbstractVector,
+    CT<:AbstractVector,
 } <: AbstractSaveFunc
     store_measurement::Val{SM}
     e_ops::TE
@@ -20,10 +21,11 @@ struct SaveFuncSMESolve{
     expvals::TEXPV
     m_expvals::TMEXPV
     tlist::TLT
+    dWdt_cache::CT
 end
 
 (f::SaveFuncSMESolve)(u, t, integrator) =
-    _save_func_smesolve(u, integrator, f.e_ops, f.m_ops, f.progr, f.iter, f.expvals, f.m_expvals, f.tlist)
+    _save_func_smesolve(u, integrator, f.e_ops, f.m_ops, f.progr, f.iter, f.expvals, f.m_expvals, f.tlist, f.dWdt_cache)
 (f::SaveFuncSMESolve{false,Nothing})(u, t, integrator) = _save_func(integrator, f.progr) # Common for both all solvers
 
 _get_e_ops_data(e_ops, ::Type{SaveFuncSMESolve}) = _get_e_ops_data(e_ops, SaveFuncMESolve)
@@ -33,7 +35,7 @@ _get_m_ops_data(sc_ops, ::Type{SaveFuncSMESolve}) =
 ##
 
 # When e_ops is a list of operators
-function _save_func_smesolve(u, integrator, e_ops, m_ops, progr, iter, expvals, m_expvals, tlist)
+function _save_func_smesolve(u, integrator, e_ops, m_ops, progr, iter, expvals, m_expvals, tlist, dWdt_cache)
     # 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.
@@ -47,8 +49,8 @@ function _save_func_smesolve(u, integrator, e_ops, m_ops, progr, iter, expvals,
     end
 
     if !isnothing(m_expvals) && iter[] > 1
-        _dWdt = _homodyne_dWdt(integrator, tlist, iter)
-        @. m_expvals[:, iter[]-1] = real(_expect(m_ops)) + _dWdt
+        _homodyne_dWdt!(dWdt_cache, integrator, tlist, iter)
+        @. m_expvals[:, iter[]-1] = real(_expect(m_ops)) + dWdt_cache
     end
 
     iter[] += 1

src/time_evolution/callback_helpers/ssesolve_callback_helpers.jl

@@ -11,6 +11,7 @@ struct SaveFuncSSESolve{
     TEXPV<:Union{Nothing,AbstractMatrix},
     TMEXPV<:Union{Nothing,AbstractMatrix},
     TLT<:AbstractVector,
+    CT<:AbstractVector,
 } <: AbstractSaveFunc
     store_measurement::Val{SM}
     e_ops::TE
@@ -20,10 +21,11 @@ struct SaveFuncSSESolve{
     expvals::TEXPV
     m_expvals::TMEXPV
     tlist::TLT
+    dWdt_cache::CT
 end
 
 (f::SaveFuncSSESolve)(u, t, integrator) =
-    _save_func_ssesolve(u, integrator, f.e_ops, f.m_ops, f.progr, f.iter, f.expvals, f.m_expvals, f.tlist)
+    _save_func_ssesolve(u, integrator, f.e_ops, f.m_ops, f.progr, f.iter, f.expvals, f.m_expvals, f.tlist, f.dWdt_cache)
 (f::SaveFuncSSESolve{false,Nothing})(u, t, integrator) = _save_func(integrator, f.progr) # Common for both all solvers
 
 _get_e_ops_data(e_ops, ::Type{SaveFuncSSESolve}) = get_data.(e_ops)
@@ -34,7 +36,7 @@ _get_save_callback_idx(cb, ::Type{SaveFuncSSESolve}) = 2 # The first one is the
 ##
 
 # When e_ops is a list of operators
-function _save_func_ssesolve(u, integrator, e_ops, m_ops, progr, iter, expvals, m_expvals, tlist)
+function _save_func_ssesolve(u, integrator, e_ops, m_ops, progr, iter, expvals, m_expvals, tlist, dWdt_cache)
     ψ = u
 
     _expect = op -> dot(ψ, op, ψ)
@@ -44,8 +46,8 @@ function _save_func_ssesolve(u, integrator, e_ops, m_ops, progr, iter, expvals,
     end
 
     if !isnothing(m_expvals) && iter[] > 1
-        _dWdt = _homodyne_dWdt(integrator, tlist, iter)
-        @. m_expvals[:, iter[]-1] = real(_expect(m_ops)) + _dWdt
+        _homodyne_dWdt!(dWdt_cache, integrator, tlist, iter)
+        @. m_expvals[:, iter[]-1] = real(_expect(m_ops)) + dWdt_cache
     end
 
     iter[] += 1