feat: PyomoDynamicOtpPRoblem

Author: vyudu

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) &&

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,24 +74,22 @@ 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
         cons.lhs - cons.rhs ≥ 0
     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)

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

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