return_problem keyword added to QuantumOptics integrators

.gitignore

@@ -3,3 +3,5 @@ docs/source/examples/
 docs/source/tutorial/_tutorial.rst
 .ipynb_checkpoints
 *_files/
+Manifest.toml
+/.vscode/**

src/mcwf.jl

@@ -322,14 +322,15 @@
 * `kwargs`: Further arguments are passed on to the ode solver.
 """
 function integrate_mcwf(dmcwf::T, jumpfun::J, tspan,
-                        psi0, seed, fout;
-                        display_beforeevent=false, display_afterevent=false,
-                        display_jumps=false,
-                        rng_state=nothing,
-                        save_everystep=false, callback=nothing,
-                        saveat=tspan,
-                        alg=OrdinaryDiffEq.DP5(),
-                        kwargs...) where {T, J}
+                    psi0, seed, fout;
+                    display_beforeevent=false, display_afterevent=false,
+                    display_jumps=false,
+                    rng_state=nothing,
+                    save_everystep=false, callback=nothing,
+                    saveat=tspan,
+                    alg=OrdinaryDiffEq.DP5(),
+                    return_problem=false,
+                    kwargs...) where {T,J}
 
     tspan_ = convert(Vector{float(eltype(tspan))}, tspan)
     # Display before or after events
@@ -388,7 +389,14 @@
 
     prob = OrdinaryDiffEq.ODEProblem{true}(df_, as_vector(psi0), (tspan_[1],tspan_[end]))
 
-    sol = OrdinaryDiffEq.solve(
+    if return_problem
+        if display_jumps
+            return Dict("out" => out, "jump_t" => jump_t, "jump_index" => jump_index, "prob" => prob, "alg" => alg, "solve_kwargs" => (reltol=1.0e-6, abstol=1.0e-8, save_everystep=false, save_start=false, save_end=false, callback=full_cb, kwargs...))
+        else
+            return Dict("out" => out, "prob" => prob, "alg" => alg, "solve_kwargs" => (reltol=1.0e-6, abstol=1.0e-8, save_everystep=false, save_start=false, save_end=false, callback=full_cb, kwargs...))
+        end
+    else
+        sol = OrdinaryDiffEq.solve(
                 prob,
                 alg;
                 reltol = 1.0e-6,
@@ -397,10 +405,11 @@
                 save_end = false,
                 callback=full_cb, kwargs...)
 
-    if display_jumps
-        return out.t, out.saveval, jump_t, jump_index
-    else
-        return out.t, out.saveval
+        if display_jumps
+            return out.t, out.saveval, jump_t, jump_index
+        else
+            return out.t, out.saveval
+        end
     end
 end
 

src/stochastic_base.jl

@@ -11,14 +11,15 @@ import DiffEqCallbacks, StochasticDiffEq, OrdinaryDiffEq
 Integrate using StochasticDiffEq
 """
 function integrate_stoch(tspan, df, dg, x0,
-            state, dstate, fout, n;
-            save_everystep = false, callback=nothing, saveat=tspan,
-            alg=StochasticDiffEq.EM(),
-            noise_rate_prototype = nothing,
-            noise_prototype_classical = nothing,
-            noise=nothing,
-            ncb=nothing,
-            kwargs...)
+    state, dstate, fout, n;
+    save_everystep = false, callback=nothing, saveat=tspan,
+    alg=StochasticDiffEq.EM(),
+    noise_rate_prototype = nothing,
+    noise_prototype_classical = nothing,
+    noise=nothing,
+    ncb=nothing,
+    return_problem=false,
+    kwargs...)
 
     function df_(dx, x, p, t)
         recast!(state,x)
@@ -65,10 +66,13 @@ function integrate_stoch(tspan, df, dg, x0,
     full_cb = OrdinaryDiffEq.CallbackSet(callback, ncb, scb)
 
     prob = StochasticDiffEq.SDEProblem{true}(df_, dg_, x0,(tspan[1],tspan[end]),
-                    noise=noise_,
-                    noise_rate_prototype=noise_rate_prototype)
+        noise=noise_,
+        noise_rate_prototype=noise_rate_prototype)
 
-    sol = StochasticDiffEq.solve(
+    if return_problem
+        return Dict("out" => out, "prob" => prob, "alg" => alg, "solve_kwargs" => (reltol=1.0e-3, abstol=1.0e-3, save_everystep=false, save_start=false, save_end=false, callback=full_cb, kwargs...))
+    else
+        sol = StochasticDiffEq.solve(
                 prob,
                 alg;
                 reltol = 1.0e-3,
@@ -77,7 +81,8 @@ function integrate_stoch(tspan, df, dg, x0,
                 save_end = false,
                 callback=full_cb, kwargs...)
 
-    out.t,out.saveval
+        return out.t, out.saveval
+    end
 end
 
 """

src/timeevolution_base.jl

@@ -12,10 +12,10 @@ function recast! end
 Integrate using OrdinaryDiffEq
 """
 function integrate(tspan, df::F, x0,
-            state, dstate, fout;
-            alg = OrdinaryDiffEq.DP5(),
-            steady_state = false, tol = 1e-3, save_everystep = false, saveat=tspan,
-            callback = nothing, kwargs...) where {F}
+        state, dstate, fout;
+        alg = OrdinaryDiffEq.DP5(),
+        steady_state = false, tol = 1e-3, save_everystep = false, saveat=tspan,
+        callback = nothing, return_problem=false, kwargs...) where {F}
 
     function df_(dx, x, p, t)
         recast!(state,x)
@@ -60,15 +60,19 @@ function integrate(tspan, df::F, x0,
 
     full_cb = OrdinaryDiffEq.CallbackSet(callback,cb)
 
-    sol = OrdinaryDiffEq.solve(
-                prob,
-                alg;
-                reltol = 1.0e-6,
-                abstol = 1.0e-8,
-                save_everystep = false, save_start = false,
-                save_end = false,
-                callback=full_cb, kwargs...)
-    out.t,out.saveval
+    if return_problem
+        return Dict("out" => out, "prob" => prob, "alg" => alg, "solve_kwargs" => (reltol=1.0e-6, abstol=1.0e-8, save_everystep=false, save_start=false, save_end=false, callback=full_cb, kwargs...))
+    else
+        sol = OrdinaryDiffEq.solve(
+            prob,
+            alg;
+            reltol=1.0e-6,
+            abstol=1.0e-8,
+            save_everystep=false, save_start=false,
+            save_end=false,
+            callback=full_cb, kwargs...)
+        return out.t, out.saveval
+    end
 end
 
 function integrate(tspan, df, x0,

test/test_stochastic_schroedinger.jl

@@ -49,9 +49,15 @@ function fstoch_2(t, psi)
 end
 
 # Non-dynamic Schrödinger
-tout, ψt4 = stochastic.schroedinger(T, ψ0, H, [zero_op, zero_op]; dt=0.1dt)
+tout, ψt4 = stochastic.schroedinger(T, ψ0, H, [zero_op, zero_op]; dt=0.1dt, seed=1)
+rslt4 = stochastic.schroedinger(T, ψ0, H, [zero_op, zero_op]; dt=0.1dt, seed=1, return_problem=true)
+sol = StochasticDiffEq.solve(rslt4["prob"], rslt4["alg"]; rslt4["solve_kwargs"]...)
+@test ψt4 == rslt4["out"].saveval
 # Dynamic Schrödinger
-tout, ψt1 = stochastic.schroedinger_dynamic(T, ψ0, fdeterm, fstoch_1; dt=0.1dt)
+tout, ψt1 = stochastic.schroedinger_dynamic(T, ψ0, fdeterm, fstoch_1; dt=0.1dt, seed=1)
+rslt1 = stochastic.schroedinger(T, ψ0, H, [zero_op, zero_op]; dt=0.1dt, seed=1, return_problem=true)
+sol = StochasticDiffEq.solve(rslt1["prob"], rslt1["alg"]; rslt1["solve_kwargs"]...)
+@test ψt1 == rslt1["out"].saveval
 tout, ψt2 = stochastic.schroedinger_dynamic(T, ψ0, fdeterm, fstoch_2; noise_processes=3, dt=0.1dt)
 
 # Test equivalence to Schrödinger equation with zero noise

test/test_timeevolution_master.jl

@@ -1,6 +1,7 @@
 using Test
 using QuantumOptics
 using LinearAlgebra
+using OrdinaryDiffEq
 
 @testset "master" begin
 
@@ -57,6 +58,11 @@ end
 @test tracedistance(L*ρ₀, ρ) < 1e-10
 
 # Test master
+tout, ρt = timeevolution.master(T, Ψ₀, Hdense, Jdense; reltol=1e-6)
+rslt = timeevolution.master(T, Ψ₀, Hdense, Jdense; reltol=1e-6, return_problem=true)
+sol = OrdinaryDiffEq.solve(rslt["prob"], rslt["alg"]; rslt["solve_kwargs"]...)
+@test ρt == rslt["out"].saveval
+
 @test timeevolution.check_master(ρ₀, Hdense, Jdense, dagger.(Jdense), nothing)
 tout, ρt = timeevolution.master(T, ρ₀, Hdense, Jdense; reltol=1e-7)
 ρ = ρt[end]

test/test_timeevolution_mcwf.jl

@@ -1,6 +1,7 @@
 using Test
 using QuantumOptics
 using Random, LinearAlgebra
+using OrdinaryDiffEq
 
 @testset "mcwf" begin
 
@@ -46,6 +47,14 @@ Jlazy = [LazyTensor(basis, 1, sqrt(γ)*sm), LazyTensor(basis, 2, sqrt(κ)*destro
 
 
 # Test mcwf
+@test timeevolution.check_mcwf(Ψ₀, Hdense, Jdense, dagger.(Jdense), nothing)
+tout, Ψt = timeevolution.mcwf(T, Ψ₀, Hdense, Jdense; seed=UInt(1), reltol=1e-7)
+rslt = timeevolution.mcwf(T, Ψ₀, Hdense, Jdense; seed=UInt(1), reltol=1e-7, return_problem=true)
+sol = OrdinaryDiffEq.solve(rslt["prob"], rslt["alg"]; rslt["solve_kwargs"]...)
+tout2 = rslt["out"].t
+Ψt2 = rslt["out"].saveval
+@test Ψt == Ψt2
+
 @test timeevolution.check_mcwf(Ψ₀, Hdense, Jdense, dagger.(Jdense), nothing)
 tout, Ψt = timeevolution.mcwf(T, Ψ₀, Hdense, Jdense; seed=UInt(1), reltol=1e-7)
 tout2, Ψt2 = timeevolution.mcwf(T, Ψ₀, Hdense, Jdense; seed=UInt(1), reltol=1e-7)

test/test_timeevolution_schroedinger.jl

@@ -1,5 +1,6 @@
 using Test
 using QuantumOptics
+using OrdinaryDiffEq
 
 @testset "schroedinger" begin
 
@@ -109,6 +110,9 @@ sz = sigmaz(basis)
 f(t, psi) = sx * π
 tspan = 0:1.0
 t, u = timeevolution.schroedinger(tspan, u0, π * sx)
+rslt = timeevolution.schroedinger(tspan, u0, π * sx, return_problem=true)
+sol = OrdinaryDiffEq.solve(rslt["prob"], rslt["alg"]; rslt["solve_kwargs"]...)
+@test u == rslt["out"].saveval
 
 # I think the tolerance on the differential equation is 1e-6, we expect the operator to be essentially the identity
 @test abs(expect(sz, u[end] * su)) - abs(expect(sz, u0 * su)) < 1e-6