Implement OBABO, remains some error on weak order

Author: Hadrien

src/caches/dynamical_caches.jl

@@ -2,30 +2,29 @@ abstract type StochasticDynamicalEqConstantCache <: StochasticDiffEqConstantCach
 abstract type StochasticDynamicalEqMutableCache <: StochasticDiffEqMutableCache end
 
 
-mutable struct BAOABConstantCache{uType,uEltypeNoUnits,uCoeffType} <: StochasticDynamicalEqConstantCache
+mutable struct BAOABConstantCache{uType,uEltypeNoUnits,uCoeffType, uCoeffMType} <: StochasticDynamicalEqConstantCache
   k::uType
   half::uEltypeNoUnits
   c1::uCoeffType
-  c2::uCoeffType
+  c2::uCoeffMType
 end
-@cache struct BAOABCache{uType,uEltypeNoUnits,rateNoiseType,uCoeffType,uTypeCombined} <: StochasticDynamicalEqMutableCache
+@cache struct BAOABCache{uType,uEltypeNoUnits,rateNoiseType,uCoeffType, uCoeffMType,uTypeCombined} <: StochasticDynamicalEqMutableCache
   utmp::uType
+  dumid::uType
   dutmp::uType
+  dunoise::uType
   k::uType
   gtmp::rateNoiseType
   noise::uType
   half::uEltypeNoUnits
   c1::uCoeffType
