don't specialize on scalar noise, only handle diagonal noise

SciML · Jul 21, 2024 · 489be39 · 489be39
1 parent 950a6a2
commit 489be39
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 13 deletions.
diff --git a/src/systems/systems.jl b/src/systems/systems.jl
@@ -126,14 +126,16 @@ function __structural_simplify(sys::AbstractSystem, io = nothing; simplify = fal
             @views copyto!(sorted_g_rows[i, :], g[g_row, :])
         end
         # Fix for https://github.com/SciML/ModelingToolkit.jl/issues/2490
-        noise_eqs = if all(row -> count(!iszero, row) == 1, eachrow(sorted_g_rows)) # Does each row have only one non-zero entry?
-            # then the noise is actually diagonal! make vector of non-zero entries.
-            # This happens when the user uses multiple `@brownian`s but never mixes them
-            map(eachrow(sorted_g_rows)) do row
-                only(filter(!iszero, row))
-            end
+        noise_eqs = if isdiag(sorted_g_rows)
+            # If the noise matrix is diagonal, then we just give solver just takes a vector column of equations
+            # and it interprets that as diagonal noise.
+            diag(sorted_g_rows)
         elseif sorted_g_rows isa AbstractMatrix && size(sorted_g_rows, 2) == 1
-            sorted_g_rows[:, 1] # Take a vector slice so solver knows there's no mixing
+            ##-------------------------------------------------------------------------------
+            ## TODO: re-enable this code once we add a way to signal that the noise is scalar
+            # sorted_g_rows[:, 1] # Take a vector slice so solver knows there's no mixing
+            ##-------------------------------------------------------------------------------
+            sorted_g_rows
         else
             sorted_g_rows
         end

diff --git a/test/sdesystem.jl b/test/sdesystem.jl
@@ -679,18 +679,18 @@ let
         β => 26.0,
         ρ => 2.33
     ]
-
     prob = SDEProblem(de, u0map, (0.0, 100.0), parammap)
-    @test solve(prob, SOSRI()).retcode == ReturnCode.Success
+    # TODO: re-enable this when we support scalar noise
+    @test_broken solve(prob, SOSRI()).retcode == ReturnCode.Success
 end
 
 let
     @parameters σ ρ β
     @variables x(t) y(t) z(t)
-    @brownian a b
-    eqs = [D(x) ~ σ * (y - x) + 0.1b * x,
-        D(y) ~ x * (ρ - z) - y + 0.1a * y,
-        D(z) ~ x * y - β * z + 0.1b * z]
+    @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)