fixed a bug with smesolve!

Author: Gavin-Rockwood

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 ST == Operator
+    if prob.states_type == 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)

src/time_evolution/smesolve.jl

@@ -1,7 +1,14 @@
 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)
+#_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
 
 function _smesolve_update_coeff(u, p, t, op_vec)
     return 2 * real(dot(op_vec, u)) #this is Tr[Sn * ρ + ρ * Sn']
@@ -146,7 +153,7 @@ function smesolveProblem(
         kwargs4...,
     )
 
-    return TimeEvolutionProblem(prob, tlist,StateOpType(), dims)
+    return TimeEvolutionProblem(prob, tlist,StateOpType(), dims, ())
 end
 
 @doc raw"""
@@ -274,8 +281,9 @@ 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.dimensions,
-        merge(prob_sme.kwargs, (progr = _output_func[2], channel = _output_func[3])),
+        (progr = _output_func[2], channel = _output_func[3]),
     )
 
     return ensemble_prob
@@ -422,7 +430,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.kwargs.isoperket), eachindex(sol)),
+        map(i -> _smesolve_generate_state.(sol[:, i].u, Ref(dims), [ens_prob.states_type]), eachindex(sol)),
         dims = 1,
     )