Optimal control interface

Project.toml

@@ -59,7 +59,7 @@ Random = "1.10"
 RecursiveArrayTools = "3.31.2"
 ReTestItems = "1.29"
 Reexport = "1.2"
-SciMLBase = "2.108.0"
+SciMLBase = "2.120.0"
 Sparspak = "0.3.11"
 StaticArrays = "1.9.8"
 Test = "1.10"

docs/src/tutorials/optimal_control.md

@@ -0,0 +1,101 @@
+# Solve Optimal Control problem
+
+A classical optimal control problem is the rocket launching problem(aka [Goddard Rocket problem](https://en.wikipedia.org/wiki/Goddard_problem)). Say we have a rocket with limited fuel and is launched vertically. And we want to control the final altitude of this rocket so that we can make the best of the limited fuel in rocket to get to the highest altitude. In this optimal control problem, the state variables are:
+
+  - Velocity of the rocket: $x_v(t)$
+  - Altitude of the rocket: $x_h(t)$
+  - Mass of the rocket and the fuel: $x_m(t)$
+
+The control variable is
+
+  - Thrust of the rocket: $u_t(t)$
+
+The dynamics of the launching can be formulated with three differential equations:
+
+$$
+\left\{\begin{aligned}
+&\frac{dx_v}{dt}=\frac{u_t-drag(x_h,x_v)}{x_m}-g(x_h)\\
+&\frac{dx_h}{dt}=x_v\\
+&\frac{dx_m}{dt}=-\frac{u_t}{c}
+\end{aligned}\right.
+$$
+
+where the drag $D(x_h,x_v)$ is a function of altitude and velocity:
+
+$$
+D(x_h,x_v)=D_c\cdot x_v^2\cdot\exp^{h_c(\frac{x_h-x_h(0)}{x_h(0)})}
+$$
+
+gravity $g(x_h)$ is a function of altitude:
+
+$$
+g(x_h)=g_0\cdot (\frac{x_h(0)}{x_h})^2
+$$
+
+$c$ is a constant. Suppose the final time is $T$, we here want to maximize the final altitude $x_h(T)$:
+
+$$
+\max x_h(T)
+$$
+
+The inequality constraints for the state variables and control variables are:
+
+$$
+\left\{\begin{aligned}
+&x_v>0\\
+&x_h>0\\
+&m_T<x_m<m_0\\
+&0<u_t<u_{t\text{max}}
+\end{aligned}\right.
+$$
+
+Similar solving for such optimal control problem can be found on JuMP.jl and InfiniteOpt.jl. The detailed parameters are taken from [COPS](https://www.mcs.anl.gov/%7Emore/cops/cops3.pdf).
+
+```julia
+using BoundaryValueDiffEqMIRK, OptimizationMOI, Ipopt
+h_0 = 1                      # Initial height
+v_0 = 0                      # Initial velocity
+m_0 = 1.0                    # Initial mass
+m_T = 0.6                    # Final mass
+g_0 = 1                      # Gravity at the surface
+h_c = 500                    # Used for drag
+c = 0.5 * sqrt(g_0 * h_0)    # Thrust-to-fuel mass
+D_c = 0.5 * 620 * m_0 / g_0  # Drag scaling
+u_t_max = 3.5 * g_0 * m_0    # Maximum thrust
+T_max = 0.2                  # Number of seconds
+T = 1_000                    # Number of time steps
+Δt = 0.2 / T;                # Time per discretized step
+
+tspan = (0.0, 0.2)
+drag(x_h, x_v) = D_c * x_v^2 * exp(-h_c * (x_h - h_0) / h_0)
+g(x_h) = g_0 * (h_0 / x_h)^2
+function rocket_launch!(du, u, p, t)
+    # u_t is the control variable (thrust)
+    x_v, x_h, x_m, u_t = u[1], u[2], u[3], u[4]
+    du[1] = (u_t-drag(x_h, x_v))/x_m - g(x_h)
+    du[2] = x_v
+    du[3] = -u_t/c
+end
+function rocket_launch_bc!(res, u, p, t)
+    res[1] = u(0.0)[1] - v_0
+    res[2] = u(0.0)[2] - h_0
+    res[3] = u(0.2)[3] - m_T
+    res[4] = u(0.2)[4] - 0.0
+end
+function constraints!(res, u, p)
+    res[1] = u[1]
+    res[2] = u[2]
+    res[3] = u[3]
+    res[4] = u[4]
+end
+cost_fun(u, p) = -u[end - 2] #Final altitude x_h. To minimize, only temporary, need to use temporary solution interpolation here similar to what we do in boundary condition evaluations.
+#cost_fun(sol, p) = -sol(0.2)[2]
+u0 = [v_0, h_0, m_0, 0.0]
+rocket_launch_fun = BVPFunction(rocket_launch!, rocket_launch_bc!; cost = cost_fun,
+    inequality = constraints!, f_prototype = zeros(3))
+rocket_launch_prob = BVProblem(rocket_launch_fun, u0, tspan; lcons = [0.0, h_0, m_T, 0.0],
+    ucons = [Inf, Inf, m_0, u_t_max])
+sol = solve(rocket_launch_prob, MIRK4(; optimize = Ipopt.Optimizer()); dt = 0.002)
+```
+
+Similar optimal control problem solving can also be deployed in JuMP.jl and InfiniteOpt.jl.

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -629,16 +629,28 @@ end
 
 @inline function __extract_lcons_ucons(prob::AbstractBVProblem, ::Type{T}, M, N) where {T}
     lcons = if isnothing(prob.lcons)
-        zeros(T, N*M)
+        zeros(T, N*M) #TODO: handle carefully when NLLS
     else
-        lcons_length = length(prob.lcons)
-        vcat(prob.lcons, zeros(T, N*M - lcons_length))
+        length_f_prototype = length(prob.f.f_prototype)
+        if !(isnothing(prob.f.equality) && isnothing(prob.f.inequality))
+            # When there are additional equality or inequality constraints
+            vcat(repeat(prob.lcons, N), zeros(T, M + (N - 1)*length_f_prototype))
+        else
+            lcons_length = length(prob.lcons)
+            vcat(prob.lcons, zeros(T, N*M - lcons_length))
+        end
     end
     ucons = if isnothing(prob.ucons)
-        zeros(T, N*M)
+        zeros(T, N*M) #TODO: handle carefully when NLLS
     else
-        ucons_length = length(prob.ucons)
-        vcat(prob.ucons, zeros(T, N*M - ucons_length))
+        length_f_prototype = length(prob.f.f_prototype)
+        if !(isnothing(prob.f.equality) && isnothing(prob.f.inequality))
+            # When there are additional equality or inequality constraints
+            vcat(repeat(prob.ucons, N), zeros(T, M + (N - 1)*length_f_prototype))
+        else
+            ucons_length = length(prob.ucons)
+            vcat(prob.ucons, zeros(T, N*M - ucons_length))
+        end
     end
     return lcons, ucons
 end

lib/BoundaryValueDiffEqFIRK/src/algorithms.jl

@@ -121,7 +121,6 @@ for stage in (2, 3, 4, 5)
           - `optimize`: Internal Optimization solver. Any solver which conforms to the SciML
             `OptimizationProblem` interface can be used. Note that any autodiff argument for
             the solver will be ignored and a custom jacobian algorithm will be used.
-
           - `jac_alg`: Jacobian Algorithm used for the nonlinear solver. Defaults to
             `BVPJacobianAlgorithm()`, which automatically decides the best algorithm to
             use based on the input types and problem type.
@@ -132,7 +131,6 @@ for stage in (2, 3, 4, 5)
                 `nonbc_diffmode` defaults to `AutoSparse(AutoForwardDiff)` if possible else
                 `AutoSparse(AutoFiniteDiff)`. For `bc_diffmode`, defaults to `AutoForwardDiff` if
                 possible else `AutoFiniteDiff`.
-
           - `nested_nlsolve`: Whether or not to use a nested nonlinear solve for the
             implicit FIRK step. Defaults to `false`. If set to `false`, the FIRK stages are
             solved as a part of the global residual. The general recommendation is to choose
@@ -150,6 +148,7 @@ for stage in (2, 3, 4, 5)
         ## References
 
         Reference for Lobatto and Radau methods:
+
         ```bibtex
         @Inbook{Jay2015,
             author="Jay, Laurent O.",
@@ -168,7 +167,9 @@ for stage in (2, 3, 4, 5)
             year = {2015},
         }
         ```
+
         References for implementation of defect control, based on the `bvp5c` solver in MATLAB:
+
         ```bibtex
         @article{shampine_solving_nodate,
             title = {Solving {Boundary} {Value} {Problems} for {Ordinary} {Differential} {Equations} in {Matlab} with bvp4c

lib/BoundaryValueDiffEqMIRK/src/collocation.jl

@@ -1,10 +1,43 @@
-function Φ!(residual, cache::MIRKCache, y, u, trait)
-    return Φ!(residual, cache.fᵢ_cache, cache.k_discrete, cache.f, cache.TU,
-        y, u, cache.p, cache.mesh, cache.mesh_dt, cache.stage, trait)
+function Φ!(residual, cache::MIRKCache, y, u, trait, constraint)
+    return Φ!(residual, cache.fᵢ_cache, cache.k_discrete, cache.f, cache.TU, y, u,
+        cache.p, cache.mesh, cache.mesh_dt, cache.stage, trait, constraint)
 end
 
-@views function Φ!(residual, fᵢ_cache, k_discrete, f!, TU::MIRKTableau,
-        y, u, p, mesh, mesh_dt, stage::Int, ::DiffCacheNeeded)
+@views function Φ!(residual, fᵢ_cache, k_discrete, f!, TU::MIRKTableau, y, u, p, mesh,
+        mesh_dt, stage::Int, ::DiffCacheNeeded, constraint::Val{true})
+    (; c, v, x, b) = TU
+
+    ttt = get_tmp(fᵢ_cache, u)
+    tmp = copy(ttt[1:3])
+
+    length_control = length(last(residual))
+
+    T = eltype(u)
+    for i in eachindex(k_discrete)
+        K = get_tmp(k_discrete[i], u)
+        residᵢ = residual[i]
+        h = mesh_dt[i]
+
+        yᵢ = get_tmp(y[i], u)
+        yᵢ₊₁ = get_tmp(y[i + 1], u)
+
+        yᵢ, uᵢ = yᵢ[1:(end - length_control)], yᵢ[(end - length_control + 1):end]
+        yᵢ₊₁, uᵢ₊₁ = yᵢ₊₁[1:(end - length_control)], yᵢ₊₁[(end - length_control + 1):end]
+
+        for r in 1:stage
+            @. tmp = (1 - v[r]) * yᵢ + v[r] * yᵢ₊₁
+            __maybe_matmul!(tmp, K[:, 1:(r - 1)], x[r, 1:(r - 1)], h, T(1))
+            f!(K[:, r], vcat(tmp, uᵢ), p, mesh[i] + c[r] * h)
+        end
+
+        # Update residual
+        @. residᵢ = yᵢ₊₁ - yᵢ
+        __maybe_matmul!(residᵢ, K[:, 1:stage], b[1:stage], -h, T(1))
+    end
+end
+
+@views function Φ!(residual, fᵢ_cache, k_discrete, f!, TU::MIRKTableau, y, u, p, mesh,
+        mesh_dt, stage::Int, ::DiffCacheNeeded, constraint::Val{false})
     (; c, v, x, b) = TU
 
     tmp = get_tmp(fᵢ_cache, u)
@@ -29,8 +62,8 @@ end
     end
 end
 
-@views function Φ!(residual, fᵢ_cache, k_discrete, f!, TU::MIRKTableau, y,
-        u, p, mesh, mesh_dt, stage::Int, ::NoDiffCacheNeeded)
+@views function Φ!(residual, fᵢ_cache, k_discrete, f!, TU::MIRKTableau, y, u, p, mesh,
+        mesh_dt, stage::Int, ::NoDiffCacheNeeded, constraint::Val{false})
     (; c, v, x, b) = TU
 
     tmp = similar(fᵢ_cache)