# feb, 2022,  Numerical Modeling in the Geosciences,  Humphrey
# homework problem in transient heat flow, using an implicit matrix formulation
# 
# ***************************************************************************
# Problem: Transient temperature field in a 100m thick layer of fresh lava,
# overlying rock with a temperature of 10C
# BCs are surface temperature of 0 degrees and a base rock temperature of
# 10degreees.  Initial lava conditions are 1000 degrees: the
# 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 100m of lava, and the underlying 200m 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), in years, thermal diffusivity, for this problem make same everywhere
n = 50   			    # set the number of nodes
thickness = 300         # set thickness of lava + rock layer in vertical, in meters
upperBC = 0 		    # set temp BCs, in degrees
lowerBC = 10            # the lower BC will also be used for intial condition on lower rock temperature
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 ***********
delz    = thickness/(n-1)      # delz used for plotting				               
z = np.linspace(0,thickness,n) # nodal locations
                               
A   = np.zeros([n,n])          # coefficient matrix, A, 
                               # now make intial condition vector of temperatures, To - Told, Tn - Tnew
To = np.append( np.ones(nodes_in_magma)*magma_temperature, np.ones(nodes_in_rock)*lowerBC ) 
Tn  = np.zeros(n)              # also make a place to keep our old values of temperature before each tstep
dz2 = delz*delz                # the 2nd derivative term is divided by dz2, 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/dz2 
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)  

# now add boundary condition indexes to the A matrix
A[0,0] = 1            # adjust the A matrix for the top BC 
A[0,1] = 0
A[-1,-2] = 0          # bottom BC
A[-1,-1] = 1          # 
# and add the BCs to the Tn vector (the right hand side vector)
To[0] = upperBC             # 
To[-1]= lowerBC       # you can view the A and Tn arrays to see the position of the BCs and diagonals

# make our basic plotting canvas
fig = plt.figure()
ax1 = fig.add_subplot(1,1,1)
ax1.grid(True)
ax1.set_title("Lava Flow Temperature Profile over time")
ax1.set_xlabel("Temperature in C")  
ax1.set_ylabel("Depth meters")  
ax1.invert_yaxis()                   # always good in Geoplots to make the surface at the top
ax1.text(400,150,'Red curves are 25 years apart');

plt.show()          # show the basic figure, we add to it below

# now for the time loop, this is a fully implicit stepping in time
for t in np.arange(0,tmax,delt):     # this is the time loop (in t) to take one step forward in time
 
    Tn = np.linalg.solve(A,To)       # direct solver in numpy, Tn now contains the solution of temperature
    print(t)                         # print into the console, so we know it is working
    line = ax1.plot(To, 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])    # To is the curve at t, Tn is at t+delt, so we are actually plotting the 'old' temperatures
    else:
        line[0].set_color('r')       # left over lines are colored red
    To = Tn                          # move the new temperature to the old temperature, to start the time step