From 45a0945d43d85b8a9b456031f7742382a73401ac Mon Sep 17 00:00:00 2001 From: Alberto Mercurio Date: Sat, 22 Feb 2025 17:56:37 +0100 Subject: [PATCH 1/2] [no ci] Add support for `OperatorKet` input for `mesolve` and `smesolve` --- src/time_evolution/mesolve.jl | 23 ++++++++++++++++------- src/time_evolution/smesolve.jl | 28 +++++++++++++++++----------- test/core-test/time_evolution.jl | 7 +++++++ 3 files changed, 40 insertions(+), 18 deletions(-) diff --git a/src/time_evolution/mesolve.jl b/src/time_evolution/mesolve.jl index ee78e7eac..ebdcb4fd9 100644 --- a/src/time_evolution/mesolve.jl +++ b/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, diff --git a/src/time_evolution/smesolve.jl b/src/time_evolution/smesolve.jl index 4890030d2..6968c7f8e 100644 --- a/src/time_evolution/smesolve.jl +++ b/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)) diff --git a/test/core-test/time_evolution.jl b/test/core-test/time_evolution.jl index bd5be390a..54f564b38 100644 --- a/test/core-test/time_evolution.jl +++ b/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" * From c950058b355947f2ea1adc50ae91d41ecfaf1cc7 Mon Sep 17 00:00:00 2001 From: Alberto Mercurio Date: Sat, 22 Feb 2025 17:59:36 +0100 Subject: [PATCH 2/2] Make changelog and format --- CHANGELOG.md | 3 +++ src/time_evolution/mesolve.jl | 4 ++-- src/time_evolution/smesolve.jl | 4 ++-- test/core-test/time_evolution.jl | 30 ++++++++++++++++++++++++++---- 4 files changed, 33 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 301d6c673..51f3946c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased](https://github.com/qutip/QuantumToolbox.jl/tree/main) +- Add support for `OperatorKet` state input for `mesolve` and `smesolve`. ([#423]) + ## [v0.28.0] Release date: 2025-02-22 @@ -171,3 +173,4 @@ Release date: 2024-11-13 [#419]: https://github.com/qutip/QuantumToolbox.jl/issues/419 [#420]: https://github.com/qutip/QuantumToolbox.jl/issues/420 [#421]: https://github.com/qutip/QuantumToolbox.jl/issues/421 +[#423]: https://github.com/qutip/QuantumToolbox.jl/issues/423 diff --git a/src/time_evolution/mesolve.jl b/src/time_evolution/mesolve.jl index ebdcb4fd9..f0796f0c4 100644 --- a/src/time_evolution/mesolve.jl +++ b/src/time_evolution/mesolve.jl @@ -79,7 +79,7 @@ function mesolveProblem( ρ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)) + to_dense(_CType(T), mat2vec(ket2dm(ψ0).data)) end L = L_evo.data @@ -89,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, (isoperket=Val(isoperket(ψ0)),)) + return TimeEvolutionProblem(prob, tlist, L_evo.dimensions, (isoperket = Val(isoperket(ψ0)),)) end @doc raw""" diff --git a/src/time_evolution/smesolve.jl b/src/time_evolution/smesolve.jl index 6968c7f8e..ae28e4716 100644 --- a/src/time_evolution/smesolve.jl +++ b/src/time_evolution/smesolve.jl @@ -104,7 +104,7 @@ function smesolveProblem( ρ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)) + to_dense(_CType(T), mat2vec(ket2dm(ψ0).data)) end progr = ProgressBar(length(tlist), enable = getVal(progress_bar)) @@ -148,7 +148,7 @@ function smesolveProblem( kwargs3..., ) - return TimeEvolutionProblem(prob, tlist, dims, (isoperket=Val(isoperket(ψ0)),)) + return TimeEvolutionProblem(prob, tlist, dims, (isoperket = Val(isoperket(ψ0)),)) end @doc raw""" diff --git a/test/core-test/time_evolution.jl b/test/core-test/time_evolution.jl index 54f564b38..6d5e401fa 100644 --- a/test/core-test/time_evolution.jl +++ b/test/core-test/time_evolution.jl @@ -159,9 +159,29 @@ 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)) + 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) @@ -208,7 +228,9 @@ @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 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" *