import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from JSAnimation import IPython_display

def aliasanim(n=300, nper=2):
    """Function that shows an animation of the aliasing problem.
    It samples a cosine function with increasing time steps.
    n: number of initial sampling points per period of the signal
    nper: number of periods shown for the real signal (suggested: n=300, nper=2)
    when launched, the sampling frequency is decreased."""

    twopi = 2 * np.pi
    k = np.arange(0, nper * n)
    # Initial signal, nicely sampled if n is large.
    # di = omega Delta t
    di = twopi / (nper*n) * nper
    t = di * k
    z = np.cos(t)

    fig = plt.figure()
    ax = plt.axes(xlim=(0, twopi*nper), ylim=(-1, 1))

    # Non-animated.
    ax.plot(t, z, 'k')
    ax.set_xlabel(r'$\omega t$')

    # Animated.
    line, = ax.plot([], [], 'k:o', alpha=0.5)
    text = ax.text(1, 1.05, '')
    line.set_data([], [])

    def init():
        return line, text

    def animate(i):
        # Create a sampling (ts and zs) with decreasing frequency.
        ts, zs = [], []
        di = twopi / (nper*n) * nper * i
        ts = di * k
        zs = np.cos(ts)
        line.set_data(ts, zs)
        if di / np.pi< 1:
            text.set_text(r'$\omega$ $\Delta t / \pi$ = %1.3f : no aliasing' % (di/np.pi))
        else:
            text.set_text(r'$\omega$ $\Delta t / \pi$ = %1.3f : ALIASING' % (di/np.pi))
        return line, text

    return animation.FuncAnimation(fig, animate, init_func=init,
                                   frames=n, interval=100)

aliasanim(n=300, nper=2)

HTML(html)