add tests

qutip · Aug 31, 2024 · a15e236 · a15e236
1 parent 5730f8c
commit a15e236
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 14 deletions.
diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl
@@ -66,7 +66,7 @@ end
         ψ0::QuantumObject,
         tlist::AbstractVector;
         sc_ops::Vector{QuantumObject{Tc, OperatorQuantumObject}}=QuantumObject{Matrix, OperatorQuantumObject}[];
-        alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm=Tsit5()
+        alg::StochasticDiffEq.StochasticDiffEqAlgorithm=SRA1()
         e_ops::Union{Nothing,AbstractVector} = nothing,
         H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
         params::NamedTuple=NamedTuple(),
@@ -100,7 +100,7 @@ Above, `C_n` is the `n`-th collapse operator and  `dW_j(t)` is the real Wiener i
 - `ψ0::QuantumObject`: The initial state of the system ``|\psi(0)\rangle``.
 - `tlist::AbstractVector`: The time list of the evolution.
 - `sc_ops::Vector`: List of stochastic collapse operators ``\{\hat{C}_n\}_n``.
-- `alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm`: The algorithm used for the time evolution.
+- `alg::StochasticDiffEq.StochasticDiffEqAlgorithm`: The algorithm used for the time evolution.
 - `e_ops::Union{Nothing,AbstractVector}`: The list of operators to be evaluated during the evolution.
 - `H_t::Union{Nothing,Function,TimeDependentOperatorSum}`: The time-dependent Hamiltonian of the system. If `nothing`, the Hamiltonian is time-independent.
 - `params::NamedTuple`: The parameters of the system.
