refactor: add interface functions for CasADi

Author: vyudu

ext/MTKCasADiDynamicOptExt.jl

@@ -6,10 +6,8 @@ using UnPack
 using NaNMath
 const MTK = ModelingToolkit
 
-# NaNMath
 for ff in [acos, log1p, acosh, log2, asin, tan, atanh, cos, log, sin, log10, sqrt]
     f = nameof(ff)
-    # These need to be defined so that JuMP can trace through functions built by Symbolics
     @eval NaNMath.$f(x::CasadiSymbolicObject) = Base.$f(x)
 end
 
@@ -76,75 +74,47 @@ function MTK.CasADiDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
         dt = nothing,
         steps = nothing,
         guesses = Dict(), kwargs...)
-    MTK.warn_overdetermined(sys, u0map)
-    _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
-    f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
-        t = tspan !== nothing ? tspan[1] : tspan, output_type = MX, kwargs...)
-
-    pmap = Dict{Any, Any}(pmap)
-    steps, is_free_t = MTK.process_tspan(tspan, dt, steps)
-    model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
-
-    CasADiDynamicOptProblem(f, u0, tspan, p, model, kwargs...)
+    process_DynamicOptProblem(CasADiDynamicOptProblem, CasADiModel, sys, u0map, tspan, pmap; dt, steps, guesses, kwargs...)
 end
 
 MTK.generate_internal_model(::Type{CasADiModel}) = CasADi.opti()
-MTK.generate_state_variable(model, u0, ns, nt)
-MTK.generate_input_variable(model, c0, nc, nt) = 1
-MTK.generate_timescale(model, dims) = 1
 
-function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
-    ctrls = MTK.unbound_inputs(sys)
-    states = unknowns(sys)
-    opti = CasADi.Opti()
+function MTK.generate_state_variable(model::Opti, u0, ns, nt, tsteps)
+    U = CasADi.variable!(model, ns, nt)
+    set_initial!(opti, U, DM(repeat(u0, 1, steps)))
+    MXLinearInterpolation(U, tsteps, tsteps[2] - tsteps[1])
+end
+
+function MTK.generate_input_variable(model::Opti, c0, nc, nt, tsteps)
+    V = CasADi.variable!(model, nc, nt)
+    !isempty(c0) && set_initial!(opti, V, DM(repeat(c0, 1, steps)))
+    MXLinearInterpolation(V, tsteps, tsteps[2] - tsteps[1])
+end
 
+function MTK.generate_timescale(model::Opti, guess, is_free_t)
     if is_free_t
-        (ts_sym, te_sym) = tspan
-        MTK.symbolic_type(ts_sym) !== MTK.NotSymbolic() &&
-            error("Free initial time problems are not currently supported in CasADiDynamicOptProblem.")
-        tₛ = variable!(opti)
-        set_initial!(opti, tₛ, pmap[te_sym])
-        subject_to!(opti, tₛ >= ts_sym)
-        hasbounds(te_sym) && begin
-            lo, hi = getbounds(te_sym)
-            subject_to!(opti, tₛ >= lo)
-            subject_to!(opti, tₛ >= hi)
-        end
-        pmap[te_sym] = tₛ
-        tsteps = LinRange(0, 1, steps)
+        tₛ = variable!(model)
+        set_initial!(model, tₛ, guess)
+        subject_to!(model, tₛ >= 0)
+        tₛ
     else
-        tₛ = MX(1)
-        tsteps = LinRange(tspan[1], tspan[2], steps)
+        MX(1)
     end
+end
 
-    U = CasADi.variable!(opti, length(states), steps)
-    V = CasADi.variable!(opti, length(ctrls), steps)
-    set_initial!(opti, U, DM(repeat(u0, 1, steps)))
-    c0 = MTK.value.([pmap[c] for c in ctrls])
-    !isempty(c0) && set_initial!(opti, V, DM(repeat(c0, 1, steps)))
-
-    U_interp = MXLinearInterpolation(U, tsteps, tsteps[2] - tsteps[1])
-    V_interp = MXLinearInterpolation(V, tsteps, tsteps[2] - tsteps[1])
-    for (i, ct) in enumerate(ctrls)
-        pmap[ct] = V[i, :]
+function MTK.add_constraint!(model::CasADiModel, expr)
+    @unpack opti = model
+    if cons isa Equation
+        subject_to!(opti, expr.lhs - expr.rhs == 0)
+    elseif cons.relational_op === Symbolics.geq
+        subject_to!(opti, expr.lhs - expr.rhs ≥ 0)
+    else
+        subject_to!(opti, expr.lhs - expr.rhs ≤ 0)
     end
-
-    model = CasADiModel(opti, U_interp, V_interp, tₛ)
-
-    set_casadi_bounds!(model, sys, pmap)
-    add_cost_function!(model, sys, tspan, pmap; is_free_t)
-    add_user_constraints!(model, sys, tspan, pmap; is_free_t)
-
-    stidxmap = Dict([v => i for (i, v) in enumerate(states)])
-    u0map = Dict([MTK.default_toterm(MTK.value(k)) => v for (k, v) in u0map])
-    u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) :
-              [stidxmap[MTK.default_toterm(k)] for (k, v) in u0map]
-    add_initial_constraints!(model, u0, u0_idxs)
-
-    model
 end
+MTK.set_objective!(model::CasADiModel, expr) = minimize!(model.opti, MX(expr))
 
-function set_casadi_bounds!(model, sys, pmap)
+function MTK.set_variable_bounds!(model, sys, pmap, tf)
     @unpack opti, U, V = model
     for (i, u) in enumerate(unknowns(sys))
         if MTK.hasbounds(u)
@@ -160,75 +130,53 @@ function set_casadi_bounds!(model, sys, pmap)
             subject_to!(opti, V.u[i, :] <= Symbolics.fixpoint_sub(hi, pmap))
         end
     end
+    if MTK.symbolic_type(tf) === MTK.ScalarSymbolic() && hasbounds(tf)
+        lo, hi = MTK.getbounds(tf)
+        subject_to!(opti, model.tₛ >= lo)
+        subject_to!(opti, model.tₛ <= hi)
+    end
 end
 
-function add_initial_constraints!(model::CasADiModel, u0, u0_idxs)
+function MTK.add_initial_constraints!(model::CasADiModel, u0, u0_idxs)
     @unpack opti, U = model
     for i in u0_idxs
         subject_to!(opti, U.u[i, 1] == u0[i])
     end
 end
 
-function add_user_constraints!(model::CasADiModel, sys, tspan, pmap; is_free_t)
+function MTK.substitute_model_vars(
+        model::CasADiModel, sys, pmap, exprs; auxmap::Dict = Dict(), is_free_t)
     @unpack opti, U, V, tₛ = model
-
     iv = MTK.get_iv(sys)
-    conssys = MTK.get_constraintsystem(sys)
-    jconstraints = isnothing(conssys) ? nothing : MTK.get_constraints(conssys)
-    (isnothing(jconstraints) || isempty(jconstraints)) && return nothing
-
-    stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
-    ctidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))])
-    cons_unknowns = map(MTK.default_toterm, unknowns(conssys))
-
-    auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in unknowns(conssys)])
-    jconstraints = substitute_casadi_vars(model, sys, pmap, jconstraints; is_free_t, auxmap)
-    # Manually substitute fixed-t variables
-    for (i, cons) in enumerate(jconstraints)
-        consvars = MTK.vars(cons)
-        for st in consvars
-            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
-            cons = Symbolics.substitute(cons, Dict(x(t) => cv(t)[idx]))
-        end
+    sts = unknowns(sys)
+    cts = MTK.unbound_inputs(sys)
 
-        if cons isa Equation
-            subject_to!(opti, cons.lhs - cons.rhs == 0)
-        elseif cons.relational_op === Symbolics.geq
-            subject_to!(opti, cons.lhs - cons.rhs ≥ 0)
-        else
-            subject_to!(opti, cons.lhs - cons.rhs ≤ 0)
-        end
-    end
-end
+    x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
+    c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
 
