refactor: get System to precompile in a trivial case

AayushSabharwal · AayushSabharwal · commit 5cad99ff2924 · 2025-10-01T18:19:42.000+05:30
diff --git a/src/systems/callbacks.jl b/src/systems/callbacks.jl
@@ -568,6 +568,9 @@ conditions(cb::AbstractCallback) = cb.conditions
 function conditions(cbs::Vector{<:AbstractCallback})
     reduce(vcat, conditions(cb) for cb in cbs; init = [])
 end
+function conditions(cbs::Vector{SymbolicContinuousCallback})
+    mapreduce(conditions, vcat, cbs; init = Equation[])
+end
 equations(cb::AbstractCallback) = conditions(cb)
 equations(cb::Vector{<:AbstractCallback}) = conditions(cb)
 
diff --git a/src/systems/system.jl b/src/systems/system.jl
@@ -46,7 +46,7 @@ struct System <: IntermediateDeprecationSystem
     this noise matrix is diagonal. Diagonal noise can be specified by providing an `N`
     length vector. If this field is `nothing`, the system does not have noise.
     """
-    noise_eqs::Union{Nothing, AbstractVector, AbstractMatrix}
+    noise_eqs::Union{Nothing, Vector{SymbolicT}, Matrix{SymbolicT}}
     """
     Jumps associated with the system. Each jump can be a `VariableRateJump`,
     `ConstantRateJump` or `MassActionJump`. See `JumpProcesses.jl` for more information.
