# -*- coding: utf-8 -*-
"""
Created on Mon Feb 17 11:07:09 2020

@author: neil
"""

# Numerical Modeling in the Geosciences,  Humphrey 2022
# 2nd problem, include radioactive heat generation, 
# ***************************************************************************
# Problem: steady state temperature field in a 20000 m thick layer of rock,
# with a surface temperature of 0degrees and a heat flux of 65 mW per square meter
# This includes reasonable radioactive heating
# We use the governing eqn Q = -K dT/dz, but this requires varying Q(z) to account for radiogenic heat
# A better approach is to use d2T/dz2=q, where q is the local radiogenic heat, 
# illustrated in the other example code on web page
# ***************************************************************************

import numpy as np
import matplotlib.pyplot as plt

K = 3               # thermal conductivity 'granite', J/(m s C)
n = 9   			# number of nodes we need to calculate, total nodes (including BCs) =n+1
thickness = 20000	# thickness of rock layer in vertical, in meters
Rq= .0000015        # nominal radioactive heat generation in upper crust J/(s m^3)
CummulativeRq = Rq*thickness # Note this is mixing an integral equation with our dirivative eqn, since we are integrating
                             # the heat sources over the depth and making assumptions, such
                             # as no heat leaks downwards
upperq  = 6.5e-2	         # W/m2 upper typical geothermal heat flux (this is a gradient condition)
if CummulativeRq > upperq :
    print("Error, cummulative radiogenic heat greater than surface heat flux")


delx = thickness/(n+1)	# delx used for plotting
upperBC = 0 		    # temp BCs, in degrees

# actual program
T = np.zeros([n+1])                                # There are n+1 nodes
T[0] = upperBC
T[1] = T[0] + (upperq)*delx/K                      # apply the gradient condition by calc the diff between T[0] and T[1]
CummulativeHeat = upperq - Rq*delx                 # since we are starting beow the sfc, subtract the radiogenic in layer

for j in np.arange(2,n+1):                         # integrate from the surface down
    T[j] = T[j-1] + CummulativeHeat*delx/K         # this the the Euler step (could use R-K for better error propagation)
    CummulativeHeat -= Rq*delx                     # as we go down, there is less radioactive heat from below, flowing up

#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 crust (with radiogenic heat)');
ax.set_ylabel('Depth from surface in meters')
ax.set_xlabel('Temperature in degrees')
ax.grid()

plt.show()