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

f = lambda x : 9 + 3*cos(pi*x) + 5*cos(4*pi*x)
analytic = lambda x,t : 9 + 3*exp(-pi**2*t)*cos(pi*x) + 5*exp(-16*pi**2*t)*cos(4*pi*x)

def maxError(n):
  end_t = 1/100

  dx = 1/(n+1)
  dt = dx**2/2
  r = dt/dx**2

  x = np.linspace(dx,1-dx,n)
  v = [f(x)]

  steps = round(end_t/dt)
  assert(abs(dt*steps-end_t) < 1e-8)

  for m in range(steps):
    new_v = (1-2*r)*v[-1]

    
    new_v[ 1:] += r*v[-1][:-1]
    new_v[:-1] += r*v[-1][ 1:]

    if 0: # Forward differences
      new_v[0] += r*v[-1][0]
      new_v[-1] += r*v[-1][-1]
    else: # 3-point formula
      new_v[0] += r*(4*v[-1][0]-v[-1][1])/3
      new_v[-1] += r*(4*v[-1][-1]-v[-1][-2])/3

    v.append(new_v)

  u = analytic(x, end_t)
  return np.max(np.abs(u-v[-1]))

n_list = np.array([2**i*50+49 for i in range(0,7)]) # Make sure t = 1/100 is a grid point
err_list = [maxError(n) for n in n_list]

print(f'Rate: ', [log(err_list[i]/err_list[i-1])/log(n_list[i]/n_list[i-1]) for i in range(len(n_list)-1)])