# Pendulum animation
# install ffmpeg (allows us to create videos)
!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
from mpl_toolkits import mplot3d
# 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)
z_vec = np.zeros(vec_size)
dz_vec = np.zeros(vec_size)
l_vec = np.zeros(vec_size)
x_vec[0] = 1.0
dx_vec[0] = 0.0
y_vec[0] = 0.0
dy_vec[0] = 1.0
z_vec[0] = 1.0
dz_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, 0.0, -2 * x_vec[i - 1]],
[0.0, m, 0.0, -2 * y_vec[i - 1]],
[0, 0, m, 1],
[-x_vec[i - 1], -y_vec[i - 1], 0.5, 0.0],
]
)
B = np.array([0.0, 0.0, -m * g, dx_vec[i - 1] ** 2 + dy_vec[i - 1] ** 2])
[ddx, ddy, ddz, 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
z_vec[i] = z_vec[i - 1] + dz_vec[i - 1] * dt
dx_vec[i] = dx_vec[i - 1] + ddx * dt
dy_vec[i] = dy_vec[i - 1] + ddy * dt
dz_vec[i] = dz_vec[i - 1] + ddz * 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("z(m)")
plt.title("z(t)")
plt.plot(t_vec, z_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 = plt.figure()
ax = plt.axes(projection="3d")
def f(x, y):
return x ** 2 + y ** 2
x = np.linspace(-1, 1, 10)
y = np.linspace(-1, 1, 10)
X, Y = np.meshgrid(x, y)
Z = f(X, Y)
ax = plt.axes(projection="3d")
ax.contour3D(X, Y, Z, 50, cmap="binary")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
ax.view_init(30, 35)
ax.set_title("problem 3 - bead on a paraboloid") # Plot Title
video_title = "simulation" # name of the mp4
# c = Circle((0,0), radius=0.1, color = 'cornflowerblue')
# ax.add_patch(c)
FPS = 20
sample_rate = int(1 / (FPS * dt))
dpi = 300 # quality of the video
writerObj = FFMpegWriter(fps=FPS)
simulation_size = len(t_vec)
x_pendulum_arm = np.zeros(simulation_size)
y_pendulum_arm = np.zeros(simulation_size)
z_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]
z_pendulum_arm[i] = z_vec[i]
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):
scat = ax.scatter3D(
x_pendulum_arm[i], y_pendulum_arm[i], z_pendulum_arm[i], c="green"
)
writerObj.grab_frame()
scat.remove()
from IPython.display import Video
Video("/work/simulation.mp4", embed=True, width=640, height=480)