#!/usr/bin/env python
# coding: utf-8

# # Latitude-dependent grey radiation

# Here is a quick example of using the `climlab.GreyRadiationModel` with a latitude dimension and seasonally varying insolation.

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import numpy as np
import matplotlib.pyplot as plt
import climlab


# In[2]:


model = climlab.GreyRadiationModel(num_lev=30, num_lat=90)
print model


# In[3]:


print model.lat


# In[4]:


insolation = climlab.radiation.DailyInsolation(domains=model.Ts.domain)


# In[5]:


model.add_subprocess('insolation', insolation)
model.subprocess.SW.flux_from_space = insolation.insolation


# In[6]:


print model


# In[7]:


model.compute_diagnostics()


# In[8]:


plt.plot(model.lat, model.SW_down_TOA)


# In[9]:


model.Tatm.shape


# In[10]:


model.integrate_years(1)


# In[11]:


plt.plot(model.lat, model.Ts)


# In[12]:


model.integrate_years(1)


# In[13]:


plt.plot(model.lat, model.timeave['Ts'])


# In[14]:


def plot_temp_section(model, timeave=True):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    if timeave:
        field = model.timeave['Tatm'].transpose()
    else:
        field = model.Tatm.transpose()
    cax = ax.contourf(model.lat, model.lev, field)
    ax.invert_yaxis()
    ax.set_xlim(-90,90)
    ax.set_xticks([-90, -60, -30, 0, 30, 60, 90])
    fig.colorbar(cax)


# In[15]:


plot_temp_section(model)


# In[16]:


model2 = climlab.RadiativeConvectiveModel(num_lev=30, num_lat=90)
insolation = climlab.radiation.DailyInsolation(domains=model2.Ts.domain)
model2.add_subprocess('insolation', insolation)
model2.subprocess.SW.flux_from_space = insolation.insolation


# In[17]:


model2.integrate_years(1)


# In[18]:


model2.integrate_years(1)


# In[19]:


plot_temp_section(model2)


# ## Testing out multi-dimensional Band Models

# In[20]:


#  Put in some ozone
import netCDF4 as nc

datapath = "http://ramadda.atmos.albany.edu:8080/repository/opendap/latest/Top/Users/BrianRose/CESM_runs/"
endstr = "/entry.das"

ozone = nc.Dataset( datapath + 'som_input/ozone_1.9x2.5_L26_2000clim_c091112.nc' + endstr )

#  Dimensions of the ozone file
lat = ozone.variables['lat'][:]
lon = ozone.variables['lon'][:]
lev = ozone.variables['lev'][:]

# Taking annual, zonal average of the ozone data
O3_zon = np.mean( ozone.variables['O3'],axis=(0,3) )


# In[21]:


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


# In[22]:


#  Set the ozone mixing ratio
O3_trans = np.transpose(O3_zon)
#  model and ozone data are on the same grid, after the transpose.
print O3_trans.shape
print lev
print model3.lev


# In[23]:


# Put in the ozone
model3.absorber_vmr['O3'] = O3_trans


# In[24]:


print model3.absorber_vmr['O3'].shape
print model3.Tatm.shape


# In[25]:


model3.step_forward()


# In[26]:


model3.integrate_years(1.)


# In[27]:


model3.integrate_years(1.)


# In[28]:


#plt.contour(model3.lat, model3.lev, model3.Tatm.transpose())
plot_temp_section(model3)


# This is now working. Will need to do some model tuning.
# 
# And start to add dynamics!

# ## Adding meridional diffusion!

# In[29]:


print model2


# In[30]:


diffmodel = climlab.process_like(model2)


# In[31]:


# thermal diffusivity in W/m**2/degC
D = 0.05
# meridional diffusivity in 1/s
K = D / diffmodel.Tatm.domain.heat_capacity[0]
print K


# In[32]:


d = climlab.dynamics.MeridionalDiffusion(K=K, state={'Tatm': diffmodel.Tatm}, **diffmodel.param)


# In[33]:


diffmodel.add_subprocess('diffusion', d)


# In[34]:


print diffmodel


# In[35]:


diffmodel.step_forward()


# In[36]:


diffmodel.integrate_years(1)


# In[37]:


diffmodel.integrate_years(1)


# In[38]:


plot_temp_section(model2)
plot_temp_section(diffmodel)


# This works as long as K is a constant.
# 
# The diffusion operation is broadcast over all vertical levels without any special code.

# In[39]:


def inferred_heat_transport( energy_in, lat_deg ):
    '''Returns the inferred heat transport (in PW) by integrating the net energy imbalance from pole to pole.'''
    from scipy import integrate
    from climlab import constants as const
    lat_rad = np.deg2rad( lat_deg )
    return ( 1E-15 * 2 * np.math.pi * const.a**2 * integrate.cumtrapz( np.cos(lat_rad)*energy_in,
            x=lat_rad, initial=0. ) )


# In[40]:


#  Plot the northward heat transport in this model
Rtoa = np.squeeze(diffmodel.timeave['ASR'] - diffmodel.timeave['OLR'])
plt.plot(diffmodel.lat, inferred_heat_transport(Rtoa, diffmodel.lat))


# ### Band model with diffusion

# In[41]:


diffband = climlab.process_like(model3)


# In[42]:


# thermal diffusivity in W/m**2/degC
D = 0.05
# meridional diffusivity in 1/s
K = D / diffband.Tatm.domain.heat_capacity[0]
print K


# In[43]:


d = climlab.dynamics.MeridionalDiffusion(K=K, state={'Tatm': diffband.Tatm}, **diffband.param)
diffband.add_subprocess('diffusion', d)
print diffband


# In[44]:


diffband.integrate_years(1)


# In[45]:


diffband.integrate_years(1)


# In[46]:


plot_temp_section(model3)
plot_temp_section(diffband)


# In[47]:


plt.plot(diffband.lat, diffband.timeave['ASR'] - diffband.timeave['OLR'])


# In[48]:


#  Plot the northward heat transport in this model
Rtoa = np.squeeze(diffband.timeave['ASR'] - diffband.timeave['OLR'])
plt.plot(diffband.lat, inferred_heat_transport(Rtoa, diffband.lat))


# In[ ]: