# example solution to our steady state heat flow problem, using finite elements.
# heat flow in the earth with radiogenic heating
# Eqn is d2T/dz2 + 3 = 0, BC T(z=0) = 0, dT(z=5)/dz = 0, z in kilometers
# This program illustrates a long winded way to assemble the Matrices and vectors
# 'z' array #1 uses equally spaced elements to show the similarity of the 
# 'A' matrix to the Finite Difference approach, (compare with FD program)

#import matplotlib.pyplot as plt
import numpy as np
import matplotlib.pyplot as plt

# set problem parameters here
Tupper = 0        # upper temperature BC
LowerG = 0        # lower gradient BC
Px     = 3        # the value of the forcing function, (constant radiogenic heat per km)
# make a mesh of nodal positions (3 different discretizations as tests)
#z = np.linspace(0,4,5)   # nodal locations, 5 nodes evenly spaced to compare with FD
#z = np.array([0,1,1.5,2,2.5,3,3.25,3.5,3.7,3.9,4])   # nodal locations suited to problem
z = np.linspace(0,4.,500)  # overkill nodes, to compare to analytic soln) 
                          
Le = np.diff(z)             # make a list of element lengths, Le is a vector
Le_inv = 1/Le               # and their inverses

## *********everything below here is written with only variables, no numbers ***********                               
# coefficient matrix, A, constructed VERY piecemeal to show it explicitly
Alower = -np.diag(Le_inv,-1)   #lower diagonal of A
Aupper = -np.diag(Le_inv,1)    #upper diagonal
Avec1  = np.append(Le_inv,[0]) #first integral of main diag
Avec2  = np.append([0],Le_inv) # 2nd of main diag
Amid   = np.diag(Avec1+Avec2,0)
A      = Amid + Aupper + Alower# put it all together

# forcing and BC vector (f), which is often the only vector that varies between problems
# and usually requires the most thought by the programer
# Note the 'dT/dz * phi' term is only non-zero at the boundaries, it is replaced with BCs
f1 = np.append(Px*Le/2,[0])    # with a constant Px, 'f' is only integrals over
f2 = np.append([0],Px*Le/2)    # single basis functions (over 2 elements)
f = f1+f2
f[0]  = Tupper                 # set the upper BC in 'f'
f[-1] = LowerG                 # set the lower gradient

# finally add the boundary condition eqns to the A matrix
# we do this by replacing the upper row with the value of T = upperBC
# and the lower row with the gradient condition dT/dz = lowerG
A[0,0] = 1            # upper BC
A[0,1] = 0
A[-1,-1] = 1          # lower BC
A[-1,-2] = -1         # this is a gradient condition

T = np.linalg.solve(A,f)

depth = z
# make our basic plotting canvas, with a title in the header bar
fig = plt.figure()
fig.canvas.set_window_title ("Temperature In the Crust")
ax1 = fig.add_subplot(1,1,1)
ax1.grid(True)
ax1.plot(T,depth, 'r-+')  
ax1.set_title("Geothermal Temperature Profile, FE solution")
ax1.set_xlabel("Temperature red-curve, analytic soln black-curve")  
ax1.set_ylabel("Depth [kilometers]")  
ax1.invert_yaxis()  

# show analytic soln
z = np.linspace(0,4.,100) 
ax1.plot((12*z - 3*z*z/2),z,'k-')
plt.show()