#This script computes pure radiative equilibrium by timestepping.
#It can use a variety of different radiation codes, and different
#models of the transmission function.

#This version is different from RadConvEq because it
# (a) incorporates stellar heating and (b) Time-steps
# the surface temperature so that one finds the temperature
# which is in equilibrium with a given solar forcing.  With
# solar absorption, one can't do the problem by just holding
# Tg fixed and finding the corresponding atmosphere that's in
# equilibrium with that

#This is an experimental version, originally modified for the AGU2008 talk.
#It has been modified to compute absorption of stellar flux
#impinging at TOA. (mainly for the M-dwarf problem, but useful
#for the discussion of solar heating in the scattering chapter)
#Eventually, this should probably change the zenith angle
#representation, and also incorporate scattering.
#As it stands, the code is mostly useful for absorption of nir.

#Warning:  Because of the way the surface temperature stepping is done
#to achieve TOA balance , when you turn off the stellar heat computation
#using the doStellarAbs flag, it does not affect the radiative flux
#(only the heating rate, which is what's used in the computation).
#If you want to do a computation with an explicit surface budget
#incorporated, you need to modify the way the time-stepping
#of surface temperature is done, and also exclude the stellar
#absorption term from the flux.
#
#Note: The stellar flux past the shortwave cutoff had to be
#added in to the calculation of incoming flux. This needs
#to be put in the main code-tree.

#ToDo:  After the data files are moved to the Workbook Dataset
#directory, update the band-data load section so it can find them

#ToDo:
#     *Make it possible to use ccm radiation as an option, so a separate
#      script isn't needed for ccm computations.
#         ->Need radcomp in ccmradFunctions. Return heating in W/kg
#     *make it possible to use the fancy version of miniclimt

#     *Handle the contribution of water vapor to surface pressure
#      in setting up the pressure grid.  (Tricky, because
#      temperature is changing!)  That's important when
#      surface temperatures are much over 300K.  The changing
#      surface pressure makes the time stepping rather tricky.

#Data on section of text which this script is associated with
Chapter = '5'
Section = '**'
Figure = '**'
#

from ClimateUtilities import *
import phys
import miniClimt as radmodel
import planets

#Path to the workbook datasets
datapath = '/Users/rtp1/Havsornen/PlanetaryClimateBook/WorkbookDatasets/'
#Path to the exponential sum tables
EsumPath = datapath + 'Chapter4Data/ExpSumTables/'

#--------------------Read in band data,set radmodel options-------------------
#
#Data for principal band of CO2. Good up to about 10% CO2 in 1bar air.
#radmodel.bandData = radmodel.loadExpSumTable(EsumPath+'CO2TablePrinBand.260K.100mb.self.data')
#
#Data for the full CO2 spectrum out to 15000 cm**-1. Use
#for calculations with near-IR solar heating
radmodel.bandData = radmodel.loadExpSumTable(EsumPath+'CO2TableExtendedWave.260K.100mb.self.data')


radmodel.BackgroundGas = phys.N2 #The transparent background gas (**Fix self-absorption issue)
radmodel.GHG =phys.CO2 #The greenhouse gas
#
radmodel.pref = 1.e4 #Reference pressure at which band data is given
radmodel.Trans = radmodel.TransEsums
#radmodel.Trans = TransMalkmus   (Don't use this with solar absorption)
#radmodel.Trans = radmodel.TransOobleck (Don't use this with solar absorption

#Set the continuum, if there is one
radmodel.Continuum = radmodel.CO2Continuum
#radmodel.Continuum = radmodel.NoContinuum

