# Numerical Modeling in the Geosciences,  Humphrey 2018
# First homework problem in heat flow 
# ***************************************************************************
# Problem: steady state temperature field in a 1000 m thick layer of rock,
# with a surface temperature of 0degrees and a heat flux of 65 mW per square meter
# BC applied at the top of the problem, 
# The governing eqn is q = -K dT/dz
# ***************************************************************************

import numpy as np
import matplotlib.pyplot as plt

K = 3                   # thermal conductivity 'granite',
n = 9   		# number of internal nodes, total nodes (including BCs) =n+1
thickness = 1000	# thickness of rock layer in vertical, in meters

delx = thickness/(n+1)	# delx used for plotting
upperBC = 0 		# surface temperature BC, in degrees
upperq  = 6.5e-2	# W/m2 upper typical geothermal heat flux (this is used as the derivative, or gradient condition at the surface)

# actual program
T = np.zeros([n+1])
T[0] = upperBC
T[1] = T[0] + upperq*delx/K    # apply the gradient condition by calc the diff between T[0] and T[1]

for z in np.arange(2,n+1):     # the derivative is just a constant!!, so very little calc!!
    T[z] = T[z-1] + upperq*delx/K

#plot result (I always like to plot the results to see if there are mistakes
fig = plt.figure()                    #create a place to plot
ax  = fig.add_subplot(111)            #create axes.  you can create multiple plots on the same figure with subplot(2,3,1) example

ax.plot(T, np.linspace(0,thickness,len(T)),'r-o')
#ax.set_ylim(T[0],T[-1])
ax.invert_yaxis()
ax.set_title('Steady State temperature in a rock slab')
ax.grid()
ax.set_ylabel('Depth from surface in meters')
ax.set_xlabel('Temperature in degrees')

plt.show()