import numpy as np import matplotlib.pyplot as plt import sympy from matplotlib.animation import FFMpegWriter from matplotlib.patches import Circle !apt update -y !apt install ffmpeg -y

# Pendulum Parameters: m_1 = 1 # kg m_2 = 1 # kg m_3 = 1 # kg l_1 = 1 # m l_2 = 1 # m l_3 = 1 # m g = 9.8 # m/s^2 # Create Symbols for Time: t = sympy.Symbol('t') # Creates symbolic variable t tau1 = sympy.Symbol('tau1') tau2 = sympy.Symbol('tau2') tau3 = sympy.Symbol('tau3') # Create Generalized Coordinates as a function of time: q = [theta_1, theta_2, theta_3] th1 = sympy.Function('th1')(t) th2 = sympy.Function('th2')(t) th3 = sympy.Function('th3')(t) # Position Equation: r = [x, y] r1 = np.array([l_1 * sympy.sin(th1), -l_1 * sympy.cos(th1)]) # Position of first pendulum r2 = np.array([l_2 * sympy.sin(th2) + r1[0], -l_2 * sympy.cos(th2) + r1[1]]) # Position of second pendulum r3 = np.array([l_3 * sympy.sin(th3) + r2[0], -l_3 * sympy.cos(th3) + r2[1]]) # Position of third pendulum # Velocity Equation: d/dt(r) = [dx/dt, dy/dt] v1 = np.array([r1[0].diff(t), r1[1].diff(t)]) # Velocity of first pendulum v2 = np.array([r2[0].diff(t), r2[1].diff(t)]) # Velocity of second pendulum v3 = np.array([r3[0].diff(t), r3[1].diff(t)]) # Velocity of third pendulum # Energy Equations: T = (1/2 * m_1 * np.dot(v1, v1)) + (1/2 * m_2 * np.dot(v2, v2)) + (1/2 * m_3 * np.dot(v3, v3)) # Kinetic Energy V = m_1 * g * r1[1] + m_2 * g * r2[1] + m_3 * g * r3[1] # Potential Energy L = T - V # Lagrangian # Lagrange Terms: dL_dth1 = L.diff(th1) dL_dth1_dt = L.diff(th1.diff(t)).diff(t) dL_dth2 = L.diff(th2) dL_dth2_dt = L.diff(th2.diff(t)).diff(t) dL_dth3 = L.diff(th3) dL_dth3_dt = L.diff(th3.diff(t)).diff(t) # Euler-Lagrange Equations: dL/dq - d/dt(dL/ddq) = 0 th1_eqn = dL_dth1 - dL_dth1_dt - tau1 th2_eqn = dL_dth2 - dL_dth2_dt - tau2 th3_eqn = dL_dth3 - dL_dth3_dt - tau3 # Replace Time Derivatives and Functions with Symbolic Variables: ddth1 = sympy.Symbol('ddth1') ddth2 = sympy.Symbol('ddth2') ddth3 = sympy.Symbol('ddth3') replacements = [(th1.diff(t).diff(t), sympy.Symbol('ddth1')), (th1.diff(t), sympy.Symbol('dth1')), (th1, sympy.Symbol('th1')), (th2.diff(t).diff(t), sympy.Symbol('ddth2')), (th2.diff(t), sympy.Symbol('dth2')), (th2, sympy.Symbol('th2')), (th3.diff(t).diff(t), sympy.Symbol('ddth3')), (th3.diff(t), sympy.Symbol('dth3')), (th3, sympy.Symbol('th3'))] th1_eqn = th1_eqn.subs(replacements) th2_eqn = th2_eqn.subs(replacements) th3_eqn = th3_eqn.subs(replacements) r1 = r1[0].subs(replacements), r1[1].subs(replacements) r2 = r2[0].subs(replacements), r2[1].subs(replacements) r3 = r3[0].subs(replacements), r3[1].subs(replacements) # Simplfiy then Force SymPy to Cancel factorization: [Sometimes needed to use .coeff()] th1_eqn = sympy.simplify(th1_eqn) th2_eqn = sympy.simplify(th2_eqn) th3_eqn = sympy.simplify(th3_eqn) th1_eqn = th1_eqn.cancel() th2_eqn = th2_eqn.cancel() th3_eqn = th3_eqn.cancel() # Solve for control torques tau1 = sympy.solve(th1_eqn,tau1) tau2 = sympy.solve(th2_eqn,tau2) tau3 = sympy.solve(th3_eqn,tau3) # Solve for accelerations ddth1 = sympy.solve(th1_eqn,ddth1) ddth2 = sympy.solve(th2_eqn,ddth2) ddth3 = sympy.solve(th3_eqn,ddth3) # Generate Lambda Functions for Torque and Position Equations: replacements = (sympy.Symbol('th1'), sympy.Symbol('dth1'), sympy.Symbol('ddth1'), sympy.Symbol('th2'), sympy.Symbol('dth2'), sympy.Symbol('ddth2'), sympy.Symbol('th3'), sympy.Symbol('dth3'), sympy.Symbol('ddth3')) tau1 = sympy.utilities.lambdify(replacements, tau1, "numpy") tau2 = sympy.utilities.lambdify(replacements, tau2, "numpy") tau3 = sympy.utilities.lambdify(replacements, tau3, "numpy") replacements = (sympy.Symbol('th1'), sympy.Symbol('dth1'), sympy.Symbol('ddth1'), sympy.Symbol('tau1'), sympy.Symbol('th2'), sympy.Symbol('dth2'), sympy.Symbol('ddth2'), sympy.Symbol('tau2'), sympy.Symbol('th3'), sympy.Symbol('dth3'), sympy.Symbol('ddth3'), sympy.Symbol('tau3')) ddth1 = sympy.utilities.lambdify(replacements, ddth1, "numpy") ddth2 = sympy.utilities.lambdify(replacements, ddth2, "numpy") ddth3 = sympy.utilities.lambdify(replacements, ddth3, "numpy") replacements = (sympy.Symbol('th1'), sympy.Symbol('th2'), sympy.Symbol('th3')) r1 = sympy.utilities.lambdify(replacements, r1, "numpy") r2 = sympy.utilities.lambdify(replacements, r2, "numpy") r3 = sympy.utilities.lambdify(replacements, r3, "numpy")

