#Script for illustrating simple models of
#ice-albedo feedback
#
#
#Data on section of text which this script is associated with
Chapter = '3.**'
Figure = '**'
#

from ClimateUtilities import *
import phys

#Define albedo and OLR functions here

#Radiating pressure pRad, in mb, is passed as the
#optional second parameter. It can be used to
#pass parameters of other OLR expressions, e.g.
#CO2 concentrations
def OLR(T,param=None):
    pRad = param
    return phys.sigma * (T**4.)*(pRad/1000.)**(4.*2./7.)
    #Comment out the preceding line and replace with the
    #following one, to use a linear OLR fit
    #return 113. + 2.177*(T-220.)

#Parameters of albedo function are globals
def albedo(T):
    if T < T1:
        return alpha_ice
    elif (T >= T1)&(T<=T2):
        r = (T-T2)**2/(T2-T1)**2
        return alpha_0 + (alpha_ice - alpha_0)*r
    else:
        return alpha_0


#Main body of script starts here
L = 1.*1370. #1.2 times modern give marginal case w/o greenhouse
alpha_ice = .6
alpha_0 = .2
T1 = 260.
T2 = 290.
pRad = 1000.

Tlist = [200.+ 2.*i for i in range(70)]
SabsList = []
OLRlist = []
NetFluxlist = []
Glist = []

for T in Tlist:
    SabsList.append((L/4.)*(1.-albedo(T)))
    OLRlist.append(OLR(T,pRad))
    NetFluxlist.append((L/4.)*(1.-albedo(T)) - OLR(T,pRad))
    Glist.append(4.*OLR(T,pRad)/(1.-albedo(T)))

#Plot the results
c1 = Curve()
c1.addCurve(Tlist)
c1.addCurve(SabsList,'Sabs','Absorbed Solar')
c1.addCurve(OLRlist,'OLR','Outgoing Longwave')
c1.Xlabel = 'Temperature (K)'
c1.Ylabel = 'Flux (W/m**2)'
c1.PlotTitle = 'Energy Balance Diagram'
w1 = plot(c1)

c2 = Curve()
c2.addCurve(Tlist)
c2.addCurve(NetFluxlist,'Net','Net Flux')
c2.addCurve([0. for i in range(len(Glist))],'Zero','Equilibrium')
c2.Xlabel = 'Temperature (K)'
c2.Ylabel = 'Net Flux (W/m**2)'
c2.PlotTitle = 'Stability diagram'
w2 = plot(c2)

c3 = Curve()
c3.addCurve(Tlist)
c3.addCurve(Glist,'G','Balance Function')
c3.addCurve([L for i in range(len(Glist))],'S','Incident Solar')
c3.switchXY = True #Switches axis so it looks like a hysteresis plot
c3.PlotTitle = 'Bifurcation diagram, pRad = %f mb'%pRad
c3.Xlabel = 'Surface Temperature (K)'
c3.Ylabel = 'Solar Constant (W/m**2)'
w3 = plot(c3)

#Make a plot of the energy balance diagram with multiple
#different solar constants
Llist = [.9*L,L,1.1*L,1.7*L]
c4 = Curve()
c4.addCurve(Tlist,'T','Surface Temperature')
for L1 in Llist:
    Solar = [(L1/4.)*(1.-albedo(T)) for T in Tlist]
    c4.addCurve(Solar,'%f'%(L1),'L = %f'%(L1))
c4.Xlabel = 'Temperature (K)'
c4.Ylabel = 'Flux (W/m**2)'
c4.PlotTitle = 'Energy Balance Diagram'
c4.addCurve(OLRlist,'OLR','Outgoing Longwave')

w4 = plot(c4)

#Now make a bifurcation diagram with pRad as the control variable
#Note that if you have changed the OLR function in such a way that
#it is independent of its parameter, this calculation will fail,
#so you should comment it out.  The technique will work, though,
#for pretty much any parameter characterizing the greenhouse
#effect, e.g. CO2 concentration.

#Note that since the control parameter is a single-valued function
#of the temperature T, we can actually make the graph much more
#swiftly by solving for, say, L in terms of T analytically, and
#then plotting L(T) instead of T(L). Since L(T) = OLR(T)/(1-albedo(T)),
#this curve can always be generated from arbitrary OLR and
#albedo functions. (it's the same curve as in the equilibrium
#ice-albedo script, and indeed perhaps part of the hysteresis stuff
#should be moved to that script).
#
#For hysteresis curves over a greenhouse parameter, one can
#proceed similarly.  For the simple OLR above, one can solve
#for pRad analytically, but in the general case one
#first computes, for a specified T, F = (1. - albedo(T))*L/4
#then solves OLR(T,GreenhouseParam) = F, to find GreenhouseParam.
#Since OLR is generally monotonically decreasing in GreenhouseParam,
#this procedure essentially always yields a single-valued function.

def radPress(flux,T):
    def g(p,const):
        return OLR(const.T,p) - const.flux
    root = newtSolve(g)
    const= Dummy()
    const.T = T
    const.flux = flux
    root.setParams(const)
    return root(500.)


L = 1370.
TList = [220.+i for i in range(131)]
gList = [radPress((1.-albedo(T))*L/4.,T) for T in TList]
cG = Curve()
cG.addCurve(gList,'pRad')
cG.addCurve(TList,'T')
#Just for variety, here I've shown that instead of
# switching axes, you can just put the data into
# the Curve in a different order
cG.PlotTitle = 'Bifurcation diagram, L = %f W/m**2'%L
cG.Ylabel = 'Surface Temperature (K)'
cG.Xlabel = 'Radiating pressure (mb)'
cG.reverseX = True #Reverse pressure axis so warmer is to the right
plot(cG)