From f5562dc215102895a6c0296877be4a886321067f Mon Sep 17 00:00:00 2001
From: Herman Sletmoen <hermansletmoen@gmail.com>
Date: Wed, 8 Jan 2025 17:07:35 +0100
Subject: [PATCH] Toggle extra non-finite value error with keyword argument

---
 src/debugging.jl              | 13 +++++++------
 src/systems/abstractsystem.jl | 12 +++++++-----
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/debugging.jl b/src/debugging.jl
index 830d32afce..c8b7372c5d 100644
--- a/src/debugging.jl
+++ b/src/debugging.jl
@@ -6,6 +6,7 @@ end
 struct LoggedFun{F}
     f::F
     args::Any
+    error_nonfinite::Bool
 end
 LoggedFunctionException(lf::LoggedFun, args, msg) = LoggedFunctionException(
     "Function $(lf.f)($(join(lf.args, ", "))) " * msg * " with input" *
@@ -20,22 +21,22 @@ function (lf::LoggedFun)(args...)
     catch err
         throw(LoggedFunctionException(lf, args, "errors")) # Julia automatically attaches original error message
     end
-    if !isfinite(val)
+    if lf.error_nonfinite && !isfinite(val)
         throw(LoggedFunctionException(lf, args, "output non-finite value $val"))
     end
     return val
 end
 
-function logged_fun(f, args...)
+function logged_fun(f, args...; error_nonfinite = true) # remember to update error_nonfinite in debug_system() docstring
     # Currently we don't really support complex numbers
-    term(LoggedFun(f, args), args..., type = Real)
+    term(LoggedFun(f, args, error_nonfinite), args..., type = Real)
 end
 
-debug_sub(eq::Equation) = debug_sub(eq.lhs) ~ debug_sub(eq.rhs)
-function debug_sub(ex)
+debug_sub(eq::Equation; kw...) = debug_sub(eq.lhs; kw...) ~ debug_sub(eq.rhs; kw...)
+function debug_sub(ex; kw...)
     iscall(ex) || return ex
     f = operation(ex)
     args = map(debug_sub, arguments(ex))
-    f in LOGGED_FUN ? logged_fun(f, args...) :
+    f in LOGGED_FUN ? logged_fun(f, args...; kw...) :
     maketerm(typeof(ex), f, args, metadata(ex))
 end
diff --git a/src/systems/abstractsystem.jl b/src/systems/abstractsystem.jl
index 492ded28ee..c3525c542b 100644
--- a/src/systems/abstractsystem.jl
+++ b/src/systems/abstractsystem.jl
@@ -2260,10 +2260,12 @@ macro mtkbuild(exprs...)
 end
 
 """
-$(SIGNATURES)
+    debug_system(sys::AbstractSystem; error_nonfinite = true)
 
 Replace functions with singularities with a function that errors with symbolic
-information. E.g.
+information. If `error_nonfinite`, debugged functions that output nonfinite values
+(like `Inf` or `NaN`) also display errors, even though the raw function itself
+does not throw an exception (like `1/0`). For example:
 
 ```julia-repl
 julia> sys = debug_system(complete(sys))
@@ -2276,15 +2278,15 @@ ERROR: Function /(1, sin(P(t))) output non-finite value Inf with input
   sin(P(t)) => 0.0
 ```
 """
-function debug_system(sys::AbstractSystem)
+function debug_system(sys::AbstractSystem; kw...)
     if has_systems(sys) && !isempty(get_systems(sys))
         error("debug_system(sys) only works on systems with no sub-systems! Consider flattening it with flatten(sys) or structural_simplify(sys) first.")
     end
     if has_eqs(sys)
-        @set! sys.eqs = debug_sub.(equations(sys))
+        @set! sys.eqs = debug_sub.(equations(sys); kw...)
     end
     if has_observed(sys)
-        @set! sys.observed = debug_sub.(observed(sys))
+        @set! sys.observed = debug_sub.(observed(sys); kw...)
     end
     return sys
 end