# Simulate System: # Initial: th1, dth1, th2, dth2, th3, dth3 x0 = 1.0, 0.0, 0.0, 1.0, 0.0, 0.0 dt = 0.001 sim_time = 10 time = np.arange(0, sim_time, dt) sim_length = len(time) # Desired thetas th1_des = np.pi/2 th2_des = 4*np.pi/4 th3_des = np.pi/2 dth1_des = 0 dth2_des = 0 dth3_des = 0 # Control gains Kp = 10 Kd = 5 # Initialize Arrays: th1_vec = np.zeros(sim_length) dth1_vec = np.zeros(sim_length) th2_vec = np.zeros(sim_length) dth2_vec = np.zeros(sim_length) th3_vec = np.zeros(sim_length) dth3_vec = np.zeros(sim_length) x1_vec = np.zeros(sim_length) y1_vec = np.zeros(sim_length) x2_vec = np.zeros(sim_length) y2_vec = np.zeros(sim_length) x3_vec = np.zeros(sim_length) y3_vec = np.zeros(sim_length) # Evaluate Initial Conditions: th1_vec[0] = x0[0] dth1_vec[0] = x0[1] th2_vec[0] = x0[2] dth2_vec[0] = x0[3] th3_vec[0] = x0[4] dth3_vec[0] = x0[5] x1_vec[0], y1_vec[0] = r1(x0[0], x0[2], x0[4]) x2_vec[0], y2_vec[0] = r2(x0[0], x0[2], x0[4]) x3_vec[0], y3_vec[0] = r3(x0[0], x0[2], x0[4]) # Inverse Dynamics: for i in range(1, sim_length): # PD Control of accelerations ddtheta1 = Kp*(th1_des-th1_vec[i-1]) + Kd*(dth1_des-dth1_vec[i-1]) ddtheta2 = Kp*(th2_des-th2_vec[i-1]) + Kd*(dth2_des-dth2_vec[i-1]) ddtheta3 = Kp*(th3_des-th3_vec[i-1]) + Kd*(dth3_des-dth3_vec[i-1]) # Solving for torques t1 = tau1(th1_vec[i-1], dth1_vec[i-1], ddtheta1, th2_vec[i-1], dth2_vec[i-1], ddtheta2, th3_vec[i-1], dth3_vec[i-1], ddtheta3) t2 = tau2(th1_vec[i-1], dth1_vec[i-1], ddtheta1, th2_vec[i-1], dth2_vec[i-1], ddtheta2, th3_vec[i-1], dth3_vec[i-1], ddtheta3) t3 = tau3(th1_vec[i-1], dth1_vec[i-1], ddtheta1, th2_vec[i-1], dth2_vec[i-1], ddtheta2, th3_vec[i-1], dth3_vec[i-1], ddtheta3) # Update theta from Euler step th1_vec[i] = th1_vec[i-1] + dth1_vec[i-1]*dt th2_vec[i] = th2_vec[i-1] + dth2_vec[i-1]*dt th3_vec[i] = th3_vec[i-1] + dth3_vec[i-1]*dt # Informed update of velcoity value from EOM and tau accth1 = ddth1(th1_vec[i-1], dth1_vec[i-1], ddtheta1, t1[0], th2_vec[i-1], dth2_vec[i-1], ddtheta2, t2[0], th3_vec[i-1], dth3_vec[i-1], ddtheta3, t3[0]) accth2 = ddth2(th1_vec[i-1], dth1_vec[i-1], ddtheta1, t1[0], th2_vec[i-1], dth2_vec[i-1], ddtheta2, t2[0], th3_vec[i-1], dth3_vec[i-1], ddtheta3, t3[0]) accth3 = ddth3(th1_vec[i-1], dth1_vec[i-1], ddtheta1, t1[0], th2_vec[i-1], dth2_vec[i-1], ddtheta2, t2[0], th3_vec[i-1], dth3_vec[i-1], ddtheta3, t3[0]) dth1_vec[i] = dth1_vec[i-1] + accth1[0]*dt dth2_vec[i] = dth2_vec[i-1] + accth2[0]*dt dth3_vec[i] = dth3_vec[i-1] + accth3[0]*dt # Animation States: x1_vec[i], y1_vec[i] = r1(th1_vec[i-1], th2_vec[i-1], th3_vec[i-1]) x2_vec[i], y2_vec[i] = r2(th1_vec[i-1], th2_vec[i-1], th3_vec[i-1]) x3_vec[i], y3_vec[i] = r3(th1_vec[i-1], th2_vec[i-1], th3_vec[i-1])

