Merge pull request #4017 from SebastianM-C/smc/dynopt

Add tune_parameters for dynamic optimization

Author: AayushSabharwal

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"

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.

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

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