Working case of ssesolve measurement

Author: albertomercurio

src/QuantumToolbox.jl

@@ -101,6 +101,7 @@ include("time_evolution/callback_helpers/sesolve_callback_helpers.jl")
 include("time_evolution/callback_helpers/mesolve_callback_helpers.jl")
 include("time_evolution/callback_helpers/mcsolve_callback_helpers.jl")
 include("time_evolution/callback_helpers/ssesolve_callback_helpers.jl")
+include("time_evolution/callback_helpers/smesolve_callback_helpers.jl")
 include("time_evolution/mesolve.jl")
 include("time_evolution/lr_mesolve.jl")
 include("time_evolution/sesolve.jl")

src/time_evolution/callback_helpers/callback_helpers.jl

@@ -13,6 +13,34 @@ function _generate_se_me_kwargs(e_ops, progress_bar, tlist, kwargs, method)
 end
 _generate_se_me_kwargs(e_ops::Nothing, progress_bar::Val{false}, tlist, kwargs, method) = kwargs
 
+function _generate_stochastic_kwargs(
+    e_ops,
+    sc_ops,
+    progress_bar,
+    tlist,
+    store_measurement,
+    kwargs,
+    method::Type{SF},
+) where {SF<:AbstractSaveFunc}
+    cb_save = _generate_stochastic_save_callback(e_ops, sc_ops, tlist, store_measurement, progress_bar, method)
+
+    if SF === SaveFuncSSESolve
+        cb_normalize = _ssesolve_generate_normalize_cb()
+        return _merge_kwargs_with_callback(kwargs, CallbackSet(cb_normalize, cb_save))
+    end
+
+    return _merge_kwargs_with_callback(kwargs, cb_save)
+end
+_generate_stochastic_kwargs(
+    e_ops::Nothing,
+    sc_ops,
+    progress_bar::Val{false},
+    tlist,
+    store_measurement::Val{false},
+    kwargs,
+    method::Type{SF},
+) where {SF<:AbstractSaveFunc} = kwargs
+
 function _merge_kwargs_with_callback(kwargs, cb)
     kwargs2 =
         haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(cb, kwargs.callback),)) :
@@ -32,6 +60,19 @@ function _generate_save_callback(e_ops, tlist, progress_bar, method)
     return PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
 end
 
+function _generate_stochastic_save_callback(e_ops, sc_ops, tlist, store_measurement, progress_bar, method)
+    e_ops_data = e_ops isa Nothing ? nothing : _get_e_ops_data(e_ops, method)
+    m_ops_data = _get_m_ops_data(sc_ops, method)
+
+    progr = getVal(progress_bar) ? ProgressBar(length(tlist), enable = getVal(progress_bar)) : nothing
+
+    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_affect! = method(store_measurement, e_ops_data, m_ops_data, progr, Ref(1), expvals, m_expvals)
+    return PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false))
+end
+
 ##
 
 # When e_ops is Nothing. Common for all solvers
@@ -49,6 +90,18 @@ end
 
 ##
 
+#=
+    To extract the measurement outcomes of a stochastic solver
+=#
+function _get_m_expvals(integrator::AbstractODESolution, method::Type{SF}) where {SF<:AbstractSaveFunc}
+    cb = _get_save_callback(integrator, method)
+    if cb isa Nothing
+        return nothing
+    else
+        return cb.affect!.m_expvals
+    end
+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).

src/time_evolution/callback_helpers/smesolve_callback_helpers.jl

@@ -0,0 +1 @@
+

src/time_evolution/callback_helpers/ssesolve_callback_helpers.jl

@@ -2,45 +2,71 @@
 Helper functions for the ssesolve callbacks. Equal to the sesolve case, but with an additional normalization before saving the expectation values.
 =#
 
-struct SaveFuncSSESolve{TE,PT<:Union{Nothing,ProgressBar},IT,TEXPV<:Union{Nothing,AbstractMatrix}} <: AbstractSaveFunc
+struct SaveFuncSSESolve{
+    SM,
+    TE,
+    TME,
+    PT<:Union{Nothing,ProgressBar},
+    IT,
+    TEXPV<:Union{Nothing,AbstractMatrix},
+    TMEXPV<:Union{Nothing,AbstractMatrix},
+} <: AbstractSaveFunc
+    store_measurement::Val{SM}
     e_ops::TE
+    m_ops::TME
     progr::PT
     iter::IT
     expvals::TEXPV
+    m_expvals::TMEXPV
 end
 
-(f::SaveFuncSSESolve)(integrator) = _save_func_ssesolve(integrator, f.e_ops, f.progr, f.iter, f.expvals)
-(f::SaveFuncSSESolve{Nothing})(integrator) = _save_func(integrator, f.progr) # Common for both all solvers
+(f::SaveFuncSSESolve)(integrator) =
+    _save_func_ssesolve(integrator, f.e_ops, f.m_ops, f.progr, f.iter, f.expvals, f.m_expvals)
+(f::SaveFuncSSESolve{false,Nothing})(integrator) = _save_func(integrator, f.progr) # Common for both all solvers
 
 _get_e_ops_data(e_ops, ::Type{SaveFuncSSESolve}) = get_data.(e_ops)
