From ce6bbc04189d1dffe61a026a447aa89c6b2231ae Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 1 Apr 2025 11:22:13 -0400
Subject: [PATCH 01/41] init: JuMPProblem

---
 ext/MTKCASADIExt.jl | 11 +++++++++++
 ext/MTKJuMPExt.jl   | 19 +++++++++++++++++++
 2 files changed, 30 insertions(+)
 create mode 100644 ext/MTKCASADIExt.jl
 create mode 100644 ext/MTKJuMPExt.jl

diff --git a/ext/MTKCASADIExt.jl b/ext/MTKCASADIExt.jl
new file mode 100644
index 0000000000..b41ece6eb7
--- /dev/null
+++ b/ext/MTKCASADIExt.jl
@@ -0,0 +1,11 @@
+"""
+an ODESystem with constraints to a JuMPProblem for optimal control solving.
+"""
+function CASADIProblem(sys::ODESystem)
+    
+end
+
+
+function CASADIProblem(prob::ODEProblem) 
+    
+end
diff --git a/ext/MTKJuMPExt.jl b/ext/MTKJuMPExt.jl
new file mode 100644
index 0000000000..b8e23a0f12
--- /dev/null
+++ b/ext/MTKJuMPExt.jl
@@ -0,0 +1,19 @@
+using JuMP, InfiniteOpt
+
+"""
+Convert an ODESystem with constraints to a JuMPProblem for optimal control solving.
+"""
+function JuMPProblem(sys::ODESystem, u0, tspan, p; dt = error("dt must be provided for JuMPProblem."))
+    steps = tspan[1]:dt:tspan[2]
+    nsteps = length(steps)
+    cost = sys.cost
+    coalesce = sys.coalesce
+
+    @variables
+    @variables
+    @variables
+end
+
+function symbolic_to_jump_constraint(cons::Union{Num, BasicSymbolic})
+    
+end

From 95d1d56329fac07f18dc99bfd3166e0f37ccb814 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 3 Apr 2025 18:25:31 -0400
Subject: [PATCH 02/41] up

---
 ext/MTKJuMPExt.jl | 37 ++++++++++++++++++++++++++++---------
 1 file changed, 28 insertions(+), 9 deletions(-)

diff --git a/ext/MTKJuMPExt.jl b/ext/MTKJuMPExt.jl
index b8e23a0f12..c90a866aaa 100644
--- a/ext/MTKJuMPExt.jl
+++ b/ext/MTKJuMPExt.jl
@@ -4,16 +4,35 @@ using JuMP, InfiniteOpt
 Convert an ODESystem with constraints to a JuMPProblem for optimal control solving.
 """
 function JuMPProblem(sys::ODESystem, u0, tspan, p; dt = error("dt must be provided for JuMPProblem."))
-    steps = tspan[1]:dt:tspan[2]
-    nsteps = length(steps)
-    cost = sys.cost
-    coalesce = sys.coalesce
-
-    @variables
-    @variables
-    @variables
+    ts = tspan[1]
+    te = tspan[2]
+    steps = ts:dt:te
+    costs = get_costs(sys)
+    consolidate = get_consolidate(sys)
+    ctrls = get_ctrls(sys)
+    states = unknowns(sys)
+    constraints = get_constraints(get_constraintsystem(sys))
+
+    model = Model()
+
+    @infinite_parameter(model, t in [tspan[1],tspan[2]], num_supports = length(steps), derivative_method = OrthogonalCollocation(2))
+    @variables(model, U[1:length(states)], Infinite(t), start = ts)
+    @variables(model, V[1:length(ctrls)], Infinite(t), start = ts)
+    @variables(model, K)
+
+    jcost = generate_jump_cost_function(sys)
+    @objective
+
+    constraints = generate_jump_constraints(constraints)
+    @constraints
+end
+
+function generate_jump_cost_function(costs, tsteps)
+end
+
+function generate_jump_constraints(constraints, jump_vars, jump_ps)
 end
 
-function symbolic_to_jump_constraint(cons::Union{Num, BasicSymbolic})
+function t_to_tstep()
     
 end

From 831ca0efc14c16cf7dbfbd00f102b7d9686b82e4 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 7 Apr 2025 11:20:04 -0400
Subject: [PATCH 03/41] add solver getter

---
 ext/MTKJuMPExt.jl | 135 +++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 117 insertions(+), 18 deletions(-)

diff --git a/ext/MTKJuMPExt.jl b/ext/MTKJuMPExt.jl
index c90a866aaa..a0a913a59a 100644
--- a/ext/MTKJuMPExt.jl
+++ b/ext/MTKJuMPExt.jl
@@ -1,38 +1,137 @@
+module MTKJuMPControlExt
+using ModelingToolkit
 using JuMP, InfiniteOpt
+using DiffEqDevTools, DiffEqBase
+
+struct JuMPProblem{uType, tType, isinplace, P, F, K} <:
+       AbstractODEProblem{uType, tType, isinplace}
+       f::F
+       u0::uType
+       tspan
+       p
+       model
+       kwargs
+end
 
 """
-Convert an ODESystem with constraints to a JuMPProblem for optimal control solving.
+    JuMPProblem(sys::ODESystem, u0, tspan, p; dt)
+
+Convert an ODESystem representing an optimal control system into a JuMP model
+for solving using optimization. Must provide `dt` for determining the length 
+of the interpolation arrays.
+
+The optimization variables:
+- a vector-of-vectors U representing the unknowns as an interpolation array
+- 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 solver constraints that encode the time-stepping used by the solver
 """
-function JuMPProblem(sys::ODESystem, u0, tspan, p; dt = error("dt must be provided for JuMPProblem."))
+function JuMPProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPProblem."), solver = :Tsit5)
     ts = tspan[1]
     te = tspan[2]
     steps = ts:dt:te
-    costs = get_costs(sys)
-    consolidate = get_consolidate(sys)
     ctrls = get_ctrls(sys)
     states = unknowns(sys)
-    constraints = get_constraints(get_constraintsystem(sys))
 
-    model = Model()
+    if !isnothing(constraintsys)
+        (length(constraints(constraintsys)) + length(u0map) > length(sts)) &&
+            @warn "The BVProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The BVP solvers will default to doing a nonlinear least-squares optimization."
+    end
 
-    @infinite_parameter(model, t in [tspan[1],tspan[2]], num_supports = length(steps), derivative_method = OrthogonalCollocation(2))
-    @variables(model, U[1:length(states)], Infinite(t), start = ts)
-    @variables(model, V[1:length(ctrls)], Infinite(t), start = ts)
-    @variables(model, K)
+    model = InfiniteModel()
+    @infinite_parameter(model, t in [ts, te], num_supports = length(steps), derivative_method = OrthogonalCollocation(2))
+    @variable(model, U[1:length(states)], Infinite(t), start = ts)
+    @variable(model, V[1:length(ctrls)], Infinite(t), start = ts)
+    @variable(model, K)
+
+    f, u0, p = process_SciMLProblem(ODEFunction{iip, specialize}, sys, _u0map, parammap;
+        t = tspan !== nothing ? tspan[1] : tspan, guesses,
+        check_length, warn_initialize_determined, eval_expression, eval_module, kwargs...)
+
+    add_jump_cost_function!(model, sys)
+    add_user_constraints!(model, sys)
+    add_solve_constraints!(model)
+
+    JuMPProblem{iip}(f, u0, tspan, p, model; kwargs...)
+end
+
+function add_jump_cost_function!(model, sys)
+    jcosts = get_costs(sys)
+    consolidate = get_consolidate(sys)
+    iv = get_iv(sys)
 
-    jcost = generate_jump_cost_function(sys)
-    @objective
+    stidxmap = Dict([v => i for (i, v) in enumerate(get_unknowns(sys))])
+    cidxmap = Dict([v => i for (i, v) in enumerate(get_ctrls(sys))])
 
-    constraints = generate_jump_constraints(constraints)
-    @constraints
+    for st in get_unknowns(sys)
+        x = operation(st)
+        t = only(arguments(st))
+        idx = stidxmap[x(iv)]
+        jcosts = Symbolics.substitute(costs, Dict(x(t) => model[:U][idx](t)))
+    end
+    
+    for ct in get_ctrls(sys)
+        p = operation(ct)
+        t = only(arguments(ct))
+        idx = cidxmap[p(iv)]
+        jcosts = Symbolics.substitute(costs, Dict(p(t) => model[:V][idx](t)))
+    end
+    
+    @objective(model, Min, consolidate(jcosts))
 end
 
-function generate_jump_cost_function(costs, tsteps)
+function add_user_constraints!(model, sys, u0map)
+    jconstraints = get_constraints(get_constraintsystem(sys))
+    iv = get_iv(sys)
+
+    stidxmap = Dict([v => i for (i, v) in enumerate(get_unknowns(sys))])
+    cidxmap = Dict([v => i for (i, v) in enumerate(get_ctrls(sys))])
+
+    for st in get_unknowns(sys)
+        x = operation(st)
+        t = only(arguments(st))
+        idx = stidxmap[x(iv)]
+        subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
+        jconstraints = Symbolics.substitute(constraints, Dict(x(t) => subval))
+    end
+
+    for ct in get_ctrls(sys)
+        p = operation(ct)
+        t = only(arguments(ct))
+        idx = cidxmap[p(iv)]
+        subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
+        jconstraints = Symbolics.substitute(constraints, Dict(p(t) => subval))
+    end
+
+    for (i, cons) in enumerate(jconstraints)
+        if cons isa Equation
+            @constraint(model, user[i], cons.lhs - cons.rhs == 0)
+        elseif cons.relational_op === Symbolics.geq 
+            @constraint(model, user[i], cons.lhs - cons.rhs ≥ 0)
+        else
+            @constraint(model, user[i], cons.lhs - cons.rhs ≤ 0)
+        end
+    end
+
+    # Add initial constraints.
 end
 
-function generate_jump_constraints(constraints, jump_vars, jump_ps)
+function add_solve_constraints!(model, tsteps, solver)
+    tableau = fetch_tableau(solver)
+
+    for (i, t) in collect(enumerate(tsteps))
+    end
 end
 
-function t_to_tstep()
-    
+"""
+Solve JuMPProblem. Takes in a symbol representing the solver.
+"""
+function solve(prob::JuMPProblem, solver_sym::Symbol)
+    model = prob.model
+    tableau_getter = Symbol(:construct, solver)
+    @eval tableau = $tableau_getter()
+    add_solve_constraints!(model, tableau)
+end
 end

From bbf71f17508712de9ee749e2953a6b2318d4bc1b Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 7 Apr 2025 11:23:13 -0400
Subject: [PATCH 04/41] add test/project

---
 Project.toml            | 1 +
 test/extensions/JuMP.jl | 3 +++
 2 files changed, 4 insertions(+)
 create mode 100644 test/extensions/JuMP.jl

diff --git a/Project.toml b/Project.toml
index 35e03e2924..04fc3d67ee 100644
--- a/Project.toml
+++ b/Project.toml
@@ -76,6 +76,7 @@ MTKChainRulesCoreExt = "ChainRulesCore"
 MTKDeepDiffsExt = "DeepDiffs"
 MTKFMIExt = "FMI"
 MTKInfiniteOptExt = "InfiniteOpt"
+MTKJuMP = "JuMP"
 MTKLabelledArraysExt = "LabelledArrays"
 
 [compat]
diff --git a/test/extensions/JuMP.jl b/test/extensions/JuMP.jl
new file mode 100644
index 0000000000..b6fc2b1fbc
--- /dev/null
+++ b/test/extensions/JuMP.jl
@@ -0,0 +1,3 @@
+import ModelingToolkit as MTK
+using JuMP, InfiniteOpt
+using DiffEqDevTools, DiffEqBase

From 63cbbee1f994bc969f59fa9e2e9b572dae8c9174 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 7 Apr 2025 16:39:10 -0400
Subject: [PATCH 05/41] Implement solver tableau

---
 ext/MTKJuMPExt.jl | 71 +++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 59 insertions(+), 12 deletions(-)

diff --git a/ext/MTKJuMPExt.jl b/ext/MTKJuMPExt.jl
index a0a913a59a..83390461b2 100644
--- a/ext/MTKJuMPExt.jl
+++ b/ext/MTKJuMPExt.jl
@@ -3,18 +3,18 @@ using ModelingToolkit
 using JuMP, InfiniteOpt
 using DiffEqDevTools, DiffEqBase
 
-struct JuMPProblem{uType, tType, isinplace, P, F, K} <:
+struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <:
        AbstractODEProblem{uType, tType, isinplace}
        f::F
        u0::uType
-       tspan
+       tspan::tType
        p
        model
        kwargs
 end
 
 """
-    JuMPProblem(sys::ODESystem, u0, tspan, p; dt)
+    JuMPControlProblem(sys::ODESystem, u0, tspan, p; dt)
 
 Convert an ODESystem representing an optimal control system into a JuMP model
 for solving using optimization. Must provide `dt` for determining the length 
@@ -28,7 +28,7 @@ The constraints are:
 - The set of user constraints passed to the ODESystem via `constraints`
 - The solver constraints that encode the time-stepping used by the solver
 """
-function JuMPProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPProblem."), solver = :Tsit5)
+function JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPProblem."), solver = :Tsit5)
     ts = tspan[1]
     te = tspan[2]
     steps = ts:dt:te
@@ -54,7 +54,7 @@ function JuMPProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be
     add_user_constraints!(model, sys)
     add_solve_constraints!(model)
 
-    JuMPProblem{iip}(f, u0, tspan, p, model; kwargs...)
+    JuMPControlProblem{iip}(f, u0, tspan, p, model; kwargs...)
 end
 
 function add_jump_cost_function!(model, sys)
@@ -118,20 +118,67 @@ function add_user_constraints!(model, sys, u0map)
     # Add initial constraints.
 end
 
-function add_solve_constraints!(model, tsteps, solver)
-    tableau = fetch_tableau(solver)
-
-    for (i, t) in collect(enumerate(tsteps))
+function add_solve_constraints!(prob, talbeau, f, tsteps)
+    A = tableau.A
+    α = tableau.α
+    c = tableau.c
+    model = prob.model
+    p = prob.p
+    dt = step(tsteps)
+
+    if is_explicit(tableau)
+        K = Any[]
+        for t in tsteps
+            for (i, h) in enumerate(c)
+                ΔU = sum([A[i, j] * K[j] for j in 1:i-1])
+                Kₙ = f(U + ΔU*dt, p, t + h*dt) 
+                push!(K, Kₙ)
+            end
+            @constraint(model, U(t) + dot(α, K) == U(t + dt))
+            empty!(K)
+        end
+    else
+        @variable(model, K[1:length(a)], Infinite(t), start = tsteps[1])
+        for t in tsteps
+            ΔUs = A * K(t)
+            for (i, h) in enumerate(c)
+                ΔU = ΔUs[i]
+                @constraint(model, K[i](t) == f(U + ΔU*dt, p, t + h*dt))
+            end
+            @constraint(model, U(t) + dot(α, K(t)) == U(t + dt))
+        end
     end
 end
 
+is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
+
 """
-Solve JuMPProblem. Takes in a symbol representing the solver.
 """
-function solve(prob::JuMPProblem, solver_sym::Symbol)
+struct JuMPControlSolution
+    model
+    sol::ODESolution
+end
+
+"""
+Solve JuMPProblem. Takes in a symbol representing the solver. Acceptable solvers may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/.
+Note that the symbol may be different than the typical
+name of the solver, e.g. :Tsitouras5 rather than Tsit5.
+"""
+function solve(prob::JuMPProblem, jump_solver, ode_solver::Symbol)
     model = prob.model
+    f = prob.f
     tableau_getter = Symbol(:construct, solver)
     @eval tableau = $tableau_getter()
-    add_solve_constraints!(model, tableau)
+    ts = prob.tspan[1]:dt:prob.tspan[2]
+    add_solve_constraints!(model, ts, tableau, f)
+
+    set_optimizer(model, solver)
+    optimize!(model)
+
+    if is_solved_and_feasible(model)
+        sol = DiffEqBase.build_solution(prob, ode_solver, ts, value(U))
+        JuMPControlSolution(model, sol)
+    end
 end
+
 end

From 80bc3d0e2051258d901a4b568ec9763dd90c9952 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 7 Apr 2025 19:23:28 -0400
Subject: [PATCH 06/41] up

---
 ext/MTKJuMPExt.jl | 49 +++++++++++++++++++++++++++--------------------
 1 file changed, 28 insertions(+), 21 deletions(-)

diff --git a/ext/MTKJuMPExt.jl b/ext/MTKJuMPExt.jl
index 83390461b2..6cfcf432ef 100644
--- a/ext/MTKJuMPExt.jl
+++ b/ext/MTKJuMPExt.jl
@@ -8,9 +8,9 @@ struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <:
        f::F
        u0::uType
        tspan::tType
-       p
-       model
-       kwargs
+       p::P
+       model::Model
+       kwargs::K
 end
 
 """
@@ -28,7 +28,7 @@ The constraints are:
 - The set of user constraints passed to the ODESystem via `constraints`
 - The solver constraints that encode the time-stepping used by the solver
 """
-function JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPProblem."), solver = :Tsit5)
+function JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPControlProblem."), guesses, eval_expression, eval_module)
     ts = tspan[1]
     te = tspan[2]
     steps = ts:dt:te
@@ -37,7 +37,7 @@ function JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt m
 
     if !isnothing(constraintsys)
         (length(constraints(constraintsys)) + length(u0map) > length(sts)) &&
-            @warn "The BVProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The BVP solvers will default to doing a nonlinear least-squares optimization."
+            @warn "The JuMPControlProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
     end
 
     model = InfiniteModel()
@@ -52,9 +52,12 @@ function JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt m
 
     add_jump_cost_function!(model, sys)
     add_user_constraints!(model, sys)
-    add_solve_constraints!(model)
 
-    JuMPControlProblem{iip}(f, u0, tspan, p, model; kwargs...)
+    stidxmap = Dict([v => i for (i, v) in enumerate(sts)])
+    u0_idxs = has_alg_eqs(sys) ? collect(1:length(sts)) : [stidxmap[k] for (k, v) in u0map]
+    add_initial_constraints!(model, u0, u0_idxs, tspan)
+
+    JuMPControlProblem{iip}(f, u0, tspan, p, model, kwargs...)
 end
 
 function add_jump_cost_function!(model, sys)
@@ -114,11 +117,16 @@ function add_user_constraints!(model, sys, u0map)
             @constraint(model, user[i], cons.lhs - cons.rhs ≤ 0)
         end
     end
+end
 
-    # Add initial constraints.
+function add_initial_constraints!(model, u0, u0_idxs, tspan)
+    ts = tspan[1]
+    @constraint(model, init_u0_idx[i in u0_idxs], U[i](ts) == u0[i])
 end
 
-function add_solve_constraints!(prob, talbeau, f, tsteps)
+is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
+
+function add_solve_constraints!(prob, tableau, f, tsteps)
     A = tableau.A
     α = tableau.α
     c = tableau.c
@@ -126,32 +134,31 @@ function add_solve_constraints!(prob, talbeau, f, tsteps)
     p = prob.p
     dt = step(tsteps)
 
+    U = model[:U]
     if is_explicit(tableau)
         K = Any[]
-        for t in tsteps
+        for τ in tsteps
             for (i, h) in enumerate(c)
                 ΔU = sum([A[i, j] * K[j] for j in 1:i-1])
-                Kₙ = f(U + ΔU*dt, p, t + h*dt) 
+                Kₙ = f(U + ΔU*dt, p, τ + h*dt)
                 push!(K, Kₙ)
             end
-            @constraint(model, U(t) + dot(α, K) == U(t + dt))
+            @constraint(model, U(τ) + dot(α, K) == U(τ + dt))
             empty!(K)
         end
     else
         @variable(model, K[1:length(a)], Infinite(t), start = tsteps[1])
-        for t in tsteps
-            ΔUs = A * K(t)
+        for τ in tsteps
+            ΔUs = A * K(τ)
             for (i, h) in enumerate(c)
                 ΔU = ΔUs[i]
-                @constraint(model, K[i](t) == f(U + ΔU*dt, p, t + h*dt))
+                @constraint(model, K[i](τ) == f(U(τ) + ΔU*dt, p, τ + h*dt))
             end
-            @constraint(model, U(t) + dot(α, K(t)) == U(t + dt))
+            @constraint(model, U(τ) + dot(α, K(τ)) == U(τ + dt))
         end
     end
 end
 
-is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
-
 """
 """
 struct JuMPControlSolution
@@ -167,16 +174,16 @@ name of the solver, e.g. :Tsitouras5 rather than Tsit5.
 function solve(prob::JuMPProblem, jump_solver, ode_solver::Symbol)
     model = prob.model
     f = prob.f
-    tableau_getter = Symbol(:construct, solver)
+    tableau_getter = Symbol(:construct, ode_solver)
     @eval tableau = $tableau_getter()
     ts = prob.tspan[1]:dt:prob.tspan[2]
     add_solve_constraints!(model, ts, tableau, f)
 
-    set_optimizer(model, solver)
+    set_optimizer(model, jump_solver)
     optimize!(model)
 
     if is_solved_and_feasible(model)
-        sol = DiffEqBase.build_solution(prob, ode_solver, ts, value(U))
+        sol = DiffEqBase.build_solution(prob, ode_solver, ts, value.(U))
         JuMPControlSolution(model, sol)
     end
 end

From 42e0048344f0e4ba878273222756de9e5839760a Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 8 Apr 2025 00:40:15 -0400
Subject: [PATCH 07/41] feat: solver tableau debugging

---
 Project.toml                    |   6 +-
 ext/MTKJuMPControlExt.jl        | 218 ++++++++++++++++++++++++++++++++
 ext/MTKJuMPExt.jl               | 191 ----------------------------
 src/ModelingToolkit.jl          |   3 +
 test/extensions/JuMP.jl         |   3 -
 test/extensions/Project.toml    |   4 +
 test/extensions/jump_control.jl |  57 +++++++++
 7 files changed, 287 insertions(+), 195 deletions(-)
 create mode 100644 ext/MTKJuMPControlExt.jl
 delete mode 100644 ext/MTKJuMPExt.jl
 delete mode 100644 test/extensions/JuMP.jl
 create mode 100644 test/extensions/jump_control.jl

diff --git a/Project.toml b/Project.toml
index 04fc3d67ee..e8831ec161 100644
--- a/Project.toml
+++ b/Project.toml
@@ -31,6 +31,7 @@ FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 JumpProcesses = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
@@ -66,8 +67,10 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
+DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 FMI = "14a09403-18e3-468f-ad8a-74f8dda2d9ac"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 
 [extensions]
@@ -76,7 +79,7 @@ MTKChainRulesCoreExt = "ChainRulesCore"
 MTKDeepDiffsExt = "DeepDiffs"
 MTKFMIExt = "FMI"
 MTKInfiniteOptExt = "InfiniteOpt"
-MTKJuMP = "JuMP"
+MTKJuMPControlExt = ["JuMP", "DiffEqDevTools"]
 MTKLabelledArraysExt = "LabelledArrays"
 
 [compat]
@@ -98,6 +101,7 @@ DeepDiffs = "1"
 DelayDiffEq = "5.50"
 DiffEqBase = "6.165.1"
 DiffEqCallbacks = "2.16, 3, 4"
+DiffEqDevTools = "2.48.0"
 DiffEqNoiseProcess = "5"
 DiffRules = "0.1, 1.0"
 DifferentiationInterface = "0.6.47"
diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
new file mode 100644
index 0000000000..1e788b05ec
--- /dev/null
+++ b/ext/MTKJuMPControlExt.jl
@@ -0,0 +1,218 @@
+module MTKJuMPControlExt
+using ModelingToolkit
+using JuMP, InfiniteOpt
+using DiffEqDevTools, DiffEqBase, SciMLBase
+using LinearAlgebra
+const MTK = ModelingToolkit
+
+struct JuMPControlProblem{uType, tType, P, F, K}
+       f::F
+       u0::uType
+       tspan::tType
+       p::P
+       model::InfiniteModel
+       kwargs::K
+
+       function JuMPControlProblem(f, u0, tspan, p, model; kwargs...) 
+           new{typeof(u0), typeof(tspan), typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
+       end
+end
+
+"""
+    JuMPControlProblem(sys::ODESystem, u0, tspan, p; dt)
+
+Convert an ODESystem representing an optimal control system into a JuMP model
+for solving using optimization. Must provide `dt` for determining the length 
+of the interpolation arrays.
+
+The optimization variables:
+- a vector-of-vectors U representing the unknowns as an interpolation array
+- 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 solver constraints that encode the time-stepping used by the solver
+"""
+function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPControlProblem."), kwargs...)
+    ts = tspan[1]
+    te = tspan[2]
+    steps = ts:dt:te
+    ctrls = controls(sys)
+    states = unknowns(sys)
+    constraintsys = MTK.get_constraintsystem(sys)
+
+    if !isnothing(constraintsys)
+        (length(constraints(constraintsys)) + length(u0map) > length(sts)) &&
+            @warn "The JuMPControlProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
+    end
+
+    f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, u0map, pmap;
+        t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
+
+    model = InfiniteModel()
+    @infinite_parameter(model, t in [ts, te], num_supports = length(steps), derivative_method = OrthogonalCollocation(2))
+    @variable(model, U[1:length(states)], Infinite(t), start = ts)
+    @variable(model, V[1:length(ctrls)], Infinite(t), start = ts)
+
+    add_jump_cost_function!(model, sys)
+    add_user_constraints!(model, sys)
+
+    stidxmap = Dict([v => i for (i, v) in enumerate(states)])
+    u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) : [stidxmap[k] for (k, v) in u0map]
+    add_initial_constraints!(model, u0, u0_idxs, tspan)
+
+    JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
+end
+
+function add_jump_cost_function!(model, sys)
+    jcosts = MTK.get_costs(sys)
+    consolidate = MTK.get_consolidate(sys)
+    if isnothing(consolidate)
+        @objective(model, Min, 0)
+        return
+    end
+    iv = MTK.get_iv(sys)
+
+    stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
+    cidxmap = Dict([v => i for (i, v) in enumerate(controls(sys))])
+
+    for st in unknowns(sys)
+        x = operation(st)
+        t = only(arguments(st))
+        idx = stidxmap[x(iv)]
+        subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
+        jcosts = Symbolics.substitute(jcosts, Dict(x(t) => subval))
+    end
+    
+    for ct in controls(sys)
+        p = operation(ct)
+        t = only(arguments(ct))
+        idx = cidxmap[p(iv)]
+        subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
+        jcosts = Symbolics.substitute(jcosts, Dict(x(t) => subval))
+    end
+    
+    @objective(model, Min, consolidate(jcosts))
+end
+
+function add_user_constraints!(model, sys)
+    jconstraints = if !(csys = MTK.get_constraintsystem(sys) isa Nothing)
+        MTK.get_constraints(csys)
+    else
+        nothing
+    end
+    isnothing(jconstraints) && return nothing
+
+    iv = MTK.get_iv(sys)
+    stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
+    cidxmap = Dict([v => i for (i, v) in enumerate(controls(sys))])
+
+    for st in unknowns(sys)
+        x = operation(st)
+        t = only(arguments(st))
+        idx = stidxmap[x(iv)]
+        subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
+        jconstraints = Symbolics.substitute(jconstraints, Dict(x(t) => subval))
+    end
+
+    for ct in controls(sys)
+        p = operation(ct)
+        t = only(arguments(ct))
+        idx = cidxmap[p(iv)]
+        subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
+        jconstraints = Symbolics.substitute(jconstraints, Dict(p(t) => subval))
+    end
+
+    for (i, cons) in enumerate(jconstraints)
+        if cons isa Equation
+            @constraint(model, user[i], cons.lhs - cons.rhs == 0)
+        elseif cons.relational_op === Symbolics.geq 
+            @constraint(model, user[i], cons.lhs - cons.rhs ≥ 0)
+        else
+            @constraint(model, user[i], cons.lhs - cons.rhs ≤ 0)
+        end
+    end
+end
+
+function add_initial_constraints!(model, u0, u0_idxs, tspan)
+    ts = tspan[1]
+    @constraint(model, init_u0_idx[i in u0_idxs], model[:U][i](ts) == u0[i])
+end
+
+is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
+
+function add_solve_constraints!(prob, tableau)
+    A = tableau.A
+    α = tableau.α
+    c = tableau.c
+    model = prob.model
+    f = prob.f
+    p = prob.p
+    tsteps = supports(model[:t])
+    pop!(tsteps)
+    dt = tsteps[2] - tsteps[1]
+
+    U = model[:U]
+    nᵤ = length(U)
+    if is_explicit(tableau)
+        K = Any[]
+        for τ in tsteps
+            for (i, h) in enumerate(c)
+                ΔU = sum([A[i, j] * K[j] for j in 1:i-1], init = zeros(nᵤ))
+                Uₙ = [U[i](τ) + ΔU[i]*dt for i in 1:nᵤ]
+                Kₙ = f(Uₙ, p, τ + h*dt)
+                push!(K, Kₙ)
+            end
+            ΔU = sum([α[i] * K[i] for i in 1:length(α)])
+            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU[n] == U[n](τ + dt))
+            empty!(K)
+        end
+    else
+        @variable(model, K[1:length(a), 1:nᵤ], Infinite(t), start = tsteps[1])
+        for τ in tsteps
+            ΔUs = [A * K(τ)]
+            for (i, h) in enumerate(c)
+                ΔU = ΔUs[i]
+                Uₙ = [U[j](τ) + ΔU[j](τ)*dt for j in 1:nᵤ]
+                @constraint(model, K[i](τ) == f(Uₙ, p, τ + h*dt))
+            end
+            ΔU = sum([α[i] * K[i] for i in 1:length(α)])
+            @constraint(model, U(τ) + dot(α, K(τ)) == U(τ + dt))
+        end
+    end
+end
+
+"""
+"""
+struct JuMPControlSolution
+    model::InfiniteModel
+    sol::ODESolution
+end
+
+"""
+Solve JuMPControlProblem. Arguments:
+- prob: a JumpControlProblem
+- jump_solver: a LP solver such as HiGHS
+- ode_solver: Takes in a symbol representing the solver. Acceptable solvers may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. Note that the symbol may be different than the typical name of the solver, e.g. :Tsitouras5 rather than Tsit5.
+"""
+function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Symbol)
+    model = prob.model
+    tableau_getter = Symbol(:construct, ode_solver)
+    @eval tableau = $tableau_getter()
+    ts = supports(model[:t])
+    add_solve_constraints!(prob, tableau)
+
+    set_optimizer(model, jump_solver)
+    optimize!(model)
+
+    if is_solved_and_feasible(model)
+        sol = DiffEqBase.build_solution(prob, ode_solver, ts, value.(U))
+        JuMPControlSolution(model, sol)
+    else
+        sol = DiffEqBase.build_solution(prob, ode_solver, ts, value.(U))
+        sol = SciMLBase.solution_new_retcode(sol, SciMLBase.ReturnCode.ConvergenceFailure)
+        JuMPControlSolution(model, sol)
+    end
+end
+
+end
diff --git a/ext/MTKJuMPExt.jl b/ext/MTKJuMPExt.jl
deleted file mode 100644
index 6cfcf432ef..0000000000
--- a/ext/MTKJuMPExt.jl
+++ /dev/null
@@ -1,191 +0,0 @@
-module MTKJuMPControlExt
-using ModelingToolkit
-using JuMP, InfiniteOpt
-using DiffEqDevTools, DiffEqBase
-
-struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <:
-       AbstractODEProblem{uType, tType, isinplace}
-       f::F
-       u0::uType
-       tspan::tType
-       p::P
-       model::Model
-       kwargs::K
-end
-
-"""
-    JuMPControlProblem(sys::ODESystem, u0, tspan, p; dt)
-
-Convert an ODESystem representing an optimal control system into a JuMP model
-for solving using optimization. Must provide `dt` for determining the length 
-of the interpolation arrays.
-
-The optimization variables:
-- a vector-of-vectors U representing the unknowns as an interpolation array
-- 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 solver constraints that encode the time-stepping used by the solver
-"""
-function JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPControlProblem."), guesses, eval_expression, eval_module)
-    ts = tspan[1]
-    te = tspan[2]
-    steps = ts:dt:te
-    ctrls = get_ctrls(sys)
-    states = unknowns(sys)
-
-    if !isnothing(constraintsys)
-        (length(constraints(constraintsys)) + length(u0map) > length(sts)) &&
-            @warn "The JuMPControlProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
-    end
-
-    model = InfiniteModel()
-    @infinite_parameter(model, t in [ts, te], num_supports = length(steps), derivative_method = OrthogonalCollocation(2))
-    @variable(model, U[1:length(states)], Infinite(t), start = ts)
-    @variable(model, V[1:length(ctrls)], Infinite(t), start = ts)
-    @variable(model, K)
-
-    f, u0, p = process_SciMLProblem(ODEFunction{iip, specialize}, sys, _u0map, parammap;
-        t = tspan !== nothing ? tspan[1] : tspan, guesses,
-        check_length, warn_initialize_determined, eval_expression, eval_module, kwargs...)
-
-    add_jump_cost_function!(model, sys)
-    add_user_constraints!(model, sys)
-
-    stidxmap = Dict([v => i for (i, v) in enumerate(sts)])
-    u0_idxs = has_alg_eqs(sys) ? collect(1:length(sts)) : [stidxmap[k] for (k, v) in u0map]
-    add_initial_constraints!(model, u0, u0_idxs, tspan)
-
-    JuMPControlProblem{iip}(f, u0, tspan, p, model, kwargs...)
-end
-
-function add_jump_cost_function!(model, sys)
-    jcosts = get_costs(sys)
-    consolidate = get_consolidate(sys)
-    iv = get_iv(sys)
-
-    stidxmap = Dict([v => i for (i, v) in enumerate(get_unknowns(sys))])
-    cidxmap = Dict([v => i for (i, v) in enumerate(get_ctrls(sys))])
-
-    for st in get_unknowns(sys)
-        x = operation(st)
-        t = only(arguments(st))
-        idx = stidxmap[x(iv)]
-        jcosts = Symbolics.substitute(costs, Dict(x(t) => model[:U][idx](t)))
-    end
-    
-    for ct in get_ctrls(sys)
-        p = operation(ct)
-        t = only(arguments(ct))
-        idx = cidxmap[p(iv)]
-        jcosts = Symbolics.substitute(costs, Dict(p(t) => model[:V][idx](t)))
-    end
-    
-    @objective(model, Min, consolidate(jcosts))
-end
-
-function add_user_constraints!(model, sys, u0map)
-    jconstraints = get_constraints(get_constraintsystem(sys))
-    iv = get_iv(sys)
-
-    stidxmap = Dict([v => i for (i, v) in enumerate(get_unknowns(sys))])
-    cidxmap = Dict([v => i for (i, v) in enumerate(get_ctrls(sys))])
-
-    for st in get_unknowns(sys)
-        x = operation(st)
-        t = only(arguments(st))
-        idx = stidxmap[x(iv)]
-        subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
-        jconstraints = Symbolics.substitute(constraints, Dict(x(t) => subval))
-    end
-
-    for ct in get_ctrls(sys)
-        p = operation(ct)
-        t = only(arguments(ct))
-        idx = cidxmap[p(iv)]
-        subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
-        jconstraints = Symbolics.substitute(constraints, Dict(p(t) => subval))
-    end
-
-    for (i, cons) in enumerate(jconstraints)
-        if cons isa Equation
-            @constraint(model, user[i], cons.lhs - cons.rhs == 0)
-        elseif cons.relational_op === Symbolics.geq 
-            @constraint(model, user[i], cons.lhs - cons.rhs ≥ 0)
-        else
-            @constraint(model, user[i], cons.lhs - cons.rhs ≤ 0)
-        end
-    end
-end
-
-function add_initial_constraints!(model, u0, u0_idxs, tspan)
-    ts = tspan[1]
-    @constraint(model, init_u0_idx[i in u0_idxs], U[i](ts) == u0[i])
-end
-
-is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
-
-function add_solve_constraints!(prob, tableau, f, tsteps)
-    A = tableau.A
-    α = tableau.α
-    c = tableau.c
-    model = prob.model
-    p = prob.p
-    dt = step(tsteps)
-
-    U = model[:U]
-    if is_explicit(tableau)
-        K = Any[]
-        for τ in tsteps
-            for (i, h) in enumerate(c)
-                ΔU = sum([A[i, j] * K[j] for j in 1:i-1])
-                Kₙ = f(U + ΔU*dt, p, τ + h*dt)
-                push!(K, Kₙ)
-            end
-            @constraint(model, U(τ) + dot(α, K) == U(τ + dt))
-            empty!(K)
-        end
-    else
-        @variable(model, K[1:length(a)], Infinite(t), start = tsteps[1])
-        for τ in tsteps
-            ΔUs = A * K(τ)
-            for (i, h) in enumerate(c)
-                ΔU = ΔUs[i]
-                @constraint(model, K[i](τ) == f(U(τ) + ΔU*dt, p, τ + h*dt))
-            end
-            @constraint(model, U(τ) + dot(α, K(τ)) == U(τ + dt))
-        end
-    end
-end
-
-"""
-"""
-struct JuMPControlSolution
-    model
-    sol::ODESolution
-end
-
-"""
-Solve JuMPProblem. Takes in a symbol representing the solver. Acceptable solvers may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/.
-Note that the symbol may be different than the typical
-name of the solver, e.g. :Tsitouras5 rather than Tsit5.
-"""
-function solve(prob::JuMPProblem, jump_solver, ode_solver::Symbol)
-    model = prob.model
-    f = prob.f
-    tableau_getter = Symbol(:construct, ode_solver)
-    @eval tableau = $tableau_getter()
-    ts = prob.tspan[1]:dt:prob.tspan[2]
-    add_solve_constraints!(model, ts, tableau, f)
-
-    set_optimizer(model, jump_solver)
-    optimize!(model)
-
-    if is_solved_and_feasible(model)
-        sol = DiffEqBase.build_solution(prob, ode_solver, ts, value.(U))
-        JuMPControlSolution(model, sol)
-    end
-end
-
-end
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index d0f427bd14..5f8c2d7610 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -347,4 +347,7 @@ export AnalysisPoint, get_sensitivity_function, get_comp_sensitivity_function,
        open_loop
 function FMIComponent end
 
+function JuMPControlProblem end
+export JuMPControlProblem
+
 end # module
diff --git a/test/extensions/JuMP.jl b/test/extensions/JuMP.jl
deleted file mode 100644
index b6fc2b1fbc..0000000000
--- a/test/extensions/JuMP.jl
+++ /dev/null
@@ -1,3 +0,0 @@
-import ModelingToolkit as MTK
-using JuMP, InfiniteOpt
-using DiffEqDevTools, DiffEqBase
diff --git a/test/extensions/Project.toml b/test/extensions/Project.toml
index 5b0de73cdf..a6d548f86f 100644
--- a/test/extensions/Project.toml
+++ b/test/extensions/Project.toml
@@ -2,7 +2,10 @@
 BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
@@ -12,6 +15,7 @@ ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Nemo = "2edaba10-b0f1-5616-af89-8c11ac63239a"
 NonlinearSolveHomotopyContinuation = "2ac3b008-d579-4536-8c91-a1a5998c2f8b"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
new file mode 100644
index 0000000000..c0a071e31d
--- /dev/null
+++ b/test/extensions/jump_control.jl
@@ -0,0 +1,57 @@
+using ModelingToolkit
+using JuMP, InfiniteOpt
+using DiffEqDevTools, DiffEqBase
+using SimpleDiffEq
+using HiGHS
+const M = ModelingToolkit
+
+@testset "ODE Solution, no cost" begin
+    # Test solving without anything attached.
+    @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
+    M.@variables x(..) y(..)
+    t = M.t_nounits; D = M.D_nounits
+
+    eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t),
+        D(y(t)) ~ -γ * y(t) + δ * x(t) * y(t)]
+
+    tspan = (0.0, 1.0)
+    u0map = [x(t) => 4.0, y(t) => 2.0]
+    parammap = [α => 7.5, β => 4, γ => 8.0, δ => 5.0]
+    @mtkbuild sys = ODESystem(eqs, t)
+
+    # Test explicit method.
+    jprob = JuMPControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
+    @test num_constraints(jprob.model) == 2 # initials
+    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    oprob = ODEProblem(sys, u0map, tspan, parammap)
+    osol = solve(oprob, SimpleTsit5(), adaptive = false)
+    @test jsol.sol.u ≈ osol.u
+
+    # Implicit method.
+    jsol2 = solve(prob, Ipopt.Optimizer, :RK4)
+    osol2 = solve(oprob, SimpleRK4(), adaptive = false)
+    @test jsol2.sol.u ≈ osol2.u
+
+    # With a constraint
+    @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
+    @variables x(..) y(..)
+
+    eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t),
+        D(y(t)) ~ -γ * y(t) + δ * x(t) * y(t)]
+
+    u0map = []
+    tspan = (0.0, 1.0)
+    guess = [x(t) => 4.0, y(t) => 2.0]
+    constr = [x(0.6) ~ 3.5, x(0.3) ~ 7.0]
+    @mtkbuild lksys = ODESystem(eqs, t; constraints = constr)
+
+    jprob = JuMPControlProblem(sys, u0map, tspan, parammap; guesses, dt = 0.01)
+    @test num_constraints(jprob.model) == 2 == num_variables(jprob.model) == 2
+    jsol = solve(prob, HiGHS.Optimizer, :Tsitouras5)
+    sol = jsol.sol
+    @test sol(0.6)[1] ≈ 3.5
+    @test sol(0.3)[1] ≈ 7.0
+end
+
+@testset "Optimal control problem" begin
+end

From 2626f5aba4ed7dee9c456fa78b4a8bf735cfd204 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 8 Apr 2025 14:50:08 -0400
Subject: [PATCH 08/41] fix: use dt in the constraints

---
 Project.toml                    |  1 -
 ext/MTKJuMPControlExt.jl        | 65 ++++++++++++++++++++++-----------
 test/extensions/Project.toml    |  2 +-
 test/extensions/ad.jl           |  2 +-
 test/extensions/jump_control.jl | 28 +++++++-------
 5 files changed, 59 insertions(+), 39 deletions(-)

diff --git a/Project.toml b/Project.toml
index e8831ec161..2779b88b83 100644
--- a/Project.toml
+++ b/Project.toml
@@ -31,7 +31,6 @@ FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 JumpProcesses = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 1e788b05ec..75260d4996 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -5,7 +5,7 @@ using DiffEqDevTools, DiffEqBase, SciMLBase
 using LinearAlgebra
 const MTK = ModelingToolkit
 
-struct JuMPControlProblem{uType, tType, P, F, K}
+struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <: SciMLBase.AbstractODEProblem{uType, tType, isinplace}
        f::F
        u0::uType
        tspan::tType
@@ -14,7 +14,7 @@ struct JuMPControlProblem{uType, tType, P, F, K}
        kwargs::K
 
        function JuMPControlProblem(f, u0, tspan, p, model; kwargs...) 
-           new{typeof(u0), typeof(tspan), typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
+           new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f), typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
        end
 end
 
@@ -50,9 +50,9 @@ function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
     model = InfiniteModel()
-    @infinite_parameter(model, t in [ts, te], num_supports = length(steps), derivative_method = OrthogonalCollocation(2))
-    @variable(model, U[1:length(states)], Infinite(t), start = ts)
-    @variable(model, V[1:length(ctrls)], Infinite(t), start = ts)
+    @infinite_parameter(model, t in [ts, te], num_supports = length(steps))
+    @variable(model, U[i = 1:length(states)], Infinite(t))
+    @variable(model, V[1:length(ctrls)], Infinite(t))
 
     add_jump_cost_function!(model, sys)
     add_user_constraints!(model, sys)
@@ -136,7 +136,8 @@ end
 
 function add_initial_constraints!(model, u0, u0_idxs, tspan)
     ts = tspan[1]
-    @constraint(model, init_u0_idx[i in u0_idxs], model[:U][i](ts) == u0[i])
+    U = model[:U]
+    @constraint(model, initial[i in u0_idxs], U[i](ts) == u0[i])
 end
 
 is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
@@ -148,6 +149,7 @@ function add_solve_constraints!(prob, tableau)
     model = prob.model
     f = prob.f
     p = prob.p
+    t = model[:t]
     tsteps = supports(model[:t])
     pop!(tsteps)
     dt = tsteps[2] - tsteps[1]
@@ -163,21 +165,21 @@ function add_solve_constraints!(prob, tableau)
                 Kₙ = f(Uₙ, p, τ + h*dt)
                 push!(K, Kₙ)
             end
-            ΔU = sum([α[i] * K[i] for i in 1:length(α)])
-            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU[n] == U[n](τ + dt))
+            ΔU = dt*sum([α[i] * K[i] for i in 1:length(α)])
+            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU[n] == U[n](τ + dt), base_name = "solve_time_$τ")
             empty!(K)
         end
     else
-        @variable(model, K[1:length(a), 1:nᵤ], Infinite(t), start = tsteps[1])
+        @variable(model, K[1:length(α), 1:nᵤ], Infinite(t), start = tsteps[1])
         for τ in tsteps
-            ΔUs = [A * K(τ)]
+            ΔUs = A * K
             for (i, h) in enumerate(c)
-                ΔU = ΔUs[i]
-                Uₙ = [U[j](τ) + ΔU[j](τ)*dt for j in 1:nᵤ]
-                @constraint(model, K[i](τ) == f(Uₙ, p, τ + h*dt))
+                ΔU = ΔUs[i, :]
+                Uₙ = [U[j] + ΔU[j]*dt for j in 1:nᵤ]
+                @constraint(model, [j in 1:nᵤ], K[i, j] == f(Uₙ, p, τ + h*dt)[j], DomainRestrictions(t => τ), base_name = "solve_K($τ)")
             end
-            ΔU = sum([α[i] * K[i] for i in 1:length(α)])
-            @constraint(model, U(τ) + dot(α, K(τ)) == U(τ + dt))
+            ΔU = dt*sum([α[i] * K[i, :] for i in 1:length(α)])
+            @constraint(model, [n = 1:nᵤ], U[n] + ΔU[n] == U[n](τ + dt), DomainRestrictions(t => τ), base_name = "solve_U($τ)")
         end
     end
 end
@@ -194,25 +196,46 @@ Solve JuMPControlProblem. Arguments:
 - prob: a JumpControlProblem
 - jump_solver: a LP solver such as HiGHS
 - ode_solver: Takes in a symbol representing the solver. Acceptable solvers may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. Note that the symbol may be different than the typical name of the solver, e.g. :Tsitouras5 rather than Tsit5.
+
+Returns a JuMPControlSolution, which contains both the model and the ODE solution.
 """
 function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Symbol)
     model = prob.model
     tableau_getter = Symbol(:construct, ode_solver)
     @eval tableau = $tableau_getter()
     ts = supports(model[:t])
