refactor: format

Author: AayushSabharwal

ext/MTKCasADiDynamicOptExt.jl

@@ -17,7 +17,9 @@ struct MXLinearInterpolation
     t::Vector{Float64}
     dt::Float64
 end
-Base.getindex(m::MXLinearInterpolation, i...) = length(i) == length(size(m.u)) ? m.u[i...] : m.u[i..., :]
+function Base.getindex(m::MXLinearInterpolation, i...)
+    length(i) == length(size(m.u)) ? m.u[i...] : m.u[i..., :]
+end
 
 mutable struct CasADiModel
     model::Opti
@@ -55,7 +57,7 @@ function (M::MXLinearInterpolation)(τ)
     (i > length(M.t) || i < 1) && error("Cannot extrapolate past the tspan.")
     colons = ntuple(_ -> (:), length(size(M.u)) - 1)
     if i < length(M.t)
-        M.u[colons..., i] + Δ*(M.u[colons..., i+1] - M.u[colons..., i])
+        M.u[colons..., i] + Δ * (M.u[colons..., i + 1] - M.u[colons..., i])
     else
         M.u[colons..., i]
     end
@@ -65,7 +67,8 @@ function MTK.CasADiDynamicOptProblem(sys::System, op, tspan;
         dt = nothing,
         steps = nothing,
         guesses = Dict(), kwargs...)
-    prob, _ = MTK.process_DynamicOptProblem(CasADiDynamicOptProblem, CasADiModel, sys, op, tspan; dt, steps, guesses, kwargs...)
+    prob, _ = MTK.process_DynamicOptProblem(
+        CasADiDynamicOptProblem, CasADiModel, sys, op, tspan; dt, steps, guesses, kwargs...)
     prob
 end
 
@@ -127,10 +130,10 @@ function MTK.lowered_integral(model::CasADiModel, expr, lo, hi)
     for (i, t) in enumerate(model.U.t)
         if lo < t < hi
             Δt = min(dt, t - lo)
-            total += (0.5*Δt*(expr[i] + expr[i-1]))
+            total += (0.5 * Δt * (expr[i] + expr[i - 1]))
         elseif t >= hi && (t - dt < hi)
             Δt = hi - t + dt
-            total += (0.5*Δt*(expr[i] + expr[i-1]))
+            total += (0.5 * Δt * (expr[i] + expr[i - 1]))
         end
     end
     model.tₛ * total
@@ -186,9 +189,13 @@ struct CasADiCollocation <: AbstractCollocation
     tableau::DiffEqBase.ODERKTableau
 end
 
-MTK.CasADiCollocation(solver, tableau = MTK.constructDefault()) = CasADiCollocation(solver, tableau)
+function MTK.CasADiCollocation(solver, tableau = MTK.constructDefault())
+    CasADiCollocation(solver, tableau)
+end
 
-function MTK.prepare_and_optimize!(prob::CasADiDynamicOptProblem, solver::CasADiCollocation; verbose = false, solver_options = Dict(), plugin_options = Dict(), kwargs...)
+function MTK.prepare_and_optimize!(
+        prob::CasADiDynamicOptProblem, solver::CasADiCollocation; verbose = false,
+        solver_options = Dict(), plugin_options = Dict(), kwargs...)
     solver_opti = add_solve_constraints!(prob, solver.tableau)
     verbose || (solver_options["print_level"] = 0)
     solver!(solver_opti, "$(solver.solver)", plugin_options, solver_options)
@@ -211,7 +218,7 @@ end
 function MTK.get_V_values(model::CasADiModel)
     value_getter = MTK.successful_solve(model) ? CasADi.debug_value : CasADi.value
     (nu, nt) = size(model.V.u)
