import matplotlib.pyplot as plt
import numpy as np

#gitter     (prøv x0=0,xf=1,t0=0,tf=0.1,n=36,m=11,u(x,0)=sin(pi*x) og sammenlign med eksplisitt)
x0 = 0
xf = 1
t0 = 0
tf = 0.1
n = 36  #x-punkter
m = 11  #t-punkter
h = (xf-x0)/(n-1)   #x-gitterbredde
k = (tf-t0)/(m-1)   #t-gitterbredde
r = k/(h**2)
x = np.linspace(x0,xf,n)
t = np.linspace(t0,tf,m)


#initialkrav
u = np.zeros((m,n))
for i in range(1,n-1):
    u[0][i] = np.sin(np.pi*h*i)

#matrise
A = np.zeros((n,n))
for i in range(1,n-1):
    A[i][i] = 1 + 2*r
    if i != 1:
        A[i][i-1] = -r
    if i != (n-2):
        A[i][i+1] = -r

#fylle inn u-verdier for t>0
for i in range(1,m):
    u[i] = np.linalg.lstsq(A,u[i-1])[0]

#plotting
plt.figure(dpi=150,figsize=(11,7))
for i in range(m):
    plt.plot(x,u[i],label=f't={round(i*k,6)}')
plt.legend(loc='upper right')
plt.grid()