+
+    # Unregister current solver constraints
+    for con in all_constraints(model)
+        if occursin("solve", JuMP.name(con))
+            unregister(model, Symbol(JuMP.name(con)))
+            delete(model, con)
+        end
+    end
+    for var in all_variables(model)
+        @show JuMP.name(var)
+        if occursin("K", JuMP.name(var))
+            unregister(model, Symbol(JuMP.name(var)))
+            delete(model, var)
+        end
+    end
     add_solve_constraints!(prob, tableau)
 
     set_optimizer(model, jump_solver)
     optimize!(model)
 
-    if is_solved_and_feasible(model)
-        sol = DiffEqBase.build_solution(prob, ode_solver, ts, value.(U))
-        JuMPControlSolution(model, sol)
-    else
-        sol = DiffEqBase.build_solution(prob, ode_solver, ts, value.(U))
+    tstatus = termination_status(model)
+    pstatus = primal_status(model)
+    !has_values(model) && error("Model not solvable; please report this to github.com/SciML/ModelingToolkit.jl.")
+
+    U_vals = value.(model[:U])
+    U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:length(ts)]
+    sol = DiffEqBase.build_solution(prob, ode_solver, ts, U_vals)
+
+    if !(pstatus === FEASIBLE_POINT && (tstatus === OPTIMAL || tstatus === LOCALLY_SOLVED || tstatus === ALMOST_OPTIMAL || tstatus === ALMOST_LOCALLY_SOLVED))
         sol = SciMLBase.solution_new_retcode(sol, SciMLBase.ReturnCode.ConvergenceFailure)
-        JuMPControlSolution(model, sol)
     end
+    JuMPControlSolution(model, sol)
 end
 
 end
diff --git a/test/extensions/Project.toml b/test/extensions/Project.toml
index a6d548f86f..9ce343202d 100644
--- a/test/extensions/Project.toml
+++ b/test/extensions/Project.toml
@@ -14,10 +14,10 @@ LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Nemo = "2edaba10-b0f1-5616-af89-8c11ac63239a"
 NonlinearSolveHomotopyContinuation = "2ac3b008-d579-4536-8c91-a1a5998c2f8b"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
+SimpleDiffEq = "05bca326-078c-5bf0-a5bf-ce7c7982d7fd"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
diff --git a/test/extensions/ad.jl b/test/extensions/ad.jl
index adaf6117c6..a456263655 100644
--- a/test/extensions/ad.jl
+++ b/test/extensions/ad.jl
@@ -3,7 +3,7 @@ using ModelingToolkit: t_nounits as t, D_nounits as D
 using Zygote
 using SymbolicIndexingInterface
 using SciMLStructures
-using OrdinaryDiffEq
+using OrdinaryDiffEqTsit5
 using NonlinearSolve
 using SciMLSensitivity
 using ForwardDiff
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index c0a071e31d..a08c95b29b 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -2,7 +2,8 @@ using ModelingToolkit
 using JuMP, InfiniteOpt
 using DiffEqDevTools, DiffEqBase
 using SimpleDiffEq
-using HiGHS
+using OrdinaryDiffEqSDIRK
+using Ipopt 
 const M = ModelingToolkit
 
 @testset "ODE Solution, no cost" begin
@@ -22,36 +23,33 @@ const M = ModelingToolkit
     # Test explicit method.
     jprob = JuMPControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
     @test num_constraints(jprob.model) == 2 # initials
-    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    jsol = solve(jprob, Ipopt.Optimizer, :RK4)
     oprob = ODEProblem(sys, u0map, tspan, parammap)
-    osol = solve(oprob, SimpleTsit5(), adaptive = false)
+    osol = solve(oprob, SimpleRK4(), dt = 0.01)
     @test jsol.sol.u ≈ osol.u
 
     # Implicit method.
-    jsol2 = solve(prob, Ipopt.Optimizer, :RK4)
-    osol2 = solve(oprob, SimpleRK4(), adaptive = false)
-    @test jsol2.sol.u ≈ osol2.u
+    jsol2 = solve(jprob, Ipopt.Optimizer, :ImplicitEuler)
+    osol2 = solve(oprob, ImplicitEuler(), dt = 0.01, adaptive = false)
+    @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001)
 
     # With a constraint
-    @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
-    @variables x(..) y(..)
-
-    eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t),
-        D(y(t)) ~ -γ * y(t) + δ * x(t) * y(t)]
-
-    u0map = []
-    tspan = (0.0, 1.0)
+    u0map = Pair[]
     guess = [x(t) => 4.0, y(t) => 2.0]
     constr = [x(0.6) ~ 3.5, x(0.3) ~ 7.0]
     @mtkbuild lksys = ODESystem(eqs, t; constraints = constr)
 
     jprob = JuMPControlProblem(sys, u0map, tspan, parammap; guesses, dt = 0.01)
     @test num_constraints(jprob.model) == 2 == num_variables(jprob.model) == 2
-    jsol = solve(prob, HiGHS.Optimizer, :Tsitouras5)
+    jsol = solve(prob, Ipopt.Optimizer, :Tsitouras5)
     sol = jsol.sol
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0
 end
 
 @testset "Optimal control problem" begin
+    # Investing
+    
+   
+    # Bang-bang control
 end

From 9e5808226c2c2f01c12f441e800d290424cc7e15 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 8 Apr 2025 14:51:02 -0400
Subject: [PATCH 09/41] test: add to runtests

---
 test/runtests.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/runtests.jl b/test/runtests.jl
index 37c738eec9..08e7425a3c 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -141,5 +141,6 @@ end
         @safetestset "LabelledArrays Test" include("labelledarrays.jl")
         @safetestset "BifurcationKit Extension Test" include("extensions/bifurcationkit.jl")
         @safetestset "InfiniteOpt Extension Test" include("extensions/test_infiniteopt.jl")
+        @safetestset "JuMPControl Extension Test" include("extensions/jump_control.jl")
     end
 end

From a09fb875fffc4e426a43d38fcdca07f368cbd8cf Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 8 Apr 2025 14:53:24 -0400
Subject: [PATCH 10/41] remove cassadi file

---
 ext/MTKCASADIExt.jl | 11 -----------
 1 file changed, 11 deletions(-)
 delete mode 100644 ext/MTKCASADIExt.jl

diff --git a/ext/MTKCASADIExt.jl b/ext/MTKCASADIExt.jl
deleted file mode 100644
index b41ece6eb7..0000000000
--- a/ext/MTKCASADIExt.jl
+++ /dev/null
@@ -1,11 +0,0 @@
-"""
-an ODESystem with constraints to a JuMPProblem for optimal control solving.
-"""
-function CASADIProblem(sys::ODESystem)
-    
-end
-
-
-function CASADIProblem(prob::ODEProblem) 
-    
-end

From 8abd12fc3a94ba3248fd28e69b8503e6223c450f Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 8 Apr 2025 18:19:30 -0400
Subject: [PATCH 11/41] up?

---
 ext/MTKJuMPControlExt.jl         | 47 ++++++++++++++++++--------------
 src/systems/diffeqs/odesystem.jl |  2 ++
 test/extensions/jump_control.jl  | 36 +++++++++++++++++-------
 3 files changed, 55 insertions(+), 30 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 75260d4996..81a6fd2f74 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -33,7 +33,7 @@ The constraints are:
 - The set of user constraints passed to the ODESystem via `constraints`
 - The solver constraints that encode the time-stepping used by the solver
 """
-function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPControlProblem."), kwargs...)
+function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPControlProblem."), guesses = Dict(), kwargs...)
     ts = tspan[1]
     te = tspan[2]
     steps = ts:dt:te
@@ -42,11 +42,12 @@ function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("
     constraintsys = MTK.get_constraintsystem(sys)
 
     if !isnothing(constraintsys)
-        (length(constraints(constraintsys)) + length(u0map) > length(sts)) &&
+        (length(constraints(constraintsys)) + length(u0map) > length(states)) &&
             @warn "The JuMPControlProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
     end
 
-    f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, u0map, pmap;
+    _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
+    f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
     model = InfiniteModel()
@@ -67,7 +68,7 @@ end
 function add_jump_cost_function!(model, sys)
     jcosts = MTK.get_costs(sys)
     consolidate = MTK.get_consolidate(sys)
-    if isnothing(consolidate)
+    if isnothing(jcosts)
         @objective(model, Min, 0)
         return
     end
@@ -81,7 +82,7 @@ function add_jump_cost_function!(model, sys)
         t = only(arguments(st))
         idx = stidxmap[x(iv)]
         subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
-        jcosts = Symbolics.substitute(jcosts, Dict(x(t) => subval))
+        jcosts = map(c -> Symbolics.substitute(c, Dict(x(t) => subval)), jcosts)
     end
     
     for ct in controls(sys)
@@ -89,30 +90,27 @@ function add_jump_cost_function!(model, sys)
         t = only(arguments(ct))
         idx = cidxmap[p(iv)]
         subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
-        jcosts = Symbolics.substitute(jcosts, Dict(x(t) => subval))
+        jcosts = map(c -> Symbolics.substitute(c, Dict(p(t) => subval)), jcosts)
     end
     
     @objective(model, Min, consolidate(jcosts))
 end
 
 function add_user_constraints!(model, sys)
-    jconstraints = if !(csys = MTK.get_constraintsystem(sys) isa Nothing)
-        MTK.get_constraints(csys)
-    else
-        nothing
-    end
-    isnothing(jconstraints) && return nothing
+    conssys = MTK.get_constraintsystem(sys)
+    jconstraints = isnothing(conssys) ? nothing : MTK.get_constraints(conssys)
+    (isnothing(jconstraints) || isempty(jconstraints)) && return nothing
 
     iv = MTK.get_iv(sys)
     stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
     cidxmap = Dict([v => i for (i, v) in enumerate(controls(sys))])
 
-    for st in unknowns(sys)
+    for st in unknowns(conssys)
         x = operation(st)
         t = only(arguments(st))
         idx = stidxmap[x(iv)]
         subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
-        jconstraints = Symbolics.substitute(jconstraints, Dict(x(t) => subval))
+        jconstraints = map(c -> Symbolics.substitute(c, Dict(x(t) => subval)), jconstraints)
     end
 
     for ct in controls(sys)
@@ -120,16 +118,16 @@ function add_user_constraints!(model, sys)
         t = only(arguments(ct))
         idx = cidxmap[p(iv)]
         subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
-        jconstraints = Symbolics.substitute(jconstraints, Dict(p(t) => subval))
+        jconstraints = map(c -> Symbolics.substitute(jconstraints, Dict(p(t) => subval)), jconstriants)
     end
 
     for (i, cons) in enumerate(jconstraints)
         if cons isa Equation
-            @constraint(model, user[i], cons.lhs - cons.rhs == 0)
+            @constraint(model, cons.lhs - cons.rhs == 0, base_name = "user[$i]")
         elseif cons.relational_op === Symbolics.geq 
-            @constraint(model, user[i], cons.lhs - cons.rhs ≥ 0)
+            @constraint(model, cons.lhs - cons.rhs ≥ 0, base_name = "user[$i]")
         else
-            @constraint(model, user[i], cons.lhs - cons.rhs ≤ 0)
+            @constraint(model, cons.lhs - cons.rhs ≤ 0, base_name = "user[$i]")
         end
     end
 end
@@ -189,6 +187,7 @@ end
 struct JuMPControlSolution
     model::InfiniteModel
     sol::ODESolution
+    input_sol::Union{Nothing, ODESolution}
 end
 
 """
@@ -213,7 +212,6 @@ function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Sym
         end
     end
     for var in all_variables(model)
-        @show JuMP.name(var)
         if occursin("K", JuMP.name(var))
             unregister(model, Symbol(JuMP.name(var)))
             delete(model, var)
@@ -232,10 +230,19 @@ function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Sym
     U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:length(ts)]
     sol = DiffEqBase.build_solution(prob, ode_solver, ts, U_vals)
 
+    input_sol = nothing
+    if !isempty(model[:V])
+        V_vals = value.(model[:V])
+        V_vals = [[V_vals[i][j] for i in 1:length(V_vals)] for j in 1:length(ts)]
+        input_sol = DiffEqBase.build_solution(prob, ode_solver, ts, V_vals)
+    end
+
     if !(pstatus === FEASIBLE_POINT && (tstatus === OPTIMAL || tstatus === LOCALLY_SOLVED || tstatus === ALMOST_OPTIMAL || tstatus === ALMOST_LOCALLY_SOLVED))
         sol = SciMLBase.solution_new_retcode(sol, SciMLBase.ReturnCode.ConvergenceFailure)
+        !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode(input_sol, SciMLBase.ReturnCode.ConvergenceFailure))
     end
-    JuMPControlSolution(model, sol)
+
+    JuMPControlSolution(model, sol, input_sol)
 end
 
 end
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index 01b0ca5fbb..3acf36ec67 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -336,6 +336,8 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
 
     if length(costs) > 1 && isnothing(consolidate)
         error("Must specify a consolidation function for the costs vector.")
+    elseif isnothing(consolidate)
+        consolidate(u) = u[1]
     end
 
     assertions = Dict{BasicSymbolic, Any}(unwrap(k) => v for (k, v) in assertions)
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index a08c95b29b..62c0ea7dbd 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -17,7 +17,7 @@ const M = ModelingToolkit
 
     tspan = (0.0, 1.0)
     u0map = [x(t) => 4.0, y(t) => 2.0]
-    parammap = [α => 7.5, β => 4, γ => 8.0, δ => 5.0]
+    parammap = [α => 1.5, β => 1.0, γ => 3.0, δ => 1.0]
     @mtkbuild sys = ODESystem(eqs, t)
 
     # Test explicit method.
@@ -39,17 +39,33 @@ const M = ModelingToolkit
     constr = [x(0.6) ~ 3.5, x(0.3) ~ 7.0]
     @mtkbuild lksys = ODESystem(eqs, t; constraints = constr)
 
-    jprob = JuMPControlProblem(sys, u0map, tspan, parammap; guesses, dt = 0.01)
-    @test num_constraints(jprob.model) == 2 == num_variables(jprob.model) == 2
-    jsol = solve(prob, Ipopt.Optimizer, :Tsitouras5)
+    jprob = JuMPControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
+    @test num_constraints(jprob.model) == 2
+    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
     sol = jsol.sol
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0
 end
 
-@testset "Optimal control problem" begin
-    # Investing
-    
-   
-    # Bang-bang control
-end
+#@testset "Optimal control: bees" begin
+#    # Example from Lawrence Evans' notes
+#    M.@variables w(..) q(..) 
+#    M.@parameters α(t) [bounds = [0, 1]] b c μ s ν
+#
+#    tspan = (0, 4)
+#    eqs = [D(w(t)) ~ -μ*w(t) + b*s*α*w(t),
+#           D(q(t)) ~ -ν*q(t) + c*(1 - α)*s*w(t)]
+#    costs = [-q(tspan[2])]
+#   
+#    @mtkbuild beesys = ODESystem(eqs, t; costs)
+#    u0map = [w(0) => 40, q(0) => 2]
+#    pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1]
+#
+#    jprob = JuMPControlProblem(beesys, u0map, tspan, pmap)
+#    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+#    control_sol = jsol.control_sol
+#    # Bang-bang control
+#end
+#
+#@testset "Constrained optimal control problems" begin
+#end

From 916a531dcd8c7467212099520c728087e1ffa086 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 8 Apr 2025 18:38:53 -0400
Subject: [PATCH 12/41] fix: consolidate method

---
 src/systems/diffeqs/odesystem.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index 3acf36ec67..7c1acbdf91 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -337,7 +337,7 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
     if length(costs) > 1 && isnothing(consolidate)
         error("Must specify a consolidation function for the costs vector.")
     elseif isnothing(consolidate)
-        consolidate(u) = u[1]
+        consolidate = u -> u[1]
     end
 
     assertions = Dict{BasicSymbolic, Any}(unwrap(k) => v for (k, v) in assertions)

From 8935b229f8ea33a62e3567d2768f8bce6a881faa Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Wed, 9 Apr 2025 15:51:53 -0400
Subject: [PATCH 13/41] feat: InfiniteOptControlProblem

---
 ext/MTKJuMPControlExt.jl | 105 +++++++++++++++++++++++++++++++++------
 1 file changed, 90 insertions(+), 15 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 81a6fd2f74..1e775b1953 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -5,7 +5,9 @@ using DiffEqDevTools, DiffEqBase, SciMLBase
 using LinearAlgebra
 const MTK = ModelingToolkit
 
-struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <: SciMLBase.AbstractODEProblem{uType, tType, isinplace}
+abstract type AbstractOptimalControlProblem{uType, tType, isinplace} <: SciMLBase.AbstractODEProblem{uType, tType, isinplace} end
+
+struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <: AbstractOptimalControlProblem{uType, tType, isinplace}
        f::F
        u0::uType
        tspan::tType
@@ -18,6 +20,19 @@ struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <: SciMLBase.Abstrac
        end
 end
 
+struct InfiniteOptControlProblem{uType, tType, isinplace, P, F, K} <: SciMLBase.AbstractODEProblem{uType, tType, isinplace}
+       f::F
+       u0::uType
+       tspan::tType
+       p::P
+       model::InfiniteModel
+       kwargs::K
+
+       function InfiniteOptControlProblem(f, u0, tspan, p, model; kwargs...) 
+           new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f), typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
+       end
+end
+
 """
     JuMPControlProblem(sys::ODESystem, u0, tspan, p; dt)
 
@@ -34,24 +49,52 @@ The constraints are:
 - The solver constraints that encode the time-stepping used by the solver
 """
 function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPControlProblem."), guesses = Dict(), kwargs...)
-    ts = tspan[1]
-    te = tspan[2]
-    steps = ts:dt:te
-    ctrls = controls(sys)
-    states = unknowns(sys)
     constraintsys = MTK.get_constraintsystem(sys)
-
     if !isnothing(constraintsys)
         (length(constraints(constraintsys)) + length(u0map) > length(states)) &&
-            @warn "The JuMPControlProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
+            @warn "The control problem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
+    end
+
+    _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
+    f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
+        t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
+    model = init_model(sys, tspan[1]:dt:tspan[2], u0map)
+
+    JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
+end
+
+"""
+    InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap; dt)
+
+Convert an ODESystem 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 `JuMPControlProblem`, but directly adds the differential equations
+of the system as derivative constraints, rather than using a solver tableau.
+"""
+function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for InfiniteOptControlProblem."), guesses = Dict(), kwargs...)
+    constraintsys = MTK.get_constraintsystem(sys)
+    if !isnothing(constraintsys)
+        (length(constraints(constraintsys)) + length(u0map) > length(unknowns(sys))) &&
+            @warn "The control problem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
     end
 
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
     f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
+    model = init_model(sys, tspan[1]:dt:tspan[2], u0map)
+    add_infopt_solve_constraints!(model, sys, pmap)
+    InfiniteOptControlProblem(f, u0, tspan, p, model, kwargs...)
+end
+
+function init_model(sys, tsteps, u0map)
+    ctrls = controls(sys)
+    states = unknowns(sys)
+
     model = InfiniteModel()
-    @infinite_parameter(model, t in [ts, te], num_supports = length(steps))
+    @infinite_parameter(model, t in [tsteps[1], tsteps[end]], num_supports = length(tsteps))
     @variable(model, U[i = 1:length(states)], Infinite(t))
     @variable(model, V[1:length(ctrls)], Infinite(t))
 
@@ -61,8 +104,6 @@ function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
     u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) : [stidxmap[k] for (k, v) in u0map]
     add_initial_constraints!(model, u0, u0_idxs, tspan)
-
-    JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
 end
 
 function add_jump_cost_function!(model, sys)
@@ -140,7 +181,29 @@ end
 
 is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
 
-function add_solve_constraints!(prob, tableau)
+function add_infopt_solve_constraints!(model, sys, pmap)
+    iv = get_iv(sys)
+    t = model[:t]
+    U = model[:U]
+    V = model[:V]
+
+    stmap = Dict([v => U[i] for (i, v) in enumerate(unknowns(sys))])
+    ctrlmap = Dict([v => V[i] for (i, v) in enumerate(controls(sys))])
+    submap = merge(stmap, ctrlmap, pmap)
+
+    @register_symbolic _D(x) = ∂(x, t)
+    # Differential equations
+    diff_eqs = diff_equations(sys)
+    diff_eqs = map(e -> Symbolics.substitute(e, submap, Differential(iv) => _D), diff_eqs)
+    @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs == diff_eqs[i].rhs)
+
+    # Algebraic equations
+    alg_eqs = alg_equations(sys)
+    alg_eqs = map(e -> Symbolics.substitute(e, submap), alg_eqs)
+    @constraint(model, A[i = 1:length(alg_eqs)], alg_eqs[i].lhs == alg_eqs[i].rhs)
+end
+
+function add_jump_solve_constraints!(prob, tableau)
     A = tableau.A
     α = tableau.α
     c = tableau.c
@@ -202,7 +265,6 @@ function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Sym
     model = prob.model
     tableau_getter = Symbol(:construct, ode_solver)
     @eval tableau = $tableau_getter()
-    ts = supports(model[:t])
 
     # Unregister current solver constraints
     for con in all_constraints(model)
@@ -218,7 +280,19 @@ function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Sym
         end
     end
     add_solve_constraints!(prob, tableau)
+    _solve(prob, jump_solver, ode_solver)
+end
 
+"""
+`derivative_method` kwarg refers to the method used by InfiniteOpt to compute derivatives. The list of possible options can be found at https://infiniteopt.github.io/InfiniteOpt.jl/stable/guide/derivative/. Defaults to FiniteDifference(Backward()).
+"""
+function DiffEqBase.solve(prob::InfiniteOptControlProblem, jump_solver; derivative_method = InfiniteOpt.FiniteDifference(Backward()))
+    set_derivative_method(prob.model[:t], derivative_method)
+    _solve(prob, jump_solver, derivative_method)
+end
+
+function _solve(prob::AbstractOptimalControlProblem, jump_solver, solver)
+    model = prob.model
     set_optimizer(model, jump_solver)
     optimize!(model)
 
@@ -226,15 +300,16 @@ function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Sym
     pstatus = primal_status(model)
     !has_values(model) && error("Model not solvable; please report this to github.com/SciML/ModelingToolkit.jl.")
 
+    ts = supports(model[:t])
     U_vals = value.(model[:U])
     U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:length(ts)]
-    sol = DiffEqBase.build_solution(prob, ode_solver, ts, U_vals)
+    sol = DiffEqBase.build_solution(prob, solver, ts, U_vals)
 
     input_sol = nothing
     if !isempty(model[:V])
         V_vals = value.(model[:V])
         V_vals = [[V_vals[i][j] for i in 1:length(V_vals)] for j in 1:length(ts)]
-        input_sol = DiffEqBase.build_solution(prob, ode_solver, ts, V_vals)
+        input_sol = DiffEqBase.build_solution(prob, solver, ts, V_vals)
     end
 
     if !(pstatus === FEASIBLE_POINT && (tstatus === OPTIMAL || tstatus === LOCALLY_SOLVED || tstatus === ALMOST_OPTIMAL || tstatus === ALMOST_LOCALLY_SOLVED))

From 7b0a64a95fd787dced6f465f7050e7a5c6394846 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 10 Apr 2025 10:09:26 -0400
Subject: [PATCH 14/41] add InfiniteOPt dep

---
 Project.toml                    | 2 +-
 test/extensions/Project.toml    | 1 -
 test/extensions/jump_control.jl | 2 ++
 3 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index 2779b88b83..763c3931b9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -78,7 +78,7 @@ MTKChainRulesCoreExt = "ChainRulesCore"
 MTKDeepDiffsExt = "DeepDiffs"
 MTKFMIExt = "FMI"
 MTKInfiniteOptExt = "InfiniteOpt"
-MTKJuMPControlExt = ["JuMP", "DiffEqDevTools"]
+MTKJuMPControlExt = ["JuMP", "DiffEqDevTools", "InfiniteOpt"]
 MTKLabelledArraysExt = "LabelledArrays"
 
 [compat]
diff --git a/test/extensions/Project.toml b/test/extensions/Project.toml
index 9ce343202d..5800fe612f 100644
--- a/test/extensions/Project.toml
+++ b/test/extensions/Project.toml
@@ -5,7 +5,6 @@ ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index 62c0ea7dbd..7cd6b025b6 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -32,6 +32,8 @@ const M = ModelingToolkit
     jsol2 = solve(jprob, Ipopt.Optimizer, :ImplicitEuler)
     osol2 = solve(oprob, ImplicitEuler(), dt = 0.01, adaptive = false)
     @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001)
+    iprob = InfiniteOptControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
+    isol = solve(iprob, Ipopt.Optimizer, derivative_method = FiniteDifference(Backward()))
 
     # With a constraint
     u0map = Pair[]

From 1ea1885290bcaf90936a3684c17a6311ab1fdb24 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 10 Apr 2025 13:47:28 -0400
Subject: [PATCH 15/41] feat: add InfiniteOptControlProblem

---
 ext/MTKJuMPControlExt.jl         | 52 +++++++++++++++-----------------
 src/ModelingToolkit.jl           |  2 ++
 src/systems/diffeqs/odesystem.jl |  2 +-
 test/extensions/jump_control.jl  | 15 ++++++---
 4 files changed, 39 insertions(+), 32 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 1e775b1953..55c9dfc3c5 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -20,7 +20,7 @@ struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <: AbstractOptimalCo
        end
 end
 
-struct InfiniteOptControlProblem{uType, tType, isinplace, P, F, K} <: SciMLBase.AbstractODEProblem{uType, tType, isinplace}
+struct InfiniteOptControlProblem{uType, tType, isinplace, P, F, K} <: AbstractOptimalControlProblem{uType, tType, isinplace}
        f::F
        u0::uType
        tspan::tType
@@ -49,16 +49,10 @@ The constraints are:
 - The solver constraints that encode the time-stepping used by the solver
 """
 function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPControlProblem."), guesses = Dict(), kwargs...)
-    constraintsys = MTK.get_constraintsystem(sys)
-    if !isnothing(constraintsys)
-        (length(constraints(constraintsys)) + length(u0map) > length(states)) &&
-            @warn "The control problem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
-    end
-
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
     f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
-    model = init_model(sys, tspan[1]:dt:tspan[2], u0map)
+    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, u0)
 
     JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
 end
@@ -74,25 +68,24 @@ Related to `JuMPControlProblem`, but directly adds the differential equations
 of the system as derivative constraints, rather than using a solver tableau.
 """
 function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for InfiniteOptControlProblem."), guesses = Dict(), kwargs...)
