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

Non-diagonal problems with a single noise channel were routed into the
diagonal code paths by length(W.dW) == 1 shortcuts and crashed with
MethodError (oop) or DimensionMismatch (iip). Routing now depends only
on is_diagonal_noise and the W.dW type.

The general-noise paths drew Rademacher variables through W.rng, which
does not exist on NoiseWrapper. They now go through an init_chi! helper
that finds the generator on the wrapped source process, and the
commented J_qr cross terms from Abdulle, Vilmart and Zygalakis (2013),
eq. (17), are activated so SROCK2 keeps weak order 2 for m > 1.

The KomBurSROCK2 stage corrections from the original PR SciML#3468 already
landed in SciML#3883 and are unchanged here.

Author: singhharsh1708

lib/StochasticDiffEqROCK/Project.toml

@@ -1,7 +1,7 @@
 name = "StochasticDiffEqROCK"
 uuid = "db241ea8-0e6b-4abc-8f2d-1adff2294fd9"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "2.0.2"
+version = "2.0.3"
 
 [deps]
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
@@ -16,6 +16,7 @@ StochasticDiffEqCore = "19c5a474-6cd1-4a5f-be79-46dc34e54d7f"
 
 [extras]
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
+DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SDEProblemLibrary = "c72e72a9-a271-4b2b-8966-303ed956772e"
@@ -40,14 +41,15 @@ SciMLBase = "3.39"
 StaticArrays = "1.9.18"
 StochasticDiffEqCore = "2"
 DiffEqDevTools = "3"
+DiffEqNoiseProcess = "5.30"
 Random = "1"
 SDEProblemLibrary = "0.1, 1"
 Test = "<0.0.1, 1"
 SafeTestsets = "0.1.0"
 julia = "1.10"
 
 [targets]
-test = ["Test", "Pkg", "SafeTestsets", "DiffEqDevTools", "SDEProblemLibrary", "Random", "SciMLTesting"]
+test = ["Test", "Pkg", "SafeTestsets", "DiffEqDevTools", "DiffEqNoiseProcess", "SDEProblemLibrary", "Random", "SciMLTesting"]
 
 [sources.OrdinaryDiffEqCore]
 path = "../OrdinaryDiffEqCore"

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ᵢ)
@@ -300,7 +302,11 @@ end
         for i in 1:length(W.dW)
             WikJ = W.dW[i]
             WikJ2 = vec_χ[i]
-            WikRange = 1 // 2 .* (W.dW .* WikJ .- (1:length(W.dW) .== i) .* abs(dt)) #.- (1:length(W.dW) .> i) .* dt .* vec_χ .+ (1:length(W.dW) .< i) .* dt .* WikJ2)
+            WikRange = 1 // 2 .* (
+                W.dW .* WikJ .- (1:length(W.dW) .== i) .* abs(dt) .-
+                    (1:length(W.dW) .> i) .* abs(dt) .* vec_χ .+
+                    (1:length(W.dW) .< i) .* abs(dt) .* WikJ2
+            )
             uₓ = Gₛ * WikRange
             WikRange = 1 // 2 .* (1:length(W.dW) .== i)
             uᵢ₋₂ = uᵢ + uₓ
@@ -332,8 +338,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 +367,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 +422,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ᵢ)
@@ -458,7 +462,12 @@ end
             WikJ2 = vec_χ[i]
             dwrange = 1:length(W.dW)
             abs_dt = abs(dt)
-            @.. WikRange = 1 // 2 * (W.dW * WikJ - (dwrange == i) * abs_dt) #+ (dwrange < i) * dt * WikJ2 - (dwrange > i) * dt * vec_χ)
+            @.. WikRange = 1 // 2 *
+                (
+                W.dW * WikJ - (dwrange == i) * abs_dt -
+                    (dwrange > i) * abs_dt * vec_χ +
+                    (dwrange < i) * abs_dt * WikJ2
+            )
             mul!(uₓ, Gₛ, WikRange)
             @.. uᵢ₋₂ = uᵢ + uₓ
             @.. WikRange = 1 // 2 * (dwrange == i)
@@ -542,14 +551,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 +642,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 +991,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 +1177,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 +1236,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 +1253,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 +1300,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 +1442,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 +1469,9 @@ 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 +1514,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ₛ₋₃
@@ -1711,7 +1723,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
 
@@ -1806,7 +1818,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
 
@@ -1840,3 +1852,18 @@ end
 
     integrator.u = u
 end
