fix: compile symbolic affects after mtkcompile in complete

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -646,6 +646,12 @@ function complete(
         if add_initial_parameters
             sys = add_initialization_parameters(sys; split)
         end
+        if has_continuous_events(sys) && is_time_dependent(sys)
+            @set! sys.continuous_events = complete.(get_continuous_events(sys); iv = get_iv(sys), alg_eqs = [alg_equations(sys); observed(sys)])
+        end
+        if has_discrete_events(sys) && is_time_dependent(sys)
+            @set! sys.discrete_events = complete.(get_discrete_events(sys); iv = get_iv(sys), alg_eqs = [alg_equations(sys); observed(sys)])
+        end
     end
     if split && has_index_cache(sys)
         @set! sys.index_cache = IndexCache(sys)

src/systems/callbacks.jl

@@ -4,6 +4,15 @@ function has_functional_affect(cb)
     affects(cb) isa ImperativeAffect
 end
 
+struct SymbolicAffect{A, K}
+    affect::A
+    kwargs::K
+end
+
+SymbolicAffect(affect::Vector{Equation}; kwargs...) = SymbolicAffect(affect, kwargs)
+SymbolicAffect(affect::SymbolicAffect; kwargs...) = SymbolicAffect(affect.affect; affect.kwargs..., kwargs...)
+SymbolicAffect(affect; kwargs...) = affect
+
 struct AffectSystem
     """The internal implicit discrete system whose equations are solved to obtain values after the affect."""
     system::AbstractSystem
@@ -15,6 +24,71 @@ struct AffectSystem
     discretes::Vector
 end
 
+function AffectSystem(spec::SymbolicAffect; iv = nothing, alg_eqs = Equation[], kwargs...)
+    AffectSystem(spec.affect; iv = something(iv, get(spec.kwargs, :iv, nothing), Some(nothing)), alg_eqs = vcat(get(spec.kwargs, :alg_eqs, Equation[]), alg_eqs), kwargs...)
+end
+
+function AffectSystem(affect::Vector{Equation}; discrete_parameters = Any[],
+        iv = nothing, alg_eqs::Vector{Equation} = Equation[], warn_no_algebraic = true, kwargs...)
+    isempty(affect) && return nothing
+    if isnothing(iv)
+        iv = t_nounits
+        @warn "No independent variable specified. Defaulting to t_nounits."
+    end
+
+    discrete_parameters isa AbstractVector || (discrete_parameters = [discrete_parameters])
+    discrete_parameters = unwrap.(discrete_parameters)
+
+    for p in discrete_parameters
+        occursin(unwrap(iv), unwrap(p)) ||
+            error("Non-time dependent parameter $p passed in as a discrete. Must be declared as @parameters $p(t).")
+    end
+
+    dvs = OrderedSet()
+    params = OrderedSet()
+    _varsbuf = Set()
+    for eq in affect
+        if !haspre(eq) && !(symbolic_type(eq.rhs) === NotSymbolic() ||
+             symbolic_type(eq.lhs) === NotSymbolic())
+            @warn "Affect equation $eq has no `Pre` operator. As such it will be interpreted as an algebraic equation to be satisfied after the callback. If you intended to use the value of a variable x before the affect, use Pre(x). Errors may be thrown if there is no `Pre` and the algebraic equation is unsatisfiable, such as X ~ X + 1."
+        end
+        collect_vars!(dvs, params, eq, iv; op = Pre)
+        empty!(_varsbuf)
+        vars!(_varsbuf, eq; op = Pre)
+        filter!(x -> iscall(x) && operation(x) isa Pre, _varsbuf)
+        union!(params, _varsbuf)
+        diffvs = collect_applied_operators(eq, Differential)
+        union!(dvs, diffvs)
+    end
+    for eq in alg_eqs
+        collect_vars!(dvs, params, eq, iv)
+    end
+    pre_params = filter(haspre ∘ value, params)
+    sys_params = collect(setdiff(params, union(discrete_parameters, pre_params)))
+    discretes = map(tovar, discrete_parameters)
+    dvs = collect(dvs)
+    _dvs = map(default_toterm, dvs)
+
+    rev_map = Dict(zip(discrete_parameters, discretes))
+    subs = merge(rev_map, Dict(zip(dvs, _dvs)))
+    affect = Symbolics.fast_substitute(affect, subs)
+    alg_eqs = Symbolics.fast_substitute(alg_eqs, subs)
+
+    @named affectsys = System(
+        vcat(affect, alg_eqs), iv, collect(union(_dvs, discretes)),
+        collect(union(pre_params, sys_params)); is_discrete = true)
+    affectsys = mtkcompile(affectsys; fully_determined = nothing)
+    # get accessed parameters p from Pre(p) in the callback parameters
+    accessed_params = Vector{Any}(filter(isparameter, map(unPre, collect(pre_params))))
+    union!(accessed_params, sys_params)
+
+    # add scalarized unknowns to the map.
+    _dvs = reduce(vcat, map(scalarize, _dvs), init = Any[])
+
+    AffectSystem(affectsys, collect(_dvs), collect(accessed_params),
+        collect(discrete_parameters))
+end
+
 system(a::AffectSystem) = a.system
 discretes(a::AffectSystem) = a.discretes
 unknowns(a::AffectSystem) = a.unknowns
@@ -159,40 +233,40 @@ will run as soon as the solver starts, while finalization affects will be execut
 """
 struct SymbolicContinuousCallback <: AbstractCallback
     conditions::Vector{Equation}
-    affect::Union{Affect, Nothing}
-    affect_neg::Union{Affect, Nothing}
-    initialize::Union{Affect, Nothing}
-    finalize::Union{Affect, Nothing}
+    affect::Union{Affect, SymbolicAffect, Nothing}
+    affect_neg::Union{Affect, SymbolicAffect, Nothing}
+    initialize::Union{Affect, SymbolicAffect, Nothing}
+    finalize::Union{Affect, SymbolicAffect, Nothing}
     rootfind::Union{Nothing, SciMLBase.RootfindOpt}
     reinitializealg::SciMLBase.DAEInitializationAlgorithm
+end
 
-    function SymbolicContinuousCallback(
-            conditions::Union{Equation, Vector{Equation}},
-            affect = nothing;
-            affect_neg = affect,
-            initialize = nothing,
-            finalize = nothing,
-            rootfind = SciMLBase.LeftRootFind,
-            reinitializealg = nothing,
-            kwargs...)
-        conditions = (conditions isa AbstractVector) ? conditions : [conditions]
-
-        if isnothing(reinitializealg)
-            if any(a -> a isa ImperativeAffect,
-                [affect, affect_neg, initialize, finalize])
-                reinitializealg = SciMLBase.CheckInit()
-            else
-                reinitializealg = SciMLBase.NoInit()
-            end
+function SymbolicContinuousCallback(
+        conditions::Union{Equation, Vector{Equation}},
+        affect = nothing;
+        affect_neg = affect,
+        initialize = nothing,
+        finalize = nothing,
+        rootfind = SciMLBase.LeftRootFind,
+        reinitializealg = nothing,
+        kwargs...)
+    conditions = (conditions isa AbstractVector) ? conditions : [conditions]
+
+    if isnothing(reinitializealg)
+        if any(a -> a isa ImperativeAffect,
+            [affect, affect_neg, initialize, finalize])
+            reinitializealg = SciMLBase.CheckInit()
+        else
+            reinitializealg = SciMLBase.NoInit()
         end
+    end
 
-        new(conditions, make_affect(affect; kwargs...),
-            make_affect(affect_neg; kwargs...),
-            make_affect(initialize; kwargs...), make_affect(
-                finalize; kwargs...),
-            rootfind, reinitializealg)
-    end # Default affect to nothing
-end
+    SymbolicContinuousCallback(conditions, SymbolicAffect(affect; kwargs...),
+        SymbolicAffect(affect_neg; kwargs...),
+        SymbolicAffect(initialize; kwargs...), SymbolicAffect(
+            finalize; kwargs...),
+        rootfind, reinitializealg)
+end # Default affect to nothing
 
 function SymbolicContinuousCallback(p::Pair, args...; kwargs...)
     SymbolicContinuousCallback(p[1], p[2], args...; kwargs...)
@@ -207,72 +281,18 @@ function SymbolicContinuousCallback(cb::Tuple, args...; kwargs...)
     end
 end
 
+function complete(cb::SymbolicContinuousCallback; kwargs...)
+    SymbolicContinuousCallback(cb.conditions, make_affect(cb.affect; kwargs...),
+        make_affect(cb.affect_neg; kwargs...), make_affect(cb.initialize; kwargs...),
+        make_affect(cb.finalize; kwargs...), cb.rootfind, cb.reinitializealg)
+end
+
+make_affect(affect::SymbolicAffect; kwargs...) = AffectSystem(affect; kwargs...)
 make_affect(affect::Nothing; kwargs...) = nothing
 make_affect(affect::Tuple; kwargs...) = ImperativeAffect(affect...)
 make_affect(affect::NamedTuple; kwargs...) = ImperativeAffect(; affect...)
 make_affect(affect::Affect; kwargs...) = affect
 
-function make_affect(affect::Vector{Equation}; discrete_parameters = Any[],
-        iv = nothing, alg_eqs::Vector{Equation} = Equation[], warn_no_algebraic = true, kwargs...)
-    isempty(affect) && return nothing
-    if isnothing(iv)
-        iv = t_nounits
-        @warn "No independent variable specified. Defaulting to t_nounits."
-    end
-
-    discrete_parameters isa AbstractVector || (discrete_parameters = [discrete_parameters])
-    discrete_parameters = unwrap.(discrete_parameters)
-
-    for p in discrete_parameters
-        occursin(unwrap(iv), unwrap(p)) ||
-            error("Non-time dependent parameter $p passed in as a discrete. Must be declared as @parameters $p(t).")
-    end
-
-    dvs = OrderedSet()
-    params = OrderedSet()
-    _varsbuf = Set()
-    for eq in affect
-        if !haspre(eq) && !(symbolic_type(eq.rhs) === NotSymbolic() ||
-             symbolic_type(eq.lhs) === NotSymbolic())
-            @warn "Affect equation $eq has no `Pre` operator. As such it will be interpreted as an algebraic equation to be satisfied after the callback. If you intended to use the value of a variable x before the affect, use Pre(x). Errors may be thrown if there is no `Pre` and the algebraic equation is unsatisfiable, such as X ~ X + 1."
-        end
-        collect_vars!(dvs, params, eq, iv; op = Pre)
-        empty!(_varsbuf)
-        vars!(_varsbuf, eq; op = Pre)
-        filter!(x -> iscall(x) && operation(x) isa Pre, _varsbuf)
-        union!(params, _varsbuf)
-        diffvs = collect_applied_operators(eq, Differential)
-        union!(dvs, diffvs)
-    end
-    for eq in alg_eqs
-        collect_vars!(dvs, params, eq, iv)
-    end
-    pre_params = filter(haspre ∘ value, params)
-    sys_params = collect(setdiff(params, union(discrete_parameters, pre_params)))
-    discretes = map(tovar, discrete_parameters)
-    dvs = collect(dvs)
-    _dvs = map(default_toterm, dvs)
-
-    rev_map = Dict(zip(discrete_parameters, discretes))
-    subs = merge(rev_map, Dict(zip(dvs, _dvs)))
-    affect = Symbolics.fast_substitute(affect, subs)
-    alg_eqs = Symbolics.fast_substitute(alg_eqs, subs)
-
-    @named affectsys = System(
-        vcat(affect, alg_eqs), iv, collect(union(_dvs, discretes)),
-        collect(union(pre_params, sys_params)); is_discrete = true)
-    affectsys = mtkcompile(affectsys; fully_determined = nothing)
-    # get accessed parameters p from Pre(p) in the callback parameters
-    accessed_params = Vector{Any}(filter(isparameter, map(unPre, collect(pre_params))))
-    union!(accessed_params, sys_params)
-
-    # add scalarized unknowns to the map.
-    _dvs = reduce(vcat, map(scalarize, _dvs), init = Any[])
-
-    AffectSystem(affectsys, collect(_dvs), collect(accessed_params),
-        collect(discrete_parameters))
-end
-
 function make_affect(affect; kwargs...)
     error("Malformed affect $(affect). This should be a vector of equations or a tuple specifying a functional affect.")
 end
@@ -374,30 +394,30 @@ Arguments:
 """
 struct SymbolicDiscreteCallback <: AbstractCallback
     conditions::Union{Number, Vector{<:Number}, Symbolic{Bool}}
-    affect::Union{Affect, Nothing}
-    initialize::Union{Affect, Nothing}
-    finalize::Union{Affect, Nothing}
+    affect::Union{Affect, SymbolicAffect, Nothing}
+    initialize::Union{Affect, SymbolicAffect, Nothing}
+    finalize::Union{Affect, SymbolicAffect, Nothing}
     reinitializealg::SciMLBase.DAEInitializationAlgorithm
+end
 
-    function SymbolicDiscreteCallback(
-            condition::Union{Symbolic{Bool}, Number, Vector{<:Number}}, affect = nothing;
-            initialize = nothing, finalize = nothing,
-            reinitializealg = nothing, kwargs...)
-        c = is_timed_condition(condition) ? condition : value(scalarize(condition))
-
-        if isnothing(reinitializealg)
-            if any(a -> a isa ImperativeAffect,
-                [affect, initialize, finalize])
-                reinitializealg = SciMLBase.CheckInit()
-            else
-                reinitializealg = SciMLBase.NoInit()
-            end
+function SymbolicDiscreteCallback(
+        condition::Union{Symbolic{Bool}, Number, Vector{<:Number}}, affect = nothing;
+        initialize = nothing, finalize = nothing,
+        reinitializealg = nothing, kwargs...)
+    c = is_timed_condition(condition) ? condition : value(scalarize(condition))
+
+    if isnothing(reinitializealg)
+        if any(a -> a isa ImperativeAffect,
+            [affect, initialize, finalize])
+            reinitializealg = SciMLBase.CheckInit()
+        else
+            reinitializealg = SciMLBase.NoInit()
         end
-        new(c, make_affect(affect; kwargs...),
-            make_affect(initialize; kwargs...),
-            make_affect(finalize; kwargs...), reinitializealg)
-    end # Default affect to nothing
-end
+    end
+    SymbolicDiscreteCallback(c, SymbolicAffect(affect; kwargs...),
+        SymbolicAffect(initialize; kwargs...),
+        SymbolicAffect(finalize; kwargs...), reinitializealg)
+end # Default affect to nothing
 
 function SymbolicDiscreteCallback(p::Pair, args...; kwargs...)
     SymbolicDiscreteCallback(p[1], p[2], args...; kwargs...)
@@ -412,6 +432,10 @@ function SymbolicDiscreteCallback(cb::Tuple, args...; kwargs...)
     end
 end
 
+function complete(cb::SymbolicDiscreteCallback; kwargs...)
+    SymbolicDiscreteCallback(cb.conditions, make_affect(cb.affect; kwargs...), make_affect(cb.initialize; kwargs...), make_affect(cb.finalize; kwargs...), cb.reinitializealg)
+end
+
 function is_timed_condition(condition::T) where {T}
     if T === Num
         false
@@ -1060,12 +1084,8 @@ end
 """
 Process the symbolic events of a system.
 """
-function create_symbolic_events(cont_events, disc_events, sys_eqs, iv)
-    alg_eqs = filter(eq -> eq.lhs isa Union{Symbolic, Number} && !is_diff_equation(eq),
-        sys_eqs)
-    cont_callbacks = to_cb_vector(cont_events; CB_TYPE = SymbolicContinuousCallback,
-        iv = iv, alg_eqs = alg_eqs, warn_no_algebraic = false)
-    disc_callbacks = to_cb_vector(disc_events; CB_TYPE = SymbolicDiscreteCallback,
-        iv = iv, alg_eqs = alg_eqs, warn_no_algebraic = false)
+function create_symbolic_events(cont_events, disc_events)
+    cont_callbacks = to_cb_vector(cont_events; CB_TYPE = SymbolicContinuousCallback)
+    disc_callbacks = to_cb_vector(disc_events; CB_TYPE = SymbolicDiscreteCallback)
     cont_callbacks, disc_callbacks
 end

src/systems/system.jl

@@ -389,7 +389,7 @@ function System(eqs::Vector{Equation}, iv, dvs, ps, brownians = [];
     end
     continuous_events,
     discrete_events = create_symbolic_events(
-        continuous_events, discrete_events, eqs, iv)
+        continuous_events, discrete_events)
 
     if iv === nothing && (!isempty(continuous_events) || !isempty(discrete_events))
         throw(EventsInTimeIndependentSystemError(continuous_events, discrete_events))

test/symbolic_events.jl

@@ -1378,3 +1378,28 @@ end
     @test SciMLBase.successful_retcode(sol)
     @test sol[x, end]≈1.0 atol=1e-6
 end
+
+@testset "Symbolic affects are compiled in `complete`" begin
+    @parameters g
+    @variables x(t) [state_priority = 10.0] y(t) [guess = 1.0]
+    @variables λ(t) [guess = 1.0]
+    eqs = [D(D(x)) ~ λ * x
+           D(D(y)) ~ λ * y - g
+           x^2 + y^2 ~ 1]
+    cevts = [[x ~ 0.0] => [D(x) ~ Pre(D(x)) + 1sign(Pre(D(x)))]]
+    @named pend = System(eqs, t; continuous_events = cevts)
+
+    scc = only(continuous_events(pend))
+    @test scc.affect isa ModelingToolkit.SymbolicAffect
+
+    pend = mtkcompile(pend)
+
+    scc = only(continuous_events(pend))
+    @test scc.affect isa ModelingToolkit.AffectSystem
+    @test length(ModelingToolkit.all_equations(scc.affect)) == 5 # 1 affect, 3 algebraic, 1 observed
+
+    u0 = [x => -1/2, D(x) => 1/2, g => 1]
+    prob = ODEProblem(pend, u0, (0.0, 5.0))
+    sol = solve(prob, FBDF())
+    @test SciMLBase.successful_retcode(sol)
+end