#---------------------Description--------------------------------------
#
#Computes and plots the surface energy balance diagram,
#which determines the surface temperature. Also computes and plots
#melting rates, as well as evaporation/sublimation rates as a function
#of air temperature. 
#
#
#-----------------------------------------------------------------------

#Data on section of text which this script is associated with
Chapter = '6'
Figure = 'fig:EarthSurfBalance,fig:PotentialMelt,fig:EvsT'
Table = ''
#
import math
import phys
from ClimateUtilities import *

#Tables of ea, for interpolation. These values are
#for 300ppmv CO2, 80% bdd layer relative humidity and 50% aloft,
#on the moist adiabat.  These would need to be changed for
#a planet other than Earth, or for different CO2 concentations.
eaList = [0.2866619245126335, 0.27846028129475953, 0.28133712781516951, 0.2932545463920273, 0.31195466217527701, 0.33493410138862478, 0.36033410293042534, 0.38704395861140806, 0.41466350840361466, 0.44311972709863801, 0.4724710424039088, 0.50324682652330455, 0.53649068478483475, 0.57090053258952678, 0.60658109861462173, 0.6437994820571683, 0.68307368263004697, 0.7250500365717123, 0.77019682655511357, 0.81895136589326822, 0.86878654042337067, 0.91215747338007425, 0.94180397982617692, 0.95785486707999867, 0.96566224250186905]
Tlist_ea = [200.0, 205.0, 210.0, 215.0, 220.0, 225.0, 230.0, 235.0, 240.0, 245.0, 250.0, 255.0, 260.0, 265.0, 270.0, 275.0, 280.0, 285.0, 290.0, 295.0, 300.0, 305.0, 310.0, 315.0, 320.0]
#Interpolated ea function. Uses polint, from ClimateUtilities
ea = interp(Tlist_ea,eaList) #Definnes interpolation function, from ClimateUtilities

#Physical constants
#
#Turbulent drag coefficient (for constant-Cd calculations)
#
Cd0 = 1.5e-03
#
#
#Specification of atmospheric composition
condensible = phys.water
noncon = noncondensible = phys.air
pa = 1.e5 #Background air pressure
#
psat = phys.satvps_function(condensible)
R = noncondensible.R
Rw = condensible.R

#Parameters needed for surface layer calculations
zstar = .00033 #Roughness height (meters)
               #.00033 gives Cd = 1.5e-3 in neutral case
z = 10. # Top of surface layer (meters)
g = 9.8 # Acceleration of gfavity (m/s**2)


#
#--------------Functions---------------------------------
#
#Drag coefficient computation including buoyancy
#Computes Richardson number
def Ri(Tg,Ta,rh,dry =False,z=10.):
    eta = rh*psat(Ta)/(pa + psat(Ta))
    etag = psat(Tg)/(pa + psat(Tg))
    if dry:
        eta = etag = 0.
    eps = condensible.MolecularWeight/noncon.MolecularWeight
    beta = g*(1.-(Ta/Tg)*(1.+(eps-1.)*etag)/(1.+(eps-1.)*eta))
    return -z*beta/(u*u)


#Computes CD as function of Ri in stable case, using
#Monin-Obukhov
def CdMO(Ri,z=10.,z0=zstar):
    Ric = .2
    if Ri > Ric:
        return 0.
    if Ri < 0.:
        Ri = 0. #Use neutral case if unstable
    return ((.4)**2/(math.log(z/z0)**2)) * (1.- Ri/Ric)**2
#Computes the boundary layer density,
#taking into account mixture of condensible
#and noncondensible
def rho(Tbar):
    rhoa = pa/(noncondensible.R * Tbar)
    rhow = relhum*psat(Tbar)/(condensible.R *Tbar)
    return rhoa+rhow  

#Computes specific heat, taking into account
#mix of condensible and noncondensible
def cp(Tbar):   
    rhoa = pa/(noncondensible.R * Tbar)
    rhow = relhum*psat(Tbar)/(condensible.R *Tbar)
    q = rhow/(rhoa+rhow)
    cpw = condensible.cp
    cpa = noncondensible.cp
    return q*cpw + (1.-q)*cpa 

    
def Fsens(Ta,Tg):
    Tbar = .5*(Ta+Tg)
    if neutral:
        Cd = Cd0
    else:
        Cd = CdMO(Ri(Tg,Ta,relhum))
    return rho(Tbar)*cp(Tbar)*u*Cd*(Tg-Ta)

#
#Latent heat flux (evaporation)
def E(Ta,Tg):
    if neutral:
        Cd = Cd0
    else:
        Cd = CdMO(Ri(Tg,Ta,relhum))
    Tbar = .5*(Ta+Tg)
    if Tg < condensible.TriplePointT:
        L = condensible.L_sublimation
    else:
        L = condensible.L_vaporization
    return (L/(Rw*Tbar))*u*Cd*(psat(Tg)- relhum*psat(Ta))

#Atmospheric back radiation into surface
#   ea is a global, and should be set according
#   to atmospheric composition and Ta
#
def IRdown(Ta):
    if Ta > 320.:
        ea1 = ea(320.)
    else:
        ea1 = ea(Ta)
    return ea1 * phys.sigma*Ta**4

#----------Main script-------------------
#Choose whether to use constant Cd or Monin-Obukhov
neutral = True
#
#Boundary layer relative humidity
relhum = .8
#Air temperature
Ta = 265. #Arctic case: 265 K; Tropical case: 300K
#
#Wind speed (m/s)
u = 5.

#Absorbed Solar flux
#S = 320. #(for tropics)
#S = 100. #For winter high latitudes, high albedo
S = 100.



TgList = [ Ta-20. + i for i in range(40)] #Adjust range as needed

netRadList = []
evapList = []
sensList = []
for Tg in TgList:
    netRad = S + IRdown(Ta) - phys.sigma*Tg**4
    netRadList.append(netRad)
    evapList.append(-E(Ta,Tg))
    sensList.append(-Fsens(Ta,Tg))

c = Curve()
c.addCurve(TgList,'Tg')
c.Xlabel = 'Ground Temperature'
c.addCurve(netRadList,'netRad')
c.addCurve(sensList,'sens')
c.addCurve(evapList,'evap')
c['netRad+sens'] = c['netRad'] + c['sens']
c['netRad+sens+evap'] = c['netRad'] + c['sens'] + c['evap']
plot(c)


#------------------------Melting----------------------------------
#Now plot potential melt as a function of temperature, for
#various values of absorbed solar radiation.  Note that this
#calculation is only valid when there is a net flux of
#energy into the surface when the surface is assumed to be
#at the melting temperature. If the net flux is outward under
#these conditions, then the surface temperature falls below
#the melting point and there is no melting. In that case one
#needs to compute the ice surface temperature in order to
#get the sublimation right. The sublimations computed below
#are valid only when the surface is melting.
TaList = [250.+i for i in range(40)]
relhum = .8
cMelt = Curve()
cMelt.addCurve(TaList,'Tsa','Air Temperature')
cMelt.Xlabel = 'Air Temperature'
def doMelt(S):
    meltList = []
    abList = []
    for Ta in TaList:
        evap = E(Ta,273.15)
        sens = Fsens(Ta,273.15)
        dF = S + IRdown(Ta) - sens - evap - phys.sigma*273.15**4
        Melt = (dF/phys.water.L_fusion)*24.*3600.*30. #mm/month
        Melt = max(Melt,0.)
        Sublimation = (evap/phys.water.L_sublimation)*24.*3600.*30. #mm/month
        netAblation = Melt + Sublimation #Valid only when melt is nonzero
        meltList.append(Melt)
        abList.append(netAblation)
    return meltList
neutral = True
cMelt.addCurve(doMelt(100.),'Melt100neutral')
cMelt.addCurve(doMelt(200.),'Melt200neutral')
cMelt.addCurve(doMelt(300.),'Melt300neutral')
neutral = False
cMelt.addCurve(doMelt(100.),'Melt100Stable')
cMelt.addCurve(doMelt(200.),'Melt200Stable')
cMelt.addCurve(doMelt(300.),'Melt300Stable')
cMelt.Xlabel = 'Surface Air Temperature'
cMelt.Ylabel = 'Melt Ablation, mm/mo'
plot(cMelt)

#------------------------Precipitation/Temperature----------------------
def dF(Tg,params):
        Ta = params.Ta
        S = params.S
        evap = E(Ta,Tg)
        sens = Fsens(Ta,Tg)
        dF = S + IRdown(Ta) - sens - evap - phys.sigma*Tg**4
        return(dF)
m = newtSolve(dF)
params = Dummy()
m.setParams(params)
relhum = .8
TaList = [220.+i for i in range(160)]
def doEvap(S):
    evapList = []
    RiList = []
    RiDryList = []
    params.S = S
    TgPrev = TaList[0]
    for Ta in TaList:
        params.Ta = Ta
        Tg = m(TgPrev) #find the ground temperature
        TgPrev = Tg #Save result for use as next initial guess
        evap = E(Ta,Tg)
        evapList.append(evap)
        RiList.append(Ri(Tg,Ta,relhum))
        RiDryList.append(Ri(Tg,Ta,relhum,True))
    return evapList,RiList,RiDryList

cEvap = Curve()
cEvap.addCurve(TaList,'Tsa')
cEvap.Xlabel = 'Air Temperature'

cRichardson = Curve()
cRichardson.addCurve(TaList,'Tsa')
cRichardson.Ylabel = 'Richardson Number'
cRichardson.Xlabel = 'Air Temperature'

neutral = True
S = 200.
print "S = %f,neutral"%S
evap,RiL,RiDryL = doEvap(S)
cRichardson.addCurve(RiL,'Ri')
cRichardson.addCurve(RiDryL,'RiDry')
cEvap.addCurve(evap,'Eneutral')
print "S = %f,MO"%S
neutral = False
evap,RiL,RiDryL = doEvap(S)
cEvap.addCurve(evap,'EMoninObukhov')
cEvap.YlogAxis = True
plot(cEvap)
plot(cRichardson)

#To see the effect of a stable surface layer, try relhum = .5 and
#S = 100.  This gives a strong enough inversion for the
#surface layer to become stable until water vapor buoyancy kicks in
#at higher T.  (Workbook problem).









    





