Merge pull request #2882 from MasonProtter/detect-diagonal-noise

Detect diagonal noise in `SDESystem`

Project.toml

@@ -12,6 +12,7 @@ ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
+DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -72,6 +73,7 @@ DataStructures = "0.17, 0.18"
 DeepDiffs = "1"
 DiffEqBase = "6.103.0"
 DiffEqCallbacks = "2.16, 3"
+DiffEqNoiseProcess = "5"
 DiffRules = "0.1, 1.0"
 Distributed = "1"
 Distributions = "0.23, 0.24, 0.25"

src/ModelingToolkit.jl

@@ -22,6 +22,7 @@ using DiffEqCallbacks
 using Graphs
 import ExprTools: splitdef, combinedef
 import OrderedCollections
+using DiffEqNoiseProcess: DiffEqNoiseProcess, WienerProcess
 
 using SymbolicIndexingInterface
 using LinearAlgebra, SparseArrays, LabelledArrays

src/systems/abstractsystem.jl

@@ -655,7 +655,8 @@ for prop in [:eqs
              :solved_unknowns
              :split_idxs
              :parent
-             :index_cache]
+             :index_cache
+             :is_scalar_noise]
     fname_get = Symbol(:get_, prop)
     fname_has = Symbol(:has_, prop)
     @eval begin

src/systems/diffeqs/sdesystem.jl

