Add smesolve solver (#389)

Co-authored-by: Yi-Te Huang <[email protected]>

.typos.toml

@@ -1,2 +1,3 @@
 [default.extend-words]
-ket = "ket"
+ket = "ket"
+sme = "sme"

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

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

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)`.

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

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

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,46 +390,14 @@ 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(),
     ntraj::Int = 1,
     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]

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()`.