#---------------------Description--------------------------------------
#
#   This script solves the Walker, Hayes and Kasting
#   CO2 regulation model.  As an alternative to the
#   original formulation, you can substitute in different
#   forms of the function relating temperature to CO2.
#   The present version uses a fixed albedo. Try experimenting
#   with an alternate formulation incorporating ice-albedo
#   feedback!
#
#-----------------------------------------------------------------------
#
#Change log
#  7/11/2012: Improved comments and introduced some additional
#             names for constants to make the script easier to
#             modify or extend.


#Data on section of text which this script is associated with
Chapter = '8'
Section = ''
Figure = 'fig:WHAKabiotic'
Table = ''
#

import math
from ClimateUtilities import *
import phys

#OLR coefficients, and polynomial OLR fit. The coefficients are
#hard-coded in here for convenience so that OLRPoly doesn't need
#to be imported.  Also, to keep the temperature calculation simple,
#the calculation of T here is based on a linearized dependence
#of OLR on temperature.  The coefficients are
#based on 270K to 300K temperature range, moist adiabat,
#50% relative humidity, Earth gravity. For other gravity and planets,
#you can use ccmradFunctions.py or the homebrew code (see Chapter 4
#scripts) to replace these coefficient tables.
#
#The fit is done to psi = log(CO2/CO2Fit0), with CO2 and CO2Fit0 in ppmv. The
#coefficient functions a(p) and b(p) are written as functions of
#p = CO2/CO2Fit0, since the whole WHAK model is written in terms
#of values relative to (or normalized by) a baseline value.
#
#The fit is done with a base temperature of Tfit0, specified below,
#so that when T = Tfit0, the OLR is just a(p).  Tfit0 does
#not actually enter into the silicate weathering equilibrium calculation,
#since the temperature appears only in the form of a difference with
#the base-case temperature. Tfit0 is only used for output purposes,
#to translate dT into the actual temperature, so one can see how close
#to freezing one gets.
CO2Fit0 = 300. #CO2 normalization value, ppmv
Tfit0 = 270. 
aCoeff = [258.283,253.341,244.395,232.74,220.23,204.7,179.9]
bCoeff = [2.443,2.476,2.498,2.466,2.37,2.216,1.95]
psi = [math.log(.1/CO2Fit0),math.log(1./CO2Fit0),math.log(10./CO2Fit0),
       math.log(100./CO2Fit0),math.log(1000./CO2Fit0),math.log(10000./CO2Fit0),
       math.log(100000./CO2Fit0)]
def a(p):
    return polint(psi,aCoeff,math.log(p))
def b(p):
    return polint(psi,bCoeff,math.log(p))

#Weathering function based on Berner formulation
#   The returned value, V is the weathering rate normalized
#   by the baseline weathering rate corresponding to p = 1
#   and dT = 0. Typically, the baseline climate is taken as
#   Earth's pre-industrial climate, but other choices are possible.
#
#Arguments:
#   dT = temperature difference between the baseline climate
#        and the climate corresponding to some altered conditions,
#        (e.g. by changing the solar constant)
#        
#   p = pCO2/pCO2(0) , with pCO2(0) being the CO2 partial pressure in
#       the baseline climate. As currently written, pCO2(0) is
#       taken to be equal to the value CO2Fit0 used in defining
#       the OLR coefficients. If for some reason you want to change
#       the CO2 concentration corresponding to p=1, you can just
#       change the value of CO2Fit0 above, and there is no need
#       to change the value arrays aCoeff and bCoeff.
#
#Internal constants:
#   aWHAK is the runoff scaling exponent.
#   a_ro is proportionality
#      giving increase of runoff with temperature.
#   bWHAK is direct CO2 exponent.
#       It is expected to be nonzero only in the absence of vegetation.
#       With vegetation, the corresponding scaling factor replaced by
#       a fixed CO2 indep. vegetation factor instead.
def V(dT,p):
    precip = 1.+a_ro*dT # Linearized form
    #precip = math.exp(-(5400./300.**2)*dT)#Alternate Clausius-Clapeyron form
    
    #Put precipitation limiter here, if desired, to incorporate
    #the limitations on precip due to the surface energy budget
    #(i.e., it's hard for precip to exceed the equivalent of the
    #surface absorbed solar radiation)
    
    precip = max(precip,.001) #To avoid negative precip
    return (precip)**aWHAK * p**bWHAK * math.exp(dT/T_U)


