In [ ]:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from JSAnimation import IPython_display

In [ ]:
# Wave 1:
k1, omega1 = 5., 5.
c1 = omega1 / k1
# Wave 2:
k2, omega2 = 3., 4.5
c2 = omega2 / k2

cg = (omega2 - omega1) / (k2 - k1)
cgc = (omega2 + omega1) / (k2 + k1)

deltg = 10 if cg <= 0 else 10.

In [ ]:
def basic_animation(frames=200, interval=30, dt=0.1):
    fig = plt.figure(figsize=(16, 4))
    ax = plt.axes(xlim=(-15, 15), ylim=(-2, 2))
    
    # Animated line.
    kw = dict(alpha=0.5, linestyle='none', marker='o')
    text = ax.text([], 0, 'Cg')
    line0, = ax.plot([], [], linestyle=':', color='gray')
    line1, = ax.plot([], [], 'r')
    line2, = ax.plot([], [], 'g')
    line3, = ax.plot([], [], 'k', alpha=0.5)
    line4, = ax.plot([], [], 'k', alpha=0.5)
    point0, = ax.plot([], [], 'r', **kw)
    point1, = ax.plot([], [], 'g', **kw)
    point2, = ax.plot([], [], 'b', **kw)

    # Non animated lines.
    x = np.arange(-15, 15, 0.05)  # Distance domain.
    ax.set_xlabel('Distance')
    ax.set_ylabel('Amplitude')
    
    def init():
        return line0, line1, line2, line3, line4, point0, point1, point2, text

    def animate(k):
        t = k * dt
        arg1, arg2 = k1 * x - omega1 * t, k2 * x - omega2 * t
        s1, s2 = np.cos(arg1), np.cos(arg2)
        summ = s1 + s2
        sg = 2 * np.cos((k2 - k1) * x / 2. - (omega2 - omega1) * t / 2.)
        
        line0.set_data(x, summ)
        line1.set_data(x,   s1)
        line2.set_data(x,   s2)
        line3.set_data(x,   sg)
        line4.set_data(x,  -sg)
        point0.set_data(t *  c1 - 10, 0)
        point1.set_data(t *  c2 - 10, 0)
        point2.set_data(t * cgc - 10, 0)
        text.set_x(t * cg - deltg)
        return line0, line1, line2, line3, line4, point0, point1, point2, text

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

In [ ]:
dt = 0.1  # Time step
nstep = 100  # Number of time steps to take.
basic_animation(frames=60, dt=dt)

In [ ]: