# Start writing code here... !apt update -y !apt install ffmpeg -y import math 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 #Parameters: m1 = 1 # kg - Mass of Pendulu m2 = 2 # kg - Mass of Block L = 1 # m - Lenght of Pendulum arm k = 10 # N/m - Spring stiffness h = 1.5 # m - Height from center of block to pendulum hinge L0 = 1 # m - Rest length of spring g = 9.8 # m/s^2 # 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 # math.radians(57.2958) dtheta_vec[0] = 0 # Initial Pole End effector Location: #y_offset = 0 # The Y Location of where the Cart and Pole are connected x_pole[0] = L * np.sin(theta_vec[0]) y_pole[0] = -L * np.cos(theta_vec[0]) # Euler Simulation: Using Matrix Form (A * x = B) # Initialize A and B: A = np.array([[m2, 0], [0, m1 * L**2]]) B = np.array([[0], [0]]) for i in range(1, vec_size): # Only the B matrix needs to be Updated: B[0] = -k * (L*np.sin(theta_vec[i - 1]) + x_vec[i - 1]) B[1] = -m1*g*L*np.sin(theta_vec[i-1]) - (1/2)*k*( -2*L*np.cos(theta_vec[i-1])*(-L*np.sin(theta_vec[i-1] + x_vec[i-1])) + 2*L*np.sin(theta_vec[i-1])*(h - L*np.cos(theta_vec[i-1]))) [ddx, ddtheta] = np.linalg.solve(A, B) # Use ddx and ddtheta to solve: x_vec[i] = x_vec[i - 1] + dx_vec[i - 1] * dt dx_vec[i] = dx_vec[i - 1] + ddx * dt theta_vec[i] = theta_vec[i - 1] + dtheta_vec[i - 1] * 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] = -L * np.cos(theta_vec[i]) #print(x_vec[i]) #print(theta_vec[i]) #print(ddx) plt.plot(t_vec, x_vec) plt.title("x position vs. time") plt.show() plt.plot(t_vec, theta_vec) plt.title("theta vs. time") plt.show() # Create Animation: # Setup Figure: fig, ax = plt.subplots() p, = ax.plot([], [], color='royalblue') min_lim = -2 max_lim = 2 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('Cartpole Simulation:') title = "simulation" # Setup Animation Writer: FPS = 20 sample_rate = int(1 / (dt * FPS)) # Real Time Playback dpi = 300 writerObj = FFMpegWriter(fps=FPS) # Initialize Pendulum bob c = Circle((0, 0), radius=0.1, color='cornflowerblue') # Draws a circle of radius 0.1 at (X, Y) location (0, 0) ax.add_patch(c) # Add the patch to the axes # We need to calculate the cartesian location of the pendulum: # Initialize Array: simulation_size = len(t_vec) # We use this alot, so lets store it. x_pendulum_arm = np.zeros(simulation_size) y_pendulum_arm = np.zeros(simulation_size) origin = (0, 0) # Pin Joint (X, Y) Location # Find the X and Y trajectories of the pendulum: for i in range(0, simulation_size): x_pendulum_arm[i] = L * np.sin(theta_vec[i]) + origin[0] y_pendulum_arm[i] = origin[1] - L * np.cos(theta_vec[i]) # Initialize Patch: (Block) width = 0.40 # Width of Cart height = width / 2 # Height of Cart xy_cart = (x_vec[0] - width / 2, -h - height / 2) # Bottom Left Corner of Cart block = Rectangle(xy_cart, width, height, color='green') # Rectangle Patch ax.add_patch(block) # Add Patch to Plot # Initialize Patch: (Spring) xspring = xy_cart[0] + width / 2, x_pendulum_arm[0] yspring = xy_cart[1] + height / 2, y_pendulum_arm[0] spring = plt.Line2D(xspring, yspring, lw = 1) ax.add_patch(spring) # Draw the Ground: ground = ax.hlines(-height / 2 - h, min_lim, max_lim, colors='black') height_hatch = 0.25 width_hatch = max_lim - min_lim xy_hatch = (min_lim, -height / 2 - height_hatch - h) ground_hatch = Rectangle(xy_hatch, width_hatch, height_hatch, facecolor='None', linestyle='None', hatch='/') ax.add_patch(ground_hatch) # Plot and Create Animation: with writerObj.saving(fig, title +".mp4", dpi): # We want to create video that represents the simulation # So we need to sample only a few of the frames: for i in range(0, simulation_size, sample_rate): # Update Pendulum Arm: x_data_points = [origin[0], x_pendulum_arm[i]] y_data_points = [origin[1], y_pendulum_arm[i]] # We want to avoid creating new plots to make an animation (Very Slow) # Instead lets take the plot we made earlier and just update it with new data. p.set_data(x_data_points, y_data_points) # Update plot with set_data # Update Pendulum Patch: patch_center = x_pendulum_arm[i], y_pendulum_arm[i] c.center = patch_center # Same idea here, instead of drawing a new patch update its location # Update Spring Patch: xspring = x_vec[i] + width/2, x_pendulum_arm[i] yspring = xy_cart[1] + height / 2, y_pendulum_arm[i] spring.set_data(xspring,yspring) #spring.set_data() # Update Cart Patch: block.set(xy=(x_vec[i] , -h - height / 2 )) # Update Drawing: fig.canvas.draw() # Update the figure with the new changes # Grab and Save Frame: writerObj.grab_frame() from IPython.display import Video Video("/work/simulation.mp4", embed=True, width=640, height=480)