-    if nu*nt != 0
+    if nu * nt != 0
         V_vals = value_getter(model.solver_opti, model.V.u)
         size(V_vals, 2) == 1 && (V_vals = V_vals')
         V_vals = [[V_vals[i, j] for i in 1:nu] for j in 1:nt]
@@ -224,9 +231,11 @@ function MTK.get_t_values(model::CasADiModel)
     value_getter = MTK.successful_solve(model) ? CasADi.debug_value : CasADi.value
     ts = value_getter(model.solver_opti, model.tₛ) .* model.U.t
 end
-MTK.objective_value(model::CasADiModel) = CasADi.pyconvert(Float64, model.solver_opti.py.value(model.solver_opti.py.f))
+function MTK.objective_value(model::CasADiModel)
+    CasADi.pyconvert(Float64, model.solver_opti.py.value(model.solver_opti.py.f))
+end
 
-function MTK.successful_solve(m::CasADiModel) 
+function MTK.successful_solve(m::CasADiModel)
     isnothing(m.solver_opti) && return false
     retcode = CasADi.return_status(m.solver_opti)
     retcode == "Solve_Succeeded" || retcode == "Solved_To_Acceptable_Level"

ext/MTKInfiniteOptExt.jl

@@ -48,26 +48,32 @@ struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
 end
 
 MTK.generate_internal_model(m::Type{InfiniteOptModel}) = InfiniteModel()
-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])
+function MTK.generate_time_variable!(m::InfiniteModel, tspan, tsteps)
+    @infinite_parameter(m, t in [tspan[1], tspan[2]], num_supports=length(tsteps))
+end
+function MTK.generate_state_variable!(m::InfiniteModel, u0::Vector, ns, ts)
+    @variable(m, U[i = 1:ns], Infinite(m[:t]), start=u0[i])
+end
+function MTK.generate_input_variable!(m::InfiniteModel, c0, nc, ts)
+    @variable(m, V[i = 1:nc], Infinite(m[:t]), start=c0[i])
+end
 
 function MTK.generate_timescale!(m::InfiniteModel, guess, is_free_t)
-    @variable(m, tₛ ≥ 0, start = guess)
+    @variable(m, tₛ≥0, start=guess)
     if !is_free_t
-        fix(tₛ, 1, force=true)
+        fix(tₛ, 1, force = true)
         set_start_value(tₛ, 1)
     end
     tₛ
 end
 
-function MTK.add_constraint!(m::InfiniteOptModel, expr::Union{Equation, Inequality}) 
+function MTK.add_constraint!(m::InfiniteOptModel, expr::Union{Equation, Inequality})
     if expr isa Equation
-        @constraint(m.model, expr.lhs - expr.rhs == 0)
+        @constraint(m.model, expr.lhs - expr.rhs==0)
     elseif expr.relational_op === Symbolics.geq
-        @constraint(m.model, expr.lhs - expr.rhs ≥ 0)
+        @constraint(m.model, expr.lhs - expr.rhs≥0)
     else
-        @constraint(m.model, expr.lhs - expr.rhs ≤ 0)
+        @constraint(m.model, expr.lhs - expr.rhs≤0)
     end
 end
 MTK.set_objective!(m::InfiniteOptModel, expr) = @objective(m.model, Min, expr)
@@ -76,20 +82,25 @@ function MTK.JuMPDynamicOptProblem(sys::System, op, tspan;
         dt = nothing,
         steps = nothing,
         guesses = Dict(), kwargs...)
-    prob, _ = MTK.process_DynamicOptProblem(JuMPDynamicOptProblem, InfiniteOptModel, sys, op, tspan; dt, steps, guesses, kwargs...)
+    prob, _ = MTK.process_DynamicOptProblem(JuMPDynamicOptProblem, InfiniteOptModel, sys,
+        op, tspan; dt, steps, guesses, kwargs...)
     prob
 end
 
 function MTK.InfiniteOptDynamicOptProblem(sys::System, op, tspan;
         dt = nothing,
         steps = nothing,
         guesses = Dict(), kwargs...)
