# Geol 5470, Humphrey 2018
# Our first, very straightforward 2D 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
# the only difference here is the plotting in 3D

import matplotlib.pyplot as plt
import numpy as np

#***** our problem assumes the block is 1500m wide, 1000m deep, with a sfc
# temperature of 1, base temperature of 7 degrees, right side of 3, and 
# left side of 6 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.

# discretize the domain
lengthx = 1000      # width
lengthz = 1000      # depth
nodes_x =  25        # in 2D, you rapidly run out of memory!, try 200 nodes!!
delx    = lengthx/(nodes_x-1)
pdelx    = lengthx/(nodes_x) # note this is for plotting only!!
# the solution to the memory problem is to use SPARSE arrays in scipy 
nodes_z = int(nodes_x * (lengthz/lengthx))

# BCs 
BCtop    = 1000
BCbottom = 0
BCright  = 500
BCleft   = 700

nodes   = nodes_z*nodes_x   # nodes is the total length of the Temperature field vector
                            # including all the unknowns and the BC nodes

# Our 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

# 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

# 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 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.

# make an A matrix with '1's on the diagonal, zeros elsewhere
A = np.identity(nodes)

for i in np.arange(nodes_z,nodes-nodes_z):   #this loops over all the nodes, 
    # except the left and right hand sides (we are counting column-wise)
    # the ith row of the A matrix is the equation for the ith node
    # check if we are at a boundary node, if not then do:
    if (((i+1)%nodes_z) !=0 ) and ((i%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[i,i]=alpha
        A[i,i-1]=beta
        A[i,i+1]=beta
        A[i,i-nodes_z]=beta      #the node to the left of a node is i-nodes_z
        A[i,i+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 
#C[-1]                               =BCbottom      #improves symetry of plot (nothing else)

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


# NOTE in an array the 1st index is the row number, the 2nd the col (in 3D the 3rd is the col,
# the 2nd is the row, and the 1st is the plane)
# Thinking in 2 and 3D you need to be very careful, so below we want to make an array
# since we numbered down the rows as our index, we are numbering in the WRONG direction
B = Tc.reshape(nodes_x,nodes_z)   # reshaping is very useful for plotting, 
# this produces an array with the col numbers going down and the rows across, in other words
# it is the transpose of the array we want
Bt = B.T   # arrays have a '.T' method to transpose


# This is another useful figure, a contour plot.  Here we make the contours stand out as black
fig4 = plt.figure()
ax4  = fig4.add_subplot(111)
cs = ax4.contour(Bt,colors ='k')   # contour just plots the contours
ax4.contourf(Bt)                   # contourf fills the contours with colors
ax4.clabel(cs,inline=1,fontsize=12)      # clabel puts labels on the contours
ax4.invert_yaxis()
ax4.set_ylabel("Depth")
ax4.set_xlabel("Width in Meters")
ax4.set_title("Temperature field as a contour")