%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import netCDF4 as nc
import climlab

"""
Explanation: ENV / ATM 415: Climate Laboratory
Exercises on Radiative-Convective Equilibrium with climlab
Tuesday April 5, 2016
End of explanation
"""


ncep_filename = 'air.mon.1981-2010.ltm.nc'
# This will try to read the data over the internet.
#ncep_url = "http://www.esrl.noaa.gov/psd/thredds/dodsC/Datasets/ncep.reanalysis.derived/"
#ncep_air = nc.Dataset( ncep_url + 'pressure/' + ncep_filename )
# Or to read from local disk
ncep_air = nc.Dataset( ncep_filename )

level = ncep_air.variables['level'][:]
lat = ncep_air.variables['lat'][:]
# A log-pressure height coordinate
zstar = -np.log(level/1000)

#  Take averages of the temperature data
Tzon = np.mean(ncep_air.variables['air'][:],axis=(0,3))
Tglobal = np.average( Tzon , weights=np.cos(np.deg2rad(lat)), axis=1) + climlab.constants.tempCtoK
#  Note the useful conversion factor. climlab.constants has lots of commonly used constant pre-defined

#  Here we are plotting with respect to log(pressure) but labeling the axis in pressure units
fig = plt.figure( figsize=(8,6) )
ax = fig.add_subplot(111)
ax.plot( Tglobal , zstar )
yticks = np.array([1000., 750., 500., 250., 100., 50., 20., 10.])
ax.set_yticks(-np.log(yticks/1000.))
ax.set_yticklabels(yticks)
ax.set_xlabel('Temperature (K)', fontsize=16)
ax.set_ylabel('Pressure (hPa)', fontsize=16 )
ax.set_title('Global, annual mean sounding from NCEP Reanalysis', fontsize = 24)
ax.grid()

#  Repeating code from previous notebook ...  set up a column model with observed temperatures.

#  initialize a grey radiation model with 30 levels
col = climlab.GreyRadiationModel()
# interpolate to 30 evenly spaced pressure levels
lev = col.lev
Tinterp = np.interp(lev, np.flipud(level), np.flipud(Tglobal))
# Initialize model with observed temperatures
col.Ts[:] = Tglobal[0]
col.Tatm[:] = Tinterp
#  The tuned value of absorptivity
eps = 0.0534
#  set it
col.subprocess.LW.absorptivity = eps

#  Pure radiative equilibrium

#  Make a clone of our first model
re = climlab.process_like(col)
#  Run out to equilibrium
re.integrate_years(2.)
#  Check for energy balance
re.ASR - re.OLR

#  And set up a RadiativeConvective model, 
rce = climlab.RadiativeConvectiveModel(adj_lapse_rate=6.)
# Set our tuned absorptivity value
rce.subprocess.LW.absorptivity = eps
#  Run out to equilibrium
rce.integrate_years(2.)
#  Check for energy balance
rce.ASR - rce.OLR

#  A handy re-usable routine for making a plot of the temperature profiles
#  We will plot temperatures with respect to log(pressure) to get a height-like coordinate
def plot_sounding(collist):
    color_cycle=['r', 'g', 'b', 'y', 'm']
    # col is either a column model object or a list of column model objects
    if isinstance(collist, climlab.Process):
        # make a list with a single item
        collist = [collist]
    fig = plt.figure()
    ax = fig.add_subplot(111)
    for i, col in enumerate(collist):
        zstar = -np.log(col.lev/climlab.constants.ps)
        ax.plot(col.Tatm, zstar, color=color_cycle[i])
        ax.plot(col.Ts, 0, 'o', markersize=12, color=color_cycle[i])
    #ax.invert_yaxis()
    yticks = np.array([1000., 750., 500., 250., 100., 50., 20., 10.])
    ax.set_yticks(-np.log(yticks/1000.))
    ax.set_yticklabels(yticks)
    ax.set_xlabel('Temperature (K)')
    ax.set_ylabel('Pressure (hPa)')
    ax.grid()
    return ax

#  Make a plot to compare observations and Radiative-Convective Equilibrium
plot_sounding([col, re, rce])

"""
Explanation: First, let's just repeat the calculations we did in the previous notebook RadiativeConvectiveEquilibrium.ipynb
(Just by copying and pasting the relevant code)
End of explanation
"""