@@ -128,13 +128,18 @@ struct SDESystem <: AbstractODESystem
     The hierarchical parent system before simplification.
     """
     parent::Any
+    """
+    Signal for whether the noise equations should be treated as a scalar process. This should only
+    be `true` when `noiseeqs isa Vector`. 
+    """
+    is_scalar_noise::Bool
 
     function SDESystem(tag, deqs, neqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed,
             tgrad,
             jac,
             ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connector_type,
             cevents, devents, parameter_dependencies, metadata = nothing, gui_metadata = nothing,
-            complete = false, index_cache = nothing, parent = nothing;
+            complete = false, index_cache = nothing, parent = nothing, is_scalar_noise = false;
             checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckComponents) > 0
             check_independent_variables([iv])
@@ -146,6 +151,9 @@ struct SDESystem <: AbstractODESystem
                 throw(ArgumentError("Noise equations ill-formed. Number of rows must match number of drift equations. size(neqs,1) = $(size(neqs,1)) != length(deqs) = $(length(deqs))"))
             end
             check_equations(equations(cevents), iv)
+            if is_scalar_noise && neqs isa AbstractMatrix
+                throw(ArgumentError("Noise equations ill-formed. Received a matrix of noise equations of size $(size(neqs)), but `is_scalar_noise` was set to `true`. Scalar noise is only compatible with an `AbstractVector` of noise equations."))
+            end
         end
         if checks == true || (checks & CheckUnits) > 0
             u = __get_unit_type(dvs, ps, iv)
@@ -154,7 +162,7 @@ struct SDESystem <: AbstractODESystem
         new(tag, deqs, neqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac,
             Wfact, Wfact_t, name, systems, defaults, connector_type, cevents, devents,
-            parameter_dependencies, metadata, gui_metadata, complete, index_cache, parent)
+            parameter_dependencies, metadata, gui_metadata, complete, index_cache, parent, is_scalar_noise)
     end
 end
 
@@ -173,7 +181,11 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs::AbstractArray, iv, dv
         discrete_events = nothing,
         parameter_dependencies = nothing,
         metadata = nothing,
-        gui_metadata = nothing)
+        gui_metadata = nothing,
+        complete = false,
+        index_cache = nothing,
+        parent = nothing,
+        is_scalar_noise = false)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     iv′ = value(iv)
@@ -210,7 +222,8 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs::AbstractArray, iv, dv
     SDESystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
         deqs, neqs, iv′, dvs′, ps′, tspan, var_to_name, ctrl′, observed, tgrad, jac,
         ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connector_type,
-        cont_callbacks, disc_callbacks, parameter_dependencies, metadata, gui_metadata; checks = checks)
+        cont_callbacks, disc_callbacks, parameter_dependencies, metadata, gui_metadata,
+        complete, index_cache, parent, is_scalar_noise; checks = checks)
 end
 
 function SDESystem(sys::ODESystem, neqs; kwargs...)
@@ -225,6 +238,7 @@ function Base.:(==)(sys1::SDESystem, sys2::SDESystem)
         isequal(nameof(sys1), nameof(sys2)) &&
         isequal(get_eqs(sys1), get_eqs(sys2)) &&
         isequal(get_noiseeqs(sys1), get_noiseeqs(sys2)) &&
+        isequal(get_is_scalar_noise(sys1), get_is_scalar_noise(sys2)) &&
         _eq_unordered(get_unknowns(sys1), get_unknowns(sys2)) &&
         _eq_unordered(get_ps(sys1), get_ps(sys2)) &&
         all(s1 == s2 for (s1, s2) in zip(get_systems(sys1), get_systems(sys2)))
@@ -616,16 +630,24 @@ function DiffEqBase.SDEProblem{iip, specialize}(
     sparsenoise === nothing && (sparsenoise = get(kwargs, :sparse, false))
 
     noiseeqs = get_noiseeqs(sys)
+    is_scalar_noise = get_is_scalar_noise(sys)
     if noiseeqs isa AbstractVector
         noise_rate_prototype = nothing
+        if is_scalar_noise
+            noise = WienerProcess(0.0, 0.0, 0.0)
+        else
+            noise = nothing
+        end
     elseif sparsenoise
         I, J, V = findnz(SparseArrays.sparse(noiseeqs))
         noise_rate_prototype = SparseArrays.sparse(I, J, zero(eltype(u0)))
+        noise = nothing
     else
         noise_rate_prototype = zeros(eltype(u0), size(noiseeqs))
+        noise = nothing
     end
 
-    SDEProblem{iip}(f, u0, tspan, p; callback = cbs,
+    SDEProblem{iip}(f, u0, tspan, p; callback = cbs, noise,
         noise_rate_prototype = noise_rate_prototype, kwargs...)
 end
 
@@ -693,22 +715,32 @@ function SDEProblemExpr{iip}(sys::SDESystem, u0map, tspan,
     sparsenoise === nothing && (sparsenoise = get(kwargs, :sparse, false))
 
     noiseeqs = get_noiseeqs(sys)
+    is_scalar_noise = get_is_scalar_noise(sys)
     if noiseeqs isa AbstractVector
         noise_rate_prototype = nothing
+        if is_scalar_noise
+            noise = WienerProcess(0.0, 0.0, 0.0)
+        else
+            noise = nothing
+        end
     elseif sparsenoise
         I, J, V = findnz(SparseArrays.sparse(noiseeqs))
         noise_rate_prototype = SparseArrays.sparse(I, J, zero(eltype(u0)))
+        noise = nothing
     else
         T = u0 === nothing ? Float64 : eltype(u0)
         noise_rate_prototype = zeros(T, size(get_noiseeqs(sys)))
+        noise = nothing
     end
     ex = quote
         f = $f
         u0 = $u0
         tspan = $tspan
         p = $p
         noise_rate_prototype = $noise_rate_prototype
-        SDEProblem(f, u0, tspan, p; noise_rate_prototype = noise_rate_prototype,
+        noise = $noise
+        SDEProblem(
+            f, u0, tspan, p; noise_rate_prototype = noise_rate_prototype, noise = noise,
             $(kwargs...))
     end
     !linenumbers ? Base.remove_linenums!(ex) : ex

src/systems/systems.jl

@@ -127,9 +127,23 @@ function __structural_simplify(sys::AbstractSystem, io = nothing; simplify = fal
             g_row > size(g, 1) && continue
             @views copyto!(sorted_g_rows[i, :], g[g_row, :])
         end
-
-        return SDESystem(full_equations(ode_sys), sorted_g_rows,
+        # Fix for https://github.com/SciML/ModelingToolkit.jl/issues/2490
+        if sorted_g_rows isa AbstractMatrix && size(sorted_g_rows, 2) == 1
+            # If there's only one brownian variable referenced across all the equations,
+            # we get a Nx1 matrix of noise equations, which is a special case known as scalar noise
+            noise_eqs = sorted_g_rows[:, 1]
+            is_scalar_noise = true
+        elseif isdiag(sorted_g_rows)
+            # If the noise matrix is diagonal, then the solver just takes a vector column of equations
+            # and it interprets that as diagonal noise.
+            noise_eqs = diag(sorted_g_rows)
+            is_scalar_noise = false
+        else
+            noise_eqs = sorted_g_rows
+            is_scalar_noise = false
+        end
+        return SDESystem(full_equations(ode_sys), noise_eqs,
             get_iv(ode_sys), unknowns(ode_sys), parameters(ode_sys);
-            name = nameof(ode_sys))
+            name = nameof(ode_sys), is_scalar_noise)
     end
 end

test/dde.jl

@@ -78,7 +78,7 @@ sol = solve(prob, RKMil())
 eqs = [D(x(t)) ~ a * x(t) + b * x(t - τ) + c + (α * x(t) + γ) * η]
 @mtkbuild sys = System(eqs, t)
 @test equations(sys) == [D(x(t)) ~ a * x(t) + b * x(t - τ) + c]
-@test isequal(ModelingToolkit.get_noiseeqs(sys), [α * x(t) + γ;;])
+@test isequal(ModelingToolkit.get_noiseeqs(sys), [α * x(t) + γ])
 prob_mtk = SDDEProblem(sys, [x(t) => 1.0 + t], tspan; constant_lags = (τ,));
 @test_nowarn sol_mtk = solve(prob_mtk, RKMil())
 

test/sdesystem.jl

@@ -657,3 +657,68 @@ ps = @SVector[p => 5.0, d => 0.5]
 sprob = SDEProblem(sys, u0, tspan, ps)
 @test sprob.f.g(sprob.u0, sprob.p, sprob.tspan[1]) isa SVector{2, Float64}
 @test_nowarn solve(sprob, ImplicitEM())
+
+let
+    @parameters σ ρ β
+    @variables x(t) y(t) z(t)
+    @brownian a
+    eqs = [D(x) ~ σ * (y - x) + 0.1a * x,
+        D(y) ~ x * (ρ - z) - y + 0.1a * y,
+        D(z) ~ x * y - β * z + 0.1a * z]
+
+    @mtkbuild de = System(eqs, t)
+
+    u0map = [
+        x => 1.0,
+        y => 0.0,
+        z => 0.0
+    ]
+
+    parammap = [
+        σ => 10.0,
+        β => 26.0,
+        ρ => 2.33
+    ]
+    prob = SDEProblem(de, u0map, (0.0, 100.0), parammap)
+    # TODO: re-enable this when we support scalar noise
+    @test solve(prob, SOSRI()).retcode == ReturnCode.Success
+end
+
+let # test to make sure that scalar noise always receive the same kicks
+    @variables x(t) y(t)
+    @brownian a
+    eqs = [D(x) ~ a,
+        D(y) ~ a]
+
+    @mtkbuild de = System(eqs, t)
+    prob = SDEProblem(de, [x => 0, y => 0], (0.0, 10.0), [])
+    sol = solve(prob, SOSRI())
+    @test sol[end][1] == sol[end][2]
+end
+
+let # test that diagonal noise is correctly handled
+    @parameters σ ρ β
+    @variables x(t) y(t) z(t)
+    @brownian a b c
+    eqs = [D(x) ~ σ * (y - x) + 0.1a * x,
+        D(y) ~ x * (ρ - z) - y + 0.1b * y,
+        D(z) ~ x * y - β * z + 0.1c * z]
+
+    @mtkbuild de = System(eqs, t)
+
+    u0map = [
+        x => 1.0,
+        y => 0.0,
+        z => 0.0
+    ]
+
+    parammap = [
+        σ => 10.0,
+        β => 26.0,
+        ρ => 2.33
+    ]
+
+    prob = SDEProblem(de, u0map, (0.0, 100.0), parammap)
+    # SOSRI only works for diagonal and scalar noise
+    @test solve(prob, SOSRI()).retcode == ReturnCode.Success
+end