-    constraintsys = MTK.get_constraintsystem(sys)
-    if !isnothing(constraintsys)
-        (length(constraints(constraintsys)) + length(u0map) > length(unknowns(sys))) &&
-            @warn "The control problem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
-    end
-
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
     f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    model = init_model(sys, tspan[1]:dt:tspan[2], u0map)
+    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, u0)
     add_infopt_solve_constraints!(model, sys, pmap)
     InfiniteOptControlProblem(f, u0, tspan, p, model, kwargs...)
 end
 
-function init_model(sys, tsteps, u0map)
+function init_model(sys, tsteps, u0map, u0)
+    constraintsys = MTK.get_constraintsystem(sys)
+    if !isnothing(constraintsys)
+        (length(constraints(constraintsys)) + length(u0map) > length(unknowns(sys))) &&
+            @warn "The control problem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
+    end
+
     ctrls = controls(sys)
     states = unknowns(sys)
-
     model = InfiniteModel()
     @infinite_parameter(model, t in [tsteps[1], tsteps[end]], num_supports = length(tsteps))
     @variable(model, U[i = 1:length(states)], Infinite(t))
@@ -103,13 +96,14 @@ function init_model(sys, tsteps, u0map)
 
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
     u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) : [stidxmap[k] for (k, v) in u0map]
-    add_initial_constraints!(model, u0, u0_idxs, tspan)
+    add_initial_constraints!(model, u0, u0_idxs, tsteps[1])
+    return model
 end
 
 function add_jump_cost_function!(model, sys)
     jcosts = MTK.get_costs(sys)
     consolidate = MTK.get_consolidate(sys)
-    if isnothing(jcosts)
+    if isnothing(jcosts) || isempty(jcosts)
         @objective(model, Min, 0)
         return
     end
@@ -173,8 +167,7 @@ function add_user_constraints!(model, sys)
     end
 end
 
-function add_initial_constraints!(model, u0, u0_idxs, tspan)
-    ts = tspan[1]
+function add_initial_constraints!(model, u0, u0_idxs, ts)
     U = model[:U]
     @constraint(model, initial[i in u0_idxs], U[i](ts) == u0[i])
 end
@@ -182,19 +175,24 @@ end
 is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
 
 function add_infopt_solve_constraints!(model, sys, pmap)
-    iv = get_iv(sys)
+    iv = MTK.get_iv(sys)
     t = model[:t]
     U = model[:U]
     V = model[:V]
 
     stmap = Dict([v => U[i] for (i, v) in enumerate(unknowns(sys))])
     ctrlmap = Dict([v => V[i] for (i, v) in enumerate(controls(sys))])
-    submap = merge(stmap, ctrlmap, pmap)
+    submap = merge(stmap, ctrlmap, Dict(pmap))
+    @show submap
 
-    @register_symbolic _D(x) = ∂(x, t)
     # Differential equations
     diff_eqs = diff_equations(sys)
-    diff_eqs = map(e -> Symbolics.substitute(e, submap, Differential(iv) => _D), diff_eqs)
+    D = Differential(iv)
+    diffsubmap = Dict([D(U[i]) => ∂(U[i], t) for i in 1:length(U)])
+    for u in unknowns(sys)
+        diff_eqs = map(e -> Symbolics.substitute(e, submap), diff_eqs)
+        diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs)
+    end
     @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs == diff_eqs[i].rhs)
 
     # Algebraic equations
@@ -273,13 +271,13 @@ function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Sym
             delete(model, con)
         end
     end
+    unregister(model, :K) 
     for var in all_variables(model)
         if occursin("K", JuMP.name(var))
-            unregister(model, Symbol(JuMP.name(var)))
             delete(model, var)
         end
     end
-    add_solve_constraints!(prob, tableau)
+    add_jump_solve_constraints!(prob, tableau)
     _solve(prob, jump_solver, ode_solver)
 end
 
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 5f8c2d7610..87ca36ef60 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -349,5 +349,7 @@ function FMIComponent end
 
 function JuMPControlProblem end
 export JuMPControlProblem
+function InfiniteOptControlProblem end
+export InfiniteOptControlProblem
 
 end # module
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index 7c1acbdf91..7e3f11b58d 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -336,7 +336,7 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
 
     if length(costs) > 1 && isnothing(consolidate)
         error("Must specify a consolidation function for the costs vector.")
-    elseif isnothing(consolidate)
+    elseif length(costs) == 1 && isnothing(consolidate)
         consolidate = u -> u[1]
     end
 
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index 7cd6b025b6..2b01e612e3 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -4,6 +4,7 @@ using DiffEqDevTools, DiffEqBase
 using SimpleDiffEq
 using OrdinaryDiffEqSDIRK
 using Ipopt 
+using BenchmarkTools
 const M = ModelingToolkit
 
 @testset "ODE Solution, no cost" begin
@@ -29,11 +30,11 @@ const M = ModelingToolkit
     @test jsol.sol.u ≈ osol.u
 
     # Implicit method.
-    jsol2 = solve(jprob, Ipopt.Optimizer, :ImplicitEuler)
-    osol2 = solve(oprob, ImplicitEuler(), dt = 0.01, adaptive = false)
+    jsol2 = @btime solve($jprob, Ipopt.Optimizer, :ImplicitEuler) # 63.031 ms, 26.49 MiB
+    osol2 = @btime solve($oprob, ImplicitEuler(), dt = 0.01, adaptive = false) # 129.375 μs, 61.91 KiB
     @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001)
     iprob = InfiniteOptControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
-    isol = solve(iprob, Ipopt.Optimizer, derivative_method = FiniteDifference(Backward()))
+    isol = @btime solve($iprob, Ipopt.Optimizer, derivative_method = FiniteDifference(Backward())) # 11.540 ms, 4.00 MiB
 
     # With a constraint
     u0map = Pair[]
@@ -43,10 +44,16 @@ const M = ModelingToolkit
 
     jprob = JuMPControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     @test num_constraints(jprob.model) == 2
-    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    jsol = @btime solve($jprob, Ipopt.Optimizer, :Tsitouras5) # 12.190 s, 9.68 GiB
     sol = jsol.sol
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0
+
+    iprob = InfiniteOptControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
+    isol = @btime solve($iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3)) # 48.564 ms, 9.58 MiB
+    sol = isol.sol
+    @test sol(0.6)[1] ≈ 3.5
+    @test sol(0.3)[1] ≈ 7.0
 end
 
 #@testset "Optimal control: bees" begin

From 92c0d045de16d34a43187ee987c02679ef18db8c Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 18 Apr 2025 05:25:46 -0400
Subject: [PATCH 16/41] fix merge

---
 ext/MTKJuMPControlExt.jl        | 125 ++++++++++++++++++--------------
 test/extensions/jump_control.jl |  14 ++--
 2 files changed, 81 insertions(+), 58 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 55c9dfc3c5..50bbc05151 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -5,32 +5,37 @@ using DiffEqDevTools, DiffEqBase, SciMLBase
 using LinearAlgebra
 const MTK = ModelingToolkit
 
-abstract type AbstractOptimalControlProblem{uType, tType, isinplace} <: SciMLBase.AbstractODEProblem{uType, tType, isinplace} end
-
-struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <: AbstractOptimalControlProblem{uType, tType, isinplace}
-       f::F
-       u0::uType
-       tspan::tType
-       p::P
-       model::InfiniteModel
-       kwargs::K
-
-       function JuMPControlProblem(f, u0, tspan, p, model; kwargs...) 
-           new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f), typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
-       end
+abstract type AbstractOptimalControlProblem{uType, tType, isinplace} <:
+              SciMLBase.AbstractODEProblem{uType, tType, isinplace} end
+
+struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <:
+       AbstractOptimalControlProblem{uType, tType, isinplace}
+    f::F
+    u0::uType
+    tspan::tType
+    p::P
+    model::InfiniteModel
+    kwargs::K
+
+    function JuMPControlProblem(f, u0, tspan, p, model; kwargs...)
+        new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f),
+            typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
+    end
 end
 
-struct InfiniteOptControlProblem{uType, tType, isinplace, P, F, K} <: AbstractOptimalControlProblem{uType, tType, isinplace}
-       f::F
-       u0::uType
-       tspan::tType
-       p::P
-       model::InfiniteModel
-       kwargs::K
-
-       function InfiniteOptControlProblem(f, u0, tspan, p, model; kwargs...) 
-           new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f), typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
-       end
+struct InfiniteOptControlProblem{uType, tType, isinplace, P, F, K} <:
+       AbstractOptimalControlProblem{uType, tType, isinplace}
+    f::F
+    u0::uType
+    tspan::tType
+    p::P
+    model::InfiniteModel
+    kwargs::K
+
+    function InfiniteOptControlProblem(f, u0, tspan, p, model; kwargs...)
+        new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f),
+            typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
+    end
 end
 
 """
@@ -48,7 +53,9 @@ The constraints are:
 - The set of user constraints passed to the ODESystem via `constraints`
 - The solver constraints that encode the time-stepping used by the solver
 """
-function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPControlProblem."), guesses = Dict(), kwargs...)
+function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap;
+        dt = error("dt must be provided for JuMPControlProblem."),
+        guesses = Dict(), kwargs...)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
     f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
@@ -67,7 +74,9 @@ of the interpolation arrays.
 Related to `JuMPControlProblem`, but directly adds the differential equations
 of the system as derivative constraints, rather than using a solver tableau.
 """
-function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for InfiniteOptControlProblem."), guesses = Dict(), kwargs...)
+function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
+        dt = error("dt must be provided for InfiniteOptControlProblem."),
+        guesses = Dict(), kwargs...)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
     f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
@@ -87,7 +96,7 @@ function init_model(sys, tsteps, u0map, u0)
     ctrls = controls(sys)
     states = unknowns(sys)
     model = InfiniteModel()
-    @infinite_parameter(model, t in [tsteps[1], tsteps[end]], num_supports = length(tsteps))
+    @infinite_parameter(model, t in [tsteps[1], tsteps[end]], num_supports=length(tsteps))
     @variable(model, U[i = 1:length(states)], Infinite(t))
     @variable(model, V[1:length(ctrls)], Infinite(t))
 
@@ -95,7 +104,8 @@ function init_model(sys, tsteps, u0map, u0)
     add_user_constraints!(model, sys)
 
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
-    u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) : [stidxmap[k] for (k, v) in u0map]
+    u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) :
+              [stidxmap[k] for (k, v) in u0map]
     add_initial_constraints!(model, u0, u0_idxs, tsteps[1])
     return model
 end
@@ -119,7 +129,7 @@ function add_jump_cost_function!(model, sys)
         subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
         jcosts = map(c -> Symbolics.substitute(c, Dict(x(t) => subval)), jcosts)
     end
-    
+
     for ct in controls(sys)
         p = operation(ct)
         t = only(arguments(ct))
@@ -127,7 +137,7 @@ function add_jump_cost_function!(model, sys)
         subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
         jcosts = map(c -> Symbolics.substitute(c, Dict(p(t) => subval)), jcosts)
     end
-    
+
     @objective(model, Min, consolidate(jcosts))
 end
 
@@ -153,23 +163,24 @@ function add_user_constraints!(model, sys)
         t = only(arguments(ct))
         idx = cidxmap[p(iv)]
         subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
-        jconstraints = map(c -> Symbolics.substitute(jconstraints, Dict(p(t) => subval)), jconstriants)
+        jconstraints = map(
+            c -> Symbolics.substitute(jconstraints, Dict(p(t) => subval)), jconstriants)
     end
 
     for (i, cons) in enumerate(jconstraints)
         if cons isa Equation
-            @constraint(model, cons.lhs - cons.rhs == 0, base_name = "user[$i]")
-        elseif cons.relational_op === Symbolics.geq 
-            @constraint(model, cons.lhs - cons.rhs ≥ 0, base_name = "user[$i]")
+            @constraint(model, cons.lhs - cons.rhs==0, base_name="user[$i]")
+        elseif cons.relational_op === Symbolics.geq
+            @constraint(model, cons.lhs - cons.rhs≥0, base_name="user[$i]")
         else
-            @constraint(model, cons.lhs - cons.rhs ≤ 0, base_name = "user[$i]")
+            @constraint(model, cons.lhs - cons.rhs≤0, base_name="user[$i]")
         end
     end
 end
 
 function add_initial_constraints!(model, u0, u0_idxs, ts)
     U = model[:U]
-    @constraint(model, initial[i in u0_idxs], U[i](ts) == u0[i])
+    @constraint(model, initial[i in u0_idxs], U[i](ts)==u0[i])
 end
 
 is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
@@ -193,12 +204,12 @@ function add_infopt_solve_constraints!(model, sys, pmap)
         diff_eqs = map(e -> Symbolics.substitute(e, submap), diff_eqs)
         diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs)
     end
-    @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs == diff_eqs[i].rhs)
+    @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs==diff_eqs[i].rhs)
 
     # Algebraic equations
     alg_eqs = alg_equations(sys)
     alg_eqs = map(e -> Symbolics.substitute(e, submap), alg_eqs)
-    @constraint(model, A[i = 1:length(alg_eqs)], alg_eqs[i].lhs == alg_eqs[i].rhs)
+    @constraint(model, A[i = 1:length(alg_eqs)], alg_eqs[i].lhs==alg_eqs[i].rhs)
 end
 
 function add_jump_solve_constraints!(prob, tableau)
@@ -219,26 +230,29 @@ function add_jump_solve_constraints!(prob, tableau)
         K = Any[]
         for τ in tsteps
             for (i, h) in enumerate(c)
-                ΔU = sum([A[i, j] * K[j] for j in 1:i-1], init = zeros(nᵤ))
-                Uₙ = [U[i](τ) + ΔU[i]*dt for i in 1:nᵤ]
-                Kₙ = f(Uₙ, p, τ + h*dt)
+                ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = zeros(nᵤ))
+                Uₙ = [U[i](τ) + ΔU[i] * dt for i in 1:nᵤ]
+                Kₙ = f(Uₙ, p, τ + h * dt)
                 push!(K, Kₙ)
             end
-            ΔU = dt*sum([α[i] * K[i] for i in 1:length(α)])
-            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU[n] == U[n](τ + dt), base_name = "solve_time_$τ")
+            ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)])
+            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU[n]==U[n](τ + dt),
+                base_name="solve_time_$τ")
             empty!(K)
         end
     else
-        @variable(model, K[1:length(α), 1:nᵤ], Infinite(t), start = tsteps[1])
+        @variable(model, K[1:length(α), 1:nᵤ], Infinite(t), start=tsteps[1])
         for τ in tsteps
             ΔUs = A * K
             for (i, h) in enumerate(c)
                 ΔU = ΔUs[i, :]
-                Uₙ = [U[j] + ΔU[j]*dt for j in 1:nᵤ]
-                @constraint(model, [j in 1:nᵤ], K[i, j] == f(Uₙ, p, τ + h*dt)[j], DomainRestrictions(t => τ), base_name = "solve_K($τ)")
+                Uₙ = [U[j] + ΔU[j] * dt for j in 1:nᵤ]
+                @constraint(model, [j in 1:nᵤ], K[i, j]==f(Uₙ, p, τ + h * dt)[j],
+                    DomainRestrictions(t => τ), base_name="solve_K($τ)")
             end
-            ΔU = dt*sum([α[i] * K[i, :] for i in 1:length(α)])
-            @constraint(model, [n = 1:nᵤ], U[n] + ΔU[n] == U[n](τ + dt), DomainRestrictions(t => τ), base_name = "solve_U($τ)")
+            ΔU = dt * sum([α[i] * K[i, :] for i in 1:length(α)])
+            @constraint(model, [n = 1:nᵤ], U[n] + ΔU[n]==U[n](τ + dt),
+                DomainRestrictions(t => τ), base_name="solve_U($τ)")
         end
     end
 end
@@ -271,7 +285,7 @@ function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Sym
             delete(model, con)
         end
     end
-    unregister(model, :K) 
+    unregister(model, :K)
     for var in all_variables(model)
         if occursin("K", JuMP.name(var))
             delete(model, var)
@@ -284,7 +298,8 @@ end
 """
 `derivative_method` kwarg refers to the method used by InfiniteOpt to compute derivatives. The list of possible options can be found at https://infiniteopt.github.io/InfiniteOpt.jl/stable/guide/derivative/. Defaults to FiniteDifference(Backward()).
 """
-function DiffEqBase.solve(prob::InfiniteOptControlProblem, jump_solver; derivative_method = InfiniteOpt.FiniteDifference(Backward()))
+function DiffEqBase.solve(prob::InfiniteOptControlProblem, jump_solver;
+        derivative_method = InfiniteOpt.FiniteDifference(Backward()))
     set_derivative_method(prob.model[:t], derivative_method)
     _solve(prob, jump_solver, derivative_method)
 end
@@ -296,7 +311,8 @@ function _solve(prob::AbstractOptimalControlProblem, jump_solver, solver)
 
     tstatus = termination_status(model)
     pstatus = primal_status(model)
-    !has_values(model) && error("Model not solvable; please report this to github.com/SciML/ModelingToolkit.jl.")
+    !has_values(model) &&
+        error("Model not solvable; please report this to github.com/SciML/ModelingToolkit.jl.")
 
     ts = supports(model[:t])
     U_vals = value.(model[:U])
@@ -310,9 +326,12 @@ function _solve(prob::AbstractOptimalControlProblem, jump_solver, solver)
         input_sol = DiffEqBase.build_solution(prob, solver, ts, V_vals)
     end
 
-    if !(pstatus === FEASIBLE_POINT && (tstatus === OPTIMAL || tstatus === LOCALLY_SOLVED || tstatus === ALMOST_OPTIMAL || tstatus === ALMOST_LOCALLY_SOLVED))
+    if !(pstatus === FEASIBLE_POINT &&
+         (tstatus === OPTIMAL || tstatus === LOCALLY_SOLVED || tstatus === ALMOST_OPTIMAL ||
+          tstatus === ALMOST_LOCALLY_SOLVED))
         sol = SciMLBase.solution_new_retcode(sol, SciMLBase.ReturnCode.ConvergenceFailure)
-        !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode(input_sol, SciMLBase.ReturnCode.ConvergenceFailure))
+        !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode(
+            input_sol, SciMLBase.ReturnCode.ConvergenceFailure))
     end
 
     JuMPControlSolution(model, sol, input_sol)
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index 2b01e612e3..7dd225aae0 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -3,7 +3,7 @@ using JuMP, InfiniteOpt
 using DiffEqDevTools, DiffEqBase
 using SimpleDiffEq
 using OrdinaryDiffEqSDIRK
-using Ipopt 
+using Ipopt
 using BenchmarkTools
 const M = ModelingToolkit
 
@@ -11,7 +11,8 @@ const M = ModelingToolkit
     # Test solving without anything attached.
     @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
     M.@variables x(..) y(..)
-    t = M.t_nounits; D = M.D_nounits
+    t = M.t_nounits
+    D = M.D_nounits
 
     eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t),
         D(y(t)) ~ -γ * y(t) + δ * x(t) * y(t)]
@@ -34,7 +35,8 @@ const M = ModelingToolkit
     osol2 = @btime solve($oprob, ImplicitEuler(), dt = 0.01, adaptive = false) # 129.375 μs, 61.91 KiB
     @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001)
     iprob = InfiniteOptControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
-    isol = @btime solve($iprob, Ipopt.Optimizer, derivative_method = FiniteDifference(Backward())) # 11.540 ms, 4.00 MiB
+    isol = @btime solve(
+        $iprob, Ipopt.Optimizer, derivative_method = FiniteDifference(Backward())) # 11.540 ms, 4.00 MiB
 
     # With a constraint
     u0map = Pair[]
@@ -49,8 +51,10 @@ const M = ModelingToolkit
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0
 
-    iprob = InfiniteOptControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
-    isol = @btime solve($iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3)) # 48.564 ms, 9.58 MiB
+    iprob = InfiniteOptControlProblem(
+        lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
+    isol = @btime solve(
+        $iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3)) # 48.564 ms, 9.58 MiB
     sol = isol.sol
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0

From a4570299ab1c6aeca6a93d310b04c89b00cefa43 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 10 Apr 2025 17:51:58 -0400
Subject: [PATCH 17/41] refactor: add optimal control interface file

---
 ext/MTKJuMPControlExt.jl                 | 31 ++++++------------------
 src/ModelingToolkit.jl                   |  7 +++---
 src/systems/optimal_control_interface.jl | 21 ++++++++++++++++
 3 files changed, 32 insertions(+), 27 deletions(-)
 create mode 100644 src/systems/optimal_control_interface.jl

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 50bbc05151..091391caaa 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -1,13 +1,10 @@
 module MTKJuMPControlExt
 using ModelingToolkit
 using JuMP, InfiniteOpt
-using DiffEqDevTools, DiffEqBase, SciMLBase
+using DiffEqDevTools, DiffEqBase
 using LinearAlgebra
 const MTK = ModelingToolkit
 
-abstract type AbstractOptimalControlProblem{uType, tType, isinplace} <:
-              SciMLBase.AbstractODEProblem{uType, tType, isinplace} end
-
 struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <:
        AbstractOptimalControlProblem{uType, tType, isinplace}
     f::F
@@ -56,6 +53,7 @@ The constraints are:
 function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap;
         dt = error("dt must be provided for JuMPControlProblem."),
         guesses = Dict(), kwargs...)
+    MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
     f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
@@ -77,6 +75,7 @@ of the system as derivative constraints, rather than using a solver tableau.
 function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
         dt = error("dt must be provided for InfiniteOptControlProblem."),
         guesses = Dict(), kwargs...)
+    MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
     f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
@@ -87,12 +86,6 @@ function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
 end
 
 function init_model(sys, tsteps, u0map, u0)
-    constraintsys = MTK.get_constraintsystem(sys)
-    if !isnothing(constraintsys)
-        (length(constraints(constraintsys)) + length(u0map) > length(unknowns(sys))) &&
-            @warn "The control problem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
-    end
-
     ctrls = controls(sys)
     states = unknowns(sys)
     model = InfiniteModel()
@@ -110,7 +103,7 @@ function init_model(sys, tsteps, u0map, u0)
     return model
 end
 
-function add_jump_cost_function!(model, sys)
+function add_jump_cost_function!(model::InfiniteModel, sys)
     jcosts = MTK.get_costs(sys)
     consolidate = MTK.get_consolidate(sys)
     if isnothing(jcosts) || isempty(jcosts)
@@ -141,7 +134,7 @@ function add_jump_cost_function!(model, sys)
     @objective(model, Min, consolidate(jcosts))
 end
 
-function add_user_constraints!(model, sys)
+function add_user_constraints!(model::InfiniteModel, sys)
     conssys = MTK.get_constraintsystem(sys)
     jconstraints = isnothing(conssys) ? nothing : MTK.get_constraints(conssys)
     (isnothing(jconstraints) || isempty(jconstraints)) && return nothing
@@ -178,14 +171,14 @@ function add_user_constraints!(model, sys)
     end
 end
 
-function add_initial_constraints!(model, u0, u0_idxs, ts)
+function add_initial_constraints!(model::InfiniteModel, u0, u0_idxs, ts)
     U = model[:U]
     @constraint(model, initial[i in u0_idxs], U[i](ts)==u0[i])
 end
 
 is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
 
-function add_infopt_solve_constraints!(model, sys, pmap)
+function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap)
     iv = MTK.get_iv(sys)
     t = model[:t]
     U = model[:U]
@@ -257,14 +250,6 @@ function add_jump_solve_constraints!(prob, tableau)
     end
 end
 
-"""
-"""
-struct JuMPControlSolution
-    model::InfiniteModel
-    sol::ODESolution
-    input_sol::Union{Nothing, ODESolution}
-end
-
 """
 Solve JuMPControlProblem. Arguments:
 - prob: a JumpControlProblem
@@ -334,7 +319,7 @@ function _solve(prob::AbstractOptimalControlProblem, jump_solver, solver)
             input_sol, SciMLBase.ReturnCode.ConvergenceFailure))
     end
 
-    JuMPControlSolution(model, sol, input_sol)
+    OptimalControlSolution(model, sol, input_sol)
 end
 
 end
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 87ca36ef60..0c83bbea10 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -347,9 +347,8 @@ export AnalysisPoint, get_sensitivity_function, get_comp_sensitivity_function,
        open_loop
 function FMIComponent end
 
-function JuMPControlProblem end
-export JuMPControlProblem
-function InfiniteOptControlProblem end
-export InfiniteOptControlProblem
+include("src/systems/optimal_control_interface.jl")
+export JuMPControlProblem, InfiniteOptControlProblem, PyomoControlProblem, CasADiControlProblem
+export OptimalControlSolution
 
 end # module
diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
new file mode 100644
index 0000000000..d93cb4cfb7
--- /dev/null
+++ b/src/systems/optimal_control_interface.jl
@@ -0,0 +1,21 @@
+abstract type AbstractOptimalControlProblem{uType, tType, isinplace} <:
+              SciMLBase.AbstractODEProblem{uType, tType, isinplace} end
+
+struct OptimalControlSolution
+    model::Any
+    sol::ODESolution
+    input_sol::Union{Nothing, ODESolution}
+end
+
+function JuMPControlProblem end
+function InfiniteOptControlProblem end
+function CasADiControlProblem end
+function PyomoControlProblem end
+
+function warn_overdetermined(sys, u0map)
+    constraintsys = get_constraintsystem(sys)
+    if !isnothing(constraintsys)
+        (length(constraints(constraintsys)) + length(u0map) > length(unknowns(sys))) &&
+            @warn "The control problem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
+    end
+end

From 565517f75d2149e7be56dc49d0e70ea2528882b4 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 10 Apr 2025 18:03:55 -0400
Subject: [PATCH 18/41] add set_silent option

---
 ext/MTKJuMPControlExt.jl        | 5 ++++-
 src/ModelingToolkit.jl          | 4 ++--
 test/extensions/jump_control.jl | 8 ++++----
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 091391caaa..f4b7298945 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -255,13 +255,15 @@ Solve JuMPControlProblem. Arguments:
 - prob: a JumpControlProblem
 - jump_solver: a LP solver such as HiGHS
 - ode_solver: Takes in a symbol representing the solver. Acceptable solvers may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. Note that the symbol may be different than the typical name of the solver, e.g. :Tsitouras5 rather than Tsit5.
+- silent: set the model silent (suppress model output)
 
 Returns a JuMPControlSolution, which contains both the model and the ODE solution.
 """
-function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Symbol)
+function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Symbol; silent = false)
     model = prob.model
     tableau_getter = Symbol(:construct, ode_solver)
     @eval tableau = $tableau_getter()
+    silent && set_silent(model)
 
     # Unregister current solver constraints
     for con in all_constraints(model)
@@ -285,6 +287,7 @@ end
 """
 function DiffEqBase.solve(prob::InfiniteOptControlProblem, jump_solver;
         derivative_method = InfiniteOpt.FiniteDifference(Backward()))
+    silent && set_silent(model)
     set_derivative_method(prob.model[:t], derivative_method)
     _solve(prob, jump_solver, derivative_method)
 end
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 0c83bbea10..323b719aa7 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -347,8 +347,8 @@ export AnalysisPoint, get_sensitivity_function, get_comp_sensitivity_function,
        open_loop
 function FMIComponent end
 
-include("src/systems/optimal_control_interface.jl")
-export JuMPControlProblem, InfiniteOptControlProblem, PyomoControlProblem, CasADiControlProblem
+include("systems/optimal_control_interface.jl")
+export AbstractOptimalControlProblem, JuMPControlProblem, InfiniteOptControlProblem, PyomoControlProblem, CasADiControlProblem
 export OptimalControlSolution
 
 end # module
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index 7dd225aae0..7e63edae24 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -31,12 +31,12 @@ const M = ModelingToolkit
     @test jsol.sol.u ≈ osol.u
 
     # Implicit method.
-    jsol2 = @btime solve($jprob, Ipopt.Optimizer, :ImplicitEuler) # 63.031 ms, 26.49 MiB
+    jsol2 = @btime solve($jprob, Ipopt.Optimizer, :ImplicitEuler, silent = true) # 63.031 ms, 26.49 MiB
     osol2 = @btime solve($oprob, ImplicitEuler(), dt = 0.01, adaptive = false) # 129.375 μs, 61.91 KiB
     @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001)
     iprob = InfiniteOptControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
     isol = @btime solve(
-        $iprob, Ipopt.Optimizer, derivative_method = FiniteDifference(Backward())) # 11.540 ms, 4.00 MiB
+        $iprob, Ipopt.Optimizer, derivative_method = FiniteDifference(Backward()), silent = true) # 11.540 ms, 4.00 MiB
 
     # With a constraint
     u0map = Pair[]
@@ -46,7 +46,7 @@ const M = ModelingToolkit
 
     jprob = JuMPControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     @test num_constraints(jprob.model) == 2
-    jsol = @btime solve($jprob, Ipopt.Optimizer, :Tsitouras5) # 12.190 s, 9.68 GiB
+    jsol = @btime solve($jprob, Ipopt.Optimizer, :Tsitouras5, silent = true) # 12.190 s, 9.68 GiB
     sol = jsol.sol
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0
@@ -54,7 +54,7 @@ const M = ModelingToolkit
     iprob = InfiniteOptControlProblem(
         lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     isol = @btime solve(
-        $iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3)) # 48.564 ms, 9.58 MiB
+        $iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
     sol = isol.sol
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0

From a276a9c70e5328abda12ccc41dbe65b4320b81ff Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 10 Apr 2025 18:04:13 -0400
Subject: [PATCH 19/41] format

