Merge pull request #142 from SciML/faster_results

Faster equation recovery for implict and explicit sindy results

Author: AlCap23

src/sindy/isindy.jl

@@ -146,7 +146,7 @@ function ImplicitSparseIdentificationResult(coeff::AbstractArray, equations::Bas
     ps = [p; p_]
 
     Ŷ = b_(X, ps, t)
-    training_error = norm.(eachrow(Y-Ŷ), 2)
+    training_error = [norm(Y[i, :]-Ŷ[i,:], 2) for i in 1:size(Ŷ, 1)]
     aicc = similar(training_error)
 
     for i in 1:length(aicc)
@@ -158,49 +158,53 @@ end
 
 
 function derive_implicit_parameterized_eqs(Ξ::AbstractArray{T, 2}, b::Basis) where T <: Real
+    
+    sparsity = Int64(norm(Ξ, 0)) # Overall sparsity
+    @parameters p[1:sparsity]
+    size_b = length(b)
 
-    sparsity = Int64(norm(Ξ, 0))
+    inds = @. ! iszero(Ξ)
 
-    @parameters p[1:sparsity]
-    p_ = zeros(eltype(Ξ), sparsity)
+    pinds = Int64.(norm.(eachcol(inds), 0))
+    
+    pinds_d = Int64.(norm.(eachcol(inds[1:size_b,:]), 0))
+    pinds_n = Int64.(norm.(eachcol(inds[size_b+1:end, :]), 0))
+    
+    p_ = similar(Ξ[inds])
+
+    eq = zeros(Operation, sum([i>0 for i in pinds]))
     cnt = 1
+    p_cnt = 1
+    eq_n = ModelingToolkit.Constant(0)
+    eq_d = ModelingToolkit.Constant(0)
 
-    b_ = Basis(Operation[], variables(b), parameters = [parameters(b)...; p...])
-
-    for i=1:size(Ξ, 2)
-        eq_d = nothing
-        eq_n = nothing
-        # Denominator
-        for j = 1:length(b)
-            if !iszero(Ξ[j,i])
-                if eq_d === nothing
-                    eq_d = p[cnt]*b[j]
-                else
-                    eq_d += p[cnt]*b[j]
+    @views for i=1:size(Ξ, 2)
+        # Numerator
+        eq_n = ModelingToolkit.Constant(0)
+        eq_d = ModelingToolkit.Constant(0)
+        if iszero(pinds_n[i]) || iszero(pinds_d[i])
+            continue
+        else
+            for j in 1:size_b
+                if inds[j, i]
+                    eq_d +=  p[p_cnt] * b.basis[j]
+                    p_[p_cnt] = Ξ[j, i]
+                    p_cnt += 1
                 end
-                p_[cnt] = Ξ[j,i]
-                cnt += 1
-            end
-        end
 
-        # Numerator
-        for j = 1:length(b)
-            if !iszero(Ξ[j+length(b),i])
-                if eq_n === nothing
-                    eq_n = p[cnt]*b[j]
-                else
-                    eq_n += p[cnt]*b[j]
+                if inds[j+size_b, i]
+                    eq_n +=  p[p_cnt] * b.basis[j]
+                    p_[p_cnt] = Ξ[j+size_b, i]
+                    p_cnt += 1
                 end
-                p_[cnt] = Ξ[j+length(b),i]
-                cnt += 1
             end
-        end
 
-        if !(isnothing(eq_d) || isnothing(eq_n))
-            push!(b_, ModelingToolkit.simplify(-eq_n ./ eq_d))
+            eq[cnt] = - eq_n / eq_d
         end
-
+        cnt += 1
     end
-    
+
+    b_ = Basis(eq, variables(b), parameters = vcat(parameters(b), p), iv = independent_variable(b))
+
     b_, p_
 end

src/sindy/results.jl

@@ -89,27 +89,29 @@ end
 
 function derive_parameterized_eqs(Ξ::AbstractArray{T, 2}, b::Basis, sparsity::Int64) where T <: Real
     @parameters p[1:sparsity]
-    p_ = zeros(eltype(Ξ), sparsity)
+    
+    inds = @. ~ iszero.(Ξ)
+    p_ = Ξ[inds]
+    pinds = Int64.(norm.(eachcol(inds), 0))
+    
     cnt = 1
-    b_ = Basis(Operation[], variables(b), parameters = [parameters(b)...; p...], iv = independent_variable(b))
-
-    for i=1:size(Ξ, 2)
-        eq = nothing
-        for j = 1:size(Ξ, 1)
-            if !iszero(Ξ[j,i])
-                if eq === nothing
-                    eq = p[cnt]*b[j]
-                else
-                    eq += p[cnt]*b[j]
-                end
-                p_[cnt] = Ξ[j,i]
-                cnt += 1
-            end
-        end
-        if !isnothing(eq)
-            push!(b_, eq)
+    eq = zeros(Operation, sum([i>0 for i in pinds]))
+
+    @inbounds for i=1:size(Ξ, 2)
+        if iszero(pinds[i])
+            continue
+        elseif i == 1 
+            eq[cnt] = sum(p[1:pinds[i]] .* b.basis[inds[:, i]])
+            cnt += 1
+        else
+            eq[cnt] = sum(p[sum(pinds[1:i-1])+1:pinds[i]+sum(pinds[1:i-1])] .* b.basis[inds[:, i]])
+            cnt += 1
         end
+        
+        
     end
+    b_ = Basis(eq, variables(b), parameters = vcat(parameters(b),p), iv = independent_variable(b))
+
     b_, p_
 end
 

test/isindy.jl

@@ -54,7 +54,7 @@
     estimator = ODEProblem(dudt, u0, tspan, ps)
     sol_ = solve(estimator, Tsit5(), saveat = 0.1)
     @test sol[:,:] ≈ sol_[:,:]
-    @test abs.(ps) ≈ abs.(Float64[-1/3 ; -1/3 ; -1.00 ; 2/3; 1.00 ;0.5 ;0.5 ; 1.0; 1.0; -1.0; 1.0])
+    @test abs.(ps) ≈ abs.(Float64[1/3, 1.0, 1/3, 2/3, 1, 0.5, 0.5, 1, 1, 1, 1]) atol = 1e-2
 
 
     @info "Michaelis-Menten-Kinetics"
@@ -90,7 +90,7 @@
     estimator = ODEProblem(dudt, u0, tspan, ps)
     sol_ = solve(estimator, Tsit5(), saveat = 0.1)
     @test isapprox(sol_[:,:], solution_1[:,:], atol = 1e-1)
-    @test abs.(ps) ≈ [1/3; 1.0; 0.2; 0.92] atol = 1e-1
+    @test abs.(ps) ≈ [1/3; 0.2; 1.0; 0.92] atol = 1e-1
 
     Ψ = ISINDy(X, DX, basis, STRRidge(1e-2), maxiter = 100, normalize = true)
     print_equations(Ψ, show_parameter = true)
@@ -101,7 +101,7 @@
     estimator = ODEProblem(dudt, u0, tspan, ps)
     sol_ = solve(estimator, Tsit5(), saveat = 0.1)
     @test isapprox(sol_[:,:], solution_1[:,:], atol = 1e-1)
-    @test abs.(ps) ≈ [1/3; 1.0; 0.2; 0.92] atol = 1e-1
+    @test abs.(ps) ≈ [1/3; 0.2; 1.0; 0.92] atol = 1e-1
 
     λs = exp10.(-3:0.1:-1)
     Ψ = ISINDy(X, DX, basis, λs ,STRRidge(1e-2), maxiter = 100, normalize = false)
@@ -113,5 +113,5 @@
     estimator = ODEProblem(dudt, u0, tspan, ps)
     sol_ = solve(estimator, Tsit5(), saveat = 0.1)
     @test isapprox(sol_[:,:], solution_1[:,:], atol = 1e-1)
-    @test abs.(ps) ≈ [1/3; 1.0; 0.2; 0.92] atol = 1e-1
+    @test abs.(ps) ≈ [1/3; 0.2; 1.0; 0.92] atol = 1e-1
 end