Updated ssesolve, smesolve and mcsolve to be minmial working changes.

Author: Gavin-Rockwood

src/time_evolution/mcsolve.jl

@@ -20,9 +20,9 @@ function _mcsolve_output_func(sol, i)
     return (sol, false)
 end
 
-function _normalize_state!(u, dims, normalize_states, type)
+function _normalize_state!(u, dims, normalize_states)
     getVal(normalize_states) && normalize!(u)
-    return QuantumObject(u, type, dims)
+    return QuantumObject(u, Ket(), dims)
 end
 
 function _mcsolve_make_Heff_QobjEvo(H::QuantumObject, c_ops)
@@ -110,15 +110,15 @@ If the environmental measurements register a quantum jump, the wave function und
 """
 function mcsolveProblem(
     H::Union{AbstractQuantumObject{Operator},Tuple},
-    ψ0::QuantumObject{ST},
+    ψ0::QuantumObject{Ket},
     tlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
     params = NullParameters(),
     rng::AbstractRNG = default_rng(),
     jump_callback::TJC = ContinuousLindbladJumpCallback(),
     kwargs...,
-) where {ST<:Union{Ket,Operator}, TJC<:LindbladJumpCallbackType}
+) where {TJC<:LindbladJumpCallbackType}
     haskey(kwargs, :save_idxs) &&
         throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
 
@@ -221,7 +221,7 @@ If the environmental measurements register a quantum jump, the wave function und
 """
 function mcsolveEnsembleProblem(
     H::Union{AbstractQuantumObject{Operator},Tuple},
-    ψ0::QuantumObject{ST},
+    ψ0::QuantumObject{Ket},
     tlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
@@ -234,7 +234,7 @@ function mcsolveEnsembleProblem(
     prob_func::Union{Function,Nothing} = nothing,
     output_func::Union{Tuple,Nothing} = nothing,
     kwargs...,
-) where {ST<:Union{Ket,Operator}, TJC<:LindbladJumpCallbackType}
+) where {TJC<:LindbladJumpCallbackType}
     _prob_func = isnothing(prob_func) ? _ensemble_dispatch_prob_func(rng, ntraj, tlist, _mcsolve_prob_func) : prob_func
     _output_func =
         output_func isa Nothing ?
