# May, 2021,  Numerical Modeling in the Geosciences,  Humphrey
# Transient 1D heat flow, outlines the construction of the F.E. matrices.
# Uses implicit Finite Difference for the time stepping
# ***************************************************************************
# Problem: Transient temperature field in top 5m of a granite mass
# subject to daily temperature swings.
# BCs are surface temperature of 0 degrees and a constant 5m temperature of
# 0degreees.  Initial conditions are 0 degrees: the surface BC is a +/- sinusoid
# of 10C amplitude, representing a diurnal swing
# This code uses a constant K, rho and Cp
# It uses a value of K/(rho*Cp) of 0.1 (m^2 per day) to give realistic
# values for rock.  This ratio is called the thermal diffusivity (k).

# our eqn is  dT/dt = k d2T/dz2
# ***************************************************************************
import matplotlib.pyplot as plt
import numpy as np

# set problem parameters here
delt = .01                # time step in fractional days
tmax = 4                  # max time in days
k    = .1                 # K/(rho*Cp), thermal diffusivity (days meters)

# create the nodes in z and therefore the 1D elements
#z = np.array([0,.1,.2,.4,5.]) # a short discretization for testing
z = np.array([0,.025,.05,.1,.15,.2,.3,.4,.6,.8,1.6,3.2,5])
n = len(z)
# Initial Conditions
To = np.zeros_like(z)
Tn = np.zeros_like(To) # and a place to put changed temperatures

# set temp BCs, in degrees
lowerBC  = 0   # could probably be a gradient condition instead
# upperBC we write as a 'lambda' function 'upperBC()'
upperBC = lambda t: 10*np.sin(2*np.pi *t)  # called as upperBC(t) and returns the BC(t)
# lambda functions are in-line functions (shorthand form of a 'def' block)

## *********everything below here is written with only variables, no numbers ***********
# make a list of element lengths
Le = np.diff(z)
Le_inv = 1/Le     # and 1/length of elements
                               
# coefficient matrix, A, constructed VERY piecemeal to show it explicitly
Alower = -np.diag(Le_inv,-1)   #lower diagonal of A
Aupper = -np.diag(Le_inv,1)    #upper diagonal
Avec1  = np.append(Le_inv,[0]) #first integral of main diag
Avec2  = np.append([0],Le_inv) # 2nd integral of main diag
Amid   = np.diag(Avec1+Avec2,0)
A      = Amid + Aupper + Alower   # put it all together
A      = k*A                   # and remember the 'k'

# coefficient matrix, B, similar to A
Blower = np.diag(Le/6,-1)
Bupper = np.diag(Le/6,1)
Bvec1  = np.append(Le/3,[0])
Bvec2  = np.append([0],Le/3)
Bmid   = np.diag(Bvec1+Bvec2,0)
B      = Bmid + Bupper + Blower
B      = B* 1/(delt)          # remember the 1/delt

# forcing vector (f), which in this problem is 0, until we FD in time
f = np.zeros(n)

Astar = A + B              # add the 2 matrices to get our A star matrix
Astar[0,0] = 1             # set the upper and lower rows of Astar as BC rows
Astar[0,1] = 0
Astar[-1,-1] = 1
Astar[-1,-2] = 0 

# make our basic plotting canvas, with a title in the header bar
fig = plt.figure()
fig.canvas.set_window_title ("Diurnal Temperature in exposed granite")
ax1 = fig.add_subplot(1,1,1)
ax1.grid(True)
ax1.set_title("Diurnal Temperature Profile over time")
ax1.set_xlabel("Temperature in Granite in C")  
ax1.set_ylabel("Depth below surface, meters")  
ax1.set_xlim(-12,12)
ax1.set_ylim(0,1)   # only show the upper meter, very little change below 1m
ax1.invert_yaxis()  # always good in Geoplots to make the surface at the top
plt.show()          # show the basic figure, we add to it below

# now for the time loop, this is a Finite Difference fully implicit stepping in time
n=0                 # n is used in plotting to remove most lines, only keeping some
for t in np.arange(0,tmax,delt):     # this is the time loop (in t) to take one step forward in time
    f = np.dot(B,To)                 # remember that B*To means element-wise multi, not matrix multi!
    f[0] = upperBC(t)                # put upper and lower BCs in f, upperBC changes every time step
    f[-1]= lowerBC
    Tn = np.linalg.solve(Astar,f)        # direct solver in numpy, Tn now contains the solution of temperature
    if t > 1.:                           # let it run for a year to 'spin up' the model from poor ICs
        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  n%12 > 0:                    # leave every 12th curve on the screen
            ax1.lines.remove(line[0])    # To is the curve at t, Tn is at t+delt
        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  
    n +=1