import numpy as np
from numpy import sin, cos, pi, exp, log
import matplotlib.pyplot as plt


end_t = 1/10

n = 100
dx = 1/(n+1)
dt = dx**2/4.25

# Make sure end_t is a grid point
steps = int(end_t/dt+1)
dt = end_t / steps

x = np.linspace(dx,1-dx,n)

def shift(x, k):
  # return [(x[i+k] if 0 <= i+k < len(x) else 0) for i in range(len(x))] #Equivalent but slower
  r = np.zeros(len(x))
  if k > 0: r[:-k] = x[k:]
  else: r[-k:] = x[:k]
  return r

for eps in [1/8,1/16,1/32,1/64, 1/100]:
  alpha = lambda u : 2+cos(u)*eps
  v_list = [sin(3*pi*x)]
  for _ in range(steps):
    v = v_list[-1]
    new_v = v + dt/dx**2 * (1/2*(alpha(shift(v,1))+alpha(      v    )) * (shift(v, 1) - v) - 
                            1/2*(alpha(      v   )+alpha(shift(v,-1))) * (v - shift(v,-1)))
    v_list.append(new_v)
  plt.plot(x, v_list[-1], label='$\epsilon = \\frac{1}{%d}$'%round(1/eps))

plt.plot(x, sin(3*pi*x) * exp(-2*(3*pi)**2*end_t), label='Explicit $\epsilon = 0$')
plt.legend()
plt.show()