# -*- coding: utf-8 -*-
"""
Created on Tue Feb 23 17:30:43 2022

@author: neil
"""
# program to allow numerical experiments on temperature in the lithosphere
# This version includes radiogenic heat sources, as well as changing the thermal properties
# Illustrates the use of the implicit method, and the matrix eqn solver
# Both the gradient BC and the set Temp are at the top of the solution 

# We use 3 different K layers to mimic the increasing conductivity as the 
# density and temperature increases in the crust
# Here we attempt to mimic the 'rigid' part of the continental Lithosphere, which is about
# 100 Kilometers or more thick.  We have a fairly good estimate of the radiogenic heat in the upper
# 10Km.  We also know the temperature of the Moho (approx 30km deep) is about 600C
# and the temperature at about 100Km is about 700-900C. We will say 800C.  And we set the surface
# heat flow at 0.65 watts/m2. 
# We can do a numeric experiment  to see if all these ideas collectively agree with each other.
# can we find a combination of reasonable K and Radiogenic heat to match these temperatures
# (Note we have 6 unknowns in 4 knowns, so your solution will not be unique)

import numpy as np
import matplotlib.pyplot as plt

# Using a Matrix implicit finite difference approach
upperTbc = 0
q = .065           # heat flux, watts per square meter, in this version applied at top
                   # since that is where it is observed
K = 3.0            # conductivity (reasonable for granite)
Rq= .0000015       # radiogenic heat in Watts/m3 (reasonable for upper 10km)
Z = 100000         # thickness of problem, 100 kilometers
n = 51             # number of nodes in Z (use more for increased accuracy maybe 51)
delz = Z/(n-1)     # delta z, length of problem / number of nodes (minus 1 for the BCs)
z = np.linspace(0,Z,n) # 'z' is a vector for plotting, and a template for making other arrays

# We are going to assume there are 3 layers in the Crust, the top is down to the Moho (30km)
# Then a layer of elastic lithosphere down to 60km, and a layer of lower crust or upper mantle
# that is still fairly elastic down to the more fluid rock at about 100km
# construct the K vector, to illustrate using variable coefficients
K1 = K* .8         # conductivity is thought to increase with depth
K2 = K* 1.1          # these are our estimates of thermal conductivites of the 3 layers
K3 = K* 1.2           # unrealistic!
                    # the radiogenic heat probably falls off with depth 
Rq1 = Rq* .8       # we know the near sfc radiogenic, but not the full 30km radiogenic.
Rq2 = Rq* .4        # These are our radiogenic estimates of the 3 layers
Rq3 = Rq* .2
# we need the K transition depths, these are numbers to make sure they don't lie exactly on a node
zK12 = 29999.9       # approx depth of Moho, in meters, depth of transion from K1 to K2
zK23 = 59999.9       # somewhat arbitrary transition from elastic Lithosphere to a more ductile

# **************************************************************************************
# below here there are no hidden numbers *********************************************
Kvec = np.zeros(n-1)   # there is one less layer than nodes, for both K and radiogenic
Rqvec= np.zeros(n-1)   # make a vector for the radiogenic heat production
                       # even though we have little data, we can try varying this
# this is easy to read but not very efficient, but assigns values to Kvec (and Rqvec if needed)
# there are fancy 'pythonic' ways of doing this using 'slicing', but for us:

for j in np.arange(len(Kvec)):        # look at the depth of each node
    Kvec[j] = K2
    Rqvec[j]= Rq2
    if z[j] < zK12: 
        Kvec[j] = K1                  # and assign a K to the layers below the K1 to K2 node
        Rqvec[j]= Rq1
    if z[j] > zK23: 
        Kvec[j] = K3             
        Rqvec[j]= Rq3
# Kvec is one less than nodes, and each K applies to the layer below the node

# pythonic way of doing this would be: ( although it is a bit difficult to read!)
#Kvec = [ K1 if i < zK12 else K2 if i < zK23 else K3 for i in z[:-1]]
#Kvec = np.array(Kvec)                # note kvec is used as an array, not a list later (cant add lists!)

# matrix construction, for the implicit FD, first construct the main diagonal 'a'
# then the upper and lower diagonals 'b' and 'c'
# we only construct the non BC diagonals here
# note this is actually quite subtle to get the correct placement
# since both BCs are at the top, the diagonals are shifted down one row 
b       = np.zeros(len(z)-1)       # length one more than Kvec
b[1:]   = Kvec[1:]
a       = np.zeros(len(z))    # be careful to make new vectors, not just new names
a[1:-1] = -Kvec[:-1] -Kvec[1:]   # making the A matrix, 'a' is -K[i-1]-K[i]
c       = np.zeros(len(z)-1)
c[:-1]  = Kvec[:-1]                  # make b and c the correct length
# because the b and c vectors are placed as below, this works
# np.diag() is a numpy shorthand for creating square matrices with non-zero diagonal entries
A = np.diag(a,0) + np.diag(b,1) + np.diag(c,-1);  # b and c end up as 'b' K[i], 'c' K[i-1]

# add boundary conditions to the A matrix
A[0,0] = 1            # adjust the A matrix for the top BC 
A[0,1] = 0

A[-1,0] = -1         # gradient BC sets the slope between node 0 and 1
A[-1,1] = 1          # gradient BC
# view the A and C arrays to see the difference between the position of the BCs

C = np.ones(n);     # make the BC vector for the right hand side
# Putting Rqvec into the C vector has to jog down one position since there are
# 2 BCs at the top, so the matrix main diagonal is shifted down one.
C[1:-1]= -Rqvec[1:]*delz*delz  # the first Rqvec[0] is lost in the gradient BC condition
C[0] = upperTbc              # upper BC is 0 degrees
C[-1] = delz*q/Kvec[0]        # gradient BC on the 2 top nodes

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

# make our basic plotting canvas, with a title in the header bar
depth = z/1000
fig = plt.figure()
ax1 = fig.add_subplot(1,1,1)
ax1.grid(True)
ax1.plot(T,depth, 'r-+')  
ax1.plot([0,T[-1]],[0, 0], 'g')
ax1.text(0,24,"K={0:4.2f}, Radio_q={1:5.2e}".format(K1,Rq1)) # put notations on the plot for K and Rad
ax1.plot([0,T[-1]],[zK12/1000, zK12/1000], 'b:')
ax1.text(0,45,"K={0:4.2f}, Radio_q={1:5.2e}".format(K2,Rq2))
ax1.plot([0,T[-1]],[zK23/1000, zK23/1000], 'b:')
ax1.text(0,85,"K={0:4.2f}, Radio_q={1:5.2e}".format(K3,Rq3))
ax1.set_title("Geothermal Temperature Profile, Radiogenic heat sources")
ax1.set_xlabel("Temperature (C)")  
ax1.set_ylabel("Depth [kilometers]")  
ax1.invert_yaxis()  # this is a hard to find, but useful pyplot command

plt.show()