refactor: use lowered_var instead of manual substitution

Author: vyudu

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"

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

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)

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,99 +42,79 @@ 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
         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)
+    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)

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