# Geol 5470, Humphrey 2022
# 2D implicit Finite Difference problem using different delx and dely
# uses constant coeffiecients and is steady state
# Rectangular problem domain. full (not sparse) matrices, with BCnodes included in matrix
# Solve temperature field in a rectangular piece of lithosphere
# This version illustrates an insulated boundary on left and non symetric x and z spacing (delx delz)

import matplotlib.pyplot as plt
import numpy as np

#***** our problem assumes the block is 1500m wide, 1000m deep, with a sfc
# temperature of 0C, base temperature of 100 degrees, right side of 30
# Note, since K is constant, we can divide it out of the problem, 
# also, less obvious, we could divide the dimensions out of the problem
# if we make delx and delz the same!! (this requires that the ratio of the
# number of nodes in x over the nodes in z be the same as the ratio of the
# side lengths. 

# discretize the domain
lengthx = 1500      # width
lengthz = 1000      # depth
nodes_z = 30        # 

# BCs 
BCtop    = 0
BCbottom = 100
BCright  = 30
BCleft   = 60    # not used

nodes_x = nodes_z   # here we make the delx and delz different
delx = lengthx/(nodes_x-1)
dx2 = delx*delx     # we only use the dx**2 in the matrix
delz = lengthz/(nodes_z-1)
dz2 = delz*delz
nodes   = nodes_z*nodes_x   # nodes is the total length of the Temperature field vector
                            # including all the unknowns and the BC nodes

# Our discretized domain mesh is a grid, our numbering scheme is that the nodes are numbered
# from 0 starting in the top left corner.  Numbers increase downward
# (column-wise) since the problem is wider than it is deep.  In large
# problems it is best to number in the shortest dimension of the grid.
# Therefore the left hand side numbers from 0 to nodes_z - 1, (we use 'm' as a mesh index in this program)
# Our temperature field will be stored as a vector, But we can also think of it as a matrix of temperatures, 
# one at each node; the matrix form matches the shape of our nodal grid.

# Our discretization, and our matrix solution technique, lead to 3 different numbering schemes:
# 1- the problem domain is a mesh of nodes, the nodes are numbered using variable 'm' from 0 to (nodes-1)
# 2- this mesh can also be be thought of as an i,j grid, where the rows are numbered in i and the cols in i
# 3- finally in our solution technique, we use a matrix with dimensions 'nodes'x'nodes', and a location in
#    that matrix is given by ii, (rows), and jj, (cols),  where each row (ii) is the equation for the 'm-th' 
#    node in the discretized domain, and therefore the 'C' vector's 'm' entry is the forcing or the BC 
#    for the 'ii' row in the 'A' matrix. 
# Important note, we use the fact that 'm' equals 'ii' in assembling the A matrix!
# And also note that we use a grid of xx,yy for the 'real world' node locations int the plotting routines

# all BCs at boundaries are fixed
# steady state and constant K
# values of the diagonal and off diagonal coeff (but in 2D these appear in
# 'bands' in the matrix, not merely on the main diagonal
alphax = -2/dx2       # value of main diagonal
alphaz = -2/dz2      # this will make the values in the A matrix non-symetric in delx and delz
betax  =  1/dx2       # value of both off diagonals
betaz  =  1/dz2

# (better in large problems would be to use sparse matrices below)

# for 2D or 3D problems, most of the work is assembling the A matrix
# It is ABSOLUTELY necessary that you have a clear picture of the indexing of your nodes!
# construct the non-trivial bands in the matrix, using a 5 point FD kernel centered on each node:-
# beta*node to left, zeros, beta*node above, alpha*node, beta*node below, zeros, beta*node to right

# You can assemble the A matrix using fancy 'slicing', but for this first time we will use a very 
# explicit assembly by looping over all the nodes to check if they are boundary nodes or not,
# and if they are non-zero nodes.

# make an A matrix with '1's on the diagonal, zeros elsewhere
# the ith row of the A matrix is the equation for the ith node
A = np.identity(nodes)   # big note, this makes all eqns, BC eqns!

for m in np.arange(0,nodes_z):  # write a insulated BC condition on the left nodes, put in matrix A 
    A[m,m+nodes_z] = -1         # which is the Temp gradient in x is zero (T[i]-T[i+nodes_z]=0]
                                # note A[i,i] is already '1'

for m in np.arange(nodes_z,nodes-nodes_z):   #this loops over all the interior nodes,
    # starts away from first boundary (left) and stops at right boundary
    # check if we are at a boundary node, if not then do:
    if (((m+1)%nodes_z) !=0 ) and ((m%nodes_z) !=0):    #leave BC nodes as '1's  
        #boundary bottom nodes are divisible by nodes_z-1, while top nodes are 1 more
        A[m,m]= alphax + alphaz
        A[m,m-1]=betaz
        A[m,m+1]=betaz
        A[m,m-nodes_z]=betax      #the node to the left of a node is i-nodes_z
        A[m,m+nodes_z]=betax

# make the known vector, with BC's in correct locations
C = np.zeros(nodes)
# the following makes good use of python 'slicing'
# (see if you can figure out what is being done!)
C[:nodes_z]                         =0             # zero gradient
C[nodes_z:nodes-nodes_z-1:nodes_z]  =BCtop     
C[2*nodes_z-1::nodes_z]             =BCbottom      #bottom BC
C[nodes-nodes_z:]                   =BCright   

# plot a SPY diagram, very useful to find out if your indexing is incorect!
fig1 = plt.figure()
ax1  = fig1.add_subplot(111)
ax1.spy(A)
ax1.set_title("SPY plot of A Array")
plt.show()

#solve
Tc = np.linalg.solve(A,C)       # direct solver in numpy, Tc now contains the solution vector of temperatures

# very simple plotting, just to introduce 2D plotting
# we will improve on this as we go along, this plots in nodal coords
fig2 = plt.figure()                                                           
ax2 = fig2.add_subplot(1,1,1)
B = Tc.reshape(nodes_x,nodes_z)   # reshaping is very useful for plotting, changes T into a matrix matching our grid
ax2.pcolor(np.arange(nodes_x),np.arange(nodes_z),B.T,shading='auto')  # the '.T' is the transpose
# We have to take the transpose since we indexed downwards, but pyplot numbers across
ax2.invert_yaxis()  # always good in Geoplots to make the surface at the top
ax2.set_ylabel("Depth (nodes)")
ax2.set_xlabel("Width (nodes)")
ax2.set_title("Temperature field")

# fancy plot (there are some things here we haven't discussed) plots in 'meters' coords
xx,zz = np.meshgrid(np.arange(nodes_x)*delx,np.arange(nodes_z)*delz) #produces a complete set of x,z coords
# finally a plot showing the directions of heat flow, overlayed over the psuedo color plot
fig5 = plt.figure()
ax5  = fig5.add_subplot(111)
                      # gradient produces the difference arrays in the 2 axis directions
U,V = np.gradient(B.T) # this gives the positive gradient in x,z, returns 2 arrays with these gradients
ax5.invert_yaxis()
ax5.pcolor(np.arange(nodes_x)*delx,np.arange(nodes_z)*delz,B.T,shading='auto') # matrixes have a method '.T' to transpose
# We want the heat to flow from hot to cold, so we need negative, also again the directions are reversed
ax5.quiver(xx,zz,-V,U)
ax5.set_ylabel("Depth in meters")
ax5.set_xlabel("Width in Meters")
ax5.set_title("Temperature field with heat flow directions and Gradients")