---
 ext/MTKJuMPControlExt.jl | 3 ++-
 src/ModelingToolkit.jl   | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index f4b7298945..45e64fa8c8 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -259,7 +259,8 @@ Solve JuMPControlProblem. Arguments:
 
 Returns a JuMPControlSolution, which contains both the model and the ODE solution.
 """
-function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Symbol; silent = false)
+function DiffEqBase.solve(
+        prob::JuMPControlProblem, jump_solver, ode_solver::Symbol; silent = false)
     model = prob.model
     tableau_getter = Symbol(:construct, ode_solver)
     @eval tableau = $tableau_getter()
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 323b719aa7..0a2ea31181 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -348,7 +348,8 @@ export AnalysisPoint, get_sensitivity_function, get_comp_sensitivity_function,
 function FMIComponent end
 
 include("systems/optimal_control_interface.jl")
-export AbstractOptimalControlProblem, JuMPControlProblem, InfiniteOptControlProblem, PyomoControlProblem, CasADiControlProblem
+export AbstractOptimalControlProblem, JuMPControlProblem, InfiniteOptControlProblem,
+       PyomoControlProblem, CasADiControlProblem
 export OptimalControlSolution
 
 end # module

From 6085324bd7a66c191f8c201f2914e932c03aaa1d Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Wed, 16 Apr 2025 15:26:25 -0400
Subject: [PATCH 20/41] partial: add free final time and bounds-handling

---
 ext/MTKJuMPControlExt.jl                 | 33 +++++++--
 src/systems/optimal_control_interface.jl |  6 ++
 src/systems/problem_utils.jl             |  1 +
 test/extensions/jump_control.jl          | 88 ++++++++++++++++++------
 4 files changed, 99 insertions(+), 29 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 45e64fa8c8..a178280428 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -55,8 +55,11 @@ function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap;
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
+    @show _u0map
     f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
+
+    (f_i, f_o) = generate_control_function(sys)
     model = init_model(sys, tspan[1]:dt:tspan[2], u0map, u0)
 
     JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
@@ -86,13 +89,14 @@ function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
 end
 
 function init_model(sys, tsteps, u0map, u0)
-    ctrls = controls(sys)
+    ctrls = MTK.unbound_inputs(sys)
     states = unknowns(sys)
     model = InfiniteModel()
     @infinite_parameter(model, t in [tsteps[1], tsteps[end]], num_supports=length(tsteps))
     @variable(model, U[i = 1:length(states)], Infinite(t))
     @variable(model, V[1:length(ctrls)], Infinite(t))
 
+    set_bounds!(model, sys)
     add_jump_cost_function!(model, sys)
     add_user_constraints!(model, sys)
 
@@ -103,6 +107,22 @@ function init_model(sys, tsteps, u0map, u0)
     return model
 end
 
+function set_bounds!(model, sys)
+    U = model[:U]
+    for (i, u) in enumerate(unknowns(sys))
+        lo, hi = MTK.getbounds(u)
+        set_lower_bound(U[i], lo)
+        set_upper_bound(U[i], hi)
+    end
+
+    V = model[:V]
+    for (i, v) in enumerate(MTK.unbound_inputs(sys))
+        lo, hi = MTK.getbounds(v)
+        set_lower_bound(V[i], lo)
+        set_upper_bound(V[i], hi)
+    end
+end
+
 function add_jump_cost_function!(model::InfiniteModel, sys)
     jcosts = MTK.get_costs(sys)
     consolidate = MTK.get_consolidate(sys)
@@ -113,7 +133,7 @@ function add_jump_cost_function!(model::InfiniteModel, sys)
     iv = MTK.get_iv(sys)
 
     stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
-    cidxmap = Dict([v => i for (i, v) in enumerate(controls(sys))])
+    cidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))])
 
     for st in unknowns(sys)
         x = operation(st)
@@ -123,7 +143,7 @@ function add_jump_cost_function!(model::InfiniteModel, sys)
         jcosts = map(c -> Symbolics.substitute(c, Dict(x(t) => subval)), jcosts)
     end
 
-    for ct in controls(sys)
+    for ct in MTK.unbound_inputs(sys)
         p = operation(ct)
         t = only(arguments(ct))
         idx = cidxmap[p(iv)]
@@ -141,7 +161,7 @@ function add_user_constraints!(model::InfiniteModel, sys)
 
     iv = MTK.get_iv(sys)
     stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
-    cidxmap = Dict([v => i for (i, v) in enumerate(controls(sys))])
+    cidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))])
 
     for st in unknowns(conssys)
         x = operation(st)
@@ -151,7 +171,7 @@ function add_user_constraints!(model::InfiniteModel, sys)
         jconstraints = map(c -> Symbolics.substitute(c, Dict(x(t) => subval)), jconstraints)
     end
 
-    for ct in controls(sys)
+    for ct in MTK.unbound_inputs(sys)
         p = operation(ct)
         t = only(arguments(ct))
         idx = cidxmap[p(iv)]
@@ -185,9 +205,8 @@ function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap)
     V = model[:V]
 
     stmap = Dict([v => U[i] for (i, v) in enumerate(unknowns(sys))])
-    ctrlmap = Dict([v => V[i] for (i, v) in enumerate(controls(sys))])
+    ctrlmap = Dict([v => V[i] for (i, v) in enumerate(MTK.unbound_inputs(sys))])
     submap = merge(stmap, ctrlmap, Dict(pmap))
-    @show submap
 
     # Differential equations
     diff_eqs = diff_equations(sys)
diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index d93cb4cfb7..f8e9c75887 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -19,3 +19,9 @@ function warn_overdetermined(sys, u0map)
             @warn "The control problem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
     end
 end
+
+"""
+IntegralNorm. When applied to an expression.
+"""
+struct IntegralNorm end
+
diff --git a/src/systems/problem_utils.jl b/src/systems/problem_utils.jl
index 41c64b78c5..94a96b8959 100644
--- a/src/systems/problem_utils.jl
+++ b/src/systems/problem_utils.jl
@@ -822,6 +822,7 @@ function maybe_build_initialization_problem(
         t = zero(floatT)
     end
 
+    @show u0map
     initializeprob = ModelingToolkit.InitializationProblem{true, SciMLBase.FullSpecialize}(
         sys, t, u0map, pmap; guesses, initialization_eqs, use_scc, kwargs...)
     if state_values(initializeprob) !== nothing
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index 7e63edae24..ee9ccd3a07 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -1,26 +1,27 @@
 using ModelingToolkit
-using JuMP, InfiniteOpt
+import JuMP, InfiniteOpt
 using DiffEqDevTools, DiffEqBase
 using SimpleDiffEq
 using OrdinaryDiffEqSDIRK
 using Ipopt
 using BenchmarkTools
+using CairoMakie
 const M = ModelingToolkit
 
 @testset "ODE Solution, no cost" begin
     # Test solving without anything attached.
     @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
-    M.@variables x(..) y(..)
+    @variables x(..) y(..)
     t = M.t_nounits
     D = M.D_nounits
 
     eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t),
         D(y(t)) ~ -γ * y(t) + δ * x(t) * y(t)]
 
+    @mtkbuild sys = ODESystem(eqs, t)
     tspan = (0.0, 1.0)
     u0map = [x(t) => 4.0, y(t) => 2.0]
     parammap = [α => 1.5, β => 1.0, γ => 3.0, δ => 1.0]
-    @mtkbuild sys = ODESystem(eqs, t)
 
     # Test explicit method.
     jprob = JuMPControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
@@ -58,27 +59,70 @@ const M = ModelingToolkit
     sol = isol.sol
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0
+
+    # Test whole-interval constraints
+    constr = [x(t) > 3, y(t) > 4]
+    @mtkbuild lksys = ODESystem(eqs, t; constraints = constr)
+    iprob = InfiniteOptControlProblem(
+        lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
+    isol = @btime solve(
+        $iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
+    sol = isol.sol
+    @test all(u -> u .> [3, 4], sol.u)
+
+    jprob = JuMPControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
+    jsol = @btime solve($jprob, Ipopt.Optimizer, :RadauIA3, silent = true) # 12.190 s, 9.68 GiB
+    sol = jsol.sol
+    @test all(u -> u .> [3, 4], sol.u)
 end
 
-#@testset "Optimal control: bees" begin
-#    # Example from Lawrence Evans' notes
-#    M.@variables w(..) q(..) 
-#    M.@parameters α(t) [bounds = [0, 1]] b c μ s ν
-#
-#    tspan = (0, 4)
-#    eqs = [D(w(t)) ~ -μ*w(t) + b*s*α*w(t),
-#           D(q(t)) ~ -ν*q(t) + c*(1 - α)*s*w(t)]
-#    costs = [-q(tspan[2])]
-#   
-#    @mtkbuild beesys = ODESystem(eqs, t; costs)
-#    u0map = [w(0) => 40, q(0) => 2]
-#    pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1]
-#
-#    jprob = JuMPControlProblem(beesys, u0map, tspan, pmap)
-#    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
-#    control_sol = jsol.control_sol
-#    # Bang-bang control
-#end
+@testset "Linear systems" begin
+    function is_bangbang(input_sol, lbounds, ubounds)
+        bangbang = true
+        for v in 1:length(input_sol.u[1])
+            all(i -> i[v] ≈ bounds[v] || i[v] ≈ bounds[u], input_sol.u) || (bangbang = false)
+        end
+        bangbang
+    end
+
+    # Double integrator
+    @variables x(..) [bounds = (0., 0.25)] v(..)
+    @variables u(t) [bounds = (-1., 1.), input = true]
+    constr = [v(1.0) ~ 0.0]
+    cost = [-x(1.0)] # Optimize the final distance.
+    @named block = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u], t)
+    block, input_idxs = structural_simplify(block, ([u],[]))
+
+    u0map = [x(t) => 0., v(t) => 0.]
+    tspan = (0., 1.)
+    parammap = [u => 0.]
+    jprob = JuMPControlProblem(block, u0map, tspan, parammap; dt = 0.01)
+    jsol = solve(jprob, Ipopt.Optimizer, :Verner8)
+    # Linear systems have bang-bang controls
+    @test is_bangbang(jsol.input_sol, [-1.], [1.])
+    # Test reached final position.
+    @test jsol.sol.u[end][1] ≈ 0.25
+
+    # Cart-pole system
+
+    # Bee example (from Lawrence Evans' notes)
+    M.@variables w(..) q(..) 
+    M.@parameters α(t) [bounds = [0, 1]] b c μ s ν
+
+    tspan = (0, 4)
+    eqs = [D(w(t)) ~ -μ*w(t) + b*s*α*w(t),
+           D(q(t)) ~ -ν*q(t) + c*(1 - α)*s*w(t)]
+    costs = [-q(tspan[2])]
+   
+    @mtkbuild beesys = ODESystem(eqs, t; costs)
+    u0map = [w(0) => 40, q(0) => 2]
+    pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1]
+
+    jprob = JuMPControlProblem(beesys, u0map, tspan, pmap)
+    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    control_sol = jsol.control_sol
+    # Bang-bang control
+end
 #
 #@testset "Constrained optimal control problems" begin
 #end

From 3a1e1de000595ba9fb88471fc57e57deb5879271 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Wed, 16 Apr 2025 18:29:34 -0400
Subject: [PATCH 21/41] implement ControlFunction

---
 ext/MTKJuMPControlExt.jl                 |  16 +--
 src/inputoutput.jl                       |   8 +-
 src/systems/diffeqs/abstractodesystem.jl |   2 +-
 src/systems/optimal_control_interface.jl | 126 +++++++++++++++++++++++
 test/extensions/jump_control.jl          |   6 +-
 5 files changed, 143 insertions(+), 15 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index a178280428..701443785c 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -15,7 +15,7 @@ struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <:
     kwargs::K
 
     function JuMPControlProblem(f, u0, tspan, p, model; kwargs...)
-        new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f),
+        new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f, 5),
             typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
     end
 end
@@ -55,13 +55,10 @@ function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap;
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
-    @show _u0map
-    f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
+    f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    (f_i, f_o) = generate_control_function(sys)
     model = init_model(sys, tspan[1]:dt:tspan[2], u0map, u0)
-
     JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
 end
 
@@ -80,7 +77,7 @@ function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
-    f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
+    f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
     model = init_model(sys, tspan[1]:dt:tspan[2], u0map, u0)
@@ -237,14 +234,17 @@ function add_jump_solve_constraints!(prob, tableau)
     dt = tsteps[2] - tsteps[1]
 
     U = model[:U]
+    V = model[:V]
     nᵤ = length(U)
+    nᵥ = length(V)
     if is_explicit(tableau)
         K = Any[]
         for τ in tsteps
             for (i, h) in enumerate(c)
                 ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = zeros(nᵤ))
                 Uₙ = [U[i](τ) + ΔU[i] * dt for i in 1:nᵤ]
-                Kₙ = f(Uₙ, p, τ + h * dt)
+                Vₙ = [V[i](τ) for i in 1:nᵥ]
+                Kₙ = f(Uₙ, Vₙ, p, τ + h * dt)
                 push!(K, Kₙ)
             end
             ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)])
@@ -259,7 +259,7 @@ function add_jump_solve_constraints!(prob, tableau)
             for (i, h) in enumerate(c)
                 ΔU = ΔUs[i, :]
                 Uₙ = [U[j] + ΔU[j] * dt for j in 1:nᵤ]
-                @constraint(model, [j in 1:nᵤ], K[i, j]==f(Uₙ, p, τ + h * dt)[j],
+                @constraint(model, [j in 1:nᵤ], K[i, j]==f(Uₙ, V, p, τ + h * dt)[j],
                     DomainRestrictions(t => τ), base_name="solve_K($τ)")
             end
             ΔU = dt * sum([α[i] * K[i, :] for i in 1:length(α)])
diff --git a/src/inputoutput.jl b/src/inputoutput.jl
index 5f9420ff3a..329a76d9c4 100644
--- a/src/inputoutput.jl
+++ b/src/inputoutput.jl
@@ -208,7 +208,9 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
         inputs = [inputs; disturbance_inputs]
     end
 
-    sys, _ = io_preprocessing(sys, inputs, []; simplify, kwargs...)
+    if !iscomplete(sys)
+        sys, _ = io_preprocessing(sys, inputs, []; simplify, kwargs...)
+    end
 
     dvs = unknowns(sys)
     ps = parameters(sys; initial_parameters = true)
@@ -250,9 +252,7 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
     f = build_function_wrapper(sys, rhss, args...; p_start = 3 + implicit_dae,
         p_end = length(p) + 2 + implicit_dae)
     f = eval_or_rgf.(f; eval_expression, eval_module)
-    f = GeneratedFunctionWrapper{(
-        3 + implicit_dae, length(args) - length(p) + 1, is_split(sys))}(f...)
-    f = f, f
+    f = GeneratedFunctionWrapper{(3, length(args) - length(p) + 1, is_split(sys))}(f...)
     ps = setdiff(parameters(sys), inputs, disturbance_inputs)
     (; f, dvs, ps, io_sys = sys)
 end
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index d9e8b05eeb..e66f03a85e 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -101,7 +101,7 @@ function calculate_control_jacobian(sys::AbstractODESystem;
     end
 
     rhs = [eq.rhs for eq in full_equations(sys)]
-    ctrls = controls(sys)
+    ctrls = unbound_inputs(sys)
 
     if sparse
         jac = sparsejacobian(rhs, ctrls, simplify = simplify)
diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index f8e9c75887..fd7adfe643 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -20,8 +20,134 @@ function warn_overdetermined(sys, u0map)
     end
 end
 
+"""
+Generate the control function f(x, u, p, t) from the ODESystem. 
+Input variables are automatically inferred but can be manually specified.
+"""
+function SciMLBase.ControlFunction{iip, specialize}(sys::ODESystem,
+        dvs = unknowns(sys),
+        ps = parameters(sys), u0 = nothing,
+        inputs = unbound_inputs(sys),
+        disturbance_inputs = disturbances(sys);
+        version = nothing, tgrad = false,
+        jac = false, controljac = false, 
+        p = nothing, t = nothing,
+        eval_expression = false,
+        sparse = false, simplify = false,
+        eval_module = @__MODULE__,
+        steady_state = false,
+        checkbounds = false,
+        sparsity = false,
+        analytic = nothing,
+        split_idxs = nothing,
+        initialization_data = nothing,
+        cse = true,
+        kwargs...) where {iip, specialize}
+
+    (f), _, _ = generate_control_function(sys, inputs, disturbance_inputs; eval_expression = true, eval_module, cse, kwargs...)
+
+    if tgrad
+        tgrad_gen = generate_tgrad(sys, dvs, ps;
+            simplify = simplify,
+            expression = Val{true},
+            expression_module = eval_module, cse,
+            checkbounds = checkbounds, kwargs...)
+        tgrad_oop, tgrad_iip = eval_or_rgf.(tgrad_gen; eval_expression, eval_module)
+        _tgrad = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(tgrad_oop, tgrad_iip)
+    else
+        _tgrad = nothing
+    end
+
+    if jac
+        jac_gen = generate_jacobian(sys, dvs, ps;
+            simplify = simplify, sparse = sparse,
+            expression = Val{true},
+            expression_module = eval_module, cse,
+            checkbounds = checkbounds, kwargs...)
+        jac_oop, jac_iip = eval_or_rgf.(jac_gen; eval_expression, eval_module)
+
+        _jac = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(jac_oop, jac_iip)
+    else
+        _jac = nothing
+    end
+
+    if controljac
+        cjac_gen = generate_control_jacobian(sys, dvs, ps;
+            simplify = simplify, sparse = sparse,
+            expression = Val{true},
+            expression_module = eval_module, cse,
+            checkbounds = checkbounds, kwargs...)
+        cjac_oop, cjac_iip = eval_or_rgf.(cjac_gen; eval_expression, eval_module)
+
+        _cjac = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(cjac_oop, cjac_iip)
+    else
+        _cjac = nothing
+    end
+
+    M = calculate_massmatrix(sys)
+    _M = if sparse && !(u0 === nothing || M === I)
+        SparseArrays.sparse(M)
+    elseif u0 === nothing || M === I
+        M
+    else
+        ArrayInterface.restructure(u0 .* u0', M)
+    end
+
+    observedfun = ObservedFunctionCache(
+        sys; steady_state, eval_expression, eval_module, checkbounds, cse)
+
+    if sparse
+        uElType = u0 === nothing ? Float64 : eltype(u0)
+        W_prototype = similar(W_sparsity(sys), uElType)
+        controljac_prototype = similar(calculate_control_jacobian(sys), uElType)
+    else
+        W_prototype = nothing
+        controljac_prototype = nothing
+    end
+    
+    ControlFunction{iip, specialize}(f;
+        sys = sys,
+        jac = _jac === nothing ? nothing : _jac,
+        controljac = _cjac === nothing ? nothing : _cjac,
+        tgrad = _tgrad === nothing ? nothing : _tgrad,
+        mass_matrix = _M,
+        jac_prototype = W_prototype,
+        controljac_prototype = controljac_prototype,
+        observed = observedfun,
+        sparsity = sparsity ? W_sparsity(sys) : nothing,
+        analytic = analytic,
+        initialization_data)
+end
+
+function SciMLBase.ControlFunction(sys::AbstractODESystem, args...; kwargs...)
+    ControlFunction{true}(sys, args...; kwargs...)
+end
+
+function SciMLBase.ControlFunction{true}(sys::AbstractODESystem, args...;
+        kwargs...)
+    ControlFunction{true, SciMLBase.AutoSpecialize}(sys, args...; kwargs...)
+end
+
+function SciMLBase.ControlFunction{false}(sys::AbstractODESystem, args...;
+        kwargs...)
+    ControlFunction{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
+end
+
 """
 IntegralNorm. When applied to an expression.
 """
 struct IntegralNorm end
 
+"""
+$(SIGNATURES)
+
+Define one or more inputs.
+
+See also [`@independent_variables`](@ref), [`@variables`](@ref) and [`@constants`](@ref).
+"""
+macro inputs(xs...)
+    Symbolics._parse_vars(:inputs,
+        Real,
+        xs,
+        toparam) |> esc
+end
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index ee9ccd3a07..5ae1e98f11 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -86,6 +86,8 @@ end
     end
 
     # Double integrator
+    t = M.t_nounits
+    D = M.D_nounits
     @variables x(..) [bounds = (0., 0.25)] v(..)
     @variables u(t) [bounds = (-1., 1.), input = true]
     constr = [v(1.0) ~ 0.0]
@@ -106,8 +108,8 @@ end
     # Cart-pole system
 
     # Bee example (from Lawrence Evans' notes)
-    M.@variables w(..) q(..) 
-    M.@parameters α(t) [bounds = [0, 1]] b c μ s ν
+    @variables w(..) q(..) 
+    @parameters α(t) [bounds = [0, 1]] b c μ s ν
 
     tspan = (0, 4)
     eqs = [D(w(t)) ~ -μ*w(t) + b*s*α*w(t),

From eae1060a55f1047d7ac2404db2b0c90c79cbaad2 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Thu, 17 Apr 2025 17:51:03 -0400
Subject: [PATCH 22/41] feat: working linear control problems

---
 ext/MTKJuMPControlExt.jl                 | 112 ++++++++++-------------
 src/systems/optimal_control_interface.jl |  10 +-
 src/systems/problem_utils.jl             |   1 -
 test/extensions/jump_control.jl          |  91 ++++++++++++++----
 4 files changed, 130 insertions(+), 84 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 701443785c..054f0ecb10 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -58,7 +58,7 @@ function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap;
     f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, u0)
+    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, pmap, u0)
     JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
 end
 
@@ -80,22 +80,23 @@ function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
     f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, u0)
+    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, pmap, u0)
     add_infopt_solve_constraints!(model, sys, pmap)
     InfiniteOptControlProblem(f, u0, tspan, p, model, kwargs...)
 end
 
-function init_model(sys, tsteps, u0map, u0)
+function init_model(sys, tsteps, u0map, pmap, u0)
     ctrls = MTK.unbound_inputs(sys)
     states = unknowns(sys)
     model = InfiniteModel()
+
     @infinite_parameter(model, t in [tsteps[1], tsteps[end]], num_supports=length(tsteps))
     @variable(model, U[i = 1:length(states)], Infinite(t))
     @variable(model, V[1:length(ctrls)], Infinite(t))
 
     set_bounds!(model, sys)
-    add_jump_cost_function!(model, sys)
-    add_user_constraints!(model, sys)
+    add_jump_cost_function!(model, sys, (tsteps[1], tsteps[2]), pmap)
+    add_user_constraints!(model, sys, pmap)
 
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
     u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) :
@@ -120,63 +121,35 @@ function set_bounds!(model, sys)
     end
 end
 
-function add_jump_cost_function!(model::InfiniteModel, sys)
+function add_jump_cost_function!(model::InfiniteModel, sys, tspan, pmap)
     jcosts = MTK.get_costs(sys)
     consolidate = MTK.get_consolidate(sys)
     if isnothing(jcosts) || isempty(jcosts)
         @objective(model, Min, 0)
         return
     end
-    iv = MTK.get_iv(sys)
-
-    stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
-    cidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))])
-
-    for st in unknowns(sys)
-        x = operation(st)
-        t = only(arguments(st))
-        idx = stidxmap[x(iv)]
-        subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
-        jcosts = map(c -> Symbolics.substitute(c, Dict(x(t) => subval)), jcosts)
-    end
+    jcosts = substitute_jump_vars(model, sys, pmap, jcosts)
 
-    for ct in MTK.unbound_inputs(sys)
-        p = operation(ct)
-        t = only(arguments(ct))
-        idx = cidxmap[p(iv)]
-        subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
-        jcosts = map(c -> Symbolics.substitute(c, Dict(p(t) => subval)), jcosts)
+    # Substitute integral
+    iv = MTK.get_iv(sys)
+    jcosts = map(c -> Symbolics.substitute(c, ∫ => Symbolics.Integral(iv in tspan)), jcosts)
+    intmap = Dict()
+    
+    for int in MTK.collect_applied_operators(jcosts, Symbolics.Integral)
+        arg = only(arguments(MTK.value(int)))
+        lower_bound, upper_bound = (int.domain.domain.left, int.domain.domain.right)
+        intmap[int] = InfiniteOpt.∫(arg, iv; lower_bound, upper_bound)
     end
-
+    jcosts = map(c -> Symbolics.substitute(c, intmap), jcosts)
     @objective(model, Min, consolidate(jcosts))
 end
 
-function add_user_constraints!(model::InfiniteModel, sys)
+function add_user_constraints!(model::InfiniteModel, sys, pmap)
     conssys = MTK.get_constraintsystem(sys)
     jconstraints = isnothing(conssys) ? nothing : MTK.get_constraints(conssys)
     (isnothing(jconstraints) || isempty(jconstraints)) && return nothing
 
-    iv = MTK.get_iv(sys)
-    stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
-    cidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))])
-
-    for st in unknowns(conssys)
-        x = operation(st)
-        t = only(arguments(st))
-        idx = stidxmap[x(iv)]
-        subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
-        jconstraints = map(c -> Symbolics.substitute(c, Dict(x(t) => subval)), jconstraints)
-    end
-
-    for ct in MTK.unbound_inputs(sys)
-        p = operation(ct)
-        t = only(arguments(ct))
-        idx = cidxmap[p(iv)]
-        subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
-        jconstraints = map(
-            c -> Symbolics.substitute(jconstraints, Dict(p(t) => subval)), jconstriants)
-    end
-
+    jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints)
     for (i, cons) in enumerate(jconstraints)
         if cons isa Equation
             @constraint(model, cons.lhs - cons.rhs==0, base_name="user[$i]")
@@ -193,31 +166,41 @@ function add_initial_constraints!(model::InfiniteModel, u0, u0_idxs, ts)
     @constraint(model, initial[i in u0_idxs], U[i](ts)==u0[i])
 end
 
-is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
-
-function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap)
+function substitute_jump_vars(model, sys, pmap, exprs)
     iv = MTK.get_iv(sys)
-    t = model[:t]
+    sts = unknowns(sys)
+    cts = MTK.unbound_inputs(sys)
     U = model[:U]
     V = model[:V]
+    # 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.substitute(c, whole_interval_map), exprs)
+
+    # for variables like x(1.0)
+    x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
+    c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
+    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.substitute(c, fixed_t_map), exprs)
+
+    exprs = map(c -> Symbolics.substitute(c, Dict(pmap)), exprs)
+    exprs
+end
 
-    stmap = Dict([v => U[i] for (i, v) in enumerate(unknowns(sys))])
-    ctrlmap = Dict([v => V[i] for (i, v) in enumerate(MTK.unbound_inputs(sys))])
-    submap = merge(stmap, ctrlmap, Dict(pmap))
+is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
 
+function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap)
     # Differential equations
-    diff_eqs = diff_equations(sys)
-    D = Differential(iv)
+    U = model[:U]
+    t = model[:t]
+    D = Differential(MTK.get_iv(sys))
     diffsubmap = Dict([D(U[i]) => ∂(U[i], t) for i in 1:length(U)])
-    for u in unknowns(sys)
-        diff_eqs = map(e -> Symbolics.substitute(e, submap), diff_eqs)
-        diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs)
-    end
+
+    diff_eqs = substitute_jump_vars(model, sys, pmap, diff_equations(sys))
+    diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs)
     @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs==diff_eqs[i].rhs)
 
     # Algebraic equations
-    alg_eqs = alg_equations(sys)
-    alg_eqs = map(e -> Symbolics.substitute(e, submap), alg_eqs)
+    alg_eqs = substitute_jump_vars(model, sys, pmap, alg_equations(sys))
     @constraint(model, A[i = 1:length(alg_eqs)], alg_eqs[i].lhs==alg_eqs[i].rhs)
 end
 
@@ -306,9 +289,10 @@ end
 `derivative_method` kwarg refers to the method used by InfiniteOpt to compute derivatives. The list of possible options can be found at https://infiniteopt.github.io/InfiniteOpt.jl/stable/guide/derivative/. Defaults to FiniteDifference(Backward()).
 """
 function DiffEqBase.solve(prob::InfiniteOptControlProblem, jump_solver;
-        derivative_method = InfiniteOpt.FiniteDifference(Backward()))
+        derivative_method = InfiniteOpt.FiniteDifference(Backward()), silent = false)
+    model = prob.model
     silent && set_silent(model)
-    set_derivative_method(prob.model[:t], derivative_method)
+    set_derivative_method(model[:t], derivative_method)
     _solve(prob, jump_solver, derivative_method)
 end
 
diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index fd7adfe643..d6cc7de320 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -134,9 +134,15 @@ function SciMLBase.ControlFunction{false}(sys::AbstractODESystem, args...;
 end
 
 """
-IntegralNorm. When applied to an expression.
+IntegralNorm. When applied to an expression in a cost
+function, assumes that the integration variable is the
+iv of the system, and assumes that the bounds are the
+tspan.
+Equivalent to Integral(t in tspan) in Symbolics.
 """
-struct IntegralNorm end
+struct ∫ <: Symbolics.Operator end
+∫(x) = ∫()(x)
+Base.show(io::IO, x::∫) = print(io, "∫")
 
 """
 $(SIGNATURES)
diff --git a/src/systems/problem_utils.jl b/src/systems/problem_utils.jl
index 94a96b8959..41c64b78c5 100644
--- a/src/systems/problem_utils.jl
+++ b/src/systems/problem_utils.jl
@@ -822,7 +822,6 @@ function maybe_build_initialization_problem(
         t = zero(floatT)
     end
 
-    @show u0map
     initializeprob = ModelingToolkit.InitializationProblem{true, SciMLBase.FullSpecialize}(
         sys, t, u0map, pmap; guesses, initialization_eqs, use_scc, kwargs...)
     if state_values(initializeprob) !== nothing
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index 5ae1e98f11..47e962ab9a 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -77,10 +77,10 @@ const M = ModelingToolkit
 end
 
 @testset "Linear systems" begin
-    function is_bangbang(input_sol, lbounds, ubounds)
+    function is_bangbang(input_sol, lbounds, ubounds, rtol = 1e-4)
         bangbang = true
-        for v in 1:length(input_sol.u[1])
-            all(i -> i[v] ≈ bounds[v] || i[v] ≈ bounds[u], input_sol.u) || (bangbang = false)
+        for v in 1:length(input_sol.u[1]) - 1
+            all(i -> ≈(i[v], bounds[v]; rtol) || ≈(i[v], bounds[u]; rtol), input_sol.u) || (bangbang = false)
         end
         bangbang
     end
@@ -91,8 +91,8 @@ end
     @variables x(..) [bounds = (0., 0.25)] v(..)
     @variables u(t) [bounds = (-1., 1.), input = true]
     constr = [v(1.0) ~ 0.0]
-    cost = [-x(1.0)] # Optimize the final distance.
-    @named block = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u], t)
+    cost = [-x(1.0)] # Maximize the final distance.
+    @named block = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u], t; costs = cost, constraints = constr)
     block, input_idxs = structural_simplify(block, ([u],[]))
 
     u0map = [x(t) => 0., v(t) => 0.]
@@ -103,28 +103,85 @@ end
     # Linear systems have bang-bang controls
     @test is_bangbang(jsol.input_sol, [-1.], [1.])
     # Test reached final position.
-    @test jsol.sol.u[end][1] ≈ 0.25
+    @test ≈(jsol.sol.u[end][1], 0.25, rtol = 1e-5)
 
-    # Cart-pole system
+    iprob = InfiniteOptControlProblem(block, u0map, tspan, parammap; dt = 0.01)
+    isol = solve(iprob, Ipopt.Optimizer; silent = true)
+    @test is_bangbang(isol.input_sol, [-1.], [1.])
+    @test ≈(isol.sol.u[end][1], 0.25, rtol = 1e-5)
 
-    # Bee example (from Lawrence Evans' notes)
-    @variables w(..) q(..) 
-    @parameters α(t) [bounds = [0, 1]] b c μ s ν
+    ###################
+    ### Bee example ###
+    ###################
+    # From Lawrence Evans' notes
+    @variables w(..) q(..) α(t) [input = true, bounds = (0, 1)]
+    @parameters b c μ s ν
 
     tspan = (0, 4)
     eqs = [D(w(t)) ~ -μ*w(t) + b*s*α*w(t),
            D(q(t)) ~ -ν*q(t) + c*(1 - α)*s*w(t)]
     costs = [-q(tspan[2])]
    
-    @mtkbuild beesys = ODESystem(eqs, t; costs)
-    u0map = [w(0) => 40, q(0) => 2]
-    pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1]
+    @named beesys = ODESystem(eqs, t; costs)
+    beesys, input_idxs = structural_simplify(beesys, ([α],[]))
+    u0map = [w(t) => 40, q(t) => 2]
+    pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1, α => 1]
 
-    jprob = JuMPControlProblem(beesys, u0map, tspan, pmap)
+    jprob = JuMPControlProblem(beesys, u0map, tspan, pmap, dt = 0.01)
     jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
-    control_sol = jsol.control_sol
-    # Bang-bang control
+    @test is_bangbang(jsol.input_sol, [0.], [1.])
+    iprob = InfiniteOptControlProblem(beesys, u0map, tspan, pmap, dt = 0.01)
+    isol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    @test is_bangbang(isol.input_sol, [0.], [1.])
 end
-#
+
+@testset "Rocket launch" begin
+    t = M.t_nounits
+    D = M.D_nounits
+
+    @variables h(..) v(..) m(..) T(..) [input = true, bounds = (0, tₘ)]
+    @parameters h_c m₀ h₀ g₀ D_c c Tₘ
+    @parameters tf
+    drag(h, v) = D_c * v^2 * exp(-h_c * (h - h₀) / h₀)
+    gravity(h) = g₀ * (h₀ / h)
+
+    eqs = [D(h(t)) ~ v(t), 
+           D(v(t)) ~ (T(t) - drag(h(t), v(t))) / m(t) - gravity(t),
+           D(m(t)) ~ -T(t) / c]
+
+    costs = [-h(tf)]
+    constraints = [T(tf) ~ 0]
+    @named rocket = ODESystem(eqs, t; costs, constraints)
+    @test tf ∈ Set(parameters(rocket))
+
+    u0map = [h(t) => h₀, m(t) => m₀, v(t) => 0]
+    pmap = [g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5*√(g₀*h₀), D_C => 0.5 * 620 * m₀/g₀, Tₘ => 3.5*g₀*m₀]
+    jprob = JuMPControlProblem(rocket, u0map, (0, tf), pmap)
+    jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3)
+    @test jsol.sol.u[end][1] ≈ 1.012
+end
+
+@testset "Free final time problem" begin
+    t = M.t_nounits
+    D = M.D_nounits
+
+    @variables x(..) u(..) [input = true, bounds = (0,1)]
+    @parameters tf
+    eqs = [D(x(t)) ~ -2 + 0.5*u]
+
+    # Integral cost function
+    costs = [∫(x-u), x(tf)]
+    consolidate(u) = u[1] + u[2]
+    jprob = JuMPControlProblem(rocket, u0map, (0, tf), pmap)
+    jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3)
+    @test jsol.sol.t[end] ≈ 10.0
+    iprob = InfiniteOptControlProblem(rocket, u0map, (0, tf), pmap)
+    isol = solve(iprob, Ipopt.Optimizer, :RadauIA3)
+    @test isol.sol.t[end] ≈ 10.0
+end
+
+@testset "Cart-pole problem" begin
+end
+
 #@testset "Constrained optimal control problems" begin
 #end

From 4adad1b01f65b9448eb8c844cdde9d347346a257 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 18 Apr 2025 05:23:42 -0400
Subject: [PATCH 23/41] feat: free final time problems

---
 ext/MTKJuMPControlExt.jl                 | 127 +++++++++++++++--------
 src/ModelingToolkit.jl                   |   1 +
 src/inputoutput.jl                       |   2 +-
 src/systems/diffeqs/odesystem.jl         |   2 +-
 src/systems/optimal_control_interface.jl |  62 ++++++++---
 src/variables.jl                         |   5 +
 test/extensions/jump_control.jl          |  44 ++++----
 7 files changed, 165 insertions(+), 78 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 054f0ecb10..310a632292 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -3,6 +3,7 @@ using ModelingToolkit
 using JuMP, InfiniteOpt
 using DiffEqDevTools, DiffEqBase
 using LinearAlgebra
+using StaticArrays
 const MTK = ModelingToolkit
 
 struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <:
@@ -14,7 +15,7 @@ struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <:
     model::InfiniteModel
     kwargs::K
 
-    function JuMPControlProblem(f, u0, tspan, p, model; kwargs...)
+    function JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
         new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f, 5),
             typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
     end
@@ -51,14 +52,18 @@ The constraints are:
 - The solver constraints that encode the time-stepping used by the solver
 """
 function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap;
-        dt = error("dt must be provided for JuMPControlProblem."),
+        dt = nothing,
+        steps = nothing,
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
     f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, pmap, u0)
+    pmap = MTK.todict(pmap)
+    steps, is_free_t = MTK.process_tspan(tspan, dt, steps)
+    model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
+
     JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
 end
 
@@ -73,55 +78,78 @@ Related to `JuMPControlProblem`, but directly adds the differential equations
 of the system as derivative constraints, rather than using a solver tableau.
 """
 function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
-        dt = error("dt must be provided for InfiniteOptControlProblem."),
+        dt = nothing,
+        steps = nothing,
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
     f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, pmap, u0)
-    add_infopt_solve_constraints!(model, sys, pmap)
+    pmap = MTK.todict(pmap)
+    steps, is_free_t = MTK.process_tspan(tspan, dt, steps)
+    model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
+
+    add_infopt_solve_constraints!(model, sys, pmap; is_free_t)
     InfiniteOptControlProblem(f, u0, tspan, p, model, kwargs...)
 end
 
-function init_model(sys, tsteps, u0map, pmap, u0)
+# Initialize InfiniteOpt model.
+function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
     ctrls = MTK.unbound_inputs(sys)
     states = unknowns(sys)
     model = InfiniteModel()
 
-    @infinite_parameter(model, t in [tsteps[1], tsteps[end]], num_supports=length(tsteps))
-    @variable(model, U[i = 1:length(states)], Infinite(t))
-    @variable(model, V[1:length(ctrls)], Infinite(t))
+    if is_free_t
+        (ts_sym, te_sym) = tspan
+        @variable(model, tf, start = pmap[te_sym])
+        hasbounds(te_sym) && begin
+            lo, hi = getbounds(te_sym)
+            set_lower_bound(tf, lo)
+            set_upper_bound(tf, hi)
+        end
+        pmap[ts_sym] = 0
+        pmap[te_sym] = 1
+        tspan = (0, 1)
+    end
+        
+    @infinite_parameter(model, t in [tspan[1], tspan[2]], num_supports = steps)
+    @variable(model, U[i = 1:length(states)], Infinite(t), start = u0[i])
+    c0 = [pmap[c] for c in ctrls]
+    @variable(model, V[i = 1:length(ctrls)], Infinite(t), start = c0[i])
 
-    set_bounds!(model, sys)
-    add_jump_cost_function!(model, sys, (tsteps[1], tsteps[2]), pmap)
-    add_user_constraints!(model, sys, pmap)
+    set_jump_bounds!(model, sys, pmap)
+    add_jump_cost_function!(model, sys, (tspan[1], tspan[2]), pmap; is_free_t)
+    add_user_constraints!(model, sys, pmap; is_free_t)
 
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
     u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) :
               [stidxmap[k] for (k, v) in u0map]
-    add_initial_constraints!(model, u0, u0_idxs, tsteps[1])
+    add_initial_constraints!(model, u0, u0_idxs, tspan[1])
     return model
 end
 
-function set_bounds!(model, sys)
+function set_jump_bounds!(model, sys, pmap)
     U = model[:U]
     for (i, u) in enumerate(unknowns(sys))
-        lo, hi = MTK.getbounds(u)
-        set_lower_bound(U[i], lo)
-        set_upper_bound(U[i], hi)
+        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))
+        end
     end
 
     V = model[:V]
     for (i, v) in enumerate(MTK.unbound_inputs(sys))
-        lo, hi = MTK.getbounds(v)
-        set_lower_bound(V[i], lo)
-        set_upper_bound(V[i], hi)
+        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))
+        end
     end
 end
 
-function add_jump_cost_function!(model::InfiniteModel, sys, tspan, pmap)
+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)
@@ -129,26 +157,36 @@ function add_jump_cost_function!(model::InfiniteModel, sys, tspan, pmap)
         return
     end
     jcosts = substitute_jump_vars(model, sys, pmap, jcosts)
+    tₛ = is_free_t ? model[:tf] : 1
 
     # Substitute integral
     iv = MTK.get_iv(sys)
-    jcosts = map(c -> Symbolics.substitute(c, ∫ => Symbolics.Integral(iv in tspan)), jcosts)
+    jcosts = map(c -> Symbolics.substitute(c, MTK.∫() => Symbolics.Integral(iv in tspan)), jcosts)
+
     intmap = Dict()
-    
     for int in MTK.collect_applied_operators(jcosts, Symbolics.Integral)
+        op = MTK.operation(int)
         arg = only(arguments(MTK.value(int)))
-        lower_bound, upper_bound = (int.domain.domain.left, int.domain.domain.right)
-        intmap[int] = InfiniteOpt.∫(arg, iv; lower_bound, upper_bound)
+        lo, hi = (op.domain.domain.left, op.domain.domain.right)
+        intmap[int] = tₛ * InfiniteOpt.∫(arg, model[:t], lo, hi)
     end
     jcosts = map(c -> Symbolics.substitute(c, intmap), jcosts)
     @objective(model, Min, consolidate(jcosts))
 end
 
-function add_user_constraints!(model::InfiniteModel, sys, pmap)
+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)
     (isnothing(jconstraints) || isempty(jconstraints)) && return nothing
 
+    if is_free_t
+        for u in MTK.get_unknowns(conssys)
+            x = MTK.operation(u)
+            t = only(arguments(u))
+            MTK.symbolic_type(t) === NotSymbolic() && error("Provided specific time constraint in a free final time problem. This is not supported by the JuMP/InfiniteOpt collocation solvers. The offending variable is $u.")
+        end
+    end
+
     jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints)
     for (i, cons) in enumerate(jconstraints)
         if cons isa Equation
@@ -188,23 +226,24 @@ end
 
 is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
 
-function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap)
+function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap; is_free_t = false)
     # Differential equations
     U = model[:U]
     t = model[:t]
     D = Differential(MTK.get_iv(sys))
     diffsubmap = Dict([D(U[i]) => ∂(U[i], t) for i in 1:length(U)])
+    tₛ = is_free_t ? model[:tf] : 1
 
     diff_eqs = substitute_jump_vars(model, sys, pmap, diff_equations(sys))
     diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs)
-    @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs==diff_eqs[i].rhs)
+    @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs == tₛ * diff_eqs[i].rhs)
 
     # Algebraic equations
     alg_eqs = substitute_jump_vars(model, sys, pmap, alg_equations(sys))
-    @constraint(model, A[i = 1:length(alg_eqs)], alg_eqs[i].lhs==alg_eqs[i].rhs)
+    @constraint(model, A[i = 1:length(alg_eqs)], alg_eqs[i].lhs == alg_eqs[i].rhs)
 end
 
-function add_jump_solve_constraints!(prob, tableau)
+function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
     A = tableau.A
     α = tableau.α
     c = tableau.c
@@ -214,6 +253,7 @@ function add_jump_solve_constraints!(prob, tableau)
     t = model[:t]
     tsteps = supports(model[:t])
     pop!(tsteps)
+    tₛ = is_free_t ? model[:tf] : 1
     dt = tsteps[2] - tsteps[1]
 
     U = model[:U]
@@ -227,7 +267,7 @@ function add_jump_solve_constraints!(prob, tableau)
                 ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = zeros(nᵤ))
                 Uₙ = [U[i](τ) + ΔU[i] * dt for i in 1:nᵤ]
                 Vₙ = [V[i](τ) for i in 1:nᵥ]
-                Kₙ = f(Uₙ, Vₙ, p, τ + h * dt)
+                Kₙ = tₛ * f(Uₙ, Vₙ, p, τ + h * dt) # scale the time
                 push!(K, Kₙ)
             end
             ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)])
@@ -237,17 +277,17 @@ function add_jump_solve_constraints!(prob, tableau)
         end
     else
         @variable(model, K[1:length(α), 1:nᵤ], Infinite(t), start=tsteps[1])
+        ΔUs = A * K
+        ΔU_tot = dt * (K' * α)
         for τ in tsteps
-            ΔUs = A * K
             for (i, h) in enumerate(c)
-                ΔU = ΔUs[i, :]
-                Uₙ = [U[j] + ΔU[j] * dt for j in 1:nᵤ]
-                @constraint(model, [j in 1:nᵤ], K[i, j]==f(Uₙ, V, p, τ + h * dt)[j],
-                    DomainRestrictions(t => τ), base_name="solve_K($τ)")
+                ΔU = @view ΔUs[i, :]
+                Uₙ = U + ΔU * dt
+                @constraint(model, [j = 1:nᵤ], K[i, j](τ) == tₛ * f(Uₙ, V, p, τ + h * dt)[j],
+                            DomainRestrictions(t => τ + h*dt), base_name="solve_K($τ)")
             end
-            ΔU = dt * sum([α[i] * K[i, :] for i in 1:length(α)])
-            @constraint(model, [n = 1:nᵤ], U[n] + ΔU[n]==U[n](τ + dt),
-                DomainRestrictions(t => τ), base_name="solve_U($τ)")
+            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n] == U[n](τ + dt),
+                        DomainRestrictions(t => τ), base_name="solve_U($τ)")
         end
     end
 end
@@ -281,7 +321,7 @@ function DiffEqBase.solve(
             delete(model, var)
         end
     end
-    add_jump_solve_constraints!(prob, tableau)
+    add_jump_solve_constraints!(prob, tableau; is_free_t = haskey(model, :tf))
     _solve(prob, jump_solver, ode_solver)
 end
 
@@ -304,9 +344,10 @@ function _solve(prob::AbstractOptimalControlProblem, jump_solver, solver)
     tstatus = termination_status(model)
     pstatus = primal_status(model)
     !has_values(model) &&
-        error("Model not solvable; please report this to github.com/SciML/ModelingToolkit.jl.")
+        error("Model not solvable; please report this to github.com/SciML/ModelingToolkit.jl with a MWE.")
 
-    ts = supports(model[:t])
+    tf = haskey(model, :tf) ? value(model[:tf]) : 1
+    ts = tf * supports(model[:t])
     U_vals = value.(model[:U])
     U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:length(ts)]
     sol = DiffEqBase.build_solution(prob, solver, ts, U_vals)
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index 0a2ea31181..bb6d991578 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -351,5 +351,6 @@ include("systems/optimal_control_interface.jl")
 export AbstractOptimalControlProblem, JuMPControlProblem, InfiniteOptControlProblem,
        PyomoControlProblem, CasADiControlProblem
 export OptimalControlSolution
+export ∫
 
 end # module
diff --git a/src/inputoutput.jl b/src/inputoutput.jl
index 329a76d9c4..11a58bd30a 100644
--- a/src/inputoutput.jl
+++ b/src/inputoutput.jl
@@ -250,7 +250,7 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
         args = (ddvs, args...)
     end
     f = build_function_wrapper(sys, rhss, args...; p_start = 3 + implicit_dae,
-        p_end = length(p) + 2 + implicit_dae)
+        p_end = length(p) + 2 + implicit_dae, kwargs...)
     f = eval_or_rgf.(f; eval_expression, eval_module)
     f = GeneratedFunctionWrapper{(3, length(args) - length(p) + 1, is_split(sys))}(f...)
     ps = setdiff(parameters(sys), inputs, disturbance_inputs)
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index 7e3f11b58d..ceef916e90 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -819,7 +819,7 @@ function validate_vars_and_find_ps!(auxvars, auxps, sysvars, iv)
                 arg isa AbstractFloat ||
                 throw(ArgumentError("Invalid argument specified for variable $var. The argument of the variable should be either $iv, a parameter, or a value specifying the time that the constraint holds."))
 
-            isparameter(arg) && push!(auxps, arg)
+            (isparameter(arg) && !isequal(arg, iv)) && push!(auxps, arg)
         else
             var ∈ sts &&
                 @warn "Variable $var has no argument. It will be interpreted as $var($iv), and the constraint will apply to the entire interval."
diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index d6cc7de320..f6da3e4e70 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -7,6 +7,15 @@ struct OptimalControlSolution
     input_sol::Union{Nothing, ODESolution}
 end
 
+function Base.show(io::IO, sol::OptimalControlSolution)
+    println("retcode: ", sol.sol.retcode, "\n")
+
+    println("Optimal control solution for following model:\n")
+    show(sol.model)
+
+    print("\n\nPlease query the model using sol.model, the solution trajectory for the system using sol.sol, or the solution trajectory for the controllers using sol.input_sol.")
+end
+
 function JuMPControlProblem end
 function InfiniteOptControlProblem end
 function CasADiControlProblem end
@@ -44,7 +53,7 @@ function SciMLBase.ControlFunction{iip, specialize}(sys::ODESystem,
         cse = true,
         kwargs...) where {iip, specialize}
 
-    (f), _, _ = generate_control_function(sys, inputs, disturbance_inputs; eval_expression = true, eval_module, cse, kwargs...)
+    (f), _, _ = generate_control_function(sys, inputs, disturbance_inputs; eval_module, cse, kwargs...)
 
     if tgrad
         tgrad_gen = generate_tgrad(sys, dvs, ps;
@@ -134,7 +143,7 @@ function SciMLBase.ControlFunction{false}(sys::AbstractODESystem, args...;
 end
 
 """
