S = 1370 # Solar constant
albedo = 0.3

# The nlayer model, calculating surface temperature
def nlayer(n,S,albedo,expnt=4):
    '''input:
    n: number of layers
    S: solar constant
    albedo: fraction of reflection
    expnt: exponent in Stefan-Bolztmann's law
    
    Output:
    T: surface temperature'''
    
    sigma = 5.67e-8 #Stefan-Boltzman constant
    # solution of surface temperature
    T = ((n+1)*S*(1-albedo)/(4*sigma))**(1./expnt)
    return T

"""
Explanation: Part II: the $n$-layer model
Here we will build the n-layer model mentioned in the lecture. We know that introducing layers in the model is to consider the role of atmosphere, especially greenhouse gases. The number of layers is equivalent to the amount of greenhouse gas in the atmosphere.
End of explanation
"""


# number of layers
# it is equivalent to the amount of greenhouse gas in the atmosphere
layers = np.arange(100)
T_layer = nlayer(layers,S,albedo,4)

# plot number of layers vs. surface temperature
plt.plot(layers,T_layer)
plt.xlabel('number of layers')
plt.ylabel('Surface temperature (K)')

"""
Explanation: Let's see how the surface temperature changes as the number of layers increases.
End of explanation
"""


print('surface temperature=',nlayer(2,S,albedo,?),'K')

"""
Explanation: Question 1: How the surface temperature changes as the number of layers increases? Is this reasonable? Why?
Answer 1:
We know that the black boday radiation is proportional to the fourth power of temperature ($E=\sigma T^4$). What if we live in a world that the exponent number is not 4?
Question 2: What if the exponent number is 2 (i.e. $E=\sigma T^2$)? What is the surface temperature based on this n-layer model (use n=2)?
End of explanation
"""


# what if the power number is not 4?
n_power = np.arange(1,7,0.2)
T_power = nlayer(2,S,albedo,n_power) # choose a two-layer model

plt.semilogy(n_power,T_power)
plt.axhline(273,color='gray',linestyle='--')
plt.axhline(373,color='gray',linestyle='--')
plt.xlabel('Stefan-Boltzmann Exponent')
plt.ylabel('Surface temperature (K)')

"""
Explanation: We can plot the how surface temperature responses as the power changes
End of explanation
"""


# import climlab
from __future__ import division, print_function
import xarray as xr
import climlab

# read NCEP/NCAR data
# read temperature, level, and latitude
ncep_url = "http://www.esrl.noaa.gov/psd/thredds/dodsC/Datasets/ncep.reanalysis.derived/"
ncep_air = xr.open_dataset(ncep_url + "pressure/air.mon.1981-2010.ltm.nc")
level = ncep_air.level
lat = ncep_air.lat

"""
Explanation: Question 3: For what range of exponents do we get a "livable" temperature? (hint: plot 2 horizontal lines at the freezing point and boiling point of water, assuming standard pressure).  
Part III: A Radiative-Convective Model
0D models, such as the ones we've been playing so far, are instructive but quite limited.
In this part, we will build a 1D climate model, which can simulate the vertical profile of temperature in the atmosphere. At first, we will only consider radiative processes, and then progressively add dry and moist convection processes. This 1D climate model is part of the climlab package (http://climlab.readthedocs.io) developed by Dr. Brian Rose at the University at Albany.
1. What are we aiming for?
Let's first look at what the vertical distribution of temperature of the Earth looks like in the observations. Here we use the NCEP/NCAR reanalysis data as the observations.
End of explanation
"""


Tzon = np.mean(ncep_air.variables['air'],axis=(0,3))
Tglobal = np.sum( Tzon * np.cos(np.deg2rad(lat)), axis=1 ) / np.sum( np.cos(np.deg2rad(lat) ) )

q_obs = np.loadtxt('q_profile.txt')

"""
Explanation: Take global averages and time averages.
End of explanation
"""


def zstar(lev):
    return -np.log(lev / climlab.constants.ps)
def get_tendencies(model):
    '''Pack all the subprocess tendencies into xarray.Datasets
    and convert to units of K / day'''
    tendencies_atm = xr.Dataset()
    tendencies_sfc = xr.Dataset()
    for name, proc, top_proc in climlab.utils.walk.walk_processes(model, topdown=False):
        if proc.time_type is not 'diagnostic':
            tendencies_atm[name] = proc.tendencies['Tatm'].to_xarray()
            tendencies_sfc[name] = proc.tendencies['Ts'].to_xarray()
    for tend in [tendencies_atm, tendencies_sfc]:
        #  convert to K / day
        tend *= climlab.constants.seconds_per_day
    return tendencies_atm, tendencies_sfc

