step splitting in euler methods

SciML · Apr 1, 2018 · 384ceea · 384ceea
1 parent 74e0daf
commit 384ceea
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 12 deletions.
diff --git a/src/alg_utils.jl b/src/alg_utils.jl
@@ -111,3 +111,7 @@ function alg_mass_matrix_compatible(alg::Union{StochasticDiffEqNewtonAlgorithm,S
         error("Algorithm must be set as symplectic or theta=1 for mass matrices")
     end
 end
+
+is_split_step(::StochasticDiffEqAlgorithm) = false
+is_split_step(::EM{split}) = split
+is_split_step(::LambaEM{split}) = split
diff --git a/src/algorithms.jl b/src/algorithms.jl
@@ -13,12 +13,17 @@
 
 # Basics
 
-struct EM <: StochasticDiffEqAlgorithm end
+struct EM{split} <: StochasticDiffEqAlgorithm end
+Base.@pure EM(split=true) = EM{split}()
+
 struct SplitEM <: StochasticDiffEqAlgorithm end
 struct EulerHeun <: StochasticDiffEqAlgorithm end
 
-struct LambaEM <: StochasticDiffEqAdaptiveAlgorithm end
+struct LambaEM{split} <: StochasticDiffEqAdaptiveAlgorithm end
+Base.@pure LambaEM(split=true) = LambaEM{split}()
+
 struct LambaEulerHeun <: StochasticDiffEqAdaptiveAlgorithm end
+
 struct RKMil{interpretation} <: StochasticDiffEqAdaptiveAlgorithm end
 Base.@pure RKMil(;interpretation=:Ito) = RKMil{interpretation}()
 

diff --git a/src/perform_step/lamba.jl b/src/perform_step/lamba.jl
@@ -2,7 +2,13 @@
   @unpack t,dt,uprev,u,W,p = integrator
   du1 = integrator.f(uprev,p,t)
   K = @muladd uprev + dt*du1
-  L = integrator.g(uprev,p,t)
+
+  if is_split_step(integrator.alg)
+    L = integrator.g(uprev,p,t+dt)
+  else
+    L = integrator.g(K,p,t+dt)
+  end
+
   mil_correction = zero(u)
 
   u = K+L*W.dW
@@ -24,10 +30,16 @@ end
 @muladd function perform_step!(integrator,cache::LambaEMCache,f=integrator.f)
   @unpack du1,du2,K,tmp,L,gtmp = cache
   @unpack t,dt,uprev,u,W,p = integrator
+
   integrator.f(du1,uprev,p,t)
-  integrator.g(L,uprev,p,t)
   @. K = @muladd uprev + dt*du1
 
+  if is_split_step(integrator.alg)
+    integrator.g(L,K,p,t+dt)
+  else
+    integrator.g(L,uprev,p,t+dt)
+  end
+
   if is_diagonal_noise(integrator.sol.prob)
     @. tmp=L*W.dW
   else

diff --git a/src/perform_step/low_order.jl b/src/perform_step/low_order.jl
@@ -1,19 +1,38 @@
 @muladd function perform_step!(integrator,cache::EMConstantCache,f=integrator.f)
   @unpack t,dt,uprev,u,W,p = integrator
+
+  K = @muladd uprev + dt*integrator.f(uprev,p,t)
+
+  if is_split_step(integrator.alg)
+    u_choice = K
+  else
+    u_choice = uprev
+  end
+
   if !is_diagonal_noise(integrator.sol.prob) || typeof(W.dW) <: Number
-    noise = integrator.g(uprev,p,t)*W.dW
+    noise = integrator.g(u_choice,p,t)*W.dW
   else
-    noise = integrator.g(uprev,p,t).*W.dW
+    noise = integrator.g(u_choice,p,t).*W.dW
   end
-  u = @muladd uprev + dt*integrator.f(uprev,p,t) + noise
+
+  u = K + noise
   integrator.u = u
 end
 
 @muladd function perform_step!(integrator,cache::EMCache,f=integrator.f)
   @unpack rtmp1,rtmp2,rtmp3 = cache
   @unpack t,dt,uprev,u,W,p = integrator
   integrator.f(rtmp1,uprev,p,t)
-  integrator.g(rtmp2,uprev,p,t)
+
+  @. u = @muladd uprev + dt*rtmp1
+
+  if is_split_step(integrator.alg)
+    u_choice = u
+  else
+    u_choice = uprev
+  end
+
+  integrator.g(rtmp2,u,p,t)
 
   if is_diagonal_noise(integrator.sol.prob)
     @tight_loop_macros for i in eachindex(u)
@@ -23,10 +42,7 @@ end
     A_mul_B!(rtmp3,rtmp2,W.dW)
   end
 
-  #@. u = @muladd uprev + dt*rtmp1 + rtmp3
-  @tight_loop_macros for i in eachindex(u)
-    @inbounds u[i] = @muladd uprev[i] + dt*rtmp1[i] + rtmp3[i]
-  end
+  @. u += rtmp3
 end
 
 @muladd function perform_step!(integrator,cache::EulerHeunConstantCache,f=integrator.f)