# Geol 5470, Humphrey 2020
# Modifies the variable K 2D program to place a blob of highly radiogenic 
# material in the middle of the problem.  A source term Qrad is put in the 'C' vector
# 2D Finite Difference problem, using variable K coeffiecients but still steady state
# Rectangular problem domain.  This is similar to the basic program, but delx does not equal
# delz, K is allowed to vary, and there is a source/sink term in the forcing vector.
# 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
# But variable K increases the difficulty of indexing considerably

import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

#***** 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 not constant, we can't 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 problem, to make it easier, we will do that
# and divide delx or delz out of the problem.

# discretize the domain
lengthx = 1500      # width
lengthz = 1000      # depth
nodes_x =  45       # in 2D, you rapidly run out of memory!, try 200 nodes!!
delx    = lengthx/(nodes_x-1) # this is the correct delx (and delz)
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))  # this makes delx = delz

# BCs 
BCtop    = 1
BCbottom = 7
BCright  = 3
BCleft   = 6

# we will put a block of different conductivity in the middle
K1 = 3    # K in most of the region
K2 = 3   # for this version we don't vary the K (although the program will handle it if you want)

Qrad = .0015   # radiogenic source, highly radioactive blob in the middle of the problem

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 often don't need ii,jj
# And also note that we use a grid of xx,yy for the 'real world' node locations int the plotting routines

# prepare the size of the forcing vector C
C = np.zeros(nodes)
# We make a K matrix, of the same shape as the discretized nodal grid, and fill it with K1.
K = np.ones([nodes_z,nodes_x])*K1
# add a block of K2 into the middle 1/3 of the matrix
for i in np.arange(int(nodes_z*1/3),int( nodes_z*2/3)): # step over the middle nodes
    for j in np.arange(int(nodes_x*1/3), int(nodes_x*2/3)):
        K[i,j]=K2  # K is now an array of dimensions nodes_x, nodes_z
        m    = j*nodes_z + i          # calculate the nodal number, from the i,j matrix index
        C[m] = -Qrad * delx*delx/K2   # put the radiogenic heat per volume in the C vector
# BIG NOTE, K is an array the size of the problem ([nodes_x,nodes_z])
# BUT A is an array of size ([nodes,nodes]), much bigger!  It requires careful indexing
# to put the K values into the correct locations in the A matrix
# The crux idea, is that the K values lie in the middle of each group of 4 nodes in the discretized mesh, 
# not at the node locations. This is called a staggered grid. The index of each K[i,j] is the same as
# the Temperature node at the upper left of each group of 4 nodes.
# Another big NOTE, any K[i,j], where i is a bottom BC node, or j is righhand BC node
# are outside the problem: they are called phantom values, since they are not used in the solution.

# 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.
# 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 which K we should use at each node.

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

for m in np.arange(nodes_z+1,nodes-nodes_z-1):   #this loops over all the nodes, 
    # except the left and right hand sides (we are counting row-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 (((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
        # since we made K a matrix matching the nodal mesh, but A has size nodalmesh by nodalmesh
        # we need to find the K[ii,jj] that are needed in our kernal equation
        # In this loop 'i' is the center temperature node (vector index) running from 0 to nodes
        # we can translate that into a matrix K[ii,jj] index by dividing by nodes_z, as follows
        i = m%nodes_z        # ii is the row index into K, it runs repetivitly from 0 to nodes_z
        j = int(m/nodes_z)   # jj is the cols index into K, it runs repetivitly from 0 to nodes_x
        # so we can add all the surrounding K values to the center node (alpha)
        A[m,m]        =  -(K[i,j] + K[i-1,j] + K[i,j-1] + K[i-1,j-1])
        # this is simplified, you might expect terms like .5*(K[ii,jj] + K[ii-1,jj]), but they all cancel
        # and simplify to the above line
        # and now add all the beta terms
        A[m,m-1]      =.5*(K[i-1,j] + K[i-1,j-1])     # once we have ii,jj it is easy to assemble A
        A[m,m+1]      =.5*(K[i,j]   + K[i,j-1])         # as long as we remeber the Ks we need to use
        A[m,m-nodes_z]=.5*(K[i,j-1] + K[i-1,j-1])     # the i node to the left of a node is i-nodes_z
        A[m,m+nodes_z]=.5*(K[i,j]   + K[i-1,j])

# make the known vector, with BC's in correct locations

# 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")
"""

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

# we improve on the simple plotting below
fig2 = plt.figure()                                                           
ax2 = fig2.add_subplot(1,1,1)
# Tc is a temperature vector, we need to reshape it into an array with
# nodes_z rows and nodes_x cols

# 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 with
# 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, (there is a logic error here)
# 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
# pcolor plots a psuedo color image of the array, using the values as gradations in color
# you can give it x,y position arrays, or more easily give it x,y array eadge position vectors and
# it will figure out the x,y coords (this is called braodcasting)
# A subtlty is that the color is a patch to the right and down of the nodes, therefore yoiu need to
# plot an array that is 1 larger in each direction to see the full psuedo plot
f=ax2.pcolor(np.arange(nodes_x+1)*pdelx,np.arange(nodes_z+1)*pdelx,Bt)  # the '.T' is the transpose
# as a result of the expanded plot, the psuedo color image is too large, but we partially fix with 
# 'pdelx' , but if you look closely the interior nodes are not quite in the correct places
#ax2.grid()
ax2.invert_yaxis()  # always good in Geoplots to make the surface at the top

bar = plt.colorbar(f)  # colorbar places a scale to the colors of the psuedo plot
bar.ax.set_ylabel("Temperature in C")
ax2.set_ylabel("Depth")
ax2.set_xlabel("Width in Meters")
ax2.set_title("Temperature field")
plt.show()

# demonstrate a 3D surface plot.  Use meshgrid to make x and z location arrays
xx,zz = np.meshgrid(np.arange(nodes_x)*delx,np.arange(nodes_z)*delx)
fig3 = plt.figure()
ax3  = fig3.add_subplot(111, projection='3d')
ax3.plot_surface(xx,zz,Bt,cmap=cm.magma, rstride = 1, cstride = 1)
ax3.set_ylabel("Depth")
ax3.set_xlabel("Width in Meters")
ax3.set_title("Temperature field as a 3D surface")

# This is another useful figure, a contour plot
fig4 = plt.figure()
ax4  = fig4.add_subplot(111)
cs = ax4.contour(xx,zz,Bt,colors ='k')
ax4.contourf(xx,zz,Bt)
ax4.clabel(cs,inline=1,fontsize=12)
ax4.invert_yaxis()
ax4.set_ylabel("Depth")
ax4.set_xlabel("Width in Meters")
ax4.set_title("Temperature field as a contour")


# finally a plot showing the directions of heat flow, overlayed
fig5 = plt.figure()
ax5  = fig5.add_subplot(111)
# gradient produces the difference arrays in the 2 axis directions
U,V = np.gradient(Bt)
ax5.invert_yaxis()
ax5.pcolor(np.arange(nodes_x+1)*pdelx,np.arange(nodes_z+1)*pdelx,Bt)
# We want the heat to flow from hot to cold, not vs
ax5.quiver(xx,zz,-V,U)
ax5.set_ylabel("Depth")
ax5.set_xlabel("Width in Meters")
ax5.set_title("Temperature field with heat flow directions and Gradients")