-function add_cost_function!(model::CasADiModel, sys, tspan, pmap; is_free_t)
-    @unpack opti, U, V, tₛ = model
-    jcosts = copy(MTK.get_costs(sys))
-    consolidate = MTK.get_consolidate(sys)
-    if isnothing(jcosts) || isempty(jcosts)
-        minimize!(opti, MX(0))
-        return
+    exprs = map(c -> Symbolics.fast_substitute(c, auxmap), exprs)
+    exprs = map(c -> Symbolics.fast_substitute(c, Dict(pmap)), exprs)
+    # tf means different things in different contexts; a [tf] in a cost function
+    # should be tₛ, while a x(tf) should translate to x[1]
+    if is_free_t
+        free_t_map = Dict([[x(tₛ) => U.u[i, end] for (i, x) in enumerate(x_ops)];
+                           [c(tₛ) => V.u[i, end] for (i, c) in enumerate(c_ops)]])
+        exprs = map(c -> Symbolics.fast_substitute(c, free_t_map), exprs)
     end
 
-    iv = MTK.get_iv(sys)
-    stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
-    ctidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))])
-
-    jcosts = substitute_casadi_vars(model, sys, pmap, jcosts; is_free_t)
-    # Substitute fixed-time variables.
-    for i in 1:length(jcosts)
-        costvars = MTK.vars(jcosts[i])
-        for st in costvars
+    exprs = substitute_fixed_t_vars(exprs)
+
+    # for variables like x(t)
+    whole_interval_map = Dict([[v => U.u[i, :] for (i, v) in enumerate(sts)];
+                               [v => V.u[i, :] for (i, v) in enumerate(cts)]])
+    exprs = map(c -> Symbolics.fast_substitute(c, whole_interval_map), exprs)
+    exprs
+end
+
+function substitute_fixed_t_vars(exprs)
+    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))
@@ -240,13 +188,18 @@ function add_cost_function!(model::CasADiModel, sys, tspan, pmap; is_free_t)
                 idx = ctidxmap[x(iv)]
                 cv = V
             end
-            jcosts[i] = Symbolics.substitute(jcosts[i], Dict(x(t) => cv(t)[idx]))
+            exprs[i] = Symbolics.fast_substitute(exprs[i], Dict(x(t) => cv(t)[idx]))
         end
     end
+end
+
+MTK.substitute_differentials(model::CasADiModel, exprs, args...) = exprs
 
+function MTK.substitute_integral(model::CasADiModel, exprs)
+    @unpack U, opti = model
     dt = U.t[2] - U.t[1]
     intmap = Dict()
-    for int in MTK.collect_applied_operators(jcosts, Symbolics.Integral)
+    for int in MTK.collect_applied_operators(exprs, Symbolics.Integral)
         op = MTK.operation(int)
         arg = only(arguments(MTK.value(int)))
         lo, hi = (op.domain.domain.left, op.domain.domain.right)
@@ -255,39 +208,11 @@ function add_cost_function!(model::CasADiModel, sys, tspan, pmap; is_free_t)
         # Approximate integral as sum.
         intmap[int] = dt * tₛ * sum(arg)
     end
-    jcosts = map(c -> Symbolics.substitute(c, intmap), jcosts)
-    jcosts = MTK.value.(jcosts)
-    minimize!(opti, MX(MTK.value(consolidate(jcosts))))
+    exprs = map(c -> Symbolics.substitute(c, intmap), exprs)
+    exprs = MTK.value.(exprs)
 end
 
