Fixes

Author: Harsh Singh

lib/StochasticDiffEqROCK/Project.toml

@@ -17,6 +17,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"
@@ -45,7 +46,7 @@ SafeTestsets = "0.1.0"
 julia = "1.10"
 
 [targets]
-test = ["Test", "Pkg", "SafeTestsets", "DiffEqDevTools", "SDEProblemLibrary", "Random"]
+test = ["Test", "Pkg", "SafeTestsets", "DiffEqDevTools", "DiffEqNoiseProcess", "SDEProblemLibrary", "Random"]
 
 [sources.OrdinaryDiffEqCore]
 path = "../OrdinaryDiffEqCore"

lib/StochasticDiffEqROCK/src/perform_step/SROCK_perform_step.jl

@@ -1426,7 +1426,7 @@ end
 @muladd function perform_step!(integrator, cache::KomBurSROCK2Cache)
     (;
         utmp, uᵢ₋₁, uᵢ₋₂, k, yₛ₋₁, yₛ₋₂, yₛ₋₃, SXₛ₋₁, SXₛ₋₂,
-        SXₛ₋₃, Gₛ, Xₛ₋₁, Xₛ₋₂, Xₛ₋₃, vec_χ,
+        SXₛ₋₃, Gₛ, Xₛ₋₁, Xₛ₋₂, Xₛ₋₃, vec_χ, WikRange,
     ) = cache
     (; t, dt, uprev, u, W, p, f) = integrator
     (; recf, mσ, mτ, mδ) = cache.constantcache
@@ -1566,8 +1566,7 @@ end
             ttmp = tᵢ₋₂ + C₁
             integrator.f.g(Gₛ, utmp, p, ttmp)
             WikRange .= 1 .* (1:length(W.dW) .== i)
-            # @.. @view(Xₛ₋₂[:,i]) =  @view(Gₛ[:,i])
-            @.. Xₛ₋₂ = Gₛ * W.dW
+            @view(Xₛ₋₂[:, i]) .= @view(Gₛ[:, i])
         end
         mul!(SXₛ₋₂, Xₛ₋₂, W.dW)
         @.. u += μₛ₋₂ * yₛ₋₂ + 3 // 8 * SXₛ₋₂
@@ -1581,14 +1580,13 @@ end
             # @.. utmp = uᵢ₋₁ + μₛ₋₃*yₛ₋₃ + δ₁*yₛ₋₂ - 1//6*W.dW[i]*@view(Xₛ₋₃[:,i]) - 1//2*W.dW[i]*@view(Xₛ₋₂[:,i]) + 1//4*SXₛ₋₃ + 3//4*SXₛ₋₂
             @.. utmp = uᵢ₋₁ + μₛ₋₃ * yₛ₋₃ + δ₁ * yₛ₋₂ + 1 // 4 * SXₛ₋₃ + 3 // 4 * SXₛ₋₂
             mul!(SXₛ₋₁, Xₛ₋₃, WikRange)
-            @.. utmp += 1 // 6 * SXₛ₋₁
+            @.. utmp -= 1 // 6 * SXₛ₋₁
             mul!(SXₛ₋₁, Xₛ₋₂, WikRange)
-            @.. utmp += 1 // 2 * SXₛ₋₁
+            @.. utmp -= 1 // 2 * SXₛ₋₁
             ttmp = tᵢ₋₁ + μₛ₋₃ + δ₁
             integrator.f.g(Gₛ, utmp, p, ttmp)
             WikRange .= 1 .* (1:length(W.dW) .== i)
-            # @.. @view(Xₛ₋₁[:,i]) = @view(Gₛ[:,i])
-            @.. Xₛ₋₁ = Gₛ * WikRange
+            @view(Xₛ₋₁[:, i]) .= @view(Gₛ[:, i])
         end
         mul!(SXₛ₋₁, Xₛ₋₁, W.dW)
         @.. u += (σ - τ) * dt * yₛ₋₁ + 3 // 8 * SXₛ₋₁

lib/StochasticDiffEqROCK/test/convergence/nondiagonal_convergence.jl

