Dynamics of an N-pendulum

In this notebook we show how to use Casadi to obtain the equations of motion of an N-pendulum.

In [1]:
from casadi import *
import numpy as np
import matplotlib.pyplot as plt
import time
from os import system
%matplotlib inline
In [2]:
def generate_n_pendulum_rhs(N=2):
    """ Generates rhs for the Hamiltonian dynamics of an N-pendulum
    In:
        N (positive int): number of links in the N-pendulum
        
    Out:
        XH (Casadi Function) : Function to compute the Hamiltonian vector field,
        takes as input 2 arguments z=(q,p) a vector of dimension 2N 
        and a vector params of dimension 2N+1 containing 
        the N lengths, the N masses and the value of the gravitational constant g.
        
        hamiltonian_fct (Casadi Function): Computes the total energy, takes as 
        input [q,p,params] and outputs the scalar Hamiltonian H.
        
        position_fct (Casadi Function): Input [q, l] where l is the vector of lengths
        and Outputs x the 2N+2 dimensional vector with the positions x0,y0,...,xN,yN
    """
    q = SX.sym("q", N) #Position
    p = SX.sym("p",N) #Momentum
    l = SX.sym("l", N) #link lengths
    m = SX.sym("m", N) #masses
    g = SX.sym("g", 1) #gravitational constant
    params = vertcat(l,m,g)
    M=[]
    for k in range(N):
        M.append(m[k])
        M.append(m[k])
    M = diag(vertcat(*M))
    x0 = 0
    y0 = 0
    x = []
    for k in range(N):
        x0 += l[k]*sin(q[k])
        y0 += -l[k]*cos(q[k])
        x.append(x0)
        x.append(y0)
    x = vertcat(*x) # 
    position_fct = Function("pos", [q, l], [x])
    # Assemble Potential energy V(q)
    V = 0
    for k in range(N):
        V+=m[k]*x[2*k+1,0]
    V*=g
    # Kinetic energy K(q,p) = 1/2 p^T M(q)^{-1} p :
    J = jacobian(x, q)
    mass_mat = mtimes([J.T, M, J])
    m_inv_p = solve(mass_mat, p)
    K = 0.5*dot(p, m_inv_p)
    # Hamiltonian:
    H = V+K
    hamiltonian_fct =  Function("H", [q,p,params], [H])
    # Dynamics rhs
    rhs = vertcat(gradient(H,p), -gradient(H,q))
    z = vertcat(q,p)
    XH = Function("X_H", [z,params], [rhs])
    return XH, hamiltonian_fct, position_fct
In [6]:
#Define the number of links:
N = 6
#Get the functions for the right hand side, energy and positions
XH, hamiltonian_fct, position_fct = generate_n_pendulum_rhs(N)
rk4_step = simpleRK(XH, 1)
In [9]:
#Simulation:
def simulate_rk4(z0, param_vals, time_step, final_time):
    n_steps = int(final_time/time_step)
    print "Total number of simulation steps:",n_steps
    zs = np.empty((2*N,n_steps))
    zs[:,0] = z0.full().T
    for k in range(1, n_steps):
        z0 = rk4_step(z0,param_vals, time_step)
        zs[:,k] = z0.full().T
    times = np.arange(n_steps)*time_step
    return zs, times

#Physical parameters:
m_vals = vertcat(DM.ones(N-1), 1.) 
l_vals = DM.ones(N)
g_val = DM([9.81])
param_vals = vertcat(l_vals, m_vals, g_val)
#Initial conditions
#z0 = vertcat(2*np.pi*DM(np.random.rand(N)), DM.zeros(N))
z0 = vertcat(3*np.pi/4. * DM.ones(N), DM.zeros(N))
h = 0.01
final_time = 10.

#Simulate
zs,times = simulate_rk4(z0, param_vals, h, final_time)

#Plot results:
for k in range(N):
    plt.plot(times,zs[k,:], label = r'$\theta_%d$'%(k))
plt.legend()
plt.xlabel('time(s)')
plt.show()
Total number of simulation steps: 1000
In [10]:
import matplotlib.animation as animation
from IPython.display import HTML
endpt_x = []
endpt_y = []

fig = plt.figure()
lim = np.sum(l_vals.full())
ax = fig.add_subplot(111, autoscale_on=False, xlim=(-lim, lim), ylim=(-lim, lim))
ax.grid()

line, = ax.plot([], [], 'o-', lw=2)
traj, = ax.plot([],[], lw=1)
time_template = 'time = %.1fs'
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)


def init():
    line.set_data([], [])
    traj.set_data([], [])
    time_text.set_text('')
    return line,traj, time_text


def animate(i):
    global endpt_x, endpt_y
    xs = position_fct(zs[0:N,i], l_vals).full()[:,0]
    thisx = [0]
    thisy = [0]
    for k in range(N):
        thisx.append(xs[2*k])
        thisy.append(xs[2*k+1])
    endpt_x.append(thisx[-1])
    endpt_y.append(thisy[-1])
    line.set_data(thisx, thisy)
    traj.set_data(endpt_x, endpt_y)
    time_text.set_text(time_template % (times[i]))
    return line,traj, time_text

ani = animation.FuncAnimation(fig, animate, np.arange(1, len(times),1),
                              interval=10, blit=True, init_func=init)

HTML(ani.to_html5_video())
#ani.save('double_pendulum.mp4', fps=15)
#plt.show()
Out[10]:
In [ ]:
#Generate C code: Not able to generate c code for the linear solve in the expression
# of the hamiltonian unfortunately...
def gen_code_and_compile(fct, name):
    objname = name + ".so"
    fct.generate(name+".c")
    print "Compiling with O3 optimization: ", objname
    print "gcc -fPIC -shared -O3 " + name + ".c -o " + objname
    t1 = time.time()
    system("gcc -fPIC -shared -O3 " + name + ".c -o " + objname)
    t2 = time.time()
    print "time = ", (t2-t1), " s"
    return external(name, "./"+objname)
#name = "mass_mat"
#rhs_compiled = gen_code_and_compile(XH, name)