# example solution to our steady state heat flow problem, using finite elements.
# heat flow in the earth with radiogenic heating
# this illustrates a long winded way to assemble the Matrices and vectors

# this version calculates the Temperature gradient at the surface
# [this skeleton outlines the steps in a steady state F.E. problem]

#import matplotlib.pyplot as plt
import numpy as np

# set problem parameters here
Px = 3                         # the value of the forcing function, here a constant
# make a mesh of nodal positions, in this simple problem I put 5 nodes, one close to the upper boundary
z = np.array([0,.01,.1,1,3])   # nodal locations
                               # make a list of element lengths
Le = np.diff(z)
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
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] = 0                       # set the upper BC in 'f'

# finally add the boundary condition eqns to the A matrix
# we do this by replacing the upper row with the simple equation
# that the value of T is the BC (in 'f')
A[0,0] = 1            #upper
A[0,1] = 0
# the lower BC is a gradient, and is 0, so we don't need to add it to 'f'

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

# Here all I do is print the result, no plots
print(T)
print("Gradient is " + str((T[1]-T[0])/Le[0]))
# I point out: to obtain this level of accuracy would require ~1000 nodes in F.D.