import numpy as np
import matplotlib.pyplot as plt

from scipy.sparse import diags
def Laasonen(dt,dy,t_max,y_max,k,T1,T2):
    s = k*dt/dy**2
    y = np.arange(0,y_max+dy,dy)
    t = np.arange(0,t_max+dt,dt)
    nt = len(t)
    ny = len(y)
    T = np.zeros((ny,))
    T[0] = T1
    T[-1] = T2
    A = diags([-s, 1+2*s, -s], [-1, 0, 1], shape=(ny-2, ny-2)).toarray()
    for n in range(nt):
        Tn = T
        B = Tn[1:-1]
        B[0] = B[0]+s*T1
        B[-1] = B[-1]+s*T2
        T[1:-1] = np.linalg.solve(A,B)
    return y,t,T,s

dt = 0.01
dy = 0.0005
k = 10**(-4)
y_max = 0.04
T1 = 100
T2 = 0

for time in np.linspace(0,1.0,10):
    y,t,T,s = Laasonen(dt,dy,time,y_max,k,T1,T2)
    plt.plot(y,T)