# ------------------------------------------------------------------------------------
# Name  : Sandeep Tripathi
# Branch: Mechanical Engineering
# Temperature Distribution in a 1 Dimensional Fin using Quadratic Element 
# ------------------------------------------------------------------------------------

# Define various modules required
from __future__ import division
import numpy as np
import pylab
import math

print("-----------Temperature Distribution in a 1 Dimensional Fin using Quadratic Element-----------\n")

# Input from the user geometrical and material parameters for the fin 
L = float(raw_input("Enter the length of the fin (mm): "))
D = float(raw_input("Enter the diameter of the fin (mm): "))
k = float(raw_input("Enter the thermal conductivity of fin (W/mK): "))
# Input from the user other required parameters
Q_f = float(raw_input("Enter the heat flux at the start of the fin (W/m^2): "))
Q = float(raw_input("Enter the heat generated within the fin (W/m^3): "))
T_end = float(raw_input("Enter the temperature at the end of the fin (deg C): "))
# Input the numbers of quadratic element from the user
n = int(raw_input("Enter the number of quadratic elements to be modelled for the fin: "))

# Initializing global stiffness matrix, global heat generation vector and global force vector
K = np.zeros([(2*n + 1), (2*n + 1)])
F = np.zeros([2*n + 1, 1])
G = np.zeros([2*n + 1, 1])

# Calculate the required values of stiffness matrix and heat generation matrix
kcon = (k * ((math.pi * (D / 1000)**2) / 4)) / (3 * ((L / 1000) / n))
Qcon = Q * ((math.pi * (D / 1000)**2) / 4) * ((L / 1000) / n)
# Forming the element stiffness and element heat generation matrix
kstif = kcon *  np.matrix([[7.0, -8.0, 1.0], [-8.0, 16.0, -8.0], [1.0, -8.0, 7.0]])
Qstif = Qcon * np.matrix([[1.0/6], [2.0/3], [1.0/6]])
F[0, 0] = Q_f * ((math.pi * (D / 1000)**2) / 4)

# Assembly of global stiffness matrix, heat generation matrix and global force matrix is carried out 
for node in  range(2*n):
    if node % 2 == 0:
        grow = node
        for erow in range(3):
            G[grow, 0] += Qstif[erow, 0]
            gcol = node
            for ecol in range(3):
                K[grow, gcol] += kstif[erow, ecol]
                gcol += 1
            grow += 1

# Summation of all the force matrix is done
Ftot = F + G
# Initializing the nodal solutions and substituting the boundary conditions
nodalsol = np.zeros([2*n + 1, 1])
nodalsol[2*n] = T_end

# Defining the criteria for stopping the solution
numiter = 100
soleqn = range(2*n)

# Solving equations for computing temperature at various nodes
# Gauss Siedel Algorithm is used
for iteration in range(numiter):
    for row in soleqn:
        const = Ftot[row]
        for col in range(2*n + 1):
            if row != col:
                const = const - (K[row, col] * nodalsol[col])
        nodalsol[row] = const/K[row, row]

# Calculating nodal solutions at intermediate length using interpolation (Quadratic Shape) functions   
Xaxis = []
Yaxis = []
elem = 0
for i in range(2*n):
    if i % 2 == 0:
        for xlen in range(0, 21, 1):
            s = xlen / 20
            tempx = (elem * (L / n)) + (s * (L / n))
            Xaxis.append(tempx)
            tempy = (nodalsol[i] * 2 * (s - 1) * (s - 0.5)) - (nodalsol[i + 1] * 4 * s * (s - 1)) + (nodalsol[i + 2] * 2 * s* (s - 0.5))
            Yaxis.append(tempy)
        elem += 1

# Plotting the given results
pylab.plot(Xaxis, Yaxis,)
pylab.title('Temperature Distribution in a 1 Dimensional Fin using Quadratic Element', fontsize = 15)
pylab.xlabel('Length of Fin (mm)')
pylab.ylabel('Temperature of Fin (deg C)')
pylab.grid(True)
pylab.show()
