SHM.py

"""
Simple harmonic motion animation example, version 1
A bunch of stuff doesn't work, and will hopefully be fixed
author: Jonathan Shi
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.animation as animation

y = -.5
width = 2.5
x0 = width/2 # initial x of left side of rectangle
height = 1

m = 1 # kg
k = 5 # N/m
a = 2.5 # max amplitude, m
overlay = False # overlays the circle on the thing

fig1 = plt.figure(num="Oscillating Mass, m = 1 kg, k = 5 N/m")
ax = plt.axes(xlim=(-6,6),ylim=(-2,2))
plt.gca().axes.get_xaxis().set_visible(False)
plt.gca().axes.get_yaxis().set_visible(False)

if not overlay:
	fig2 = plt.figure(num="Circle")
	axCircle = plt.axes()
	plt.axis('square')
	plt.xlim(-1.5, 1.5)
	plt.ylim(-1.5, 1.5)
	plt.grid('on')
	plt.ylabel("Displacement in x")

fig3 = plt.figure(num="Wave")
axTrig = plt.axes()
plt.ylim(-a-1-x0, a+1+x0)
plt.grid('on')
plt.xlabel("Time (s)")
plt.ylabel("Displacement in x (m)")

rectangle = patches.Rectangle((x0, y), width, height, fill=True)
line, = ax.plot([], [], 'b')
vertical1, = ax.plot([a+x0, a+x0], [-2, 2], 'r')
ax.text(a+x0, -1, "  Maximum")
vertical2, = ax.plot([-a-x0, -a-x0], [-2, 2], 'r')
ax.text(-a-x0-1.5, -1, "Minimum")
vertical3 = ax.plot([0, 0], [-2, 2], 'k')
ax.text(.1, -1, " Equilibrium")
ax.add_patch(rectangle)

# fig2 = plt.figure("Cycle")
# axCircle = fig2.add_subplot(2, 1, 1)
# plt.xlim(-1.5, 1.5)
# plt.ylim(-1.5, 1.5)
aC = ax if overlay else axCircle

bigArc, = aC.plot([], [], 'k')
angleArc, = aC.plot([], [], 'r')
lineToAngle, = aC.plot([], [], 'k')
lineToOrigin, = aC.plot([0, 1], [0, 0], 'k')
# axTrig = fig2.add_subplot(2, 1, 2)
sineLine, = axTrig.plot([], [], 'b')

lines = [line, bigArc, angleArc, lineToAngle, lineToOrigin, sineLine]

def animate(t): # this t is saved as the current t
	# for the box
	# x = Acos(sqrt(k/m)*t)
	t = np.sqrt(k/m) * float(t) / 50
	x = a * np.cos(t) - x0
	# print t, x
	rectangle.set_xy([x, y])
	line.set_data([-6, x], [0, 0])
	lineToAngle.set_data([0, np.cos(t)], [0, np.sin(t)])

	arcT = np.linspace(0, t % (2*np.pi), 500)
	bigArcX = []
	bigArcY = []
	littleArcX = []
	littleArcY = []
	for v in arcT:
		bigArcX.append(np.cos(v))
		littleArcX.append(0.2 * np.cos(v))
		bigArcY.append(np.sin(v))
		littleArcY.append(0.2 * np.sin(v))
	# bigArc.set_data(bigArcX, bigArcY)
	angleArc.set_data(littleArcX, littleArcY)
	sinT = np.linspace(t - 6, t, 500)
	if plt.fignum_exists("Wave"):
		sineLine.set_data(sinT, a * np.cos(sinT))
		plt.sca(axTrig)
		plt.xlim(t - 6, t)
	return lines

delay = 70

aniBox = animation.FuncAnimation(fig1, animate, interval=delay, repeat=False, blit=False)
if not overlay:
	aniGraph1 = animation.FuncAnimation(fig2, animate, interval=delay, repeat=False, blit=False)
aniGraph2 = animation.FuncAnimation(fig3, animate, interval=delay, repeat=False, blit=False)

plt.show()