@@ -279,30 +279,37 @@ struct System <: IntermediateDeprecationSystem
             """))
         end
         @assert iv === nothing || symtype(iv) === Real
-        jumps = Vector{JumpType}(jumps)
-        if (checks == true || (checks & CheckComponents) > 0) && iv !== nothing
-            check_independent_variables([iv])
+        if (checks isa Bool && checks === true || checks isa Int && (checks & CheckComponents) > 0) && iv !== nothing
+            check_independent_variables((iv,))
             check_variables(unknowns, iv)
             check_parameters(ps, iv)
             check_equations(eqs, iv)
-            if noise_eqs !== nothing && size(noise_eqs, 1) != length(eqs)
-                throw(IllFormedNoiseEquationsError(size(noise_eqs, 1), length(eqs)))
+            Neq = length(eqs)
+            if noise_eqs isa Matrix{SymbolicT}
+                N1 = size(noise_eqs, 1)
+            elseif noise_eqs isa Vector{SymbolicT}
+                N1 = length(noise_eqs)
+            elseif noise_eqs === nothing
+                N1 = Neq
+            else
+                error()
             end
+            N1 == Neq || throw(IllFormedNoiseEquationsError(N1, Neq))
             check_equations(equations(continuous_events), iv)
             check_subsystems(systems)
         end
-        if checks == true || (checks & CheckUnits) > 0
-            u = __get_unit_type(unknowns, ps, iv)
-            if noise_eqs === nothing
-                check_units(u, eqs)
-            else
-                check_units(u, eqs, noise_eqs)
-            end
-            if iv !== nothing
-                check_units(u, jumps, iv)
-            end
-            isempty(constraints) || check_units(u, constraints)
-        end
+        # if checks == true || (checks & CheckUnits) > 0
+        #     u = __get_unit_type(unknowns, ps, iv)
+        #     if noise_eqs === nothing
+        #         check_units(u, eqs)
+        #     else
+        #         check_units(u, eqs, noise_eqs)
+        #     end
+        #     if iv !== nothing
+        #         check_units(u, jumps, iv)
+        #     end
+        #     isempty(constraints) || check_units(u, constraints)
+        # end
         new(tag, eqs, noise_eqs, jumps, constraints, costs,
             consolidate, unknowns, ps, brownians, iv,
             observed, parameter_dependencies, var_to_name, name, description, defaults,
@@ -321,13 +328,11 @@ function default_consolidate(costs, subcosts)
     return reduce(+, costs; init = 0.0) + reduce(+, subcosts; init = 0.0)
 end
 
-function unwrap_vars(vars::AbstractArray{SymbolicT})
-    vec(vars)
-end
-function unwrap_vars(vars)
-    result = SymbolicT[]
-    for var in vars
-        push!(result, unwrap(var))
+unwrap_vars(vars::AbstractArray{SymbolicT}) = vars
+function unwrap_vars(vars::AbstractArray)
+    result = similar(vars, SymbolicT)
+    for i in eachindex(vars)
+        result[i] = SU.Const{VartypeT}(vars[i])
     end
     return result
 end
@@ -372,29 +377,30 @@ function System(eqs::Vector{Equation}, iv, dvs, ps, brownians = SymbolicT[];
         initializesystem = nothing, is_initializesystem = false, is_discrete = false,
         preface = [], checks = true)
     name === nothing && throw(NoNameError())
+
+    if !(eqs isa Vector{Equation})
+        eqs = Equation[eqs]
+    end
+    eqs = eqs::Vector{Equation}
+
     if !isempty(parameter_dependencies)
         @invokelatest warn_pdeps()
-        eqs = Equation[eqs; parameter_dependencies]
+        append!(eqs, parameter_dependencies)
     end
 
     iv = unwrap(iv)
-    ps = unwrap_vars(ps)
-    dvs = unwrap_vars(dvs)
+    ps = vec(unwrap_vars(ps))
+    dvs = vec(unwrap_vars(dvs))
     if iv !== nothing
         filter!(!Base.Fix2(isdelay, iv), dvs)
     end
     brownians = unwrap_vars(brownians)
 
-    if !(eqs isa Vector{Equation})
-        eqs = Equation[eqs]
-    end
-    eqs = eqs::Vector{Equation}
-
     if noise_eqs !== nothing
-        noise_eqs = unwrap.(noise_eqs)
+        noise_eqs = unwrap_vars(noise_eqs)
     end
 
-    costs = unwrap_vars(costs)
+    costs = vec(unwrap_vars(costs))
 
     defaults = defsdict(defaults)
     guesses = defsdict(guesses)
@@ -421,8 +427,12 @@ function System(eqs::Vector{Equation}, iv, dvs, ps, brownians = SymbolicT[];
 
         process_variables!(var_to_name, defaults, guesses, dvs)
         process_variables!(var_to_name, defaults, guesses, ps)
-        process_variables!(var_to_name, defaults, guesses, SymbolicT[eq.lhs for eq in observed])
-        process_variables!(var_to_name, defaults, guesses, SymbolicT[eq.rhs for eq in observed])
+        buffer = SymbolicT[]
+        for eq in observed
+            push!(buffer, eq.lhs)
+            push!(buffer, eq.rhs)
+        end
+        process_variables!(var_to_name, defaults, guesses, buffer)
 
         for var in dvs
             if isinput(var)
@@ -435,10 +445,9 @@ function System(eqs::Vector{Equation}, iv, dvs, ps, brownians = SymbolicT[];
     filter!(!(Base.Fix1(===, COMMON_NOTHING) ∘ last), defaults)
     filter!(!(Base.Fix1(===, COMMON_NOTHING) ∘ last), guesses)
 
-    sysnames = nameof.(systems)
-    unique_sysnames = Set(sysnames)
-    if length(unique_sysnames) != length(sysnames)
-        throw(NonUniqueSubsystemsError(sysnames, unique_sysnames))
+
+    if !allunique(map(nameof, systems))
+        nonunique_subsystems(systems)
     end
     continuous_events,
     discrete_events = create_symbolic_events(
@@ -452,7 +461,10 @@ function System(eqs::Vector{Equation}, iv, dvs, ps, brownians = SymbolicT[];
         is_dde = _check_if_dde(eqs, iv, systems)
     end
 
-    assertions = Dict{SymbolicT, String}(unwrap(k) => v for (k, v) in assertions)
+    _assertions = Dict{SymbolicT, String}
+    for (k, v) in assertions
+        _assertions[unwrap(k)::SymbolicT] = v
+    end
 
     if isempty(metadata)
         metadata = MetadataT()
@@ -466,6 +478,7 @@ function System(eqs::Vector{Equation}, iv, dvs, ps, brownians = SymbolicT[];
         metadata = meta
     end
     metadata = refreshed_metadata(metadata)
+    jumps = Vector{JumpType}(jumps)
     System(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)), eqs, noise_eqs, jumps, constraints,
         costs, consolidate, dvs, ps, brownians, iv, observed, Equation[],
         var_to_name, name, description, defaults, guesses, systems, initialization_eqs,
@@ -475,6 +488,12 @@ function System(eqs::Vector{Equation}, iv, dvs, ps, brownians = SymbolicT[];
         initializesystem, is_initializesystem, is_discrete; checks)
 end
 
+@noinline function nonunique_subsystems(systems)
+    sysnames = nameof.(systems)
+    unique_sysnames = Set(sysnames)
+    throw(NonUniqueSubsystemsError(sysnames, unique_sysnames))
+end
+
 @noinline function warn_pdeps()
     @warn """
     The `parameter_dependencies` keyword argument is deprecated. Please provide all
@@ -749,19 +768,15 @@ differential equations.
 """
 is_dde(sys::AbstractSystem) = has_is_dde(sys) && get_is_dde(sys)
 
-function _check_if_dde(eqs, iv, subsystems)
-    is_dde = any(ModelingToolkit.is_dde, subsystems)
-    if !is_dde
-        vs = Set()
-        for eq in eqs
-            vars!(vs, eq)
-            is_dde = any(vs) do sym
-                isdelay(unwrap(sym), iv)
-            end
-            is_dde && break
-        end
+_check_if_dde(eqs::Vector{Equation}, iv::Nothing, subsystems::Vector{System}) = false
+function _check_if_dde(eqs::Vector{Equation}, iv::SymbolicT, subsystems::Vector{System})
+    any(ModelingToolkit.is_dde, subsystems) && return true
+    pred = Base.Fix2(isdelay, iv)
+    for eq in eqs
+        SU.query!(pred, eq.lhs) && return true
+        SU.query!(pred, eq.rhs) && return true
     end
-    return is_dde
+    return false
 end
 
 """
