# -*- coding: utf-8 -*-
"""
Created on Thu Mar 26 20:15:10 2020

@author: neil
"""


# Feb, 2020,  Numerical Modeling in the Geosciences,  Humphrey
# homework problem in non-linear heat flow, 3rd version
# Since this is a CONTINUUM problem (the K does not vary discretely), we try the math version
# of the governing PDE  d2T/dz2 = - 2/T * (dT/dz)^2
# 
# Non-linear 1D heat flow, steady state
# this is a simplified example to show the basic technique of iterating slowly towards
# a solution by increasing the non-linear term slowly (and hoping for the best!) 
# this illustrates a convergence criteria
# ***************************************************************************
# Problem: steady state temperature field in a 100 m thick layer of rock,
# with a surface temperature of 100degrees and a basal temperature of
# 1000degreees, with a K that depends on T^2
# This variation of K is of course unrealistic, but the resulting non-linear
# eqn actually has an analytic solution, which is not typically true!
# ***************************************************************************

import numpy as np
import matplotlib.pyplot as plt
 
# set problem parameters here
omega = .25             # omega is our 'sneak up on the answer' parameter, that gives the proportion 
                        # of the new iteration values that we should average with the old values

MaxIterations = 75      # Since convergence is not guaranteed, set a maximum number of iterations
convergenceT  = .01      # run until the temperatures don't change by some max amount
n         = 11   	    # set the number of nodes, 
thickness = 100 		# set thickness of rock layer in vertical, in meters
upperBC   = 100 	    # set temp BCs, in degrees
lowerBC   = 1000 
Ko        = 1           # important, for a non-linear problem, it is usually necessary to define
                        # starting values carefully... in this case using
                        # T=0 as a starting point does not work (why-- because K is a fn of T, and K(0)=0)
 
# *********everything below here is written with only variables, no numbers ***********
# first set up a figure to plot our results, while we calculate
fig = plt.figure()
fig.canvas.set_window_title ("Temperature In the Crust")
ax1 = fig.add_subplot(1,1,1)
ax1.grid(True)
ax1.set_title("Temperature Profile, K depends on the temperature!")

upperX  = 0             # this uses X for depth, surface is x=0
lowerX  = upperX+thickness  
depths  = np.linspace(0,thickness,n)
Tn      = np.linspace(upperBC,lowerBC,num=n)        #set the starting guess temperature field
ax1.plot(Tn,depths, 'g-+')                          #plot the starting guess temperature
ax1.invert_yaxis()
ax1.set_xlabel("Temperature")  
ax1.set_ylabel("Depth")  
plt.show()

# now set up the finite difference matrix and vectors; A*T=U				
A    = np.zeros((n,n))        #coefficient matrix
To   = np.ones_like(Tn)       #this is the old solution vector
U    = np.zeros(n)            #BCvector, note this is a column vector (important)
delx = thickness/(n-1) 	
dx2  = delx*delx 
# insert the BC conditons, these don't change
U[0]   = upperBC
U[-1]  = lowerBC 

for i in np.arange(1,n-1):       # skip top and bottom, that way we don't have reinput the BCs
    A[i,i]=-2*Ko/dx2              #Beta
    A[i,i+1]=Ko/dx2               #Alpha2
    A[i,i-1]=Ko/dx2               #Alpha1

A[0,0]  =1              # these are the BC nodes in the matrix
A[-1,-1]=1 

nn=0                    #loop counter
while nn < MaxIterations:    # make sure the loop will not go forever
    nn=nn+1 
    To = Tn 
    for i in np.arange(1,n-1):       # skip top and bottom, that way we don't have reinput the BCs
        # we have moved the non-linear terms to the right-hand forcing/BC vector U
        U[i] = -(2/To[i])*( (To[i+1]-To[i-1])/(2*delx) )**2
        #U[i] = -(1/(750*(1+To[i]/750)))*( (To[i+1]-To[i-1])/(2*delx) )**2
     
    # and solve for the temperature field
    Tn =  np.linalg.solve(A,U)      # Tn now contains the approx solution of temperaturre
    if np.max(np.abs(Tn-To)) < convergenceT:
        break                       # exit the loop if the diff Tn-To is less than out convergence Temperature

    Tn =  omega*Tn + (1-omega)*To    # this is a very good little trick for non-linear problems
                                     #  if necessary you can use less of Tn, to keep it stable
    ax1.plot(Tn,depths, 'r-+') 
    plt.pause(.1)
else:  # an 'else' may be used after a 'while' (it is like a 'finally do'), it is run if the loop exits
       # without a 'break' statement (this is somewhat sophisticated python)
    print('Convergence not achieved, Max Iterations reached')    
print('stopped after {} iterations'.format(nn))    
print('temperature at 50m depth {}'.format(Tn[len(Tn)//2]))
ax1.plot(Tn,depths, 'b-+')  #plot the last numerical solution so we can compare with the analytic

# compare our FD solution to analytic soln
T = (9.99e6*depths + 1e6)**(1/3) # analytic soln to the non-linear equation T^2 * d2T/dx2 + 2T * (dT/dx)^2 = 0
                                 # analytic solns to non-linear problems can be difficult or impossible
ax1.plot(T,depths,'k+-')         # plot analytic soln