diff --git a/.typos.toml b/.typos.toml index 73648865c..e3264a3d3 100644 --- a/.typos.toml +++ b/.typos.toml @@ -1,2 +1,3 @@ [default.extend-words] -ket = "ket" \ No newline at end of file +ket = "ket" +sme = "sme" diff --git a/CHANGELOG.md b/CHANGELOG.md index 25e8f3888..08a281c88 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Fix CUDA `sparse_to_dense`. ([#386]) - Improve pseudo inverse spectrum solver. ([#388]) +- Add `smesolve` function for stochastic master equation. ([#389]) ## [v0.25.2] Release date: 2025-02-02 @@ -109,3 +110,4 @@ Release date: 2024-11-13 [#383]: https://github.com/qutip/QuantumToolbox.jl/issues/383 [#386]: https://github.com/qutip/QuantumToolbox.jl/issues/386 [#388]: https://github.com/qutip/QuantumToolbox.jl/issues/388 +[#389]: https://github.com/qutip/QuantumToolbox.jl/issues/389 diff --git a/docs/src/resources/api.md b/docs/src/resources/api.md index 1063f21f3..d3d09c3c1 100644 --- a/docs/src/resources/api.md +++ b/docs/src/resources/api.md @@ -189,17 +189,20 @@ cosm TimeEvolutionProblem TimeEvolutionSol TimeEvolutionMCSol -TimeEvolutionSSESol +TimeEvolutionStochasticSol sesolveProblem mesolveProblem mcsolveProblem mcsolveEnsembleProblem ssesolveProblem ssesolveEnsembleProblem +smesolveProblem +smesolveEnsembleProblem sesolve mesolve mcsolve ssesolve +smesolve dfd_mesolve liouvillian liouvillian_generalized diff --git a/docs/src/users_guide/time_evolution/intro.md b/docs/src/users_guide/time_evolution/intro.md index 3b9ccf985..aabbb25a0 100644 --- a/docs/src/users_guide/time_evolution/intro.md +++ b/docs/src/users_guide/time_evolution/intro.md @@ -36,7 +36,8 @@ The following table lists the solvers provided by `QuantumToolbox` for dynamic q | Unitary evolution, Schrödinger equation | [`sesolve`](@ref) | [`sesolveProblem`](@ref) | [`TimeEvolutionSol`](@ref) | | Lindblad master eqn. or Von Neuman eqn.| [`mesolve`](@ref) | [`mesolveProblem`](@ref) | [`TimeEvolutionSol`](@ref) | | Monte Carlo evolution | [`mcsolve`](@ref) | [`mcsolveProblem`](@ref) [`mcsolveEnsembleProblem`](@ref) | [`TimeEvolutionMCSol`](@ref) | -| Stochastic Schrödinger equation | [`ssesolve`](@ref) | [`ssesolveProblem`](@ref) [`ssesolveEnsembleProblem`](@ref) | [`TimeEvolutionSSESol`](@ref) | +| Stochastic Schrödinger equation | [`ssesolve`](@ref) | [`ssesolveProblem`](@ref) [`ssesolveEnsembleProblem`](@ref) | [`TimeEvolutionStochasticSol`](@ref) | +| Stochastic master equation | [`smesolve`](@ref) | [`smesolveProblem`](@ref) [`smesolveEnsembleProblem`](@ref) | [`TimeEvolutionStochasticSol`](@ref) | !!! note "Solving dynamics with pre-defined problems" `QuantumToolbox` provides two different methods to solve the dynamics. One can use the function calls listed above by either taking all the operators (like Hamiltonian and collapse operators, etc.) as inputs directly, or generating the `prob`lems by yourself and take it as an input of the function call, e.g., `sesolve(prob)`. diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl index 3677bb178..4776c0d43 100644 --- a/src/QuantumToolbox.jl +++ b/src/QuantumToolbox.jl @@ -106,6 +106,7 @@ include("time_evolution/lr_mesolve.jl") include("time_evolution/sesolve.jl") include("time_evolution/mcsolve.jl") include("time_evolution/ssesolve.jl") +include("time_evolution/smesolve.jl") include("time_evolution/time_evolution_dynamical.jl") # Others diff --git a/src/qobj/superoperators.jl b/src/qobj/superoperators.jl index 7fe381478..affc0a980 100644 --- a/src/qobj/superoperators.jl +++ b/src/qobj/superoperators.jl @@ -164,22 +164,34 @@ function liouvillian( ) where {OpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject}} L = liouvillian(H, Id_cache) if !(c_ops isa Nothing) - # sum all the (time-independent) c_ops first - c_ops_ti = filter(op -> isa(op, QuantumObject), c_ops) - if !isempty(c_ops_ti) - L += mapreduce(op -> lindblad_dissipator(op, Id_cache), +, c_ops_ti) - end - - # sum rest of the QobjEvo together - c_ops_td = filter(op -> isa(op, QuantumObjectEvolution), c_ops) - if !isempty(c_ops_td) - L += mapreduce(op -> lindblad_dissipator(op, Id_cache), +, c_ops_td) - end + L += _sum_lindblad_dissipators(c_ops, Id_cache) end return L end +liouvillian(H::Nothing, c_ops::Union{AbstractVector,Tuple}, Id_cache::Diagonal = I(prod(c_ops[1].dims))) = + _sum_lindblad_dissipators(c_ops, Id_cache) + +liouvillian(H::Nothing, c_ops::Nothing) = 0 + liouvillian(H::AbstractQuantumObject{OperatorQuantumObject}, Id_cache::Diagonal = I(prod(H.dimensions))) = get_typename_wrapper(H)(_liouvillian(H.data, Id_cache), SuperOperator, H.dimensions) liouvillian(H::AbstractQuantumObject{SuperOperatorQuantumObject}, Id_cache::Diagonal) = H + +function _sum_lindblad_dissipators(c_ops, Id_cache::Diagonal) + D = 0 + # sum all the (time-independent) c_ops first + c_ops_ti = filter(op -> isa(op, QuantumObject), c_ops) + if !isempty(c_ops_ti) + D += mapreduce(op -> lindblad_dissipator(op, Id_cache), +, c_ops_ti) + end + + # sum rest of the QobjEvo together + c_ops_td = filter(op -> isa(op, QuantumObjectEvolution), c_ops) + if !isempty(c_ops_td) + D += mapreduce(op -> lindblad_dissipator(op, Id_cache), +, c_ops_td) + end + + return D +end diff --git a/src/time_evolution/mcsolve.jl b/src/time_evolution/mcsolve.jl index 94c0372f8..7c006f8e8 100644 --- a/src/time_evolution/mcsolve.jl +++ b/src/time_evolution/mcsolve.jl @@ -12,11 +12,6 @@ function _mcsolve_prob_func(prob, i, repeat, global_rng, seeds, tlist) return remake(prob, f = f, callback = cb) end -function _mcsolve_dispatch_prob_func(rng, ntraj, tlist) - seeds = map(i -> rand(rng, UInt64), 1:ntraj) - return (prob, i, repeat) -> _mcsolve_prob_func(prob, i, repeat, rng, seeds, tlist) -end - # Standard output function function _mcsolve_output_func(sol, i) idx = _mc_get_jump_callback(sol).affect!.jump_times_which_idx[] @@ -25,43 +20,6 @@ function _mcsolve_output_func(sol, i) return (sol, false) end -# Output function with progress bar update -function _mcsolve_output_func_progress(sol, i, progr) - next!(progr) - return _mcsolve_output_func(sol, i) -end - -# Output function with distributed channel update for progress bar -function _mcsolve_output_func_distributed(sol, i, channel) - put!(channel, true) - return _mcsolve_output_func(sol, i) -end - -function _mcsolve_dispatch_output_func(::ET, progress_bar, ntraj) where {ET<:Union{EnsembleSerial,EnsembleThreads}} - if getVal(progress_bar) - progr = ProgressBar(ntraj, enable = getVal(progress_bar)) - f = (sol, i) -> _mcsolve_output_func_progress(sol, i, progr) - return (f, progr, nothing) - else - return (_mcsolve_output_func, nothing, nothing) - end -end -function _mcsolve_dispatch_output_func( - ::ET, - progress_bar, - ntraj, -) where {ET<:Union{EnsembleSplitThreads,EnsembleDistributed}} - if getVal(progress_bar) - progr = ProgressBar(ntraj, enable = getVal(progress_bar)) - progr_channel::RemoteChannel{Channel{Bool}} = RemoteChannel(() -> Channel{Bool}(1)) - - f = (sol, i) -> _mcsolve_output_func_distributed(sol, i, progr_channel) - return (f, progr, progr_channel) - else - return (_mcsolve_output_func, nothing, nothing) - end -end - function _normalize_state!(u, dims, normalize_states) getVal(normalize_states) && normalize!(u) return QuantumObject(u, type = Ket, dims = dims) @@ -280,9 +238,10 @@ function mcsolveEnsembleProblem( output_func::Union{Tuple,Nothing} = nothing, kwargs..., ) where {TJC<:LindbladJumpCallbackType} - _prob_func = prob_func isa Nothing ? _mcsolve_dispatch_prob_func(rng, ntraj, tlist) : prob_func + _prob_func = isnothing(prob_func) ? _ensemble_dispatch_prob_func(rng, ntraj, tlist, _mcsolve_prob_func) : prob_func _output_func = - output_func isa Nothing ? _mcsolve_dispatch_output_func(ensemble_method, progress_bar, ntraj) : output_func + output_func isa Nothing ? + _ensemble_dispatch_output_func(ensemble_method, progress_bar, ntraj, _mcsolve_output_func) : output_func prob_mc = mcsolveProblem( H, @@ -431,38 +390,6 @@ function mcsolve( return mcsolve(ens_prob_mc, alg, ntraj, ensemble_method, normalize_states) end -function _mcsolve_solve_ens( - ens_prob_mc::TimeEvolutionProblem, - alg::OrdinaryDiffEqAlgorithm, - ensemble_method::ET, - ntraj::Int, -) where {ET<:Union{EnsembleSplitThreads,EnsembleDistributed}} - sol = nothing - - @sync begin - @async while take!(ens_prob_mc.kwargs.channel) - next!(ens_prob_mc.kwargs.progr) - end - - @async begin - sol = solve(ens_prob_mc.prob, alg, ensemble_method, trajectories = ntraj) - put!(ens_prob_mc.kwargs.channel, false) - end - end - - return sol -end - -function _mcsolve_solve_ens( - ens_prob_mc::TimeEvolutionProblem, - alg::OrdinaryDiffEqAlgorithm, - ensemble_method, - ntraj::Int, -) - sol = solve(ens_prob_mc.prob, alg, ensemble_method, trajectories = ntraj) - return sol -end - function mcsolve( ens_prob_mc::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit5(), @@ -470,7 +397,7 @@ function mcsolve( ensemble_method = EnsembleThreads(), normalize_states = Val(true), ) - sol = _mcsolve_solve_ens(ens_prob_mc, alg, ensemble_method, ntraj) + sol = _ensemble_dispatch_solve(ens_prob_mc, alg, ensemble_method, ntraj) dims = ens_prob_mc.dimensions _sol_1 = sol[:, 1] diff --git a/src/time_evolution/mesolve.jl b/src/time_evolution/mesolve.jl index ae1a40667..a87bc3864 100644 --- a/src/time_evolution/mesolve.jl +++ b/src/time_evolution/mesolve.jl @@ -1,7 +1,9 @@ export mesolveProblem, mesolve -_mesolve_make_L_QobjEvo(H::QuantumObject, c_ops) = QobjEvo(liouvillian(H, c_ops); type = SuperOperator) +_mesolve_make_L_QobjEvo(H::Union{QuantumObject,Nothing}, c_ops) = QobjEvo(liouvillian(H, c_ops); type = SuperOperator) _mesolve_make_L_QobjEvo(H::Union{QuantumObjectEvolution,Tuple}, c_ops) = liouvillian(QobjEvo(H), c_ops) +_mesolve_make_L_QobjEvo(H::Nothing, c_ops::Nothing) = throw(ArgumentError("Both H and +c_ops are Nothing. You are probably running the wrong function.")) @doc raw""" mesolveProblem( @@ -31,7 +33,7 @@ where # Arguments - `H`: Hamiltonian of the system ``\hat{H}``. It can be either a [`QuantumObject`](@ref), a [`QuantumObjectEvolution`](@ref), or a `Tuple` of operator-function pairs. -- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. +- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. It can be either a [`Ket`](@ref) or a [`Operator`](@ref). - `tlist`: List of times at which to save either the state or the expectation values of the system. - `c_ops`: List of collapse operators ``\{\hat{C}_n\}_n``. It can be either a `Vector` or a `Tuple`. - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`. @@ -119,7 +121,7 @@ where # Arguments - `H`: Hamiltonian of the system ``\hat{H}``. It can be either a [`QuantumObject`](@ref), a [`QuantumObjectEvolution`](@ref), or a `Tuple` of operator-function pairs. -- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. +- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. It can be either a [`Ket`](@ref) or a [`Operator`](@ref). - `tlist`: List of times at which to save either the state or the expectation values of the system. - `c_ops`: List of collapse operators ``\{\hat{C}_n\}_n``. It can be either a `Vector` or a `Tuple`. - `alg`: The algorithm for the ODE solver. The default value is `Tsit5()`. diff --git a/src/time_evolution/smesolve.jl b/src/time_evolution/smesolve.jl new file mode 100644 index 000000000..904951d42 --- /dev/null +++ b/src/time_evolution/smesolve.jl @@ -0,0 +1,375 @@ +export smesolveProblem, smesolveEnsembleProblem, smesolve + +_smesolve_generate_state(u, dims) = QuantumObject(vec2mat(u), type = Operator, dims = dims) + +function _smesolve_update_coeff(u, p, t, op_vec) + return real(dot(u, op_vec)) / 2 #this is Tr[Sn * ρ + ρ * Sn'] +end + +_smesolve_ScalarOperator(op_vec) = + ScalarOperator(one(eltype(op_vec)), (a, u, p, t) -> -_smesolve_update_coeff(u, p, t, op_vec)) + +@doc raw""" + smesolveProblem( + H::Union{AbstractQuantumObject{OperatorQuantumObject},Tuple}, + ψ0::QuantumObject{KetQuantumObject}, + tlist::AbstractVector, + c_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing; + e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + params::NamedTuple = NamedTuple(), + rng::AbstractRNG = default_rng(), + progress_bar::Union{Val,Bool} = Val(true), + kwargs..., + ) + +Generate the SDEProblem for the Stochastic Master Equation time evolution of an open quantum system. This is defined by the following stochastic differential equation: + +```math +d| \rho (t) = -i [\hat{H}, \rho(t)] dt + \sum_n \mathcal{D}[\hat{C}_n] \rho(t) dt + \sum_n \mathcal{D}[\hat{S}_n] \rho(t) dt + \sum_n \mathcal{H}[\hat{S}_n] \rho(t) dW_n(t), +``` + +where + +```math +\mathcal{D}[\hat{O}] \rho = \hat{O} \rho \hat{O}^\dagger - \frac{1}{2} \{\hat{O}^\dagger \hat{O}, \rho\}, +``` + +is the Lindblad superoperator, and + +```math +\mathcal{H}[\hat{O}] \rho = \hat{O} \rho + \rho \hat{O}^\dagger - \mathrm{Tr}[\hat{O} \rho + \rho \hat{O}^\dagger] \rho, + +Above, ``\hat{C}_n`` represent the operators related to pure dissipation, while ``\hat{S}_n`` are the measurement operators. The ``dW_n(t)`` term is the real Wiener increment associated to ``\hat{S}_n``. See [Wiseman2009Quantum](@cite) for more details. + +# Arguments + +- `H`: Hamiltonian of the system ``\hat{H}``. It can be either a [`QuantumObject`](@ref), a [`QuantumObjectEvolution`](@ref), or a `Tuple` of operator-function pairs. +- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. It can be either a [`Ket`](@ref) or a [`Operator`](@ref). +- `tlist`: List of times at which to save either the state or the expectation values of the system. +- `c_ops`: List of collapse operators ``\{\hat{C}_n\}_n``. It can be either a `Vector` or a `Tuple`. +- `sc_ops`: List of measurement collapse operators ``\{\hat{S}_n\}_n``. It can be either a `Vector` or a `Tuple`. +- `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`. +- `params`: `NamedTuple` 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. +- `kwargs`: The keyword arguments for the ODEProblem. + +# Notes + +- The states will be saved depend on the keyword argument `saveat` in `kwargs`. +- If `e_ops` is empty, the default value of `saveat=tlist` (saving the states corresponding to `tlist`), otherwise, `saveat=[tlist[end]]` (only save the final state). You can also specify `e_ops` and `saveat` separately. +- The default tolerances in `kwargs` are given as `reltol=1e-2` and `abstol=1e-2`. +- For more details about `kwargs` please refer to [`DifferentialEquations.jl` (Keyword Arguments)](https://docs.sciml.ai/DiffEqDocs/stable/basics/common_solver_opts/) + +# Returns + +- `prob`: The [`TimeEvolutionProblem`](@ref) containing the `SDEProblem` for the Stochastic Master Equation time evolution. +""" +function smesolveProblem( + H::Union{AbstractQuantumObject{OperatorQuantumObject},Tuple}, + ψ0::QuantumObject{KetQuantumObject}, + tlist::AbstractVector, + c_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing; + e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + params::NamedTuple = NamedTuple(), + rng::AbstractRNG = default_rng(), + progress_bar::Union{Val,Bool} = Val(true), + kwargs..., +) + haskey(kwargs, :save_idxs) && + throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox.")) + + isnothing(sc_ops) && + throw(ArgumentError("The list of measurement collapse operators must be provided. Use mesolveProblem instead.")) + + tlist = _check_tlist(tlist, _FType(ψ0)) + + L_evo = _mesolve_make_L_QobjEvo(H, c_ops) + _mesolve_make_L_QobjEvo(nothing, sc_ops) + check_dimensions(L_evo, ψ0) + dims = L_evo.dimensions + + T = Base.promote_eltype(L_evo, ψ0) + ρ0 = sparse_to_dense(_CType(T), mat2vec(ket2dm(ψ0).data)) # Convert it to dense vector with complex element type + + progr = ProgressBar(length(tlist), enable = getVal(progress_bar)) + + sc_ops_evo_data = Tuple(map(get_data ∘ QobjEvo, sc_ops)) + + K = get_data(L_evo) + + Id = I(prod(dims)) + D_l = map(sc_ops_evo_data) do op + # TODO: Implement the three-argument dot function for SciMLOperators.jl + # Currently, we are assuming a time-independent MatrixOperator + op_vec = mat2vec(adjoint(op.A)) + return _spre(op, Id) + _spost(op', Id) + _smesolve_ScalarOperator(op_vec) * IdentityOperator(prod(dims)^2) + end + D = DiffusionOperator(D_l) + + p = (progr = progr, times = tlist, Hdims = dims, n_sc_ops = length(sc_ops), params...) + + is_empty_e_ops = (e_ops isa Nothing) ? true : isempty(e_ops) + + saveat = is_empty_e_ops ? tlist : [tlist[end]] + default_values = (DEFAULT_SDE_SOLVER_OPTIONS..., saveat = saveat) + kwargs2 = merge(default_values, kwargs) + kwargs3 = _generate_se_me_kwargs(e_ops, makeVal(progress_bar), tlist, kwargs2, SaveFuncMESolve) + + tspan = (tlist[1], tlist[end]) + noise = + RealWienerProcess!(tlist[1], zeros(length(sc_ops)), zeros(length(sc_ops)), save_everystep = false, rng = rng) + noise_rate_prototype = similar(ρ0, length(ρ0), length(sc_ops)) + prob = SDEProblem{true}(K, D, ρ0, tspan, p; noise_rate_prototype = noise_rate_prototype, noise = noise, kwargs3...) + + return TimeEvolutionProblem(prob, tlist, dims) +end + +@doc raw""" + smesolveEnsembleProblem( + H::Union{AbstractQuantumObject{OperatorQuantumObject},Tuple}, + ψ0::QuantumObject{KetQuantumObject}, + tlist::AbstractVector, + c_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing; + e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + params::NamedTuple = NamedTuple(), + rng::AbstractRNG = default_rng(), + ntraj::Int = 1, + ensemble_method = EnsembleThreads(), + prob_func::Union{Function, Nothing} = nothing, + output_func::Union{Tuple,Nothing} = nothing, + progress_bar::Union{Val,Bool} = Val(true), + kwargs..., + ) + +Generate the SDEProblem for the Stochastic Master Equation time evolution of an open quantum system. This is defined by the following stochastic differential equation: + +```math +d| \rho (t) = -i [\hat{H}, \rho(t)] dt + \sum_n \mathcal{D}[\hat{C}_n] \rho(t) dt + \sum_n \mathcal{D}[\hat{S}_n] \rho(t) dt + \sum_n \mathcal{H}[\hat{S}_n] \rho(t) dW_n(t), +``` + +where + +```math +\mathcal{D}[\hat{O}] \rho = \hat{O} \rho \hat{O}^\dagger - \frac{1}{2} \{\hat{O}^\dagger \hat{O}, \rho\}, +``` + +is the Lindblad superoperator, and + +```math +\mathcal{H}[\hat{O}] \rho = \hat{O} \rho + \rho \hat{O}^\dagger - \mathrm{Tr}[\hat{O} \rho + \rho \hat{O}^\dagger] \rho, + +Above, ``\hat{C}_n`` represent the operators related to pure dissipation, while ``\hat{S}_n`` are the measurement operators. The ``dW_n(t)`` term is the real Wiener increment associated to ``\hat{S}_n``. See [Wiseman2009Quantum](@cite) for more details. + +# Arguments + +- `H`: Hamiltonian of the system ``\hat{H}``. It can be either a [`QuantumObject`](@ref), a [`QuantumObjectEvolution`](@ref), or a `Tuple` of operator-function pairs. +- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. It can be either a [`Ket`](@ref) or a [`Operator`](@ref). +- `tlist`: List of times at which to save either the state or the expectation values of the system. +- `c_ops`: List of collapse operators ``\{\hat{C}_n\}_n``. It can be either a `Vector` or a `Tuple`. +- `sc_ops`: List of measurement collapse operators ``\{\hat{S}_n\}_n``. It can be either a `Vector` or a `Tuple`. +- `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`. +- `params`: `NamedTuple` 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()`. +- `prob_func`: Function to use for generating the ODEProblem. +- `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. +- `kwargs`: The keyword arguments for the ODEProblem. + +# Notes + +- The states will be saved depend on the keyword argument `saveat` in `kwargs`. +- If `e_ops` is empty, the default value of `saveat=tlist` (saving the states corresponding to `tlist`), otherwise, `saveat=[tlist[end]]` (only save the final state). You can also specify `e_ops` and `saveat` separately. +- The default tolerances in `kwargs` are given as `reltol=1e-2` and `abstol=1e-2`. +- For more details about `kwargs` please refer to [`DifferentialEquations.jl` (Keyword Arguments)](https://docs.sciml.ai/DiffEqDocs/stable/basics/common_solver_opts/) + +# Returns + +- `prob`: The [`TimeEvolutionProblem`](@ref) containing the Ensemble `SDEProblem` for the Stochastic Master Equation time evolution. +""" +function smesolveEnsembleProblem( + H::Union{AbstractQuantumObject{OperatorQuantumObject},Tuple}, + ψ0::QuantumObject{KetQuantumObject}, + tlist::AbstractVector, + c_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing; + e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + params::NamedTuple = NamedTuple(), + rng::AbstractRNG = default_rng(), + ntraj::Int = 1, + ensemble_method = EnsembleThreads(), + prob_func::Union{Function,Nothing} = nothing, + output_func::Union{Tuple,Nothing} = nothing, + progress_bar::Union{Val,Bool} = Val(true), + kwargs..., +) + _prob_func = + isnothing(prob_func) ? _ensemble_dispatch_prob_func(rng, ntraj, tlist, _stochastic_prob_func) : prob_func + _output_func = + output_func isa Nothing ? + _ensemble_dispatch_output_func(ensemble_method, progress_bar, ntraj, _stochastic_output_func) : output_func + + prob_sme = smesolveProblem( + H, + ψ0, + tlist, + c_ops, + sc_ops; + e_ops = e_ops, + params = params, + rng = rng, + progress_bar = Val(false), + kwargs..., + ) + + ensemble_prob = TimeEvolutionProblem( + EnsembleProblem(prob_sme, prob_func = _prob_func, output_func = _output_func[1], safetycopy = true), + prob_sme.times, + prob_sme.dimensions, + (progr = _output_func[2], channel = _output_func[3]), + ) + + return ensemble_prob +end + +@doc raw""" + smesolve( + H::Union{AbstractQuantumObject{OperatorQuantumObject},Tuple}, + ψ0::QuantumObject{KetQuantumObject}, + tlist::AbstractVector, + c_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing; + alg::StochasticDiffEqAlgorithm = SRA1(), + e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + params::NamedTuple = NamedTuple(), + rng::AbstractRNG = default_rng(), + ntraj::Int = 1, + ensemble_method = EnsembleThreads(), + prob_func::Union{Function, Nothing} = nothing, + output_func::Union{Tuple,Nothing} = nothing, + progress_bar::Union{Val,Bool} = Val(true), + kwargs..., + ) + +Stochastic Master Equation time evolution of an open quantum system. This is defined by the following stochastic differential equation: + +```math +d| \rho (t) = -i [\hat{H}, \rho(t)] dt + \sum_n \mathcal{D}[\hat{C}_n] \rho(t) dt + \sum_n \mathcal{D}[\hat{S}_n] \rho(t) dt + \sum_n \mathcal{H}[\hat{S}_n] \rho(t) dW_n(t), +``` + +where + +```math +\mathcal{D}[\hat{O}] \rho = \hat{O} \rho \hat{O}^\dagger - \frac{1}{2} \{\hat{O}^\dagger \hat{O}, \rho\}, +``` + +is the Lindblad superoperator, and + +```math +\mathcal{H}[\hat{O}] \rho = \hat{O} \rho + \rho \hat{O}^\dagger - \mathrm{Tr}[\hat{O} \rho + \rho \hat{O}^\dagger] \rho, + +Above, ``\hat{C}_n`` represent the operators related to pure dissipation, while ``\hat{S}_n`` are the measurement operators. The ``dW_n(t)`` term is the real Wiener increment associated to ``\hat{S}_n``. See [Wiseman2009Quantum](@cite) for more details. + +# Arguments + +- `H`: Hamiltonian of the system ``\hat{H}``. It can be either a [`QuantumObject`](@ref), a [`QuantumObjectEvolution`](@ref), or a `Tuple` of operator-function pairs. +- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. It can be either a [`Ket`](@ref) or a [`Operator`](@ref). +- `tlist`: List of times at which to save either the state or the expectation values of the system. +- `c_ops`: List of collapse operators ``\{\hat{C}_n\}_n``. It can be either a `Vector` or a `Tuple`. +- `sc_ops`: List of measurement collapse operators ``\{\hat{C}_n\}_n``. It can be either a `Vector` or a `Tuple`. +- `alg`: The algorithm to use for the stochastic differential equation. Default is `SRA1()`. +- `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`. +- `params`: `NamedTuple` 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()`. +- `prob_func`: Function to use for generating the ODEProblem. +- `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. +- `kwargs`: The keyword arguments for the ODEProblem. + +# Notes + +- The states will be saved depend on the keyword argument `saveat` in `kwargs`. +- If `e_ops` is empty, the default value of `saveat=tlist` (saving the states corresponding to `tlist`), otherwise, `saveat=[tlist[end]]` (only save the final state). You can also specify `e_ops` and `saveat` separately. +- The default tolerances in `kwargs` are given as `reltol=1e-2` and `abstol=1e-2`. +- For more details about `kwargs` please refer to [`DifferentialEquations.jl` (Keyword Arguments)](https://docs.sciml.ai/DiffEqDocs/stable/basics/common_solver_opts/) + +# Returns + +- `sol::TimeEvolutionStochasticSol`: The solution of the time evolution. See [`TimeEvolutionStochasticSol`](@ref). +""" +function smesolve( + H::Union{AbstractQuantumObject{OperatorQuantumObject},Tuple}, + ψ0::QuantumObject{KetQuantumObject}, + tlist::AbstractVector, + c_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + sc_ops::Union{Nothing,AbstractVector,Tuple} = nothing; + alg::StochasticDiffEqAlgorithm = SRA1(), + e_ops::Union{Nothing,AbstractVector,Tuple} = nothing, + params::NamedTuple = NamedTuple(), + rng::AbstractRNG = default_rng(), + ntraj::Int = 1, + ensemble_method = EnsembleThreads(), + prob_func::Union{Function,Nothing} = nothing, + output_func::Union{Tuple,Nothing} = nothing, + progress_bar::Union{Val,Bool} = Val(true), + kwargs..., +) + ensemble_prob = smesolveEnsembleProblem( + H, + ψ0, + tlist, + c_ops, + sc_ops; + e_ops = e_ops, + params = params, + rng = rng, + ntraj = ntraj, + ensemble_method = ensemble_method, + prob_func = prob_func, + output_func = output_func, + progress_bar = progress_bar, + kwargs..., + ) + + return smesolve(ensemble_prob, alg, ntraj, ensemble_method) +end + +function smesolve( + ens_prob::TimeEvolutionProblem, + alg::StochasticDiffEqAlgorithm = SRA1(), + ntraj::Int = 1, + ensemble_method = EnsembleThreads(), +) + sol = _ensemble_dispatch_solve(ens_prob, alg, ensemble_method, ntraj) + + _sol_1 = sol[:, 1] + _expvals_sol_1 = _se_me_sse_get_expvals(_sol_1) + + 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_all isa Nothing ? nothing : stack(_expvals_all) + states = map(i -> _smesolve_generate_state.(sol[:, i].u, Ref(dims)), eachindex(sol)) + + expvals = + _se_me_sse_get_expvals(_sol_1) isa Nothing ? nothing : + dropdims(sum(expvals_all, dims = 3), dims = 3) ./ length(sol) + + return TimeEvolutionStochasticSol( + ntraj, + ens_prob.times, + states, + expvals, + expvals_all, + sol.converged, + _sol_1.alg, + _sol_1.prob.kwargs[:abstol], + _sol_1.prob.kwargs[:reltol], + ) +end diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl index e827db017..0b8ec2146 100644 --- a/src/time_evolution/ssesolve.jl +++ b/src/time_evolution/ssesolve.jl @@ -1,5 +1,25 @@ export ssesolveProblem, ssesolveEnsembleProblem, ssesolve +# TODO: Merge this with _stochastic_prob_func +function _ssesolve_prob_func(prob, i, repeat) + internal_params = prob.p + + global_rng = internal_params.global_rng + seed = internal_params.seeds[i] + traj_rng = typeof(global_rng)() + seed!(traj_rng, seed) + + noise = RealWienerProcess!( + prob.tspan[1], + zeros(internal_params.n_sc_ops), + zeros(internal_params.n_sc_ops), + save_everystep = false, + rng = traj_rng, + ) + + return remake(prob, noise = noise, seed = seed) +end + #= struct DiffusionOperator @@ -46,41 +66,18 @@ function _ssesolve_update_coeff(u, p, t, op) return real(dot(u, op.A, u)) #this is en/2: /2 = Re end -function _ssesolve_prob_func(prob, i, repeat) - internal_params = prob.p - - global_rng = internal_params.global_rng - seed = internal_params.seeds[i] - traj_rng = typeof(global_rng)() - seed!(traj_rng, seed) - - noise = RealWienerProcess!( - prob.tspan[1], - zeros(internal_params.n_sc_ops), - zeros(internal_params.n_sc_ops), - save_everystep = false, - rng = traj_rng, - ) - - return remake(prob, noise = noise, seed = seed) -end - -# Standard output function -_ssesolve_output_func(sol, i) = (sol, false) - # Output function with progress bar update function _ssesolve_output_func_progress(sol, i) next!(sol.prob.p.progr) - return _ssesolve_output_func(sol, i) + return _stochastic_output_func(sol, i) end # Output function with distributed channel update for progress bar function _ssesolve_output_func_distributed(sol, i) put!(sol.prob.p.progr_channel, true) - return _ssesolve_output_func(sol, i) + return _stochastic_output_func(sol, i) end -_ssesolve_dispatch_output_func() = _ssesolve_output_func _ssesolve_dispatch_output_func(::ET) where {ET<:Union{EnsembleSerial,EnsembleThreads}} = _ssesolve_output_func_progress _ssesolve_dispatch_output_func(::EnsembleDistributed) = _ssesolve_output_func_distributed @@ -121,7 +118,7 @@ and e_n = \langle \hat{C}_n + \hat{C}_n^\dagger \rangle. ``` -Above, `\hat{C}_n` is the `n`-th collapse operator and `dW_j(t)` is the real Wiener increment associated to `\hat{C}_n`. See [Wiseman2009Quantum](@cite) for more details. +Above, ``\hat{C}_n`` is the `n`-th collapse operator and ``dW_n(t)`` is the real Wiener increment associated to ``\hat{C}_n``. See [Wiseman2009Quantum](@cite) for more details. # Arguments @@ -240,7 +237,7 @@ and e_n = \langle \hat{C}_n + \hat{C}_n^\dagger \rangle. ``` -Above, `\hat{C}_n` is the `n`-th collapse operator and `dW_j(t)` is the real Wiener increment associated to `\hat{C}_n`. See [Wiseman2009Quantum](@cite) for more details. +Above, ``\hat{C}_n`` is the `n`-th collapse operator and ``dW_n(t)`` is the real Wiener increment associated to ``\hat{C}_n``. See [Wiseman2009Quantum](@cite) for more details. # Arguments @@ -363,7 +360,7 @@ and e_n = \langle \hat{C}_n + \hat{C}_n^\dagger \rangle. ``` -Above, `\hat{C}_n` is the `n`-th collapse operator and `dW_j(t)` is the real Wiener increment associated to `\hat{C}_n`. See [Wiseman2009Quantum](@cite) for more details. +Above, ``\hat{C}_n`` is the `n`-th collapse operator and ``dW_n(t)`` is the real Wiener increment associated to ``\hat{C}_n``. See [Wiseman2009Quantum](@cite) for more details. # Arguments @@ -393,7 +390,7 @@ Above, `\hat{C}_n` is the `n`-th collapse operator and `dW_j(t)` is the real Wie # Returns -- `sol::TimeEvolutionSSESol`: The solution of the time evolution. See also [`TimeEvolutionSSESol`](@ref). +- `sol::TimeEvolutionStochasticSol`: The solution of the time evolution. See [`TimeEvolutionStochasticSol`](@ref). """ function ssesolve( H::Union{AbstractQuantumObject{OperatorQuantumObject},Tuple}, @@ -458,7 +455,7 @@ function ssesolve( _se_me_sse_get_expvals(_sol_1) isa Nothing ? nothing : dropdims(sum(expvals_all, dims = 3), dims = 3) ./ length(sol) - return TimeEvolutionSSESol( + return TimeEvolutionStochasticSol( ntraj, _sol_1.prob.p.times, states, diff --git a/src/time_evolution/time_evolution.jl b/src/time_evolution/time_evolution.jl index 39d5d1641..722463101 100644 --- a/src/time_evolution/time_evolution.jl +++ b/src/time_evolution/time_evolution.jl @@ -1,4 +1,4 @@ -export TimeEvolutionSol, TimeEvolutionMCSol, TimeEvolutionSSESol +export TimeEvolutionSol, TimeEvolutionMCSol, TimeEvolutionStochasticSol export liouvillian_floquet, liouvillian_generalized @@ -149,9 +149,9 @@ function Base.show(io::IO, sol::TimeEvolutionMCSol) end @doc raw""" - struct TimeEvolutionSSESol + struct TimeEvolutionStochasticSol -A structure storing the results and some information from solving trajectories of the Stochastic Shrodinger equation time evolution. +A structure storing the results and some information from solving trajectories of the Stochastic time evolution. # Fields (Attributes) @@ -165,7 +165,7 @@ A structure storing the results and some information from solving trajectories o - `abstol::Real`: The absolute tolerance which is used during the solving process. - `reltol::Real`: The relative tolerance which is used during the solving process. """ -struct TimeEvolutionSSESol{ +struct TimeEvolutionStochasticSol{ TT<:AbstractVector{<:Real}, TS<:AbstractVector, TE<:Union{AbstractMatrix,Nothing}, @@ -185,8 +185,8 @@ struct TimeEvolutionSSESol{ reltol::RT end -function Base.show(io::IO, sol::TimeEvolutionSSESol) - print(io, "Solution of quantum trajectories\n") +function Base.show(io::IO, sol::TimeEvolutionStochasticSol) + print(io, "Solution of stochastic quantum trajectories\n") print(io, "(converged: $(sol.converged))\n") print(io, "--------------------------------\n") print(io, "num_trajectories = $(sol.ntraj)\n") @@ -202,6 +202,11 @@ function Base.show(io::IO, sol::TimeEvolutionSSESol) return nothing end +####################################### +#= + Callbacks for Monte Carlo quantum trajectories +=# + abstract type LindbladJumpCallbackType end struct ContinuousLindbladJumpCallback <: LindbladJumpCallbackType @@ -225,6 +230,116 @@ function _check_tlist(tlist, T::Type) return tlist2 end +####################################### +#= +Helpers for handling output of ensemble problems. +This is very useful especially for dispatching which method to use to update the progress bar. +=# + +# Output function with progress bar update +function _ensemble_output_func_progress(sol, i, progr, output_func) + next!(progr) + return output_func(sol, i) +end + +# Output function with distributed channel update for progress bar +function _ensemble_output_func_distributed(sol, i, channel, output_func) + put!(channel, true) + return output_func(sol, i) +end + +function _ensemble_dispatch_output_func( + ::ET, + progress_bar, + ntraj, + output_func, +) where {ET<:Union{EnsembleSerial,EnsembleThreads}} + if getVal(progress_bar) + progr = ProgressBar(ntraj, enable = getVal(progress_bar)) + f = (sol, i) -> _ensemble_output_func_progress(sol, i, progr, output_func) + return (f, progr, nothing) + else + return (output_func, nothing, nothing) + end +end +function _ensemble_dispatch_output_func( + ::ET, + progress_bar, + ntraj, + output_func, +) where {ET<:Union{EnsembleSplitThreads,EnsembleDistributed}} + if getVal(progress_bar) + progr = ProgressBar(ntraj, enable = getVal(progress_bar)) + progr_channel::RemoteChannel{Channel{Bool}} = RemoteChannel(() -> Channel{Bool}(1)) + + f = (sol, i) -> _ensemble_output_func_distributed(sol, i, progr_channel, output_func) + return (f, progr, progr_channel) + else + return (output_func, nothing, nothing) + end +end + +function _ensemble_dispatch_prob_func(rng, ntraj, tlist, prob_func) + seeds = map(i -> rand(rng, UInt64), 1:ntraj) + return (prob, i, repeat) -> prob_func(prob, i, repeat, rng, seeds, tlist) +end + +function _ensemble_dispatch_solve( + ens_prob_mc::TimeEvolutionProblem, + alg::Union{<:OrdinaryDiffEqAlgorithm,<:StochasticDiffEqAlgorithm}, + ensemble_method::ET, + ntraj::Int, +) where {ET<:Union{EnsembleSplitThreads,EnsembleDistributed}} + sol = nothing + + @sync begin + @async while take!(ens_prob_mc.kwargs.channel) + next!(ens_prob_mc.kwargs.progr) + end + + @async begin + sol = solve(ens_prob_mc.prob, alg, ensemble_method, trajectories = ntraj) + put!(ens_prob_mc.kwargs.channel, false) + end + end + + return sol +end +function _ensemble_dispatch_solve( + ens_prob_mc::TimeEvolutionProblem, + alg::Union{<:OrdinaryDiffEqAlgorithm,<:StochasticDiffEqAlgorithm}, + ensemble_method, + ntraj::Int, +) + sol = solve(ens_prob_mc.prob, alg, ensemble_method, trajectories = ntraj) + return sol +end + +####################################### +#= + Stochastic funcs +=# +function _stochastic_prob_func(prob, i, repeat, rng, seeds, tlist) + params = prob.prob.p + + seed = seeds[i] + traj_rng = typeof(rng)() + seed!(traj_rng, seed) + + noise = RealWienerProcess!( + prob.prob.tspan[1], + zeros(params.n_sc_ops), + zeros(params.n_sc_ops), + save_everystep = false, + rng = traj_rng, + ) + + return remake(prob.prob, noise = noise, seed = seed) +end + +# Standard output function +_stochastic_output_func(sol, i) = (sol, false) + ####################################### function liouvillian_floquet( diff --git a/test/core-test/time_evolution.jl b/test/core-test/time_evolution.jl index a1b852e6a..e09690365 100644 --- a/test/core-test/time_evolution.jl +++ b/test/core-test/time_evolution.jl @@ -18,6 +18,10 @@ e_ops = [a' * a, σz] c_ops = [sqrt(γ * (1 + nth)) * a, sqrt(γ * nth) * a', sqrt(γ * (1 + nth)) * σm, sqrt(γ * nth) * σm'] + sme_η = 0.7 # Efficiency of the homodyne detector for smesolve + c_ops_sme = [sqrt(1 - sme_η) * op for op in c_ops] + sc_ops_sme = [sqrt(sme_η) * op for op in c_ops] + ψ0_int = Qobj(round.(Int, real.(ψ0.data)), dims = ψ0.dims) # Used for testing the type inference @testset "sesolve" begin @@ -93,7 +97,7 @@ end end - @testset "mesolve, mcsolve, and ssesolve" begin + @testset "mesolve, mcsolve, ssesolve and smesolve" begin tlist = range(0, 10 / γ, 100) prob_me = mesolveProblem(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false)) @@ -124,6 +128,7 @@ 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)) ρ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) @@ -137,13 +142,15 @@ sol_mc_string = sprint((t, s) -> show(t, "text/plain", s), sol_mc) sol_mc_string_states = sprint((t, s) -> show(t, "text/plain", s), sol_mc_states) sol_sse_string = sprint((t, s) -> show(t, "text/plain", s), sol_sse) + sol_sme_string = sprint((t, s) -> show(t, "text/plain", s), sol_sme) @test prob_me.prob.f.f isa MatrixOperator @test prob_mc.prob.f.f isa MatrixOperator - @test sum(abs.(sol_mc.expect .- sol_me.expect)) / length(tlist) < 0.1 - @test sum(abs.(sol_mc2.expect .- sol_me.expect)) / length(tlist) < 0.1 - @test sum(abs.(vec(expect_mc_states_mean) .- vec(sol_me.expect[1, :]))) / length(tlist) < 0.1 - @test sum(abs.(vec(expect_mc_states_mean2) .- vec(sol_me.expect[1, :]))) / length(tlist) < 0.1 - @test sum(abs.(sol_sse.expect .- sol_me.expect)) / length(tlist) < 0.1 + @test sum(abs, sol_mc.expect .- sol_me.expect) / length(tlist) < 0.1 + @test sum(abs, sol_mc2.expect .- sol_me.expect) / length(tlist) < 0.1 + @test sum(abs, vec(expect_mc_states_mean) .- vec(sol_me.expect[1, :])) / length(tlist) < 0.1 + @test sum(abs, vec(expect_mc_states_mean2) .- vec(sol_me.expect[1, :])) / length(tlist) < 0.1 + @test sum(abs, sol_sse.expect .- sol_me.expect) / length(tlist) < 0.1 + @test sum(abs, sol_sme.expect .- sol_me.expect) / length(tlist) < 0.1 @test length(sol_me.times) == length(tlist) @test length(sol_me.states) == 1 @test size(sol_me.expect) == (length(e_ops), length(tlist)) @@ -159,6 +166,8 @@ @test sol_mc_states.expect === nothing @test length(sol_sse.times) == length(tlist) @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 sol_me_string == "Solution of time evolution\n" * "(return code: $(sol_me.retcode))\n" * @@ -189,7 +198,7 @@ "abstol = $(sol_mc_states.abstol)\n" * "reltol = $(sol_mc_states.reltol)\n" @test sol_sse_string == - "Solution of quantum trajectories\n" * + "Solution of stochastic quantum trajectories\n" * "(converged: $(sol_sse.converged))\n" * "--------------------------------\n" * "num_trajectories = $(sol_sse.ntraj)\n" * @@ -198,6 +207,16 @@ "SDE alg.: $(sol_sse.alg)\n" * "abstol = $(sol_sse.abstol)\n" * "reltol = $(sol_sse.reltol)\n" + @test sol_sme_string == + "Solution of stochastic quantum trajectories\n" * + "(converged: $(sol_sme.converged))\n" * + "--------------------------------\n" * + "num_trajectories = $(sol_sme.ntraj)\n" * + "num_states = $(length(sol_sme.states[1]))\n" * + "num_expect = $(size(sol_sme.expect, 1))\n" * + "SDE alg.: $(sol_sme.alg)\n" * + "abstol = $(sol_sme.abstol)\n" * + "reltol = $(sol_sme.reltol)\n" tlist1 = Float64[] tlist2 = [0, 0.2, 0.1] @@ -211,6 +230,9 @@ @test_throws ArgumentError ssesolve(H, ψ0, tlist1, c_ops, progress_bar = Val(false)) @test_throws ArgumentError ssesolve(H, ψ0, tlist2, c_ops, progress_bar = Val(false)) @test_throws ArgumentError ssesolve(H, ψ0, tlist3, c_ops, progress_bar = Val(false)) + @test_throws ArgumentError smesolve(H, ψ0, tlist1, c_ops_sme, sc_ops_sme, progress_bar = Val(false)) + @test_throws ArgumentError smesolve(H, ψ0, tlist2, c_ops_sme, sc_ops_sme, progress_bar = Val(false)) + @test_throws ArgumentError smesolve(H, ψ0, tlist3, c_ops_sme, sc_ops_sme, progress_bar = Val(false)) # Time-Dependent Hamiltonian # ssesolve is slow to be run on CI. It is not removed from the test because it may be useful for testing in more powerful machines. @@ -364,6 +386,52 @@ @test allocs_tot < 570000 # TODO: Fix this high number of allocations end + @testset "Memory Allocations (smesolve)" begin + allocs_tot = @allocations smesolve( + H, + ψ0, + tlist, + c_ops_sme, + sc_ops_sme, + e_ops = e_ops, + ntraj = 100, + progress_bar = Val(false), + ) # Warm-up + allocs_tot = @allocations smesolve( + H, + ψ0, + tlist, + c_ops_sme, + sc_ops_sme, + e_ops = e_ops, + ntraj = 100, + progress_bar = Val(false), + ) + @test allocs_tot < 2710000 # TODO: Fix this high number of allocations + + allocs_tot = @allocations smesolve( + H, + ψ0, + tlist, + c_ops_sme, + sc_ops_sme, + ntraj = 100, + saveat = [tlist[end]], + progress_bar = Val(false), + ) # Warm-up + allocs_tot = @allocations smesolve( + H, + ψ0, + tlist, + c_ops_sme, + sc_ops_sme, + ntraj = 100, + saveat = [tlist[end]], + progress_bar = Val(false), + ) + @test allocs_tot < 570000 # TODO: Fix this high number of allocations + end + @testset "Type Inference mesolve" begin coef(p, t) = exp(-t) ad_t = QobjEvo(a', coef) @@ -465,17 +533,125 @@ ) end - @testset "mcsolve and ssesolve reproducibility" begin + @testset "Type Inference smesolve" begin + c_ops_sme_tuple = Tuple(c_ops_sme) # To avoid type instability, we must have a Tuple instead of a Vector + sc_ops_sme_tuple = Tuple(sc_ops_sme) # To avoid type instability, we must have a Tuple instead of a Vector + @inferred smesolveEnsembleProblem( + H, + ψ0, + tlist, + c_ops_sme_tuple, + sc_ops_sme_tuple, + ntraj = 5, + e_ops = e_ops, + progress_bar = Val(false), + rng = rng, + ) + @inferred smesolve( + H, + ψ0, + tlist, + c_ops_sme_tuple, + sc_ops_sme_tuple, + ntraj = 5, + e_ops = e_ops, + progress_bar = Val(false), + rng = rng, + ) + @inferred smesolve( + H, + ψ0, + tlist, + c_ops_sme_tuple, + sc_ops_sme_tuple, + ntraj = 5, + progress_bar = Val(true), + rng = rng, + ) + @inferred smesolve( + H, + ψ0, + [0, 10], + c_ops_sme_tuple, + sc_ops_sme_tuple, + ntraj = 5, + progress_bar = Val(false), + rng = rng, + ) + @inferred smesolve( + H, + ψ0_int, + tlist, + c_ops_sme_tuple, + sc_ops_sme_tuple, + ntraj = 5, + progress_bar = Val(false), + rng = rng, + ) + @inferred smesolve( + H, + ψ0, + tlist, + c_ops_sme_tuple, + sc_ops_sme_tuple, + ntraj = 5, + e_ops = (a' * a, a'), + progress_bar = Val(false), + rng = rng, + ) # We test the type inference for Tuple of different types + end + + @testset "mcsolve, ssesolve and smesolve reproducibility" begin rng = MersenneTwister(1234) sol_mc1 = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng) + rng = MersenneTwister(1234) sol_sse1 = ssesolve(H, ψ0, tlist, c_ops, ntraj = 50, e_ops = e_ops, progress_bar = Val(false), rng = rng) + rng = MersenneTwister(1234) + sol_sme1 = smesolve( + H, + ψ0, + tlist, + c_ops_sme, + sc_ops_sme, + ntraj = 50, + e_ops = e_ops, + progress_bar = Val(false), + rng = rng, + ) rng = MersenneTwister(1234) sol_mc2 = mcsolve(H, ψ0, tlist, c_ops, ntraj = 500, e_ops = e_ops, progress_bar = Val(false), rng = rng) + rng = MersenneTwister(1234) sol_sse2 = ssesolve(H, ψ0, tlist, c_ops, ntraj = 50, e_ops = e_ops, progress_bar = Val(false), rng = rng) + rng = MersenneTwister(1234) + sol_sme2 = smesolve( + H, + ψ0, + tlist, + c_ops_sme, + sc_ops_sme, + ntraj = 50, + e_ops = e_ops, + progress_bar = Val(false), + rng = rng, + ) rng = MersenneTwister(1234) sol_mc3 = mcsolve(H, ψ0, tlist, c_ops, ntraj = 510, e_ops = e_ops, progress_bar = Val(false), rng = rng) + rng = MersenneTwister(1234) + sol_sse3 = ssesolve(H, ψ0, tlist, c_ops, ntraj = 60, e_ops = e_ops, progress_bar = Val(false), rng = rng) + rng = MersenneTwister(1234) + sol_sme3 = smesolve( + H, + ψ0, + tlist, + c_ops_sme, + sc_ops_sme, + ntraj = 60, + e_ops = e_ops, + progress_bar = Val(false), + rng = rng, + ) @test sol_mc1.expect ≈ sol_mc2.expect atol = 1e-10 @test sol_mc1.expect_all ≈ sol_mc2.expect_all atol = 1e-10 @@ -486,6 +662,13 @@ @test sol_sse1.expect ≈ sol_sse2.expect atol = 1e-10 @test sol_sse1.expect_all ≈ sol_sse2.expect_all atol = 1e-10 + + @test sol_sse1.expect_all ≈ sol_sse3.expect_all[:, :, 1:50] atol = 1e-10 + + @test sol_sme1.expect ≈ sol_sme2.expect atol = 1e-10 + @test sol_sme1.expect_all ≈ sol_sme2.expect_all atol = 1e-10 + + @test sol_sme1.expect_all ≈ sol_sme3.expect_all[:, :, 1:50] atol = 1e-10 end end