[RFC] Add functional guesses

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.

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)

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

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)

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

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)
 

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

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