-    prob, pmap = MTK.process_DynamicOptProblem(InfiniteOptDynamicOptProblem, InfiniteOptModel, sys, op, tspan; dt, steps, guesses, kwargs...)
+    prob, pmap = MTK.process_DynamicOptProblem(
+        InfiniteOptDynamicOptProblem, InfiniteOptModel,
+        sys, op, tspan; dt, steps, guesses, kwargs...)
     MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan)
     prob
 end
 
-MTK.lowered_integral(model::InfiniteOptModel, expr, lo, hi) = model.tₛ * InfiniteOpt.∫(expr, model.model[:t], lo, hi)
+function MTK.lowered_integral(model::InfiniteOptModel, expr, lo, hi)
+    model.tₛ * InfiniteOpt.∫(expr, model.model[:t], lo, hi)
+end
 MTK.lowered_derivative(model::InfiniteOptModel, i) = ∂(model.U[i], model.model[:t])
 
 function MTK.process_integral_bounds(model::InfiniteOptModel, integral_span, tspan)
@@ -125,7 +136,7 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
     nᵥ = length(V)
     if MTK.is_explicit(tableau)
         K = Any[]
-        for τ in tsteps[1:end-1]
+        for τ in tsteps[1:(end - 1)]
             for (i, h) in enumerate(c)
                 ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = zeros(nᵤ))
                 Uₙ = [U[i](τ) + ΔU[i] * dt for i in 1:nᵤ]
@@ -142,14 +153,15 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
         K = @variable(model, K[1:length(α), 1:nᵤ], Infinite(model[:t]))
         ΔUs = A * K
         ΔU_tot = dt * (K' * α)
-        for τ in tsteps[1:end-1]
+        for τ in tsteps[1:(end - 1)]
             for (i, h) in enumerate(c)
                 ΔU = @view ΔUs[i, :]
                 Uₙ = U + ΔU * dt
                 @constraint(model, [j = 1:nᵤ], K[i, j]==(tₛ * f(Uₙ, V, p, τ + h * dt)[j]),
                     DomainRestrictions(t => τ), base_name="solve_K$i($τ)")
             end
-            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tsteps[end])),
+            @constraint(model,
+                [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tsteps[end])),
                 DomainRestrictions(t => τ), base_name="solve_U($τ)")
         end
     end
@@ -159,15 +171,21 @@ struct JuMPCollocation <: AbstractCollocation
     solver::Any
     tableau::DiffEqBase.ODERKTableau
 end
-MTK.JuMPCollocation(solver, tableau = MTK.constructDefault()) = JuMPCollocation(solver, tableau)
+function MTK.JuMPCollocation(solver, tableau = MTK.constructDefault())
+    JuMPCollocation(solver, tableau)
+end
 
 struct InfiniteOptCollocation <: AbstractCollocation
     solver::Any
     derivative_method::InfiniteOpt.AbstractDerivativeMethod
 end
-MTK.InfiniteOptCollocation(solver, derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward())) = InfiniteOptCollocation(solver, derivative_method)
+function MTK.InfiniteOptCollocation(
+        solver, derivative_method = InfiniteOpt.FiniteDifference(InfiniteOpt.Backward()))
+    InfiniteOptCollocation(solver, derivative_method)
+end
 
-function MTK.prepare_and_optimize!(prob::JuMPDynamicOptProblem, solver::JuMPCollocation; verbose = false, kwargs...)
+function MTK.prepare_and_optimize!(
+        prob::JuMPDynamicOptProblem, solver::JuMPCollocation; verbose = false, kwargs...)
     model = prob.wrapped_model.model
     verbose || set_silent(model)
     # Unregister current solver constraints
