# Import Libraries:
!apt update -y
!apt install ffmpeg -y
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter
from matplotlib.patches import Circle
# Pendulum Parameters
g = 9.8 # m/s^2
m1 = 1 # kg
m2 = 2 # kg
dt = 0.001 # time increments
k = 10 # N/m
h = 1.5 # m
L = 1 # m
t_vec = np.arange(0,10,dt) # time vector
# Initialize a vector of zeros
x_vec = np.zeros(len(t_vec))
theta_vec = np.zeros(len(t_vec))
dx_vec = np.zeros(len(t_vec))
dtheta_vec = np.zeros(len(t_vec))
ddx_vec = np.zeros(len(t_vec))
ddtheta_vec = np.zeros(len(t_vec))
distance = np.zeros(len(t_vec))
# Set our initial condition
x_vec[0] = 0 # initial x position
theta_vec[0] = 1 # initial y postion
dx_vec[0] = 0 # initial x velocity
dtheta_vec[0] = 0 # initial y velocity
# Euler's Method (approximately integrates the differential equation)
for i in range(1, len(t_vec)):
# X direction
x_vec[i] = x_vec[i-1] + dx_vec[i-1]*dt
ddx_vec[i-1] = (k*(np.sqrt((L*np.sin(theta_vec[i-1]) - x_vec[i-1])**2
+ (h - L*np.cos(theta_vec[i-1]))**2) - L)*(L*np.sin(theta_vec[i-1])
- x_vec[i-1]))/(m2*np.sqrt((L*np.sin(theta_vec[i-1]) - x_vec[i-1])**2
+ (h - L*np.cos(theta_vec[i-1]))**2))
dx_vec[i] = dx_vec[i-1] + ddx_vec[i-1]*dt
# Y direction
theta_vec[i] = theta_vec[i-1] + dtheta_vec[i-1]*dt
ddtheta_vec[i-1] = -((g*np.sin(theta_vec[i-1]))/(L))
- (k*(np.sqrt((L*np.sin(theta_vec[i-1])
- x_vec[i-1])**2 + (h - L*np.cos(theta_vec[i-1]))**2) - L)
*((L*np.sin(theta_vec[i-1])
- x_vec[i-1])*(L*np.cos(theta_vec[i-1]))
+ (h - L*np.cos(theta_vec[i-1]))
*(L*np.sin(theta_vec[i-1]))))
/(m1*L**2*np.sqrt((L*np.sin(theta_vec[i-1])
- x_vec[i-1])**2 + (h - L*np.cos(theta_vec[i-1]))**2))
dtheta_vec[i] = dtheta_vec[i-1] + ddtheta_vec[i-1]*dt
distance[i] = -h
plt.plot(t_vec,x_vec)
plt.title('X Positon')
plt.xlabel('Time')
plt.ylabel('X')
plt.show()
plt.plot(t_vec,theta_vec)
plt.title('Theta Angle')
plt.xlabel('Time')
plt.ylabel('Theta')
plt.show()
# Setup Figure: Initialize Figure / Axe Handles
fig, ax = plt.subplots() # Initialize Figure and Axes
p, = ax.plot([], [], color='black') # Initialize Empty Plot
q, = ax.plot([], [], color='black')
ax.axis('equal')
ax.set_xlim([-3, 3]) # X Lim
ax.set_ylim([-3, 3]) # Y Lim
ax.set_xlabel('X') # X Label
ax.set_ylabel('Y') # Y Label
ax.set_title('Penedulum Simulation:') # Plot Title
video_title = "simulation2" # name of the mp4
# Initialize Patch:
c = Circle((0, -0.65), radius=0.05, color='red')
ax.add_patch(c) # Add the patch to the axes
b = Circle((0, -1.5),radius = 0.1, color = 'red')
ax.add_patch(b)
# Setup Animation Writer:
FPS = 20 # Frames per second for the video
sample_rate = int(1 / (dt * FPS)) # Rate at which we sample the simulation to visualize
dpi = 300 # Quality of the Video
writerObj = FFMpegWriter(fps=FPS) # Video Object
# 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])
# Plot and Create Animation:
with writerObj.saving(fig, video_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]]
x1_data_points = [x_vec[i], x_pendulum_arm[i]]
y1_data_points = [distance[i], 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
q.set_data(x1_data_points, y1_data_points)
# 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
red_center = x_vec[i], distance[i]
b.center = red_center
# 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/simulation2.mp4", embed=True, width=640, height=480)