Make saveat kwarg

qojulia · Apr 14, 2020 · 9aa3ff8 · 9aa3ff8
1 parent 9933083
commit 9aa3ff8
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/src/mcwf.jl b/src/mcwf.jl
@@ -246,6 +246,7 @@ function integrate_mcwf(dmcwf::Function, jumpfun::Function, tspan,
                         display_beforeevent=false, display_afterevent=false,
                         display_jumps=false,
                         save_everystep=false, callback=nothing,
+                        saveat=tspan,
                         alg=OrdinaryDiffEq.DP5(),
                         kwargs...) where T
 
@@ -282,7 +283,7 @@ function integrate_mcwf(dmcwf::Function, jumpfun::Function, tspan,
     dstate = copy(psi0)
     out_type = pure_inference(fout, Tuple{eltype(tspan),typeof(state)})
     out = DiffEqCallbacks.SavedValues(Float64,out_type)
-    scb = DiffEqCallbacks.SavingCallback(fout_,out,saveat=tspan,
+    scb = DiffEqCallbacks.SavingCallback(fout_,out,saveat=saveat,
                                          save_everystep=save_everystep,
                                          save_start = false)
 

diff --git a/src/stochastic_base.jl b/src/stochastic_base.jl
@@ -12,7 +12,7 @@ Integrate using StochasticDiffEq
 """
 function integrate_stoch(tspan::Vector{Float64}, df::Function, dg::Function, x0::Vector{ComplexF64},
             state::T, dstate::T, fout::Function, n::Int;
-            save_everystep = false, callback=nothing,
+            save_everystep = false, callback=nothing, saveat=tspan,
             alg::StochasticDiffEq.StochasticDiffEqAlgorithm=StochasticDiffEq.EM(),
             noise_rate_prototype = nothing,
             noise_prototype_classical = nothing,
@@ -58,7 +58,7 @@ function integrate_stoch(tspan::Vector{Float64}, df::Function, dg::Function, x0:
 
     out = DiffEqCallbacks.SavedValues(Float64,out_type)
 
-    scb = DiffEqCallbacks.SavingCallback(fout_,out,saveat=tspan,
+    scb = DiffEqCallbacks.SavingCallback(fout_,out,saveat=saveat,
                                          save_everystep=save_everystep,
                                          save_start = false)
 

diff --git a/src/timeevolution_base.jl b/src/timeevolution_base.jl
@@ -16,7 +16,7 @@ Integrate using OrdinaryDiffEq
 function integrate(tspan::Vector{Float64}, df::Function, x0::DiffArray,
             state::T, dstate::T, fout::Function;
             alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm = OrdinaryDiffEq.DP5(),
-            steady_state = false, tol = 1e-3, save_everystep = false,
+            steady_state = false, tol = 1e-3, save_everystep = false, saveat=tspan,
             callback = nothing, kwargs...) where T
 
     function df_(dx::DiffArray, x::DiffArray, p, t)
@@ -34,7 +34,7 @@ function integrate(tspan::Vector{Float64}, df::Function, x0::DiffArray,
 
     out = DiffEqCallbacks.SavedValues(Float64,out_type)
 
-    scb = DiffEqCallbacks.SavingCallback(fout_,out,saveat=tspan,
+    scb = DiffEqCallbacks.SavingCallback(fout_,out,saveat=saveat,
                                          save_everystep=save_everystep,
                                          save_start = false)