Add error handling for JuMPDynamicOptProblem

When we have more collocation points than expected
the constraints are wrong

Author: SebastianM-C

ext/MTKCasADiDynamicOptExt.jl

@@ -28,10 +28,11 @@ mutable struct CasADiModel{T}
     P::T
     tₛ::MX
     is_free_final::Bool
+    tsteps::LinRange{Float64, Int}
     solver_opti::Union{Nothing, Opti}
 
-    function CasADiModel(opti, U, V, P, tₛ, is_free_final, solver_opti = nothing)
-        new{typeof(P)}(opti, U, V, P, tₛ, is_free_final, solver_opti)
+    function CasADiModel(opti, U, V, P, tₛ, is_free_final, tsteps, solver_opti = nothing)
+        new{typeof(P)}(opti, U, V, P, tₛ, is_free_final, tsteps, solver_opti)
     end
 end
 
@@ -160,7 +161,10 @@ function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau)
     solver_opti = copy(model)
 
     tsteps = U.t
-    dt = tsteps[2] - tsteps[1]
+    dt = (tsteps[end] - tsteps[1]) / (length(tsteps) - 1)
+
+    # CasADi uses linear interpolation for cost evaluations that are not on the collocation points
+    @assert tsteps == wrapped_model.tsteps "Collocation points mismatch."
 
     nᵤ = size(U.u, 1)
     nᵥ = size(V.u, 1)

ext/MTKInfiniteOptExt.jl

@@ -17,6 +17,7 @@ struct InfiniteOptModel
     P::Vector{<:AbstractVariableRef}
     tₛ::AbstractVariableRef
     is_free_final::Bool
+    tsteps::LinRange{Float64, Int}
 end
 
 struct JuMPDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
@@ -139,6 +140,10 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
     t = model[:t]
     tsteps = supports(t)
 
+    # InfiniteOpt can introduce additional collocation points
+    # Make sure that the collocation points are correct.
+    MTK.check_collocation_time_mismatch(prob.f.sys, wrapped_model.tsteps, tsteps)
+
     dt = (tsteps[end] - tsteps[1]) / (length(tsteps) - 1)
 
     nᵤ = length(U)

ext/MTKPyomoDynamicOptExt.jl

@@ -26,16 +26,17 @@ struct PyomoDynamicOptModel{T}
     P::T
     tₛ::PyomoVar
     is_free_final::Bool
+    tsteps::LinRange{Float64, Int}
     solver_model::Union{Nothing, ConcreteModel}
     dU::PyomoVar
     model_sym::Union{Num, Symbolics.BasicSymbolic}
     t_sym::Union{Num, Symbolics.BasicSymbolic}
     dummy_sym::Union{Num, Symbolics.BasicSymbolic}
 
-    function PyomoDynamicOptModel(model, U, V, P, tₛ, is_free_final)
+    function PyomoDynamicOptModel(model, U, V, P, tₛ, is_free_final, tsteps)
         @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel) T_SYM DUMMY_SYM
         model.dU = dae.DerivativeVar(U, wrt = model.t, initialize = 0)