#  To read from internet
#datapath = "http://ramadda.atmos.albany.edu:8080/repository/opendap/latest/Top/Users/Brian+Rose/CESM+runs/som_input/"
#endstr = "/entry.das"

ozone_filename = 'ozone_1.9x2.5_L26_2000clim_c091112.nc'
datapath = ''
endstr = ''

ozone = nc.Dataset( datapath + ozone_filename + endstr )


print ozone.variables['O3']

lat = ozone.variables['lat'][:]
lon = ozone.variables['lon'][:]
lev = ozone.variables['lev'][:]


"""
Explanation: Stratospheric ozone
End of explanation
"""


print lev

"""
Explanation: The pressure levels in this dataset are:
End of explanation
"""


O3_zon = np.mean( ozone.variables['O3'],axis=(0,3) )
O3_global = np.sum( O3_zon * np.cos(np.deg2rad(lat)), axis=1 ) / np.sum( np.cos(np.deg2rad(lat) ) )


O3_global.shape

fig = plt.figure(figsize=(15,5))
ax1 = fig.add_subplot(1,2,1)
cax = ax1.contourf(lat, np.log(lev/climlab.constants.ps), O3_zon * 1.E6)
ax1.invert_yaxis()
ax1.set_xlabel('Latitude', fontsize=16)
ax1.set_ylabel('Pressure (hPa)', fontsize=16 )
yticks = np.array([1000., 500., 250., 100., 50., 20., 10., 5.])
ax1.set_yticks( np.log(yticks/1000.) )
ax1.set_yticklabels( yticks )
ax1.set_title('Ozone concentration (annual mean)', fontsize = 16);
plt.colorbar(cax)

ax2 = fig.add_subplot(1,2,2)
ax2.plot( O3_global * 1.E6, np.log(lev/climlab.constants.ps) )
ax2.invert_yaxis()
ax2.set_xlabel('Ozone (ppm)', fontsize=16)
ax2.set_ylabel('Pressure (hPa)', fontsize=16 )
yticks = np.array([1000., 500., 250., 100., 50., 20., 10., 5.])
ax2.set_yticks( np.log(yticks/1000.) )
ax2.set_yticklabels( yticks )
ax2.set_title('Global mean ozone concentration', fontsize = 16);

"""
Explanation: Take the global average of the ozone climatology, and plot it as a function of pressure (or height)
End of explanation
"""


#  Create the column with appropriate vertical coordinate, surface albedo and convective adjustment
band1 = climlab.BandRCModel(lev=lev, adj_lapse_rate=6)
print band1

"""
Explanation: This shows that most of the ozone is indeed in the stratosphere, and peaks near the top of the stratosphere.
The BandRCModel: a radiative model with individual spectral bands
Here is a brief introduction to the climlab.BandRCModel process.
This is a model that divides the spectrum into 7 distinct bands: three shortwave and four longwave.
As we will see, the process works much like the familiar climlab.RadiativeConvectiveModel.
About the spectra
The shortwave is divided into three channels:

Channel 0 is the Hartley and Huggins band (extreme UV, 200 - 340 nm, 1% of total flux, strong ozone absorption)
Channel 1 is Chappuis band (450 - 800 nm, 27% of total flux, moderate ozone absorption)
Channel 2 is remaining radiation (72% of total flux, largely in the visible range, no ozone absorption)

The longwave is divided into four bands:

Band 0 is the window region (between 8.5 and 11 $\mu$m), 17% of total flux.
Band 1 is the CO2 absorption channel (the band of strong absorption by CO2 around 15 $\mu$m), 15% of total flux
Band 2 is a weak water vapor absorption channel, 35% of total flux
Band 3 is a strong water vapor absorption channel, 33% of total flux

The longwave decomposition is not as easily related to specific wavelengths, as in reality there is a lot of overlap between H$_2$O and CO$_2$ absorption features (as well as absorption by other greenhouse gases such as CH$_4$ and N$_2$O that we are not representing).
Create a Radiative-Convective model using these spectral bands, and put in some ozone
Now create a new column model object on the same pressure levels as the ozone data. We are also going set an adjusted lapse rate of 6 K / km, and tune the longwave absorption
End of explanation
"""


band1.absorber_vmr

"""
Explanation: Check out the list of subprocesses.
We now have a process called H2O, in addition to things we've seen before.
This model keeps track of water vapor. More on that later!
About the radiatively active gases
The Band model is aware of three different absorbing gases: O3 (ozone), CO2, and H2O (water vapor). The abundances of these gases are stored in a dictionary of arrays as follows:
End of explanation
"""


