Direct single shooting with CasADi

direct_single_shooting.py

from casadi import *
from numpy import *
import matplotlib.pyplot as plt

N = 20      # Control discretization
T = 10.0    # End time

# Declare variables (use scalar graph)
u  = SX.sym("u")    # control
x  = SX.sym("x",2)  # states

# System dynamics
xdot = vertcat( [(1 - x[1]**2)*x[0] - x[1] + u, x[0]] )
qdot = x[0]**2 + x[1]**2 + u**2
f = SXFunction([x,u],[xdot,qdot])
f.init()

# RK4 with M steps
U = MX.sym("U")
X = MX.sym("X",2)
M = 10; DT = T/(N*M)
XF = X
QF = 0
for j in range(M):
    [k1, k1_q] = f([XF,             U])
    [k2, k2_q] = f([XF + DT/2 * k1, U])
    [k3, k3_q] = f([XF + DT/2 * k2, U])
    [k4, k4_q] = f([XF + DT   * k3, U])
    XF += DT/6*(k1   + 2*k2   + 2*k3   + k4)
    QF += DT/6*(k1_q + 2*k2_q + 2*k3_q + k4_q)
F = MXFunction([X,U],[XF,QF])
F.init()

# Formulate NLP (use matrix graph)
nv = N
v = MX.sym("v",nv)

# Objective function
J=0

# Get an expression for the cost and state at end
X = MX([0,1])
for k in range(N):
    [X, qf] = F([X,v[k]])
    J += qf

# Terminal constraints: x_0(T)=x_1(T)=0
g = X

# Allocate an NLP solver
nlp = MXFunction(nlpIn(x=v),nlpOut(f=J,g=g))
solver = NlpSolver("ipopt", nlp)
solver.init()

# Set bounds and initial guess
solver.setInput( 0., "x0")
solver.setInput(-1., "lbx")
solver.setInput( 1., "ubx")
solver.setInput( 0., "lbg")
solver.setInput( 0., "ubg")

# Solve the problem
solver.evaluate()

# Retrieve the solution
u_opt = solver.getOutput("x")

# Plot the results
plt.figure(1)
plt.clf()
plt.step(linspace(0,T,N),u_opt,'-.')
plt.title("Van der Pol optimization - single shooting")
plt.xlabel('time')
plt.legend(['u trajectory'])
plt.grid()
plt.show()