-        new{typeof(P)}(model, U, V, P, tₛ, is_free_final, nothing,
+        new{typeof(P)}(model, U, V, P, tₛ, is_free_final, tsteps, nothing,
             PyomoVar(model.dU), MODEL_SYM, T_SYM, DUMMY_SYM)
     end
 end

src/systems/optimal_control_interface.jl

@@ -222,6 +222,56 @@ function process_tspan(tspan, dt, steps)
     end
 end
 
+function get_discrete_time_evaluations(expr)
+    vars = Symbolics.get_variables(expr)
+
+    # Group by base variable
+    result = Dict()
+
+    for v in vars
+        if iscall(v)
+            args = arguments(v)
+            if length(args) == 1 && args[1] isa Number
+                base_var = operation(v)
+                time_point = args[1]
+
+                if !haskey(result, base_var)
+                    result[base_var] = Float64[]
+                end
+                push!(result[base_var], time_point)
+            end
+        end
+    end
+
+    # Sort and unique the time points
+    for (var, times) in result
+        result[var] = sort!(unique!(times))
+    end
+
+    return result
+end
+
+function check_collocation_time_mismatch(sys, expected_tsteps, tsteps)
+    if collect(expected_tsteps)≠tsteps
+        eval_times = get_discrete_time_evaluations(cost(sys))
+        for (var, ts) in eval_times
+            tdiff = setdiff(ts, expected_tsteps)
+            @info tdiff
+            if !isempty(tdiff)
+                error("$var is evaluated inside the cost function at $(length(tdiff)) points " *
+                    "that are not in the $(length(expected_tsteps)) collocation points. " *
+                    "Cost evaluation points must align with the collocation grid. "*
+                    "Adjust the dt to match the time points used in the cost function.")
+            end
+        end
+        if length(expected_tsteps) ≠ length(tsteps)
+            error("Found extra $(abs(length(expected_tsteps) - length(tsteps))) collocation points.")
+        elseif maximum(abs.(expected_tsteps .- tsteps)) > 1e-10
+            error("The collocation points differ from the expected ones by more than 1e-10.")
+        end
+    end
+end
+
 ##########################
 ### MODEL CONSTRUCTION ###
 ##########################
@@ -274,7 +324,7 @@ function process_DynamicOptProblem(
     P_backend = needs_individual_tunables(model) ? P_syms : P
 
     tₛ = generate_timescale!(model, get(pmap, tspan[2], tspan[2]), is_free_t)
-    fullmodel = model_type(model, U, V, P_backend, tₛ, is_free_t)
+    fullmodel = model_type(model, U, V, P_backend, tₛ, is_free_t, tsteps)
 
     merge!(pmap, Dict(tunable_params .=> P_syms))
 
@@ -427,7 +477,6 @@ function process_integral_bounds end
 function lowered_integral end
 function lowered_derivative end
 function lowered_var end
-function lowered_param end
 function fixed_t_map end
 
 function add_user_constraints!(model, sys, tspan, pmap)

test/extensions/dynamic_optimization.jl

@@ -487,10 +487,8 @@ end
     # test with different time stepping
 
     jprob = JuMPDynamicOptProblem(sys′, u0map, tspan; dt=1/120, tune_parameters=true)
-    jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5()))
-
-    @test jsol.sol.ps[δ] ≈ 1.8 rtol=1e-4 broken=true
-    @test jsol.sol.ps[α] ≈ 2.5 rtol=1e-4 broken=true
+    err_msg = "x is evaluated inside the cost function at 40 points that are not in the 121 collocation points."
+    @test_throws err_msg solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5()))
 
     iprob = InfiniteOptDynamicOptProblem(sys′, u0map, tspan, dt = 1/50, tune_parameters=true)
     isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3)))
@@ -503,8 +501,8 @@ end
     iprob = InfiniteOptDynamicOptProblem(sys′, u0map, tspan, dt = 1/120, tune_parameters=true)
     isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3)))
 
-    @test isol.sol.ps[δ] ≈ 1.8 rtol=1e-3
-    @test isol.sol.ps[α] ≈ 2.5 rtol=1e-3
+    @test isol.sol.ps[δ] ≈ 1.8 rtol=1e-4
+    @test isol.sol.ps[α] ≈ 2.5 rtol=1e-4
 
     if ENABLE_CASADI
         cprob = CasADiDynamicOptProblem(sys′, u0map, tspan; dt = 1/50, tune_parameters=true)
@@ -516,8 +514,8 @@ end
 
         cprob = CasADiDynamicOptProblem(sys′, u0map, tspan; dt = 1/120, tune_parameters=true)
         csol = solve(cprob, CasADiCollocation("ipopt", constructRK4()))
-        @test csol.sol.ps[δ] ≈ 1.8 rtol=1e-4 broken=false
-        @test csol.sol.ps[α] ≈ 2.5 rtol=1e-4 broken=true
+        @test csol.sol.ps[δ] ≈ 1.8 rtol=1e-4
+        @test csol.sol.ps[α] ≈ 2.5 rtol=1e-3
     end
 
     pprob = PyomoDynamicOptProblem(sys′, u0map, tspan, dt = 1/50, tune_parameters=true)