# -*- coding: utf-8 -*-
"""
Created on Mon Mar 30 12:50:54 2020

@author: neil
"""

# Geol 5470, Humphrey 2020
#  this is  the 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
# However, we place a gradient condition on the bottom

import matplotlib.pyplot as plt
import numpy as np


#***** our problem assumes the block is 100m wide, 100m deep, with a sfc
# temperature of 2, right side of 4, and 
# left side of 6 degrees. With a 0 gradient at the bottom
# Note, since K is constant, we divide it out of the problem, 
# also, less obvious, if we make delx and delz the same, we can divide delx delz out

# discretize the domain
lengthx = 100      # width
lengthz = 100      # depth
nodes_x =  6       # 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    = 2
BCbottom = 0    # this is a gradient conditon
BCright  = 4
BCleft   = 6

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

# 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+1,nodes-nodes_z):   #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
    if ((i+1)%nodes_z) ==0 :    # this chooses only the index of lower boundary nodes
        A[i,i-1] = -1           # this is the gradient condition, make A[i,i-1]=-1 and A[i,i]=1
        
    # 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, while top nodes are 1 more
        A[i,i]        =  -4
        A[i,i-1]      =1     # 
        A[i,i+1]      =1     # 
        A[i,i-nodes_z]=1     # the i node to the left of a node is i-nodes_z
        A[i,i+nodes_z]=1
    

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

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


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