diff --git a/docs/src/tutorials/dynamic_optimization.md b/docs/src/tutorials/dynamic_optimization.md index 40e83f8f38..25f62766e1 100644 --- a/docs/src/tutorials/dynamic_optimization.md +++ b/docs/src/tutorials/dynamic_optimization.md @@ -68,6 +68,46 @@ axislegend(ax2) fig ``` +### Providing an initial trajectory + +By default every state variable is seeded with its constant value from `u0map` at +each collocation point. When that starting point is a poor one, the solver can be +given a guess of the whole trajectory instead. The `initial_trajectory` keyword +takes a map from states to symbolic expressions in the independent variable: + +```@example dynamic_opt +jprob_guess = JuMPDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001, + initial_trajectory = Dict(h(t) => 1 + t, v(t) => 1.0)) +jsol_guess = solve(jprob_guess, JuMPCollocation(Ipopt.Optimizer)); +``` + +Each expression is compiled to a function and evaluated at the collocation points to +produce the start values handed to the optimizer. Only the states listed are affected +— the rest keep their constant seed. This changes where the solve starts from, not the +optimum it converges to. + +Expressions may reference parameters, which are resolved from the operating point, so +the guess can be written in terms of the same quantities as the model: + +```@example dynamic_opt +jprob_guess2 = JuMPDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001, + initial_trajectory = Dict(h(t) => h₀ + t, v(t) => t)) +jsol_guess2 = solve(jprob_guess2, JuMPCollocation(Ipopt.Optimizer)); +``` + +Only parameters of the compiled system can be referenced this way; anything else — +another state, or a quantity `mtkcompile` has eliminated that does not reduce to time +and parameters — raises an `ArgumentError` naming the unresolved quantity. + +`initial_trajectory` is supported by the JuMP, InfiniteOpt, and CasADi backends. +Passing a non-empty map to Pyomo raises an `ArgumentError`. + +!!! note + + For free final time problems (see below) the collocation grid is normalized to + `[0, 1]`, so the trajectory expressions are evaluated at normalized time rather + than physical time. + ### Free final time problems There are additionally a class of dynamic optimization problems where we would like to know how to control our system to achieve something in the least time. Such problems are called free final time problems, since the final time is unknown. To model these problems in ModelingToolkit, we declare the final time as a parameter. diff --git a/lib/ModelingToolkitBase/ext/MTKCasADiDynamicOptExt.jl b/lib/ModelingToolkitBase/ext/MTKCasADiDynamicOptExt.jl index 48be495474..6b9e88e153 100644 --- a/lib/ModelingToolkitBase/ext/MTKCasADiDynamicOptExt.jl +++ b/lib/ModelingToolkitBase/ext/MTKCasADiDynamicOptExt.jl @@ -101,12 +101,12 @@ function MTK.CasADiDynamicOptProblem( dt = nothing, steps = nothing, tune_parameters = false, - guesses = Dict(), + guesses = Dict(), initial_trajectory = Dict(), bounds = Dict(), kwargs... ) prob, _ = MTK.process_DynamicOptProblem( - CasADiDynamicOptProblem, CasADiModel, sys, op, tspan; dt, steps, tune_parameters, guesses, bounds, kwargs... + CasADiDynamicOptProblem, CasADiModel, sys, op, tspan; dt, steps, tune_parameters, guesses, initial_trajectory, bounds, kwargs... ) return prob end @@ -121,6 +121,14 @@ function MTK.generate_state_variable!(model::Opti, u0, ns, tsteps) return MXLinearInterpolation(U, tsteps, tsteps[2] - tsteps[1]) end +function MTK.set_initial_trajectory!(m::Opti, U, idx, traj) + # The collocation grid is fixed when the variables are created, so the + # trajectory is sampled onto it. This overrides the constant `u0` seed set + # in `generate_state_variable!` for the entries of state `idx`. + t_samples = traj.(U.t) + return set_initial!(m, U[idx], DM(t_samples)) +end + function MTK.generate_input_variable!(model::Opti, c0, nc, tsteps) nt = length(tsteps) V = CasADi.variable!(model, nc, nt) diff --git a/lib/ModelingToolkitBase/ext/MTKInfiniteOptExt.jl b/lib/ModelingToolkitBase/ext/MTKInfiniteOptExt.jl index 20c98912c4..68f46382af 100644 --- a/lib/ModelingToolkitBase/ext/MTKInfiniteOptExt.jl +++ b/lib/ModelingToolkitBase/ext/MTKInfiniteOptExt.jl @@ -101,6 +101,10 @@ function MTK.generate_tunable_params!(m::InfiniteModel, p0, np) return @variable(m, P[i = 1:np], start = p0[i]) end +function MTK.set_initial_trajectory!(m::InfiniteModel, U, idx, traj) + return set_start_value(U[idx], traj) +end + function MTK.generate_timescale!(m::InfiniteModel, guess, is_free_t) @variable(m, tₛ ≥ 0, start = guess) if !is_free_t @@ -146,13 +150,13 @@ function MTK.JuMPDynamicOptProblem( dt = nothing, steps = nothing, tune_parameters = false, - guesses = Dict(), + guesses = Dict(), initial_trajectory = Dict(), bounds = Dict(), kwargs... ) prob, _ = MTK.process_DynamicOptProblem( JuMPDynamicOptProblem, InfiniteOptModel, sys, - op, tspan; dt, steps, tune_parameters, guesses, bounds, kwargs... + op, tspan; dt, steps, tune_parameters, guesses, initial_trajectory, bounds, kwargs... ) return prob end @@ -162,13 +166,13 @@ function MTK.InfiniteOptDynamicOptProblem( dt = nothing, steps = nothing, tune_parameters = false, - guesses = Dict(), + guesses = Dict(), initial_trajectory = Dict(), bounds = Dict(), kwargs... ) prob, pmap = MTK.process_DynamicOptProblem( InfiniteOptDynamicOptProblem, InfiniteOptModel, - sys, op, tspan; dt, steps, tune_parameters, guesses, bounds, kwargs... + sys, op, tspan; dt, steps, tune_parameters, guesses, initial_trajectory, bounds, kwargs... ) MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) return prob diff --git a/lib/ModelingToolkitBase/ext/MTKPyomoDynamicOptExt.jl b/lib/ModelingToolkitBase/ext/MTKPyomoDynamicOptExt.jl index 307d451565..e9e0734e41 100644 --- a/lib/ModelingToolkitBase/ext/MTKPyomoDynamicOptExt.jl +++ b/lib/ModelingToolkitBase/ext/MTKPyomoDynamicOptExt.jl @@ -97,13 +97,13 @@ end function MTK.PyomoDynamicOptProblem( sys::System, op, tspan; dt = nothing, steps = nothing, tune_parameters = false, - guesses = Dict(), + guesses = Dict(), initial_trajectory = Dict(), bounds = Dict(), kwargs... ) prob, pmap = MTK.process_DynamicOptProblem( PyomoDynamicOptProblem, PyomoDynamicOptModel, - sys, op, tspan; dt, steps, tune_parameters, guesses, bounds, kwargs... + sys, op, tspan; dt, steps, tune_parameters, guesses, initial_trajectory, bounds, kwargs... ) conc_model = prob.wrapped_model.model MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan) diff --git a/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl b/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl index a08c8a54bd..4cad9db4b1 100644 --- a/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl +++ b/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl @@ -339,6 +339,7 @@ export calculate_jacobian, generate_jacobian, generate_rhs, generate_custom_func export calculate_control_jacobian, generate_control_jacobian export calculate_tgrad, generate_tgrad export generate_cost, calculate_cost_gradient, generate_cost_gradient +export generate_trajectory export calculate_cost_hessian, generate_cost_hessian export calculate_massmatrix, generate_diffusion_function export stochastic_integral_transform diff --git a/lib/ModelingToolkitBase/src/systems/codegen.jl b/lib/ModelingToolkitBase/src/systems/codegen.jl index 767c42ed4f..d871fbed82 100644 --- a/lib/ModelingToolkitBase/src/systems/codegen.jl +++ b/lib/ModelingToolkitBase/src/systems/codegen.jl @@ -497,6 +497,39 @@ function generate_history(sys::System, u0, opts::GeneratedFunctionOptions) ) end +""" + $(TYPEDSIGNATURES) + +Generate a function `f(p, t)` which evaluates the symbolic expression `expr` at time `t` +with the parameter object `p`. `expr` may involve the independent variable, parameters, +observed variables and bound parameters of `sys` (the latter two are inlined +symbolically), but not its unknowns. The `f(p, t)` signature matches the initial guess +function convention used across the SciML ecosystem, e.g. by `BVProblem`. + +# Keyword Arguments + +$GENERATE_X_KWARGS +""" +function generate_trajectory(sys::System, expr, opts::GeneratedFunctionOptions) + (; eval_expression, eval_module, compiler_options) = opts + expression = expression_val(opts) + wrap_gfw = wrap_gfw_val(opts) + p = reorder_parameters(sys) + res = build_function_wrapper( + sys, expr, [p; Any[get_iv(sys)]], BuildFunctionWrapperOptions(; + p_start = 1, p_end = length(p), wrap_delays = false, + codegen_function_options = opts.codegen + ) + ) + if !(expr isa AbstractArray || symbolic_type(expr) == ArraySymbolic()) + res = res[1] + end + return maybe_compile_function( + expression, wrap_gfw, (1, 2, is_split(sys)), res; + compiler_options, eval_expression, eval_module + ) +end + """ $(TYPEDSIGNATURES) diff --git a/lib/ModelingToolkitBase/src/systems/optimal_control_interface.jl b/lib/ModelingToolkitBase/src/systems/optimal_control_interface.jl index ca3b632728..8aa9c50ecb 100644 --- a/lib/ModelingToolkitBase/src/systems/optimal_control_interface.jl +++ b/lib/ModelingToolkitBase/src/systems/optimal_control_interface.jl @@ -323,8 +323,10 @@ function process_DynamicOptProblem( dt = nothing, steps = nothing, tune_parameters = false, - guesses = Dict(), - bounds = Dict(), kwargs... + guesses = Dict(), initial_trajectory = Dict(), + bounds = Dict(), + eval_expression = false, eval_module = @__MODULE__, + kwargs... ) warn_overdetermined(sys, op) ctrls = inputs(sys) @@ -333,6 +335,7 @@ function process_DynamicOptProblem( stidxmap = Dict([v => i for (i, v) in enumerate(states)]) op = Dict([default_toterm(value(k)) => v for (k, v) in op]) + initial_trajectory = Dict([default_toterm(value(k)) => v for (k, v) in initial_trajectory]) bounds = Dict([default_toterm(value(k)) => v for (k, v) in bounds]) u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) : [stidxmap[default_toterm(k)] for (k, v) in op if haskey(stidxmap, k)] @@ -341,7 +344,7 @@ function process_DynamicOptProblem( f, u0, p = process_SciMLProblem( ODEInputFunction, sys, _op; - t = tspan !== nothing ? tspan[1] : tspan, kwargs... + t = tspan !== nothing ? tspan[1] : tspan, eval_expression, eval_module, kwargs... ) model_tspan, steps, is_free_t = process_tspan(tspan, dt, steps) warn_overdetermined(sys, op) @@ -370,6 +373,12 @@ function process_DynamicOptProblem( model = generate_internal_model(model_type) generate_time_variable!(model, model_tspan, tsteps) U = generate_state_variable!(model, u0, length(states), tsteps) + # Apply start trajectories, compiling the symbolic expressions to callables + for (var, traj) in initial_trajectory + idx = get(stidxmap, var, nothing) + idx === nothing && continue + set_initial_trajectory!(model, U, idx, build_trajectory_function(sys, var, traj, p; eval_expression, eval_module)) + end 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 @@ -393,6 +402,60 @@ function process_DynamicOptProblem( return prob_type(f, u0, tspan, p, fullmodel; kwargs...), pmap end +""" + build_trajectory_function(sys, var, traj, p; eval_expression = false, eval_module = @__MODULE__) + +Compile an `initial_trajectory` entry for `var` into a callable of the independent +variable of `sys`, closing over the parameter object `p`. + +`traj` is a symbolic expression in the independent variable and parameters; observed +variables and parameter bindings it references are inlined symbolically, and parameter +values are read from `p` when the trajectory is evaluated. A constant is a valid +trajectory; callables are not supported and raise an `ArgumentError`. + +The result is a `Function`, which is what backends such as InfiniteOpt require of +`JuMP.set_start_value`. +""" +function build_trajectory_function( + sys, var, traj, p; eval_expression = false, eval_module = @__MODULE__ + ) + if symbolic_type(traj) === NotSymbolic() + # A guess that does not vary in time is still a valid trajectory. + traj isa Number && return Returns(traj) + throw( + ArgumentError( + "Only symbolic trajectories are supported for `initial_trajectory`, " * + "got a $(nameof(typeof(traj))) for $var." + ) + ) + end + + iv = get_iv(sys) + expr = get_ir_info(sys).obs_subber(unwrap(traj)) + + unresolved = filter( + v -> !isequal(v, unwrap(iv)) && !is_parameter(sys, v), + Symbolics.get_variables(expr) + ) + isempty(unresolved) || throw( + ArgumentError( + "The `initial_trajectory` for $var may only depend on the independent " * + "variable $iv and on parameters, but after inlining observed variables " * + "and parameter bindings it also depends on $(join(unresolved, ", "))." + ) + ) + + opts = GeneratedFunctionOptions(; expression = Val{false}, eval_expression, eval_module) + rgf = generate_trajectory(sys, expr, opts) + + return Base.Fix1(rgf, p) +end + +# Set the start values of state `idx` from a trajectory callable. +# Backends without a method for their model type do not support this. +function set_initial_trajectory!(model, U, idx, traj) + throw(ArgumentError("The `initial_trajectory` keyword argument is not supported by the $(nameof(typeof(model))) backend.")) +end function generate_time_variable! end function generate_internal_model end function generate_state_variable! end diff --git a/lib/ModelingToolkitBase/test/optimization/dynamic_optimization.jl b/lib/ModelingToolkitBase/test/optimization/dynamic_optimization.jl index 9a25a7ef79..5262032c6b 100644 --- a/lib/ModelingToolkitBase/test/optimization/dynamic_optimization.jl +++ b/lib/ModelingToolkitBase/test/optimization/dynamic_optimization.jl @@ -852,3 +852,77 @@ end jsol_tf = solve(jprob_tf, JuMPCollocation(Ipopt.Optimizer, ExplicitTableaus.RK4())) @test 2 * jsol_tf.sol[x][end] ≤ 0.8 end + +struct UnsupportedTrajectoryBackend end + +@testset "Expression-valued initial trajectories" begin + @variables x(..) v(..) + @variables u(..) [bounds = (-1.0, 1.0), input = true] + constr = [v(1.0) ~ 0.0] + cost = [-x(1.0)] + + @named block = System( + [D(x(t)) ~ v(t), D(v(t)) ~ u(t)], t; costs = cost, constraints = constr + ) + block = mtkcompile(block; inputs = [u(t)]) + + u0map = [x(t) => 0.0, v(t) => 0.0] + tspan = (0.0, 1.0) + parammap = [u(t) => 0.0] + + # Trajectories are symbolic expressions in the independent variable + traj = Dict(x(t) => 0.125 * t^2, v(t) => 0.25 * t) + + iprob = InfiniteOptDynamicOptProblem( + block, [u0map; parammap], tspan; dt = 0.01, + initial_trajectory = traj + ) + isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer)) + @test ≈(isol.sol[x(t)][end], 0.25, rtol = 1.0e-3) + + # Without trajectory should also work + iprob2 = InfiniteOptDynamicOptProblem(block, [u0map; parammap], tspan; dt = 0.01) + isol2 = solve(iprob2, InfiniteOptCollocation(Ipopt.Optimizer)) + @test ≈(isol2.sol[x(t)][end], 0.25, rtol = 1.0e-3) + + if ENABLE_CASADI + cprob = CasADiDynamicOptProblem( + block, [u0map; parammap], tspan; dt = 0.01, + initial_trajectory = traj + ) + csol = solve(cprob, CasADiCollocation("ipopt")) + @test ≈(csol.sol[x(t)][end], 0.25, rtol = 1.0e-3) + end + + # Backends without a `set_initial_trajectory!` method report it clearly + @test_throws ArgumentError M.set_initial_trajectory!( + UnsupportedTrajectoryBackend(), nothing, 1, identity + ) + + # Expressions are compiled to callables of the independent variable + p_test = MTKParameters(block, parammap) + fx = M.build_trajectory_function(block, x(t), 0.125 * t^2, p_test) + @test fx isa Function + @test fx(2.0) ≈ 0.5 + + # A constant expression is still a valid trajectory + fc = M.build_trajectory_function(block, x(t), 3.0, p_test) + @test fc isa Function + @test fc(2.0) == 3.0 + + # Callables are rejected: trajectories are symbolic-only for now + g = τ -> 7.0 + @test_throws ArgumentError M.build_trajectory_function(block, x(t), g, p_test) + + # Parameters stay symbolic and are read from the parameter object when the + # trajectory is evaluated + @parameters a + @named psys = System([D(x(t)) ~ a * v(t), D(v(t)) ~ 0.0], t) + psys = mtkcompile(psys) + fp = M.build_trajectory_function(psys, x(t), a * t, MTKParameters(psys, [a => 4.0])) + @test fp(2.0) ≈ 8.0 + + # Anything that does not reduce to time and parameters is reported + @parameters b + @test_throws ArgumentError M.build_trajectory_function(block, x(t), b * t, p_test) +end