diff --git a/CHANGELOG.md b/CHANGELOG.md index 8eb59f8d1..253fba8fd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Align some attributes of `mcsolve`, `ssesolve` and `smesolve` results with `QuTiP`. ([#402]) - Improve ensemble generation of `ssesolve` and change parameters handling on stochastic processes. ([#403]) - Set default trajectories to 500 and rename the keyword argument `ensemble_method` to `ensemblealg`. ([#405]) +- Introduce measurement on `ssesolve` and `smesolve`. ([#404]) ## [v0.26.0] Release date: 2025-02-09 @@ -133,4 +134,5 @@ Release date: 2024-11-13 [#398]: https://github.com/qutip/QuantumToolbox.jl/issues/398 [#402]: https://github.com/qutip/QuantumToolbox.jl/issues/402 [#403]: https://github.com/qutip/QuantumToolbox.jl/issues/403 +[#404]: https://github.com/qutip/QuantumToolbox.jl/issues/404 [#405]: https://github.com/qutip/QuantumToolbox.jl/issues/405 diff --git a/docs/src/users_guide/time_evolution/stochastic.md b/docs/src/users_guide/time_evolution/stochastic.md index 3687001ce..852ae089d 100644 --- a/docs/src/users_guide/time_evolution/stochastic.md +++ b/docs/src/users_guide/time_evolution/stochastic.md @@ -105,12 +105,16 @@ sse_sol = ssesolve( sc_ops, e_ops = [x], ntraj = ntraj, + store_measurement = Val(true), ) +measurement_avg = sum(sse_sol.measurement, dims=2) / size(sse_sol.measurement, 2) +measurement_avg = dropdims(measurement_avg, dims=2) + # plot by CairoMakie.jl fig = Figure(size = (500, 350)) ax = Axis(fig[1, 1], xlabel = "Time") -#lines!(ax, tlist, real(sse_sol.xxxxxx), label = L"J_x", color = :red, linestyle = :solid) TODO: add this in the future +lines!(ax, tlist[1:end-1], real(measurement_avg[1,:]), label = L"J_x", color = :red, linestyle = :solid) lines!(ax, tlist, real(sse_sol.expect[1,:]), label = L"\langle x \rangle", color = :black, linestyle = :solid) axislegend(ax, position = :rt) @@ -134,12 +138,16 @@ sme_sol = smesolve( sc_ops, e_ops = [x], ntraj = ntraj, + store_measurement = Val(true), ) +measurement_avg = sum(sme_sol.measurement, dims=2) / size(sme_sol.measurement, 2) +measurement_avg = dropdims(measurement_avg, dims=2) + # plot by CairoMakie.jl fig = Figure(size = (500, 350)) ax = Axis(fig[1, 1], xlabel = "Time") -#lines!(ax, tlist, real(sme_sol.xxxxxx), label = L"J_x", color = :red, linestyle = :solid) TODO: add this in the future +lines!(ax, tlist[1:end-1], real(measurement_avg[1,:]), label = L"J_x", color = :red, linestyle = :solid) lines!(ax, tlist, real(sme_sol.expect[1,:]), label = L"\langle x \rangle", color = :black, linestyle = :solid) axislegend(ax, position = :rt) diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl index 7f512d685..f1bb81b95 100644 --- a/src/QuantumToolbox.jl +++ b/src/QuantumToolbox.jl @@ -97,11 +97,12 @@ include("qobj/block_diagonal_form.jl") # time evolution include("time_evolution/time_evolution.jl") +include("time_evolution/callback_helpers/callback_helpers.jl") 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/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") diff --git a/src/time_evolution/callback_helpers/callback_helpers.jl b/src/time_evolution/callback_helpers/callback_helpers.jl index 6f4103069..4170635a8 100644 --- a/src/time_evolution/callback_helpers/callback_helpers.jl +++ b/src/time_evolution/callback_helpers/callback_helpers.jl @@ -4,6 +4,8 @@ This file contains helper functions for callbacks. The affect! function are defi ## +abstract type AbstractSaveFunc end + # Multiple dispatch depending on the progress_bar and e_ops types function _generate_se_me_kwargs(e_ops, progress_bar, tlist, kwargs, method) cb = _generate_save_callback(e_ops, tlist, progress_bar, method) @@ -11,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),)) : @@ -30,38 +60,29 @@ function _generate_save_callback(e_ops, tlist, progress_bar, method) return PresetTimeCallback(tlist, _save_affect!, save_positions = (false, false)) end -_get_e_ops_data(e_ops, ::Type{SaveFuncSESolve}) = get_data.(e_ops) -_get_e_ops_data(e_ops, ::Type{SaveFuncMESolve}) = [_generate_mesolve_e_op(op) for op in e_ops] # Broadcasting generates type instabilities on Julia v1.10 -_get_e_ops_data(e_ops, ::Type{SaveFuncSSESolve}) = get_data.(e_ops) - -_generate_mesolve_e_op(op) = mat2vec(adjoint(get_data(op))) - -#= -This function add the normalization callback to the kwargs. It is needed to stabilize the integration when using the ssesolve method. -=# -function _ssesolve_add_normalize_cb(kwargs) - _condition = (u, t, integrator) -> true - _affect! = (integrator) -> normalize!(integrator.u) - cb = DiscreteCallback(_condition, _affect!; save_positions = (false, false)) - # return merge(kwargs, (callback = CallbackSet(kwargs[:callback], cb),)) +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) - cb_set = haskey(kwargs, :callback) ? CallbackSet(kwargs[:callback], cb) : cb + progr = getVal(progress_bar) ? ProgressBar(length(tlist), enable = getVal(progress_bar)) : nothing - kwargs2 = merge(kwargs, (callback = cb_set,)) + 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 - return kwargs2 + _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 both mesolve and sesolve +# When e_ops is Nothing. Common for all solvers function _save_func(integrator, progr) next!(progr) u_modified!(integrator, false) return nothing end -# When progr is Nothing. Common for both mesolve and sesolve +# When progr is Nothing. Common for all solvers function _save_func(integrator, progr::Nothing) u_modified!(integrator, false) return nothing @@ -69,9 +90,34 @@ 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). +=# +function _get_e_ops(integrator::AbstractODEIntegrator, method::Type{SF}) where {SF<:AbstractSaveFunc} + cb = _get_save_callback(integrator, method) + if cb isa Nothing + return nothing + else + return cb.affect!.e_ops + end +end + # Get the e_ops from a given AbstractODESolution. Valid for `sesolve`, `mesolve` and `ssesolve`. -function _se_me_sse_get_expvals(sol::AbstractODESolution) - cb = _se_me_sse_get_save_callback(sol) +function _get_expvals(sol::AbstractODESolution, method::Type{SF}) where {SF<:AbstractSaveFunc} + cb = _get_save_callback(sol, method) if cb isa Nothing return nothing else @@ -79,28 +125,46 @@ function _se_me_sse_get_expvals(sol::AbstractODESolution) end end -function _se_me_sse_get_save_callback(sol::AbstractODESolution) +#= + _get_save_callback + +Return the Callback that is responsible for saving the expectation values of the system. +=# +function _get_save_callback(sol::AbstractODESolution, method::Type{SF}) where {SF<:AbstractSaveFunc} kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple to support Zygote.jl if hasproperty(kwargs, :callback) - return _se_me_sse_get_save_callback(kwargs.callback) + return _get_save_callback(kwargs.callback, method) else return nothing end end -_se_me_sse_get_save_callback(integrator::AbstractODEIntegrator) = _se_me_sse_get_save_callback(integrator.opts.callback) -function _se_me_sse_get_save_callback(cb::CallbackSet) +_get_save_callback(integrator::AbstractODEIntegrator, method::Type{SF}) where {SF<:AbstractSaveFunc} = + _get_save_callback(integrator.opts.callback, method) +function _get_save_callback(cb::CallbackSet, method::Type{SF}) where {SF<:AbstractSaveFunc} cbs_discrete = cb.discrete_callbacks if length(cbs_discrete) > 0 - _cb = cb.discrete_callbacks[1] - return _se_me_sse_get_save_callback(_cb) + idx = _get_save_callback_idx(cb, method) + _cb = cb.discrete_callbacks[idx] + return _get_save_callback(_cb, method) else return nothing end end -function _se_me_sse_get_save_callback(cb::DiscreteCallback) - if typeof(cb.affect!) <: Union{SaveFuncSESolve,SaveFuncMESolve,SaveFuncSSESolve} +function _get_save_callback(cb::DiscreteCallback, ::Type{SF}) where {SF<:AbstractSaveFunc} + if typeof(cb.affect!) <: AbstractSaveFunc return cb end return nothing end -_se_me_sse_get_save_callback(cb::ContinuousCallback) = nothing +_get_save_callback(cb::ContinuousCallback, ::Type{SF}) where {SF<:AbstractSaveFunc} = nothing + +_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]) + + return _dWdt +end diff --git a/src/time_evolution/callback_helpers/mcsolve_callback_helpers.jl b/src/time_evolution/callback_helpers/mcsolve_callback_helpers.jl index cc3e3598d..b0095e705 100644 --- a/src/time_evolution/callback_helpers/mcsolve_callback_helpers.jl +++ b/src/time_evolution/callback_helpers/mcsolve_callback_helpers.jl @@ -2,7 +2,7 @@ Helper functions for the mcsolve callbacks. =# -struct SaveFuncMCSolve{TE,IT,TEXPV} +struct SaveFuncMCSolve{TE,IT,TEXPV} <: AbstractSaveFunc e_ops::TE iter::IT expvals::TEXPV @@ -10,6 +10,9 @@ end (f::SaveFuncMCSolve)(integrator) = _save_func_mcsolve(integrator, f.e_ops, f.iter, f.expvals) +_get_save_callback_idx(cb, ::Type{SaveFuncMCSolve}) = _mcsolve_has_continuous_jump(cb) ? 1 : 2 + +## struct LindbladJump{ T1, T2, @@ -167,37 +170,6 @@ _mcsolve_discrete_condition(u, t, integrator) = ## -#= - _mc_get_save_callback - -Return the Callback that is responsible for saving the expectation values of the system. -=# -function _mc_get_save_callback(sol::AbstractODESolution) - kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple to support Zygote.jl - return _mc_get_save_callback(kwargs.callback) # There is always the Jump callback -end -_mc_get_save_callback(integrator::AbstractODEIntegrator) = _mc_get_save_callback(integrator.opts.callback) -function _mc_get_save_callback(cb::CallbackSet) - cbs_discrete = cb.discrete_callbacks - - if length(cbs_discrete) > 0 - idx = _mcsolve_has_continuous_jump(cb) ? 1 : 2 - _cb = cb.discrete_callbacks[idx] - return _mc_get_save_callback(_cb) - else - return nothing - end -end -_mc_get_save_callback(cb::DiscreteCallback) = - if cb.affect! isa SaveFuncMCSolve - return cb - else - return nothing - end -_mc_get_save_callback(cb::ContinuousCallback) = nothing - -## - function _mc_get_jump_callback(sol::AbstractODESolution) kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple to support Zygote.jl return _mc_get_jump_callback(kwargs.callback) # There is always the Jump callback @@ -215,8 +187,8 @@ _mc_get_jump_callback(cb::DiscreteCallback) = cb ## #= -With this function we extract the c_ops and c_ops_herm from the LindbladJump `affect!` function of the callback of the integrator. -This callback can be a DiscreteLindbladJumpCallback or a ContinuousLindbladJumpCallback. + With this function we extract the c_ops and c_ops_herm from the LindbladJump `affect!` function of the callback of the integrator. + This callback can be a DiscreteLindbladJumpCallback or a ContinuousLindbladJumpCallback. =# function _mcsolve_get_c_ops(integrator::AbstractODEIntegrator) cb = _mc_get_jump_callback(integrator) @@ -227,28 +199,6 @@ function _mcsolve_get_c_ops(integrator::AbstractODEIntegrator) 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). -=# -function _mcsolve_get_e_ops(integrator::AbstractODEIntegrator) - cb = _mc_get_save_callback(integrator) - if cb isa Nothing - return nothing - else - return cb.affect!.e_ops - end -end - -function _mcsolve_get_expvals(sol::AbstractODESolution) - cb = _mc_get_save_callback(sol) - if cb isa Nothing - return nothing - else - return cb.affect!.expvals - end -end - #= _mcsolve_initialize_callbacks(prob, tlist) diff --git a/src/time_evolution/callback_helpers/mesolve_callback_helpers.jl b/src/time_evolution/callback_helpers/mesolve_callback_helpers.jl index 449e2645a..ee253d765 100644 --- a/src/time_evolution/callback_helpers/mesolve_callback_helpers.jl +++ b/src/time_evolution/callback_helpers/mesolve_callback_helpers.jl @@ -2,7 +2,7 @@ Helper functions for the mesolve callbacks. =# -struct SaveFuncMESolve{TE,PT<:Union{Nothing,ProgressBar},IT,TEXPV<:Union{Nothing,AbstractMatrix}} +struct SaveFuncMESolve{TE,PT<:Union{Nothing,ProgressBar},IT,TEXPV<:Union{Nothing,AbstractMatrix}} <: AbstractSaveFunc e_ops::TE progr::PT iter::IT @@ -12,6 +12,8 @@ end (f::SaveFuncMESolve)(integrator) = _save_func_mesolve(integrator, f.e_ops, f.progr, f.iter, f.expvals) (f::SaveFuncMESolve{Nothing})(integrator) = _save_func(integrator, f.progr) +_get_e_ops_data(e_ops, ::Type{SaveFuncMESolve}) = [_generate_mesolve_e_op(op) for op in e_ops] # Broadcasting generates type instabilities on Julia v1.10 + ## # When e_ops is a list of operators @@ -29,7 +31,7 @@ function _save_func_mesolve(integrator, e_ops, progr, iter, expvals) end function _mesolve_callbacks_new_e_ops!(integrator::AbstractODEIntegrator, e_ops) - cb = _se_me_sse_get_save_callback(integrator) + cb = _get_save_callback(integrator, SaveFuncMESolve) if cb isa Nothing return nothing else @@ -37,3 +39,5 @@ function _mesolve_callbacks_new_e_ops!(integrator::AbstractODEIntegrator, e_ops) return nothing end end + +_generate_mesolve_e_op(op) = mat2vec(adjoint(get_data(op))) diff --git a/src/time_evolution/callback_helpers/sesolve_callback_helpers.jl b/src/time_evolution/callback_helpers/sesolve_callback_helpers.jl index 54f0945f9..e5205ffc6 100644 --- a/src/time_evolution/callback_helpers/sesolve_callback_helpers.jl +++ b/src/time_evolution/callback_helpers/sesolve_callback_helpers.jl @@ -2,7 +2,7 @@ Helper functions for the sesolve callbacks. =# -struct SaveFuncSESolve{TE,PT<:Union{Nothing,ProgressBar},IT,TEXPV<:Union{Nothing,AbstractMatrix}} +struct SaveFuncSESolve{TE,PT<:Union{Nothing,ProgressBar},IT,TEXPV<:Union{Nothing,AbstractMatrix}} <: AbstractSaveFunc e_ops::TE progr::PT iter::IT @@ -12,6 +12,8 @@ end (f::SaveFuncSESolve)(integrator) = _save_func_sesolve(integrator, f.e_ops, f.progr, f.iter, f.expvals) (f::SaveFuncSESolve{Nothing})(integrator) = _save_func(integrator, f.progr) # Common for both mesolve and sesolve +_get_e_ops_data(e_ops, ::Type{SaveFuncSESolve}) = get_data.(e_ops) + ## # When e_ops is a list of operators diff --git a/src/time_evolution/callback_helpers/smesolve_callback_helpers.jl b/src/time_evolution/callback_helpers/smesolve_callback_helpers.jl new file mode 100644 index 000000000..ab14bf81b --- /dev/null +++ b/src/time_evolution/callback_helpers/smesolve_callback_helpers.jl @@ -0,0 +1,55 @@ +#= +Helper functions for the smesolve callbacks. Almost equal to the mesolve case, but with an additional possibility to store the measurement operators expectation values. +=# + +struct SaveFuncSMESolve{ + SM, + TE, + TE2, + PT<:Union{Nothing,ProgressBar}, + IT, + TEXPV<:Union{Nothing,AbstractMatrix}, + TMEXPV<:Union{Nothing,AbstractMatrix}, +} <: AbstractSaveFunc + store_measurement::Val{SM} + e_ops::TE + m_ops::TE2 + progr::PT + iter::IT + expvals::TEXPV + m_expvals::TMEXPV +end + +(f::SaveFuncSMESolve)(integrator) = + _save_func_smesolve(integrator, f.e_ops, f.m_ops, f.progr, f.iter, f.expvals, f.m_expvals) +(f::SaveFuncSMESolve{false,Nothing})(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) +_get_m_ops_data(sc_ops, ::Type{SaveFuncSMESolve}) = + map(op -> _generate_mesolve_e_op(op) + _generate_mesolve_e_op(op'), sc_ops) + +## + +# When e_ops is a list of operators +function _save_func_smesolve(integrator, e_ops, m_ops, progr, iter, expvals, m_expvals) + # 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. + + ρ = integrator.u + + _expect = op -> dot(op, ρ) + + 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 diff --git a/src/time_evolution/callback_helpers/ssesolve_callback_helpers.jl b/src/time_evolution/callback_helpers/ssesolve_callback_helpers.jl index 69fca4b11..bd20d2d2c 100644 --- a/src/time_evolution/callback_helpers/ssesolve_callback_helpers.jl +++ b/src/time_evolution/callback_helpers/ssesolve_callback_helpers.jl @@ -1,25 +1,65 @@ #= -Helper functions for the ssesolve callbacks. Equal to the sesolve case, but with an additional normalization before saving the expectation values. +Helper functions for the ssesolve callbacks. Almost equal to the sesolve case, but with an additional possibility to store the measurement operators expectation values. Also, this callback is not the first one, but the second one, due to the presence of the normalization callback. =# -struct SaveFuncSSESolve{TE,PT<:Union{Nothing,ProgressBar},IT,TEXPV<:Union{Nothing,AbstractMatrix}} +struct SaveFuncSSESolve{ + SM, + TE, + TE2, + PT<:Union{Nothing,ProgressBar}, + IT, + TEXPV<:Union{Nothing,AbstractMatrix}, + TMEXPV<:Union{Nothing,AbstractMatrix}, +} <: AbstractSaveFunc + store_measurement::Val{SM} e_ops::TE + m_ops::TE2 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 mesolve and sesolve +(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}) = 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 + +## + +#= + This function generates the normalization callback. It is needed to stabilize the integration when using the ssesolve method. +=# +function _ssesolve_generate_normalize_cb() + _condition = (u, t, integrator) -> true + _affect! = (integrator) -> normalize!(integrator.u) + cb = DiscreteCallback(_condition, _affect!; save_positions = (false, false)) + + return cb +end diff --git a/src/time_evolution/mcsolve.jl b/src/time_evolution/mcsolve.jl index 172e8e099..210cf37a7 100644 --- a/src/time_evolution/mcsolve.jl +++ b/src/time_evolution/mcsolve.jl @@ -398,9 +398,10 @@ function mcsolve( dims = ens_prob_mc.dimensions _sol_1 = sol[:, 1] - _expvals_sol_1 = _mcsolve_get_expvals(_sol_1) + _expvals_sol_1 = _get_expvals(_sol_1, SaveFuncMCSolve) - _expvals_all = _expvals_sol_1 isa Nothing ? nothing : map(i -> _mcsolve_get_expvals(sol[:, i]), eachindex(sol)) + _expvals_all = + _expvals_sol_1 isa Nothing ? nothing : map(i -> _get_expvals(sol[:, i], SaveFuncMCSolve), eachindex(sol)) 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)) col_times = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.col_times, eachindex(sol)) diff --git a/src/time_evolution/mesolve.jl b/src/time_evolution/mesolve.jl index 7db098bdd..ee78e7eac 100644 --- a/src/time_evolution/mesolve.jl +++ b/src/time_evolution/mesolve.jl @@ -178,7 +178,7 @@ function mesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit return TimeEvolutionSol( prob.times, ρt, - _se_me_sse_get_expvals(sol), + _get_expvals(sol, SaveFuncMESolve), sol.retcode, sol.alg, NamedTuple(sol.prob.kwargs).abstol, diff --git a/src/time_evolution/sesolve.jl b/src/time_evolution/sesolve.jl index 368d46241..571b1a396 100644 --- a/src/time_evolution/sesolve.jl +++ b/src/time_evolution/sesolve.jl @@ -154,7 +154,7 @@ function sesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit return TimeEvolutionSol( prob.times, ψt, - _se_me_sse_get_expvals(sol), + _get_expvals(sol, SaveFuncSESolve), sol.retcode, sol.alg, NamedTuple(sol.prob.kwargs).abstol, diff --git a/src/time_evolution/smesolve.jl b/src/time_evolution/smesolve.jl index 741c0d045..32f02be70 100644 --- a/src/time_evolution/smesolve.jl +++ b/src/time_evolution/smesolve.jl @@ -20,6 +20,7 @@ _smesolve_ScalarOperator(op_vec) = params = NullParameters(), rng::AbstractRNG = default_rng(), progress_bar::Union{Val,Bool} = Val(true), + store_measurement::Union{Val, Bool} = Val(false), kwargs..., ) @@ -54,6 +55,7 @@ Above, ``\hat{C}_i`` represent the collapse operators related to pure dissipatio - `params`: `NullParameters` of parameters to pass to the solver. - `rng`: Random number generator for reproducibility. - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities. +- `store_measurement`: Whether to store the measurement expectation values. Default is `Val(false)`. - `kwargs`: The keyword arguments for the ODEProblem. # Notes @@ -77,6 +79,7 @@ function smesolveProblem( params = NullParameters(), rng::AbstractRNG = default_rng(), progress_bar::Union{Val,Bool} = Val(true), + store_measurement::Union{Val,Bool} = Val(false), kwargs..., ) where {StateOpType<:Union{KetQuantumObject,OperatorQuantumObject}} haskey(kwargs, :save_idxs) && @@ -111,11 +114,24 @@ function smesolveProblem( 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, SaveFuncMESolve) + kwargs3 = _generate_stochastic_kwargs( + e_ops, + sc_ops, + makeVal(progress_bar), + tlist, + makeVal(store_measurement), + kwargs2, + SaveFuncSMESolve, + ) tspan = (tlist[1], tlist[end]) - noise = - RealWienerProcess!(tlist[1], zeros(length(sc_ops)), zeros(length(sc_ops)), save_everystep = false, rng = rng) + noise = RealWienerProcess!( + tlist[1], + zeros(length(sc_ops)), + zeros(length(sc_ops)), + save_everystep = getVal(store_measurement), + rng = rng, + ) noise_rate_prototype = similar(ρ0, length(ρ0), length(sc_ops)) prob = SDEProblem{true}( K, @@ -141,11 +157,12 @@ end e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, params = NullParameters(), rng::AbstractRNG = default_rng(), - ntraj::Int = 1, - ensemble_method = EnsembleThreads(), + ntraj::Int = 500, + ensemblealg::EnsembleAlgorithm = EnsembleThreads(), 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..., ) @@ -179,11 +196,12 @@ Above, ``\hat{C}_i`` represent the collapse operators related to pure dissipatio - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`. - `params`: `NullParameters` of parameters to pass to the solver. - `rng`: Random number generator for reproducibility. -- `ntraj`: Number of trajectories to use. -- `ensemble_method`: Ensemble method to use. Default to `EnsembleThreads()`. +- `ntraj`: Number of trajectories to use. Default is `500`. +- `ensemblealg`: Ensemble method to use. Default to `EnsembleThreads()`. - `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 expectation values. Default is `Val(false)`. - `kwargs`: The keyword arguments for the ODEProblem. # Notes @@ -206,19 +224,27 @@ function smesolveEnsembleProblem( e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, params = NullParameters(), rng::AbstractRNG = default_rng(), - ntraj::Int = 1, - ensemble_method = EnsembleThreads(), + ntraj::Int = 500, + ensemblealg::EnsembleAlgorithm = EnsembleThreads(), 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..., ) where {StateOpType<:Union{KetQuantumObject,OperatorQuantumObject}} _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 + _ensemble_dispatch_output_func(ensemblealg, progress_bar, ntraj, _stochastic_output_func) : output_func prob_sme = smesolveProblem( H, @@ -230,6 +256,7 @@ function smesolveEnsembleProblem( params = params, rng = rng, progress_bar = Val(false), + store_measurement = makeVal(store_measurement), kwargs..., ) @@ -254,11 +281,12 @@ end e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, params = NullParameters(), rng::AbstractRNG = default_rng(), - ntraj::Int = 1, - ensemble_method = EnsembleThreads(), + ntraj::Int = 500, + ensemblealg::EnsembleAlgorithm = EnsembleThreads(), 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..., ) @@ -293,11 +321,12 @@ Above, ``\hat{C}_i`` represent the collapse operators related to pure dissipatio - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`. - `params`: `NullParameters` of parameters to pass to the solver. - `rng`: Random number generator for reproducibility. -- `ntraj`: Number of trajectories to use. -- `ensemble_method`: Ensemble method to use. Default to `EnsembleThreads()`. +- `ntraj`: Number of trajectories to use. Default is `500`. +- `ensemblealg`: Ensemble method to use. Default to `EnsembleThreads()`. - `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 expectation values. Default is `Val(false)`. - `kwargs`: The keyword arguments for the ODEProblem. # Notes @@ -321,11 +350,12 @@ function smesolve( e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, params = NullParameters(), rng::AbstractRNG = default_rng(), - ntraj::Int = 1, - ensemble_method = EnsembleThreads(), + ntraj::Int = 500, + ensemblealg::EnsembleAlgorithm = EnsembleThreads(), 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..., ) where {StateOpType<:Union{KetQuantumObject,OperatorQuantumObject}} ensemble_prob = smesolveEnsembleProblem( @@ -338,34 +368,41 @@ function smesolve( params = params, rng = rng, ntraj = ntraj, - ensemble_method = ensemble_method, + ensemblealg = ensemblealg, prob_func = prob_func, output_func = output_func, progress_bar = progress_bar, + store_measurement = makeVal(store_measurement), kwargs..., ) - return smesolve(ensemble_prob, alg, ntraj, ensemble_method) + return smesolve(ensemble_prob, alg, ntraj, ensemblealg) end function smesolve( ens_prob::TimeEvolutionProblem, alg::StochasticDiffEqAlgorithm = SRA1(), - ntraj::Int = 1, - ensemble_method = EnsembleThreads(), + ntraj::Int = 500, + ensemblealg::EnsembleAlgorithm = EnsembleThreads(), ) - sol = _ensemble_dispatch_solve(ens_prob, alg, ensemble_method, ntraj) + sol = _ensemble_dispatch_solve(ens_prob, alg, ensemblealg, ntraj) _sol_1 = sol[:, 1] - _expvals_sol_1 = _se_me_sse_get_expvals(_sol_1) + _expvals_sol_1 = _get_expvals(_sol_1, SaveFuncMESolve) + _m_expvals_sol_1 = _get_m_expvals(_sol_1, SaveFuncSMESolve) dims = ens_prob.dimensions - _expvals_all = _expvals_sol_1 isa Nothing ? nothing : map(i -> _se_me_sse_get_expvals(sol[:, i]), eachindex(sol)) + _expvals_all = + _expvals_sol_1 isa Nothing ? nothing : map(i -> _get_expvals(sol[:, i], SaveFuncMESolve), eachindex(sol)) expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all, dims = 2) # Stack on dimension 2 to align with QuTiP states = map(i -> _smesolve_generate_state.(sol[:, i].u, Ref(dims)), eachindex(sol)) + _m_expvals = + _m_expvals_sol_1 isa Nothing ? nothing : map(i -> _get_m_expvals(sol[:, i], SaveFuncSMESolve), eachindex(sol)) + m_expvals = _m_expvals isa Nothing ? nothing : stack(_m_expvals, dims = 2) + expvals = - _se_me_sse_get_expvals(_sol_1) isa Nothing ? nothing : + _get_expvals(_sol_1, SaveFuncMESolve) isa Nothing ? nothing : dropdims(sum(expvals_all, dims = 2), dims = 2) ./ length(sol) return TimeEvolutionStochasticSol( @@ -375,6 +412,7 @@ function smesolve( expvals, expvals, # This is average_expect expvals_all, + m_expvals, # Measurement expectation values sol.converged, _sol_1.alg, _sol_1.prob.kwargs[:abstol], diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl index d137d5352..c428b1547 100644 --- a/src/time_evolution/ssesolve.jl +++ b/src/time_evolution/ssesolve.jl @@ -22,6 +22,7 @@ _ScalarOperator_e2_2(op, f = +) = params = NullParameters(), rng::AbstractRNG = default_rng(), progress_bar::Union{Val,Bool} = Val(true), + store_measurement::Union{Val, Bool} = Val(false), kwargs..., ) @@ -56,6 +57,7 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is th - `params`: `NullParameters` of parameters to pass to the solver. - `rng`: Random number generator for reproducibility. - `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 @@ -78,6 +80,7 @@ function ssesolveProblem( params = NullParameters(), rng::AbstractRNG = default_rng(), progress_bar::Union{Val,Bool} = Val(true), + store_measurement::Union{Val,Bool} = Val(false), kwargs..., ) haskey(kwargs, :save_idxs) && @@ -111,12 +114,24 @@ 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!(tlist[1], zeros(length(sc_ops)), zeros(length(sc_ops)), save_everystep = false, rng = rng) + noise = RealWienerProcess!( + tlist[1], + zeros(length(sc_ops)), + zeros(length(sc_ops)), + save_everystep = getVal(store_measurement), + rng = rng, + ) noise_rate_prototype = similar(ψ0, length(ψ0), length(sc_ops)) prob = SDEProblem{true}( K, @@ -126,7 +141,7 @@ function ssesolveProblem( params; noise_rate_prototype = noise_rate_prototype, noise = noise, - kwargs4..., + kwargs3..., ) return TimeEvolutionProblem(prob, tlist, dims) @@ -141,11 +156,12 @@ end e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, params = NullParameters(), rng::AbstractRNG = default_rng(), - ntraj::Int = 1, - ensemble_method = EnsembleThreads(), + ntraj::Int = 500, + ensemblealg::EnsembleAlgorithm = EnsembleThreads(), 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..., ) @@ -179,12 +195,13 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is t - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`. - `params`: `NullParameters` of parameters to pass to the solver. - `rng`: Random number generator for reproducibility. -- `ntraj`: Number of trajectories to use. -- `ensemble_method`: Ensemble method to use. Default to `EnsembleThreads()`. +- `ntraj`: Number of trajectories to use. Default is `500`. +- `ensemblealg`: Ensemble method to use. Default to `EnsembleThreads()`. - `jump_callback`: The Jump Callback type: Discrete or Continuous. The default is `ContinuousLindbladJumpCallback()`, which is more precise. - `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 @@ -206,19 +223,27 @@ function ssesolveEnsembleProblem( e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, params = NullParameters(), rng::AbstractRNG = default_rng(), - ntraj::Int = 1, - ensemble_method = EnsembleThreads(), + ntraj::Int = 500, + ensemblealg::EnsembleAlgorithm = EnsembleThreads(), 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 + _ensemble_dispatch_output_func(ensemblealg, progress_bar, ntraj, _stochastic_output_func) : output_func prob_sme = ssesolveProblem( H, @@ -229,6 +254,7 @@ function ssesolveEnsembleProblem( params = params, rng = rng, progress_bar = Val(false), + store_measurement = makeVal(store_measurement), kwargs..., ) @@ -252,11 +278,12 @@ end e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, params = NullParameters(), rng::AbstractRNG = default_rng(), - ntraj::Int = 1, - ensemble_method = EnsembleThreads(), + ntraj::Int = 500, + ensemblealg::EnsembleAlgorithm = EnsembleThreads(), 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..., ) @@ -294,11 +321,12 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is th - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`. - `params`: `NullParameters` of parameters to pass to the solver. - `rng`: Random number generator for reproducibility. -- `ntraj`: Number of trajectories to use. -- `ensemble_method`: Ensemble method to use. Default to `EnsembleThreads()`. +- `ntraj`: Number of trajectories to use. Default is `500`. +- `ensemblealg`: Ensemble method to use. Default to `EnsembleThreads()`. - `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 @@ -322,11 +350,12 @@ function ssesolve( e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, params = NullParameters(), rng::AbstractRNG = default_rng(), - ntraj::Int = 1, - ensemble_method = EnsembleThreads(), + ntraj::Int = 500, + ensemblealg::EnsembleAlgorithm = EnsembleThreads(), 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( @@ -338,35 +367,42 @@ function ssesolve( params = params, rng = rng, ntraj = ntraj, - ensemble_method = ensemble_method, + ensemblealg = ensemblealg, prob_func = prob_func, output_func = output_func, progress_bar = progress_bar, + store_measurement = makeVal(store_measurement), kwargs..., ) - return ssesolve(ens_prob, alg, ntraj, ensemble_method) + return ssesolve(ens_prob, alg, ntraj, ensemblealg) end function ssesolve( ens_prob::TimeEvolutionProblem, alg::StochasticDiffEqAlgorithm = SRA1(), - ntraj::Int = 1, - ensemble_method = EnsembleThreads(), + ntraj::Int = 500, + ensemblealg::EnsembleAlgorithm = EnsembleThreads(), ) - sol = _ensemble_dispatch_solve(ens_prob, alg, ensemble_method, ntraj) + sol = _ensemble_dispatch_solve(ens_prob, alg, ensemblealg, ntraj) _sol_1 = sol[:, 1] - _expvals_sol_1 = _se_me_sse_get_expvals(_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 - _expvals_all = _expvals_sol_1 isa Nothing ? nothing : map(i -> _se_me_sse_get_expvals(sol[:, i]), eachindex(sol)) + _expvals_all = + _expvals_sol_1 isa Nothing ? nothing : map(i -> _get_expvals(sol[:, i], SaveFuncSSESolve), eachindex(sol)) 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 = - _se_me_sse_get_expvals(_sol_1) isa Nothing ? nothing : + _get_expvals(_sol_1, SaveFuncSSESolve) isa Nothing ? nothing : dropdims(sum(expvals_all, dims = 2), dims = 2) ./ length(sol) return TimeEvolutionStochasticSol( @@ -376,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], diff --git a/src/time_evolution/time_evolution.jl b/src/time_evolution/time_evolution.jl index c32805524..737837358 100644 --- a/src/time_evolution/time_evolution.jl +++ b/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 @@ -349,7 +351,7 @@ function _stochastic_prob_func(prob, i, repeat, rng, seeds, tlist; kwargs...) prob.prob.tspan[1], zeros(kwargs[:n_sc_ops]), zeros(kwargs[:n_sc_ops]), - save_everystep = false, + save_everystep = getVal(kwargs[:store_measurement]), rng = traj_rng, ) diff --git a/src/time_evolution/time_evolution_dynamical.jl b/src/time_evolution/time_evolution_dynamical.jl index 77969751e..c320946e8 100644 --- a/src/time_evolution/time_evolution_dynamical.jl +++ b/src/time_evolution/time_evolution_dynamical.jl @@ -245,7 +245,7 @@ function dfd_mesolve( return TimeEvolutionSol( dfd_prob.times, ρt, - _se_me_sse_get_expvals(sol), + _get_expvals(sol, SaveFuncMESolve), sol.retcode, sol.alg, NamedTuple(sol.prob.kwargs).abstol, @@ -522,7 +522,7 @@ function _DSF_mcsolve_Affect!(integrator) # e_ops0 = params.e_ops # c_ops0 = params.c_ops - e_ops0 = _mcsolve_get_e_ops(integrator) + e_ops0 = _get_e_ops(integrator, SaveFuncMCSolve) c_ops0, c_ops0_herm = _mcsolve_get_c_ops(integrator) copyto!(ψt, integrator.u) diff --git a/test/core-test/time_evolution.jl b/test/core-test/time_evolution.jl index 2b9d0afe7..419c17677 100644 --- a/test/core-test/time_evolution.jl +++ b/test/core-test/time_evolution.jl @@ -24,7 +24,7 @@ ψ0_int = Qobj(round.(Int, real.(ψ0.data)), dims = ψ0.dims) # Used for testing the type inference - @testset "sesolve" begin + @testset "sesolve" verbose = true begin tlist = range(0, 20 * 2π / g, 1000) saveat_idxs = 500:900 saveat = tlist[saveat_idxs] @@ -100,7 +100,7 @@ end end - @testset "mesolve, mcsolve, ssesolve and smesolve" begin + @testset "mesolve, mcsolve, ssesolve and smesolve" verbose = true begin tlist = range(0, 10 / γ, 100) saveat_idxs = 50:90 saveat = tlist[saveat_idxs] @@ -130,8 +130,29 @@ progress_bar = Val(false), jump_callback = DiscreteLindbladJumpCallback(), ) - sol_sse = ssesolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false)) - sol_sme = smesolve(H, ψ0, tlist, c_ops_sme, sc_ops_sme, ntraj = 500, e_ops = e_ops, progress_bar = Val(false)) + sol_sse = ssesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false)) + sol_sse2 = ssesolve( + H, + ψ0, + tlist, + c_ops, + e_ops = e_ops, + ntraj = 20, + progress_bar = Val(false), + store_measurement = Val(true), + ) + sol_sme = smesolve(H, ψ0, tlist, c_ops_sme, sc_ops_sme, e_ops = e_ops, progress_bar = Val(false)) + sol_sme2 = smesolve( + H, + ψ0, + tlist, + c_ops_sme, + sc_ops_sme, + e_ops = e_ops, + ntraj = 20, + progress_bar = Val(false), + store_measurement = Val(true), + ) ρt_mc = [ket2dm.(normalize.(states)) for states in sol_mc_states.states] expect_mc_states = mapreduce(states -> expect.(Ref(e_ops[1]), states), hcat, ρt_mc) @@ -172,6 +193,11 @@ @test size(sol_sse.expect) == (length(e_ops), length(tlist)) @test length(sol_sme.times) == length(tlist) @test size(sol_sme.expect) == (length(e_ops), length(tlist)) + @test isnothing(sol_sse.measurement) + @test isnothing(sol_sme.measurement) + @test size(sol_sse2.measurement) == (length(c_ops), 20, length(tlist) - 1) + @test size(sol_sme2.measurement) == (length(sc_ops_sme), 20, length(tlist) - 1) + @test sol_me_string == "Solution of time evolution\n" * "(return code: $(sol_me.retcode))\n" * @@ -255,7 +281,7 @@ sol_se = sesolve(H_dr_fr, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false)) sol_me = mesolve(H_dr_fr, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false)) sol_mc = mcsolve(H_dr_fr, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), rng = rng) - # sol_sse = ssesolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng) + # sol_sse = ssesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), rng = rng) # Time Evolution in the lab frame @@ -268,7 +294,7 @@ sol_se_td = sesolve(H_td, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false), params = p) sol_me_td = mesolve(H_td, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p) sol_mc_td = mcsolve(H_td, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng) - # sol_sse_td = ssesolve(H_td, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng) + # sol_sse_td = ssesolve(H_td, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng) @test sol_se.expect ≈ sol_se_td.expect atol = 1e-6 * length(tlist) @test sol_me.expect ≈ sol_me_td.expect atol = 1e-6 * length(tlist) @@ -282,7 +308,7 @@ sol_me_td2 = mesolve(L_td, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p) sol_mc_td2 = mcsolve(H_td2, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng) # sol_sse_td2 = - # ssesolve(H_td2, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng) + # ssesolve(H_td2, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false), params = p, rng = rng) @test sol_se.expect ≈ sol_se_td2.expect atol = 1e-6 * length(tlist) @test sol_me.expect ≈ sol_me_td2.expect atol = 1e-6 * length(tlist)