#The "climate function" which gives the temperature anomaly vs pCO2/pCO2(0).
#This is the temperature difference with the unperturbed case that is
#in equilibrium with absorbed solar S=S0 and 1XCO2(0) (p=1).
#
#You could include ice-albedo feedback here.
def dT(p):
    return (S0 + dS  - a(p)  )/b(p) - (S0 - a(1.))/b(1.)


#Objective function for computing what pressure causes
#the weathering rate to equal the outgassing. Note that
#the outgassing rate, Vtarget, is normalized to the
#weathering rate in the baseline climate, because that is
#the way the weathering function V(dT,p) is defined. 
def f(p,Vtarget):
    return V(dT(p),p) - Vtarget

#Instantiate a Newton's method root finder
m = newtSolve(f)

#Set constants
T_U = 10. #Berner uses 11.11, but we rounded to 10
aWHAK = .65 #From Berner, .65
bWHAK = .5 #From Berner, .5 . Applies only without vegetation
           #Try lower values to see how weakening this dependence
           # makes a tighter thermostat.
a_ro = .03  #Runoff vs. temperature coefficient

alpha = .2 #Albedo (fixed in this calculation)
S0 = (1.-alpha)*1370./4.  #Baseline absorbed solar radiation 
T0 = Tfit0 + (S0 - a(1.))/b(1.) #Baseline temperature
#
pCO2 = 1. #Initial guess

#Lists in which we will store results
SList = []
TList = []
TuList = []
CO2List = []

#Find the equilibrium climate for various outgassing rates.
m.params = 1. #Set outgassing rate (relative to present,
              #or other baseline used)

#This loop does the calculation. For more reliable convergence,
#it is best to do two separate calculations, starting from
#the baseline state (S0,T0,pCO2(0)) and going to lower or
#higher S. In the text, the results from the two loops
#were spliced together into a single graph. dSfact gives the
#proportional change in the absorbed solar radiation relative
#to the baseline. 

for dSfact in [.01*i for i in range(31)]: #Past dimmer sun
#for dSfact in [-.01*i for i in range(50)]: #Future brighter sun
    dS = -dSfact*S0
    pCO2 = m(pCO2)
    dT1 = dT(pCO2)
    Tu = T0 + dT(1.) #unthermostatted value
    T = T0 + dT1
    SList.append(S0+dS)
    TList.append(T)
    TuList.append(Tu)
    CO2List.append(pCO2*CO2Fit0*1e-3) #Convert to mb
    
    print S0 + dS,pCO2*CO2Fit0*1e-3,Tu,T,m.params

#Put the results in Curve objects for output and plotting
cT = Curve()
cT.addCurve(SList,'S')
cT.addCurve(TList,'T')
cT.addCurve(TuList,'T(unthermostatted)')
cT.PlotTitle = 'Temperature vs Absorbed Solar'
cT.Xlabel = 'Absorbed Solar (W/m**2)'
cT.Ylabel = 'Temperature (K)'

cCO2 = Curve()
cCO2.addCurve(SList,'S')
cCO2.addCurve(CO2List,'CO2(mb)')
cCO2.YlogAxis = True
cCO2.PlotTitle = 'CO2 vs Absorbed Solar'
cCO2.Xlabel = 'Absorbed Solar (W/m**2)'
cCO2.Ylabel = 'pCO2 (mb)'