Support better logging for singular cases in integration

lib/OrdinaryDiffEqCore/Project.toml

@@ -26,6 +26,7 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -37,12 +38,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"
 
@@ -77,6 +80,7 @@ Pkg = "1"
 Polyester = "0.7"
 PrecompileTools = "1.2.1, 1.3"
 Preferences = "1.5.0"
+Printf = "1.9"
 Random = "<0.0.1, 1"
 RecursiveArrayTools = "4.2.0"
 Reexport = "1.2.2"

lib/OrdinaryDiffEqCore/ext/OrdinaryDiffEqCoreSparseArraysExt.jl

@@ -1,7 +1,7 @@
 module OrdinaryDiffEqCoreSparseArraysExt
 
 using SparseArrays: SparseMatrixCSC
-import OrdinaryDiffEqCore: _isdiag, find_algebraic_vars_eqs
+import OrdinaryDiffEqCore: _isdiag, find_algebraic_vars_eqs, _find_large_jac_entries!
 
 # Efficient O(nnz) isdiag check for sparse matrices.
 # Standard isdiag is O(n²) which is prohibitively slow for large sparse matrices.
@@ -22,6 +22,20 @@ function _isdiag(A::SparseMatrixCSC)
     return true
 end
 
+# only look at nonzero vals
+function _find_large_jac_entries!(rows::Set{Int}, cols::Set{Int}, entries::Vector, jac::SparseMatrixCSC)
+    @inbounds for j in axes(jac, 2)
+        for k in jac.colptr[j]:(jac.colptr[j + 1] - 1)
+            val = jac.nzval[k]
+            if !isfinite(val) || abs(val) > 1e6
+                i = jac.rowval[k]
+                push!(rows, i)
+                push!(cols, j)
+                push!(entries, (i, j, val))
+            end
+        end
+    end
+end
 """
     find_algebraic_vars_eqs(M::SparseMatrixCSC)
 

lib/OrdinaryDiffEqCore/ext/OrdinaryDiffEqModelingToolkitExt.jl

@@ -0,0 +1,121 @@
+module OrdinaryDiffEqModelingToolkitExt
+
+using OrdinaryDiffEqCore, ModelingToolkit
+using Printf: @sprintf
+
+function OrdinaryDiffEqCore.system_singularity_rootcause(sys, u, uprev)
+    diagnosis = String[]
+
+    #check for assertion failures
+    unks = unknowns(sys)
+    curr_substitution_map = Dict(zip(unks, u))
+    prev_substitution_map = Dict(zip(unknowns(sys), uprev))
+
+    for (cond, msg) in ModelingToolkit.assertions(sys)
+        subclauses = String[]
+        find_failing_subterms(cond, prev_substitution_map, curr_substitution_map, subclauses)
+        if !isempty(subclauses)
+            push!(diagnosis, "\n\nAssertion violated: $cond - \"$msg\"")
+            append!(diagnosis, subclauses)
+        end
+    end
+
+    #find singularity causes in equations
+    singularities = String[]
+    for eq in equations(sys)
+        find_singular_subterms(eq, eq.rhs, prev_substitution_map, singularities)
+    end
+    if !isempty(singularities)
+        push!(diagnosis, "\nSymbolic Analysis of MTK System:")
+        append!(diagnosis, singularities)
+    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
+
+function find_failing_subterms(cond, prev_map, curr_map, diagnosis)
+    c = Symbolics.unwrap(cond)
+    !SymbolicUtils.iscall(c) && return diagnosis
+    op = SymbolicUtils.operation(c)
+    args = SymbolicUtils.arguments(c)
+
+    if (op === (<) || op === (>) || op === (<=) || op === (>=)) && length(args) == 2
+        #compare using previous non-nan values to find violating subclauses, then output current values
+        lhs = Symbolics.value(Symbolics.substitute(args[1], prev_map))
+        rhs = Symbolics.value(Symbolics.substitute(args[2], prev_map))
+        if lhs isa Number && rhs isa Number
+            # small margin -> violated
+            margin = (op === (<) || op === (<=)) ? rhs - lhs : lhs - rhs
+            if margin <= 1e-6
+                push!(diagnosis, "   subclause `$c` violated: $(clause_values(c, curr_map))")
+            end
+        end
+    elseif op === (!=) && length(args) == 2
+        lhs = Symbolics.value(Symbolics.substitute(args[1], prev_map))
+        rhs = Symbolics.value(Symbolics.substitute(args[2], prev_map))
+        if lhs isa Number && rhs isa Number && abs(lhs - rhs) <= 1e-6
+            push!(diagnosis, "   subclause `$c` violated: $(clause_values(c, curr_map))")
+        end
+    elseif op === (==) && length(args) == 2
+        lhs = Symbolics.value(Symbolics.substitute(args[1], prev_map))
+        rhs = Symbolics.value(Symbolics.substitute(args[2], prev_map))
+        if lhs isa Number && rhs isa Number && abs(lhs - rhs) > 1e-6
+            push!(diagnosis, "   subclause `$c` violated: $(clause_values(c, curr_map))")
+        end
+    else #recurse
+        for arg in args
+            find_failing_subterms(arg, prev_map, curr_map, diagnosis)
+        end
+    end
+    return diagnosis
+end
+
+function clause_values(c, curr_map)
+    parts = String[]
+    for v in Symbolics.get_variables(c)
+        val = Symbolics.value(Symbolics.substitute(v, curr_map))
+        push!(parts, val isa Number ? "$v = $(@sprintf("%.4g", val))" : "$v = $val")
+    end
+    return join(parts, ", ")
+end
+
+end

lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl

@@ -42,6 +42,7 @@ import SciMLOperators: AbstractSciMLOperator, MatrixOperator, FunctionOperator,
     isconstant
 
 import Random
+import Printf: @sprintf
 
 import RecursiveArrayTools: recursivecopy!, recursivecopy, recursive_bottom_eltype, recursive_unitless_bottom_eltype, recursive_unitless_eltype, copyat_or_push!, DiffEqArray
 
@@ -84,7 +85,7 @@ using SciMLBase: SciMLBase, CallbackSet, ContinuousCallback, DAEProblem,
 using SciMLOperators: SciMLOperators
 using CommonSolve: solve
 
-import SciMLBase: AbstractNonlinearProblem, alg_order, LinearAliasSpecifier
+import SciMLBase: AbstractNonlinearProblem, alg_order, LinearAliasSpecifier, log_instability
 
 import SciMLBase: islinear
 # `calculate_residuals`/`calculate_residuals!` are unused here but re-exported for

lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl

@@ -601,6 +601,7 @@ function increment_reject!(stats)
     return stats.nreject += 1
 end
 
+
 function log_step!(progress_name, progress_id, progress_message, dt, u, p, t, tspan)
     t1, t2 = tspan
     return @logmsg(
@@ -611,6 +612,160 @@ function log_step!(progress_name, progress_id, progress_message, dt, u, p, t, ts
     )
 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, uprev) = ""
+
+function SciMLBase.log_instability(integrator::ODEIntegrator)
+    W = _get_W(integrator)
+    u = integrator.u
+    u0 = integrator.sol.prob.u0
+
+    # state analysis: NaN/Inf components, and components that have blown up
+    nan_inf_idxs = findall(!isfinite, u)
+    blown_idxs = Int[]
+    if length(u) == length(u0)
+        for i in eachindex(u)
+            ref = max(abs(u0[i]), oneunit(eltype(u)))
+            abs(u[i]) > 1.0e6 * ref && push!(blown_idxs, i)
+        end
+    end
+
+    # jacobian analysis over rows and columns for large values
+    jac = if W !== nothing && hasproperty(W, :J)
+        #rosenbrock
+        W.J
+    elseif hasproperty(integrator.cache, :J)
+        #radau
+        get_fresh_jacobian(integrator, integrator.cache)
+    elseif hasproperty(integrator.cache, :nlsolver) &&
+            hasproperty(integrator.cache.nlsolver.cache, :J)
+        #BDF
+        integrator.cache.nlsolver.cache.J
+    else
+        nothing
+    end
+
+    bad_entries = nothing
+    singularity_rows = nothing
+    singularity_cols = nothing
+    if jac !== nothing
+        rows = Set{Int}()
+        cols = Set{Int}()
+        entries = Tuple{Int, Int, eltype(jac)}[]
+        _find_large_jac_entries!(rows, cols, entries, jac)
+
+        # keep only entries within 10 orders of magnitude of the largest finite entry,
+        # plus any non-finite entries. filters out large-but-normal model parameters 
+        max_finite = 0.0
+        for (_, _, v) in entries
+            if isfinite(v)
+                max_finite = max(max_finite, abs(v))
+            end
+        end
+        cutoff = max_finite * 1e-10
+        filter!(t -> !isfinite(t[3]) || abs(t[3]) >= cutoff, entries) #only keep those vals within 1e10 of max or inf/nan
+        sort!(entries, by = t -> (!isfinite(t[3]), abs(t[3])), rev = true)
+
+        # derive rows and columns from remaining entries
+        row_set = Set{Int}()
+        col_set = Set{Int}()
+        for (i, j, _) in entries
+            push!(row_set, i)
+            push!(col_set, j)
+        end
+        bad_entries = entries
+        singularity_rows = sort!(collect(row_set))
+        singularity_cols = sort!(collect(col_set))
+    end
+
+    # trace diagnostics to symbolic system if present
+    f = integrator.sol.prob.f
+    sys = (hasproperty(f, :sys) && f.sys !== nothing) ? f.sys : nothing
+    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
+    symbolic_analysis = system_singularity_rootcause(sys, u, integrator.uprev)
+
+    # diagnostic message construction
+    diagnostic = String[]
+    if !isempty(nan_inf_idxs) #state vars
+        if u isa AbstractArray
+            n_nan = length(nan_inf_idxs)
+            n_total = length(u)
+            if n_nan == n_total
+                push!(diagnostic, "All $n_total state variables are non-finite (NaN/Inf)")
+            elseif n_nan > 3
+                push!(diagnostic, "$n_nan of $n_total state variables are non-finite (NaN/Inf): indices $nan_inf_idxs")
+            else
+                for i in nan_inf_idxs
+                    push!(diagnostic, "u[$i] = $(u[i]) is non-finite (NaN/Inf)")
+                end
+            end
+        else
+            push!(diagnostic, "u = $u is non-finite (NaN/Inf)")
+        end
+    elseif !isempty(blown_idxs)
+        if u isa AbstractArray
+            for i in blown_idxs
+                push!(diagnostic, "u[$i] = $(@sprintf("%.4g", u[i])) has grown >1e6× its initial value")
+            end
+        else
+            push!(diagnostic, "u = $(@sprintf("%.4g", u)) has grown >1e6× its initial value")
+        end
+    end
+
+    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
+            isfinite(v) ? (has_large = true) : (has_nonfinite = true)
+        end
+        entry_desc = if has_nonfinite && has_large
+            "non-finite and large"
+        elseif has_nonfinite
+            "non-finite"
+        else
+            "unusually large"
+        end
+
+        example_strs = String[]
+        for (i, j, v) in first(bad_entries, 5)
+            push!(example_strs, "J[$i,$j] = $(@sprintf("%.4g", v))")
+        end
+        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)
+                    push!(diagnostic, "  row $row corresponds to equation: $(sym_eqs[row])") #trace rows back to symbolic eqs
+                end
+            end
+        end
+        # jac cols
+        if !isempty(singularity_cols)
+            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)
+                        push!(diagnostic, "  col $col corresponds to variable: $(sym_vars[col])") #trace cols back to symbolic vars
+                    end
+                end
+            end
+        end
+    end
+
+    diagnostic = isempty(diagnostic) ? "" : "\n\nDiagnostics:\n" * join(diagnostic, "\n\n") * "."
+
+    if !isempty(symbolic_analysis)
+        diagnostic *= join(symbolic_analysis, "\n")
+    end
+
+    return diagnostic
+end
+
 function fixed_t_for_tstop_error!(integrator, ttmp)
     if _get_next_step_tstop(integrator)
         _set_tstop_flag!(integrator, false)

lib/OrdinaryDiffEqCore/src/misc_utils.jl

@@ -156,6 +156,18 @@ end
 # Sparse specialization is provided in OrdinaryDiffEqCoreSparseArraysExt
 _isdiag(A::AbstractMatrix) = isdiag(A)
 
+# Dense fallback to find large Jacobian entries. 
+# Sparse specialization is provided in OrdinaryDiffEqCoreSparseArraysExt
+function _find_large_jac_entries!(rows::Set{Int}, cols::Set{Int}, entries::Vector, jac::AbstractMatrix)
+    for i in axes(jac, 1), j in axes(jac, 2)
+        val = jac[i, j]
+        if !isfinite(val) || abs(val) > 1e6
+            push!(rows, i)
+            push!(cols, j)
+            push!(entries, (i, j, val))
+        end
+    end
+end
 """
     find_algebraic_vars_eqs(M)
 

lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl

@@ -40,7 +40,7 @@ using OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, OrdinaryDiffEqAdaptiveImplici
     TryAgain,
     Divergence, constvalue, @SciMLMessage
 
-import OrdinaryDiffEqCore: get_chunksize, resize_J_W!, alg_autodiff
+import OrdinaryDiffEqCore: get_chunksize, resize_J_W!, alg_autodiff, get_fresh_jacobian
 
 import ConstructionBase