feat: add assertions to function regardless of debug_system

add logging in `debug_system`

Author: AayushSabharwal

src/debugging.jl

@@ -43,6 +43,17 @@ function debug_sub(ex, funcs; kw...)
     maketerm(typeof(ex), f, args, metadata(ex))
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+A function which returns `NaN` if `condition` fails, and `0.0` otherwise.
+"""
+function _nan_condition(condition::Bool)
+    condition ? 0.0 : NaN
+end
+
+@register_symbolic _nan_condition(condition::Bool)
+
 """
     $(TYPEDSIGNATURES)
 
@@ -51,9 +62,10 @@ and zero if it is true. In case the condition is false and `log == true`,
 `message` will be logged as an `@error`.
 """
 function _debug_assertion(expr::Bool, message::String, log::Bool)
-    expr && return 0.0
+    value = _nan_condition(expr)
+    isnan(value) || return value
     log && @error message
-    return NaN
+    return value
 end
 
 @register_symbolic _debug_assertion(expr::Bool, message::String, log::Bool)
@@ -67,13 +79,22 @@ const ASSERTION_LOG_VARIABLE = only(@parameters __log_assertions_ₘₜₖ::Bool
 """
     $(TYPEDSIGNATURES)
 
-Get a symbolic expression as per the requirement of `debug_system` for all the assertions
-in `assertions`. `is_split` denotes whether the corresponding system is a split system.
+Get a symbolic expression for all the assertions in `sys`. The expression returns `NaN`
+if any of the assertions fail, and `0.0` otherwise. If `ASSERTION_LOG_VARIABLE` is a
+parameter in the system, it will control whether the message associated with each
+assertion is logged when it fails.
 """
-function get_assertions_expr(assertions::Dict{BasicSymbolic, String}, is_split::Bool)
+function get_assertions_expr(sys::AbstractSystem)
+    asserts = assertions(sys)
     term = 0
-    for (k, v) in assertions
-        term += _debug_assertion(k, "Assertion $k failed:\n$v", ASSERTION_LOG_VARIABLE)
+    if is_parameter(sys, ASSERTION_LOG_VARIABLE)
+        for (k, v) in asserts
+            term += _debug_assertion(k, "Assertion $k failed:\n$v", ASSERTION_LOG_VARIABLE)
+        end
+    else
+        for (k, v) in asserts
+            term += _nan_condition(k)
+        end
     end
     return term
 end

src/systems/abstractsystem.jl

@@ -2303,12 +2303,11 @@ ERROR: Function /(1, sin(P(t))) output non-finite value Inf with input
   sin(P(t)) => 0.0
 ```
 
-Additionally, all assertions in the system are validated in the equations. If any of
-the conditions are false, the right hand side of at least one of the equations of
-the system will evaluate to `NaN`. A new parameter is also added to the system which
-controls whether the message associated with each assertion will be logged when the
-assertion fails. This parameter defaults to `true` and can be toggled by
-symbolic indexing with `ModelingToolkit.ASSERTION_LOG_VARIABLE`. For example,
+Additionally, all assertions in the system are optionally logged when they fail.
+A new parameter is also added to the system which controls whether the message associated
+with each assertion will be logged when the assertion fails. This parameter defaults to
+`true` and can be toggled by symbolic indexing with
+`ModelingToolkit.ASSERTION_LOG_VARIABLE`. For example,
 `prob.ps[ModelingToolkit.ASSERTION_LOG_VARIABLE] = false` will disable logging.
 """
 function debug_system(
@@ -2321,18 +2320,16 @@ function debug_system(
     end
     if has_eqs(sys)
         eqs = debug_sub.(equations(sys), Ref(functions); kw...)
-        expr = get_assertions_expr(assertions(sys))
-        eqs[end] = eqs[end].lhs ~ eqs[end].rhs + expr
         @set! sys.eqs = eqs
         @set! sys.ps = unique!([get_ps(sys); ASSERTION_LOG_VARIABLE])
         @set! sys.defaults = merge(get_defaults(sys), Dict(ASSERTION_LOG_VARIABLE => true))
-        if iscomplete(sys)
-            sys = complete(sys; split = is_split(sys))
-        end
     end
     if has_observed(sys)
         @set! sys.observed = debug_sub.(observed(sys), Ref(functions); kw...)
     end
+    if iscomplete(sys)
+        sys = complete(sys; split = is_split(sys))
+    end
     return sys
 end
 

src/systems/diffeqs/abstractodesystem.jl

@@ -168,6 +168,10 @@ function generate_function(sys::AbstractODESystem, dvs = unknowns(sys),
     rhss = implicit_dae ? [_iszero(eq.lhs) ? eq.rhs : eq.rhs - eq.lhs for eq in eqs] :
            [eq.rhs for eq in eqs]
 
+    if !isempty(assertions(sys))
+        rhss[end] += unwrap(get_assertions_expr(sys))
+    end
+
     # TODO: add an optional check on the ordering of observed equations
     u = dvs
     p = reorder_parameters(sys, ps)

test/debugging.jl

@@ -8,12 +8,23 @@ using ModelingToolkit: t_nounits as t, D_nounits as D, ASSERTION_LOG_VARIABLE
 sys_ode = structural_simplify(inner_ode)
 sys_sde = structural_simplify(inner_sde)
 
-@testset "`debug_system` adds assertions" begin
+@testset "assertions are present in generated `f`" begin
+    @testset "$(typeof(sys))" for (Problem, sys, alg) in [
+        (ODEProblem, sys_ode, Tsit5()), (SDEProblem, sys_sde, ImplicitEM())]
+        @test !is_parameter(sys, ASSERTION_LOG_VARIABLE)
+        prob = Problem(sys, [x => 0.1], (0.0, 5.0))
+        sol = solve(prob, alg)
+        @test !SciMLBase.successful_retcode(sol)
+        @test isnan(prob.f.f([0.0], prob.p, sol.t[end])[1])
+    end
+end
+
+@testset "`debug_system` adds logging" begin
     @testset "$(typeof(sys))" for (Problem, sys, alg) in [
         (ODEProblem, sys_ode, Tsit5()), (SDEProblem, sys_sde, ImplicitEM())]
         dsys = debug_system(sys; functions = [])
         @test is_parameter(dsys, ASSERTION_LOG_VARIABLE)
-        prob = Problem(dsys, [x => 1.0], (0.0, 5.0))
+        prob = Problem(dsys, [x => 0.1], (0.0, 5.0))
         sol = solve(prob, alg)
         @test !SciMLBase.successful_retcode(sol)
         prob.ps[ASSERTION_LOG_VARIABLE] = true
@@ -29,7 +40,7 @@ end
         @mtkbuild outer = ctor(Equation[], t; systems = [inner])
         dsys = debug_system(outer; functions = [])
         @test is_parameter(dsys, ASSERTION_LOG_VARIABLE)
-        prob = Problem(dsys, [inner.x => 1.0], (0.0, 5.0))
+        prob = Problem(dsys, [inner.x => 0.1], (0.0, 5.0))
         sol = solve(prob, alg)
         @test !SciMLBase.successful_retcode(sol)
         prob.ps[ASSERTION_LOG_VARIABLE] = true