# 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
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)