Pendulum and Block: HW 1.2
Import Libraries
# Install ffmpeg (allows us to create videos)
!apt update -y
!apt install ffmpeg -y
import numpy as np
import matplotlib.pyplot as plt
import sympy
from matplotlib.animation import FFMpegWriter
from matplotlib.patches import Circle, Rectangle
Generate Equations of Motion with SymPy
# Parameters
L = 1 # length of pendulum (m)
m_1 = 1 # mass of pendulum (kg)
k = 10 # spring stiffness (N/m)
m_2 = 2 # mass of block (kg)
h = 1.5 # height of pendulum hinge (m)
r_0 = 1 # rest length of spring (m)
g = 9.8 # gravitational constant (m/s^2)
# Create Symbols for Time:
t = sympy.Symbol('t') # Creates symbolic variable t
# Create Generalized Coordinates as a function of time: q = [theta, x]
th = sympy.Function('th')(t)
x = sympy.Function('x')(t)
# Position Equation: r = [x, y]
r1 = np.array([L * sympy.sin(th), -L * sympy.cos(th)]) # Position of pendulum
r2 = np.array([x,-h]) # Position of block
# Velocity Equation: d/dt(r) = [dx/dt, dy/dt]
v1 = np.array([r1[0].diff(t), r1[1].diff(t)]) # Velocity of pendulum
v2 = np.array([r2[0].diff(t), 0]) # Velocity of block
# Length of spring squared
r_squared = (r1[0]-r2[0])**2 + (r1[1]-r2[1])**2
r = r_squared**0.5
r0 = 1
# Energy Equations:
T = 1/2 * m_1 * np.dot(v1, v1) + 1/2 * m_2 * np.dot(v2, v2) # Kinetic Energy
# V = m_1 * g * r1[1] + m_2 * g * r2[1] + 1/2 * k * r_squared# Potential Energy
V = m_1 * g * r1[1] + m_2 * g * r2[1] + 1/2 * k * (r0-r)**2# Potential Energy
L = T - V # Lagrangian
# Lagrange Terms:
dL_dth = L.diff(th)
dL_dth_dt = L.diff(th.diff(t)).diff(t)
dL_dx = L.diff(x)
dL_dx_dt = L.diff(x.diff(t)).diff(t)
# Euler-Lagrange Equations: dL/dq - d/dt(dL/ddq) = 0
th_eqn = dL_dth - dL_dth_dt
x_eqn = dL_dx - dL_dx_dt
replacements = [(th.diff(t).diff(t), sympy.Symbol('ddth')), (th.diff(t), sympy.Symbol('dth')), (th, sympy.Symbol('th')),
(x.diff(t).diff(t), sympy.Symbol('ddx')), (x.diff(t), sympy.Symbol('dx')), (x, sympy.Symbol('x'))]
th_eqn = th_eqn.subs(replacements)
x_eqn = x_eqn.subs(replacements)
print(th_eqn)
print(x_eqn)
r1 = r1[0].subs(replacements), r1[1].subs(replacements)
r2 = r2[0].subs(replacements), -h
# Solve for Coefficients for A * x = B where x = [ddth1 ddth2]
A1 = th_eqn.coeff(sympy.Symbol('ddth'))
A2 = th_eqn.coeff(sympy.Symbol('ddx'))
A3 = x_eqn.coeff(sympy.Symbol('ddth'))
A4 = x_eqn.coeff(sympy.Symbol('ddx'))
print(A1)
# Multiply remaining terms by -1 to switch to other side of equation: A * x - B = 0 -> A * x = B
remainder = [(sympy.Symbol('ddth'), 0), (sympy.Symbol('ddx'), 0)]
B1 = -1 * th_eqn.subs(remainder)
B2 = -1 * x_eqn.subs(remainder)
# Generate Lambda Functions for A and B and Position Equations:
replacements = (sympy.Symbol('th'), sympy.Symbol('dth'), sympy.Symbol('x'), sympy.Symbol('dx'))
# A1 = sympy.utilities.lambdify(replacements, A1, "numpy")
# A2 = sympy.utilities.lambdify(replacements, A2, "numpy")
# A3 = sympy.utilities.lambdify(replacements, A3, "numpy")
# A4 = sympy.utilities.lambdify(replacements, A4, "numpy")
# B1 = sympy.utilities.lambdify(replacements, B1, "numpy")
# B2 = sympy.utilities.lambdify(replacements, B2, "numpy")
r1 = sympy.utilities.lambdify(replacements, r1, "numpy")
r2 = sympy.utilities.lambdify(replacements, r2, "numpy")
ddth = B1/A1
ddx = B2/A4
ddth_func = sympy.utilities.lambdify(replacements, ddth, "numpy")
ddx_func = sympy.utilities.lambdify(replacements, ddx, "numpy")
Simulation of Pendulum and Block
# Simulate System:
x0 = 1, 0, 0, 0 # th, dth, x, dx
dt = 0.001
sim_time = 10
time = np.arange(0, sim_time, dt)
sim_length = len(time)
# Initialize Arrays:
th_vec = np.zeros(sim_length)
dth_vec = np.zeros(sim_length)
x_vec = np.zeros(sim_length)
dx_vec = np.zeros(sim_length)
x1_vec = np.zeros(sim_length)
y1_vec = np.zeros(sim_length)
x2_vec = np.zeros(sim_length)
y2_vec = np.zeros(sim_length)
# Evaluate Initial Conditions:
th_vec[0] = x0[0]
dth_vec[0] = x0[1]
x_vec[0] = x0[2]
dx_vec[0] = x0[3]
x1_vec[0], y1_vec[0] = r1(x0[0], x0[1], x0[2], x0[3])
x2_vec[0], y2_vec[0] = r2(x0[0], x0[1], x0[2], x0[3])
# Initialize A and B:
A = np.array([[0, 0], [0, 0]])
B = np.array([0, 0])
# Euler Integration:
for i in range(1, sim_length):
# Evaluate Dynamics:
ddth = ddth_func(th_vec[i-1], dth_vec[i-1], x_vec[i-1], dx_vec[i-1])
ddx = ddx_func(th_vec[i-1], dth_vec[i-1], x_vec[i-1], dx_vec[i-1])
# Euler Step Integration:
th_vec[i] = th_vec[i-1] + dth_vec[i-1] * dt
dth_vec[i] = dth_vec[i-1] + ddth * dt
x_vec[i] = x_vec[i-1] + dx_vec[i-1] * dt
dx_vec[i] = dx_vec[i-1] + ddx * dt
# Animation States:
x1_vec[i], y1_vec[i] = r1(th_vec[i-1], dth_vec[i-1], x_vec[i-1], dx_vec[i-1])
x2_vec[i], y2_vec[i] = r2(th_vec[i-1], dth_vec[i-1], x_vec[i-1], dx_vec[i-1])
plt.plot(time,th_vec)
plt.title("Theta over time")
plt.show()
plt.plot(time,x_vec)
plt.title("X over time")
plt.show()
Animate the Pendulum and Block
# Create Animation:
# Setup Figure: Initialize Figure / Axe Handles
fig, ax = plt.subplots()
p1, = ax.plot([], [], color='black', linewidth=2)
p2, = ax.plot([], [], color='black', linewidth=2)
lb, ub = -2.5, 2.5
ax.axis('equal')
ax.set_xlim([lb, ub])
ax.set_xlabel('X') # X Label
ax.set_ylabel('Y') # Y Label
ax.set_title('Pendulum Block Simulation:')
video_title = "simulation"
# Setup Animation Writer:
FPS = 20
sample_rate = int(1 / (dt * FPS))
dpi = 300
writerObj = FFMpegWriter(fps=FPS)
# Initialize Patch: Pendulum 1 and 2
pendulum = Circle((0, 0), radius=0.1, color='cornflowerblue', zorder=10)
width = 0.25
height = 0.15
block = Rectangle((0,0), width, height, color='cornflowerblue', zorder=10)
ax.add_patch(pendulum)
ax.add_patch(block)
# Draw Static Objects:
pin_joint = Circle((0, 0), radius=0.05, color='black', zorder=10)
ax.add_patch(pin_joint)
# Plot and Create Animation:
with writerObj.saving(fig, video_title+".mp4", dpi):
for i in range(0, sim_length, sample_rate):
# Draw Pendulum Arm:
x_pendulum_arm = [0, x1_vec[i], x2_vec[i]]
y_pendulum_arm = [0, y1_vec[i], y2_vec[i]]
p1.set_data(x_pendulum_arm, y_pendulum_arm)
# Update Pendulum Patches:
pendulum_center = x1_vec[i], y1_vec[i]
block_center = x2_vec[i], y2_vec[i]
pendulum.center = pendulum_center
block.set(xy=(x_vec[i] - width / 2, -h-height/2))
# Update Drawing:
fig.canvas.draw()
# Grab and Save Frame:
writerObj.grab_frame()