@@ -190,7 +208,8 @@ function MTK.prepare_and_optimize!(prob::JuMPDynamicOptProblem, solver::JuMPColl
     model
 end
 
-function MTK.prepare_and_optimize!(prob::InfiniteOptDynamicOptProblem, solver::InfiniteOptCollocation; verbose = false, kwargs...)
+function MTK.prepare_and_optimize!(prob::InfiniteOptDynamicOptProblem,
+        solver::InfiniteOptCollocation; verbose = false, kwargs...)
     model = prob.wrapped_model.model
     verbose || set_silent(model)
     set_derivative_method(model[:t], solver.derivative_method)
@@ -223,8 +242,8 @@ function MTK.successful_solve(model::InfiniteModel)
         error("Model not solvable; please report this to github.com/SciML/ModelingToolkit.jl with a MWE.")
 
     pstatus === FEASIBLE_POINT &&
-         (tstatus === OPTIMAL || tstatus === LOCALLY_SOLVED || tstatus === ALMOST_OPTIMAL ||
-          tstatus === ALMOST_LOCALLY_SOLVED)
+        (tstatus === OPTIMAL || tstatus === LOCALLY_SOLVED || tstatus === ALMOST_OPTIMAL ||
+         tstatus === ALMOST_LOCALLY_SOLVED)
 end
 
 import InfiniteOpt: JuMP, GeneralVariableRef

ext/MTKPyomoDynamicOptExt.jl

@@ -108,7 +108,8 @@ function MTK.add_constraint!(pmodel::PyomoDynamicOptModel, cons; n_idxs = 1)
     else
         cons.lhs - cons.rhs ≤ 0
     end
-    expr = Symbolics.substitute(Symbolics.unwrap(expr), SPECIAL_FUNCTIONS_DICT, fold = false)
+    expr = Symbolics.substitute(
+        Symbolics.unwrap(expr), SPECIAL_FUNCTIONS_DICT, fold = false)
 
     cons_sym = Symbol("cons", hash(cons))
     if occursin(Symbolics.unwrap(t_sym), expr)
@@ -141,17 +142,17 @@ end
 function MTK.lowered_integral(m::PyomoDynamicOptModel, arg, lo, hi)
     @unpack model, model_sym, t_sym, dummy_sym = m
     total = 0
-    dt = Pyomo.pyconvert(Float64, (model.t.at(-1) - model.t.at(1))/(model.steps - 1))
+    dt = Pyomo.pyconvert(Float64, (model.t.at(-1) - model.t.at(1)) / (model.steps - 1))
     f = Symbolics.build_function(arg, model_sym, t_sym, expression = false)
     for (i, t) in enumerate(model.t)
         if Bool(lo < t) && Bool(t < hi)
-            t_p = model.t.at(i-1)
+            t_p = model.t.at(i - 1)
             Δt = min(t - lo, t - t_p)
-            total += 0.5*Δt*(f(model, t) + f(model, t_p))
+            total += 0.5 * Δt * (f(model, t) + f(model, t_p))
         elseif Bool(t >= hi) && Bool(t - dt < hi)
-            t_p = model.t.at(i-1)
+            t_p = model.t.at(i - 1)
             Δt = hi - t + dt
-            total += 0.5*Δt*(f(model, t) + f(model, t_p))
+            total += 0.5 * Δt * (f(model, t) + f(model, t_p))
         end
     end
     PyomoVar(model.tₛ * total)

src/systems/optimal_control_interface.jl

@@ -93,7 +93,7 @@ function PyomoCollocation end
 
 function warn_overdetermined(sys, op)
     cstrs = constraints(sys)
-    init_conds = filter(x -> value(x) ∈ Set(unknowns(sys)), [k for (k,v) in op])
+    init_conds = filter(x -> value(x) ∈ Set(unknowns(sys)), [k for (k, v) in op])
     if !isempty(cstrs)
         (length(cstrs) + length(init_conds) > length(unknowns(sys))) &&
             @warn "The control problem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by op) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
@@ -227,19 +227,19 @@ end
 ##########################
 ### MODEL CONSTRUCTION ###
 ##########################
-function process_DynamicOptProblem(prob_type::Type{<:AbstractDynamicOptProblem}, model_type, sys::System, op, tspan;
-    dt = nothing,
-    steps = nothing,
-    guesses = Dict(), kwargs...)
-
+function process_DynamicOptProblem(
+        prob_type::Type{<:AbstractDynamicOptProblem}, model_type, sys::System, op, tspan;
+        dt = nothing,
+        steps = nothing,
+        guesses = Dict(), kwargs...)
     warn_overdetermined(sys, op)
     ctrls = unbound_inputs(sys)
     states = unknowns(sys)
 
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
     op = Dict([default_toterm(value(k)) => v for (k, v) in op])
     u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) :
-        [stidxmap[default_toterm(k)] for (k, v) in op if haskey(stidxmap, k)]
+              [stidxmap[default_toterm(k)] for (k, v) in op if haskey(stidxmap, k)]
 
     _op = has_alg_eqs(sys) ? op : merge(Dict(op), Dict(guesses))
     f, u0, p = process_SciMLProblem(ODEInputFunction, sys, _op;
@@ -302,7 +302,7 @@ function set_variable_bounds!(m, sys, pmap, tf)
     end
 end
 
-is_free_final(model) = model.is_free_final 
+is_free_final(model) = model.is_free_final
 
 function add_cost_function!(model, sys, tspan, pmap)
     jcosts = cost(sys)
@@ -335,14 +335,15 @@ function substitute_integral(model, expr, tspan)
     Symbolics.substitute(expr, intmap)
 end
 
-function process_integral_bounds(model, integral_span, tspan) 
+function process_integral_bounds(model, integral_span, tspan)
     if is_free_final(model) && isequal(integral_span, tspan)
         integral_span = (0, 1)
     elseif is_free_final(model)
         error("Free final time problems cannot handle partial timespans.")
     else
         (lo, hi) = integral_span
-        (lo < tspan[1] || hi > tspan[2]) && error("Integral bounds are beyond the timespan.")
+        (lo < tspan[1] || hi > tspan[2]) &&
+            error("Integral bounds are beyond the timespan.")
         integral_span
     end
 end
@@ -353,12 +354,14 @@ function substitute_model_vars(model, sys, exprs, tspan)
     c_ops = [operation(unwrap(ct)) for ct in unbound_inputs(sys)]
     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, whole_t_map(model, t, 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, free_t_map(model, tf, x_ops, c_ops)), exprs)
         exprs = map(c -> Symbolics.fast_substitute(c, Dict(tf => _tf)), exprs)
     end
     exprs = map(c -> Symbolics.fast_substitute(c, fixed_t_map(model, x_ops, c_ops)), exprs)
@@ -392,7 +395,8 @@ function fixed_t_map end
 function add_user_constraints!(model, sys, tspan, pmap)
     jconstraints = get_constraints(sys)
     (isnothing(jconstraints) || isempty(jconstraints)) && return nothing
-    cons_dvs, cons_ps = process_constraint_system(jconstraints, Set(unknowns(sys)), parameters(sys), get_iv(sys); validate = false)
+    cons_dvs, cons_ps = process_constraint_system(
+        jconstraints, Set(unknowns(sys)), parameters(sys), get_iv(sys); validate = false)
 
     is_free_final(model) && check_constraint_vars(cons_dvs)
 
@@ -421,12 +425,13 @@ function add_equational_constraints!(model, sys, pmap, tspan)
 end
 
 function set_objective! end
-objective_value(sol::DynamicOptSolution) = objective_value(sol.model) 
+objective_value(sol::DynamicOptSolution) = objective_value(sol.model)
 
 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))])
+    diffsubmap = Dict([D(lowered_var(model, :U, i, t)) => lowered_derivative(model, i)
+                       for i in 1:length(unknowns(sys))])
     eqs = map(c -> Symbolics.substitute(c, diffsubmap), eqs)
 end
 
@@ -466,9 +471,10 @@ 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...)
+function DiffEqBase.solve(prob::AbstractDynamicOptProblem,
+        solver::AbstractCollocation; verbose = false, kwargs...)
     solved_model = prepare_and_optimize!(prob, solver; verbose, kwargs...)
-    
+
     ts = get_t_values(solved_model)
     Us = get_U_values(solved_model)
     Vs = get_V_values(solved_model)