-IntegralNorm. When applied to an expression in a cost
+Integral operator. When applied to an expression in a cost
 function, assumes that the integration variable is the
 iv of the system, and assumes that the bounds are the
 tspan.
@@ -143,17 +152,46 @@ Equivalent to Integral(t in tspan) in Symbolics.
 struct ∫ <: Symbolics.Operator end
 ∫(x) = ∫()(x)
 Base.show(io::IO, x::∫) = print(io, "∫")
+Base.nameof(::∫) = :∫
 
-"""
-$(SIGNATURES)
+function (I::∫)(x)
+    Term{symtype(x)}(I, Any[x])
+end
+
+function (I::∫)(x::Num)
+    v = value(x)
+    Num(I(v))
+end
 
-Define one or more inputs.
+SymbolicUtils.promote_symtype(::Int, t) = t
 
-See also [`@independent_variables`](@ref), [`@variables`](@ref) and [`@constants`](@ref).
-"""
-macro inputs(xs...)
-    Symbolics._parse_vars(:inputs,
-        Real,
-        xs,
-        toparam) |> esc
+# returns the JuMP timespan, the number of steps, and whether it is a free time problem.
+function process_tspan(tspan, dt, steps)
+    is_free_time = false
+    if isnothing(dt) && isnothing(steps)
+        error("Must provide either the dt or the number of intervals to the collocation solvers (JuMP, InfiniteOpt, CasADi).")
+    elseif symbolic_type(tspan[1]) === ScalarSymbolic() || symbolic_type(tspan[2]) === ScalarSymbolic()
+        isnothing(steps) && error("Free final time problems require specifying the number of steps, rather than dt.")
+        isnothing(dt) || @warn "Specified dt for free final time problem. This will be ignored; dt will be determined by the number of timesteps."
+
+        return steps, true
+    else
+        isnothing(steps) || @warn "Specified number of steps for problem with concrete tspan. This will be ignored; number of steps will be determined by dt."
+
+        return length(tspan[1]:dt:tspan[2]), false 
+    end
 end
+
+#"""
+#$(SIGNATURES)
+#
+#Define one or more inputs.
+#
+#See also [`@independent_variables`](@ref), [`@variables`](@ref) and [`@constants`](@ref).
+#"""
+#macro inputs(xs...)
+#    Symbolics._parse_vars(:inputs,
+#        Real,
+#        xs,
+#        toparam) |> esc
+#end
diff --git a/src/variables.jl b/src/variables.jl
index f3dd16819d..939a68c706 100644
--- a/src/variables.jl
+++ b/src/variables.jl
@@ -332,6 +332,11 @@ function hasbounds(x)
     any(isfinite.(b[1]) .|| isfinite.(b[2]))
 end
 
+function setbounds(x::Num, bounds)
+    (lb, ub) = bounds
+    setmetadata(x, VariableBounds, (lb, ub))
+end
+
 ## Disturbance =================================================================
 struct VariableDisturbance end
 Symbolics.option_to_metadata_type(::Val{:disturbance}) = VariableDisturbance
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index 47e962ab9a..4a6788b008 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -139,26 +139,26 @@ end
     t = M.t_nounits
     D = M.D_nounits
 
-    @variables h(..) v(..) m(..) T(..) [input = true, bounds = (0, tₘ)]
-    @parameters h_c m₀ h₀ g₀ D_c c Tₘ
-    @parameters tf
+    @parameters h_c m₀ h₀ g₀ D_c c Tₘ m_c
+    @variables h(..) v(..) m(..) [bounds = (m_c, 1)] T(..) [input = true, bounds = (0, Tₘ)]
     drag(h, v) = D_c * v^2 * exp(-h_c * (h - h₀) / h₀)
     gravity(h) = g₀ * (h₀ / h)
 
     eqs = [D(h(t)) ~ v(t), 
-           D(v(t)) ~ (T(t) - drag(h(t), v(t))) / m(t) - gravity(t),
+           D(v(t)) ~ (T(t) - drag(h(t), v(t))) / m(t) - gravity(h(t)),
            D(m(t)) ~ -T(t) / c]
 
-    costs = [-h(tf)]
-    constraints = [T(tf) ~ 0]
+    (ts, te) = (0., 0.2)
+    costs = [-h(te)]
+    constraints = [T(te) ~ 0]
     @named rocket = ODESystem(eqs, t; costs, constraints)
-    @test tf ∈ Set(parameters(rocket))
+    rocket, input_idxs = structural_simplify(rocket, ([T(t)], []))
 
     u0map = [h(t) => h₀, m(t) => m₀, v(t) => 0]
-    pmap = [g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5*√(g₀*h₀), D_C => 0.5 * 620 * m₀/g₀, Tₘ => 3.5*g₀*m₀]
-    jprob = JuMPControlProblem(rocket, u0map, (0, tf), pmap)
+    pmap = [g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5*√(g₀*h₀), D_c => 0.5 * 620 * m₀/g₀, Tₘ => 3.5*g₀*m₀, T(t) => 0., h₀ => 1, m_c => 0.6]
+    jprob = JuMPControlProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
     jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3)
-    @test jsol.sol.u[end][1] ≈ 1.012
+    @test jsol.sol.u[end][1] > 1.012
 end
 
 @testset "Free final time problem" begin
@@ -167,20 +167,22 @@ end
 
     @variables x(..) u(..) [input = true, bounds = (0,1)]
     @parameters tf
-    eqs = [D(x(t)) ~ -2 + 0.5*u]
-
+    eqs = [D(x(t)) ~ -2 + 0.5*u(t)]
     # Integral cost function
-    costs = [∫(x-u), x(tf)]
+    costs = [-∫(x(t)-u(t)), -x(tf)]
     consolidate(u) = u[1] + u[2]
-    jprob = JuMPControlProblem(rocket, u0map, (0, tf), pmap)
-    jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3)
-    @test jsol.sol.t[end] ≈ 10.0
-    iprob = InfiniteOptControlProblem(rocket, u0map, (0, tf), pmap)
-    isol = solve(iprob, Ipopt.Optimizer, :RadauIA3)
-    @test isol.sol.t[end] ≈ 10.0
-end
+    @named rocket = ODESystem(eqs, t; costs, consolidate)
+    rocket, input_idxs = structural_simplify(rocket, ([u(t)], []))
+
+    u0map = [x(t) => 17.5]
+    pmap = [u(t) => 0., tf => 8]
+    jprob = JuMPControlProblem(rocket, u0map, (0, tf), pmap; steps = 201)
+    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    @test isapprox(jsol.sol.t[end], 10.0, rtol = 1e-3)
 
-@testset "Cart-pole problem" begin
+    iprob = InfiniteOptControlProblem(rocket, u0map, (0, tf), pmap; steps = 200)
+    isol = solve(iprob, Ipopt.Optimizer)
+    @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3)
 end
 
 #@testset "Constrained optimal control problems" begin

From cd9b0392b49cee75e36332b395edcd3688d8a351 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 18 Apr 2025 05:23:58 -0400
Subject: [PATCH 24/41] format

---
 ext/MTKJuMPControlExt.jl                 | 34 +++++++-------
 src/systems/optimal_control_interface.jl | 22 ++++++----
 test/extensions/jump_control.jl          | 56 +++++++++++++-----------
 3 files changed, 62 insertions(+), 50 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 310a632292..24dc05efd2 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -102,7 +102,7 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
 
     if is_free_t
         (ts_sym, te_sym) = tspan
-        @variable(model, tf, start = pmap[te_sym])
+        @variable(model, tf, start=pmap[te_sym])
         hasbounds(te_sym) && begin
             lo, hi = getbounds(te_sym)
             set_lower_bound(tf, lo)
@@ -112,11 +112,11 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
         pmap[te_sym] = 1
         tspan = (0, 1)
     end
-        
-    @infinite_parameter(model, t in [tspan[1], tspan[2]], num_supports = steps)
-    @variable(model, U[i = 1:length(states)], Infinite(t), start = u0[i])
+
+    @infinite_parameter(model, t in [tspan[1], tspan[2]], num_supports=steps)
+    @variable(model, U[i = 1:length(states)], Infinite(t), start=u0[i])
     c0 = [pmap[c] for c in ctrls]
-    @variable(model, V[i = 1:length(ctrls)], Infinite(t), start = c0[i])
+    @variable(model, V[i = 1:length(ctrls)], Infinite(t), start=c0[i])
 
     set_jump_bounds!(model, sys, pmap)
     add_jump_cost_function!(model, sys, (tspan[1], tspan[2]), pmap; is_free_t)
@@ -161,7 +161,8 @@ function add_jump_cost_function!(model::InfiniteModel, sys, tspan, pmap; is_free
 
     # Substitute integral
     iv = MTK.get_iv(sys)
-    jcosts = map(c -> Symbolics.substitute(c, MTK.∫() => Symbolics.Integral(iv in tspan)), jcosts)
+    jcosts = map(
+        c -> Symbolics.substitute(c, MTK.∫() => Symbolics.Integral(iv in tspan)), jcosts)
 
     intmap = Dict()
     for int in MTK.collect_applied_operators(jcosts, Symbolics.Integral)
@@ -183,7 +184,8 @@ function add_user_constraints!(model::InfiniteModel, sys, pmap; is_free_t = fals
         for u in MTK.get_unknowns(conssys)
             x = MTK.operation(u)
             t = only(arguments(u))
-            MTK.symbolic_type(t) === NotSymbolic() && error("Provided specific time constraint in a free final time problem. This is not supported by the JuMP/InfiniteOpt collocation solvers. The offending variable is $u.")
+            MTK.symbolic_type(t) === NotSymbolic() &&
+                error("Provided specific time constraint in a free final time problem. This is not supported by the JuMP/InfiniteOpt collocation solvers. The offending variable is $u.")
         end
     end
 
@@ -211,13 +213,15 @@ function substitute_jump_vars(model, sys, pmap, exprs)
     U = model[:U]
     V = model[:V]
     # 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)]])
+    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.substitute(c, whole_interval_map), exprs)
 
     # for variables like x(1.0)
     x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
     c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
-    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)]])
+    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.substitute(c, fixed_t_map), exprs)
 
     exprs = map(c -> Symbolics.substitute(c, Dict(pmap)), exprs)
@@ -236,11 +240,11 @@ function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap; is_free_
 
     diff_eqs = substitute_jump_vars(model, sys, pmap, diff_equations(sys))
     diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs)
-    @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs == tₛ * diff_eqs[i].rhs)
+    @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs==tₛ * diff_eqs[i].rhs)
 
     # Algebraic equations
     alg_eqs = substitute_jump_vars(model, sys, pmap, alg_equations(sys))
-    @constraint(model, A[i = 1:length(alg_eqs)], alg_eqs[i].lhs == alg_eqs[i].rhs)
+    @constraint(model, A[i = 1:length(alg_eqs)], alg_eqs[i].lhs==alg_eqs[i].rhs)
 end
 
 function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
@@ -283,11 +287,11 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
             for (i, h) in enumerate(c)
                 ΔU = @view ΔUs[i, :]
                 Uₙ = U + ΔU * dt
-                @constraint(model, [j = 1:nᵤ], K[i, j](τ) == tₛ * f(Uₙ, V, p, τ + h * dt)[j],
-                            DomainRestrictions(t => τ + h*dt), base_name="solve_K($τ)")
+                @constraint(model, [j = 1:nᵤ], K[i, j](τ)==tₛ * f(Uₙ, V, p, τ + h * dt)[j],
+                    DomainRestrictions(t => τ + h * dt), base_name="solve_K($τ)")
             end
-            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n] == U[n](τ + dt),
-                        DomainRestrictions(t => τ), base_name="solve_U($τ)")
+            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](τ + dt),
+                DomainRestrictions(t => τ), base_name="solve_U($τ)")
         end
     end
 end
diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index f6da3e4e70..a2201d008e 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -39,7 +39,7 @@ function SciMLBase.ControlFunction{iip, specialize}(sys::ODESystem,
         inputs = unbound_inputs(sys),
         disturbance_inputs = disturbances(sys);
         version = nothing, tgrad = false,
-        jac = false, controljac = false, 
+        jac = false, controljac = false,
         p = nothing, t = nothing,
         eval_expression = false,
         sparse = false, simplify = false,
@@ -52,8 +52,8 @@ function SciMLBase.ControlFunction{iip, specialize}(sys::ODESystem,
         initialization_data = nothing,
         cse = true,
         kwargs...) where {iip, specialize}
-
-    (f), _, _ = generate_control_function(sys, inputs, disturbance_inputs; eval_module, cse, kwargs...)
+    (f), _, _ = generate_control_function(
+        sys, inputs, disturbance_inputs; eval_module, cse, kwargs...)
 
     if tgrad
         tgrad_gen = generate_tgrad(sys, dvs, ps;
@@ -113,7 +113,7 @@ function SciMLBase.ControlFunction{iip, specialize}(sys::ODESystem,
         W_prototype = nothing
         controljac_prototype = nothing
     end
-    
+
     ControlFunction{iip, specialize}(f;
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
@@ -170,15 +170,19 @@ function process_tspan(tspan, dt, steps)
     is_free_time = false
     if isnothing(dt) && isnothing(steps)
         error("Must provide either the dt or the number of intervals to the collocation solvers (JuMP, InfiniteOpt, CasADi).")
-    elseif symbolic_type(tspan[1]) === ScalarSymbolic() || symbolic_type(tspan[2]) === ScalarSymbolic()
-        isnothing(steps) && error("Free final time problems require specifying the number of steps, rather than dt.")
-        isnothing(dt) || @warn "Specified dt for free final time problem. This will be ignored; dt will be determined by the number of timesteps."
+    elseif symbolic_type(tspan[1]) === ScalarSymbolic() ||
+           symbolic_type(tspan[2]) === ScalarSymbolic()
+        isnothing(steps) &&
+            error("Free final time problems require specifying the number of steps, rather than dt.")
+        isnothing(dt) ||
+            @warn "Specified dt for free final time problem. This will be ignored; dt will be determined by the number of timesteps."
 
         return steps, true
     else
-        isnothing(steps) || @warn "Specified number of steps for problem with concrete tspan. This will be ignored; number of steps will be determined by dt."
+        isnothing(steps) ||
+            @warn "Specified number of steps for problem with concrete tspan. This will be ignored; number of steps will be determined by dt."
 
-        return length(tspan[1]:dt:tspan[2]), false 
+        return length(tspan[1]:dt:tspan[2]), false
     end
 end
 
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index 4a6788b008..dcf029a40f 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -79,8 +79,9 @@ end
 @testset "Linear systems" begin
     function is_bangbang(input_sol, lbounds, ubounds, rtol = 1e-4)
         bangbang = true
-        for v in 1:length(input_sol.u[1]) - 1
-            all(i -> ≈(i[v], bounds[v]; rtol) || ≈(i[v], bounds[u]; rtol), input_sol.u) || (bangbang = false)
+        for v in 1:(length(input_sol.u[1]) - 1)
+            all(i -> ≈(i[v], bounds[v]; rtol) || ≈(i[v], bounds[u]; rtol), input_sol.u) ||
+                (bangbang = false)
         end
         bangbang
     end
@@ -88,26 +89,27 @@ end
     # Double integrator
     t = M.t_nounits
     D = M.D_nounits
-    @variables x(..) [bounds = (0., 0.25)] v(..)
-    @variables u(t) [bounds = (-1., 1.), input = true]
+    @variables x(..) [bounds = (0.0, 0.25)] v(..)
+    @variables u(t) [bounds = (-1.0, 1.0), input = true]
     constr = [v(1.0) ~ 0.0]
     cost = [-x(1.0)] # Maximize the final distance.
-    @named block = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u], t; costs = cost, constraints = constr)
-    block, input_idxs = structural_simplify(block, ([u],[]))
+    @named block = ODESystem(
+        [D(x(t)) ~ v(t), D(v(t)) ~ u], t; costs = cost, constraints = constr)
+    block, input_idxs = structural_simplify(block, ([u], []))
 
-    u0map = [x(t) => 0., v(t) => 0.]
-    tspan = (0., 1.)
-    parammap = [u => 0.]
+    u0map = [x(t) => 0.0, v(t) => 0.0]
+    tspan = (0.0, 1.0)
+    parammap = [u => 0.0]
     jprob = JuMPControlProblem(block, u0map, tspan, parammap; dt = 0.01)
     jsol = solve(jprob, Ipopt.Optimizer, :Verner8)
     # Linear systems have bang-bang controls
-    @test is_bangbang(jsol.input_sol, [-1.], [1.])
+    @test is_bangbang(jsol.input_sol, [-1.0], [1.0])
     # Test reached final position.
     @test ≈(jsol.sol.u[end][1], 0.25, rtol = 1e-5)
 
     iprob = InfiniteOptControlProblem(block, u0map, tspan, parammap; dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer; silent = true)
-    @test is_bangbang(isol.input_sol, [-1.], [1.])
+    @test is_bangbang(isol.input_sol, [-1.0], [1.0])
     @test ≈(isol.sol.u[end][1], 0.25, rtol = 1e-5)
 
     ###################
@@ -118,21 +120,21 @@ end
     @parameters b c μ s ν
 
     tspan = (0, 4)
-    eqs = [D(w(t)) ~ -μ*w(t) + b*s*α*w(t),
-           D(q(t)) ~ -ν*q(t) + c*(1 - α)*s*w(t)]
+    eqs = [D(w(t)) ~ -μ * w(t) + b * s * α * w(t),
+        D(q(t)) ~ -ν * q(t) + c * (1 - α) * s * w(t)]
     costs = [-q(tspan[2])]
-   
+
     @named beesys = ODESystem(eqs, t; costs)
-    beesys, input_idxs = structural_simplify(beesys, ([α],[]))
+    beesys, input_idxs = structural_simplify(beesys, ([α], []))
     u0map = [w(t) => 40, q(t) => 2]
     pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1, α => 1]
 
     jprob = JuMPControlProblem(beesys, u0map, tspan, pmap, dt = 0.01)
     jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
-    @test is_bangbang(jsol.input_sol, [0.], [1.])
+    @test is_bangbang(jsol.input_sol, [0.0], [1.0])
     iprob = InfiniteOptControlProblem(beesys, u0map, tspan, pmap, dt = 0.01)
     isol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
-    @test is_bangbang(isol.input_sol, [0.], [1.])
+    @test is_bangbang(isol.input_sol, [0.0], [1.0])
 end
 
 @testset "Rocket launch" begin
@@ -144,18 +146,20 @@ end
     drag(h, v) = D_c * v^2 * exp(-h_c * (h - h₀) / h₀)
     gravity(h) = g₀ * (h₀ / h)
 
-    eqs = [D(h(t)) ~ v(t), 
-           D(v(t)) ~ (T(t) - drag(h(t), v(t))) / m(t) - gravity(h(t)),
-           D(m(t)) ~ -T(t) / c]
+    eqs = [D(h(t)) ~ v(t),
+        D(v(t)) ~ (T(t) - drag(h(t), v(t))) / m(t) - gravity(h(t)),
+        D(m(t)) ~ -T(t) / c]
 
-    (ts, te) = (0., 0.2)
+    (ts, te) = (0.0, 0.2)
     costs = [-h(te)]
     constraints = [T(te) ~ 0]
     @named rocket = ODESystem(eqs, t; costs, constraints)
     rocket, input_idxs = structural_simplify(rocket, ([T(t)], []))
 
     u0map = [h(t) => h₀, m(t) => m₀, v(t) => 0]
-    pmap = [g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5*√(g₀*h₀), D_c => 0.5 * 620 * m₀/g₀, Tₘ => 3.5*g₀*m₀, T(t) => 0., h₀ => 1, m_c => 0.6]
+    pmap = [
+        g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5 * √(g₀ * h₀), D_c => 0.5 * 620 * m₀ / g₀,
+        Tₘ => 3.5 * g₀ * m₀, T(t) => 0.0, h₀ => 1, m_c => 0.6]
     jprob = JuMPControlProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
     jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3)
     @test jsol.sol.u[end][1] > 1.012
@@ -165,17 +169,17 @@ end
     t = M.t_nounits
     D = M.D_nounits
 
-    @variables x(..) u(..) [input = true, bounds = (0,1)]
+    @variables x(..) u(..) [input = true, bounds = (0, 1)]
     @parameters tf
-    eqs = [D(x(t)) ~ -2 + 0.5*u(t)]
+    eqs = [D(x(t)) ~ -2 + 0.5 * u(t)]
     # Integral cost function
-    costs = [-∫(x(t)-u(t)), -x(tf)]
+    costs = [-∫(x(t) - u(t)), -x(tf)]
     consolidate(u) = u[1] + u[2]
     @named rocket = ODESystem(eqs, t; costs, consolidate)
     rocket, input_idxs = structural_simplify(rocket, ([u(t)], []))
 
     u0map = [x(t) => 17.5]
-    pmap = [u(t) => 0., tf => 8]
+    pmap = [u(t) => 0.0, tf => 8]
     jprob = JuMPControlProblem(rocket, u0map, (0, tf), pmap; steps = 201)
     jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
     @test isapprox(jsol.sol.t[end], 10.0, rtol = 1e-3)

From 6988707d667cdcb7ce1dc04eb03396a8d937d4c8 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 22 Apr 2025 16:47:44 -0400
Subject: [PATCH 25/41] test: add trasncription tests

---
 ext/MTKJuMPControlExt.jl        |  5 +--
 test/extensions/jump_control.jl | 62 +++++++++++++++++++++++++++------
 2 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 24dc05efd2..4c90cc8448 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -40,8 +40,9 @@ end
     JuMPControlProblem(sys::ODESystem, u0, tspan, p; dt)
 
 Convert an ODESystem representing an optimal control system into a JuMP model
-for solving using optimization. Must provide `dt` for determining the length 
-of the interpolation arrays.
+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 `nsteps`.
 
 The optimization variables:
 - a vector-of-vectors U representing the unknowns as an interpolation array
diff --git a/test/extensions/jump_control.jl b/test/extensions/jump_control.jl
index dcf029a40f..d285e86886 100644
--- a/test/extensions/jump_control.jl
+++ b/test/extensions/jump_control.jl
@@ -6,6 +6,7 @@ using OrdinaryDiffEqSDIRK
 using Ipopt
 using BenchmarkTools
 using CairoMakie
+using DataInterpolations
 const M = ModelingToolkit
 
 @testset "ODE Solution, no cost" begin
@@ -76,21 +77,26 @@ const M = ModelingToolkit
     @test all(u -> u .> [3, 4], sol.u)
 end
 
-@testset "Linear systems" begin
-    function is_bangbang(input_sol, lbounds, ubounds, rtol = 1e-4)
-        bangbang = true
-        for v in 1:(length(input_sol.u[1]) - 1)
-            all(i -> ≈(i[v], bounds[v]; rtol) || ≈(i[v], bounds[u]; rtol), input_sol.u) ||
-                (bangbang = false)
-        end
-        bangbang
+function is_bangbang(input_sol, lbounds, ubounds, rtol = 1e-4)
+    for v in 1:(length(input_sol.u[1]) - 1)
+        all(i -> ≈(i[v], bounds[v]; rtol) || ≈(i[v], bounds[u]; rtol), input_sol.u) ||
+            return false
     end
+    true
+end
+
+function ctrl_to_spline(inputsol, splineType)
+    us = reduce(vcat, inputsol.u)
+    ts = reduce(vcat, inputsol.t)
+    splineType(us, ts)
+end
 
+@testset "Linear systems" begin
     # Double integrator
     t = M.t_nounits
     D = M.D_nounits
     @variables x(..) [bounds = (0.0, 0.25)] v(..)
-    @variables u(t) [bounds = (-1.0, 1.0), input = true]
+    @variables u(..) [bounds = (-1.0, 1.0), input = true]
     constr = [v(1.0) ~ 0.0]
     cost = [-x(1.0)] # Maximize the final distance.
     @named block = ODESystem(
@@ -99,18 +105,27 @@ end
 
     u0map = [x(t) => 0.0, v(t) => 0.0]
     tspan = (0.0, 1.0)
-    parammap = [u => 0.0]
+    parammap = [u(t) => 0.0]
     jprob = JuMPControlProblem(block, u0map, tspan, parammap; dt = 0.01)
     jsol = solve(jprob, Ipopt.Optimizer, :Verner8)
     # Linear systems have bang-bang controls
     @test is_bangbang(jsol.input_sol, [-1.0], [1.0])
     # Test reached final position.
     @test ≈(jsol.sol.u[end][1], 0.25, rtol = 1e-5)
+    # Test dynamics
+    @parameters (u_interp::LinearInterpolation)(..)
+    block_ode = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u_interp(t)], t)
+    spline = ctrl_to_spline(jsol.input_sol, LinearInterpolation)
+    oprob = ODEProblem(block, u0map, tspan, [u_interp => spline])
+    osol = solve(oprob, Vern8())
+    @test jsol.sol.u ≈ osol.u
 
     iprob = InfiniteOptControlProblem(block, u0map, tspan, parammap; dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer; silent = true)
     @test is_bangbang(isol.input_sol, [-1.0], [1.0])
     @test ≈(isol.sol.u[end][1], 0.25, rtol = 1e-5)
+    osol = solve(oprob, ImplicitEuler())
+    @test isol.sol.u ≈ osol.u
 
     ###################
     ### Bee example ###
@@ -133,8 +148,16 @@ end
     jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
     @test is_bangbang(jsol.input_sol, [0.0], [1.0])
     iprob = InfiniteOptControlProblem(beesys, u0map, tspan, pmap, dt = 0.01)
-    isol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    isol = solve(iprob, Ipopt.Optimizer; silent = true)
     @test is_bangbang(isol.input_sol, [0.0], [1.0])
+
+    @parameters (α_interp::LinearInterpolation)(..)
+    eqs = [D(w(t)) ~ -μ * w(t) + b * s * α_interp(t) * w(t),
+           D(q(t)) ~ -ν * q(t) + c * (1 - α_interp(t)) * s * w(t)]
+    beesys_ode = ODESystem(eqs, t)
+    oprob = ODEProblem(beesys_ode, u0map, tspan, [α_interp => ctrl_to_spline(jsol.input_sol, LinearInterpolation)])
+    osol = solve(oprob, Tsit5())
+    @test osol.u ≈ jsol.sol.u
 end
 
 @testset "Rocket launch" begin
@@ -163,6 +186,17 @@ end
     jprob = JuMPControlProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
     jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3)
     @test jsol.sol.u[end][1] > 1.012
+    
+    # Test solution
+    @parameters (T_interp::CubicSpline)(..)
+    eqs = [D(h(t)) ~ v(t),
+        D(v(t)) ~ (T_interp(t) - drag(h(t), v(t))) / m(t) - gravity(h(t)),
+        D(m(t)) ~ -T_interp(t) / c]
+    rocket_ode = ODESystem(eqs, t)
+    interpmap = Dict(T_interp => ctrl_to_spline(jsol.inputsol, CubicSpline))
+    oprob = ODEProblem(rocket_ode, u0map, tspan, merge(pmap, interpmap))
+    osol = solve(oprob, RadauIA3())
+    @test jsol.sol.u ≈ osol.u
 end
 
 @testset "Free final time problem" begin
@@ -189,5 +223,11 @@ end
     @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3)
 end
 
+using JuliaSimCompiler
+using Multibody.PlanarMechanics
+
+@testset "Cart-pole problem" begin
+end
+
 #@testset "Constrained optimal control problems" begin
 #end

From 353f12bbff5a16751a9b7bdb0973f18f062d01e9 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 22 Apr 2025 19:48:20 -0400
Subject: [PATCH 26/41] init new project for optimal control tests

---
 src/systems/optimal_control_interface.jl         | 16 ++++++++--------
 test/dynamic_optimization/Project.toml           |  5 +++++
 .../jump_control.jl                              |  0
 test/extensions/Project.toml                     |  3 ---
 4 files changed, 13 insertions(+), 11 deletions(-)
 create mode 100644 test/dynamic_optimization/Project.toml
 rename test/{extensions => dynamic_optimization}/jump_control.jl (100%)

diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index a2201d008e..c6d46bb0ba 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -33,7 +33,7 @@ end
 Generate the control function f(x, u, p, t) from the ODESystem. 
 Input variables are automatically inferred but can be manually specified.
 """
-function SciMLBase.ControlFunction{iip, specialize}(sys::ODESystem,
+function SciMLBase.ODEInputFunction{iip, specialize}(sys::ODESystem,
         dvs = unknowns(sys),
         ps = parameters(sys), u0 = nothing,
         inputs = unbound_inputs(sys),
@@ -114,7 +114,7 @@ function SciMLBase.ControlFunction{iip, specialize}(sys::ODESystem,
         controljac_prototype = nothing
     end
 
-    ControlFunction{iip, specialize}(f;
+    ODEInputFunction{iip, specialize}(f;
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
         controljac = _cjac === nothing ? nothing : _cjac,
@@ -128,18 +128,18 @@ function SciMLBase.ControlFunction{iip, specialize}(sys::ODESystem,
         initialization_data)
 end
 
-function SciMLBase.ControlFunction(sys::AbstractODESystem, args...; kwargs...)
-    ControlFunction{true}(sys, args...; kwargs...)
+function SciMLBase.ODEInputFunction(sys::AbstractODESystem, args...; kwargs...)
+    ODEInputFunction{true}(sys, args...; kwargs...)
 end
 
-function SciMLBase.ControlFunction{true}(sys::AbstractODESystem, args...;
+function SciMLBase.ODEInputFunction{true}(sys::AbstractODESystem, args...;
         kwargs...)
-    ControlFunction{true, SciMLBase.AutoSpecialize}(sys, args...; kwargs...)
+    ODEInputFunction{true, SciMLBase.AutoSpecialize}(sys, args...; kwargs...)
 end
 
-function SciMLBase.ControlFunction{false}(sys::AbstractODESystem, args...;
+function SciMLBase.ODEInputFunction{false}(sys::AbstractODESystem, args...;
         kwargs...)
-    ControlFunction{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
+    ODEInputFunction{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
 end
 
 """
diff --git a/test/dynamic_optimization/Project.toml b/test/dynamic_optimization/Project.toml
new file mode 100644
index 0000000000..ae0890688b
--- /dev/null
+++ b/test/dynamic_optimization/Project.toml
@@ -0,0 +1,5 @@
+[deps]
+DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
+InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
+Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
+Multibody = "e1cad5d1-98ef-44f9-a79a-9ca4547f95b9"
diff --git a/test/extensions/jump_control.jl b/test/dynamic_optimization/jump_control.jl
similarity index 100%
rename from test/extensions/jump_control.jl
rename to test/dynamic_optimization/jump_control.jl
diff --git a/test/extensions/Project.toml b/test/extensions/Project.toml
index 5800fe612f..fe2189b169 100644
--- a/test/extensions/Project.toml
+++ b/test/extensions/Project.toml
@@ -3,12 +3,9 @@ BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
-DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
-Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
-JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Nemo = "2edaba10-b0f1-5616-af89-8c11ac63239a"

