# program to solve temperature in the lithosphere
# This version has no heat sources But it changes the thermal properties
# Mainly illustrates the use of the implicit method, and the matrix eqn solver
# This version places the gradient BC at the base of the solution space
# We use 3 different layers to mimic the increasing conductivity as the 
# density and temperature increases in the crust

# Big Note: placement of nodes on the boundaries between K layers is crucial
# to accurately modeling the temperature kinks that occur at real geologic boundaries

import numpy as np
import matplotlib.pyplot as plt

# Using a Matrix implicit 2nd order finite difference approach
q = .065           # heat flux, watts per square meter, in this version applied at base
K = 3.0            # conductivity (reasonable for granite)
Z = 100000         # thickness of problem, 100 kilometers
n = 11             # number of nodes in Z
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

# construct the K vector, to illustrate using variable coefficients
# first we define the values of the 3 layers, K1, K2, K3
K1 = K
K2 = K*1.4 
K3 = K*2.2       # K increases in steps with depth, ie heat transfer gets easier with depth
# And we need the transition depths, we don't list depths at the exact nodal locations since
# that might cause ambiguity with our 'if' statements
zK12 = 29999
zK23 = 59999       # transition depths in meters

# Note it is good practice to keep any parameters in the header section
# burying important numbers (such as the value of conductivity or whatever) in the code
# IS BAD FORM
# below here there are no hidden numbers *********************************************
Kvec = np.zeros(n-1)                  # there is one less inbetween K value than nodes
# this is easy to read but not a very efficient way to make a vector of conductivities
# there are fancy 'pythonic' ways of doing this, but for us:

for j in np.arange(len(Kvec)):        # look at the depth of each node
    if z[j] < zK12: 
        Kvec[j] = K1                  # and assign a K to the layer below the node
    elif z[j] < zK23:
        Kvec[j] = K2
    else: 
        Kvec[j] = K3             
# Kvec is one less than nodes, and each K applies to the layer below the node
# in other words, K[i] lies inbetween T[i] and T[i+1]
# It is important that the K values apply to the whole region inbetween nodes and does not 
# change values inbetween nodes, since that invalidates our abstraction of variable K

# matrix construction, for the implicit FD, first construct the main diagonal 'a'
# then the upper and lower diagonals 'b' and 'c'
# we could construct the matrix cell by cell, but numpy has a nice way of making diagonal arrays
# called np.diag(a,n), where a is a vector, which makes a square array of size len(a) by len(a) ,
# and the vector is placed at the main diagonal, plus or minus n.
# we only construct the non BC diagonals here
# note this is actually quite subtle to get the correct placement
# since the K values have to apply to the regions between nodes (see notes)

a       = np.zeros(len(z))       # make a vector that will be the length of the diagonal of the array
a[1:-1] = -(Kvec[:-1] + Kvec[1:])#  note, if you add vectors or arrays in numpy, you are adding
                                 # cell by cell ie. for each i, a[i] = -(Kvec[i-1] + Kvec[i])
b = c   = Kvec                   # make b and c the correct length, note the Python shorthand!
# because the b and c vectors are placed offset as below, this works
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]
# if you don't understand this, try printing out the matrix
# What is happening is that numpy makes an array A, by adding 3 arrays created by the 3 calls to 
# np.diag.  The first diag call places 'a' vector on the main diagonal, then you are adding
# the b and c diagonals above and below the main diagonal.

# add boundary conditions to the A matrix, note the direct addressing; A[row,col]
A[0,0] = 1            # adjust the A matrix for the top BC 
A[0,1] = 0            # this sets a proscribed value to the top Temperature
# NOTE the negative index below counts back from end!
A[-1,-2] = -1         # gradient BC sets the slope between node 0 and 1
A[-1,-1] = 1          # as a temperature gradient BC
# view the A and C arrays to see the difference between the position of the BCs

C=np.zeros(n);       # make the BC vector for the right hand side
C[0] = 0             # upper BC is 0 degrees
C[-1] = delz*q/K3    # gradient BC on the 2 bottom nodes

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

depth = z/1000       # convert depths to Kilometers for plotting
# 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([0,T[-1]],[0, 0], 'g')
ax1.plot([0,T[-1]],[zK12/1000, zK12/1000], 'b:')
ax1.plot([0,T[-1]],[zK23/1000, zK23/1000], 'b:')
ax1.set_title("Geothermal Temperature Profile, no heat sources")
ax1.set_xlabel("Temperature")  
ax1.set_ylabel("Depth [kilometers]")  
ax1.invert_yaxis()  

plt.show()