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

@author: neil
"""
# minor modification to place both a fixed (Dirichlet) and gradient condition (Neumann)
# at the surface, and no condition at the bottom
# 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 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)

import numpy as np
import matplotlib.pyplot as plt
# header section
Tupper = 0
Tgrad  = -0.065    # W/m^2 at surface, negative since Z is positive down, and heat flow is up
K      = 3
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)-1)   # size of square matrix is n by n, off-diagonal is one less
b     = np.ones(len(z)  )      # make b and c the correct length
c     = np.ones(len(z)-2)
# np.diag(x,0) makes a square matrix with 'x' on the diagonal, 2nd arg places off the diagonal
A = np.diag(a,-1) + np.diag(b,0) + np.diag(c,-2);  # 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]=Tlower
A[0,0] = 1         # this is the dirichlet condition
A[0,1] = 0
A[1,0]= 1          # this is the Neuman condition
A[1,1]= -1

C=np.zeros(n);     # make the 'BC vector' for the right hand side
C[0] = Tupper      # upper BC is 0 degrees
C[1] = Tgrad*delz/K      # gradient condition (moved K and delz to right hand vector)

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.set_title("Geothermal Temperature Profile, no heat sources")
ax1.set_xlabel("Temperature")  
ax1.set_ylabel("Depth [kilometers]")  
ax1.invert_yaxis()  

plt.show()