From ac08f9ac93b477d34e0cd8947519c3e4af8acc39 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Wed, 23 Apr 2025 19:09:18 -0400
Subject: [PATCH 27/41] test: more test fixes

---
 ext/MTKJuMPControlExt.jl                  | 39 ++++++----
 src/systems/diffeqs/odesystem.jl          |  7 +-
 src/systems/systems.jl                    |  5 ++
 test/dynamic_optimization/Project.toml    | 11 ++-
 test/dynamic_optimization/jump_control.jl | 90 ++++++++++++++++-------
 test/runtests.jl                          | 11 +++
 6 files changed, 122 insertions(+), 41 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 4c90cc8448..463431e235 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -30,7 +30,7 @@ struct InfiniteOptControlProblem{uType, tType, isinplace, P, F, K} <:
     model::InfiniteModel
     kwargs::K
 
-    function InfiniteOptControlProblem(f, u0, tspan, p, model; kwargs...)
+    function InfiniteOptControlProblem(f, u0, tspan, p, model, kwargs...)
         new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f),
             typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
     end
@@ -58,7 +58,7 @@ function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap;
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
-    f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
+    f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
     pmap = MTK.todict(pmap)
@@ -84,7 +84,7 @@ function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
-    f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
+    f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
     pmap = MTK.todict(pmap)
@@ -116,7 +116,7 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
 
     @infinite_parameter(model, t in [tspan[1], tspan[2]], num_supports=steps)
     @variable(model, U[i = 1:length(states)], Infinite(t), start=u0[i])
-    c0 = [pmap[c] for c in ctrls]
+    c0 = MTK.value.([pmap[c] for c in ctrls])
     @variable(model, V[i = 1:length(ctrls)], Infinite(t), start=c0[i])
 
     set_jump_bounds!(model, sys, pmap)
@@ -124,8 +124,9 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
     add_user_constraints!(model, sys, pmap; is_free_t)
 
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
+    u0map = Dict([MTK.default_toterm(MTK.value(k)) => v for (k, v) in u0map])
     u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) :
-              [stidxmap[k] for (k, v) in u0map]
+        [stidxmap[MTK.default_toterm(k)] for (k, v) in u0map]
     add_initial_constraints!(model, u0, u0_idxs, tspan[1])
     return model
 end
@@ -190,7 +191,10 @@ function add_user_constraints!(model::InfiniteModel, sys, pmap; is_free_t = fals
         end
     end
 
-    jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints)
+    auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in unknowns(conssys)])
+    jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints; auxmap)
+    
+    # Substitute to-term'd variables
     for (i, cons) in enumerate(jconstraints)
         if cons isa Equation
             @constraint(model, cons.lhs - cons.rhs==0, base_name="user[$i]")
@@ -207,25 +211,28 @@ function add_initial_constraints!(model::InfiniteModel, u0, u0_idxs, ts)
     @constraint(model, initial[i in u0_idxs], U[i](ts)==u0[i])
 end
 
-function substitute_jump_vars(model, sys, pmap, exprs)
+function substitute_jump_vars(model, sys, pmap, exprs; auxmap = Dict())
     iv = MTK.get_iv(sys)
     sts = unknowns(sys)
     cts = MTK.unbound_inputs(sys)
     U = model[:U]
     V = model[:V]
+    exprs = map(c -> Symbolics.fixpoint_sub(c, auxmap), exprs)
+
     # 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.substitute(c, whole_interval_map), exprs)
+    exprs = map(c -> Symbolics.fixpoint_sub(c, whole_interval_map), exprs)
 
     # for variables like x(1.0)
     x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
     c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
     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.substitute(c, fixed_t_map), exprs)
 
-    exprs = map(c -> Symbolics.substitute(c, Dict(pmap)), exprs)
+    exprs = map(c -> Symbolics.fixpoint_sub(c, fixed_t_map), exprs)
+
+    exprs = map(c -> Symbolics.fixpoint_sub(c, Dict(pmap)), exprs)
     exprs
 end
 
@@ -255,8 +262,10 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
     model = prob.model
     f = prob.f
     p = prob.p
+
     t = model[:t]
-    tsteps = supports(model[:t])
+    tsteps = supports(t)
+    tmax = tsteps[end]
     pop!(tsteps)
     tₛ = is_free_t ? model[:tf] : 1
     dt = tsteps[2] - tsteps[1]
@@ -280,6 +289,7 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
                 base_name="solve_time_$τ")
             empty!(K)
         end
+        @show num_variables(model)
     else
         @variable(model, K[1:length(α), 1:nᵤ], Infinite(t), start=tsteps[1])
         ΔUs = A * K
@@ -288,10 +298,10 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
             for (i, h) in enumerate(c)
                 ΔU = @view ΔUs[i, :]
                 Uₙ = U + ΔU * dt
-                @constraint(model, [j = 1:nᵤ], K[i, j](τ)==tₛ * f(Uₙ, V, p, τ + h * dt)[j],
-                    DomainRestrictions(t => τ + h * dt), base_name="solve_K($τ)")
+                @constraint(model, [j = 1:nᵤ], K[i, j](τ)==(tₛ * f(Uₙ, V, p, τ + h * dt)[j]),
+                            DomainRestrictions(t => min(τ + h * dt, tmax)), base_name="solve_K($τ)")
             end
-            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](τ + dt),
+            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tmax)),
                 DomainRestrictions(t => τ), base_name="solve_U($τ)")
         end
     end
