#------------------------------------------------------------
#Script to solve the two-stream equations with scattering,
#by numerical integration as an ODE problem. This version
#will break down when the atmosphere becomes very optically thick.
#
#As currently written, it illustrates the numerical solution in
#its simplest form, and doesn't provide for making the scattering
#coefficients (specifically the single-scatter albedo) a function
#of optical depth. The atmosphere is also assumed isothermal.
#
#The script does the calculation two different ways: first starting
#the integration from the top and working downward, then starting
#from the bottom and working upwards.  Both are plotted. The
#two calculations should give the same result, unless roundoff
#problems start to become significant (which eventually becomes
#the case when the atmosphere becomes very optically thick).
#
#Note that this script plots the diffuse flux only.  To get
#the total flux, the direct beam term should be added in
#------------------------------------------------------------
#Data on section of text which this script is associated with
Chapter = '5.**'
Figure = 'None'
#

#
#ToDo:
#


from ClimateUtilities import *
import phys

import numpy,math


#Planck functions in wavenumber space (cm**-1)
def Planck(wavenum,T):
    return 100.*math.pi*phys.c*phys.B(100.*wavenum*phys.c,T)

#Function returning the derivatives of the
#upward and downward flux, for use with the
#integrator. I = [Iplus,Iminus].
#
def dIdtau(tau,I,source):
    #Scattering coefficients
    #   These would eventually be functions of depth,
    #   in cases where the single-scatter albedo omg0
    #   depends on depth.  In a general calculation,
    #   the temperature Ta also would depend on depth
    gamma1 = gamma + gammaPrime*(1.-omg0) #Assumes ghat=0
    gamma2 = gamma - gammaPrime*(1.-omg0)
    gammaB = gamma1-gamma2
    gammaPlus = gammaMinus = .5*omg0 #Assumes ghat = 0
    #
    #Direct beam source term
    Fsun = Lsun*Planck(wave,Tsun)/(phys.sigma*Tsun**4)
    DirectBeam = Fsun*math.exp(-(tauInf-tau)/cosz)
    #
    Iplus = I[0]
    Iminus = I[1]
    dIpdtau = -gamma1*Iplus + gamma2*Iminus + \
              source*(gammaB*Planck(wave,Ta) + gammaPlus*DirectBeam)
    dImdtau = gamma1*Iminus - gamma2*Iplus - \
              source*(gammaB*Planck(wave,Ta)+gammaMinus*DirectBeam)
    return numpy.array([dIpdtau,dImdtau])

#Set parameters-----------------------------------------------
#Coefficients for the two-stream approx
gamma = 1.
gammaPrime = 1.
#
omg0 = .9 #Single scatter albedo
cosz  = .5 #Cosine of zenith angle
Tsun = 6000 #Blackbody temperature of star
            #(Used for spectrum of direct beam illumination)
Lsun = 1370. #Stellar constant
wave = 600.
Tg = 320.
Ta = 300.
#
dtau = .1
tauInf = 20.
#Done setting parameters -------------------------------------

#Construct a homogeneous solution satisfying
#Iminus = 0 at top of atmosphere and Iplus = 1
#at bottom.  This will be added to a particular
#solution satisfying the upper boundary condition,
#in order to satisfy the lower one. We
#integrate from the top downward

#Initial conditions
Istart = numpy.array([1.,0.])
m = integrator(dIdtau,tauInf,Istart,-dtau)
m.setParams(0.) #Homogeneous. No source
tau = tauInf

tauList = [tauInf]
IList = [Istart]
while tau > 0.:
    if tau - dtau < 0.:
        break
    tau,I = m.next()
    tauList.append(tau)
    IList.append(I.copy())

Ip0 = IList[-1][0]
Ihom = [I/Ip0 for I in IList] #Normalize

#Now construct the particular solution
#Initial conditions
Istart = numpy.array([1.,0.])
m = integrator(dIdtau,tauInf,Istart,-dtau)
m.setParams(1.) #Inhomogeneous.
tau = tauInf

tauList = [tauInf]
IList = [Istart]
while tau > 0.:
    if tau - dtau < 0.:
        break
    tau,I = m.next()
    tauList.append(tau)
    IList.append(I.copy())

#Find correct superposition of particular and
#homogeneous solution to satisfy lower bc
Ip0 = IList[-1][0]
coeff = (Planck(wave,Tg)-Ip0)
for i in range(len(IList)):
    IList[i] += coeff*Ihom[i]


#Plot
c = Curve()
c.addCurve(tauList,'tau')
c.addCurve([I[0] for I in IList],'Iplus')
c.addCurve([I[1] for I in IList],'Iminus')
c.Xlabel = 'tau'
c.switchXY = True
plot(c)

#Alternate approach: Start from bottom and work up
#Initial conditions
Istart = numpy.array([0.,1.])
m = integrator(dIdtau,0.,Istart,dtau)
m.setParams(0.) #Homogeneous. No source
tau = 0.

tauList = [tauInf]
IList = [Istart]
while tau < tauInf:
    if tau + dtau > tauInf:
        break
    tau,I = m.next()
    tauList.append(tau)
    IList.append(I.copy())

ImTop = IList[-1][1]
Ihom = [I/ImTop for I in IList] #Normalize

#Now construct the particular solution
#Initial conditions
Istart = numpy.array([Planck(wave,Tg),0.])
m = integrator(dIdtau,0.,Istart,dtau)
m.setParams(1.) #Inhomogeneous.
tau = 0.

tauList = [0.]
IList = [Istart]
while tau < tauInf:
    if tau + dtau > tauInf:
        break
    tau,I = m.next()
    tauList.append(tau)
    IList.append(I.copy())

#Find correct superposition of particular and
#homogeneous solution to satisfy lower bc
ImTop = IList[-1][1]
for i in range(len(IList)):
    IList[i] -= ImTop*Ihom[i]

#Plot
c = Curve()
c.addCurve(tauList,'tau')
c.addCurve([I[0] for I in IList],'Iplus')
c.addCurve([I[1] for I in IList],'Iminus')
c.Xlabel = 'tau'
c.switchXY = True
plot(c)