diff --git a/src/utils.jl b/src/utils.jl
@@ -117,34 +117,31 @@ const CheckUnits = 1 << 2
 
 function check_independent_variables(ivs)
     for iv in ivs
-        isparameter(iv) ||
-            @warn "Independent variable $iv should be defined with @independent_variables $iv."
+        isparameter(iv) || @invokelatest warn_indepvar(iv)
     end
 end
 
+@noinline function warn_indepvar(iv::SymbolicT)
+    @warn "Independent variable $iv should be defined with @independent_variables $iv."
+end
+
 function check_parameters(ps, iv)
     for p in ps
         isequal(iv, p) &&
             throw(ArgumentError("Independent variable $iv not allowed in parameters."))
     end
 end
 
-function is_delay_var(iv, var)
-    if Symbolics.isarraysymbolic(var)
-        return is_delay_var(iv, first(collect(var)))
-    end
-    args = nothing
-    try
-        args = arguments(var)
-    catch
-        return false
+function is_delay_var(iv::SymbolicT, var::SymbolicT)
+    Moshi.Match.@match var begin
+        BSImpl.Term(; f, args) => begin
+            length(args) > 1 && return false
+            arg = args[1]
+            isequal(arg, iv) && return false
+            return symtype(arg) <: Real
+        end
+        _ => false
     end
-    length(args) > 1 && return false
-    isequal(first(args), iv) && return false
-    delay = iv - first(args)
-    delay isa Integer ||
-        delay isa AbstractFloat ||
-        (delay isa Num && isreal(value(delay)))
 end
 
 function check_variables(dvs, iv)
@@ -187,20 +184,35 @@ function collect_ivs(eqs, op = Differential)
     return ivs
 end
 
+struct IndepvarCheckPredicate
+    iv::SymbolicT
+end
+
+function (icp::IndepvarCheckPredicate)(ex::SymbolicT)
+    Moshi.Match.@match ex begin
+        BSImpl.Term(; f) && if f isa Differential end => begin
+            f = f::Differential
+            isequal(f.x, icp.iv) || throw_multiple_iv(icp.iv, f.x)
+            return false
+        end
+        _ => false
+    end
+end
+
+@noinline function throw_multiple_iv(iv, newiv)
+    throw(ArgumentError("Differential w.r.t. variable ($newiv) other than the independent variable ($iv) are not allowed."))
+end
+
 """
     check_equations(eqs, iv)
 
 Assert that equations are well-formed when building ODE, i.e., only containing a single independent variable.
 """
-function check_equations(eqs, iv)
-    ivs = collect_ivs(eqs)
-    display = collect(ivs)
-    length(ivs) <= 1 ||
-        throw(ArgumentError("Differential w.r.t. multiple variables $display are not allowed."))
-    if length(ivs) == 1
-        single_iv = pop!(ivs)
-        isequal(single_iv, iv) ||
-            throw(ArgumentError("Differential w.r.t. variable ($single_iv) other than the independent variable ($iv) are not allowed."))
+function check_equations(eqs::Vector{Equation}, iv::SymbolicT)
+    icp = IndepvarCheckPredicate(iv)
+    for eq in eqs
+        SU.query!(icp, eq.lhs)
+        SU.query!(icp, eq.rhs)
     end
 end
 
@@ -211,10 +223,12 @@ Assert that the subsystems have the appropriate namespacing behavior.
 """
 function check_subsystems(systems)
     idxs = findall(!does_namespacing, systems)
-    if !isempty(idxs)
-        names = join("  " .* string.(nameof.(systems[idxs])), "\n")
-        throw(ArgumentError("All subsystems must have namespacing enabled. The following subsystems do not perform namespacing:\n$(names)"))
-    end
+    isempty(idxs) || throw_bad_namespacing(systems, idxs)
+end
+
+@noinline function throw_bad_namespacing(systems, idxs)
+    names = join("  " .* string.(nameof.(systems[idxs])), "\n")
+    throw(ArgumentError("All subsystems must have namespacing enabled. The following subsystems do not perform namespacing:\n$(names)"))
 end
 
 """
@@ -626,6 +640,7 @@ function collect_vars!(unknowns, parameters, expr, iv; depth = 0, op = Symbolics
     if issym(expr)
         return collect_var!(unknowns, parameters, expr, iv; depth)
     end
+    SymbolicUtils.isconst(expr) && return
     for var in vars(expr; op)
         while iscall(var) && operation(var) isa op
             validate_operator(operation(var), arguments(var), iv; context = expr)
