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

@author: neil
"""

# 2021,  Numerical Modeling in the Geosciences,  Humphrey
# problem in non-linear heat flow
# 
# Non-linear 1D heat flow, steady state, (program 1)
# programs 1 and 2 try to solve the eqn as d( K*T^2 dT/dz)/dz = 0
# 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!) 
# ***************************************************************************
# 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, especially the step size for iteration to an answer
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

iterations = 25         # for this simple problem we set the number of interations to get our answer
                        # for more complex problems, you wan to check for
                        # convergence of your solution, (see the next example program)
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
                                                    # for a guess we use a linear temperature spread
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 theoretically a column vector
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 to see how many times we go thru the loop
while( nn < iterations):
    nn=nn+1 
    To = Tn 
    # reform the A matrix with new values of K(T*T)
    for i in np.arange(1,n-1):       # skip top and bottom rows, 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
        T2i = To[i]  *To[i]          # T2i is the square of To[i], while T2m is square of To above To[i]
        T2p = To[i+1]*To[i+1]        # and T2p is the square of the To value below (to the right) of To[i]
        h_mean1 = Ko*(2*T2m*T2i/(T2m+T2i)) 	  #note this is where we make K=Ko*T*T
        h_mean2 = Ko*(2*T2i*T2p/(T2i+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 an approx solution of temperaturre 
    
    # the core of the iteration, is to add a little of the new estimate of T to the old and hope it 
    # improves the answer. NOTE using to much of the new values usually causes oscillations or divergence
    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)
    
ax1.plot(Tn,depths, 'b-+')  #plot the last numerical solution so we can compare with the analytic

# compare our FD solution to analytic soln (this problem was chosen since it has an 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