#Solves the diffusion equation with a flux boundary
#condition at the surface

Chapter = "7"
Figure = None

from ClimateUtilities import *


N = 200 # Number of points to use in the vertical
dz = .05 # Grid spacing in meters
dt = 10. # Time step in seconds
nsmooth = 100 # How often to do smoothing step
kCond = 2.24# Thermal conductivity, = rho cp D
D = 1.16e-6 # Thermal diffusivity

#OLR coefficients for quadratic fit
#  Based on 150K stratosphere, rh=.5, c02 = 300ppmv
Tbase,a,b,c = 180.,48.461,1.5866,.0029663

Tstart = 300.
v = numpy.array((1.,-2.,1.))/(dz*dz)
#Function to return second derivative of temperature array
#Note that this does not set the upper boundary condition,
#but it does set the bottom boundary condition
def Tzz(T):
    temp = numpy.convolve(T,v,1)
    temp[-1] = temp[-2] #No flux bottom B.C.
    return temp

#Smoothing function, to deal with the leapfrog oscillation
def smooth(T2,T3):
    Tbar = .5*(T2+T3)
    return Tbar,Tbar

#OLR function
def OLR(Ts):
    return a + b*(Ts-Tbase) + c*(Ts-Tbase)*(Ts-Tbase)


#Main script starts here

# Set up arrays
T1 = numpy.zeros(N,numpy.float)
T2 = numpy.zeros(N,numpy.float)
T3 = numpy.zeros(N,numpy.float)

#Initialize
T1[:] = Tstart
T2[:] = Tstart

#Time step
t = 0.
tlist = []
Tslist = []
cDepth = Curve() # Curve to hold depth data
cDepth.addCurve( [dz*i for i in range(N)],'depth') # Depth data
for i in range(600000):
    T3 = T1 + 2.*dt*D * Tzz(T2)
    #Do flux boundary condition
    T3[0] = T3[1] - dz*OLR(T3[0])/kCond
    # Save results for plotting
    t += dt
    if i%1000 == 0:
        print t/(24.*3600.),T3[0]
        tlist.append(t/(24.*3600.))
        Tslist .append(T3[0])
    if i%(5*24*3600/dt) == 0:
        cDepth.addCurve(T2,'T(%f,z)'%(t/(24.*3600.)))
    #Cycle the data to get ready for the next step
    T2,T1 = T3,T2
    #Smooth if necessary
    if i%nsmooth == 0:
        T2,T1 = smooth(T2,T1)
cTime = Curve()
cTime.addCurve(tlist,'t')
cTime.addCurve(Tslist,'Ts')