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

lib/StochasticDiffEqROCK/Project.toml

@@ -1,13 +1,14 @@
 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"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
@@ -16,8 +17,8 @@ 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"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -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", "SciMLTesting"]
 
 [sources.OrdinaryDiffEqCore]
 path = "../OrdinaryDiffEqCore"

lib/StochasticDiffEqROCK/src/StochasticDiffEqROCK.jl

@@ -19,6 +19,8 @@ import DiffEqBase: full_cache, rand_cache, ratenoise_cache
 import MuladdMacro: @muladd
 import SciMLBase
 
+import Random
+
 using LinearAlgebra
 using StaticArrays
 using RecursiveArrayTools

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,9 @@ 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)
+    vec_χ = gen_prob ? init_χ!(similar(W.dW), W) : nothing
 
     alg = unwrap_alg(integrator, true)
     alg.eigen_est === nothing ? maxeig!(integrator, cache) : alg.eigen_est(integrator)
@@ -265,7 +264,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 +299,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 +335,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 +364,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 +419,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 +459,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 +548,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 +639,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 +988,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 +1174,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 +1233,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 +1250,7 @@ end
     τ = mτ[deg_index]
 
     sqrt_dt = sqrt(abs(dt))
-    (gen_prob) && (vec_χ = 2 .* floor.(1 // 2 .+ false .* W.dW .+ rand(length(W.dW))) .- 1)
+    vec_χ = gen_prob ? init_χ!(similar(W.dW), W) : nothing
 
     tᵢ₋₂ = t
     uᵢ₋₂ = uprev
@@ -1289,7 +1294,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 +1436,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 +1463,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 +1508,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 +1717,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 +1812,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 +1846,24 @@ end
 
     integrator.u = u
 end
+
+# Fill `vec_χ` with independent Rademacher (±1) samples drawn from the RNG of the
+# integrator's noise process. Drawn in bulk and mapped by broadcast rather than by a
+# scalar loop, so the routine does not scalar-index `vec_χ` and stays usable for array
+# types that forbid it.
+function init_χ!(vec_χ, W)
+    Random.rand!(rng(W), vec_χ)
+    one_χ = one(eltype(vec_χ))
+    vec_χ .= ifelse.(vec_χ .< 1 // 2, -one_χ, one_χ)
+    return vec_χ
+end
+
+# `NoiseWrapper` does not carry its own `rng`; the generator lives on the wrapped
+# source process, which may itself be a wrapper (#3188). Processes that replay a
+# recorded path (`NoiseGrid`) have no generator at all, so fall back to the global
+# one rather than erroring.
+function rng(W)
+    hasfield(typeof(W), :rng) && return W.rng
+    hasfield(typeof(W), :source) && return rng(W.source)
+    return Random.default_rng()
+end

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()