multi-link pendulum - python simulation and animation of a multi-link pendulum

This jupyter notebook uses Lagrange formalism to derive the dynamics of a multi-link pendulum, simulates it and shows an animation of resulting solution. Unfortunately, it seems to be not very time-performant, so the best I could do in a reasonable amount of time was 6 links. That is why the dynamic equations are also exported to be simulated later in rust.

[1] 1994 - Murray, Li, Sastry: A Mathematical Introduction to Robotic Manipulation (https://www.ce.cit.tum.de/fileadmin/w00cgn/rm/pdf/murray-li-sastry-94-complete.pdf)

Notes:

• indexing starts with zero
• angular positions and velocities are defined absolutely (in general joint based measurements of a link orientation is relative to the previous link)

(Benjamin Jahn, TU Ilmenau, 2021)

Initilization

initialization: clear workspace, define latex use for display

# clear workspace
%reset -f

# latex math display
from IPython.display import Math, Latex

def latexDisplay(latex_obj): display(Math(latex_obj))
# def latexDisplay(latex_obj): display(Latex(r'$$' + (latex_obj) + r'$$'))

Kinematics and Dynamics

symbolic computation of the dynamics of a multi-link pendulum: derive dynamics using Lagrange equations of the second kind [1, p.168ff]

# import sympy for symbolic computations
from sympy import *

# NUMBER OF LINKS (PUT THE DESIRED NUMBER OF LINKS HERE)
N = 6

# generalized coordinates
q = Matrix([Symbol(r'q_' + str(i), real=True) for i in range(0,N)])
dq = Matrix([Symbol(r'\dot{q}_' + str(i), real=True) for i in range(0,N)])
ddq = Matrix([Symbol(r'\ddot{q}_' + str(i), real=True) for i in range(0,N)])

# time dependency and 
t = symbols(r't', real=True, positive=True); # time

# generalized coordinates as function of time
q_t = Matrix([Function(r'q_' + str(i), real=True)(t) for i in range(0,N)])
dq_t = Matrix([Function(r'\dot{q}_' + str(i), real=True)(t) for i in range(0,N)])
ddq_t = Matrix([Function(r'\ddot{q}_' + str(i), real=True)(t) for i in range(0,N)])

# dictionary for substitution
q_t2s = dict(zip(q_t,q)); q_s2t = dict(zip(q,q_t))
dq_t2s = dict(zip(dq_t,dq)); dq_s2t = dict(zip(dq,dq_t))
ddq_t2s = dict(zip(ddq_t,ddq)); ddq_s2t = dict(zip(ddq,ddq_t))

# substitute derivatives
q2dq = dict(zip([diff(q,t) for q in q_t], dq_t))
dq2ddq = dict(zip([diff(dq,t) for dq in dq_t], ddq_t))

# kinematic parameter
l = MatrixSymbol('l', N, 1); a = MatrixSymbol('a', N, 1)

# dynamic parameter
J = MatrixSymbol('J', N, 1); m = MatrixSymbol('m', N, 1); d = MatrixSymbol('d', N, 1)
grav = symbols('g', real=True, positive=True) # gravitational constant

# planar rotation
def rot(q): return Matrix([[cos(q), -sin(q)],[sin(q), cos(q)]])

# center of mass positions
p = zeros(2,N)
for i in range(0,N):
  for j in range(0, i):
    p[:,i] = p[:,i] + rot(q_t[j])*Matrix([0, l[j]])
  p[:,i] = p[:,i]  + rot(q_t[i])*Matrix([0, a[i]])
# latexDisplay(r'p =' + latex(p.subs(q_t2s)))

# center of mass velocities
dp = diff(p, t).subs(q2dq)
#latexDisplay(r'\dot{p} =' + latex(dp.subs(q_t2s).subs(dq_t2s)))

# kinetic and potential energy
T = sum([(Rational(1,2)*dp.col(i).T*m[i]*dp.col(i))[0] for i in range(0,N)]) + (Rational(1,2)*q.T*diag(*J)*q)[0] #; T = simplify(T)
V = sum([(m[i]*Matrix([0, grav]).T*p.col(i))[0] for i in range(0,N)]) #; V = simplify(V)
#latexDisplay(r'T =' + latex(T.subs(q_t2s).subs(dq_t2s)))
#latexDisplay(r'V =' + latex(V.subs(q_t2s).subs(dq_t2s)))

# Lagrangian --> ODE
L = T-V
Lode = diff(Matrix([L]).jacobian(dq_t),t) - Matrix([L]).jacobian(q_t); Lode = Lode.T
Lode = Lode.subs(q2dq).subs(dq2ddq)

# rephrase as robotic equation of motion
M = Lode.jacobian(ddq_t)
M = factor(M)
C = Rational(1, 2) * (Lode - M * ddq_t).jacobian(dq_t)
g = Lode - M * ddq_t - C * dq_t
g = simplify(g)
D = (diag(*d) + diag(*[*d[1:], 0])
    + Matrix(N, N, lambda i, j: -d[i] if i == j + 1 else 0) + Matrix(N, N, lambda i, j: -d[i + 1] if i + 1 == j else 0))
# ode = M.LUsolve(-C*dq_t-D*dq_t-g); ode = ode.subs(q_t2s).subs(dq_t2s)

Simulation

# numerics imports
import numpy as np
from math import ceil, floor
from scipy.integrate import solve_ivp

# numeric values for parameters
l_num = np.linspace(0.5, 0.35, N) # length of links [m]
a_num = l_num*1/2 # distance to center-of-mass for each link [m]
m_num = np.linspace(1,0.1, N)/N*1 # mass of each link [kg]
J_num = np.linspace(0.02, 0.01, N)/N*3 # moment-of-inertia for each link [kg m^2]
d_num = np.ones(N)*0.005 # viscous friction coefficient [Nm s]
grav_num = 9.81; # gravitational constant [m^3 (kg s^2)^-1]

substitutions = dict(zip([*l, *a, *m, *J, *d, grav],[*l_num, *a_num, *m_num, *J_num, *d_num, grav_num]))
#ode = ode.subs(substitutions)
M_num = M.subs(substitutions)
C_num = C.subs(substitutions)
g_num = g.subs(substitutions)
D_num = D.subs(substitutions)   

Code Generation for Rust (for faster simulation later)

from sympy.utilities.codegen import codegen
q_codegen = ["q_" + str(idx) for idx in range(0,N)]; dq_codegen = ["dq_" + str(idx) for idx in range(0,N)]
subs_codegen = dict(zip(q,q_codegen)) | dict(zip(dq,dq_codegen)) | q_t2s | dq_t2s

[(_, M_code)] = codegen([("M_flat", Matrix(flatten(M_num.subs(q_t2s).subs(dq_t2s).subs(subs_codegen))))], "Rust", header=False, empty=False)
C_code = M_code.replace("0, ", "0.0 ,"); C_code = C_code.replace("0]", "0.0]")
M_code_tmp = M_code.splitlines()
M_code_tmp[0] = "".join(["fn M_flat(q: [f64; " + str(N) + "]) -> " + " [f64; ", str(N**2), "] {"])
for idx in range(N-1,-1,-1): 
    M_code_tmp.insert(1, "    let q_" + str(idx) + " = " + "q[" + str(idx) + "]; " + "let q_" + str(idx) + "s = " + "q_" + str(idx) + ".sin(); " + "let q_" + str(idx) + "c = " + "q_" + str(idx) + ".cos(); ")
    M_code_tmp[-3] = M_code_tmp[-3].replace("q_" + str(idx) + ".sin()", "q_" + str(idx) + "s")
    M_code_tmp[-3] = M_code_tmp[-3].replace("q_" + str(idx) + ".cos()", "q_" + str(idx) + "c")
M_code = "\r\n".join(M_code_tmp)
# print(M_code)

[(_, C_code)] = codegen([("C_flat", Matrix(flatten(C_num.subs(q_t2s).subs(dq_t2s).subs(subs_codegen))))], "Rust", header=False, empty=False)
C_code = C_code.replace("0, ", "0.0 ,"); C_code = C_code.replace("0]", "0.0]")
C_code_tmp = C_code.splitlines()
C_code_tmp[0] = "".join(["fn C_flat(q: [f64; " + str(N) + "], dq: [f64; " + str(N) + "]) -> " + " [f64; ", str(N**2), "] {"])
for idx in range(N-1,-1,-1): 
    C_code_tmp.insert(1, "    let q_" + str(idx) + " = " + "q[" + str(idx) + "]; " + "let dq_" + str(idx) + " = " + "dq[" + str(idx) + "]; " + "let q_" + str(idx) + "s = " + "q_" + str(idx) + ".sin(); " + "let q_" + str(idx) + "c = " + "q_" + str(idx) + ".cos(); ")
    C_code_tmp[-3] = C_code_tmp[-3].replace("q_" + str(idx) + ".sin()", "q_" + str(idx) + "s")
    C_code_tmp[-3] = C_code_tmp[-3].replace("q_" + str(idx) + ".cos()", "q_" + str(idx) + "c")
C_code = "\r\n".join(C_code_tmp)
# print(C_code)

[(_, D_code)] = codegen([("D_flat", Matrix(flatten(D_num.subs(q_t2s).subs(dq_t2s).subs(subs_codegen))))], "Rust", header=False, empty=False)
D_code = D_code.replace("0, ", "0.0 ,"); D_code = D_code.replace("0]", "0.0]")
D_code_tmp = D_code.splitlines()
D_code_tmp[0] = "".join([D_code_tmp[0].split("->")[0], "->", " [f64; ", str(N**2), "] {"])
D_code = "\r\n".join(D_code_tmp)
# print(D_code)

[(_, g_code)] = codegen([("g_flat", Matrix(flatten(g_num.subs(q_t2s).subs(dq_t2s).subs(subs_codegen))))], "Rust", header=False, empty=False)
g_code = g_code.replace("0, ", "0.0 ,"); g_code = g_code.replace("0]", "0.0]")
g_code_tmp = g_code.splitlines()
g_code_tmp[0] = "".join(["fn g_flat(q: [f64; " + str(N) + "]) -> " + " [f64; ", str(N), "] {"])
for idx in range(N-1,-1,-1): 
    g_code_tmp.insert(1, "    let q_" + str(idx) + " = " + "q[" + str(idx) + "]; " + "let q_" + str(idx) + "s = " + "q_" + str(idx) + ".sin(); " + "let q_" + str(idx) + "c = " + "q_" + str(idx) + ".cos(); ")
    g_code_tmp[-3] = g_code_tmp[-3].replace("q_" + str(idx) + ".sin()", "q_" + str(idx) + "s")
    g_code_tmp[-3] = g_code_tmp[-3].replace("q_" + str(idx) + ".cos()", "q_" + str(idx) + "c")
g_code = "\r\n".join(g_code_tmp)
# print(g_code)

f = open("./rust-assisted/src/multi_link_pendulum_ode_" + str(N) + ".rs", "w")
preembel = "use ndarray::*;" + "\r\n" + "use ndarray_linalg::*;" "\r\n\r\n" + "const DOF:usize = " + str(N) + ";" + "\r\n" + "\r\n"
main_fct = "pub fn robotic_system(t:f64, x:Array1::<f64>) -> Array1::<f64> {\n"+"let mut q_vec = [0.0; DOF];\n"+"let mut dq_vec = [0.0; DOF];\n"+"for idx in 0..DOF {\n"+"q_vec[idx] = x[idx];\n"+"dq_vec[idx] = x[idx+DOF];\n"+"}\n"+"let q:Array1::<f64> = Array::from_vec(q_vec.to_vec());\n"+"let dq:Array1::<f64> = Array::from_vec(dq_vec.to_vec());\n"+"let n:usize = q.dim();\n"+"let dq_:Array1::<f64> = dq.clone();\n"+"let M: Array2::<f64> = Array::from_shape_vec((n,n), M_flat(q_vec).to_vec()).unwrap();\n"+"let C: Array2::<f64> = Array::from_shape_vec((n,n), C_flat(q_vec, dq_vec).to_vec()).unwrap();\n"+"let D: Array2::<f64> = Array::from_shape_vec((n,n), D_flat().to_vec()).unwrap();\n"+"let g: Array1::<f64> = Array::from_vec(g_flat(q_vec).to_vec());\n"+"let ddq_:Array1::<f64> = M.solve_into(-C.dot(&dq) - D.dot(&dq) - g).unwrap();\n"+"return concatenate(Axis(0), &[dq_.view(), ddq_.view()]).unwrap();\n"+"}"
result = preembel + M_code + "\r\n" + C_code  + "\r\n" + D_code + "\r\n" + g_code + "\r\n" + "\r\n" + main_fct
f.write(result)
f.close()

!!! Stop here if you only want to generate rust functions, as the simulations takes a long time for higher link numbers in python !!!

# prepared ode for solve_ivp
ode = M_num.LUsolve(-C_num*dq_t-D_num*dq_t-g_num); ode = ode.subs(q_t2s).subs(dq_t2s)
# ode = simplify(ode)
ode_sys = Matrix([*dq, *ode]) # express odes of order two as system of odes of order one
ode_f = lambdify([*q, *dq], [*ode_sys])

# execute simulation
T = 16; Ts = 1e-3
x0 = np.zeros(2*N); x0[-1] = 0.05
t_ = np.linspace(0,T, ceil(T/Ts+1), endpoint=True)

# solve ode without any input and nominal parameters
sol = solve_ivp(lambda t,x: ode_f(*x), [0,t_[-1]], x0, 'RK45', t_eval = t_, rtol=1e-5); 

Plots

# plot imports
import plotly.graph_objs as go
import plotly.io as pio
pio.templates.default = "none"
pio.renderers.default = "notebook_connected"

# define generic plot function for angular position of triple pendulum
from colour import Color

def plot_pendulum(sol):
    start_color = '#0000ff'; end_color = '#00ff80'
    colorscale = [x.hex for x in list(Color(start_color).range_to(Color(end_color), N))]

    fig = go.Figure()
    for idx in range(0, N):
        fig.add_trace(go.Scatter(x=sol.t, y=sol.y[idx, :] / np.pi, name=r"$q_" + str(idx) + "$", 
            line=dict(width=2, color=colorscale[idx])))

    fig.update_layout(height=600, width=900, margin=go.layout.Margin(t=20), 
        xaxis=dict(title=r"$\text{time in seconds}$", mirror=True, ticks='outside', showline=True),
        yaxis=dict(title=r"$\text{angular position in } \pi \text{ rad}$", 
            tickmode="linear", tick0=0.0, dtick=1, mirror=True, ticks='outside', showline=True))

    fig.show()

# plot angular positions
plot_pendulum(sol)

Animation

## import matplotlib for animations
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import animation, rc, colormaps
matplotlib.rcParams['animation.embed_limit'] = 2**128
from math import *
rc('animation', html='jshtml')

# define pendulum animation function
def animate_pendulum(sol, fps=30, colormap='winter'):

    # colors = cm.get_cmap(colormap, N)
    colors = colormaps[colormap].resampled(N)

    # frame number and index distance for given fps (tried to avoid interpolation to not slow down animation even further)
    T2anim = sol.t[-1]
    frameNumber = fps*T2anim
    didx = floor(len(sol.t)/frameNumber)

    # create and setup figure
    px = 1/plt.rcParams['figure.dpi']  # pixel in inches
    fig, ax = plt.subplots(figsize=(900*px,540*px))
    ax.set_aspect('equal')
    # plt.axis('equal'); 
    plt.close()
    l_sum = sum(l_num)
    ax.set_xlim((-l_sum*1.65, l_sum*1.65)); ax.set_ylim((-l_sum*1.35, l_sum*1.35))

    # create line objects
    line_base, = ax.plot([], [], lw=3, c=(0, 0, 0))
    lines = [None]*N
    for idx in range(0,N):
        lines[idx], = ax.plot([], [], lw=4, c=colors(idx/N))

    # initialization function: plot the background of each frame
    def init():
        # draw rail
        line_base.set_data(np.linspace(-l_sum*1.65, l_sum*1.65, 10), 0*np.linspace(-2, 2, 10)) 
        
        return (line_base,*lines, ) # return line objects

    # animation function: this is called sequentially
    def animate(i):

        # current configuration of pendulum
        p = np.zeros((2,N+1))
        for idx in range(1,N+1):
            p[:,idx] = p[:,idx-1] + np.array([-l_num[idx-1]*sin(sol.y[idx-1,i*didx]), l_num[idx-1]*cos(sol.y[idx-1,i*didx])])

        for idx in range(0,N):
            lines[idx].set_data([p[0,idx], p[0, idx+1]],[p[1,idx], p[1, idx+1]])

        return (*lines,) # return line objects

    # return animation object
    return animation.FuncAnimation(fig, animate, init_func=init, frames=floor(frameNumber), interval=1/fps*1000, blit=True)

anim = animate_pendulum(sol) # no control input --> cart fixed in position
anim 

Once Loop Reflect

Check out the simulation assisted by rust, for up to 32 links !!!: multi_link_pendulum_rust.ipynb