# Create Animation: # Setup Figure: Initialize Figure / Axe Handles fig, ax = plt.subplots() p1, = ax.plot([], [], color='black', linewidth=2) p2, = ax.plot([], [], color='black', linewidth=2) p3, = ax.plot([], [], color='black', linewidth=2) lb, ub = -5, 5 ax.axis('equal') ax.set_xlim([lb, ub]) ax.set_xlabel('X') # X Label ax.set_ylabel('Y') # Y Label ax.set_title('Triple Pendulum Simulation:') video_title = "simulation" # Setup Animation Writer: FPS = 20 sample_rate = int(1 / (dt * FPS)) dpi = 300 writerObj = FFMpegWriter(fps=FPS) # Initialize Patch: Pendulum 1, 2, and 3 pendulum_1 = Circle((0, 0), radius=0.1, color='cornflowerblue', zorder=10) pendulum_2 = Circle((0, 0), radius=0.1, color='cornflowerblue', zorder=10) pendulum_3 = Circle((0, 0), radius=0.1, color='cornflowerblue', zorder=10) ax.add_patch(pendulum_1) ax.add_patch(pendulum_2) ax.add_patch(pendulum_3) # Draw Static Objects: pin_joint = Circle((0, 0), radius=0.05, color='black', zorder=10) ax.add_patch(pin_joint) # Plot and Create Animation: with writerObj.saving(fig, video_title+".mp4", dpi): for i in range(0, sim_length, sample_rate): # Draw Pendulum Arm: x_pendulum_arm = [0, x1_vec[i], x2_vec[i], x3_vec[i]] y_pendulum_arm = [0, y1_vec[i], y2_vec[i], y3_vec[i]] p1.set_data(x_pendulum_arm, y_pendulum_arm) # Update Pendulum Patches: pendulum_1_center = x1_vec[i], y1_vec[i] pendulum_2_center = x2_vec[i], y2_vec[i] pendulum_3_center = x3_vec[i], y3_vec[i] pendulum_1.center = pendulum_1_center pendulum_2.center = pendulum_2_center pendulum_3.center = pendulum_3_center # Update Drawing: fig.canvas.draw() # Grab and Save Frame: writerObj.grab_frame() from IPython.display import Video Video("/work/"+video_title+".mp4", embed=True, width=640, height=480)