+_get_m_ops_data(sc_ops, ::Type{SaveFuncSSESolve}) = map(op -> Hermitian(get_data(op) + get_data(op)'), sc_ops)
 
-_get_save_callback_idx(cb, ::Type{SaveFuncSSESolve}) = 1 # The first one is the normalization callback
+_get_save_callback_idx(cb, ::Type{SaveFuncSSESolve}) = 2 # The first one is the normalization callback
 
 ##
 
 # When e_ops is a list of operators
-function _save_func_ssesolve(integrator, e_ops, progr, iter, expvals)
-    ψ = normalize!(integrator.u)
+function _save_func_ssesolve(integrator, e_ops, m_ops, progr, iter, expvals, m_expvals)
+    ψ = integrator.u
+
     _expect = op -> dot(ψ, op, ψ)
-    @. expvals[:, iter[]] = _expect(e_ops)
+
+    if !isnothing(e_ops)
+        @. expvals[:, iter[]] = _expect(e_ops)
+    end
+
+    if !isnothing(m_expvals) && iter[] > 1
+        _dWdt = _homodyne_dWdt(integrator)
+        @. m_expvals[:, iter[]-1] = real(_expect(m_ops)) + _dWdt
+    end
+
     iter[] += 1
 
     return _save_func(integrator, progr)
 end
 
+# 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])
+
+    return _dWdt
+end
+
 ##
 
 #=
-This function adds the normalization callback to the kwargs. It is needed to stabilize the integration when using the ssesolve method.
+    This function generates the normalization callback. It is needed to stabilize the integration when using the ssesolve method.
 =#
-function _ssesolve_add_normalize_cb(kwargs)
+function _ssesolve_generate_normalize_cb()
     _condition = (u, t, integrator) -> true
     _affect! = (integrator) -> normalize!(integrator.u)
     cb = DiscreteCallback(_condition, _affect!; save_positions = (false, false))
 
-    cb_set = haskey(kwargs, :callback) ? CallbackSet(kwargs[:callback], cb) : cb
-
-    kwargs2 = merge(kwargs, (callback = cb_set,))
-
-    return kwargs2
+    return cb
 end

src/time_evolution/smesolve.jl

@@ -376,6 +376,7 @@ function smesolve(
         expvals,
         expvals, # This is average_expect
         expvals_all,
+        nothing, # Measurement expectation values
         sol.converged,
         _sol_1.alg,
         _sol_1.prob.kwargs[:abstol],

src/time_evolution/ssesolve.jl

@@ -114,8 +114,15 @@ function ssesolveProblem(
     D = DiffusionOperator(D_l)
 
     kwargs2 = _merge_saveat(tlist, e_ops, DEFAULT_SDE_SOLVER_OPTIONS; kwargs...)
-    kwargs3 = _generate_se_me_kwargs(e_ops, makeVal(progress_bar), tlist, kwargs2, SaveFuncSSESolve)
-    kwargs4 = _ssesolve_add_normalize_cb(kwargs3)
+    kwargs3 = _generate_stochastic_kwargs(
+        e_ops,
+        sc_ops,
+        makeVal(progress_bar),
+        tlist,
+        makeVal(store_measurement),
+        kwargs2,
+        SaveFuncSSESolve,
+    )
 
     tspan = (tlist[1], tlist[end])
     noise = RealWienerProcess!(
@@ -134,7 +141,7 @@ function ssesolveProblem(
         params;
         noise_rate_prototype = noise_rate_prototype,
         noise = noise,
-        kwargs4...,
+        kwargs3...,
     )
 
     return TimeEvolutionProblem(prob, tlist, dims)
@@ -154,6 +161,7 @@ end
         prob_func::Union{Function, Nothing} = nothing,
         output_func::Union{Tuple,Nothing} = nothing,
         progress_bar::Union{Val,Bool} = Val(true),
+        store_measurement::Union{Val,Bool} = Val(false),
         kwargs...,
     )
 
@@ -193,6 +201,7 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and  ``dW_n(t)`` is t
 - `prob_func`: Function to use for generating the SDEProblem.
 - `output_func`: a `Tuple` containing the `Function` to use for generating the output of a single trajectory, the (optional) `ProgressBar` object, and the (optional) `RemoteChannel` object.
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
+- `store_measurement`: Whether to store the measurement results. Default is `Val(false)`.
 - `kwargs`: The keyword arguments for the ODEProblem.
 
 # Notes
@@ -219,11 +228,19 @@ function ssesolveEnsembleProblem(
     prob_func::Union{Function,Nothing} = nothing,
     output_func::Union{Tuple,Nothing} = nothing,
     progress_bar::Union{Val,Bool} = Val(true),
+    store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
 )
     _prob_func =
         isnothing(prob_func) ?
-        _ensemble_dispatch_prob_func(rng, ntraj, tlist, _stochastic_prob_func; n_sc_ops = length(sc_ops)) : prob_func
+        _ensemble_dispatch_prob_func(
+            rng,
+            ntraj,
+            tlist,
+            _stochastic_prob_func;
+            n_sc_ops = length(sc_ops),
+            store_measurement = makeVal(store_measurement),
+        ) : prob_func
     _output_func =
         output_func isa Nothing ?
         _ensemble_dispatch_output_func(ensemble_method, progress_bar, ntraj, _stochastic_output_func) : output_func
@@ -237,6 +254,7 @@ function ssesolveEnsembleProblem(
         params = params,
         rng = rng,
         progress_bar = Val(false),
+        store_measurement = makeVal(store_measurement),
         kwargs...,
     )
 
@@ -265,6 +283,7 @@ end
         prob_func::Union{Function, Nothing} = nothing,
         output_func::Union{Tuple,Nothing} = nothing,
         progress_bar::Union{Val,Bool} = Val(true),
+        store_measurement::Union{Val,Bool} = Val(false),
         kwargs...,
     )
 
@@ -307,6 +326,7 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is th
 - `prob_func`: Function to use for generating the SDEProblem.
 - `output_func`: a `Tuple` containing the `Function` to use for generating the output of a single trajectory, the (optional) `ProgressBar` object, and the (optional) `RemoteChannel` object.
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
+- `store_measurement`: Whether to store the measurement results. Default is `Val(false)`.
 - `kwargs`: The keyword arguments for the ODEProblem.
 
 # Notes
@@ -335,6 +355,7 @@ function ssesolve(
     prob_func::Union{Function,Nothing} = nothing,
     output_func::Union{Tuple,Nothing} = nothing,
     progress_bar::Union{Val,Bool} = Val(true),
+    store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
 )
     ens_prob = ssesolveEnsembleProblem(
@@ -350,6 +371,7 @@ function ssesolve(
         prob_func = prob_func,
         output_func = output_func,
         progress_bar = progress_bar,
+        store_measurement = makeVal(store_measurement),
         kwargs...,
     )
 
@@ -366,6 +388,7 @@ function ssesolve(
 
     _sol_1 = sol[:, 1]
     _expvals_sol_1 = _get_expvals(_sol_1, SaveFuncSSESolve)
+    _m_expvals_sol_1 = _get_m_expvals(_sol_1, SaveFuncSSESolve)
 
     normalize_states = Val(false)
     dims = ens_prob.dimensions
@@ -374,6 +397,10 @@ function ssesolve(
     expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all, dims = 2) # Stack on dimension 2 to align with QuTiP
     states = map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol))
 
+    _m_expvals =
+        _m_expvals_sol_1 isa Nothing ? nothing : map(i -> _get_m_expvals(sol[:, i], SaveFuncSSESolve), eachindex(sol))
+    m_expvals = _m_expvals isa Nothing ? nothing : stack(_m_expvals, dims = 2)
+
     expvals =
         _get_expvals(_sol_1, SaveFuncSSESolve) isa Nothing ? nothing :
         dropdims(sum(expvals_all, dims = 2), dims = 2) ./ length(sol)
@@ -385,6 +412,7 @@ function ssesolve(
         expvals,
         expvals, # This is average_expect
         expvals_all,
+        m_expvals, # Measurement expectation values
         sol.converged,
         _sol_1.alg,
         _sol_1.prob.kwargs[:abstol],

src/time_evolution/time_evolution.jl

@@ -173,6 +173,7 @@ struct TimeEvolutionStochasticSol{
     TS<:AbstractVector,
     TE<:Union{AbstractMatrix,Nothing},
     TEA<:Union{AbstractArray,Nothing},
+    TEM<:Union{AbstractArray,Nothing},
     AlgT<:StochasticDiffEqAlgorithm,
     AT<:Real,
     RT<:Real,
@@ -183,6 +184,7 @@ struct TimeEvolutionStochasticSol{
     expect::TE
     average_expect::TE # Currently just a synonym for `expect`
     runs_expect::TEA
+    measurement::TEM
     converged::Bool
     alg::AlgT
     abstol::AT
@@ -345,11 +347,13 @@ function _stochastic_prob_func(prob, i, repeat, rng, seeds, tlist; kwargs...)
     traj_rng = typeof(rng)()
     seed!(traj_rng, seed)
 
+    store_measurement = haskey(kwargs, :store_measurement) ? getVal(kwargs[:store_measurement]) : false
+
     noise = RealWienerProcess!(
         prob.prob.tspan[1],
         zeros(kwargs[:n_sc_ops]),
         zeros(kwargs[:n_sc_ops]),
-        save_everystep = false,
+        save_everystep = store_measurement,
         rng = traj_rng,
     )