Updated ssesolve, smesolve and mcsolve to be minimal changes along with a few other changes.

Author: Gavin-Rockwood

src/time_evolution/mcsolve.jl

@@ -118,7 +118,7 @@ function mcsolveProblem(
     rng::AbstractRNG = default_rng(),
     jump_callback::TJC = ContinuousLindbladJumpCallback(),
     kwargs...,
-) where {TJC<:LindbladJumpCallbackType,ST<:Union{Ket,Operator}}
+) where {ST<:Union{Ket,Operator}, TJC<:LindbladJumpCallbackType}
     haskey(kwargs, :save_idxs) &&
         throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
 
@@ -234,7 +234,7 @@ function mcsolveEnsembleProblem(
     prob_func::Union{Function,Nothing} = nothing,
     output_func::Union{Tuple,Nothing} = nothing,
     kwargs...,
-) where {TJC<:LindbladJumpCallbackType,ST<:Union{Ket,Operator}}
+) where {ST<:Union{Ket,Operator}, 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 ?
@@ -261,7 +261,7 @@ function mcsolveEnsembleProblem(
     ensemble_prob = TimeEvolutionProblem(
         EnsembleProblem(prob_mc.prob, prob_func = _prob_func, output_func = _output_func[1], safetycopy = false),
         prob_mc.times,
-        ST(),
+        prob_mc.states_type,
         prob_mc.dimensions,
         (progr = _output_func[2], channel = _output_func[3]),
     )
@@ -375,7 +375,7 @@ function mcsolve(
     keep_runs_results::Union{Val,Bool} = Val(false),
     normalize_states::Union{Val,Bool} = Val(true),
     kwargs...,
-) where {TJC<:LindbladJumpCallbackType} where {ST<:Union{Ket,Operator}}
+) where {ST<:Union{Ket,Operator}, TJC<:LindbladJumpCallbackType}
     ens_prob_mc = mcsolveEnsembleProblem(
         H,
         ψ0,

src/time_evolution/mesolve.jl

@@ -114,16 +114,16 @@ function mesolveProblem(
     # end
     if isoper(ψ0)
         ρ0 = to_dense(_complex_float_type(T), mat2vec(ψ0.data))
-        state_type = Operator()
+        states_type = Operator()
     elseif isoperket(ψ0)
         ρ0 = to_dense(_complex_float_type(T), copy(ψ0.data))
-        state_type = OperatorKet()
+        states_type = OperatorKet()
     elseif isket(ψ0)
         ρ0 = to_dense(_complex_float_type(T), mat2vec(ket2dm(ψ0).data))
-        state_type = Operator()
+        states_type = Operator()
     elseif issuper(ψ0)
         ρ0 = to_dense(_complex_float_type(T), copy(ψ0.data))
-        state_type = SuperOperator()
+        states_type = SuperOperator()
     end
 
     L = cache_operator(L_evo.data, ρ0)
@@ -136,7 +136,7 @@ function mesolveProblem(
 
     prob = ODEProblem{getVal(inplace),FullSpecialize}(L, ρ0, tspan, params; kwargs4...)
 
-    return TimeEvolutionProblem(prob, tlist, state_type, L_evo.dimensions)
+    return TimeEvolutionProblem(prob, tlist, states_type, L_evo.dimensions)
 end
 
 @doc raw"""
@@ -398,7 +398,7 @@ function mesolve_map(
     ens_prob = TimeEvolutionProblem(
         EnsembleProblem(prob.prob, prob_func = _prob_func, output_func = _output_func[1], safetycopy = false),
         prob.times,
-        StateOpType(),
+        prob.states_type,
         prob.dimensions,
         (progr = _output_func[2], channel = _output_func[3]),
     )

src/time_evolution/sesolve.jl

@@ -73,6 +73,7 @@ function sesolveProblem(
         throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
 
     tlist = _check_tlist(tlist, _float_type(ψ0))
+    states_type = ψ0.type
 
     H_evo = _sesolve_make_U_QobjEvo(H) # Multiply by -i
     isoper(H_evo) || throw(ArgumentError("The Hamiltonian must be an Operator."))
@@ -90,7 +91,7 @@ function sesolveProblem(
 
     prob = ODEProblem{getVal(inplace),FullSpecialize}(U, ψ0, tspan, params; kwargs4...)
 
-    return TimeEvolutionProblem(prob, tlist, ST(), H_evo.dimensions)
+    return TimeEvolutionProblem(prob, tlist, states_type, H_evo.dimensions)
 end
 
 @doc raw"""

src/time_evolution/smesolve.jl

@@ -108,10 +108,12 @@ function smesolveProblem(
     dims = L_evo.dimensions
 
     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))
+    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()
     else
-        to_dense(_complex_float_type(T), mat2vec(ket2dm(ψ0).data))
+        ρ0 = to_dense(_complex_float_type(T), mat2vec(ket2dm(ψ0).data))
+        states_type = Operator()
     end
 
     sc_ops_evo_data = Tuple(map(get_data ∘ QobjEvo, sc_ops_list))
@@ -153,7 +155,7 @@ function smesolveProblem(
         kwargs4...,
     )
 
-    return TimeEvolutionProblem(prob, tlist,StateOpType(), dims, ())
+    return TimeEvolutionProblem(prob, tlist, states_type, dims, ())
 end
 
 @doc raw"""
@@ -281,7 +283,7 @@ function smesolveEnsembleProblem(
     ensemble_prob = TimeEvolutionProblem(
         EnsembleProblem(prob_sme, prob_func = _prob_func, output_func = _output_func[1], safetycopy = true),
         prob_sme.times,
-        StateOpType(),
+        prob_sme.states_type,
         prob_sme.dimensions,
         (progr = _output_func[2], channel = _output_func[3]),
     )

src/time_evolution/ssesolve.jl

@@ -95,6 +95,7 @@ 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."))
@@ -142,7 +143,7 @@ function ssesolveProblem(
         kwargs4...,
     )
 
-    return TimeEvolutionProblem(prob, tlist, ST(), dims)
+    return TimeEvolutionProblem(prob, tlist, states_type, dims)
 end
 
 @doc raw"""
@@ -252,7 +253,7 @@ function ssesolveEnsembleProblem(
             progr_desc = "[ssesolve] ",
         ) : output_func
 
-    prob_sme = ssesolveProblem(
+    prob_sse = ssesolveProblem(
         H,
         ψ0,
         tlist,
@@ -266,10 +267,10 @@ function ssesolveEnsembleProblem(
     )
 
     ensemble_prob = TimeEvolutionProblem(
-        EnsembleProblem(prob_sme, prob_func = _prob_func, output_func = _output_func[1], safetycopy = true),
-        prob_sme.times,
-        ST(),
-        prob_sme.dimensions,
+        EnsembleProblem(prob_sse, prob_func = _prob_func, output_func = _output_func[1], safetycopy = true),
+        prob_sse.times,
+        prob_sse.states_type,
+        prob_sse.dimensions,
         (progr = _output_func[2], channel = _output_func[3]),
     )