yticks = np.array([1000., 750., 500., 250., 100., 50., 20., 10.])

plt.style.use('ggplot')
plt.figure( figsize=(8,8) )
plt.plot( Tglobal + 273.15, zstar(level),label='observations',color='C2',linestyle=':',linewidth=4)
plt.yticks(zstar(yticks),yticks)
plt.xlabel('Temperature (K)')
plt.ylabel('Pressure (hPa)')
plt.title('Global, annual mean sounding from NCEP Reanalysis')

"""
Explanation: Define functions:
(1) Get the log scale axis of pressure
(2) Get the temperature tendency variables from the model
End of explanation
"""


#  Choose the surface albedo
alb = 0.2
#  State variables (Air and surface temperature)
state = climlab.column_state(num_lev=50)
#  Parent model process
rad_model = climlab.TimeDependentProcess(state=state)
#  Set the radiation model
rad = climlab.radiation.RRTMG(state=state, specific_humidity=q_obs, albedo=alb)
#  Couple the radiation model
rad_model.add_subprocess('Radiation', rad)

"""
Explanation: Question 3: Please describe the vertical temperature distribution.
Answer 3:
2. Radiative equilibrium
Now we only use the radiation model (RRTMG) to simulate the temperature profile and compare it with the observation. This model can consider the radiation effect of different gases, at different wavelengths (like CO$_2$ mainly absorbs long-wave radiation, but O$_3$ mainly absorbs short-wave radiation).
End of explanation
"""


#  Volumetric mixing ratios
rad.absorber_vmr

"""
Explanation: Let's look at what radiatively active gases are in this model.
End of explanation
"""


#  Start from isothermal state
rad_model.state.Tatm[:] = rad_model.state.Ts

# Integrate the model for 4 years
rad_model.integrate_years(4)

# Check whether the model reaches the equilibrium
# usually it does when the number smaller than 0.1
rad_model.ASR - rad_model.OLR

"""
Explanation: Question 4: What gases are considered in the radiation model?
Answer 4:
Let's run this model until it reaches equilibrium (i.e. energy fluxes balance each other out, so the temperature is constant). We need to specify an initial state, which will be an isothermal one (the atmosphere starts out with the same temperature everywhere - obviously not realistic, but instructive to see how radiation changes that situation).
End of explanation
"""


tendencies_atm, tendencies_sfc = get_tendencies(rad_model)

# plot the temperature profile and temperature tendencies

# prepare the legend names of different temperature tendencies
plot_names = {'top': 'Total', 'LW': 'LW radiation',
              'SW': 'SW radiation',
              'Radiation': 'Radiation (net)',
              'Convection': 'Convection'}

plt.style.use('ggplot')
plt.figure( figsize=(12,6) )
# plot the first subplot
plt.subplot(1, 2, 1)
plt.plot( Tglobal + 273.15, zstar(level),label='observations',color='C2',linestyle=':',linewidth=4 )
plt.plot(rad_model.Tatm,zstar(rad_model.lev),label='pure radiation',color='C0')
plt.legend()
plt.yticks(zstar(yticks),yticks)
plt.xlabel('Temperature (K)')
plt.ylabel('Pressure (hPa)')
plt.title('Temperature profile')

# plot the second subplot
plt.subplot(1, 2, 2)
for i, name in enumerate(tendencies_atm.data_vars):
    plt.plot(tendencies_atm[name], zstar(rad_model.lev),label=plot_names[name])
plt.legend()
plt.yticks(zstar(yticks),yticks)
plt.xlabel('Temperature tendency (K/day)')
plt.ylabel('Pressure (hPa)')
plt.title('Temperature tendency of the pure radiation case')

"""
Explanation: Now let's get the temperature tendencies. A tendency is the rate of change of a quantity due to different processes; in this case radiation (shortwave and longwave) and (later on) convection. This allows us to quantify the contribution of different physical processes to the evolution of temperature: when a process contriutes to warming the atmosphere, its tendency is positive. When it contributes to cooling the atmosphere, its tendency is negative.
End of explanation
"""


state2 = climlab.column_state(num_lev=50)
#  Parent model process
rcm_dry = climlab.TimeDependentProcess(state=state2)
#  Set the radiation model
rad_dry = climlab.radiation.RRTMG(state=state2, specific_humidity=q_obs, albedo=alb)
#  Convective adjustment
conv_dry = climlab.convection.ConvectiveAdjustment(state=state2, adj_lapse_rate=9.8)
#  Couple everything together
rcm_dry.add_subprocess('Radiation', rad_dry)
rcm_dry.add_subprocess('Convection', conv_dry)

