push in the right direction, still some issues with adaptivity

Author: frankschae

src/caches/implicit_split_step_caches.jl

@@ -64,7 +64,7 @@ function alg_cache(alg::ISSEulerHeun, prob, u, ΔW, ΔZ, p, rate_prototype, nois
   gtmp2 = zero(rate_prototype)
 
   if is_diagonal_noise(prob)
-    gtmp3 = gtmp2
+    gtmp3 = copy(gtmp2)
     dW_cache = nothing
   else
     gtmp3 = zero(noise_rate_prototype)

src/perform_step/implicit_split_step.jl

@@ -1,25 +1,24 @@
-@muladd function perform_step!(integrator,
-                               cache::Union{ISSEMConstantCache,
-                                            ISSEulerHeunConstantCache})
-  @unpack t,dt,uprev,u,p,f = integrator
+@muladd function perform_step!(integrator, cache::Union{ISSEMConstantCache, ISSEulerHeunConstantCache})
+
+  @unpack t, dt, uprev, u, p, f = integrator
   @unpack nlsolver = cache
   alg = unwrap_alg(integrator, true)
   theta = alg.theta
-  alg.symplectic ? a = dt/2 : a = theta*dt
+  alg.symplectic ? a = dt / 2 : a = theta * dt
   OrdinaryDiffEq.markfirststage!(nlsolver)
 
   # TODO: Stochastic extrapolants?
   u = uprev
 
   repeat_step = false
 
-  L = integrator.g(uprev,p,t)
-  ftmp = integrator.f(uprev,p,t)
+  L = integrator.g(uprev, p, t)
+  ftmp = integrator.f(uprev, p, t)
 
   if alg.symplectic
     z = zero(u) # constant extrapolation, justified by ODE IM
   else
-    z = dt*ftmp # linear extrapolation
+    z = dt * ftmp # linear extrapolation
   end
   nlsolver.z = z
 
@@ -29,7 +28,7 @@
     #u = uprev + z/2
     tmp = uprev
   else
-    tmp = uprev + dt*(1-theta)*ftmp
+    tmp = uprev + dt * (1 - theta) * ftmp
   end
   nlsolver.tmp = tmp
 
@@ -39,56 +38,71 @@
   if alg.symplectic
     u = tmp + z
   else
-    u = tmp + theta*z
+    u = tmp + theta * z
   end
 
-  gtmp = L.*integrator.W.dW
+  if !is_diagonal_noise(integrator.sol.prob)
+    gtmp = L * integrator.W.dW
+  else
+    gtmp = L .* integrator.W.dW
+  end
 
   if typeof(cache) <: ISSEulerHeunConstantCache
     utilde = u + gtmp
-    gtmp = ((integrator.g(utilde,p,t) + L)/2)*integrator.W.dW
+    if !is_diagonal_noise(integrator.sol.prob)
+      gtmp = ((integrator.g(utilde, p, t) + L) / 2) * integrator.W.dW
+    else
+      gtmp = ((integrator.g(utilde, p, t) + L) / 2) .* integrator.W.dW
+    end
   end
 
   u += gtmp
 
   if integrator.opts.adaptive
     if has_Wfact(f)
       # This means the Jacobian was never computed!
-      J = f.jac(uprev,p,t)
+      J = f.jac(uprev, p, t)
     else
       J = OrdinaryDiffEq.calc_J(integrator, nlsolver.cache)
     end
-    Ed = dt*(dt*J*ftmp)/2
+    Ed = dt * (dt * J * ftmp) / 2
 
     if typeof(cache) <: ISSEMConstantCache
       K = @.. uprev + dt * ftmp
-      utilde =  @.. K + integrator.sqdt * L
-      ggprime = (integrator.g(utilde,p,t).-L)./(integrator.sqdt)
-      En = ggprime .* (integrator.W.dW.^2 .- dt)./2
+      utilde = @.. K + integrator.sqdt * L
+      ggprime = (integrator.g(utilde, p, t) .- L) ./ (integrator.sqdt)
+      if !is_diagonal_noise(integrator.sol.prob)
+        En = ggprime * (integrator.W.dW .^ 2 .- dt) ./ 2
+      else
+        En = ggprime .* (integrator.W.dW .^ 2 .- dt) ./ 2
+      end
     elseif typeof(cache) <: ISSEulerHeunConstantCache
-      utilde = uprev + L*integrator.sqdt
-      ggprime = (integrator.g(utilde,p,t).-L)./(integrator.sqdt)
-      En = ggprime.*(integrator.W.dW.^2)./2
+      utilde = @.. uprev + L * integrator.sqdt
+      ggprime = (integrator.g(utilde, p, t) .- L) ./ (integrator.sqdt)
+      if !is_diagonal_noise(integrator.sol.prob)
+        En = ggprime * (integrator.W.dW .^ 2) ./ 2
+      else
+        En = ggprime .* (integrator.W.dW .^ 2) ./ 2
+      end
     end
 
     resids = calculate_residuals(Ed, En, uprev, u, integrator.opts.abstol,
-                                 integrator.opts.reltol, integrator.opts.delta,
-                                 integrator.opts.internalnorm,t)
-    integrator.EEst = integrator.opts.internalnorm(resids,t)
+      integrator.opts.reltol, integrator.opts.delta,
+      integrator.opts.internalnorm, t)
+    integrator.EEst = integrator.opts.internalnorm(resids, t)
   end
 
   integrator.u = u
 end
 
