# [ATM 623: Climate Modeling](../index.ipynb)

[Brian E. J. Rose](http://www.atmos.albany.edu/facstaff/brose/index.html), University at Albany

# Lecture 7: Grey radiation modeling with climlab

### About these notes:

This document uses the interactive [`Jupyter notebook`](https://jupyter.org) format. The notes can be accessed in several different ways:

- The interactive notebooks are hosted on `github` at https://github.com/brian-rose/ClimateModeling_courseware
- The latest versions can be viewed as static web pages [rendered on nbviewer](http://nbviewer.ipython.org/github/brian-rose/ClimateModeling_courseware/blob/master/index.ipynb)
- A complete snapshot of the notes as of May 2017 (end of spring semester) are [available on Brian's website](http://www.atmos.albany.edu/facstaff/brose/classes/ATM623_Spring2017/Notes/index.html).

[Also here is a legacy version from 2015](http://www.atmos.albany.edu/facstaff/brose/classes/ATM623_Spring2015/Notes/index.html).

Many of these notes make use of the `climlab` package, available at https://github.com/brian-rose/climlab

```python
#  Ensure compatibility with Python 2 and 3
from __future__ import print_function, division
```

## Contents

1. [Introducing `climlab`](#section1)
2. [Using `climlab` to implement the two-layer leaky greenhouse model](#section2)
3. [The observed annual, global mean temperature profile](#section3)
4. [A 30-layer model using the observed temperatures](#section4)
5. [Radiative forcing in the 30-layer model](#section5)
6. [Radiative equilibrium in the 30-layer model](#section6)
7. [Radiative-Convective Equilibrium in the 30-layer model](#section7)
8. [Putting stratospheric ozone in the grey-gas model](#section8)

____________

## 1. Introducing `climlab`
____________

``climlab`` is a flexible engine for process-oriented climate modeling.
It is based on a very general concept of a model as a collection of individual,
interacting processes. ``climlab`` defines a base class called ``Process``, which
can contain an arbitrarily complex tree of sub-processes (each also some
sub-class of ``Process``). Every climate process (radiative, dynamical,
physical, turbulent, convective, chemical, etc.) can be simulated as a stand-alone
process model given appropriate input, or as a sub-process of a more complex model.
New classes of model can easily be defined and run interactively by putting together an
appropriate collection of sub-processes.

``climlab`` is a work-in-progress, and the code base will evolve substantially over the course of this semester.
The latest code can always be found on ``github``:

https://github.com/brian-rose/climlab

You are strongly encouraged to clone the ``climlab`` repository and use ``git`` to keep your local copy up-to-date.

Running this notebook requires that ``climlab`` is already installed on your system.

```python
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import xarray as xr
from xarray.ufuncs import cos, deg2rad, log
import climlab
```

____________

## 2. Using `climlab` to implement the two-layer leaky greenhouse model
____________

One of the things that ``climlab`` is set up to do is the grey-radiation modeling we have already been discussing.

Since we already derived a [complete analytical solution to the two-layer leaky greenhouse model](Lecture06 -- Elementary greenhouse models.ipynb), we will use this to validate the `climlab` code.

### Validation

We want to verify that the model reproduces the observed OLR given observed temperatures, and the absorptivity that we tuned in the analytical model. The target numbers are:

\begin{align}
T_s &= 288 \text{ K} \\
T_0 &= 275 \text{ K} \\
T_1 &= 230 \text{ K} \\
\end{align}

$$ \epsilon = 0.586  $$

$$ OLR = 238.5 \text{ W m}^{-2} $$

### Initialize a model in `climlab`
The first thing we do is create a new model.

The following example code is sparsely commented but will hopefully orient you on the basics of defining and working with a `climlab Process` object.

```python
#  Test in a 2-layer atmosphere
col = climlab.GreyRadiationModel(num_lev=2)
print( col)
```

    climlab Process of type <class 'climlab.model.column.GreyRadiationModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (2,)
    The subprocess tree:
    top: <class 'climlab.model.column.GreyRadiationModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>

```python
col.subprocess
```

    {'LW': <climlab.radiation.greygas.GreyGas at 0x1a0be75898>,
     'SW': <climlab.radiation.greygas.GreyGasSW at 0x106d19358>,
     'insolation': <climlab.radiation.insolation.FixedInsolation at 0x30aa01d68>}

Every item in the above dictionary is itself an instance of the `climlab.Process` object:

```python
print( col.subprocess['LW'])
```

    climlab Process of type <class 'climlab.radiation.greygas.GreyGas'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (2,)
    The subprocess tree:
    top: <class 'climlab.radiation.greygas.GreyGas'>

The `state` dictionary holds the state variables of the model. In this case, temperatures:

```python
col.state
```

    {'Tatm': Field([ 200.,  278.]), 'Ts': Field([ 288.])}

Access these either through dictionary methods or as attributes of the model object:

```python
print( col.state['Ts'])
print( col.Ts)
col.Ts is col.state['Ts']
```

    [ 288.]
    [ 288.]

    True

Now we are assigning the "observed" temperatures to our model state:

```python
col.Ts[:] = 288.
col.Tatm[:] = np.array([230., 275.])
col.state
```

    {'Tatm': Field([ 230.,  275.]), 'Ts': Field([ 288.])}

```python
LW = col.subprocess['LW']
print( LW)
```

    climlab Process of type <class 'climlab.radiation.greygas.GreyGas'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (2,)
    The subprocess tree:
    top: <class 'climlab.radiation.greygas.GreyGas'>

```python
LW.absorptivity
```

    Field([ 0.47737425,  0.47737425])

```python
#  copying the tuned value of epsilon from Lecture 6 notes
LW.absorptivity = 0.586
LW.absorptivity
```

    Field([ 0.586,  0.586])

```python
#  This does all the calculations that would be performed at each time step,
#  but doesn't actually update the temperatures
col.compute_diagnostics()
#  Print out the dictionary
col.diagnostics
```

    {'ASR': array([ 239.2513]),
     'LW_absorbed_atm': array([ 20.02990881, -96.98387764]),
     'LW_absorbed_sfc': Field([-161.57076227]),
     'LW_down_sfc': array([ 228.53426769]),
     'LW_emission': Field([  92.98664086,  190.03779837]),
     'LW_up_sfc': Field([ 390.10502995]),
     'OLR': array([ 238.5247311]),
     'SW_absorbed_atm': array([ 0.,  0.]),
     'SW_absorbed_sfc': Field([ 239.2513]),
     'SW_down_TOA': array([ 341.3]),
     'SW_up_TOA': array([ 102.0487]),
     'SW_up_sfc': Field([ 102.0487]),
     'absorbed': array([ 0.,  0.]),
     'absorbed_total': 0.0,
     'coszen': Field([ 1.]),
     'emission': Field([ 0.,  0.]),
     'emission_sfc': Field([ 0.]),
     'flux_from_sfc': Field([ 102.0487]),
     'flux_reflected_up': array([   0.    ,    0.    ,  102.0487]),
     'flux_to_sfc': array([ 341.3]),
     'flux_to_space': array([ 102.0487]),
     'insolation': array([ 341.3]),
     'planetary_albedo': array([ 0.299])}

```python
#  Check OLR against our analytical solution
col.OLR
```

    array([ 238.5247311])

```python
# Like the state variables, the diagnostics can also be accessed in two different ways
col.diagnostics['OLR']
```

    array([ 238.5247311])

```python
col.state
```

    {'Tatm': Field([ 230.,  275.]), 'Ts': Field([ 288.])}

```python
# perform a single time step
col.step_forward()
```

```python
col.state
```

    {'Tatm': Field([ 230.33784312,  273.3641795 ]), 'Ts': Field([ 289.60514636])}

We just stepped forward one discreet unit in time. Because we didn't specify a timestep when we created the model, it is set to a default value:

```python
col.timestep
```

    86400.0

which is 1 day (expressed in seconds).

Now we will integrate the model out to equilibrium.

We could easily write a loop to call the `step_forward()` method many times.

Or use a handy shortcut that allows us to specify the integration length in physical time units:

```python
# integrate out to radiative equilibrium
col.integrate_years(2.)
```

    Integrating for 730 steps, 730.4844 days, or 2.0 years.
    Total elapsed time is 2.0014116660123062 years.

```python
# Check for equilibrium
col.ASR - col.OLR
```

    array([ -2.39951930e-07])

```python
#  The temperatures at radiative equilibrium
col.state
```

    {'Tatm': Field([ 233.72131685,  262.28540231]), 'Ts': Field([ 296.38447748])}

Compare these to the analytical solutions for radiative equilibrium with $\epsilon = 0.58$:

\begin{align}
T_s &= 296.4 \text{ K} \\
T_0 &= 262.3 \text{ K} \\
T_1 &= 233.8 \text{ K} \\
\end{align}

So it looks like `climlab` agrees with our analytical results to within 0.1 K. That's good.

____________

## 3. The observed annual, global mean temperature profile
____________

We want to model the OLR in a column whose temperatures match observations. As we've done before, we'll calculate the global, annual mean air temperature from the NCEP Reanalysis data.

```python
# This will try to read the data over the internet.
ncep_filename = 'air.mon.1981-2010.ltm.nc'
#  to read over internet
ncep_url = "http://www.esrl.noaa.gov/psd/thredds/dodsC/Datasets/ncep.reanalysis.derived/pressure/"
path = ncep_url
#  Open handle to data
ncep_air = xr.open_dataset( path + ncep_filename, decode_times=False )
```

```python
#  Take global, annual average and convert to Kelvin
weight = cos(deg2rad(ncep_air.lat)) / cos(deg2rad(ncep_air.lat)).mean(dim='lat')
Tglobal = (ncep_air.air * weight).mean(dim=('lat','lon','time'))
print( Tglobal)
```

    <xarray.DataArray (level: 17)>
    array([ 15.179082,  11.207002,   7.838327,   0.219941,  -6.448343, -14.888844,
           -25.570467, -39.369685, -46.797908, -53.652235, -60.563551, -67.006048,
           -65.532927, -61.486637, -55.853581, -51.593945, -43.219982])
    Coordinates:
      * level    (level) float32 1000.0 925.0 850.0 700.0 600.0 500.0 400.0 ...

We're going to convert this to degrees Kelvin, using a handy list of pre-defined constants in `climlab.constants`

```python
climlab.constants.tempCtoK
```

    273.15

```python
Tglobal += climlab.constants.tempCtoK
print( Tglobal)
```

    <xarray.DataArray (level: 17)>
    array([ 288.329082,  284.357002,  280.988327,  273.369941,  266.701657,
            258.261156,  247.579533,  233.780315,  226.352092,  219.497765,
            212.586449,  206.143952,  207.617073,  211.663363,  217.296419,
            221.556055,  229.930018])
    Coordinates:
      * level    (level) float32 1000.0 925.0 850.0 700.0 600.0 500.0 400.0 ...

```python
#  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 zstar(lev):
    return -np.log(lev / climlab.constants.ps)

def plot_soundings(result_list, name_list, plot_obs=True, fixed_range=True):
    color_cycle=['r', 'g', 'b', 'y']
    # col is either a column model object or a list of column model objects
    #if isinstance(state_list, climlab.Process):
    #    # make a list with a single item
    #    collist = [collist]
    fig, ax = plt.subplots(figsize=(9,9))
    if plot_obs:
        ax.plot(Tglobal, zstar(Tglobal.level), color='k', label='Observed')
    for i, state in enumerate(result_list):
        Tatm = state['Tatm']
        lev = Tatm.domain.axes['lev'].points
        Ts = state['Ts']
        ax.plot(Tatm, zstar(lev), color=color_cycle[i], label=name_list[i])
        ax.plot(Ts, 0, 'o', markersize=12, color=color_cycle[i])
    #ax.invert_yaxis()
    yticks = np.array([1000., 750., 500., 250., 100., 50., 20., 10., 5.])
    ax.set_yticks(-np.log(yticks/1000.))
    ax.set_yticklabels(yticks)
    ax.set_xlabel('Temperature (K)', fontsize=14)
    ax.set_ylabel('Pressure (hPa)', fontsize=14)
    ax.grid()
    ax.legend()
    if fixed_range:
        ax.set_xlim([200, 300])
        ax.set_ylim(zstar(np.array([1000., 5.])))
    #ax2 = ax.twinx()

    return ax
```

```python
plot_soundings([],[] );
```

____________

## 4. A 30-layer model using the observed temperatures
____________

```python
#  initialize a grey radiation model with 30 levels
col = climlab.GreyRadiationModel()
print( col)
```

    climlab Process of type <class 'climlab.model.column.GreyRadiationModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (30,)
    The subprocess tree:
    top: <class 'climlab.model.column.GreyRadiationModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>

```python
col.lev
```

    array([  16.66666667,   50.        ,   83.33333333,  116.66666667,
            150.        ,  183.33333333,  216.66666667,  250.        ,
            283.33333333,  316.66666667,  350.        ,  383.33333333,
            416.66666667,  450.        ,  483.33333333,  516.66666667,
            550.        ,  583.33333333,  616.66666667,  650.        ,
            683.33333333,  716.66666667,  750.        ,  783.33333333,
            816.66666667,  850.        ,  883.33333333,  916.66666667,
            950.        ,  983.33333333])

```python
 col.lev_bounds
```

    array([    0.        ,    33.33333333,    66.66666667,   100.        ,
             133.33333333,   166.66666667,   200.        ,   233.33333333,
             266.66666667,   300.        ,   333.33333333,   366.66666667,
             400.        ,   433.33333333,   466.66666667,   500.        ,
             533.33333333,   566.66666667,   600.        ,   633.33333333,
             666.66666667,   700.        ,   733.33333333,   766.66666667,
             800.        ,   833.33333333,   866.66666667,   900.        ,
             933.33333333,   966.66666667,  1000.        ])

```python
# interpolate to 30 evenly spaced pressure levels
lev = col.lev
Tinterp = np.interp(lev, np.flipud(Tglobal.level), np.flipud(Tglobal))
Tinterp
#  Need to 'flipud' because the interpolation routine
#  needs the pressure data to be in increasing order
```

    array([ 224.34737604,  211.66336272,  206.96235273,  208.29145078,
            212.5864485 ,  217.1939926 ,  221.78254026,  226.35209151,
            231.30424019,  236.08018424,  240.67992366,  245.27966307,
            249.35980325,  252.92034421,  256.48088517,  259.66790594,
            262.4814065 ,  265.29490707,  267.81303803,  270.03579937,
            272.25856072,  274.21642867,  275.90940323,  277.60237779,
            279.29535236,  280.98832692,  282.48551584,  283.98270476,
            285.68102879,  287.44639785])

```python
# Initialize model with observed temperatures
col.Ts[:] = Tglobal[0]
col.Tatm[:] = Tinterp
```

```python
# This should look just like the observations
result_list = [col.state]
name_list = ['Observed, interpolated']
plot_soundings(result_list, name_list);
```

### Tune absorptivity to get observed OLR

```python
col.compute_diagnostics()
col.OLR
```

    array([ 263.15004392])

```python
# Need to tune absorptivity to get OLR = 238.5
epsarray = np.linspace(0.01, 0.1, 100)
OLRarray = np.zeros_like(epsarray)
```

```python
for i in range(epsarray.size):
    col.subprocess['LW'].absorptivity = epsarray[i]
    col.compute_diagnostics()
    OLRarray[i] = col.OLR

plt.plot(epsarray, OLRarray)
plt.grid()
plt.xlabel('epsilon')
plt.ylabel('OLR')
```

The necessary value seems to lie near 0.055 or so.

We can be more precise with a numerical root-finder.

```python
def OLRanom(eps):
    col.subprocess['LW'].absorptivity = eps
    col.compute_diagnostics()
    return col.OLR - 238.5
```

```python
# Use numerical root-finding to get the equilibria
from scipy.optimize import brentq
# brentq is a root-finding function
#  Need to give it a function and two end-points
#  It will look for a zero of the function between those end-points
eps = brentq(OLRanom, 0.01, 0.1)
print( eps)
```

    0.05369075475256376

```python
col.subprocess.LW.absorptivity = eps
col.subprocess.LW.absorptivity
```

    Field([ 0.05369075,  0.05369075,  0.05369075,  0.05369075,  0.05369075,
            0.05369075,  0.05369075,  0.05369075,  0.05369075,  0.05369075,
            0.05369075,  0.05369075,  0.05369075,  0.05369075,  0.05369075,
            0.05369075,  0.05369075,  0.05369075,  0.05369075,  0.05369075,
            0.05369075,  0.05369075,  0.05369075,  0.05369075,  0.05369075,
            0.05369075,  0.05369075,  0.05369075,  0.05369075,  0.05369075])

```python
col.compute_diagnostics()
col.OLR
```

    array([ 238.5])

____________

## 5. Radiative forcing in the 30-layer model
____________

Let's compute radiative forcing for a **2% increase in absorptivity**.

```python
#  clone our model using a built-in climlab function
col2 = climlab.process_like(col)
print( col2)
```

    climlab Process of type <class 'climlab.model.column.GreyRadiationModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (30,)
    The subprocess tree:
    top: <class 'climlab.model.column.GreyRadiationModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>

```python
col2.subprocess['LW'].absorptivity *= 1.02
col2.subprocess['LW'].absorptivity
```

    Field([ 0.05476457,  0.05476457,  0.05476457,  0.05476457,  0.05476457,
            0.05476457,  0.05476457,  0.05476457,  0.05476457,  0.05476457,
            0.05476457,  0.05476457,  0.05476457,  0.05476457,  0.05476457,
            0.05476457,  0.05476457,  0.05476457,  0.05476457,  0.05476457,
            0.05476457,  0.05476457,  0.05476457,  0.05476457,  0.05476457,
            0.05476457,  0.05476457,  0.05476457,  0.05476457,  0.05476457])

```python
#  Radiative forcing by definition is the change in TOA radiative flux,
# HOLDING THE TEMPERATURES FIXED.
col2.Ts - col.Ts
```

    Field([ 0.])

```python
col2.Tatm - col.Tatm
```

    Field([ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.])

```python
col2.compute_diagnostics()
col2.OLR
```

    array([ 236.6538411])

The OLR decreased after we added the extra absorbers, as we expect. Now we can calculate the Radiative Forcing:

```python
RF = -(col2.OLR - col.OLR)
print( 'The radiative forcing is %.2f W/m2.' %RF)
```

    The radiative forcing is 1.85 W/m2.

____________

## 6. Radiative equilibrium in the 30-layer model
____________

```python
re = climlab.process_like(col)
```

```python
#  To get to equilibrium, we just time-step the model forward long enough
re.integrate_years(1.)
```

    Integrating for 365 steps, 365.2422 days, or 1.0 years.
    Total elapsed time is 0.9993368783782377 years.

```python
#  Check for energy balance
print( 'The net downward radiative flux at TOA is %.4f W/m2.' %(re.ASR - re.OLR))
```

    The net downward radiative flux at TOA is -0.0016 W/m2.

```python
result_list.append(re.state)
name_list.append('Radiative equilibrium (grey gas)')
plot_soundings(result_list, name_list)
```

Some properties of the **radiative equilibrium** temperature profile:

- The surface is warmer than observed.
- The lower troposphere is colder than observed.
- Very cold air is sitting immediately above the warm surface.
- There is no tropopause, no stratosphere.

____________

## 7. Radiative-Convective Equilibrium in the 30-layer model
____________

We recognize that the large drop in temperature just above the surface is unphysical. Parcels of air in direct contact with the ground will be warmed by mechansisms other than radiative transfer.

These warm air parcels will then become buoyant, and will convect upward, mixing their heat content with the environment.

We **parameterize** the statistical effects of this mixing through a **convective adjustment**.

At each timestep, our model checks for any locations at which the **lapse rate** exceeds some threshold. Unstable layers are removed through an energy-conserving mixing formula.

This process is assumed to be fast relative to radiative heating. In the model, it is instantaneous.

### Add the convective adjustment as an additional subprocess

```python
#  Here is the existing model
print( re)
```

    climlab Process of type <class 'climlab.model.column.GreyRadiationModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (30,)
    The subprocess tree:
    top: <class 'climlab.model.column.GreyRadiationModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>

```python
#  First we make a new clone
rce = climlab.process_like(re)
#  Then create a new ConvectiveAdjustment process
conv = climlab.convection.ConvectiveAdjustment(state=rce.state,
                                               adj_lapse_rate=6.)
#  And add it to our model
rce.add_subprocess('Convective Adjustment', conv)
print( rce)
```

    climlab Process of type <class 'climlab.model.column.GreyRadiationModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (30,)
    The subprocess tree:
    top: <class 'climlab.model.column.GreyRadiationModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>
       Convective Adjustment: <class 'climlab.convection.convadj.ConvectiveAdjustment'>

This model is exactly like our previous models, except for one additional subprocess called ``Convective Adjustment``.

We passed a parameter ``adj_lapse_rate`` (in K / km) that sets the neutrally stable lapse rate -- in this case, 6 K / km.

This number is chosed to very loosely represent the net effect of **moist convection**.

```python
#  Run out to equilibrium
rce.integrate_years(1.)
```

    Integrating for 365 steps, 365.2422 days, or 1.0 years.
    Total elapsed time is 1.9986737567564754 years.

```python
#  Check for energy balance
rce.ASR - rce.OLR
```

    array([ 0.00079773])

```python
result_list.append(rce.state)
name_list.append('Radiatve-Convective equilibrium (grey gas)')
```

```python
plot_soundings(result_list, name_list)
```

Introducing convective adjustment into the model cools the surface quite a bit (compared to Radiative Equilibrium, in green here) -- and warms the lower troposphere. It gives us a MUCH better fit to observations.

But of course we still have no stratosphere.

____________

## 8. Putting stratospheric ozone in the grey-gas model
____________

Our model has no equivalent of the stratosphere, where temperature increases with height. That's because our model has been completely transparent to shortwave radiation up until now.

We can load the observed ozone climatology from the input files for the CESM model:

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

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

```python
print( ozone)
```

    <xarray.Dataset>
    Dimensions:  (ilev: 27, lat: 96, lev: 26, lon: 144, slat: 95, slon: 144, time: 12)
    Coordinates:
      * ilev     (ilev) float64 2.194 4.895 9.882 18.05 29.84 44.62 61.61 78.51 ...
      * lat      (lat) float64 -90.0 -88.11 -86.21 -84.32 -82.42 -80.53 -78.63 ...
      * lev      (lev) float64 3.545 7.389 13.97 23.94 37.23 53.11 70.06 85.44 ...
      * lon      (lon) float64 0.0 2.5 5.0 7.5 10.0 12.5 15.0 17.5 20.0 22.5 ...
      * slat     (slat) float64 -89.05 -87.16 -85.26 -83.37 -81.47 -79.58 -77.68 ...
      * slon     (slon) float64 -1.25 1.25 3.75 6.25 8.75 11.25 13.75 16.25 ...
      * time     (time) datetime64[ns] 2000-01-15 2000-02-15 2000-03-15 ...
    Data variables:
        P0       float64 ...
        date     (time) int32 ...
        datesec  (time) int32 ...
        gw       (lat) float64 ...
        hyai     (ilev) float64 ...
        hyam     (lev) float64 ...
        hybi     (ilev) float64 ...
        hybm     (lev) float64 ...
        w_stag   (slat) float64 ...
        O3       (time, lev, lat, lon) float32 ...
        PS       (time, lat, lon) float32 ...
    Attributes:
        Conventions:                     CF-1.0
        source:                          CAM
        case:                            ar5_cam_1850-2000_03
        title:
        logname:                         lamar
        host:                            be0207en.ucar.ed
        Version:                         $Name$
        revision_Id:                     $Id$
        initial_file:                    /ptmp/lamar/cam-inputs/ar5_cam_1850_06.c...
        topography_file:                 /fs/cgd/csm/inputdata/atm/cam/topo/USGS-...
        history:                         Wed Jan  5 14:56:28 2011: ncks -d time,1...
        nco_openmp_thread_number:        1
        NCO:                             4.0.5
        DODS_EXTRA.Unlimited_Dimension:  time

The pressure levels in this dataset are:

```python
print( ozone.lev)
```

    <xarray.DataArray 'lev' (lev: 26)>
    array([   3.544638,    7.388814,   13.967214,   23.944625,   37.23029 ,
             53.114605,   70.05915 ,   85.439115,  100.514695,  118.250335,
            139.115395,  163.66207 ,  192.539935,  226.513265,  266.481155,
            313.501265,  368.81798 ,  433.895225,  510.455255,  600.5242  ,
            696.79629 ,  787.70206 ,  867.16076 ,  929.648875,  970.55483 ,
            992.5561  ])
    Coordinates:
      * lev      (lev) float64 3.545 7.389 13.97 23.94 37.23 53.11 70.06 85.44 ...
    Attributes:
        long_name:      hybrid level at midpoints (1000*(A+B))
        units:          level
        positive:       down
        standard_name:  atmosphere_hybrid_sigma_pressure_coordinate
        formula_terms:  a: hyam b: hybm p0: P0 ps: PS

### Take the global average of the ozone climatology, and plot it as a function of pressure (or height)

```python
#  Take global, annual average and convert to Kelvin
weight_ozone = cos(deg2rad(ozone.lat)) / cos(deg2rad(ozone.lat)).mean(dim='lat')
O3_global = (ozone.O3 * weight_ozone).mean(dim=('lat','lon','time'))
print( O3_global)
```

    <xarray.DataArray (lev: 26)>
    array([  7.827929e-06,   8.641505e-06,   7.589400e-06,   5.245671e-06,
             3.177616e-06,   1.823200e-06,   9.807570e-07,   6.228705e-07,
             4.476205e-07,   3.344812e-07,   2.625703e-07,   2.078981e-07,
             1.570746e-07,   1.124255e-07,   8.060050e-08,   6.278265e-08,
             5.429906e-08,   4.995061e-08,   4.600757e-08,   4.229778e-08,
             3.805591e-08,   3.387686e-08,   3.121716e-08,   2.978071e-08,
             2.879810e-08,   2.754299e-08])
    Coordinates:
      * lev      (lev) float64 3.545 7.389 13.97 23.94 37.23 53.11 70.06 85.44 ...

```python
ax = plt.figure(figsize=(10,8)).add_subplot(111)
ax.plot( O3_global * 1.E6, -np.log(ozone.lev/climlab.constants.ps) )
ax.set_xlabel('Ozone (ppm)', fontsize=16)
ax.set_ylabel('Pressure (hPa)', fontsize=16 )
yticks = np.array([1000., 750., 500., 250., 100., 50., 20., 10., 5.])
ax.set_yticks(-np.log(yticks/1000.))
ax.set_yticklabels(yticks)
ax.grid()
ax.set_title('Global, annual mean ozone concentration', fontsize = 24);
```

This shows that most of the ozone is indeed in the stratosphere, and peaks near the top of the stratosphere.

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.

```python
# the RadiativeConvectiveModel is pre-defined in climlab
#  It contains the same components are our previous model
#   But here we are specifying a different set of vertical levels.
oz_col = climlab.RadiativeConvectiveModel(lev = ozone.lev, adj_lapse_rate=6)
print( oz_col)
```

    climlab Process of type <class 'climlab.model.column.RadiativeConvectiveModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (26,)
    The subprocess tree:
    top: <class 'climlab.model.column.RadiativeConvectiveModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>
       convective adjustment: <class 'climlab.convection.convadj.ConvectiveAdjustment'>

Now we will do something new: let the column absorb some shortwave radiation. We will assume that the shortwave absorptivity is proportional to the ozone concentration we plotted above.

Now we need to weight the absorptivity by the pressure (mass) of each layer.

```python
# This number is an arbitrary parameter that scales how absorptive we are making the ozone
# in our grey gas model
ozonefactor = 75
dp = oz_col.Tatm.domain.lev.delta
epsSW = O3_global.values * dp * ozonefactor
```

We want to use the field `epsSW` as the absorptivity for our SW radiation model.

Let's see what the absorptivity is current set to:

```python
print( oz_col.subprocess['SW'].absorptivity)
```

    [ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.
      0.  0.  0.  0.  0.  0.  0.  0.]

It defaults to zero.

Before changing this (putting in the ozone), let's take a look at the shortwave absorption in the column:

```python
oz_col.compute_diagnostics()
```

```python
oz_col.diagnostics['SW_absorbed_atm']
```

    array([ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.])

Let's now put in the ozone:

```python
oz_col.subprocess['SW'].absorptivity = epsSW
print( oz_col.subprocess['SW'].absorptivity)
```

    [  3.20948549e-03   3.37750296e-03   4.71182551e-03   4.57614201e-03
       3.47591203e-03   2.24450924e-03   1.18884331e-03   7.11369788e-04
       5.50761612e-04   4.84170273e-04   4.47141487e-04   4.16507314e-04
       3.70212131e-04   3.11733073e-04   2.62922861e-04   2.40936639e-04
       2.45147939e-04   2.65307555e-04   2.87482272e-04   2.95567946e-04
       2.67120872e-04   2.16427978e-04   1.66169127e-04   1.15468088e-04
       6.79353134e-05   3.81013280e-05]

Let's check how this changes the SW absorption:

```python
oz_col.compute_diagnostics()
oz_col.SW_absorbed_atm
```

    array([ 1.40571621,  1.47671285,  2.05685574,  1.99236565,  1.50916587,
            0.97239284,  0.51429089,  0.30750441,  0.23797743,  0.20913838,
            0.19309134,  0.17981756,  0.1597929 ,  0.13452308,  0.11343944,
            0.10393806,  0.10574109,  0.11442201,  0.12396838,  0.12743553,
            0.11515129,  0.09328389,  0.07161233,  0.04975714,  0.02927232,
            0.01641659])

It is now non-zero, and largest near the top of the column (also top of the array) where the ozone concentration is highest.

Now it's time to run the model out to radiative-convective equilibrium

```python
oz_col.integrate_years(1.)
```

    Integrating for 365 steps, 365.2422 days, or 1.0 years.
    Total elapsed time is 0.9993368783782377 years.

```python
print( oz_col.ASR - oz_col.OLR)
```

    [-0.0040524]

And let's now see what we got!

```python
result_list.append(oz_col.state)
name_list.append('Radiative-Convective equilibrium with O3')
```

```python
#  Make a plot to compare observations, Radiative Equilibrium, Radiative-Convective Equilibrium, and RCE with ozone!
plot_soundings(result_list, name_list)
```

And we finally have something that looks looks like the tropopause, with temperature increasing above at approximately the correct rate.

There are still plenty of discrepancies between this model solution and the observations, including:

- Tropopause temperature is too warm, by about 15 degrees.
- Surface temperature is too cold

There are a number of parameters we might adjust if we wanted to improve the fit, including:

- Longwave absorptivity
- Surface albedo

Feel free to experiment! (That's what models are for, after all).

### The take home message

The dominant effect of stratospheric ozone is to vastly increase the radiative equilibrium temperature in the ozone layer. The temperature needs to be higher so that the longwave emission can balance the shortwave absorption.

Without ozone to absorb incoming solar radiation, the **temperature does not increase with height**.

This simple grey-gas model illustrates this principle very clearly.

[Back to ATM 623 notebook home](../index.ipynb)

____________
## Version information
____________

```python
%load_ext version_information
%version_information numpy, scipy, matplotlib, xarray, climlab
```

    Loading extensions from ~/.ipython/extensions is deprecated. We recommend managing extensions like any other Python packages, in site-packages.

Software	Version

Python	3.6.2 64bit [GCC 4.2.1 Compatible Clang 4.0.1 (tags/RELEASE_401/final)]

IPython	6.1.0

OS	Darwin 16.7.0 x86_64 i386 64bit

numpy	1.13.1

scipy	0.19.1

matplotlib	2.0.2

xarray	0.9.6

climlab	0.6.0.dev10

Wed Oct 11 13:35:42 2017 EDT

____________

## Credits

The author of this notebook is [Brian E. J. Rose](http://www.atmos.albany.edu/facstaff/brose/index.html), University at Albany.

It was developed in support of [ATM 623: Climate Modeling](http://www.atmos.albany.edu/facstaff/brose/classes/ATM623_Spring2015/), a graduate-level course in the [Department of Atmospheric and Envionmental Sciences](http://www.albany.edu/atmos/index.php)

Development of these notes and the [climlab software](https://github.com/brian-rose/climlab) is partially supported by the National Science Foundation under award AGS-1455071 to Brian Rose. Any opinions, findings, conclusions or recommendations expressed here are mine and do not necessarily reflect the views of the National Science Foundation.
____________

```python

```

Lectures/Lecture07 -- Grey radiation modeling with climlab.ipynb


## Neuromatch Academy 2020 -- RL Day

# Tutorial 1 - Learning to predict

Please execute the cell below to initialize the notebook environment

```python
#@ title Introduction
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="", width=854, height=480, fs=1)
print("Video available at https://youtu.be/" + video.id)
video
```

    Video available at https://youtu.be/

```python
import numpy as np
import matplotlib.pyplot as plt
import ipywidgets as widgets       # import widgets
```

```python
# @title Figure Settings
%matplotlib inline
fig_w, fig_h = (8, 8)
plt.rcParams.update({'figure.figsize': (fig_w, fig_h)})
%config InlineBackend.figure_format = 'retina'
```

```python
# @title Helper functions

# Parameters for simulations
reward_time = 10
reward_magnitude = 10

n_trials = 20000
n_steps = 40

def plot_value_function(V, ax=None, show=True):
  """Plot V(s), the value function"""
  if not ax:
    fig, ax = plt.subplots()

  ax.stem(V, use_line_collection=True)
  ax.set_ylabel('Value')
  ax.set_xlabel('State')
  ax.set_title("Value function: $V(s)$")

  if show:
    plt.show()

def plot_tde_trace(TDE, ax=None, show=True, skip=400):
  """Plot the TD Error across trials"""
  if not ax:
    fig, ax = plt.subplots()

  indx = np.arange(0, TDE.shape[1], skip)
  im = ax.imshow(TDE[:,indx])
  ax.set_title('TD-error over learning')
  ax.set_ylabel('State')
  ax.set_xlabel('Iterations')
  ax.set_xticklabels([f"{int(skip * x)}" for x in ax.get_xticks()])
  ax.figure.colorbar(im)
  if show:
    plt.show()

def learning_summary_plot(V, TDE):
  """Summary plot for Ex1"""
  fig, (ax1, ax2) = plt.subplots(nrows = 2, gridspec_kw={'height_ratios': [1, 2]})

  plot_value_function(V, ax=ax1, show=False)
  plot_tde_trace(TDE, ax=ax2, show=False)
  plt.tight_layout()

def reward_guesser_title_hint(r1, r2):
  """"Provide a mildly obfuscated hint for a demo."""
  if (r1==14 and r2==6) or (r1==6 and r2==14):
    return "Technically correct...."

  if not (((r1+r2) & 5) & 7) - 4: # Don't spoil the fun :-)
    return "Congratulations! You solved it!"

  return "Keep trying...."

#@title Default title text
class ClassicalConditioning:

    def __init__(self, n_steps, reward_magnitude, reward_time):

        # Task variables
        self.n_steps = n_steps
        self.n_actions = 0
        self.cs_time = int(n_steps/4) - 1

        # Reward variables
        self.reward_state = [0,0]
        self.reward_magnitude = None
        self.reward_probability = None
        self.reward_time = None

        self.set_reward(reward_magnitude, reward_time)

        # Time step at which the conditioned stimulus is presented

        # Create a state dictionary
        self._create_state_dictionary()

    def set_reward(self, reward_magnitude, reward_time):

        """
        Determine reward state and magnitude of reward
        """
        if reward_time >= self.n_steps - self.cs_time:
            self.reward_magnitude = 0

        else:
            self.reward_magnitude = reward_magnitude
            self.reward_state = [1, reward_time]

    def get_outcome(self, current_state):

        """
        Determine next state and reward
        """
        # Update state
        if current_state < self.n_steps - 1:
            next_state = current_state + 1
        else:
            next_state = 0

        # Check for reward
        if self.reward_state == self.state_dict[current_state]:
            reward = self.reward_magnitude
        else:
            reward = 0

        return next_state, reward

    def _create_state_dictionary(self):

        """
        This dictionary maps number of time steps/ state identities
        in each episode to some useful state attributes:

        state      - 0 1 2 3 4 5 (cs) 6 7 8 9 10 11 12 ...
        is_delay   - 0 0 0 0 0 0 (cs) 1 1 1 1  1  1  1 ...
        t_in_delay - 0 0 0 0 0 0 (cs) 1 2 3 4  5  6  7 ...
        """
        d = 0

        self.state_dict = {}
        for s in range(self.n_steps):
            if s <= self.cs_time:
                self.state_dict[s] = [0,0]
            else:
                d += 1 # Time in delay
                self.state_dict[s] = [1,d]

class MultiRewardCC(ClassicalConditioning):
  """Classical conditioning paradigm, except that one randomly selected reward,
    magnitude, from a list, is delivered of a single fixed reward."""
  def __init__(self, n_steps, reward_magnitudes, reward_time=None):
    """"Build a multi-reward classical conditioning environment
      Args:
        - nsteps: Maximum number of steps
        - reward_magnitudes: LIST of possible reward magnitudes.
        - reward_time: Single fixed reward time
      Uses numpy global random state.
      """
    super().__init__(n_steps, 1, reward_time)
    self.reward_magnitudes = reward_magnitudes

  def get_outcome(self, current_state):
    next_state, reward = super().get_outcome(current_state)
    if reward:
      reward=np.random.choice(self.reward_magnitudes)
    return next_state, reward

class ProbabilisticCC(ClassicalConditioning):
  """Classical conditioning paradigm, except that rewards are stochastically omitted."""
  def __init__(self, n_steps, reward_magnitude, reward_time=None, p_reward=0.75):
    """"Build a multi-reward classical conditioning environment
      Args:
        - nsteps: Maximum number of steps
        - reward_magnitudes: Reward magnitudes.
        - reward_time: Single fixed reward time.
        - p_reward: probability that reward is actually delivered in rewarding state
      Uses numpy global random state.
      """
    super().__init__(n_steps, reward_magnitude, reward_time)
    self.p_reward = p_reward

  def get_outcome(self, current_state):
    next_state, reward = super().get_outcome(current_state)
    if reward:
      reward*= int(np.random.uniform(size=1)[0] < self.p_reward)
    return next_state, reward

```

---

## Tutorial objectives

In this tutorial, we will learn how to estimate state-value functions in a classical conditioning paradigm using Temporal Difference (TD) learning and examine TD-errors at the presentation of the conditioned and unconditioned stimulus (CS and US) under different CS-US contingencies. These exercises will provide you with an understanding of both how RPEs behave in classical conditioning and what we should expect if Dopamine represents a "canonical" model-free RPE.

* You will learn to use the standard tapped delay line conditioning model
* You will understand how RPEs move to CS
* You will understand how variability in reward size effects RPEs
* You will understand how differences in US-CS timing effect RPEs
* Advanced exercise
---

__Environment:__

- The agent experiences the environment in episodes or trials.
- Episodes terminate by transitioning to the inter-trial-interval (ITI) state and they are initiated from the ITI state as well. We clamp the value of the terminal/ITI states to zero.
- The classical conditioning environment is composed of a sequence of states that the agent deterministically transitions through. Starting at State 0, the agent moves to State 1 in the first step, from State 1 to Start 2 in the second, and so on.  These states represent time in the tapped delay line representation
- Within each episode, the agent is presented a CS and US (reward).
- For each exercise, we use a different CS-US contingency.
- The agent's goal is to learn to predict expected rewards from each state in the trial.

__Definitions:__

1. Returns:
\begin{align}
G_{t} = r_{t+1} + \gamma r_{t+2} + \gamma^2 r_{t+3} + ... = \sum \limits_{k = 1}^{\infty} \gamma^{k-1} r_{t+k}
\end{align}

2. Value:
\begin{align}
V(s_{t}) = \mathbb{E} [ G_{t} | s_{t}] = \mathbb{E} [r_{t+1} + \gamma V_{t+1} | s_{t}]
\end{align}

3. TD-error:
\begin{align}
\delta_{t} = r_{t+1} + \gamma V(s_{t+1}) - V(s_{t})
\end{align}

4. Value updates:
\begin{align}
V(s_{t}) \leftarrow V(s_{t}) + \alpha \delta_{t}
\end{align}

**Run the following code for your implementation:**

---
# EXERCISE 1: TD-learning with guaranteed rewards.

Implement TD-learning to estimate the state-value function in the classical-conditioning world with guaranteed rewards, with a fixed magnitude, at a fixed delay after the CS. Save TD-errors over learning so we can visualize them--you're going to need to compute them anyway.

Use the provided code to estimate the value function.

```python
env = ClassicalConditioning(n_steps, reward_magnitude=reward_magnitude, reward_time=reward_time) # Initialize the environment with n_steps; CS is presented ~1/4 along the way (e.g. for n_steps=10, CS is presented at t=9)

def td_learner(env, n_trials, gamma=0.98, alpha=0.001):
  V = np.zeros(env.n_steps) # Array to store values over states (time)
  TDE = np.zeros((env.n_steps, n_trials)) # Array to store TD errors

  for n in range(n_trials):

    state = 0 # Initial state
    for t in range(env.n_steps):

        next_state, reward = env.get_outcome(state) # Get next state and next reward
        is_delay  = env.state_dict[state][0] # Is the current state in the delay period (after CS)?
        ########################################################################
        ## Insert code to calculate the TD-error
        # Hint: this is the reward prediction *error*. When would we expect
        #  Comment out the last line in this block when you're done.
        raise NotImplementedError("Student excercise: need to implement RPE")
        #################################################################################

        ########################################################################
        ## Insert code to update the Value Function V(st) below.
        #  Comment out the last line in this block when you're done.
        #  Hint: use the TD error you calculated above.
        raise NotImplementedError("Student excercise: need to implement Value update")
        #################################################################################
        state = next_state # Update state

  return V, TDE

#V, TDE = td_learning_guaranteed_reward()
#learning_summary_plot(V, TDE)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D5_ReinforcementLearning/solutions/W2D5_Tutorial1_Solution_55115845.py)

*Example output:*

Interactive Demo 1: US to CS Transfer
---
During classical conditioning, the subject's behavioral response (e.g., salivating) transfer from the unconditioned stimulus (US) (e.g., the smell of tasty food) to the conditioned stimulus (CS) that predicts it. Reward prediction errors play an important role in this process by adjusting the value of states according to their expected, discounted return.

Use the widget below to examine how reward prediction errors change over time. Before training (red line), the only the reward state has high reward prediction error. As training progresses (blue, slider), the reward prediction errors shift to the conditioned stimulus, where they end up when the trial is complete (green).

Dopamine neurons, which are thought to carry reward prediction errors _in vivo_, show exactly the same behavior!

```python
#@title
@widgets.interact
def plot_tde_by_trial(trial = widgets.IntSlider(value=5000, min=0, max=n_trials-1 , step=1, description="Trial #")):
  if 'TDE' not in globals():
    print("Complete Exercise 1 to enable this interactive demo!")
  else:

    fig, ax = plt.subplots()
    ax.axhline(0, color='k') # Use this + basefmt=' ' to keep the legend clean.
    ax.stem(TDE[:, 0], linefmt='r-', markerfmt='rd', basefmt=' ',
            label="Before Learning (Trial 0)",
            use_line_collection=True)
    ax.stem(TDE[:, -1], linefmt='g-', markerfmt='gs', basefmt=' ',
            label="After Learning (Trial $\infty$)",
            use_line_collection=True)
    ax.stem(TDE[:, trial], linefmt='b-', markerfmt='bo', basefmt=' ',
            label=f"Trial {trial}",
            use_line_collection=True)

    ax.set_xlabel("State in trial")
    ax.set_ylabel("TD Error")
    ax.set_title("Temporal Difference Error by Trial")
    ax.legend()

    plt.show()
```

Interactive Demo 2 : Learning Rates and Discount Factors
---

Our TD-learning agent has two parameter that control how it learns: $\alpha$, the learning rate, and $\gamma$, the discount factor. In Exercise 1, we set these parameters to $\alpha=0.001$ and $\gamma=0.98$ for you. Here, you'll investigate how changing these parameters alters the model that TD-learning learns.

Before enabling the interactive demo below, take a moment to think about the functions of these two parameters. $\alpha$ controls the size of the Value function updates produced by each TD-error. In our simple, deterministic world, will this affect the final model we learn? Is a larger $\alpha$ necessarily better in more complex, realistic environments?

The discount rate $\gamma$ applies an exponentially-decaying weight to returns occuring in the future, rather than the present timestep. How does this affect the model we learn? What happens when $\gamma=0$ or $\gamma \geq 1$?

```python
#@title

@widgets.interact
def plot_summary_alpha_gamma(alpha = widgets.FloatSlider(value=0.0001, min=0.001, max=0.1, step=0.0001, description="alpha"),
                             gamma = widgets.FloatSlider(value=0.980, min=0, max=1.1, step=0.010, description="gamma")):
  env = ClassicalConditioning(n_steps, reward_magnitude, reward_time)
  try:
    V_params, TDE_params = td_learner(env, n_trials, gamma=gamma, alpha=alpha)
  except NotImplementedError:
    print("Finish Exercise 1 to enable this interactive demo")

  learning_summary_plot(V_params,TDE_params)

```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D5_ReinforcementLearning/solutions/W2D5_Tutorial1_Solution_a94dedaf.py)

# EXERCISE 2: TD-learning with varying reward magnitudes

In the previous exercise, the environment was as simple as possible. On every trial, the CS predicted the same reward, at the same time, with 100% certainty. In the next few exercises, we will make the environment more progressively more complicated and examine the TD-learner's behavior.

First, will replace the environment with one that dispenses multiple rewards. Shown below is the final value function $V$ for a TD learner that was trained in an enviroment where the CS predicted rewards of 6 or 14 units (both equally likely). Can you find another pair of rewards, both equally likely, that exactly match this value function?

Hints:
* Carefully consider the definition of the value function $V$. This can be solved analytically.
* There is no need to change $\alpha$ or $\gamma$.
* Due to the randomness, there may be a small amount of variation.

```python
#@title Interactive Demo 3: Match the Value Functions

rng_state = np.random.get_state()
env = MultiRewardCC(40, [6, 14], reward_time)
V_multi, TDE_multi = td_learner(env, n_trials, gamma=0.98, alpha=0.001)

@widgets.interact
def reward_guesser_interaction(r1 = widgets.IntText(value=0,  description="Reward 1"),
                               r2 = widgets.IntText(value=0, description="Reward 2")):
  try:
    env2 = MultiRewardCC(40, [r1, r2], reward_time)
    V_guess, _ = td_learner(env2, n_trials, gamma=0.98, alpha=0.001)
    fig, ax = plt.subplots()
    m, l, _ = ax.stem(V_multi, linefmt='y-', markerfmt='yo', basefmt=' ', label="Target",
            use_line_collection=True)
    m.set_markersize(15)
    m.set_markerfacecolor('none')
    l.set_linewidth(4)
    m, _, _ = ax.stem(V_guess, linefmt='r', markerfmt='rx', basefmt=' ', label="Guess",
                      use_line_collection=True)
    m.set_markersize(15)

    ax.set_xlabel("State")
    ax.set_ylabel("Value")
    ax.set_title("Guess V(s)\n" + reward_guesser_title_hint(r1, r2))
    ax.legend()
  except NotImplementedError:
    print("Please finish Exercise 1 first!")
```

__

Examining the TD Error
---
Run the cell below to plot the TD errors from our multi-reward environment. A new feature appears in this plot? What is it? Why does it happen?

```python
plot_tde_trace(TDE_multi)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D5_ReinforcementLearning/solutions/W2D5_Tutorial1_Solution_2efbc169.py)

---
# EXERCISE 3: TD-learning with probabilistic rewards

In this environment, we'll return to delivering a single reward of ten units. However, it will be delivered intermittently: on 20 percent of trials, the CS will be shown but the agent will not receive the usual reward; the remaining 80% will proceed as usual.

 Run the cell below to simulate. How does this compare with the previous experiment?

Earlier in the notebook, we saw that changing $\alpha$ had little effect on learning in a deterministic environment. What happens if you set it to an large value, like 1, in this noisier scenario? Does it seem like it will _ever_ converge?

```python
np.random.set_state(rng_state) # Resynchronize everyone's notebooks

try:
  env = ProbabilisticCC(n_steps, reward_magnitude=10, reward_time=10, p_reward=0.8)
  V_stochastic, TDE_stochastic = td_learner(env, n_trials*2, alpha=1)
  learning_summary_plot(V_stochastic,TDE_stochastic)
except NotImplementedError:
  print("Please finish Exercise 1 first")
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D5_ReinforcementLearning/solutions/W2D5_Tutorial1_Solution_251e1fdb.py)

# ADVANCED Exercise: TD-reset

In this exercise we will implement a commonly used heuristic used in modeling activity of dopamine neurons, TD-reset.
Implement TD-learning as in previous exercises, but set TD-error to zero on all steps after reward (US).

1. Plot value function and TD-errors.
2. Can you explain how the reset is changing the TD-errors and value function?

```python
env = ProbabilisticCC(n_steps, reward_magnitude=10, reward_time=10, p_reward=0.8)
def td_reset_learner(env, n_trials, alpha=0.001, gamma=0.98):
  V = np.zeros(env.n_steps)
  TDE_reset = np.zeros((env.n_steps, n_trials))

  for n in range(n_trials):
    state = 0
    reset = False

    for t in range(env.n_steps):

        next_state, reward = env.get_outcome(state)
        is_delay  = env.state_dict[state][0]

        ########################################################################
        # Insert your code here to convert your TD learner into TD-reset
        # Hint: you can reuse much of your previous work.
        #TDE_reset[state] =....
        #
        #
        #
        #V[state] += ....
        # Comment out the line below when you are done
        raise NotImplementedError("Student excercise: need to implement TD-reset update")
        ########################################################################

        state = next_state

  return V, TDE

# Uncomment these two lines to visualize your results
#V_reset, TDE_reset = td_reset_learner(env, n_trials)
#learning_summary_plot(V_reset, TDE_reset)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D5_ReinforcementLearning/solutions/W2D5_Tutorial1_Solution_584e22ec.py)

*Example output:*

# Advanced: Removing the CS
---
In Exercise 1, you (should have) included a term that depends on the conditioned stimulus. Remove it and see what happens. Do you understand why?
This phenomena often fools people attempting to train animals--beware!

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D5_ReinforcementLearning/solutions/W2D5_Tutorial1_Solution_e34e8483.py)

Conclusion
___
In this notebook, we have developed a simple TD Learner and examined how its state representations and reward prediction errors evolve during training. By manipualting its environment and parameters ($\alpha$, $\gamma$), you developed an intuition for how it behaves.

This simple model closely resembles the behavior of subjects undergoing classical conditioning tasks and the dopamine neurons that may underlie that behavior. You may have implemented TD-reset or used the model to recreate a common experimental error. The update rule used here has been extensively studied for [more than 70 years](https://www.pnas.org/content/108/Supplement_3/15647) as a possible explanation for artificial and biological learning.

However, you may have noticed that something is missing from this notebook. We carefully calculated the value of each state, but did not use it to actually do anything. Using values to plan _**Actions**_ is coming up next!


# The ARMI Material Library

While *nuclides* are the microscopic building blocks of nature, their collection into *materials* is what we interact with at the engineering scale. The ARMI Framework provides a `Material` class, which has a composition (how many of each nuclide are in the material), and a variety of thermomechanical properties (many of which are temperature dependent), such as:

* Mass density
* Heat capacity
* Linear or volumetric thermal expansion
* Thermal conductivity
* Solidus/liquidus temperature

and so on.

Many of these properties are widely available in the literature for fresh materials. As materials are irradiated, the properties tend to change in complex ways. Material objects can be extended to account for such changes.

The ARMI Framework comes with a small set of example material definitions. These are generally quite incomplete (often missing temperature dependence), and are of academic quality at best. To do engineering design calculations, users of ARMI are expected to make or otherwise prepare materials. As the ecosystem grows, we hope the material library will mature.

In any case, here we will explore the use of `Material`s. Let's get an instance of the Uranium Oxide material.

```python
from armi.materials import uraniumOxide
uo2 = uraniumOxide.UO2()
density500 = uo2.density(Tc=500)
print(f"The density of UO2 @ T = 500C is {density500:.2f} g/cc")
```

Taking a look at the composition

```python
print(uo2.p.massFrac)
```

The mass fractions of a material, plus its mass density, fully define the composition. Conversions between number density/fraction and mass density/fraction are handled on the next level up (on `Component`s), which we will explore soon.

ARMI automatically thermally-expands materials based on their coefficients of linear expansion. For instance, a piece of Uranium Oxide that's 10 cm at room temperature would be longer at 500 C according to the formula:

\begin{equation}
\frac{\Delta L}{L_0} = \alpha \Delta T
\end{equation}

On the reactor model, this all happens behind the scenes. But here at the material library level, we can see it in detail.

```python
L0 = 10.0
dLL = uo2.linearExpansionFactor(500,25)
L = L0 * (1+dLL)
print(f"Hot length is {L:.4f} cm")

```

Let's plot the heat capacity as a function of temperature in K.

```python
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
Tk = np.linspace(300,2000)
heatCapacity = [uo2.heatCapacity(Tk=ti) for ti in Tk]
plt.plot(Tk, heatCapacity)
plt.title("$UO_2$ heat capacity vs. temperature")
plt.xlabel("Temperature (K)")
plt.ylabel("Heat capacity (J/kg-K)")
plt.grid(ls='--',alpha=0.3)
```

Different physics plugins require different properties to be defined. For pure neutronics runs, mass density and composition is enough. But for thermal/hydraulics runs, heat capacity and thermal conductivity is needed for solids, and more is needed for coolants. As irradiation models are investigated, creep, corrosion, porosity, swelling, and other factors will be necessary.


## Introduction

This notebook will go over optimization of multiple response variables using the MRUplift Framework. Readers are encouraged to read the [example on single responses](https://github.com/Ibotta/ibotta_uplift/blob/master/examples/mr_uplift_one_response_example.ipynb) if they'd like a general introduction to uplift models. This will go over:

1. Hypothetical Problem and The Data Generating Process

2. Building / Gridsearching the Uplift Model

3. Evaluating the Tradeoffs with the out-of-sample ERUPT metric

4. Assign Optimal Treatments for New Observations

```python
import numpy as np
import pandas as pd

from mr_uplift.dataset.data_simulation import get_simple_uplift_data
from mr_uplift.mr_uplift import MRUplift
from ggplot import *
```

    /Users/samweiss/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:34: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
      from ._conv import register_converters as _register_converters
    Using TensorFlow backend.
    /Users/samweiss/anaconda3/lib/python3.6/site-packages/ggplot/utils.py:81: FutureWarning: pandas.tslib is deprecated and will be removed in a future version.
    You can access Timestamp as pandas.Timestamp
      pd.tslib.Timestamp,
    /Users/samweiss/anaconda3/lib/python3.6/site-packages/ggplot/stats/smoothers.py:4: FutureWarning: The pandas.lib module is deprecated and will be removed in a future version. These are private functions and can be accessed from pandas._libs.lib instead
      from pandas.lib import Timestamp
    /Users/samweiss/anaconda3/lib/python3.6/site-packages/statsmodels/compat/pandas.py:56: FutureWarning: The pandas.core.datetools module is deprecated and will be removed in a future version. Please use the pandas.tseries module instead.
      from pandas.core import datetools

### Business Problem

Imagine we are data scientists working for a startup that would like to grow the business more but in an efficient way.

As a tactic to increase user activity the company gives all users a potentially expensive treatment. In order to reduce costs we were assigned the task of using data to find a subset of users that should continue receiving the costly treatment while continue increasing growth as measured by revenue.

The company would rather have and increase in revenue even if it decreases overall profitiability. However, we do not know the relative tradeoff of growth in revenue vs profitibility the business would like to take.

If we knew the relative tradeoffs we could create a new variable $y_{new} = (1-a)*revenue - a*costs$ and maximize $y_{new}$ as a normal uplift problem (see example 1). Since we don't know the value of $a$ we will have to calculate the tradeoffs first and then decide the optimal value.

This problem therefore entails both finding users and evaluating tradeoffs among different response variables.

We are given some explanatory variables for users $x$, a random treatment of whether a users recieved marketing campaign or not $t$, and response variable of profitibility $y$.

We can use MRUplift to find the tradeoff between increases in revenue vs the increase in costs. Once we have found a suitable tradeoff we can then use the model to target users that should recieve the treatment given our new objectives.

### Data Generating Process

Below is the data generating process of the data we have. (I have included a noise response variable to demonstrate what happens when we try to model noise.)

\begin{equation}
x_1  \sim runif(0,1)
\end{equation}

\begin{equation}
x_2 \sim runif(0,1)
\end{equation}

\begin{equation}
e_1 \sim rnorm(0,1)
\end{equation}

\begin{equation}
e_2 \sim rnorm(0,1)
\end{equation}

\begin{equation}
t \sim rbinom(.5)
\end{equation}

\begin{equation}
noise \sim rnorm(0,1)
\end{equation}

\begin{equation}
revenue = x_1*t + e_1
\end{equation}

\begin{equation}
costs = x_2*t + e_2
\end{equation}

\begin{equation}
profit = revenue - costs
\end{equation}

```python
y, x, t = get_simple_uplift_data(10000)

y = pd.DataFrame(y)
y.columns = ['revenue','cost', 'noise']
y['profit'] = y['revenue'] - y['cost']
```

### Model Building / Gridsearch
After instantiating the MRUplift class the `.fit` function will build the model. It first seperates the data into a train / test split. It builds standard scaler transformerd on all variables $x, y, t$.

Then it builds and runs grisdesarch using keras neural network model that minimizes the mean squared error of the form $y = f(t,x)$. The user can input a custom parameter grid.

```python
uplift_model = MRUplift()
param_grid = dict(num_nodes=[8], dropout=[.1,.5], activation=[
                          'relu'], num_layers=[1,2], epochs=[25], batch_size=[30])

uplift_model.fit(x, y, t.reshape(-1,1), param_grid = param_grid, n_jobs = 1)

```

    /Users/samweiss/anaconda3/lib/python3.6/site-packages/sklearn/utils/validation.py:475: DataConversionWarning: Data with input dtype int64 was converted to float64 by StandardScaler.
      warnings.warn(msg, DataConversionWarning)

### Evaluating Tradeoffs using out-of-sample ERUPT Curves

Using the test dataset MRUplift will then evaluate the model using the ERUPT metrics with different objective weights with the `get_erupt_curves` function.

Continuing the notation from the [single response example](https://github.com/Ibotta/ibotta_uplift/blob/master/examples/ibotta_uplift_multiple_response_example.ipynb) the policy assignment is expanded to include weights $W$ below:

\begin{equation}
    \pi(x_i, W) =argmax \:_{t \in T} \sum_j w_j *E[y_{j,i} | X=x_i, T=t]
\end{equation}

Each objective weights will consist of an array equal to the number of responses. With these weights it calculates a new variable based on predictions to maximize $y_{new} = weight_0*\hat{y_0} + weight_1*\hat{y_1} + ... + weight_{num responses}*\hat{y_{num responses}}$.

Varying the objective weights allows us to see what happens when we change the focus of model from a revenue maximizing problem to a cost minimizing problem.

The `get_erupt_curves` function outputs two dataframes:

1) The first is the ERUPT metric and standard deviation for all response variables given a set of model weights. In this example the weights are set up to go from a cost minimizing to a revenue maximimization objective. We see both costs and revenue go up but at different rates. Profit is maximized when the revenue weight is .5 and cost weight is .5. Finally we can also see the noise response variable is not effected as expected.

2) The second pandas df shows the distribution of treatments under the optimal assignment for a set of objective weights. In this example we can see those assigned the treatment increase as we weight the objective function towards to be more revenue maximization.

```python
erupt_curves, dists = uplift_model.get_erupt_curves()

#I create a new variable that shows the weight of the revenue variable
dists['weights_1'] = [np.float(x.split(',')[0]) for x in dists['weights']]
erupt_curves['weights_1'] = [np.float(x.split(',')[0]) for x in erupt_curves['weights']]

ggplot(aes(x='weights_1', y='mean', group = 'assignment', colour = 'assignment'), data=erupt_curves) +\
    geom_line()+\
    geom_point()+facet_wrap("response_var_names")
```

```python
ggplot(aes(x='weights_1', y='num_observations'), data=dists) +\
    geom_line()+\
    geom_point()+facet_wrap('tmt')
```

### Assigning Optimal Treatments for new observations
After building and evaluating an uplift model the modeller may deem it worthy of production. To assign new users the optimal treatment one can use the `predict_optimal_treatments` function to do so below but now requires a weights assignment.
I have chosen a weights assignment of .6 for revenue and -.4 for costs to be more alligned with the business objective.

```python
#generate 5 new observation
_, x_new ,_  = get_simple_uplift_data(5)
uplift_model.predict_optimal_treatments(x_new, weights = np.array([.6,-.4,0,0]).reshape(1,-1))
```

    /Users/samweiss/anaconda3/lib/python3.6/site-packages/sklearn/utils/validation.py:475: DataConversionWarning: Data with input dtype int64 was converted to float64 by StandardScaler.
      warnings.warn(msg, DataConversionWarning)

    array([[1],
           [1],
           [1],
           [1],
           [0]])


# Tutorial 1: Reverse non-equilibrium simulation of a LJ fluid

In this tutorial, we will use the azplugins implementation of the [reverse non-equilibrium MD method](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.59.4894) by Müller-Plathe to generate shear flow in a simple Lennard-Jones (LJ) fluid. It is recommend to read this paper before you continue with the tutorial to familarize yourself with the method. The main benefit of this method over other non-equilibrium techniques to generate shear is that it is relatively straightforward to compute the viscosity of the fluid. The basic idea is to exchange momentum between two areas of the simulation box in an unphysical way by swapping momentum of particles. Then, the physical response of the system is recorded, the velocity gradient generated by of the physical momentum flux. Commonly, non-equilibrium simulations are done the other way around: a velocity gradient is generated (by e.g. Lees-Edwards boundary conditions or SLLOD) and the stress/momentum is measured. Often, it is difficult to measure stress in the simulation, but it generally very easy to measure a velocity gradient.

The simulation box setup in azplugins is shown below. Hoomd-blue convention is that the box dimensions go from $-L/2$ to $L/2$ in each dimension, such that the origin is in the center of the simulation box. For the reverse perturbation plugin, the $x$ direction is the flow direction and the $z$ is the gradient direction. The locations of the two layers where the momentum exchanges are performed are at $\pm L_z/4$. Each layer has a thickness of `width=`$w$. I recommend using cubic boxes or boxes elongated in the gradient direction only. For details please see this [publication](https://arxiv.org/pdf/1811.04097.pdf) and references therein, as well as the follow up [publication](https://arxiv.org/pdf/1907.07086.pdf).

A triangular shaped flow is established from swapping momentum between particles in the two exchange regions. Because of periodic boundary conditions, the box contains two linear shear regions between the two exchange regions.

To decide which particle pairs momentum to swap, the original publication uses the particles in the exchange region which have the largest momentum in the opposing direction of the flow. For more control over the flow, the plugin follows this [publication](https://doi.org/10.1063/1.3276454) and allows you to specify a target momentum with the variable ``target_momentum`` instead of picking the largest opposing pair. This way, the algorithm uses the pair with momentum closest to $\pm$``target_momentum``, resulting in a more consistent, predictable momentum swap each time.

### System setup

First, we need to set up a simple equlibrated LJ fluid simulation. If you are not familiar with this, check out the [hoomd-blue tutorial](https://github.com/joaander/hoomd-examples/blob/master/Tutorial%20-%20MD%20-%20Lennard%20Jones.ipynb). We saved the equilibrated LJ fluid configuration in a gsd snapshot. For this tutorial, we apply the reverse perturbation method to a standard Lennard-Jones system at a reduced temperature of $k_\text{B}T=0.722$ and number density $\rho =0.848$ in a cubic box with $L = 10$. These values are close to the triple point of the LJ fluid and are the same as used in the original publication.

In principle, the reverse perturbation algorithm can be applied to many different systems. We have found it to work best with fairly homogeneous fluid-like systems, strongly ordered or gel-like networks might not work as nicely. It is strongly suggested to do extensive testing for your system of choice if you wish to apply the reverse perturbation method to a new system. Start with a pure solvent, confirming known fluid properties and check the flow profile in all dimensions extensively.

The reverse perturbation method conserves both momentum and kinetic energy and should not need a thermostat. In fact, it does not work well if you apply a global thermostat, because the thermostat will not take the gradient into account when calculating the temperature. If your system needs thermostatting, you may want to either use a local cell based thermostat or thermostat only the non-swapped dimensions.

The code below imports all needed modules,  initializes a hoomd simulation context and loads the initial snapshot from a gsd file.

```Python
import numpy as np
import sys
import hoomd
import azplugins
from hoomd import md
from hoomd import data

L = 10

hoomd.context.initialize()
hoomd.context.SimulationContext()

system = hoomd.init.read_gsd(filename='tutorial_reverse_perturbation_1_LJ_init.gsd')

```
Now, we just need to set the correct interactions for this system. In this case, we are using the standard LJ pair potential with a cutoff of `r_cut=3.0`, $\epsilon=$ `epsilon=1.0`, and $\sigma=$`sigma=1.0`.  For convinience, we are also defining a particle group ``all`` containing all particles in the simulation.
```Python
nl = hoomd.md.nlist.cell()
lj = hoomd.md.pair.lj(r_cut=3.0, nlist=nl)
lj.pair_coeff.set('A', 'A', epsilon=1.0, sigma=1.0)

all = hoomd.group.all()
```

### Measuring quantities during the non-equilibrium simulation

For measuring the flow, we will need to at least measure the $x$ velocity component in $z$ direction, $v_x(z)$. For this purpose, ``azplugins.flow.FlowProfiler()`` can be used to record velocity histograms. In a real simulation, you should  record all quantities of interest, including density, temperature and other velocity components in all spatial directions. You can use the ``hoomd.analyze.callback()`` with a custom class to do so. We have observed all kinds of gradients, including unwanted ones when using this method, you are strongly encuraged you to thorougly test this for your system/solvent of choice.

```Python
measure_vel = azplugins.flow.FlowProfiler(system=system, axis=2, bins=100, range=(-L/2,L/2), area=L**2)
analyze = hoomd.analyze.callback(measure_vel, period=1e2)
```
Here, ``axis`` is direction for binning (0=`x`, 1=`y`, 2=`z`).

The next part is setting up all properties we want to record during the shear flow simulation, we use hoomds ``hoomd.analyze.log`` to record the momentum exchanged during the simulation. The reverse perturbation azplugin registers a variable called ``rp_momentum`` to hoomd. It contains the value of the exchanged momentum for each execution of the timestep before. So, value reported on timestep 0 will always be zero, the value at timestep 1 will be the momentum exchanged on timestep 0, and so on. This is because of the order how things are executed within hoomd. The variable only records the value of the current execution. Meaning, if you want to know the exact amount of momentum exchanged during the simulation, you have to call
 ``hoomd.analyze.log`` with the same period as the swapping period in ``azplugins.flow.reverse_perturbation``. Ideally, the simulation reaches a steady state and the average over a long run can be used to calculate the momentum exchange. If the simulation does not reach a steady state, calculating the viscosity is not possible in a reliable way.

```Python
# log the exchanged momentum during the simulation
log = hoomd.analyze.log(filename="tutorial_reverse_perturbation_1_LJ.log",
                        quantities=['rp_momentum'],
                        period=1e2,overwrite=True)

# write out a trajectory during the simulation
hoomd.dump.gsd(filename="tutorial_reverse_perturbation_1_LJ_trajectory.gsd",
               overwrite=True, period=1e2, group=all)
```

### Run the non-equilibrium simulation

The last task is to disable the previous thermostat set up the NVE integrator. Then we can setup and enable the reverse perturbation plugin ``azplugins.flow.reverse_perturbation`` and  let the simulation run for another 500,000 timesteps. Read about the ``azplugins.flow.reverse_perturbation`` parameters in the next section.

```Python
# define timestep and integrator
hoomd.md.integrate.mode_standard(dt=0.005)

nve = hoomd.md.integrate.nve(group=all)

# enable the reverse perturbation and run the simulation
f = azplugins.flow.reverse_perturbation(group=all,width=1.0,Nswap=10,period=1,target_momentum=0.5)
hoomd.run(5e5)
np.savetxt('tutorial_reverse_perturbation_1_LJ_vx.hist', np.column_stack((measure_vel.centers, measure_vel.number_velocity[:,0])))
```

### azplugins.flow.reverse_perturbation parameters

The plugin has a `group`, enabling you to choose which particles you would like to apply this plugin to. For example if you have solutes in solvent, you can only swap momentum of the solvent particles and leave the solutes untouched. Simply create a "solvents" group and assign the plugin to only operate on it and not all particles.

The `width` parameter sets the width of the exchange regions. What you should choose here depends on the density and system size. Because the swapping only happens between pairs in the exchange regions, they need to be big enough to contain some particles to pick from. So, the actual thing you care about is the average number of particles in the regions. The velocity profile will be deviating from the linear behavior in the excange regions and will not result in two perfect triangles, in case of the LJ fluid it results in rounded-off tips.
So, as a result aim for as small as possible and as big as needed. The easiest way to test it is to do a quick check of the average number of particles in the exchange regions and run various small/short simulations, recording the exchanged momentum for every swap. See if it is constant and doesn't have zero recordings in between, because zero exchanged momentum usually means that there was no pair at all to exchange momentum between, so `width` was too small. The width also limits the maximum achivable shear to a certain degree, again, if no pairs are left to swap, no momentum can be exchanged.

The following parameters determine the value of the momentum exchanged and therefore the shear flow which is generated:

- The parameter `Nswap` sets how many pairs are supposed to be swapped per execution. It gives the maximum, not the minimum. If the algorithm doesn't find suitable pairs, nothing will be swapped.

- The next parameter, `period` just gives the period of execution of the plugin.  It has been noted in the [literature](https://doi.org/10.1063/1.3276454) , that it is better to try to achive a more constant momentum exchange for a more steady state. What this means is, that it might be beneficial to try, for example `Nswap=2, period=1` instead of `Nswap=10, period=5`. While both combinations result in the same number of pairs swapped on average, the first one is more continous.

- The last parameter `target_momentum` tunes what the momentum swap per particle should be. This parameter, in conjunction with the other ones, can be used to set the overall momentum exchanged over time. If you choose a very large value, the algorihm is unlikely to find pairs with that momentum and will just swap the two biggest ones. For estimating realistic values I recommend to look at the Boltzmann distribution of the velocities and see what momentum values are present in the system at that temperature. If you record `rp_momentum` you should get values close to $2\cdot$`target_momentum`, because the particle in the lower exchange region had a momentum close to `+target_momentum` and the particle in the upper echchange region `-target_momentum`.

There is no easy way or standard recipe for picking the parameters. You should do some simple estimates about how much momentum you roughly want (based on the expected viscosity of your system) to be exchanged and vary the parameters extensively to test.

### Analyzing the results

First, have a look at the generated trajectory. Use either [vmd](https://www.ks.uiuc.edu/Research/vmd/) with the [gsd plugin](https://github.com/mphoward/gsd-vmd), [ovito](https://www.ovito.org/), or your favorite configuration viewer. You should see a homogeneous fluid. The parameters in this tutorial generate a shear rate which is probably hard to see with the naked eye, but if you increase the shear rate, then you'll be able to see the flow. Check for inhomogeneites, bubbles or other unexpected features. For a fluid which can phase separate, this might indicate unwanted temperature/density gradients.

Even if we can't see it with the naked eye, we can check the $v_x(z)$ histogram we recorded to see what the velocity gradient in the system actually is:

```python
import numpy as np
import matplotlib.pyplot as plt

# reading the recorded velocity profile
velocity_profile = np.genfromtxt('tutorial_reverse_perturbation_1_LJ_vx.hist', skip_header=True)
# plotting the velocity profile
plt.figure(figsize=(4,2.2))
plt.plot(velocity_profile[:,0], velocity_profile[:,1])
plt.xlabel('$z$')
plt.ylabel('$v_x(z)$')
plt.show()
```

You should see a triangular shaped flow profile, with a maximum of $\approx 0.07$ in the upper exchange region at $L_z/4=10/4=2.5$ and a minimum of $\approx -0.07$ in the lower exchange region at $-2.5$. Because the simulation wasn't very long, the histogram isn't very smooth, which can be improved with runtime. Also note that we also recorded the startup of the flow from the quiescent equilibrated liquid. Because the simulation was shorter, this effect is probably noticable. To avoid this for your real simulation start recording after the shear flow is established. Simply run a short shear simulation after equilibration, before recording.  The slope of the linear parts in the velocity profile is the shear rate.

This method of generating shear has a upper limit as well as a lower limit for the shear rate. You will see that the velocity profiles eventually will not be triangular anymore if you increase the shear rate by swapping more and more pairs, swapping more frequently or increasing the target momentum. This indicates that you've reached the upper level of the linear shear response regime of your fluid. For a LJ fluid, the shear profiles become somewhat sine/S-shaped. In that case, the slope is not constant, so the shear rate is also not constant, preventing the calculation of the viscosity. If your shear rate is too low, the slope will be very difficult to measure and essentially resemble noise, making your viscosity measurment very inaccurate.

We can also check what momentum we exchange during the simulation:

```python
# reading momentum exchange
data = np.genfromtxt('tutorial_reverse_perturbation_1_LJ.log', skip_header=True)

# plotting
plt.figure(figsize=(4,2.2))
# skipping first entry, because that one will be zero in any case
plt.plot(data[1:,0], data[1:,1])
plt.xlabel('timestep')
plt.ylabel('momentum')
plt.show()
```

Because we used `Nswap=1` and `target_momentum=0.5`, each time the plugin is executed it exchanges momentum of roughly 1.You can see that with the chosen parameters we get a nice constant momentum swap during the entire simulation run.  Note that the sum of `data[1:,2]` does not give the total exchanged momentum, because we only recorded every `period=1e2` timesteps.

### Viscosity calculation

Now, the viscosity $\mu$ in linear shear flow is simply the constant between the shear rate $\dot \gamma$ and stress $\tau$,
\begin{align}
\tau = \mu \dot \gamma \quad \text{or}\quad \mu = \frac{\tau}{\dot \gamma}
\end{align}
Meaning, that if we can determine both shear rate $\dot \gamma$ and stress $\tau$ from our simulation, we can easily calculate the viscosity!
The shear stress in this method is given by:
\begin{align}
\tau = \frac{\langle \Delta p_x\rangle}{2L_xL_y\Delta t}
\end{align}
where $\langle \Delta p_x\rangle$  denotes the average amount of momentum exchanged during the interval $\Delta t$. The original publication uses the total exchanged momentum during the whole simulation time. Because we do not want to write to file that frequently and record large files, we record less frequently and use the average instead. In the end, for a well behaved steady-state simulation this should be identical. The shear rate $\dot \gamma$ can be extracted from the linear part of velocity profile with a linear fit:

```python
L = 10
w = 1

linear_part = velocity_profile[np.abs(velocity_profile[:,0])<L/4. - w/2.]
plt.figure(figsize=(4,2.2))
plt.plot(velocity_profile[:,0], velocity_profile[:,1])
plt.plot(linear_part[:,0], linear_part[:,1])

gamma_dot,b = np.polyfit(linear_part[:,0], linear_part[:,1], 1)
a = np.arange(-L/4.,L/4.,0.1)
plt.plot(a, gamma_dot*a+b)
plt.xlabel('$z$')
plt.ylabel('$v_x(z)$')
plt.show()
```

You can fit both linear parts, or fit a full triangle to increase accuracy, or numerically determine the derivative, as shown below. It is important to exclude the exchange regions regardless of method used, because the velocity profile is usually not linear inside the exchange regions.

```python
plt.figure(figsize=(4,2.2))
dx = velocity_profile[1,0]-velocity_profile[0,0]
grad = np.gradient(velocity_profile[:,1],dx)
plt.plot(velocity_profile[:,0], grad)
b = np.arange(-L/2.,-L/4.,0.1)
c = np.arange(L/4.,L/2.,0.1)
plt.plot(a, gamma_dot*np.ones(len(a)))
plt.plot(b, -gamma_dot*np.ones(len(b)))
plt.plot(c, -gamma_dot*np.ones(len(c)))

plt.xlabel('$z$')
plt.ylabel('$\dot \gamma(z)$')
plt.show()
```

We exchanged a momentum of roughly 1 every 10 timesteps, so 1 per $0.05$ time, since the timestep during the simulation was `dt=0.005`. We can check by calculating the average, which should be close to 1. Then, we simply have to divide $\tau$ by $\dot \gamma$:

```python
# skipping first entry, because that one will be zero in any case
mom_av = np.average(data[1:,1])
tau = mom_av/(2*L*L*0.05)
viscosity = tau/gamma_dot
print(viscosity)
```

    3.2514264848390932

The original publication reports viscosity values around 3.2-3.3, so this very simple and short simulation gives a rough good first estimate already. For increasing the accuracy of this estimate, you should determine the viscosity at different shear rates and plot your results $\tau$ vs $\dot \gamma$.

We have observed unexpected secondary flows in elongated boxes (see this [publication](https://arxiv.org/pdf/1811.04097.pdf) for details), so it is worthwhile to carefully check for box size effects.
You should test the smalles and biggest dimensions you are intending to use in your simulations.

Also vary the parameters of the reverse perturbation to check for systematic effects. Sweep a reasonable set of shear rates to determine the upper and lower limit of the acsessible linear response regime.

Look for unwanted gradients in temperature, density or velocity along all spatial dimensions.

### Future reading

- Original publication: Florian Mueller-Plathe. Reversing the perturbation innonequilibrium
    molecular dynamics:  An easy way to calculate the shear viscosity of
    fluids. Phys. Rev. E, 59:4894-4898, May 1999.
    <http://dx.doi.org/10.1103/PhysRevE.59.4894>
- Follow-up publications, describe the algorithm in more detail:
    * Müller-Plathe, F., & Bordat, P. (2004). Reverse Non-equilibrium Molecular Dynamics. Lecture Notes in Physics, 310–326. <http://dx.doi.org/10.1007/978-3-540-39895-0_10>
    * Müller-Plathe, F., & Reith, D. (1999). Cause and effect reversed in non-equilibrium molecular dynamics: an easy route to transport coefficients. Computational and Theoretical Polymer Science, 9(3-4), 203–209. <http://dx.doi.org/10.1016/s1089-3156(99)00006-9>
- Our own publications about the method, describe some of the observed problems of the algorithm:
    * Instability of Shear Flows in Spatially Periodic Domains, MP Howard, A Statt, HA Stone, TM Truskett, arXiv preprint arXiv:1907.07086 <https://arxiv.org/pdf/1907.07086.pdf>
    * Unexpected secondary flows in reverse nonequilibrium shear flow simulations, A Statt, MP Howard, AZ Panagiotopoulos, Physical Review Fluids 4 (4), 043905 <https://arxiv.org/pdf/1811.04097.pdf>
- This is the documentation page of the lammps implementation of the same algorithm: https://lammps.sandia.gov/doc/fix_viscosity.html Comparing against a different simulation package can be a useful for finding bugs.


Just some scrap paper

```python
%matplotlib inline
%load_ext autoreload
%autoreload 2
from pylab import *
from spylind import spyIVP as sivp
from spylind import spylind as spl
import sympy as sm

from spylind import spylind as spl # master equation useful bits
import qutip as q
sm.init_printing()

# For testing/debuggin:
if 1:
    import MT
    from MT import liouvillian as L
    import pdb
```

    The autoreload extension is already loaded. To reload it, use:
      %reload_ext autoreload

## Copying pyMBE

```python
#
H0, [EinS, H1], [DeltaS, Hdet], c_opL=sld.makeBlochOperators(gsEnergies= [0,1], esEnergies=[0], osc_strengths=np.array([[1],[1]]) ,
                        T1_opt=2)
H_desc = [H0, (EinS, dest), (sm.conjugate(EinS), dest.H), (DeltaS, Hdet) ]

```

```python
H1
```

    matrix([[0, 0, 0.5*conjugate(Ef)],
            [0, 0, 0.5*conjugate(Ef)],
            [0.5*Ef, 0.5*Ef, 0]], dtype=object)

```python
H1a
```

    matrix([[0, 0, 0],
            [0, 0, 0],
            [0.500000000000000, 0.500000000000000, 0]], dtype=object)

```python
zS = sm.symbols("z, ")
H1a = np.matrix(sm.Matrix(np.tril(H1, 0)).subs(EinS,1))
H1b = H1a.T.conj()
H_desc = [H0, [EinS, H1a], [ sm.conjugate(EinS), H1b], (DeltaS, Hdet) ]
display(H_desc[0], c_opL, H1a, H1b)
```

    matrix([[ 0.        +0.j,  0.        +0.j,  0.        +0.j],
            [ 0.        +0.j, -6.28318531+0.j,  0.        +0.j],
            [ 0.        +0.j,  0.        +0.j,  0.        +0.j]])

    [matrix([[0, 0, 0.707106781186548],
             [0, 0, 0],
             [0, 0, 0]], dtype=object),
     matrix([[0, 0, 0],
             [0, 0, 0.707106781186548],
             [0, 0, 0]], dtype=object)]

    matrix([[0, 0, 0],
            [0, 0, 0],
            [0.500000000000000, 0.500000000000000, 0]], dtype=object)

    matrix([[0, 0, 0.500000000000000],
            [0, 0, 0.500000000000000],
            [0, 0, 0]], dtype=object)

```python
eq, ex_vals=spl.makeMESymb(H_desc, c_opL=c_opL, e_opL = [H1a], bReturnMatrixEquation=True)
display(eq, ex_vals)
```

```python

```


Dinámica de Fluidos Computacional Aplicada - IM0642

Sesión 03: Formulación diferencial para las ecuaciones de fluidos

ESCUELA DE INGENIERÍA

Especialización en Turbomáquinas

MEDELLÍN - COLOMBIA

2021/01

 Text provided under a Creative Commons Attribution license, CC-BY. All code is made available under the FSF-approved MIT license.(c) Carlos Alberto Alvarez Henao

***

***Docente:*** Carlos Alberto Álvarez Henao, I.C. D.Sc.

***e-mail:*** calvar52@eafit.edu.co

***skype:*** carlos.alberto.alvarez.henao

***Herramienta:*** [Jupyter](http://jupyter.org/)

***Kernel:*** Python 3.8

***

Tabla de contenido

1  Ecuaciones para CFD

1.1  Conservación de la masa  (2D):

1.2  Conservación del Momentum (2D)

1.2.1  Forma conservativa

1.2.2  Ecuaciones de Conservación del Momentum Forma conservativa (2D)

1.3  Conservación de la Energía (2D)

2  Ecuaciones de la Dinámica de Fluidos

2.1  Ecuaciones del Movimiento

2.2  Ecuación termodinámica

2.3  Ecuación de continuidad (de la masa)

## Ecuaciones para CFD

[Volver a la Tabla de Contenido](#TOC)

### Conservación de la masa  (*2D*):

La anterior figura implica que *la variación temporal de la cantidad de masa en el elemento es igual a la descarga resultante a través de las fronteras del elemento*. Matemáticamente se expresa:

\begin{equation}
\frac{\delta \rho}{\delta t}\delta x \delta y = - \left[ \frac{\delta (\rho u)}{\delta x} \delta x \right] \delta y - \left[ \frac{\delta (\rho v)}{\delta y} \delta y \right] \delta x
\end{equation}

Como se trata de elementos de tamaño diferencial, la multiplicación del valor de dos elementos de tamaño muy pequeño darán un valor mucho menor, entonces, el valor $\delta x \delta y \rightarrow 0$, y reemplazando en la ecuación anterior:

$$\frac{\partial \rho}{\partial t} + \left [ \frac{\partial (\rho u)}{\partial x} + \frac{\partial (\rho v)}{\partial x} \right ] =0$$

El primer término es la variación temporal de la densidad, y el segundo término, entre corchetes, es la tasa de variación de masa por unidad de volumen.

[Volver a la Tabla de Contenido](#TOC)

### Conservación del Momentum (*2D*)

de la ***2a Ley de Newton***

Variación de la velocidad entre dos puntos

$$\delta u = \frac{\partial u}{\partial x}\delta x+\frac{\partial u}{\partial y}\delta y+\frac{\partial u}{\partial t}\delta t$$

dividiendo por $\delta t$

\begin{equation}
\begin{split}
\frac{\delta u}{\delta t} &= \frac{\partial u}{\partial x}\frac{\delta x}{\delta t}+\frac{\partial u}{\partial y}\frac{\delta y}{\delta t}+\frac{\partial u}{\partial t} \\
 &=\frac{\partial u}{\partial t}+u\frac{\partial u}{\partial x}+v\frac{\partial u}{\partial y}
\end{split}
\end{equation}

el primer término de la ecuación anterior es la acelaración del elemento de fluido.

[Volver a la Tabla de Contenido](#TOC)

#### Forma conservativa

Los coeficientes de las derivadas son constantes, o si son variables, sus derivadas no aparecen en la *EDP*!

\begin{equation}
\begin{split}
\rho \frac{Du}{Dt}+u \underbrace{\left [ \frac{\partial \rho}{\partial t}+\frac{\partial (\rho u)}{\partial x}+\frac{\partial (\rho v)}{\partial y} \right ]}_{\text{Ecuación de continuidad=0}} \\
& =\frac{\partial (\rho u)}{\partial t}+\frac{\partial (\rho u^2)}{\partial x}+\frac{\partial (\rho uv)}{\partial y} \\
& =\frac{\partial (\rho u)}{\partial t}+\nabla \cdot (\rho u \boldsymbol{V})
\end{split}
\end{equation}

[Volver a la Tabla de Contenido](#TOC)

#### Ecuaciones de Conservación del Momentum Forma conservativa (*2D*)

$$\frac{\partial (\rho u)}{\partial t}+\frac{\partial (\rho u^2)}{\partial x}+\frac{\partial (\rho uv)}{\partial y}=-\frac{\partial p}{\partial x}+\frac{\partial \tau_{xx}}{\partial x}+\frac{\partial \tau_{yx}}{\partial y}+\rho F_{x}$$

$$\frac{\partial (\rho v)}{\partial t}+\frac{\partial (\rho v^2)}{\partial y}+\frac{\partial (\rho uv)}{\partial x}=-\frac{\partial p}{\partial y}+\frac{\partial \tau_{yy}}{\partial y}+\frac{\partial \tau_{xy}}{\partial y}+\rho F_{y}$$

[Volver a la Tabla de Contenido](#TOC)

### Conservación de la Energía (*2D*)

$$\rho \frac{DE}{Dt}=S+W$$

Suma de tres componentes:

- ***Energía Cinética*** de traslación del fluido debido a su movimiento;

- ***Energía interna*** del fluido debido a la vibración natural de las moléculas que lo componen;

- ***Energía Potencial Gravitacional***.

con:

$$E=e+\frac{V^2}{2}$$

donde $\frac{V^2}{2}$ es la energía por unidad de masa y $V=\sqrt{u^2+v^2}$

\begin{equation}
\begin{split}
S & =\rho \frac{\partial Q}{\partial t} + \frac{\partial}{\partial x} \left ( \kappa \frac{\partial T}{\partial x} \right ) +\frac{\partial}{\partial y} \left ( \kappa \frac{\partial T}{\partial y} \right )\\
 &=\rho \frac{\partial Q}{\partial t} + \nabla \cdot \dot{\boldsymbol{q}}
\end{split}
\end{equation}

contiene los efectos de las fuentes de calor sobre la energía total del fluido.

- $\dot{\boldsymbol{q}}=\kappa \nabla T$: Flujo de calor a través de las paredes del volumen de control.

- $\kappa$: coeficiente de conductividad térmica del fluido

- $\rho \frac{\partial Q}{\partial t}$: Tasa de calor por unidad de masa producido por agentes externos, como absorción de calor por el fluido, o internos como reacciones químicas.

- $\nabla \cdot \dot{\boldsymbol{q}}$: Transferencia de calor por conducción a través de las paredes del elemento de fluido, debido a la presencia de gradientes de temperatura (radiación).

Trabajo realizado sobre el fluido por las fuerzas externas, incluyendo la contribución debido a la variación de la energía potencial:

$$W=- \underbrace{ \left[ \frac{\partial (up)}{\partial x} + \frac{\partial (vp)}{\partial y} + \frac{\partial (u \tau_{xx})}{\partial x} + \frac{\partial (u \tau_{yx})}{\partial y} + \frac{\partial (v \tau_{xy})}{\partial x} + \frac{\partial (v \tau_{yy})}{\partial y}\right]}_{\text{Tensiones viscosas y presión}} + \underbrace{\rho \boldsymbol{F} \cdot \boldsymbol{V}}_{\text{Fuerzas de campo}}$$

[Volver a la Tabla de Contenido](#TOC)

## Ecuaciones de la Dinámica de Fluidos

### Ecuaciones del Movimiento

$$\frac{\partial u}{\partial t} = \underbrace{-u\frac{\partial u}{\partial x}-v\frac{\partial u}{\partial y}}_{convección}-\underbrace{\theta \frac{\partial \pi'}{\partial x}}_{presión} + \underbrace{fv}_{coriolis}+\underbrace{\frac{\partial}{\partial x} \left( K_m\frac{\partial u}{\partial x}\right)+\frac{\partial}{\partial y} \left( K_m\frac{\partial u}{\partial y}\right)}_{difusión}$$

$$\frac{\partial v}{\partial t} = -u\frac{\partial v}{\partial x}-v\frac{\partial v}{\partial y}-\theta \frac{\partial \pi'}{\partial x} + fu+\frac{\partial}{\partial x} \left( K_m\frac{\partial v}{\partial x}\right)+\frac{\partial}{\partial y} \left( K_m\frac{\partial v}{\partial y}\right)$$

[Volver a la Tabla de Contenido](#TOC)

### Ecuación termodinámica

$$\frac{\partial \theta _{il}}{\partial t}=-u\frac{\partial \theta_{il}}{\partial x}-v\frac{\partial \theta_{il}}{\partial y}+\frac{\partial}{\partial x} \left( K_h \frac{\partial \theta_{il}}{\partial t}\right)_{rad}+\frac{\partial}{\partial x}\left( K_h \frac{\partial \theta_{il}}{\partial x}\right)+\frac{\partial}{\partial y}\left( K_h \frac{\partial \theta_{il}}{\partial y}\right)$$

[Volver a la Tabla de Contenido](#TOC)

### Ecuación de continuidad (de la masa)

$$\frac{\partial \pi'}{\partial t}=-\frac{R\pi_0}{c_v\rho_0\theta_0}\left(\frac{\partial\rho_0\theta_0u}{\partial x}+\frac{\partial\rho_0\theta_0v}{\partial y}\right)$$

[Volver a la Tabla de Contenido](#TOC)