detect fewer brownians than equations, but noise still diag

Author: MasonProtter

src/systems/diffeqs/sdesystem.jl

@@ -245,11 +245,11 @@ function Base.:(==)(sys1::SDESystem, sys2::SDESystem)
         all(s1 == s2 for (s1, s2) in zip(get_systems(sys1), get_systems(sys2)))
 end
 
-function __num_isdiag(mat)
-    for i in axes(mat, 1), j in axes(mat, 2)
-        i == j || isequal(mat[i, j], 0) || return false
-    end
-    return true
+function __num_isdiag_noise(mat)
+    all(col -> count(!iszero, col) <= 1, eachcol(mat))
+end
+function __get_num_diag_noise(mat)
+    vec(sum(mat; dims = 2))
 end
 
 function generate_diffusion_function(sys::SDESystem, dvs = unknowns(sys),
@@ -258,7 +258,7 @@ function generate_diffusion_function(sys::SDESystem, dvs = unknowns(sys),
     if isdde
         eqs = delay_to_function(sys, eqs)
     end
-    if eqs isa AbstractMatrix && __num_isdiag(eqs)
+    if eqs isa AbstractMatrix && __num_isdiag_noise(eqs)
         eqs = diag(eqs)
     end
     u = map(x -> time_varying_as_func(value(x), sys), dvs)

src/systems/systems.jl

@@ -133,10 +133,11 @@ function __structural_simplify(sys::AbstractSystem, io = nothing; simplify = fal
             # 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)
+        elseif __num_isdiag_noise(sorted_g_rows)
+            # If each column of the noise matrix has either 0 or 1 non-zero entry, then this is "diagonal noise".
+            # In this case, the solver just takes a vector column of equations and it interprets that to
+            # mean that each noise process is independant
+            noise_eqs = __get_num_diag_noise(sorted_g_rows)
             is_scalar_noise = false
         else
             noise_eqs = sorted_g_rows

test/sdesystem.jl

@@ -725,12 +725,12 @@ end
 
 @testset "Non-diagonal noise check" begin
     @parameters σ ρ β
-    @variables x(t) y(t) z(t)
+    @variables x(tt) y(tt) z(tt)
     @brownian a b c
     eqs = [D(x) ~ σ * (y - x) + 0.1a * x + 0.1b * y,
         D(y) ~ x * (ρ - z) - y + 0.1b * y,
         D(z) ~ x * y - β * z + 0.1c * z]
-    @mtkbuild de = System(eqs, t)
+    @mtkbuild de = System(eqs, tt)
 
     u0map = [
         x => 1.0,
@@ -746,5 +746,32 @@ end
 
     prob = SDEProblem(de, u0map, (0.0, 100.0), parammap)
     # SOSRI only works for diagonal and scalar noise
+    @test_throws ErrorException solve(prob, SOSRI()).retcode==ReturnCode.Success
+    # ImplictEM does work for non-diagonal noise
     @test solve(prob, ImplicitEM()).retcode == ReturnCode.Success
 end
+
+@testset "Diagonal noise, less brownians than equations" begin
+    @parameters σ ρ β
+    @variables x(tt) y(tt) z(tt)
+    @brownian a b
+    eqs = [D(x) ~ σ * (y - x) + 0.1a * x,     # One brownian
+        D(y) ~ x * (ρ - z) - y + 0.1b * y, # Another brownian
+        D(z) ~ x * y - β * z]              # no brownians -- still diagonal
+    @mtkbuild de = System(eqs, tt)
+
+    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)
+    @test solve(prob, SOSRI()).retcode == ReturnCode.Success
+end