Fix noise derivative definition

Author: albertomercurio

src/time_evolution/callback_helpers/callback_helpers.jl

@@ -24,12 +24,17 @@ function _generate_stochastic_kwargs(
 ) where {SF<:AbstractSaveFunc}
     cb_save = _generate_stochastic_save_callback(e_ops, sc_ops, tlist, store_measurement, progress_bar, method)
 
+    # Ensure that the noise is stored in tlist. # TODO: Fix this directly in DiffEqNoiseProcess.jl
+    # See https://github.com/SciML/DiffEqNoiseProcess.jl/issues/214 for example
+    tstops = haskey(kwargs, :tstops) ? sort!(vcat(tlist, kwargs.tstops)) : tlist
+    kwargs2 = merge(kwargs, (tstops = tstops,))
+
     if SF === SaveFuncSSESolve
         cb_normalize = _ssesolve_generate_normalize_cb()
-        return _merge_kwargs_with_callback(kwargs, CallbackSet(cb_normalize, cb_save))
+        return _merge_kwargs_with_callback(kwargs2, CallbackSet(cb_normalize, cb_save))
     end
 
-    return _merge_kwargs_with_callback(kwargs, cb_save)
+    return _merge_kwargs_with_callback(kwargs2, cb_save)
 end
 _generate_stochastic_kwargs(
     e_ops::Nothing,
@@ -69,7 +74,7 @@ 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)
+    _save_func = method(store_measurement, e_ops_data, m_ops_data, progr, Ref(1), expvals, m_expvals, tlist)
     return FunctionCallingCallback(_save_func, funcat = tlist)
 end
 
@@ -163,8 +168,12 @@ _get_save_callback_idx(cb, method) = 1
 # %% ------------ Noise Measurement Helpers ------------ %%
 
 # TODO: Add some cache mechanism to avoid memory allocations
-function _homodyne_dWdt(integrator)
-    @inbounds _dWdt = (integrator.W.u[end] .- integrator.W.u[end-1]) ./ (integrator.W.t[end] - integrator.W.t[end-1])
+# TODO: To improve. See https://github.com/SciML/DiffEqNoiseProcess.jl/issues/214
+function _homodyne_dWdt(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])
 
     return _dWdt
 end

src/time_evolution/callback_helpers/smesolve_callback_helpers.jl

@@ -10,6 +10,7 @@ struct SaveFuncSMESolve{
     IT,
     TEXPV<:Union{Nothing,AbstractMatrix},
     TMEXPV<:Union{Nothing,AbstractMatrix},
+    TLT<:AbstractVector,
 } <: AbstractSaveFunc
     store_measurement::Val{SM}
     e_ops::TE
@@ -18,10 +19,11 @@ struct SaveFuncSMESolve{
     iter::IT
     expvals::TEXPV
     m_expvals::TMEXPV
+    tlist::TLT
 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)
+    _save_func_smesolve(u, integrator, f.e_ops, f.m_ops, f.progr, f.iter, f.expvals, f.m_expvals, f.tlist)
 (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)
@@ -31,7 +33,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)
+function _save_func_smesolve(u, integrator, e_ops, m_ops, progr, iter, expvals, m_expvals, tlist)
     # 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.
@@ -45,7 +47,7 @@ function _save_func_smesolve(u, integrator, e_ops, m_ops, progr, iter, expvals,
     end
 
     if !isnothing(m_expvals) && iter[] > 1
-        _dWdt = _homodyne_dWdt(integrator)
+        _dWdt = _homodyne_dWdt(integrator, tlist, iter)
         @. m_expvals[:, iter[]-1] = real(_expect(m_ops)) + _dWdt
     end
 

src/time_evolution/callback_helpers/ssesolve_callback_helpers.jl

@@ -10,6 +10,7 @@ struct SaveFuncSSESolve{
     IT,
     TEXPV<:Union{Nothing,AbstractMatrix},
     TMEXPV<:Union{Nothing,AbstractMatrix},
+    TLT<:AbstractVector,
 } <: AbstractSaveFunc
     store_measurement::Val{SM}
     e_ops::TE
@@ -18,10 +19,11 @@ struct SaveFuncSSESolve{
     iter::IT
     expvals::TEXPV
     m_expvals::TMEXPV
+    tlist::TLT
 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)
+    _save_func_ssesolve(u, integrator, f.e_ops, f.m_ops, f.progr, f.iter, f.expvals, f.m_expvals, f.tlist)
 (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)
@@ -32,7 +34,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)
+function _save_func_ssesolve(u, integrator, e_ops, m_ops, progr, iter, expvals, m_expvals, tlist)
     ψ = u
 
     _expect = op -> dot(ψ, op, ψ)
@@ -42,7 +44,7 @@ function _save_func_ssesolve(u, integrator, e_ops, m_ops, progr, iter, expvals,
     end
 
     if !isnothing(m_expvals) && iter[] > 1
-        _dWdt = _homodyne_dWdt(integrator)
+        _dWdt = _homodyne_dWdt(integrator, tlist, iter)
         @. m_expvals[:, iter[]-1] = real(_expect(m_ops)) + _dWdt
     end