# Install ffmpeg (allows us to create videos) !apt update -y !apt install ffmpeg -y import numpy as np import matplotlib.pyplot as plt import sympy from matplotlib.animation import FFMpegWriter from matplotlib.patches import Circle, Rectangle

# Parameters L = 1 # length of pendulum (m) m_1 = 1 # mass of pendulum (kg) k = 10 # spring stiffness (N/m) m_2 = 2 # mass of block (kg) h = 1.5 # height of pendulum hinge (m) r_0 = 1 # rest length of spring (m) g = 9.8 # gravitational constant (m/s^2) # Create Symbols for Time: t = sympy.Symbol('t') # Creates symbolic variable t # Create Generalized Coordinates as a function of time: q = [theta, x] th = sympy.Function('th')(t) x = sympy.Function('x')(t) # Position Equation: r = [x, y] r1 = np.array([L * sympy.sin(th), -L * sympy.cos(th)]) # Position of pendulum r2 = np.array([x,-h]) # Position of block # Velocity Equation: d/dt(r) = [dx/dt, dy/dt] v1 = np.array([r1[0].diff(t), r1[1].diff(t)]) # Velocity of pendulum v2 = np.array([r2[0].diff(t), 0]) # Velocity of block # Length of spring squared r_squared = (r1[0]-r2[0])**2 + (r1[1]-r2[1])**2 r = r_squared**0.5 r0 = 1 # Energy Equations: T = 1/2 * m_1 * np.dot(v1, v1) + 1/2 * m_2 * np.dot(v2, v2) # Kinetic Energy # V = m_1 * g * r1[1] + m_2 * g * r2[1] + 1/2 * k * r_squared# Potential Energy V = m_1 * g * r1[1] + m_2 * g * r2[1] + 1/2 * k * (r0-r)**2# Potential Energy L = T - V # Lagrangian # Lagrange Terms: dL_dth = L.diff(th) dL_dth_dt = L.diff(th.diff(t)).diff(t) dL_dx = L.diff(x) dL_dx_dt = L.diff(x.diff(t)).diff(t) # Euler-Lagrange Equations: dL/dq - d/dt(dL/ddq) = 0 th_eqn = dL_dth - dL_dth_dt x_eqn = dL_dx - dL_dx_dt replacements = [(th.diff(t).diff(t), sympy.Symbol('ddth')), (th.diff(t), sympy.Symbol('dth')), (th, sympy.Symbol('th')), (x.diff(t).diff(t), sympy.Symbol('ddx')), (x.diff(t), sympy.Symbol('dx')), (x, sympy.Symbol('x'))] th_eqn = th_eqn.subs(replacements) x_eqn = x_eqn.subs(replacements) print(th_eqn) print(x_eqn) r1 = r1[0].subs(replacements), r1[1].subs(replacements) r2 = r2[0].subs(replacements), -h # Solve for Coefficients for A * x = B where x = [ddth1 ddth2] A1 = th_eqn.coeff(sympy.Symbol('ddth')) A2 = th_eqn.coeff(sympy.Symbol('ddx')) A3 = x_eqn.coeff(sympy.Symbol('ddth')) A4 = x_eqn.coeff(sympy.Symbol('ddx')) print(A1) # Multiply remaining terms by -1 to switch to other side of equation: A * x - B = 0 -> A * x = B remainder = [(sympy.Symbol('ddth'), 0), (sympy.Symbol('ddx'), 0)] B1 = -1 * th_eqn.subs(remainder) B2 = -1 * x_eqn.subs(remainder) # Generate Lambda Functions for A and B and Position Equations: replacements = (sympy.Symbol('th'), sympy.Symbol('dth'), sympy.Symbol('x'), sympy.Symbol('dx')) # A1 = sympy.utilities.lambdify(replacements, A1, "numpy") # A2 = sympy.utilities.lambdify(replacements, A2, "numpy") # A3 = sympy.utilities.lambdify(replacements, A3, "numpy") # A4 = sympy.utilities.lambdify(replacements, A4, "numpy") # B1 = sympy.utilities.lambdify(replacements, B1, "numpy") # B2 = sympy.utilities.lambdify(replacements, B2, "numpy") r1 = sympy.utilities.lambdify(replacements, r1, "numpy") r2 = sympy.utilities.lambdify(replacements, r2, "numpy") ddth = B1/A1 ddx = B2/A4 ddth_func = sympy.utilities.lambdify(replacements, ddth, "numpy") ddx_func = sympy.utilities.lambdify(replacements, ddx, "numpy")

# Simulate System: x0 = 1, 0, 0, 0 # th, dth, x, dx dt = 0.001 sim_time = 10 time = np.arange(0, sim_time, dt) sim_length = len(time) # Initialize Arrays: th_vec = np.zeros(sim_length) dth_vec = np.zeros(sim_length) x_vec = np.zeros(sim_length) dx_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) # Evaluate Initial Conditions: th_vec[0] = x0[0] dth_vec[0] = x0[1] x_vec[0] = x0[2] dx_vec[0] = x0[3] x1_vec[0], y1_vec[0] = r1(x0[0], x0[1], x0[2], x0[3]) x2_vec[0], y2_vec[0] = r2(x0[0], x0[1], x0[2], x0[3]) # Initialize A and B: A = np.array([[0, 0], [0, 0]]) B = np.array([0, 0]) # Euler Integration: for i in range(1, sim_length): # Evaluate Dynamics: ddth = ddth_func(th_vec[i-1], dth_vec[i-1], x_vec[i-1], dx_vec[i-1]) ddx = ddx_func(th_vec[i-1], dth_vec[i-1], x_vec[i-1], dx_vec[i-1]) # Euler Step Integration: th_vec[i] = th_vec[i-1] + dth_vec[i-1] * dt dth_vec[i] = dth_vec[i-1] + ddth * dt x_vec[i] = x_vec[i-1] + dx_vec[i-1] * dt dx_vec[i] = dx_vec[i-1] + ddx * dt # Animation States: x1_vec[i], y1_vec[i] = r1(th_vec[i-1], dth_vec[i-1], x_vec[i-1], dx_vec[i-1]) x2_vec[i], y2_vec[i] = r2(th_vec[i-1], dth_vec[i-1], x_vec[i-1], dx_vec[i-1]) plt.plot(time,th_vec) plt.title("Theta over time") plt.show() plt.plot(time,x_vec) plt.title("X over time") plt.show()

# 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) lb, ub = -2.5, 2.5 ax.axis('equal') ax.set_xlim([lb, ub]) ax.set_xlabel('X') # X Label ax.set_ylabel('Y') # Y Label ax.set_title('Pendulum Block 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 and 2 pendulum = Circle((0, 0), radius=0.1, color='cornflowerblue', zorder=10) width = 0.25 height = 0.15 block = Rectangle((0,0), width, height, color='cornflowerblue', zorder=10) ax.add_patch(pendulum) ax.add_patch(block) # 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]] y_pendulum_arm = [0, y1_vec[i], y2_vec[i]] p1.set_data(x_pendulum_arm, y_pendulum_arm) # Update Pendulum Patches: pendulum_center = x1_vec[i], y1_vec[i] block_center = x2_vec[i], y2_vec[i] pendulum.center = pendulum_center block.set(xy=(x_vec[i] - width / 2, -h-height/2)) # Update Drawing: fig.canvas.draw() # Grab and Save Frame: writerObj.grab_frame()