diff --git a/Project.toml b/Project.toml index 1dd887a4e9..7911ed181c 100644 --- a/Project.toml +++ b/Project.toml @@ -123,7 +123,7 @@ FunctionWrappers = "1.1" FunctionWrappersWrappers = "0.1" Graphs = "1.5.2" ImplicitDiscreteSolve = "0.1.2, 1" -InfiniteOpt = "0.5" +InfiniteOpt = "0.6" InteractiveUtils = "1" JuliaFormatter = "1.0.47, 2" JumpProcesses = "9.19" @@ -152,7 +152,7 @@ RecursiveArrayTools = "3.26" Reexport = "0.2, 1" RuntimeGeneratedFunctions = "0.5.9" SCCNonlinearSolve = "1.4.0" -SciMLBase = "2.115.0" +SciMLBase = "2.125.0" SciMLPublic = "1.0.0" SciMLStructures = "1.7" Serialization = "1" diff --git a/docs/src/tutorials/dynamic_optimization.md b/docs/src/tutorials/dynamic_optimization.md index ff1ffab3dc..5ba9598cba 100644 --- a/docs/src/tutorials/dynamic_optimization.md +++ b/docs/src/tutorials/dynamic_optimization.md @@ -101,7 +101,7 @@ The `tf` mapping in the parameter map is treated as an initial guess. Please note that, at the moment, free final time problems cannot support constraints defined at definite time values, like `x(3) ~ 2`. !!! warning - + The Pyomo collocation methods (LagrangeRadau, LagrangeLegendre) currently are bugged for free final time problems. Strongly suggest using BackwardEuler() for such problems when using Pyomo as the backend. When declaring the problem in this case we need to provide the number of steps, since dt can't be known in advanced. Let's solve plot our final solution and the controller for the block, using InfiniteOpt as the backend: @@ -126,3 +126,90 @@ axislegend(ax1) axislegend(ax2) fig ``` + +### Parameter estimation + +The dynamic optimization framework can also be used for parameter estimation. In this approach, we treat unknown parameters as tunable variables and minimize the difference between model predictions and observed data. + +Let's demonstrate this with the Lotka-Volterra equations. First, we'll generate some synthetic data by solving the system with known parameter values: + +```@example dynamic_opt +@parameters α = 1.5 β = 1.0 [tunable=false] γ = 3.0 δ = 1.0 +@variables x_pe(t) y_pe(t) + +eqs_pe = [D(x_pe) ~ α * x_pe - β * x_pe * y_pe, + D(y_pe) ~ -γ * y_pe + δ * x_pe * y_pe] + +@mtkcompile sys0_pe = System(eqs_pe, t) +tspan_pe = (0.0, 1.0) +u0map_pe = [x_pe => 4.0, y_pe => 2.0] + +# True parameter values (these are what we'll try to recover) +parammap_pe = [α => 2.5, δ => 1.8] + +oprob_pe = ODEProblem(sys0_pe, [u0map_pe; parammap_pe], tspan_pe) +osol_pe = solve(oprob_pe, Tsit5()) + +# Generate synthetic data at 51 time points +ts_pe = range(tspan_pe..., length=51) +data_pe = osol_pe(ts_pe, idxs=x_pe).u +``` + +Now we'll set up the parameter estimation problem. We use `EvalAt` to evaluate the state at specific time points and construct a least-squares cost function: + +```@example dynamic_opt +costs_pe = [abs2(EvalAt(ti)(x_pe) - data_pe[i]) for (i, ti) in enumerate(ts_pe)] + +@mtkcompile sys_pe = System(eqs_pe, t; costs = costs_pe) +``` + +By default the cost values are sumed up, if a different behaviour is desired, the `consolidate` keyword can be set in the `System` definition. + +Next, we select which parameters to tune using `subset_tunables`. Here we'll estimate `α` and `δ` while keeping `β` and `γ` fixed: + +```@example dynamic_opt +sys_pe′ = subset_tunables(sys_pe, [α, δ]) +``` + +Now we can solve the parameter estimation problem. Note the `tune_parameters=true` flag: + +```@example dynamic_opt +iprob_pe = InfiniteOptDynamicOptProblem(sys_pe′, u0map_pe, tspan_pe; dt=1/50, tune_parameters=true) +isol_pe = solve(iprob_pe, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3))) + +println("Estimated α = ", isol_pe.sol.ps[α], " (true value: 2.5)") +println("Estimated δ = ", isol_pe.sol.ps[δ], " (true value: 1.8)") +``` + +Let's visualize the fit: + +```@example dynamic_opt +fig = Figure(resolution = (800, 400)) +ax = Axis(fig[1, 1], title = "Parameter Estimation Results", xlabel = "Time", ylabel = "Prey Population") +scatter!(ax, ts_pe, data_pe, label = "Data", markersize = 8) +lines!(ax, isol_pe.sol.t, isol_pe.sol[x_pe], label = "Fitted Model", linewidth = 2) +axislegend(ax) +fig +``` + +!!! note "Time Alignment for Cost Evaluation" + When using `EvalAt` for parameter estimation, different backends handle the case when evaluation times don't align with collocation points differently: + + - **JuMP**: Will throw an error asking you to adjust `dt` if evaluation times don't match collocation points exactly. + - **CasADi**: Uses linear interpolation between collocation points for cost evaluations at intermediate times. + - **InfiniteOpt**: Automatically adds support points for the evaluation times, handling mismatched grids gracefully. + + For example, InfiniteOpt can use a different `dt` than what the data spacing requires: + + ```@example dynamic_opt + # With InfiniteOpt, dt doesn't need to match the data points: + iprob_pe2 = InfiniteOptDynamicOptProblem(sys_pe′, u0map_pe, tspan_pe, + dt = 1/120, tune_parameters=true) + isol_pe2 = solve(iprob_pe2, InfiniteOptCollocation(Ipopt.Optimizer, + InfiniteOpt.OrthogonalCollocation(3))) + + println("With dt=1/120: Estimated α = ", isol_pe2.sol.ps[α], " (true value: 2.5)") + println("With dt=1/120: Estimated δ = ", isol_pe2.sol.ps[δ], " (true value: 1.8)") + ``` + + This flexibility makes InfiniteOpt particularly convenient for parameter estimation when your data points don't naturally align with a uniform collocation grid. diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl index addc478d98..f61cec4274 100644 --- a/ext/MTKCasADiDynamicOptExt.jl +++ b/ext/MTKCasADiDynamicOptExt.jl @@ -21,21 +21,23 @@ function Base.getindex(m::MXLinearInterpolation, i...) length(i) == length(size(m.u)) ? m.u[i...] : m.u[i..., :] end -mutable struct CasADiModel +mutable struct CasADiModel{T} model::Opti U::MXLinearInterpolation V::MXLinearInterpolation + P::T tₛ::MX is_free_final::Bool + tsteps::LinRange{Float64, Int} solver_opti::Union{Nothing, Opti} - function CasADiModel(opti, U, V, tₛ, is_free_final, solver_opti = nothing) - new(opti, U, V, 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 struct CasADiDynamicOptProblem{uType, tType, isinplace, P, F, K} <: - AbstractDynamicOptProblem{uType, tType, isinplace} + SciMLBase.AbstractDynamicOptProblem{uType, tType, isinplace} f::F u0::uType tspan::tType @@ -66,10 +68,11 @@ end function MTK.CasADiDynamicOptProblem(sys::System, op, tspan; dt = nothing, steps = nothing, + tune_parameters = false, guesses = Dict(), kwargs...) prob, _ = MTK.process_DynamicOptProblem( - CasADiDynamicOptProblem, CasADiModel, sys, op, tspan; dt, steps, guesses, kwargs...) + CasADiDynamicOptProblem, CasADiModel, sys, op, tspan; dt, steps, tune_parameters, guesses, kwargs...) prob end @@ -90,6 +93,14 @@ function MTK.generate_input_variable!(model::Opti, c0, nc, tsteps) MXLinearInterpolation(V, tsteps, tsteps[2] - tsteps[1]) end +function MTK.generate_tunable_params!(model::Opti, p0, np) + P = CasADi.variable!(model, np) + for i in 1:np + set_initial!(model, P[i], p0[i]) + end + P +end + function MTK.generate_timescale!(model::Opti, guess, is_free_t) if is_free_t tₛ = variable!(model) @@ -140,14 +151,20 @@ function MTK.lowered_integral(model::CasADiModel, expr, lo, hi) model.tₛ * total end +MTK.needs_individual_tunables(::Opti) = true +MTK.get_param_for_pmap(::Opti, P, i) = P[i] + function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau) @unpack A, α, c = tableau @unpack wrapped_model, f, p = prob - @unpack model, U, V, tₛ = wrapped_model + @unpack model, U, V, P, tₛ = wrapped_model 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) @@ -160,7 +177,7 @@ function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau) ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = MX(zeros(nᵤ))) Uₙ = U.u[:, k] + ΔU * dt Vₙ = V.u[:, k] - Kₙ = tₛ * f(Uₙ, Vₙ, p, τ + h * dt) # scale the time + Kₙ = tₛ * MTK.f_wrapper(f, Uₙ, Vₙ, p, P, τ + h * dt) # scale the time push!(K, Kₙ) end ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)]) @@ -176,7 +193,7 @@ function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau) ΔU = ΔUs[i, :]' Uₙ = U.u[:, k] + ΔU * dt Vₙ = V.u[:, k] - subject_to!(solver_opti, Kᵢ[:, i] == tₛ * f(Uₙ, Vₙ, p, τ + h * dt)) + subject_to!(solver_opti, Kᵢ[:, i] == tₛ * MTK.f_wrapper(f, Uₙ, Vₙ, p, P, τ + h * dt)) end ΔU_tot = dt * (Kᵢ * α) subject_to!(solver_opti, U.u[:, k] + ΔU_tot == U.u[:, k + 1]) @@ -228,6 +245,11 @@ function MTK.get_V_values(model::CasADiModel) end end +function MTK.get_P_values(model::CasADiModel) + value_getter = MTK.successful_solve(model) ? CasADi.debug_value : CasADi.value + [value_getter(model.solver_opti, model.P[i]) for i in eachindex(model.P)] +end + 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 diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index e0f02c0436..bcfb7ce597 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -2,6 +2,7 @@ module MTKInfiniteOptExt using ModelingToolkit using InfiniteOpt using DiffEqBase +using SciMLStructures using LinearAlgebra using StaticArrays using UnPack @@ -13,12 +14,14 @@ struct InfiniteOptModel model::InfiniteModel U::Vector{<:AbstractVariableRef} V::Vector{<:AbstractVariableRef} + P::Vector{<:AbstractVariableRef} tₛ::AbstractVariableRef is_free_final::Bool + tsteps::LinRange{Float64, Int} end struct JuMPDynamicOptProblem{uType, tType, isinplace, P, F, K} <: - AbstractDynamicOptProblem{uType, tType, isinplace} + SciMLBase.AbstractDynamicOptProblem{uType, tType, isinplace} f::F u0::uType tspan::tType @@ -33,7 +36,7 @@ struct JuMPDynamicOptProblem{uType, tType, isinplace, P, F, K} <: end struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <: - AbstractDynamicOptProblem{uType, tType, isinplace} + SciMLBase.AbstractDynamicOptProblem{uType, tType, isinplace} f::F u0::uType tspan::tType @@ -57,6 +60,9 @@ 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_tunable_params!(m::InfiniteModel, p0, np) + @variable(m, P[i=1:np], start=p0[i]) +end function MTK.generate_timescale!(m::InfiniteModel, guess, is_free_t) @variable(m, tₛ ≥ 0, start = guess) @@ -81,20 +87,22 @@ MTK.set_objective!(m::InfiniteOptModel, expr) = @objective(m.model, Min, expr) function MTK.JuMPDynamicOptProblem(sys::System, op, tspan; dt = nothing, steps = nothing, + tune_parameters = false, guesses = Dict(), kwargs...) prob, _ = MTK.process_DynamicOptProblem(JuMPDynamicOptProblem, InfiniteOptModel, sys, - op, tspan; dt, steps, guesses, kwargs...) + op, tspan; dt, steps, tune_parameters, guesses, kwargs...) prob end function MTK.InfiniteOptDynamicOptProblem(sys::System, op, tspan; dt = nothing, steps = nothing, + tune_parameters = false, guesses = Dict(), kwargs...) prob, pmap = MTK.process_DynamicOptProblem(InfiniteOptDynamicOptProblem, InfiniteOptModel, - sys, op, tspan; dt, steps, guesses, kwargs...) + sys, op, tspan; dt, steps, tune_parameters, guesses, kwargs...) MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) prob end @@ -128,10 +136,15 @@ end function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau) @unpack A, α, c = tableau @unpack wrapped_model, f, p = prob - @unpack tₛ, U, V, model = wrapped_model + @unpack tₛ, U, V, P, model = wrapped_model t = model[:t] tsteps = supports(t) - dt = tsteps[2] - tsteps[1] + + # 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) nᵥ = length(V) @@ -142,7 +155,7 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau) Δ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ᵤ] Vₙ = [V[i](τ) for i in 1:nᵥ] - Kₙ = tₛ * f(Uₙ, Vₙ, p, τ + h * dt) + Kₙ = tₛ * MTK.f_wrapper(f, Uₙ, Vₙ, p, P, τ + h * dt) push!(K, Kₙ) end ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)]) @@ -158,12 +171,12 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau) 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($τ)") + @constraint(model, [j = 1:nᵤ], K[i, j]==(tₛ * MTK.f_wrapper(f, Uₙ, V, p, P, τ + h * dt)[j]), + DomainRestriction(==(τ), t), base_name="solve_K$i($τ)") end @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tsteps[end])), - DomainRestrictions(t => τ), base_name="solve_U($τ)") + DomainRestriction(==(τ), t), base_name="solve_U($τ)") end end end @@ -233,6 +246,7 @@ function MTK.get_U_values(m::InfiniteModel) U_vals = value.(m[:U]) U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:nt] end +MTK.get_P_values(m::InfiniteModel) = value(m[:P]) MTK.get_t_values(m::InfiniteModel) = value(m[:tₛ]) * supports(m[:t]) MTK.objective_value(m::InfiniteModel) = InfiniteOpt.objective_value(m) @@ -257,11 +271,11 @@ end # JuMP variables and Symbolics variables never compare equal. When tracing through dynamics, a function argument can be either a JuMP variable or A Symbolics variable, it can never be both. function Base.isequal(::SymbolicUtils.Symbolic, - ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}) + ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, JuMP.GenericNonlinearExpr}) false end function Base.isequal( - ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}, + ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, JuMP.GenericNonlinearExpr}, ::SymbolicUtils.Symbolic) false end diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl index 5b4e9e7a1c..a2a05d297a 100644 --- a/ext/MTKPyomoDynamicOptExt.jl +++ b/ext/MTKPyomoDynamicOptExt.jl @@ -16,30 +16,33 @@ const SPECIAL_FUNCTIONS_DICT = Dict([acos => Pyomo.py_acos, log => Pyomo.py_log, sin => Pyomo.py_sin, sqrt => Pyomo.py_sqrt, - exp => Pyomo.py_exp]) + exp => Pyomo.py_exp, + abs2 => (x -> x^2)]) -struct PyomoDynamicOptModel +struct PyomoDynamicOptModel{T} model::ConcreteModel U::PyomoVar V::PyomoVar + 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, 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(model, U, V, 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 struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <: - AbstractDynamicOptProblem{uType, tType, isinplace} + SciMLBase.AbstractDynamicOptProblem{uType, tType, isinplace} f::F u0::uType tspan::tType @@ -60,11 +63,11 @@ end _getproperty(s, name::Val{fieldname}) where {fieldname} = getproperty(s, fieldname) function MTK.PyomoDynamicOptProblem(sys::System, op, tspan; - dt = nothing, steps = nothing, + dt = nothing, steps = nothing, tune_parameters = false, guesses = Dict(), kwargs...) prob, pmap = MTK.process_DynamicOptProblem(PyomoDynamicOptProblem, PyomoDynamicOptModel, - sys, op, tspan; dt, steps, guesses, kwargs...) + sys, op, tspan; dt, steps, tune_parameters, guesses, kwargs...) conc_model = prob.wrapped_model.model MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) prob @@ -94,6 +97,13 @@ function MTK.generate_input_variable!(m::ConcreteModel, c0, nc, ts) PyomoVar(m.V) end +function MTK.generate_tunable_params!(m::ConcreteModel, p0, np) + m.p_idxs = pyomo.RangeSet(1, np) + init_f = Pyomo.pyfunc((m, i) -> (p0[Pyomo.pyconvert(Int, i)])) + m.P = pyomo.Var(m.p_idxs, initialize = init_f) + PyomoVar(m.P) +end + function MTK.generate_timescale!(m::ConcreteModel, guess, is_free_t) m.tₛ = is_free_t ? pyomo.Var(initialize = guess, bounds = (0, Inf)) : Pyomo.Py(1) PyomoVar(m.tₛ) @@ -169,6 +179,18 @@ function MTK.lowered_var(m::PyomoDynamicOptModel, uv, i, t) Symbolics.unwrap(var) end +# For Pyomo, we need to create symbolic representations for pmap +# This is needed because Pyomo uses symbolic building for constraints +function MTK.get_param_for_pmap(m::ConcreteModel, P::PyomoVar, i) + # Create a symbolic variable that will be used in the pmap + # The actual PyomoVar will be accessed via the symbolic representation + @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel) + P_sym = Symbolics.value(pysym_getproperty(MODEL_SYM, :P)) + Symbolics.unwrap(P_sym[i]) +end + +MTK.needs_individual_tunables(m::ConcreteModel) = true + struct PyomoCollocation <: AbstractCollocation solver::Union{String, Symbol} derivative_method::Pyomo.DiscretizationMethod @@ -208,6 +230,10 @@ function MTK.get_V_values(output::PyomoOutput) m = output.model [[Pyomo.pyconvert(Float64, pyomo.value(m.V[i, t])) for i in m.v_idxs] for t in m.t] end +function MTK.get_P_values(output::PyomoOutput) + m = output.model + [Pyomo.pyconvert(Float64, pyomo.value(m.P[i])) for i in m.p_idxs] +end function MTK.get_t_values(output::PyomoOutput) m = output.model Pyomo.pyconvert(Float64, pyomo.value(m.tₛ)) * [Pyomo.pyconvert(Float64, t) for t in m.t] diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl index 5a0ddbf8d5..5375ba6c87 100644 --- a/src/systems/optimal_control_interface.jl +++ b/src/systems/optimal_control_interface.jl @@ -1,6 +1,3 @@ -abstract type AbstractDynamicOptProblem{uType, tType, isinplace} <: - SciMLBase.AbstractODEProblem{uType, tType, isinplace} end - abstract type AbstractCollocation end struct DynamicOptSolution @@ -21,9 +18,9 @@ end """ JuMPDynamicOptProblem(sys::System, op, tspan; dt, steps, guesses, kwargs...) -Convert an 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 +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`. To construct the problem, please load InfiniteOpt along with ModelingToolkit. @@ -32,8 +29,8 @@ function JuMPDynamicOptProblem end """ InfiniteOptDynamicOptProblem(sys::System, op, tspan; dt) -Convert an System representing an optimal control system into a InfiniteOpt model -for solving using optimization. Must provide `dt` for determining the length +Convert a 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 @@ -45,9 +42,9 @@ function InfiniteOptDynamicOptProblem end """ CasADiDynamicOptProblem(sys::System, op, tspan; dt, steps, guesses, kwargs...) -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 +Convert a 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`. To construct the problem, please load CasADi along with ModelingToolkit. @@ -56,9 +53,9 @@ function CasADiDynamicOptProblem end """ PyomoDynamicOptProblem(sys::System, op, tspan; dt, steps) -Convert an System 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 +Convert a System 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`. To construct the problem, please load Pyomo along with ModelingToolkit. @@ -225,17 +222,69 @@ 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 ### ########################## function process_DynamicOptProblem( - prob_type::Type{<:AbstractDynamicOptProblem}, model_type, sys::System, op, tspan; + prob_type::Type{<:SciMLBase.AbstractDynamicOptProblem}, model_type, sys::System, op, tspan; dt = nothing, steps = nothing, + tune_parameters = false, guesses = Dict(), kwargs...) warn_overdetermined(sys, op) ctrls = unbound_inputs(sys) states = unknowns(sys) + tunable_params = tune_parameters ? tunable_parameters(sys) : [] stidxmap = Dict([v => i for (i, v) in enumerate(states)]) op = Dict([default_toterm(value(k)) => v for (k, v) in op]) @@ -249,18 +298,35 @@ function process_DynamicOptProblem( model_tspan, steps, is_free_t = process_tspan(tspan, dt, steps) warn_overdetermined(sys, op) - pmap = filter(p -> (first(p) ∉ Set(unknowns(sys))), op) - pmap = recursive_unwrap(AnyDict(pmap)) - evaluate_varmap!(pmap, keys(pmap)) + # Build pmap for symbolic substitution in costs/constraints/bounds + all_parameters = default_toterm.(parameters(sys)) + # Extract all parameter values from processed p (which has defaults filled in) + getter = SymbolicIndexingInterface.getp(sys, all_parameters) + pmap = AnyDict(all_parameters .=> getter(p)) + + # Filter out tunable parameters - they should remain symbolic + tunable_set = Set(default_toterm.(tunable_params)) + pmap = filter(kvp -> first(kvp) ∉ tunable_set, pmap) + c0 = value.([pmap[c] for c in ctrls]) + p0, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p) tsteps = LinRange(model_tspan[1], model_tspan[2], steps) model = generate_internal_model(model_type) generate_time_variable!(model, model_tspan, tsteps) U = generate_state_variable!(model, u0, length(states), tsteps) V = generate_input_variable!(model, c0, length(ctrls), tsteps) + P = generate_tunable_params!(model, p0, length(tunable_params)) + # Add the symbolic representation of the tunable parameters to the map + # The order of the Tunable section is given by the tunable_parameters(sys) + # Some backends need symbolic accessors instead of raw variables + P_syms = [get_param_for_pmap(model, P, i) for i in eachindex(tunable_params)] + 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, tₛ, is_free_t) + fullmodel = model_type(model, U, V, P_backend, tₛ, is_free_t, tsteps) + + merge!(pmap, Dict(tunable_params .=> P_syms)) set_variable_bounds!(fullmodel, sys, pmap, tspan[2]) add_cost_function!(fullmodel, sys, tspan, pmap) @@ -274,9 +340,28 @@ function generate_time_variable! end function generate_internal_model end function generate_state_variable! end function generate_input_variable! end +function generate_tunable_params! end function generate_timescale! end function add_initial_constraints! end function add_constraint! end +# Default: return P[i] directly. Symbolic backends (like Pyomo) can override this. +get_param_for_pmap(model, P, i) = P isa AbstractArray ? P[i] : P +# Some backends need symbolic accessors instead of raw variables (CasADi in particular) +needs_individual_tunables(model) = false + +function f_wrapper(f, Uₙ, Vₙ, p, P, t) + if isempty(P) + # no tunable parameters + return f(Uₙ, Vₙ, p, t) + end + if SciMLStructures.isscimlstructure(p) + _, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p) + p′ = repack(P) + f(Uₙ, Vₙ, p′, t) + else + f(Uₙ, Vₙ, P, t) + end +end function set_variable_bounds!(m, sys, pmap, tf) @unpack model, U, V, tₛ = m @@ -443,8 +528,8 @@ function substitute_toterm(vars, exprs) exprs = map(c -> Symbolics.fast_substitute(c, toterm_map), exprs) end -function substitute_params(pmap, exprs) - exprs = map(c -> Symbolics.fixpoint_sub(c, Dict(pmap)), exprs) +function substitute_params(pmap::Dict, exprs) + exprs = map(c -> Symbolics.fixpoint_sub(c, pmap), exprs) end function check_constraint_vars(vars) @@ -467,6 +552,7 @@ function prepare_and_optimize! end function get_t_values end function get_U_values end function get_V_values end +function get_P_values end function successful_solve end """ @@ -474,17 +560,25 @@ 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, +function CommonSolve.solve(prob::SciMLBase.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) + Ps = get_P_values(solved_model) is_free_final(prob.wrapped_model) && (ts .+ prob.tspan[1]) - ode_sol = DiffEqBase.build_solution(prob, solver, ts, Us) - input_sol = isnothing(Vs) ? nothing : DiffEqBase.build_solution(prob, solver, ts, Vs) + # update the parameters with the ones in the solved_model + if !isempty(Ps) + new_p = SciMLStructures.replace(SciMLStructures.Tunable(), prob.p, Ps) + new_prob = remake(prob, p=new_p) + else + new_prob = prob + end + ode_sol = SciMLBase.build_solution(new_prob, solver, ts, Us) + input_sol = isnothing(Vs) ? nothing : SciMLBase.build_solution(new_prob, solver, ts, Vs) if !successful_solve(solved_model) ode_sol = SciMLBase.solution_new_retcode( diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl index eebdc0873c..f6b6d75382 100644 --- a/test/extensions/dynamic_optimization.jl +++ b/test/extensions/dynamic_optimization.jl @@ -1,4 +1,5 @@ using ModelingToolkit +using ModelingToolkit: t_nounits as t, D_nounits as D import InfiniteOpt using DiffEqDevTools, DiffEqBase using SimpleDiffEq @@ -7,6 +8,7 @@ using Ipopt using DataInterpolations using CasADi using Pyomo +using Test import DiffEqBase: solve const M = ModelingToolkit @@ -17,8 +19,6 @@ const ENABLE_CASADI = VERSION >= v"1.11" # Test solving without anything attached. @parameters α=1.5 β=1.0 γ=3.0 δ=1.0 @variables x(..) y(..) - t = M.t_nounits - D = M.D_nounits eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t), D(y(t)) ~ -γ * y(t) + δ * x(t) * y(t)] @@ -133,8 +133,6 @@ end @testset "Linear systems" begin # Double integrator - t = M.t_nounits - D = M.D_nounits @variables x(..) v(..) @variables u(..) [bounds = (-1.0, 1.0), input = true] constr = [v(1.0) ~ 0.0] @@ -239,52 +237,50 @@ end end @testset "Rocket launch" begin - t = M.t_nounits - D = M.D_nounits - @parameters h_c m₀ h₀ g₀ D_c c Tₘ m_c - @variables h(..) v(..) m(..) = m₀ [bounds = (m_c, 1)] T(..) [input = true, bounds = (0, Tₘ)] + ps = @parameters h_c m₀ h₀ g₀ D_c c Tₘ m_c + vars = @variables h(t) v(t) m(t) = m₀ [bounds = (m_c, 1)] T(t) [input = true, bounds = (0, Tₘ)] drag(h, v) = D_c * v^2 * exp(-h_c * (h - h₀) / h₀) gravity(h) = g₀ * (h₀ / h) - eqs = [D(h(t)) ~ v(t), - D(v(t)) ~ (T(t) - drag(h(t), v(t))) / m(t) - gravity(h(t)), - D(m(t)) ~ -T(t) / c] + eqs = [D(h) ~ v, + D(v) ~ (T - drag(h, v)) / m - gravity(h), + D(m) ~ -T / c] (ts, te) = (0.0, 0.2) - costs = [-h(te)] - cons = [T(te) ~ 0, m(te) ~ m_c] - @named rocket = System(eqs, t; costs, constraints = cons) - rocket = mtkcompile(rocket; inputs = [T(t)]) + costs = [-EvalAt(te)(h)] + cons = [EvalAt(te)(T) ~ 0, EvalAt(te)(m) ~ m_c] + @named rocket = System(eqs, t, vars, ps; costs, constraints = cons) + rocket = mtkcompile(rocket; inputs = [T]) - u0map = [h(t) => h₀, m(t) => m₀, v(t) => 0] + u0map = [h => h₀, m => m₀, v => 0] pmap = [ g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5 * √(g₀ * h₀), D_c => 0.5 * 620 * m₀ / g₀, - Tₘ => 3.5 * g₀ * m₀, T(t) => 0.0, h₀ => 1, m_c => 0.6] + Tₘ => 3.5 * g₀ * m₀, T => 0.0, h₀ => 1, m_c => 0.6] jprob = JuMPDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001, cse = false) jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIIA5())) - @test jsol.sol[h(t)][end] > 1.012 + @test jsol.sol[h][end] > 1.012 if ENABLE_CASADI cprob = CasADiDynamicOptProblem( rocket, [u0map; pmap], (ts, te); dt = 0.001, cse = false) csol = solve(cprob, CasADiCollocation("ipopt")) - @test csol.sol[h(t)][end] > 1.012 + @test csol.sol[h][end] > 1.012 end iprob = InfiniteOptDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001) isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) - @test isol.sol[h(t)][end] > 1.012 + @test isol.sol[h][end] > 1.012 pprob = PyomoDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001, cse = false) psol = solve(pprob, PyomoCollocation("ipopt", LagrangeRadau(4))) - @test psol.sol[h(t)][end] > 1.012 + @test psol.sol[h][end] > 1.012 # Test solution @parameters (T_interp::CubicSpline)(..) - eqs = [D(h(t)) ~ v(t), - D(v(t)) ~ (T_interp(t) - drag(h(t), v(t))) / m(t) - gravity(h(t)), - D(m(t)) ~ -T_interp(t) / c] + eqs = [D(h) ~ v, + D(v) ~ (T_interp(t) - drag(h, v)) / m - gravity(h), + D(m) ~ -T_interp(t) / c] @mtkcompile rocket_ode = System(eqs, t) interpmap = Dict(T_interp => ctrl_to_spline(jsol.input_sol, CubicSpline)) oprob = ODEProblem(rocket_ode, merge(Dict(u0map), Dict(pmap), interpmap), (ts, te)) @@ -306,8 +302,6 @@ end end @testset "Free final time problems" begin - t = M.t_nounits - D = M.D_nounits @variables x(..) u(..) [input = true, bounds = (0, 1)] @parameters tf @@ -422,3 +416,119 @@ end psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(4))) @test psol.sol.u[end] ≈ [π, 0, 0, 0] end + +@testset "Parameter defaults usage" begin + # Test that parameter defaults are used when not explicitly provided in op + @parameters α=1.5 β=1.0 γ=3.0 δ=1.0 + @variables x(t) y(t) + + eqs = [D(x) ~ α * x - β * x * y + D(y) ~ -γ * y + δ * x * y] + + sys = mtkcompile(System(eqs, t, name=:sys)) + tspan = (0.0, 1.0) + u0map = [x => 4.0, y => 2.0] + + # Only provide initial conditions, rely on parameter defaults + jprob = JuMPDynamicOptProblem(sys, u0map, tspan, dt = 0.01) + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRK4())) + + # Compare with ODEProblem that also uses defaults + oprob = ODEProblem(sys, u0map, tspan) + osol = solve(oprob, SimpleRK4(), dt = 0.01) + + @test jsol.sol.u ≈ osol.u + + iprob = InfiniteOptDynamicOptProblem(sys, u0map, tspan, dt = 0.01) + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3))) + @test isol.sol.u ≈ osol.u rtol=1e-4 + + if ENABLE_CASADI + cprob = CasADiDynamicOptProblem(sys, u0map, tspan, dt = 0.01) + csol = solve(cprob, CasADiCollocation("ipopt", constructRK4())) + @test csol.sol.u ≈ osol.u + end + + pprob = PyomoDynamicOptProblem(sys, u0map, tspan, dt = 0.01) + psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler())) + + @test psol.sol.u ≈ osol.u rtol=1e-2 +end + +@testset "Parameter estimation" begin + @parameters α = 1.5 β = 1.0 [tunable=false] γ = 3.0 δ = 1.0 + @variables x(t) y(t) + + eqs = [D(x) ~ α * x - β * x * y, + D(y) ~ -γ * y + δ * x * y] + + @mtkcompile sys0 = System(eqs, t) + tspan = (0.0, 1.0) + u0map = [x => 4.0, y => 2.0] + parammap = [α => 2.5, β => 1.0, γ => 3.0, δ => 1.8] + + oprob = ODEProblem(sys0, [u0map; parammap], tspan) + osol = solve(oprob, Tsit5()) + ts = range(tspan..., length=51) + data = osol(ts, idxs=x).u + + costs = [abs2(EvalAt(t)(x)-data[i]) for (i, t) in enumerate(ts)] + consolidate(u, sub) = sum(u) + + @mtkcompile sys = System(eqs, t; costs, consolidate) + + sys′ = subset_tunables(sys, [δ, α]) + jprob = JuMPDynamicOptProblem(sys′, u0map, tspan; dt=1/50, tune_parameters=true) + jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5())) + + @test jsol.sol.ps[δ] ≈ 1.8 rtol=1e-4 + @test jsol.sol.ps[α] ≈ 2.5 rtol=1e-4 + + # test with different time stepping + + jprob = JuMPDynamicOptProblem(sys′, u0map, tspan; dt=1/120, tune_parameters=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))) + + @test isol.sol.ps[δ] ≈ 1.8 rtol=1e-3 + @test isol.sol.ps[α] ≈ 2.5 rtol=1e-3 + + # test with different time stepping + + 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-4 + @test isol.sol.ps[α] ≈ 2.5 rtol=1e-4 + + if ENABLE_CASADI + cprob = CasADiDynamicOptProblem(sys′, u0map, tspan; dt = 1/50, tune_parameters=true) + csol = solve(cprob, CasADiCollocation("ipopt", constructRK4())) + @test csol.sol.ps[δ] ≈ 1.8 rtol=1e-4 + @test csol.sol.ps[α] ≈ 2.5 rtol=1e-4 + + # test with different time stepping + + 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 + @test csol.sol.ps[α] ≈ 2.5 rtol=1e-3 + end + + pprob = PyomoDynamicOptProblem(sys′, u0map, tspan, dt = 1/50, tune_parameters=true) + psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(4))) + + @test psol.sol.ps[δ] ≈ 1.8 rtol=1e-4 + @test psol.sol.ps[α] ≈ 2.5 rtol=1e-4 + + # test with different time stepping + + # pprob = PyomoDynamicOptProblem(sys′, u0map, tspan, dt = 1/120, tune_parameters=true) + # psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(4))) + + # @test psol.sol.ps[δ] ≈ 1.8 rtol=1e-4 + # @test psol.sol.ps[α] ≈ 2.5 rtol=1e-4 +end