@@ -122,7 +122,7 @@ function ssesolveProblem(
     ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
     tlist::AbstractVector,
     sc_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
-    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
     e_ops::Union{Nothing,AbstractVector} = nothing,
     H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
@@ -138,7 +138,7 @@ function ssesolveProblem(
 
     progress_bar_val = makeVal(progress_bar)
 
-    t_l = convert(Vector{Float64}, tlist) # Convert it into Float64 to avoid type instabilities for OrdinaryDiffEq.jl
+    t_l = convert(Vector{Float64}, tlist) # Convert it into Float64 to avoid type instabilities for StochasticDiffEq.jl
 
     ϕ0 = get_data(ψ0)
 
@@ -202,7 +202,7 @@ end
         ψ0::QuantumObject,
         tlist::AbstractVector;
         sc_ops::Vector{QuantumObject{Tc, OperatorQuantumObject}}=QuantumObject{Matrix, OperatorQuantumObject}[];
-        alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm=Tsit5()
+        alg::StochasticDiffEq.StochasticDiffEqAlgorithm=SRA1()
         e_ops::Union{Nothing,AbstractVector} = nothing,
         H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
         params::NamedTuple=NamedTuple(),
@@ -237,7 +237,7 @@ Above, `C_n` is the `n`-th collapse operator and  `dW_j(t)` is the real Wiener i
 - `ψ0::QuantumObject`: The initial state of the system ``|\psi(0)\rangle``.
 - `tlist::AbstractVector`: The time list of the evolution.
 - `sc_ops::Vector`: List of stochastic collapse operators ``\{\hat{C}_n\}_n``.
-- `alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm`: The algorithm used for the time evolution.
+- `alg::StochasticDiffEq.StochasticDiffEqAlgorithm`: The algorithm used for the time evolution.
 - `e_ops::Union{Nothing,AbstractVector}`: The list of operators to be evaluated during the evolution.
 - `H_t::Union{Nothing,Function,TimeDependentOperatorSum}`: The time-dependent Hamiltonian of the system. If `nothing`, the Hamiltonian is time-independent.
 - `params::NamedTuple`: The parameters of the system.
@@ -260,7 +260,7 @@ function ssesolveEnsembleProblem(
     ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
     tlist::AbstractVector,
     sc_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
-    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
     e_ops::Union{Nothing,AbstractVector} = nothing,
     H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
@@ -275,13 +275,12 @@ function ssesolveEnsembleProblem(
     return ensemble_prob
 end
 
-
 @doc raw"""
     ssesolve(H::QuantumObject,
         ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
         tlist::AbstractVector,
         sc_ops::Vector{QuantumObject{Tc, OperatorQuantumObject}}=QuantumObject{Matrix, OperatorQuantumObject}[];
-        alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm=Tsit5(),
+        alg::StochasticDiffEq.StochasticDiffEqAlgorithm=SRA1(),
         e_ops::Vector{QuantumObject{Te, OperatorQuantumObject}}=QuantumObject{Matrix, OperatorQuantumObject}[],
         H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
         params::NamedTuple=NamedTuple(),
@@ -321,7 +320,7 @@ Above, `C_n` is the `n`-th collapse operator and  `dW_j(t)` is the real Wiener i
 - `ψ0::QuantumObject`: Initial state of the system ``|\psi(0)\rangle``.
 - `tlist::AbstractVector`: List of times at which to save the state of the system.
 - `sc_ops::Vector`: List of stochastic collapse operators ``\{\hat{C}_n\}_n``.
-- `alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm`: Algorithm to use for the time evolution.
+- `alg::StochasticDiffEq.StochasticDiffEqAlgorithm`: Algorithm to use for the time evolution.
 - `e_ops::Vector`: List of operators for which to calculate expectation values.
 - `H_t::Union{Nothing,Function,TimeDependentOperatorSum}`: Time-dependent part of the Hamiltonian.
 - `params::NamedTuple`: Dictionary of parameters to pass to the solver.
@@ -350,7 +349,7 @@ function ssesolve(
     ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
     tlist::AbstractVector,
     sc_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
-    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
     e_ops::Union{Nothing,AbstractVector} = nothing,
     H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
@@ -379,7 +378,7 @@ end
 
 function ssesolve(
     ens_prob::EnsembleProblem;
-    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
     n_traj::Int = 1,
     ensemble_method = EnsembleThreads(),
 )

diff --git a/src/time_evolution/time_evolution.jl b/src/time_evolution/time_evolution.jl
@@ -120,7 +120,7 @@ function Base.show(io::IO, sol::TimeEvolutionSSESol)
     print(io, "(converged: $(sol.converged))\n")
     print(io, "--------------------------------\n")
     print(io, "num_trajectories = $(sol.n_traj)\n")
-    # print(io, "num_states = $(length(sol.states[1]))\n")
+    print(io, "num_states = $(length(sol.states[1]))\n")
     print(io, "num_expect = $(size(sol.expect, 1))\n")
     print(io, "SDE alg.: $(sol.alg)\n")
     print(io, "abstol = $(sol.abstol)\n")

diff --git a/test/time_evolution_and_partial_trace.jl b/test/time_evolution_and_partial_trace.jl
@@ -44,7 +44,7 @@
         end
     end
 
-    @testset "mesolve and mcsolve" begin
+    @testset "mesolve, mcsolve, and ssesolve" begin
         N = 10
         a = destroy(N)
         a_d = a'
@@ -57,9 +57,12 @@
         sol_me2 = mesolve(H, psi0, t_l, c_ops, progress_bar = Val(false))
         sol_me3 = mesolve(H, psi0, t_l, c_ops, e_ops = e_ops, saveat = t_l, progress_bar = Val(false))
         sol_mc = mcsolve(H, psi0, t_l, c_ops, n_traj = 500, e_ops = e_ops, progress_bar = Val(false))
+        sol_sse = ssesolve(H, psi0, t_l, c_ops, n_traj = 500, e_ops = e_ops, progress_bar = Val(false))
         sol_me_string = sprint((t, s) -> show(t, "text/plain", s), sol_me)
         sol_mc_string = sprint((t, s) -> show(t, "text/plain", s), sol_mc)
+        sol_sse_string = sprint((t, s) -> show(t, "text/plain", s), sol_sse)
         @test sum(abs.(sol_mc.expect .- sol_me.expect)) / length(t_l) < 0.1
+        @test sum(abs.(sol_sse.expect .- sol_me.expect)) / length(t_l) < 0.1
         @test length(sol_me.states) == 1
         @test size(sol_me.expect) == (length(e_ops), length(t_l))
         @test length(sol_me2.states) == length(t_l)
@@ -85,6 +88,16 @@
               "ODE alg.: $(sol_mc.alg)\n" *
               "abstol = $(sol_mc.abstol)\n" *
               "reltol = $(sol_mc.reltol)\n"
+        @test sol_sse_string ==
+              "Solution of quantum trajectories\n" *
+              "(converged: $(sol_sse.converged))\n" *
+              "--------------------------------\n" *
+              "num_trajectories = $(sol_sse.n_traj)\n" *
+              "num_states = $(length(sol_sse.states[1]))\n" *
+              "num_expect = $(size(sol_sse.expect, 1))\n" *
+              "SDE alg.: $(sol_sse.alg)\n" *
+              "abstol = $(sol_sse.abstol)\n" *
+              "reltol = $(sol_sse.reltol)\n"
 
         @testset "Type Inference mesolve" begin
             if VERSION >= v"1.10"
@@ -110,6 +123,22 @@
                 @inferred mcsolve(H, psi0, t_l, c_ops, n_traj = 500, progress_bar = Val(true))
             end
         end
+
+        @testset "Type Inference ssesolve" begin
+            if VERSION >= v"1.10"
+                @inferred ssesolveEnsembleProblem(
+                    H,
+                    psi0,
+                    t_l,
+                    c_ops,
+                    n_traj = 500,
+                    e_ops = e_ops,
+                    progress_bar = Val(false),
+                )
+                @inferred ssesolve(H, psi0, t_l, c_ops, n_traj = 500, e_ops = e_ops, progress_bar = Val(false))
+                @inferred ssesolve(H, psi0, t_l, c_ops, n_traj = 500, progress_bar = Val(true))
+            end
+        end
     end
 
     @testset "exceptions" begin