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

@author: neil
"""
# Toy program using a (Matrix) implicit 2nd order finite difference approach
# program to solve temperature in the lithosphere, steady state, (using d2T/dx2 = 0)
# For the situation where the BCs are on different boundaries: Here Tupper and Tlower
# Mainly illustrates the use of the implicit method, and the matrix eqn solver
# This version puts the BC values into the right-hand side ('forcing') vector
# and therefore deletes the first and last rows of the A matrix
# If you are unfamiliar with using linear algebra to solve a system of equations
# See the note on linear algebra on the web pages (week 3)

import numpy as np
import matplotlib.pyplot as plt
# header section
Tupper = 0
Tlower = 25
Z = 1000          # thickness of problem, 1 kilometers
n = 5             # number of nodes in Z, including the 2 BC nodes (only 3 unknown nodes)
                  # 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)-2)   # size of square matrix, excluding 2 BC nodes
b = c = np.ones(len(z)-3)      # 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

C=np.zeros(n-2);     # make the 'BC vector' for the right hand side
C[0] = -Tupper       # upper BC is 0 degrees
C[-1]= -Tlower      # gradient BC on the 2 bottom nodes

# make an array for holding the solution
T = np.zeros(n)
T[0] = Tupper      # place the known (BC) temperatures
T[-1]= Tlower

T[1:-1] = np.linalg.solve(A,C)   # direct solver in numpy

depth = z/1000
# 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, no heat sources")
ax1.set_xlabel("Temperature")  
ax1.set_ylabel("Depth [kilometers]")  
ax1.invert_yaxis()  

plt.show()