bouncing ball

bouncr.py

'''
This example uses a MovieWriter directly to grab individual frames and
write them to a file. This avoids any event loop integration, but has
the advantage of working with even the Agg backend. This is not recommended
for use in an interactive setting.

This version avoids the ugly stateful plt.* expressions, and the 
set_data method. Instead, it uses the more generic axis methods. This is slower,
but easier to remember, I hope.
'''
import matplotlib
matplotlib.use("Agg")
import kevrekidis as kv
import numpy as np
import matplotlib.animation as manimation

FFMpegWriter = manimation.writers['ffmpeg']
writer = FFMpegWriter(fps=15)

fig, ax = kv.fa(proj3d=False)


x, y, z = 0, 0, 0


azim = 0
T = 0
irange = 300
stepsRange = 8
tmax = irange * stepsRange
dpi = 300
y = y0 = 10.0
v = 0.0
r = 1.0
Y = [y]
V = [v]
dt = .05
k = 60.0
m = 4.0
def circ(y, ax, yrad=None, rad=r, color='black'):
    if yrad is None:
        yrad = r
    xrad = r
    xrad += r - yrad
    T = np.arange(0,np.pi*2+.1,.1)
    X = xrad*np.cos(T)
    Y = yrad*np.sin(T) + y
    ax.plot(X, Y, color=color)
    return X, Y
with writer.saving(fig, "bouncr.mp4", dpi):
    for i in range(irange):
        y += v * dt
        f = -9.8 * m
        if y < 0:
            y = .01
            v = 0
        v *= .99 # air resistance
        if y < r * 1.05:
            v *= .9 # fur? I dunno.
        if y < r and y != 0:
            f += min(k * 1 / y, 1000)
            v *= .8  # bouncing energy loss
        v += f/m * dt
        print i, irange, y, v, f
        ax.cla()
        ax.set_xlim((-y0/2.0, y0/2.0))
        ax.set_ylim((0, y0))
        ax.set_xticks([])
        ax.set_xticklabels([])
        if y < r:
            yrad = y
        else:
            yrad = r
        circ(y, ax, yrad=yrad)
#         ax.scatter(0, y)
        writer.grab_frame()

from os import system
system("notify-send 'plot done' &")
system("mplayer bouncr.mp4 &")