Elena's bug

pendulum_debug.py

# Pendulum collocation
# w position of cart
# w1,y1 position of pendulum

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


nk = 50    # Control discretization
tf = 3.0  # End time

m = 0.1
l = 1.
M = 1.
gr = 9.81

# Dimensions
nu = 1

# -----------------------------------------------------------------------------
# Collocation setup
# -----------------------------------------------------------------------------



# Degree of interpolating polynomial
d = 6 #Fails for high order no f%#"@ clue why...

# Choose collocation points
tau_root = collocationPoints(d,"radau")

# Size of the finite elements
h = tf/nk

# Coefficients of the collocation equation
C = NP.zeros((d+1,d+1))

# Coefficients of the continuity equation
D = NP.zeros(d+1)

# Dimensionless time inside one control interval
tau = SX.sym("tau")

# All collocation time points
T = NP.zeros((nk,d+1))
for k in range(nk):
    for j in range(d+1):
        T[k,j] = (k + tau_root[j])


# For all collocation points
for j in range(d+1):
    # Construct Lagrange polynomials to get the polynomial basis at the collocation point
    L = 1
    for r in range(d+1):
        if r != j:
            L *= (tau-tau_root[r])/(tau_root[j]-tau_root[r])
    lfcn = SXFunction('lfcn', [tau],[L])

    # Evaluate the polynomial at the final time to get the coefficients of the continuity equation
    [D[j]] = lfcn([1.0])

    # Evaluate the time derivative of the polynomial at all collocation points to get the coefficients of the continuity equation
    tfcn = lfcn.tangent()
    for r in range(d+1):
        C[j][r], _ = tfcn([tau_root[r]])

# -----------------------------------------------------------------------------
# Model Pendulum
# -----------------------------------------------------------------------------


# Declare variables (use scalar graph)
t  = SX.sym('t')          # time
u  = SX.sym('u')          # control
xa = SX.sym('xa')    # algebraic state
p  = SX.sym('p',0,1)      # parameters


state_labels = ['w0','w','y','dw0','dw','dy']  # cart, mass x, mass y
xd    = struct_symSX(state_labels)
xddot = struct_symSX([ 'd'+name for name in state_labels ])

res = vertcat([xddot['dw0'] - xd['dw0'],\
               xddot['dw'] -  xd['dw'],\
               xddot['dy'] -  xd['dy'],\
               M*xddot['ddw0'] + (xd['w0']-xd['w'])*xa + u, \
               m*xddot['ddw']  + (xd['w']-xd['w0'])*xa , \
               m*xddot['ddy']  +         (xd['y']*xa) + (m*gr) ,\
               (xd['w'] - xd['w0'])*(xddot['ddw'] - xddot['ddw0']) + xd['y']*xddot['ddy'] + xd['dy']**2 + (xd['dw0']-xd['dw'])**2 ])

# System dynamics (implicit formulation)
ffcn = SXFunction('ffcn', [t,xd,xddot,xa,u,p],[res])

# Objective function
xref = xd()
xref['y'] = l
lagr = 1e2*(u-0)**2

LagrangeTerm = SXFunction('lagrange', [t,xd,xa,u,p],[lagr])


# Differential state bounds
#Path bounds
xD_min = xd(-inf)
xD_max = xd( inf)
#Initial bounds
xDi_min = xd()
xDi_max = xd()
xDi_min['y'] = xDi_max['y'] = -l
#Final bounds
xDf_min = xd(-inf)
xDf_max = xd( inf)
for name in ['y','w0','dw0','dy']:
    xDf_min[name] = xDf_max[name] = xref[name]


# Structure holding NLP variables
V = struct_symMX([
                  (
                   entry('Xd',repeat=[nk,d+1],struct=xd),
                   entry('XA',repeat=[nk,d],struct=xa),
                   entry('U',repeat=[nk],shape=nu)
                   )
                  ])

P = MX.sym('P',0,1)



# Constraint function for the NLP
coll_cstr = []
continuity_cstr = []
# For all finite elements
for k in range(nk):

    # For all collocation points
    for j in range(1,d+1):

        # Get an expression for the state derivative at the collocation point
        xp_jk = 0
        for r in range (d+1):
            xp_jk += C[r,j]*V["Xd",k,r]

        # Add collocation equations to the NLP
        [fk] = ffcn([h*T[k,j],V["Xd",k,j],xp_jk/h, V["XA",k,j-1], V["U",k], P])
        coll_cstr.append(fk)

   # Get an expression for the state at the end of the finite element
    xf_k = 0
    for r in range(d+1):
        xf_k += D[r]*V["Xd",k,r]

    # Add continuity equation to NLP
    if k < nk-1:
        continuity_cstr.append(V["Xd",k+1,0] - xf_k)

xtf = 0
for r in range(d+1):
    xtf += D[r]*V["Xd",-1,r]

g = struct_MX(
              [
                entry('collocation',expr=coll_cstr),
                entry('continuity',expr=continuity_cstr),
                entry('finaltime', expr = xtf)
                ]
              )


# Objective function
f = 0
for k in range(nk):
    [fk] = LagrangeTerm([h*T[k,j],V["Xd",k,0], V["XA",k,0], V["U",k], P])
    f += fk


# Initialise NLP - solver
nlp = MXFunction('nlp', nlpIn(x=V, p=P),nlpOut(f=f,g=g))

# -----------------------------------------------------------------------------
## SOLVE THE NLP
# -----------------------------------------------------------------------------


# Concatenate constraints
lbg = g()
ubg = g()

for k,name in enumerate(state_labels):
    lbg['finaltime',k] = xDf_min[name]
    ubg['finaltime',k] = xDf_max[name]

# default bound -inf +inf
vars_lb   = V(-inf)
vars_ub   = V( inf)
vars_init = V()


# Set states and its bounds
for k in range(nk):
    vars_init["Xd",k,:,'w0']  = 0
    vars_init["Xd",k,:,'w']   =  l*sin(pi*k/float(nk))
    vars_init["Xd",k,:,'y']   = -l*cos(pi*k/float(nk))
    vars_init["Xd",k,:,'dy']  = l*sin(pi*k/float(nk))*pi/nk/h
    vars_init["Xd",k,:,'dw']  = l*cos(pi*k/float(nk))*pi/nk/h
    vars_init["Xd",k,:,'dw0'] = 0

for name in xd.keys():
    vars_lb["Xd",:,:,name]   = xD_min[name]
    vars_ub["Xd",:,:,name]   = xD_max[name]

vars_init["XA",:,:] = np.array([-sign(xDi_min['y'])*9.81*m])

# Set controls and its bounds
vars_init["U",:] = np.array([0])
vars_lb["U",:]   = np.array([-20])
vars_ub["U",:]   = np.array([20])

# State at initial time
for name in xd.keys():
    vars_lb["Xd",0,0,name] = xDi_min[name]
    vars_ub["Xd",0,0,name] = xDi_max[name]


# ALLOCATE NLP SOLVER

opts = {}
opts["expand"] = True
opts["max_iter"] = 1000
opts["linear_solver"] = 'ma57'
# opts['hessian_approximation'] = 'limited-memory'
# opts["linear_solver"] = 'ma27'

solver = NlpSolver("solver", "ipopt", nlp, opts)


# Initial condition
arg = {}
arg["x0"] = vars_init

# Bounds on x
arg["lbx"] = vars_lb
arg["ubx"] = vars_ub

# Bounds on g
arg["lbg"] = lbg
arg["ubg"] = ubg


# Solve the problem
res = solver(arg)


# Print the optimal cost
print "optimal cost: ", float(res["f"])