[no ci] Add support for OperatorKet input for mesolve and smesolve

Author: albertomercurio

src/time_evolution/mesolve.jl

@@ -33,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``. It can be either a [`Ket`](@ref) or a [`Operator`](@ref).
+- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. It can be either a [`Ket`](@ref), [`Operator`](@ref) or [`OperatorKet`](@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`.
@@ -65,7 +65,7 @@ function mesolveProblem(
     kwargs...,
 ) where {
     HOpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject},
-    StateOpType<:Union{KetQuantumObject,OperatorQuantumObject},
+    StateOpType<:Union{KetQuantumObject,OperatorQuantumObject,OperatorKetQuantumObject},
 }
     haskey(kwargs, :save_idxs) &&
         throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
@@ -76,7 +76,11 @@ function mesolveProblem(
     check_dimensions(L_evo, ψ0)
 
     T = Base.promote_eltype(L_evo, ψ0)
-    ρ0 = to_dense(_CType(T), mat2vec(ket2dm(ψ0).data)) # Convert it to dense vector with complex element type
+    ρ0 = if isoperket(ψ0) # Convert it to dense vector with complex element type
+        to_dense(_CType(T), copy(ψ0.data))
+    else
+        to_dense(_CType(T), mat2vec(ket2dm(ψ0).data)) 
+    end
     L = L_evo.data
 
     kwargs2 = _merge_saveat(tlist, e_ops, DEFAULT_ODE_SOLVER_OPTIONS; kwargs...)
@@ -85,7 +89,7 @@ function mesolveProblem(
     tspan = (tlist[1], tlist[end])
     prob = ODEProblem{getVal(inplace),FullSpecialize}(L, ρ0, tspan, params; kwargs3...)
 
-    return TimeEvolutionProblem(prob, tlist, L_evo.dimensions)
+    return TimeEvolutionProblem(prob, tlist, L_evo.dimensions, (isoperket=Val(isoperket(ψ0)),))
 end
 
 @doc raw"""
@@ -117,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``. It can be either a [`Ket`](@ref) or a [`Operator`](@ref).
+- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. It can be either a [`Ket`](@ref), [`Operator`](@ref) or [`OperatorKet`](@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()`.
@@ -152,7 +156,7 @@ function mesolve(
     kwargs...,
 ) where {
     HOpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject},
-    StateOpType<:Union{KetQuantumObject,OperatorQuantumObject},
+    StateOpType<:Union{KetQuantumObject,OperatorQuantumObject,OperatorKetQuantumObject},
 }
     prob = mesolveProblem(
         H,
@@ -173,7 +177,12 @@ end
 function mesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit5())
     sol = solve(prob.prob, alg)
 
-    ρt = map(ϕ -> QuantumObject(vec2mat(ϕ), type = Operator, dims = prob.dimensions), sol.u)
+    # No type instabilities since `isoperket` is a Val, and so it is known at compile time
+    if getVal(prob.kwargs.isoperket)
+        ρt = map(ϕ -> QuantumObject(ϕ, type = OperatorKet, dims = prob.dimensions), sol.u)
+    else
+        ρt = map(ϕ -> QuantumObject(vec2mat(ϕ), type = Operator, dims = prob.dimensions), sol.u)
+    end
 
     return TimeEvolutionSol(
         prob.times,

src/time_evolution/smesolve.jl

@@ -1,6 +1,7 @@
 export smesolveProblem, smesolveEnsembleProblem, smesolve
 
-_smesolve_generate_state(u, dims) = QuantumObject(vec2mat(u), type = Operator, 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_update_coeff(u, p, t, op_vec)
     return 2 * real(dot(op_vec, u)) #this is Tr[Sn * ρ + ρ * Sn']
@@ -47,7 +48,7 @@ Above, ``\hat{C}_i`` represent the collapse operators related to pure dissipatio
 # 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).
+- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. It can be either a [`Ket`](@ref), [`Operator`](@ref) or [`OperatorKet`](@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}_i\}_i``. It can be either a `Vector` or a `Tuple`.
 - `sc_ops`: List of stochastic collapse operators ``\{\hat{S}_n\}_n``. It can be either a `Vector`, a `Tuple` or a [`AbstractQuantumObject`](@ref). It is recommended to use the last case when only one operator is provided.
@@ -84,7 +85,7 @@ function smesolveProblem(
     progress_bar::Union{Val,Bool} = Val(true),
     store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
-) where {StateOpType<:Union{KetQuantumObject,OperatorQuantumObject}}
+) where {StateOpType<:Union{KetQuantumObject,OperatorQuantumObject,OperatorKetQuantumObject}}
     haskey(kwargs, :save_idxs) &&
         throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
 
