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

f = lambda x : 2*np.minimum(x,1-x)
def analytic(x, t):
  s = np.zeros_like(x)
  for k in range(1,201):
    s += 8/pi**2 * (sin(k*pi/2)/k**2) * exp(-(k*pi)**2*t) * sin(k*pi*x)
  return s

end_t = 1/10

n = 50
dx = 1/(n+1)

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


for r in [1/2, 1, 2, 10]:
  dt = dx**2 * r

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

  theta = 0.5
  I = scipy.sparse.diags([1], offsets=[0], shape=(n,n), format='csc')
  A = scipy.sparse.diags([-1,2,-1], offsets=[-1,0,1], shape=(n,n), format='csc') / dx**2

  v_list = [f(x)]
  for _ in range(steps):
    new_v = scipy.sparse.linalg.spsolve(I + theta*dt*A, (I - (1-theta)*dt*A) @ v_list[-1])
    v_list.append(new_v)
  plt.plot(x, v_list[-1], label='$\\frac{\Delta t}{\Delta x^2} = %s$'%str(r))

plt.plot(x, analytic(x, end_t), label='Analytic')
plt.legend()
plt.show()