# feb, 2020,  Numerical Modeling in the Geosciences,  Humphrey
# homework problem in transient heat flow, using an implicit matrix formulation
# This version tries to adjust for the low heat transfer to the atmosphere
# ***************************************************************************
# Problem: Transient temperature field in a 100m thick layer of fresh magma,
# overlying rock with a temperature of 0C
# BCs are surface temperature of 0 degrees and an initial rock temperature of
# 0degreees.  Initial magma conditions are 1000 degrees: 
# the surface condition is a gradient condition, we start with zero flux
# BC's are steady in time.  This code uses a constant K, rho and Cp
# It uses a value of K/(rho*Cp) of 30 (m^2 per year) to give realistic
# values for rock.  This ratio is called the thermal diffusivity (k).
# We model both the 1Km of magma, and the underlying 1km of initially cool
# rock, using the same diffusivity for both.
# our eqn is  dT/dt = k d2T/dz2
# ***************************************************************************
import matplotlib.pyplot as plt
import numpy as np

# set problem parameters here
delt = 1                # time step in years
tmax = 250              # max time in years
k  = 30                 # K/(rho*Cp), thermal diffusivity
n = 50   			    # set the number of nodes
thickness = 300         # set thickness of magma + rock layer in vertical, in meters
delx    = thickness/(n-1)      # delx used for plotting (should really be z
lowerBC = 0 		    # set temp BCs, in degrees
upperBC = 0             # this gradient condition says there is no heat flow (insulated)
#upperBC = 2.2*delx             # this is a gradient condition, assume that the air has little conductivity
magma_temperature = 1000
nodes_in_magma = int(n/3)
nodes_in_rock = n-int(n/3)

# *********everything below here is written with only variables, no numbers ***********
				            
z = np.linspace(0,thickness,n) # nodal locations
                               
A   = np.zeros([n,n])          # coefficient matrix, A, 
                               # now make intial condition vector of temperatures
Tn = np.append( np.ones(nodes_in_magma)*magma_temperature, np.ones(nodes_in_rock)*lowerBC ) 
To  = np.zeros(n)              # also make a place to keep our old values of temperature before each tstep
dx2 = delx*delx                # the 2nd derivative term is divided by dx2, so calc here to save computing

# we can assemble the A matrix outside the time loop since in this simple case it is a constant
# the A matrix is simply a tri-diagonal with beta on the main, and alpha on the off diagonals
alpha = -delt*k/dx2 
beta  = -2*alpha + 1
A     = np.diag(np.ones(n)*beta,0) + np.diag(np.ones(n-1)*alpha,1) + np.diag(np.ones(n-1)*alpha,-1)  
# note efficient adding of arrays

# add boundary conditions to the A matrix
A[0,0] = -1            # adjust the A matrix for the top BC, this is a gradient condition 
A[0,1] = 1
# NOTE the negative index below counts back from end!
A[-1,-2] = 0          # bottom BC
A[-1,-1] = 1          # 
# and add the BCs to the Tn vector (the right hand side vector)
Tn[0] = upperBC             # 
Tn[-1]= lowerBC
# you can view the A and Tn arrays to see the difference between the position of the BCs

# 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.set_title("Magma Flow, Temperature Profile over time")
ax1.set_xlabel("Temperature, degrees C")  
ax1.set_ylabel("Depth meters")  
ax1.invert_yaxis()  
ax1.text(400,150,'Red curves are 25 years apart');

plt.show()

# now for the time loop
# note this illustrrates a simple method of 'animation', although matplotlib gives a spurious error msg!
for t in np.arange(0,tmax,delt):     # this is the time loop to take one step forward in time
    Tn[0] = upperBC          # the gradient BC needs to be reladed each time because the temperature at Tn[-1]
                        # does not stay at zero!
    To = Tn             # move the new temperature to the old temperature, to start the time step    
    Tn = np.linalg.solve(A,To)   # direct solver in numpy			# Tn now contains the solution of temperature

    line = ax1.plot(Tn, z,'-k')      # 'line' is a list of lines (although here it is actually just one)
    plt.pause(0.01)                  # put the line up long enough to see it
    if  t%25 > 0:                    # leave every 25 year curve on the screen
        ax1.lines.remove(line[0])
    else:
        line[0].set_color('r')       # left over lines are colored red