band1.state

band1.integrate_years(2)

# Check for energy balance
band1.ASR - band1.OLR

#  Add another line to our graph!
plot_sounding([col, re, rce, band1])

"""
Explanation: Ozone and CO2 are both specified in the model. The default, as you see above, is zero ozone, and constant (well-mixed) CO2 at a volume mixing ratio of 3.8E-4 or 380 ppm.
Water vapor is handled differently: it is determined by the model at each timestep. We make the following assumptions, following a classic paper on radiative-convective equilibrium by Manabe and Wetherald (J. Atmos. Sci. 1967):

the relative humidity just above the surface is fixed at 77% (can be changed of course... see the parameter col1.relative_humidity
water vapor drops off linearly with pressure
there is a small specified amount of water vapor in the stratosphere.

Radiative-convective equilibrium in the band model without ozone
End of explanation
"""


band2 = climlab.process_like(band1)
print band2

band2.absorber_vmr['O3'] = O3_global
band2.absorber_vmr

#  Run the model out to equilibrium!
band2.integrate_years(2.)

#  Add another line to our graph!
plot_sounding([col, re, rce, band1, band2])

"""
Explanation: Now put in the ozone
End of explanation
"""


band2.absorber_vmr['CO2']

# Let's double CO2 and calculate radiative forcing
band3 = climlab.process_like(band2)
band3.absorber_vmr['CO2'] *= 2.
band3.absorber_vmr['CO2']

"""
Explanation: Once we include ozone we get a well-defined stratosphere. 
Things to consider / try:

Here we used the global annual mean Q = 341.3 W m$^{-2}$. We might want to consider latitudinal or seasonal variations in Q.
We also used the global annual mean ozone profile! Ozone varies tremendously in latitude and by season. That information is all contained in the ozone data file we opened above. We might explore the effects of those variations.
We can calculate climate sensitivity in this model by doubling the CO2 concentration and re-running out to the new equilibrium. Does the amount of ozone affect the climate sensitivity?  (example below)
An important shortcoming of the model: there are no clouds! (that would be the next step in the hierarchy of column models)
Clouds would act both in the shortwave (increasing the albedo, cooling the climate) and in the longwave (greenhouse effect, warming the climate). Which effect is stronger depends on the vertical structure of the clouds (high or low clouds) and their optical properties (e.g. thin cirrus clouds are nearly transparent to solar radiation but are good longwave absorbers).

Climate sensitivity to doubling CO2
We now actually have a model that is aware of the amount of atmospheric CO2:
End of explanation
"""


band3.compute_diagnostics()
print 'The radiative forcing for doubling CO2 is %f W/m2.' % (band2.OLR - band3.OLR)

#  and make another copy, which we will integrate out to equilibrium
band4 = climlab.process_like(band3)
band4.integrate_years(3)
band4.ASR - band4.OLR

DeltaT = band4.Ts - band2.Ts
print 'The Equilibrium Climate Sensitivity is %f K.' % DeltaT

"""
Explanation: We've just increased CO2 from 380 ppm to 760 ppm.
Radiative forcing is the instantaneous change in OLR...
End of explanation
"""


#  We multiply the H2O mixing ratio by 1000 to get units of g / kg
#  (amount of water vapor per mass of air)
plt.plot( band2.absorber_vmr['H2O'] *1000, lev, label='before 2xCO2')
plt.plot( band4.absorber_vmr['H2O'] * 1000., lev, label='after 2xCO2')
plt.xlabel('g/kg H2O')
plt.ylabel('pressure (hPa)')
plt.legend(loc='upper right')
plt.grid()
#  This reverses the axis so pressure decreases upward
plt.gca().invert_yaxis()

"""
Explanation: Investigating the role of water vapor feedback in climate sensitivity
Let's make some plots of the vertical profiles of water vapor before and after the climate change.
End of explanation
"""


#  First, make a new clone
noh2o = climlab.process_like(band2)

#  See what absorbing gases are currently present
noh2o.absorber_vmr

#  double the CO2 !
noh2o.absorber_vmr['CO2'] *= 2

#  Check out the list of subprocesses
print noh2o

#  Remove the process that changes the H2O
noh2o.remove_subprocess('H2O')
print noh2o