#Define function to do time integration for n steps
def steps(Tg,T,nSteps,dtime):
    for i in range(nSteps):
        #Tg = T[-1] (No longer needed. Tg computed by iteration)
        #Do smoothing
        if i%10 == 0:
            for j in range(1,len(T)-1):
                T[j] = .25*T[j-1] + .5*T[j] + .25*T[j+1]
        #
        fluxLW,heatLW = radmodel.radCompLW(p,T,Tg,q)
        fluxStellar,heatStellar = radmodel.radCompStellar(p,T,Tg,q)
        fluxStellar += -OutBandStellarFlux #Stellar flux past shortwave cutoff
        flux = fluxLW+fluxStellar
        if doStellarAbs:    
            heat = heatLW+heatStellar
        else:
            heat = heatLW
        dT = heat*dtime
        #Limit the temperature change per step
        dT = Numeric.where(dT>5.,5.,dT)
        dT = Numeric.where(dT<-5.,-5.,dT)
        #Midpoint method time stepping
        fluxLW,heatLW = radmodel.radCompLW(p,T+.5*dT,Tg+.5*dT[-1],q)
        fluxStellar,heatStellar = radmodel.radCompStellar(p,T+.5*dT,Tg+.5*dT[-1],q)
        fluxStellar += -OutBandStellarFlux #Stellar flux past shortwave cutoff
        flux = fluxLW+fluxStellar
        if doStellarAbs:    
            heat = heatLW+heatStellar
        else:
            heat = heatLW
        dT = heat*dtime
        #Limit the temperature change per step
        dT = Numeric.where(dT>5.,5.,dT)
        dT = Numeric.where(dT<-5.,-5.,dT)
        T += dT
        #
        dTmax = max(abs(dT)) #To keep track of convergence
        
        #Using flux[0] in the following is not a mistake.
        #This is just a trick to relax towards satisfying
        #the TOA balance. The formulation assumes that
        #turbulent fluxes keep the low level air temperature
        #near the ground temperature.

        #Estimate the surface temperature change needed to
        #bring the TOA energy budget into balance
        Trad = (fluxLW[0]/phys.sigma)**.25
        dTg = -flux[0]/(4.*phys.sigma*Trad**3)
        #Relax partway towards that value
        Tg = Tg + .1*dTg
        #
        Tad = Tg*(p/p[-1])**Rcp
        T = Numeric.where(T<Tad,Tad,T)
        #
        if i%1 == 0:
            print i,T[0],T[n/2],T[-1],Tg,-flux[0],dTmax/dtime
    return Tg,Tad,T,flux,fluxStellar,fluxLW,heat,heatStellar,heatLW
    
def transmissionSpectrumStellar(p,T,Tg,q):
    n = len(p)
    nuList = []
    fluxList = []
    for band in radmodel.bandData:
        nuList.append((band.nu1+band.nu2)/2.)
        flux,heat = radmodel.stellarRadCompBand(p,T,Tg,q,band)
        fluxList.append(-flux[-1]/(band.nu2-band.nu1))
    c = Curve()
    c.addCurve(nuList,'nu')
    c.addCurve(fluxList,'flux')
    return c

#Computation of out-of range stellar flux
mPlanck = romberg(radmodel.Planck)
def CompOutBandStellar():
    nu2 = radmodel.bandData[-1].nu2
    return .25*(radmodel.Lstellar/(phys.sigma*radmodel.Tstellar**4))*mPlanck([nu2,10.*nu2],radmodel.Tstellar)

#------------History function, to write output--------------------------        
def history(n,tag):
    suffix = tag+'%d'%n+'.txt'
    c1 = Curve()
    c1.addCurve(p)
    c1.addCurve(flux,'Flux')
    c1.addCurve(fluxStellar,'Stellar')
    c1.addCurve(fluxLW,'LW')
    c1.switchXY=c1.reverseY = True
    c1.PlotTitle = 'Flux'
    c1.dump(outputDir+'Flux'+suffix)


    c2 = Curve()
    c2.addCurve(p)
    c2.addCurve(heat,'Heating')
    c2.addCurve(heatStellar,'Stellar')
    c2.addCurve(heatLW,'LW')
    c2.switchXY=c2.reverseY = True
    c2.PlotTitle = 'Heating Rate'
    c2.dump(outputDir+'Heat'+suffix)

    c3 = Curve()
    c3.addCurve(p,'p')
    c3.addCurve(T,'T')
    c3.addCurve(Tad,'Tadiabat')
    c3.switchXY=c3.reverseY = True
    c3.dump(outputDir+'T'+suffix)