-@muladd function perform_step!(integrator, cache::Union{ISSEMCache,
-                                                        ISSEulerHeunCache})
-  @unpack t,dt,uprev,u,p,f = integrator
-  @unpack gtmp,gtmp2,dW_cache,nlsolver,k,dz = cache
-  @unpack z,tmp = nlsolver
+@muladd function perform_step!(integrator, cache::Union{ISSEMCache, ISSEulerHeunCache})
+  @unpack t, dt, uprev, u, p, f = integrator
+  @unpack gtmp, gtmp2, dW_cache, nlsolver, k, dz = cache
+  @unpack z, tmp = nlsolver
 
   J = (OrdinaryDiffEq.isnewton(nlsolver) ? nlsolver.cache.J : nothing)
   alg = unwrap_alg(integrator, true)
-  alg.symplectic ? a = dt/2 : a = alg.theta*dt
+  alg.symplectic ? a = dt / 2 : a = alg.theta * dt
   dW = integrator.W.dW
   mass_matrix = integrator.f.mass_matrix
   theta = alg.theta
@@ -97,20 +111,28 @@ end
   repeat_step = false
 
   if integrator.success_iter > 0 && !integrator.u_modified && alg.extrapolant == :interpolant
-    current_extrapolant!(u,t+dt,integrator)
+    current_extrapolant!(u, t + dt, integrator)
   elseif alg.extrapolant == :linear
-    @.. u = uprev + integrator.fsalfirst*dt
+    @.. u = uprev + integrator.fsalfirst * dt
   else # :constant
-    copyto!(u,uprev)
+    copyto!(u, uprev)
   end
 
-  integrator.f(tmp,uprev,p,t)
+  integrator.f(tmp, uprev, p, t)
   integrator.g(gtmp, uprev, p, t)
 
   if alg.symplectic
     @.. z = zero(eltype(u)) # Justified by ODE solvers, constrant extrapolation when IM
   else
-    @.. z = dt*tmp # linear extrapolation
+    @.. z = dt * tmp # linear extrapolation
+  end
+
+  # Handle noise computations
+
+  if is_diagonal_noise(integrator.sol.prob)
+    @.. gtmp2 = gtmp * dW
+  else
+    mul!(gtmp2, gtmp, dW)
   end
 
   ###
@@ -144,8 +166,8 @@ end
     if typeof(cache) <: ISSEMCache
 
       if !is_diagonal_noise(integrator.sol.prob)
-        integrator.g(gtmp2, z, p, t)
-        g_sized2 = norm(gtmp2, 2)
+        integrator.g(gtmp, z, p, t)
+        g_sized2 = norm(gtmp, 2)
         @.. dW_cache = dW .^ 2 - dt
         diff_tmp = integrator.opts.internalnorm(dW_cache, t)
         En = (g_sized2 - g_sized) / (2integrator.sqdt) * diff_tmp
@@ -191,18 +213,9 @@ end
   end
 
   ##############################################################################
-
-  # Handle noise computations
-
-  if is_diagonal_noise(integrator.sol.prob)
-    @.. gtmp2 = gtmp * dW
-  else
-    mul!(gtmp2, gtmp, dW)
-  end
-
   if typeof(cache) <: ISSEulerHeunCache
     gtmp3 = cache.gtmp3
-    @.. z = uprev + gtmp2
+    @.. z = u + gtmp2
     integrator.g(gtmp3, z, p, t)
     @.. gtmp = (gtmp3 + gtmp) / 2
     if is_diagonal_noise(integrator.sol.prob)
@@ -211,7 +224,6 @@ end
       mul!(gtmp2, gtmp, dW)
     end
   end
-
   @.. u += gtmp2
 
   ##############################################################################