!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
# Define Pendulum Parameters:
g = 9.8 # gravity (m/s^2)
L = 1 # Length (m)
m1 = 1 # blob (kg)
m2 = 2 # block (kg)
k = 10
# 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)
x_vec[0] = 0.0
dx_vec[0] = 0.0
theta_vec[0] = 1.0
dtheta_vec[0] = 0.0
# Euler Simulation: Using Matrix Form (A * x = B)
# Initialize A and B:
A = np.array([[m2, 0.0], [0, m1 * L ** 2]])
B = np.array([0.0, 0.0])
for i in range(1, vec_size):
# b must be updated every iteration:
sq = (
((x_vec[i - 1]) - (L * np.sin(theta_vec[i - 1]))) ** 2
+ (((1.5) - (L * np.cos(theta_vec[i - 1]))) ** 2)
) ** (-0.5)
B[0] = (-k) * (x_vec[i - 1] - (L * np.sin(theta_vec[i - 1]))) * (1 - sq)
B[1] = (-m1 * g * L * np.sin(theta_vec[i - 1])) - (
-(k * L * x_vec[i - 1] * np.cos(theta_vec[i - 1]))
+ (k * L * 1.5 * np.sin(theta_vec[i - 1]))
) * (1 - sq)
[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
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
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(r"$\theta$(rad)")
plt.title(r"$\theta$(t)")
plt.plot(t_vec, theta_vec)
plt.show()
# Create Animation:
# Setup Figure:
from matplotlib.patches import Circle
fig, ax = plt.subplots()
(p,) = ax.plot([], [], color="royalblue")
min_lim = -5
max_lim = 5
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("problem 2 - pendulum and block connected by a spring")
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, 0 - 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
# 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, 0 - height / 2 - height_hatch)
ground_hatch = Rectangle(
xy_hatch, width_hatch, height_hatch, facecolor="None", linestyle="None", hatch="/"
)
ax.add_patch(ground_hatch)
c1 = Circle((0, 0), radius=0.1, color="cornflowerblue")
ax.add_patch(c1)
simulation_size = len(t_vec)
x1_pendulum_arm = np.zeros(simulation_size)
y1_pendulum_arm = np.zeros(simulation_size)
x2_pendulum_arm = np.zeros(simulation_size)
y2_pendulum_arm = np.zeros(simulation_size)
for i in range(0, simulation_size):
x1_pendulum_arm[i] = L * np.sin(theta_vec[i])
y1_pendulum_arm[i] = 1.5 - L * np.cos(theta_vec[i])
x2_pendulum_arm[i] = x_vec[i]
y2_pendulum_arm[i] = 0
# Animate:
with writerObj.saving(fig, title + ".mp4", dpi):
for i in range(0, vec_size, sample_rate):
# Update Pendulum Arm(rod)
x_data_points1 = [0, x1_pendulum_arm[i], x1_pendulum_arm[i], x2_pendulum_arm[i]]
y_data_points1 = [
1.5,
y1_pendulum_arm[i],
y1_pendulum_arm[i],
y2_pendulum_arm[i],
]
p.set_data(x_data_points1, y_data_points1)
patch_center1 = x1_pendulum_arm[i], y1_pendulum_arm[i]
c1.center = patch_center1
# Update Cart Patch:
r.set(xy=(x_vec[i] - width / 2, 0 - height / 2))
# Update Drawing:
fig.canvas.draw()
# 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)