@@ -0,0 +1,70 @@
+using StochasticDiffEqROCK, DiffEqDevTools, DiffEqNoiseProcess, Test, Random
+
+# Non-diagonal SDE: 2-component system driven by 1 Brownian motion (n=2, m=1)
+# du_i = μ u_i dt + σ u_i dW    (same scalar Brownian motion for both components)
+# Weak solution: E[u_i(T)] = u_i(0) exp(μ T)
+
+const μ_test = -0.5
+const σ_test = 0.1
+
+f_nd_oop(u, p, t) = μ_test .* u
+g_nd_oop(u, p, t) = σ_test .* reshape(u, length(u), 1)
+
+f_nd_iip!(du, u, p, t) = (du .= μ_test .* u)
+function g_nd_iip!(du, u, p, t)
+    for i in axes(du, 1)
+        du[i, 1] = σ_test * u[i]
+    end
+end
+
+# Analytic (Itô): u_i(t) = u_i(0) exp((μ - σ²/2)t + σ W(t))
+analytic_nd(u0, p, t, W) = u0 .* exp.((μ_test - σ_test^2 / 2) * t .+ σ_test .* W[1])
+
+u0 = [1.0, 1.0]
+tspan = (0.0, 1.0)
+
+prob_oop = SDEProblem(
+    SDEFunction(f_nd_oop, g_nd_oop; analytic = analytic_nd),
+    u0, tspan;
+    noise_rate_prototype = zeros(2, 1)
+)
+prob_iip = SDEProblem(
+    SDEFunction(f_nd_iip!, g_nd_iip!; analytic = analytic_nd),
+    u0, tspan;
+    noise_rate_prototype = zeros(2, 1)
+)
+
+Random.seed!(100)
+dts = 1 .// 2 .^ (6:-1:2)
+
+@testset "SROCK2 non-diagonal OOP weak convergence (order ~2)" begin
+    sim = test_convergence(dts, prob_oop, SROCK2(), trajectories = Int(5e4),
+                           save_everystep = false, weak_timeseries_errors = false)
+    @test abs(sim.𝒪est[:weak_final] - 2.0) < 0.5
+end
+
+@testset "SROCK2 non-diagonal IIP weak convergence (order ~2)" begin
+    sim = test_convergence(dts, prob_iip, SROCK2(), trajectories = Int(5e4),
+                           save_everystep = false, weak_timeseries_errors = false)
+    @test abs(sim.𝒪est[:weak_final] - 2.0) < 0.5
+end
+
+# KomBurSROCK2 smoke test: non-diagonal noise must not crash with DimensionMismatch
+@testset "KomBurSROCK2 non-diagonal IIP does not crash" begin
+    prob_smoke = SDEProblem(f_nd_iip!, g_nd_iip!, u0, tspan;
+                            noise_rate_prototype = zeros(2, 1))
+    sol = solve(prob_smoke, KomBurSROCK2(), dt = 0.01)
+    @test sol.retcode == ReturnCode.Success
+end
+
+# Smoke test: NoiseWrapper must not crash
+@testset "SROCK2 NoiseWrapper does not crash" begin
+    prob_base = SDEProblem(f_nd_iip!, g_nd_iip!, u0, tspan;
+                           noise_rate_prototype = zeros(2, 1))
+    sol_base = solve(prob_base, SROCK2(), dt = 0.01, save_noise = true)
+    W_wrap = NoiseWrapper(sol_base.W)
+    prob_wrap = SDEProblem(f_nd_iip!, g_nd_iip!, u0, tspan;
+                           noise = W_wrap, noise_rate_prototype = zeros(2, 1))
+    sol_wrap = solve(prob_wrap, SROCK2(), dt = 0.01)
+    @test sol_wrap.retcode == ReturnCode.Success
+end

lib/StochasticDiffEqROCK/test/runtests.jl

@@ -22,3 +22,9 @@ if TEST_GROUP == "ALL" || TEST_GROUP == "SROCKC2WeakConvergence"
         include("weak_convergence/weak_srockc2.jl")
     end
 end
+
+if TEST_GROUP == "ALL" || TEST_GROUP == "NonDiagonalConvergence"
+    @time @safetestset "Non-Diagonal Noise Convergence Tests (SROCK2, #3188, #3170)" begin
+        include("convergence/nondiagonal_convergence.jl")
+    end
+end

lib/StochasticDiffEqROCK/test/test_groups.toml

@@ -6,5 +6,8 @@ versions = ["1"]
 runner = ["self-hosted", "Linux", "X64", "high-memory"]
 timeout = 240
 
+[NonDiagonalConvergence]
+versions = ["1"]
+
 [QA]
 versions = ["1"]