add rootcause analysis of symbolic system

Author: Shreyas-Ekanathan

lib/OrdinaryDiffEqCore/Project.toml

@@ -39,12 +39,14 @@ SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
 
 [weakdeps]
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [extensions]
 OrdinaryDiffEqCoreMooncakeExt = "Mooncake"
+OrdinaryDiffEqModelingToolkitExt = "ModelingToolkit"
 OrdinaryDiffEqCorePolyesterExt = "Polyester"
 OrdinaryDiffEqCoreSparseArraysExt = "SparseArrays"
 

lib/OrdinaryDiffEqCore/ext/OrdinaryDiffEqModelingToolkitExt.jl

@@ -0,0 +1,54 @@
+module OrdinaryDiffEqModelingToolkitExt
+
+using OrdinaryDiffEqCore, ModelingToolkit
+using Printf: @sprintf
+
+function OrdinaryDiffEqCore.system_singularity_rootcause(sys, u)
+    substitution_map = Dict(zip(unknowns(sys), u))
+    diagnosis = String[]
+    for eq in equations(sys)
+        find_singular_subterms(eq, eq.rhs, substitution_map, diagnosis)
+    end
+    return diagnosis
+end
+
+function find_singular_subterms(eq, expr, sub_map, diagnosis)
+    expr = Symbolics.unwrap(expr)
+    !SymbolicUtils.iscall(expr) && return diagnosis
+    op = SymbolicUtils.operation(expr)
+    args = SymbolicUtils.arguments(expr)
+
+    if op === (/) #division, singular if we divide by small thing
+        d = Symbolics.value(Symbolics.substitute(args[2], sub_map))
+        if d isa Number && abs(d) < 1e-10
+            push!(diagnosis, "in equation $eq: division by very small value $(args[2]) ≈ $(@sprintf("%.4g", d)) leads to singularity.")
+        end
+    elseif op === log #singular if we log small thing
+        x = Symbolics.value(Symbolics.substitute(args[1], sub_map))
+        if x isa Number && x <= 1e-10
+            push!(diagnosis, "in equation $eq: log of $(args[1]) = $(@sprintf("%.4g", x)) near/at singularity (derivative blows up).")
+        end
+    elseif op === sqrt 
+        x = Symbolics.value(Symbolics.substitute(args[1], sub_map))
+        if x isa Number && x < 1e-10
+            push!(diagnosis, "in equation $eq: sqrt of $(args[1]) = $(@sprintf("%.4g", x)) near/at singularity (derivative blows up).")
+        end
+    elseif op === (^)
+        e = Symbolics.value(Symbolics.substitute(args[2], sub_map))
+        b = Symbolics.value(Symbolics.substitute(args[1], sub_map))
+        if e isa Number && b isa Number #two cases
+            if e < 0 && abs(b) < 1e-10
+                push!(diagnosis, "in equation $eq: ($(args[1])) raised to power $e with base ≈ $(@sprintf("%.4g", b)) going to 0; result diverges.")
+            elseif e > 0 && abs(b) > 1
+                push!(diagnosis, "in equation $eq: ($(args[1]) ≈ $(@sprintf("%.4g", b))) raised to power $e - base magnitude is large and being amplified.")
+            end
+        end
+    end
+
+    for arg in args
+        find_singular_subterms(eq, arg, sub_map, diagnosis)
+    end
+    return diagnosis
+end
+
+end

lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl

@@ -615,6 +615,8 @@ end
 # overrides this with a method that calls calc_J to get a fresh Jacobian.
 get_fresh_jacobian(integrator, cache) = cache.J
 
+system_singularity_rootcause(sys, u) = ""
+
 function SciMLBase.log_instability(integrator::ODEIntegrator)
     W = _get_W(integrator)
     u = integrator.u
@@ -684,6 +686,11 @@ function SciMLBase.log_instability(integrator::ODEIntegrator)
     sym_eqs = (sys !== nothing && hasfield(typeof(sys), :eqs)) ? getfield(sys, :eqs) : nothing
     sym_vars = (sys !== nothing && hasfield(typeof(sys), :unknowns)) ? getfield(sys, :unknowns) : nothing
 
+    # symbolic analysis
+    #skip jac analysis if this isn't empty
+    u_vals = length(nan_inf_idxs) == 0 ? u : integrator.uprev
+    symbolic_analysis = system_singularity_rootcause(sys, u_vals)
+
     # diagnostic message construction
     diagnostic = String[]
     if !isempty(nan_inf_idxs) #state vars
@@ -712,7 +719,7 @@ function SciMLBase.log_instability(integrator::ODEIntegrator)
         end
     end
 
-    if bad_entries !== nothing && !isempty(bad_entries) #Jacobian analysis
+    if bad_entries !== nothing && !isempty(bad_entries) && isempty(symbolic_analysis) #Jacobian analysis (skipped if we have symbolic analysis)
         has_nonfinite = false
         has_large = false
         for (_, _, v) in bad_entries
@@ -730,7 +737,7 @@ function SciMLBase.log_instability(integrator::ODEIntegrator)
         for (i, j, v) in first(bad_entries, 5)
             push!(example_strs, "J[$i,$j] = $(@sprintf("%.4g", v))")
         end
-        push!(diagnostic, "Jacobian row(s) $singularity_rows have $entry_desc entries (e.g. $(join(example_strs, ", "))), suggesting a singularity in those equation(s)")
+        push!(diagnostic, "\nJacobian row(s) $singularity_rows have $entry_desc entries (e.g. $(join(example_strs, ", "))), suggesting a singularity in those equation(s)")
         if sym_eqs !== nothing
             for row in singularity_rows
                 if row <= length(sym_eqs)
@@ -740,7 +747,7 @@ function SciMLBase.log_instability(integrator::ODEIntegrator)
         end
         # jac cols
         if !isempty(singularity_cols)
-            push!(diagnostic, "Jacobian column(s) $singularity_cols have $entry_desc entries, suggesting those state component(s) are diverging")
+            push!(diagnostic, "\nJacobian column(s) $singularity_cols have $entry_desc entries, suggesting those state component(s) are diverging")
             if sym_vars !== nothing
                 for col in singularity_cols
                     if col <= length(sym_vars)
@@ -750,7 +757,13 @@ function SciMLBase.log_instability(integrator::ODEIntegrator)
             end
         end
     end
-    diagnostic = isempty(diagnostic) ? "" : "\nDiagnostics:\n" * join(diagnostic, "\n") * "."
+
+    diagnostic = isempty(diagnostic) ? "" : "\n\nDiagnostics:\n" * join(diagnostic, "\n") * "."
+
+    if !isempty(symbolic_analysis)
+        diagnostic *= "\n\nSymbolic Analysis of MTK System:\n" * join(symbolic_analysis, "\n")
+    end
+
     return diagnostic
 end