feat: add assertions field to ODESystem

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -983,6 +983,7 @@ for prop in [:eqs
              :gui_metadata
              :discrete_subsystems
              :parameter_dependencies
+             :assertions
              :solved_unknowns
              :split_idxs
              :parent
@@ -1468,6 +1469,24 @@ end
 """
 $(TYPEDSIGNATURES)
 
+Get the assertions for a system `sys` and its subsystems.
+"""
+function assertions(sys::AbstractSystem)
+    has_assertions(sys) || return Dict{BasicSymbolic, String}()
+
+    asserts = get_assertions(sys)
+    systems = get_systems(sys)
+    namespaced_asserts = mapreduce(
+        merge!, systems; init = Dict{BasicSymbolic, String}()) do subsys
+        Dict{BasicSymbolic, String}(namespace_expr(k, subsys) => v
+        for (k, v) in assertions(subsys))
+    end
+    return merge(asserts, namespaced_asserts)
+end
+
+"""
+$(TYPEDSIGNATURES)
+
 Get the guesses for variables in the initialization system of the system `sys` and its subsystems.
 
 See also [`initialization_equations`](@ref) and [`ModelingToolkit.get_guesses`](@ref).
@@ -3036,6 +3055,11 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem;
         kwargs = merge(kwargs, (initialization_eqs = ieqs, guesses = guesses))
     end
 
+    if has_assertions(basesys)
+        kwargs = merge(
+            kwargs, (; assertions = merge(get_assertions(basesys), get_assertions(sys))))
+    end
+
     return T(args...; kwargs...)
 end
 

src/systems/diffeqs/odesystem.jl

@@ -137,6 +137,11 @@ struct ODESystem <: AbstractODESystem
     """
     parameter_dependencies::Vector{Equation}
     """
+    Mapping of conditions which should be true throughout the solve to corresponding error
+    messages. These will be added to the equations when calling `debug_system`.
+    """
+    assertions::Dict{BasicSymbolic, String}
+    """
     Metadata for the system, to be used by downstream packages.
     """
     metadata::Any
@@ -190,7 +195,7 @@ struct ODESystem <: AbstractODESystem
             jac, ctrl_jac, Wfact, Wfact_t, name, description, systems, defaults, guesses,
             torn_matching, initializesystem, initialization_eqs, schedule,
             connector_type, preface, cevents,
-            devents, parameter_dependencies,
+            devents, parameter_dependencies, assertions = Dict{BasicSymbolic, String}(),
             metadata = nothing, gui_metadata = nothing, is_dde = false,
             tstops = [], tearing_state = nothing,
             substitutions = nothing, complete = false, index_cache = nothing,
@@ -210,7 +215,7 @@ struct ODESystem <: AbstractODESystem
         new(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac, Wfact, Wfact_t, name, description, systems, defaults, guesses, torn_matching,
             initializesystem, initialization_eqs, schedule, connector_type, preface,
-            cevents, devents, parameter_dependencies, metadata,
+            cevents, devents, parameter_dependencies, assertions, metadata,
             gui_metadata, is_dde, tstops, tearing_state, substitutions, complete, index_cache,
             discrete_subsystems, solved_unknowns, split_idxs, parent)
     end
@@ -235,6 +240,7 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         continuous_events = nothing,
         discrete_events = nothing,
         parameter_dependencies = Equation[],
+        assertions = Dict(),
         checks = true,
         metadata = nothing,
         gui_metadata = nothing,
@@ -286,12 +292,13 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
     if is_dde === nothing
         is_dde = _check_if_dde(deqs, iv′, systems)
     end
+    assertions = Dict{BasicSymbolic, Any}(unwrap(k) => v for (k, v) in assertions)
     ODESystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
         deqs, iv′, dvs′, ps′, tspan, var_to_name, ctrl′, observed, tgrad, jac,
         ctrl_jac, Wfact, Wfact_t, name, description, systems,
         defaults, guesses, nothing, initializesystem,
         initialization_eqs, schedule, connector_type, preface, cont_callbacks,
-        disc_callbacks, parameter_dependencies,
+        disc_callbacks, parameter_dependencies, assertions,
         metadata, gui_metadata, is_dde, tstops, checks = checks)
 end
 
@@ -364,6 +371,7 @@ function flatten(sys::ODESystem, noeqs = false)
             name = nameof(sys),
             description = description(sys),
             initialization_eqs = initialization_equations(sys),
+            assertions = assertions(sys),
             is_dde = is_dde(sys),
             tstops = symbolic_tstops(sys),
             metadata = get_metadata(sys),