Solution, exercise 4

solution_exercise4.py

#import casadi
import numpy
from casadi import *

# Initial and final position
px0 = 0.0;  py0 = 1.5
pxF = 20.0; pyF = 0.0

# Time horizon and discretization
T = 3.0
N = 20
DT = T/float(N)

# Friction term
alpha = 0.02
CVODES = False

# Guess for v0
vx0_guess = 5.0;  vy0_guess = 5.0
v0_guess = [vx0_guess,vy0_guess]

# ODE right hand side
def f(x):
   px = x[0]; vx = x[1]
   py = x[2]; vy = x[3]
   v = sqrt(vx*vx + vy*vy)
   return numpy.array([vx,-alpha*v*vx,vy,-alpha*v*vy - 9.81])

# Integrate with RK4
x = numpy.array([px0, vx0_guess, py0, vy0_guess])
for k in range(N):
    k1 = f(x)
    k2 = f(x + DT/2*k1)
    k3 = f(x + DT/2*k2)
    k4 = f(x + DT*k3)
    x = x + DT/6*(k1 + 2*k2 + 2*k3 + k4)
print "rk4: " + str(x)

# ODE right hand side as CasADi function
x = SX.sym('x',4)
px = x[0]; vx = x[1]
py = x[2]; vy = x[3]
v = sqrt(vx*vx + vy*vy)
xdot = vertcat([vx,-alpha*v*vx,vy,-alpha*v*vy - 9.81])
f = SXFunction([x], [xdot])
f.setOption("name","f")
f.init()

# CasADi integrator
dae = SXFunction(daeIn(x=x), daeOut(ode=xdot))
integrator = Integrator("cvodes",dae)
integrator.setOption("tf",T)
integrator.init()

# Get an expression for the position at the end time
v0 = MX.sym('v0',2)
vx0 = v0[0]
vy0 = v0[1]
X = vertcat([px0, vx0, py0, vy0])
if CVODES:
   integrator_in = integratorIn(x0=x)
   integrator_out = integrator(integrator_in)
   [x] = integratorOut(integrator_out,"xf")
else:
    for k in range(N):
       [k1] = f([X          ])
       [k2] = f([X + DT/2*k1])
       [k3] = f([X + DT/2*k2])
       [k4] = f([X + DT  *k3])       
       X = X + DT/6*(k1 + 2*k2 + 2*k3 + k4)

# Create function for evaluating the state the end time
pf = X[0::2]
F = MXFunction([v0],[pf])
F.setOption("name","F")
F.init()
print "F([v0]): ", F([v0_guess])

# Formulate the single shooting NLP
nlp = MXFunction( nlpIn(x=v0), nlpOut(f=0, g=pf))
solver = NlpSolver("ipopt",nlp)
solver.setOption("max_iter",30)
solver.init()

# Solve the NLP
solver.setInput([5,5], "x0")
solver.setInput([pxF,pyF], "lbg")
solver.setInput([pxF,pyF], "ubg")
solver.evaluate()
print "ipopt:         ", str(solver.getOutput("x"))

# Newton scheme, root-finding function
R = MXFunction([v0],[pf - MX([pxF,pyF])])
R.init()

# Calculate Jacobian
J = R.jacobian()
J.init()

# Newton's method
v = numpy.array(v0_guess)
for it in range(10):
   J.setInput(v)
   J.evaluate()
   Jk = J.getOutput(0)
   Rk = J.getOutput(1)
   v = v - solve(Jk, Rk)
   print "iteration ", it, ", || Rk || = ", norm_2(Rk)
print "Gauss-Newton: ", v