#!/usr/bin/env python3
#
# heatimplicit.py
#
# Implicit numerical solution of the heat equation with a time
# dependent boundary condition (example from numerical computing with
# python).
#
# This python code solves the heat equation dT/dt = k d^2T/dx^2 on the
# x-interval [0,1], subject to time dependent initial conditions
# T(x=0,t)=sin(10t) and T(x=1,t) = 0. The initial condition is T(0,x)=0.
#
# Ramses van Zon, November 26, 2018

import numpy as np
import scipy.linalg as la
import matplotlib.pyplot as plt

k = 0.2
t1 = 0
t2 = 2.5
x1 = 0
x2 = 1
nx = 100
x = np.linspace(x1,x2,nx+1)
dx = x[1]-x[0]
   
for dt in [ 0.00005, 0.0005 , 0.005, 0.05 ]:

   nt = int((t2-t1)/dt)
   alpha = dt*k/dx**2

   # allocate the temparature array
   T = np.zeros(nx)

   # allocate an array to plot the image
   TasImage = np.zeros([nx,nt+1])

   # build matrix
   A = np.zeros([nx,nx])
   for j in range(1,nx-1):
       A[j,j-1] = -alpha
       A[j,j] = 1 + 2*alpha
       A[j,j+1] = -alpha
   # first and last element are different because of the boundary conditions
   A[0,0] = 1
   A[nx-1,nx-1] = 1

   TasImage[:,0] = np.flip(T).copy()

   for i in range(1,nt):
       t = t1 + i*dt
       # implicit solve
       b = T.copy()
       T = la.solve(A,b)
       # re-adjust boundary condition
       T[0] = np.sin(10*t)
       T[nx-1] = 0
       TasImage[:,i] = np.flip(T).copy()
       print ("t =",t)

   plt.style.use('dark_background')
   plt.imshow(TasImage, aspect='equal', extent=[t1,t2,x1,x2])
   plt.axes().set_aspect((t2-t1)/(x2-x1))
   plt.xlabel('time [s]')
   plt.ylabel('distance [m]')
   plt.colorbar()
   plt.text(t1+0.06*(t2-t1),0.9*(x2-x1)+x1,r"$\alpha="+str(alpha)+"$",size=16,color='k')
   plt.savefig("alpha"+str(alpha)+".pdf",transparent=True)
   plt.clf()