Solution, Exercise 8 (first part)

solution_exercise8.py

import numpy as NP
from casadi import *

# Choose collocation points
d = 4  # Degree
tau_root = collocationPoints(d,"legendre")

# Calculate coefficients
C = DMatrix.nan(d+1,d+1)  # For the collocation equation
D = DMatrix.nan(d+1)      # For the continuity equation
for j in range(d+1):
  # Lagrange polynomial
  tau = SX.sym("tau")
  L = 1
  for r in range(d+1):
    if r != j:
      L *= (tau-tau_root[r])/(tau_root[j]-tau_root[r])

  # Evaluate the polynomial at the final time to get the coefficients of the 
  # continuity equation
  lfcn = SXFunction([tau],[L])
  lfcn.init()
  [D[j]] = lfcn([1.0])

  # Evaluate the time derivative of the polynomial at all collocation points to
  # get the coefficients of the collocation equation
  dL = gradient(L,tau)
  dfcn = SXFunction([tau],[dL])
  dfcn.init()
  for r in range(d+1):
    [C[j,r]] = dfcn([tau_root[r]])

# End time
tf = 3.

# Number of time steps
N = 20

# Integrator step size
DT = tf/N

# Declare variables
nx = 4
x  = SX.sym("x",nx)
px = x[0]; vx = x[1]; py = x[2]; vy = x[3]

# Constants
alpha = 0.02
g0 = 9.81
h = 1.5

# ODE right hand side function
xdot = vertcat([vx,
                -alpha*vx*sqrt(vx**2 + vy**2),
                vy,
                -alpha*vy*sqrt(vx**2 + vy**2) - g0])
f = SXFunction([x],[xdot])
f.init()

# NLP (root-finding) variables
nw = nx + d*nx + nx
w = MX.sym("w",nw)
X = vertsplit(w,nx) # split up in chunks of length nx
# equivalent to: X = [w[nx*k:nx*(k+1)] for k in range(d+2)]


# Get the collocation equations (that define V)
g = []
for j in range(1,d+1):
  # Expression for the state derivative at the collocation point
  xdot_j = 0
  for r in range (d+1):
    xdot_j += C[r,j]*X[r]

  # Append collocation equations
  [f_j] = f([X[j]])
  g.append(DT*f_j - xdot_j)

# Get an expression for the state at the end of the finite element
xf = 0
for r in range(d+1):
  xf += D[r]*X[r]

# Append continuity equation
g.append(xf - X[d+1])

# Concatenate equations
g = vertcat(g)

# Bounds on w
lbw = -DMatrix.inf(nw)
ubw =  DMatrix.inf(nw)

# Create NLP solver
nlp = MXFunction(nlpIn(x=w),nlpOut(f=0,g=g))
solver = NlpSolver('ipopt',nlp)
solver.setOption("print_level",0)
solver.setOption("print_time",False)
solver.init()

# Initial value
x = DMatrix([0., 5., h, 5.])

# Initial guess
wk = repmat(x,1+d+1,1)

# Integrate
for k in range(N):
    # Update bounds on w
    ubw[:nx] = lbw[:nx] = x

    # Solve NLP
    solver.setInput(wk, "x0")
    solver.setInput(lbw, "lbx")
    solver.setInput(ubw, "ubx")
    solver.setInput(0, "lbg")
    solver.setInput(0, "ubg")
    solver.evaluate()
    wk = solver.getOutput('x')
    
    # Get state at end
    x = wk[-nx:] # MATLAB: wk(end-nx:end)

    print "Step %2d (%s, %2d iterations): %s" % \
        (k, solver.getStat('return_status'),solver.getStat('iter_count'),x)