Merge pull request SciML#8 from tkf/DynamicSS-tspan

Add tspan option to DynamicSS
benneti · Oct 18, 2018 · f3bf5ef · f3bf5ef
2 parents 2cfbdec + b155816
commit f3bf5ef
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/src/algorithms.jl b/src/algorithms.jl
@@ -5,9 +5,14 @@ struct SSRootfind{F} <: SteadyStateDiffEqAlgorithm
 end
 SSRootfind(;nlsolve=(f,u0,abstol) -> (res=NLsolve.nlsolve(f,u0,ftol = abstol);res.zero)) = SSRootfind(nlsolve)
 
-struct DynamicSS{A,AT,RT} <: SteadyStateDiffEqAlgorithm
+struct DynamicSS{A,AT,RT,TS} <: SteadyStateDiffEqAlgorithm
   alg::A
   abstol::AT
   reltol::RT
+  tspan::TS
 end
-DynamicSS(alg;abstol = 1e-8, reltol = 1e-6) = DynamicSS(alg,abstol,reltol)
+DynamicSS(alg; abstol = 1e-8, reltol = 1e-6, tspan = Inf) =
+  DynamicSS(alg, abstol, reltol, tspan)
+
+# Backward compatibility:
+DynamicSS(alg, abstol, reltol) = DynamicSS(alg; abstol=abstol, reltol=reltol)
diff --git a/src/solve.jl b/src/solve.jl
@@ -43,7 +43,12 @@ end
 function DiffEqBase.__solve(prob::DiffEqBase.AbstractSteadyStateProblem,
                             alg::DynamicSS,args...;kwargs...)
 
-  _prob = ODEProblem(prob.f,prob.u0,(0.0,Inf),prob.p)
-  solve(_prob,alg.alg,args...;kwargs...,
-        callback=TerminateSteadyState(alg.abstol,alg.reltol))
+  tspan = alg.tspan isa Tuple ? alg.tspan : (zero(alg.tspan), alg.tspan)
+  _prob = ODEProblem(prob.f,prob.u0,tspan,prob.p)
+  sol = solve(_prob,alg.alg,args...;kwargs...,
+              callback=TerminateSteadyState(alg.abstol,alg.reltol))
+  if sol.t[end] == _prob.tspan[end]
+    sol = DiffEqBase.solution_new_retcode(sol, :Failure)
+  end
+  sol
 end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -9,6 +9,7 @@ u0 = zeros(2)
 prob = SteadyStateProblem(f,u0)
 abstol = 1e-8
 sol = solve(prob,SSRootfind())
+@test sol.retcode == :Success
 p = nothing
 
 du = zeros(2)
@@ -18,22 +19,29 @@ f(du,sol.u,nothing,0)
 prob = ODEProblem(f,u0,(0.0,1.0))
 prob = SteadyStateProblem(prob)
 sol = solve(prob,SSRootfind(nlsolve = (f,u0,abstol) -> (res=NLsolve.nlsolve(f,u0,autodiff=:forward,method=:newton,iterations=Int(1e6),ftol=abstol);res.zero) ))
+@test sol.retcode == :Success
 
 f(du,sol.u,nothing,0)
 @test du == [0,0]
 
 # Use Sundials
 sol = solve(prob,SSRootfind(nlsolve = (f,u0,abstol) -> (res=Sundials.kinsol(f,u0)) ))
+@test sol.retcode == :Success
 f(du,sol.u,nothing,0)
 @test du == [0,0]
 
 using OrdinaryDiffEq
 sol = solve(prob,DynamicSS(Rodas5()))
+@test sol.retcode == :Success
 
 f(du,sol.u[end],p,0)
 @test du ≈ [0,0] atol = 1e-7
 
+sol = solve(prob,DynamicSS(Rodas5(),tspan=1e-3))
+@test sol.retcode != :Success
+
 sol = solve(prob,DynamicSS(CVODE_BDF()),dt=1.0)
+@test sol.retcode == :Success
 
 f(du,sol.u[end],p,0)
 @test du ≈ [0,0] atol = 1e-6