+
+# Fill `vec_χ` with independent Rademacher (±1) samples drawn from the RNG of the
+# integrator's noise process. Kept as a scalar loop so it works for both `Vector`
+# and `MArray` storage without requiring the Random stdlib as a dependency.
+function init_χ!(vec_χ, W)
+    r = rng(W)
+    for i in eachindex(vec_χ)
+        vec_χ[i] = 2 * (rand(r) < 1 // 2) - 1
+    end
+    return vec_χ
+end
+
+# `NoiseWrapper` does not carry its own `rng`; the generator lives on the wrapped
+# source process, so accessing `W.rng` directly crashes (#3188).
+rng(W) = hasfield(typeof(W), :rng) ? W.rng : W.source.rng

lib/StochasticDiffEqROCK/test/runtests.jl

@@ -24,6 +24,10 @@ if TEST_GROUP == "ALL" || TEST_GROUP == "Core"
     @time @safetestset "KomBurSROCK2 non-diagonal noise" begin
         include("kombursrock2_nondiag_tests.jl")
     end
+
+    @time @safetestset "SROCK non-diagonal noise regressions" begin
+        include("srock_nondiag_tests.jl")
+    end
 end
 
 if TEST_GROUP == "ALL" || TEST_GROUP == "SROCKC2WeakConvergence"
@@ -32,6 +36,12 @@ if TEST_GROUP == "ALL" || TEST_GROUP == "SROCKC2WeakConvergence"
     end
 end
 
+if TEST_GROUP == "ALL" || TEST_GROUP == "SROCK2NonDiagonalConvergence"
+    @time @safetestset "SROCK2 Non-Diagonal Weak Convergence Tests" begin
+        include("weak_convergence/weak_srock2_nondiag.jl")
+    end
+end
+
 # Run QA tests (Aqua, JET) - skip on pre-release Julia
 if (TEST_GROUP == "QA" || TEST_GROUP == "ALL") && isempty(VERSION.prerelease)
     activate_qa_env()

lib/StochasticDiffEqROCK/test/srock_nondiag_tests.jl

@@ -0,0 +1,55 @@
+using StochasticDiffEqROCK, DiffEqNoiseProcess, Test, Random
+import SciMLBase
+
+# Non-square noise (n = 2 states, m = 1 Brownian motion) used to crash with
+# MethodError / DimensionMismatch because a length-one W.dW shortcut routed these
+# problems into the diagonal-noise code path, see SciML/OrdinaryDiffEq.jl#3170.
+@testset "m=1 non-diagonal noise does not crash" begin
+    n, m = 2, 1
+    f_oop(u, p, t) = zeros(n)
+    g_oop(u, p, t) = ones(n, m)
+    f_iip(du, u, p, t) = (du .= 0.0)
+    g_iip(du, u, p, t) = (du .= 1.0)
+    u0 = zeros(n)
+    nrp = ones(n, m)
+    prob_oop = SDEProblem(f_oop, g_oop, u0, (0.0, 1.0); noise_rate_prototype = nrp)
+    prob_iip = SDEProblem(f_iip, g_iip, u0, (0.0, 1.0); noise_rate_prototype = nrp)
+
+    @testset "$(nameof(typeof(alg)))" for alg in (
+            SROCK2(), SROCKEM(), SROCKC2(), KomBurSROCK2(), TangXiaoSROCK2(),
+        )
+        for prob in (prob_oop, prob_iip)
+            Random.seed!(100)
+            sol = solve(prob, alg, dt = 1 / 2^4)
+            @test SciMLBase.successful_retcode(sol)
+            @test all(isfinite, sol.u[end])
+        end
+    end
+end
+
+# NoiseWrapper keeps its RNG on the wrapped source process, so the former direct
+# `W.rng` access crashed for non-diagonal problems, see SciML/OrdinaryDiffEq.jl#3188.
+@testset "SROCK2 with NoiseWrapper (non-diagonal noise)" begin
+    f_iip(du, u, p, t) = (du .= -0.5 .* u)
+    function g_iip(du, u, p, t)
+        du[1, 1] = 0.1 * u[1]
+        du[1, 2] = 0.05 * u[2]
+        du[2, 1] = 0.05 * u[1]
+        du[2, 2] = 0.1 * u[2]
+        return nothing
+    end
+    u0 = [1.0, 1.0]
+    tspan = (0.0, 1.0)
+    Random.seed!(100)
+    prob = SDEProblem(f_iip, g_iip, u0, tspan; noise_rate_prototype = zeros(2, 2))
+    sol = solve(prob, SROCK2(), dt = 1 / 2^6, save_noise = true)
+    @test SciMLBase.successful_retcode(sol)
+
+    prob_wrapped = SDEProblem(
+        f_iip, g_iip, u0, tspan;
+        noise = NoiseWrapper(sol.W), noise_rate_prototype = zeros(2, 2)
+    )
+    sol_wrapped = solve(prob_wrapped, SROCK2(), dt = 1 / 2^6)
+    @test SciMLBase.successful_retcode(sol_wrapped)
+    @test all(isfinite, sol_wrapped.u[end])
+end

lib/StochasticDiffEqROCK/test/test_groups.toml

@@ -9,5 +9,10 @@ versions = ["1"]
 timeout = 240
 local_only = true
 
+[SROCK2NonDiagonalConvergence]
+versions = ["1"]
+timeout = 240
+local_only = true
+
 [QA]
 versions = ["lts", "1"]