add tune_parameters for the rest of the backends

Co-authored-by: Claude <[email protected]>

Author: SebastianM-C

ext/MTKCasADiDynamicOptExt.jl

@@ -21,16 +21,17 @@ 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
     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, solver_opti = nothing)
+        new{typeof(P)}(opti, U, V, P, tₛ, is_free_final, solver_opti)
     end
 end
 
@@ -66,10 +67,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 +92,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)
@@ -143,7 +153,7 @@ end
 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
@@ -160,7 +170,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 +186,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 +238,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)
+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

@@ -97,10 +97,11 @@ 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
@@ -131,16 +132,6 @@ function MTK.lowered_var(m::InfiniteOptModel, uv, i, t)
     t isa Union{Num, Symbolics.Symbolic} ? X[i] : X[i](t)
 end
 
-function f_wrapper(f, Uₙ, Vₙ, p, P, t)
-    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 add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
     @unpack A, α, c = tableau
     @unpack wrapped_model, f, p = prob
@@ -159,7 +150,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_wrapper(f, Uₙ, Vₙ, p, 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(α)])
@@ -175,12 +166,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_wrapper(f, Uₙ, V, p, 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

ext/MTKPyomoDynamicOptExt.jl

@@ -16,12 +16,14 @@ 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
     model::ConcreteModel
     U::PyomoVar
     V::PyomoVar
+    P::PyomoVar
     tₛ::PyomoVar
     is_free_final::Bool
     solver_model::Union{Nothing, ConcreteModel}
@@ -30,10 +32,10 @@ struct PyomoDynamicOptModel
     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)
         @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(model, U, V, P, tₛ, is_free_final, nothing,
             PyomoVar(model.dU), MODEL_SYM, T_SYM, DUMMY_SYM)
     end
 end
@@ -60,11 +62,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 +96,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 +178,16 @@ function MTK.lowered_var(m::PyomoDynamicOptModel, uv, i, t)
     Symbolics.unwrap(var)
 end
 
+function MTK.lowered_param(m::PyomoDynamicOptModel, i)
+    P = Symbolics.value(pysym_getproperty(m.model_sym, :P))
+    Symbolics.unwrap(P[i])
+end
+
+# For Pyomo, return symbolic accessors instead of raw PyomoVar
+function MTK.get_param_for_pmap(m::PyomoDynamicOptModel, P, i)
+    MTK.lowered_param(m, i)
+end
+
 struct PyomoCollocation <: AbstractCollocation
     solver::Union{String, Symbol}
     derivative_method::Pyomo.DiscretizationMethod
@@ -208,6 +227,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]