Fix noise kwarg propagation for SDEProblem(f::SDESystem)

Project.toml

@@ -58,6 +58,7 @@ SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"

src/problems/sdeproblem.jl

@@ -77,7 +77,20 @@ end
         t = tspan !== nothing ? tspan[1] : tspan, check_length, eval_expression,
         eval_module, check_compatibility, sparse, expression, kwargs...)
 
-    noise, noise_rate_prototype = calculate_noise_and_rate_prototype(sys, u0; sparsenoise)
+    # Only calculate noise and noise_rate_prototype if not provided by user
+    if !haskey(kwargs, :noise) && !haskey(kwargs, :noise_rate_prototype)
+        noise, noise_rate_prototype = calculate_noise_and_rate_prototype(sys, u0; sparsenoise)
+    elseif !haskey(kwargs, :noise)
+        noise, _ = calculate_noise_and_rate_prototype(sys, u0; sparsenoise)
+        noise_rate_prototype = kwargs[:noise_rate_prototype]
+    elseif !haskey(kwargs, :noise_rate_prototype)
+        _, noise_rate_prototype = calculate_noise_and_rate_prototype(sys, u0; sparsenoise)
+        noise = kwargs[:noise]
+    else
+        noise = kwargs[:noise]
+        noise_rate_prototype = kwargs[:noise_rate_prototype]
+    end
+
     kwargs = process_kwargs(sys; expression, callback, eval_expression, eval_module,
         op, kwargs...)
 

test/sdesystem.jl

@@ -1,5 +1,6 @@
 using ModelingToolkit, StaticArrays, LinearAlgebra
 using StochasticDiffEq, OrdinaryDiffEq, SparseArrays
+using DiffEqNoiseProcess: NoiseWrapper
 using Random, Test
 using Setfield
 using Statistics
@@ -953,3 +954,47 @@ end
     @test ModelingToolkit.isbrownian(p)
     @test ModelingToolkit.isbrownian(q)
 end
+
+@testset "noise kwarg propagation (issue #3664)" begin
+    @parameters σ ρ β
+    @variables x(tt) y(tt) z(tt)
+
+    u0 = [1.0, 0.0, 0.0]
+    T = (0.0, 5.0)
+
+    eqs = [D(x) ~ σ * (y - x),
+        D(y) ~ x * (ρ - z) - y,
+        D(z) ~ x * y - β * z]
+    noiseeqs = [3.0,
+        3.0,
+        3.0]
+    @mtkbuild sde_lorentz = SDESystem(eqs, noiseeqs, tt, [x, y, z], [σ, ρ, β])
+    parammap = [σ, ρ, β] .=> [10, 28.0, 8 / 3]
+
+    # Test that user-provided noise is respected
+    Random.seed!(1)
+    noise1 = StochasticDiffEq.RealWienerProcess(0.0, 0.0, 0.0; save_everystep = true)
+    u0_dict = Dict(unknowns(sde_lorentz) .=> u0)
+    prob1 = SDEProblem(sde_lorentz, merge(u0_dict, Dict(parammap)), T; noise = noise1)
+    sol1 = solve(prob1, SRIW1())
+
+    # Verify noise was actually used (curW should be modified)
+    @test noise1.curW != 0.0
+
+    # Test that using the same noise via NoiseWrapper gives deterministic results
+    noise2 = NoiseWrapper(noise1)
+    prob2 = SDEProblem(sde_lorentz, merge(u0_dict, Dict(parammap)), T; noise = noise2)
+    sol2 = solve(prob2, SRIW1())
+
+    # Same noise should give same results
+    @test sol1.u[end] ≈ sol2.u[end]
+
+    # Test that without providing noise, different results are obtained
+    Random.seed!(1)
+    prob3 = SDEProblem(sde_lorentz, merge(u0_dict, Dict(parammap)), T)
+    Random.seed!(2)
+    prob4 = SDEProblem(sde_lorentz, merge(u0_dict, Dict(parammap)), T)
+    sol3 = solve(prob3, SRIW1(), seed = 1)
+    sol4 = solve(prob4, SRIW1(), seed = 2)
+    @test !(sol3.u[end] ≈ sol4.u[end])
+end