# -*- coding: utf-8 -*-
"""
Created on Thu Feb 18 21:03:57 2021

@author: neil
"""
# Program using a (Matrix) implicit 2nd order finite difference approach
# program to solve temperature in the lithosphere, steady state, (using d2T/dx2 +3 = 0)
# includes unrealistic radiogenic heating to match 1st FE problem
# Mainly illustrates the use of the implicit method, and the matrix eqn solver
# This version includes all the nodes, including BC nodes in the matrix

import numpy as np
import matplotlib.pyplot as plt
# header section
Tupper = 0        # upper temperature BC
LowerG = 0        # lower gradient BC

Z = 4             # thickness of problem, 1 kilometers
n = 5             # number of nodes in Z, including the 2 BC nodes (only 3 unknown nodes)
#n = 41            # number of nodes in Z, including the 2 BC nodes 
#n = 401
                  # this uses a small matrix so you can print out and observe the results
delz = Z/(n-1)    # delta z, length of problem / number of nodes, we are using z positive down
z = np.linspace(0,Z,n) # 'z' is a vector for plotting, and a template for making other arrays

# matrix construction, for the implicit FD, first construct the main diagonal 'a'
# then the upper and lower diagonals 'b' and 'c'
a     = -2*np.ones(len(z))   # size of square matrix, excluding 2 BC nodes
b = c = np.ones(len(z)-1)      # make b and c the correct length, note the Python shorthand!
# np.diag(x,0) makes a square matrix with 'x' on the diagonal, 2nd arg places off the diagonal
A = np.diag(a,0) + np.diag(b,1) + np.diag(c,-1);  # makes a square matrix with 3 diagonals
# if you don't understand this, try printing out the matrix
# add the BC nodes to the matrix, essentially saying T[0]=Tupper,T[-1] is a gradient conditions of 0
A[0,0] = 1
A[0,1] = 0
A[-1,-1]= 1
A[-1,-2]= -1

C=np.ones(n)*-3*delz*delz;     # make the 'BC vector' for the right hand side, the '3' is radiogenic heat
C[0] = Tupper        # upper BC is 0 degrees
C[-1]= LowerG        # gradient BC on the 2 bottom nodes

T = np.linalg.solve(A,C)   # direct solver in numpy

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, FD solution")
ax1.set_xlabel("Temperature, black line is analytic soln")  
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()