Table of Contents
Decay
This example showcases : – using a generator to drive an animation, – changing axes limits during an animation.
import itertools
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
def data_gen():
for cnt in itertools.count():
t = cnt / 10
yield t, np.sin(2*np.pi*t) * np.exp(-t/10)
def init():
ax.set_ylim(-1.1, 1.1)
ax.set_xlim(0, 10)
del xdata[:]
del ydata[:]
line.set_data(xdata, ydata)
return line,
fig, ax = plt.subplots()
line = ax.plot([], [], lw=2)
ax.grid()
xdata, ydata = [], []
def run(data):
# update the data
t, y = data
xdata.append(t)
ydata.append(y)
xmin, xmax = ax.get_xlim()
if t >= xmax:
ax.set_xlim(xmin, 2*xmax)
ax.figure.canyas.draw()
line.set_data(xdata, ydata)
return line,
ani = animation.FuncAnimation(fig, run, data_gen, interval=10,init_func=init)
plt.show()Code language: PHP (php)output:
The Bayes update
This animation displays the posterior estimate updates as it is refitted when new data arrives. The vertical line represents the theoretical value to which the plotted distribution should converge.
import math
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
def beta_pdf(x, a, b):
return (x**(a-1) * (1-x)**(b-1) * math.gamma(a + b)
/ (math.gamma(a) * math.gamma(b)))
class UpdateDist:
def __init__(self, ax, prob=0.5):
self.success = 0
self.prob = prob
self.line, = ax.plot([], [], 'k-')
self.x = np.linspace(0, 1, 200)
self.ax = ax
# Set up plot parameters
self.ax.set_xlim(0, 1)
self.ax.set_ylim(0, 10)
self.ax.grid(True)
# This vertical line represents the theoretical value, to
# which the plotted distribution should converge.
self.ax.axvline(prob, linestyle='--', color='black')
def __call__(self, i):
# This way the plot can continuously run and we just keep
# watching new realizations of the process
if i == 0:
self.success = 0
self.line.set_data([], [])
return self.line,
# Choose success based on exceed a threshold with a uniform pick
if np.random.rand(1,) < self.prob:
self.success += 1
y = beta_pdf(self.x, self.success + 1, (i - self.success) + 1)
self.line.set_data(self.x, y)
return self.line,
# Fixing random state for reproducibility
np.random.seed(19680801)
fig, ax = plt.subplots()
ud = UpdateDist(ax, prob=0.7)
anim = FuncAnimation(fig, ud, frames=100, interval=100, blit=True)
plt.show()Code language: HTML, XML (xml)output:
The double pendulum problem
from numpy import sin, cos
import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as integrate
import matplotlib.animation as animation
G = 9.8 # acceleration due to gravity, in m/s^2
L1 = 1.0 # length of pendulum 1 in m
L2 = 1.0 # length of pendulum 2 in m
M1 = 1.0 # mass of pendulum 1 in kg
M2 = 1.0 # mass of pendulum 2 in kg
t_stop = 5 # how many seconds to simulate
def derivs(state, t):
dydx = np.zeros_like(state)
dydx[0] = state[1]
delta = state[2] - state[0]
den1 = (M1+M2) * L1 - M2 * L1 * cos(delta) * cos(delta)
dydx[1] = ((M2 * L1 * state[1] * state[1] * sin(delta) * cos(delta)
+ M2 * G * sin(state[2]) * cos(delta)
+ M2 * L2 * state[3] * state[3] * sin(delta)
- (M1+M2) * G * sin(state[0]))
/ den1)
dydx[2] = state[3]
den2 = (L2/L1) * den1
dydx[3] = ((- M2 * L2 * state[3] * state[3] * sin(delta) * cos(delta)
+ (M1+M2) * G * sin(state[0]) * cos(delta)
- (M1+M2) * L1 * state[1] * state[1] * sin(delta)
- (M1+M2) * G * sin(state[2]))
/ den2)
return dydx
# create a time array from 0..100 sampled at 0.05 second steps
dt = 0.05
t = np.arange(0, t_stop, dt)
# th1 and th2 are the initial angles (degrees)
# w10 and w20 are the initial angular velocities (degrees per second)
th1 = 120.0
w1 = 0.0
th2 = -10.0
w2 = 0.0
# initial state
state = np.radians([th1, w1, th2, w2])
# integrate your ODE using scipy.integrate.
y = integrate.odeint(derivs, state, t)
x1 = L1*sin(y[:, 0])
y1 = -L1*cos(y[:, 0])
x2 = L2*sin(y[:, 2]) + x1
y2 = -L2*cos(y[:, 2]) + y1
fig = plt.figure(figsize=(5, 4))
ax = fig.add_subplot(111, autoscale_on=False, xlim=(-2, 2), ylim=(-2, 1))
ax.set_aspect('equal')
ax.grid()
line, = ax.plot([], [], 'o-', lw=2)
time_template = 'time = %.1fs'
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
def animate(i):
thisx = [0, x1[i], x2[i]]
thisy = [0, y1[i], y2[i]]
line.set_data(thisx, thisy)
time_text.set_text(time_template % (i*dt))
return line, time_text
ani = animation.FuncAnimation(
fig, animate, len(y), interval=dt*1000, blit=True)
plt.show()Code language: PHP (php)output:
Oscilloscope
import numpy as np from matplotlib.lines import Line2D import matplotlib.pyplot as plt import matplotlib.animation as animation class Scope: def __init__(self, ax, maxt=2, dt=0.02): self.ax = ax self.dt = dt self.maxt = maxt self.tdata = [0] self.ydata = [0] self.line = Line2D(self.tdata, self.ydata) self.ax.add_line(self.line) self.ax.set_ylim(-.1, 1.1) self.ax.set_xlim(0, self.maxt) def update(self, y): lastt = self.tdata[-1] if lastt > self.tdata[0] + self.maxt: # reset the arrays self.tdata = [self.tdata[-1]] self.ydata = [self.ydata[-1]] self.ax.set_xlim(self.tdata[0], self.tdata[0] + self.maxt) self.ax.figure.canvas.draw() t = self.tdata[-1] + self.dt self.tdata.append(t) self.ydata.append(y) self.line.set_data(self.tdata, self.ydata) return self.line, def emitter(p=0.1): """Return a random value in [0, 1) with probability p, else 0.""" while True: v = np.random.rand(1) if v > p: yield 0. else: yield np.random.rand(1) # Fixing random state for reproducibility np.random.seed(19680801 // 10) fig, ax = plt.subplots() scope = Scope(ax) # pass a generator in "emitter" to produce data for the update func ani = animation.FuncAnimation(fig, scope.update, emitter, interval=50, blit=True) plt.show()
