3rd order Rosenbrocks joining the fun.

SciML · Jul 4, 2017 · f530a40 · f530a40
1 parent 73999f0
commit f530a40
Show file tree

Hide file tree

Showing 6 changed files with 550 additions and 7 deletions.
diff --git a/src/OrdinaryDiffEq.jl b/src/OrdinaryDiffEq.jl
@@ -111,7 +111,7 @@ module OrdinaryDiffEq
          Feagin10, Feagin12, Feagin14, CompositeAlgorithm
 
   export Rosenbrock23, Rosenbrock32, RosShamp4, Veldd4, Velds4, GRK4T, GRK4A,
-         Ros4LStab
+         Ros4LStab, ROS3P, Rodas3
 
   export IIF1, IIF2
 

diff --git a/src/alg_utils.jl b/src/alg_utils.jl
@@ -10,6 +10,8 @@ isfsal(alg::Tsit5) = true
 isfsal(alg::Vern6) = true
 isfsal(alg::Rosenbrock23) = true
 isfsal(alg::Rosenbrock32) = true
+isfsal(alg::ROS3P) = true
+isfsal(alg::Rodas3) = true
 isfsal(alg::RosShamp4) = true
 isfsal(alg::Veldd4) = true
 isfsal(alg::Velds4) = true
@@ -85,6 +87,8 @@ qmax_default(alg::DP8) = 6
 get_chunksize(alg::OrdinaryDiffEqAlgorithm) = error("This algorithm does not have a chunk size defined.")
 get_chunksize{CS,AD}(alg::Rosenbrock23{CS,AD}) = CS
 get_chunksize{CS,AD}(alg::Rosenbrock32{CS,AD}) = CS
+get_chunksize{CS,AD}(alg::ROS3P{CS,AD}) = CS
+get_chunksize{CS,AD}(alg::Rodas3{CS,AD}) = CS
 get_chunksize{CS,AD}(alg::RosShamp4{CS,AD}) = CS
 get_chunksize{CS,AD}(alg::Veldd4{CS,AD}) = CS
 get_chunksize{CS,AD}(alg::Velds4{CS,AD}) = CS
@@ -100,6 +104,8 @@ alg_extrapolates(alg::Trapezoid) = true
 alg_autodiff(alg::OrdinaryDiffEqAlgorithm) = error("This algorithm does not have an autodifferentiation option defined.")
 alg_autodiff{CS,AD}(alg::Rosenbrock23{CS,AD}) = AD
 alg_autodiff{CS,AD}(alg::Rosenbrock32{CS,AD}) = AD
+alg_autodiff{CS,AD}(alg::ROS3P{CS,AD}) = AD
+alg_autodiff{CS,AD}(alg::Rodas3{CS,AD}) = AD
 alg_autodiff{CS,AD}(alg::RosShamp4{CS,AD}) = AD
 alg_autodiff{CS,AD}(alg::Veldd4{CS,AD}) = AD
 alg_autodiff{CS,AD}(alg::Velds4{CS,AD}) = AD
@@ -164,6 +170,8 @@ alg_order(alg::Feagin14) = 14
 
 alg_order(alg::Rosenbrock23) = 2
 alg_order(alg::Rosenbrock32) = 3
+alg_order(alg::ROS3P) = 3
+alg_order(alg::Rodas3) = 3
 alg_order(alg::RosShamp4) = 4
 alg_order(alg::Veldd4) = 4
 alg_order(alg::Velds4) = 4
@@ -193,6 +201,8 @@ alg_adaptive_order(alg::Feagin14) = 12
 
 alg_adaptive_order(alg::Rosenbrock23) = 3
 alg_adaptive_order(alg::Rosenbrock32) = 2
+alg_adaptive_order(alg::ROS3P) = 2
+alg_adaptive_order(alg::Rodas3) = 2
 alg_adaptive_order(alg::RosShamp4) = 3
 alg_adaptive_order(alg::Veldd4) = 3
 alg_adaptive_order(alg::Velds4) = 3

