From cc6abdef19e33335905149d2536f6a16d5ee7a9e Mon Sep 17 00:00:00 2001
From: Mason Protter <mason.protter@icloud.com>
Date: Thu, 1 Aug 2024 12:52:40 +0200
Subject: [PATCH] fix iteration order, add test that catches catalyst case

---
 src/systems/diffeqs/sdesystem.jl | 40 ++++++++++++++++++++++++++------
 test/sdesystem.jl                |  9 +++----
 2 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
index cefa5dc485..6a47af2868 100644
--- a/src/systems/diffeqs/sdesystem.jl
+++ b/src/systems/diffeqs/sdesystem.jl
@@ -246,30 +246,56 @@ function Base.:(==)(sys1::SDESystem, sys2::SDESystem)
 end
 
 function __num_isdiag_noise(mat)
-    for j in axes(mat, 2)
+    for i in axes(mat, 1)
         nnz = 0
-        for i in axes(mat, 1)
+        for j in axes(mat, 2)
             if !isequal(mat[i, j], 0)
                 nnz += 1
             end
         end
         if nnz > 1
-            return false
+            return (false)
         end
     end
     true
 end
 function __get_num_diag_noise(mat)
-    map(axes(mat, 2)) do j
-        for i in axes(mat, 1)
-            if !isequal(mat[i, j], 0)
-                return mat[i, j]
+    map(axes(mat, 1)) do i
+        for j in axes(mat, 2)
+            mij = mat[i, j]
+            if !isequal(mij, 0)
+                return mij
             end
         end
         0
     end
 end
 
+# function __num_isdiag_noise(mat)
+#     for j in axes(mat, 2)
+#         nnz = 0
+#         for i in axes(mat, 1)
+#             if !isequal(mat[i, j], 0)
+#                 nnz += 1
+#             end
+#         end
+#         if nnz > 1
+#             return false
+#         end
+#     end
+#     true
+# end
+# function __get_num_diag_noise(mat)
+#     map(axes(mat, 2)) do j
+#         for i in axes(mat, 1)
+#             if !isequal(mat[i, j], 0)
+#                 return mat[i, j]
+#             end
+#         end
+#         0
+#     end
+# end
+
 function generate_diffusion_function(sys::SDESystem, dvs = unknowns(sys),
         ps = full_parameters(sys); isdde = false, kwargs...)
     eqs = get_noiseeqs(sys)
diff --git a/test/sdesystem.jl b/test/sdesystem.jl
index 439b61f6b9..6e83d40d5f 100644
--- a/test/sdesystem.jl
+++ b/test/sdesystem.jl
@@ -726,10 +726,10 @@ end
 @testset "Non-diagonal noise check" begin
     @parameters σ ρ β
     @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]
+    @brownian a b c d e f
+    eqs = [D(x) ~ σ * (y - x) + 0.1a * x + d,
+        D(y) ~ x * (ρ - z) - y + 0.1b * y + e,
+        D(z) ~ x * y - β * z + 0.1c * z + f]
     @mtkbuild de = System(eqs, tt)
 
     u0map = [
@@ -749,6 +749,7 @@ end
     @test_throws ErrorException solve(prob, SOSRI()).retcode==ReturnCode.Success
     # ImplicitEM does work for non-diagonal noise
     @test solve(prob, ImplicitEM()).retcode == ReturnCode.Success
+    @test size(ModelingToolkit.get_noiseeqs(de)) == (3, 6)
 end
 
 @testset "Diagonal noise, less brownians than equations" begin