Aliasing animation plot

aliasing-animation.py

import numpy as np

import matplotlib
import matplotlib.pyplot as plt
from matplotlib import cycler

from matplotlib.animation import FuncAnimation

colors = [_['color'] for _ in list(matplotlib.rcParams['axes.prop_cycle'])]

## Animation of aliasing
frame_rate = 20

# full wave should stay on screen for 2 seconds
# full wave should take 2 seconds to wipe out
# one second between waves
# each wave = 7 seconds * 4 waves = 28 seconds of animation

anim_duration = 28
n_frames = anim_duration * frame_rate

sched_0 = np.zeros(n_frames) + 3
sched_0[0*frame_rate:2*frame_rate] = np.linspace(-2, 0, endpoint=False, num=2*frame_rate)
sched_0[2*frame_rate:4*frame_rate] = 0
sched_0[4*frame_rate:6*frame_rate] = np.linspace(0, 2, endpoint=False, num=2*frame_rate)

sched_1 = np.zeros(n_frames) + 3
sched_1[7*frame_rate:9*frame_rate] = np.linspace(-2, 0, endpoint=False, num=2*frame_rate)
sched_1[9*frame_rate:11*frame_rate] = 0
sched_1[11*frame_rate:13*frame_rate] = np.linspace(0, 2, endpoint=False, num=2*frame_rate)

sched_2 = np.zeros(n_frames) + 3
sched_2[14*frame_rate:16*frame_rate] = np.linspace(-2, 0, endpoint=False, num=2*frame_rate)
sched_2[16*frame_rate:18*frame_rate] = 0
sched_2[18*frame_rate:20*frame_rate] = np.linspace(0, 2, endpoint=False, num=2*frame_rate)

sched_3 = np.zeros(n_frames) + 3
sched_3[21*frame_rate:23*frame_rate] = np.linspace(-2, 0, endpoint=False, num=2*frame_rate)
sched_3[23*frame_rate:25*frame_rate] = 0
sched_3[25*frame_rate:27*frame_rate] = np.linspace(0, 2, endpoint=False, num=2*frame_rate)

# --- #

fig = plt.figure(figsize=(8, 6))

fs_real = 1000
fs = 5
duration = 2

f0 = 1

f1 = f0 + fs
f2 = f0 + 2 * fs
f3 = f0 - fs

# Plot the continuous time curves
t = np.linspace(0, duration, num=duration * fs_real, endpoint=False)
x0 = np.cos(2 * np.pi * f0 * t)
x1 = np.cos(2 * np.pi * f1 * t)
x2 = np.cos(2 * np.pi * f2 * t)
x3 = np.cos(2 * np.pi * f3 * t)

# Plot the samples
N = int(duration * fs)
xsamp = np.cos(2 * np.pi * f0 * np.arange(N) / fs)
tsamp = np.arange(N) / fs

ax = plt.gca()

p0 = plt.plot([], [], label=r'$f={}$ Hz'.format(f0), color=colors[0])[0]
p1 = plt.plot([], [], label=r'$f={}$ Hz'.format(f1), color=colors[2])[0]
p2 = plt.plot([], [], label=r'$f={}$ Hz'.format(f2), color=colors[3])[0]
p3 = plt.plot([], [], label=r'$f={}$ Hz'.format(f3), color=colors[4])[0]

plt.plot(tsamp, xsamp, linestyle='', marker='.', color=colors[1], label=r'$x[n]$', zorder=5)

plt.legend(ncol=2, loc='upper right', fontsize='large')
plt.title(r'Aliasing example at $f_s={}$ Hz'.format(fs))
plt.xlabel('Time $t$ [seconds]')
plt.ylim([-1.25, 1.25])
plt.xlim([0, 2.0])
plt.yticks(np.linspace(-1, 1, num=9))

def init():
    return p0, p1, p2, p3

def animate(i):
    time = i / float(frame_rate)
    
    if time <= 7:
        n = (t >= sched_0[i]) & (t <= sched_0[i] + 2)
        p0.set_data(t[n], x0[n])
        return (p0,)
    elif time <= 14:
        n = (t >= sched_1[i]) & (t <= sched_1[i] + 2)
        p1.set_data(t[n], x1[n])
        return (p1,)
    elif time <= 21:
        n = (t >= sched_2[i]) & (t <= sched_2[i] + 2)
        p2.set_data(t[n], x2[n])
        return (p2,)
    else:
        n = (t >= sched_3[i]) & (t <= sched_3[i] + 2)
        p3.set_data(t[n], x3[n])
        return (p3,)

animation = FuncAnimation(fig, animate,
                               init_func=init,
                               frames=n_frames,
                               interval=1000./frame_rate,
                               blit=True)