From 9a556d7d99daf69f367e175f83440610e16535b5 Mon Sep 17 00:00:00 2001
From: Sikorski <sikorski@zib.de>
Date: Fri, 30 Sep 2022 11:34:07 +0200
Subject: [PATCH 1/6] fix SROCK2 for NoiseWrapper noise

fixes #505
---
 src/perform_step/SROCK_perform_step.jl | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/perform_step/SROCK_perform_step.jl b/src/perform_step/SROCK_perform_step.jl
index fafb4e825..4908b6a5a 100644
--- a/src/perform_step/SROCK_perform_step.jl
+++ b/src/perform_step/SROCK_perform_step.jl
@@ -332,8 +332,12 @@ end
 
   sqrt_dt   = sqrt(abs(dt))
   if gen_prob
-    vec_χ .= 1//2 .+ oftype(W.dW, rand(W.rng, length(W.dW)))
-    @.. vec_χ = 2*floor(vec_χ) - 1
+    if hasfield(W, :rng)
+        rand!(W.rng, vec_χ)
+    else
+        rand!(vec_χ)
+    end
+    @.. vec_χ = 2*floor(vec_χ + 1//2) - 1
   end
 
   μ = recf[start]  # here κ = 0

From 31265446599504151173a7e31aee672de03fcb4b Mon Sep 17 00:00:00 2001
From: Sikorski <sikorski@zib.de>
Date: Fri, 30 Sep 2022 12:01:25 +0200
Subject: [PATCH 2/6] use W.rng for vec_\chi in all SROCK solvers

---
 src/perform_step/SROCK_perform_step.jl | 35 ++++++++++++++++++--------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/perform_step/SROCK_perform_step.jl b/src/perform_step/SROCK_perform_step.jl
index 4908b6a5a..867ac3adc 100644
--- a/src/perform_step/SROCK_perform_step.jl
+++ b/src/perform_step/SROCK_perform_step.jl
@@ -183,7 +183,10 @@ end
   @unpack recf, recf2, mα, mσ, mτ = cache
 
   gen_prob = !((is_diagonal_noise(integrator.sol.prob)) || (typeof(W.dW) <: Number) || (length(W.dW) == 1))
-  gen_prob && (vec_χ = 2 .* floor.(false .* W.dW .+ 1//2 .+ oftype(W.dW, rand(W.rng,length(W.dW)))) .- true)
+  if gen_prob
+    vec_χ = similar(W.dW)
+    init_χ!(vec_χ, W)
+  end
 
   alg = unwrap_alg(integrator, true)
   alg.eigen_est === nothing ? maxeig!(integrator, cache) : alg.eigen_est(integrator)
@@ -316,6 +319,9 @@ end
   @unpack recf, recf2, mα, mσ, mτ = cache.constantcache
   ccache = cache.constantcache
   gen_prob = !((is_diagonal_noise(integrator.sol.prob)) || (typeof(W.dW) <: Number) || (length(W.dW) == 1))
+  if gen_prob
+    init_χ!(vec_χ, W)
+  end
 
   alg = unwrap_alg(integrator, true)
   alg.eigen_est === nothing ? maxeig!(integrator, cache) : alg.eigen_est(integrator)
@@ -331,14 +337,6 @@ end
   τ = 1//2*((1-α)^2) + 2*α*(1-α)*mσ[deg_index] + (α^2)*(mσ[deg_index]*(mσ[deg_index]+mτ[deg_index]))
 
   sqrt_dt   = sqrt(abs(dt))
-  if gen_prob
-    if hasfield(W, :rng)
-        rand!(W.rng, vec_χ)
-    else
-        rand!(vec_χ)
-    end
-    @.. vec_χ = 2*floor(vec_χ + 1//2) - 1
-  end
 
   μ = recf[start]  # here κ = 0
   tᵢ = t + α*dt*μ
@@ -1110,6 +1108,10 @@ end
   @unpack recf, mσ, mτ, mδ = cache
 
   gen_prob = !((is_diagonal_noise(integrator.sol.prob)) || (typeof(W.dW) <: Number) || (length(W.dW) == 1))
+  if gen_prob
+    vec_χ = similar(W.dW)
+    init_χ!(vec_χ, W)
+  end
 
   alg = unwrap_alg(integrator, true)
   alg.eigen_est === nothing ? maxeig!(integrator, cache) : alg.eigen_est(integrator)
@@ -1125,7 +1127,6 @@ end
   τ = mτ[deg_index]
 
   sqrt_dt   = sqrt(abs(dt))
-  (gen_prob) && (vec_χ = 2 .* floor.( 1//2 .+ false .* W.dW .+ rand(length(W.dW))) .- 1)
 
   tᵢ₋₂ = t
   uᵢ₋₂ = uprev
@@ -1299,6 +1300,9 @@ end
 
   ccache = cache.constantcache
   gen_prob = !((is_diagonal_noise(integrator.sol.prob)) || (typeof(W.dW) <: Number) || (length(W.dW) == 1))
+  if gen_prob
+    init_χ!(vec_χ, W)
+  end
 
   alg = unwrap_alg(integrator, true)
   alg.eigen_est === nothing ? maxeig!(integrator, cache) : alg.eigen_est(integrator)
@@ -1314,7 +1318,7 @@ end
   τ = mτ[deg_index]
 
   sqrt_dt   = sqrt(abs(dt))
-  (gen_prob) && (vec_χ .= 2 .* floor.(1//2 .+ false .* vec_χ .+ rand(length(vec_χ))) .- 1)
+
 
   tᵢ₋₂ = t
   @.. uᵢ₋₂ = uprev
@@ -1680,3 +1684,12 @@ end
 
   integrator.u = u
 end
+
+function init_χ!(vec_χ, W)
+    if hasfield(W, :rng)
+        rand!(W.rng, vec_χ)
+    else
+        rand!(vec_χ)
+    end
+    @.. vec_χ = 2*floor(vec_χ + 1//2) - 1
+end

From 0fba9cf600e9cdba3e3256f0490ccfd44ef26232 Mon Sep 17 00:00:00 2001
From: Sikorski <sikorski@zib.de>
Date: Fri, 30 Sep 2022 12:12:55 +0200
Subject: [PATCH 3/6] extract rng from NoiseWrapper

---
 src/perform_step/SROCK_perform_step.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/perform_step/SROCK_perform_step.jl b/src/perform_step/SROCK_perform_step.jl
index 867ac3adc..486e1adfc 100644
--- a/src/perform_step/SROCK_perform_step.jl
+++ b/src/perform_step/SROCK_perform_step.jl
@@ -1688,6 +1688,8 @@ end
 function init_χ!(vec_χ, W)
     if hasfield(W, :rng)
         rand!(W.rng, vec_χ)
+    elseif isa(W, NoiseWrapper)
+        rand!(W.W.rng, vec_χ)
     else
         rand!(vec_χ)
     end

From b67ed07ba2e8fea71c1815b554823577266db585 Mon Sep 17 00:00:00 2001
From: Sikorski <sikorski@zib.de>
Date: Fri, 30 Sep 2022 14:45:21 +0200
Subject: [PATCH 4/6] add forgotten typeof

---
 src/perform_step/SROCK_perform_step.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/perform_step/SROCK_perform_step.jl b/src/perform_step/SROCK_perform_step.jl
index 486e1adfc..e8e2bf5b9 100644
--- a/src/perform_step/SROCK_perform_step.jl
+++ b/src/perform_step/SROCK_perform_step.jl
@@ -1686,7 +1686,7 @@ end
 end
 
 function init_χ!(vec_χ, W)
-    if hasfield(W, :rng)
+    if hasfield(typeof(W), :rng)
         rand!(W.rng, vec_χ)
     elseif isa(W, NoiseWrapper)
         rand!(W.W.rng, vec_χ)

From 21fa59440fcc6a12ec22ecc7f45bbc38599a3108 Mon Sep 17 00:00:00 2001
From: Sikorski <sikorski@zib.de>
Date: Fri, 30 Sep 2022 14:56:03 +0200
Subject: [PATCH 5/6] access  instead of

---
 src/perform_step/SROCK_perform_step.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/perform_step/SROCK_perform_step.jl b/src/perform_step/SROCK_perform_step.jl
index e8e2bf5b9..88d92f42d 100644
--- a/src/perform_step/SROCK_perform_step.jl
+++ b/src/perform_step/SROCK_perform_step.jl
@@ -1689,7 +1689,7 @@ function init_χ!(vec_χ, W)
     if hasfield(typeof(W), :rng)
         rand!(W.rng, vec_χ)
     elseif isa(W, NoiseWrapper)
-        rand!(W.W.rng, vec_χ)
+        rand!(W.source.rng, vec_χ)
     else
         rand!(vec_χ)
     end

From 9650ab7d254b91e2a0057f23cbb2e9730a5a5dea Mon Sep 17 00:00:00 2001
From: Sikorski <sikorski@zib.de>
Date: Fri, 30 Sep 2022 15:45:44 +0200
Subject: [PATCH 6/6] switch to dispatch on rng(<:AbstractNoiseProcess)

---
 src/perform_step/SROCK_perform_step.jl | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/src/perform_step/SROCK_perform_step.jl b/src/perform_step/SROCK_perform_step.jl
index 88d92f42d..9d70d5279 100644
--- a/src/perform_step/SROCK_perform_step.jl
+++ b/src/perform_step/SROCK_perform_step.jl
@@ -1686,12 +1686,9 @@ end
 end
 
 function init_χ!(vec_χ, W)
-    if hasfield(typeof(W), :rng)
-        rand!(W.rng, vec_χ)
-    elseif isa(W, NoiseWrapper)
-        rand!(W.source.rng, vec_χ)
-    else
-        rand!(vec_χ)
-    end
+    rand!(rng(W), vec_χ)
     @.. vec_χ = 2*floor(vec_χ + 1//2) - 1
 end
+
+rng(W::DiffEqNoiseProcess.AbstractNoiseProcess) = W.rng
+rng(W::NoiseWrapper) = W.source.rng