[WIP] feat: reduce reliance on metadata in structural_simplify

ext/MTKBifurcationKitExt.jl

@@ -23,7 +23,7 @@ struct ObservableRecordFromSolution{S, T}
     # A Vector of pairs (Symbolic => value) with the default values of all system variables and parameters.
     subs_vals::T
 
-    function ObservableRecordFromSolution(nsys::NonlinearSystem,
+    function ObservableRecordFromSolution(nsys::System,
             plot_var,
             bif_idx,
             u0_vals,
@@ -82,7 +82,7 @@ end
 ### Creates BifurcationProblem Overloads ###
 
 # When input is a NonlinearSystem.
-function BifurcationKit.BifurcationProblem(nsys::NonlinearSystem,
+function BifurcationKit.BifurcationProblem(nsys::System,
         u0_bif,
         ps,
         bif_par,
@@ -92,7 +92,7 @@ function BifurcationKit.BifurcationProblem(nsys::NonlinearSystem,
         jac = true,
         kwargs...)
     if !ModelingToolkit.iscomplete(nsys)
-        error("A completed `NonlinearSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `BifurcationProblem`")
+        error("A completed `System` is required. Call `complete` or `structural_simplify` on the system before creating a `BifurcationProblem`")
     end
     @set! nsys.index_cache = nothing # force usage of a parameter vector instead of `MTKParameters`
     # Creates F and J functions.
@@ -144,11 +144,11 @@ function BifurcationKit.BifurcationProblem(nsys::NonlinearSystem,
 end
 
 # When input is a ODESystem.
-function BifurcationKit.BifurcationProblem(osys::ODESystem, args...; kwargs...)
+function BifurcationKit.BifurcationProblem(osys::System, args...; kwargs...)
     if !ModelingToolkit.iscomplete(osys)
         error("A completed `ODESystem` is required. Call `complete` or `structural_simplify` on the system before creating a `BifurcationProblem`")
     end
-    nsys = NonlinearSystem([0 ~ eq.rhs for eq in full_equations(osys)],
+    nsys = System([0 ~ eq.rhs for eq in full_equations(osys)],
         unknowns(osys),
         parameters(osys);
         observed = observed(osys),

ext/MTKCasADiDynamicOptExt.jl

@@ -56,9 +56,9 @@ function (M::MXLinearInterpolation)(τ)
 end
 
 """
-    CasADiDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt, steps)
+    CasADiDynamicOptProblem(sys::System, u0, tspan, p; dt, steps)
 
-Convert an ODESystem representing an optimal control system into a CasADi model
+Convert an System representing an optimal control system into a CasADi model
 for solving using optimization. Must provide either `dt`, the timestep between collocation 
 points (which, along with the timespan, determines the number of points), or directly 
 provide the number of points as `steps`.
@@ -68,10 +68,10 @@ The optimization variables:
 - a vector-of-vectors V representing the controls as an interpolation array
 
 The constraints are:
-- The set of user constraints passed to the ODESystem via `constraints`
+- The set of user constraints passed to the System via `constraints`
 - The solver constraints that encode the time-stepping used by the solver
 """
-function MTK.CasADiDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
+function MTK.CasADiDynamicOptProblem(sys::System, u0map, tspan, pmap;
         dt = nothing,
         steps = nothing,
         guesses = Dict(), kwargs...)
@@ -80,7 +80,8 @@ function MTK.CasADiDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
     f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, output_type = MX, kwargs...)
 
-    pmap = Dict{Any, Any}(pmap)
+    pmap = MTK.recursive_unwrap(MTK.AnyDict(pmap))
+    MTK.evaluate_varmap!(pmap, keys(pmap))
     steps, is_free_t = MTK.process_tspan(tspan, dt, steps)
     model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
 
@@ -143,15 +144,15 @@ function set_casadi_bounds!(model, sys, pmap)
     for (i, u) in enumerate(unknowns(sys))
         if MTK.hasbounds(u)
             lo, hi = MTK.getbounds(u)
-            subject_to!(opti, Symbolics.fixpoint_sub(lo, pmap) <= U.u[i, :])
-            subject_to!(opti, U.u[i, :] <= Symbolics.fixpoint_sub(hi, pmap))
+            subject_to!(opti, Symbolics.fast_substitute(lo, pmap) <= U.u[i, :])
+            subject_to!(opti, U.u[i, :] <= Symbolics.fast_substitute(hi, pmap))
         end
     end
     for (i, v) in enumerate(MTK.unbound_inputs(sys))
         if MTK.hasbounds(v)
             lo, hi = MTK.getbounds(v)
-            subject_to!(opti, Symbolics.fixpoint_sub(lo, pmap) <= V.u[i, :])
-            subject_to!(opti, V.u[i, :] <= Symbolics.fixpoint_sub(hi, pmap))
+            subject_to!(opti, Symbolics.fast_substitute(lo, pmap) <= V.u[i, :])
+            subject_to!(opti, V.u[i, :] <= Symbolics.fast_substitute(hi, pmap))
         end
     end
 end
@@ -167,15 +168,15 @@ function add_user_constraints!(model::CasADiModel, sys, tspan, pmap; is_free_t)
     @unpack opti, U, V, tₛ = model
 
     iv = MTK.get_iv(sys)
-    conssys = MTK.get_constraintsystem(sys)
-    jconstraints = isnothing(conssys) ? nothing : MTK.get_constraints(conssys)
+    jconstraints = MTK.get_constraints(sys)
     (isnothing(jconstraints) || isempty(jconstraints)) && return nothing
 
     stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
     ctidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))])
-    cons_unknowns = map(MTK.default_toterm, unknowns(conssys))
+    cons_dvs, cons_ps = MTK.process_constraint_system(
+        jconstraints, Set(unknowns(sys)), parameters(sys), iv; validate = false)
 
-    auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in unknowns(conssys)])
+    auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in cons_dvs])
     jconstraints = substitute_casadi_vars(model, sys, pmap, jconstraints; is_free_t, auxmap)
     # Manually substitute fixed-t variables
     for (i, cons) in enumerate(jconstraints)
