remove alloc in perform_step for nonscalar noise

Author: rmsrosa

src/caches/basic_method_caches.jl

@@ -73,19 +73,20 @@ function alg_cache(alg::RandomEM,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prot
 end
 
 struct RandomHeunConstantCache <: StochasticDiffEqConstantCache end
-@cache struct RandomHeunCache{uType,rateType} <: StochasticDiffEqMutableCache
+@cache struct RandomHeunCache{uType,rateType,randType} <: StochasticDiffEqMutableCache
   u::uType
   uprev::uType
   tmp::uType
   rtmp1::rateType
   rtmp2::rateType
+  wtmp::randType
 end
 
 alg_cache(alg::RandomHeun,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} = RandomHeunConstantCache()
 
 function alg_cache(alg::RandomHeun,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}
-  tmp = zero(u); rtmp1 = zero(rate_prototype); rtmp2 = zero(rate_prototype)
-  RandomHeunCache(u,uprev,tmp,rtmp1,rtmp2)
+  tmp = zero(u); rtmp1 = zero(rate_prototype); rtmp2 = zero(rate_prototype); wtmp = zero(ΔW)
+  RandomHeunCache(u,uprev,tmp,rtmp1,rtmp2,wtmp)
 end
 
 struct SimplifiedEMConstantCache <: StochasticDiffEqConstantCache end

src/perform_step/low_order.jl

@@ -124,17 +124,23 @@ end
 @muladd function perform_step!(integrator,cache::RandomHeunConstantCache)
     @unpack t,dt,uprev,u,W,p,f = integrator
     ftmp = integrator.f(uprev,p,t,W.curW)
-    tmp = @.. uprev + dt * ftmp 
-    u = uprev .+ (dt/2) .* (ftmp .+ integrator.f(tmp,p,t+dt, W.curW .+ W.dW))
+    tmp = @.. uprev + dt * ftmp
+    wtmp = @.. W.curW + W.dW
+    u = uprev .+ (dt/2) .* (ftmp .+ integrator.f(tmp,p,t+dt, wtmp))
     integrator.u = u
 end
 
 @muladd function perform_step!(integrator,cache::RandomHeunCache)
-    @unpack tmp, rtmp1, rtmp2 = cache
+    @unpack tmp, rtmp1, rtmp2, wtmp = cache
     @unpack t,dt,uprev,u,W,p,f = integrator
     integrator.f(rtmp1,uprev,p,t,W.curW)
     @.. tmp = uprev + dt * rtmp1
-    integrator.f(rtmp2,tmp,p,t+dt,W.curW+W.dW)
+    if W.dW isa Number
+        wtmp = W.curW + W.dW
+    else
+        @.. wtmp = W.curW + W.dW
+    end
+    integrator.f(rtmp2,tmp,p,t+dt,wtmp)
     @.. u = uprev + (dt/2) * (rtmp1 + rtmp2)
 end