Fix definition of noise derivative in stochastic solvers (#453)

Author: albertomercurio

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Support different length for `to` and `from` on GeneralDimensions. ([#448])
 - Extend the `Makie.jl` extension to all the other available backends. ([#450])
+- Fix definition of noise derivative in stochastic solvers. ([#453])
 
 ## [v0.30.0]
 Release date: 2025-04-12
@@ -206,3 +207,4 @@ Release date: 2024-11-13
 [#443]: https://github.com/qutip/QuantumToolbox.jl/issues/443
 [#448]: https://github.com/qutip/QuantumToolbox.jl/issues/448
 [#450]: https://github.com/qutip/QuantumToolbox.jl/issues/450
+[#453]: https://github.com/qutip/QuantumToolbox.jl/issues/453

src/qobj/quantum_object_evo.jl

@@ -468,6 +468,9 @@ function _promote_to_scimloperator(α::Number, data::ScaledOperator)
     isconstant(data.λ) && return ScaledOperator(α * data.λ, data.L)
     return ScaledOperator(data.λ, _promote_to_scimloperator(α, data.L)) # Try to propagate the rule
 end
+function _promote_to_scimloperator(α::Number, data::AddedOperator)
+    return AddedOperator(_promote_to_scimloperator.(α, data.ops)) # Try to propagate the rule
+end
 function _promote_to_scimloperator(α::Number, data::AbstractSciMLOperator)
     return α * data # Going back to the generic case
 end

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) ? unique!(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,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)
+    _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
 
@@ -162,9 +169,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!(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_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

@@ -10,6 +10,8 @@ struct SaveFuncSMESolve{
     IT,
     TEXPV<:Union{Nothing,AbstractMatrix},
     TMEXPV<:Union{Nothing,AbstractMatrix},
+    TLT<:AbstractVector,
+    CT<:AbstractVector,
 } <: AbstractSaveFunc
     store_measurement::Val{SM}
     e_ops::TE
@@ -18,10 +20,12 @@ struct SaveFuncSMESolve{
     iter::IT
     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)
+    _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)
@@ -31,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)
+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.
@@ -45,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)
-        @. 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

@@ -10,6 +10,8 @@ struct SaveFuncSSESolve{
     IT,
     TEXPV<:Union{Nothing,AbstractMatrix},
     TMEXPV<:Union{Nothing,AbstractMatrix},
+    TLT<:AbstractVector,
+    CT<:AbstractVector,
 } <: AbstractSaveFunc
     store_measurement::Val{SM}
     e_ops::TE
@@ -18,10 +20,12 @@ struct SaveFuncSSESolve{
     iter::IT
     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)
+    _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)
@@ -32,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)
+function _save_func_ssesolve(u, integrator, e_ops, m_ops, progr, iter, expvals, m_expvals, tlist, dWdt_cache)
     ψ = u
 
     _expect = op -> dot(ψ, op, ψ)
@@ -42,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)
-        @. 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/ssesolve.jl

@@ -113,7 +113,7 @@ function ssesolveProblem(
         sc_ops_evo_data,
     )
 
-    K = -1im * get_data(H_eff_evo) + K_l
+    K = get_data(QobjEvo(H_eff_evo, -1im)) + K_l
 
     D_l = map(op -> op + _ScalarOperator_e(op, -) * IdentityOperator(prod(dims)), sc_ops_evo_data)
     D = DiffusionOperator(D_l)

src/time_evolution/time_evolution.jl

@@ -3,7 +3,7 @@ export TimeEvolutionSol, TimeEvolutionMCSol, TimeEvolutionStochasticSol
 export liouvillian_floquet, liouvillian_generalized
 
 const DEFAULT_ODE_SOLVER_OPTIONS = (abstol = 1e-8, reltol = 1e-6, save_everystep = false, save_end = true)
-const DEFAULT_SDE_SOLVER_OPTIONS = (abstol = 1e-2, reltol = 1e-2, save_everystep = false, save_end = true)
+const DEFAULT_SDE_SOLVER_OPTIONS = (abstol = 1e-3, reltol = 2e-3, save_everystep = false, save_end = true)
 const COL_TIMES_WHICH_INIT_SIZE = 200
 
 @doc raw"""