#==============Now do the calculation====================================

#--------------Initializations-------------------------------------------
    
#----------------Set initial time step--------------------------------
dtime = 50. #Timestep in days; 5 days is the usual for Earthlike case
            #For the radiative convective case, you can get away with
            #using 50 for the first few hundred time steps, then
            #reducing to 5 or less later as the solution converges
#
heatfact = 24.*3600./1000. #Rough conversion to K/day
                           #If we are only interested in the steady
                           #state, and not in accurately reproducing
                           #the detailed time dependence, it doesn't
                           #matter if we use the exact specific heat
                           #here to convert the heating rate in W/kg
                           #into Kelvins/day
dtime = dtime*heatfact
#----------------------------------------------------------------------


#-------Set up the pressure array------------------------------
#   We typically use 60 points with extra resolution near the
#   ground when doing pure radiative equilibrium, where there
#   is a radiative boundary layer near the ground that needs to
#   be resolved.  For radiative-convective equilibrium, the adiabat
#   takes over near the ground, and setpLin(...), which gives uniform
#   resolution, does pretty well with as little as 20 points. If the
#   tropopause gets too high up, though, you may need extra resolution
#   in the stratosphere, and in that case using a logarithmic grid
#   can be advantageous.
#
#   Caution: If you go to surface temperatures much above 300K,
#   then you need to keep track of the total surface pressure
#   vs. the surface pressure of dry air. That's not done here.
#--------------------------------------------------------------
n = 20 
#ps = 600.
#ptop = 1. #Use 100 except for Venus, where we use 1.e4
ps,ptop = 2.e5,100.
#p = radmodel.setpExtraGroundRes(ps,ptop,n)
p = radmodel.setpLin(ps,ptop,n)
#p = radmodel.setpLog(ps,ptop,n) #Puts extra resolution in stratosphere.
#-------------------------------------------------------------------------


#------------Set the greenhouse gas concentration------------------------
#q = 10.e-4 #Makes optical depth 10 for Oobleck
#q = (44./29.)*300.e-6 #Earthlike CO2
q = 1. #Use for Mars and Venus. When q=1, the background gas isn't used
       #In this case, remember to use the self-broadened table,
       #not the foreign-broadened table.
Rcp = radmodel.GHG.Rcp #Change if the GHG isn't dominant
#------------------------------------------------------------------------

#Where to put the output
outputDir = 'TestResults/'
caseTag = 'Test'

#--------------Other parameters-------------------------------------------
doStellarAbs = True
radmodel.g = 10. #Set the gravity
#radmodel.g = planets.Mars.g #Mars case
radmodel.Tstellar = 6000.
#radmodel.Lstellar = 1000. #Tuned for Mars summer afternoon sounding
radmodel.Lstellar = 1400.
OutBandStellarFlux = CompOutBandStellar()
print OutBandStellarFlux
#
#To speed up the calculation, specify a cutoff wavenumber
#beyond which the longwave (thermal) flux calculation won't
#be done.  If the atmosphere gets very hot, though, this
#cutoff may need to be increased. For cool atmospheres, you
#could make it smaller.  To figure out what value to use,
#examine the Planck function for the maximum temperature in
#the atmosphere.
radmodel.LWCutoff = 3000. #In 1/cm
#
dtime = .5**3*dtime # .5**3 is standard
#
#Initialize the temperature
#c = readTable(outputDir+'TTestTrop80.txt')
#p = c['p']
#T = c['T']
#Tg = T[-1]
Tg = 250.
T = Tg*(p/p[-1])**Rcp
#

#---------------------------------------------------------
#--------------Initializations Done-----------------------
#--------------Now do the time stepping-------------------
#---------------------------------------------------------
for i in range(0,20):
    nout = 10*i
    if i%5 == 0:
         dtime = .5*dtime
    Tg,Tad,T,flux,fluxStellar,fluxLW,heat,heatStellar,heatLW = steps(Tg,T,10,dtime)
    print Tg,T[-1],T[-5]
    history(nout,caseTag)