diff --git a/src/caches/rosenbrock_caches.jl b/src/caches/rosenbrock_caches.jl
@@ -1,3 +1,8 @@
+
+################################################################################
+
+# Shampine's Low-order Rosenbrocks
+
 type Rosenbrock23Cache{uType,uArrayType,rateType,du2Type,LinuType,vecuType,JType,TabType,TFType,UFType,F,JCType} <: OrdinaryDiffEqMutableCache
   u::uType
   uprev::uType
@@ -165,13 +170,17 @@ function alg_cache(alg::Rosenbrock32,u,rate_prototype,uEltypeNoUnits,tTypeNoUnit
   Rosenbrock32ConstantCache(uEltypeNoUnits,tf,uf)
 end
 
+################################################################################
+
+### 3rd order specialized Rosenbrocks
+
 immutable Rosenbrock33ConstantCache{TF,UF,Tab} <: OrdinaryDiffEqConstantCache
   tf::TF
   uf::UF
   tab::Tab
 end
 
-type Rosenbrock33ConstantCache{uType,uArrayType,rateType,du2Type,LinuType,vecuType,JType,TabType,TFType,UFType,F,JCType} <: OrdinaryDiffEqMutableCache
+type Rosenbrock33Cache{uType,uArrayType,rateType,du2Type,LinuType,vecuType,JType,TabType,TFType,UFType,F,JCType} <: OrdinaryDiffEqMutableCache
   u::uType
   uprev::uType
   k1::rateType
@@ -200,10 +209,10 @@ type Rosenbrock33ConstantCache{uType,uArrayType,rateType,du2Type,LinuType,vecuTy
   jac_config::JCType
 end
 
-u_cache(c::Rosenbrock33ConstantCache) = (c.dT,c.tmp)
-du_cache(c::Rosenbrock33ConstantCache) = (c.k₁,c.k₂,c.k₃,c.du1,c.du2,c.f₁,c.fsalfirst,c.fsallast,c.linsolve_tmp)
-jac_cache(c::Rosenbrock33ConstantCache) = (c.J,c.W)
-vecu_cache(c::Rosenbrock33ConstantCache) = (c.vectmp,c.vectmp2,c.vectmp3)
+u_cache(c::Rosenbrock33Cache) = (c.dT,c.tmp)
+du_cache(c::Rosenbrock33Cache) = (c.k₁,c.k₂,c.k₃,c.du1,c.du2,c.f₁,c.fsalfirst,c.fsallast,c.linsolve_tmp)
+jac_cache(c::Rosenbrock33Cache) = (c.J,c.W)
+vecu_cache(c::Rosenbrock33Cache) = (c.vectmp,c.vectmp2,c.vectmp3)
 
 function alg_cache(alg::ROS3P,u,rate_prototype,uEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,::Type{Val{true}})
   k1 = zeros(rate_prototype)
@@ -235,11 +244,101 @@ function alg_cache(alg::ROS3P,u,rate_prototype,uEltypeNoUnits,tTypeNoUnits,uprev
   else
     jac_config = nothing
   end
-  Rosenbrock33ConstantCache(u,uprev,k1,k2,k3,k4,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
+  Rosenbrock33Cache(u,uprev,k1,k2,k3,k4,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                     fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                     linsolve_tmp_vec,alg.linsolve,jac_config)
 end
 
+function alg_cache(alg::ROS3P,u,rate_prototype,uEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,::Type{Val{false}})
+  tf = TimeDerivativeWrapper(f,u)
+  uf = UDerivativeWrapper(f,t)
+  Rosenbrock33ConstantCache(tf,uf,ROS3PConstantCache(realtype(uEltypeNoUnits),realtype(tTypeNoUnits)))
+end
+
+type Rosenbrock34Cache{uType,uArrayType,rateType,du2Type,LinuType,vecuType,JType,TabType,TFType,UFType,F,JCType} <: OrdinaryDiffEqMutableCache
+  u::uType
+  uprev::uType
+  k1::rateType
+  k2::rateType
+  k3::rateType
+  k4::rateType
+  du::rateType
+  du1::rateType
+  du2::du2Type
+  vectmp::vecuType
+  vectmp2::vecuType
+  vectmp3::vecuType
+  vectmp4::vecuType
+  fsalfirst::rateType
+  fsallast::rateType
+  dT::uArrayType
+  J::JType
+  W::JType
+  tmp::uArrayType
+  tab::TabType
+  tf::TFType
+  uf::UFType
+  linsolve_tmp::LinuType
+  linsolve_tmp_vec::vecuType
+  linsolve::F
+  jac_config::JCType
+end
+
+u_cache(c::Rosenbrock34Cache) = (c.dT,c.tmp)
+du_cache(c::Rosenbrock34Cache) = (c.k₁,c.k₂,c.k₃,c.du1,c.du2,c.f₁,c.fsalfirst,c.fsallast,c.linsolve_tmp)
+jac_cache(c::Rosenbrock34Cache) = (c.J,c.W)
+vecu_cache(c::Rosenbrock34Cache) = (c.vectmp,c.vectmp2,c.vectmp3)
+
+function alg_cache(alg::Rodas3,u,rate_prototype,uEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,::Type{Val{true}})
+  k1 = zeros(rate_prototype)
+  k2 = zeros(rate_prototype)
+  k3 = zeros(rate_prototype)
+  k4 = zeros(rate_prototype)
+  du = zeros(rate_prototype)
+  du1 = zeros(rate_prototype)
+  du2 = zeros(rate_prototype)
+  vectmp = vec(similar(u,indices(u)))
+  vectmp2 = vec(similar(u,indices(u)))
+  vectmp3 = vec(similar(u,indices(u)))
+  vectmp4 = vec(similar(u,indices(u)))
+  fsalfirst = zeros(rate_prototype)
+  fsallast = zeros(rate_prototype)
+  dT = similar(u,indices(u))
+  J = zeros(uEltypeNoUnits,length(u),length(u)) # uEltype?
+  W = similar(J);
+  tmp = similar(u,indices(u))
+  tab = Rodas3ConstantCache(realtype(uEltypeNoUnits),realtype(tTypeNoUnits))
+  vf = VectorF(f,size(u))
+  vfr = VectorFReturn(f,size(u))
+  tf = TimeGradientWrapper(vf,uprev)
+  uf = UJacobianWrapper(vfr,t)
+  linsolve_tmp = similar(u,indices(u))
+  linsolve_tmp_vec = vec(linsolve_tmp)
+  if alg_autodiff(alg)
+    jac_config = ForwardDiff.JacobianConfig(uf,vec(du1),vec(uprev),ForwardDiff.Chunk{determine_chunksize(u,alg)}())
+  else
+    jac_config = nothing
+  end
+  Rosenbrock34Cache(u,uprev,k1,k2,k3,k4,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
+                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
+                    linsolve_tmp_vec,alg.linsolve,jac_config)
+end
+
+immutable Rosenbrock34ConstantCache{TF,UF,Tab} <: OrdinaryDiffEqConstantCache
+  tf::TF
+  uf::UF
+  tab::Tab
+end
+
+function alg_cache(alg::Rodas3,u,rate_prototype,uEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,::Type{Val{false}})
+  tf = TimeDerivativeWrapper(f,u)
+  uf = UDerivativeWrapper(f,t)
+  Rosenbrock34ConstantCache(tf,uf,Rodas3ConstantCache(realtype(uEltypeNoUnits),realtype(tTypeNoUnits)))
+end
+
+
+################################################################################
+
 immutable Rosenbrock4ConstantCache{TF,UF,Tab} <: OrdinaryDiffEqConstantCache
   tf::TF
   uf::UF