Fix SROCK2 NoiseWrapper and non-square noise (#3188, #3170)

- Add init_χ! helper and use correct RNG dispatch for NoiseWrapper
- Remove broken commutative noise opt-out (length=1) for non-diagonal noise
- This ensures matrix noise shapes correctly route to generalized noise path
- Add DiffEqNoiseProcess to deps

Author: Harsh Singh

lib/StochasticDiffEqROCK/src/caches/SROCK_caches.jl

@@ -161,7 +161,7 @@ function alg_cache(
     uᵢ₋₁ = zero(u)
     uᵢ₋₂ = zero(u)
     Gₛ = zero(noise_rate_prototype)
-    if (!alg.strong_order_1 || is_diagonal_noise(prob) || ΔW isa Number || length(ΔW) == 1)
+    if (!alg.strong_order_1 || is_diagonal_noise(prob) || ΔW isa Number )
         Gₛ₁ = Gₛ
     else
         Gₛ₁ = zero(noise_rate_prototype)
@@ -295,7 +295,7 @@ function alg_cache(
     uᵢ₋₁ = zero(u)
     uᵢ₋₂ = zero(u)
     Gₛ = zero(noise_rate_prototype)
-    if ΔW isa Number || length(ΔW) == 1 || is_diagonal_noise(prob)
+    if ΔW isa Number  || is_diagonal_noise(prob)
         Gₛ₁ = Gₛ
     else
         Gₛ₁ = zero(noise_rate_prototype)
@@ -377,7 +377,7 @@ function alg_cache(
     Xₛ₋₃ = zero(noise_rate_prototype)
     vec_χ = false .* vec(ΔW)
     WikRange = false .* vec(ΔW)
-    if ΔW isa Number || length(ΔW) == 1 || is_diagonal_noise(prob)
+    if ΔW isa Number  || is_diagonal_noise(prob)
         Gₛ = Xₛ₋₁
         SXₛ₋₁ = utmp
         SXₛ₋₂ = utmp

lib/StochasticDiffEqROCK/src/perform_step/SROCK_perform_step.jl

@@ -194,10 +194,12 @@ end
     (; recf, recf2, mα, mσ, mτ) = cache
 
     gen_prob = !(
-        (is_diagonal_noise(integrator.sol.prob)) || (W.dW isa Number) ||
-            (length(W.dW) == 1)
+        (is_diagonal_noise(integrator.sol.prob)) || (W.dW isa Number) 
     )
-    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)
@@ -265,7 +267,7 @@ end
     # Now uᵢ₋₂ = uₛ₋₂, uᵢ₋₁ = uₛ₋₁, uᵢ = uₛ
     # Similarly tᵢ₋₂ = tₛ₋₂, tᵢ₋₁ = tₛ₋₁, tᵢ = tₛ
 
-    if (W.dW isa Number) || (length(W.dW) == 1) || is_diagonal_noise(integrator.sol.prob)
+    if (W.dW isa Number)  || is_diagonal_noise(integrator.sol.prob)
         Gₛ = integrator.f.g(uᵢ₋₁, p, tᵢ₋₁)
         u += Gₛ .* W.dW
         Gₛ = integrator.f.g(uᵢ, p, tᵢ)
@@ -332,8 +334,7 @@ end
     (; recf, recf2, mα, mσ, mτ) = cache.constantcache
     ccache = cache.constantcache
     gen_prob = !(
-        (is_diagonal_noise(integrator.sol.prob)) || (W.dW isa Number) ||
-            (length(W.dW) == 1)
+        (is_diagonal_noise(integrator.sol.prob)) || (W.dW isa Number) 
     )
 
     alg = unwrap_alg(integrator, true)
@@ -362,8 +363,7 @@ 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
+        init_χ!(vec_χ, W)
     end
 
     μ = recf[start]  # here κ = 0
@@ -418,7 +418,7 @@ end
     # Now uᵢ₋₂ = uₛ₋₂, uᵢ₋₁ = uₛ₋₁, uᵢ = uₛ
     # Similarly tᵢ₋₂ = tₛ₋₂, tᵢ₋₁ = tₛ₋₁, tᵢ = tₛ
 
-    if (W.dW isa Number) || (length(W.dW) == 1) || is_diagonal_noise(integrator.sol.prob)
+    if (W.dW isa Number)  || is_diagonal_noise(integrator.sol.prob)
         integrator.f.g(Gₛ, uᵢ₋₁, p, tᵢ₋₁)
         @.. u += Gₛ * W.dW
         integrator.f.g(Gₛ, uᵢ, p, tᵢ)
@@ -542,14 +542,14 @@ end
     end
 
     Gₛ = integrator.f.g(u, p, tᵢ)
-    if (W.dW isa Number) || (length(W.dW) == 1) || is_diagonal_noise(integrator.sol.prob)
+    if (W.dW isa Number)  || is_diagonal_noise(integrator.sol.prob)
         u += Gₛ .* W.dW
     else
         u += Gₛ * W.dW
     end
 
     if integrator.alg.strong_order_1
-        if (W.dW isa Number) || (length(W.dW) == 1) ||
+        if (W.dW isa Number)  ||
                 (is_diagonal_noise(integrator.sol.prob))
             uᵢ₋₂ = @. 1 // 2 * Gₛ * (W.dW^2 - abs(dt))
             tmp = @. u + uᵢ₋₂
@@ -633,15 +633,15 @@ end
     end
 
     integrator.f.g(Gₛ, u, p, tᵢ)
-    if (W.dW isa Number) || (length(W.dW) == 1) || is_diagonal_noise(integrator.sol.prob)
+    if (W.dW isa Number)  || is_diagonal_noise(integrator.sol.prob)
         @.. u += Gₛ * W.dW
     else
         mul!(uᵢ₋₁, Gₛ, W.dW)
         u += uᵢ₋₁
     end
 
     if integrator.alg.strong_order_1
-        if (W.dW isa Number) || (length(W.dW) == 1) ||
+        if (W.dW isa Number)  ||
                 (is_diagonal_noise(integrator.sol.prob))
             @.. uᵢ₋₂ = 1 // 2 * Gₛ * (W.dW^2 - abs(dt))
             @.. tmp = u + uᵢ₋₂
@@ -982,7 +982,7 @@ end
         end
     end
 
-    if (W.dW isa Number) || (length(W.dW) == 1)
+    if (W.dW isa Number) 
         Gₛ = integrator.f.g(Û₁, p, t̂₁)
         uₓ += Gₛ * W.dW
 
@@ -1168,7 +1168,7 @@ end
         end
     end
 
-    if (W.dW isa Number) || (length(W.dW) == 1) || is_diagonal_noise(integrator.sol.prob)
+    if (W.dW isa Number)  || is_diagonal_noise(integrator.sol.prob)
         integrator.f.g(Gₛ, Û₁, p, t̂₁)
         @.. uₓ += Gₛ * W.dW
 
@@ -1227,8 +1227,7 @@ end
     (; recf, mσ, mτ, mδ) = cache
 
     gen_prob = !(
-        (is_diagonal_noise(integrator.sol.prob)) || (W.dW isa Number) ||
-            (length(W.dW) == 1)
+        (is_diagonal_noise(integrator.sol.prob)) || (W.dW isa Number) 
     )
 
     alg = unwrap_alg(integrator, true)
@@ -1245,7 +1244,10 @@ end
     τ = mτ[deg_index]
 
     sqrt_dt = sqrt(abs(dt))
-    (gen_prob) && (vec_χ = 2 .* floor.(1 // 2 .+ false .* W.dW .+ rand(length(W.dW))) .- 1)
+    if gen_prob
+        vec_χ = similar(W.dW)
+        init_χ!(vec_χ, W)
+    end
 
     tᵢ₋₂ = t
     uᵢ₋₂ = uprev
@@ -1289,7 +1291,7 @@ end
     tᵢ₋₁ += θₛ₋₃ * (tᵢ₋₁ - tᵢ₋₂)
     tᵢ₋₂ = ttmp
 
-    if W.dW isa Number || length(W.dW) == 1 || is_diagonal_noise(integrator.sol.prob)
+    if W.dW isa Number  || is_diagonal_noise(integrator.sol.prob)
         # stage s-3
         yₛ₋₃ = integrator.f(uᵢ₋₁, p, tᵢ₋₁)
         utmp = uᵢ₋₁ + μₛ₋₃ * yₛ₋₃
@@ -1431,8 +1433,7 @@ end
 
     ccache = cache.constantcache
     gen_prob = !(
-        (is_diagonal_noise(integrator.sol.prob)) || (W.dW isa Number) ||
-            (length(W.dW) == 1)
+        (is_diagonal_noise(integrator.sol.prob)) || (W.dW isa Number) 
     )
 
     alg = unwrap_alg(integrator, true)
@@ -1459,7 +1460,7 @@ end
     τ = mτ[deg_index]
 
     sqrt_dt = sqrt(abs(dt))
-    (gen_prob) && (vec_χ .= 2 .* floor.(1 // 2 .+ false .* vec_χ .+ rand(length(vec_χ))) .- 1)
+    if gen_prob; init_χ!(vec_χ, W); end
 
     tᵢ₋₂ = t
     @.. uᵢ₋₂ = uprev
@@ -1502,7 +1503,7 @@ end
     tᵢ₋₁ += θₛ₋₃ * (tᵢ₋₁ - tᵢ₋₂)
     tᵢ₋₂ = ttmp
 
-    if W.dW isa Number || length(W.dW) == 1 || is_diagonal_noise(integrator.sol.prob)
+    if W.dW isa Number  || is_diagonal_noise(integrator.sol.prob)
         # stage s-3
         integrator.f(yₛ₋₃, uᵢ₋₁, p, tᵢ₋₁)
         @.. utmp = uᵢ₋₁ + μₛ₋₃ * yₛ₋₃
@@ -1713,7 +1714,7 @@ end
     uᵢ₋₂ = integrator.f(uᵢ₋₂, p, tᵢ₋₂)
     u += dt * (σ + τ) * uᵢ₋₂
 
-    if (W.dW isa Number) || (length(W.dW) == 1) || is_diagonal_noise(integrator.sol.prob)
+    if (W.dW isa Number)  || is_diagonal_noise(integrator.sol.prob)
         Gₛ = integrator.f.g(uᵢ₋₁, p, tᵢ₋₁)
         u += Gₛ .* W.dW
 
@@ -1808,7 +1809,7 @@ end
     integrator.f(k, uᵢ₋₂, p, tᵢ₋₂)
     @.. u += dt * (σ + τ) * k
 
-    if (W.dW isa Number) || (length(W.dW) == 1) || is_diagonal_noise(integrator.sol.prob)
+    if (W.dW isa Number)  || is_diagonal_noise(integrator.sol.prob)
         integrator.f.g(Gₛ, uᵢ₋₁, p, tᵢ₋₁)
         @.. u += Gₛ * W.dW
 
@@ -1842,3 +1843,11 @@ end
 
     integrator.u = u
 end
+
+function init_χ!(vec_χ, W)
+    rand!(rng(W), vec_χ)
+    @.. vec_χ = 2 * floor(vec_χ + 1 // 2) - 1
+end
+
+rng(W::DiffEqNoiseProcess.AbstractNoiseProcess) = W.rng
+rng(W::DiffEqNoiseProcess.NoiseWrapper) = W.source.rng