# -*- coding: utf-8 -*-
"""
Created on Tue Mar 24 18:30:47 2020

@author: neil
"""

# 2021 Numerical Modeling in the Geosciences,  Humphrey
# homework problem in non-linear heat flow, (2nd program)
# 
# 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 using a convergence criteria
# (basically the same as program 1 until the iteration loop)
# ***************************************************************************
# 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, the main improvement here is the next line
convergenceT  = .1      # run until the temperatures don't change by some max amount
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 = 50      # Since convergence is not guaranteed, set a maximum number of iterations

n         = 101   	    # 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 
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
        # remember that averaging K values works best if we use the harmonic mean, as below
        T2m = To[i-1]*To[i-1]        # since the K(i) values are based on T**2, produce the squares
        T2  = To[i]  *To[i]
        T2p = To[i+1]*To[i+1]
        h_mean1 = Ko*(2*T2m*T2/(T2m+T2)) 	#note this is where we make K=Ko*T*T
        h_mean2 = Ko*(2*T2*T2p/(T2+T2p))    #this is the harmonic means of the K from i to i+1

        A[i,i]   =  -(h_mean1+h_mean2)/dx2    #Beta
        A[i,i+1] =  h_mean2/dx2               #alpha2
        A[i,i-1] =  h_mean1/dx2               #alpha1
    
    # and solve for the temperature field
    Tn =  np.linalg.solve(A,U)      # Tn now contains the approx solution of temperaturre 
    
    # here we break out of the loop if the temperatures stop changing
    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))    
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