"""
This scripts simulate the moviment of the first N Fourier modes of the solution of the wave equation
(with L=1 and c=1):
    u_tt = u_xx, u(0,t)=u(1,t)=0, u(x,0) = f(x), u_t(x,0) = 0
where f is a triangle function, with the top at x=a. 
"""
import numpy as np
import matplotlib.pyplot as plt
plt.ion()
plt.grid('on')

x = np.linspace(0,1,101)
tspan = np.linspace(0,4,101)    
N = 100     # The number of Fourier modes

a = 0.5     # Location of the top point of the triangle    
Bn = [np.sin(a*n*np.pi)/(n*np.pi)**2/a/(1-a) for n in range(1,N+1)]     # Fourier coefficients 

# Make the truncated Fourier series S_N
s = np.zeros_like(x)
for n in range(1,N+1):
    s += Bn[n-1]*np.sin(n*np.pi*x)

plt.figure(1)
plt.clf()
plt.plot(x,s)
plt.grid(True)
plt.axis([0,1,-0.6,0.6])
plt.xlabel('x')
plt.ylabel('u(x,t)')
plt.title(f'time={0:.2f}')

input('Press a key to continue')

# Main loop
for t in tspan:
    plt.clf()
    s = np.zeros_like(x)
    for n in range(1,N+1):
        s += Bn[n-1]*np.sin(n*np.pi*x)*np.cos(n*np.pi*t)

    plt.plot(x,s)
    plt.grid(True)
    plt.axis([0,1,-0.6,0.6])
    plt.xlabel('x')
    plt.ylabel('u(x,t)')
    plt.title(f'time={t:.2f}')
    plt.pause(0.001)

# Plot the amplitude spectrum B_n as a function of n. 
# Only the first 20 are shown here (the rest is very small)
plt.figure(2)
plt.plot(range(1,21),Bn[:20],'o')
plt.xlabel('n')
plt.ylabel('B_n')
plt.title('Amplitude spectrum (Fourier coefficients B_n)')