# April, 2018,  Numerical Modeling in the Geosciences,  Humphrey
# homework problem in transient heat flow, using an implicit matrix formulation
# This outlines the construction of the F.E. matrices, it leaves the solution to you.
# ***************************************************************************
# Problem: Transient temperature field in a 3km thick layer of ground,
# overlying a injected dyke with a temperature of 800C
# BCs are surface temperature of 0 degrees and a constant dyke temperature of
# 800degreees.  Initial conditions are 0 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 only the 3Km of earth.
# 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 (yrs meters)
n = 30   			    # set the number of nodes
thickness = 3000        # set thickness of rock layer in vertical, in meters
upperBC = 0 		    # set temp BCs, in degrees
lowerBC = 800 

# make a mesh of nodal positions, in this simple problem I put 3 nodes close to the lower boundary
# and make them evensized to the surface
z1 = np.linspace(0,thickness-200,n-3)          # nodal locations
z2 = np.linspace(thickness-125,thickness,3)    # nodal locations
z = np.append(z1,z2)
# make a list of element lengths
Le = np.diff(z)
Le_inv = 1/Le

## *********everything below here is written with only variables, no numbers ***********
                               
# 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 of main diag
Amid   = np.diag(Avec1+Avec2,0)
A   = Amid + Aupper + Alower   # put it all together

# coefficient matrix, B
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

# forcing vector (f)
f = np.zeros(n)
f[-1] = 800

# finally add the boundary condition eqns to the A and B matrices
# we do this by replacing the upper and lower rows with the simple equation
# that the value of T is the BC in 'f'
A[0,0] = 1            #upper
B[0,0] = 0
A[0,1] = 0
B[0,1] = 0

A[-1,-1] = 1          #lower
B[-1,-1] = 0
A[-1,-2] = 0
B[-1,-2] = 0

# you can proceed from this point using your Finite Difference code since you have
# all the necessary matrices and vectors