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

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

for div in [2,6]:
  def maxError(n):
    end_t = 1/10

    dx = 1/(n+1)
    dt = dx**2/div
    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:]
      v.append(new_v)

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

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

  print(f'dt = dx**2/{div}, 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)])