@@ -359,7 +359,7 @@ If the environmental measurements register a quantum jump, the wave function und
 """
 function mcsolve(
     H::Union{AbstractQuantumObject{Operator},Tuple},
-    ψ0::QuantumObject{ST},
+    ψ0::QuantumObject{Ket},
     tlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     alg::AbstractODEAlgorithm = DP5(),
@@ -375,7 +375,7 @@ function mcsolve(
     keep_runs_results::Union{Val,Bool} = Val(false),
     normalize_states::Union{Val,Bool} = Val(true),
     kwargs...,
-) where {ST<:Union{Ket,Operator}, TJC<:LindbladJumpCallbackType}
+) where {TJC<:LindbladJumpCallbackType}
     ens_prob_mc = mcsolveEnsembleProblem(
         H,
         ψ0,
@@ -415,11 +415,8 @@ function mcsolve(
         _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_all = stack(
-        map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states, [ens_prob_mc.states_type]), eachindex(sol)), # Unsure why ens_prob_mc.states_type needs to be in an array but the other two arguments don't!
-        dims = 1,
-    )
+    # stack to transform Vector{Vector{QuantumObject}} -> Matrix{QuantumObject}
+    states_all = stack(map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol)), dims = 1)
 
     col_times = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.col_times, eachindex(sol))
     col_which = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.col_which, eachindex(sol))
@@ -439,4 +436,4 @@ function mcsolve(
         kwargs.abstol,
         kwargs.reltol,
     )
-end
+end

src/time_evolution/mesolve.jl

@@ -7,7 +7,7 @@ _mesolve_make_L_QobjEvo(H::Nothing, c_ops::Nothing) = throw(ArgumentError("Both
 c_ops are Nothing. You are probably running the wrong function."))
 
 function _gen_mesolve_solution(sol, prob::TimeEvolutionProblem{ST}) where {ST<:Union{Operator,OperatorKet,SuperOperator}}
-    if prob.states_type == Operator
+    if prob.states_type isa Operator
         ρt = map(ϕ -> QuantumObject(vec2mat(ϕ), type = prob.states_type, dims = prob.dimensions), sol.u)
     else
         ρt = map(ϕ -> QuantumObject(ϕ, type = prob.states_type, dims = prob.dimensions), sol.u)
@@ -66,6 +66,7 @@ where
 
 # Notes
 
+- The initial state can also be [`SuperOperator`](@ref) (such as a super-identity). This is useful for simulating many density matrices simultaneously or calculating process matrices. Currently must be Square. 
 - 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.
 - If `H` is an [`Operator`](@ref), `ψ0` is a [`Ket`](@ref) and `c_ops` is `Nothing`, the function will call [`sesolveProblem`](@ref) instead.
@@ -106,24 +107,15 @@ function mesolveProblem(
     L_evo = _mesolve_make_L_QobjEvo(H, c_ops)
     check_dimensions(L_evo, ψ0)
 
-    T = Base.promote_eltype(L_evo, ψ0)
-    # ρ0 = if isoperket(ψ0) # Convert it to dense vector with complex element type
-    #     to_dense(_complex_float_type(T), copy(ψ0.data))
-    # else
-    #     to_dense(_complex_float_type(T), mat2vec(ket2dm(ψ0).data))
-    # end
-    if isoper(ψ0)
-        ρ0 = to_dense(_complex_float_type(T), mat2vec(ψ0.data))
-        states_type = Operator()
-    elseif isoperket(ψ0)
-        ρ0 = to_dense(_complex_float_type(T), copy(ψ0.data))
-        states_type = OperatorKet()
-    elseif isket(ψ0)
-        ρ0 = to_dense(_complex_float_type(T), mat2vec(ket2dm(ψ0).data))
+    # Convert to dense vector with complex element type
+
+    T = _complex_float_type(Base.promote_eltype(L_evo, ψ0))
+    if isoperket(ψ0) || issuper(ψ0)
+        ρ0 = to_dense(T, copy(ψ0.data))
+        states_type = ψ0.type
+    else
+        ρ0 = to_dense(T, mat2vec(ket2dm(ψ0).data))
         states_type = Operator()
-    elseif issuper(ψ0)
-        ρ0 = to_dense(_complex_float_type(T), copy(ψ0.data))
-        states_type = SuperOperator()
     end
 
     L = cache_operator(L_evo.data, ρ0)
@@ -180,6 +172,7 @@ where
 
 # Notes
 
+- The initial state can also be [`SuperOperator`](@ref) (such as a super-identity). This is useful for simulating many density matrices simultaneously or calculating process matrices. Currently must be Square. 
 - 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.
 - If `H` is an [`Operator`](@ref), `ψ0` is a [`Ket`](@ref) and `c_ops` is `Nothing`, the function will call [`sesolve`](@ref) instead.
@@ -292,6 +285,7 @@ for each combination in the ensemble.
 
 # Notes
 
+- The initial state can also be [`SuperOperator`](@ref) (such as a super-identity). This is useful for simulating many density matrices simultaneously or calculating process matrices. Currently must be Square. 
 - The function returns an array of solutions with dimensions matching the Cartesian product of initial states and parameter sets.
 - If `ψ0` is a vector of `m` states and `params = (p1, p2, ...)` where `p1` has length `n1`, `p2` has length `n2`, etc., the output will be of size `(m, n1, n2, ...)`.
 - If `H` is an [`Operator`](@ref), `ψ0` is a [`Ket`](@ref) and `c_ops` is `Nothing`, the function will call [`sesolve_map`](@ref) instead.
@@ -329,14 +323,10 @@ function mesolve_map(
     # Convert to appropriate format based on state type
     ψ0_iter = map(ψ0) do state
         T = _complex_float_type(eltype(state))
-        if isoper(state)
-            to_dense(_complex_float_type(T), mat2vec(state.data))
-        elseif isoperket(state)
-            to_dense(_complex_float_type(T), copy(state.data))
-        elseif isket(state)
-            to_dense(_complex_float_type(T), mat2vec(ket2dm(state).data))
-        elseif issuper(state)
-            to_dense(_complex_float_type(T), copy(state.data))
+        if isoperket(state) || issuper(state)
+            to_dense(T, copy(state.data))
+        else
+            to_dense(T, mat2vec(ket2dm(state).data))
         end
     end
     if params isa NullParameters

src/time_evolution/sesolve.jl

@@ -50,6 +50,7 @@ Generate the ODEProblem for the Schrödinger time evolution of a quantum system:
 
 # Notes
 
+- Initial state can also be [`Operator`](@ref)s where each column represents a state vector, such as the Identity operator. This can be used, for example, to calculate the propagator.
 - 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-6` and `abstol=1e-8`.
@@ -127,6 +128,7 @@ Time evolution of a closed quantum system using the Schrödinger equation:
 
 # Notes
 
+- Initial state can also be [`Operator`](@ref)s where each column represents a state vector, such as the Identity operator. This can be used, for example, to calculate the propagator.
 - 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-6` and `abstol=1e-8`.
@@ -215,9 +217,10 @@ for each combination in the ensemble.
 - `kwargs`: The keyword arguments for the ODEProblem.
 
 # Notes
-
+- Initial state can also be [`Operator`](@ref)s where each column represents a state vector, such as the Identity operator. This can be used, for example, to calculate the propagator.
 - The function returns an array of solutions with dimensions matching the Cartesian product of initial states and parameter sets.
 - If `ψ0` is a vector of `m` states and `params = (p1, p2, ...)` where `p1` has length `n1`, `p2` has length `n2`, etc., the output will be of size `(m, n1, n2, ...)`.
+- Similarly, the initial state(s) can also be `Operator`s where each column represents a state vector, such as the Identity operator. This can be used, for example, to calculate many propagators.
 - See [`sesolve`](@ref) for more details.
 
 # Returns
@@ -239,7 +242,7 @@ function sesolve_map(
 
     ψ0 = map(to_dense, ψ0) # Convert all initial states to dense vectors
 
-    ψ0_iter = map(get_data, ψ0)
+      ψ0_iter = map(state -> to_dense(_complex_float_type(eltype(state)), copy(state.data)), ψ0) 
     if params isa NullParameters
         iter = collect(Iterators.product(ψ0_iter, [params])) |> vec # convert nx1 Matrix into Vector
     else

src/time_evolution/smesolve.jl

@@ -1,14 +1,7 @@
 export smesolveProblem, smesolveEnsembleProblem, smesolve
 
-#_smesolve_generate_state(u, dims, isoperket::Val{false}) = QuantumObject(vec2mat(u), type = Operator(), dims = dims)
-#_smesolve_generate_state(u, dims, isoperket::Val{true}) = QuantumObject(u, type = OperatorKet(), dims = dims)
-function _smesolve_generate_state(u, dims, type)
-    if type == OperatorKet
-        return QuantumObject(u, type = type, dims = dims)
-    else
-        return QuantumObject(vec2mat(u), type = Operator(), dims = dims)
-    end
-end
+_smesolve_generate_state(u, dims, isoperket::Val{false}) = QuantumObject(vec2mat(u), type = Operator(), dims = dims)
+_smesolve_generate_state(u, dims, isoperket::Val{true}) = QuantumObject(u, type = OperatorKet(), dims = dims)
 
 function _smesolve_update_coeff(u, p, t, op_vec)
     return 2 * real(dot(op_vec, u)) #this is Tr[Sn * ρ + ρ * Sn']
@@ -110,10 +103,10 @@ function smesolveProblem(
     T = Base.promote_eltype(L_evo, ψ0)
     if isoperket(ψ0) # Convert it to dense vector with complex element type
         ρ0 = to_dense(_complex_float_type(T), copy(ψ0.data))
-        states_type = OperatorKet()
+        state_type = OperatorKet()
     else
         ρ0 = to_dense(_complex_float_type(T), mat2vec(ket2dm(ψ0).data))
-        states_type = Operator()
+        state_type = Operator()
     end
 
     sc_ops_evo_data = Tuple(map(get_data ∘ QobjEvo, sc_ops_list))
@@ -155,7 +148,7 @@ function smesolveProblem(
         kwargs4...,
     )
 
-    return TimeEvolutionProblem(prob, tlist, states_type, dims, ())
+    return TimeEvolutionProblem(prob, tlist, state_type, dims, (isoperket = Val(isoperket(ψ0)),))
 end
 
 @doc raw"""
