# This code block installs ffmpeg to your Deepnote Machine # Allows Deepnote to write MP4 !apt update -y !apt install ffmpeg -y # Cart Pole Code: import numpy as np from numpy.linalg import solve import matplotlib.pyplot as plt from matplotlib.animation import FFMpegWriter from matplotlib.patches import Rectangle from matplotlib.patches import Circle # Define Parameters: g = 9.8 # gravity (m/s^2) L = 1 # Pendulum Length (m) k = 10 # Spring Constant (N/m) m1 = 1 # Pendulum mass (kg) m2 = 2 # Cart mass (kg) h = 1.5 # Pendulum height (m) l_o = 1 # Spring free length (m) # Define Time Vector for Simulation dt = 0.001 # Time Step sim_time = 10 # Length of Simulation t_vec = np.arange(0, sim_time, dt) # Initialize State Vectors: vec_size = len(t_vec) # We use this value alot so lets store it x_vec = np.zeros(vec_size) dx_vec = np.zeros(vec_size) theta_vec = np.zeros(vec_size) dtheta_vec = np.zeros(vec_size) # Pole End effector Location for Animation: x_pole = np.zeros(vec_size) y_pole = np.zeros(vec_size) # Assign Intial Values: x_vec[0] = 0 dx_vec[0] = 0 theta_vec[0] = 1 dtheta_vec[0] = 0 # Initial Pole End effector Location: x_pole[0] = L*np.sin(theta_vec[0]) y_pole[0] = h - L*np.cos(theta_vec[0]) # Euler Simulation: Using Matrix Form (A * x = B) # Initialize A and B: A = np.array([[L**2*m1, 0], [0, m2]]) B = np.array([(k*(2*L*np.sin(theta_vec[0])*(h-L*np.cos(theta_vec[0]))-2*L*np.cos(theta_vec[0])*(x_vec[0]-L*np.sin(theta_vec[0])))*(l_o-((h-L*np.cos(theta_vec[0]))**2+(x_vec[0]-L*np.sin(theta_vec[0]))**2)**(1/2)))/(2*((h-L*np.cos(theta_vec[0]))**2+(x_vec[0]-L*np.sin(theta_vec[0]))**2)**(1/2)) - m1*g*L*np.sin(theta_vec[0]), (k*(l_o-((h-L*np.cos(theta_vec[0]))**2+(x_vec[0]-L*np.sin(theta_vec[0]))**2)**(1/2))*(2*x_vec[0]-2*L*np.sin(theta_vec[0])))/(2*((h-L*np.cos(theta_vec[0]))**2+(x_vec[0]-L*np.sin(theta_vec[0]))**2)**(1/2))]) for i in range(1, vec_size): # B must be updated every iteration: B[0] = (k*(2*L*np.sin(theta_vec[i-1])*(h-L*np.cos(theta_vec[i-1]))-2*L*np.cos(theta_vec[i-1])*(x_vec[i-1]-L*np.sin(theta_vec[i-1])))*(l_o-((h-L*np.cos(theta_vec[i-1]))**2+(x_vec[i-1]-L*np.sin(theta_vec[i-1]))**2)**(1/2)))/(2*((h-L*np.cos(theta_vec[i-1]))**2+(x_vec[i-1]-L*np.sin(theta_vec[i-1]))**2)**(1/2)) - m1*g*L*np.sin(theta_vec[i-1]) B[1] = (k*(l_o-((h-L*np.cos(theta_vec[i-1]))**2+(x_vec[i-1]-L*np.sin(theta_vec[i-1]))**2)**(1/2))*(2*x_vec[i-1]-2*L*np.sin(theta_vec[i-1])))/(2*((h-L*np.cos(theta_vec[i-1]))**2+(x_vec[i-1]-L*np.sin(theta_vec[i-1]))**2)**(1/2)) [ddtheta, ddx] = np.linalg.solve(A, B) # Use ddx and ddtheta to solve: x_vec[i] = x_vec[i-1] + dx_vec[i-1]*dt theta_vec[i] = theta_vec[i-1] + dtheta_vec[i-1]*dt dx_vec[i] = dx_vec[i-1] + ddx*dt dtheta_vec[i] = dtheta_vec[i-1] + ddtheta*dt # Extra States for Animation: x_pole[i] = L*np.sin(theta_vec[i]) y_pole[i] = h - L*np.cos(theta_vec[i]) # Plotting x(t) and theta(t) plt.plot(t_vec, x_vec) plt.title("X Cart Position versus Time") plt.xlabel("Time [s]") plt.ylabel("X Position [m]") plt.show() plt.plot(t_vec, theta_vec) plt.title("Theta Pendulum Position versus Time") plt.xlabel("Time [s]") plt.ylabel("Theta [rad]") plt.show() # Create Animation: # Setup Figure: fig, ax = plt.subplots() p, = ax.plot([], [], color='black') min_lim = -3 max_lim = 3 ax.axis('equal') ax.set_xlim([min_lim, max_lim]) ax.set_ylim([min_lim, max_lim]) ax.set_xlabel('X') ax.set_ylabel('Y') ax.set_title('Pendulum with Spring and Cart Simulation:') title = "simulation" # Setup Animation Writer: FPS = 20 sample_rate = int(1/(dt*FPS)) # Real Time Playback dpi = 300 writerObj = FFMpegWriter(fps=FPS) # Initialize Patch: (Cart) width = 1 # Width of Cart height = width/2 # Height of Cart xy_cart = (x_vec[0]-width/2, -height/2) # Bottom Left Corner of Cart r = Rectangle(xy_cart, width, height, color='cornflowerblue') # Rectangle Patch ax.add_patch(r) # Add Patch to Plot # Initialize Patch: (Pendulum mass) c = Circle((x_pole[0], y_pole[0]), radius=0.1, color='red') ax.add_patch(c) # Draw the Ground: ground = ax.hlines(-height/2, min_lim, max_lim, colors='black') height_hatch = 0.25 width_hatch = max_lim - min_lim xy_hatch = (min_lim, -height/2-height_hatch) ground_hatch = Rectangle(xy_hatch, width_hatch, height_hatch, facecolor='None', linestyle='None', hatch='/') ax.add_patch(ground_hatch) # Animate: with writerObj.saving(fig, title + ".mp4", dpi): for i in range(0, vec_size, sample_rate): # Update Pendulum Arm and Spring: x_pole_arm = [0, x_pole[i], x_vec[i]] y_pole_arm = [h, y_pole[i], 0] p.set_data(x_pole_arm, y_pole_arm) # Update pendulum mass patch: patch_center = x_pole[i], y_pole[i] c.center = patch_center # Update Cart Patch: r.set(xy=(x_vec[i]-width/2, -height/2)) # Update Drawing: fig.canvas.draw() # Save Frame: writerObj.grab_frame()