'''
This script computes the moist adiabat for an atmosphere
which is a mixture of a noncondensing and a condensing
component.  It also computes the vertical profile of
the mass specific concentration of the condensible
component, and does plots

When doing plots using the plot(...) function
implemented with Ngl, the exponents on the axes are sometimes
hard to read; in particular, the powers of 10 on the concentration
axis are negative, though you need good eyes to see it.

Note that a version of the MoistAdiabat function, analogous
to the satvps_function used for saturation vapor pressures,
is located in the phys module.  Type help(phys.MoistAdiabat)
for more information. 
'''
#Data on section of text which this script is associated with
Chapter = '2.**'
Figure = '**'


import phys
from ClimateUtilities import *
from math import *


#Choose your mix of gases here
condensible = phys.H2O
noncon = phys.air


#Thermodynamic constants (needed for dry adiabat
#and one-component saturated adiabat computations)
L = condensible.L_vaporization
Ra = noncon.R
Rc = condensible.R
cpa = noncon.cp
Tr = condensible.TriplePointT
psat_ref = condensible.TriplePointP

#Dry and one-component adiabat functions, for comparison
#Note that the use of Numeric.log lets this take
#an array as an argument for p. The rest of the arithmetic,
#including the exponentiation in the dry adiabat, works
#equally well for a Numeric array as for a regular number
#
#**ToDo: Note that the formula for Tdry uses R/cp for dry air for
#comparison. This is what was used in the first edition
#of the Principles of Planetary Climate. However for high
#temperatures, the condensible concentration alters the
#value of R/cp, so it would be more appropriate to use the value
#of R/cp based on the actual mix of condensible and noncondensible.
#This may not matter much, though, since the main point of the comparison
#is to show the convergence to the dry adiabat at low temperatures.
def Tsat(p):
    return(1./(1./Tr - (Rc/L)*Numeric.log(p/psat_ref)))
def Tdry(Ts,p):
    return Ts*(p/p[0])**(Ra/cpa)

#**ToDo: Add controls on resolution, top of atmosphere, etc.
def MoistAdiabat(ps,Ts,condensible=phys.H2O,noncondensible = phys.air):
    #Set up saturation vapor pressure function
    satvp = phys.satvps_function(condensible)
    #Set up thermodynamic constants
    eps = condensible.MolecularWeight/noncon.MolecularWeight
    L = condensible.L_vaporization
    Ra = noncon.R
    Rc = condensible.R
    cpa = noncon.cp
    cpc = condensible.cp
    #
    #Set up derivative function for integrator
    def slope(logpa,logT):
        pa = exp(logpa)
        T = exp(logT)
        qsat = eps*(satvp(T)/pa)
        num = (1. + (L/(Ra*T))*qsat)*Ra
        den = cpa + (cpc + (L/(Rc*T) - 1.)*(L/T))*qsat
        return num/den
    #Initial conditions
    step = -.05  #Step size for integration
    ptop = 1000. #Where to stop integratoin
    #
    #
    logpa = log(ps)
    logT = log(Ts)
    ad = integrator(slope,logpa,logT,step )

    #Initialize lists to save results
    pL = [exp(logpa) + satvp(exp(logT))]
    molarConL = [satvp(exp(logT))/pL[0]]
    TL = [exp(logT)]
    #Integration loop
    p = 1.e30 #Dummy initial value, to get started
    while p > ptop:
            ans = ad.next()
            pa = exp(ans[0])
            T = exp(ans[1])
            p = pa+satvp(T)
            pL.append(p)
            molarConL.append(satvp(T)/p)
            TL.append(T)
    pL = Numeric.array(pL)
    TL = Numeric.array(TL)
    molarConL = Numeric.array(molarConL)
    #Compute the mass specific concentration
    Mc = condensible.MolecularWeight
    Mnc = noncon.MolecularWeight
    Mbar = molarConL*Mc +(1.-molarConL)*Mnc
    qL = (Mc/Mbar)*molarConL     
    return pL,TL,molarConL,qL

ps = 1.e5
Ts = 325.
p,T,molarCon,q = MoistAdiabat(ps,Ts,condensible,noncon)

#Plot temperature
c = Curve()
c.addCurve(p,'p','Pressure (Pa)')
c.addCurve(T,'T','Temperature (K)')
c.addCurve(Tdry(Ts,p),'Tdry','Dry adiabat')
c.addCurve(Tsat(p),'Tad1','One component adiabat')
c.switchXY = c.reverseY = c.YlogAxis = True
c.PlotTitle = "Temperature Profile"
c.Ylabel = "Temperature (K)"
c.Xlabel = "Pressure (Pa)"
plot(c)

#Plot Molar Concentration
c1 = Curve()
c1.switchXY = c1.reverseY = c1.YlogAxis = c1.XlogAxis = True
c1.addCurve(p,'p','Pressure (Pa)')
c1.addCurve(molarCon+1.e-30,'molarCon','Molar concentration of condensible')
c1.PlotTitle = "Molar Concentration Profile"
c1.Ylabel = "Concentration (fraction)"
c1.Xlabel = "Pressure (Pa)"
plot(c1)

#Plot Mass Concentration
c2 = Curve()
c2.switchXY = c2.reverseY = c2.YlogAxis = c2.XlogAxis = True
c2.addCurve(p,'p','Pressure (Pa)')
c2.addCurve(q+1.e-30,'q','Mass specific concentration of condensible')
c2.PlotTitle = "Mass Concentration Profile"
c2.Ylabel = "Concentration (fraction)"
c2.Xlabel = "Pressure (Pa)"
plot(c2)