@@ -207,9 +208,8 @@ end
 
 function add_cost_function!(model::CasADiModel, sys, tspan, pmap; is_free_t)
     @unpack opti, U, V, tₛ = model
-    jcosts = copy(MTK.get_costs(sys))
-    consolidate = MTK.get_consolidate(sys)
-    if isnothing(jcosts) || isempty(jcosts)
+    jcosts = cost(sys)
+    if Symbolics._iszero(jcosts)
         minimize!(opti, MX(0))
         return
     end
@@ -218,24 +218,22 @@ function add_cost_function!(model::CasADiModel, sys, tspan, pmap; is_free_t)
     stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
     ctidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))])
 
-    jcosts = substitute_casadi_vars(model, sys, pmap, jcosts; is_free_t)
+    jcosts = substitute_casadi_vars(model, sys, pmap, [jcosts]; is_free_t)[1]
     # Substitute fixed-time variables.
-    for i in 1:length(jcosts)
-        costvars = MTK.vars(jcosts[i])
-        for st in costvars
-            MTK.iscall(st) || continue
-            x = operation(st)
-            t = only(arguments(st))
-            MTK.symbolic_type(t) === MTK.NotSymbolic() || continue
-            if haskey(stidxmap, x(iv))
-                idx = stidxmap[x(iv)]
-                cv = U
-            else
-                idx = ctidxmap[x(iv)]
-                cv = V
-            end
-            jcosts[i] = Symbolics.substitute(jcosts[i], Dict(x(t) => cv(t)[idx]))
+    costvars = MTK.vars(jcosts)
+    for st in costvars
+        MTK.iscall(st) || continue
+        x = operation(st)
+        t = only(arguments(st))
+        MTK.symbolic_type(t) === MTK.NotSymbolic() || continue
+        if haskey(stidxmap, x(iv))
+            idx = stidxmap[x(iv)]
+            cv = U
+        else
+            idx = ctidxmap[x(iv)]
+            cv = V
         end
+        jcosts = Symbolics.substitute(jcosts, Dict(x(t) => cv(t)[idx]))
     end
 
     dt = U.t[2] - U.t[1]
@@ -249,9 +247,9 @@ function add_cost_function!(model::CasADiModel, sys, tspan, pmap; is_free_t)
         # Approximate integral as sum.
         intmap[int] = dt * tₛ * sum(arg)
     end
