Support better logging for singular cases in integration

lib/DiffEqBase/src/verbosity.jl

@@ -2,6 +2,7 @@
     sub_specifiers = (:linear_verbosity, :nonlinear_verbosity)
     toggles = (
         :dt_NaN, :init_NaN, :dense_output_saveat, :max_iters, :dt_min_unstable, :instability,
+        :symbolic_diagnostic,
         :newton_convergence, :step_rejected, :step_accepted, :convergence_limit,
         :alg_switch, :stiff_detection, :mismatched_input_output_type, :jacobian_update,
         :w_factorization, :newton_iterations,
@@ -29,6 +30,7 @@
             max_iters = Silent(),
             dt_min_unstable = Silent(),
             instability = Silent(),
+            symbolic_diagnostic = Silent(),
             newton_convergence = Silent(),
             step_rejected = Silent(),
             step_accepted = Silent(),
@@ -67,6 +69,7 @@
             max_iters = WarnLevel(),
             dt_min_unstable = WarnLevel(),
             instability = WarnLevel(),
+            symbolic_diagnostic = Silent(),
             newton_convergence = WarnLevel(),
             step_rejected = Silent(),
             step_accepted = Silent(),
@@ -105,6 +108,7 @@
             max_iters = WarnLevel(),
             dt_min_unstable = WarnLevel(),
             instability = WarnLevel(),
+            symbolic_diagnostic = WarnLevel(),
             newton_convergence = Silent(),
             step_rejected = Silent(),
             step_accepted = Silent(),
@@ -143,6 +147,7 @@
             max_iters = WarnLevel(),
             dt_min_unstable = WarnLevel(),
             instability = WarnLevel(),
+            symbolic_diagnostic = WarnLevel(),
             newton_convergence = WarnLevel(),
             step_rejected = Silent(),
             step_accepted = Silent(),
@@ -181,6 +186,7 @@
             max_iters = WarnLevel(),
             dt_min_unstable = WarnLevel(),
             instability = WarnLevel(),
+            symbolic_diagnostic = WarnLevel(),
             newton_convergence = WarnLevel(),
             step_rejected = InfoLevel(),
             step_accepted = InfoLevel(),
@@ -215,7 +221,7 @@
     groups = (
         error_control = (
             :dt_NaN, :init_NaN, :dense_output_saveat, :max_iters, :dt_min_unstable,
-            :instability, :newton_convergence, :step_rejected, :step_accepted, :convergence_limit,
+            :instability, :symbolic_diagnostic, :newton_convergence, :step_rejected, :step_accepted, :convergence_limit,
         ),
         performance = (
             :alg_switch, :stiff_detection, :mismatched_input_output_type, :jacobian_update,

lib/OrdinaryDiffEqCore/Project.toml

@@ -28,6 +28,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"
@@ -43,6 +44,7 @@ Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
+
 [extensions]
 OrdinaryDiffEqCoreMooncakeExt = "Mooncake"
 OrdinaryDiffEqCorePolyesterExt = "Polyester"
@@ -82,6 +84,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/src/OrdinaryDiffEqCore.jl

@@ -50,6 +50,7 @@ import SciMLOperators: 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
 
@@ -92,7 +93,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_numerical_instability, has_mtk_sys
 
 import SciMLOperators: islinear
 # `calculate_residuals`/`calculate_residuals!` are unused here but re-exported for

lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl

@@ -697,6 +697,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(
@@ -707,6 +708,213 @@ 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
+
+SciMLBase.has_mtk_sys(integrator::ODEIntegrator) = hasproperty(integrator.sol.prob.f, :sys) && integrator.sol.prob.f.sys !== nothing
+
+#get atmp values by cache, for use in diagnostics
+error_estimate_residuals(cache) = hasfield(typeof(cache), :atmp) ? getfield(cache, :atmp) : nothing
+error_estimate_residuals(cache::CompositeCache) = error_estimate_residuals(@inbounds cache.caches[cache.current])
+function error_estimate_residuals(cache::DefaultCache)
+    1 <= cache.current <= 6 || return nothing
+    name = (:cache1, :cache2, :cache3, :cache4, :cache5, :cache6)[cache.current]
+    return isdefined(cache, name) ? error_estimate_residuals(getfield(cache, name)) : nothing
+end
+
+function SciMLBase.log_numerical_instability(integrator::ODEIntegrator; jacobian_logging = true)
+    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
+        # keep only components within 20 orders of magnitude of the largest
+        if !isempty(blown_idxs)
+            max_blown = maximum(abs(u[i]) for i in blown_idxs)
+            cutoff = max_blown * 1e-20
+            filter!(i -> abs(u[i]) >= cutoff, blown_idxs)
+            sort!(blown_idxs, by = i -> abs(u[i]), rev = true)
+        end
+    end
+
+    # Jacobian analysis: rows and columns holding non-finite or unusually large entries
+    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 #no jac to analyze
+        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 Jacobian rows/cols back to equations/variables
+    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
+
+    # each analysis gets its own section
+    state_analysis = String[]
+    jacobian_analysis = String[]
+    error_analysis = String[]
+
+    # state diagnostics message
+    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!(state_analysis, "All $n_total state variables are non-finite (NaN/Inf)")
+            elseif n_nan > 3
+                push!(state_analysis, "$n_nan of $n_total state variables are non-finite (NaN/Inf): indices $nan_inf_idxs")
+            else
+                for i in nan_inf_idxs
+                    push!(state_analysis, "u[$i] = $(u[i]) is non-finite (NaN/Inf)")
+                end
+            end
+        else
+            push!(state_analysis, "u = $u is non-finite (NaN/Inf)")
+        end
+    elseif !isempty(blown_idxs)
+        if u isa AbstractArray
+            for i in blown_idxs
+                push!(state_analysis, "u[$i] = $(@sprintf("%.4g", u[i])) has grown >1e6× its initial value")
+            end
+        else
+            push!(state_analysis, "u = $(@sprintf("%.4g", u)) has grown >1e6× its initial value")
+        end
+    end
+
+    # Jacobian diagnostics message
+    if jacobian_logging && bad_entries !== nothing && !isempty(bad_entries)
+        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!(jacobian_analysis, "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!(jacobian_analysis, "  row $row corresponds to equation: $(sym_eqs[row])") #trace rows back to symbolic eqs
+                end
+            end
+        end
+        # jac cols
+        if !isempty(singularity_cols)
+            push!(jacobian_analysis, "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!(jacobian_analysis, "  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") * "."
+    # error estimate analysis
+    if integrator.opts.adaptive
+        push!(error_analysis, "step error estimate EEst = $(@sprintf("%.4g", get_EEst(integrator))) (a step is accepted when EEst <= 1)")
+        atmp = error_estimate_residuals(integrator.cache)
+        if atmp isa AbstractArray && !isempty(atmp) && eltype(atmp) <: Number
+            nonfinite = count(!isfinite, atmp)
+            nonfinite > 0 && push!(error_analysis, "$nonfinite of $(length(atmp)) weighted residuals are non-finite (NaN/Inf)")
+            idxs = collect(eachindex(atmp))
+            n = min(3, length(idxs))
+            # sort NaN residuals along Inf
+            partialsort!(idxs, 1:n, by = i -> (v = abs(atmp[i]); isnan(v) ? typemax(v) : v), rev = true)
+
+            with_state = u isa AbstractArray && eachindex(u) == eachindex(atmp)
+            contributors = String[]
+            for i in idxs[1:n]
+                line = "  atmp[$i] = $(@sprintf("%.4g", atmp[i]))"
+                if with_state
+                    line *= ", u[$i] = $(@sprintf("%.4g", u[i]))"
+                    line *= ", uprev[$i] = $(@sprintf("%.4g", integrator.uprev[i]))"
+                end
+                push!(contributors, line)
+            end
+            push!(error_analysis, "largest contributors to EEst = internalnorm(atmp), where atmp is the tolerance-weighted local error per state component:\n" * join(contributors, "\n"))
+        end
+    end
+
+    # assemble the message, one titled section per non-empty analysis
+    sections = (
+        ("State Analysis", state_analysis),
+        ("Jacobian Analysis", jacobian_analysis),
+        ("Error Analysis", error_analysis)
+    )
+    all(isempty(msgs) for (_, msgs) in sections) && return ""
+
+    diagnostic = "\n\nDiagnostics:"
+    for (title, msgs) in sections
+        isempty(msgs) && continue
+        body = join(("  " * replace(msg, "\n" => "\n  ") for msg in msgs), "\n")
+        diagnostic *= "\n\n$title:\n$body"
+    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

@@ -229,6 +229,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

@@ -41,7 +41,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
 

lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl

@@ -354,6 +354,8 @@ function calc_J(integrator, cache, next_step::Bool = false)
     return J
 end
 
+get_fresh_jacobian(integrator, cache::OrdinaryDiffEqCache) = calc_J(integrator, cache)
+
 """
     calc_J!(J, integrator, cache, next_step::Bool = false) -> J