fix: fix constraint and cost transcription for pyomo

Author: vyudu

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)

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

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
 

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