rcm_dry.state.Tatm[:] = rcm_dry.state.Ts
rcm_dry.integrate_years(4)
rcm_dry.ASR - rcm_dry.OLR

"""
Explanation: Question 5: Compare the temperature profile simulated by the radiation process. What are the main discrepancies? Forcus in particular on the near-surface temperature and lapse rate. 
Answer 5:
3. adding dry convection
Obviously radiation is not the only important process in the atmosphere. When the air near the surface gets too hot, it wants to rise, which transports heat upwards, thereby reducing the lapse rate. Let's add this dry convection process to the radiative processes seen above, and see how that affects the temperature profile. What follows is some model gobbledygook that specifies convection to work at a dry adibatic lapse rate of 9.8 degrees/km, as seen in class.
End of explanation
"""


tendencies_atm, tendencies_sfc = get_tendencies(rcm_dry)

plt.style.use('ggplot')
plt.figure( figsize=(12,6) )
plt.subplot(1, 2, 1)
plt.plot( Tglobal + 273.15, zstar(level),label='observations',color='C2',linestyle=':',linewidth=4 )
plt.plot(rad_model.Tatm,zstar(rad_model.lev),label='pure radiation',color='C0')
plt.plot(rcm_dry.Tatm,zstar(rad_model.lev),label='dry adjust',color='C3',linestyle='--')
plt.legend()
plt.yticks(zstar(yticks),yticks)
plt.xlabel('Temperature (K)')
plt.ylabel('Pressure (hPa)')
plt.title('Temperature profile')

plt.subplot(1, 2, 2)
for i, name in enumerate(tendencies_atm.data_vars):
    plt.plot(tendencies_atm[name], zstar(rad_model.lev),label=plot_names[name])
plt.legend()
plt.yticks(zstar(yticks),yticks)
plt.xlabel('Temperature tendency (K/day)')
plt.ylabel('Pressure (hPa)')
plt.title('Temperature tendency of the dry adjustment case')

"""
Explanation: OK, great, the model reaches equilibrium. Now let's see what the temperature tendencies look like:
End of explanation
"""


state3 = climlab.column_state(num_lev=50)
#  Parent model process
rcm = climlab.TimeDependentProcess(state=state3)
#  Fixed relative humidity
h2o = climlab.radiation.ManabeWaterVapor(state=state3)
#  Couple water vapor to radiation
rad = climlab.radiation.RRTMG(state=state3, specific_humidity=h2o.q, albedo=alb)
#  Convective adjustment
conv = climlab.convection.ConvectiveAdjustment(state=state3, adj_lapse_rate=6)
#  Couple everything together
rcm.add_subprocess('Radiation', rad)
rcm.add_subprocess('WaterVapor', h2o)
rcm.add_subprocess('Convection', conv)

rcm.state.Tatm[:] = rcm.state.Ts
rcm.integrate_years(4)
rcm.ASR - rcm.OLR

tendencies_atm, tendencies_sfc = get_tendencies(rcm)

plt.style.use('ggplot')
plt.figure( figsize=(12,6) )
plt.subplot(1, 2, 1)
plt.plot( Tglobal + 273.15, zstar(level),label='observations',color='C2',linestyle=':',linewidth=4 )
plt.plot(rad_model.Tatm,zstar(rad_model.lev),label='pure radiation',color='C0')
plt.plot(rcm_dry.Tatm,zstar(rad_model.lev),label='dry adjust',color='C3',linestyle='--')
plt.plot(rcm.Tatm,zstar(rcm.lev),label='moist adjust',color='C1')
plt.legend()
plt.yticks(zstar(yticks),yticks)
plt.xlabel('Temperature (K)')
plt.ylabel('Pressure (hPa)')
plt.title('Temperature profile')

plt.subplot(1, 2, 2)
for i, name in enumerate(tendencies_atm.data_vars):
    plt.plot(tendencies_atm[name], zstar(rad_model.lev),label=plot_names[name])
plt.legend()
plt.yticks(zstar(yticks),yticks)
plt.xlabel('Temperature tendency (K/day)')
plt.ylabel('Pressure (hPa)')
plt.title('Temperature tendency of the moist adjustment case')

"""
Explanation: Question 6: Does adding the dry convection process improve the simulation of vertical temperature profile? Look at the temperature tendencies, what's the role of convection process? At what level does this role stop, and why do you think that is?
Answer 6:
4. Adding moist convection
Dry convection is fine and good, but in the real world there is often moisture, which condenses once the temperature gets cold enough, so that changes the game. How so? Let us see. We start with some model gobbledygook to set things up proper, then we run the model for 4 years, so it has a chance to reach equilibrium.
End of explanation
"""


