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

"""
Explanation: ENV / ATM 415: Climate Laboratory
The planetary energy budget in CESM simulations
Tuesday April 19 and Thursday April 21, 2016
End of explanation
"""


##  To read data over the internet
control_filename = 'som_1850_f19.cam.h0.clim.nc'
datapath = 'http://ramadda.atmos.albany.edu:8080/repository/opendap/latest/Top/Users/Brian+Rose/CESM+runs/'
endstr = '/entry.das'
control = nc.Dataset( datapath + 'som_1850_f19/' + control_filename + endstr )

##  To read from a local copy of the file 
## (just a small subset of the total list of variables, to save disk space)
#ontrol_filename = 'som_1850_f19.cam.h0.clim_subset.nc'
#control = nc.Dataset( control_filename )

"""
Explanation: Open the output from our control simulation with the slab ocean version of the CESM:
End of explanation
"""


for v in control.variables: print v

"""
Explanation: The full file from the online server contains many many variables, describing all aspects of the model climatology.
Whether we see a long list or a short list in the following code block depends on whether we are reading the full output file or the much smaller subset:
End of explanation
"""


for field in ['TS', 'FLNT', 'FSNT', 'FLNTC', 'FSNTC']:
    print control.variables[field]

"""
Explanation: Today we need just a few of these variables:

TS: the surface temperature
FLNT: the longwave radiation at the top of the atmosphere (i.e. what we call the OLR)
FSNT: the net shortwave radiation at the top of the atmosphere (i.e. what we call the ASR)
FLNTC: the clear-sky OLR
FSNTC: the clear-sky ASR

Take a look at some of the meta-data for these fields:
End of explanation
"""


print control.variables['lat']

"""
Explanation: Each one of these variables has dimensions (12, 96, 144), which corresponds to time (12 months), latitude and longitude.
Take a look at one of the coordinate variables:
End of explanation
"""


lat = control.variables['lat'][:]
lon = control.variables['lon'][:]
print lat

"""
Explanation: Now let's load in the coordinate data, to use later for plotting:
End of explanation
"""


#  A re-usable function to make a map of a 2d field on a latitude / longitude grid
def make_map(field_2d):
    #  Make a filled contour plot
    fig = plt.figure(figsize=(10,5))
    cax = plt.contourf(lon, lat, field_2d)
    #  draw a single contour to outline the continents
    plt.contour( lon, lat, control.variables['LANDFRAC'][0,:,:], [0.5], colors='k')
    plt.xlabel('Longitude (degrees east)')
    plt.ylabel('Latitude (degrees north)')
    plt.colorbar(cax)

#  Here is a convenient function that takes the name of a variable in our CESM output
#  and make a map of its annual average
def map_this(fieldname, dataset=control):
    field = dataset.variables[fieldname][:]
    field_annual = np.mean(field, axis=0)
    make_map(field_annual)

#  Use this function to make a quick map of the annual average surface temperature:
map_this('TS')

"""
Explanation: Surface temperature in the control simulation
End of explanation
"""


#  The lat/lon dimensions after taking the time average:
TS_annual = np.mean(control.variables['TS'][:], axis=0)
TS_annual.shape

"""
Explanation: Computing a global average
End of explanation
"""


def global_mean(field_2d):
    '''This function takes a 2D array on a regular latitude-longitude grid 
       and returns the global area-weighted average'''
    zonal_mean = np.mean(field_2d, axis=1)
    return np.average(zonal_mean, weights=np.cos(np.deg2rad(lat)))

#  Again, a convenience function that takes just the name of the model output field
#  and returns its time and global average
def global_mean_this(fieldname, dataset=control):
    field = dataset.variables[fieldname][:]
    field_annual = np.mean(field, axis=0)
    return global_mean(field_annual)

"""
Explanation: Define a little re-usable function to take the global average of any of these fields:
End of explanation
"""


global_mean_this('TS')

"""
Explanation: Now compute the global average surface temperature in the simulation:
End of explanation
"""


#  This field is not included in the small subset file
#  so this will only work if you are reading the full file from the online server
control.variables['CLOUD']

"""
Explanation: Cloud cover in the control simulation
The model simulates cloud amount in every grid box. The cloud field is thus 4-dimensional:
End of explanation
"""


control.variables['CLDTOT']

map_this('CLDTOT')

"""
Explanation: To simplify things we can just look at the total cloud cover, integrated from the surface to the top of the atmosphere:
End of explanation
"""


#  To get you started, here is the ASR
map_this('FSNT')

map_this('FLNT')

net_radiation = np.mean(control.variables['FSNT'][:] - control.variables['FLNT'][:], axis=0)
make_map(net_radiation)

"""
Explanation: Which parts of Earth are cloudy and which are not? (at least in this simulation)
Exercise 1: Make three maps: ASR, OLR, and the net radiation ASR-OLR (all annual averages)
What interesting features do you see on these maps?
End of explanation
"""


global_mean(net_radiation)

"""
Explanation: Exercise 2: Calculate the global average net radiation.
Is it close to zero? What does that mean?
End of explanation
"""


#  The meta-data:
control.variables['co2vmr']

# The data themselves, expressed in ppm:
control.variables['co2vmr'][:] * 1E6

"""
Explanation: Exercise 3: Make maps of the clear-sky ASR and clear-sky OLR
These diagnostics have been calculated by the GCM. Basically at every timestep, the GCM calculates the radiation twice: once with the clouds and once without the clouds.
Exercise 4: Make a map of the Cloud Radiative Effect
Recall that we define $CRE$ as
$$ CRE = \left( ASR - ASR_{clear} \right) - \left( OLR - OLR_{clear} \right) $$
This quantity is positive where the clouds have a net warming effect on the climate.
Exercise 5: in the global average, are the clouds warming or cooling the climate in the CESM control simulation?
Climate sensitivity in the CESM: the effects of doubling CO2
How much CO2 was in the atmosphere for the control simulation?
This information is available in the full output file (this won't work with the local subset file):
End of explanation
"""


##  To read data over the internet
# doubleCO2_filename = 'som_1850_2xCO2.cam.h0.clim.nc'
# doubleCO2 = nc.Dataset( datapath + 'som_1850_f19/' + doubleCO2_filename + endstr )

##  To read from a local copy of the file 
## (just a small subset of the total list of variables, to save disk space)
doubleCO2_filename = 'som_1850_2xCO2.cam.h0.clim_subset.nc'
doubleCO2 = nc.Dataset( doubleCO2_filename )

"""
Explanation: Answer: the CO2 concentration is 284.7 ppm in the control simulation.
Now we want to see how the climate changes in the CESM when we double CO2 and run it out to equilibrium.
I have done this. Because we are using a slab ocean model, it reaches equilibrium after just a few decades.
Let's now open up the output file from the 2xCO2 scenario:
End of explanation
"""


doubleCO2.variables['co2vmr'][:] * 1E6

"""
Explanation: This file has all the same fields as control, but they reflect the new equilibrium climate after doubling CO2.
Let's verify the CO2 amount:
End of explanation
"""