@@ -285,7 +278,7 @@ function smesolveEnsembleProblem(
         prob_sme.times,
         prob_sme.states_type,
         prob_sme.dimensions,
-        (progr = _output_func[2], channel = _output_func[3]),
+        merge(prob_sme.kwargs, (progr = _output_func[2], channel = _output_func[3])),
     )
 
     return ensemble_prob
@@ -432,7 +425,7 @@ function smesolve(
 
     # stack to transform Vector{Vector{QuantumObject}} -> Matrix{QuantumObject}
     states_all = stack(
-        map(i -> _smesolve_generate_state.(sol[:, i].u, Ref(dims), [ens_prob.states_type]), eachindex(sol)),
+        map(i -> _smesolve_generate_state.(sol[:, i].u, Ref(dims), ens_prob.kwargs.isoperket), eachindex(sol)),
         dims = 1,
     )
 
@@ -454,4 +447,4 @@ function smesolve(
         kwargs.abstol,
         kwargs.reltol,
     )
-end
+end

src/time_evolution/ssesolve.jl

@@ -76,7 +76,7 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is th
 """
 function ssesolveProblem(
     H::Union{AbstractQuantumObject{Operator},Tuple},
-    ψ0::QuantumObject{ST},
+    ψ0::QuantumObject{Ket},
     tlist::AbstractVector,
     sc_ops::Union{Nothing,AbstractVector,Tuple,AbstractQuantumObject} = nothing;
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
@@ -85,7 +85,7 @@ function ssesolveProblem(
     progress_bar::Union{Val,Bool} = Val(true),
     store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
-) where {ST<:Union{Ket,Operator}}
+)
     haskey(kwargs, :save_idxs) &&
         throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
 
@@ -95,13 +95,13 @@ function ssesolveProblem(
     sc_ops_isa_Qobj = sc_ops isa AbstractQuantumObject # We can avoid using non-diagonal noise if sc_ops is just an AbstractQuantumObject
 
     tlist = _check_tlist(tlist, _float_type(ψ0))
-    states_type = ψ0.type
 
     H_eff_evo = _mcsolve_make_Heff_QobjEvo(H, sc_ops_list)
     isoper(H_eff_evo) || throw(ArgumentError("The Hamiltonian must be an Operator."))
     check_dimensions(H_eff_evo, ψ0)
     dims = H_eff_evo.dimensions
 
+    states_type = ψ0.type
     ψ0 = to_dense(_complex_float_type(ψ0), get_data(ψ0))
 
     sc_ops_evo_data = Tuple(map(get_data ∘ QobjEvo, sc_ops_list))
@@ -219,7 +219,7 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and  ``dW_n(t)`` is t
 """
 function ssesolveEnsembleProblem(
     H::Union{AbstractQuantumObject{Operator},Tuple},
-    ψ0::QuantumObject{ST},
+    ψ0::QuantumObject{Ket},
     tlist::AbstractVector,
     sc_ops::Union{Nothing,AbstractVector,Tuple,AbstractQuantumObject} = nothing;
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
@@ -232,7 +232,7 @@ function ssesolveEnsembleProblem(
     progress_bar::Union{Val,Bool} = Val(true),
     store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
-) where {ST<:Union{Ket,Operator}}
+)
     _prob_func =
         isnothing(prob_func) ?
         _ensemble_dispatch_prob_func(
@@ -253,7 +253,7 @@ function ssesolveEnsembleProblem(
             progr_desc = "[ssesolve] ",
         ) : output_func
 
-    prob_sse = ssesolveProblem(
+    prob_sme = ssesolveProblem(
         H,
         ψ0,
         tlist,
@@ -267,10 +267,10 @@ function ssesolveEnsembleProblem(
     )
 
     ensemble_prob = TimeEvolutionProblem(
-        EnsembleProblem(prob_sse, prob_func = _prob_func, output_func = _output_func[1], safetycopy = true),
-        prob_sse.times,
-        prob_sse.states_type,
-        prob_sse.dimensions,
+        EnsembleProblem(prob_sme, prob_func = _prob_func, output_func = _output_func[1], safetycopy = true),
+        prob_sme.times,
+        prob_sme.states_type,
+        prob_sme.dimensions,
         (progr = _output_func[2], channel = _output_func[3]),
     )
 
@@ -357,7 +357,7 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is th
 """
 function ssesolve(
     H::Union{AbstractQuantumObject{Operator},Tuple},
-    ψ0::QuantumObject{ST},
+    ψ0::QuantumObject{Ket},
     tlist::AbstractVector,
     sc_ops::Union{Nothing,AbstractVector,Tuple,AbstractQuantumObject} = nothing;
     alg::Union{Nothing,AbstractSDEAlgorithm} = nothing,
@@ -372,7 +372,7 @@ function ssesolve(
     keep_runs_results::Union{Val,Bool} = Val(false),
     store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
-) where {ST<:Union{Ket,Operator}}
+)
     ens_prob = ssesolveEnsembleProblem(
         H,
         ψ0,
@@ -419,7 +419,7 @@ function ssesolve(
     expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all, dims = 2) # Stack on dimension 2 to align with QuTiP
 
     # stack to transform Vector{Vector{QuantumObject}} -> Matrix{QuantumObject}
-    states_all = stack(map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states, [ens_prob.states_type]), eachindex(sol)), dims = 1)
+    states_all = stack(map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol)), dims = 1)
 
     _m_expvals =
         _m_expvals_sol_1 isa Nothing ? nothing : map(i -> _get_m_expvals(sol[:, i], SaveFuncSSESolve), eachindex(sol))
@@ -439,4 +439,4 @@ function ssesolve(
         kwargs.abstol,
         kwargs.reltol,
     )
-end
+end