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