!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
# Pendulum Parameters
g = 9.8
m = 1
# 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)
y_vec = np.zeros(vec_size)
dy_vec = np.zeros(vec_size)
l_vec = np.zeros(vec_size)
x_vec[0] = -1.0
dx_vec[0] = 0.0
y_vec[0] = 1.0
dy_vec[0] = 0.0
l_vec[0] = 0.0
# Euler Simulation: Using Matrix Form (A * x = B)
# Initialize A and B:
for i in range(1, vec_size):
A = np.array([[m, 0.0, -2*x_vec[i-1]], [0.0, m, 1.0], [-x_vec[i-1], 0.5, 0.0]])
B = np.array([0.0, -m*g, dx_vec[i-1]**2])
[ddx, ddy, l] = np.linalg.solve(A, B)
# Use ddx and ddtheta to solve:
x_vec[i] = x_vec[i - 1] + dx_vec[i - 1] * dt
y_vec[i] = y_vec[i - 1] + dy_vec[i - 1] * dt
dx_vec[i] = dx_vec[i - 1] + ddx * dt
dy_vec[i] = dy_vec[i - 1] + ddy * dt
l_vec[i] = l
plt.xlabel('t(sec)')
plt.ylabel('x(m)')
plt.title('x(t)')
plt.plot(t_vec,x_vec)
plt.show()
plt.xlabel('t(sec)')
plt.ylabel('y(m)')
plt.title('y(t)')
plt.plot(t_vec,y_vec)
plt.show()
plt.xlabel('t(sec)')
plt.ylabel(r'$\lambda$(N)')
plt.title(r'$\lambda$(t)')
plt.plot(t_vec,l_vec)
plt.show()
# set up our figure for the drawing of the pendulum
fig, ax = plt.subplots()
# create a plot on those axes which is currently empty
p, = ax.plot([], [], color='cornflowerblue') # Initialize Empty Plot
ax.set_xlim([-3,3]) # x lim
ax.set_ylim([-10,10]) # y lim
ax.axis('equal')
ax.set_xlabel('X') # x label
ax.set_ylabel('Y') # y label
ax.set_title('problem 1 - bead on a parabola') # Plot Title
video_title = "simulation" # name of the mp4
# now we want to add stuff on the plot
c = Circle((0,0), radius=0.1, color = 'cornflowerblue')
ax.add_patch(c)
# define information of our animation
FPS = 20
sample_rate = int(1/(FPS*dt))
dpi = 300 # quality of the video
writerObj = FFMpegWriter(fps = FPS)
# now we are putting the rod and bob at the right place
# initialize an array containing the position of pendulum over time
simulation_size = len(t_vec)
x_pendulum_arm = np.zeros(simulation_size)
y_pendulum_arm = np.zeros(simulation_size)
for i in range(0,simulation_size):
x_pendulum_arm[i] = x_vec[i]
y_pendulum_arm[i] = y_vec[i]
# we have computed all the pendulum positions
xx = np.linspace(-1.5, 1.5, 100)
yy = xx**2
ax.plot(xx, yy)
# now we want to update the plot and save frames and create the video
# Plot and Create Animation:
with writerObj.saving(fig, video_title+".mp4", dpi):
# Loop thru the frames we want to animate
for i in range(0, simulation_size, sample_rate):
# Update Pendulum Arm(rod)
#x_data_points = [0, x_pendulum_arm[i]]
#y_data_points = [0, y_pendulum_arm[i]]
# update the data in plotting object
#p.set_data(x_data_points, y_data_points)
# Update Pendulum bob:
patch_center = x_pendulum_arm[i], y_pendulum_arm[i]
# note: commas without brackets create list
c.center = patch_center
# Update Drawing:
fig.canvas.draw() # Update the figure with the new changes
# Grab and Save Frame:
writerObj.grab_frame()
# We can also run it inline with the following block of code:
from IPython.display import Video
Video("/work/simulation.mp4", embed=True, width=640, height=480)
# You can find the video's path by right clicking on it in Notebooks & Files
# and selecting "Copy path to clipboard"