#
# Solution of unsteady 1D diffusion equation withinin the domain [0,1] using
# explicit second order finite difference method.
# Dirichlet and Neumann boundary conditions are implemented.
#
# This is a part of teaching materials provided for the course
# "Computational Heat & Fluid Flow (ME 605)" taught at IIT Goa
# during the winter semester of 2018
#
# Author: Y Sudhakar, IIT Goa
# email: sudhakar@iitgoa.ac.in
#
# Tested with python 3.6.3 on 16 August 2018
#
# Assignment for the students
# 1. Check solution at large time steps -- recall stability condition
# 2. Implement Robin boundary condition
# 3. Implement one-sided 2nd order approximation for Neumann BC
# 4. Write the iteration number and error-norm in a file
# 5. Introduce a source term and solve the problem
#

# This function applies boundary conditions on the given array
def apply_bc(u):
    u[0] = 100.0
    if bc == 'Dirichlet':
        u[num_mesh-1] = 50.0
    elif bc == 'Neumann':
        u[num_mesh-1] = u[num_mesh-2]
    else:
        sys.exit('Set correct boundary condition at the right side boundary')

# ============== Program begins here ==================================
import numpy as np
import math as m
import matplotlib.pyplot as plt
import sys

# --------------- Set the following input parameters ------------------
num_mesh = 51                  # number of mesh points
del_t    = 0.00001               # size of the time step
alpha    = 1.0                 # Thermal diffusivity
bc       = 'Dirichlet'         # Boundary condition at the right boundary; set either 'Dirichlet', 'Neumann' or 'Robin'
# ---------------------------------------------------------------------

del_x    	= 1.0/(num_mesh-1.0)        # mesh size (Delta_x)
xmesh    	= np.zeros(num_mesh)        # Mesh points
u_old    	= np.zeros(num_mesh)        # solution at the previous time step
u_new    	= np.zeros(num_mesh)        # Solution at the new time step
steady_term	= np.zeros(num_mesh)        # RHS of the discretized governing equation (approaches zero at steady state)

# Compute the location of mesh points
for i in range(0, len(xmesh)):
    xmesh[i] = i * del_x

# apply boundary conditions on the initial field
apply_bc(u_new)
apply_bc(u_old)

# Solve the discretized equation until steady state is achieved
norm = 100.0   # measure of difference between values at two consequetive time steps
iter = 0
converged = False
while converged==False:
    # the discretized equation
    for i in range(1, len(u_new)-1):
        u_new[i] = u_old[i] + alpha * del_t / (del_x**2) * (u_old[i-1]+u_old[i+1]-2.0*u_old[i])
    apply_bc(u_new)
    
    for i in range(1, len(u_new)-1):
        steady_term[i] = abs(alpha*(u_new[i+1] - 2.0*u_new[i] + u_new[i-1]) / (del_x**2.0))
    norm = max(steady_term)

    iter = iter + 1
    print(iter, norm)
    if norm <0.00001:
        converged = True
    # Soution not converged -- copy current value into old
    for i in range(0, len(u_new)):
        u_old[i] = u_new[i]

print(u_new)

# plot the converged results
plt.plot(xmesh,u_new,'x-')
plt.xlabel('x')
plt.ylabel('Temperature')
plt.show()
#plt.savefig('unsteady-diffusion.pdf')