From 94576dbec700682eef80054dce76f1dcf6c78518 Mon Sep 17 00:00:00 2001
From: Herman Sletmoen <hermansletmoen@gmail.com>
Date: Wed, 8 Jan 2025 17:27:32 +0100
Subject: [PATCH] Let user pass list of functions that should be debugged

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

diff --git a/src/debugging.jl b/src/debugging.jl
index c8b7372c5d..f6652183cb 100644
--- a/src/debugging.jl
+++ b/src/debugging.jl
@@ -1,5 +1,3 @@
-const LOGGED_FUN = Set([log, sqrt, (^), /, inv])
-
 struct LoggedFunctionException <: Exception
     msg::String
 end
@@ -32,11 +30,11 @@ function logged_fun(f, args...; error_nonfinite = true) # remember to update err
     term(LoggedFun(f, args, error_nonfinite), args..., type = Real)
 end
 
-debug_sub(eq::Equation; kw...) = debug_sub(eq.lhs; kw...) ~ debug_sub(eq.rhs; kw...)
-function debug_sub(ex; kw...)
+debug_sub(eq::Equation, funcs; kw...) = debug_sub(eq.lhs, funcs; kw...) ~ debug_sub(eq.rhs, funcs; kw...)
+function debug_sub(ex, funcs; kw...)
     iscall(ex) || return ex
     f = operation(ex)
-    args = map(debug_sub, arguments(ex))
-    f in LOGGED_FUN ? logged_fun(f, args...; kw...) :
+    args = map(ex -> debug_sub(ex, funcs; kw...), arguments(ex))
+    f in funcs ? 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 c3525c542b..796ff73eba 100644
--- a/src/systems/abstractsystem.jl
+++ b/src/systems/abstractsystem.jl
@@ -2260,11 +2260,11 @@ macro mtkbuild(exprs...)
 end
 
 """
-    debug_system(sys::AbstractSystem; error_nonfinite = true)
+    debug_system(sys::AbstractSystem; functions = [log, sqrt, (^), /, inv, asin, acos], error_nonfinite = true)
 
-Replace functions with singularities with a function that errors with symbolic
-information. If `error_nonfinite`, debugged functions that output nonfinite values
-(like `Inf` or `NaN`) also display errors, even though the raw function itself
+Wrap `functions` in `sys` so any error thrown in them shows helpful symbolic-numeric
+information about its input. If `error_nonfinite`, 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
@@ -2278,15 +2278,18 @@ ERROR: Function /(1, sin(P(t))) output non-finite value Inf with input
   sin(P(t)) => 0.0
 ```
 """
-function debug_system(sys::AbstractSystem; kw...)
+function debug_system(sys::AbstractSystem; functions = [log, sqrt, (^), /, inv, asin, acos], kw...)
+    if !(functions isa Set)
+        functions = Set(functions) # more efficient "in" lookup
+    end
     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); kw...)
+        @set! sys.eqs = debug_sub.(equations(sys), Ref(functions); kw...)
     end
     if has_observed(sys)
-        @set! sys.observed = debug_sub.(observed(sys); kw...)
+        @set! sys.observed = debug_sub.(observed(sys), Ref(functions); kw...)
     end
     return sys
 end