# 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()