"""
Explanation: Water vapor decreases from the surface upward mostly because the temperature decreases!
Our model uses a parameterization that holds the relative humidity to a fixed profile. Since the saturation point increases strongly with temperature, the actual amount of water vapor increases as the climate warms to preserve the relative humidity.
We can see this increase in the graph above.
Removing the fixed-relative-humidity water vapor subprocess
Make a new model that does NOT have a water vapor feedback. Hold the H2O amount fixed as we warm.
End of explanation
"""


noh2o.absorber_vmr

"""
Explanation: Notice that the subprocess labelled 'H2O' is gone from the list.
This was the process that modified the water vapor. But note that the model still contains H2O:
End of explanation
"""


noh2o.integrate_years(3)
noh2o.ASR - noh2o.OLR

#  Repeat the same plot of water vapor profiles, but add a third curve
#  We'll plot the new model as a dashed line to make it easier to see.
plt.plot( band2.absorber_vmr['H2O'] *1000, lev, label='before 2xCO2')
plt.plot( band4.absorber_vmr['H2O'] * 1000., lev, label='after 2xCO2')
plt.plot( noh2o.absorber_vmr['H2O'] * 1000., lev, linestyle='--', label='No H2O feedback')
plt.xlabel('g/kg H2O')
plt.ylabel('pressure (hPa)')
plt.legend(loc='upper right')
plt.grid()
#  This reverses the axis so pressure decreases upward
plt.gca().invert_yaxis()

"""
Explanation: But this will be held fixed now as the climate changes in noh2o.
Let's go ahead and run it out to equilibrium.
End of explanation
"""


DeltaT_noh2o = noh2o.Ts - band2.Ts
print 'The Equilibrium Climate Sensitivity with water vapor feedback is %f K.' % DeltaT_noh2o

"""
Explanation: Indeed, the water vapor is identical in the new equilibrium climate to the old pre-CO2-increase model. 
So this model does NOT have a water vapor feedback.
How does this affect the climate sensivity?
End of explanation
"""


DeltaT - DeltaT_noh2o

"""
Explanation: So the effect of the water vapor feedback on the climate sensitivity to doubled CO2 is:
End of explanation
"""


#  Create the column with appropriate vertical coordinate, surface albedo and convective adjustment
tuneband = climlab.BandRCModel(lev=lev, adj_lapse_rate=6, albedo_sfc=0.1)
print tuneband

tuneband.absorber_vmr['O3'] = O3_global

tuneband.absorber_vmr

tuneband.compute_diagnostics()

tuneband.absorber_vmr

tuneband.subprocess.LW.absorption_cross_section

#  This set of parameters gives correct Ts and ASR
#   But seems weird... the radiative forcing for doubling CO2 is tiny
#tuneband.subprocess.SW.reflectivity[15] = 0.255
#tuneband.subprocess.LW.absorption_cross_section['CO2'][1] = 2.
#tuneband.subprocess.LW.absorption_cross_section['H2O'][2] = 0.45

#  Just tune the surface albedo, won't get correct ASR but get correct Ts
tuneband = climlab.BandRCModel(lev=lev, adj_lapse_rate=6, albedo_sfc=0.22)

tuneband.integrate_converge()

tuneband.Ts

tuneband.ASR

tuneband.OLR

tuneband.param

tband2 = climlab.process_like(tuneband)

tband2.subprocess['LW'].absorber_vmr['CO2'] *= 2

tband2.compute_diagnostics()

tband2.OLR

tband2.step_forward()

tband2.OLR

tband2.integrate_converge()

tband2.Ts


ncep_air

ncep_air.variables.keys()

latmodel = climlab.BandRCModel(lat=lat, lev=lev, adj_lapse_rate=6, albedo_sfc=0.22)
print latmodel

latmodel.Tatm.shape

ncep_air.variables['air'].shape

tuneband.Ts

tuneband.subprocess.SW.albedo_sfc = 0.4

tuneband.integrate_converge()

tuneband.Ts

tuneband.

#  make a model on the same grid as the ozone
model = climlab.BandRCModel(lev=lev, lat=lat, albedo_sfc=0.22)
insolation = climlab.radiation.insolation.AnnualMeanInsolation(domains=model.Ts.domain)
model.add_subprocess('insolation', insolation)
model.subprocess.SW.flux_from_space = model.subprocess.insolation.insolation
print model

model.subprocess.insolation.insolation

"""
Explanation: We get about an additional degree of warming from the water vapor increase.
MORE MESSING AROUND
End of explanation
"""