-  c2::uCoeffType
+  c2::uCoeffMType
   tmp::uTypeCombined
 end
 
 function alg_cache(alg::BAOAB,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{false}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
   k = zero(rate_prototype.x[1])
-  if typeof(alg.gamma) <: Number
-      c1 = exp.(-alg.gamma*dt)
-      c2 = sqrt.(1 .- alg.scale_noise*c1.^2)# if scale_noise == false, c2 = 1
-  elseif typeof(alg.gamma) <: AbstractMatrix
+  if typeof(alg.gamma) <: AbstractMatrix
       c1 = exp(-alg.gamma*dt)
       c2 = cholesky(I - alg.scale_noise*c1*transpose(c1)).U# if scale_noise == false, c2 = 1
   else
@@ -36,18 +35,18 @@ function alg_cache(alg::BAOAB,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototy
 end
 
 function alg_cache(alg::BAOAB,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{true}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
+  dumid = zero(u.x[1])
   dutmp = zero(u.x[1])
+  dunoise = zero(u.x[1])
   utmp = zero(u.x[2])
   k = zero(rate_prototype.x[1])
 
   gtmp = zero(noise_rate_prototype)
   noise = zero(rate_prototype.x[1])
 
   half = uEltypeNoUnits(1//2)
-  if typeof(alg.gamma) <: Number
-      c1 = exp.(-alg.gamma*dt)
-      c2 = sqrt.(1 .- alg.scale_noise*c1.^2)# if scale_noise == false, c2 = 1
-  elseif typeof(alg.gamma) <: AbstractMatrix
+
+  if typeof(alg.gamma) <: AbstractMatrix
       c1 = exp(-alg.gamma*dt)
       c2 = cholesky(I - alg.scale_noise*c1*transpose(c1)).U# if scale_noise == false, c2 = 1
   else
@@ -57,36 +56,35 @@ function alg_cache(alg::BAOAB,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototy
 
   tmp = zero(u)
 
-  BAOABCache(utmp, dutmp, k, gtmp, noise, half, c1, c2, tmp)
+  BAOABCache(utmp, dumid, dutmp, dunoise, k, gtmp, noise, half, c1, c2, tmp)
 end
 
 
 
-mutable struct ABOBAConstantCache{uType,uEltypeNoUnits, uCoeffType} <: StochasticDynamicalEqConstantCache
+mutable struct ABOBAConstantCache{uType,uEltypeNoUnits, uCoeffType, uCoeffMType} <: StochasticDynamicalEqConstantCache
   k::uType
   half::uEltypeNoUnits
   c₂::uCoeffType
-  σ::uCoeffType
+  σ::uCoeffMType
 end
-@cache struct ABOBACache{uType,uEltypeNoUnits,rateNoiseType,uCoeffType,uTypeCombined} <: StochasticDynamicalEqMutableCache
+@cache struct ABOBACache{uType,uEltypeNoUnits,rateNoiseType,uCoeffType, uCoeffMType,uTypeCombined} <: StochasticDynamicalEqMutableCache
   utmp::uType
+  dumid::uType
   dutmp::uType
+  dunoise::uType
   k::uType
   gtmp::rateNoiseType
   noise::uType
   half::uEltypeNoUnits
   c₂::uCoeffType
-  σ::uCoeffType
+  σ::uCoeffMType
   tmp::uTypeCombined
 end
 
 function alg_cache(alg::ABOBA,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{false}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
   k = zero(rate_prototype.x[1])
 
-  if typeof(alg.gamma) <: Number
-      c₂ = exp.(-alg.gamma*dt)
-      σ = sqrt.(1 .- alg.scale_noise*c₂.^2)
-  elseif typeof(alg.gamma) <: AbstractMatrix
+  if typeof(alg.gamma) <: AbstractMatrix
       c₂ = exp(-alg.gamma*dt)
       σ = cholesky(I - alg.scale_noise*c₂*transpose(c₂)).U
   else
@@ -99,6 +97,8 @@ end
 
 function alg_cache(alg::ABOBA,prob,u,ΔW,ΔZ,p,rate_prototype, noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{true}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
   dutmp = zero(u.x[1])
+  dumid = zero(u.x[1])
+  dunoise = zero(u.x[1])
   utmp = zero(u.x[2])
   k = zero(rate_prototype.x[1])
 
@@ -107,10 +107,7 @@ function alg_cache(alg::ABOBA,prob,u,ΔW,ΔZ,p,rate_prototype, noise_rate_protot
 
   half = uEltypeNoUnits(1//2)
 
-  if typeof(alg.gamma) <: Number
-      c₂ = exp.(-alg.gamma*dt)
-      σ = sqrt.(1 .- alg.scale_noise*c₂.^2)
-  elseif typeof(alg.gamma) <: AbstractMatrix
+  if typeof(alg.gamma) <: AbstractMatrix
       c₂ = exp(-alg.gamma*dt)
       σ = cholesky(I - alg.scale_noise*c₂*transpose(c₂)).U
   else
@@ -120,64 +117,64 @@ function alg_cache(alg::ABOBA,prob,u,ΔW,ΔZ,p,rate_prototype, noise_rate_protot
 
   tmp = zero(u)
 
-  ABOBACache(utmp, dutmp, k, gtmp, noise, half, c₂, σ, tmp)
+  ABOBACache(utmp, dumid, dutmp, dunoise, k, gtmp, noise, half, c₂, σ, tmp)
 end
 
 
 
 
-mutable struct OBABOConstantCache{uType,uEltypeNoUnits, uCoeffType} <: StochasticDynamicalEqConstantCache
+mutable struct OBABOConstantCache{uType,rateNoiseType, uEltypeNoUnits, uCoeffType, uCoeffMType} <: StochasticDynamicalEqConstantCache
   k::uType
+  sig::rateNoiseType
   half::uEltypeNoUnits
   c₂::uCoeffType
-  σ::uCoeffType
+  σ::uCoeffMType
 end
-@cache struct OBABOCache{uType,uEltypeNoUnits,rateNoiseType,uCoeffType,uTypeCombined} <: StochasticDynamicalEqMutableCache
+
+@cache struct OBABOCache{uType,uEltypeNoUnits,rateNoiseType,uCoeffType, uCoeffMType,uTypeCombined} <: StochasticDynamicalEqMutableCache
   utmp::uType
+  dumid::uType
   dutmp::uType
+  dunoise::uType
   k::uType
   gtmp::rateNoiseType
   noise::uType
-  noisetmp::uType
   half::uEltypeNoUnits
   c₂::uCoeffType
-  σ::uCoeffType
+  σ::uCoeffMType
   tmp::uTypeCombined
 end
 
 function alg_cache(alg::OBABO,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{false}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
   k = zero(rate_prototype.x[1])
+  sig = zero(noise_rate_prototype)
   half=uEltypeNoUnits(1//2)
-  if typeof(alg.gamma) <: Number
-      c₂ = exp.(-alg.gamma*half*dt)
-      σ = sqrt.(1 .- alg.scale_noise*c₂.^2)
-  elseif typeof(alg.gamma) <: AbstractMatrix
+
+  if typeof(alg.gamma) <: AbstractMatrix
       c₂ = exp(-alg.gamma*half*dt)
       σ = cholesky(I - alg.scale_noise*c₂*transpose(c₂)).U
   else
       c₂ = exp.(-alg.gamma*half*dt)
       σ = sqrt.(1 .- alg.scale_noise*c₂.^2)
   end
    # if scale_noise == false, c2 = 1
-  OBABOConstantCache(k, half, c₂, σ)
+  OBABOConstantCache(k, sig, half, c₂, σ)
 end
 
 function alg_cache(alg::OBABO,prob,u,ΔW,ΔZ,p,rate_prototype, noise_rate_prototype,jump_rate_prototype,::Type{uEltypeNoUnits},::Type{uBottomEltypeNoUnits},::Type{tTypeNoUnits},uprev,f,t,dt,::Type{Val{true}}) where {uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits}
   dutmp = zero(u.x[1])
+  dumid = zero(u.x[1])
+  dunoise = zero(u.x[1])
   utmp = zero(u.x[2])
   k = zero(rate_prototype.x[1])
-  noise = zero(rate_prototype.x[1])
 
   gtmp = zero(noise_rate_prototype)
-  noisetmp = zero(rate_prototype.x[1])
+  noise = zero(rate_prototype.x[1])
 
 
   half = uEltypeNoUnits(1//2)
 
-  if typeof(alg.gamma) <: Number
-      c₂ = exp.(-alg.gamma*half*dt)
-      σ = sqrt.(1 .- alg.scale_noise*c₂.^2)
-  elseif typeof(alg.gamma) <: AbstractMatrix
+  if typeof(alg.gamma) <: AbstractMatrix
       c₂ = exp(-alg.gamma*half*dt)
       σ = cholesky(I - alg.scale_noise*c₂*transpose(c₂)).U
   else
@@ -187,5 +184,5 @@ function alg_cache(alg::OBABO,prob,u,ΔW,ΔZ,p,rate_prototype, noise_rate_protot
 
   tmp = zero(u)
 
-  OBABOCache(utmp, dutmp, k, noise, gtmp, noisetmp, half, c₂, σ, tmp)
+  OBABOCache(utmp, dumid, dutmp, dunoise, k, gtmp, noise, half, c₂, σ, tmp)
 end

src/perform_step/dynamical.jl

@@ -49,7 +49,6 @@ end
     du3 = c1.*du2 + c2.*noise
   end
 
-
   # A
   u = u2 + half*dt*du3
 
@@ -62,25 +61,25 @@ end
 
 @muladd function perform_step!(integrator,cache::BAOABCache)
   @unpack t,dt,sqdt,uprev,u,p,W,f = integrator
-  @unpack utmp, dutmp, k, gtmp, noise, half, c1, c2 = cache
+  @unpack utmp, dumid, dutmp, dunoise, k, gtmp, noise, half, c1, c2 = cache
   du1 = uprev.x[1]
   u1 = uprev.x[2]
 
   # B
-  @.. dutmp = du1 + half*dt*k
+  @.. dumid = du1 + half*dt*k
 
   # A
-  @.. utmp = u1 + half*dt*dutmp
+  @.. utmp = u1 + half*dt*dumid
 
   # O
   integrator.g(gtmp,utmp,p,t+dt*half)
   @.. noise = gtmp*W.dW / sqdt
   if typeof(c2) <: AbstractMatrix
-      mul!(dutmp,c1,dutmp)
-      mul!(noise,c2,noise)
-      @.. dutmp+= noise
+      mul!(dutmp,c1,dumid)
+      mul!(dunoise,c2,noise)
+      @.. dutmp+= dunoise
   else
-      @.. dutmp = c1*dutmp + c2*noise
+      @.. dutmp = c1*dumid + c2*noise
   end
 
   # A
@@ -119,7 +118,7 @@ end
 
 @muladd function perform_step!(integrator,cache::ABOBACache)
   @unpack t,dt,sqdt,uprev,u,p,W,f = integrator
-  @unpack utmp, dutmp, k, gtmp, noise, half, c₂, σ = cache
+  @unpack utmp, dumid, dutmp, dunoise, k, gtmp, noise, half, c₂, σ = cache
   du1 = uprev.x[1]
   u1 = uprev.x[2]
 
@@ -129,26 +128,26 @@ end
 
   # B
   f.f1(k,du1,utmp,p,t+dt)
-  @.. dutmp = du1 + half*dt*k
+  @.. dumid = du1 + half*dt*k
 
   # O
   integrator.g(gtmp,utmp,p,t+dt*half)
   @.. noise = gtmp*W.dW / sqdt
 
   if typeof(σ) <: AbstractMatrix
-      mul!(dutmp,c₂,dutmp)
-      mul!(noise,σ,noise)
-      @.. dutmp+=noise
+      mul!(dutmp,c₂,dumid)
+      mul!(dunoise,σ,noise)
+      @.. dutmp+=dunoise
   else
-      @.. dutmp = c₂*dutmp + σ*noise
+      @.. dutmp = c₂*dumid + σ*noise
   end
 
 
   # B
   @.. u.x[1] = dutmp + half*dt*k
 
   # A: xt+1
-  @.. u.x[2] = utmp + half*dt*dutmp
+  @.. u.x[2] = utmp + half*dt*u.x[1]
 end
 
 
@@ -161,6 +160,7 @@ function initialize!(integrator, cache::OBABOConstantCache)
 
   verify_f2(integrator.f.f2, du1, u1, p, t, integrator, cache)
   cache.k = integrator.f.f1(du1,u1,p,t)
+  cache.sig = integrator.g(u1,p,t)
 end
 
 function initialize!(integrator, cache::OBABOCache)
@@ -170,6 +170,87 @@ function initialize!(integrator, cache::OBABOCache)
 
   verify_f2(integrator.f.f2, cache.k, du1, u1, p, t, integrator, cache)
   integrator.f.f1(cache.k,du1,u1,p,t)
-
   integrator.g(cache.gtmp,u1,p,t)
 end
+
+
+@muladd function perform_step!(integrator,cache::OBABOConstantCache)
+  @unpack t,dt,sqdt,uprev,u,p,W,f = integrator
+  @unpack half, c₂, σ = cache
+  du1 = uprev.x[1]
+  u1 = uprev.x[2]
+
+  # O
+  noise = cache.sig.*W.dW ./ sqdt
+  if typeof(σ) <: AbstractMatrix || typeof(noise) <: Number
+    du2 = c₂*du1 + σ*noise
+  else
+    du2 = c₂.*du1 + σ.*noise
+  end
+
+  # B
+  dumid = du2 + half*dt*cache.k
+
+  # A
+  u = u1 + dt*dumid
+
+  cache.k = f.f1(dumid,u,p,t+dt)
+  # B
+  du3 = dumid + half*dt*cache.k
+
+  # O
+  cache.sig = integrator.g(u,p,t+dt)
+  noise = cache.sig.*W.dW ./ sqdt # That should be a second noise
+  if typeof(σ) <: AbstractMatrix || typeof(noise) <: Number
+    du = c₂*du3 + σ*noise
+  else
+    du = c₂.*du3 + σ.*noise
+  end
+
+  integrator.u = ArrayPartition((du, u))
+end
+
+
+@muladd function perform_step!(integrator,cache::OBABOCache)
+  @unpack t,dt,sqdt,uprev,u,p,W,f = integrator
+  @unpack utmp, dumid, dutmp, dunoise, k, gtmp, noise, half, c₂, σ = cache
+  du1 = uprev.x[1]
+  u1 = uprev.x[2]
+
+  # O
+  @.. noise = gtmp*W.dW / sqdt
+
+  if typeof(σ) <: AbstractMatrix
+      mul!(dutmp,c₂,du1)
+      mul!(dunoise,σ,noise)
+      @.. dutmp+=dunoise
+  else
+      @.. dutmp = c₂*du1 + σ*noise
+  end
+
+  # B
+
+  @.. dumid = dutmp + half*dt*k
+
+  # A: xt+1
+  @.. u.x[2] = u1 + dt*dumid
+
+
+  # B
+  f.f1(k,dumid,u.x[2],p,t+dt)
+  @..  dutmp = dumid + half*dt*k
+
+  # O
+  integrator.g(gtmp,u.x[2],p,t+dt)
+  @.. noise = gtmp*W.dW / sqdt  # That should be a second noise
+
+  if typeof(σ) <: AbstractMatrix
+      mul!(u.x[1],c₂,dutmp)
+      mul!(dunoise,σ,noise)
+      @.. u.x[1]+=dunoise
+  else
+      @.. u.x[1] = c₂*dutmp + σ*noise
+  end
+
+
+end