from numpy import empty
import matplotlib.pyplot as pl 

# Constants
L = 0.01      # Thickness of steel in meters
D = 4.25e-6   # Thermal diffusivity of steel
N = 100       # Number of divisions in grid
a = L/N       # Grid spacing
h = 1.0e-4      # Time-step
epsilon = h/10

Tlo = 0.0     # Low temperature in Celsius
Tmid = 30.0   # Intermediate temperature in Celsius
Thi = 60.0    # Hi temperature in Celsius

tdisp = [0.01, 0.1, 0.4, 1.0, 10.0]     # Display times, in ascending order
tend = tdisp[-1] + epsilon  # Continue only to just past last display time

# Create arrays
T = empty(N+1,float)    # N divisions, so N+1 points
T[0] = Thi
T[N] = Tlo
T[1:N] = Tmid
Tp = empty(N+1,float)
Tp[0] = Thi
Tp[N] = Tlo

# Main loop
t = 0.0
c = h*D/(a*a)
t_idx = 0   # index of display time list
while t<tend:

    # Calculate the new values of T
    Tp[1:N] = T[1:N] + c*(T[2:N+1]+T[0:N-1]-2.0*T[1:N])     # array in one swoop
    # Elements 0 and N have fixed temperatures
    T = Tp
    t += h  # increment time

    # Make plots at the given display times
    if abs(t-tdisp[t_idx]) < epsilon:
        pl.plot(T, label=str(tdisp[t_idx])+" s")
        if t_idx < len(tdisp)-1:  # don't allow index to go out of range
            t_idx += 1

pl.xlabel(r"$x (\times$ "+str(a)+" m)")
pl.ylabel("$T$, Celsius")
pl.title("Temperature profile of "+str(L)+"-m bar")
pl.legend()
pl.show()
print("Done.")