# Geol 5470, Humphrey 2022
# Our first, very straightforward 2D implicit Finite Difference problem
# uses constant coeffiecients and is steady state
# Rectangular problem domain.  Simplest approach
# full (not sparse) matrices, with BCnodes included in matrix
# Solve temperature field in a rectangular piece of lithosphere, assuming
# we know all the bounding temperatures.  So it is steady state, fixed BCs.

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, and 
# left side of 60 degrees.
# 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.  For this first problem, to make it trivial, we will do that
# and divide delx or delz out of the problem.

# this first program just illustrates the techniques, the plot needs improvement 
# since they don't plot the actual dimensions for the 'real world' problem

# discretize the domain
lengthx = 1500      # width, (in this program we don't even use these x,z dimensions)
lengthz = 1000      # depth
nodes_z = 3         # in 2D, you rapidly run out of memory!, try 200 nodes!!
# (advanced Note, the solution to the memory problem is to use SPARSE arrays in scipy) 
nodes_x = int(nodes_z * (lengthx/lengthz))

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

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, we use the fact that 'm' equals 'ii' in assembling the A matrix! often ii,jj indexing is not needed

# 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
alpha = -4       # value of main diagonal
beta  =  1       # value of both off diagonals

# 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.
# so a row in the A matrix is: 
# zeros, beta*node to left, zeros, beta*node above, alpha*node, beta*node below, zeros, beta*node to right, zeros

# 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 m-th row of the A matrix is the equation for the m-th node
A = np.identity(nodes)   # big note, this makes all eqns, BC eqns! (but rows are changed below)

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 (((m+1)%nodes_z) !=0 ) and ((m%nodes_z) !=0):    #leave BC nodes as '1's  
        # fill in the rows that represent interior nodes
        # boundary bottom nodes are divisible by nodes_z-1, while top nodes are 1 more
        A[m,m]=alpha
        A[m,m-1]=beta
        A[m,m+1]=beta
        A[m,m-nodes_z]=beta      #the node to the left of a node is i-nodes_z
        A[m,m+nodes_z]=beta

# 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]                         =BCleft        # 
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
fig2 = plt.figure()                                                           
ax2 = fig2.add_subplot(1,1,1)
B = Tc.reshape(nodes_x,nodes_z)   # reshaping is very useful for plotting, 
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
# see the third program for more notes on this
#B = Tc.reshape(nodes_z,nodes_x,order='F') #reduces this to one command, but is hard to understand!
#ax2.pcolor(np.arange(nodes_x),np.arange(nodes_z),B)
#ax2.grid()
ax2.invert_yaxis()  # always good in Geoplots to make the surface at the top
ax2.set_ylabel("Depth")
ax2.set_title("Temperature field")