|
| 1 | +module MTKPyomoDynamicOptExt |
| 2 | +using ModelingToolkit |
| 3 | +using PythonCall |
| 4 | +using DiffEqBase |
| 5 | +using UnPack |
| 6 | +using NaNMath |
| 7 | +const MTK = ModelingToolkit |
| 8 | + |
| 9 | +# import pyomo |
| 10 | +const pyomo = PythonCall.pynew() |
| 11 | +PythonCall.pycopy!(pyomo, pyimport("pyomo.environ")) |
| 12 | + |
| 13 | +struct PyomoDAEVar |
| 14 | + v::Py |
| 15 | +end |
| 16 | +(v::PyomoDAEVar)(t) = v.v[:, t] |
| 17 | +getindex(v::PyomoDAEVar, i::Union{Num, Symbolic}, t::Union{Num, Symbolic}) = wrap(Term{symeltype(A)}(getindex, [A, unwrap(i), unwrap(t)])) |
| 18 | + |
| 19 | +for ff in [acos, log1p, acosh, log2, asin, tan, atanh, cos, log, sin, log10, sqrt] |
| 20 | + f = nameof(ff) |
| 21 | + @eval NaNMath.$f(x::PyomoDAEVar) = Base.$f(x) |
| 22 | +end |
| 23 | + |
| 24 | +const SymbolicConcreteModel = Symbolics.symstruct(ConcreteModel) |
| 25 | + |
| 26 | +struct PyomoModel |
| 27 | + model::ConcreteModel |
| 28 | + U |
| 29 | + V |
| 30 | + tₛ::Union{Int} |
| 31 | + is_free_final::Bool |
| 32 | + model_sym::SymbolicConcreteModel |
| 33 | + t_sym::Union{Num, BasicSymbolic} |
| 34 | + idx_sym::Union{Num, BasicSymbolic} |
| 35 | + |
| 36 | + function PyomoModel(model, U, V, tₛ, is_free_final) |
| 37 | + @variables MODEL_SYM::SymbolicConcreteModel IDX_SYM::Int T_SYM |
| 38 | + PyomoModel(model, U, V, tₛ, is_free_final, MODEL_SYM, T_SYM, INDEX_SYM) |
| 39 | + end |
| 40 | +end |
| 41 | + |
| 42 | +struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <: |
| 43 | + AbstractDynamicOptProblem{uType, tType, isinplace} |
| 44 | + f::F |
| 45 | + u0::uType |
| 46 | + tspan::tType |
| 47 | + p::P |
| 48 | + wrapped_model::ConcreteModel |
| 49 | + kwargs::K |
| 50 | + |
| 51 | + function PyomoDynamicOptProblem(f, u0, tspan, p, model, kwargs...) |
| 52 | + new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f, 5), |
| 53 | + typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs) |
| 54 | + end |
| 55 | +end |
| 56 | + |
| 57 | +""" |
| 58 | + PyomoDynamicOptProblem(sys::ODESystem, u0, tspan, p; dt, steps) |
| 59 | +
|
| 60 | +Convert an ODESystem representing an optimal control system into a Pyomo model |
| 61 | +for solving using optimization. Must provide either `dt`, the timestep between collocation |
| 62 | +points (which, along with the timespan, determines the number of points), or directly |
| 63 | +provide the number of points as `steps`. |
| 64 | +""" |
| 65 | +function MTK.PyomoDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap; |
| 66 | + dt = nothing, |
| 67 | + steps = nothing, |
| 68 | + guesses = Dict(), kwargs...) |
| 69 | + prob = MTK.process_DynamicOptProblem(PyomoDynamicOptProblem, PyomoModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...) |
| 70 | + MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) |
| 71 | + prob |
| 72 | +end |
| 73 | + |
| 74 | +MTK.generate_internal_model(m::Type{PyomoModel}) = pyomo.ConcreteModel() |
| 75 | +function MTK.generate_time_variable!(m::ConcreteModel, tspan, tsteps) |
| 76 | + m.t = pyomo.ContinuousSet(initialize = collect(tsteps), bounds = tspan) |
| 77 | +end |
| 78 | + |
| 79 | +function MTK.generate_state_variable!(m::ConcreteModel, u0, ns, ts) |
| 80 | + m.u_idxs = pyomo.RangeSet(1, ns) |
| 81 | + pyomo.Var(m.u_idxs, m.t) |
| 82 | +end |
| 83 | + |
| 84 | +function MTK.generate_input_variable!(m::ConcreteModel, u0, nc, ts) |
| 85 | + m.v_idxs = pyomo.RangeSet(1, nc) |
| 86 | + pyomo.Var(m.v_idxs, m.t) |
| 87 | +end |
| 88 | + |
| 89 | +function MTK.generate_timescale(m::ConcreteModel, guess, is_free_t) |
| 90 | + m.tₛ = is_free_t ? pyomo.Var(initialize = guess, bounds = (0, Inf)) : 1 |
| 91 | +end |
| 92 | + |
| 93 | +function MTK.add_constraint!(pmodel::PyomoModel, cons) |
| 94 | + @unpack model, model_sym, idx_sym, t_sym = pmodel |
| 95 | + expr = if cons isa Equation |
| 96 | + cons.lhs - cons.rhs == 0 |
| 97 | + elseif cons.relational_op === Symbolics.geq |
| 98 | + cons.lhs - cons.rhs ≥ 0 |
| 99 | + else |
| 100 | + cons.lhs - cons.rhs ≤ 0 |
| 101 | + end |
| 102 | + constraint_f = Symbolics.build_function(expr, model_sym, idx_sym, t_sym) |
| 103 | + pyomo.Constraint(rule = constraint_f) |
| 104 | +end |
| 105 | + |
| 106 | +function MTK.set_objective!(m::PyomoModel, expr) = pyomo.Objective(expr = expr) |
| 107 | + |
| 108 | +function add_initial_constraints!(model::PyomoModel, u0, u0_idxs) |
| 109 | + for i in u0_idxs |
| 110 | + model.U[i, 0].fix(u0[i]) |
| 111 | + end |
| 112 | +end |
| 113 | + |
| 114 | +function substitute_fixed_t_vars!(model::PyomoModel, sys, exprs) |
| 115 | + stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))]) |
| 116 | + ctidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))]) |
| 117 | + iv = MTK.get_iv(sys) |
| 118 | + |
| 119 | + for cons in jconstraints |
| 120 | + consvars = MTK.vars(cons) |
| 121 | + for st in consvars |
| 122 | + MTK.iscall(st) || continue |
| 123 | + x = MTK.operation(st) |
| 124 | + t = only(MTK.arguments(st)) |
| 125 | + MTK.symbolic_type(t) === MTK.NotSymbolic() || continue |
| 126 | + if haskey(stidxmap, x(iv)) |
| 127 | + idx = stidxmap[x(iv)] |
| 128 | + cv = :U |
| 129 | + else |
| 130 | + idx = ctidxmap[x(iv)] |
| 131 | + cv = :V |
| 132 | + end |
| 133 | + model.t.add(t) |
| 134 | + cons = Symbolics.substitute(cons, Dict(x(t) => model.cv[idx, t])) |
| 135 | + end |
| 136 | + end |
| 137 | +end |
| 138 | + |
| 139 | +function MTK.substitute_model_vars(pmodel::PyomoModel, sys, pmap, exprs, tspan) |
| 140 | + @unwrap model, model_sym, idx_sym, t_sym = pmodel |
| 141 | + x_ops = [MTK.operation(MTK.unwrap(st)) for st in unknowns(sys)] |
| 142 | + c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in MTK.unbound_inputs(sys)] |
| 143 | + mU = Symbolics.symbolic_getproperty(model_sym, :U) |
| 144 | + mV = Symbolics.symbolic_getproperty(model_sym, :V) |
| 145 | + |
| 146 | + (ti, tf) = tspan |
| 147 | + if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() |
| 148 | + _tf = model.tₛ + ti |
| 149 | + exprs = map(c -> Symbolics.fast_substitute(c, Dict(tf => _tf)), exprs) |
| 150 | + free_t_map = Dict([[x(tₛ) => mU[i, end] for (i, x) in enumerate(x_ops)]; |
| 151 | + [c(tₛ) => mV[i, end] for (i, c) in enumerate(c_ops)]]) |
| 152 | + exprs = map(c -> Symbolics.fixpoint_sub(c, free_t_map), exprs) |
| 153 | + end |
| 154 | + |
| 155 | + whole_interval_map = Dict([[v => mU[i, t_sym] for (i, v) in enumerate(sts)]; |
| 156 | + [v => mV[i, t_sym] for (i, v) in enumerate(cts)]]) |
| 157 | + exprs = map(c -> Symbolics.fixpoint_sub(c, whole_interval_map), exprs) |
| 158 | + exprs |
| 159 | +end |
| 160 | + |
| 161 | +function MTK.substitute_integral!(model::PyomoModel, exprs, tspan) |
| 162 | + intmap = Dict() |
| 163 | + for int in MTK.collect_applied_operators(exprs, Symbolics.Integral) |
| 164 | + op = MTK.operation(int) |
| 165 | + arg = only(arguments(MTK.value(int))) |
| 166 | + lo, hi = MTK.value.((op.domain.domain.left, op.domain.domain.right)) |
| 167 | + if MTK.is_free_final(model) && isequal((lo, hi), tspan) |
| 168 | + (lo, hi) = (0, 1) |
| 169 | + elseif MTK.is_free_final(model) |
| 170 | + error("Free final time problems cannot handle partial timespans.") |
| 171 | + end |
| 172 | + intmap[int] = model.tₛ * InfiniteOpt.∫(arg, model.model[:t], lo, hi) |
| 173 | + end |
| 174 | + exprs = map(c -> Symbolics.substitute(c, intmap), exprs) |
| 175 | +end |
| 176 | + |
| 177 | +function MTK.substitute_differentials(model::PyomoModel, sys, eqs) |
| 178 | + pmodel = prob.model |
| 179 | + @unpack model, model_sym, t_sym, idx_sym = pmodel |
| 180 | + model.dU = pyomo.DerivativeVar(model.U, wrt = model.t) |
| 181 | + |
| 182 | + mdU = Symbolics.symbolic_getproperty(model_sym, :dU) |
| 183 | + mU = Symbolics.symbolic_getproperty(model_sym, :U) |
| 184 | + mtₛ = Symbolics.symbolic_getproperty(model_sym, :tₛ) |
| 185 | + diffsubmap = Dict([D(mU[i, t_sym]) => mdU[i, t_sym] for i in 1:length(unknowns(sys))]) |
| 186 | + |
| 187 | + diff_eqs = substitute_model_vars(model, sys, pmap, diff_equations(sys)) |
| 188 | + diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs) |
| 189 | + [mtₛ * eq.rhs - eq.lhs == 0 for eq in diff_eqs] |
| 190 | +end |
| 191 | + |
| 192 | +struct PyomoCollocation <: AbstractCollocation |
| 193 | + solver::Any |
| 194 | + derivative_method |
| 195 | +end |
| 196 | +MTK.PyomoCollocation(solver, derivative_method = 1) = PyomoCollocation(solver, derivative_method) |
| 197 | + |
| 198 | +function MTK.prepare_and_optimize!() |
| 199 | +end |
| 200 | +function MTK.get_U_values() |
| 201 | +end |
| 202 | +function MTK.get_V_values() |
| 203 | +end |
| 204 | +function MTK.get_t_values() |
| 205 | +end |
| 206 | +function MTK.successful_solve() |
| 207 | +end |
| 208 | + |
| 209 | +end |
0 commit comments