@@ -323,6 +333,7 @@ function DiffEqBase.solve(
     unregister(model, :K)
     for var in all_variables(model)
         if occursin("K", JuMP.name(var))
+            unregister(model, Symbol(JuMP.name(var)))
             delete(model, var)
         end
     end
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index ceef916e90..11fdf984ab 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -406,7 +406,7 @@ function ODESystem(eqs, iv; constraints = Equation[], costs = Num[], kwargs...)
     end
     costs = wrap.(costs)
 
-    return ODESystem(eqs, iv, collect(Iterators.flatten((diffvars, algevars, consvars))),
+    return ODESystem(eqs, iv, collect(Iterators.flatten((diffvars, algevars))),
         collect(new_ps); constraintsystem, costs, kwargs...)
 end
 
@@ -769,7 +769,9 @@ function process_constraint_system(
     constraintps = OrderedSet()
     for cons in constraints
         collect_vars!(constraintsts, constraintps, cons, iv)
+        union!(constraintsts, collect_applied_operators(cons, Differential))
     end
+    @show constraintsts
 
     # Validate the states.
     validate_vars_and_find_ps!(constraintsts, constraintps, sts, iv)
@@ -811,6 +813,9 @@ function validate_vars_and_find_ps!(auxvars, auxps, sysvars, iv)
         elseif length(arguments(var)) > 1
             throw(ArgumentError("Too many arguments for variable $var."))
         elseif length(arguments(var)) == 1
+            if iscall(var) && operation(var) isa Differential
+                var = only(arguments(var))
+            end
             arg = only(arguments(var))
             operation(var)(iv) ∈ sts ||
                 throw(ArgumentError("Variable $var is not a variable of the ODESystem. Called variables must be variables of the ODESystem."))
diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index 52f93afb9b..9da7249300 100644
--- a/src/systems/systems.jl
+++ b/src/systems/systems.jl
@@ -163,6 +163,11 @@ function __structural_simplify(
     end
 end
 
+function toterm_auxsystems(system::ODESystem)
+    constraints = system.constraintsystem.constraints
+
+end
+
 """
     $(TYPEDSIGNATURES)
 
diff --git a/test/dynamic_optimization/Project.toml b/test/dynamic_optimization/Project.toml
index ae0890688b..a5eaf3144b 100644
--- a/test/dynamic_optimization/Project.toml
+++ b/test/dynamic_optimization/Project.toml
@@ -1,5 +1,14 @@
 [deps]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
-Multibody = "e1cad5d1-98ef-44f9-a79a-9ca4547f95b9"
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
+OrdinaryDiffEqFIRK = "5960d6e9-dd7a-4743-88e7-cf307b64f125"
+OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
+OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2"
+SimpleDiffEq = "05bca326-078c-5bf0-a5bf-ce7c7982d7fd"
diff --git a/test/dynamic_optimization/jump_control.jl b/test/dynamic_optimization/jump_control.jl
index d285e86886..1565cf387f 100644
--- a/test/dynamic_optimization/jump_control.jl
+++ b/test/dynamic_optimization/jump_control.jl
@@ -2,7 +2,7 @@ using ModelingToolkit
 import JuMP, InfiniteOpt
 using DiffEqDevTools, DiffEqBase
 using SimpleDiffEq
-using OrdinaryDiffEqSDIRK
+using OrdinaryDiffEqSDIRK, OrdinaryDiffEqVerner, OrdinaryDiffEqTsit5, OrdinaryDiffEqFIRK
 using Ipopt
 using BenchmarkTools
 using CairoMakie
@@ -100,8 +100,8 @@ end
     constr = [v(1.0) ~ 0.0]
     cost = [-x(1.0)] # Maximize the final distance.
     @named block = ODESystem(
-        [D(x(t)) ~ v(t), D(v(t)) ~ u], t; costs = cost, constraints = constr)
-    block, input_idxs = structural_simplify(block, ([u], []))
+                             [D(x(t)) ~ v(t), D(v(t)) ~ u(t)], t; costs = cost, constraints = constr)
+    block, input_idxs = structural_simplify(block, ([u(t)], []))
 
     u0map = [x(t) => 0.0, v(t) => 0.0]
     tspan = (0.0, 1.0)
@@ -113,19 +113,19 @@ end
     # Test reached final position.
     @test ≈(jsol.sol.u[end][1], 0.25, rtol = 1e-5)
     # Test dynamics
-    @parameters (u_interp::LinearInterpolation)(..)
-    block_ode = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u_interp(t)], t)
-    spline = ctrl_to_spline(jsol.input_sol, LinearInterpolation)
-    oprob = ODEProblem(block, u0map, tspan, [u_interp => spline])
-    osol = solve(oprob, Vern8())
-    @test jsol.sol.u ≈ osol.u
+    @parameters (u_interp::ConstantInterpolation)(..)
+    @mtkbuild block_ode = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u_interp(t)], t)
+    spline = ctrl_to_spline(jsol.input_sol, ConstantInterpolation)
+    oprob = ODEProblem(block_ode, u0map, tspan, [u_interp => spline])
+    osol = solve(oprob, Vern8(), dt = 0.01, adaptive = false)
+    @test ≈(jsol.sol.u, osol.u, rtol = 0.05)
 
     iprob = InfiniteOptControlProblem(block, u0map, tspan, parammap; dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer; silent = true)
     @test is_bangbang(isol.input_sol, [-1.0], [1.0])
     @test ≈(isol.sol.u[end][1], 0.25, rtol = 1e-5)
-    osol = solve(oprob, ImplicitEuler())
-    @test isol.sol.u ≈ osol.u
+    osol = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false)
+    @test ≈(isol.sol.u, osol.u, rtol = 0.05)
 
     ###################
     ### Bee example ###
@@ -154,10 +154,12 @@ end
     @parameters (α_interp::LinearInterpolation)(..)
     eqs = [D(w(t)) ~ -μ * w(t) + b * s * α_interp(t) * w(t),
            D(q(t)) ~ -ν * q(t) + c * (1 - α_interp(t)) * s * w(t)]
-    beesys_ode = ODESystem(eqs, t)
-    oprob = ODEProblem(beesys_ode, u0map, tspan, [α_interp => ctrl_to_spline(jsol.input_sol, LinearInterpolation)])
-    osol = solve(oprob, Tsit5())
-    @test osol.u ≈ jsol.sol.u
+    @mtkbuild beesys_ode = ODESystem(eqs, t)
+    oprob = ODEProblem(beesys_ode, u0map, tspan, merge(Dict(pmap), Dict(α_interp => ctrl_to_spline(jsol.input_sol, LinearInterpolation))))
+    osol = solve(oprob, Tsit5(); dt = 0.01, adaptive = false)
+    @test ≈(osol.u, jsol.sol.u, rtol = 0.01)
+    osol2 = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false)
+    @test ≈(osol2.u, isol.sol.u, rtol = 0.01)
 end
 
 @testset "Rocket launch" begin
@@ -175,8 +177,8 @@ end
 
     (ts, te) = (0.0, 0.2)
     costs = [-h(te)]
-    constraints = [T(te) ~ 0]
-    @named rocket = ODESystem(eqs, t; costs, constraints)
+    cons = [T(te) ~ 0]
+    @named rocket = ODESystem(eqs, t; costs, constraints = cons)
     rocket, input_idxs = structural_simplify(rocket, ([T(t)], []))
 
     u0map = [h(t) => h₀, m(t) => m₀, v(t) => 0]
@@ -184,18 +186,22 @@ end
         g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5 * √(g₀ * h₀), D_c => 0.5 * 620 * m₀ / g₀,
         Tₘ => 3.5 * g₀ * m₀, T(t) => 0.0, h₀ => 1, m_c => 0.6]
     jprob = JuMPControlProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
-    jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3)
+    jsol = solve(jprob, Ipopt.Optimizer, :RadauIIA3)
     @test jsol.sol.u[end][1] > 1.012
+
+    iprob = InfiniteOptControlProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
+    isol = solve(iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3))
+    @test isol.sol.u[end][1] > 1.012
     
     # Test solution
     @parameters (T_interp::CubicSpline)(..)
     eqs = [D(h(t)) ~ v(t),
         D(v(t)) ~ (T_interp(t) - drag(h(t), v(t))) / m(t) - gravity(h(t)),
         D(m(t)) ~ -T_interp(t) / c]
-    rocket_ode = ODESystem(eqs, t)
-    interpmap = Dict(T_interp => ctrl_to_spline(jsol.inputsol, CubicSpline))
-    oprob = ODEProblem(rocket_ode, u0map, tspan, merge(pmap, interpmap))
-    osol = solve(oprob, RadauIA3())
+    @mtkbuild rocket_ode = ODESystem(eqs, t)
+    interpmap = Dict(T_interp => ctrl_to_spline(jsol.input_sol, CubicSpline))
+    oprob = ODEProblem(rocket_ode, u0map, (ts, te), merge(Dict(pmap), interpmap))
+    osol = solve(oprob, RadauIIA3())
     @test jsol.sol.u ≈ osol.u
 end
 
@@ -223,10 +229,44 @@ end
     @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3)
 end
 
-using JuliaSimCompiler
-using Multibody.PlanarMechanics
-
 @testset "Cart-pole problem" begin
+    # gravity, length, moment of Inertia, drag coeff
+    @parameters g l mₚ mₖ
+    @variables x(..) θ(..) u(t) [input = true, bounds = (-10, 10)]
+
+    s = sin(θ(t))
+    c = cos(θ(t))
+    H = [mₖ+mₚ  mₚ*l*c
+         mₚ*l*c mₚ*l^2]
+    C = [0 -mₚ*D(θ(t))*l*s
+         0  0]
+    qd = [D(x(t)), D(θ(t))]
+    G = [0, mₚ*g*l*s]
+    B = [1, 0]
+
+    tf = 5
+    rhss = -H \ Vector(C*qd + G - B*u)
+    eqs = [D(D(x(t))) ~ rhss[1], D(D(θ(t))) ~ rhss[2]]
+    cons = [θ(tf) ~ π, x(tf) ~ 0, D(θ(tf)) ~ 0, D(x(tf)) ~ 0]
+    costs = [∫(u^2)]
+    tspan = (0, tf)
+
+    @named cartpole = ODESystem(eqs, t; costs, constraints = cons)
+    cartpole, input_idxs = structural_simplify(cartpole, ([u], []))
+
+    u0map = [D(x(t)) => 0., D(θ(t)) => 0., θ(t) => 0., x(t) => 0.]
+    pmap = [mₖ => 1., mₚ => 0.2, l => 0.5, g => 9.81, u => 0]
+    jprob = JuMPControlProblem(cartpole, u0map, tspan, pmap; dt = 0.04)
+    jsol = solve(jprob, Ipopt.Optimizer, :RK4)
+    @test jsol.sol.u[end] ≈ [π, 0, 0, 0]
+
+    iprob = InfiniteOptControlProblem(cartpole, u0map, tspan, pmap; dt = 0.04)
+    isol = solve(iprob, Ipopt.Optimizer)
+    @test isol.sol.u[end] ≈ [π, 0, 0, 0]
+end
+
+# RC Circuit
+@testset "MTK Components" begin
 end
 
 #@testset "Constrained optimal control problems" begin
diff --git a/test/runtests.jl b/test/runtests.jl
index 08e7425a3c..0e0ea29a2c 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -22,6 +22,12 @@ function activate_downstream_env()
     Pkg.instantiate()
 end
 
+function activate_dynamic_optimization_env()
+    Pkg.activate("dynamic_optimization")
+    Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
+    Pkg.instantiate()
+end
+
 @time begin
     if GROUP == "All" || GROUP == "InterfaceI"
         @testset "InterfaceI" begin
@@ -143,4 +149,9 @@ end
         @safetestset "InfiniteOpt Extension Test" include("extensions/test_infiniteopt.jl")
         @safetestset "JuMPControl Extension Test" include("extensions/jump_control.jl")
     end
+
+    if GROUP == "All" || GROUP == "Dynamic Optimization"
+        activate_dynamic_optimization_env()
+        @safetestset "JuMP Collocation Solvers" include("dynamic_optimization/jump_control")
+    end
 end

From 943181e1bb95f8f9ecd0191aa59882627341d4cc Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 25 Apr 2025 14:56:55 -0400
Subject: [PATCH 28/41] more test fixes

---
 ext/MTKJuMPControlExt.jl                  |  9 ++-
 src/systems/diffeqs/odesystem.jl          |  3 +-
 test/dynamic_optimization/jump_control.jl | 78 +++++++++++++----------
 3 files changed, 51 insertions(+), 39 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 463431e235..45cc9e4365 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -289,17 +289,16 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
                 base_name="solve_time_$τ")
             empty!(K)
         end
-        @show num_variables(model)
     else
-        @variable(model, K[1:length(α), 1:nᵤ], Infinite(t), start=tsteps[1])
+        @variable(model, K[1:length(α), 1:nᵤ], Infinite(t))
         ΔUs = A * K
         ΔU_tot = dt * (K' * α)
         for τ in tsteps
             for (i, h) in enumerate(c)
                 ΔU = @view ΔUs[i, :]
-                Uₙ = U + ΔU * dt
-                @constraint(model, [j = 1:nᵤ], K[i, j](τ)==(tₛ * f(Uₙ, V, p, τ + h * dt)[j]),
-                            DomainRestrictions(t => min(τ + h * dt, tmax)), base_name="solve_K($τ)")
+                Uₙ = U + ΔU * h * dt
+                @constraint(model, [j = 1:nᵤ], K[i, j]==(tₛ * f(Uₙ, V, p, τ + h * dt)[j]),
+                            DomainRestrictions(t => τ), base_name="solve_K$i($τ)")
             end
             @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tmax)),
                 DomainRestrictions(t => τ), base_name="solve_U($τ)")
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index 11fdf984ab..62efcc3f3f 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -762,7 +762,7 @@ end
 Build the constraint system for the ODESystem.
 """
 function process_constraint_system(
-        constraints::Vector{Equation}, sts, ps, iv; consname = :cons)
+        constraints::Vector, sts, ps, iv; consname = :cons)
     isempty(constraints) && return nothing
 
     constraintsts = OrderedSet()
@@ -771,7 +771,6 @@ function process_constraint_system(
         collect_vars!(constraintsts, constraintps, cons, iv)
         union!(constraintsts, collect_applied_operators(cons, Differential))
     end
-    @show constraintsts
 
     # Validate the states.
     validate_vars_and_find_ps!(constraintsts, constraintps, sts, iv)
diff --git a/test/dynamic_optimization/jump_control.jl b/test/dynamic_optimization/jump_control.jl
index 1565cf387f..43ccedf33e 100644
--- a/test/dynamic_optimization/jump_control.jl
+++ b/test/dynamic_optimization/jump_control.jl
@@ -4,10 +4,8 @@ using DiffEqDevTools, DiffEqBase
 using SimpleDiffEq
 using OrdinaryDiffEqSDIRK, OrdinaryDiffEqVerner, OrdinaryDiffEqTsit5, OrdinaryDiffEqFIRK
 using Ipopt
-using BenchmarkTools
-using CairoMakie
 using DataInterpolations
-const M = ModelingToolkit
+#const M = ModelingToolkit
 
 @testset "ODE Solution, no cost" begin
     # Test solving without anything attached.
@@ -26,19 +24,19 @@ const M = ModelingToolkit
 
     # Test explicit method.
     jprob = JuMPControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
-    @test num_constraints(jprob.model) == 2 # initials
+    @test JuMP.num_constraints(jprob.model) == 2 # initials
     jsol = solve(jprob, Ipopt.Optimizer, :RK4)
     oprob = ODEProblem(sys, u0map, tspan, parammap)
     osol = solve(oprob, SimpleRK4(), dt = 0.01)
     @test jsol.sol.u ≈ osol.u
 
     # Implicit method.
-    jsol2 = @btime solve($jprob, Ipopt.Optimizer, :ImplicitEuler, silent = true) # 63.031 ms, 26.49 MiB
-    osol2 = @btime solve($oprob, ImplicitEuler(), dt = 0.01, adaptive = false) # 129.375 μs, 61.91 KiB
+    jsol2 = solve(jprob, Ipopt.Optimizer, :ImplicitEuler, silent = true) # 63.031 ms, 26.49 MiB
+    osol2 = solve(oprob, ImplicitEuler(), dt = 0.01, adaptive = false) # 129.375 μs, 61.91 KiB
     @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001)
     iprob = InfiniteOptControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
-    isol = @btime solve(
-        $iprob, Ipopt.Optimizer, derivative_method = FiniteDifference(Backward()), silent = true) # 11.540 ms, 4.00 MiB
+    isol = solve(iprob, Ipopt.Optimizer, derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward()), silent = true) # 11.540 ms, 4.00 MiB
+    @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001)
 
     # With a constraint
     u0map = Pair[]
@@ -47,34 +45,29 @@ const M = ModelingToolkit
     @mtkbuild lksys = ODESystem(eqs, t; constraints = constr)
 
     jprob = JuMPControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
-    @test num_constraints(jprob.model) == 2
-    jsol = @btime solve($jprob, Ipopt.Optimizer, :Tsitouras5, silent = true) # 12.190 s, 9.68 GiB
-    sol = jsol.sol
-    @test sol(0.6)[1] ≈ 3.5
-    @test sol(0.3)[1] ≈ 7.0
+    @test JuMP.num_constraints(jprob.model) == 2
+    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5, silent = true) # 12.190 s, 9.68 GiB
+    @test jsol.sol(0.6)[1] ≈ 3.5
+    @test jsol.sol(0.3)[1] ≈ 7.0
 
     iprob = InfiniteOptControlProblem(
         lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
-    isol = @btime solve(
-        $iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
+    isol = solve(iprob, Ipopt.Optimizer, derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
     sol = isol.sol
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0
 
     # Test whole-interval constraints
-    constr = [x(t) > 3, y(t) > 4]
+    constr = [x(t) ≳ 1, y(t) ≳ 1]
     @mtkbuild lksys = ODESystem(eqs, t; constraints = constr)
     iprob = InfiniteOptControlProblem(
         lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
-    isol = @btime solve(
-        $iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
-    sol = isol.sol
-    @test all(u -> u .> [3, 4], sol.u)
+    isol = solve(iprob, Ipopt.Optimizer, derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
+    @test all(u -> u > [1, 1], isol.sol.u)
 
     jprob = JuMPControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
-    jsol = @btime solve($jprob, Ipopt.Optimizer, :RadauIA3, silent = true) # 12.190 s, 9.68 GiB
-    sol = jsol.sol
-    @test all(u -> u .> [3, 4], sol.u)
+    jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3, silent = true) # 12.190 s, 9.68 GiB
+    @test all(u -> u > [1, 1], jsol.sol.u)
 end
 
 function is_bangbang(input_sol, lbounds, ubounds, rtol = 1e-4)
@@ -186,11 +179,11 @@ end
         g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5 * √(g₀ * h₀), D_c => 0.5 * 620 * m₀ / g₀,
         Tₘ => 3.5 * g₀ * m₀, T(t) => 0.0, h₀ => 1, m_c => 0.6]
     jprob = JuMPControlProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
-    jsol = solve(jprob, Ipopt.Optimizer, :RadauIIA3)
+    jsol = solve(jprob, Ipopt.Optimizer, :RadauIIA5, silent = true)
     @test jsol.sol.u[end][1] > 1.012
 
     iprob = InfiniteOptControlProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
-    isol = solve(iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3))
+    isol = solve(iprob, Ipopt.Optimizer, derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true)
     @test isol.sol.u[end][1] > 1.012
     
     # Test solution
@@ -201,8 +194,8 @@ end
     @mtkbuild rocket_ode = ODESystem(eqs, t)
     interpmap = Dict(T_interp => ctrl_to_spline(jsol.input_sol, CubicSpline))
     oprob = ODEProblem(rocket_ode, u0map, (ts, te), merge(Dict(pmap), interpmap))
-    osol = solve(oprob, RadauIIA3())
-    @test jsol.sol.u ≈ osol.u
+    osol = solve(oprob, RadauIIA5(); adaptive = false, dt = 0.005)
+    @test ≈(jsol.sol.u, osol.u, rtol = 0.02)
 end
 
 @testset "Free final time problem" begin
@@ -266,8 +259,29 @@ end
 end
 
 # RC Circuit
-@testset "MTK Components" begin
-end
-
-#@testset "Constrained optimal control problems" begin
-#end
+# using ModelingToolkitStandardLibrary.Electrical
+# @testset "MTK Components" begin
+#     @mtkmodel RL begin
+#         @parameters begin
+#             R = 1.0
+#             L = 1.0
+#         end
+#         @components begin
+#             resistor = Resistor(R = R)
+#             inductor = Inductor(L = L)
+#             source = Voltage()
+#             ground = Ground()
+#         end
+#         @equations begin
+#             connect(source.p, resistor.p)
+#             connect(resistor.n, inductor.p)
+#             connect(inductor.n, source.n, ground.g)
+#         end
+#     end
+# 
+#     costs = []
+#     coalesce = sum
+#     @named sys = RL()
+#     sys, _ = structural_simplify(sys, inputs = [sys.source.V.u])
+#     @parameters tf λ i₀
+# end

From facb5c02001510e7e2fda02c7344e87ff8f3f6fc Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 25 Apr 2025 15:01:21 -0400
Subject: [PATCH 29/41] clean up comments and tests

---
 src/systems/optimal_control_interface.jl  | 14 ------------
 src/systems/systems.jl                    |  5 ----
 test/dynamic_optimization/jump_control.jl | 28 -----------------------
 test/extensions/Project.toml              |  2 --
 4 files changed, 49 deletions(-)

diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index c6d46bb0ba..fbf04c1b99 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -185,17 +185,3 @@ function process_tspan(tspan, dt, steps)
         return length(tspan[1]:dt:tspan[2]), false
     end
 end
-
-#"""
-#$(SIGNATURES)
-#
-#Define one or more inputs.
-#
-#See also [`@independent_variables`](@ref), [`@variables`](@ref) and [`@constants`](@ref).
-#"""
-#macro inputs(xs...)
-#    Symbolics._parse_vars(:inputs,
-#        Real,
-#        xs,
-#        toparam) |> esc
-#end
diff --git a/src/systems/systems.jl b/src/systems/systems.jl
index 9da7249300..52f93afb9b 100644
--- a/src/systems/systems.jl
+++ b/src/systems/systems.jl
@@ -163,11 +163,6 @@ function __structural_simplify(
     end
 end
 
-function toterm_auxsystems(system::ODESystem)
-    constraints = system.constraintsystem.constraints
-
-end
-
 """
     $(TYPEDSIGNATURES)
 
diff --git a/test/dynamic_optimization/jump_control.jl b/test/dynamic_optimization/jump_control.jl
index 43ccedf33e..284316d6d7 100644
--- a/test/dynamic_optimization/jump_control.jl
+++ b/test/dynamic_optimization/jump_control.jl
@@ -257,31 +257,3 @@ end
     isol = solve(iprob, Ipopt.Optimizer)
     @test isol.sol.u[end] ≈ [π, 0, 0, 0]
 end
-
-# RC Circuit
-# using ModelingToolkitStandardLibrary.Electrical
-# @testset "MTK Components" begin
-#     @mtkmodel RL begin
-#         @parameters begin
-#             R = 1.0
-#             L = 1.0
-#         end
-#         @components begin
-#             resistor = Resistor(R = R)
-#             inductor = Inductor(L = L)
-#             source = Voltage()
-#             ground = Ground()
-#         end
-#         @equations begin
-#             connect(source.p, resistor.p)
-#             connect(resistor.n, inductor.p)
-#             connect(inductor.n, source.n, ground.g)
-#         end
-#     end
-# 
-#     costs = []
-#     coalesce = sum
-#     @named sys = RL()
-#     sys, _ = structural_simplify(sys, inputs = [sys.source.V.u])
-#     @parameters tf λ i₀
-# end
diff --git a/test/extensions/Project.toml b/test/extensions/Project.toml
index fe2189b169..03e65b4978 100644
--- a/test/extensions/Project.toml
+++ b/test/extensions/Project.toml
@@ -2,7 +2,6 @@
 BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
@@ -13,7 +12,6 @@ NonlinearSolveHomotopyContinuation = "2ac3b008-d579-4536-8c91-a1a5998c2f8b"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
-SimpleDiffEq = "05bca326-078c-5bf0-a5bf-ce7c7982d7fd"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

From 8a2f618748a7e94d38d8d651cf274ea595a9d566 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 25 Apr 2025 15:04:09 -0400
Subject: [PATCH 30/41] format

---
 ext/MTKJuMPControlExt.jl                  |  6 ++--
 test/dynamic_optimization/jump_control.jl | 40 ++++++++++++++---------
 2 files changed, 28 insertions(+), 18 deletions(-)

diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPControlExt.jl
index 45cc9e4365..30486f845d 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPControlExt.jl
@@ -126,7 +126,7 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
     u0map = Dict([MTK.default_toterm(MTK.value(k)) => v for (k, v) in u0map])
     u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) :
-        [stidxmap[MTK.default_toterm(k)] for (k, v) in u0map]
+              [stidxmap[MTK.default_toterm(k)] for (k, v) in u0map]
     add_initial_constraints!(model, u0, u0_idxs, tspan[1])
     return model
 end
@@ -193,7 +193,7 @@ function add_user_constraints!(model::InfiniteModel, sys, pmap; is_free_t = fals
 
     auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in unknowns(conssys)])
     jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints; auxmap)
-    
+
     # Substitute to-term'd variables
     for (i, cons) in enumerate(jconstraints)
         if cons isa Equation
@@ -298,7 +298,7 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
                 ΔU = @view ΔUs[i, :]
                 Uₙ = U + ΔU * h * dt
                 @constraint(model, [j = 1:nᵤ], K[i, j]==(tₛ * f(Uₙ, V, p, τ + h * dt)[j]),
-                            DomainRestrictions(t => τ), base_name="solve_K$i($τ)")
+                    DomainRestrictions(t => τ), base_name="solve_K$i($τ)")
             end
             @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tmax)),
                 DomainRestrictions(t => τ), base_name="solve_U($τ)")
diff --git a/test/dynamic_optimization/jump_control.jl b/test/dynamic_optimization/jump_control.jl
index 284316d6d7..587afa8ecf 100644
--- a/test/dynamic_optimization/jump_control.jl
+++ b/test/dynamic_optimization/jump_control.jl
@@ -35,7 +35,9 @@ using DataInterpolations
     osol2 = solve(oprob, ImplicitEuler(), dt = 0.01, adaptive = false) # 129.375 μs, 61.91 KiB
     @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001)
     iprob = InfiniteOptControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
-    isol = solve(iprob, Ipopt.Optimizer, derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward()), silent = true) # 11.540 ms, 4.00 MiB
+    isol = solve(iprob, Ipopt.Optimizer,
+        derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward()),
+        silent = true) # 11.540 ms, 4.00 MiB
     @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001)
 
     # With a constraint
@@ -52,7 +54,8 @@ using DataInterpolations
 
     iprob = InfiniteOptControlProblem(
         lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
-    isol = solve(iprob, Ipopt.Optimizer, derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
+    isol = solve(iprob, Ipopt.Optimizer,
+        derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
     sol = isol.sol
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0
@@ -62,7 +65,8 @@ using DataInterpolations
     @mtkbuild lksys = ODESystem(eqs, t; constraints = constr)
     iprob = InfiniteOptControlProblem(
         lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
-    isol = solve(iprob, Ipopt.Optimizer, derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
+    isol = solve(iprob, Ipopt.Optimizer,
+        derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
     @test all(u -> u > [1, 1], isol.sol.u)
 
     jprob = JuMPControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
@@ -93,7 +97,7 @@ end
     constr = [v(1.0) ~ 0.0]
     cost = [-x(1.0)] # Maximize the final distance.
     @named block = ODESystem(
-                             [D(x(t)) ~ v(t), D(v(t)) ~ u(t)], t; costs = cost, constraints = constr)
+        [D(x(t)) ~ v(t), D(v(t)) ~ u(t)], t; costs = cost, constraints = constr)
     block, input_idxs = structural_simplify(block, ([u(t)], []))
 
     u0map = [x(t) => 0.0, v(t) => 0.0]
@@ -146,9 +150,13 @@ end
 
     @parameters (α_interp::LinearInterpolation)(..)
     eqs = [D(w(t)) ~ -μ * w(t) + b * s * α_interp(t) * w(t),
-           D(q(t)) ~ -ν * q(t) + c * (1 - α_interp(t)) * s * w(t)]
+        D(q(t)) ~ -ν * q(t) + c * (1 - α_interp(t)) * s * w(t)]
     @mtkbuild beesys_ode = ODESystem(eqs, t)
-    oprob = ODEProblem(beesys_ode, u0map, tspan, merge(Dict(pmap), Dict(α_interp => ctrl_to_spline(jsol.input_sol, LinearInterpolation))))
+    oprob = ODEProblem(beesys_ode,
+        u0map,
+        tspan,
+        merge(Dict(pmap),
+            Dict(α_interp => ctrl_to_spline(jsol.input_sol, LinearInterpolation))))
     osol = solve(oprob, Tsit5(); dt = 0.01, adaptive = false)
     @test ≈(osol.u, jsol.sol.u, rtol = 0.01)
     osol2 = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false)
@@ -182,10 +190,12 @@ end
     jsol = solve(jprob, Ipopt.Optimizer, :RadauIIA5, silent = true)
     @test jsol.sol.u[end][1] > 1.012
 
-    iprob = InfiniteOptControlProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
-    isol = solve(iprob, Ipopt.Optimizer, derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true)
+    iprob = InfiniteOptControlProblem(
+        rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
+    isol = solve(iprob, Ipopt.Optimizer,
+        derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true)
     @test isol.sol.u[end][1] > 1.012
-    
+
     # Test solution
     @parameters (T_interp::CubicSpline)(..)
     eqs = [D(h(t)) ~ v(t),
@@ -229,16 +239,16 @@ end
 
     s = sin(θ(t))
     c = cos(θ(t))
-    H = [mₖ+mₚ  mₚ*l*c
+    H = [mₖ+mₚ mₚ*l*c
          mₚ*l*c mₚ*l^2]
     C = [0 -mₚ*D(θ(t))*l*s
-         0  0]
+         0 0]
     qd = [D(x(t)), D(θ(t))]
-    G = [0, mₚ*g*l*s]
+    G = [0, mₚ * g * l * s]
     B = [1, 0]
 
     tf = 5
-    rhss = -H \ Vector(C*qd + G - B*u)
+    rhss = -H \ Vector(C * qd + G - B * u)
     eqs = [D(D(x(t))) ~ rhss[1], D(D(θ(t))) ~ rhss[2]]
     cons = [θ(tf) ~ π, x(tf) ~ 0, D(θ(tf)) ~ 0, D(x(tf)) ~ 0]
     costs = [∫(u^2)]
@@ -247,8 +257,8 @@ end
     @named cartpole = ODESystem(eqs, t; costs, constraints = cons)
     cartpole, input_idxs = structural_simplify(cartpole, ([u], []))
 
-    u0map = [D(x(t)) => 0., D(θ(t)) => 0., θ(t) => 0., x(t) => 0.]
-    pmap = [mₖ => 1., mₚ => 0.2, l => 0.5, g => 9.81, u => 0]
+    u0map = [D(x(t)) => 0.0, D(θ(t)) => 0.0, θ(t) => 0.0, x(t) => 0.0]
+    pmap = [mₖ => 1.0, mₚ => 0.2, l => 0.5, g => 9.81, u => 0]
     jprob = JuMPControlProblem(cartpole, u0map, tspan, pmap; dt = 0.04)
     jsol = solve(jprob, Ipopt.Optimizer, :RK4)
     @test jsol.sol.u[end] ≈ [π, 0, 0, 0]

From 022987c60ec5e111cae4608e1c3c3ac4a0e09edd Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 25 Apr 2025 17:49:52 -0400
Subject: [PATCH 31/41] rename Control -> DynamicOpt

---
 Project.toml                                  |   2 +-
 ...PControlExt.jl => MTKJuMPDynamicOptExt.jl} |  42 ++--
 src/ModelingToolkit.jl                        |   6 +-
 src/systems/optimal_control_interface.jl      |  14 +-
 test/downstream/Project.toml                  |   1 +
 test/dynamic_optimization/jump_control.jl     |  32 +--
 test/extensions/ad.jl                         |   1 +
 test/runtests.jl                              | 182 +++++++++---------
 8 files changed, 141 insertions(+), 139 deletions(-)
 rename ext/{MTKJuMPControlExt.jl => MTKJuMPDynamicOptExt.jl} (90%)

diff --git a/Project.toml b/Project.toml
index 763c3931b9..e5adfbf35b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -78,7 +78,7 @@ MTKChainRulesCoreExt = "ChainRulesCore"
 MTKDeepDiffsExt = "DeepDiffs"
 MTKFMIExt = "FMI"
 MTKInfiniteOptExt = "InfiniteOpt"
-MTKJuMPControlExt = ["JuMP", "DiffEqDevTools", "InfiniteOpt"]
+MTKJuMPDynamicOptExt = ["JuMP", "DiffEqDevTools", "InfiniteOpt"]
 MTKLabelledArraysExt = "LabelledArrays"
 
 [compat]
diff --git a/ext/MTKJuMPControlExt.jl b/ext/MTKJuMPDynamicOptExt.jl
similarity index 90%
rename from ext/MTKJuMPControlExt.jl
rename to ext/MTKJuMPDynamicOptExt.jl
index 30486f845d..882456dab3 100644
--- a/ext/MTKJuMPControlExt.jl
+++ b/ext/MTKJuMPDynamicOptExt.jl
@@ -1,4 +1,4 @@
-module MTKJuMPControlExt
+module MTKJuMPDynamicOptExt
 using ModelingToolkit
 using JuMP, InfiniteOpt
 using DiffEqDevTools, DiffEqBase
@@ -6,8 +6,8 @@ using LinearAlgebra
 using StaticArrays
 const MTK = ModelingToolkit
 
-struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <:
-       AbstractOptimalControlProblem{uType, tType, isinplace}
+struct JuMPDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
+       AbstractDynamicOptProblem{uType, tType, isinplace}
     f::F
     u0::uType
     tspan::tType
@@ -15,14 +15,14 @@ struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <:
     model::InfiniteModel
     kwargs::K
 
-    function JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
+    function JuMPDynamicOptProblem(f, u0, tspan, p, model, kwargs...)
         new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f, 5),
             typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
     end
 end
 
-struct InfiniteOptControlProblem{uType, tType, isinplace, P, F, K} <:
-       AbstractOptimalControlProblem{uType, tType, isinplace}
+struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
+       AbstractDynamicOptProblem{uType, tType, isinplace}
     f::F
     u0::uType
     tspan::tType
@@ -30,14 +30,14 @@ struct InfiniteOptControlProblem{uType, tType, isinplace, P, F, K} <:
     model::InfiniteModel
     kwargs::K
 
-    function InfiniteOptControlProblem(f, u0, tspan, p, model, kwargs...)
+    function InfiniteOptDynamicOptProblem(f, u0, tspan, p, model, kwargs...)
         new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f),
             typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
     end
 end
 
 """
-    JuMPControlProblem(sys::ODESystem, u0, tspan, p; dt)
+    JuMPDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt)
 
 Convert an ODESystem representing an optimal control system into a JuMP model
 for solving using optimization. Must provide either `dt`, the timestep between collocation 
@@ -52,7 +52,7 @@ The constraints are:
 - The set of user constraints passed to the ODESystem via `constraints`
 - The solver constraints that encode the time-stepping used by the solver
 """
-function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap;
+function MTK.JuMPDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
         dt = nothing,
         steps = nothing,
         guesses = Dict(), kwargs...)
@@ -65,20 +65,20 @@ function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap;
     steps, is_free_t = MTK.process_tspan(tspan, dt, steps)
     model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
 
-    JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
+    JuMPDynamicOptProblem(f, u0, tspan, p, model, kwargs...)
 end
 
 """
-    InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap; dt)
+    InfiniteOptDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; dt)
 
 Convert an ODESystem 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 `JuMPControlProblem`, but directly adds the differential equations
+Related to `JuMPDynamicOptProblem`, but directly adds the differential equations
 of the system as derivative constraints, rather than using a solver tableau.
 """
-function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
+function MTK.InfiniteOptDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
         dt = nothing,
         steps = nothing,
         guesses = Dict(), kwargs...)
@@ -92,7 +92,7 @@ function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
     model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
 
     add_infopt_solve_constraints!(model, sys, pmap; is_free_t)
-    InfiniteOptControlProblem(f, u0, tspan, p, model, kwargs...)
+    InfiniteOptDynamicOptProblem(f, u0, tspan, p, model, kwargs...)
 end
 
 # Initialize InfiniteOpt model.
@@ -307,16 +307,16 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
 end
 
 """
-Solve JuMPControlProblem. Arguments:
-- prob: a JumpControlProblem
+Solve JuMPDynamicOptProblem. Arguments:
+- prob: a JumpDynamicOptProblem
 - jump_solver: a LP solver such as HiGHS
 - ode_solver: Takes in a symbol representing the solver. Acceptable solvers may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. Note that the symbol may be different than the typical name of the solver, e.g. :Tsitouras5 rather than Tsit5.
 - silent: set the model silent (suppress model output)
 
-Returns a JuMPControlSolution, which contains both the model and the ODE solution.
+Returns a DynamicOptSolution, which contains both the model and the ODE solution.
 """
 function DiffEqBase.solve(
-        prob::JuMPControlProblem, jump_solver, ode_solver::Symbol; silent = false)
+        prob::JuMPDynamicOptProblem, jump_solver, ode_solver::Symbol; silent = false)
     model = prob.model
     tableau_getter = Symbol(:construct, ode_solver)
     @eval tableau = $tableau_getter()
@@ -343,7 +343,7 @@ end
 """
 `derivative_method` kwarg refers to the method used by InfiniteOpt to compute derivatives. The list of possible options can be found at https://infiniteopt.github.io/InfiniteOpt.jl/stable/guide/derivative/. Defaults to FiniteDifference(Backward()).
 """
-function DiffEqBase.solve(prob::InfiniteOptControlProblem, jump_solver;
+function DiffEqBase.solve(prob::InfiniteOptDynamicOptProblem, jump_solver;
         derivative_method = InfiniteOpt.FiniteDifference(Backward()), silent = false)
     model = prob.model
     silent && set_silent(model)
@@ -351,7 +351,7 @@ function DiffEqBase.solve(prob::InfiniteOptControlProblem, jump_solver;
     _solve(prob, jump_solver, derivative_method)
 end
 
-function _solve(prob::AbstractOptimalControlProblem, jump_solver, solver)
+function _solve(prob::AbstractDynamicOptProblem, jump_solver, solver)
     model = prob.model
     set_optimizer(model, jump_solver)
     optimize!(model)
@@ -382,7 +382,7 @@ function _solve(prob::AbstractOptimalControlProblem, jump_solver, solver)
             input_sol, SciMLBase.ReturnCode.ConvergenceFailure))
     end
 
-    OptimalControlSolution(model, sol, input_sol)
+    DynamicOptSolution(model, sol, input_sol)
 end
 
 end
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index bb6d991578..ff75b520f2 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -348,9 +348,9 @@ export AnalysisPoint, get_sensitivity_function, get_comp_sensitivity_function,
 function FMIComponent end
 
 include("systems/optimal_control_interface.jl")
-export AbstractOptimalControlProblem, JuMPControlProblem, InfiniteOptControlProblem,
-       PyomoControlProblem, CasADiControlProblem
-export OptimalControlSolution
+export AbstractDynamicOptProblem, JuMPDynamicOptProblem, InfiniteOptDynamicOptProblem,
+       PyomoDynamicOptProblem, CasADiDynamicOptProblem
+export DynamicOptSolution 
 export ∫
 
 end # module
diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index fbf04c1b99..1183d04db6 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -1,13 +1,13 @@
-abstract type AbstractOptimalControlProblem{uType, tType, isinplace} <:
+abstract type AbstractDynamicOptProblem{uType, tType, isinplace} <:
               SciMLBase.AbstractODEProblem{uType, tType, isinplace} end
 
-struct OptimalControlSolution
+struct DynamicOptSolution
     model::Any
     sol::ODESolution
     input_sol::Union{Nothing, ODESolution}
 end
 
-function Base.show(io::IO, sol::OptimalControlSolution)
+function Base.show(io::IO, sol::DynamicOptSolution)
     println("retcode: ", sol.sol.retcode, "\n")
 
     println("Optimal control solution for following model:\n")
@@ -16,10 +16,10 @@ function Base.show(io::IO, sol::OptimalControlSolution)
     print("\n\nPlease query the model using sol.model, the solution trajectory for the system using sol.sol, or the solution trajectory for the controllers using sol.input_sol.")
 end
 
-function JuMPControlProblem end
-function InfiniteOptControlProblem end
-function CasADiControlProblem end
-function PyomoControlProblem end
+function JuMPDynamicOptProblem end
+function InfiniteOptDynamicOptProblem end
+function CasADiDynamicOptProblem end
+function PyomoDynamicOptProblem end
 
 function warn_overdetermined(sys, u0map)
     constraintsys = get_constraintsystem(sys)
diff --git a/test/downstream/Project.toml b/test/downstream/Project.toml
index ade09e797b..aa58095744 100644
--- a/test/downstream/Project.toml
+++ b/test/downstream/Project.toml
@@ -4,6 +4,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 
 [compat]
diff --git a/test/dynamic_optimization/jump_control.jl b/test/dynamic_optimization/jump_control.jl
index 587afa8ecf..431491ddcd 100644
--- a/test/dynamic_optimization/jump_control.jl
+++ b/test/dynamic_optimization/jump_control.jl
@@ -23,7 +23,7 @@ using DataInterpolations
     parammap = [α => 1.5, β => 1.0, γ => 3.0, δ => 1.0]
 
     # Test explicit method.
-    jprob = JuMPControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
+    jprob = JuMPDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01)
     @test JuMP.num_constraints(jprob.model) == 2 # initials
     jsol = solve(jprob, Ipopt.Optimizer, :RK4)
     oprob = ODEProblem(sys, u0map, tspan, parammap)
@@ -34,7 +34,7 @@ using DataInterpolations
     jsol2 = solve(jprob, Ipopt.Optimizer, :ImplicitEuler, silent = true) # 63.031 ms, 26.49 MiB
     osol2 = solve(oprob, ImplicitEuler(), dt = 0.01, adaptive = false) # 129.375 μs, 61.91 KiB
     @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001)
-    iprob = InfiniteOptControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
+    iprob = InfiniteOptDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer,
         derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward()),
         silent = true) # 11.540 ms, 4.00 MiB
@@ -46,13 +46,13 @@ using DataInterpolations
     constr = [x(0.6) ~ 3.5, x(0.3) ~ 7.0]
     @mtkbuild lksys = ODESystem(eqs, t; constraints = constr)
 
-    jprob = JuMPControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
+    jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     @test JuMP.num_constraints(jprob.model) == 2
     jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5, silent = true) # 12.190 s, 9.68 GiB
     @test jsol.sol(0.6)[1] ≈ 3.5
     @test jsol.sol(0.3)[1] ≈ 7.0
 
-    iprob = InfiniteOptControlProblem(
+    iprob = InfiniteOptDynamicOptProblem(
         lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer,
         derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
@@ -63,13 +63,13 @@ using DataInterpolations
     # Test whole-interval constraints
     constr = [x(t) ≳ 1, y(t) ≳ 1]
     @mtkbuild lksys = ODESystem(eqs, t; constraints = constr)
-    iprob = InfiniteOptControlProblem(
+    iprob = InfiniteOptDynamicOptProblem(
         lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer,
         derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
     @test all(u -> u > [1, 1], isol.sol.u)
 
-    jprob = JuMPControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
+    jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3, silent = true) # 12.190 s, 9.68 GiB
     @test all(u -> u > [1, 1], jsol.sol.u)
 end
@@ -103,7 +103,7 @@ end
     u0map = [x(t) => 0.0, v(t) => 0.0]
     tspan = (0.0, 1.0)
     parammap = [u(t) => 0.0]
-    jprob = JuMPControlProblem(block, u0map, tspan, parammap; dt = 0.01)
+    jprob = JuMPDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01)
     jsol = solve(jprob, Ipopt.Optimizer, :Verner8)
     # Linear systems have bang-bang controls
     @test is_bangbang(jsol.input_sol, [-1.0], [1.0])
@@ -117,7 +117,7 @@ end
     osol = solve(oprob, Vern8(), dt = 0.01, adaptive = false)
     @test ≈(jsol.sol.u, osol.u, rtol = 0.05)
 
-    iprob = InfiniteOptControlProblem(block, u0map, tspan, parammap; dt = 0.01)
+    iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer; silent = true)
     @test is_bangbang(isol.input_sol, [-1.0], [1.0])
     @test ≈(isol.sol.u[end][1], 0.25, rtol = 1e-5)
@@ -141,10 +141,10 @@ end
     u0map = [w(t) => 40, q(t) => 2]
     pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1, α => 1]
 
-    jprob = JuMPControlProblem(beesys, u0map, tspan, pmap, dt = 0.01)
+    jprob = JuMPDynamicOptProblem(beesys, u0map, tspan, pmap, dt = 0.01)
     jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
     @test is_bangbang(jsol.input_sol, [0.0], [1.0])
-    iprob = InfiniteOptControlProblem(beesys, u0map, tspan, pmap, dt = 0.01)
+    iprob = InfiniteOptDynamicOptProblem(beesys, u0map, tspan, pmap, dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer; silent = true)
     @test is_bangbang(isol.input_sol, [0.0], [1.0])
 
@@ -186,11 +186,11 @@ end
     pmap = [
         g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5 * √(g₀ * h₀), D_c => 0.5 * 620 * m₀ / g₀,
         Tₘ => 3.5 * g₀ * m₀, T(t) => 0.0, h₀ => 1, m_c => 0.6]
-    jprob = JuMPControlProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
+    jprob = JuMPDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
     jsol = solve(jprob, Ipopt.Optimizer, :RadauIIA5, silent = true)
     @test jsol.sol.u[end][1] > 1.012
 
-    iprob = InfiniteOptControlProblem(
+    iprob = InfiniteOptDynamicOptProblem(
         rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
     isol = solve(iprob, Ipopt.Optimizer,
         derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true)
@@ -223,11 +223,11 @@ end
 
     u0map = [x(t) => 17.5]
     pmap = [u(t) => 0.0, tf => 8]
-    jprob = JuMPControlProblem(rocket, u0map, (0, tf), pmap; steps = 201)
+    jprob = JuMPDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 201)
     jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
     @test isapprox(jsol.sol.t[end], 10.0, rtol = 1e-3)
 
-    iprob = InfiniteOptControlProblem(rocket, u0map, (0, tf), pmap; steps = 200)
+    iprob = InfiniteOptDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 200)
     isol = solve(iprob, Ipopt.Optimizer)
     @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3)
 end
@@ -259,11 +259,11 @@ end
 
     u0map = [D(x(t)) => 0.0, D(θ(t)) => 0.0, θ(t) => 0.0, x(t) => 0.0]
     pmap = [mₖ => 1.0, mₚ => 0.2, l => 0.5, g => 9.81, u => 0]
-    jprob = JuMPControlProblem(cartpole, u0map, tspan, pmap; dt = 0.04)
+    jprob = JuMPDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04)
     jsol = solve(jprob, Ipopt.Optimizer, :RK4)
     @test jsol.sol.u[end] ≈ [π, 0, 0, 0]
 
-    iprob = InfiniteOptControlProblem(cartpole, u0map, tspan, pmap; dt = 0.04)
+    iprob = InfiniteOptDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04)
     isol = solve(iprob, Ipopt.Optimizer)
     @test isol.sol.u[end] ≈ [π, 0, 0, 0]
 end
diff --git a/test/extensions/ad.jl b/test/extensions/ad.jl
index a456263655..14649b6bb6 100644
--- a/test/extensions/ad.jl
+++ b/test/extensions/ad.jl
@@ -4,6 +4,7 @@ using Zygote
 using SymbolicIndexingInterface
 using SciMLStructures
 using OrdinaryDiffEqTsit5
+using OrdinaryDiffEqNonlinearSolve
 using NonlinearSolve
 using SciMLSensitivity
 using ForwardDiff
diff --git a/test/runtests.jl b/test/runtests.jl
index 0e0ea29a2c..5f27f8bb88 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -29,102 +29,102 @@ function activate_dynamic_optimization_env()
 end
 
 @time begin
-    if GROUP == "All" || GROUP == "InterfaceI"
-        @testset "InterfaceI" begin
-            @safetestset "Linear Algebra Test" include("linalg.jl")
-            @safetestset "AbstractSystem Test" include("abstractsystem.jl")
-            @safetestset "Variable Scope Tests" include("variable_scope.jl")
-            @safetestset "Symbolic Parameters Test" include("symbolic_parameters.jl")
-            @safetestset "Parsing Test" include("variable_parsing.jl")
-            @safetestset "Simplify Test" include("simplify.jl")
-            @safetestset "Direct Usage Test" include("direct.jl")
-            @safetestset "System Linearity Test" include("linearity.jl")
-            @safetestset "Input Output Test" include("input_output_handling.jl")
-            @safetestset "Clock Test" include("clock.jl")
-            @safetestset "ODESystem Test" include("odesystem.jl")
-            @safetestset "Dynamic Quantities Test" include("dq_units.jl")
-            @safetestset "Unitful Quantities Test" include("units.jl")
-            @safetestset "Mass Matrix Test" include("mass_matrix.jl")
-            @safetestset "Reduction Test" include("reduction.jl")
-            @safetestset "Split Parameters Test" include("split_parameters.jl")
-            @safetestset "StaticArrays Test" include("static_arrays.jl")
-            @safetestset "Components Test" include("components.jl")
-            @safetestset "Model Parsing Test" include("model_parsing.jl")
-            @safetestset "Error Handling" include("error_handling.jl")
-            @safetestset "StructuralTransformations" include("structural_transformation/runtests.jl")
-            @safetestset "Basic transformations" include("basic_transformations.jl")
-            @safetestset "State Selection Test" include("state_selection.jl")
-            @safetestset "Symbolic Event Test" include("symbolic_events.jl")
-            @safetestset "Stream Connect Test" include("stream_connectors.jl")
-            @safetestset "Domain Connect Test" include("domain_connectors.jl")
-            @safetestset "Lowering Integration Test" include("lowering_solving.jl")
-            @safetestset "Dependency Graph Test" include("dep_graphs.jl")
-            @safetestset "Function Registration Test" include("function_registration.jl")
-            @safetestset "Precompiled Modules Test" include("precompile_test.jl")
-            @safetestset "DAE Jacobians Test" include("dae_jacobian.jl")
-            @safetestset "Jacobian Sparsity" include("jacobiansparsity.jl")
-            @safetestset "Modelingtoolkitize Test" include("modelingtoolkitize.jl")
-            @safetestset "FuncAffect Test" include("funcaffect.jl")
-            @safetestset "Constants Test" include("constants.jl")
-            @safetestset "Parameter Dependency Test" include("parameter_dependencies.jl")
-            @safetestset "Equation Type Accessors Test" include("equation_type_accessors.jl")
-            @safetestset "System Accessor Functions Test" include("accessor_functions.jl")
-            @safetestset "Equations with complex values" include("complex.jl")
-        end
-    end
+    #if GROUP == "All" || GROUP == "InterfaceI"
+    #    @testset "InterfaceI" begin
+    #        @safetestset "Linear Algebra Test" include("linalg.jl")
+    #        @safetestset "AbstractSystem Test" include("abstractsystem.jl")
+    #        @safetestset "Variable Scope Tests" include("variable_scope.jl")
+    #        @safetestset "Symbolic Parameters Test" include("symbolic_parameters.jl")
+    #        @safetestset "Parsing Test" include("variable_parsing.jl")
+    #        @safetestset "Simplify Test" include("simplify.jl")
+    #        @safetestset "Direct Usage Test" include("direct.jl")
+    #        @safetestset "System Linearity Test" include("linearity.jl")
+    #        @safetestset "Input Output Test" include("input_output_handling.jl")
+    #        @safetestset "Clock Test" include("clock.jl")
+    #        @safetestset "ODESystem Test" include("odesystem.jl")
+    #        @safetestset "Dynamic Quantities Test" include("dq_units.jl")
+    #        @safetestset "Unitful Quantities Test" include("units.jl")
+    #        @safetestset "Mass Matrix Test" include("mass_matrix.jl")
+    #        @safetestset "Reduction Test" include("reduction.jl")
+    #        @safetestset "Split Parameters Test" include("split_parameters.jl")
+    #        @safetestset "StaticArrays Test" include("static_arrays.jl")
+    #        @safetestset "Components Test" include("components.jl")
+    #        @safetestset "Model Parsing Test" include("model_parsing.jl")
+    #        @safetestset "Error Handling" include("error_handling.jl")
+    #        @safetestset "StructuralTransformations" include("structural_transformation/runtests.jl")
+    #        @safetestset "Basic transformations" include("basic_transformations.jl")
+    #        @safetestset "State Selection Test" include("state_selection.jl")
+    #        @safetestset "Symbolic Event Test" include("symbolic_events.jl")
+    #        @safetestset "Stream Connect Test" include("stream_connectors.jl")
+    #        @safetestset "Domain Connect Test" include("domain_connectors.jl")
+    #        @safetestset "Lowering Integration Test" include("lowering_solving.jl")
+    #        @safetestset "Dependency Graph Test" include("dep_graphs.jl")
+    #        @safetestset "Function Registration Test" include("function_registration.jl")
+    #        @safetestset "Precompiled Modules Test" include("precompile_test.jl")
+    #        @safetestset "DAE Jacobians Test" include("dae_jacobian.jl")
+    #        @safetestset "Jacobian Sparsity" include("jacobiansparsity.jl")
+    #        @safetestset "Modelingtoolkitize Test" include("modelingtoolkitize.jl")
+    #        @safetestset "FuncAffect Test" include("funcaffect.jl")
+    #        @safetestset "Constants Test" include("constants.jl")
+    #        @safetestset "Parameter Dependency Test" include("parameter_dependencies.jl")
+    #        @safetestset "Equation Type Accessors Test" include("equation_type_accessors.jl")
+    #        @safetestset "System Accessor Functions Test" include("accessor_functions.jl")
+    #        @safetestset "Equations with complex values" include("complex.jl")
+    #    end
+    #end
 
-    if GROUP == "All" || GROUP == "Initialization"
-        @safetestset "Guess Propagation" include("guess_propagation.jl")
-        @safetestset "Hierarchical Initialization Equations" include("hierarchical_initialization_eqs.jl")
-        @safetestset "InitializationSystem Test" include("initializationsystem.jl")
-        @safetestset "Initial Values Test" include("initial_values.jl")
-    end
+    #if GROUP == "All" || GROUP == "Initialization"
+    #    @safetestset "Guess Propagation" include("guess_propagation.jl")
+    #    @safetestset "Hierarchical Initialization Equations" include("hierarchical_initialization_eqs.jl")
+    #    @safetestset "InitializationSystem Test" include("initializationsystem.jl")
+    #    @safetestset "Initial Values Test" include("initial_values.jl")
+    #end
 
-    if GROUP == "All" || GROUP == "InterfaceII"
-        @testset "InterfaceII" begin
-            @safetestset "Code Generation Test" include("code_generation.jl")
-            @safetestset "IndexCache Test" include("index_cache.jl")
-            @safetestset "Variable Utils Test" include("variable_utils.jl")
-            @safetestset "Variable Metadata Test" include("test_variable_metadata.jl")
-            @safetestset "OptimizationSystem Test" include("optimizationsystem.jl")
-            @safetestset "Discrete System" include("discrete_system.jl")
-            @safetestset "Implicit Discrete System" include("implicit_discrete_system.jl")
-            @safetestset "SteadyStateSystem Test" include("steadystatesystems.jl")
-            @safetestset "SDESystem Test" include("sdesystem.jl")
-            @safetestset "DDESystem Test" include("dde.jl")
-            @safetestset "NonlinearSystem Test" include("nonlinearsystem.jl")
-            @safetestset "SCCNonlinearProblem Test" include("scc_nonlinear_problem.jl")
-            @safetestset "PDE Construction Test" include("pdesystem.jl")
-            @safetestset "JumpSystem Test" include("jumpsystem.jl")
-            @safetestset "Optimal Control + Constraints Tests" include("optimal_control.jl")
-            @safetestset "print_tree" include("print_tree.jl")
-            @safetestset "Constraints Test" include("constraints.jl")
-            @safetestset "IfLifting Test" include("if_lifting.jl")
-            @safetestset "Analysis Points Test" include("analysis_points.jl")
-            @safetestset "Causal Variables Connection Test" include("causal_variables_connection.jl")
-            @safetestset "Debugging Test" include("debugging.jl")
-            @safetestset "Namespacing test" include("namespacing.jl")
-            @safetestset "Subsystem replacement" include("substitute_component.jl")
-        end
-    end
+    #if GROUP == "All" || GROUP == "InterfaceII"
+    #    @testset "InterfaceII" begin
+    #        @safetestset "Code Generation Test" include("code_generation.jl")
+    #        @safetestset "IndexCache Test" include("index_cache.jl")
+    #        @safetestset "Variable Utils Test" include("variable_utils.jl")
+    #        @safetestset "Variable Metadata Test" include("test_variable_metadata.jl")
+    #        @safetestset "OptimizationSystem Test" include("optimizationsystem.jl")
+    #        @safetestset "Discrete System" include("discrete_system.jl")
+    #        @safetestset "Implicit Discrete System" include("implicit_discrete_system.jl")
+    #        @safetestset "SteadyStateSystem Test" include("steadystatesystems.jl")
+    #        @safetestset "SDESystem Test" include("sdesystem.jl")
+    #        @safetestset "DDESystem Test" include("dde.jl")
+    #        @safetestset "NonlinearSystem Test" include("nonlinearsystem.jl")
+    #        @safetestset "SCCNonlinearProblem Test" include("scc_nonlinear_problem.jl")
+    #        @safetestset "PDE Construction Test" include("pdesystem.jl")
+    #        @safetestset "JumpSystem Test" include("jumpsystem.jl")
+    #        @safetestset "Optimal Control + Constraints Tests" include("optimal_control.jl")
+    #        @safetestset "print_tree" include("print_tree.jl")
+    #        @safetestset "Constraints Test" include("constraints.jl")
+    #        @safetestset "IfLifting Test" include("if_lifting.jl")
+    #        @safetestset "Analysis Points Test" include("analysis_points.jl")
+    #        @safetestset "Causal Variables Connection Test" include("causal_variables_connection.jl")
+    #        @safetestset "Debugging Test" include("debugging.jl")
+    #        @safetestset "Namespacing test" include("namespacing.jl")
+    #        @safetestset "Subsystem replacement" include("substitute_component.jl")
+    #    end
+    #end
 
-    if GROUP == "All" || GROUP == "SymbolicIndexingInterface"
-        @safetestset "SymbolicIndexingInterface test" include("symbolic_indexing_interface.jl")
-        @safetestset "SciML Problem Input Test" include("sciml_problem_inputs.jl")
-        @safetestset "MTKParameters Test" include("mtkparameters.jl")
-    end
+    #if GROUP == "All" || GROUP == "SymbolicIndexingInterface"
+    #    @safetestset "SymbolicIndexingInterface test" include("symbolic_indexing_interface.jl")
+    #    @safetestset "SciML Problem Input Test" include("sciml_problem_inputs.jl")
+    #    @safetestset "MTKParameters Test" include("mtkparameters.jl")
+    #end
 
-    if GROUP == "All" || GROUP == "Extended"
-        @safetestset "Test Big System Usage" include("bigsystem.jl")
-        println("C compilation test requires gcc available in the path!")
-        @safetestset "C Compilation Test" include("ccompile.jl")
-        @testset "Distributed Test" include("distributed.jl")
-        @testset "Serialization" include("serialization.jl")
-    end
+    #if GROUP == "All" || GROUP == "Extended"
+    #    @safetestset "Test Big System Usage" include("bigsystem.jl")
+    #    println("C compilation test requires gcc available in the path!")
+    #    @safetestset "C Compilation Test" include("ccompile.jl")
+    #    @testset "Distributed Test" include("distributed.jl")
+    #    @testset "Serialization" include("serialization.jl")
+    #end
 
-    if GROUP == "All" || GROUP == "RegressionI"
-        @safetestset "Latexify recipes Test" include("latexify.jl")
-    end
+    #if GROUP == "All" || GROUP == "RegressionI"
+    #    @safetestset "Latexify recipes Test" include("latexify.jl")
+    #end
 
     if GROUP == "All" || GROUP == "Downstream"
         activate_downstream_env()

From 9766abae8f3772bb53c177e74b2ace66ccfb51ce Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 25 Apr 2025 23:06:54 -0400
Subject: [PATCH 32/41] test fixes

---
 src/inputoutput.jl                        |  4 ++--
 src/structural_transformation/utils.jl    |  1 +
 src/systems/diffeqs/odesystem.jl          |  6 +++---
 test/downstream/test_disturbance_model.jl | 12 ++++++------
 test/extensions/Project.toml              |  1 +
 test/input_output_handling.jl             | 20 ++++++++++----------
 test/odesystem.jl                         | 12 +++++-------
 7 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/src/inputoutput.jl b/src/inputoutput.jl
index 11a58bd30a..8c1a084c9b 100644
--- a/src/inputoutput.jl
+++ b/src/inputoutput.jl
@@ -430,7 +430,7 @@ function add_input_disturbance(sys, dist::DisturbanceModel, inputs = nothing; kw
     augmented_sys = ODESystem(eqs, t, systems = [dsys], name = gensym(:outer))
     augmented_sys = extend(augmented_sys, sys)
 
-    (f_oop, f_ip), dvs, p, io_sys = generate_control_function(augmented_sys, all_inputs,
+    f, dvs, p, io_sys = generate_control_function(augmented_sys, all_inputs,
         [d]; kwargs...)
-    (f_oop, f_ip), augmented_sys, dvs, p, io_sys
+    f, augmented_sys, dvs, p, io_sys
 end
diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 14628f2958..15a46531d9 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -96,6 +96,7 @@ function check_consistency(state::TransformationState, orig_inputs; nothrow = fa
     fullvars = get_fullvars(state)
     neqs = n_concrete_eqs(state)
     @unpack graph, var_to_diff = state.structure
+    @show equations(state.sys)
     highest_vars = computed_highest_diff_variables(complete!(state.structure))
     n_highest_vars = 0
     for (v, h) in enumerate(highest_vars)
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index 62efcc3f3f..5d1994fc12 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -803,12 +803,12 @@ Return the set of additional parameters found in the system, e.g. in x(p) ~ 3 th
 parameter of the system.
 """
 function validate_vars_and_find_ps!(auxvars, auxps, sysvars, iv)
-    sts = sysvars
+    sts = Set(sysvars)
+    @show sts
 
     for var in auxvars
         if !iscall(var)
-            occursin(iv, var) && (var ∈ sts ||
-             throw(ArgumentError("Time-dependent variable $var is not an unknown of the system.")))
+            var ∈ sts || throw(ArgumentError("Time-independent variable $var is not an unknown of the system."))
         elseif length(arguments(var)) > 1
             throw(ArgumentError("Too many arguments for variable $var."))
         elseif length(arguments(var)) == 1
diff --git a/test/downstream/test_disturbance_model.jl b/test/downstream/test_disturbance_model.jl
index 97276437e2..c96830efa3 100644
--- a/test/downstream/test_disturbance_model.jl
+++ b/test/downstream/test_disturbance_model.jl
@@ -149,10 +149,10 @@ sol = solve(prob, Tsit5())
 ## Generate function for an augmented Unscented Kalman Filter =====================
 # temp = open_loop(model_with_disturbance, :d)
 outputs = [P.inertia1.phi, P.inertia2.phi, P.inertia1.w, P.inertia2.w]
-(f_oop1, f_ip), x_sym, p_sym, io_sys = ModelingToolkit.generate_control_function(
+f, x_sym, p_sym, io_sys = ModelingToolkit.generate_control_function(
     model_with_disturbance, [:u], [:d1, :d2, :dy], split = false)
 
-(f_oop2, f_ip2), x_sym, p_sym, io_sys = ModelingToolkit.generate_control_function(
+f, x_sym, p_sym, io_sys = ModelingToolkit.generate_control_function(
     model_with_disturbance, [:u], [:d1, :d2, :dy],
     disturbance_argument = true, split = false)
 
@@ -168,22 +168,22 @@ x0, p = ModelingToolkit.get_u0_p(io_sys, op, op)
 x = zeros(5)
 u = zeros(1)
 d = zeros(3)
-@test f_oop2(x, u, p, t, d) == zeros(5)
+@test f(x, u, p, t, d) == zeros(5)
 @test measurement(x, u, p, 0.0) == [0, 0, 0, 0]
 @test measurement2(x, u, p, 0.0, d) == [0]
 
 # Add to the integrating disturbance input
 d = [1, 0, 0]
-@test sort(f_oop2(x, u, p, 0.0, d)) == [0, 0, 0, 1, 1] # Affects disturbance state and one velocity
+@test sort(f(x, u, p, 0.0, d)) == [0, 0, 0, 1, 1] # Affects disturbance state and one velocity
 @test measurement2(x, u, p, 0.0, d) == [0]
 
 d = [0, 1, 0]
-@test sort(f_oop2(x, u, p, 0.0, d)) == [0, 0, 0, 0, 1] # Affects one velocity
+@test sort(f(x, u, p, 0.0, d)) == [0, 0, 0, 0, 1] # Affects one velocity
 @test measurement(x, u, p, 0.0) == [0, 0, 0, 0]
 @test measurement2(x, u, p, 0.0, d) == [0]
 
 d = [0, 0, 1]
-@test sort(f_oop2(x, u, p, 0.0, d)) == [0, 0, 0, 0, 0] # Affects nothing
+@test sort(f(x, u, p, 0.0, d)) == [0, 0, 0, 0, 0] # Affects nothing
 @test measurement(x, u, p, 0.0) == [0, 0, 0, 0]
 @test measurement2(x, u, p, 0.0, d) == [1] # We have now disturbed the output
 
diff --git a/test/extensions/Project.toml b/test/extensions/Project.toml
index 03e65b4978..9e58e972d0 100644
--- a/test/extensions/Project.toml
+++ b/test/extensions/Project.toml
@@ -9,6 +9,7 @@ LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Nemo = "2edaba10-b0f1-5616-af89-8c11ac63239a"
 NonlinearSolveHomotopyContinuation = "2ac3b008-d579-4536-8c91-a1a5998c2f8b"
+OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
diff --git a/test/input_output_handling.jl b/test/input_output_handling.jl
index 693a00b9ad..2c16c89320 100644
--- a/test/input_output_handling.jl
+++ b/test/input_output_handling.jl
@@ -173,7 +173,7 @@ end
         p = [rand()]
         x = [rand()]
         u = [rand()]
-        @test f[1](x, u, p, 1) ≈ -x + u
+        @test f(x, u, p, 1) ≈ -x + u
 
         # With disturbance inputs
         @variables x(t)=0 u(t)=0 [input = true] d(t)=0
@@ -191,7 +191,7 @@ end
         p = [rand()]
         x = [rand()]
         u = [rand()]
-        @test f[1](x, u, p, 1) ≈ -x + u
+        @test f(x, u, p, 1) ≈ -x + u
 
         ## With added d argument
         @variables x(t)=0 u(t)=0 [input = true] d(t)=0
@@ -210,7 +210,7 @@ end
         x = [rand()]
         u = [rand()]
         d = [rand()]
-        @test f[1](x, u, p, t, d) ≈ -x + u + [d[]^2]
+        @test f(x, u, p, t, d) ≈ -x + u + [d[]^2]
     end
 end
 
@@ -273,7 +273,7 @@ x = ModelingToolkit.varmap_to_vars(
     merge(ModelingToolkit.defaults(model),
         Dict(D.(unknowns(model)) .=> 0.0)), dvs)
 u = [rand()]
-out = f[1](x, u, p, 1)
+out = f(x, u, p, 1)
 i = findfirst(isequal(u[1]), out)
 @test i isa Int
 @test iszero(out[[1:(i - 1); (i + 1):end]])
@@ -333,7 +333,7 @@ model_outputs = [model.inertia1.w, model.inertia2.w, model.inertia1.phi, model.i
 @named dmodel = Blocks.StateSpace([0.0], [1.0], [1.0], [0.0]) # An integrating disturbance
 
 @named dist = ModelingToolkit.DisturbanceModel(model.torque.tau.u, dmodel)
-(f_oop, f_ip), outersys, dvs, p, io_sys = ModelingToolkit.add_input_disturbance(model, dist)
+f, outersys, dvs, p, io_sys = ModelingToolkit.add_input_disturbance(model, dist)
 
 @unpack u, d = outersys
 matrices, ssys = linearize(outersys, [u, d], model_outputs)
@@ -348,8 +348,8 @@ x0 = randn(5)
 x1 = copy(x0) + x_add # add disturbance state perturbation
 u = randn(1)
 pn = MTKParameters(io_sys, [])
-xp0 = f_oop(x0, u, pn, 0)
-xp1 = f_oop(x1, u, pn, 0)
+xp0 = f(x0, u, pn, 0)
+xp1 = f(x1, u, pn, 0)
 
 @test xp0 ≈ matrices.A * x0 + matrices.B * [u; 0]
 @test xp1 ≈ matrices.A * x1 + matrices.B * [u; 0]
@@ -402,7 +402,7 @@ outs = collect(complete(model).output.u)
 disturbed_input = ins[1]
 @named dist_integ = DisturbanceModel(disturbed_input, integrator)
 
-(f_oop, f_ip), augmented_sys, dvs, p = ModelingToolkit.add_input_disturbance(model,
+f, augmented_sys, dvs, p = ModelingToolkit.add_input_disturbance(model,
     dist_integ,
     ins)
 
@@ -447,7 +447,7 @@ end
     @named sys = ODESystem(eqs, t, [x], [])
 
     f, dvs, ps, io_sys = ModelingToolkit.generate_control_function(sys, simplify = true)
-    @test f[1]([0.5], nothing, MTKParameters(io_sys, []), 0.0) ≈ [1.0]
+    @test f([0.5], nothing, MTKParameters(io_sys, []), 0.0) ≈ [1.0]
 end
 
 @testset "With callable symbolic" begin
@@ -459,5 +459,5 @@ end
     p = MTKParameters(io_sys, [])
     u = [1.0]
     x = [1.0]
-    @test_nowarn f[1](x, u, p, 0.0)
+    @test_nowarn f(x, u, p, 0.0)
 end
diff --git a/test/odesystem.jl b/test/odesystem.jl
index afb9e6e440..9219fca5fb 100644
--- a/test/odesystem.jl
+++ b/test/odesystem.jl
@@ -472,14 +472,12 @@ sys = complete(sys)
 
 # check inputs
 let
-    @parameters f k d
-    @variables x(t) ẋ(t)
-    δ = D
+    @parameters k d
+    @variables x(t) ẋ(t) f(t) [input = true]
 
-    eqs = [δ(x) ~ ẋ, δ(ẋ) ~ f - k * x - d * ẋ]
-    @named sys = ODESystem(eqs, t, [x, ẋ], [f, d, k]; controls = [f])
-
-    calculate_control_jacobian(sys)
+    eqs = [D(x) ~ ẋ, D(ẋ) ~ f - k * x - d * ẋ]
+    @named sys = ODESystem(eqs, t, [x, ẋ], [d, k])
+    sys, _ = structural_simplify(sys, ([f], []))
 
     @test isequal(calculate_control_jacobian(sys),
         reshape(Num[0, 1], 2, 1))

From 86ace9e8b54b9dadb74fd593f94105850f259686 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Fri, 25 Apr 2025 23:07:30 -0400
Subject: [PATCH 33/41] format

---
 src/ModelingToolkit.jl           | 2 +-
 src/systems/diffeqs/odesystem.jl | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
index ff75b520f2..6f91730582 100644
--- a/src/ModelingToolkit.jl
+++ b/src/ModelingToolkit.jl
@@ -350,7 +350,7 @@ function FMIComponent end
 include("systems/optimal_control_interface.jl")
 export AbstractDynamicOptProblem, JuMPDynamicOptProblem, InfiniteOptDynamicOptProblem,
        PyomoDynamicOptProblem, CasADiDynamicOptProblem
-export DynamicOptSolution 
+export DynamicOptSolution
 export ∫
 
 end # module
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index 5d1994fc12..9264e77235 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -808,7 +808,8 @@ function validate_vars_and_find_ps!(auxvars, auxps, sysvars, iv)
 
     for var in auxvars
         if !iscall(var)
-            var ∈ sts || throw(ArgumentError("Time-independent variable $var is not an unknown of the system."))
+            var ∈ sts ||
+                throw(ArgumentError("Time-independent variable $var is not an unknown of the system."))
         elseif length(arguments(var)) > 1
             throw(ArgumentError("Too many arguments for variable $var."))
         elseif length(arguments(var)) == 1

From 9af446c665c4a34ae61a68d95f313b0a19741553 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 28 Apr 2025 13:58:41 -0400
Subject: [PATCH 34/41] add JuMP to extensions Project.toml

---
 test/extensions/Project.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/extensions/Project.toml b/test/extensions/Project.toml
index 9e58e972d0..f5cedf49fe 100644
--- a/test/extensions/Project.toml
+++ b/test/extensions/Project.toml
@@ -5,6 +5,7 @@ ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Nemo = "2edaba10-b0f1-5616-af89-8c11ac63239a"

From 8400ed387d396f94363028219a0fbd5edf309f9d Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 28 Apr 2025 14:28:52 -0400
Subject: [PATCH 35/41] fix tests

---
 test/extensions/test_infiniteopt.jl | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/test/extensions/test_infiniteopt.jl b/test/extensions/test_infiniteopt.jl
index e45aa0f2fd..833b9f3275 100644
--- a/test/extensions/test_infiniteopt.jl
+++ b/test/extensions/test_infiniteopt.jl
@@ -22,13 +22,14 @@ using ModelingToolkit: D_nounits as D, t_nounits as t, varmap_to_vars
 end
 @named model = Pendulum()
 model = complete(model)
-
 inputs = [model.τ]
-(f_oop, f_ip), dvs, psym, io_sys = ModelingToolkit.generate_control_function(
-    model, inputs, split = false)
-
 outputs = [model.y]
-f_obs = ModelingToolkit.build_explicit_observed_function(io_sys, outputs; inputs = inputs)
+model, _ = structural_simplify(model, (inputs, outputs))
+
+f, dvs, psym, io_sys = ModelingToolkit.generate_control_function(
+    model, split = false)
+
+f_obs = ModelingToolkit.build_explicit_observed_function(io_sys, outputs; inputs)
 
 expected_state_order = [model.θ, model.ω]
 permutation = [findfirst(isequal(x), expected_state_order) for x in dvs] # This maps our expected state order to the actual state order
@@ -64,7 +65,7 @@ InfiniteOpt.@variables(m,
 # Trace the dynamics
 x0, p = ModelingToolkit.get_u0_p(io_sys, [model.θ => 0, model.ω => 0], [model.L => L])
 
-xp = f_oop(x, u, p, τ)
+xp = f(x, u, p, τ)
 cp = f_obs(x, u, p, τ) # Test that it's possible to trace through an observed function
 
 @objective(m, Min, tf)

From 887c3e914e9bb38751c695b4644d909a9f5ce52a Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 28 Apr 2025 14:41:14 -0400
Subject: [PATCH 36/41] move jump control to downsteram

---
 test/downstream/Project.toml                       | 11 ++++++++++-
 test/downstream/analysis_points.jl                 |  2 +-
 test/downstream/inversemodel.jl                    |  2 +-
 .../jump_control.jl                                |  0
 test/downstream/linearize.jl                       |  2 +-
 test/downstream/test_disturbance_model.jl          |  2 +-
 test/dynamic_optimization/Project.toml             | 14 --------------
 7 files changed, 14 insertions(+), 19 deletions(-)
 rename test/{dynamic_optimization => downstream}/jump_control.jl (100%)
 delete mode 100644 test/dynamic_optimization/Project.toml

diff --git a/test/downstream/Project.toml b/test/downstream/Project.toml
index aa58095744..f7f6885b84 100644
--- a/test/downstream/Project.toml
+++ b/test/downstream/Project.toml
@@ -1,10 +1,19 @@
 [deps]
 ControlSystemsMTK = "687d7614-c7e5-45fc-bfc3-9ee385575c88"
+DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
+DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
+InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
+Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqFIRK = "5960d6e9-dd7a-4743-88e7-cf307b64f125"
 OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
+OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce"
+OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
+OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2"
+SimpleDiffEq = "05bca326-078c-5bf0-a5bf-ce7c7982d7fd"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 
 [compat]
diff --git a/test/downstream/analysis_points.jl b/test/downstream/analysis_points.jl
index 29b9aad512..58c4a97a8d 100644
--- a/test/downstream/analysis_points.jl
+++ b/test/downstream/analysis_points.jl
@@ -1,4 +1,4 @@
-using ModelingToolkit, OrdinaryDiffEq, LinearAlgebra, ControlSystemsBase
+using ModelingToolkit, OrdinaryDiffEqRosenbrock, LinearAlgebra, ControlSystemsBase
 using ModelingToolkitStandardLibrary.Mechanical.Rotational
 using ModelingToolkitStandardLibrary.Blocks
 using ModelingToolkit: connect, t_nounits as t, D_nounits as D
diff --git a/test/downstream/inversemodel.jl b/test/downstream/inversemodel.jl
index 32d5ee87ec..7bed0bc1e8 100644
--- a/test/downstream/inversemodel.jl
+++ b/test/downstream/inversemodel.jl
@@ -1,7 +1,7 @@
 using ModelingToolkit
 using ModelingToolkitStandardLibrary
 using ModelingToolkitStandardLibrary.Blocks
-using OrdinaryDiffEq
+using OrdinaryDiffEqRosenbrock
 using SymbolicIndexingInterface
 using Test
 using ControlSystemsMTK: tf, ss, get_named_sensitivity, get_named_comp_sensitivity
diff --git a/test/dynamic_optimization/jump_control.jl b/test/downstream/jump_control.jl
similarity index 100%
rename from test/dynamic_optimization/jump_control.jl
rename to test/downstream/jump_control.jl
diff --git a/test/downstream/linearize.jl b/test/downstream/linearize.jl
index 16df29f834..d82bd696dd 100644
--- a/test/downstream/linearize.jl
+++ b/test/downstream/linearize.jl
@@ -206,7 +206,7 @@ lsys, ssys = linearize(sat, [u], [y]; op = Dict(u => 2))
 @test lsys.D[] == 0
 
 # Test case when unknowns in system do not have equations in initialization system
-using ModelingToolkit, OrdinaryDiffEq, LinearAlgebra
+using ModelingToolkit, LinearAlgebra
 using ModelingToolkitStandardLibrary.Mechanical.Rotational
 using ModelingToolkitStandardLibrary.Blocks: Add, Sine, PID, SecondOrder, Step, RealOutput
 using ModelingToolkit: connect
diff --git a/test/downstream/test_disturbance_model.jl b/test/downstream/test_disturbance_model.jl
index c96830efa3..cd9d769e99 100644
--- a/test/downstream/test_disturbance_model.jl
+++ b/test/downstream/test_disturbance_model.jl
@@ -3,7 +3,7 @@ This file implements and tests a typical workflow for state estimation with dist
 The primary subject of the tests is the analysis-point features and the
 analysis-point specific method for `generate_control_function`.
 =#
-using ModelingToolkit, OrdinaryDiffEq, LinearAlgebra, Test
+using ModelingToolkit, OrdinaryDiffEqTsit5, LinearAlgebra, Test
 using ModelingToolkitStandardLibrary.Mechanical.Rotational
 using ModelingToolkitStandardLibrary.Blocks
 using ModelingToolkit: connect
diff --git a/test/dynamic_optimization/Project.toml b/test/dynamic_optimization/Project.toml
deleted file mode 100644
index a5eaf3144b..0000000000
--- a/test/dynamic_optimization/Project.toml
+++ /dev/null
@@ -1,14 +0,0 @@
-[deps]
-BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
-CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
-DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
-DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
-InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
-Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
-ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
-ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
-OrdinaryDiffEqFIRK = "5960d6e9-dd7a-4743-88e7-cf307b64f125"
-OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
-OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
-OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2"
-SimpleDiffEq = "05bca326-078c-5bf0-a5bf-ce7c7982d7fd"

From 80d5d2b3401846cb962f55f1556339da0d5903ac Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 28 Apr 2025 15:20:43 -0400
Subject: [PATCH 37/41] fix more tests

---
 src/structural_transformation/utils.jl |  1 -
 src/systems/diffeqs/odesystem.jl       |  1 -
 test/downstream/Project.toml           |  1 +
 test/downstream/jump_control.jl        |  8 +++++---
 test/runtests.jl                       | 12 +-----------
 5 files changed, 7 insertions(+), 16 deletions(-)

diff --git a/src/structural_transformation/utils.jl b/src/structural_transformation/utils.jl
index 15a46531d9..14628f2958 100644
--- a/src/structural_transformation/utils.jl
+++ b/src/structural_transformation/utils.jl
@@ -96,7 +96,6 @@ function check_consistency(state::TransformationState, orig_inputs; nothrow = fa
     fullvars = get_fullvars(state)
     neqs = n_concrete_eqs(state)
     @unpack graph, var_to_diff = state.structure
-    @show equations(state.sys)
     highest_vars = computed_highest_diff_variables(complete!(state.structure))
     n_highest_vars = 0
     for (v, h) in enumerate(highest_vars)
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index 9264e77235..77d8ed589d 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -804,7 +804,6 @@ parameter of the system.
 """
 function validate_vars_and_find_ps!(auxvars, auxps, sysvars, iv)
     sts = Set(sysvars)
-    @show sts
 
     for var in auxvars
         if !iscall(var)
diff --git a/test/downstream/Project.toml b/test/downstream/Project.toml
index f7f6885b84..1192dda074 100644
--- a/test/downstream/Project.toml
+++ b/test/downstream/Project.toml
@@ -4,6 +4,7 @@ DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
diff --git a/test/downstream/jump_control.jl b/test/downstream/jump_control.jl
index 431491ddcd..303b18bbeb 100644
--- a/test/downstream/jump_control.jl
+++ b/test/downstream/jump_control.jl
@@ -5,7 +5,7 @@ using SimpleDiffEq
 using OrdinaryDiffEqSDIRK, OrdinaryDiffEqVerner, OrdinaryDiffEqTsit5, OrdinaryDiffEqFIRK
 using Ipopt
 using DataInterpolations
-#const M = ModelingToolkit
+const M = ModelingToolkit
 
 @testset "ODE Solution, no cost" begin
     # Test solving without anything attached.
@@ -108,7 +108,7 @@ end
     # Linear systems have bang-bang controls
     @test is_bangbang(jsol.input_sol, [-1.0], [1.0])
     # Test reached final position.
-    @test ≈(jsol.sol.u[end][1], 0.25, rtol = 1e-5)
+    @test ≈(jsol.sol.u[end][2], 0.25, rtol = 1e-5)
     # Test dynamics
     @parameters (u_interp::ConstantInterpolation)(..)
     @mtkbuild block_ode = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u_interp(t)], t)
@@ -120,7 +120,7 @@ end
     iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer; silent = true)
     @test is_bangbang(isol.input_sol, [-1.0], [1.0])
-    @test ≈(isol.sol.u[end][1], 0.25, rtol = 1e-5)
+    @test ≈(isol.sol.u[end][2], 0.25, rtol = 1e-5)
     osol = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false)
     @test ≈(isol.sol.u, osol.u, rtol = 0.05)
 
@@ -233,6 +233,8 @@ end
 end
 
 @testset "Cart-pole problem" begin
+    t = M.t_nounits
+    D = M.D_nounits
     # gravity, length, moment of Inertia, drag coeff
     @parameters g l mₚ mₖ
     @variables x(..) θ(..) u(t) [input = true, bounds = (-10, 10)]
diff --git a/test/runtests.jl b/test/runtests.jl
index 5f27f8bb88..3031e855a3 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -22,12 +22,6 @@ function activate_downstream_env()
     Pkg.instantiate()
 end
 
-function activate_dynamic_optimization_env()
-    Pkg.activate("dynamic_optimization")
-    Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
-    Pkg.instantiate()
-end
-
 @time begin
     #if GROUP == "All" || GROUP == "InterfaceI"
     #    @testset "InterfaceI" begin
@@ -128,6 +122,7 @@ end
 
     if GROUP == "All" || GROUP == "Downstream"
         activate_downstream_env()
+        @safetestset "JuMP Collocation Solvers" include("downstream/jump_control.jl")
         @safetestset "Linearization Tests" include("downstream/linearize.jl")
         @safetestset "Linearization Dummy Derivative Tests" include("downstream/linearization_dd.jl")
         @safetestset "Inverse Models Test" include("downstream/inversemodel.jl")
@@ -149,9 +144,4 @@ end
         @safetestset "InfiniteOpt Extension Test" include("extensions/test_infiniteopt.jl")
         @safetestset "JuMPControl Extension Test" include("extensions/jump_control.jl")
     end
-
-    if GROUP == "All" || GROUP == "Dynamic Optimization"
-        activate_dynamic_optimization_env()
-        @safetestset "JuMP Collocation Solvers" include("dynamic_optimization/jump_control")
-    end
 end

From 56caf18ad2737d9febdfb6af3831f78986c3b254 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Mon, 28 Apr 2025 16:13:22 -0400
Subject: [PATCH 38/41] rtest: emove from extensions

---
 test/runtests.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 3031e855a3..679a36f26c 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -142,6 +142,5 @@ end
         @safetestset "LabelledArrays Test" include("labelledarrays.jl")
         @safetestset "BifurcationKit Extension Test" include("extensions/bifurcationkit.jl")
         @safetestset "InfiniteOpt Extension Test" include("extensions/test_infiniteopt.jl")
-        @safetestset "JuMPControl Extension Test" include("extensions/jump_control.jl")
     end
 end

From e080243bee69123e30e1e335e1384c26d655a3ea Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 29 Apr 2025 14:14:55 -0400
Subject: [PATCH 39/41] fix: make free final time problems with constraints
 work

---
 ext/MTKJuMPDynamicOptExt.jl              | 49 +++++++++++++++++-------
 src/systems/optimal_control_interface.jl |  2 +-
 test/downstream/jump_control.jl          | 38 +++++++++++++++---
 3 files changed, 68 insertions(+), 21 deletions(-)

diff --git a/ext/MTKJuMPDynamicOptExt.jl b/ext/MTKJuMPDynamicOptExt.jl
index 882456dab3..448963dc60 100644
--- a/ext/MTKJuMPDynamicOptExt.jl
+++ b/ext/MTKJuMPDynamicOptExt.jl
@@ -42,7 +42,7 @@ end
 Convert an ODESystem 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 `nsteps`.
+provide the number of points as `steps`.
 
 The optimization variables:
 - a vector-of-vectors U representing the unknowns as an interpolation array
@@ -61,7 +61,7 @@ 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 = MTK.todict(pmap)
+    pmap = Dict{Any, Any}(pmap)
     steps, is_free_t = MTK.process_tspan(tspan, dt, steps)
     model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
 
@@ -87,7 +87,7 @@ 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 = MTK.todict(pmap)
+    pmap = Dict{Any, Any}(pmap)
     steps, is_free_t = MTK.process_tspan(tspan, dt, steps)
     model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
 
@@ -103,14 +103,15 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
 
     if is_free_t
         (ts_sym, te_sym) = tspan
+        MTK.symbolic_type(ts_sym) !== MTK.NotSymbolic() && error("Free initial time problems are not currently supported.")
         @variable(model, tf, start=pmap[te_sym])
+        set_lower_bound(tf, ts_sym)
         hasbounds(te_sym) && begin
             lo, hi = getbounds(te_sym)
             set_lower_bound(tf, lo)
             set_upper_bound(tf, hi)
         end
-        pmap[ts_sym] = 0
-        pmap[te_sym] = 1
+        pmap[te_sym] = model[:tf]
         tspan = (0, 1)
     end
 
@@ -118,6 +119,9 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
     @variable(model, U[i = 1:length(states)], Infinite(t), start=u0[i])
     c0 = MTK.value.([pmap[c] for c in ctrls])
     @variable(model, V[i = 1:length(ctrls)], Infinite(t), start=c0[i])
+    for (i, ct) in enumerate(ctrls)
+        pmap[ct] = model[:V][i]
+    end
 
     set_jump_bounds!(model, sys, pmap)
     add_jump_cost_function!(model, sys, (tspan[1], tspan[2]), pmap; is_free_t)
@@ -131,6 +135,13 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
     return model
 end
 
+"""
+Modify the pmap by replacing controls with V[i](t), and tf with the model's final time variable for free final time problems.
+"""
+function modified_pmap(model, sys, pmap)
+    pmap = Dict{Any, Any}(pmap)
+end
+
 function set_jump_bounds!(model, sys, pmap)
     U = model[:U]
     for (i, u) in enumerate(unknowns(sys))
@@ -158,7 +169,7 @@ function add_jump_cost_function!(model::InfiniteModel, sys, tspan, pmap; is_free
         @objective(model, Min, 0)
         return
     end
-    jcosts = substitute_jump_vars(model, sys, pmap, jcosts)
+    jcosts = substitute_jump_vars(model, sys, pmap, jcosts; is_free_t)
     tₛ = is_free_t ? model[:tf] : 1
 
     # Substitute integral
@@ -186,13 +197,14 @@ function add_user_constraints!(model::InfiniteModel, sys, pmap; is_free_t = fals
         for u in MTK.get_unknowns(conssys)
             x = MTK.operation(u)
             t = only(arguments(u))
-            MTK.symbolic_type(t) === NotSymbolic() &&
-                error("Provided specific time constraint in a free final time problem. This is not supported by the JuMP/InfiniteOpt collocation solvers. The offending variable is $u.")
+            if (MTK.symbolic_type(t) === MTK.NotSymbolic())
+                error("Provided specific time constraint in a free final time problem. This is not supported by the JuMP/InfiniteOpt collocation solvers. The offending variable is $u. Specific-time constraints can only be specified at the beginning or end of the timespan.")
+            end
         end
     end
 
     auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in unknowns(conssys)])
-    jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints; auxmap)
+    jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints; auxmap, is_free_t)
 
     # Substitute to-term'd variables
     for (i, cons) in enumerate(jconstraints)
@@ -211,13 +223,22 @@ function add_initial_constraints!(model::InfiniteModel, u0, u0_idxs, ts)
     @constraint(model, initial[i in u0_idxs], U[i](ts)==u0[i])
 end
 
-function substitute_jump_vars(model, sys, pmap, exprs; auxmap = Dict())
+function substitute_jump_vars(model, sys, pmap, exprs; auxmap = Dict(), is_free_t = false)
     iv = MTK.get_iv(sys)
     sts = unknowns(sys)
     cts = MTK.unbound_inputs(sys)
     U = model[:U]
     V = model[:V]
+    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)
+    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)
+    end
 
     # for variables like x(t)
     whole_interval_map = Dict([[v => U[i] for (i, v) in enumerate(sts)];
@@ -225,14 +246,10 @@ function substitute_jump_vars(model, sys, pmap, exprs; auxmap = Dict())
     exprs = map(c -> Symbolics.fixpoint_sub(c, whole_interval_map), exprs)
 
     # for variables like x(1.0)
-    x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
-    c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
     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.fixpoint_sub(c, Dict(pmap)), exprs)
     exprs
 end
 
@@ -246,8 +263,12 @@ function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap; is_free_
     diffsubmap = Dict([D(U[i]) => ∂(U[i], t) for i in 1:length(U)])
     tₛ = is_free_t ? model[:tf] : 1
 
+    @show diff_equations(sys)
+    @show pmap
     diff_eqs = substitute_jump_vars(model, sys, pmap, diff_equations(sys))
+    @show diff_eqs
     diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs)
+    @show diff_eqs
     @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs==tₛ * diff_eqs[i].rhs)
 
     # Algebraic equations
diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index 1183d04db6..be303192e4 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -173,7 +173,7 @@ function process_tspan(tspan, dt, steps)
     elseif symbolic_type(tspan[1]) === ScalarSymbolic() ||
            symbolic_type(tspan[2]) === ScalarSymbolic()
         isnothing(steps) &&
-            error("Free final time problems require specifying the number of steps, rather than dt.")
+            error("Free final time problems require specifying the number of steps using the keyword arg `steps`, rather than dt.")
         isnothing(dt) ||
             @warn "Specified dt for free final time problem. This will be ignored; dt will be determined by the number of timesteps."
 
diff --git a/test/downstream/jump_control.jl b/test/downstream/jump_control.jl
index 303b18bbeb..be36f8bbbc 100644
--- a/test/downstream/jump_control.jl
+++ b/test/downstream/jump_control.jl
@@ -178,7 +178,7 @@ end
 
     (ts, te) = (0.0, 0.2)
     costs = [-h(te)]
-    cons = [T(te) ~ 0]
+    cons = [T(te) ~ 0, m(te) ~ m_c]
     @named rocket = ODESystem(eqs, t; costs, constraints = cons)
     rocket, input_idxs = structural_simplify(rocket, ([T(t)], []))
 
@@ -186,14 +186,14 @@ end
     pmap = [
         g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5 * √(g₀ * h₀), D_c => 0.5 * 620 * m₀ / g₀,
         Tₘ => 3.5 * g₀ * m₀, T(t) => 0.0, h₀ => 1, m_c => 0.6]
-    jprob = JuMPDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
+    jprob = JuMPDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false)
     jsol = solve(jprob, Ipopt.Optimizer, :RadauIIA5, silent = true)
     @test jsol.sol.u[end][1] > 1.012
 
     iprob = InfiniteOptDynamicOptProblem(
-        rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
+        rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false)
     isol = solve(iprob, Ipopt.Optimizer,
-        derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true)
+                 derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true)
     @test isol.sol.u[end][1] > 1.012
 
     # Test solution
@@ -204,11 +204,16 @@ end
     @mtkbuild rocket_ode = ODESystem(eqs, t)
     interpmap = Dict(T_interp => ctrl_to_spline(jsol.input_sol, CubicSpline))
     oprob = ODEProblem(rocket_ode, u0map, (ts, te), merge(Dict(pmap), interpmap))
-    osol = solve(oprob, RadauIIA5(); adaptive = false, dt = 0.005)
+    osol = solve(oprob, RadauIIA5(); adaptive = false, dt = 0.001)
     @test ≈(jsol.sol.u, osol.u, rtol = 0.02)
+
+    interpmap1 = Dict(T_interp => ctrl_to_spline(isol.input_sol, CubicSpline))
+    oprob1 = ODEProblem(rocket_ode, u0map, (ts, te), merge(Dict(pmap), interpmap1))
+    osol1 = solve(oprob1, RadauIIA5(); adaptive = false, dt = 0.001)
+    @test ≈(isol.sol.u, osol1.u, rtol = 0.02)
 end
 
-@testset "Free final time problem" begin
+@testset "Free final time problems" begin
     t = M.t_nounits
     D = M.D_nounits
 
@@ -230,6 +235,27 @@ end
     iprob = InfiniteOptDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 200)
     isol = solve(iprob, Ipopt.Optimizer)
     @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3)
+
+    @variables x(..) v(..)
+    @variables u(..) [bounds = (-1.0, 1.0), input = true]
+    @parameters tf
+    constr = [v(tf) ~ 0, x(tf) ~ 0]
+    cost = [tf] # Minimize time
+
+    @named block = ODESystem(
+        [D(x(t)) ~ v(t), D(v(t)) ~ u(t)], t; costs = cost, constraints = constr)
+    block, input_idxs = structural_simplify(block, ([u(t)], []))
+    
+    u0map = [x(t) => 1.0, v(t) => 0.0]
+    tspan = (0.0, tf)
+    parammap = [u(t) => 0.0, tf => 1.0]
+    jprob = JuMPDynamicOptProblem(block, u0map, tspan, parammap; steps = 51)
+    jsol = solve(jprob, Ipopt.Optimizer, :Verner8, silent = true)
+    @test isapprox(jsol.sol.t[end], 2.0, atol = 1e-5)
+
+    iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; steps = 51)
+    isol = solve(iprob, Ipopt.Optimizer, silent = true)
+    @test isapprox(isol.sol.t[end], 2.0, atol = 1e-5)
 end
 
 @testset "Cart-pole problem" begin

From bcbd593023db80fa0d979d0c3ee794f3f601fb2a Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Tue, 29 Apr 2025 14:17:34 -0400
Subject: [PATCH 40/41] format: format

---
 ext/MTKJuMPDynamicOptExt.jl     | 6 ++++--
 test/downstream/jump_control.jl | 4 ++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/ext/MTKJuMPDynamicOptExt.jl b/ext/MTKJuMPDynamicOptExt.jl
index 448963dc60..0e86064ff4 100644
--- a/ext/MTKJuMPDynamicOptExt.jl
+++ b/ext/MTKJuMPDynamicOptExt.jl
@@ -103,7 +103,8 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
 
     if is_free_t
         (ts_sym, te_sym) = tspan
-        MTK.symbolic_type(ts_sym) !== MTK.NotSymbolic() && error("Free initial time problems are not currently supported.")
+        MTK.symbolic_type(ts_sym) !== MTK.NotSymbolic() &&
+            error("Free initial time problems are not currently supported.")
         @variable(model, tf, start=pmap[te_sym])
         set_lower_bound(tf, ts_sym)
         hasbounds(te_sym) && begin
@@ -236,7 +237,8 @@ function substitute_jump_vars(model, sys, pmap, exprs; auxmap = Dict(), is_free_
     exprs = map(c -> Symbolics.fixpoint_sub(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)]])
+        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)
     end
 
diff --git a/test/downstream/jump_control.jl b/test/downstream/jump_control.jl
index be36f8bbbc..ca339a60eb 100644
--- a/test/downstream/jump_control.jl
+++ b/test/downstream/jump_control.jl
@@ -193,7 +193,7 @@ end
     iprob = InfiniteOptDynamicOptProblem(
         rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false)
     isol = solve(iprob, Ipopt.Optimizer,
-                 derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true)
+        derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true)
     @test isol.sol.u[end][1] > 1.012
 
     # Test solution
@@ -245,7 +245,7 @@ end
     @named block = ODESystem(
         [D(x(t)) ~ v(t), D(v(t)) ~ u(t)], t; costs = cost, constraints = constr)
     block, input_idxs = structural_simplify(block, ([u(t)], []))
-    
+
     u0map = [x(t) => 1.0, v(t) => 0.0]
     tspan = (0.0, tf)
     parammap = [u(t) => 0.0, tf => 1.0]

From d0cd1e388f364208b307977f859aa3c13c705095 Mon Sep 17 00:00:00 2001
From: vyudu <vincent.duyuan@gmail.com>
Date: Wed, 30 Apr 2025 09:45:32 -0400
Subject: [PATCH 41/41] fix: fix rocket launch test

---
 ext/MTKJuMPDynamicOptExt.jl     |  4 ----
 test/downstream/jump_control.jl | 23 +++++++++++------------
 2 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/ext/MTKJuMPDynamicOptExt.jl b/ext/MTKJuMPDynamicOptExt.jl
index 0e86064ff4..fad10b5048 100644
--- a/ext/MTKJuMPDynamicOptExt.jl
+++ b/ext/MTKJuMPDynamicOptExt.jl
@@ -265,12 +265,8 @@ function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap; is_free_
     diffsubmap = Dict([D(U[i]) => ∂(U[i], t) for i in 1:length(U)])
     tₛ = is_free_t ? model[:tf] : 1
 
-    @show diff_equations(sys)
-    @show pmap
     diff_eqs = substitute_jump_vars(model, sys, pmap, diff_equations(sys))
-    @show diff_eqs
     diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs)
-    @show diff_eqs
     @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs==tₛ * diff_eqs[i].rhs)
 
     # Algebraic equations
diff --git a/test/downstream/jump_control.jl b/test/downstream/jump_control.jl
index ca339a60eb..4b454914c9 100644
--- a/test/downstream/jump_control.jl
+++ b/test/downstream/jump_control.jl
@@ -25,7 +25,7 @@ const M = ModelingToolkit
     # Test explicit method.
     jprob = JuMPDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01)
     @test JuMP.num_constraints(jprob.model) == 2 # initials
-    jsol = solve(jprob, Ipopt.Optimizer, :RK4)
+    jsol = solve(jprob, Ipopt.Optimizer, :RK4, silent = true)
     oprob = ODEProblem(sys, u0map, tspan, parammap)
     osol = solve(oprob, SimpleRK4(), dt = 0.01)
     @test jsol.sol.u ≈ osol.u
@@ -104,7 +104,7 @@ end
     tspan = (0.0, 1.0)
     parammap = [u(t) => 0.0]
     jprob = JuMPDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01)
-    jsol = solve(jprob, Ipopt.Optimizer, :Verner8)
+    jsol = solve(jprob, Ipopt.Optimizer, :Verner8, silent = true)
     # Linear systems have bang-bang controls
     @test is_bangbang(jsol.input_sol, [-1.0], [1.0])
     # Test reached final position.
@@ -142,7 +142,7 @@ end
     pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1, α => 1]
 
     jprob = JuMPDynamicOptProblem(beesys, u0map, tspan, pmap, dt = 0.01)
-    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5, silent = true)
     @test is_bangbang(jsol.input_sol, [0.0], [1.0])
     iprob = InfiniteOptDynamicOptProblem(beesys, u0map, tspan, pmap, dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer; silent = true)
@@ -191,9 +191,8 @@ end
     @test jsol.sol.u[end][1] > 1.012
 
     iprob = InfiniteOptDynamicOptProblem(
-        rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false)
-    isol = solve(iprob, Ipopt.Optimizer,
-        derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true)
+        rocket, u0map, (ts, te), pmap; dt = 0.001)
+    isol = solve(iprob, Ipopt.Optimizer, silent = true)
     @test isol.sol.u[end][1] > 1.012
 
     # Test solution
@@ -209,8 +208,8 @@ end
 
     interpmap1 = Dict(T_interp => ctrl_to_spline(isol.input_sol, CubicSpline))
     oprob1 = ODEProblem(rocket_ode, u0map, (ts, te), merge(Dict(pmap), interpmap1))
-    osol1 = solve(oprob1, RadauIIA5(); adaptive = false, dt = 0.001)
-    @test ≈(isol.sol.u, osol1.u, rtol = 0.02)
+    osol1 = solve(oprob1, ImplicitEuler(); adaptive = false, dt = 0.001)
+    @test ≈(isol.sol.u, osol1.u, rtol = 0.01)
 end
 
 @testset "Free final time problems" begin
@@ -229,11 +228,11 @@ end
     u0map = [x(t) => 17.5]
     pmap = [u(t) => 0.0, tf => 8]
     jprob = JuMPDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 201)
-    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5, silent = true)
     @test isapprox(jsol.sol.t[end], 10.0, rtol = 1e-3)
 
     iprob = InfiniteOptDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 200)
-    isol = solve(iprob, Ipopt.Optimizer)
+    isol = solve(iprob, Ipopt.Optimizer, silent = true)
     @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3)
 
     @variables x(..) v(..)
@@ -288,10 +287,10 @@ end
     u0map = [D(x(t)) => 0.0, D(θ(t)) => 0.0, θ(t) => 0.0, x(t) => 0.0]
     pmap = [mₖ => 1.0, mₚ => 0.2, l => 0.5, g => 9.81, u => 0]
     jprob = JuMPDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04)
-    jsol = solve(jprob, Ipopt.Optimizer, :RK4)
+    jsol = solve(jprob, Ipopt.Optimizer, :RK4, silent = true)
     @test jsol.sol.u[end] ≈ [π, 0, 0, 0]
 
     iprob = InfiniteOptDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04)
-    isol = solve(iprob, Ipopt.Optimizer)
+    isol = solve(iprob, Ipopt.Optimizer, silent = true)
     @test isol.sol.u[end] ≈ [π, 0, 0, 0]
 end