%matplotlib inline
from IPython.display import HTML
import matplotlib.animation
import matplotlib.pyplot as plt
import numpy as np
class Pendulum:
def __init__(self, length, mass):
self.length = length
self.mass = mass
self.q = 0
self.q_dot = 0
def reset(self, q, q_dot):
"""
Set the angle and angular velocity of the pendulum
"""
self.q = q
self.q_dot = q_dot
def step(self, dt):
"""
Compute a forward step on the internal model.
Updates the position and velocity of the pendulum
"""
self.q_dot = self.q_dot + self._q_ddot() * dt
self.q = self.q + self.q_dot * dt
def position(self):
"""
Position of the tip of the pendulum
"""
x = self.length * np.sin(self.q)
y = - self.length * np.cos(self.q)
return x, y
def _q_ddot(self):
EARTH_GRAVITY_CONSTANT = 9.81 # m/s^2
return - (EARTH_GRAVITY_CONSTANT / self.length) * np.sin(self.q)
pendulum = Pendulum(length=2.0, mass=0.42)
pendulum.reset(1.0, 0.)
dt = 0.05
angle = [pendulum.q]
pos = [pendulum.position()]
time = [0]
for i in range(60):
pendulum.step(dt)
angle.append(pendulum.q)
pos.append(pendulum.position())
time.append(time[-1] + dt)
X = np.asarray(pos).T[:][0]
Y = np.asarray(pos).T[:][1]
%%capture
fig, ax = plt.subplots()
ax.axis([-2, 2, -2.5, 0.5])
ax.set_aspect('equal')
ax.set_xlabel('X[m]')
ax.set_ylabel('Y[m]')
l, = ax.plot([], [], '-', lw=2)
c = ax.add_patch(plt.Circle((0, 0), 0.1))
def init():
l.set_data([], [])
c.center = (X[0], Y[0])
def animate(i):
l.set_data([0, X[i]], [0, Y[i]])
c.center = (X[i], Y[i])
ani = matplotlib.animation.FuncAnimation(fig, animate, frames=len(X), init_func=init)
HTML(ani.to_jshtml())