#!/usr/bin/python3
# Heat equation with the FTCS method

import numpy as np
import matplotlib.pyplot as plt

D = 4.25E-6  # diffusion coefficient
w = .01      # width of spatial interval
T0 = 283     # T at x=0
T1 = 373     # T at x=w
N = 100      # spatial discretization
a = w/N      # lattice constant
tmax = 10    # solution sought for 0 <= t <= tmax
h = 1.E-3    # time step
t = 0
c = h * D / a**2

# initial temperature profile at t = 0
T = np.ones(N+1)
T *= T0
T[N] = T1 
tplot = [.01, .1, 1, 10]  # intermediate times for plotting

while t < tmax:
    T[1:N] = T[1:N] + c * (T[:N-1] + T[2:] - 2 * T[1:N])
    for tp in tplot: # are we (close to) where we want to plot?
        if abs(t - tp) < 0.1 * h:
            plt.plot(T - 273., label = "t = " + str(tp) + " s")
    t += h

plt.xlabel("x [10^(-4) m]")
plt.ylabel("T(x) [C]")   
plt.legend(loc = "upper left")  
plt.show()