formatting changes removed

Author: henhen724

.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 @@ Integrate a single Monte Carlo wave function trajectory.
 * `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 @@ function integrate_mcwf(dmcwf::T, jumpfun::J, tspan,
 
     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 @@ function integrate_mcwf(dmcwf::T, jumpfun::J, tspan,
                 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)
@@ -42,33 +42,37 @@ function integrate(tspan, df::F, x0,
                                          save_start = false,
                                          tdir = first(tspan)<last(tspan) ? one(eltype(tspan)) : -one(eltype(tspan)))
 
-    prob = OrdinaryDiffEq.ODEProblem{true}(df_, x0,(convert(tType, tspan[1]),convert(tType, tspan[end])))
-
-    if steady_state
-        affect! = function (integrator)
-            !save_everystep && scb.affect!(integrator,true)
-            OrdinaryDiffEq.terminate!(integrator)
-        end
-        _cb = OrdinaryDiffEq.DiscreteCallback(
-                                SteadyStateCondtion(copy(state),tol,state),
-                                affect!;
-                                save_positions = (false,false))
-        cb = OrdinaryDiffEq.CallbackSet(_cb,scb)
-    else
-        cb = scb
-    end
+                                         prob = OrdinaryDiffEq.ODEProblem{true}(df_, x0,(convert(tType, tspan[1]),convert(tType, tspan[end])))
+
+                                         if steady_state
+                                            affect! = function (integrator)
+                                                !save_everystep && scb.affect!(integrator,true)
+                                                OrdinaryDiffEq.terminate!(integrator)
+                                            end
+                                            _cb = OrdinaryDiffEq.DiscreteCallback(
+                                                                    SteadyStateCondtion(copy(state),tol,state),
+                                                                    affect!;
+                                                                    save_positions = (false,false))
+                                            cb = OrdinaryDiffEq.CallbackSet(_cb,scb)
+                                        else
+                                            cb = scb
+                                        end
 
     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