-    jcosts = map(c -> Symbolics.substitute(c, intmap), jcosts)
-    jcosts = MTK.value.(jcosts)
-    minimize!(opti, MX(MTK.value(consolidate(jcosts))))
+    jcosts = Symbolics.substitute(jcosts, intmap)
+    jcosts = MTK.value(jcosts)
+    minimize!(opti, MX(jcosts))
 end
 
 function substitute_casadi_vars(
@@ -264,20 +262,20 @@ function substitute_casadi_vars(
     x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
     c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
 
-    exprs = map(c -> Symbolics.fixpoint_sub(c, auxmap), exprs)
-    exprs = map(c -> Symbolics.fixpoint_sub(c, Dict(pmap)), exprs)
+    exprs = map(c -> Symbolics.fast_substitute(c, auxmap), exprs)
+    exprs = map(c -> Symbolics.fast_substitute(c, Dict(pmap)), exprs)
     # tf means different things in different contexts; a [tf] in a cost function
     # should be tₛ, while a x(tf) should translate to x[1]
     if is_free_t
         free_t_map = Dict([[x(tₛ) => U.u[i, end] for (i, x) in enumerate(x_ops)];
                            [c(tₛ) => V.u[i, end] for (i, c) in enumerate(c_ops)]])
-        exprs = map(c -> Symbolics.fixpoint_sub(c, free_t_map), exprs)
+        exprs = map(c -> Symbolics.fast_substitute(c, free_t_map), exprs)
     end
 
     # for variables like x(t)
     whole_interval_map = Dict([[v => U.u[i, :] for (i, v) in enumerate(sts)];
                                [v => V.u[i, :] for (i, v) in enumerate(cts)]])
-    exprs = map(c -> Symbolics.fixpoint_sub(c, whole_interval_map), exprs)
+    exprs = map(c -> Symbolics.fast_substitute(c, whole_interval_map), exprs)
     exprs
 end
 

ext/MTKFMIExt.jl

@@ -289,7 +289,7 @@ function MTK.FMIComponent(::Val{Ver}; fmu = nothing, tolerance = 1e-6,
     end
 
     eqs = [observed; diffeqs]
-    return ODESystem(eqs, t, states, params; parameter_dependencies, defaults = defs,
+    return System(eqs, t, states, params; parameter_dependencies, defaults = defs,
         discrete_events = [instance_management_callback], name, initialization_eqs)
 end
 

ext/MTKInfiniteOptExt.jl

@@ -39,9 +39,9 @@ struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
 end
 
 """
-    JuMPDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt)
+    JuMPDynamicOptProblem(sys::System, u0, tspan, p; dt)
 
-Convert an ODESystem representing an optimal control system into a JuMP model
+Convert a System representing an optimal control system into a JuMP model
 for solving using optimization. Must provide either `dt`, the timestep between collocation 
 points (which, along with the timespan, determines the number of points), or directly 
 provide the number of points as `steps`.
@@ -51,10 +51,10 @@ The optimization variables:
 - a vector-of-vectors V representing the controls as an interpolation array
 
 The constraints are:
-- The set of user constraints passed to the ODESystem via `constraints`
+- The set of user constraints passed to the System via `constraints`
 - The solver constraints that encode the time-stepping used by the solver
 """
-function MTK.JuMPDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
+function MTK.JuMPDynamicOptProblem(sys::System, u0map, tspan, pmap;
         dt = nothing,
         steps = nothing,
         guesses = Dict(), kwargs...)
@@ -63,24 +63,25 @@ function MTK.JuMPDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
     f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    pmap = Dict{Any, Any}(pmap)
+    pmap = MTK.recursive_unwrap(MTK.AnyDict(pmap))
+    MTK.evaluate_varmap!(pmap, keys(pmap))
     steps, is_free_t = MTK.process_tspan(tspan, dt, steps)
     model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
 
     JuMPDynamicOptProblem(f, u0, tspan, p, model, kwargs...)
 end
 
 """
-    InfiniteOptDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; dt)
+    InfiniteOptDynamicOptProblem(sys::System, u0map, tspan, pmap; dt)
 
-Convert an ODESystem representing an optimal control system into a InfiniteOpt model
+Convert System representing an optimal control system into a InfiniteOpt model
 for solving using optimization. Must provide `dt` for determining the length 
 of the interpolation arrays.
 
 Related to `JuMPDynamicOptProblem`, but directly adds the differential equations
 of the system as derivative constraints, rather than using a solver tableau.
 """
-function MTK.InfiniteOptDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
+function MTK.InfiniteOptDynamicOptProblem(sys::System, u0map, tspan, pmap;
         dt = nothing,
         steps = nothing,
         guesses = Dict(), kwargs...)
@@ -89,7 +90,8 @@ function MTK.InfiniteOptDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
     f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    pmap = Dict{Any, Any}(pmap)
+    pmap = MTK.recursive_unwrap(MTK.AnyDict(pmap))
+    MTK.evaluate_varmap!(pmap, keys(pmap))
     steps, is_free_t = MTK.process_tspan(tspan, dt, steps)
     model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
 
@@ -150,29 +152,28 @@ function set_jump_bounds!(model, sys, pmap)
     for (i, u) in enumerate(unknowns(sys))
         if MTK.hasbounds(u)
             lo, hi = MTK.getbounds(u)
-            set_lower_bound(U[i], Symbolics.fixpoint_sub(lo, pmap))
-            set_upper_bound(U[i], Symbolics.fixpoint_sub(hi, pmap))
+            set_lower_bound(U[i], Symbolics.fast_substitute(lo, pmap))
+            set_upper_bound(U[i], Symbolics.fast_substitute(hi, pmap))
         end
     end
 
     V = model[:V]
     for (i, v) in enumerate(MTK.unbound_inputs(sys))
         if MTK.hasbounds(v)
             lo, hi = MTK.getbounds(v)
-            set_lower_bound(V[i], Symbolics.fixpoint_sub(lo, pmap))
-            set_upper_bound(V[i], Symbolics.fixpoint_sub(hi, pmap))
+            set_lower_bound(V[i], Symbolics.fast_substitute(lo, pmap))
+            set_upper_bound(V[i], Symbolics.fast_substitute(hi, pmap))
         end
     end
 end
 
 function add_jump_cost_function!(model::InfiniteModel, sys, tspan, pmap; is_free_t = false)
-    jcosts = MTK.get_costs(sys)
-    consolidate = MTK.get_consolidate(sys)
-    if isnothing(jcosts) || isempty(jcosts)
+    jcosts = cost(sys)
+    if Symbolics._iszero(jcosts)
         @objective(model, Min, 0)
         return
     end
-    jcosts = substitute_jump_vars(model, sys, pmap, jcosts; is_free_t)
+    jcosts = substitute_jump_vars(model, sys, pmap, [jcosts]; is_free_t)[1]
     tₛ = is_free_t ? model[:tf] : 1
 
     # Substitute integral
@@ -187,17 +188,18 @@ function add_jump_cost_function!(model::InfiniteModel, sys, tspan, pmap; is_free
         hi = haskey(pmap, hi) ? 1 : MTK.value(hi)
         intmap[int] = tₛ * InfiniteOpt.∫(arg, model[:t], lo, hi)
     end
-    jcosts = map(c -> Symbolics.substitute(c, intmap), jcosts)
-    @objective(model, Min, consolidate(jcosts))
+    jcosts = Symbolics.substitute(jcosts, intmap)
+    @objective(model, Min, MTK.value(jcosts))
 end
 
 function add_user_constraints!(model::InfiniteModel, sys, pmap; is_free_t = false)
-    conssys = MTK.get_constraintsystem(sys)
-    jconstraints = isnothing(conssys) ? nothing : MTK.get_constraints(conssys)
+    jconstraints = MTK.get_constraints(sys)
     (isnothing(jconstraints) || isempty(jconstraints)) && return nothing
+    cons_dvs, cons_ps = MTK.process_constraint_system(
+        jconstraints, Set(unknowns(sys)), parameters(sys), MTK.get_iv(sys); validate = false)
 
     if is_free_t
-        for u in MTK.get_unknowns(conssys)
+        for u in cons_dvs
             x = MTK.operation(u)
             t = only(arguments(u))
             if (MTK.symbolic_type(t) === MTK.NotSymbolic())
@@ -206,7 +208,7 @@ function add_user_constraints!(model::InfiniteModel, sys, pmap; is_free_t = fals
         end
     end
 
-    auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in unknowns(conssys)])
+    auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in cons_dvs])
     jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints; auxmap, is_free_t)
 
     # Substitute to-term'd variables
@@ -235,25 +237,25 @@ function substitute_jump_vars(model, sys, pmap, exprs; auxmap = Dict(), is_free_
     x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
     c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
 
-    exprs = map(c -> Symbolics.fixpoint_sub(c, auxmap), exprs)
-    exprs = map(c -> Symbolics.fixpoint_sub(c, Dict(pmap)), exprs)
+    exprs = map(c -> Symbolics.fast_substitute(c, auxmap), exprs)
+    exprs = map(c -> Symbolics.fast_substitute(c, Dict(pmap)), exprs)
     if is_free_t
         tf = model[:tf]
         free_t_map = Dict([[x(tf) => U[i](1) for (i, x) in enumerate(x_ops)];
                            [c(tf) => V[i](1) for (i, c) in enumerate(c_ops)]])
-        exprs = map(c -> Symbolics.fixpoint_sub(c, free_t_map), exprs)
+        exprs = map(c -> Symbolics.fast_substitute(c, free_t_map), exprs)
     end
 
     # for variables like x(t)
     whole_interval_map = Dict([[v => U[i] for (i, v) in enumerate(sts)];
                                [v => V[i] for (i, v) in enumerate(cts)]])
-    exprs = map(c -> Symbolics.fixpoint_sub(c, whole_interval_map), exprs)
+    exprs = map(c -> Symbolics.fast_substitute(c, whole_interval_map), exprs)
 
     # for variables like x(1.0)
     fixed_t_map = Dict([[x_ops[i] => U[i] for i in 1:length(U)];
                         [c_ops[i] => V[i] for i in 1:length(V)]])
 
-    exprs = map(c -> Symbolics.fixpoint_sub(c, fixed_t_map), exprs)
+    exprs = map(c -> Symbolics.fast_substitute(c, fixed_t_map), exprs)
     exprs
 end