From c4a5616971f1bf25891aa264b508d3fea7c3d6f9 Mon Sep 17 00:00:00 2001 From: Fredrik Bagge Carlson Date: Tue, 19 Nov 2024 09:27:50 +0100 Subject: [PATCH 1/3] add capability to trace MTK dynamics with InfiniteOpt --- Project.toml | 3 + ext/MTKInfiniteOptExt.jl | 19 ++++++ test/extensions/Project.toml | 3 + test/extensions/test_infiniteopt.jl | 102 ++++++++++++++++++++++++++++ test/runtests.jl | 1 + 5 files changed, 128 insertions(+) create mode 100644 ext/MTKInfiniteOptExt.jl create mode 100644 test/extensions/test_infiniteopt.jl diff --git a/Project.toml b/Project.toml index 18b4d55992..5848298418 100644 --- a/Project.toml +++ b/Project.toml @@ -64,6 +64,7 @@ BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6" HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327" +InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57" LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800" [extensions] @@ -72,6 +73,7 @@ MTKChainRulesCoreExt = "ChainRulesCore" MTKDeepDiffsExt = "DeepDiffs" MTKHomotopyContinuationExt = "HomotopyContinuation" MTKLabelledArraysExt = "LabelledArrays" +MTKInfiniteOptExt = "InfiniteOpt" [compat] AbstractTrees = "0.3, 0.4" @@ -104,6 +106,7 @@ FunctionWrappers = "1.1" FunctionWrappersWrappers = "0.1" Graphs = "1.5.2" HomotopyContinuation = "2.11" +InfiniteOpt = "0.5" InteractiveUtils = "1" JuliaFormatter = "1.0.47" JumpProcesses = "9.13.1" diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl new file mode 100644 index 0000000000..88bf724d14 --- /dev/null +++ b/ext/MTKInfiniteOptExt.jl @@ -0,0 +1,19 @@ +module MTKInfiniteOptExt +import ModelingToolkit +import SymbolicUtils +import NaNMath +import InfiniteOpt + +# This file contains method definitions to make it possible to trace through functions generated by MTK using JuMP variables + +for ff in [acos, log1p, acosh, log2, asin, lgamma, 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::GeneralVariableRef) = Base.$f(x) +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. +Base.isequal(::SymbolicUtils.BasicSymbolic{Real}, ::InfiniteOpt.GeneralVariableRef) = false +Base.isequal(::InfiniteOpt.GeneralVariableRef, ::SymbolicUtils.BasicSymbolic{Real}) = false + +end diff --git a/test/extensions/Project.toml b/test/extensions/Project.toml index 37614d9bfb..5f7afe222a 100644 --- a/test/extensions/Project.toml +++ b/test/extensions/Project.toml @@ -4,6 +4,9 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327" +InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57" +Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" +JuMP = "4076af6c-e467-56ae-b986-b466b2749572" LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" Nemo = "2edaba10-b0f1-5616-af89-8c11ac63239a" diff --git a/test/extensions/test_infiniteopt.jl b/test/extensions/test_infiniteopt.jl new file mode 100644 index 0000000000..e45aa0f2fd --- /dev/null +++ b/test/extensions/test_infiniteopt.jl @@ -0,0 +1,102 @@ +using ModelingToolkit, InfiniteOpt, JuMP, Ipopt +using ModelingToolkit: D_nounits as D, t_nounits as t, varmap_to_vars + +@mtkmodel Pendulum begin + @parameters begin + g = 9.8 + L = 0.4 + K = 1.2 + m = 0.3 + end + @variables begin + θ(t) # state + ω(t) # state + τ(t) = 0 # input + y(t) # output + end + @equations begin + D(θ) ~ ω + D(ω) ~ -g / L * sin(θ) - K / m * ω + τ / m / L^2 + y ~ θ * 180 / π + end +end +@named model = Pendulum() +model = complete(model) + +inputs = [model.τ] +(f_oop, f_ip), dvs, psym, io_sys = ModelingToolkit.generate_control_function( + model, inputs, split = false) + +outputs = [model.y] +f_obs = ModelingToolkit.build_explicit_observed_function(io_sys, outputs; inputs = inputs) + +expected_state_order = [model.θ, model.ω] +permutation = [findfirst(isequal(x), expected_state_order) for x in dvs] # This maps our expected state order to the actual state order + +## + +ub = varmap_to_vars([model.θ => 2pi, model.ω => 10], dvs) +lb = varmap_to_vars([model.θ => -2pi, model.ω => -10], dvs) +xf = varmap_to_vars([model.θ => pi, model.ω => 0], dvs) +nx = length(dvs) +nu = length(inputs) +ny = length(outputs) + +## +m = InfiniteModel(optimizer_with_attributes(Ipopt.Optimizer, + "print_level" => 0, "acceptable_tol" => 1e-3, "constr_viol_tol" => 1e-5, "max_iter" => 1000, + "tol" => 1e-5, "mu_strategy" => "monotone", "nlp_scaling_method" => "gradient-based", + "alpha_for_y" => "safer-min-dual-infeas", "bound_mult_init_method" => "mu-based", "print_user_options" => "yes")); + +@infinite_parameter(m, τ in [0, 1], num_supports=51, + derivative_method=OrthogonalCollocation(4)) # Time variable +guess_xs = [t -> pi, t -> 0.1][permutation] +guess_us = [t -> 0.1] +InfiniteOpt.@variables(m, + begin + # state variables + (lb[i] <= x[i = 1:nx] <= ub[i], Infinite(τ), start = guess_xs[i]) # state variables + -10 <= u[i = 1:nu] <= 10, Infinite(τ), (start = guess_us[i]) # control variables + 0 <= tf <= 10, (start = 5) # Final time + 0.2 <= L <= 0.6, (start = 0.4) # Length parameter + end) + +# Trace the dynamics +x0, p = ModelingToolkit.get_u0_p(io_sys, [model.θ => 0, model.ω => 0], [model.L => L]) + +xp = f_oop(x, u, p, τ) +cp = f_obs(x, u, p, τ) # Test that it's possible to trace through an observed function + +@objective(m, Min, tf) +@constraint(m, [i = 1:nx], x[i](0)==x0[i]) # Initial condition +@constraint(m, [i = 1:nx], x[i](1)==xf[i]) # Terminal state + +x_scale = varmap_to_vars([model.θ => 1 + model.ω => 1], dvs) + +# Add dynamics constraints +@constraint(m, [i = 1:nx], (∂(x[i], τ) - tf * xp[i]) / x_scale[i]==0) + +optimize!(m) + +# Extract the optimal solution +opt_tf = value(tf) +opt_time = opt_tf * value(τ) +opt_x = [value(x[i]) for i in permutation] +opt_u = [value(u[i]) for i in 1:nu] +opt_L = value(L) + +# Plot the results +# using Plots +# plot(opt_time, opt_x[1], label = "θ", xlabel = "Time [s]", layout=3) +# plot!(opt_time, opt_x[2], label = "ω", sp=2) +# plot!(opt_time, opt_u[1], label = "τ", sp=3) + +using Test +@test opt_x[1][end]≈pi atol=1e-3 +@test opt_x[2][end]≈0 atol=1e-3 + +@test opt_x[1][1]≈0 atol=1e-3 +@test opt_x[2][1]≈0 atol=1e-3 + +@test opt_L≈0.2 atol=1e-3 # Smallest permissible length is optimal diff --git a/test/runtests.jl b/test/runtests.jl index 192ead935c..25b692f678 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -113,5 +113,6 @@ end @safetestset "Auto Differentiation Test" include("extensions/ad.jl") @safetestset "LabelledArrays Test" include("labelledarrays.jl") @safetestset "BifurcationKit Extension Test" include("extensions/bifurcationkit.jl") + @safetestset "InfiniteOpt Extension Test" include("extensions/test_infiniteopt.jl") end end From 00bfd1d539977588416cfb6de8cdcbcbf8ce2c24 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Tue, 19 Nov 2024 09:36:52 -0100 Subject: [PATCH 2/3] Update MTKInfiniteOptExt.jl --- ext/MTKInfiniteOptExt.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index 88bf724d14..79463f7e71 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -15,5 +15,6 @@ 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. Base.isequal(::SymbolicUtils.BasicSymbolic{Real}, ::InfiniteOpt.GeneralVariableRef) = false Base.isequal(::InfiniteOpt.GeneralVariableRef, ::SymbolicUtils.BasicSymbolic{Real}) = false - +Base.isequal(::SymbolicUtils.Symbolic, ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}) = false +Base.isequal(::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}, ::SymbolicUtils.Symbolic) = false end From 72009993b81bc0e7908e35c20775b0c03589beb1 Mon Sep 17 00:00:00 2001 From: Fredrik Bagge Carlson Date: Wed, 20 Nov 2024 18:39:08 +0100 Subject: [PATCH 3/3] rm lgamma --- ext/MTKInfiniteOptExt.jl | 16 +++++++++++----- test/odesystem.jl | 2 +- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl index 79463f7e71..3b482a1f03 100644 --- a/ext/MTKInfiniteOptExt.jl +++ b/ext/MTKInfiniteOptExt.jl @@ -3,18 +3,24 @@ import ModelingToolkit import SymbolicUtils import NaNMath import InfiniteOpt +import InfiniteOpt: JuMP, GeneralVariableRef # This file contains method definitions to make it possible to trace through functions generated by MTK using JuMP variables -for ff in [acos, log1p, acosh, log2, asin, lgamma, tan, atanh, cos, log, sin, log10, sqrt] +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::GeneralVariableRef) = Base.$f(x) 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. -Base.isequal(::SymbolicUtils.BasicSymbolic{Real}, ::InfiniteOpt.GeneralVariableRef) = false -Base.isequal(::InfiniteOpt.GeneralVariableRef, ::SymbolicUtils.BasicSymbolic{Real}) = false -Base.isequal(::SymbolicUtils.Symbolic, ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}) = false -Base.isequal(::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}, ::SymbolicUtils.Symbolic) = false +function Base.isequal(::SymbolicUtils.Symbolic, + ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}) + false +end +function Base.isequal( + ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr}, + ::SymbolicUtils.Symbolic) + false +end end diff --git a/test/odesystem.jl b/test/odesystem.jl index 90d1c4c578..02eb80ef85 100644 --- a/test/odesystem.jl +++ b/test/odesystem.jl @@ -1521,6 +1521,6 @@ end @parameters p[1:2] = [1.0, 2.0] @mtkbuild sys = ODESystem([D(x) ~ x, y^2 ~ x + sum(p)], t) prob = DAEProblem(sys, [D(x) => x, D(y) => D(x) / 2y], [], (0.0, 1.0)) - sol = solve(prob, DFBDF(), abstol=1e-8, reltol=1e-8) + sol = solve(prob, DFBDF(), abstol = 1e-8, reltol = 1e-8) @test sol[x]≈sol[y^2 - sum(p)] atol=1e-5 end