@@ -100,7 +101,11 @@ function smesolveProblem(
     dims = L_evo.dimensions
 
     T = Base.promote_eltype(L_evo, ψ0)
-    ρ0 = to_dense(_CType(T), mat2vec(ket2dm(ψ0).data)) # Convert it to dense vector with complex element type
+    ρ0 = if isoperket(ψ0) # Convert it to dense vector with complex element type
+        to_dense(_CType(T), copy(ψ0.data))
+    else
+        to_dense(_CType(T), mat2vec(ket2dm(ψ0).data)) 
+    end
 
     progr = ProgressBar(length(tlist), enable = getVal(progress_bar))
 
@@ -143,7 +148,7 @@ function smesolveProblem(
         kwargs3...,
     )
 
-    return TimeEvolutionProblem(prob, tlist, dims)
+    return TimeEvolutionProblem(prob, tlist, dims, (isoperket=Val(isoperket(ψ0)),))
 end
 
 @doc raw"""
@@ -188,7 +193,7 @@ Above, ``\hat{C}_i`` represent the collapse operators related to pure dissipatio
 # 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).
+- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. It can be either a [`Ket`](@ref), [`Operator`](@ref) or [`OperatorKet`](@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}_i\}_i``. It can be either a `Vector` or a `Tuple`.
 - `sc_ops`: List of stochastic collapse operators ``\{\hat{S}_n\}_n``. It can be either a `Vector`, a `Tuple` or a [`AbstractQuantumObject`](@ref). It is recommended to use the last case when only one operator is provided.
@@ -233,7 +238,7 @@ function smesolveEnsembleProblem(
     progress_bar::Union{Val,Bool} = Val(true),
     store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
-) where {StateOpType<:Union{KetQuantumObject,OperatorQuantumObject}}
+) where {StateOpType<:Union{KetQuantumObject,OperatorQuantumObject,OperatorKetQuantumObject}}
     _prob_func =
         isnothing(prob_func) ?
         _ensemble_dispatch_prob_func(
@@ -266,7 +271,7 @@ function smesolveEnsembleProblem(
         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]),
+        merge(prob_sme.kwargs, (progr = _output_func[2], channel = _output_func[3])),
     )
 
     return ensemble_prob
@@ -315,7 +320,7 @@ Above, ``\hat{C}_i`` represent the collapse operators related to pure dissipatio
 # 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).
+- `ψ0`: Initial state of the system ``|\psi(0)\rangle``. It can be either a [`Ket`](@ref), [`Operator`](@ref) or [`OperatorKet`](@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}_i\}_i``. It can be either a `Vector` or a `Tuple`.
 - `sc_ops`: List of stochastic collapse operators ``\{\hat{S}_n\}_n``. It can be either a `Vector`, a `Tuple` or a [`AbstractQuantumObject`](@ref). It is recommended to use the last case when only one operator is provided.
@@ -362,7 +367,7 @@ function smesolve(
     progress_bar::Union{Val,Bool} = Val(true),
     store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
-) where {StateOpType<:Union{KetQuantumObject,OperatorQuantumObject}}
+) where {StateOpType<:Union{KetQuantumObject,OperatorQuantumObject,OperatorKetQuantumObject}}
     ensemble_prob = smesolveEnsembleProblem(
         H,
         ψ0,
@@ -406,7 +411,8 @@ function smesolve(
     _expvals_all =
         _expvals_sol_1 isa Nothing ? nothing : map(i -> _get_expvals(sol[:, i], SaveFuncMESolve), eachindex(sol))
     expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all, dims = 2) # Stack on dimension 2 to align with QuTiP
-    states = map(i -> _smesolve_generate_state.(sol[:, i].u, Ref(dims)), eachindex(sol))
+
+    states = map(i -> _smesolve_generate_state.(sol[:, i].u, Ref(dims), ens_prob.kwargs.isoperket), eachindex(sol))
 
     _m_expvals =
         _m_expvals_sol_1 isa Nothing ? nothing : map(i -> _get_m_expvals(sol[:, i], SaveFuncSMESolve), eachindex(sol))

test/core-test/time_evolution.jl

@@ -158,6 +158,11 @@
         )
         sol_sme3 = smesolve(H, ψ0, tlist, c_ops_sme2, sc_ops_sme2, e_ops = e_ops, progress_bar = Val(false))
 
+        # For testing the `OperatorKet` input
+        sol_me4 = mesolve(H, operator_to_vector(ket2dm(ψ0)), tlist, c_ops, saveat=saveat, progress_bar = Val(false))
+        sol_sme4 = smesolve(H, ψ0, tlist, c_ops_sme, sc_ops_sme, saveat=saveat, ntraj=10, progress_bar = Val(false), rng = MersenneTwister(12))
+        sol_sme5 = smesolve(H, operator_to_vector(ket2dm(ψ0)), tlist, c_ops_sme, sc_ops_sme, saveat=saveat, ntraj=10, progress_bar = Val(false), rng = MersenneTwister(12))
+
         ρ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)
         expect_mc_states_mean = sum(expect_mc_states, dims = 2) / size(expect_mc_states, 2)
@@ -190,6 +195,7 @@
         @test length(sol_me3.states) == length(saveat)
         @test size(sol_me3.expect) == (length(e_ops), length(tlist))
         @test sol_me3.expect[1, saveat_idxs] ≈ expect(e_ops[1], sol_me3.states) atol = 1e-6
+        @test all([sol_me3.states[i] ≈ vector_to_operator(sol_me4.states[i]) for i in eachindex(saveat)])
         @test length(sol_mc.times) == length(tlist)
         @test size(sol_mc.expect) == (length(e_ops), length(tlist))
         @test length(sol_mc_states.times) == length(tlist)
@@ -202,6 +208,7 @@
         @test isnothing(sol_sme.measurement)
         @test size(sol_sse2.measurement) == (length(c_ops), 20, length(tlist) - 1)
         @test size(sol_sme2.measurement) == (length(sc_ops_sme), 20, length(tlist) - 1)
+        @test all([sol_sme4.states[j][i] ≈ vector_to_operator(sol_sme5.states[j][i]) for i in eachindex(saveat), j in 1:10])
 
         @test sol_me_string ==
               "Solution of time evolution\n" *