#!/usr/bin/python3
# (1+1)-dimensional wave equation with FTCS (unstable!)

import numpy as np
import matplotlib.pyplot as plt

L = .3       # string length
c = 100.     # phase velocity
N = 100      # number of spatial slices
a = L/N      # lattice constant
tmax = .01   # solution sought for 0 <= t <= tmax
h = 1.E-6    # time increment
t = 0
alpha = h * c**2 / a**2

# starting profile
phi = np.zeros(N+1)
for k in range(1, N+1):
    phi[k] = np.sin(k / N * np.pi) * k / N / 100
psi = np.zeros(N+1)
tplot = [0, .001, .004, .006]  # intermediate times for plotting

while t < tmax:
    phi[1:N], psi[1:N] = phi[1:N] + h * psi[1:N], \
                         psi[1:N] + alpha * (phi[:N-1] + phi[2:] - 2 * phi[1:N])
    for tp in tplot:
        if abs(t - tp) < 0.1 * h:
            plt.plot(phi, label = "t = " + str(tp) + " s")
    t += h

plt.xlabel("x")
plt.ylabel("phi(x)")   
plt.legend(loc = "upper left")  
plt.show()