-function substitute_casadi_vars(
-        model::CasADiModel, sys, pmap, exprs; auxmap::Dict = Dict(), is_free_t)
-    @unpack opti, U, V, tₛ = model
-    iv = MTK.get_iv(sys)
-    sts = unknowns(sys)
-    cts = MTK.unbound_inputs(sys)
-
-    x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
-    c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
-
-    exprs = map(c -> Symbolics.fixpoint_sub(c, auxmap), exprs)
-    exprs = map(c -> Symbolics.fixpoint_sub(c, Dict(pmap)), exprs)
-    # tf means different things in different contexts; a [tf] in a cost function
-    # should be tₛ, while a x(tf) should translate to x[1]
-    if is_free_t
-        free_t_map = Dict([[x(tₛ) => U.u[i, end] for (i, x) in enumerate(x_ops)];
-                           [c(tₛ) => V.u[i, end] for (i, c) in enumerate(c_ops)]])
-        exprs = map(c -> Symbolics.fixpoint_sub(c, free_t_map), exprs)
-    end
-
-    # for variables like x(t)
-    whole_interval_map = Dict([[v => U.u[i, :] for (i, v) in enumerate(sts)];
-                               [v => V.u[i, :] for (i, v) in enumerate(cts)]])
-    exprs = map(c -> Symbolics.fixpoint_sub(c, whole_interval_map), exprs)
-    exprs
-end
-
-function add_solve_constraints(prob, tableau)
+function add_solve_constraints!(prob, tableau)
     @unpack A, α, c = tableau
     @unpack model, f, p = prob
     @unpack opti, U, V, tₛ = model
@@ -332,57 +257,29 @@ function add_solve_constraints(prob, tableau)
     solver_opti
 end
 
-"""
-    solve(prob::CasADiDynamicOptProblem, casadi_solver, ode_solver; plugin_options, solver_options, silent)
-
-`plugin_options` and `solver_options` get propagated to the Opti object in CasADi.
-
-NOTE: the solver should be passed in as a string to CasADi. "ipopt"
-"""
-function DiffEqBase.solve(
-        prob::CasADiDynamicOptProblem, solver::Union{String, Symbol} = "ipopt",
-        tableau_getter = MTK.constructDefault; plugin_options::Dict = Dict(),
-        solver_options::Dict = Dict(), silent = false)
-    @unpack model, u0, p, tspan, f = prob
-    tableau = tableau_getter()
-    @unpack opti, U, V, tₛ = model
-
+function MTK.prepare_solver()
     opti = add_solve_constraints(prob, tableau)
-    silent && (solver_options["print_level"] = 0)
     solver!(opti, "$solver", plugin_options, solver_options)
+end
+function MTK.get_U_values()
+    U_vals = value_getter(U.u)
+    size(U_vals, 2) == 1 && (U_vals = U_vals')
+    U_vals = [[U_vals[i, j] for i in 1:size(U_vals, 1)] for j in 1:length(ts)]
+end
+function MTK.get_V_values()
+end
+function MTK.get_t_values()
+    ts = value_getter(tₛ) * U.t
+end
 
-    failed = false
-    value_getter = nothing
-    sol = nothing
+function MTK.optimize_model!()
     try
         sol = CasADi.solve!(opti)
         value_getter = x -> CasADi.value(sol, x)
     catch ErrorException
         value_getter = x -> CasADi.debug_value(opti, x)
         failed = true
     end
-
-    ts = value_getter(tₛ) * U.t
-    U_vals = value_getter(U.u)
-    size(U_vals, 2) == 1 && (U_vals = U_vals')
-    U_vals = [[U_vals[i, j] for i in 1:size(U_vals, 1)] for j in 1:length(ts)]
-    ode_sol = DiffEqBase.build_solution(prob, tableau_getter, ts, U_vals)
-
-    input_sol = nothing
-    if prod(size(V.u)) != 0
-        V_vals = value_getter(V.u)
-        size(V_vals, 2) == 1 && (V_vals = V_vals')
-        V_vals = [[V_vals[i, j] for i in 1:size(V_vals, 1)] for j in 1:length(ts)]
-        input_sol = DiffEqBase.build_solution(prob, tableau_getter, ts, V_vals)
-    end
-
-    if failed
-        ode_sol = SciMLBase.solution_new_retcode(
-            ode_sol, SciMLBase.ReturnCode.ConvergenceFailure)
-        !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode(
-            input_sol, SciMLBase.ReturnCode.ConvergenceFailure))
-    end
-
-    DynamicOptSolution(model, ode_sol, input_sol)
 end
+MTK.successful_solve() = true
 end