Skip to content

Instantly share code, notes, and snippets.

@jaeandersson
Created September 24, 2014 09:40
Show Gist options
  • Save jaeandersson/9c52ef3f0f3529d77f62 to your computer and use it in GitHub Desktop.
Save jaeandersson/9c52ef3f0f3529d77f62 to your computer and use it in GitHub Desktop.
Solution, Exercise 8 (first part)
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment