# -*- 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,
# 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 includes all the nodes, including BC nodes in the 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)
# this is modified to include 2 K layers, layer 1 has K=1, layer2 has K=2

import numpy as np
import matplotlib.pyplot as plt
# header section
Tupper = 0
Tlower = 25
Z = 1000          # thickness of problem, 1 kilometers
n =7           # number of nodes in Z, needs to be odd
                  # 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

K = np.ones(n-1)  # Make a vector of K values, K[i] is either value 1 or 2
for i in np.arange(n//2,n-1): #note that '//' is integer divide
    K[i] = 2                  # this makes a vector, half = 1 and half = 2
    
# matrix construction, for the implicit FD, first construct the main diagonal 'a'
# then the upper and lower diagonals 'b' and 'c'
a     = np.ones(len(z))     # size of square matrix, including 2 BC nodes
b     = np.ones(len(z)-1)   # make b and c the correct length
c     = np.ones(len(z)-1)   # c will be the lower diagonal
for i in np.arange(1,n-1):  # make the main diagonal 'a', and the upper 'b' and lower 'c' diags
    a[i] = -K[i]-K[i-1]
    b[i] = K[i]
    c[i] = K[i]     # this is subtle, 'c' gets loaded into A, with c[0] going into
                    # 'A' row 1, not 0!. In effect c[i] is therefore K[i-1]!
    
A = np.diag(a,0) + np.diag(b,1) + np.diag(c,-1);  # makes a square matrix with 3 diagonals
# np.diag() makes a square matrix with a given vector on a diagonal
# 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]=Tlower
A[0,0] = 1
A[0,1] = 0
A[-1,-1]= 1
A[-1,-2]= 0

C=np.zeros(n);     # 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

T = 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.plot(T,depth, 'k*') 
ax1.plot([0,25],[.5,.5],'g')
ax1.set_title("Geothermal Temperature Profile, no heat sources")
ax1.set_xlabel("Temperature")  
ax1.set_ylabel("Depth [kilometers]")  
ax1.invert_yaxis()  

plt.show()