chore: move the docstrings to the interface file

Author: vyudu

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

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

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