From 2ccdbbaf09bf051d14713dab75a949a3402222c7 Mon Sep 17 00:00:00 2001 From: vyudu Date: Thu, 15 May 2025 12:39:14 -0400 Subject: [PATCH 01/22] refactor: add definitions of shared functions to optimal_control_interface.jl --- ext/MTKCasADiDynamicOptExt.jl | 5 + src/systems/optimal_control_interface.jl | 147 +++++++++++++++++++++++ 2 files changed, 152 insertions(+) diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index 8f6b2c345e..e07791f4ab 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -90,6 +90,11 @@ function MTK.CasADiDynamicOptProblem(sys::System, u0map, tspan, pmap; CasADiDynamicOptProblem(f, u0, tspan, p, model, kwargs...) end +MTK.generate_U(model, dims) = 1 +MTK.generate_V(model, dims) = 1 +MTK.generate_timescale(model, dims) = 1 +MTK.generate_internal_model(::Type{CasADiModel}) = CasADi.opti() + function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false) ctrls = MTK.unbound_inputs(sys) states = unknowns(sys) diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index 2b460f77a2..04809f5686 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -1,6 +1,8 @@ abstract type AbstractDynamicOptProblem{uType, tType, isinplace} <: SciMLBase.AbstractODEProblem{uType, tType, isinplace} end +abstract type AbstractCollocation end + struct DynamicOptSolution model::Any sol::ODESolution @@ -148,3 +150,148 @@ function process_tspan(tspan, dt, steps) return length(tspan[1]:dt:tspan[2]), false end end + +function process_DynamicOptProblem(prob_type::AbstractDynamicOptProblem, model_type, sys::ODESystem, u0map, tspan, pmap; + 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(ODEInputFunction, sys, _u0map, pmap; + t = tspan !== nothing ? tspan[1] : tspan, kwargs...) + + 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] + pmap = Dict{Any, Any}(pmap) + steps, is_free_t = MTK.process_tspan(tspan, dt, steps) + + ctrls = MTK.unbound_inputs(sys) + states = unknowns(sys) + + model = generate_internal_model(model_type) + U = generate_U(model, u0) + V = generate_V() + tₛ = generate_timescale() + fullmodel = model_type(model, U, V, tₛ) + + set_variable_bounds!(fullmodel, sys, pmap) + add_cost_function!(fullmodel, sys, tspan, pmap; is_free_t) + add_user_constraints!(fullmodel, sys, tspan, pmap; is_free_t) + add_initial_constraints!(fullmodel, u0, u0_idxs) + + prob_type(f, u0, tspan, p, fullmodel, kwargs...) +end + +function add_cost_function!() + jcosts = copy(MTK.get_costs(sys)) + consolidate = MTK.get_consolidate(sys) + if isnothing(jcosts) || isempty(jcosts) + minimize!(opti, MX(0)) + return + end + + jcosts = substitute_model_vars(model, sys, pmap, jcosts; is_free_t) + jcosts = substitute_free_final_vars(model, sys, pmap, jcosts; is_free_t) + jcosts = substitute_fixed_t_vars(model, sys, pmap, jcosts; is_free_t) + jcosts = substitute_integral() +end + +function add_user_constraints!() + conssys = MTK.get_constraintsystem(sys) + jconstraints = isnothing(conssys) ? nothing : MTK.get_constraints(conssys) + (isnothing(jconstraints) || isempty(jconstraints)) && return nothing + + auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in unknowns(conssys)]) + jconstraints = substitute_model_vars(model, sys, pmap, jconstraints; auxmap, is_free_t) + + for c in jconstraints + if cons isa Equation + add_constraint!() + elseif cons.relational_op === Symbolics.geq + add_constraint!() + else + add_constraint!() + end + end +end + +function generate_U end +function generate_V end +function generate_timescale end + +function add_initial_constraints! end +function add_constraint! end + +function add_collocation_solve_constraints!(prob, tableau) + nᵤ = size(U.u, 1) + nᵥ = size(V.u, 1) + + if is_explicit(tableau) + K = MX[] + for k in 1:(length(tsteps) - 1) + τ = tsteps[k] + for (i, h) in enumerate(c) + ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = MX(zeros(nᵤ))) + Uₙ = U.u[:, k] + ΔU * dt + Vₙ = V.u[:, k] + 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(α)]) + subject_to!(solver_opti, U.u[:, k] + ΔU == U.u[:, k + 1]) + empty!(K) + end + else + for k in 1:(length(tsteps) - 1) + τ = tsteps[k] + # Kᵢ = generate_K() + Kᵢ = variable!(solver_opti, nᵤ, length(α)) + ΔUs = A * Kᵢ' # the stepsize at each stage of the implicit method + for (i, h) in enumerate(c) + ΔU = ΔUs[i, :]' + Uₙ = U.u[:, k] + ΔU * dt + Vₙ = V.u[:, k] + subject_to!(solver_opti, Kᵢ[:, i] == tₛ * f(Uₙ, Vₙ, p, τ + h * dt)) + end + ΔU_tot = dt * (Kᵢ * α) + subject_to!(solver_opti, U.u[:, k] + ΔU_tot == U.u[:, k + 1]) + end + end +end + +function add_equational_solve_constraints!() + diff_eqs = substitute_differentials() + add_constraint!() + + alg_eqs = substitute_model_vars() + add_constraint!() +end + +""" +Add the solve constraints, set the solver (Ipopt, e.g.) +""" +function prepare_solver end + +function DiffEqBase.solve(prob::AbstractDynamicOptProblem, solver::AbstractCollocation) + #add_solve_constraints!(prob, solver) + solver = prepare_solver(prob, solver) + sol = solve_prob(prob, solver) + + ts = get_t_values(sol) + Us = get_U_values(sol) + Vs = get_V_values(sol) + + ode_sol = DiffEqBase.build_solution(prob, solver, ts, Us) + input_sol = DiffEqBase.build_solution(prob, solver, ts, Vs) + + if successful_solve(model) + ode_sol = SciMLBase.solution_new_retcode( + ode_sol, SciMLBase.ReturnCode.ConvergenceFailure) + !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode( + input_sol, SciMLBase.ReturnCode.ConvergenceFailure)) + end + DynamicOptSolution(model, ode_sol, input_sol) +end From df1a0e496fb0d8e7afb8fd034595bcbf251bb591 Mon Sep 17 00:00:00 2001 From: vyudu Date: Fri, 16 May 2025 17:05:45 -0400 Subject: [PATCH 02/22] refactor: aduse optimal control itnerface --- ext/MTKCasADiDynamicOptExt.jl | 9 +- ext/MTKInfiniteOptExt.jl | 342 ++++++++--------------- src/systems/optimal_control_interface.jl | 181 ++++++------ test/extensions/dynamic_optimization.jl | 6 +- 4 files changed, 223 insertions(+), 315 deletions(-) diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index e07791f4ab..e0c8fbb5ea 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -25,6 +25,7 @@ struct CasADiModel U::MXLinearInterpolation V::MXLinearInterpolation tₛ::MX + is_free_final::Bool end struct CasADiDynamicOptProblem{uType, tType, isinplace, P, F, K} <: @@ -90,10 +91,10 @@ function MTK.CasADiDynamicOptProblem(sys::System, u0map, tspan, pmap; CasADiDynamicOptProblem(f, u0, tspan, p, model, kwargs...) end -MTK.generate_U(model, dims) = 1 -MTK.generate_V(model, dims) = 1 -MTK.generate_timescale(model, dims) = 1 MTK.generate_internal_model(::Type{CasADiModel}) = CasADi.opti() +MTK.generate_state_variable(model, u0, ns, nt) +MTK.generate_input_variable(model, c0, nc, nt) = 1 +MTK.generate_timescale(model, dims) = 1 function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false) ctrls = MTK.unbound_inputs(sys) @@ -317,7 +318,7 @@ function add_solve_constraints(prob, tableau) for k in 1:(length(tsteps) - 1) τ = tsteps[k] Kᵢ = variable!(solver_opti, nᵤ, length(α)) - ΔUs = A * Kᵢ' # the stepsize at each stage of the implicit method + ΔUs = A * Kᵢ' for (i, h) in enumerate(c) ΔU = ΔUs[i, :]' Uₙ = U.u[:, k] + ΔU * dt diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index 40eb6f5264..6dd80232f8 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -4,6 +4,7 @@ using InfiniteOpt using DiffEqBase using LinearAlgebra using StaticArrays +using UnPack import SymbolicUtils import NaNMath const MTK = ModelingToolkit @@ -38,6 +39,29 @@ struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <: end end +struct InfiniteOptModel + model::InfiniteModel + U::AbstractVariableRef + V::AbstractVariableRef + tₛ::Union + is_free_final::Bool +end + +MTK.generate_internal_model(m::Type{InfiniteOptModel}) = InfiniteModel() +MTK.generate_state_variable!(m::InfiniteModel, u0::Vector, ns, nt) = @variable(m, U[i = 1:nt], Infinite(m[:t]), start=u0[i]) +MTK.generate_input_variable!(m::InfiniteModel, c0, nc, nt) = @variable(m, V[i = 1:nc], Infinite(m[:t]), start=c0[i]) + +function MTK.generate_timescale!(m::InfiniteModel, guess, is_free_t) + @variable(m, tₛ ≥ 0, start = guess) + if !is_free_t + fix(m[:tₛ], 1) + set_start_value(m[:tₛ], 1) + end +end + +MTK.add_constraint!(m::InfiniteOptModel, expr) = @constraint(m.model, expr) +MTK.set_objective!(m::InfiniteOptModel, expr) = @objective(m.model, Min, expr) + """ JuMPDynamicOptProblem(sys::System, u0, tspan, p; dt) @@ -58,19 +82,7 @@ function MTK.JuMPDynamicOptProblem(sys::System, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - MTK.warn_overdetermined(sys, u0map) - _u0map = has_alg_eqs(sys) ? MTK.to_varmap(u0map, unknowns(sys)) : - merge(Dict(u0map), Dict(guesses)) - pmap = MTK.to_varmap(pmap, parameters(sys)) - f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, merge(_u0map, pmap); - t = tspan !== nothing ? tspan[1] : tspan, kwargs...) - - pmap = MTK.recursive_unwrap(MTK.AnyDict(pmap)) - MTK.evaluate_varmap!(pmap, keys(pmap)) - steps, is_free_t = MTK.process_tspan(tspan, dt, steps) - model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t) - - JuMPDynamicOptProblem(f, u0, tspan, p, model, kwargs...) + process_DynamicOptProblem(JuMPDynamicOptProblem, InfiniteOptModel, u0map, tspan, pmap; dt, steps, guesses, kwargs...) end """ @@ -87,216 +99,87 @@ function MTK.InfiniteOptDynamicOptProblem(sys::System, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - MTK.warn_overdetermined(sys, u0map) - _u0map = has_alg_eqs(sys) ? MTK.to_varmap(u0map, unknowns(sys)) : - merge(Dict(u0map), Dict(guesses)) - pmap = MTK.to_varmap(pmap, parameters(sys)) - f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, merge(_u0map, pmap); - t = tspan !== nothing ? tspan[1] : tspan, kwargs...) - - pmap = MTK.recursive_unwrap(MTK.AnyDict(pmap)) - MTK.evaluate_varmap!(pmap, keys(pmap)) - steps, is_free_t = MTK.process_tspan(tspan, dt, steps) - model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t) - - add_infopt_solve_constraints!(model, sys, pmap; is_free_t) - InfiniteOptDynamicOptProblem(f, u0, tspan, p, model, kwargs...) -end - -# 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() - - 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[te_sym] = model[:tf] - 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 = 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) - 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[MTK.default_toterm(k)] for (k, v) in u0map] - add_initial_constraints!(model, u0, u0_idxs, tspan[1]) - return model + prob = process_DynamicOptProblem(InfiniteOptDynamicOptProblem, InfiniteOptModel, u0map, tspan, pmap; dt, steps, guesses, kwargs...) 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] +function MTK.set_variable_bounds!(model, sys, pmap, tf) for (i, u) in enumerate(unknowns(sys)) if MTK.hasbounds(u) lo, hi = MTK.getbounds(u) - set_lower_bound(U[i], Symbolics.fast_substitute(lo, pmap)) - set_upper_bound(U[i], Symbolics.fast_substitute(hi, pmap)) + set_lower_bound(model.U[i], Symbolics.fixpoint_sub(lo, pmap)) + set_upper_bound(model.U[i], Symbolics.fixpoint_sub(hi, pmap)) end end - V = model[:V] for (i, v) in enumerate(MTK.unbound_inputs(sys)) if MTK.hasbounds(v) lo, hi = MTK.getbounds(v) - set_lower_bound(V[i], Symbolics.fast_substitute(lo, pmap)) - set_upper_bound(V[i], Symbolics.fast_substitute(hi, pmap)) + set_lower_bound(model.V[i], Symbolics.fixpoint_sub(lo, pmap)) + set_upper_bound(model.V[i], Symbolics.fixpoint_sub(hi, pmap)) end end -end -function add_jump_cost_function!(model::InfiniteModel, sys, tspan, pmap; is_free_t = false) - jcosts = cost(sys) - if Symbolics._iszero(jcosts) - @objective(model, Min, 0) - return + if symbolic_type(tf) === ScalarSymbolic() && hasbounds(tf) + lo, hi = MTK.getbounds(tf) + set_lower_bound(model.tₛ, lo) + set_upper_bound(model.tₛ, hi) end - jcosts = substitute_jump_vars(model, sys, pmap, [jcosts]; is_free_t)[1] - tₛ = is_free_t ? model[:tf] : 1 - - # Substitute integral - iv = MTK.get_iv(sys) +end - intmap = Dict() +function MTK.substitute_integral(model, jcosts) for int in MTK.collect_applied_operators(jcosts, Symbolics.Integral) op = MTK.operation(int) arg = only(arguments(MTK.value(int))) lo, hi = (op.domain.domain.left, op.domain.domain.right) lo = MTK.value(lo) hi = haskey(pmap, hi) ? 1 : MTK.value(hi) - intmap[int] = tₛ * InfiniteOpt.∫(arg, model[:t], lo, hi) - end - jcosts = Symbolics.substitute(jcosts, intmap) - @objective(model, Min, MTK.value(jcosts)) -end - -function add_user_constraints!(model::InfiniteModel, sys, pmap; is_free_t = false) - jconstraints = MTK.get_constraints(sys) - (isnothing(jconstraints) || isempty(jconstraints)) && return nothing - cons_dvs, cons_ps = MTK.process_constraint_system( - jconstraints, Set(unknowns(sys)), parameters(sys), MTK.get_iv(sys); validate = false) - - if is_free_t - for u in cons_dvs - x = MTK.operation(u) - t = only(arguments(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 cons_dvs]) - jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints; auxmap, is_free_t) - - # 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]") - 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]") - end + intmap[int] = model.tₛ * InfiniteOpt.∫(arg, model.model.t, lo, hi) end + jcosts = map(c -> Symbolics.substitute(c, intmap), jcosts) end -function add_initial_constraints!(model::InfiniteModel, u0, u0_idxs, ts) - U = model[:U] +function MTK.add_initial_constraints!(model::InfiniteOptModel, u0, u0_idxs, ts) + U = model.U @constraint(model, initial[i in u0_idxs], U[i](ts)==u0[i]) end -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.fast_substitute(c, auxmap), exprs) - exprs = map(c -> Symbolics.fast_substitute(c, Dict(pmap)), exprs) - if is_free_t - tf = model[:tf] +function MTK.substitute_model_vars(model, sys, exprs; tf = nothing) + whole_interval_map = Dict([[v => model.U[i] for (i, v) in enumerate(unknowns(sys))]; + [v => model.V[i] for (i, v) in enumerate(MTK.unbound_inputs(sys))]]) + exprs = map(c -> Symbolics.fast_substitute(c, whole_interval_map), exprs) + + x_ops = [MTK.operation(MTK.unwrap(st)) for st in unknowns(sys)] + c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in MTK.unbound_inputs(sys)] + + if symbolic_type(tf) === ScalarSymbolic() 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.fast_substitute(c, free_t_map), exprs) end - # for variables like x(t) - whole_interval_map = Dict([[v => U[i] for (i, v) in enumerate(sts)]; - [v => V[i] for (i, v) in enumerate(cts)]]) - exprs = map(c -> Symbolics.fast_substitute(c, whole_interval_map), exprs) - # for variables like x(1.0) fixed_t_map = Dict([[x_ops[i] => U[i] for i in 1:length(U)]; [c_ops[i] => V[i] for i in 1:length(V)]]) - exprs = map(c -> Symbolics.fast_substitute(c, fixed_t_map), exprs) - exprs end -function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap; is_free_t = false) - # Differential equations - U = model[:U] - t = model[:t] +function MTK.substitute_differentials(model::InfiniteOptModel, eqs) + U = model.U + t = model.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==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) + map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs) end -function add_jump_solve_constraints!(prob, tableau; is_free_t = false) - A = tableau.A - α = tableau.α - c = tableau.c - model = prob.model - f = prob.f - p = prob.p - - t = model[:t] - tsteps = supports(t) - tmax = tsteps[end] - pop!(tsteps) - tₛ = is_free_t ? model[:tf] : 1 +function add_solve_constraints!(prob::JuMPDynamicOptProblem, solver) + @unpack A, α, c = solver.tableau + @unpack model, f, p = prob + tsteps = supports(model.model[:t]) dt = tsteps[2] - tsteps[1] - U = model[:U] - V = model[:V] + tₛ = model.tₛ + U = model.U + V = model.V nᵤ = length(U) nᵥ = length(V) if MTK.is_explicit(tableau) @@ -306,7 +189,7 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false) Δ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ₙ = tₛ * f(Uₙ, Vₙ, p, τ + h * dt) # scale the time + Kₙ = tₛ * f(Uₙ, Vₙ, p, τ + h * dt) push!(K, Kₙ) end ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)]) @@ -325,27 +208,39 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false) @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)), + @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tsteps[end])), DomainRestrictions(t => τ), base_name="solve_U($τ)") end end end """ -Solve JuMPDynamicOptProblem. Arguments: -- prob: a JumpDynamicOptProblem -- jump_solver: a LP solver such as HiGHS -- tableau_getter: Takes in a function to fetch a tableau. Tableau loaders look like `constructRK4` and may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. If this argument is not passed in, the solver will default to Radau second order. -- silent: set the model silent (suppress model output) +JuMP Collocation solver. +- solver: a optimization solver such as Ipopt +- tableau: An ODE RK tableau. Load a tableau by calling a function like `constructRK4` and may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. If this argument is not passed in, the solver will default to Radau second order. Returns a DynamicOptSolution, which contains both the model and the ODE solution. """ -function DiffEqBase.solve( - prob::JuMPDynamicOptProblem, jump_solver, tableau_getter = MTK.constructDefault; silent = false) - model = prob.model - tableau = tableau_getter() - silent && set_silent(model) +struct JuMPCollocation + solver::Any + tableau::DiffEqBase.ODERKTableau +end +JuMPCollocation(solver; tableau = MTK.constructDefault()) = JuMPCollocation(solver, tableau) + +""" +InfiniteOpt Collocation solver. +- solver: an optimization solver such as Ipopt +- `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()). +""" +struct InfiniteOptCollocation + solver::Any + derivative_method::InfiniteOpt.AbstractDerivativeMethod +end +InfiniteOptCollocation(solver; derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward())) = InfiniteOptCollocation(solver, derivative_method) +function MTK.prepare_solver!(prob::JuMPDynamicOptProblem, solver::JuMPCollocation; verbose = false, kwargs...) + model = prob.model.model + verbose || set_silent(model) # Unregister current solver constraints for con in all_constraints(model) if occursin("solve", JuMP.name(con)) @@ -360,53 +255,48 @@ function DiffEqBase.solve( delete(model, var) end end - add_jump_solve_constraints!(prob, tableau; is_free_t = haskey(model, :tf)) - _solve(prob, jump_solver, tableau_getter) + add_collocation_solve_constraints!(model, solver.tableau) + set_optimizer(model, solver.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::InfiniteOptDynamicOptProblem, jump_solver; - derivative_method = InfiniteOpt.FiniteDifference(Backward()), silent = false) - model = prob.model - silent && set_silent(model) - set_derivative_method(model[:t], derivative_method) - _solve(prob, jump_solver, derivative_method) +function MTK.prepare_solver!(prob::InfiniteOptDynamicOptProblem, solver::InfiniteOptCollocation; verbose = false, kwargs...) + model = prob.model.model + verbose || set_silent(model) + add_equational_constraints!(model, prob.f.sys, prob.tspan) + set_derivative_method(model[:t], solver.derivative_method) + set_optimizer(model, solver.solver) end -function _solve(prob::AbstractDynamicOptProblem, jump_solver, solver) - model = prob.model - set_optimizer(model, jump_solver) - optimize!(model) +function MTK.optimize_model!(prob::Union{InfiniteOptDynamicOptProblem, JuMPDynamicOptProblem}, solver) + optimize!(prob.model.model) + prob.model +end +function MTK.get_V_values(m::InfiniteOptModel) + if !isempty(m.V) + V_vals = value.(m.V) + V_vals = [[V_vals[i][j] for i in 1:length(V_vals)] for j in 1:length(ts)] + else + nothing + end +end +function MTK.get_U_values(m::InfiniteOptModel) + U_vals = value.(m.U) + U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:length(ts)] +end + +MTK.get_t_values(model) = prob.tspan[1] .+ model.tₛ * supports(model.model[:t]) + +function MTK.successful_solve(m::InfiniteOptModel) + model = m.model tstatus = termination_status(model) pstatus = primal_status(model) !has_values(model) && error("Model not solvable; please report this to github.com/SciML/ModelingToolkit.jl with a MWE.") - 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) - - 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, solver, ts, V_vals) - end - - if !(pstatus === FEASIBLE_POINT && + 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 - - DynamicOptSolution(model, sol, input_sol) + tstatus === ALMOST_LOCALLY_SOLVED) end import InfiniteOpt: JuMP, GeneralVariableRef diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index 04809f5686..04c1a0c74b 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -22,6 +22,10 @@ function JuMPDynamicOptProblem end function InfiniteOptDynamicOptProblem end function CasADiDynamicOptProblem end +function JuMPCollocation end +function InfiniteOptCollocation end +function CasADiCollocation end + function warn_overdetermined(sys, u0map) cstrs = constraints(sys) if !isempty(cstrs) @@ -133,6 +137,9 @@ end # 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 + symbolic_type(tspan[1]) !== NotSymbolic() && + error("Free initial time problems are not currently supported by the collocation solvers.") + 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() || @@ -142,16 +149,19 @@ function process_tspan(tspan, dt, steps) 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 + return (0, 1), 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 + return tspan, length(tspan[1]:dt:tspan[2]), false end end -function process_DynamicOptProblem(prob_type::AbstractDynamicOptProblem, model_type, sys::ODESystem, u0map, tspan, pmap; +########################## +### MODEL CONSTRUCTION ### +########################## +function process_DynamicOptProblem(prob_type::Type{<:AbstractDynamicOptProblem}, model_type, sys::ODESystem, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) @@ -166,128 +176,135 @@ function process_DynamicOptProblem(prob_type::AbstractDynamicOptProblem, model_t u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) : [stidxmap[MTK.default_toterm(k)] for (k, v) in u0map] pmap = Dict{Any, Any}(pmap) - steps, is_free_t = MTK.process_tspan(tspan, dt, steps) + model_tspan, steps, is_free_t = MTK.process_tspan(tspan, dt, steps) ctrls = MTK.unbound_inputs(sys) states = unknowns(sys) + c0 = MTK.value.([pmap[c] for c in ctrls]) model = generate_internal_model(model_type) - U = generate_U(model, u0) - V = generate_V() - tₛ = generate_timescale() - fullmodel = model_type(model, U, V, tₛ) + generate_time_variable!(model, model_tspan, steps) + U = generate_state_variable!(model, u0, length(states), length(steps)) + V = generate_input_variable!(model, c0, length(ctrls), length(steps)) + tₛ = generate_timescale!(model, get(pmap, tspan[2], tspan[2]), is_free_t) + fullmodel = model_type(model, U, V, tₛ, is_free_t) - set_variable_bounds!(fullmodel, sys, pmap) - add_cost_function!(fullmodel, sys, tspan, pmap; is_free_t) - add_user_constraints!(fullmodel, sys, tspan, pmap; is_free_t) - add_initial_constraints!(fullmodel, u0, u0_idxs) + set_variable_bounds!(fullmodel, sys, pmap, tspan[2]) + add_cost_function!(fullmodel, sys, tspan, pmap) + add_user_constraints!(fullmodel, sys, tspan, pmap) + add_initial_constraints!(fullmodel, u0, u0_idxs, model_tspan[1]) prob_type(f, u0, tspan, p, fullmodel, kwargs...) end -function add_cost_function!() +function generate_internal_model end +function generate_state_variable! end +function generate_input_variable! end +function generate_timescale! end +function set_variable_bounds! end +function add_initial_constraints! end +function add_constraint! end +is_free_final(model) = model.is_free_final + +function add_cost_function!(model, sys, tspan, pmap) jcosts = copy(MTK.get_costs(sys)) consolidate = MTK.get_consolidate(sys) if isnothing(jcosts) || isempty(jcosts) - minimize!(opti, MX(0)) + set_objective!(model, 0) return end - - jcosts = substitute_model_vars(model, sys, pmap, jcosts; is_free_t) - jcosts = substitute_free_final_vars(model, sys, pmap, jcosts; is_free_t) - jcosts = substitute_fixed_t_vars(model, sys, pmap, jcosts; is_free_t) - jcosts = substitute_integral() + jcosts = substitute_model_vars(model, sys, jcosts; tf = tspan[2]) + jcosts = substitute_params(pmap, jcosts) + jcosts = substitute_integral(model, jcosts) + set_objective!(model, consolidate(jcosts)) end -function add_user_constraints!() +function add_user_constraints!(model, sys, tspan, pmap) conssys = MTK.get_constraintsystem(sys) jconstraints = isnothing(conssys) ? nothing : MTK.get_constraints(conssys) (isnothing(jconstraints) || isempty(jconstraints)) && return nothing + consvars = MTK.get_unknowns(conssys) + is_free_final(model) && check_constraint_vars(consvars) - auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in unknowns(conssys)]) - jconstraints = substitute_model_vars(model, sys, pmap, jconstraints; auxmap, is_free_t) + jconstraints = substitute_model_vars(model, sys, jcosts; tf = tspan[2]) + jconstraints = substitute_toterm(consvars, jconstraints) + jconstraints = substitute_params(pmap, jconstraints) for c in jconstraints if cons isa Equation - add_constraint!() + add_constraint!(model, c.lhs - c.rhs == 0) elseif cons.relational_op === Symbolics.geq - add_constraint!() + add_constraint!(model, c.lhs - c.rhs ≥ 0) else - add_constraint!() + add_constraint!(model, c.lhs - c.rhs ≤ 0) end end end -function generate_U end -function generate_V end -function generate_timescale end - -function add_initial_constraints! end -function add_constraint! end +function add_equational_constraints!(model, sys, tspan) + model = model.model + diff_eqs = substitute_model_vars(model, sys, diff_equations(sys); tf = tspan[2]) + diff_eqs = substitute_differentials(model, sys, diff_eqs) + for eq in diff_eqs + add_constraint!(model, eq.lhs == eq.rhs * model.tₛ) + end -function add_collocation_solve_constraints!(prob, tableau) - nᵤ = size(U.u, 1) - nᵥ = size(V.u, 1) - - if is_explicit(tableau) - K = MX[] - for k in 1:(length(tsteps) - 1) - τ = tsteps[k] - for (i, h) in enumerate(c) - ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = MX(zeros(nᵤ))) - Uₙ = U.u[:, k] + ΔU * dt - Vₙ = V.u[:, k] - 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(α)]) - subject_to!(solver_opti, U.u[:, k] + ΔU == U.u[:, k + 1]) - empty!(K) - end - else - for k in 1:(length(tsteps) - 1) - τ = tsteps[k] - # Kᵢ = generate_K() - Kᵢ = variable!(solver_opti, nᵤ, length(α)) - ΔUs = A * Kᵢ' # the stepsize at each stage of the implicit method - for (i, h) in enumerate(c) - ΔU = ΔUs[i, :]' - Uₙ = U.u[:, k] + ΔU * dt - Vₙ = V.u[:, k] - subject_to!(solver_opti, Kᵢ[:, i] == tₛ * f(Uₙ, Vₙ, p, τ + h * dt)) - end - ΔU_tot = dt * (Kᵢ * α) - subject_to!(solver_opti, U.u[:, k] + ΔU_tot == U.u[:, k + 1]) - end + alg_eqs = substitute_model_vars(model, sys, alg_equations(sys); tf = tspan[2]) + for eq in alg_eqs + add_constraint!(model, eq.lhs == eq.rhs * model.tₛ) end end -function add_equational_solve_constraints!() - diff_eqs = substitute_differentials() - add_constraint!() +function set_objective! end +"""Substitute variables like x(1.5) with the corresponding model variables.""" +function substitute_model_vars end +function substitute_integral end +function substitute_differentials end + +function substitute_toterm(vars, exprs) + toterm_map = Dict([u => MTK.default_toterm(MTK.value(u)) for u in vars]) + exprs = map(c -> Symbolics.fast_substitute(c, toterm_map), exprs) +end + +function substitute_params(pmap, exprs) + exprs = map(c -> Symbolics.fast_substitute(c, Dict(pmap)), exprs) +end - alg_eqs = substitute_model_vars() - add_constraint!() +function check_constraint_vars(vars) + for u in vars + x = operation(u) + t = only(arguments(u)) + if (symbolic_type(t) === NotSymbolic()) + error("Provided specific time constraint in a free final time problem. This is not supported by the collocation solvers at the moment. The offending variable is $u. Specific-time user constraints can only be specified at the end of the timespan.") + end + end end +######################## +### SOLVER UTILITIES ### +######################## """ -Add the solve constraints, set the solver (Ipopt, e.g.) +Add the solve constraints, set the solver (Ipopt, e.g.) and solver options. """ -function prepare_solver end - -function DiffEqBase.solve(prob::AbstractDynamicOptProblem, solver::AbstractCollocation) - #add_solve_constraints!(prob, solver) - solver = prepare_solver(prob, solver) - sol = solve_prob(prob, solver) +function prepare_solver! end +function optimize_model! end +function get_t_values end +function get_U_values end +function get_V_values end +function successful_solve end + +function DiffEqBase.solve(prob::AbstractDynamicOptProblem, solver::AbstractCollocation; verbose = false, kwargs...) + solver = prepare_solver!(prob, solver) + model = optimize_model!(prob, solver) - ts = get_t_values(sol) - Us = get_U_values(sol) - Vs = get_V_values(sol) + ts = get_t_values(model) + Us = get_U_values(model) + Vs = get_V_values(model) ode_sol = DiffEqBase.build_solution(prob, solver, ts, Us) - input_sol = DiffEqBase.build_solution(prob, solver, ts, Vs) + input_sol = isnothing(Vs) ? nothing : DiffEqBase.build_solution(prob, solver, ts, Vs) - if successful_solve(model) + if !successful_solve(model) ode_sol = SciMLBase.solution_new_retcode( ode_sol, SciMLBase.ReturnCode.ConvergenceFailure) !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode( diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index dd2368b558..03aeebf1bc 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -1,5 +1,5 @@ using ModelingToolkit -import JuMP, InfiniteOpt +import InfiniteOpt using DiffEqDevTools, DiffEqBase using SimpleDiffEq using OrdinaryDiffEqSDIRK, OrdinaryDiffEqVerner, OrdinaryDiffEqTsit5, OrdinaryDiffEqFIRK @@ -27,7 +27,7 @@ const M = ModelingToolkit # Test explicit method. jprob = JuMPDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01) - @test JuMP.num_constraints(jprob.model) == 2 # initials + @test InfiniteOpt.num_constraints(jprob.model) == 2 # initials jsol = solve(jprob, Ipopt.Optimizer, constructRK4, silent = true) oprob = ODEProblem(sys, [u0map; parammap], tspan) osol = solve(oprob, SimpleRK4(), dt = 0.01) @@ -56,7 +56,7 @@ const M = ModelingToolkit @mtkcompile lksys = System(eqs, t; constraints = constr) jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) - @test JuMP.num_constraints(jprob.model) == 2 + @test InfiniteOpt.num_constraints(jprob.model) == 2 jsol = solve(jprob, Ipopt.Optimizer, constructTsitouras5, silent = true) # 12.190 s, 9.68 GiB @test jsol.sol(0.6; idxs = x(t)) ≈ 3.5 @test jsol.sol(0.3; idxs = x(t)) ≈ 7.0 From 912d71b318424c0d3737cb22d75d3048c4c89097 Mon Sep 17 00:00:00 2001 From: vyudu Date: Fri, 16 May 2025 19:05:05 -0400 Subject: [PATCH 03/22] refactor: add interface functions for CasADi --- ext/MTKCasADiDynamicOptExt.jl | 284 ++++++++--------------- ext/MTKInfiniteOptExt.jl | 107 +++++---- src/ModelingToolkit.jl | 2 + src/systems/optimal_control_interface.jl | 62 ++--- test/extensions/dynamic_optimization.jl | 58 +++-- 5 files changed, 211 insertions(+), 302 deletions(-) diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index e0c8fbb5ea..554d9faec4 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -6,10 +6,8 @@ using UnPack using NaNMath const MTK = ModelingToolkit -# NaNMath for ff in [acos, log1p, acosh, log2, asin, tan, atanh, cos, log, sin, log10, sqrt] f = nameof(ff) - # These need to be defined so that JuMP can trace through functions built by Symbolics @eval NaNMath.$f(x::CasadiSymbolicObject) = Base.$f(x) end @@ -76,78 +74,47 @@ function MTK.CasADiDynamicOptProblem(sys::System, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - MTK.warn_overdetermined(sys, u0map) - _u0map = has_alg_eqs(sys) ? MTK.to_varmap(u0map, unknowns(sys)) : - merge(Dict(u0map), Dict(guesses)) - pmap = MTK.to_varmap(pmap, parameters(sys)) - f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, merge(_u0map, pmap); - t = tspan !== nothing ? tspan[1] : tspan, output_type = MX, kwargs...) - - pmap = MTK.recursive_unwrap(MTK.AnyDict(pmap)) - MTK.evaluate_varmap!(pmap, keys(pmap)) - steps, is_free_t = MTK.process_tspan(tspan, dt, steps) - model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t) - - CasADiDynamicOptProblem(f, u0, tspan, p, model, kwargs...) + process_DynamicOptProblem(CasADiDynamicOptProblem, CasADiModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) end MTK.generate_internal_model(::Type{CasADiModel}) = CasADi.opti() -MTK.generate_state_variable(model, u0, ns, nt) -MTK.generate_input_variable(model, c0, nc, nt) = 1 -MTK.generate_timescale(model, dims) = 1 -function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false) - ctrls = MTK.unbound_inputs(sys) - states = unknowns(sys) - opti = CasADi.Opti() +function MTK.generate_state_variable(model::Opti, u0, ns, nt, tsteps) + U = CasADi.variable!(model, ns, nt) + set_initial!(opti, U, DM(repeat(u0, 1, steps))) + MXLinearInterpolation(U, tsteps, tsteps[2] - tsteps[1]) +end +function MTK.generate_input_variable(model::Opti, c0, nc, nt, tsteps) + V = CasADi.variable!(model, nc, nt) + !isempty(c0) && set_initial!(opti, V, DM(repeat(c0, 1, steps))) + MXLinearInterpolation(V, tsteps, tsteps[2] - tsteps[1]) +end + +function MTK.generate_timescale(model::Opti, guess, is_free_t) 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 in CasADiDynamicOptProblem.") - tₛ = variable!(opti) - set_initial!(opti, tₛ, pmap[te_sym]) - subject_to!(opti, tₛ >= ts_sym) - hasbounds(te_sym) && begin - lo, hi = getbounds(te_sym) - subject_to!(opti, tₛ >= lo) - subject_to!(opti, tₛ >= hi) - end - pmap[te_sym] = tₛ - tsteps = LinRange(0, 1, steps) + tₛ = variable!(model) + set_initial!(model, tₛ, guess) + subject_to!(model, tₛ >= 0) + tₛ else - tₛ = MX(1) - tsteps = LinRange(tspan[1], tspan[2], steps) + MX(1) end +end - U = CasADi.variable!(opti, length(states), steps) - V = CasADi.variable!(opti, length(ctrls), steps) - set_initial!(opti, U, DM(repeat(u0, 1, steps))) - c0 = MTK.value.([pmap[c] for c in ctrls]) - !isempty(c0) && set_initial!(opti, V, DM(repeat(c0, 1, steps))) - - U_interp = MXLinearInterpolation(U, tsteps, tsteps[2] - tsteps[1]) - V_interp = MXLinearInterpolation(V, tsteps, tsteps[2] - tsteps[1]) - for (i, ct) in enumerate(ctrls) - pmap[ct] = V[i, :] +function MTK.add_constraint!(model::CasADiModel, expr) + @unpack opti = model + if cons isa Equation + subject_to!(opti, expr.lhs - expr.rhs == 0) + elseif cons.relational_op === Symbolics.geq + subject_to!(opti, expr.lhs - expr.rhs ≥ 0) + else + subject_to!(opti, expr.lhs - expr.rhs ≤ 0) end - - model = CasADiModel(opti, U_interp, V_interp, tₛ) - - set_casadi_bounds!(model, sys, pmap) - add_cost_function!(model, sys, tspan, pmap; is_free_t) - add_user_constraints!(model, sys, tspan, 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[MTK.default_toterm(k)] for (k, v) in u0map] - add_initial_constraints!(model, u0, u0_idxs) - - model end +MTK.set_objective!(model::CasADiModel, expr) = minimize!(model.opti, MX(expr)) -function set_casadi_bounds!(model, sys, pmap) +function MTK.set_variable_bounds!(model, sys, pmap, tf) @unpack opti, U, V = model for (i, u) in enumerate(unknowns(sys)) if MTK.hasbounds(u) @@ -163,36 +130,56 @@ function set_casadi_bounds!(model, sys, pmap) subject_to!(opti, V.u[i, :] <= Symbolics.fast_substitute(hi, pmap)) end end + if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() && hasbounds(tf) + lo, hi = MTK.getbounds(tf) + subject_to!(opti, model.tₛ >= lo) + subject_to!(opti, model.tₛ <= hi) + end end -function add_initial_constraints!(model::CasADiModel, u0, u0_idxs) +function MTK.add_initial_constraints!(model::CasADiModel, u0, u0_idxs) @unpack opti, U = model for i in u0_idxs subject_to!(opti, U.u[i, 1] == u0[i]) end end -function add_user_constraints!(model::CasADiModel, sys, tspan, pmap; is_free_t) +function MTK.substitute_model_vars( + model::CasADiModel, sys, pmap, exprs; auxmap::Dict = Dict(), is_free_t) @unpack opti, U, V, tₛ = model - iv = MTK.get_iv(sys) - jconstraints = MTK.get_constraints(sys) - (isnothing(jconstraints) || isempty(jconstraints)) && return nothing - - stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))]) - ctidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))]) - cons_dvs, cons_ps = MTK.process_constraint_system( - jconstraints, Set(unknowns(sys)), parameters(sys), iv; validate = false) - - auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in cons_dvs]) - jconstraints = substitute_casadi_vars(model, sys, pmap, jconstraints; is_free_t, auxmap) - # Manually substitute fixed-t variables - for (i, cons) in enumerate(jconstraints) - consvars = MTK.vars(cons) - for st in consvars + sts = unknowns(sys) + cts = MTK.unbound_inputs(sys) + + 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.fast_substitute(c, auxmap), exprs) + exprs = map(c -> Symbolics.fast_substitute(c, Dict(pmap)), exprs) + # tf means different things in different contexts; a [tf] in a cost function + # should be tₛ, while a x(tf) should translate to x[1] + if is_free_t + free_t_map = Dict([[x(tₛ) => U.u[i, end] for (i, x) in enumerate(x_ops)]; + [c(tₛ) => V.u[i, end] for (i, c) in enumerate(c_ops)]]) + exprs = map(c -> Symbolics.fast_substitute(c, free_t_map), exprs) + end + + exprs = substitute_fixed_t_vars(exprs) + + # for variables like x(t) + whole_interval_map = Dict([[v => U.u[i, :] for (i, v) in enumerate(sts)]; + [v => V.u[i, :] for (i, v) in enumerate(cts)]]) + exprs = map(c -> Symbolics.fast_substitute(c, whole_interval_map), exprs) + exprs +end + +function substitute_fixed_t_vars(exprs) + for i in 1:length(exprs) + subvars = MTK.vars(exprs[i]) + for st in subvars MTK.iscall(st) || continue - x = MTK.operation(st) - t = only(MTK.arguments(st)) + x = operation(st) + t = only(arguments(st)) MTK.symbolic_type(t) === MTK.NotSymbolic() || continue if haskey(stidxmap, x(iv)) idx = stidxmap[x(iv)] @@ -201,52 +188,19 @@ function add_user_constraints!(model::CasADiModel, sys, tspan, pmap; is_free_t) idx = ctidxmap[x(iv)] cv = V end - cons = Symbolics.substitute(cons, Dict(x(t) => cv(t)[idx])) - end - - if cons isa Equation - subject_to!(opti, cons.lhs - cons.rhs == 0) - elseif cons.relational_op === Symbolics.geq - subject_to!(opti, cons.lhs - cons.rhs ≥ 0) - else - subject_to!(opti, cons.lhs - cons.rhs ≤ 0) + exprs[i] = Symbolics.fast_substitute(exprs[i], Dict(x(t) => cv(t)[idx])) end + jcosts = Symbolics.substitute(jcosts, Dict(x(t) => cv(t)[idx])) end end -function add_cost_function!(model::CasADiModel, sys, tspan, pmap; is_free_t) - @unpack opti, U, V, tₛ = model - jcosts = cost(sys) - if Symbolics._iszero(jcosts) - minimize!(opti, MX(0)) - return - end - - iv = MTK.get_iv(sys) - stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))]) - ctidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))]) - - jcosts = substitute_casadi_vars(model, sys, pmap, [jcosts]; is_free_t)[1] - # Substitute fixed-time variables. - costvars = MTK.vars(jcosts) - for st in costvars - MTK.iscall(st) || continue - x = operation(st) - t = only(arguments(st)) - MTK.symbolic_type(t) === MTK.NotSymbolic() || continue - if haskey(stidxmap, x(iv)) - idx = stidxmap[x(iv)] - cv = U - else - idx = ctidxmap[x(iv)] - cv = V - end - jcosts = Symbolics.substitute(jcosts, Dict(x(t) => cv(t)[idx])) - end +MTK.substitute_differentials(model::CasADiModel, exprs, args...) = exprs +function MTK.substitute_integral(model::CasADiModel, exprs) + @unpack U, opti = model dt = U.t[2] - U.t[1] intmap = Dict() - for int in MTK.collect_applied_operators(jcosts, Symbolics.Integral) + for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) op = MTK.operation(int) arg = only(arguments(MTK.value(int))) lo, hi = (op.domain.domain.left, op.domain.domain.right) @@ -255,39 +209,11 @@ function add_cost_function!(model::CasADiModel, sys, tspan, pmap; is_free_t) # Approximate integral as sum. intmap[int] = dt * tₛ * sum(arg) end - jcosts = Symbolics.substitute(jcosts, intmap) - jcosts = MTK.value(jcosts) - minimize!(opti, MX(jcosts)) -end - -function substitute_casadi_vars( - model::CasADiModel, sys, pmap, exprs; auxmap::Dict = Dict(), is_free_t) - @unpack opti, U, V, tₛ = model - iv = MTK.get_iv(sys) - sts = unknowns(sys) - cts = MTK.unbound_inputs(sys) - - 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.fast_substitute(c, auxmap), exprs) - exprs = map(c -> Symbolics.fast_substitute(c, Dict(pmap)), exprs) - # tf means different things in different contexts; a [tf] in a cost function - # should be tₛ, while a x(tf) should translate to x[1] - if is_free_t - free_t_map = Dict([[x(tₛ) => U.u[i, end] for (i, x) in enumerate(x_ops)]; - [c(tₛ) => V.u[i, end] for (i, c) in enumerate(c_ops)]]) - exprs = map(c -> Symbolics.fast_substitute(c, free_t_map), exprs) - end - - # for variables like x(t) - whole_interval_map = Dict([[v => U.u[i, :] for (i, v) in enumerate(sts)]; - [v => V.u[i, :] for (i, v) in enumerate(cts)]]) - exprs = map(c -> Symbolics.fast_substitute(c, whole_interval_map), exprs) - exprs + exprs = map(c -> Symbolics.substitute(c, intmap), exprs) + exprs = MTK.value.(exprs) end -function add_solve_constraints(prob, tableau) +function add_solve_constraints!(prob, tableau) @unpack A, α, c = tableau @unpack model, f, p = prob @unpack opti, U, V, tₛ = model @@ -332,28 +258,22 @@ function add_solve_constraints(prob, tableau) solver_opti end -""" - solve(prob::CasADiDynamicOptProblem, casadi_solver, ode_solver; plugin_options, solver_options, silent) - -`plugin_options` and `solver_options` get propagated to the Opti object in CasADi. - -NOTE: the solver should be passed in as a string to CasADi. "ipopt" -""" -function DiffEqBase.solve( - prob::CasADiDynamicOptProblem, solver::Union{String, Symbol} = "ipopt", - tableau_getter = MTK.constructDefault; plugin_options::Dict = Dict(), - solver_options::Dict = Dict(), silent = false) - @unpack model, u0, p, tspan, f = prob - tableau = tableau_getter() - @unpack opti, U, V, tₛ = model - +function MTK.prepare_solver() opti = add_solve_constraints(prob, tableau) - silent && (solver_options["print_level"] = 0) solver!(opti, "$solver", plugin_options, solver_options) +end +function MTK.get_U_values() + U_vals = value_getter(U.u) + size(U_vals, 2) == 1 && (U_vals = U_vals') + U_vals = [[U_vals[i, j] for i in 1:size(U_vals, 1)] for j in 1:length(ts)] +end +function MTK.get_V_values() +end +function MTK.get_t_values() + ts = value_getter(tₛ) * U.t +end - failed = false - value_getter = nothing - sol = nothing +function MTK.optimize_model!() try sol = CasADi.solve!(opti) value_getter = x -> CasADi.value(sol, x) @@ -361,28 +281,6 @@ function DiffEqBase.solve( value_getter = x -> CasADi.debug_value(opti, x) failed = true end - - ts = value_getter(tₛ) * U.t - U_vals = value_getter(U.u) - size(U_vals, 2) == 1 && (U_vals = U_vals') - U_vals = [[U_vals[i, j] for i in 1:size(U_vals, 1)] for j in 1:length(ts)] - ode_sol = DiffEqBase.build_solution(prob, tableau_getter, ts, U_vals) - - input_sol = nothing - if prod(size(V.u)) != 0 - V_vals = value_getter(V.u) - size(V_vals, 2) == 1 && (V_vals = V_vals') - V_vals = [[V_vals[i, j] for i in 1:size(V_vals, 1)] for j in 1:length(ts)] - input_sol = DiffEqBase.build_solution(prob, tableau_getter, ts, V_vals) - end - - if failed - ode_sol = SciMLBase.solution_new_retcode( - ode_sol, SciMLBase.ReturnCode.ConvergenceFailure) - !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode( - input_sol, SciMLBase.ReturnCode.ConvergenceFailure)) - end - - DynamicOptSolution(model, ode_sol, input_sol) end +MTK.successful_solve() = true end diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index 6dd80232f8..6e4a24ecae 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -9,13 +9,21 @@ import SymbolicUtils import NaNMath const MTK = ModelingToolkit +struct InfiniteOptModel + model::InfiniteModel + U::Vector{<:AbstractVariableRef} + V::Vector{<:AbstractVariableRef} + tₛ::AbstractVariableRef + is_free_final::Bool +end + struct JuMPDynamicOptProblem{uType, tType, isinplace, P, F, K} <: AbstractDynamicOptProblem{uType, tType, isinplace} f::F u0::uType tspan::tType p::P - model::InfiniteModel + model::InfiniteOptModel kwargs::K function JuMPDynamicOptProblem(f, u0, tspan, p, model, kwargs...) @@ -30,7 +38,7 @@ struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <: u0::uType tspan::tType p::P - model::InfiniteModel + model::InfiniteOptModel kwargs::K function InfiniteOptDynamicOptProblem(f, u0, tspan, p, model, kwargs...) @@ -39,27 +47,29 @@ struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <: end end -struct InfiniteOptModel - model::InfiniteModel - U::AbstractVariableRef - V::AbstractVariableRef - tₛ::Union - is_free_final::Bool -end - MTK.generate_internal_model(m::Type{InfiniteOptModel}) = InfiniteModel() -MTK.generate_state_variable!(m::InfiniteModel, u0::Vector, ns, nt) = @variable(m, U[i = 1:nt], Infinite(m[:t]), start=u0[i]) +MTK.generate_time_variable!(m::InfiniteModel, tspan, steps) = @infinite_parameter(m, t in [tspan[1], tspan[2]], num_supports = steps) +MTK.generate_state_variable!(m::InfiniteModel, u0::Vector, ns, nt) = @variable(m, U[i = 1:ns], Infinite(m[:t]), start=u0[i]) MTK.generate_input_variable!(m::InfiniteModel, c0, nc, nt) = @variable(m, V[i = 1:nc], Infinite(m[:t]), start=c0[i]) function MTK.generate_timescale!(m::InfiniteModel, guess, is_free_t) @variable(m, tₛ ≥ 0, start = guess) if !is_free_t - fix(m[:tₛ], 1) - set_start_value(m[:tₛ], 1) + fix(tₛ, 1, force=true) + set_start_value(tₛ, 1) end + tₛ end -MTK.add_constraint!(m::InfiniteOptModel, expr) = @constraint(m.model, expr) +function MTK.add_constraint!(m::InfiniteOptModel, expr::Union{Equation, Inequality}) + if expr isa Equation + @constraint(m.model, expr.lhs - expr.rhs == 0) + elseif expr.relational_op === Symbolics.geq + @constraint(m.model, expr.lhs - eq.rhs ≥ 0) + else + @constraint(m.model, expr.lhs - eq.rhs ≤ 0) + end +end MTK.set_objective!(m::InfiniteOptModel, expr) = @objective(m.model, Min, expr) """ @@ -82,7 +92,7 @@ function MTK.JuMPDynamicOptProblem(sys::System, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - process_DynamicOptProblem(JuMPDynamicOptProblem, InfiniteOptModel, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + MTK.process_DynamicOptProblem(JuMPDynamicOptProblem, InfiniteOptModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) end """ @@ -99,7 +109,7 @@ function MTK.InfiniteOptDynamicOptProblem(sys::System, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - prob = process_DynamicOptProblem(InfiniteOptDynamicOptProblem, InfiniteOptModel, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + MTK.process_DynamicOptProblem(InfiniteOptDynamicOptProblem, InfiniteOptModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) end function MTK.set_variable_bounds!(model, sys, pmap, tf) @@ -119,28 +129,28 @@ function MTK.set_variable_bounds!(model, sys, pmap, tf) end end - if symbolic_type(tf) === ScalarSymbolic() && hasbounds(tf) + if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() && hasbounds(tf) lo, hi = MTK.getbounds(tf) set_lower_bound(model.tₛ, lo) set_upper_bound(model.tₛ, hi) end end -function MTK.substitute_integral(model, jcosts) - for int in MTK.collect_applied_operators(jcosts, Symbolics.Integral) +function MTK.substitute_integral(model, exprs) + intmap = Dict() + for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) op = MTK.operation(int) arg = only(arguments(MTK.value(int))) - lo, hi = (op.domain.domain.left, op.domain.domain.right) - lo = MTK.value(lo) - hi = haskey(pmap, hi) ? 1 : MTK.value(hi) - intmap[int] = model.tₛ * InfiniteOpt.∫(arg, model.model.t, lo, hi) + lo, hi = MTK.value.((op.domain.domain.left, op.domain.domain.right)) + hi = (MTK.symbolic_type(hi) === MTK.ScalarSymbolic()) ? 1 : hi + intmap[int] = model.tₛ * InfiniteOpt.∫(arg, model.model[:t], lo, hi) end - jcosts = map(c -> Symbolics.substitute(c, intmap), jcosts) + exprs = map(c -> Symbolics.substitute(c, intmap), exprs) end -function MTK.add_initial_constraints!(model::InfiniteOptModel, u0, u0_idxs, ts) - U = model.U - @constraint(model, initial[i in u0_idxs], U[i](ts)==u0[i]) +function MTK.add_initial_constraints!(m::InfiniteOptModel, u0, u0_idxs, ts) + @show m.U + @constraint(m.model, initial[i in u0_idxs], m.U[i](ts)==u0[i]) end function MTK.substitute_model_vars(model, sys, exprs; tf = nothing) @@ -151,15 +161,15 @@ function MTK.substitute_model_vars(model, sys, exprs; tf = nothing) x_ops = [MTK.operation(MTK.unwrap(st)) for st in unknowns(sys)] c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in MTK.unbound_inputs(sys)] - if symbolic_type(tf) === ScalarSymbolic() - 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)]]) + if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() + free_t_map = Dict([[x(tf) => model.U[i](1) for (i, x) in enumerate(x_ops)]; + [c(tf) => model.V[i](1) for (i, c) in enumerate(c_ops)]]) exprs = map(c -> Symbolics.fast_substitute(c, free_t_map), exprs) end # for variables like x(1.0) - fixed_t_map = Dict([[x_ops[i] => U[i] for i in 1:length(U)]; - [c_ops[i] => V[i] for i in 1:length(V)]]) + fixed_t_map = Dict([[x_ops[i] => model.U[i] for i in 1:length(model.U)]; + [c_ops[i] => model.V[i] for i in 1:length(model.V)]]) exprs = map(c -> Symbolics.fast_substitute(c, fixed_t_map), exprs) end @@ -171,8 +181,8 @@ function MTK.substitute_differentials(model::InfiniteOptModel, eqs) map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs) end -function add_solve_constraints!(prob::JuMPDynamicOptProblem, solver) - @unpack A, α, c = solver.tableau +function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau) + @unpack A, α, c = tableau @unpack model, f, p = prob tsteps = supports(model.model[:t]) dt = tsteps[2] - tsteps[1] @@ -184,7 +194,7 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, solver) nᵥ = length(V) if MTK.is_explicit(tableau) K = Any[] - for τ in tsteps + for τ in tsteps[1:end-1] 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ᵤ] @@ -193,7 +203,7 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, solver) 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), + @constraint(model.model, [n = 1:nᵤ], U[n](τ) + ΔU[n]==U[n](τ + dt), base_name="solve_time_$τ") empty!(K) end @@ -201,14 +211,14 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, solver) @variable(model, K[1:length(α), 1:nᵤ], Infinite(t)) ΔUs = A * K ΔU_tot = dt * (K' * α) - for τ in tsteps + for τ in tsteps[1:end-1] for (i, h) in enumerate(c) Δ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]), + @constraint(model.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, tsteps[end])), + @constraint(model.model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tsteps[end])), DomainRestrictions(t => τ), base_name="solve_U($τ)") end end @@ -221,22 +231,22 @@ JuMP Collocation solver. Returns a DynamicOptSolution, which contains both the model and the ODE solution. """ -struct JuMPCollocation +struct JuMPCollocation <: AbstractCollocation solver::Any tableau::DiffEqBase.ODERKTableau end -JuMPCollocation(solver; tableau = MTK.constructDefault()) = JuMPCollocation(solver, tableau) +MTK.JuMPCollocation(solver, tableau = MTK.constructDefault()) = JuMPCollocation(solver, tableau) """ InfiniteOpt Collocation solver. - solver: an optimization solver such as Ipopt - `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()). """ -struct InfiniteOptCollocation +struct InfiniteOptCollocation <: AbstractCollocation solver::Any derivative_method::InfiniteOpt.AbstractDerivativeMethod end -InfiniteOptCollocation(solver; derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward())) = InfiniteOptCollocation(solver, derivative_method) +MTK.InfiniteOptCollocation(solver, derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward())) = InfiniteOptCollocation(solver, derivative_method) function MTK.prepare_solver!(prob::JuMPDynamicOptProblem, solver::JuMPCollocation; verbose = false, kwargs...) model = prob.model.model @@ -255,7 +265,7 @@ function MTK.prepare_solver!(prob::JuMPDynamicOptProblem, solver::JuMPCollocatio delete(model, var) end end - add_collocation_solve_constraints!(model, solver.tableau) + add_solve_constraints!(prob, solver.tableau) set_optimizer(model, solver.solver) end @@ -273,19 +283,20 @@ function MTK.optimize_model!(prob::Union{InfiniteOptDynamicOptProblem, JuMPDynam end function MTK.get_V_values(m::InfiniteOptModel) + nt = length(supports(m.model[:t])) if !isempty(m.V) V_vals = value.(m.V) - V_vals = [[V_vals[i][j] for i in 1:length(V_vals)] for j in 1:length(ts)] + V_vals = [[V_vals[i][j] for i in 1:length(V_vals)] for j in 1:nt] else nothing end end function MTK.get_U_values(m::InfiniteOptModel) + nt = length(supports(m.model[:t])) U_vals = value.(m.U) - U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:length(ts)] + U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:nt] end - -MTK.get_t_values(model) = prob.tspan[1] .+ model.tₛ * supports(model.model[:t]) +MTK.get_t_values(model) = model.tₛ * supports(model.model[:t]) function MTK.successful_solve(m::InfiniteOptModel) model = m.model diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl index 0e9962c241..9608c6083d 100644 --- a/src/ModelingToolkit.jl +++ b/src/ModelingToolkit.jl @@ -357,6 +357,8 @@ function FMIComponent end include("systems/optimal_control_interface.jl") export AbstractDynamicOptProblem, JuMPDynamicOptProblem, InfiniteOptDynamicOptProblem, CasADiDynamicOptProblem +export AbstractCollocation, JuMPCollocation, InfiniteOptCollocation, + CasADiCollocation export DynamicOptSolution @public apply_to_variables diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index 04c1a0c74b..3fb0fc91d9 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -165,23 +165,25 @@ function process_DynamicOptProblem(prob_type::Type{<:AbstractDynamicOptProblem}, dt = nothing, steps = nothing, guesses = Dict(), kwargs...) + warn_overdetermined(sys, u0map) + ctrls = unbound_inputs(sys) + states = unknowns(sys) - MTK.warn_overdetermined(sys, u0map) _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses)) - f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap; - t = tspan !== nothing ? tspan[1] : tspan, kwargs...) - stidxmap = Dict([v => i for (i, v) in enumerate(states)]) - u0map = Dict([MTK.default_toterm(MTK.value(k)) => v for (k, v) in u0map]) + u0map = Dict([default_toterm(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] - pmap = Dict{Any, Any}(pmap) - model_tspan, steps, is_free_t = MTK.process_tspan(tspan, dt, steps) + [stidxmap[default_toterm(k)] for (k, v) in u0map] - ctrls = MTK.unbound_inputs(sys) - states = unknowns(sys) - c0 = MTK.value.([pmap[c] for c in ctrls]) + f, u0, p = process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap; + t = tspan !== nothing ? tspan[1] : tspan, kwargs...) + model_tspan, steps, is_free_t = process_tspan(tspan, dt, steps) + pmap = recursive_unwrap(AnyDict(pmap)) + evaluate_varmap!(pmap, keys(pmap)) + c0 = value.([pmap[c] for c in ctrls]) + + tsteps = LinRange(model_tspan[1], model_tspan[2], steps) model = generate_internal_model(model_type) generate_time_variable!(model, model_tspan, steps) U = generate_state_variable!(model, u0, length(states), length(steps)) @@ -197,6 +199,7 @@ function process_DynamicOptProblem(prob_type::Type{<:AbstractDynamicOptProblem}, prob_type(f, u0, tspan, p, fullmodel, kwargs...) end +function generate_time_variable! end function generate_internal_model end function generate_state_variable! end function generate_input_variable! end @@ -207,8 +210,8 @@ function add_constraint! end is_free_final(model) = model.is_free_final function add_cost_function!(model, sys, tspan, pmap) - jcosts = copy(MTK.get_costs(sys)) - consolidate = MTK.get_consolidate(sys) + jcosts = copy(get_costs(sys)) + consolidate = get_consolidate(sys) if isnothing(jcosts) || isempty(jcosts) set_objective!(model, 0) return @@ -220,24 +223,19 @@ function add_cost_function!(model, sys, tspan, pmap) end function add_user_constraints!(model, sys, tspan, pmap) - conssys = MTK.get_constraintsystem(sys) - jconstraints = isnothing(conssys) ? nothing : MTK.get_constraints(conssys) + conssys = get_constraintsystem(sys) + jconstraints = isnothing(conssys) ? nothing : get_constraints(conssys) (isnothing(jconstraints) || isempty(jconstraints)) && return nothing - consvars = MTK.get_unknowns(conssys) + consvars = get_unknowns(conssys) is_free_final(model) && check_constraint_vars(consvars) - jconstraints = substitute_model_vars(model, sys, jcosts; tf = tspan[2]) jconstraints = substitute_toterm(consvars, jconstraints) jconstraints = substitute_params(pmap, jconstraints) + jconstraints = substitute_model_vars(model, sys, jconstraints; tf = tspan[2]) for c in jconstraints - if cons isa Equation - add_constraint!(model, c.lhs - c.rhs == 0) - elseif cons.relational_op === Symbolics.geq - add_constraint!(model, c.lhs - c.rhs ≥ 0) - else - add_constraint!(model, c.lhs - c.rhs ≤ 0) - end + @show c + add_constraint!(model, c) end end @@ -246,12 +244,12 @@ function add_equational_constraints!(model, sys, tspan) diff_eqs = substitute_model_vars(model, sys, diff_equations(sys); tf = tspan[2]) diff_eqs = substitute_differentials(model, sys, diff_eqs) for eq in diff_eqs - add_constraint!(model, eq.lhs == eq.rhs * model.tₛ) + add_constraint!(model, eq.lhs ~ eq.rhs * model.tₛ) end alg_eqs = substitute_model_vars(model, sys, alg_equations(sys); tf = tspan[2]) for eq in alg_eqs - add_constraint!(model, eq.lhs == eq.rhs * model.tₛ) + add_constraint!(model, eq.lhs ~ eq.rhs * model.tₛ) end end @@ -262,7 +260,7 @@ function substitute_integral end function substitute_differentials end function substitute_toterm(vars, exprs) - toterm_map = Dict([u => MTK.default_toterm(MTK.value(u)) for u in vars]) + toterm_map = Dict([u => default_toterm(value(u)) for u in vars]) exprs = map(c -> Symbolics.fast_substitute(c, toterm_map), exprs) end @@ -293,18 +291,24 @@ function get_U_values end function get_V_values end function successful_solve end +""" + solve(prob::AbstractDynamicOptProblem, solver::AbstractCollocation; verbose = false, kwargs...) + +- kwargs are used for other options. For example, the `plugin_options` and `solver_options` will propagated to the Opti object in CasADi. +""" function DiffEqBase.solve(prob::AbstractDynamicOptProblem, solver::AbstractCollocation; verbose = false, kwargs...) - solver = prepare_solver!(prob, solver) + solver = prepare_solver!(prob, solver; verbose, kwargs...) model = optimize_model!(prob, solver) ts = get_t_values(model) Us = get_U_values(model) Vs = get_V_values(model) + is_free_final(model) && (ts .+ tspan[1]) ode_sol = DiffEqBase.build_solution(prob, solver, ts, Us) input_sol = isnothing(Vs) ? nothing : DiffEqBase.build_solution(prob, solver, ts, Vs) - if !successful_solve(model) + if !successful_solve(model) ode_sol = SciMLBase.solution_new_retcode( ode_sol, SciMLBase.ReturnCode.ConvergenceFailure) !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode( diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index 03aeebf1bc..0d8ad71931 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -27,26 +27,22 @@ const M = ModelingToolkit # Test explicit method. jprob = JuMPDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01) - @test InfiniteOpt.num_constraints(jprob.model) == 2 # initials - jsol = solve(jprob, Ipopt.Optimizer, constructRK4, silent = true) + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRK4())) oprob = ODEProblem(sys, [u0map; parammap], tspan) osol = solve(oprob, SimpleRK4(), dt = 0.01) cprob = CasADiDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01) - csol = solve(cprob, "ipopt", constructRK4) + csol = solve(cprob, CasADiCollocation("ipopt", constructRK4())) @test jsol.sol.u ≈ osol.u @test csol.sol.u ≈ osol.u # Implicit method. - jsol2 = solve(jprob, Ipopt.Optimizer, constructImplicitEuler, silent = true) # 63.031 ms, 26.49 MiB - osol2 = solve(oprob, ImplicitEuler(), dt = 0.01, adaptive = false) # 129.375 μs, 61.91 KiB - jsol2 = solve(jprob, Ipopt.Optimizer, constructImplicitEuler, silent = true) # 63.031 ms, 26.49 MiB + osol2 = solve(oprob, ImplicitEuler(), dt = 0.01, adaptive = false) + jsol2 = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructImplicitEuler())) @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001) 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 + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test ≈(isol.sol.u, osol2.u, rtol = 0.001) - csol2 = solve(cprob, "ipopt", constructImplicitEuler, silent = true) + csol2 = solve(cprob, CasADiCollocation("ipopt", constructImplicitEuler())) @test ≈(csol2.sol.u, osol2.u, rtol = 0.001) # With a constraint @@ -57,20 +53,19 @@ const M = ModelingToolkit jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) @test InfiniteOpt.num_constraints(jprob.model) == 2 - jsol = solve(jprob, Ipopt.Optimizer, constructTsitouras5, silent = true) # 12.190 s, 9.68 GiB + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5())) @test jsol.sol(0.6; idxs = x(t)) ≈ 3.5 @test jsol.sol(0.3; idxs = x(t)) ≈ 7.0 cprob = CasADiDynamicOptProblem( lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) - csol = solve(cprob, "ipopt", constructTsitouras5, silent = true) + csol = solve(cprob, CasADiCollocation("ipopt", constructTsitouras5())) @test csol.sol(0.6; idxs = x(t)) ≈ 3.5 @test csol.sol(0.3; idxs = x(t)) ≈ 7.0 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 + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3))) # 48.564 ms, 9.58 MiB sol = isol.sol @test sol(0.6; idxs = x(t)) ≈ 3.5 @test sol(0.3; idxs = x(t)) ≈ 7.0 @@ -80,17 +75,16 @@ const M = ModelingToolkit @mtkcompile lksys = System(eqs, t; constraints = constr) iprob = InfiniteOptDynamicOptProblem( lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) - isol = solve(iprob, Ipopt.Optimizer, - derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true) + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3))) @test all(u -> u > [1, 1], isol.sol.u) jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) - jsol = solve(jprob, Ipopt.Optimizer, constructRadauIA3, silent = true) # 12.190 s, 9.68 GiB + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIA3())) @test all(u -> u > [1, 1], jsol.sol.u) cprob = CasADiDynamicOptProblem( lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) - csol = solve(cprob, "ipopt", constructRadauIA3, silent = true) + csol = solve(cprob, CasADiCollocation("ipopt", constructRadauIA3())) @test all(u -> u > [1, 1], csol.sol.u) end @@ -124,14 +118,14 @@ 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, constructVerner8, silent = true) + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructVerner8())) # Linear systems have bang-bang controls @test is_bangbang(jsol.input_sol, [-1.0], [1.0]) # Test reached final position. @test ≈(jsol.sol[x(t)][end], 0.25, rtol = 1e-5) cprob = CasADiDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01) - csol = solve(cprob, "ipopt", constructVerner8, silent = true) + csol = solve(cprob, CasADiCollocation("ipopt", constructVerner8())) # Linear systems have bang-bang controls @test is_bangbang(csol.input_sol, [-1.0], [1.0]) # Test reached final position. @@ -147,7 +141,7 @@ end @test ≈(csol.sol.u, osol.u, rtol = 0.05) iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01) - isol = solve(iprob, Ipopt.Optimizer; silent = true) + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test is_bangbang(isol.input_sol, [-1.0], [1.0]) @test ≈(isol.sol[x(t)][end], 0.25, rtol = 1e-5) osol = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false) @@ -171,13 +165,13 @@ 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, constructTsitouras5, silent = true) + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5())) @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) + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test is_bangbang(isol.input_sol, [0.0], [1.0]) cprob = CasADiDynamicOptProblem(beesys, u0map, tspan, pmap; dt = 0.01) - csol = solve(cprob, "ipopt", constructTsitouras5, silent = true) + csol = solve(cprob, CasADiCollocation("ipopt", constructTsitouras5())) @test is_bangbang(csol.input_sol, [0.0], [1.0]) @parameters (α_interp::LinearInterpolation)(..) @@ -265,15 +259,15 @@ 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, constructTsitouras5, silent = true) + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5())) @test isapprox(jsol.sol.t[end], 10.0, rtol = 1e-3) cprob = CasADiDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 201) - csol = solve(cprob, "ipopt", constructTsitouras5, silent = true) + csol = solve(cprob, CasADiCollocation("ipopt", constructTsitouras5())) @test isapprox(csol.sol.t[end], 10.0, rtol = 1e-3) iprob = InfiniteOptDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 200) - isol = solve(iprob, Ipopt.Optimizer, silent = true) + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3) @variables x(..) v(..) @@ -290,15 +284,15 @@ end 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, constructVerner8, silent = true) + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructVerner8())) @test isapprox(jsol.sol.t[end], 2.0, atol = 1e-5) cprob = CasADiDynamicOptProblem(block, u0map, (0, tf), parammap; steps = 51) - csol = solve(cprob, "ipopt", constructVerner8, silent = true) + csol = solve(cprob, CasADiCollocation("ipopt", constructVerner8())) @test isapprox(csol.sol.t[end], 2.0, atol = 1e-5) iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; steps = 51) - isol = solve(iprob, Ipopt.Optimizer, silent = true) + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isapprox(isol.sol.t[end], 2.0, atol = 1e-5) end @@ -332,7 +326,7 @@ 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, constructRK4, silent = true) + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRK4())) @test jsol.sol.u[end] ≈ [π, 0, 0, 0] cprob = CasADiDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04) @@ -340,6 +334,6 @@ end @test csol.sol.u[end] ≈ [π, 0, 0, 0] iprob = InfiniteOptDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04) - isol = solve(iprob, Ipopt.Optimizer, silent = true) + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isol.sol.u[end] ≈ [π, 0, 0, 0] end From 4ca744a72b52485c728d1d0fc84add4294b2ba58 Mon Sep 17 00:00:00 2001 From: vyudu Date: Sat, 17 May 2025 17:16:34 -0400 Subject: [PATCH 04/22] correctly implement interface --- ext/MTKCasADiDynamicOptExt.jl | 175 +++++++++++++---------- ext/MTKInfiniteOptExt.jl | 59 ++++---- src/systems/optimal_control_interface.jl | 48 ++++--- test/extensions/dynamic_optimization.jl | 9 +- 4 files changed, 164 insertions(+), 127 deletions(-) diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index 554d9faec4..1dbb776615 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -18,12 +18,17 @@ struct MXLinearInterpolation dt::Float64 end -struct CasADiModel - opti::Opti +mutable struct CasADiModel + model::Opti U::MXLinearInterpolation V::MXLinearInterpolation tₛ::MX is_free_final::Bool + solver_opti::Union{Nothing, Opti} + + function CasADiModel(opti, U, V, tₛ, is_free_final, solver_opti = nothing) + new(opti, U, V, tₛ, is_free_final, solver_opti) + end end struct CasADiDynamicOptProblem{uType, tType, isinplace, P, F, K} <: @@ -74,24 +79,27 @@ function MTK.CasADiDynamicOptProblem(sys::System, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - process_DynamicOptProblem(CasADiDynamicOptProblem, CasADiModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + MTK.process_DynamicOptProblem(CasADiDynamicOptProblem, CasADiModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) end -MTK.generate_internal_model(::Type{CasADiModel}) = CasADi.opti() +MTK.generate_internal_model(::Type{CasADiModel}) = CasADi.Opti() +MTK.generate_time_variable!(opti::Opti, args...) = nothing -function MTK.generate_state_variable(model::Opti, u0, ns, nt, tsteps) +function MTK.generate_state_variable!(model::Opti, u0, ns, tsteps) + nt = length(tsteps) U = CasADi.variable!(model, ns, nt) - set_initial!(opti, U, DM(repeat(u0, 1, steps))) + set_initial!(model, U, DM(repeat(u0, 1, nt))) MXLinearInterpolation(U, tsteps, tsteps[2] - tsteps[1]) end -function MTK.generate_input_variable(model::Opti, c0, nc, nt, tsteps) +function MTK.generate_input_variable!(model::Opti, c0, nc, tsteps) + nt = length(tsteps) V = CasADi.variable!(model, nc, nt) - !isempty(c0) && set_initial!(opti, V, DM(repeat(c0, 1, steps))) + !isempty(c0) && set_initial!(model, V, DM(repeat(c0, 1, nt))) MXLinearInterpolation(V, tsteps, tsteps[2] - tsteps[1]) end -function MTK.generate_timescale(model::Opti, guess, is_free_t) +function MTK.generate_timescale!(model::Opti, guess, is_free_t) if is_free_t tₛ = variable!(model) set_initial!(model, tₛ, guess) @@ -102,78 +110,73 @@ function MTK.generate_timescale(model::Opti, guess, is_free_t) end end -function MTK.add_constraint!(model::CasADiModel, expr) - @unpack opti = model - if cons isa Equation - subject_to!(opti, expr.lhs - expr.rhs == 0) - elseif cons.relational_op === Symbolics.geq - subject_to!(opti, expr.lhs - expr.rhs ≥ 0) +function MTK.add_constraint!(m::CasADiModel, expr) + if expr isa Equation + subject_to!(m.model, expr.lhs - expr.rhs == 0) + elseif expr.relational_op === Symbolics.geq + subject_to!(m.model, expr.lhs - expr.rhs ≥ 0) else - subject_to!(opti, expr.lhs - expr.rhs ≤ 0) + subject_to!(m.model, expr.lhs - expr.rhs ≤ 0) end end -MTK.set_objective!(model::CasADiModel, expr) = minimize!(model.opti, MX(expr)) +MTK.set_objective!(m::CasADiModel, expr) = minimize!(m.model, MX(expr)) -function MTK.set_variable_bounds!(model, sys, pmap, tf) - @unpack opti, U, V = model +function MTK.set_variable_bounds!(m::CasADiModel, sys, pmap, tf) + @unpack model, U, tₛ, V = m for (i, u) in enumerate(unknowns(sys)) if MTK.hasbounds(u) lo, hi = MTK.getbounds(u) - subject_to!(opti, Symbolics.fast_substitute(lo, pmap) <= U.u[i, :]) - subject_to!(opti, U.u[i, :] <= Symbolics.fast_substitute(hi, pmap)) + subject_to!(model, Symbolics.fixpoint_sub(lo, pmap) <= U.u[i, :]) + subject_to!(model, U.u[i, :] <= Symbolics.fixpoint_sub(hi, pmap)) end end for (i, v) in enumerate(MTK.unbound_inputs(sys)) if MTK.hasbounds(v) lo, hi = MTK.getbounds(v) - subject_to!(opti, Symbolics.fast_substitute(lo, pmap) <= V.u[i, :]) - subject_to!(opti, V.u[i, :] <= Symbolics.fast_substitute(hi, pmap)) + subject_to!(model, Symbolics.fixpoint_sub(lo, pmap) <= V.u[i, :]) + subject_to!(model, V.u[i, :] <= Symbolics.fixpoint_sub(hi, pmap)) end end if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() && hasbounds(tf) lo, hi = MTK.getbounds(tf) - subject_to!(opti, model.tₛ >= lo) - subject_to!(opti, model.tₛ <= hi) + subject_to!(model, tₛ >= lo) + subject_to!(model, tₛ <= hi) end end -function MTK.add_initial_constraints!(model::CasADiModel, u0, u0_idxs) - @unpack opti, U = model +function MTK.add_initial_constraints!(m::CasADiModel, u0, u0_idxs, args...) + @unpack model, U = m for i in u0_idxs - subject_to!(opti, U.u[i, 1] == u0[i]) + subject_to!(model, U.u[i, 1] == u0[i]) end end -function MTK.substitute_model_vars( - model::CasADiModel, sys, pmap, exprs; auxmap::Dict = Dict(), is_free_t) - @unpack opti, U, V, tₛ = model +function MTK.substitute_model_vars(m::CasADiModel, sys, exprs, tspan) + @unpack model, U, V, tₛ = m iv = MTK.get_iv(sys) sts = unknowns(sys) cts = MTK.unbound_inputs(sys) - 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.fast_substitute(c, auxmap), exprs) - exprs = map(c -> Symbolics.fast_substitute(c, Dict(pmap)), exprs) - # tf means different things in different contexts; a [tf] in a cost function - # should be tₛ, while a x(tf) should translate to x[1] - if is_free_t - free_t_map = Dict([[x(tₛ) => U.u[i, end] for (i, x) in enumerate(x_ops)]; - [c(tₛ) => V.u[i, end] for (i, c) in enumerate(c_ops)]]) + (ti, tf) = tspan + if MTK.is_free_final(m) + _tf = tₛ + ti + exprs = map(c -> Symbolics.fast_substitute(c, Dict(tf => _tf)), exprs) + free_t_map = Dict([[x(_tf) => U.u[i, end] for (i, x) in enumerate(x_ops)]; + [c(_tf) => V.u[i, end] for (i, c) in enumerate(c_ops)]]) exprs = map(c -> Symbolics.fast_substitute(c, free_t_map), exprs) end - exprs = substitute_fixed_t_vars(exprs) - - # for variables like x(t) + exprs = substitute_fixed_t_vars(m, sys, exprs) whole_interval_map = Dict([[v => U.u[i, :] for (i, v) in enumerate(sts)]; [v => V.u[i, :] for (i, v) in enumerate(cts)]]) exprs = map(c -> Symbolics.fast_substitute(c, whole_interval_map), exprs) - exprs end -function substitute_fixed_t_vars(exprs) +function substitute_fixed_t_vars(model::CasADiModel, sys, exprs) + stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))]) + ctidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))]) + iv = MTK.get_iv(sys) for i in 1:length(exprs) subvars = MTK.vars(exprs[i]) for st in subvars @@ -183,27 +186,28 @@ function substitute_fixed_t_vars(exprs) MTK.symbolic_type(t) === MTK.NotSymbolic() || continue if haskey(stidxmap, x(iv)) idx = stidxmap[x(iv)] - cv = U + cv = model.U else idx = ctidxmap[x(iv)] - cv = V + cv = model.V end exprs[i] = Symbolics.fast_substitute(exprs[i], Dict(x(t) => cv(t)[idx])) end jcosts = Symbolics.substitute(jcosts, Dict(x(t) => cv(t)[idx])) end + exprs end -MTK.substitute_differentials(model::CasADiModel, exprs, args...) = exprs +MTK.substitute_differentials(model::CasADiModel, sys, eqs) = exprs -function MTK.substitute_integral(model::CasADiModel, exprs) - @unpack U, opti = model +function MTK.substitute_integral(m::CasADiModel, exprs, tspan) + @unpack U, model, tₛ = m dt = U.t[2] - U.t[1] intmap = Dict() for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) op = MTK.operation(int) arg = only(arguments(MTK.value(int))) - lo, hi = (op.domain.domain.left, op.domain.domain.right) + lo, hi = MTK.value.((op.domain.domain.left, op.domain.domain.right)) !isequal((lo, hi), tspan) && error("Non-whole interval bounds for integrals are not currently supported for CasADiDynamicOptProblem.") # Approximate integral as sum. @@ -213,11 +217,11 @@ function MTK.substitute_integral(model::CasADiModel, exprs) exprs = MTK.value.(exprs) end -function add_solve_constraints!(prob, tableau) +function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau) @unpack A, α, c = tableau @unpack model, f, p = prob - @unpack opti, U, V, tₛ = model - solver_opti = copy(opti) + @unpack model, U, V, tₛ = model + solver_opti = copy(model) tsteps = U.t dt = tsteps[2] - tsteps[1] @@ -258,29 +262,56 @@ function add_solve_constraints!(prob, tableau) solver_opti end -function MTK.prepare_solver() - opti = add_solve_constraints(prob, tableau) - solver!(opti, "$solver", plugin_options, solver_options) +""" +CasADi Collocation solver. +- solver: an optimization solver such as Ipopt. Should be given as a string or symbol in all lowercase, e.g. "ipopt" +- tableau: An ODE RK tableau. Load a tableau by calling a function like `constructRK4` and may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. If this argument is not passed in, the solver will default to Radau second order. +""" +struct CasADiCollocation <: AbstractCollocation + solver::Union{String, Symbol} + tableau::DiffEqBase.ODERKTableau +end +MTK.CasADiCollocation(solver, tableau = MTK.constructDefault()) = CasADiCollocation(solver, tableau) + +function MTK.prepare_and_optimize!(prob::CasADiDynamicOptProblem, solver::CasADiCollocation; verbose = false, solver_options = Dict(), plugin_options = Dict(), kwargs...) + solver_opti = add_solve_constraints!(prob, solver.tableau) + verbose || (solver_options["print_level"] = 0) + solver!(solver_opti, "$(solver.solver)", plugin_options, solver_options) + try + CasADi.solve!(solver_opti) + catch ErrorException + end + prob.model.solver_opti = solver_opti end -function MTK.get_U_values() - U_vals = value_getter(U.u) + +function MTK.get_U_values(model::CasADiModel) + value_getter = MTK.successful_solve(model) ? CasADi.debug_value : CasADi.value + (nu, nt) = size(model.U.u) + U_vals = value_getter(model.solver_opti, model.U.u) size(U_vals, 2) == 1 && (U_vals = U_vals') - U_vals = [[U_vals[i, j] for i in 1:size(U_vals, 1)] for j in 1:length(ts)] + U_vals = [[U_vals[i, j] for i in 1:nu] for j in 1:nt] end -function MTK.get_V_values() + +function MTK.get_V_values(model::CasADiModel) + value_getter = MTK.successful_solve(model) ? CasADi.debug_value : CasADi.value + (nu, nt) = size(model.V.u) + if nu*nt != 0 + V_vals = value_getter(model.solver_opti, model.V.u) + size(V_vals, 2) == 1 && (V_vals = V_vals') + V_vals = [[V_vals[i, j] for i in 1:nu] for j in 1:nt] + else + nothing + end end -function MTK.get_t_values() - ts = value_getter(tₛ) * U.t + +function MTK.get_t_values(model::CasADiModel) + value_getter = MTK.successful_solve(model) ? CasADi.debug_value : CasADi.value + ts = value_getter(model.solver_opti, model.tₛ) .* model.U.t end -function MTK.optimize_model!() - try - sol = CasADi.solve!(opti) - value_getter = x -> CasADi.value(sol, x) - catch ErrorException - value_getter = x -> CasADi.debug_value(opti, x) - failed = true - end +function MTK.successful_solve(m::CasADiModel) + isnothing(m.solver_opti) && return false + retcode = CasADi.return_status(m.solver_opti) + retcode == "Solve_Succeeded" || retcode == "Solved_To_Acceptable_Level" end -MTK.successful_solve() = true end diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index 6e4a24ecae..4a2f025f00 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -49,8 +49,8 @@ end MTK.generate_internal_model(m::Type{InfiniteOptModel}) = InfiniteModel() MTK.generate_time_variable!(m::InfiniteModel, tspan, steps) = @infinite_parameter(m, t in [tspan[1], tspan[2]], num_supports = steps) -MTK.generate_state_variable!(m::InfiniteModel, u0::Vector, ns, nt) = @variable(m, U[i = 1:ns], Infinite(m[:t]), start=u0[i]) -MTK.generate_input_variable!(m::InfiniteModel, c0, nc, nt) = @variable(m, V[i = 1:nc], Infinite(m[:t]), start=c0[i]) +MTK.generate_state_variable!(m::InfiniteModel, u0::Vector, ns, ts) = @variable(m, U[i = 1:ns], Infinite(m[:t]), start=u0[i]) +MTK.generate_input_variable!(m::InfiniteModel, c0, nc, ts) = @variable(m, V[i = 1:nc], Infinite(m[:t]), start=c0[i]) function MTK.generate_timescale!(m::InfiniteModel, guess, is_free_t) @variable(m, tₛ ≥ 0, start = guess) @@ -65,9 +65,9 @@ function MTK.add_constraint!(m::InfiniteOptModel, expr::Union{Equation, Inequali if expr isa Equation @constraint(m.model, expr.lhs - expr.rhs == 0) elseif expr.relational_op === Symbolics.geq - @constraint(m.model, expr.lhs - eq.rhs ≥ 0) + @constraint(m.model, expr.lhs - expr.rhs ≥ 0) else - @constraint(m.model, expr.lhs - eq.rhs ≤ 0) + @constraint(m.model, expr.lhs - expr.rhs ≤ 0) end end MTK.set_objective!(m::InfiniteOptModel, expr) = @objective(m.model, Min, expr) @@ -109,10 +109,12 @@ function MTK.InfiniteOptDynamicOptProblem(sys::System, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - MTK.process_DynamicOptProblem(InfiniteOptDynamicOptProblem, InfiniteOptModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + prob = MTK.process_DynamicOptProblem(InfiniteOptDynamicOptProblem, InfiniteOptModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + MTK.add_equational_constraints!(prob.model, sys, pmap, tspan) + prob end -function MTK.set_variable_bounds!(model, sys, pmap, tf) +function MTK.set_variable_bounds!(model::InfiniteOptModel, sys, pmap, tf) for (i, u) in enumerate(unknowns(sys)) if MTK.hasbounds(u) lo, hi = MTK.getbounds(u) @@ -136,24 +138,27 @@ function MTK.set_variable_bounds!(model, sys, pmap, tf) end end -function MTK.substitute_integral(model, exprs) +function MTK.substitute_integral(model, exprs, tspan) intmap = Dict() for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) op = MTK.operation(int) arg = only(arguments(MTK.value(int))) lo, hi = MTK.value.((op.domain.domain.left, op.domain.domain.right)) - hi = (MTK.symbolic_type(hi) === MTK.ScalarSymbolic()) ? 1 : hi + if MTK.is_free_final(model) && isequal((lo, hi), tspan) + (lo, hi) = (0, 1) + elseif MTK.is_free_final(model) + error("Free final time problems cannot handle partial timespans.") + end intmap[int] = model.tₛ * InfiniteOpt.∫(arg, model.model[:t], lo, hi) end exprs = map(c -> Symbolics.substitute(c, intmap), exprs) end function MTK.add_initial_constraints!(m::InfiniteOptModel, u0, u0_idxs, ts) - @show m.U @constraint(m.model, initial[i in u0_idxs], m.U[i](ts)==u0[i]) end -function MTK.substitute_model_vars(model, sys, exprs; tf = nothing) +function MTK.substitute_model_vars(model::InfiniteOptModel, sys, exprs, tspan) whole_interval_map = Dict([[v => model.U[i] for (i, v) in enumerate(unknowns(sys))]; [v => model.V[i] for (i, v) in enumerate(MTK.unbound_inputs(sys))]]) exprs = map(c -> Symbolics.fast_substitute(c, whole_interval_map), exprs) @@ -161,9 +166,12 @@ function MTK.substitute_model_vars(model, sys, exprs; tf = nothing) x_ops = [MTK.operation(MTK.unwrap(st)) for st in unknowns(sys)] c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in MTK.unbound_inputs(sys)] + (ti, tf) = tspan if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() - free_t_map = Dict([[x(tf) => model.U[i](1) for (i, x) in enumerate(x_ops)]; - [c(tf) => model.V[i](1) for (i, c) in enumerate(c_ops)]]) + _tf = model.tₛ + ti + exprs = map(c -> Symbolics.fast_substitute(c, Dict(tf => _tf)), exprs) + free_t_map = Dict([[x(_tf) => model.U[i](1) for (i, x) in enumerate(x_ops)]; + [c(_tf) => model.V[i](1) for (i, c) in enumerate(c_ops)]]) exprs = map(c -> Symbolics.fast_substitute(c, free_t_map), exprs) end @@ -173,18 +181,19 @@ function MTK.substitute_model_vars(model, sys, exprs; tf = nothing) exprs = map(c -> Symbolics.fast_substitute(c, fixed_t_map), exprs) end -function MTK.substitute_differentials(model::InfiniteOptModel, eqs) +function MTK.substitute_differentials(model::InfiniteOptModel, sys, eqs) U = model.U t = model.model[:t] D = Differential(MTK.get_iv(sys)) diffsubmap = Dict([D(U[i]) => ∂(U[i], t) for i in 1:length(U)]) - map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs) + map(e -> Symbolics.substitute(e, diffsubmap), eqs) end function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau) @unpack A, α, c = tableau @unpack model, f, p = prob - tsteps = supports(model.model[:t]) + t = model.model[:t] + tsteps = supports(t) dt = tsteps[2] - tsteps[1] tₛ = model.tₛ @@ -208,7 +217,7 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau) empty!(K) end else - @variable(model, K[1:length(α), 1:nᵤ], Infinite(t)) + @variable(model.model, K[1:length(α), 1:nᵤ], Infinite(t)) ΔUs = A * K ΔU_tot = dt * (K' * α) for τ in tsteps[1:end-1] @@ -228,8 +237,6 @@ end JuMP Collocation solver. - solver: a optimization solver such as Ipopt - tableau: An ODE RK tableau. Load a tableau by calling a function like `constructRK4` and may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. If this argument is not passed in, the solver will default to Radau second order. - -Returns a DynamicOptSolution, which contains both the model and the ODE solution. """ struct JuMPCollocation <: AbstractCollocation solver::Any @@ -240,7 +247,7 @@ MTK.JuMPCollocation(solver, tableau = MTK.constructDefault()) = JuMPCollocation( """ InfiniteOpt Collocation solver. - solver: an optimization solver such as Ipopt -- `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()). +- `derivative_method`: 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()). """ struct InfiniteOptCollocation <: AbstractCollocation solver::Any @@ -248,7 +255,7 @@ struct InfiniteOptCollocation <: AbstractCollocation end MTK.InfiniteOptCollocation(solver, derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward())) = InfiniteOptCollocation(solver, derivative_method) -function MTK.prepare_solver!(prob::JuMPDynamicOptProblem, solver::JuMPCollocation; verbose = false, kwargs...) +function MTK.prepare_and_optimize!(prob::JuMPDynamicOptProblem, solver::JuMPCollocation; verbose = false, kwargs...) model = prob.model.model verbose || set_silent(model) # Unregister current solver constraints @@ -267,19 +274,15 @@ function MTK.prepare_solver!(prob::JuMPDynamicOptProblem, solver::JuMPCollocatio end add_solve_constraints!(prob, solver.tableau) set_optimizer(model, solver.solver) + optimize!(model) end -function MTK.prepare_solver!(prob::InfiniteOptDynamicOptProblem, solver::InfiniteOptCollocation; verbose = false, kwargs...) +function MTK.prepare_and_optimize!(prob::InfiniteOptDynamicOptProblem, solver::InfiniteOptCollocation; verbose = false, kwargs...) model = prob.model.model verbose || set_silent(model) - add_equational_constraints!(model, prob.f.sys, prob.tspan) set_derivative_method(model[:t], solver.derivative_method) set_optimizer(model, solver.solver) -end - -function MTK.optimize_model!(prob::Union{InfiniteOptDynamicOptProblem, JuMPDynamicOptProblem}, solver) - optimize!(prob.model.model) - prob.model + optimize!(model) end function MTK.get_V_values(m::InfiniteOptModel) @@ -296,7 +299,7 @@ function MTK.get_U_values(m::InfiniteOptModel) U_vals = value.(m.U) U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:nt] end -MTK.get_t_values(model) = model.tₛ * supports(model.model[:t]) +MTK.get_t_values(model) = value(model.tₛ) * supports(model.model[:t]) function MTK.successful_solve(m::InfiniteOptModel) model = m.model diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index 3fb0fc91d9..92e5467980 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -186,8 +186,8 @@ function process_DynamicOptProblem(prob_type::Type{<:AbstractDynamicOptProblem}, tsteps = LinRange(model_tspan[1], model_tspan[2], steps) model = generate_internal_model(model_type) generate_time_variable!(model, model_tspan, steps) - U = generate_state_variable!(model, u0, length(states), length(steps)) - V = generate_input_variable!(model, c0, length(ctrls), length(steps)) + U = generate_state_variable!(model, u0, length(states), tsteps) + V = generate_input_variable!(model, c0, length(ctrls), tsteps) tₛ = generate_timescale!(model, get(pmap, tspan[2], tspan[2]), is_free_t) fullmodel = model_type(model, U, V, tₛ, is_free_t) @@ -216,9 +216,9 @@ function add_cost_function!(model, sys, tspan, pmap) set_objective!(model, 0) return end - jcosts = substitute_model_vars(model, sys, jcosts; tf = tspan[2]) + jcosts = substitute_model_vars(model, sys, jcosts, tspan) jcosts = substitute_params(pmap, jcosts) - jcosts = substitute_integral(model, jcosts) + jcosts = substitute_integral(model, jcosts, tspan) set_objective!(model, consolidate(jcosts)) end @@ -230,24 +230,24 @@ function add_user_constraints!(model, sys, tspan, pmap) is_free_final(model) && check_constraint_vars(consvars) jconstraints = substitute_toterm(consvars, jconstraints) + jconstraints = substitute_model_vars(model, sys, jconstraints, tspan) jconstraints = substitute_params(pmap, jconstraints) - jconstraints = substitute_model_vars(model, sys, jconstraints; tf = tspan[2]) for c in jconstraints - @show c add_constraint!(model, c) end end -function add_equational_constraints!(model, sys, tspan) - model = model.model - diff_eqs = substitute_model_vars(model, sys, diff_equations(sys); tf = tspan[2]) +function add_equational_constraints!(model, sys, pmap, tspan) + diff_eqs = substitute_model_vars(model, sys, diff_equations(sys), tspan) + diff_eqs = substitute_params(pmap, diff_eqs) diff_eqs = substitute_differentials(model, sys, diff_eqs) for eq in diff_eqs add_constraint!(model, eq.lhs ~ eq.rhs * model.tₛ) end - alg_eqs = substitute_model_vars(model, sys, alg_equations(sys); tf = tspan[2]) + alg_eqs = substitute_model_vars(model, sys, alg_equations(sys), tspan) + alg_eqs = substitute_params(pmap, alg_eqs) for eq in alg_eqs add_constraint!(model, eq.lhs ~ eq.rhs * model.tₛ) end @@ -256,6 +256,12 @@ end function set_objective! end """Substitute variables like x(1.5) with the corresponding model variables.""" function substitute_model_vars end +""" +Substitute integrals. For an integral from (ts, te): +- Free final time problems should transcribe this to (0, 1) in the case that (ts, te) is the original timespan. Free final time problems cannot handle partial timespans. +- CasADi cannot handle partial timespans, even for non-free-final time problems. +time problems and unchanged otherwise. +""" function substitute_integral end function substitute_differentials end @@ -265,7 +271,7 @@ function substitute_toterm(vars, exprs) end function substitute_params(pmap, exprs) - exprs = map(c -> Symbolics.fast_substitute(c, Dict(pmap)), exprs) + exprs = map(c -> Symbolics.fixpoint_sub(c, Dict(pmap)), exprs) end function check_constraint_vars(vars) @@ -282,10 +288,9 @@ end ### SOLVER UTILITIES ### ######################## """ -Add the solve constraints, set the solver (Ipopt, e.g.) and solver options. +Add the solve constraints, set the solver (Ipopt, e.g.) and solver options, optimize the model. """ -function prepare_solver! end -function optimize_model! end +function prepare_and_optimize! end function get_t_values end function get_U_values end function get_V_values end @@ -297,22 +302,21 @@ function successful_solve end - kwargs are used for other options. For example, the `plugin_options` and `solver_options` will propagated to the Opti object in CasADi. """ function DiffEqBase.solve(prob::AbstractDynamicOptProblem, solver::AbstractCollocation; verbose = false, kwargs...) - solver = prepare_solver!(prob, solver; verbose, kwargs...) - model = optimize_model!(prob, solver) + prepare_and_optimize!(prob, solver; verbose, kwargs...) - ts = get_t_values(model) - Us = get_U_values(model) - Vs = get_V_values(model) - is_free_final(model) && (ts .+ tspan[1]) + ts = get_t_values(prob.model) + Us = get_U_values(prob.model) + Vs = get_V_values(prob.model) + is_free_final(prob.model) && (ts .+ prob.tspan[1]) ode_sol = DiffEqBase.build_solution(prob, solver, ts, Us) input_sol = isnothing(Vs) ? nothing : DiffEqBase.build_solution(prob, solver, ts, Vs) - if !successful_solve(model) + if !successful_solve(prob.model) ode_sol = SciMLBase.solution_new_retcode( ode_sol, SciMLBase.ReturnCode.ConvergenceFailure) !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode( input_sol, SciMLBase.ReturnCode.ConvergenceFailure)) end - DynamicOptSolution(model, ode_sol, input_sol) + DynamicOptSolution(prob.model.model, ode_sol, input_sol) end diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index 0d8ad71931..2d24f92d1f 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -52,7 +52,6 @@ const M = ModelingToolkit @mtkcompile lksys = System(eqs, t; constraints = constr) jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) - @test InfiniteOpt.num_constraints(jprob.model) == 2 jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5())) @test jsol.sol(0.6; idxs = x(t)) ≈ 3.5 @test jsol.sol(0.3; idxs = x(t)) ≈ 7.0 @@ -213,16 +212,16 @@ 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 = JuMPDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) - jsol = solve(jprob, Ipopt.Optimizer, constructRadauIIA5, silent = true) + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIIA5())) @test jsol.sol[h(t)][end] > 1.012 cprob = CasADiDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) - csol = solve(cprob, "ipopt"; silent = true) + csol = solve(cprob, CasADiCollocation("ipopt")) @test csol.sol[h(t)][end] > 1.012 iprob = InfiniteOptDynamicOptProblem( rocket, u0map, (ts, te), pmap; dt = 0.001) - isol = solve(iprob, Ipopt.Optimizer, silent = true) + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isol.sol[h(t)][end] > 1.012 # Test solution @@ -330,7 +329,7 @@ end @test jsol.sol.u[end] ≈ [π, 0, 0, 0] cprob = CasADiDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04) - csol = solve(cprob, "ipopt", constructRK4, silent = true) + csol = solve(cprob, CasADiCollocation("ipopt", constructRK4())) @test csol.sol.u[end] ≈ [π, 0, 0, 0] iprob = InfiniteOptDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04) From 8d268e9f1064f92f415a3a5426b9c9ec677bd442 Mon Sep 17 00:00:00 2001 From: vyudu Date: Mon, 19 May 2025 14:54:06 -0400 Subject: [PATCH 05/22] refactor: move set_variable_bounds to interface function --- ext/MTKCasADiDynamicOptExt.jl | 37 ++++------------- ext/MTKInfiniteOptExt.jl | 52 ++++++------------------ src/systems/optimal_control_interface.jl | 36 +++++++++++++--- 3 files changed, 51 insertions(+), 74 deletions(-) diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index 1dbb776615..2dbd184683 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -17,6 +17,7 @@ struct MXLinearInterpolation t::Vector{Float64} dt::Float64 end +Base.getindex(m::MXLinearInterpolation, i...) = length(i) == length(size(m.u)) ? m.u[i...] : m.u[i..., :] mutable struct CasADiModel model::Opti @@ -37,7 +38,7 @@ struct CasADiDynamicOptProblem{uType, tType, isinplace, P, F, K} <: u0::uType tspan::tType p::P - model::CasADiModel + wrapped_model::CasADiModel kwargs::K function CasADiDynamicOptProblem(f, u0, tspan, p, model, kwargs...) @@ -52,10 +53,11 @@ function (M::MXLinearInterpolation)(τ) Δ = nt - i + 1 (i > length(M.t) || i < 1) && error("Cannot extrapolate past the tspan.") + colons = ntuple(_ -> (:), length(size(M.u)) - 1) if i < length(M.t) - M.u[:, i] + Δ * (M.u[:, i + 1] - M.u[:, i]) + M.u[colons..., i] + Δ*(M.u[colons..., i+1] - M.u[colons..., i]) else - M.u[:, i] + M.u[colons..., i] end end @@ -121,29 +123,6 @@ function MTK.add_constraint!(m::CasADiModel, expr) end MTK.set_objective!(m::CasADiModel, expr) = minimize!(m.model, MX(expr)) -function MTK.set_variable_bounds!(m::CasADiModel, sys, pmap, tf) - @unpack model, U, tₛ, V = m - for (i, u) in enumerate(unknowns(sys)) - if MTK.hasbounds(u) - lo, hi = MTK.getbounds(u) - subject_to!(model, Symbolics.fixpoint_sub(lo, pmap) <= U.u[i, :]) - subject_to!(model, U.u[i, :] <= Symbolics.fixpoint_sub(hi, pmap)) - end - end - for (i, v) in enumerate(MTK.unbound_inputs(sys)) - if MTK.hasbounds(v) - lo, hi = MTK.getbounds(v) - subject_to!(model, Symbolics.fixpoint_sub(lo, pmap) <= V.u[i, :]) - subject_to!(model, V.u[i, :] <= Symbolics.fixpoint_sub(hi, pmap)) - end - end - if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() && hasbounds(tf) - lo, hi = MTK.getbounds(tf) - subject_to!(model, tₛ >= lo) - subject_to!(model, tₛ <= hi) - end -end - function MTK.add_initial_constraints!(m::CasADiModel, u0, u0_idxs, args...) @unpack model, U = m for i in u0_idxs @@ -219,8 +198,8 @@ end function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau) @unpack A, α, c = tableau - @unpack model, f, p = prob - @unpack model, U, V, tₛ = model + @unpack wrapped_model, f, p = prob + @unpack model, U, V, tₛ = wrapped_model solver_opti = copy(model) tsteps = U.t @@ -281,7 +260,7 @@ function MTK.prepare_and_optimize!(prob::CasADiDynamicOptProblem, solver::CasADi CasADi.solve!(solver_opti) catch ErrorException end - prob.model.solver_opti = solver_opti + prob.wrapped_model.solver_opti = solver_opti end function MTK.get_U_values(model::CasADiModel) diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index 4a2f025f00..db9eb3f3a0 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -23,7 +23,7 @@ struct JuMPDynamicOptProblem{uType, tType, isinplace, P, F, K} <: u0::uType tspan::tType p::P - model::InfiniteOptModel + wrapped_model::InfiniteOptModel kwargs::K function JuMPDynamicOptProblem(f, u0, tspan, p, model, kwargs...) @@ -38,7 +38,7 @@ struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <: u0::uType tspan::tType p::P - model::InfiniteOptModel + wrapped_model::InfiniteOptModel kwargs::K function InfiniteOptDynamicOptProblem(f, u0, tspan, p, model, kwargs...) @@ -110,34 +110,10 @@ function MTK.InfiniteOptDynamicOptProblem(sys::System, u0map, tspan, pmap; steps = nothing, guesses = Dict(), kwargs...) prob = MTK.process_DynamicOptProblem(InfiniteOptDynamicOptProblem, InfiniteOptModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) - MTK.add_equational_constraints!(prob.model, sys, pmap, tspan) + MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) prob end -function MTK.set_variable_bounds!(model::InfiniteOptModel, sys, pmap, tf) - for (i, u) in enumerate(unknowns(sys)) - if MTK.hasbounds(u) - lo, hi = MTK.getbounds(u) - set_lower_bound(model.U[i], Symbolics.fixpoint_sub(lo, pmap)) - set_upper_bound(model.U[i], Symbolics.fixpoint_sub(hi, pmap)) - end - end - - for (i, v) in enumerate(MTK.unbound_inputs(sys)) - if MTK.hasbounds(v) - lo, hi = MTK.getbounds(v) - set_lower_bound(model.V[i], Symbolics.fixpoint_sub(lo, pmap)) - set_upper_bound(model.V[i], Symbolics.fixpoint_sub(hi, pmap)) - end - end - - if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() && hasbounds(tf) - lo, hi = MTK.getbounds(tf) - set_lower_bound(model.tₛ, lo) - set_upper_bound(model.tₛ, hi) - end -end - function MTK.substitute_integral(model, exprs, tspan) intmap = Dict() for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) @@ -191,14 +167,12 @@ end function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau) @unpack A, α, c = tableau - @unpack model, f, p = prob - t = model.model[:t] + @unpack wrapped_model, f, p = prob + @unpack tₛ, U, V, model = wrapped_model + t = model[:t] tsteps = supports(t) dt = tsteps[2] - tsteps[1] - tₛ = model.tₛ - U = model.U - V = model.V nᵤ = length(U) nᵥ = length(V) if MTK.is_explicit(tableau) @@ -212,22 +186,22 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau) push!(K, Kₙ) end ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)]) - @constraint(model.model, [n = 1:nᵤ], U[n](τ) + ΔU[n]==U[n](τ + dt), + @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU[n]==U[n](τ + dt), base_name="solve_time_$τ") empty!(K) end else - @variable(model.model, K[1:length(α), 1:nᵤ], Infinite(t)) + K = @variable(model, K[1:length(α), 1:nᵤ], Infinite(model[:t])) ΔUs = A * K ΔU_tot = dt * (K' * α) for τ in tsteps[1:end-1] for (i, h) in enumerate(c) ΔU = @view ΔUs[i, :] - Uₙ = U + ΔU * h * dt - @constraint(model.model, [j = 1:nᵤ], K[i, j]==(tₛ * f(Uₙ, V, p, τ + h * dt)[j]), + Uₙ = U + ΔU * 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.model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tsteps[end])), + @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tsteps[end])), DomainRestrictions(t => τ), base_name="solve_U($τ)") end end @@ -256,7 +230,7 @@ end MTK.InfiniteOptCollocation(solver, derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward())) = InfiniteOptCollocation(solver, derivative_method) function MTK.prepare_and_optimize!(prob::JuMPDynamicOptProblem, solver::JuMPCollocation; verbose = false, kwargs...) - model = prob.model.model + model = prob.wrapped_model.model verbose || set_silent(model) # Unregister current solver constraints for con in all_constraints(model) @@ -278,7 +252,7 @@ function MTK.prepare_and_optimize!(prob::JuMPDynamicOptProblem, solver::JuMPColl end function MTK.prepare_and_optimize!(prob::InfiniteOptDynamicOptProblem, solver::InfiniteOptCollocation; verbose = false, kwargs...) - model = prob.model.model + model = prob.wrapped_model.model verbose || set_silent(model) set_derivative_method(model[:t], solver.derivative_method) set_optimizer(model, solver.solver) diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index 92e5467980..b91c76182f 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -207,6 +207,30 @@ function generate_timescale! end function set_variable_bounds! end function add_initial_constraints! end function add_constraint! end + +function set_variable_bounds!(m, sys, pmap, tf) + @unpack model, U, V, tₛ = m + for (i, u) in enumerate(unknowns(sys)) + if hasbounds(u) + lo, hi = getbounds(u) + add_constraint!(m, U[i] ≳ Symbolics.fixpoint_sub(lo, pmap)) + add_constraint!(m, U[i] ≲ Symbolics.fixpoint_sub(hi, pmap)) + end + end + for (i, v) in enumerate(unbound_inputs(sys)) + if hasbounds(v) + lo, hi = getbounds(v) + add_constraint!(m, V[i] ≳ Symbolics.fixpoint_sub(lo, pmap)) + add_constraint!(m, V[i] ≲ Symbolics.fixpoint_sub(hi, pmap)) + end + end + if symbolic_type(tf) === ScalarSymbolic() && hasbounds(tf) + lo, hi = getbounds(tf) + set_lower_bound(tₛ, Symbolics.fixpoint_sub(lo, pmap)) + set_upper_bound(tₛ, Symbolics.fixpoint_sub(hi, pmap)) + end +end + is_free_final(model) = model.is_free_final function add_cost_function!(model, sys, tspan, pmap) @@ -304,19 +328,19 @@ function successful_solve end function DiffEqBase.solve(prob::AbstractDynamicOptProblem, solver::AbstractCollocation; verbose = false, kwargs...) prepare_and_optimize!(prob, solver; verbose, kwargs...) - ts = get_t_values(prob.model) - Us = get_U_values(prob.model) - Vs = get_V_values(prob.model) - is_free_final(prob.model) && (ts .+ prob.tspan[1]) + ts = get_t_values(prob.wrapped_model) + Us = get_U_values(prob.wrapped_model) + Vs = get_V_values(prob.wrapped_model) + is_free_final(prob.wrapped_model) && (ts .+ prob.tspan[1]) ode_sol = DiffEqBase.build_solution(prob, solver, ts, Us) input_sol = isnothing(Vs) ? nothing : DiffEqBase.build_solution(prob, solver, ts, Vs) - if !successful_solve(prob.model) + if !successful_solve(prob.wrapped_model) ode_sol = SciMLBase.solution_new_retcode( ode_sol, SciMLBase.ReturnCode.ConvergenceFailure) !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode( input_sol, SciMLBase.ReturnCode.ConvergenceFailure)) end - DynamicOptSolution(prob.model.model, ode_sol, input_sol) + DynamicOptSolution(prob.wrapped_model.model, ode_sol, input_sol) end From b5b0f4456da373a2c864279ccc170684225469de Mon Sep 17 00:00:00 2001 From: vyudu Date: Mon, 19 May 2025 18:00:34 -0400 Subject: [PATCH 06/22] refactor: centralize substitute_differentiasl and substitute_integral --- ext/MTKCasADiDynamicOptExt.jl | 59 ++----- ext/MTKInfiniteOptExt.jl | 66 +++---- ext/MTKPyomoDynamicOptExt.jl | 209 +++++++++++++++++++++++ src/systems/optimal_control_interface.jl | 61 +++++-- 4 files changed, 299 insertions(+), 96 deletions(-) create mode 100644 ext/MTKPyomoDynamicOptExt.jl diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index 2dbd184683..f70b69a59d 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -130,32 +130,20 @@ function MTK.add_initial_constraints!(m::CasADiModel, u0, u0_idxs, args...) end end -function MTK.substitute_model_vars(m::CasADiModel, sys, exprs, tspan) - @unpack model, U, V, tₛ = m - iv = MTK.get_iv(sys) - sts = unknowns(sys) - cts = MTK.unbound_inputs(sys) - x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts] - c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts] - (ti, tf) = tspan - if MTK.is_free_final(m) - _tf = tₛ + ti - exprs = map(c -> Symbolics.fast_substitute(c, Dict(tf => _tf)), exprs) - free_t_map = Dict([[x(_tf) => U.u[i, end] for (i, x) in enumerate(x_ops)]; - [c(_tf) => V.u[i, end] for (i, c) in enumerate(c_ops)]]) - exprs = map(c -> Symbolics.fast_substitute(c, free_t_map), exprs) - end +function free_t_map(model, tf, x_ops, c_ops) + Dict([[x(tf) => model.U.u[i, end] for (i, x) in enumerate(x_ops)]; + [c(tf) => model.V.u[i, end] for (i, c) in enumerate(c_ops)]]) +end + +function whole_t_map(model, sys) + Dict([[v => model.U.u[i, :] for (i, v) in enumerate(unknowns(sys))]; + [v => model.V.u[i, :] for (i, v) in enumerate(MTK.unbound_inputs(sys))]]) - exprs = substitute_fixed_t_vars(m, sys, exprs) - whole_interval_map = Dict([[v => U.u[i, :] for (i, v) in enumerate(sts)]; - [v => V.u[i, :] for (i, v) in enumerate(cts)]]) - exprs = map(c -> Symbolics.fast_substitute(c, whole_interval_map), exprs) end -function substitute_fixed_t_vars(model::CasADiModel, sys, exprs) - stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))]) - ctidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))]) - iv = MTK.get_iv(sys) +function fixed_t_map(model::CasADiModel, x_ops, c_ops, exprs) + stidxmap = Dict([v => i for (i, v) in x_ops]) + ctidxmap = Dict([v => i for (i, v) in c_ops]) for i in 1:length(exprs) subvars = MTK.vars(exprs[i]) for st in subvars @@ -163,11 +151,11 @@ function substitute_fixed_t_vars(model::CasADiModel, sys, exprs) x = operation(st) t = only(arguments(st)) MTK.symbolic_type(t) === MTK.NotSymbolic() || continue - if haskey(stidxmap, x(iv)) - idx = stidxmap[x(iv)] + if haskey(stidxmap, x) + idx = stidxmap[x] cv = model.U else - idx = ctidxmap[x(iv)] + idx = ctidxmap[x] cv = model.V end exprs[i] = Symbolics.fast_substitute(exprs[i], Dict(x(t) => cv(t)[idx])) @@ -177,24 +165,7 @@ function substitute_fixed_t_vars(model::CasADiModel, sys, exprs) exprs end -MTK.substitute_differentials(model::CasADiModel, sys, eqs) = exprs - -function MTK.substitute_integral(m::CasADiModel, exprs, tspan) - @unpack U, model, tₛ = m - dt = U.t[2] - U.t[1] - intmap = Dict() - for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) - op = MTK.operation(int) - arg = only(arguments(MTK.value(int))) - lo, hi = MTK.value.((op.domain.domain.left, op.domain.domain.right)) - !isequal((lo, hi), tspan) && - error("Non-whole interval bounds for integrals are not currently supported for CasADiDynamicOptProblem.") - # Approximate integral as sum. - intmap[int] = dt * tₛ * sum(arg) - end - exprs = map(c -> Symbolics.substitute(c, intmap), exprs) - exprs = MTK.value.(exprs) -end +MTK.lowered_integral(model, expr, args...) = model.tₛ * (model.U.t[2] - model.U.t[1]) * expr function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau) @unpack A, α, c = tableau diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index db9eb3f3a0..61f4ec70b3 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -48,7 +48,7 @@ struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <: end MTK.generate_internal_model(m::Type{InfiniteOptModel}) = InfiniteModel() -MTK.generate_time_variable!(m::InfiniteModel, tspan, steps) = @infinite_parameter(m, t in [tspan[1], tspan[2]], num_supports = steps) +MTK.generate_time_variable!(m::InfiniteModel, tspan, steps) = @infinite_parameter(m, t in [tspan[1], tspan[2]], num_supports = length(tsteps)) MTK.generate_state_variable!(m::InfiniteModel, u0::Vector, ns, ts) = @variable(m, U[i = 1:ns], Infinite(m[:t]), start=u0[i]) MTK.generate_input_variable!(m::InfiniteModel, c0, nc, ts) = @variable(m, V[i = 1:nc], Infinite(m[:t]), start=c0[i]) @@ -114,55 +114,39 @@ function MTK.InfiniteOptDynamicOptProblem(sys::System, u0map, tspan, pmap; prob end -function MTK.substitute_integral(model, exprs, tspan) - intmap = Dict() - for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) - op = MTK.operation(int) - arg = only(arguments(MTK.value(int))) - lo, hi = MTK.value.((op.domain.domain.left, op.domain.domain.right)) - if MTK.is_free_final(model) && isequal((lo, hi), tspan) - (lo, hi) = (0, 1) - elseif MTK.is_free_final(model) - error("Free final time problems cannot handle partial timespans.") - end - intmap[int] = model.tₛ * InfiniteOpt.∫(arg, model.model[:t], lo, hi) +MTK.lowered_integral(model, expr, lo, hi) = model.tₛ * InfiniteOpt.∫(arg, model.model[:t], lo, hi) + +function MTK.process_integral_bounds(model, integral_span, tspan) + if MTK.is_free_final(model) && isequal(integral_span, tspan) + integral_span = (0, 1) + elseif MTK.is_free_final(model) + error("Free final time problems cannot handle partial timespans.") + else + integral_span end - exprs = map(c -> Symbolics.substitute(c, intmap), exprs) end function MTK.add_initial_constraints!(m::InfiniteOptModel, u0, u0_idxs, ts) - @constraint(m.model, initial[i in u0_idxs], m.U[i](ts)==u0[i]) + for i in u0_idxs + fix(m.U[i], u0[i], force = true) + end end -function MTK.substitute_model_vars(model::InfiniteOptModel, sys, exprs, tspan) - whole_interval_map = Dict([[v => model.U[i] for (i, v) in enumerate(unknowns(sys))]; - [v => model.V[i] for (i, v) in enumerate(MTK.unbound_inputs(sys))]]) - exprs = map(c -> Symbolics.fast_substitute(c, whole_interval_map), exprs) - - x_ops = [MTK.operation(MTK.unwrap(st)) for st in unknowns(sys)] - c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in MTK.unbound_inputs(sys)] - - (ti, tf) = tspan - if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() - _tf = model.tₛ + ti - exprs = map(c -> Symbolics.fast_substitute(c, Dict(tf => _tf)), exprs) - free_t_map = Dict([[x(_tf) => model.U[i](1) for (i, x) in enumerate(x_ops)]; - [c(_tf) => model.V[i](1) for (i, c) in enumerate(c_ops)]]) - exprs = map(c -> Symbolics.fast_substitute(c, free_t_map), exprs) - end +MTK.lowered_derivative(model, i) = ∂(model.U[i], model.model[:t]) + +function MTK.fixed_t_map(model::InfiniteOptModel, x_ops, c_ops, exprs) + Dict([[x_ops[i] => model.U[i] for i in 1:length(model.U)]; + [c_ops[i] => model.V[i] for i in 1:length(model.V)]]) +end - # for variables like x(1.0) - fixed_t_map = Dict([[x_ops[i] => model.U[i] for i in 1:length(model.U)]; - [c_ops[i] => model.V[i] for i in 1:length(model.V)]]) - exprs = map(c -> Symbolics.fast_substitute(c, fixed_t_map), exprs) +function MTK.free_t_map(model::InfiniteOptModel, tf, x_ops, c_ops) + Dict([[x(tf) => model.U[i](1) for (i, x) in enumerate(x_ops)]; + [c(tf) => model.V[i](1) for (i, c) in enumerate(c_ops)]]) end -function MTK.substitute_differentials(model::InfiniteOptModel, sys, eqs) - U = model.U - t = model.model[:t] - D = Differential(MTK.get_iv(sys)) - diffsubmap = Dict([D(U[i]) => ∂(U[i], t) for i in 1:length(U)]) - map(e -> Symbolics.substitute(e, diffsubmap), eqs) +function MTK.whole_t_map(model::InfiniteOptModel, sys) + whole_interval_map = Dict([[v => model.U[i] for (i, v) in enumerate(unknowns(sys))]; + [v => model.V[i] for (i, v) in enumerate(MTK.unbound_inputs(sys))]]) end function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau) diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl new file mode 100644 index 0000000000..859c17d72a --- /dev/null +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -0,0 +1,209 @@ +module MTKPyomoDynamicOptExt +using ModelingToolkit +using PythonCall +using DiffEqBase +using UnPack +using NaNMath +const MTK = ModelingToolkit + +# import pyomo +const pyomo = PythonCall.pynew() +PythonCall.pycopy!(pyomo, pyimport("pyomo.environ")) + +struct PyomoDAEVar + v::Py +end +(v::PyomoDAEVar)(t) = v.v[:, t] +getindex(v::PyomoDAEVar, i::Union{Num, Symbolic}, t::Union{Num, Symbolic}) = wrap(Term{symeltype(A)}(getindex, [A, unwrap(i), unwrap(t)])) + +for ff in [acos, log1p, acosh, log2, asin, tan, atanh, cos, log, sin, log10, sqrt] + f = nameof(ff) + @eval NaNMath.$f(x::PyomoDAEVar) = Base.$f(x) +end + +const SymbolicConcreteModel = Symbolics.symstruct(ConcreteModel) + +struct PyomoModel + model::ConcreteModel + U + V + tₛ::Union{Int} + is_free_final::Bool + model_sym::SymbolicConcreteModel + t_sym::Union{Num, BasicSymbolic} + idx_sym::Union{Num, BasicSymbolic} + + function PyomoModel(model, U, V, tₛ, is_free_final) + @variables MODEL_SYM::SymbolicConcreteModel IDX_SYM::Int T_SYM + PyomoModel(model, U, V, tₛ, is_free_final, MODEL_SYM, T_SYM, INDEX_SYM) + end +end + +struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <: + AbstractDynamicOptProblem{uType, tType, isinplace} + f::F + u0::uType + tspan::tType + p::P + wrapped_model::ConcreteModel + kwargs::K + + function PyomoDynamicOptProblem(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 + +""" + PyomoDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt, steps) + +Convert an ODESystem representing an optimal control system into a Pyomo model +for solving using optimization. Must provide either `dt`, the timestep between collocation +points (which, along with the timespan, determines the number of points), or directly +provide the number of points as `steps`. +""" +function MTK.PyomoDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; + dt = nothing, + steps = nothing, + guesses = Dict(), kwargs...) + prob = MTK.process_DynamicOptProblem(PyomoDynamicOptProblem, PyomoModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) + prob +end + +MTK.generate_internal_model(m::Type{PyomoModel}) = pyomo.ConcreteModel() +function MTK.generate_time_variable!(m::ConcreteModel, tspan, tsteps) + m.t = pyomo.ContinuousSet(initialize = collect(tsteps), bounds = tspan) +end + +function MTK.generate_state_variable!(m::ConcreteModel, u0, ns, ts) + m.u_idxs = pyomo.RangeSet(1, ns) + pyomo.Var(m.u_idxs, m.t) +end + +function MTK.generate_input_variable!(m::ConcreteModel, u0, nc, ts) + m.v_idxs = pyomo.RangeSet(1, nc) + pyomo.Var(m.v_idxs, m.t) +end + +function MTK.generate_timescale(m::ConcreteModel, guess, is_free_t) + m.tₛ = is_free_t ? pyomo.Var(initialize = guess, bounds = (0, Inf)) : 1 +end + +function MTK.add_constraint!(pmodel::PyomoModel, cons) + @unpack model, model_sym, idx_sym, t_sym = pmodel + expr = if cons isa Equation + cons.lhs - cons.rhs == 0 + elseif cons.relational_op === Symbolics.geq + cons.lhs - cons.rhs ≥ 0 + else + cons.lhs - cons.rhs ≤ 0 + end + constraint_f = Symbolics.build_function(expr, model_sym, idx_sym, t_sym) + pyomo.Constraint(rule = constraint_f) +end + +function MTK.set_objective!(m::PyomoModel, expr) = pyomo.Objective(expr = expr) + +function add_initial_constraints!(model::PyomoModel, u0, u0_idxs) + for i in u0_idxs + model.U[i, 0].fix(u0[i]) + end +end + +function substitute_fixed_t_vars!(model::PyomoModel, sys, exprs) + stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))]) + ctidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))]) + iv = MTK.get_iv(sys) + + for cons in jconstraints + consvars = MTK.vars(cons) + for st in consvars + MTK.iscall(st) || continue + x = MTK.operation(st) + t = only(MTK.arguments(st)) + MTK.symbolic_type(t) === MTK.NotSymbolic() || continue + if haskey(stidxmap, x(iv)) + idx = stidxmap[x(iv)] + cv = :U + else + idx = ctidxmap[x(iv)] + cv = :V + end + model.t.add(t) + cons = Symbolics.substitute(cons, Dict(x(t) => model.cv[idx, t])) + end + end +end + +function MTK.substitute_model_vars(pmodel::PyomoModel, sys, pmap, exprs, tspan) + @unwrap model, model_sym, idx_sym, t_sym = pmodel + x_ops = [MTK.operation(MTK.unwrap(st)) for st in unknowns(sys)] + c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in MTK.unbound_inputs(sys)] + mU = Symbolics.symbolic_getproperty(model_sym, :U) + mV = Symbolics.symbolic_getproperty(model_sym, :V) + + (ti, tf) = tspan + if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() + _tf = model.tₛ + ti + exprs = map(c -> Symbolics.fast_substitute(c, Dict(tf => _tf)), exprs) + free_t_map = Dict([[x(tₛ) => mU[i, end] for (i, x) in enumerate(x_ops)]; + [c(tₛ) => mV[i, end] for (i, c) in enumerate(c_ops)]]) + exprs = map(c -> Symbolics.fixpoint_sub(c, free_t_map), exprs) + end + + whole_interval_map = Dict([[v => mU[i, t_sym] for (i, v) in enumerate(sts)]; + [v => mV[i, t_sym] for (i, v) in enumerate(cts)]]) + exprs = map(c -> Symbolics.fixpoint_sub(c, whole_interval_map), exprs) + exprs +end + +function MTK.substitute_integral!(model::PyomoModel, exprs, tspan) + intmap = Dict() + for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) + op = MTK.operation(int) + arg = only(arguments(MTK.value(int))) + lo, hi = MTK.value.((op.domain.domain.left, op.domain.domain.right)) + if MTK.is_free_final(model) && isequal((lo, hi), tspan) + (lo, hi) = (0, 1) + elseif MTK.is_free_final(model) + error("Free final time problems cannot handle partial timespans.") + end + intmap[int] = model.tₛ * InfiniteOpt.∫(arg, model.model[:t], lo, hi) + end + exprs = map(c -> Symbolics.substitute(c, intmap), exprs) +end + +function MTK.substitute_differentials(model::PyomoModel, sys, eqs) + pmodel = prob.model + @unpack model, model_sym, t_sym, idx_sym = pmodel + model.dU = pyomo.DerivativeVar(model.U, wrt = model.t) + + mdU = Symbolics.symbolic_getproperty(model_sym, :dU) + mU = Symbolics.symbolic_getproperty(model_sym, :U) + mtₛ = Symbolics.symbolic_getproperty(model_sym, :tₛ) + diffsubmap = Dict([D(mU[i, t_sym]) => mdU[i, t_sym] for i in 1:length(unknowns(sys))]) + + diff_eqs = substitute_model_vars(model, sys, pmap, diff_equations(sys)) + diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs) + [mtₛ * eq.rhs - eq.lhs == 0 for eq in diff_eqs] +end + +struct PyomoCollocation <: AbstractCollocation + solver::Any + derivative_method +end +MTK.PyomoCollocation(solver, derivative_method = 1) = PyomoCollocation(solver, derivative_method) + +function MTK.prepare_and_optimize!() +end +function MTK.get_U_values() +end +function MTK.get_V_values() +end +function MTK.get_t_values() +end +function MTK.successful_solve() +end + +end diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index b91c76182f..04d508bc11 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -243,9 +243,52 @@ function add_cost_function!(model, sys, tspan, pmap) jcosts = substitute_model_vars(model, sys, jcosts, tspan) jcosts = substitute_params(pmap, jcosts) jcosts = substitute_integral(model, jcosts, tspan) + set_objective!(model, consolidate(jcosts)) end +""" +Substitute integrals. For an integral from (ts, te): +- Free final time problems should transcribe this to (0, 1) in the case that (ts, te) is the original timespan. Free final time problems cannot handle partial timespans. +- CasADi cannot handle partial timespans, even for non-free-final time problems. +time problems and unchanged otherwise. +""" +function substitute_integral(model, exprs, tspan) + intmap = Dict() + for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) + op = MTK.operation(int) + arg = only(arguments(MTK.value(int))) + lo, hi = MTK.value.((op.domain.domain.left, op.domain.domain.right)) + lo, hi = process_integral_bounds(model, (lo, hi), tspan) + intmap[int] = lowered_integral(model, expr, lo, hi) + end + expr = map(c -> Symbolics.substitute(c, intmap), expr) + expr = value.(expr) +end + +"""Substitute variables like x(1.5) with the corresponding model variables.""" +function substitute_model_vars(model, sys, exprs) + x_ops = [MTK.operation(MTK.unwrap(st)) for st in unknowns(sys)] + c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in MTK.unbound_inputs(sys)] + + exprs = map(c -> Symbolics.fast_substitute(c, fixed_t_map(model, x_ops, c_ops, exprs)), exprs) + exprs = map(c -> Symbolics.fast_substitute(c, whole_t_map(model, sys)), exprs) + + (ti, tf) = tspan + if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() + _tf = model.tₛ + ti + exprs = map(c -> Symbolics.fast_substitute(c, Dict(tf => _tf)), exprs) + exprs = map(c -> Symbolics.fast_substitute(c, free_t_map(model, _tf, x_ops, c_ops)), exprs) + end +end + +function process_integral_bounds end +function lowered_integral end +function lowered_derivative end +function free_t_map end +function fixed_t_map end +function whole_t_map end + function add_user_constraints!(model, sys, tspan, pmap) conssys = get_constraintsystem(sys) jconstraints = isnothing(conssys) ? nothing : get_constraints(conssys) @@ -265,7 +308,7 @@ end function add_equational_constraints!(model, sys, pmap, tspan) diff_eqs = substitute_model_vars(model, sys, diff_equations(sys), tspan) diff_eqs = substitute_params(pmap, diff_eqs) - diff_eqs = substitute_differentials(model, sys, diff_eqs) + diff_eqs = substitute_differentials(model, sys, diff_eqs, tspan) for eq in diff_eqs add_constraint!(model, eq.lhs ~ eq.rhs * model.tₛ) end @@ -278,16 +321,12 @@ function add_equational_constraints!(model, sys, pmap, tspan) end function set_objective! end -"""Substitute variables like x(1.5) with the corresponding model variables.""" -function substitute_model_vars end -""" -Substitute integrals. For an integral from (ts, te): -- Free final time problems should transcribe this to (0, 1) in the case that (ts, te) is the original timespan. Free final time problems cannot handle partial timespans. -- CasADi cannot handle partial timespans, even for non-free-final time problems. -time problems and unchanged otherwise. -""" -function substitute_integral end -function substitute_differentials end + +function substitute_differentials(model, sys, eqs) + D = Differential(MTK.get_iv(sys)) + diffsubmap = Dict([D(model.U[i]) => lowered_derivative(model, i) for i in 1:length(U)]) + diff_eqs = map(c -> Symbolics.substitute(c, diffsubmap), diff_eqs) +end function substitute_toterm(vars, exprs) toterm_map = Dict([u => default_toterm(value(u)) for u in vars]) From 072824cf99562f94e8c01a6a53743e6b04cebacd Mon Sep 17 00:00:00 2001 From: vyudu Date: Mon, 19 May 2025 21:58:00 -0400 Subject: [PATCH 07/22] chore: move the docstrings to the interface file --- ext/MTKCasADiDynamicOptExt.jl | 24 +--- ext/MTKInfiniteOptExt.jl | 43 +------ ext/MTKPyomoDynamicOptExt.jl | 155 ++++++++++------------- src/systems/optimal_control_interface.jl | 68 +++++++++- 4 files changed, 141 insertions(+), 149 deletions(-) diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index f70b69a59d..d780472540 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -61,23 +61,7 @@ function (M::MXLinearInterpolation)(τ) end end -""" - CasADiDynamicOptProblem(sys::System, u0, tspan, p; dt, steps) - -Convert an System representing an optimal control system into a CasADi model -for solving using optimization. Must provide either `dt`, the timestep between collocation -points (which, along with the timespan, determines the number of points), or directly -provide the number of points as `steps`. - -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 System via `constraints` -- The solver constraints that encode the time-stepping used by the solver -""" -function MTK.CasADiDynamicOptProblem(sys::System, u0map, tspan, pmap; +function MTK.CasADiDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) @@ -212,15 +196,11 @@ function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau) solver_opti end -""" -CasADi Collocation solver. -- solver: an optimization solver such as Ipopt. Should be given as a string or symbol in all lowercase, e.g. "ipopt" -- tableau: An ODE RK tableau. Load a tableau by calling a function like `constructRK4` and may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. If this argument is not passed in, the solver will default to Radau second order. -""" struct CasADiCollocation <: AbstractCollocation solver::Union{String, Symbol} tableau::DiffEqBase.ODERKTableau end + MTK.CasADiCollocation(solver, tableau = MTK.constructDefault()) = CasADiCollocation(solver, tableau) function MTK.prepare_and_optimize!(prob::CasADiDynamicOptProblem, solver::CasADiCollocation; verbose = false, solver_options = Dict(), plugin_options = Dict(), kwargs...) diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index 61f4ec70b3..4b59714e35 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -72,40 +72,14 @@ function MTK.add_constraint!(m::InfiniteOptModel, expr::Union{Equation, Inequali end MTK.set_objective!(m::InfiniteOptModel, expr) = @objective(m.model, Min, expr) -""" - JuMPDynamicOptProblem(sys::System, u0, tspan, p; dt) - -Convert a System representing an optimal control system into a JuMP model -for solving using optimization. Must provide either `dt`, the timestep between collocation -points (which, along with the timespan, determines the number of points), or directly -provide the number of points as `steps`. - -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 System via `constraints` -- The solver constraints that encode the time-stepping used by the solver -""" -function MTK.JuMPDynamicOptProblem(sys::System, u0map, tspan, pmap; +function MTK.JuMPDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) MTK.process_DynamicOptProblem(JuMPDynamicOptProblem, InfiniteOptModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) end -""" - InfiniteOptDynamicOptProblem(sys::System, u0map, tspan, pmap; dt) - -Convert System representing an optimal control system into a InfiniteOpt model -for solving using optimization. Must provide `dt` for determining the length -of the interpolation arrays. - -Related to `JuMPDynamicOptProblem`, but directly adds the differential equations -of the system as derivative constraints, rather than using a solver tableau. -""" -function MTK.InfiniteOptDynamicOptProblem(sys::System, u0map, tspan, pmap; +function MTK.InfiniteOptDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) @@ -115,6 +89,7 @@ function MTK.InfiniteOptDynamicOptProblem(sys::System, u0map, tspan, pmap; end MTK.lowered_integral(model, expr, lo, hi) = model.tₛ * InfiniteOpt.∫(arg, model.model[:t], lo, hi) +MTK.lowered_derivative(model, i) = ∂(model.U[i], model.model[:t]) function MTK.process_integral_bounds(model, integral_span, tspan) if MTK.is_free_final(model) && isequal(integral_span, tspan) @@ -132,8 +107,6 @@ function MTK.add_initial_constraints!(m::InfiniteOptModel, u0, u0_idxs, ts) end end -MTK.lowered_derivative(model, i) = ∂(model.U[i], model.model[:t]) - function MTK.fixed_t_map(model::InfiniteOptModel, x_ops, c_ops, exprs) Dict([[x_ops[i] => model.U[i] for i in 1:length(model.U)]; [c_ops[i] => model.V[i] for i in 1:length(model.V)]]) @@ -191,22 +164,12 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau) end end -""" -JuMP Collocation solver. -- solver: a optimization solver such as Ipopt -- tableau: An ODE RK tableau. Load a tableau by calling a function like `constructRK4` and may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. If this argument is not passed in, the solver will default to Radau second order. -""" struct JuMPCollocation <: AbstractCollocation solver::Any tableau::DiffEqBase.ODERKTableau end MTK.JuMPCollocation(solver, tableau = MTK.constructDefault()) = JuMPCollocation(solver, tableau) -""" -InfiniteOpt Collocation solver. -- solver: an optimization solver such as Ipopt -- `derivative_method`: 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()). -""" struct InfiniteOptCollocation <: AbstractCollocation solver::Any derivative_method::InfiniteOpt.AbstractDerivativeMethod diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl index 859c17d72a..b822ca0203 100644 --- a/ext/MTKPyomoDynamicOptExt.jl +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -14,7 +14,8 @@ struct PyomoDAEVar v::Py end (v::PyomoDAEVar)(t) = v.v[:, t] -getindex(v::PyomoDAEVar, i::Union{Num, Symbolic}, t::Union{Num, Symbolic}) = wrap(Term{symeltype(A)}(getindex, [A, unwrap(i), unwrap(t)])) +getindex(v::PyomoDAEVar, i::Union{Num, Symbolic}, t::Union{Num, Symbolic}) = wrap(Term{symeltype(v)}(getindex, [v, unwrap(i), unwrap(t)])) +getindex(v::PyomoDAEVar, i::Int) = wrap(Term{symeltype(v)}(getindex, [v, unwrap(i), Colon()])) for ff in [acos, log1p, acosh, log2, asin, tan, atanh, cos, log, sin, log10, sqrt] f = nameof(ff) @@ -25,9 +26,9 @@ const SymbolicConcreteModel = Symbolics.symstruct(ConcreteModel) struct PyomoModel model::ConcreteModel - U - V - tₛ::Union{Int} + U::PyomoDAEVar + V::PyomoDAEVar + tₛ::Union{Int, Py} is_free_final::Bool model_sym::SymbolicConcreteModel t_sym::Union{Num, BasicSymbolic} @@ -54,36 +55,30 @@ struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <: end end -""" - PyomoDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt, steps) - -Convert an ODESystem representing an optimal control system into a Pyomo model -for solving using optimization. Must provide either `dt`, the timestep between collocation -points (which, along with the timespan, determines the number of points), or directly -provide the number of points as `steps`. -""" function MTK.PyomoDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; - dt = nothing, - steps = nothing, + dt = nothing, steps = nothing, guesses = Dict(), kwargs...) prob = MTK.process_DynamicOptProblem(PyomoDynamicOptProblem, PyomoModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + prob.wrapped_model.model.dU = pyomo.DerivativeVar(prob.wrapped_model.model.U, wrt = model.t) MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) prob end MTK.generate_internal_model(m::Type{PyomoModel}) = pyomo.ConcreteModel() + function MTK.generate_time_variable!(m::ConcreteModel, tspan, tsteps) + m.steps = length(tsteps) m.t = pyomo.ContinuousSet(initialize = collect(tsteps), bounds = tspan) end function MTK.generate_state_variable!(m::ConcreteModel, u0, ns, ts) m.u_idxs = pyomo.RangeSet(1, ns) - pyomo.Var(m.u_idxs, m.t) + PyomoDAEVar(pyomo.Var(m.u_idxs, m.t)) end function MTK.generate_input_variable!(m::ConcreteModel, u0, nc, ts) m.v_idxs = pyomo.RangeSet(1, nc) - pyomo.Var(m.v_idxs, m.t) + PyomoDAEVar(pyomo.Var(m.v_idxs, m.t)) end function MTK.generate_timescale(m::ConcreteModel, guess, is_free_t) @@ -111,99 +106,89 @@ function add_initial_constraints!(model::PyomoModel, u0, u0_idxs) end end -function substitute_fixed_t_vars!(model::PyomoModel, sys, exprs) - stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))]) - ctidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))]) - iv = MTK.get_iv(sys) - - for cons in jconstraints - consvars = MTK.vars(cons) - for st in consvars +function fixed_t_map(m::PyomoModel, sys, exprs) + stidxmap = Dict([v => i for (i, v) in x_ops]) + ctidxmap = Dict([v => i for (i, v) in c_ops]) + mU = Symbolics.symbolic_getproperty(model_sym, :U) + mV = Symbolics.symbolic_getproperty(model_sym, :V) + fixed_t_map = Dict() + for expr in exprs + vars = MTK.vars(exprs) + for st in vars MTK.iscall(st) || continue x = MTK.operation(st) t = only(MTK.arguments(st)) MTK.symbolic_type(t) === MTK.NotSymbolic() || continue - if haskey(stidxmap, x(iv)) - idx = stidxmap[x(iv)] - cv = :U - else - idx = ctidxmap[x(iv)] - cv = :V - end - model.t.add(t) - cons = Symbolics.substitute(cons, Dict(x(t) => model.cv[idx, t])) + m.model.t.add(t) + fixed_t_map[x(t)] = haskey(stidxmap, x) ? mU[stidxmap[x], t] : mV[ctidxmap[x], t] end end + fixed_t_map end -function MTK.substitute_model_vars(pmodel::PyomoModel, sys, pmap, exprs, tspan) - @unwrap model, model_sym, idx_sym, t_sym = pmodel - x_ops = [MTK.operation(MTK.unwrap(st)) for st in unknowns(sys)] - c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in MTK.unbound_inputs(sys)] +function MTK.free_t_map(model, x_ops, c_ops) mU = Symbolics.symbolic_getproperty(model_sym, :U) mV = Symbolics.symbolic_getproperty(model_sym, :V) + Dict([[x(tₛ) => mU[i, end] for (i, x) in enumerate(x_ops)]; + [c(tₛ) => mV[i, end] for (i, c) in enumerate(c_ops)]]) +end - (ti, tf) = tspan - if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() - _tf = model.tₛ + ti - exprs = map(c -> Symbolics.fast_substitute(c, Dict(tf => _tf)), exprs) - free_t_map = Dict([[x(tₛ) => mU[i, end] for (i, x) in enumerate(x_ops)]; - [c(tₛ) => mV[i, end] for (i, c) in enumerate(c_ops)]]) - exprs = map(c -> Symbolics.fixpoint_sub(c, free_t_map), exprs) - end +function MTK.whole_t_map(model) + mU = Symbolics.symbolic_getproperty(model_sym, :U) + mV = Symbolics.symbolic_getproperty(model_sym, :V) + Dict([[v => mU[i, t_sym] for (i, v) in enumerate(sts)]; + [v => mV[i, t_sym] for (i, v) in enumerate(cts)]]) +end - whole_interval_map = Dict([[v => mU[i, t_sym] for (i, v) in enumerate(sts)]; - [v => mV[i, t_sym] for (i, v) in enumerate(cts)]]) - exprs = map(c -> Symbolics.fixpoint_sub(c, whole_interval_map), exprs) - exprs -end - -function MTK.substitute_integral!(model::PyomoModel, exprs, tspan) - intmap = Dict() - for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) - op = MTK.operation(int) - arg = only(arguments(MTK.value(int))) - lo, hi = MTK.value.((op.domain.domain.left, op.domain.domain.right)) - if MTK.is_free_final(model) && isequal((lo, hi), tspan) - (lo, hi) = (0, 1) - elseif MTK.is_free_final(model) - error("Free final time problems cannot handle partial timespans.") - end - intmap[int] = model.tₛ * InfiniteOpt.∫(arg, model.model[:t], lo, hi) - end - exprs = map(c -> Symbolics.substitute(c, intmap), exprs) +function MTK.lowered_integral(m::PyomoModel, arg) + @unpack model, model_sym, t_sym = m + arg_f = Symbolics.build_function(arg, model_sym, t_sym) + Integral(model.t, wrt = model.t, rule=arg_f) end -function MTK.substitute_differentials(model::PyomoModel, sys, eqs) - pmodel = prob.model - @unpack model, model_sym, t_sym, idx_sym = pmodel - model.dU = pyomo.DerivativeVar(model.U, wrt = model.t) +MTK.process_integral_bounds(model, integral_span, tspan) = integral_span +function MTK.lowered_derivative(m::PyomoModel, i) mdU = Symbolics.symbolic_getproperty(model_sym, :dU) - mU = Symbolics.symbolic_getproperty(model_sym, :U) - mtₛ = Symbolics.symbolic_getproperty(model_sym, :tₛ) - diffsubmap = Dict([D(mU[i, t_sym]) => mdU[i, t_sym] for i in 1:length(unknowns(sys))]) - - diff_eqs = substitute_model_vars(model, sys, pmap, diff_equations(sys)) - diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs) - [mtₛ * eq.rhs - eq.lhs == 0 for eq in diff_eqs] + mdU[i, t_sym] end struct PyomoCollocation <: AbstractCollocation - solver::Any - derivative_method + solver::Union{String, Symbol} + derivative_method::Pyomo.DiscretizationMethod end -MTK.PyomoCollocation(solver, derivative_method = 1) = PyomoCollocation(solver, derivative_method) -function MTK.prepare_and_optimize!() -end -function MTK.get_U_values() +MTK.PyomoCollocation(solver, derivative_method = LagrangeRadau(5)) = PyomoCollocation(solver, derivative_method) + +function MTK.prepare_and_optimize!(prob::PyomoDynamicOptProblem, collocation; verbose, kwargs...) + m = prob.wrapped_model.model + dm = collocation.derivative_method + discretizer = TransformationFactory(dm) + ncp = is_finite_difference(dm) ? 1 : dm.np + discretizer.apply_to(model, wrt = m.t, nfe = m.steps, ncp = ncp, scheme = scheme_string(dm)) + solver = SolverFactory(string(collocation.solver)) + solver.solve(m) end -function MTK.get_V_values() + +function MTK.get_U_values(m::PyomoModel) + [pyomo.value(model.U[i]) for i in model.U] end -function MTK.get_t_values() + +function MTK.get_V_values(m::PyomoModel) + [pyomo.value(model.V[i]) for i in model.V] end -function MTK.successful_solve() + +function MTK.get_t_values(m::PyomoModel) + [pyomo.value(model.t[i]) for i in model.t] end +function MTK.successful_solve(m::PyomoModel) + ss = m.solver.status + tc = m.solver.termination_condition + if ss == opt.SolverStatus.ok && (tc == opt.TerminationStatus.optimal || tc == opt.TerminationStatus.locallyOptimal) + return true + else + return false + end +end end diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index 04d508bc11..31dbf5e5c1 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -18,13 +18,78 @@ function Base.show(io::IO, sol::DynamicOptSolution) 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 +""" + JuMPDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt, steps, guesses, kwargs...) + +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 `steps`. + +To construct the problem, please load InfiniteOpt along with ModelingToolkit. +""" function JuMPDynamicOptProblem end +""" + 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 `JuMPDynamicOptProblem`, but directly adds the differential equations +of the system as derivative constraints, rather than using a solver tableau. + +To construct the problem, please load InfiniteOpt along with ModelingToolkit. +""" function InfiniteOptDynamicOptProblem end +""" + CasADiDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt, steps, guesses, kwargs...) + +Convert an ODESystem representing an optimal control system into a CasADi model +for solving using optimization. Must provide either `dt`, the timestep between collocation +points (which, along with the timespan, determines the number of points), or directly +provide the number of points as `steps`. + +To construct the problem, please load CasADi along with ModelingToolkit. +""" function CasADiDynamicOptProblem end +""" + PyomoDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt, steps) +Convert an ODESystem representing an optimal control system into a Pyomo model +for solving using optimization. Must provide either `dt`, the timestep between collocation +points (which, along with the timespan, determines the number of points), or directly +provide the number of points as `steps`. + +To construct the problem, please load Pyomo along with ModelingToolkit. +""" +function PyomoDynamicOptProblem end + +### Collocations +""" +JuMP Collocation solver. +- solver: a optimization solver such as Ipopt +- tableau: An ODE RK tableau. Load a tableau by calling a function like `constructRK4` and may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. If this argument is not passed in, the solver will default to Radau second order. +""" function JuMPCollocation end +""" +InfiniteOpt Collocation solver. +- solver: an optimization solver such as Ipopt +- `derivative_method`: 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 InfiniteOptCollocation end +""" +CasADi Collocation solver. +- solver: an optimization solver such as Ipopt. Should be given as a string or symbol in all lowercase, e.g. "ipopt" +- tableau: An ODE RK tableau. Load a tableau by calling a function like `constructRK4` and may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. If this argument is not passed in, the solver will default to Radau second order. +""" function CasADiCollocation end +""" +Pyomo Collocation solver. +- solver: an optimization solver such as Ipopt. Should be given as a string or symbol in all lowercase, e.g. "ipopt" +- derivative_method: a derivative method from Pyomo. The choices here are ForwardEuler, BackwardEuler, MidpointEuler, LagrangeRadau, or LagrangeLegendre. The last two should additionally have a number indicating the number of collocation points per timestep, e.g. PyomoCollocation("ipopt", LagrangeRadau(3)). Defaults to LagrangeRadau(5). +""" +function PyomoCollocation end function warn_overdetermined(sys, u0map) cstrs = constraints(sys) @@ -243,7 +308,6 @@ function add_cost_function!(model, sys, tspan, pmap) jcosts = substitute_model_vars(model, sys, jcosts, tspan) jcosts = substitute_params(pmap, jcosts) jcosts = substitute_integral(model, jcosts, tspan) - set_objective!(model, consolidate(jcosts)) end @@ -316,7 +380,7 @@ function add_equational_constraints!(model, sys, pmap, tspan) alg_eqs = substitute_model_vars(model, sys, alg_equations(sys), tspan) alg_eqs = substitute_params(pmap, alg_eqs) for eq in alg_eqs - add_constraint!(model, eq.lhs ~ eq.rhs * model.tₛ) + add_constraint!(model, eq.lhs ~ eq.rhs) end end From 7fda5b0f13a6ad689afc24a652377a97243b84ad Mon Sep 17 00:00:00 2001 From: vyudu Date: Mon, 19 May 2025 22:08:15 -0400 Subject: [PATCH 08/22] test: add Pyomo tests --- test/extensions/dynamic_optimization.jl | 48 +++++++++++++++++++++++-- 1 file changed, 46 insertions(+), 2 deletions(-) diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index 2d24f92d1f..d476b7a066 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -44,6 +44,9 @@ const M = ModelingToolkit @test ≈(isol.sol.u, osol2.u, rtol = 0.001) csol2 = solve(cprob, CasADiCollocation("ipopt", constructImplicitEuler())) @test ≈(csol2.sol.u, osol2.u, rtol = 0.001) + pprob = PyomoDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01) + psol = solve(cprob, PyomoCollocation("ipopt", BackwardEuler())) + @test psol.sol.u ≈ osol2.u # With a constraint u0map = Pair[] @@ -62,6 +65,12 @@ const M = ModelingToolkit @test csol.sol(0.6; idxs = x(t)) ≈ 3.5 @test csol.sol(0.3; idxs = x(t)) ≈ 7.0 + pprob = PyomoDynamicOptProblem( + lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(3))) + @test psol.sol(0.6)[1] ≈ 3.5 + @test psol.sol(0.3)[1] ≈ 7.0 + iprob = InfiniteOptDynamicOptProblem( lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3))) # 48.564 ms, 9.58 MiB @@ -77,10 +86,14 @@ const M = ModelingToolkit isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3))) @test all(u -> u > [1, 1], isol.sol.u) - jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + jprob = PyoDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIA3())) @test all(u -> u > [1, 1], jsol.sol.u) + pprob = PyomoDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + psol = solve(pprob, PyomoCollocation("ipopt", MidpointEuler())) + @test all(u -> u > [1, 1], psol.sol.u) + cprob = CasADiDynamicOptProblem( lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) csol = solve(cprob, CasADiCollocation("ipopt", constructRadauIA3())) @@ -125,7 +138,6 @@ end cprob = CasADiDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01) csol = solve(cprob, CasADiCollocation("ipopt", constructVerner8())) - # Linear systems have bang-bang controls @test is_bangbang(csol.input_sol, [-1.0], [1.0]) # Test reached final position. @test ≈(csol.sol[x(t)][end], 0.25, rtol = 1e-5) @@ -143,8 +155,15 @@ end isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test is_bangbang(isol.input_sol, [-1.0], [1.0]) @test ≈(isol.sol[x(t)][end], 0.25, rtol = 1e-5) + + pprob = PyomoDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01) + psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) + @test is_bangbang(psol.input_sol, [-1.0], [1.0]) + @test ≈(psol.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) + @test ≈(psol.sol.u, osol.u, rtol = 0.05) ################### ### Bee example ### @@ -172,6 +191,9 @@ end cprob = CasADiDynamicOptProblem(beesys, u0map, tspan, pmap; dt = 0.01) csol = solve(cprob, CasADiCollocation("ipopt", constructTsitouras5())) @test is_bangbang(csol.input_sol, [0.0], [1.0]) + pprob = PyomoDynamicOptProblem(beesys, u0map, tspan, pmap, dt = 0.01) + psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) + @test is_bangbang(psol.input_sol, [0.0], [1.0]) @parameters (α_interp::LinearInterpolation)(..) eqs = [D(w(t)) ~ -μ * w(t) + b * s * α_interp(t) * w(t), @@ -186,6 +208,7 @@ end @test ≈(osol.u, csol.sol.u, rtol = 0.01) osol2 = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false) @test ≈(osol2.u, isol.sol.u, rtol = 0.01) + @test ≈(osol2.u, psol.sol.u, rtol = 0.01) end @testset "Rocket launch" begin @@ -224,6 +247,10 @@ end isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isol.sol[h(t)][end] > 1.012 + pprob = PyomoCollocationDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) + psol = solve(pprob, PyomoCollocation("ipopt")) + @test psol.sol.u[end][1] > 1.012 + # Test solution @parameters (T_interp::CubicSpline)(..) eqs = [D(h(t)) ~ v(t), @@ -240,6 +267,11 @@ end oprob1 = ODEProblem(rocket_ode, merge(Dict(u0map), Dict(pmap), interpmap1), (ts, te)) osol1 = solve(oprob1, ImplicitEuler(); adaptive = false, dt = 0.001) @test ≈(isol.sol.u, osol1.u, rtol = 0.01) + + interpmap2 = Dict(T_interp => ctrl_to_spline(psol.input_sol, CubicSpline)) + oprob2 = ODEProblem(rocket_ode, u0map, (ts, te), merge(Dict(pmap), interpmap2)) + osol2 = solve(oprob2, RadauIIA5(); adaptive = false, dt = 0.001) + @test ≈(psol.sol.u, osol2.u, rtol = 0.01) end @testset "Free final time problems" begin @@ -269,6 +301,10 @@ end isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3) + pprob = PyomoDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 201) + psol = solve(pprob, PyomoCollocation("ipopt")) + @test isapprox(psol.sol.t[end], 10.0, rtol = 1e-3) + @variables x(..) v(..) @variables u(..) [bounds = (-1.0, 1.0), input = true] @parameters tf @@ -293,6 +329,10 @@ end iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; steps = 51) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isapprox(isol.sol.t[end], 2.0, atol = 1e-5) + + pprob = PyomoDynamicOptProblem(block, u0map, tspan, parammap; steps = 51) + psol = solve(pprob, PyomoCollocation("ipopt")) + @test isapprox(psol.sol.t[end], 2.0, atol = 1e-5) end @testset "Cart-pole problem" begin @@ -335,4 +375,8 @@ end iprob = InfiniteOptDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isol.sol.u[end] ≈ [π, 0, 0, 0] + + pprob = PyomoDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04) + psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(4))) + @test psol.sol.u[end] ≈ [π, 0, 0, 0] end From a275b339569a76f184fbfaa61c3c203588d7a268 Mon Sep 17 00:00:00 2001 From: vyudu Date: Thu, 22 May 2025 09:26:41 -0400 Subject: [PATCH 09/22] refactor: use lowered_var instead of manual substitution --- Project.toml | 2 + ext/MTKCasADiDynamicOptExt.jl | 39 +----- ext/MTKInfiniteOptExt.jl | 24 +--- ext/MTKPyomoDynamicOptExt.jl | 155 ++++++++--------------- src/ModelingToolkit.jl | 4 +- src/systems/optimal_control_interface.jl | 74 +++++++---- test/extensions/dynamic_optimization.jl | 6 +- 7 files changed, 123 insertions(+), 181 deletions(-) diff --git a/Project.toml b/Project.toml index 9645a412b0..215c0bcd1b 100644 --- a/Project.toml +++ b/Project.toml @@ -46,6 +46,7 @@ OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +Pyomo = "0e8e1daf-01b5-4eba-a626-3897743a3816" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47" @@ -81,6 +82,7 @@ MTKDeepDiffsExt = "DeepDiffs" MTKFMIExt = "FMI" MTKInfiniteOptExt = "InfiniteOpt" MTKLabelledArraysExt = "LabelledArrays" +MTKPyomoDynamicOptExt = "Pyomo" [compat] ADTypes = "1.14.0" diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index d780472540..d678e853e6 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -105,6 +105,7 @@ function MTK.add_constraint!(m::CasADiModel, expr) subject_to!(m.model, expr.lhs - expr.rhs ≤ 0) end end + MTK.set_objective!(m::CasADiModel, expr) = minimize!(m.model, MX(expr)) function MTK.add_initial_constraints!(m::CasADiModel, u0, u0_idxs, args...) @@ -114,42 +115,12 @@ function MTK.add_initial_constraints!(m::CasADiModel, u0, u0_idxs, args...) end end -function free_t_map(model, tf, x_ops, c_ops) - Dict([[x(tf) => model.U.u[i, end] for (i, x) in enumerate(x_ops)]; - [c(tf) => model.V.u[i, end] for (i, c) in enumerate(c_ops)]]) -end - -function whole_t_map(model, sys) - Dict([[v => model.U.u[i, :] for (i, v) in enumerate(unknowns(sys))]; - [v => model.V.u[i, :] for (i, v) in enumerate(MTK.unbound_inputs(sys))]]) - -end - -function fixed_t_map(model::CasADiModel, x_ops, c_ops, exprs) - stidxmap = Dict([v => i for (i, v) in x_ops]) - ctidxmap = Dict([v => i for (i, v) in c_ops]) - for i in 1:length(exprs) - subvars = MTK.vars(exprs[i]) - for st in subvars - MTK.iscall(st) || continue - x = operation(st) - t = only(arguments(st)) - MTK.symbolic_type(t) === MTK.NotSymbolic() || continue - if haskey(stidxmap, x) - idx = stidxmap[x] - cv = model.U - else - idx = ctidxmap[x] - cv = model.V - end - exprs[i] = Symbolics.fast_substitute(exprs[i], Dict(x(t) => cv(t)[idx])) - end - jcosts = Symbolics.substitute(jcosts, Dict(x(t) => cv(t)[idx])) - end - exprs +function MTK.lowered_var(m::CasADiModel, uv, i, t) + X = getfield(m, uv) + t isa Union{Num, Symbolics.Symbolic} ? X.u[i, :] : X(t)[i] end -MTK.lowered_integral(model, expr, args...) = model.tₛ * (model.U.t[2] - model.U.t[1]) * expr +MTK.lowered_integral(model::CasADiModel, expr, args...) = model.tₛ * (model.U.t[2] - model.U.t[1]) * sum(expr) function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau) @unpack A, α, c = tableau diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index 4b59714e35..8cd83e4323 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -48,7 +48,7 @@ struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <: end MTK.generate_internal_model(m::Type{InfiniteOptModel}) = InfiniteModel() -MTK.generate_time_variable!(m::InfiniteModel, tspan, steps) = @infinite_parameter(m, t in [tspan[1], tspan[2]], num_supports = length(tsteps)) +MTK.generate_time_variable!(m::InfiniteModel, tspan, tsteps) = @infinite_parameter(m, t in [tspan[1], tspan[2]], num_supports = length(tsteps)) MTK.generate_state_variable!(m::InfiniteModel, u0::Vector, ns, ts) = @variable(m, U[i = 1:ns], Infinite(m[:t]), start=u0[i]) MTK.generate_input_variable!(m::InfiniteModel, c0, nc, ts) = @variable(m, V[i = 1:nc], Infinite(m[:t]), start=c0[i]) @@ -88,8 +88,8 @@ function MTK.InfiniteOptDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; prob end -MTK.lowered_integral(model, expr, lo, hi) = model.tₛ * InfiniteOpt.∫(arg, model.model[:t], lo, hi) -MTK.lowered_derivative(model, i) = ∂(model.U[i], model.model[:t]) +MTK.lowered_integral(model::InfiniteOptModel, expr, lo, hi) = model.tₛ * InfiniteOpt.∫(expr, model.model[:t], lo, hi) +MTK.lowered_derivative(model::InfiniteOptModel, i) = ∂(model.U[i], model.model[:t]) function MTK.process_integral_bounds(model, integral_span, tspan) if MTK.is_free_final(model) && isequal(integral_span, tspan) @@ -103,23 +103,13 @@ end function MTK.add_initial_constraints!(m::InfiniteOptModel, u0, u0_idxs, ts) for i in u0_idxs - fix(m.U[i], u0[i], force = true) + fix(m.U[i](0), u0[i], force = true) end end -function MTK.fixed_t_map(model::InfiniteOptModel, x_ops, c_ops, exprs) - Dict([[x_ops[i] => model.U[i] for i in 1:length(model.U)]; - [c_ops[i] => model.V[i] for i in 1:length(model.V)]]) -end - -function MTK.free_t_map(model::InfiniteOptModel, tf, x_ops, c_ops) - Dict([[x(tf) => model.U[i](1) for (i, x) in enumerate(x_ops)]; - [c(tf) => model.V[i](1) for (i, c) in enumerate(c_ops)]]) -end - -function MTK.whole_t_map(model::InfiniteOptModel, sys) - whole_interval_map = Dict([[v => model.U[i] for (i, v) in enumerate(unknowns(sys))]; - [v => model.V[i] for (i, v) in enumerate(MTK.unbound_inputs(sys))]]) +function MTK.lowered_var(m::InfiniteOptModel, uv, i, t) + X = getfield(m, uv) + t isa Union{Num, Symbolics.Symbolic} ? X[i] : X[i](t) end function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau) diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl index b822ca0203..5e5087fd9f 100644 --- a/ext/MTKPyomoDynamicOptExt.jl +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -1,42 +1,26 @@ module MTKPyomoDynamicOptExt using ModelingToolkit -using PythonCall +using Pyomo using DiffEqBase using UnPack using NaNMath const MTK = ModelingToolkit -# import pyomo -const pyomo = PythonCall.pynew() -PythonCall.pycopy!(pyomo, pyimport("pyomo.environ")) - -struct PyomoDAEVar - v::Py -end -(v::PyomoDAEVar)(t) = v.v[:, t] -getindex(v::PyomoDAEVar, i::Union{Num, Symbolic}, t::Union{Num, Symbolic}) = wrap(Term{symeltype(v)}(getindex, [v, unwrap(i), unwrap(t)])) -getindex(v::PyomoDAEVar, i::Int) = wrap(Term{symeltype(v)}(getindex, [v, unwrap(i), Colon()])) - -for ff in [acos, log1p, acosh, log2, asin, tan, atanh, cos, log, sin, log10, sqrt] - f = nameof(ff) - @eval NaNMath.$f(x::PyomoDAEVar) = Base.$f(x) -end - -const SymbolicConcreteModel = Symbolics.symstruct(ConcreteModel) - -struct PyomoModel +struct PyomoDynamicOptModel model::ConcreteModel - U::PyomoDAEVar - V::PyomoDAEVar - tₛ::Union{Int, Py} + U::PyomoVar + V::PyomoVar + tₛ::Union{Int, PyomoVar} is_free_final::Bool - model_sym::SymbolicConcreteModel - t_sym::Union{Num, BasicSymbolic} - idx_sym::Union{Num, BasicSymbolic} - - function PyomoModel(model, U, V, tₛ, is_free_final) - @variables MODEL_SYM::SymbolicConcreteModel IDX_SYM::Int T_SYM - PyomoModel(model, U, V, tₛ, is_free_final, MODEL_SYM, T_SYM, INDEX_SYM) + dU::PyomoVar + model_sym::Union{Num, Symbolics.BasicSymbolic} + t_sym::Union{Num, Symbolics.BasicSymbolic} + idx_sym::Union{Num, Symbolics.BasicSymbolic} + + function PyomoDynamicOptModel(model, U, V, tₛ, is_free_final) + @variables MODEL_SYM::Symbolics.symstruct(PyomoDynamicOptModel) IDX_SYM::Int T_SYM + model.dU = dae.DerivativeVar(U, wrt = model.t, initialize = 0) + new(model, U, V, tₛ, is_free_final, PyomoVar(model.dU), MODEL_SYM, T_SYM, IDX_SYM) end end @@ -46,7 +30,7 @@ struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <: u0::uType tspan::tType p::P - wrapped_model::ConcreteModel + wrapped_model::PyomoDynamicOptModel kwargs::K function PyomoDynamicOptProblem(f, u0, tspan, p, model, kwargs...) @@ -58,35 +42,38 @@ end function MTK.PyomoDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - prob = MTK.process_DynamicOptProblem(PyomoDynamicOptProblem, PyomoModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) - prob.wrapped_model.model.dU = pyomo.DerivativeVar(prob.wrapped_model.model.U, wrt = model.t) + prob = MTK.process_DynamicOptProblem(PyomoDynamicOptProblem, PyomoDynamicOptModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + conc_model = prob.wrapped_model.model MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) prob end -MTK.generate_internal_model(m::Type{PyomoModel}) = pyomo.ConcreteModel() +MTK.generate_internal_model(m::Type{PyomoDynamicOptModel}) = ConcreteModel(pyomo.ConcreteModel()) function MTK.generate_time_variable!(m::ConcreteModel, tspan, tsteps) m.steps = length(tsteps) - m.t = pyomo.ContinuousSet(initialize = collect(tsteps), bounds = tspan) + m.t = dae.ContinuousSet(initialize = tsteps, bounds = tspan) end function MTK.generate_state_variable!(m::ConcreteModel, u0, ns, ts) m.u_idxs = pyomo.RangeSet(1, ns) - PyomoDAEVar(pyomo.Var(m.u_idxs, m.t)) + m.U = pyomo.Var(m.u_idxs, m.t, initialize = 0) + PyomoVar(m.U) end -function MTK.generate_input_variable!(m::ConcreteModel, u0, nc, ts) +function MTK.generate_input_variable!(m::ConcreteModel, c0, nc, ts) m.v_idxs = pyomo.RangeSet(1, nc) - PyomoDAEVar(pyomo.Var(m.v_idxs, m.t)) + m.V = pyomo.Var(m.v_idxs, m.t, initialize = 0) + PyomoVar(m.V) end -function MTK.generate_timescale(m::ConcreteModel, guess, is_free_t) - m.tₛ = is_free_t ? pyomo.Var(initialize = guess, bounds = (0, Inf)) : 1 +function MTK.generate_timescale!(m::ConcreteModel, guess, is_free_t) + m.tₛ = is_free_t ? PyomoVar(pyomo.Var(initialize = guess, bounds = (0, Inf))) : 1 end -function MTK.add_constraint!(pmodel::PyomoModel, cons) +function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons) @unpack model, model_sym, idx_sym, t_sym = pmodel + @show model.dU expr = if cons isa Equation cons.lhs - cons.rhs == 0 elseif cons.relational_op === Symbolics.geq @@ -94,63 +81,40 @@ function MTK.add_constraint!(pmodel::PyomoModel, cons) else cons.lhs - cons.rhs ≤ 0 end - constraint_f = Symbolics.build_function(expr, model_sym, idx_sym, t_sym) - pyomo.Constraint(rule = constraint_f) + constraint_f = Symbolics.build_function(expr, model_sym, idx_sym, t_sym, expression = Val{false}) + @show typeof(constraint_f) + @show typeof(Pyomo.pyfunc(constraint_f)) + cons_sym = gensym() + setproperty!(model, cons_sym, pyomo.Constraint(model.u_idxs, model.t, rule = Pyomo.pyfunc(constraint_f))) end -function MTK.set_objective!(m::PyomoModel, expr) = pyomo.Objective(expr = expr) +function MTK.set_objective!(m::PyomoDynamicOptModel, expr) + m.model.obj = pyomo.Objective(expr = expr) +end -function add_initial_constraints!(model::PyomoModel, u0, u0_idxs) +function MTK.add_initial_constraints!(model::PyomoDynamicOptModel, u0, u0_idxs, ts) for i in u0_idxs model.U[i, 0].fix(u0[i]) end end -function fixed_t_map(m::PyomoModel, sys, exprs) - stidxmap = Dict([v => i for (i, v) in x_ops]) - ctidxmap = Dict([v => i for (i, v) in c_ops]) - mU = Symbolics.symbolic_getproperty(model_sym, :U) - mV = Symbolics.symbolic_getproperty(model_sym, :V) - fixed_t_map = Dict() - for expr in exprs - vars = MTK.vars(exprs) - for st in vars - MTK.iscall(st) || continue - x = MTK.operation(st) - t = only(MTK.arguments(st)) - MTK.symbolic_type(t) === MTK.NotSymbolic() || continue - m.model.t.add(t) - fixed_t_map[x(t)] = haskey(stidxmap, x) ? mU[stidxmap[x], t] : mV[ctidxmap[x], t] - end - end - fixed_t_map -end - -function MTK.free_t_map(model, x_ops, c_ops) - mU = Symbolics.symbolic_getproperty(model_sym, :U) - mV = Symbolics.symbolic_getproperty(model_sym, :V) - Dict([[x(tₛ) => mU[i, end] for (i, x) in enumerate(x_ops)]; - [c(tₛ) => mV[i, end] for (i, c) in enumerate(c_ops)]]) -end - -function MTK.whole_t_map(model) - mU = Symbolics.symbolic_getproperty(model_sym, :U) - mV = Symbolics.symbolic_getproperty(model_sym, :V) - Dict([[v => mU[i, t_sym] for (i, v) in enumerate(sts)]; - [v => mV[i, t_sym] for (i, v) in enumerate(cts)]]) -end - -function MTK.lowered_integral(m::PyomoModel, arg) +function MTK.lowered_integral(m::PyomoDynamicOptModel, arg, lo, hi) @unpack model, model_sym, t_sym = m arg_f = Symbolics.build_function(arg, model_sym, t_sym) - Integral(model.t, wrt = model.t, rule=arg_f) + dae.Integral(model.t, wrt = model.t, rule=arg_f) end MTK.process_integral_bounds(model, integral_span, tspan) = integral_span -function MTK.lowered_derivative(m::PyomoModel, i) - mdU = Symbolics.symbolic_getproperty(model_sym, :dU) - mdU[i, t_sym] +function MTK.lowered_derivative(m::PyomoDynamicOptModel, i) + mdU = Symbolics.symbolic_getproperty(m.model_sym, :dU).val + Symbolics.unwrap(mdU[i, m.t_sym]) +end + +function MTK.lowered_var(m::PyomoDynamicOptModel, uv, i, t) + X = Symbolics.symbolic_getproperty(m.model_sym, uv).val + var = t isa Union{Num, Symbolics.Symbolic} ? X[i, m.t_sym] : X[i, t] + Symbolics.unwrap(var) end struct PyomoCollocation <: AbstractCollocation @@ -164,25 +128,18 @@ function MTK.prepare_and_optimize!(prob::PyomoDynamicOptProblem, collocation; ve m = prob.wrapped_model.model dm = collocation.derivative_method discretizer = TransformationFactory(dm) - ncp = is_finite_difference(dm) ? 1 : dm.np - discretizer.apply_to(model, wrt = m.t, nfe = m.steps, ncp = ncp, scheme = scheme_string(dm)) + ncp = Pyomo.is_finite_difference(dm) ? 1 : dm.np + discretizer.apply_to(m, wrt = m.t, nfe = m.steps, scheme = Pyomo.scheme_string(dm)) solver = SolverFactory(string(collocation.solver)) - solver.solve(m) + solver.solve(m, tee = true) + Main.xx[] = solver end -function MTK.get_U_values(m::PyomoModel) - [pyomo.value(model.U[i]) for i in model.U] -end - -function MTK.get_V_values(m::PyomoModel) - [pyomo.value(model.V[i]) for i in model.V] -end - -function MTK.get_t_values(m::PyomoModel) - [pyomo.value(model.t[i]) for i in model.t] -end +MTK.get_U_values(m::PyomoDynamicOptModel) = [pyomo.value(m.model.U[i]) for i in m.model.U.index_set()] +MTK.get_V_values(m::PyomoDynamicOptModel) = [pyomo.value(m.model.V[i]) for i in m.model.V.index_set()] +MTK.get_t_values(m::PyomoDynamicOptModel) = Pyomo.get_results(m.model, :t) -function MTK.successful_solve(m::PyomoModel) +function MTK.successful_solve(m::PyomoDynamicOptModel) ss = m.solver.status tc = m.solver.termination_condition if ss == opt.SolverStatus.ok && (tc == opt.TerminationStatus.optimal || tc == opt.TerminationStatus.locallyOptimal) diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl index 9608c6083d..255c5ade48 100644 --- a/src/ModelingToolkit.jl +++ b/src/ModelingToolkit.jl @@ -356,9 +356,9 @@ function FMIComponent end include("systems/optimal_control_interface.jl") export AbstractDynamicOptProblem, JuMPDynamicOptProblem, InfiniteOptDynamicOptProblem, - CasADiDynamicOptProblem + CasADiDynamicOptProblem, PyomoDynamicOptProblem export AbstractCollocation, JuMPCollocation, InfiniteOptCollocation, - CasADiCollocation + CasADiCollocation, PyomoCollocation export DynamicOptSolution @public apply_to_variables diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index 31dbf5e5c1..b094c4a24f 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -250,7 +250,7 @@ function process_DynamicOptProblem(prob_type::Type{<:AbstractDynamicOptProblem}, tsteps = LinRange(model_tspan[1], model_tspan[2], steps) model = generate_internal_model(model_type) - generate_time_variable!(model, model_tspan, steps) + generate_time_variable!(model, model_tspan, tsteps) U = generate_state_variable!(model, u0, length(states), tsteps) V = generate_input_variable!(model, c0, length(ctrls), tsteps) tₛ = generate_timescale!(model, get(pmap, tspan[2], tspan[2]), is_free_t) @@ -269,24 +269,26 @@ function generate_internal_model end function generate_state_variable! end function generate_input_variable! end function generate_timescale! end -function set_variable_bounds! end function add_initial_constraints! end function add_constraint! end function set_variable_bounds!(m, sys, pmap, tf) @unpack model, U, V, tₛ = m + t = get_iv(sys) for (i, u) in enumerate(unknowns(sys)) + var = lowered_var(m, :U, i, t) if hasbounds(u) lo, hi = getbounds(u) - add_constraint!(m, U[i] ≳ Symbolics.fixpoint_sub(lo, pmap)) - add_constraint!(m, U[i] ≲ Symbolics.fixpoint_sub(hi, pmap)) + add_constraint!(m, var ≳ Symbolics.fixpoint_sub(lo, pmap)) + add_constraint!(m, var ≲ Symbolics.fixpoint_sub(hi, pmap)) end end for (i, v) in enumerate(unbound_inputs(sys)) + var = lowered_var(m, :V, i, t) if hasbounds(v) lo, hi = getbounds(v) - add_constraint!(m, V[i] ≳ Symbolics.fixpoint_sub(lo, pmap)) - add_constraint!(m, V[i] ≲ Symbolics.fixpoint_sub(hi, pmap)) + add_constraint!(m, var ≳ Symbolics.fixpoint_sub(lo, pmap)) + add_constraint!(m, var ≲ Symbolics.fixpoint_sub(hi, pmap)) end end if symbolic_type(tf) === ScalarSymbolic() && hasbounds(tf) @@ -319,39 +321,58 @@ time problems and unchanged otherwise. """ function substitute_integral(model, exprs, tspan) intmap = Dict() - for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) - op = MTK.operation(int) - arg = only(arguments(MTK.value(int))) - lo, hi = MTK.value.((op.domain.domain.left, op.domain.domain.right)) + for int in collect_applied_operators(exprs, Symbolics.Integral) + op = operation(int) + arg = only(arguments(value(int))) + lo, hi = value.((op.domain.domain.left, op.domain.domain.right)) lo, hi = process_integral_bounds(model, (lo, hi), tspan) - intmap[int] = lowered_integral(model, expr, lo, hi) + intmap[int] = lowered_integral(model, arg, lo, hi) end - expr = map(c -> Symbolics.substitute(c, intmap), expr) - expr = value.(expr) + exprs = map(c -> Symbolics.substitute(c, intmap), exprs) + value.(exprs) end -"""Substitute variables like x(1.5) with the corresponding model variables.""" -function substitute_model_vars(model, sys, exprs) - x_ops = [MTK.operation(MTK.unwrap(st)) for st in unknowns(sys)] - c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in MTK.unbound_inputs(sys)] +"""Substitute variables like x(1.5), x(t), etc. with the corresponding model variables.""" +function substitute_model_vars(model, sys, exprs, tspan) + x_ops = [operation(unwrap(st)) for st in unknowns(sys)] + c_ops = [operation(unwrap(ct)) for ct in unbound_inputs(sys)] + t = get_iv(sys) - exprs = map(c -> Symbolics.fast_substitute(c, fixed_t_map(model, x_ops, c_ops, exprs)), exprs) - exprs = map(c -> Symbolics.fast_substitute(c, whole_t_map(model, sys)), exprs) + exprs = map(c -> Symbolics.fast_substitute(c, whole_t_map(model, t, x_ops, c_ops)), exprs) + exprs = map(c -> Symbolics.fast_substitute(c, fixed_t_map(model, x_ops, c_ops)), exprs) (ti, tf) = tspan - if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() + if symbolic_type(tf) === ScalarSymbolic() _tf = model.tₛ + ti exprs = map(c -> Symbolics.fast_substitute(c, Dict(tf => _tf)), exprs) exprs = map(c -> Symbolics.fast_substitute(c, free_t_map(model, _tf, x_ops, c_ops)), exprs) end + exprs +end + +"""Mappings for variables that depend on the final time parameter, x(tf).""" +function free_t_map(m, tf, x_ops, c_ops) + Dict([[x(tf) => lowered_var(m, :U, i, 1) for (i, x) in enumerate(x_ops)]; + [c(tf) => lowered_var(m, :V, i, 1) for (i, c) in enumerate(c_ops)]]) +end + +"""Mappings for variables that cover the whole timespan, x(t).""" +function whole_t_map(m, t, x_ops, c_ops) + Dict([[v(t) => lowered_var(m, :U, i, t) for (i, v) in enumerate(x_ops)]; + [v(t) => lowered_var(m, :V, i, t) for (i, v) in enumerate(c_ops)]]) +end + +"""Mappings for variables that cover the whole timespan, x(t).""" +function fixed_t_map(m, x_ops, c_ops) + Dict([[v => (t -> lowered_var(m, :U, i, t)) for (i, v) in enumerate(x_ops)]; + [v => (t -> lowered_var(m, :V, i, t)) for (i, v) in enumerate(c_ops)]]) end function process_integral_bounds end function lowered_integral end function lowered_derivative end -function free_t_map end +function lowered_var end function fixed_t_map end -function whole_t_map end function add_user_constraints!(model, sys, tspan, pmap) conssys = get_constraintsystem(sys) @@ -372,7 +393,7 @@ end function add_equational_constraints!(model, sys, pmap, tspan) diff_eqs = substitute_model_vars(model, sys, diff_equations(sys), tspan) diff_eqs = substitute_params(pmap, diff_eqs) - diff_eqs = substitute_differentials(model, sys, diff_eqs, tspan) + diff_eqs = substitute_differentials(model, sys, diff_eqs) for eq in diff_eqs add_constraint!(model, eq.lhs ~ eq.rhs * model.tₛ) end @@ -387,9 +408,10 @@ end function set_objective! end function substitute_differentials(model, sys, eqs) - D = Differential(MTK.get_iv(sys)) - diffsubmap = Dict([D(model.U[i]) => lowered_derivative(model, i) for i in 1:length(U)]) - diff_eqs = map(c -> Symbolics.substitute(c, diffsubmap), diff_eqs) + t = get_iv(sys) + D = Differential(t) + diffsubmap = Dict([D(lowered_var(model, :U, i, t)) => lowered_derivative(model, i) for i in 1:length(unknowns(sys))]) + map(c -> Symbolics.substitute(c, diffsubmap), eqs) end function substitute_toterm(vars, exprs) diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index d476b7a066..706eb11789 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -45,7 +45,7 @@ const M = ModelingToolkit csol2 = solve(cprob, CasADiCollocation("ipopt", constructImplicitEuler())) @test ≈(csol2.sol.u, osol2.u, rtol = 0.001) pprob = PyomoDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01) - psol = solve(cprob, PyomoCollocation("ipopt", BackwardEuler())) + psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) @test psol.sol.u ≈ osol2.u # With a constraint @@ -118,7 +118,7 @@ end # Double integrator t = M.t_nounits D = M.D_nounits - @variables x(..) [bounds = (0.0, 0.25)] v(..) + @variables x(..) v(..) @variables u(..) [bounds = (-1.0, 1.0), input = true] constr = [v(1.0) ~ 0.0] cost = [-x(1.0)] # Maximize the final distance. @@ -130,7 +130,7 @@ end tspan = (0.0, 1.0) parammap = [u(t) => 0.0] jprob = JuMPDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01) - jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructVerner8())) + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructVerner8()), verbose = true) # Linear systems have bang-bang controls @test is_bangbang(jsol.input_sol, [-1.0], [1.0]) # Test reached final position. From 3c591c057f5b75c7d4a0585650df427bef4e2cf9 Mon Sep 17 00:00:00 2001 From: vyudu Date: Mon, 26 May 2025 14:22:22 -0400 Subject: [PATCH 10/22] feat: PyomoDynamicOtpPRoblem --- ext/MTKInfiniteOptExt.jl | 21 +++++----- ext/MTKPyomoDynamicOptExt.jl | 50 +++++++++++++----------- src/systems/optimal_control_interface.jl | 15 +++---- test/extensions/dynamic_optimization.jl | 1 + 4 files changed, 47 insertions(+), 40 deletions(-) diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index 8cd83e4323..21cf117ab3 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -186,6 +186,7 @@ function MTK.prepare_and_optimize!(prob::JuMPDynamicOptProblem, solver::JuMPColl add_solve_constraints!(prob, solver.tableau) set_optimizer(model, solver.solver) optimize!(model) + model end function MTK.prepare_and_optimize!(prob::InfiniteOptDynamicOptProblem, solver::InfiniteOptCollocation; verbose = false, kwargs...) @@ -194,26 +195,26 @@ function MTK.prepare_and_optimize!(prob::InfiniteOptDynamicOptProblem, solver::I set_derivative_method(model[:t], solver.derivative_method) set_optimizer(model, solver.solver) optimize!(model) + model end -function MTK.get_V_values(m::InfiniteOptModel) - nt = length(supports(m.model[:t])) - if !isempty(m.V) - V_vals = value.(m.V) +function MTK.get_V_values(m::InfiniteModel) + nt = length(supports(m[:t])) + if !isempty(m[:V]) + V_vals = value.(m[:V]) V_vals = [[V_vals[i][j] for i in 1:length(V_vals)] for j in 1:nt] else nothing end end -function MTK.get_U_values(m::InfiniteOptModel) - nt = length(supports(m.model[:t])) - U_vals = value.(m.U) +function MTK.get_U_values(m::InfiniteModel) + nt = length(supports(m[:t])) + U_vals = value.(m[:U]) U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:nt] end -MTK.get_t_values(model) = value(model.tₛ) * supports(model.model[:t]) +MTK.get_t_values(m::InfiniteModel) = value(m[:tₛ]) * supports(m[:t]) -function MTK.successful_solve(m::InfiniteOptModel) - model = m.model +function MTK.successful_solve(model::InfiniteModel) tstatus = termination_status(model) pstatus = primal_status(model) !has_values(model) && diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl index 5e5087fd9f..69c2bdaffd 100644 --- a/ext/MTKPyomoDynamicOptExt.jl +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -4,23 +4,26 @@ using Pyomo using DiffEqBase using UnPack using NaNMath +using Setfield const MTK = ModelingToolkit struct PyomoDynamicOptModel model::ConcreteModel U::PyomoVar V::PyomoVar - tₛ::Union{Int, PyomoVar} + tₛ::PyomoVar is_free_final::Bool + solver_model::Union{Nothing, ConcreteModel} dU::PyomoVar model_sym::Union{Num, Symbolics.BasicSymbolic} t_sym::Union{Num, Symbolics.BasicSymbolic} - idx_sym::Union{Num, Symbolics.BasicSymbolic} + uidx_sym::Union{Num, Symbolics.BasicSymbolic} + vidx_sym::Union{Num, Symbolics.BasicSymbolic} function PyomoDynamicOptModel(model, U, V, tₛ, is_free_final) - @variables MODEL_SYM::Symbolics.symstruct(PyomoDynamicOptModel) IDX_SYM::Int T_SYM + @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel) U_IDX_SYM::Int V_IDX_SYM::Int T_SYM model.dU = dae.DerivativeVar(U, wrt = model.t, initialize = 0) - new(model, U, V, tₛ, is_free_final, PyomoVar(model.dU), MODEL_SYM, T_SYM, IDX_SYM) + new(model, U, V, tₛ, is_free_final, nothing, PyomoVar(model.dU), MODEL_SYM, T_SYM, U_IDX_SYM, V_IDX_SYM) end end @@ -39,6 +42,9 @@ struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <: end end +pysym_getproperty(s, name::Symbol) = Symbolics.wrap(SymbolicUtils.term(_getproperty, s, Val{name}(), type = Symbolics.Struct{PyomoVar})) +_getproperty(s, name::Val{fieldname}) where fieldname = getproperty(s, fieldname) + function MTK.PyomoDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) @@ -68,12 +74,11 @@ function MTK.generate_input_variable!(m::ConcreteModel, c0, nc, ts) end function MTK.generate_timescale!(m::ConcreteModel, guess, is_free_t) - m.tₛ = is_free_t ? PyomoVar(pyomo.Var(initialize = guess, bounds = (0, Inf))) : 1 + m.tₛ = is_free_t ? PyomoVar(pyomo.Var(initialize = guess, bounds = (0, Inf))) : PyomoVar(Pyomo.Py(1)) end -function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons) - @unpack model, model_sym, idx_sym, t_sym = pmodel - @show model.dU +function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons; n_idxs = 1) + @unpack model, model_sym, t_sym = pmodel expr = if cons isa Equation cons.lhs - cons.rhs == 0 elseif cons.relational_op === Symbolics.geq @@ -81,11 +86,10 @@ function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons) else cons.lhs - cons.rhs ≤ 0 end - constraint_f = Symbolics.build_function(expr, model_sym, idx_sym, t_sym, expression = Val{false}) - @show typeof(constraint_f) - @show typeof(Pyomo.pyfunc(constraint_f)) - cons_sym = gensym() - setproperty!(model, cons_sym, pyomo.Constraint(model.u_idxs, model.t, rule = Pyomo.pyfunc(constraint_f))) + f_expr = Symbolics.build_function(expr, model_sym, t_sym) + cons_sym = Symbol("cons", hash(cons)) + constraint_f = eval(:(cons_sym = $f_expr)) + setproperty!(model, cons_sym, pyomo.Constraint(model.t, rule = Pyomo.pyfunc(constraint_f))) end function MTK.set_objective!(m::PyomoDynamicOptModel, expr) @@ -107,12 +111,12 @@ end MTK.process_integral_bounds(model, integral_span, tspan) = integral_span function MTK.lowered_derivative(m::PyomoDynamicOptModel, i) - mdU = Symbolics.symbolic_getproperty(m.model_sym, :dU).val + mdU = Symbolics.value(pysym_getproperty(m.model_sym, :dU)) Symbolics.unwrap(mdU[i, m.t_sym]) end function MTK.lowered_var(m::PyomoDynamicOptModel, uv, i, t) - X = Symbolics.symbolic_getproperty(m.model_sym, uv).val + X = Symbolics.value(pysym_getproperty(m.model_sym, uv)) var = t isa Union{Num, Symbolics.Symbolic} ? X[i, m.t_sym] : X[i, t] Symbolics.unwrap(var) end @@ -125,21 +129,21 @@ end MTK.PyomoCollocation(solver, derivative_method = LagrangeRadau(5)) = PyomoCollocation(solver, derivative_method) function MTK.prepare_and_optimize!(prob::PyomoDynamicOptProblem, collocation; verbose, kwargs...) - m = prob.wrapped_model.model + solver_m = prob.wrapped_model.model.clone() dm = collocation.derivative_method discretizer = TransformationFactory(dm) ncp = Pyomo.is_finite_difference(dm) ? 1 : dm.np - discretizer.apply_to(m, wrt = m.t, nfe = m.steps, scheme = Pyomo.scheme_string(dm)) + discretizer.apply_to(solver_m, wrt = solver_m.t, nfe = solver_m.steps, scheme = Pyomo.scheme_string(dm)) solver = SolverFactory(string(collocation.solver)) - solver.solve(m, tee = true) - Main.xx[] = solver + solver.solve(solver_m, tee = true) + solver_m end -MTK.get_U_values(m::PyomoDynamicOptModel) = [pyomo.value(m.model.U[i]) for i in m.model.U.index_set()] -MTK.get_V_values(m::PyomoDynamicOptModel) = [pyomo.value(m.model.V[i]) for i in m.model.V.index_set()] -MTK.get_t_values(m::PyomoDynamicOptModel) = Pyomo.get_results(m.model, :t) +MTK.get_U_values(m::ConcreteModel) = [[pyomo.value(m.U[i, t]) for i in m.u_idxs] for t in m.t] +MTK.get_V_values(m::ConcreteModel) = [[pyomo.value(m.V[i, t]) for i in m.v_idxs] for t in m.t] +MTK.get_t_values(m::ConcreteModel) = [t for t in m.t] -function MTK.successful_solve(m::PyomoDynamicOptModel) +function MTK.successful_solve(m::ConcreteModel) ss = m.solver.status tc = m.solver.termination_condition if ss == opt.SolverStatus.ok && (tc == opt.TerminationStatus.optimal || tc == opt.TerminationStatus.locallyOptimal) diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index b094c4a24f..7e59d893cc 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -395,6 +395,7 @@ function add_equational_constraints!(model, sys, pmap, tspan) diff_eqs = substitute_params(pmap, diff_eqs) diff_eqs = substitute_differentials(model, sys, diff_eqs) for eq in diff_eqs + @show typeof(eq.lhs) add_constraint!(model, eq.lhs ~ eq.rhs * model.tₛ) end @@ -411,7 +412,7 @@ function substitute_differentials(model, sys, eqs) t = get_iv(sys) D = Differential(t) diffsubmap = Dict([D(lowered_var(model, :U, i, t)) => lowered_derivative(model, i) for i in 1:length(unknowns(sys))]) - map(c -> Symbolics.substitute(c, diffsubmap), eqs) + eqs = map(c -> Symbolics.substitute(c, diffsubmap), eqs) end function substitute_toterm(vars, exprs) @@ -451,21 +452,21 @@ function successful_solve end - kwargs are used for other options. For example, the `plugin_options` and `solver_options` will propagated to the Opti object in CasADi. """ function DiffEqBase.solve(prob::AbstractDynamicOptProblem, solver::AbstractCollocation; verbose = false, kwargs...) - prepare_and_optimize!(prob, solver; verbose, kwargs...) + solved_model = prepare_and_optimize!(prob, solver; verbose, kwargs...) - ts = get_t_values(prob.wrapped_model) - Us = get_U_values(prob.wrapped_model) - Vs = get_V_values(prob.wrapped_model) + ts = get_t_values(solved_model) + Us = get_U_values(solved_model) + Vs = get_V_values(solved_model) is_free_final(prob.wrapped_model) && (ts .+ prob.tspan[1]) ode_sol = DiffEqBase.build_solution(prob, solver, ts, Us) input_sol = isnothing(Vs) ? nothing : DiffEqBase.build_solution(prob, solver, ts, Vs) - if !successful_solve(prob.wrapped_model) + if !successful_solve(solved_model) ode_sol = SciMLBase.solution_new_retcode( ode_sol, SciMLBase.ReturnCode.ConvergenceFailure) !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode( input_sol, SciMLBase.ReturnCode.ConvergenceFailure)) end - DynamicOptSolution(prob.wrapped_model.model, ode_sol, input_sol) + DynamicOptSolution(solved_model, ode_sol, input_sol) end diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index 706eb11789..e7056b2399 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -6,6 +6,7 @@ using OrdinaryDiffEqSDIRK, OrdinaryDiffEqVerner, OrdinaryDiffEqTsit5, OrdinaryDi using Ipopt using DataInterpolations using CasADi +using Pyomo import DiffEqBase: solve const M = ModelingToolkit From 33ad60a68cb92f7285055d342e70e9fefed19ee2 Mon Sep 17 00:00:00 2001 From: vyudu Date: Tue, 27 May 2025 01:34:16 -0400 Subject: [PATCH 11/22] fix: fix constraint and cost transcription for pyomo --- ext/MTKCasADiDynamicOptExt.jl | 1 + ext/MTKPyomoDynamicOptExt.jl | 72 ++++++++++++++++++------ src/systems/optimal_control_interface.jl | 8 +-- test/extensions/dynamic_optimization.jl | 13 +++-- 4 files changed, 65 insertions(+), 29 deletions(-) diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index d678e853e6..12ecf31196 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -183,6 +183,7 @@ function MTK.prepare_and_optimize!(prob::CasADiDynamicOptProblem, solver::CasADi catch ErrorException end prob.wrapped_model.solver_opti = solver_opti + prob.wrapped_model end function MTK.get_U_values(model::CasADiModel) diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl index 69c2bdaffd..bff1460ceb 100644 --- a/ext/MTKPyomoDynamicOptExt.jl +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -7,6 +7,20 @@ using NaNMath using Setfield const MTK = ModelingToolkit +SPECIAL_FUNCTIONS_DICT = Dict([acos => Pyomo.py_acos, + log1p => Pyomo.py_log1p, + acosh => Pyomo.py_acosh, + log2 => Pyomo.py_log2, + asin => Pyomo.py_asin, + tan => Pyomo.py_tan, + atanh => Pyomo.py_atanh, + cos => Pyomo.py_cos, + log => Pyomo.py_log, + sin => Pyomo.py_sin, + log10 => Pyomo.py_log10, + sqrt => Pyomo.py_sqrt, + exp => Pyomo.py_exp]) + struct PyomoDynamicOptModel model::ConcreteModel U::PyomoVar @@ -17,13 +31,12 @@ struct PyomoDynamicOptModel dU::PyomoVar model_sym::Union{Num, Symbolics.BasicSymbolic} t_sym::Union{Num, Symbolics.BasicSymbolic} - uidx_sym::Union{Num, Symbolics.BasicSymbolic} - vidx_sym::Union{Num, Symbolics.BasicSymbolic} + dummy_sym::Union{Num, Symbolics.BasicSymbolic} function PyomoDynamicOptModel(model, U, V, tₛ, is_free_final) - @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel) U_IDX_SYM::Int V_IDX_SYM::Int T_SYM + @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel) T_SYM DUMMY_SYM model.dU = dae.DerivativeVar(U, wrt = model.t, initialize = 0) - new(model, U, V, tₛ, is_free_final, nothing, PyomoVar(model.dU), MODEL_SYM, T_SYM, U_IDX_SYM, V_IDX_SYM) + new(model, U, V, tₛ, is_free_final, nothing, PyomoVar(model.dU), MODEL_SYM, T_SYM, DUMMY_SYM) end end @@ -42,7 +55,7 @@ struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <: end end -pysym_getproperty(s, name::Symbol) = Symbolics.wrap(SymbolicUtils.term(_getproperty, s, Val{name}(), type = Symbolics.Struct{PyomoVar})) +pysym_getproperty(s::Union{Num, Symbolics.Symbolic}, name::Symbol) = Symbolics.wrap(SymbolicUtils.term(_getproperty, Symbolics.unwrap(s), Val{name}(), type = Symbolics.Struct{PyomoVar})) _getproperty(s, name::Val{fieldname}) where fieldname = getproperty(s, fieldname) function MTK.PyomoDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; @@ -78,7 +91,7 @@ function MTK.generate_timescale!(m::ConcreteModel, guess, is_free_t) end function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons; n_idxs = 1) - @unpack model, model_sym, t_sym = pmodel + @unpack model, model_sym, t_sym, dummy_sym = pmodel expr = if cons isa Equation cons.lhs - cons.rhs == 0 elseif cons.relational_op === Symbolics.geq @@ -86,14 +99,30 @@ function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons; n_idxs = 1) else cons.lhs - cons.rhs ≤ 0 end - f_expr = Symbolics.build_function(expr, model_sym, t_sym) + expr = Symbolics.substitute(expr, SPECIAL_FUNCTIONS_DICT) + @show expr + cons_sym = Symbol("cons", hash(cons)) - constraint_f = eval(:(cons_sym = $f_expr)) - setproperty!(model, cons_sym, pyomo.Constraint(model.t, rule = Pyomo.pyfunc(constraint_f))) + if occursin(Symbolics.unwrap(t_sym), expr) + f = eval(Symbolics.build_function(expr, model_sym, t_sym)) + setproperty!(model, cons_sym, pyomo.Constraint(model.t, rule = Pyomo.pyfunc(f))) + else + f = eval(Symbolics.build_function(expr, model_sym, dummy_sym)) + setproperty!(model, cons_sym, pyomo.Constraint(rule = Pyomo.pyfunc(f))) + end end -function MTK.set_objective!(m::PyomoDynamicOptModel, expr) - m.model.obj = pyomo.Objective(expr = expr) +function MTK.set_objective!(pmodel::PyomoDynamicOptModel, expr) + @unpack model, model_sym, t_sym, dummy_sym = pmodel + @show expr + expr = Symbolics.substitute(expr, SPECIAL_FUNCTIONS_DICT) + if occursin(Symbolics.unwrap(t_sym), expr) + f = eval(Symbolics.build_function(expr, model_sym, t_sym)) + model.obj = pyomo.Objective(model.t, rule = Pyomo.pyfunc(f)) + else + f = eval(Symbolics.build_function(expr, model_sym, dummy_sym)) + model.obj = pyomo.Objective(rule = Pyomo.pyfunc(f)) + end end function MTK.add_initial_constraints!(model::PyomoDynamicOptModel, u0, u0_idxs, ts) @@ -104,8 +133,14 @@ end function MTK.lowered_integral(m::PyomoDynamicOptModel, arg, lo, hi) @unpack model, model_sym, t_sym = m - arg_f = Symbolics.build_function(arg, model_sym, t_sym) - dae.Integral(model.t, wrt = model.t, rule=arg_f) + @show arg + @show Symbolics.build_function(arg, model_sym, t_sym) + arg_f = eval(Symbolics.build_function(arg, model_sym, t_sym)) + @show arg_f + int_sym = Symbol(:int, hash(arg)) + @show int_sym + setproperty!(model, int_sym, dae.Integral(model.t, wrt = model.t, rule=Pyomo.pyfunc(arg_f))) + PyomoVar(model.tₛ * model.int_sym) end MTK.process_integral_bounds(model, integral_span, tspan) = integral_span @@ -135,15 +170,16 @@ function MTK.prepare_and_optimize!(prob::PyomoDynamicOptProblem, collocation; ve ncp = Pyomo.is_finite_difference(dm) ? 1 : dm.np discretizer.apply_to(solver_m, wrt = solver_m.t, nfe = solver_m.steps, scheme = Pyomo.scheme_string(dm)) solver = SolverFactory(string(collocation.solver)) - solver.solve(solver_m, tee = true) - solver_m + results = solver.solve(solver_m, tee = true) + ConcreteModel(solver_m) end -MTK.get_U_values(m::ConcreteModel) = [[pyomo.value(m.U[i, t]) for i in m.u_idxs] for t in m.t] -MTK.get_V_values(m::ConcreteModel) = [[pyomo.value(m.V[i, t]) for i in m.v_idxs] for t in m.t] -MTK.get_t_values(m::ConcreteModel) = [t for t in m.t] +MTK.get_U_values(m::ConcreteModel) = [[Pyomo.pyconvert(Float64, pyomo.value(m.U[i, t])) for i in m.u_idxs] for t in m.t] +MTK.get_V_values(m::ConcreteModel) = [[Pyomo.pyconvert(Float64, pyomo.value(m.V[i, t])) for i in m.v_idxs] for t in m.t] +MTK.get_t_values(m::ConcreteModel) = [Pyomo.pyconvert(Float64, t) for t in m.t] function MTK.successful_solve(m::ConcreteModel) + return true ss = m.solver.status tc = m.solver.termination_condition if ss == opt.SolverStatus.ok && (tc == opt.TerminationStatus.optimal || tc == opt.TerminationStatus.locallyOptimal) diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index 7e59d893cc..f6cb499e91 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -328,8 +328,7 @@ function substitute_integral(model, exprs, tspan) lo, hi = process_integral_bounds(model, (lo, hi), tspan) intmap[int] = lowered_integral(model, arg, lo, hi) end - exprs = map(c -> Symbolics.substitute(c, intmap), exprs) - value.(exprs) + exprs = map(c -> Symbolics.substitute(c, intmap), value.(exprs)) end """Substitute variables like x(1.5), x(t), etc. with the corresponding model variables.""" @@ -339,14 +338,14 @@ function substitute_model_vars(model, sys, exprs, tspan) t = get_iv(sys) exprs = map(c -> Symbolics.fast_substitute(c, whole_t_map(model, t, x_ops, c_ops)), exprs) - exprs = map(c -> Symbolics.fast_substitute(c, fixed_t_map(model, x_ops, c_ops)), exprs) (ti, tf) = tspan if symbolic_type(tf) === ScalarSymbolic() _tf = model.tₛ + ti + exprs = map(c -> Symbolics.fast_substitute(c, free_t_map(model, tf, x_ops, c_ops)), exprs) exprs = map(c -> Symbolics.fast_substitute(c, Dict(tf => _tf)), exprs) - exprs = map(c -> Symbolics.fast_substitute(c, free_t_map(model, _tf, x_ops, c_ops)), exprs) end + exprs = map(c -> Symbolics.fast_substitute(c, fixed_t_map(model, x_ops, c_ops)), exprs) exprs end @@ -395,7 +394,6 @@ function add_equational_constraints!(model, sys, pmap, tspan) diff_eqs = substitute_params(pmap, diff_eqs) diff_eqs = substitute_differentials(model, sys, diff_eqs) for eq in diff_eqs - @show typeof(eq.lhs) add_constraint!(model, eq.lhs ~ eq.rhs * model.tₛ) end diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index e7056b2399..4f52b08d67 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -47,7 +47,7 @@ const M = ModelingToolkit @test ≈(csol2.sol.u, osol2.u, rtol = 0.001) pprob = PyomoDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01) psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) - @test psol.sol.u ≈ osol2.u + @test all([≈(psol.sol(t), osol2(t), rtol = 1e-3) for t in 0.:0.01:1.]) # With a constraint u0map = Pair[] @@ -69,6 +69,7 @@ const M = ModelingToolkit pprob = PyomoDynamicOptProblem( lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(3))) + @show psol.sol @test psol.sol(0.6)[1] ≈ 3.5 @test psol.sol(0.3)[1] ≈ 7.0 @@ -87,7 +88,7 @@ const M = ModelingToolkit isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3))) @test all(u -> u > [1, 1], isol.sol.u) - jprob = PyoDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIA3())) @test all(u -> u > [1, 1], jsol.sol.u) @@ -160,11 +161,11 @@ end pprob = PyomoDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01) psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) @test is_bangbang(psol.input_sol, [-1.0], [1.0]) - @test ≈(psol.sol.u[end][2], 0.25, rtol = 1e-5) + @test ≈(psol.sol.u[end][2], 0.25, rtol = 1e-3) osol = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false) @test ≈(isol.sol.u, osol.u, rtol = 0.05) - @test ≈(psol.sol.u, osol.u, rtol = 0.05) + @test all([≈(psol.sol(t), osol(t), rtol = 0.05) for t in 0.:0.01:1.]) ################### ### Bee example ### @@ -209,7 +210,7 @@ end @test ≈(osol.u, csol.sol.u, rtol = 0.01) osol2 = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false) @test ≈(osol2.u, isol.sol.u, rtol = 0.01) - @test ≈(osol2.u, psol.sol.u, rtol = 0.01) + @test all([≈(psol.sol(t), osol2(t), rtol = 0.01) for t in 0.:0.01:4.]) end @testset "Rocket launch" begin @@ -248,7 +249,7 @@ end isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isol.sol[h(t)][end] > 1.012 - pprob = PyomoCollocationDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) + pprob = PyomoDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) psol = solve(pprob, PyomoCollocation("ipopt")) @test psol.sol.u[end][1] > 1.012 From 5988c8123be5b259c64824fdb00766561a98cb4d Mon Sep 17 00:00:00 2001 From: vyudu Date: Wed, 28 May 2025 12:07:19 -0400 Subject: [PATCH 12/22] fix: pyomo needs one fewer finite element than # tsteps --- ext/MTKPyomoDynamicOptExt.jl | 49 ++++++++++++++++--------- test/extensions/dynamic_optimization.jl | 8 ++-- 2 files changed, 35 insertions(+), 22 deletions(-) diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl index bff1460ceb..bbd6dba115 100644 --- a/ext/MTKPyomoDynamicOptExt.jl +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -76,18 +76,21 @@ end function MTK.generate_state_variable!(m::ConcreteModel, u0, ns, ts) m.u_idxs = pyomo.RangeSet(1, ns) - m.U = pyomo.Var(m.u_idxs, m.t, initialize = 0) + init_f = Pyomo.pyfunc((m, i, t) -> (u0[Pyomo.pyconvert(Int, i)])) + m.U = pyomo.Var(m.u_idxs, m.t, initialize = init_f) PyomoVar(m.U) end function MTK.generate_input_variable!(m::ConcreteModel, c0, nc, ts) m.v_idxs = pyomo.RangeSet(1, nc) - m.V = pyomo.Var(m.v_idxs, m.t, initialize = 0) + init_f = Pyomo.pyfunc((m, i, t) -> (c0[Pyomo.pyconvert(Int, i)])) + m.V = pyomo.Var(m.v_idxs, m.t, initialize = init_f) PyomoVar(m.V) end function MTK.generate_timescale!(m::ConcreteModel, guess, is_free_t) - m.tₛ = is_free_t ? PyomoVar(pyomo.Var(initialize = guess, bounds = (0, Inf))) : PyomoVar(Pyomo.Py(1)) + m.tₛ = is_free_t ? pyomo.Var(initialize = guess, bounds = (0, Inf)) : Pyomo.Py(1) + PyomoVar(m.tₛ) end function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons; n_idxs = 1) @@ -100,7 +103,6 @@ function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons; n_idxs = 1) cons.lhs - cons.rhs ≤ 0 end expr = Symbolics.substitute(expr, SPECIAL_FUNCTIONS_DICT) - @show expr cons_sym = Symbol("cons", hash(cons)) if occursin(Symbolics.unwrap(t_sym), expr) @@ -133,12 +135,8 @@ end function MTK.lowered_integral(m::PyomoDynamicOptModel, arg, lo, hi) @unpack model, model_sym, t_sym = m - @show arg - @show Symbolics.build_function(arg, model_sym, t_sym) arg_f = eval(Symbolics.build_function(arg, model_sym, t_sym)) - @show arg_f int_sym = Symbol(:int, hash(arg)) - @show int_sym setproperty!(model, int_sym, dae.Integral(model.t, wrt = model.t, rule=Pyomo.pyfunc(arg_f))) PyomoVar(model.tₛ * model.int_sym) end @@ -168,21 +166,36 @@ function MTK.prepare_and_optimize!(prob::PyomoDynamicOptProblem, collocation; ve dm = collocation.derivative_method discretizer = TransformationFactory(dm) ncp = Pyomo.is_finite_difference(dm) ? 1 : dm.np - discretizer.apply_to(solver_m, wrt = solver_m.t, nfe = solver_m.steps, scheme = Pyomo.scheme_string(dm)) + discretizer.apply_to(solver_m, wrt = solver_m.t, nfe = solver_m.steps - 1, scheme = Pyomo.scheme_string(dm)) solver = SolverFactory(string(collocation.solver)) results = solver.solve(solver_m, tee = true) - ConcreteModel(solver_m) + PyomoOutput(results, solver_m) +end + +struct PyomoOutput + result::Pyomo.Py + model::Pyomo.Py end -MTK.get_U_values(m::ConcreteModel) = [[Pyomo.pyconvert(Float64, pyomo.value(m.U[i, t])) for i in m.u_idxs] for t in m.t] -MTK.get_V_values(m::ConcreteModel) = [[Pyomo.pyconvert(Float64, pyomo.value(m.V[i, t])) for i in m.v_idxs] for t in m.t] -MTK.get_t_values(m::ConcreteModel) = [Pyomo.pyconvert(Float64, t) for t in m.t] +function MTK.get_U_values(output::PyomoOutput) + m = output.model + [[Pyomo.pyconvert(Float64, pyomo.value(m.U[i, t])) for i in m.u_idxs] for t in m.t] +end +function MTK.get_V_values(output::PyomoOutput) + m = output.model + [[Pyomo.pyconvert(Float64, pyomo.value(m.V[i, t])) for i in m.v_idxs] for t in m.t] +end +function MTK.get_t_values(output::PyomoOutput) + m = output.model + Pyomo.pyconvert(Float64, pyomo.value(m.tₛ)) * [Pyomo.pyconvert(Float64, t) for t in m.t] +end -function MTK.successful_solve(m::ConcreteModel) - return true - ss = m.solver.status - tc = m.solver.termination_condition - if ss == opt.SolverStatus.ok && (tc == opt.TerminationStatus.optimal || tc == opt.TerminationStatus.locallyOptimal) +function MTK.successful_solve(output::PyomoOutput) + r = output.result + ss = r.solver.status + Main.xx[] = ss + tc = r.solver.termination_condition + if Bool(ss == opt.SolverStatus.ok) && (Bool(tc == opt.TerminationCondition.optimal) || Bool(tc == opt.TerminationCondition.locallyOptimal)) return true else return false diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index 4f52b08d67..85b06eee6f 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -250,7 +250,7 @@ end @test isol.sol[h(t)][end] > 1.012 pprob = PyomoDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) - psol = solve(pprob, PyomoCollocation("ipopt")) + psol = solve(pprob, PyomoCollocation("ipopt", MidpointEuler())) @test psol.sol.u[end][1] > 1.012 # Test solution @@ -319,7 +319,7 @@ end u0map = [x(t) => 1.0, v(t) => 0.0] tspan = (0.0, tf) - parammap = [u(t) => 0.0, tf => 1.0] + parammap = [u(t) => 1.0, tf => 1.0] jprob = JuMPDynamicOptProblem(block, u0map, tspan, parammap; steps = 51) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructVerner8())) @test isapprox(jsol.sol.t[end], 2.0, atol = 1e-5) @@ -329,11 +329,11 @@ end @test isapprox(csol.sol.t[end], 2.0, atol = 1e-5) iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; steps = 51) - isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer), verbose = true) @test isapprox(isol.sol.t[end], 2.0, atol = 1e-5) pprob = PyomoDynamicOptProblem(block, u0map, tspan, parammap; steps = 51) - psol = solve(pprob, PyomoCollocation("ipopt")) + psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) @test isapprox(psol.sol.t[end], 2.0, atol = 1e-5) end From 5f33573f83dc8df1de3fcaf8ab29ce10e9dfacb7 Mon Sep 17 00:00:00 2001 From: vyudu Date: Thu, 29 May 2025 03:29:07 -0400 Subject: [PATCH 13/22] docs: dynamic optimization docs --- docs/src/tutorials/dynamic_optimization.md | 105 +++++++++++++++++++++ ext/MTKCasADiDynamicOptExt.jl | 16 +++- ext/MTKInfiniteOptExt.jl | 3 +- ext/MTKPyomoDynamicOptExt.jl | 43 ++++++--- src/systems/optimal_control_interface.jl | 13 +++ test/extensions/dynamic_optimization.jl | 23 +++-- 6 files changed, 180 insertions(+), 23 deletions(-) create mode 100644 docs/src/tutorials/dynamic_optimization.md diff --git a/docs/src/tutorials/dynamic_optimization.md b/docs/src/tutorials/dynamic_optimization.md new file mode 100644 index 0000000000..a1a1a44384 --- /dev/null +++ b/docs/src/tutorials/dynamic_optimization.md @@ -0,0 +1,105 @@ +# Solving Dynamic Optimization Problems +Systems in ModelingToolkit.jl can be directly converted to dynamic optimization or optimal control problems. In such systems, one has one or more input variables that are externally controlled to control the dynamics of the system. A dynamic optimization solves for the optimal time trajectory of the input variables in order to maximize or minimize a desired objective function. For example, a car driver might like to know how to step on the accelerator if the goal is to finish a race while using the least gas. + +To begin, let us take a rocket launch example. The input variable here is the thrust exerted by the engine. The rocket state is described by its current height and velocity. +```julia +using ModelingToolkit +t = ModelingToolkit.t_nounits +D = ModelingToolkit.D_nounits + +@parameters h_c m₀ h₀ g₀ D_c c Tₘ m_c +@variables begin + h(..) + v(..) + m(..) [bounds = (m_c, 1)] + T(..) [input = true, bounds = (0, Tₘ)] +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] + +(ts, te) = (0.0, 0.2) +costs = [-h(te)] +cons = [T(te) ~ 0, m(te) ~ m_c] + +@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] +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] +``` +What we would like to optimize here is the final height of the rocket. We do this by providing a vector of expressions corresponding to the costs. By default, the sense of the optimization is to minimize the provided cost. So to maximize the rocket height at the final time, we write `-h(te)` as the cost. + +Now we can construct a problem and solve it. Let us use JuMP as our backend here. +```julia +jprob = JuMPDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) +jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIIA5())) +``` + +Let's plot our final solution and the controller here: +```julia +``` + +###### Free final time problems +There are additionally a class of dynamic optimization problems where we would like to know how to control our system to achieve something in the least time. Such problems are called free final time problems, since the final time is unknown. To model these problems in ModelingToolkit, we declare the final time as a parameter. + +```julia +@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] +``` + +Please note that, at the moment, free final time problems cannot support constraints defined at definite time values, like `x(3) ~ 2`. + +Let's plot our final solution and the controller for the block: +```julia +``` + +### Solvers +Currently 4 backends are exposed for solving dynamic optimization problems using collocation: JuMP, InfiniteOpt, CasADi, and Pyomo. + +Please note that there are differences in how to construct the collocation solver for the different cases. For example, the Python based ones, CasADi and Pyomo, expect the solver to be passed in as a string (CasADi and Pyomo come pre-loaded with certain solvers, but other solvers may need to be manually installed using `pip` or `conda`), while JuMP/InfiniteOpt expect the optimizer object to be passed in directly: +``` +JuMPCollocation(Ipopt.Optimizer, constructRK4()) +CasADiCollocation("ipopt", constructRK4()) +``` + +**JuMP** and **CasADi** collocation require an ODE tableau to be passed in. These can be constructed by calling the `constructX()` functions from DiffEqDevTools. If none is passed in, both solvers will default to using Radau second-order with five collocation points. + +**Pyomo** and **InfiniteOpt** each have their own built-in collocation methods. +1. **InfiniteOpt**: The list of InfiniteOpt collocation methods can be found [in the table on this page](https://infiniteopt.github.io/InfiniteOpt.jl/stable/guide/derivative/). If none is passed in, the solver defaults to `FiniteDifference(Backward())`, which is effectively implicit Euler. +2. **Pyomo**: The list of Pyomo collocation methods can be found [here](). If none is passed in, the solver defaults to a `LagrangeRadau(3)`. + +```@docs; canonical = false +JuMPCollocation +InfiniteOptCollocation +CasADiCollocation +PyomoCollocation +solve(::AbstractDynamicOptProblem) +``` + +### Problem constructors +```@docs; canonical = false +JuMPDynamicOptProblem +InfiniteOptDynamicOptProblem +CasADiDynamicOptProblem +PyomoDynamicOptProblem +``` diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index 12ecf31196..2b4494863c 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -120,7 +120,20 @@ function MTK.lowered_var(m::CasADiModel, uv, i, t) t isa Union{Num, Symbolics.Symbolic} ? X.u[i, :] : X(t)[i] end -MTK.lowered_integral(model::CasADiModel, expr, args...) = model.tₛ * (model.U.t[2] - model.U.t[1]) * sum(expr) +function MTK.lowered_integral(model::CasADiModel, expr, lo, hi) + total = MX(0) + dt = model.U.t[2] - model.U.t[1] + for (i, t) in enumerate(model.U.t) + if lo < t < hi + Δt = min(dt, t - lo) + total += (0.5*Δt*(expr[i] + expr[i-1])) + elseif t >= hi && (t - dt < hi) + Δt = hi - t + dt + total += (0.5*Δt*(expr[i] + expr[i-1])) + end + end + model.tₛ * total +end function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau) @unpack A, α, c = tableau @@ -210,6 +223,7 @@ function MTK.get_t_values(model::CasADiModel) value_getter = MTK.successful_solve(model) ? CasADi.debug_value : CasADi.value ts = value_getter(model.solver_opti, model.tₛ) .* model.U.t end +MTK.objective_value(model::CasADiModel) = CasADi.pyconvert(Float64, model.solver_opti.py.value(model.solver_opti.py.f)) function MTK.successful_solve(m::CasADiModel) isnothing(m.solver_opti) && return false diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index 21cf117ab3..31183e3761 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -91,7 +91,7 @@ end MTK.lowered_integral(model::InfiniteOptModel, expr, lo, hi) = model.tₛ * InfiniteOpt.∫(expr, model.model[:t], lo, hi) MTK.lowered_derivative(model::InfiniteOptModel, i) = ∂(model.U[i], model.model[:t]) -function MTK.process_integral_bounds(model, integral_span, tspan) +function MTK.process_integral_bounds(model::InfiniteOptModel, integral_span, tspan) if MTK.is_free_final(model) && isequal(integral_span, tspan) integral_span = (0, 1) elseif MTK.is_free_final(model) @@ -213,6 +213,7 @@ function MTK.get_U_values(m::InfiniteModel) U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:nt] end MTK.get_t_values(m::InfiniteModel) = value(m[:tₛ]) * supports(m[:t]) +MTK.objective_value(m::InfiniteModel) = InfiniteOpt.objective_value(m) function MTK.successful_solve(model::InfiniteModel) tstatus = termination_status(model) diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl index bbd6dba115..8b668001f0 100644 --- a/ext/MTKPyomoDynamicOptExt.jl +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -8,16 +8,13 @@ using Setfield const MTK = ModelingToolkit SPECIAL_FUNCTIONS_DICT = Dict([acos => Pyomo.py_acos, - log1p => Pyomo.py_log1p, acosh => Pyomo.py_acosh, - log2 => Pyomo.py_log2, asin => Pyomo.py_asin, tan => Pyomo.py_tan, atanh => Pyomo.py_atanh, cos => Pyomo.py_cos, log => Pyomo.py_log, sin => Pyomo.py_sin, - log10 => Pyomo.py_log10, sqrt => Pyomo.py_sqrt, exp => Pyomo.py_exp]) @@ -36,10 +33,21 @@ struct PyomoDynamicOptModel function PyomoDynamicOptModel(model, U, V, tₛ, is_free_final) @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel) T_SYM DUMMY_SYM model.dU = dae.DerivativeVar(U, wrt = model.t, initialize = 0) + #add_time_equation!(model, MODEL_SYM, T_SYM) new(model, U, V, tₛ, is_free_final, nothing, PyomoVar(model.dU), MODEL_SYM, T_SYM, DUMMY_SYM) end end +function add_time_equation!(model::ConcreteModel, model_sym, t_sym) + model.dtime = dae.DerivativeVar(model.time) + + mdt = Symbolics.value(pysym_getproperty(model_sym, :dtime)) + mts = Symbolics.value(pysym_getproperty(model_sym, :tₛ)) + expr = mdt[t_sym] - mts == 0 + f = Pyomo.pyfunc(eval(Symbolics.build_function(expr, model_sym, t_sym))) + model.time_eq = pyomo.Constraint(model.t, rule = f) +end + struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <: AbstractDynamicOptProblem{uType, tType, isinplace} f::F @@ -72,6 +80,7 @@ MTK.generate_internal_model(m::Type{PyomoDynamicOptModel}) = ConcreteModel(pyomo function MTK.generate_time_variable!(m::ConcreteModel, tspan, tsteps) m.steps = length(tsteps) m.t = dae.ContinuousSet(initialize = tsteps, bounds = tspan) + m.time = pyomo.Var(m.t) end function MTK.generate_state_variable!(m::ConcreteModel, u0, ns, ts) @@ -116,7 +125,6 @@ end function MTK.set_objective!(pmodel::PyomoDynamicOptModel, expr) @unpack model, model_sym, t_sym, dummy_sym = pmodel - @show expr expr = Symbolics.substitute(expr, SPECIAL_FUNCTIONS_DICT) if occursin(Symbolics.unwrap(t_sym), expr) f = eval(Symbolics.build_function(expr, model_sym, t_sym)) @@ -134,15 +142,24 @@ function MTK.add_initial_constraints!(model::PyomoDynamicOptModel, u0, u0_idxs, end function MTK.lowered_integral(m::PyomoDynamicOptModel, arg, lo, hi) - @unpack model, model_sym, t_sym = m - arg_f = eval(Symbolics.build_function(arg, model_sym, t_sym)) - int_sym = Symbol(:int, hash(arg)) - setproperty!(model, int_sym, dae.Integral(model.t, wrt = model.t, rule=Pyomo.pyfunc(arg_f))) - PyomoVar(model.tₛ * model.int_sym) + @unpack model, model_sym, t_sym, dummy_sym = m + total = 0 + dt = Pyomo.pyconvert(Float64, (model.t.at(-1) - model.t.at(1))/(model.steps - 1)) + f = Symbolics.build_function(arg, model_sym, t_sym, expression = false) + for (i, t) in enumerate(model.t) + if Bool(lo < t) && Bool(t < hi) + t_p = model.t.at(i-1) + Δt = min(t - lo, t - t_p) + total += 0.5*Δt*(f(model, t) + f(model, t_p)) + elseif Bool(t >= hi) && Bool(t - dt < hi) + t_p = model.t.at(i-1) + Δt = hi - t + dt + total += 0.5*Δt*(f(model, t) + f(model, t_p)) + end + end + PyomoVar(model.tₛ * total) end -MTK.process_integral_bounds(model, integral_span, tspan) = integral_span - function MTK.lowered_derivative(m::PyomoDynamicOptModel, i) mdU = Symbolics.value(pysym_getproperty(m.model_sym, :dU)) Symbolics.unwrap(mdU[i, m.t_sym]) @@ -167,6 +184,7 @@ function MTK.prepare_and_optimize!(prob::PyomoDynamicOptProblem, collocation; ve discretizer = TransformationFactory(dm) ncp = Pyomo.is_finite_difference(dm) ? 1 : dm.np discretizer.apply_to(solver_m, wrt = solver_m.t, nfe = solver_m.steps - 1, scheme = Pyomo.scheme_string(dm)) + solver = SolverFactory(string(collocation.solver)) results = solver.solve(solver_m, tee = true) PyomoOutput(results, solver_m) @@ -190,10 +208,11 @@ function MTK.get_t_values(output::PyomoOutput) Pyomo.pyconvert(Float64, pyomo.value(m.tₛ)) * [Pyomo.pyconvert(Float64, t) for t in m.t] end +MTK.objective_value(output::PyomoOutput) = Pyomo.pyconvert(Float64, pyomo.value(output.model.obj)) + function MTK.successful_solve(output::PyomoOutput) r = output.result ss = r.solver.status - Main.xx[] = ss tc = r.solver.termination_condition if Bool(ss == opt.SolverStatus.ok) && (Bool(tc == opt.TerminationCondition.optimal) || Bool(tc == opt.TerminationCondition.locallyOptimal)) return true diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index f6cb499e91..d0668a98ef 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -331,6 +331,18 @@ function substitute_integral(model, exprs, tspan) exprs = map(c -> Symbolics.substitute(c, intmap), value.(exprs)) end +function process_integral_bounds(model, integral_span, tspan) + if is_free_final(model) && isequal(integral_span, tspan) + integral_span = (0, 1) + elseif is_free_final(model) + error("Free final time problems cannot handle partial timespans.") + else + (lo, hi) = integral_span + (lo < tspan[1] || hi > tspan[2]) && error("Integral bounds are beyond the timespan.") + integral_span + end +end + """Substitute variables like x(1.5), x(t), etc. with the corresponding model variables.""" function substitute_model_vars(model, sys, exprs, tspan) x_ops = [operation(unwrap(st)) for st in unknowns(sys)] @@ -405,6 +417,7 @@ function add_equational_constraints!(model, sys, pmap, tspan) end function set_objective! end +objective_value(sol::DynamicOptSolution) = objective_value(sol.model) function substitute_differentials(model, sys, eqs) t = get_iv(sys) diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index 85b06eee6f..feb8ddfba8 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -47,7 +47,7 @@ const M = ModelingToolkit @test ≈(csol2.sol.u, osol2.u, rtol = 0.001) pprob = PyomoDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01) psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) - @test all([≈(psol.sol(t), osol2(t), rtol = 1e-3) for t in 0.:0.01:1.]) + @test all([≈(psol.sol(t), osol2(t), rtol = 1e-2) for t in 0.:0.01:1.]) # With a constraint u0map = Pair[] @@ -111,9 +111,9 @@ function is_bangbang(input_sol, lbounds, ubounds, rtol = 1e-4) end function ctrl_to_spline(inputsol, splineType) - us = reduce(vcat, inputsol.u) - ts = reduce(vcat, inputsol.t) - splineType(us, ts) + us = reduce(vcat, inputsol.u) + ts = reduce(vcat, inputsol.t) + splineType(us, ts) end @testset "Linear systems" begin @@ -163,7 +163,8 @@ end @test is_bangbang(psol.input_sol, [-1.0], [1.0]) @test ≈(psol.sol.u[end][2], 0.25, rtol = 1e-3) - osol = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false) + spline = ctrl_to_spline(isol.input_sol, ConstantInterpolation) + oprob = ODEProblem(block_ode, u0map, tspan, [u_interp => spline]) @test ≈(isol.sol.u, osol.u, rtol = 0.05) @test all([≈(psol.sol(t), osol(t), rtol = 0.05) for t in 0.:0.01:1.]) @@ -250,7 +251,7 @@ end @test isol.sol[h(t)][end] > 1.012 pprob = PyomoDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) - psol = solve(pprob, PyomoCollocation("ipopt", MidpointEuler())) + psol = solve(pprob, PyomoCollocation("ipopt", LagrangeRadau(4))) @test psol.sol.u[end][1] > 1.012 # Test solution @@ -273,7 +274,7 @@ end interpmap2 = Dict(T_interp => ctrl_to_spline(psol.input_sol, CubicSpline)) oprob2 = ODEProblem(rocket_ode, u0map, (ts, te), merge(Dict(pmap), interpmap2)) osol2 = solve(oprob2, RadauIIA5(); adaptive = false, dt = 0.001) - @test ≈(psol.sol.u, osol2.u, rtol = 0.01) + @test all([≈(psol.sol(t), osol2(t), rtol = 0.01) for t in 0:0.001:0.2]) end @testset "Free final time problems" begin @@ -294,18 +295,22 @@ end jprob = JuMPDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 201) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5())) @test isapprox(jsol.sol.t[end], 10.0, rtol = 1e-3) + @test ≈(M.objective_value(jsol), -92.75, atol = 0.25) cprob = CasADiDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 201) csol = solve(cprob, CasADiCollocation("ipopt", constructTsitouras5())) @test isapprox(csol.sol.t[end], 10.0, rtol = 1e-3) + @test ≈(M.objective_value(csol), -92.75, atol = 0.25) iprob = InfiniteOptDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 200) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3) + @test ≈(M.objective_value(isol), -92.75, atol = 0.25) pprob = PyomoDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 201) - psol = solve(pprob, PyomoCollocation("ipopt")) + psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) @test isapprox(psol.sol.t[end], 10.0, rtol = 1e-3) + @test ≈(M.objective_value(psol), -92.75, atol = 0.1) @variables x(..) v(..) @variables u(..) [bounds = (-1.0, 1.0), input = true] @@ -375,7 +380,7 @@ end @test csol.sol.u[end] ≈ [π, 0, 0, 0] iprob = InfiniteOptDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04) - isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(2))) @test isol.sol.u[end] ≈ [π, 0, 0, 0] pprob = PyomoDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04) From f7d246f5e082f20329f162fcfe0c15d9d85af97e Mon Sep 17 00:00:00 2001 From: vyudu Date: Thu, 29 May 2025 03:36:31 -0400 Subject: [PATCH 14/22] docs: make examples dynamic --- docs/src/tutorials/dynamic_optimization.md | 27 ++++++++++++++++------ 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/docs/src/tutorials/dynamic_optimization.md b/docs/src/tutorials/dynamic_optimization.md index a1a1a44384..8cad7746c4 100644 --- a/docs/src/tutorials/dynamic_optimization.md +++ b/docs/src/tutorials/dynamic_optimization.md @@ -2,7 +2,8 @@ Systems in ModelingToolkit.jl can be directly converted to dynamic optimization or optimal control problems. In such systems, one has one or more input variables that are externally controlled to control the dynamics of the system. A dynamic optimization solves for the optimal time trajectory of the input variables in order to maximize or minimize a desired objective function. For example, a car driver might like to know how to step on the accelerator if the goal is to finish a race while using the least gas. To begin, let us take a rocket launch example. The input variable here is the thrust exerted by the engine. The rocket state is described by its current height and velocity. -```julia + +```@example dynamic_opt using ModelingToolkit t = ModelingToolkit.t_nounits D = ModelingToolkit.D_nounits @@ -37,19 +38,25 @@ pmap = [ What we would like to optimize here is the final height of the rocket. We do this by providing a vector of expressions corresponding to the costs. By default, the sense of the optimization is to minimize the provided cost. So to maximize the rocket height at the final time, we write `-h(te)` as the cost. Now we can construct a problem and solve it. Let us use JuMP as our backend here. -```julia +```@example dynamic_opt +using InfiniteOpt jprob = JuMPDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIIA5())) ``` Let's plot our final solution and the controller here: -```julia +```@example dynamic_opt +using CairoMakie +plot( + plot(jsol.sol, title = "Rocket trajectory"), + plot(jsol.input_sol, title = "Control trajectory") +) ``` ###### Free final time problems There are additionally a class of dynamic optimization problems where we would like to know how to control our system to achieve something in the least time. Such problems are called free final time problems, since the final time is unknown. To model these problems in ModelingToolkit, we declare the final time as a parameter. -```julia +```@example dynamic_opt @variables x(..) v(..) @variables u(..) [bounds = (-1.0, 1.0), input = true] @parameters tf @@ -69,8 +76,14 @@ parammap = [u(t) => 0.0, tf => 1.0] Please note that, at the moment, free final time problems cannot support constraints defined at definite time values, like `x(3) ~ 2`. -Let's plot our final solution and the controller for the block: -```julia +!!! warning + + The Pyomo collocation methods (LagrangeRadau, LagrangeLegendre) currently are bugged for free final time problems. Strongly suggest using BackwardEuler() for such problems. + +Let's solve plot our final solution and the controller for the block, using InfiniteOpt as the backend: +```@example dynamic_opt +iprob = InfiniteOptDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) +isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) ``` ### Solvers @@ -86,7 +99,7 @@ CasADiCollocation("ipopt", constructRK4()) **Pyomo** and **InfiniteOpt** each have their own built-in collocation methods. 1. **InfiniteOpt**: The list of InfiniteOpt collocation methods can be found [in the table on this page](https://infiniteopt.github.io/InfiniteOpt.jl/stable/guide/derivative/). If none is passed in, the solver defaults to `FiniteDifference(Backward())`, which is effectively implicit Euler. -2. **Pyomo**: The list of Pyomo collocation methods can be found [here](). If none is passed in, the solver defaults to a `LagrangeRadau(3)`. +2. **Pyomo**: The list of Pyomo collocation methods can be found [at the bottom of this page](https://github.com/SciML/Pyomo.jl). If none is passed in, the solver defaults to a `LagrangeRadau(3)`. ```@docs; canonical = false JuMPCollocation From 10719ae0465300490a7f2ee2a13f146853e2f4e3 Mon Sep 17 00:00:00 2001 From: vyudu Date: Thu, 29 May 2025 12:20:57 -0400 Subject: [PATCH 15/22] fix: update ext/tests to v10 --- ext/MTKCasADiDynamicOptExt.jl | 5 +- ext/MTKInfiniteOptExt.jl | 9 +-- ext/MTKPyomoDynamicOptExt.jl | 4 +- src/systems/optimal_control_interface.jl | 68 +++++++++++---------- test/extensions/dynamic_optimization.jl | 77 ++++++++++++------------ 5 files changed, 84 insertions(+), 79 deletions(-) diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index 2b4494863c..73e4a77ed4 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -61,11 +61,12 @@ function (M::MXLinearInterpolation)(τ) end end -function MTK.CasADiDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; +function MTK.CasADiDynamicOptProblem(sys::System, op, tspan; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - MTK.process_DynamicOptProblem(CasADiDynamicOptProblem, CasADiModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + prob, _ = MTK.process_DynamicOptProblem(CasADiDynamicOptProblem, CasADiModel, sys, op, tspan; dt, steps, guesses, kwargs...) + prob end MTK.generate_internal_model(::Type{CasADiModel}) = CasADi.Opti() diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index 31183e3761..8150cd06f7 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -72,18 +72,19 @@ function MTK.add_constraint!(m::InfiniteOptModel, expr::Union{Equation, Inequali end MTK.set_objective!(m::InfiniteOptModel, expr) = @objective(m.model, Min, expr) -function MTK.JuMPDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; +function MTK.JuMPDynamicOptProblem(sys::System, op, tspan; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - MTK.process_DynamicOptProblem(JuMPDynamicOptProblem, InfiniteOptModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + prob, _ = MTK.process_DynamicOptProblem(JuMPDynamicOptProblem, InfiniteOptModel, sys, op, tspan; dt, steps, guesses, kwargs...) + prob end -function MTK.InfiniteOptDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; +function MTK.InfiniteOptDynamicOptProblem(sys::System, op, tspan; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - prob = MTK.process_DynamicOptProblem(InfiniteOptDynamicOptProblem, InfiniteOptModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + prob, pmap = MTK.process_DynamicOptProblem(InfiniteOptDynamicOptProblem, InfiniteOptModel, sys, op, tspan; dt, steps, guesses, kwargs...) MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) prob end diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl index 8b668001f0..2301a86b68 100644 --- a/ext/MTKPyomoDynamicOptExt.jl +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -66,10 +66,10 @@ end pysym_getproperty(s::Union{Num, Symbolics.Symbolic}, name::Symbol) = Symbolics.wrap(SymbolicUtils.term(_getproperty, Symbolics.unwrap(s), Val{name}(), type = Symbolics.Struct{PyomoVar})) _getproperty(s, name::Val{fieldname}) where fieldname = getproperty(s, fieldname) -function MTK.PyomoDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; +function MTK.PyomoDynamicOptProblem(sys::System, op, tspan; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - prob = MTK.process_DynamicOptProblem(PyomoDynamicOptProblem, PyomoDynamicOptModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) + prob, pmap = MTK.process_DynamicOptProblem(PyomoDynamicOptProblem, PyomoDynamicOptModel, sys, op, tspan; dt, steps, guesses, kwargs...) conc_model = prob.wrapped_model.model MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) prob diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index d0668a98ef..ea19afea86 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -19,9 +19,9 @@ function Base.show(io::IO, sol::DynamicOptSolution) end """ - JuMPDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt, steps, guesses, kwargs...) + JuMPDynamicOptProblem(sys::System, u0, tspan, p; dt, steps, guesses, kwargs...) -Convert an ODESystem representing an optimal control system into a JuMP model +Convert an System representing an optimal control system into a JuMP model for solving using optimization. Must provide either `dt`, the timestep between collocation points (which, along with the timespan, determines the number of points), or directly provide the number of points as `steps`. @@ -30,9 +30,9 @@ To construct the problem, please load InfiniteOpt along with ModelingToolkit. """ function JuMPDynamicOptProblem end """ - InfiniteOptDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; dt) + InfiniteOptDynamicOptProblem(sys::System, op, tspan; dt) -Convert an ODESystem representing an optimal control system into a InfiniteOpt model +Convert an System representing an optimal control system into a InfiniteOpt model for solving using optimization. Must provide `dt` for determining the length of the interpolation arrays. @@ -43,9 +43,9 @@ To construct the problem, please load InfiniteOpt along with ModelingToolkit. """ function InfiniteOptDynamicOptProblem end """ - CasADiDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt, steps, guesses, kwargs...) + CasADiDynamicOptProblem(sys::System, u0, tspan, p; dt, steps, guesses, kwargs...) -Convert an ODESystem representing an optimal control system into a CasADi model +Convert an System representing an optimal control system into a CasADi model for solving using optimization. Must provide either `dt`, the timestep between collocation points (which, along with the timespan, determines the number of points), or directly provide the number of points as `steps`. @@ -54,9 +54,9 @@ To construct the problem, please load CasADi along with ModelingToolkit. """ function CasADiDynamicOptProblem end """ - PyomoDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt, steps) + PyomoDynamicOptProblem(sys::System, u0, tspan, p; dt, steps) -Convert an ODESystem representing an optimal control system into a Pyomo model +Convert an System representing an optimal control system into a Pyomo model for solving using optimization. Must provide either `dt`, the timestep between collocation points (which, along with the timespan, determines the number of points), or directly provide the number of points as `steps`. @@ -91,11 +91,12 @@ Pyomo Collocation solver. """ function PyomoCollocation end -function warn_overdetermined(sys, u0map) +function warn_overdetermined(sys, op) cstrs = constraints(sys) + init_conds = filter(x -> value(x) ∈ Set(unknowns(sys)), [k for (k,v) in op]) if !isempty(cstrs) - (length(cstrs) + 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." + (length(cstrs) + length(init_conds) > length(unknowns(sys))) && + @warn "The control problem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by op) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization." end end @@ -226,24 +227,27 @@ end ########################## ### MODEL CONSTRUCTION ### ########################## -function process_DynamicOptProblem(prob_type::Type{<:AbstractDynamicOptProblem}, model_type, sys::ODESystem, u0map, tspan, pmap; +function process_DynamicOptProblem(prob_type::Type{<:AbstractDynamicOptProblem}, model_type, sys::System, op, tspan; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - warn_overdetermined(sys, u0map) + + warn_overdetermined(sys, op) ctrls = unbound_inputs(sys) states = unknowns(sys) - _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses)) stidxmap = Dict([v => i for (i, v) in enumerate(states)]) - u0map = Dict([default_toterm(value(k)) => v for (k, v) in u0map]) + op = Dict([default_toterm(value(k)) => v for (k, v) in op]) u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) : - [stidxmap[default_toterm(k)] for (k, v) in u0map] + [stidxmap[default_toterm(k)] for (k, v) in op if haskey(stidxmap, k)] - f, u0, p = process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap; + _op = has_alg_eqs(sys) ? op : merge(Dict(op), Dict(guesses)) + f, u0, p = process_SciMLProblem(ODEInputFunction, sys, _op; t = tspan !== nothing ? tspan[1] : tspan, kwargs...) model_tspan, steps, is_free_t = process_tspan(tspan, dt, steps) + warn_overdetermined(sys, op) + pmap = filter(p -> (first(p) ∉ Set(unknowns(sys))), op) pmap = recursive_unwrap(AnyDict(pmap)) evaluate_varmap!(pmap, keys(pmap)) c0 = value.([pmap[c] for c in ctrls]) @@ -261,7 +265,7 @@ function process_DynamicOptProblem(prob_type::Type{<:AbstractDynamicOptProblem}, add_user_constraints!(fullmodel, sys, tspan, pmap) add_initial_constraints!(fullmodel, u0, u0_idxs, model_tspan[1]) - prob_type(f, u0, tspan, p, fullmodel, kwargs...) + prob_type(f, u0, tspan, p, fullmodel, kwargs...), pmap end function generate_time_variable! end @@ -301,16 +305,16 @@ end is_free_final(model) = model.is_free_final function add_cost_function!(model, sys, tspan, pmap) - jcosts = copy(get_costs(sys)) - consolidate = get_consolidate(sys) - if isnothing(jcosts) || isempty(jcosts) + jcosts = cost(sys) + if Symbolics._iszero(jcosts) set_objective!(model, 0) return end - jcosts = substitute_model_vars(model, sys, jcosts, tspan) + + jcosts = substitute_model_vars(model, sys, [jcosts], tspan) jcosts = substitute_params(pmap, jcosts) - jcosts = substitute_integral(model, jcosts, tspan) - set_objective!(model, consolidate(jcosts)) + jcosts = substitute_integral(model, only(jcosts), tspan) + set_objective!(model, value(jcosts)) end """ @@ -319,16 +323,16 @@ Substitute integrals. For an integral from (ts, te): - CasADi cannot handle partial timespans, even for non-free-final time problems. time problems and unchanged otherwise. """ -function substitute_integral(model, exprs, tspan) +function substitute_integral(model, expr, tspan) intmap = Dict() - for int in collect_applied_operators(exprs, Symbolics.Integral) + for int in collect_applied_operators(expr, Symbolics.Integral) op = operation(int) arg = only(arguments(value(int))) lo, hi = value.((op.domain.domain.left, op.domain.domain.right)) lo, hi = process_integral_bounds(model, (lo, hi), tspan) intmap[int] = lowered_integral(model, arg, lo, hi) end - exprs = map(c -> Symbolics.substitute(c, intmap), value.(exprs)) + Symbolics.substitute(expr, intmap) end function process_integral_bounds(model, integral_span, tspan) @@ -386,13 +390,13 @@ function lowered_var end function fixed_t_map end function add_user_constraints!(model, sys, tspan, pmap) - conssys = get_constraintsystem(sys) - jconstraints = isnothing(conssys) ? nothing : get_constraints(conssys) + jconstraints = get_constraints(sys) (isnothing(jconstraints) || isempty(jconstraints)) && return nothing - consvars = get_unknowns(conssys) - is_free_final(model) && check_constraint_vars(consvars) + cons_dvs, cons_ps = process_constraint_system(jconstraints, Set(unknowns(sys)), parameters(sys), get_iv(sys); validate = false) + + is_free_final(model) && check_constraint_vars(cons_dvs) - jconstraints = substitute_toterm(consvars, jconstraints) + jconstraints = substitute_toterm(cons_dvs, jconstraints) jconstraints = substitute_model_vars(model, sys, jconstraints, tspan) jconstraints = substitute_params(pmap, jconstraints) diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index feb8ddfba8..0243b298d6 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -27,11 +27,11 @@ const M = ModelingToolkit parammap = [α => 1.5, β => 1.0, γ => 3.0, δ => 1.0] # Test explicit method. - jprob = JuMPDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01) + jprob = JuMPDynamicOptProblem(sys, [u0map; parammap], tspan, dt = 0.01) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRK4())) oprob = ODEProblem(sys, [u0map; parammap], tspan) osol = solve(oprob, SimpleRK4(), dt = 0.01) - cprob = CasADiDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01) + cprob = CasADiDynamicOptProblem(sys, [u0map; parammap], tspan, dt = 0.01) csol = solve(cprob, CasADiCollocation("ipopt", constructRK4())) @test jsol.sol.u ≈ osol.u @test csol.sol.u ≈ osol.u @@ -40,12 +40,12 @@ const M = ModelingToolkit osol2 = solve(oprob, ImplicitEuler(), dt = 0.01, adaptive = false) jsol2 = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructImplicitEuler())) @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001) - iprob = InfiniteOptDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01) + iprob = InfiniteOptDynamicOptProblem(sys, [u0map; parammap], tspan, dt = 0.01) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test ≈(isol.sol.u, osol2.u, rtol = 0.001) csol2 = solve(cprob, CasADiCollocation("ipopt", constructImplicitEuler())) @test ≈(csol2.sol.u, osol2.u, rtol = 0.001) - pprob = PyomoDynamicOptProblem(sys, u0map, tspan, parammap, dt = 0.01) + pprob = PyomoDynamicOptProblem(sys, [u0map; parammap], tspan, dt = 0.01) psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) @test all([≈(psol.sol(t), osol2(t), rtol = 1e-2) for t in 0.:0.01:1.]) @@ -55,26 +55,26 @@ const M = ModelingToolkit constr = [x(0.6) ~ 3.5, x(0.3) ~ 7.0] @mtkcompile lksys = System(eqs, t; constraints = constr) - jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + jprob = JuMPDynamicOptProblem(lksys, [u0map; parammap], tspan; guesses = guess, dt = 0.01) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5())) @test jsol.sol(0.6; idxs = x(t)) ≈ 3.5 @test jsol.sol(0.3; idxs = x(t)) ≈ 7.0 cprob = CasADiDynamicOptProblem( - lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + lksys, [u0map; parammap], tspan; guesses = guess, dt = 0.01) csol = solve(cprob, CasADiCollocation("ipopt", constructTsitouras5())) @test csol.sol(0.6; idxs = x(t)) ≈ 3.5 @test csol.sol(0.3; idxs = x(t)) ≈ 7.0 pprob = PyomoDynamicOptProblem( - lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + lksys, [u0map; parammap], tspan; guesses = guess, dt = 0.01) psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(3))) @show psol.sol @test psol.sol(0.6)[1] ≈ 3.5 @test psol.sol(0.3)[1] ≈ 7.0 iprob = InfiniteOptDynamicOptProblem( - lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + lksys, [u0map; parammap], tspan; guesses = guess, dt = 0.01) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3))) # 48.564 ms, 9.58 MiB sol = isol.sol @test sol(0.6; idxs = x(t)) ≈ 3.5 @@ -84,20 +84,20 @@ const M = ModelingToolkit constr = [x(t) ≳ 1, y(t) ≳ 1] @mtkcompile lksys = System(eqs, t; constraints = constr) iprob = InfiniteOptDynamicOptProblem( - lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + lksys, [u0map; parammap], tspan; guesses = guess, dt = 0.01) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3))) @test all(u -> u > [1, 1], isol.sol.u) - jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + jprob = JuMPDynamicOptProblem(lksys, [u0map; parammap], tspan; guesses = guess, dt = 0.01) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIA3())) @test all(u -> u > [1, 1], jsol.sol.u) - pprob = PyomoDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + pprob = PyomoDynamicOptProblem(lksys, [u0map; parammap], tspan; guesses = guess, dt = 0.01) psol = solve(pprob, PyomoCollocation("ipopt", MidpointEuler())) @test all(u -> u > [1, 1], psol.sol.u) cprob = CasADiDynamicOptProblem( - lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01) + lksys, [u0map; parammap], tspan; guesses = guess, dt = 0.01) csol = solve(cprob, CasADiCollocation("ipopt", constructRadauIA3())) @test all(u -> u > [1, 1], csol.sol.u) end @@ -131,14 +131,14 @@ end u0map = [x(t) => 0.0, v(t) => 0.0] tspan = (0.0, 1.0) parammap = [u(t) => 0.0] - jprob = JuMPDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01) + jprob = JuMPDynamicOptProblem(block, [u0map; parammap], tspan; dt = 0.01) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructVerner8()), verbose = true) # Linear systems have bang-bang controls @test is_bangbang(jsol.input_sol, [-1.0], [1.0]) # Test reached final position. @test ≈(jsol.sol[x(t)][end], 0.25, rtol = 1e-5) - cprob = CasADiDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01) + cprob = CasADiDynamicOptProblem(block, [u0map; parammap], tspan; dt = 0.01) csol = solve(cprob, CasADiCollocation("ipopt", constructVerner8())) @test is_bangbang(csol.input_sol, [-1.0], [1.0]) # Test reached final position. @@ -153,18 +153,18 @@ end @test ≈(jsol.sol.u, osol.u, rtol = 0.05) @test ≈(csol.sol.u, osol.u, rtol = 0.05) - iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01) + iprob = InfiniteOptDynamicOptProblem(block, [u0map; parammap], tspan; dt = 0.01) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test is_bangbang(isol.input_sol, [-1.0], [1.0]) @test ≈(isol.sol[x(t)][end], 0.25, rtol = 1e-5) - pprob = PyomoDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01) + pprob = PyomoDynamicOptProblem(block, [u0map; parammap], tspan; dt = 0.01) psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) @test is_bangbang(psol.input_sol, [-1.0], [1.0]) @test ≈(psol.sol.u[end][2], 0.25, rtol = 1e-3) spline = ctrl_to_spline(isol.input_sol, ConstantInterpolation) - oprob = ODEProblem(block_ode, u0map, tspan, [u_interp => spline]) + oprob = ODEProblem(block_ode, [u0map; u_interp => spline], tspan) @test ≈(isol.sol.u, osol.u, rtol = 0.05) @test all([≈(psol.sol(t), osol(t), rtol = 0.05) for t in 0.:0.01:1.]) @@ -185,16 +185,16 @@ end u0map = [w(t) => 40, q(t) => 2] pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1, α => 1] - jprob = JuMPDynamicOptProblem(beesys, u0map, tspan, pmap, dt = 0.01) + jprob = JuMPDynamicOptProblem(beesys, [u0map; pmap], tspan, dt = 0.01) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5())) @test is_bangbang(jsol.input_sol, [0.0], [1.0]) - iprob = InfiniteOptDynamicOptProblem(beesys, u0map, tspan, pmap, dt = 0.01) + iprob = InfiniteOptDynamicOptProblem(beesys, [u0map; pmap], tspan, dt = 0.01) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test is_bangbang(isol.input_sol, [0.0], [1.0]) - cprob = CasADiDynamicOptProblem(beesys, u0map, tspan, pmap; dt = 0.01) + cprob = CasADiDynamicOptProblem(beesys, [u0map; pmap], tspan; dt = 0.01) csol = solve(cprob, CasADiCollocation("ipopt", constructTsitouras5())) @test is_bangbang(csol.input_sol, [0.0], [1.0]) - pprob = PyomoDynamicOptProblem(beesys, u0map, tspan, pmap, dt = 0.01) + pprob = PyomoDynamicOptProblem(beesys, [u0map; pmap], tspan, dt = 0.01) psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) @test is_bangbang(psol.input_sol, [0.0], [1.0]) @@ -237,20 +237,19 @@ 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.001, cse = false) + jprob = JuMPDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001, cse = false) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIIA5())) @test jsol.sol[h(t)][end] > 1.012 - cprob = CasADiDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) + cprob = CasADiDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001, cse = false) csol = solve(cprob, CasADiCollocation("ipopt")) @test csol.sol[h(t)][end] > 1.012 - iprob = InfiniteOptDynamicOptProblem( - rocket, u0map, (ts, te), pmap; dt = 0.001) + iprob = InfiniteOptDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isol.sol[h(t)][end] > 1.012 - pprob = PyomoDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) + pprob = PyomoDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001, cse = false) psol = solve(pprob, PyomoCollocation("ipopt", LagrangeRadau(4))) @test psol.sol.u[end][1] > 1.012 @@ -272,7 +271,7 @@ end @test ≈(isol.sol.u, osol1.u, rtol = 0.01) interpmap2 = Dict(T_interp => ctrl_to_spline(psol.input_sol, CubicSpline)) - oprob2 = ODEProblem(rocket_ode, u0map, (ts, te), merge(Dict(pmap), interpmap2)) + oprob2 = ODEProblem(rocket_ode, merge(Dict(u0map), Dict(pmap), interpmap2), (ts, te)) osol2 = solve(oprob2, RadauIIA5(); adaptive = false, dt = 0.001) @test all([≈(psol.sol(t), osol2(t), rtol = 0.01) for t in 0:0.001:0.2]) end @@ -292,22 +291,22 @@ end u0map = [x(t) => 17.5] pmap = [u(t) => 0.0, tf => 8] - jprob = JuMPDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 201) + jprob = JuMPDynamicOptProblem(rocket, [u0map; pmap], (0, tf); steps = 201) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5())) @test isapprox(jsol.sol.t[end], 10.0, rtol = 1e-3) @test ≈(M.objective_value(jsol), -92.75, atol = 0.25) - cprob = CasADiDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 201) + cprob = CasADiDynamicOptProblem(rocket, [u0map; pmap], (0, tf); steps = 201) csol = solve(cprob, CasADiCollocation("ipopt", constructTsitouras5())) @test isapprox(csol.sol.t[end], 10.0, rtol = 1e-3) @test ≈(M.objective_value(csol), -92.75, atol = 0.25) - iprob = InfiniteOptDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 200) + iprob = InfiniteOptDynamicOptProblem(rocket, [u0map; pmap], (0, tf); steps = 200) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3) @test ≈(M.objective_value(isol), -92.75, atol = 0.25) - pprob = PyomoDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 201) + pprob = PyomoDynamicOptProblem(rocket, [u0map; pmap], (0, tf); steps = 201) psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) @test isapprox(psol.sol.t[end], 10.0, rtol = 1e-3) @test ≈(M.objective_value(psol), -92.75, atol = 0.1) @@ -325,19 +324,19 @@ end u0map = [x(t) => 1.0, v(t) => 0.0] tspan = (0.0, tf) parammap = [u(t) => 1.0, tf => 1.0] - jprob = JuMPDynamicOptProblem(block, u0map, tspan, parammap; steps = 51) + jprob = JuMPDynamicOptProblem(block, [u0map; parammap], tspan; steps = 51) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructVerner8())) @test isapprox(jsol.sol.t[end], 2.0, atol = 1e-5) - cprob = CasADiDynamicOptProblem(block, u0map, (0, tf), parammap; steps = 51) + cprob = CasADiDynamicOptProblem(block, [u0map; parammap], (0, tf); steps = 51) csol = solve(cprob, CasADiCollocation("ipopt", constructVerner8())) @test isapprox(csol.sol.t[end], 2.0, atol = 1e-5) - iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; steps = 51) + iprob = InfiniteOptDynamicOptProblem(block, [u0map; parammap], tspan; steps = 51) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer), verbose = true) @test isapprox(isol.sol.t[end], 2.0, atol = 1e-5) - pprob = PyomoDynamicOptProblem(block, u0map, tspan, parammap; steps = 51) + pprob = PyomoDynamicOptProblem(block, [u0map; parammap], tspan; steps = 51) psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) @test isapprox(psol.sol.t[end], 2.0, atol = 1e-5) end @@ -371,19 +370,19 @@ 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) + jprob = JuMPDynamicOptProblem(cartpole, [u0map; pmap], tspan; dt = 0.04) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRK4())) @test jsol.sol.u[end] ≈ [π, 0, 0, 0] - cprob = CasADiDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04) + cprob = CasADiDynamicOptProblem(cartpole, [u0map; pmap], tspan; dt = 0.04) csol = solve(cprob, CasADiCollocation("ipopt", constructRK4())) @test csol.sol.u[end] ≈ [π, 0, 0, 0] - iprob = InfiniteOptDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04) + iprob = InfiniteOptDynamicOptProblem(cartpole, [u0map; pmap], tspan; dt = 0.04) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(2))) @test isol.sol.u[end] ≈ [π, 0, 0, 0] - pprob = PyomoDynamicOptProblem(cartpole, u0map, tspan, pmap; dt = 0.04) + pprob = PyomoDynamicOptProblem(cartpole, [u0map; pmap], tspan; dt = 0.04) psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(4))) @test psol.sol.u[end] ≈ [π, 0, 0, 0] end From 7daecfbf3c88f4179791fbd94b39dcc11e9511b8 Mon Sep 17 00:00:00 2001 From: vyudu Date: Thu, 29 May 2025 18:32:40 -0400 Subject: [PATCH 16/22] docs: update docs to v10 --- docs/Project.toml | 3 + docs/src/tutorials/dynamic_optimization.md | 87 +++++++++++++++------- 2 files changed, 65 insertions(+), 25 deletions(-) diff --git a/docs/Project.toml b/docs/Project.toml index ef6af51d79..314a1d7f84 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -5,11 +5,14 @@ BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665" CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" ControlSystemsBase = "aaaaaaaa-a6ca-5380-bf3e-84a91bcd477e" DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" +DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" DynamicQuantities = "06fc5a27-2a28-4c7c-a15d-362465fb6821" FMI = "14a09403-18e3-468f-ad8a-74f8dda2d9ac" FMIZoo = "724179cf-c260-40a9-bd27-cccc6fe2f195" +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" diff --git a/docs/src/tutorials/dynamic_optimization.md b/docs/src/tutorials/dynamic_optimization.md index 8cad7746c4..7753516795 100644 --- a/docs/src/tutorials/dynamic_optimization.md +++ b/docs/src/tutorials/dynamic_optimization.md @@ -1,7 +1,7 @@ # Solving Dynamic Optimization Problems Systems in ModelingToolkit.jl can be directly converted to dynamic optimization or optimal control problems. In such systems, one has one or more input variables that are externally controlled to control the dynamics of the system. A dynamic optimization solves for the optimal time trajectory of the input variables in order to maximize or minimize a desired objective function. For example, a car driver might like to know how to step on the accelerator if the goal is to finish a race while using the least gas. -To begin, let us take a rocket launch example. The input variable here is the thrust exerted by the engine. The rocket state is described by its current height and velocity. +To begin, let us take a rocket launch example. The input variable here is the thrust exerted by the engine. The rocket state is described by its current height, mass, and velocity. The mass decreases as the rocket loses fuel while thrusting. ```@example dynamic_opt using ModelingToolkit @@ -12,8 +12,8 @@ D = ModelingToolkit.D_nounits @variables begin h(..) v(..) - m(..) [bounds = (m_c, 1)] - T(..) [input = true, bounds = (0, Tₘ)] + m(..), [bounds = (m_c, 1)] + T(..), [input = true, bounds = (0, Tₘ)] end drag(h, v) = D_c * v^2 * exp(-h_c * (h - h₀) / h₀) @@ -27,8 +27,8 @@ eqs = [D(h(t)) ~ v(t), costs = [-h(te)] cons = [T(te) ~ 0, m(te) ~ m_c] -@named rocket = ODESystem(eqs, t; costs, constraints = cons) -rocket, input_idxs = structural_simplify(rocket, ([T(t)], [])) +@named rocket = System(eqs, t; costs, constraints = cons) +rocket = mtkcompile(rocket, inputs = [T(t)]) u0map = [h(t) => h₀, m(t) => m₀, v(t) => 0] pmap = [ @@ -37,70 +37,107 @@ pmap = [ ``` What we would like to optimize here is the final height of the rocket. We do this by providing a vector of expressions corresponding to the costs. By default, the sense of the optimization is to minimize the provided cost. So to maximize the rocket height at the final time, we write `-h(te)` as the cost. -Now we can construct a problem and solve it. Let us use JuMP as our backend here. +Now we can construct a problem and solve it. Let us use JuMP as our backend here. Note that the package trigger is actually [InfiniteOpt](https://infiniteopt.github.io/InfiniteOpt.jl/stable/), and not JuMP - this package includes JuMP but is designed for optimization on function spaces. Additionally we need to load the solver package - we will use [Ipopt](https://github.com/jump-dev/Ipopt.jl) here (a good choice in general). + +Here we have also loaded DiffEqDevTools because we will need to construct the ODE tableau. This is only needed if one desires a custom ODE tableau for the collocation - by default the solver will use RadauIIA5. ```@example dynamic_opt -using InfiniteOpt -jprob = JuMPDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) -jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIIA5())) +using InfiniteOpt, Ipopt, DiffEqDevTools +jprob = JuMPDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001) +jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIIA5())); ``` +The solution has three fields: `jsol.sol` is the ODE solution for the states, `jsol.input_sol` is the ODE solution for the inputs, and `jsol.model` is the wrapped model that we can use to query things like objective and constraint residuals. -Let's plot our final solution and the controller here: +Let's plot the final solution and the controller here: ```@example dynamic_opt using CairoMakie -plot( - plot(jsol.sol, title = "Rocket trajectory"), - plot(jsol.input_sol, title = "Control trajectory") -) +fig = Figure(resolution = (800, 400)) +ax1 = Axis(fig[1,1], title = "Rocket trajectory", xlabel = "Time") +ax2 = Axis(fig[1,2], title = "Control trajectory", xlabel = "Time") + +for u in unknowns(rocket) + lines!(ax1, jsol.sol.t, jsol.sol[u], label = string(u)) +end +lines!(ax2, jsol.input_sol, label = "Thrust") +axislegend(ax1) +axislegend(ax2) +fig ``` -###### Free final time problems +### Free final time problems There are additionally a class of dynamic optimization problems where we would like to know how to control our system to achieve something in the least time. Such problems are called free final time problems, since the final time is unknown. To model these problems in ModelingToolkit, we declare the final time as a parameter. +Below we have a model system called the double integrator. We control the acceleration of a block in order to reach a desired destination in the least time. ```@example dynamic_opt -@variables x(..) v(..) -@variables u(..) [bounds = (-1.0, 1.0), input = true] +@variables begin + x(..) + v(..) + u(..), [bounds = (-1.0, 1.0), input = true] +end + @parameters tf constr = [v(tf) ~ 0, x(tf) ~ 0] cost = [tf] # Minimize time -@named block = ODESystem( +@named block = System( [D(x(t)) ~ v(t), D(v(t)) ~ u(t)], t; costs = cost, constraints = constr) -block, input_idxs = structural_simplify(block, ([u(t)], [])) +block = mtkcompile(block; inputs = [u(t)]) u0map = [x(t) => 1.0, v(t) => 0.0] tspan = (0.0, tf) parammap = [u(t) => 0.0, tf => 1.0] ``` +The `tf` mapping in the parameter map is treated as an initial guess. Please note that, at the moment, free final time problems cannot support constraints defined at definite time values, like `x(3) ~ 2`. !!! warning - The Pyomo collocation methods (LagrangeRadau, LagrangeLegendre) currently are bugged for free final time problems. Strongly suggest using BackwardEuler() for such problems. + The Pyomo collocation methods (LagrangeRadau, LagrangeLegendre) currently are bugged for free final time problems. Strongly suggest using BackwardEuler() for such problems when using Pyomo as the backend. -Let's solve plot our final solution and the controller for the block, using InfiniteOpt as the backend: +When declaring the problem in this case we need to provide the number of steps, since dt can't be known in advanced. Let's solve plot our final solution and the controller for the block, using InfiniteOpt as the backend: ```@example dynamic_opt -iprob = InfiniteOptDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false) -isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) +iprob = InfiniteOptDynamicOptProblem(block, [u0map; parammap], tspan; steps = 100) +isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)); +``` + +Let's plot the final solution and the controller here: +```@example dynamic_opt +fig = Figure(resolution = (800, 400)) +ax1 = Axis(fig[1,1], title = "Block trajectory", xlabel = "Time") +ax2 = Axis(fig[1,2], title = "Control trajectory", xlabel = "Time") + +for u in unknowns(block) + lines!(ax1, isol.sol.t, isol.sol[u], label = string(u)) +end +lines!(ax2, isol.input_sol, label = "Acceleration") +axislegend(ax1) +axislegend(ax2) +fig ``` ### Solvers Currently 4 backends are exposed for solving dynamic optimization problems using collocation: JuMP, InfiniteOpt, CasADi, and Pyomo. -Please note that there are differences in how to construct the collocation solver for the different cases. For example, the Python based ones, CasADi and Pyomo, expect the solver to be passed in as a string (CasADi and Pyomo come pre-loaded with certain solvers, but other solvers may need to be manually installed using `pip` or `conda`), while JuMP/InfiniteOpt expect the optimizer object to be passed in directly: +Please note that there are differences in how to construct the collocation solver for the different cases. For example, the Python based ones, CasADi and Pyomo, expect the solver to be passed in as a string (CasADi and Pyomo come pre-loaded with Ipopt, but other solvers may need to be manually installed using `pip` or `conda`), while JuMP/InfiniteOpt expect the optimizer object to be passed in directly: ``` JuMPCollocation(Ipopt.Optimizer, constructRK4()) CasADiCollocation("ipopt", constructRK4()) ``` -**JuMP** and **CasADi** collocation require an ODE tableau to be passed in. These can be constructed by calling the `constructX()` functions from DiffEqDevTools. If none is passed in, both solvers will default to using Radau second-order with five collocation points. +**JuMP** and **CasADi** collocation require an ODE tableau to be passed in. These can be constructed by calling the `constructX()` functions from DiffEqDevTools. The list of tableaus can be found [here](https://docs.sciml.ai/DiffEqDevDocs/dev/internals/tableaus/). If none is passed in, both solvers will default to using Radau second-order with five collocation points. **Pyomo** and **InfiniteOpt** each have their own built-in collocation methods. 1. **InfiniteOpt**: The list of InfiniteOpt collocation methods can be found [in the table on this page](https://infiniteopt.github.io/InfiniteOpt.jl/stable/guide/derivative/). If none is passed in, the solver defaults to `FiniteDifference(Backward())`, which is effectively implicit Euler. 2. **Pyomo**: The list of Pyomo collocation methods can be found [at the bottom of this page](https://github.com/SciML/Pyomo.jl). If none is passed in, the solver defaults to a `LagrangeRadau(3)`. +Some examples of the latter two collocations: +```julia +PyomoCollocation("ipopt", LagrangeRadau(2)) +InfiniteOptCollocation(Ipopt.Optimizer, OrthogonalCollocation(3)) +``` + ```@docs; canonical = false JuMPCollocation InfiniteOptCollocation From f33f8c70b429db71eb0ab9ce168f75557ecb72db Mon Sep 17 00:00:00 2001 From: vyudu Date: Mon, 2 Jun 2025 10:02:36 -0400 Subject: [PATCH 17/22] fix: error when using collocation on fft problems --- docs/src/API/dynamic_opt.md | 41 ++++++++++++++++++++++++ ext/MTKPyomoDynamicOptExt.jl | 60 +++++++++++++++++++++++------------- 2 files changed, 80 insertions(+), 21 deletions(-) create mode 100644 docs/src/API/dynamic_opt.md diff --git a/docs/src/API/dynamic_opt.md b/docs/src/API/dynamic_opt.md new file mode 100644 index 0000000000..1a7081154c --- /dev/null +++ b/docs/src/API/dynamic_opt.md @@ -0,0 +1,41 @@ +### Solvers + +Currently 4 backends are exposed for solving dynamic optimization problems using collocation: JuMP, InfiniteOpt, CasADi, and Pyomo. + +Please note that there are differences in how to construct the collocation solver for the different cases. For example, the Python based ones, CasADi and Pyomo, expect the solver to be passed in as a string (CasADi and Pyomo come pre-loaded with Ipopt, but other solvers may need to be manually installed using `pip` or `conda`), while JuMP/InfiniteOpt expect the optimizer object to be passed in directly: + +``` +JuMPCollocation(Ipopt.Optimizer, constructRK4()) +CasADiCollocation("ipopt", constructRK4()) +``` + +**JuMP** and **CasADi** collocation require an ODE tableau to be passed in. These can be constructed by calling the `constructX()` functions from DiffEqDevTools. The list of tableaus can be found [here](https://docs.sciml.ai/DiffEqDevDocs/dev/internals/tableaus/). If none is passed in, both solvers will default to using Radau second-order with five collocation points. + +**Pyomo** and **InfiniteOpt** each have their own built-in collocation methods. + + 1. **InfiniteOpt**: The list of InfiniteOpt collocation methods can be found [in the table on this page](https://infiniteopt.github.io/InfiniteOpt.jl/stable/guide/derivative/). If none is passed in, the solver defaults to `FiniteDifference(Backward())`, which is effectively implicit Euler. + 2. **Pyomo**: The list of Pyomo collocation methods can be found [at the bottom of this page](https://github.com/SciML/Pyomo.jl). If none is passed in, the solver defaults to a `LagrangeRadau(3)`. + +Some examples of the latter two collocations: + +```julia +PyomoCollocation("ipopt", LagrangeRadau(2)) +InfiniteOptCollocation(Ipopt.Optimizer, OrthogonalCollocation(3)) +``` + +```@docs; canonical = false +JuMPCollocation +InfiniteOptCollocation +CasADiCollocation +PyomoCollocation +solve(::AbstractDynamicOptProblem) +``` + +### Problem constructors + +```@docs; canonical = false +JuMPDynamicOptProblem +InfiniteOptDynamicOptProblem +CasADiDynamicOptProblem +PyomoDynamicOptProblem +``` diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl index 2301a86b68..01a7fa726b 100644 --- a/ext/MTKPyomoDynamicOptExt.jl +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -8,15 +8,15 @@ using Setfield const MTK = ModelingToolkit SPECIAL_FUNCTIONS_DICT = Dict([acos => Pyomo.py_acos, - acosh => Pyomo.py_acosh, - asin => Pyomo.py_asin, - tan => Pyomo.py_tan, - atanh => Pyomo.py_atanh, - cos => Pyomo.py_cos, - log => Pyomo.py_log, - sin => Pyomo.py_sin, - sqrt => Pyomo.py_sqrt, - exp => Pyomo.py_exp]) + acosh => Pyomo.py_acosh, + asin => Pyomo.py_asin, + tan => Pyomo.py_tan, + atanh => Pyomo.py_atanh, + cos => Pyomo.py_cos, + log => Pyomo.py_log, + sin => Pyomo.py_sin, + sqrt => Pyomo.py_sqrt, + exp => Pyomo.py_exp]) struct PyomoDynamicOptModel model::ConcreteModel @@ -34,7 +34,8 @@ struct PyomoDynamicOptModel @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel) T_SYM DUMMY_SYM model.dU = dae.DerivativeVar(U, wrt = model.t, initialize = 0) #add_time_equation!(model, MODEL_SYM, T_SYM) - new(model, U, V, tₛ, is_free_final, nothing, PyomoVar(model.dU), MODEL_SYM, T_SYM, DUMMY_SYM) + new(model, U, V, tₛ, is_free_final, nothing, + PyomoVar(model.dU), MODEL_SYM, T_SYM, DUMMY_SYM) end end @@ -63,19 +64,26 @@ struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <: end end -pysym_getproperty(s::Union{Num, Symbolics.Symbolic}, name::Symbol) = Symbolics.wrap(SymbolicUtils.term(_getproperty, Symbolics.unwrap(s), Val{name}(), type = Symbolics.Struct{PyomoVar})) -_getproperty(s, name::Val{fieldname}) where fieldname = getproperty(s, fieldname) +function pysym_getproperty(s::Union{Num, Symbolics.Symbolic}, name::Symbol) + Symbolics.wrap(SymbolicUtils.term( + _getproperty, Symbolics.unwrap(s), Val{name}(), type = Symbolics.Struct{PyomoVar})) +end +_getproperty(s, name::Val{fieldname}) where {fieldname} = getproperty(s, fieldname) function MTK.PyomoDynamicOptProblem(sys::System, op, tspan; dt = nothing, steps = nothing, guesses = Dict(), kwargs...) - prob, pmap = MTK.process_DynamicOptProblem(PyomoDynamicOptProblem, PyomoDynamicOptModel, sys, op, tspan; dt, steps, guesses, kwargs...) + prob, + pmap = MTK.process_DynamicOptProblem(PyomoDynamicOptProblem, PyomoDynamicOptModel, + sys, op, tspan; dt, steps, guesses, kwargs...) conc_model = prob.wrapped_model.model MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) prob end -MTK.generate_internal_model(m::Type{PyomoDynamicOptModel}) = ConcreteModel(pyomo.ConcreteModel()) +function MTK.generate_internal_model(m::Type{PyomoDynamicOptModel}) + ConcreteModel(pyomo.ConcreteModel()) +end function MTK.generate_time_variable!(m::ConcreteModel, tspan, tsteps) m.steps = length(tsteps) @@ -83,7 +91,7 @@ function MTK.generate_time_variable!(m::ConcreteModel, tspan, tsteps) m.time = pyomo.Var(m.t) end -function MTK.generate_state_variable!(m::ConcreteModel, u0, ns, ts) +function MTK.generate_state_variable!(m::ConcreteModel, u0, ns, ts) m.u_idxs = pyomo.RangeSet(1, ns) init_f = Pyomo.pyfunc((m, i, t) -> (u0[Pyomo.pyconvert(Int, i)])) m.U = pyomo.Var(m.u_idxs, m.t, initialize = init_f) @@ -113,7 +121,7 @@ function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons; n_idxs = 1) end expr = Symbolics.substitute(expr, SPECIAL_FUNCTIONS_DICT) - cons_sym = Symbol("cons", hash(cons)) + cons_sym = Symbol("cons", hash(cons)) if occursin(Symbolics.unwrap(t_sym), expr) f = eval(Symbolics.build_function(expr, model_sym, t_sym)) setproperty!(model, cons_sym, pyomo.Constraint(model.t, rule = Pyomo.pyfunc(f))) @@ -176,14 +184,21 @@ struct PyomoCollocation <: AbstractCollocation derivative_method::Pyomo.DiscretizationMethod end -MTK.PyomoCollocation(solver, derivative_method = LagrangeRadau(5)) = PyomoCollocation(solver, derivative_method) +function MTK.PyomoCollocation(solver, derivative_method = LagrangeRadau(5)) + PyomoCollocation(solver, derivative_method) +end -function MTK.prepare_and_optimize!(prob::PyomoDynamicOptProblem, collocation; verbose, kwargs...) +function MTK.prepare_and_optimize!( + prob::PyomoDynamicOptProblem, collocation; verbose, kwargs...) solver_m = prob.wrapped_model.model.clone() dm = collocation.derivative_method discretizer = TransformationFactory(dm) + if MTK.is_free_final(prob.wrapped_model) && !Pyomo.is_finite_difference(dm) + error("The Lagrange-Radau and Lagrange-Legendre collocations currently cannot be used for free final problems.") + end ncp = Pyomo.is_finite_difference(dm) ? 1 : dm.np - discretizer.apply_to(solver_m, wrt = solver_m.t, nfe = solver_m.steps - 1, scheme = Pyomo.scheme_string(dm)) + discretizer.apply_to(solver_m, wrt = solver_m.t, nfe = solver_m.steps - 1, + scheme = Pyomo.scheme_string(dm)) solver = SolverFactory(string(collocation.solver)) results = solver.solve(solver_m, tee = true) @@ -208,13 +223,16 @@ function MTK.get_t_values(output::PyomoOutput) Pyomo.pyconvert(Float64, pyomo.value(m.tₛ)) * [Pyomo.pyconvert(Float64, t) for t in m.t] end -MTK.objective_value(output::PyomoOutput) = Pyomo.pyconvert(Float64, pyomo.value(output.model.obj)) +function MTK.objective_value(output::PyomoOutput) + Pyomo.pyconvert(Float64, pyomo.value(output.model.obj)) +end function MTK.successful_solve(output::PyomoOutput) r = output.result ss = r.solver.status tc = r.solver.termination_condition - if Bool(ss == opt.SolverStatus.ok) && (Bool(tc == opt.TerminationCondition.optimal) || Bool(tc == opt.TerminationCondition.locallyOptimal)) + if Bool(ss == opt.SolverStatus.ok) && (Bool(tc == opt.TerminationCondition.optimal) || + Bool(tc == opt.TerminationCondition.locallyOptimal)) return true else return false From d5a49b0668e4170affbd028fe88b9fe456f467c0 Mon Sep 17 00:00:00 2001 From: vyudu Date: Mon, 2 Jun 2025 10:03:12 -0400 Subject: [PATCH 18/22] reemove solver info from tutorial --- docs/src/tutorials/dynamic_optimization.md | 65 +++++++--------------- 1 file changed, 19 insertions(+), 46 deletions(-) diff --git a/docs/src/tutorials/dynamic_optimization.md b/docs/src/tutorials/dynamic_optimization.md index 7753516795..ff1ffab3dc 100644 --- a/docs/src/tutorials/dynamic_optimization.md +++ b/docs/src/tutorials/dynamic_optimization.md @@ -1,4 +1,5 @@ # Solving Dynamic Optimization Problems + Systems in ModelingToolkit.jl can be directly converted to dynamic optimization or optimal control problems. In such systems, one has one or more input variables that are externally controlled to control the dynamics of the system. A dynamic optimization solves for the optimal time trajectory of the input variables in order to maximize or minimize a desired objective function. For example, a car driver might like to know how to step on the accelerator if the goal is to finish a race while using the least gas. To begin, let us take a rocket launch example. The input variable here is the thrust exerted by the engine. The rocket state is described by its current height, mass, and velocity. The mass decreases as the rocket loses fuel while thrusting. @@ -10,8 +11,8 @@ D = ModelingToolkit.D_nounits @parameters h_c m₀ h₀ g₀ D_c c Tₘ m_c @variables begin - h(..) - v(..) + h(..) + v(..) m(..), [bounds = (m_c, 1)] T(..), [input = true, bounds = (0, Tₘ)] end @@ -20,8 +21,8 @@ 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] + 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, 0.2) costs = [-h(te)] @@ -35,24 +36,28 @@ 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] ``` + What we would like to optimize here is the final height of the rocket. We do this by providing a vector of expressions corresponding to the costs. By default, the sense of the optimization is to minimize the provided cost. So to maximize the rocket height at the final time, we write `-h(te)` as the cost. Now we can construct a problem and solve it. Let us use JuMP as our backend here. Note that the package trigger is actually [InfiniteOpt](https://infiniteopt.github.io/InfiniteOpt.jl/stable/), and not JuMP - this package includes JuMP but is designed for optimization on function spaces. Additionally we need to load the solver package - we will use [Ipopt](https://github.com/jump-dev/Ipopt.jl) here (a good choice in general). Here we have also loaded DiffEqDevTools because we will need to construct the ODE tableau. This is only needed if one desires a custom ODE tableau for the collocation - by default the solver will use RadauIIA5. + ```@example dynamic_opt using InfiniteOpt, Ipopt, DiffEqDevTools jprob = JuMPDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIIA5())); ``` + The solution has three fields: `jsol.sol` is the ODE solution for the states, `jsol.input_sol` is the ODE solution for the inputs, and `jsol.model` is the wrapped model that we can use to query things like objective and constraint residuals. Let's plot the final solution and the controller here: + ```@example dynamic_opt using CairoMakie fig = Figure(resolution = (800, 400)) -ax1 = Axis(fig[1,1], title = "Rocket trajectory", xlabel = "Time") -ax2 = Axis(fig[1,2], title = "Control trajectory", xlabel = "Time") +ax1 = Axis(fig[1, 1], title = "Rocket trajectory", xlabel = "Time") +ax2 = Axis(fig[1, 2], title = "Control trajectory", xlabel = "Time") for u in unknowns(rocket) lines!(ax1, jsol.sol.t, jsol.sol[u], label = string(u)) @@ -64,12 +69,14 @@ fig ``` ### Free final time problems + There are additionally a class of dynamic optimization problems where we would like to know how to control our system to achieve something in the least time. Such problems are called free final time problems, since the final time is unknown. To model these problems in ModelingToolkit, we declare the final time as a parameter. Below we have a model system called the double integrator. We control the acceleration of a block in order to reach a desired destination in the least time. + ```@example dynamic_opt @variables begin - x(..) + x(..) v(..) u(..), [bounds = (-1.0, 1.0), input = true] end @@ -88,6 +95,7 @@ u0map = [x(t) => 1.0, v(t) => 0.0] tspan = (0.0, tf) parammap = [u(t) => 0.0, tf => 1.0] ``` + The `tf` mapping in the parameter map is treated as an initial guess. Please note that, at the moment, free final time problems cannot support constraints defined at definite time values, like `x(3) ~ 2`. @@ -97,16 +105,18 @@ Please note that, at the moment, free final time problems cannot support constra The Pyomo collocation methods (LagrangeRadau, LagrangeLegendre) currently are bugged for free final time problems. Strongly suggest using BackwardEuler() for such problems when using Pyomo as the backend. When declaring the problem in this case we need to provide the number of steps, since dt can't be known in advanced. Let's solve plot our final solution and the controller for the block, using InfiniteOpt as the backend: + ```@example dynamic_opt iprob = InfiniteOptDynamicOptProblem(block, [u0map; parammap], tspan; steps = 100) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)); ``` Let's plot the final solution and the controller here: + ```@example dynamic_opt fig = Figure(resolution = (800, 400)) -ax1 = Axis(fig[1,1], title = "Block trajectory", xlabel = "Time") -ax2 = Axis(fig[1,2], title = "Control trajectory", xlabel = "Time") +ax1 = Axis(fig[1, 1], title = "Block trajectory", xlabel = "Time") +ax2 = Axis(fig[1, 2], title = "Control trajectory", xlabel = "Time") for u in unknowns(block) lines!(ax1, isol.sol.t, isol.sol[u], label = string(u)) @@ -116,40 +126,3 @@ axislegend(ax1) axislegend(ax2) fig ``` - -### Solvers -Currently 4 backends are exposed for solving dynamic optimization problems using collocation: JuMP, InfiniteOpt, CasADi, and Pyomo. - -Please note that there are differences in how to construct the collocation solver for the different cases. For example, the Python based ones, CasADi and Pyomo, expect the solver to be passed in as a string (CasADi and Pyomo come pre-loaded with Ipopt, but other solvers may need to be manually installed using `pip` or `conda`), while JuMP/InfiniteOpt expect the optimizer object to be passed in directly: -``` -JuMPCollocation(Ipopt.Optimizer, constructRK4()) -CasADiCollocation("ipopt", constructRK4()) -``` - -**JuMP** and **CasADi** collocation require an ODE tableau to be passed in. These can be constructed by calling the `constructX()` functions from DiffEqDevTools. The list of tableaus can be found [here](https://docs.sciml.ai/DiffEqDevDocs/dev/internals/tableaus/). If none is passed in, both solvers will default to using Radau second-order with five collocation points. - -**Pyomo** and **InfiniteOpt** each have their own built-in collocation methods. -1. **InfiniteOpt**: The list of InfiniteOpt collocation methods can be found [in the table on this page](https://infiniteopt.github.io/InfiniteOpt.jl/stable/guide/derivative/). If none is passed in, the solver defaults to `FiniteDifference(Backward())`, which is effectively implicit Euler. -2. **Pyomo**: The list of Pyomo collocation methods can be found [at the bottom of this page](https://github.com/SciML/Pyomo.jl). If none is passed in, the solver defaults to a `LagrangeRadau(3)`. - -Some examples of the latter two collocations: -```julia -PyomoCollocation("ipopt", LagrangeRadau(2)) -InfiniteOptCollocation(Ipopt.Optimizer, OrthogonalCollocation(3)) -``` - -```@docs; canonical = false -JuMPCollocation -InfiniteOptCollocation -CasADiCollocation -PyomoCollocation -solve(::AbstractDynamicOptProblem) -``` - -### Problem constructors -```@docs; canonical = false -JuMPDynamicOptProblem -InfiniteOptDynamicOptProblem -CasADiDynamicOptProblem -PyomoDynamicOptProblem -``` From 8750986b434cb9f24baa22d3702d6886d0d6081e Mon Sep 17 00:00:00 2001 From: vyudu Date: Mon, 2 Jun 2025 11:04:17 -0400 Subject: [PATCH 19/22] update Project.toml, fix substitution bug --- Project.toml | 3 ++- ext/MTKPyomoDynamicOptExt.jl | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 215c0bcd1b..098757e3b9 100644 --- a/Project.toml +++ b/Project.toml @@ -46,7 +46,6 @@ OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" -Pyomo = "0e8e1daf-01b5-4eba-a626-3897743a3816" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47" @@ -74,6 +73,7 @@ DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6" FMI = "14a09403-18e3-468f-ad8a-74f8dda2d9ac" InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57" LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800" +Pyomo = "0e8e1daf-01b5-4eba-a626-3897743a3816" [extensions] MTKBifurcationKitExt = "BifurcationKit" @@ -142,6 +142,7 @@ OrdinaryDiffEqCore = "1.15.0" OrdinaryDiffEqDefault = "1.2" OrdinaryDiffEqNonlinearSolve = "1.5.0" PrecompileTools = "1" +Pyomo = "0.1.0" REPL = "1" RecursiveArrayTools = "3.26" Reexport = "0.2, 1" diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl index 01a7fa726b..ce667cde6d 100644 --- a/ext/MTKPyomoDynamicOptExt.jl +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -119,7 +119,7 @@ function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons; n_idxs = 1) else cons.lhs - cons.rhs ≤ 0 end - expr = Symbolics.substitute(expr, SPECIAL_FUNCTIONS_DICT) + expr = Symbolics.substitute(Symbolics.unwrap(expr), SPECIAL_FUNCTIONS_DICT) cons_sym = Symbol("cons", hash(cons)) if occursin(Symbolics.unwrap(t_sym), expr) From 97bb8ffacb4000819de1a3aea9745261e2d45ba9 Mon Sep 17 00:00:00 2001 From: vyudu Date: Mon, 2 Jun 2025 13:44:09 -0400 Subject: [PATCH 20/22] add Pyomo to extensions --- test/extensions/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/test/extensions/Project.toml b/test/extensions/Project.toml index 9f43e6f4a4..6a8d01a7b8 100644 --- a/test/extensions/Project.toml +++ b/test/extensions/Project.toml @@ -20,6 +20,7 @@ OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8" OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf" OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a" OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2" +Pyomo = "0e8e1daf-01b5-4eba-a626-3897743a3816" SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1" SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226" SimpleDiffEq = "05bca326-078c-5bf0-a5bf-ce7c7982d7fd" From a8181789e5cac6e2cb34fcbcb7d9400af5a88683 Mon Sep 17 00:00:00 2001 From: vyudu Date: Thu, 5 Jun 2025 10:18:56 -0400 Subject: [PATCH 21/22] fix: fix pyomo problems on lts --- ext/MTKPyomoDynamicOptExt.jl | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl index ce667cde6d..cb0aafc432 100644 --- a/ext/MTKPyomoDynamicOptExt.jl +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -7,7 +7,7 @@ using NaNMath using Setfield const MTK = ModelingToolkit -SPECIAL_FUNCTIONS_DICT = Dict([acos => Pyomo.py_acos, +const SPECIAL_FUNCTIONS_DICT = Dict([acos => Pyomo.py_acos, acosh => Pyomo.py_acosh, asin => Pyomo.py_asin, tan => Pyomo.py_tan, @@ -33,22 +33,11 @@ struct PyomoDynamicOptModel function PyomoDynamicOptModel(model, U, V, tₛ, is_free_final) @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel) T_SYM DUMMY_SYM model.dU = dae.DerivativeVar(U, wrt = model.t, initialize = 0) - #add_time_equation!(model, MODEL_SYM, T_SYM) new(model, U, V, tₛ, is_free_final, nothing, PyomoVar(model.dU), MODEL_SYM, T_SYM, DUMMY_SYM) end end -function add_time_equation!(model::ConcreteModel, model_sym, t_sym) - model.dtime = dae.DerivativeVar(model.time) - - mdt = Symbolics.value(pysym_getproperty(model_sym, :dtime)) - mts = Symbolics.value(pysym_getproperty(model_sym, :tₛ)) - expr = mdt[t_sym] - mts == 0 - f = Pyomo.pyfunc(eval(Symbolics.build_function(expr, model_sym, t_sym))) - model.time_eq = pyomo.Constraint(model.t, rule = f) -end - struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <: AbstractDynamicOptProblem{uType, tType, isinplace} f::F @@ -119,7 +108,7 @@ function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons; n_idxs = 1) else cons.lhs - cons.rhs ≤ 0 end - expr = Symbolics.substitute(Symbolics.unwrap(expr), SPECIAL_FUNCTIONS_DICT) + expr = Symbolics.substitute(Symbolics.unwrap(expr), SPECIAL_FUNCTIONS_DICT, fold = false) cons_sym = Symbol("cons", hash(cons)) if occursin(Symbolics.unwrap(t_sym), expr) @@ -133,7 +122,7 @@ end function MTK.set_objective!(pmodel::PyomoDynamicOptModel, expr) @unpack model, model_sym, t_sym, dummy_sym = pmodel - expr = Symbolics.substitute(expr, SPECIAL_FUNCTIONS_DICT) + expr = Symbolics.substitute(expr, SPECIAL_FUNCTIONS_DICT, fold = false) if occursin(Symbolics.unwrap(t_sym), expr) f = eval(Symbolics.build_function(expr, model_sym, t_sym)) model.obj = pyomo.Objective(model.t, rule = Pyomo.pyfunc(f)) From 639ed6a1c0374687ffd25f2efa0ed5b26d6a21da Mon Sep 17 00:00:00 2001 From: vyudu Date: Thu, 5 Jun 2025 10:48:03 -0400 Subject: [PATCH 22/22] fix: use symbolic indexing for solutions of pprobs --- test/extensions/dynamic_optimization.jl | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index 0243b298d6..3884bc1aa9 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -69,9 +69,8 @@ const M = ModelingToolkit pprob = PyomoDynamicOptProblem( lksys, [u0map; parammap], tspan; guesses = guess, dt = 0.01) psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(3))) - @show psol.sol - @test psol.sol(0.6)[1] ≈ 3.5 - @test psol.sol(0.3)[1] ≈ 7.0 + @test psol.sol(0.6; idxs = x(t)) ≈ 3.5 + @test psol.sol(0.3; idxs = x(t)) ≈ 7.0 iprob = InfiniteOptDynamicOptProblem( lksys, [u0map; parammap], tspan; guesses = guess, dt = 0.01) @@ -161,7 +160,7 @@ end pprob = PyomoDynamicOptProblem(block, [u0map; parammap], tspan; dt = 0.01) psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) @test is_bangbang(psol.input_sol, [-1.0], [1.0]) - @test ≈(psol.sol.u[end][2], 0.25, rtol = 1e-3) + @test ≈(psol.sol[x(t)][end], 0.25, rtol = 1e-3) spline = ctrl_to_spline(isol.input_sol, ConstantInterpolation) oprob = ODEProblem(block_ode, [u0map; u_interp => spline], tspan) @@ -251,7 +250,7 @@ end pprob = PyomoDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001, cse = false) psol = solve(pprob, PyomoCollocation("ipopt", LagrangeRadau(4))) - @test psol.sol.u[end][1] > 1.012 + @test psol.sol[h(t)][end] > 1.012 # Test solution @parameters (T_interp::CubicSpline)(..)