Lambda School Data Science

*Unit 2, Sprint 1, Module 2*

---

# Regression & Classification, Module 2
- Go from simple regression (1 feature) to multiple regression (2+ features)
- Get and plot coefficients
- Use regression metrics: MAE (Mean Absolute Error) and $R^2$ Score
- Understand how ordinary least squares regression minimizes the sum of squared errors
- Explain why overfitting is a problem. Do train/test split

### Setup

You can work locally (follow the [local setup instructions](https://lambdaschool.github.io/ds/unit2/local/)) or on Colab (run the code cell below).

```python
import os, sys
in_colab = 'google.colab' in sys.modules

# If you're in Colab...
if in_colab:
    # Pull files from Github repo
    os.chdir('/content')
    !git init .
    !git remote add origin https://github.com/LambdaSchool/DS-Unit-2-Regression-Classification.git
    !git pull origin master

    # Install required python packages
    !pip install -r requirements.txt

    # Change into directory for module
    os.chdir('module2')
```

```python
# Ignore this Numpy warning when using Plotly Express:
# FutureWarning: Method .ptp is deprecated and will be removed in a future version. Use numpy.ptp instead.
import warnings
warnings.filterwarnings(action='ignore', category=FutureWarning, module='numpy')
```

# Go from simple regression (1 feature) to multiple regression (2+ features)

## Overview: Predict Elections! 🗳

```python
import pandas as pd
import plotly.express as px

df = pd.read_csv('../data/bread_peace_voting.csv')
px.scatter(
    df,
    x='Average Recent Growth in Personal Incomes',
    y='Incumbent Party Vote Share',
    text='Year',
    title='US Presidential Elections, 1952-2016',
    trendline='ols',  # Ordinary Least Squares
)
```

#### Douglas Hibbs, [Background Information on the ‘Bread and Peace’ Model of Voting in Postwar US Presidential Elections](https://douglas-hibbs.com/background-information-on-bread-and-peace-voting-in-us-presidential-elections/)

> Aggregate two-party vote shares going to candidates of the party holding the presidency during the postwar era are well explained by just two fundamental determinants:

> (1) Positively by weighted-average growth of per capita real disposable personal income over the term.
> (2) Negatively by cumulative US military fatalities (scaled to population) owing to unprovoked, hostile deployments of American armed forces in foreign wars.

#### Data sources
- 1952-2012: Douglas Hibbs, [2014 lecture at Deakin University Melbourne](http://www.douglas-hibbs.com/HibbsArticles/HIBBS-PRESVOTE-SLIDES-MELBOURNE-Part1-2014-02-26.pdf), Slide 40
- 2016, Vote Share: [The American Presidency Project](https://www.presidency.ucsb.edu/statistics/elections)
- 2016, Recent Growth in Personal Incomes: [The 2016 election economy: the "Bread and Peace" model final forecast](https://angrybearblog.com/2016/11/the-2016-election-economy-the-bread-and-peace-model-final-forecast.html)
- 2016, US Military Fatalities: Assumption that Afghanistan War fatalities in 2012-16 occured at the same rate as 2008-12

> Fatalities denotes the cumulative number of American military fatalities per millions of US population the in Korea, Vietnam, Iraq and Afghanistan wars during the presidential terms preceding the 1952, 1964, 1968, 1976 and 2004, 2008 and 2012 elections. —[Hibbs](http://www.douglas-hibbs.com/HibbsArticles/HIBBS-PRESVOTE-SLIDES-MELBOURNE-Part1-2014-02-26.pdf), Slide 33

## Follow Along

Look at the data

```python

```

What's the average Incumbent Party Vote Share?

```python

```

Add another feature to the scatterplot

```python

```

Plot the hyperplane of best fit

```python
import itertools
import numpy as np
import plotly.express as px
import plotly.graph_objs as go
from sklearn.linear_model import LinearRegression

def regression_3d(df, x, y, z, num=100, **kwargs):
    """
    Visualize linear regression in 3D: 2 features + 1 target

    df : Pandas DataFrame
    x : string, feature 1 column in df
    y : string, feature 2 column in df
    z : string, target column in df
    num : integer, number of quantiles for each feature
    """

    # Plot data
    fig = px.scatter_3d(df, x, y, z, **kwargs)

    # Fit Linear Regression
    features = [x, y]
    target = z
    model = LinearRegression()
    model.fit(df[features], df[target])

    # Define grid of coordinates in the feature space
    xmin, xmax = df[x].min(), df[x].max()
    ymin, ymax = df[y].min(), df[y].max()
    xcoords = np.linspace(xmin, xmax, num)
    ycoords = np.linspace(ymin, ymax, num)
    coords = list(itertools.product(xcoords, ycoords))

    # Make predictions for the grid
    predictions = model.predict(coords)
    Z = predictions.reshape(num, num).T

    # Plot predictions as a 3D surface (plane)
    fig.add_trace(go.Surface(x=xcoords, y=ycoords, z=Z))

    return fig
```

```python

```

Fit Linear Regression with 2 features

- Jake VanderPlas, [_Python Data Science Handbook,_ Chapter 5.2, Introducing Scikit-Learn — Basics of the API](https://jakevdp.github.io/PythonDataScienceHandbook/05.02-introducing-scikit-learn.html#Basics-of-the-API)
- Scikit-Learn documentation, [sklearn.linear_model.LinearRegression](https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.LinearRegression.html)

```python
# 1. Import the appropriate estimator class from Scikit-Learn

# 2. Instantiate this class

# 3. Arrange X features matrix & y target vector

# 4. Fit the model
```

## Challenge

# Get and plot coefficients

## Overview

## Follow Along

What's the equation for the hyperplane?

Can you relate the intercept and coefficients to what you see in the plot?

```python

```

What if ...

Income growth = 0%, fatalities = 0

```python

```

Income growth = 1% (fatalities = 0)

```python

```

The difference between these predictions = ?

```python

```

What if... income growth = 2% (fatalities = 0)

```python

```

The difference between these predictions = ?

```python

```

What if... (income growth=2%) fatalities = 100

```python

```

The difference between these predictions = ?

```python

```

What if income growth = 3% (fatalities = 100)

```python

```

The difference between these predictions = ?

```python

```

What if (income growth = 3%) fatalities = 200

```python

```

The difference between these predictions = ?

```python

```

Plot coefficients

```python

```

## Challenge

#  Use regression metrics: MAE (Mean Absolute Error) and $R^2$ Score

## Overview

## Follow Along

How's the error?

```python

```

How does this compare to guessing?

```python

```

Plot the "residuals" (errors)

```python
import matplotlib.pyplot as plt
from sklearn.metrics import mean_absolute_error, r2_score

def regression_residuals(df, feature, target, m, b):
    """
    Visualize linear regression, with residual errors,
    in 2D: 1 feature + 1 target.

    Use the m & b parameters to "fit the model" manually.

    df : Pandas DataFrame
    feature : string, feature column in df
    target : string, target column in df
    m : numeric, slope for linear equation
    b : numeric, intercept for linear requation
    """

    # Plot data
    df.plot.scatter(feature, target)

    # Make predictions
    x = df[feature]
    y = df[target]
    y_pred = m*x + b

    # Plot predictions
    plt.plot(x, y_pred)

    # Plot residual errors
    for x, y1, y2 in zip(x, y, y_pred):
        plt.plot((x, x), (y1, y2), color='grey')

    # Print regression metrics
    mae = mean_absolute_error(y, y_pred)
    r2 = r2_score(y, y_pred)
    print('Mean Absolute Error:', mae)
    print('R^2:', r2)

```

```python

```

## Challenge

# Understand how ordinary least squares regression minimizes the sum of squared errors

## Overview

But Ordinary Least Squares Regression *doesn't* directly minimize MAE or R^2...

1. Guess & Check
2. Linear Algebra

## Follow Along

### Guess & Check

```python
from matplotlib.patches import Rectangle
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

def regression_squared_errors(df, feature, target, m, b):
    """
    Visualize linear regression, with squared errors,
    in 2D: 1 feature + 1 target.

    Use the m & b parameters to "fit the model" manually.

    df : Pandas DataFrame
    feature : string, feature column in df
    target : string, target column in df
    m : numeric, slope for linear equation
    b : numeric, intercept for linear requation
    """

    # Plot data
    fig = plt.figure(figsize=(7,7))
    ax = plt.axes()
    df.plot.scatter(feature, target, ax=ax)

    # Make predictions
    x = df[feature]
    y = df[target]
    y_pred = m*x + b

    # Plot predictions
    ax.plot(x, y_pred)

    # Plot squared errors
    xmin, xmax = ax.get_xlim()
    ymin, ymax = ax.get_ylim()
    scale = (xmax-xmin)/(ymax-ymin)
    for x, y1, y2 in zip(x, y, y_pred):
        bottom_left = (x, min(y1, y2))
        height = abs(y1 - y2)
        width = height * scale
        ax.add_patch(Rectangle(xy=bottom_left, width=width, height=height, alpha=0.1))

    # Print regression metrics
    mse = mean_squared_error(y, y_pred)
    rmse = np.sqrt(mse)
    mae = mean_absolute_error(y, y_pred)
    r2 = r2_score(y, y_pred)
    print('Mean Squared Error:', mse)
    print('Root Mean Squared Error:', rmse)
    print('Mean Absolute Error:', mae)
    print('R^2:', r2)
```

```python

```

### Linear Algebra

The same result that is found by minimizing the sum of the squared errors can be also found through a linear algebra process known as the "Least Squares Solution:"

\begin{align}
\hat{\beta} = (X^{T}X)^{-1}X^{T}y
\end{align}

Before we can work with this equation in its linear algebra form we have to understand how to set up the matrices that are involved in this equation.

### The $\beta$ vector

The $\beta$ vector represents all the parameters that we are trying to estimate, our $y$ vector and $X$ matrix values are full of data from our dataset. The $\beta$ vector holds the variables that we are solving for: $\beta_0$ and $\beta_1$

Now that we have all of the necessary parts we can set them up in the following equation:

\begin{align}
y = X \beta + \epsilon
\end{align}

Since our $\epsilon$ value represents **random** error we can assume that it will equal zero on average.

\begin{align}
y = X \beta
\end{align}

The objective now is to isolate the $\beta$ matrix. We can do this by pre-multiplying both sides by "X transpose" $X^{T}$.

\begin{align}
X^{T}y =  X^{T}X \beta
\end{align}

Since anything times its transpose will result in a square matrix, if that matrix is then an invertible matrix, then we should be able to multiply both sides by its inverse to remove it from the right hand side. (We'll talk tomorrow about situations that could lead to $X^{T}X$ not being invertible.)

\begin{align}
(X^{T}X)^{-1}X^{T}y =  (X^{T}X)^{-1}X^{T}X \beta
\end{align}

Since any matrix multiplied by its inverse results in the identity matrix, and anything multiplied by the identity matrix is itself, we are left with only $\beta$ on the right hand side:

\begin{align}
(X^{T}X)^{-1}X^{T}y = \hat{\beta}
\end{align}

We will now call it "beta hat" $\hat{\beta}$ because it now represents our estimated values for $\beta_0$ and $\beta_1$

### Lets calculate our $\beta$ coefficients with numpy!

```python
# This is NOT an objective you'll be tested on. It's just a demo.

# X is a matrix. Add constant for the intercept.
from statsmodels.api import add_constant
X = add_constant(df[feature].values)
print('X')
print(X)

# y is a column vector
y = df[target].values[:, np.newaxis]
print('y')
print(y)

# Least squares solution in code
X_transpose = X.T
X_transpose_X = X_transpose @ X
X_transpose_X_inverse = np.linalg.inv(X_transpose_X)
X_transpose_y = X_transpose @ y
beta_hat = X_transpose_X_inverse @ X_transpose_y

print('Beta Hat')
print(beta_hat)
```

```python
# Scikit-learn gives the exact same results
from sklearn.linear_model import LinearRegression
model = LinearRegression()
features = ['Average Recent Growth in Personal Incomes']
target = 'Incumbent Party Vote Share'
X = df[features]
y = df[target]
model.fit(X, y)
model.intercept_, model.coef_
```

## Challenge

# Explain why overfitting is a problem. Do train/test split

## Overview

#### Jake VanderPlas, [_Python Data Science Handbook,_ Chapter 5.3](https://jakevdp.github.io/PythonDataScienceHandbook/05.03-hyperparameters-and-model-validation.html#The-Bias-variance-trade-off)

Fundamentally, the question of "the best model" is about finding a sweet spot in the tradeoff between bias and variance. Consider the following figure, which presents two regression fits to the same dataset:

The model on the left attempts to find a straight-line fit through the data. Because the data are intrinsically more complicated than a straight line, the straight-line model will never be able to describe this dataset well. Such a model is said to _underfit_ the data: that is, it does not have enough model flexibility to suitably account for all the features in the data; another way of saying this is that the model has high _bias_.

The model on the right attempts to fit a high-order polynomial through the data. Here the model fit has enough flexibility to nearly perfectly account for the fine features in the data, but even though it very accurately describes the training data, its precise form seems to be more reflective of the particular noise properties of the data rather than the intrinsic properties of whatever process generated that data. Such a model is said to _overfit_ the data: that is, it has so much model flexibility that the model ends up accounting for random errors as well as the underlying data distribution; another way of saying this is that the model has high _variance_.

From the scores associated with these two models, we can make an observation that holds more generally:

- For high-bias models, the performance of the model on the validation set is similar to the performance on the training set.

- For high-variance models, the performance of the model on the validation set is far worse than the performance on the training set.

If we imagine that we have some ability to tune the model complexity, we would expect the training score and validation score to behave as illustrated in the following figure:

The diagram shown here is often called a validation curve, and we see the following essential features:

- The training score is everywhere higher than the validation score. This is generally the case: the model will be a better fit to data it has seen than to data it has not seen.
- For very low model complexity (a high-bias model), the training data is under-fit, which means that the model is a poor predictor both for the training data and for any previously unseen data.
- For very high model complexity (a high-variance model), the training data is over-fit, which means that the model predicts the training data very well, but fails for any previously unseen data.
- For some intermediate value, the validation curve has a maximum. This level of complexity indicates a suitable trade-off between bias and variance.

The means of tuning the model complexity varies from model to model.

## Follow Along

Wrangle New York City property sales data

```python
# Read New York City property sales data, from first 4 months of 2019.
# Dataset has 23040 rows, 21 columns.
df = pd.read_csv('../data/NYC_Citywide_Rolling_Calendar_Sales.csv')
assert df.shape == (23040, 21)

# Change column names. Replace spaces with underscores
df.columns = [col.replace(' ', '_') for col in df]

# Remove symbols from SALE_PRICE string, convert to integer
df['SALE_PRICE'] = (
    df['SALE_PRICE']
    .str.replace('$','')
    .str.replace('-','')
    .str.replace(',','')
    .astype(int)
)

# Keep subset of rows:
# Tribeca neighborhood, Condos - Elevator Apartments,
# 1 unit, sale price more than $1, less than $35 million
mask = (
    (df['NEIGHBORHOOD'].str.contains('TRIBECA')) &
    (df['BUILDING_CLASS_CATEGORY'] == '13 CONDOS - ELEVATOR APARTMENTS') &
    (df['TOTAL_UNITS'] == 1) &
    (df['SALE_PRICE'] > 0) &
    (df['SALE_PRICE'] < 35000000)
)
df = df[mask]

# Data now has 90 rows, 21 columns
assert df.shape == (90, 21)

# Convert SALE_DATE to datetime
df['SALE_DATE'] = pd.to_datetime(df['SALE_DATE'], infer_datetime_format=True)

from ipywidgets import interact
import pandas as pd
from sklearn.linear_model import LinearRegression

# Read New York City property sales data, from first 4 months of 2019.
# Dataset has 23040 rows, 21 columns.
df = pd.read_csv('../data/NYC_Citywide_Rolling_Calendar_Sales.csv')
assert df.shape == (23040, 21)

# Change column names. Replace spaces with underscores
df.columns = [col.replace(' ', '_') for col in df]

# Remove symbols from SALE_PRICE string, convert to integer
df['SALE_PRICE'] = (
    df['SALE_PRICE']
    .str.replace('$','')
    .str.replace('-','')
    .str.replace(',','')
    .astype(int)
)

# Keep subset of rows:
# Tribeca neighborhood, Condos - Elevator Apartments,
# 1 unit, sale price more than $1, less than $35 million
mask = (
    (df['NEIGHBORHOOD'].str.contains('TRIBECA')) &
    (df['BUILDING_CLASS_CATEGORY'] == '13 CONDOS - ELEVATOR APARTMENTS') &
    (df['TOTAL_UNITS'] == 1) &
    (df['SALE_PRICE'] > 0) &
    (df['SALE_PRICE'] < 35000000)
)
df = df[mask]

# Data now has 90 rows, 21 columns
assert df.shape == (90, 21)

# Convert SALE_DATE to datetime
df['SALE_DATE'] = pd.to_datetime(df['SALE_DATE'], infer_datetime_format=True)

# Arrange X features matrix & y target vector
features = ['GROSS_SQUARE_FEET']
target = 'SALE_PRICE'
X = df[features]
y = df[target]
```

Do random train/test spliit

```python
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=11)
```

Repeatedly fit increasingly complex models, and keep track of the scores

```python
from IPython.display import display, HTML
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression

# Credit: Jake VanderPlas, Python Data Science Handbook, Chapter 5.3
# https://jakevdp.github.io/PythonDataScienceHandbook/05.03-hyperparameters-and-model-validation.html#Validation-curves-in-Scikit-Learn
def PolynomialRegression(degree=2, **kwargs):
    return make_pipeline(PolynomialFeatures(degree),
                         LinearRegression(**kwargs))

polynomial_degrees = range(1, 10, 2)
train_r2s = []
test_r2s = []

for degree in polynomial_degrees:
    model = PolynomialRegression(degree)
    display(HTML(f'Polynomial degree={degree}'))

    model.fit(X_train, y_train)
    train_r2 = model.score(X_train, y_train)
    test_r2 = model.score(X_test, y_test)
    display(HTML(f'<b style="color: blue">Train R2 {train_r2:.2f}</b>'))
    display(HTML(f'<b style="color: red">Test R2 {test_r2:.2f}</b>'))

    plt.scatter(X_train, y_train, color='blue', alpha=0.5)
    plt.scatter(X_test, y_test, color='red', alpha=0.5)
    plt.xlabel(features)
    plt.ylabel(target)

    x_domain = np.linspace(X.min(), X.max())
    curve = model.predict(x_domain)
    plt.plot(x_domain, curve, color='blue')
    plt.show()
    display(HTML('<hr/>'))

    train_r2s.append(train_r2)
    test_r2s.append(test_r2)

    guess = model.predict(np.reshape(1497, (1,1)))
    display(HTML(f'Guess))

display(HTML('Validation Curve'))
plt.plot(polynomial_degrees, train_r2s, color='blue', label='Train')
plt.plot(polynomial_degrees, test_r2s, color='red', label='Test')
plt.xlabel('Model Complexity (Polynomial Degree)')
plt.ylabel('R^2 Score')
plt.legend()
plt.show()
```

## Challenge

# Review

# Sources

Jake VanderPlas, [_Python Data Science Handbook,_ Chapter 5.3](https://jakevdp.github.io/PythonDataScienceHandbook/05.03-hyperparameters-and-model-validation.html#The-Bias-variance-trade-off)

#### Douglas Hibbs, [Background Information on the ‘Bread and Peace’ Model of Voting in Postwar US Presidential Elections](https://douglas-hibbs.com/background-information-on-bread-and-peace-voting-in-us-presidential-elections/)

> Aggregate two-party vote shares going to candidates of the party holding the presidency during the postwar era are well explained by just two fundamental determinants:

> (1) Positively by weighted-average growth of per capita real disposable personal income over the term.
> (2) Negatively by cumulative US military fatalities (scaled to population) owing to unprovoked, hostile deployments of American armed forces in foreign wars.

#### Data sources
- 1952-2012: Douglas Hibbs, [2014 lecture at Deakin University Melbourne](http://www.douglas-hibbs.com/HibbsArticles/HIBBS-PRESVOTE-SLIDES-MELBOURNE-Part1-2014-02-26.pdf), Slide 40
- 2016, Vote Share: [The American Presidency Project](https://www.presidency.ucsb.edu/statistics/elections)
- 2016, Recent Growth in Personal Incomes: [The 2016 election economy: the "Bread and Peace" model final forecast](https://angrybearblog.com/2016/11/the-2016-election-economy-the-bread-and-peace-model-final-forecast.html)
- 2016, US Military Fatalities: Assumption that Afghanistan War fatalities in 2012-16 occured at the same rate as 2008-12

> Fatalities denotes the cumulative number of American military fatalities per millions of US population the in Korea, Vietnam, Iraq and Afghanistan wars during the presidential terms preceding the 1952, 1964, 1968, 1976 and 2004, 2008 and 2012 elections. —[Hibbs](http://www.douglas-hibbs.com/HibbsArticles/HIBBS-PRESVOTE-SLIDES-MELBOURNE-Part1-2014-02-26.pdf), Slide 33


# ForestDML, ForestDRLearner, OrthoForest and CausalForest: Basic Example

We depict the performance of our `ForestDML`, `ForestDRLearner`, `OrthoForest` and `CausalForest` estimators on the same data generating process as the one used in the tutorial page of the grf package (see https://github.com/grf-labs/grf#usage-examples). This is mostly for qualitative comparison and verification purposes among our implementation of variants of Causal Forests and the implementation in the grf R package.

```python
# Helper imports
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
%matplotlib inline
```

# DGP

We use the following data generating process (DGP) from [here](https://github.com/grf-labs/grf#usage-examples):

\begin{align}
X \sim& \text{Normal}(0,\, I_{p})\\
T =& \text{Binomial}(1, .4 + .2 \cdot 1\{X[0] > 0\})\\
Y =& (X[0] \cdot 1\{X[0] > 0\}) \cdot T + X[1] + X[2] \cdot 1\{X[2] < 0\} + \epsilon, &\; \epsilon \sim \text{Normal}(0, 1)\\
\end{align}

We use $p=10$ and draw $n=2000$ samples from this DGP.

```python
import numpy as np
import scipy.special
np.random.seed(123)
n = 2000
p = 10
X = np.random.normal(size=(n, p))
true_propensity = lambda x: .4 + .2 * (x[:, 0] > 0)
true_effect = lambda x: (x[:, 0] * (x[:, 0] > 0))
true_conf = lambda x: x[:, 1] + np.clip(x[:, 2], - np.inf, 0)
T = np.random.binomial(1, true_propensity(X))
Y =  true_effect(X) * T + true_conf(X) + np.random.normal(size=(n,))
```

### Cross-Validated Forest Nuisance Models

We use forest based estimators (Gradient Boosted Forests or Random Forests) as nuisance models. For the meta-learner versions of our forest based estimators, we also use a generic forest estimator even as a final model. The hyperparameters of the forest models (e.g. number of estimators, max depth, min leaf size) is chosen via cross validation. We also choose among Gradient or Random Forests via cross validation

```python
from econml.sklearn_extensions.model_selection import GridSearchCVList
from sklearn.linear_model import Lasso, LogisticRegression
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.base import clone
from econml.sklearn_extensions.linear_model import WeightedLasso

def first_stage_reg():
    return GridSearchCVList([Lasso(),
                             RandomForestRegressor(n_estimators=100, random_state=123),
                             GradientBoostingRegressor(random_state=123)],
                             param_grid_list=[{'alpha': [.001, .01, .1, 1, 10]},
                                               {'max_depth': [3, None],
                                               'min_samples_leaf': [10, 50]},
                                              {'n_estimators': [50, 100],
                                               'max_depth': [3],
                                               'min_samples_leaf': [10, 30]}],
                             cv=5,
                             scoring='neg_mean_squared_error')

def first_stage_clf():
    return GridSearchCVList([LogisticRegression(),
                             RandomForestClassifier(n_estimators=100, random_state=123),
                             GradientBoostingClassifier(random_state=123)],
                             param_grid_list=[{'C': [0.01, .1, 1, 10, 100]},
                                              {'max_depth': [3, 5],
                                               'min_samples_leaf': [10, 50]},
                                              {'n_estimators': [50, 100],
                                               'max_depth': [3],
                                               'min_samples_leaf': [10, 30]}],
                             cv=5,
                             scoring='neg_mean_squared_error')

def final_stage():
    return GridSearchCVList([WeightedLasso(),
                             RandomForestRegressor(n_estimators=100, random_state=123)],
                             param_grid_list=[{'alpha': [.001, .01, .1, 1, 10]},
                                              {'max_depth': [3, 5],
                                               'min_samples_leaf': [10, 50]}],
                             cv=5,
                             scoring='neg_mean_squared_error')
```

```python
model_y = clone(first_stage_reg().fit(X, Y).best_estimator_)
model_y
```

    GradientBoostingRegressor(min_samples_leaf=30, n_estimators=50,
                              random_state=123)

```python
model_t = clone(first_stage_clf().fit(X, T).best_estimator_)
model_t
```

    RandomForestClassifier(max_depth=3, min_samples_leaf=10, random_state=123)

# DML Estimators

```python
from econml.dml import CausalForestDML
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
from sklearn.dummy import DummyRegressor, DummyClassifier

est = CausalForestDML(model_y=model_y,
                      model_t=model_t,
                      discrete_treatment=True,
                      cv=3,
                      n_estimators=4000,
                      min_samples_leaf=10,
                      verbose=0,
                      min_weight_fraction_leaf=.005)
est.fit(Y, T, X=X)
```

    <econml.dml.causal_forest.CausalForestDML at 0x1fbaae74c08>

```python
from econml.dml import NonParamDML
est2 = NonParamDML(model_y=model_y,
                   model_t=model_t,
                   cv=3,
                   discrete_treatment=True,
                   model_final=final_stage())
est2.fit(Y, T, X=X)
```

    <econml.dml.dml.NonParamDML at 0x1fbb3693c48>

```python
X_test = np.zeros((100, p))
X_test[:, 0] = np.linspace(-2, 2, 100)
```

```python
pred = est.effect(X_test)
lb, ub = est.effect_interval(X_test, alpha=0.01)
```

```python
pred2 = est2.effect(X_test)
```

```python
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.plot(X_test[:, 0], true_effect(X_test), '--')
plt.plot(X_test[:, 0], pred2, label='nonparamdml')
plt.plot(X_test[:, 0], pred, label='forestdml (causal forest)')
plt.fill_between(X_test[:, 0], lb, ub, alpha=.4, label='honestrf_ci')
plt.legend()
plt.show()
```

### First Stage Learned Models

```python
# Model T
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.title('honestrf')
for mdl in est.models_t:
    plt.plot(X_test[:, 0], mdl.predict_proba(X_test)[:, 1])
plt.plot(X_test[:, 0], true_propensity(X_test), '--', label='truth')
plt.legend()

plt.subplot(1, 2, 2)
plt.title('rf')
for mdl in est2.models_t:
    plt.plot(X_test[:, 0], mdl.predict_proba(X_test)[:, 1])
plt.plot(X_test[:, 0], true_propensity(X_test), '--', label='truth')
plt.legend()
plt.show()
```

```python
# Model Y
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.title('honestrf')
for mdl in est.models_y:
    plt.plot(X_test[:, 0], mdl.predict(X_test))
plt.plot(X_test[:, 0], true_effect(X_test) * true_propensity(X_test) + true_conf(X_test), '--', label='truth')
plt.legend()

plt.subplot(1, 2, 2)
plt.title('rf')
for mdl in est2.models_y:
    plt.plot(X_test[:, 0], mdl.predict(X_test))
plt.plot(X_test[:, 0], true_effect(X_test) * true_propensity(X_test) + true_conf(X_test), '--', label='truth')
plt.legend()
plt.show()
```

### Interpretability of CATE Model of NonParamDML with SHAP

```python
import shap
import string

feature_names = list(string.ascii_lowercase)[:X.shape[1]]
# explain the model's predictions using SHAP values
shap_values = est.shap_values(X[:100], feature_names=feature_names)
```

     99%|===================| 99/100 [00:51<00:00]

```python
# visualize the first prediction's explanation (use matplotlib=True to avoid Javascript)
shap.force_plot(shap_values["Y0"]["T0_1"][0], matplotlib=True)
```

```python
shap.summary_plot(shap_values["Y0"]["T0_1"])
```

# DRLearner

```python
model_regression = clone(first_stage_reg().fit(np.hstack([T.reshape(-1, 1), X]), Y).best_estimator_)
model_regression
```

    GradientBoostingRegressor(min_samples_leaf=30, n_estimators=50,
                              random_state=123)

```python
from econml.dr import ForestDRLearner
from sklearn.dummy import DummyRegressor, DummyClassifier

est = ForestDRLearner(model_regression=model_y,
                      model_propensity=model_t,
                      cv=3,
                      n_estimators=4000,
                      min_samples_leaf=10,
                      verbose=0,
                      min_weight_fraction_leaf=.005)
est.fit(Y, T, X=X)
```

    <econml.dr._drlearner.ForestDRLearner at 0x1fbb4a5d688>

```python
from econml.dr import DRLearner
est2 = DRLearner(model_regression=model_y,
                 model_propensity=model_t,
                 model_final=final_stage(),
                 cv=3)
est2.fit(Y, T.reshape((-1, 1)), X=X)
```

    <econml.dr._drlearner.DRLearner at 0x1fbb7fc1c88>

```python
X_test = np.zeros((100, p))
X_test[:, 0] = np.linspace(-2, 2, 100)
```

```python
pred = est.effect(X_test)
lb, ub = est.effect_interval(X_test, alpha=0.01)
```

```python
pred2 = est2.effect(X_test)
```

```python
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.plot(X_test[:, 0], true_effect(X_test), '--')
plt.plot(X_test[:, 0], pred2, label='nonparamdml')
plt.plot(X_test[:, 0], pred, label='forestdml (causal forest)')
plt.fill_between(X_test[:, 0], lb, ub, alpha=.4, label='honestrf_ci')
plt.legend()
plt.show()
```

### First stage nuisance models

```python
# Model T
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.title('honestrf')
for mdl in est.models_propensity:
    plt.plot(X_test[:, 0], mdl.predict_proba(X_test)[:, 1])
plt.plot(X_test[:, 0], true_propensity(X_test), '--', label='truth')
plt.legend()

plt.subplot(1, 2, 2)
plt.title('rf')
for mdl in est2.models_propensity:
    plt.plot(X_test[:, 0], mdl.predict_proba(X_test)[:, 1])
plt.plot(X_test[:, 0], true_propensity(X_test), '--', label='truth')
plt.legend()
plt.show()
```

```python
# Model Y
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.title('honestrf')
for mdl in est.models_regression:
    plt.plot(X_test[:, 0], mdl.predict(np.hstack([X_test, np.ones((X_test.shape[0], 1))])))
plt.plot(X_test[:, 0], true_effect(X_test) + true_conf(X_test), '--', label='truth')
plt.legend()

plt.subplot(1, 2, 2)
plt.title('rf')
for mdl in est2.models_regression:
    plt.plot(X_test[:, 0], mdl.predict(np.hstack([X_test, np.ones((X_test.shape[0], 1))])))
plt.plot(X_test[:, 0], true_effect(X_test) + true_conf(X_test), '--', label='truth')
plt.legend()
plt.show()
```

### Interpretability of CATE Model of DRLearner with SHAP

```python
# explain the model's predictions using SHAP values
shap_values = est.shap_values(X[:100], feature_names=feature_names)
```

     99%|===================| 99/100 [01:16<00:00]

```python
# visualize the first prediction's explanation (use matplotlib=True to avoid Javascript)
shap.force_plot(shap_values["Y0"]["T0_1"][0], matplotlib=True)
```

```python
shap.summary_plot(shap_values["Y0"]["T0_1"])
```

# OrthoForest

```python
from econml.orf import DROrthoForest
from sklearn.linear_model import Lasso, LassoCV, LogisticRegression, LogisticRegressionCV
from econml.sklearn_extensions.linear_model import WeightedLassoCV

est3 = DROrthoForest(model_Y=Lasso(alpha=0.01),
                     propensity_model=LogisticRegression(C=1),
                     model_Y_final=WeightedLassoCV(cv=3),
                     propensity_model_final=LogisticRegressionCV(cv=3),
                     n_trees=1000, min_leaf_size=10)
est3.fit(Y, T, X=X)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   58.4s
    [Parallel(n_jobs=-1)]: Done 112 tasks      | elapsed:  1.1min
    [Parallel(n_jobs=-1)]: Done 272 tasks      | elapsed:  1.3min
    [Parallel(n_jobs=-1)]: Done 496 tasks      | elapsed:  1.6min
    [Parallel(n_jobs=-1)]: Done 784 tasks      | elapsed:  2.0min
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:  2.3min finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    1.5s
    [Parallel(n_jobs=-1)]: Done 112 tasks      | elapsed:    8.5s
    [Parallel(n_jobs=-1)]: Done 272 tasks      | elapsed:   20.7s
    [Parallel(n_jobs=-1)]: Done 496 tasks      | elapsed:   35.0s
    [Parallel(n_jobs=-1)]: Done 784 tasks      | elapsed:   51.7s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:  1.1min finished

    <econml.orf._ortho_forest.DROrthoForest at 0x1fbb814ee48>

```python
pred3 = est3.effect(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   51.7s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:  1.1min finished

```python
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.plot(X_test[:, 0], true_effect(X_test), '--')
plt.plot(X_test[:, 0], pred, label='forestdr')
plt.plot(X_test[:, 0], pred2, label='nonparamdr')
plt.plot(X_test[:, 0], pred3, label='discreteorf')
plt.fill_between(X_test[:, 0], lb, ub, alpha=.4, label='forest_dr_ci')
plt.legend()
plt.show()
```


### http://nlp.seas.harvard.edu/2018/04/03/attention.html

```python
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import math, copy, time
from torch.autograd import Variable
import matplotlib.pyplot as plt
import seaborn
seaborn.set_context(context="talk")
%matplotlib inline
from IPython.display import Image
```

## Background

In sequential models, the number of operations required to relate signals from two arbitrary input or output positions grows in the distance between positions, linearly for ConvS2S and logarithmically for ByteNet. This makes it more difficult to learn dependencies between distant positions.

In the Transformer this is reduced to a constant number of operations, albeit at the cost of reduced effective resolution due to averaging attention-weighted positions, an effect we counteract with Multi-Head Attention.

Self-attention, sometimes called intra-attention is an attention mechanism relating different positions of a single sequence in order to compute a representation of the sequence. Self-attention has been used successfully in a variety of tasks including reading comprehension, abstractive summarization, textual entailment and learning task-independent sentence representations.

To the best of our knowledge, however, the Transformer is the first transduction model relying entirely on self-attention to compute representations of its input and output without using sequence aligned RNNs or convolution.

## Model Architecture

Most competitive neural sequence transduction models have an encoder-decoder structure (cite). Here, the encoder maps an input sequence of symbol representations (x_1, …, x_n) to a sequence of continuous representations z = (z_1, …, z_n). Given z, the decoder then generates an output sequence (y_1, …, y_m) of symbols one element at a time. At each step the model is auto-regressive (cite), consuming the previously generated symbols as additional input when generating the next.

```python
class EncoderDecoder(nn.Module):
    """
    A standard Encoder-Decoder architecture. Base for this and many
    other models.
    """
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator

    def forward(self, src, tgt, src_mask, tgt_mask):
        "Take in and process masked src and target sequences."
        return self.decode(self.encode(src, src_mask), src_mask,
                            tgt, tgt_mask)

    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)

    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
```

```python
class Generator(nn.Module):
    "Define standard linear + softmax generation step."
    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        return F.log_softmax(self.proj(x), dim=-1)
```

The Transformer follows this overall architecture using stacked self-attention and point-wise, fully connected layers for both the encoder and decoder, shown in the left and right halves of Figure 1, respectively.

```python
Image(filename='images/ModalNet-21.png')
```

## Encoder and Decoder Stacks

### Encoder
The encoder is composed of a stack of N = 6 identical layers.

```python
def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
```

```python
class Encoder(nn.Module):
    "Core encoder is a stack of N layers"
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        "Pass the input (and mask) through each layer in turn."
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)
```

We employ a residual connection (cite) around each of the two sub-layers, followed by layer normalization (cite).

```python
class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
```

That is, the output of each sub-layer is LayerNorm(x + Sublayer(x)), where Sublayer(x) is the function implemented by the sub-layer itself. We apply dropout (cite) to the output of each sub-layer, before it is added to the sub-layer input and normalized.

To facilitate these residual connections, all sub-layers in the model, as well as the embedding layers, produce outputs of dimension d_model = 512.

```python
class SublayerConnection(nn.Module):
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    """
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        "Apply residual connection to any sublayer with the same size."
        return x + self.dropout(sublayer(self.norm(x)))
```

Each layer has two sub-layers. The first is a multi-head self-attention mechanism, and the second is a simple, position-wise fully connected feed- forward network.

```python
class EncoderLayer(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size

    def forward(self, x, mask):
        "Follow Figure 1 (left) for connections."
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)
```

### Decoder
The decoder is also composed of a stack of N = 6 identical layers.

```python
class Decoder(nn.Module):
    "Generic N layer decoder with masking."
    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)
```

In addition to the two sub-layers in each encoder layer, the decoder inserts a third sub-layer, which performs multi-head attention over the output of the encoder stack. Similar to the encoder, we employ residual connections around each of the sub-layers, followed by layer normalization.

```python
class DecoderLayer(nn.Module):
    "Decoder is made of self-attn, src-attn, and feed forward (defined below)"
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, src_mask, tgt_mask):
        "Follow Figure 1 (right) for connections."
        m = memory
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        return self.sublayer[2](x, self.feed_forward)
```

We also modify the self-attention sub-layer in the decoder stack to prevent positions from attending to subsequent positions. This masking, combined with fact that the output embeddings are offset by one position, ensures that the predictions for position i can depend only on the known outputs at positions less than i.

```python
def subsequent_mask(size):
    "Mask out subsequent positions."
    attn_shape = (1, size, size)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    return torch.from_numpy(subsequent_mask) == 0
```

Below the attention mask shows the position each tgt word (row) is allowed to look at (column). Words are blocked for attending to future words during training.

```python
plt.figure(figsize=(5,5))
plt.imshow(subsequent_mask(20)[0])
None
```

### Attention
An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.

We call our particular attention “Scaled Dot-Product Attention”. The input consists of queries and keys of dimension d_k, and values of dimension d_v. We compute the dot products of the query with all keys, divide each by √d_k, and apply a softmax function to obtain the weights on the values.

```python
Image(filename='images/ModalNet-19.png')
```

In practice, we compute the attention function on a set of queries simultaneously, packed together into a matrix
Q. The keys and values are also packed together into matrices K and V. We compute the matrix of outputs as:

$\mathrm{Attention}(Q, K, V) = \mathrm{softmax}(\frac{QK^T}{\sqrt{d_k}})V$

```python
def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn
```

The two most commonly used attention functions are additive attention (cite), and dot-product (multiplicative) attention. Dot-product attention is identical to our algorithm, except for the scaling factor of $\frac{1}{\sqrt{d_k}}$
. Additive attention computes the compatibility function using a feed-forward network with a single hidden layer. While the two are similar in theoretical complexity, dot-product attention is much faster and more space-efficient in practice, since it can be implemented using highly optimized matrix multiplication code.

While for small values of ${d_k}$ the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of ${d_k}$(cite). We suspect that for large values of ${d_k}$, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients (To illustrate why the dot products get large, assume that the components of q and k are independent random variables with mean 0 and variance 1. Then their dot product, q⋅ k=∑d_k i=1 q_i k_i, has mean 0 and variance d_k.). To counteract this effect, we scale the dot products by $\frac{1}{\sqrt{d_k}}$.

```python
Image(filename='images/ModalNet-20.png')
```

Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this.

\begin{equation}
\mathrm{MultiHead}(Q, K, V) = \mathrm{Concat}(\mathrm{head_1}, ...,
\mathrm{head_h})W^O    \\
    \text{where}~\mathrm{head_i} = \mathrm{Attention}(QW^Q_i, KW^K_i, VW^V_i)
\end{equation}

Where the projections are parameter matrices $W^Q_i \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W^K_i \in
\mathbb{R}^{d_{\text{model}} \times d_k}$, $W^V_i \in \mathbb{R}^{d_{\text{model}} \times d_v}$ and $W^O \in \mathbb{R}^{hd_v \times d_{\text{model}}}$. In this work we employ h = 8 parallel attention layers, or heads. For each of these we use d_k=d_v=d_model/h=64. Due to the reduced dimension of each head, the total computational cost is similar to that of single-head attention with full dimensionality.

```python
class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k
        query, key, value = \
            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
             for l, x in zip(self.linears, (query, key, value))]

        # 2) Apply attention on all the projected vectors in batch.
        x, self.attn = attention(query, key, value, mask=mask,
                                 dropout=self.dropout)

        # 3) "Concat" using a view and apply a final linear.
        x = x.transpose(1, 2).contiguous() \
             .view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)
```

## Applications of Attention in our Model

The Transformer uses multi-head attention in three different ways: 1) In “encoder-decoder attention” layers, the queries come from the previous decoder layer, and the memory keys and values come from the output of the encoder. This allows every position in the decoder to attend over all positions in the input sequence. This mimics the typical encoder-decoder attention mechanisms in sequence-to-sequence models such as (cite).

2) The encoder contains self-attention layers. In a self-attention layer all of the keys, values and queries come from the same place, in this case, the output of the previous layer in the encoder. Each position in the encoder can attend to all positions in the previous layer of the encoder.

3) Similarly, self-attention layers in the decoder allow each position in the decoder to attend to all positions in the decoder up to and including that position. We need to prevent leftward information flow in the decoder to preserve the auto-regressive property. We implement this inside of scaled dot- product attention by masking out (setting to $-\infty$) all values in the input of the softmax which correspond to illegal connections.

## Position-wise Feed-Forward Networks

In addition to attention sub-layers, each of the layers in our encoder and decoder contains a fully connected feed-forward network, which is applied to each position separately and identically. This consists of two linear transformations with a ReLU activation in between.

$\mathrm{FFN}(x)=\max(0, xW_1 + b_1) W_2 + b_2$

While the linear transformations are the same across different positions, they use different parameters from layer to layer. Another way of describing this is as two convolutions with kernel size 1. The dimensionality of input and output is d_model=512, and the inner-layer has dimensionality d_ff=2048.

```python
class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))
```

## Embeddings and Softmax

Similarly to other sequence transduction models, we use learned embeddings to convert the input tokens and output tokens to vectors of dimension d_model. We also use the usual learned linear transformation and softmax function to convert the decoder output to predicted next-token probabilities. In our model, we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation, similar to (cite). In the embedding layers, we multiply those weights by
√d_model.

```python
class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)
```

## Positional Encoding

Since our model contains no recurrence and no convolution, in order for the model to make use of the order of the sequence, we must inject some information about the relative or absolute position of the tokens in the sequence. To this end, we add “positional encodings” to the input embeddings at the bottoms of the encoder and decoder stacks. The positional encodings have the same dimension
d_model as the embeddings, so that the two can be summed. There are many choices of positional encodings, learned and fixed (cite).

In this work, we use sine and cosine functions of different frequencies:$PE_{(pos,2i)} = sin(pos / 10000^{2i/d_{\text{model}}})$

$PE_{(pos,2i+1)} = cos(pos / 10000^{2i/d_{\text{model}}})$ where pos is the position and i is the dimension. That is, each dimension of the positional encoding corresponds to a sinusoid. The wavelengths form a geometric progression from  2π to 10000⋅2π. We chose this function because we hypothesized it would allow the model to easily learn to attend by relative positions, since for any fixed offset k, $PE_{pos+k}$ can be represented as a linear function of $PE_{pos}$.

In addition, we apply dropout to the sums of the embeddings and the positional encodings in both the encoder and decoder stacks. For the base model, we use a rate of $P_{drop}=0.1$.

```python
class PositionalEncoding(nn.Module):
    "Implement the PE function."
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + Variable(self.pe[:, :x.size(1)],
                         requires_grad=False)
        return self.dropout(x)
```

Below the positional encoding will add in a sine wave based on position. The frequency and offset of the wave is different for each dimension.

```python
plt.figure(figsize=(15, 5))
pe = PositionalEncoding(20, 0)
y = pe.forward(Variable(torch.zeros(1, 100, 20)))
plt.plot(np.arange(100), y[0, :, 4:8].data.numpy())
plt.legend(["dim %d"%p for p in [4,5,6,7]])
None
```

We also experimented with using learned positional embeddings (cite) instead, and found that the two versions produced nearly identical results. We chose the sinusoidal version because it may allow the model to extrapolate to sequence lengths longer than the ones encountered during training.

## Full Model

Here we define a function that takes in hyperparameters and produces a full model.

```python
def make_model(src_vocab, tgt_vocab, N=6,
               d_model=512, d_ff=2048, h=8, dropout=0.1):
    "Helper: Construct a model from hyperparameters."
    c = copy.deepcopy
    attn = MultiHeadedAttention(h, d_model)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    position = PositionalEncoding(d_model, dropout)
    model = EncoderDecoder(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn),
                             c(ff), dropout), N),
        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
        Generator(d_model, tgt_vocab))

    # This was important from their code.
    # Initialize parameters with Glorot / fan_avg.
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform(p)
    return model
```

```python
# Small example model.
tmp_model = make_model(10, 10, 2)
None
```

    <ipython-input-32-8c4144c2405f>:20: UserWarning: nn.init.xavier_uniform is now deprecated in favor of nn.init.xavier_uniform_.
      nn.init.xavier_uniform(p)

## Training

This section describes the training regime for our models.

We stop for a quick interlude to introduce some of the tools needed to train a standard encoder decoder model. First we define a batch object that holds the src and target sentences for training, as well as constructing the masks.

## Batches and Masking

```python
class Batch:
    "Object for holding a batch of data with mask during training."
    def __init__(self, src, trg=None, pad=0):
        self.src = src
        self.src_mask = (src != pad).unsqueeze(-2)
        if trg is not None:
            self.trg = trg[:, :-1]
            self.trg_y = trg[:, 1:]
            self.trg_mask = \
                self.make_std_mask(self.trg, pad)
            self.ntokens = (self.trg_y != pad).data.sum()

    @staticmethod
    def make_std_mask(tgt, pad):
        "Create a mask to hide padding and future words."
        tgt_mask = (tgt != pad).unsqueeze(-2)
        tgt_mask = tgt_mask & Variable(
            subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data))
        return tgt_mask
```

Next we create a generic training and scoring function to keep track of loss. We pass in a generic loss compute function that also handles parameter updates.

## Training Loop

```python
def run_epoch(data_iter, model, loss_compute):
    "Standard Training and Logging Function"
    start = time.time()
    total_tokens = 0
    total_loss = 0
    tokens = 0
    for i, batch in enumerate(data_iter):
        out = model.forward(batch.src, batch.trg,
                            batch.src_mask, batch.trg_mask)
        loss = loss_compute(out, batch.trg_y, batch.ntokens)
        total_loss += loss
        total_tokens += batch.ntokens
        tokens += batch.ntokens
        if i % 50 == 1:
            elapsed = time.time() - start
            print("Epoch Step: %d Loss: %f Tokens per Sec: %f" %
                    (i, loss / batch.ntokens, tokens / elapsed))
            start = time.time()
            tokens = 0
    return total_loss / total_tokens
```

## Training Data and Batching

We trained on the standard WMT 2014 English-German dataset consisting of about 4.5 million sentence pairs. Sentences were encoded using byte-pair encoding, which has a shared source-target vocabulary of about 37000 tokens. For English- French, we used the significantly larger WMT 2014 English-French dataset consisting of 36M sentences and split tokens into a 32000 word-piece vocabulary.

Sentence pairs were batched together by approximate sequence length. Each training batch contained a set of sentence pairs containing approximately 25000 source tokens and 25000 target tokens.

We will use torch text for batching. This is discussed in more detail below. Here we create batches in a torchtext function that ensures our batch size padded to the maximum batchsize does not surpass a threshold (25000 if we have 8 gpus).

```python
global max_src_in_batch, max_tgt_in_batch
def batch_size_fn(new, count, sofar):
    "Keep augmenting batch and calculate total number of tokens + padding."
    global max_src_in_batch, max_tgt_in_batch
    if count == 1:
        max_src_in_batch = 0
        max_tgt_in_batch = 0
    max_src_in_batch = max(max_src_in_batch,  len(new.src))
    max_tgt_in_batch = max(max_tgt_in_batch,  len(new.trg) + 2)
    src_elements = count * max_src_in_batch
    tgt_elements = count * max_tgt_in_batch
    return max(src_elements, tgt_elements)
```


# Ice-albedo feedback and Snowball Earth

Welcome!  In this activity we will apply our simple energy balance model ("EBM") in Climlab to evaluate the climate impacts of changing ice cover.

We will study one of the most extreme cases of the ice-albedo feedback hinted at in the geologic record: Snowball Earth.

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

```python
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import climlab
```

```python
def state_plot(model, figsize=(12,5), show=True, ice_temp=True): ## define a standard plot for temperature and albedo
    """Plot the temperature and albedo at the current state of the model.
    Shade the current ice line in grey."""
    templimits = -30,35
    alimits = min(model.albedo)-0.05, max(model.albedo)+0.05
    latlimits = -90,90
    lat_ticks = np.arange(-90,90,30)

    Ts = np.array(model.Ts).flatten()
    if ice_temp:
    	Tf = float(model.param['Tf'])
    else:
    	Tf=0

    fig = plt.figure(figsize=figsize)

    ax1 = fig.add_subplot(1,2,1)
    ax1.plot(model.lat, Ts)
    ax1.set(xlim=latlimits, ylim=templimits,
           ylabel='Temperature [deg C]', xlabel='Latitude', xticks=lat_ticks)
    ax1.fill_between(model.lat, Ts, y2=Tf, where=Ts<Tf, color='LightGrey', alpha=0.5)
    ax1.grid()

    ax2 = fig.add_subplot(1,2,2)
    icerect1 = patches.Rectangle((latlimits[0], 0), width=model.icelat[0]-latlimits[0], height=1,
                            color='LightGrey', alpha=0.5)
    icerect2 = patches.Rectangle((model.icelat[1], 0), width=latlimits[1]-model.icelat[1], height=1,
                            color='LightGrey', alpha=0.5)
    ax2.add_patch(icerect1)
    ax2.add_patch(icerect2)
    ax2.plot(model.lat, model.albedo)
    ax2.set(xlim=latlimits, ylim=alimits,
           ylabel='Albedo', xlabel='Latitude', xticks=lat_ticks)
    ax2.grid()
```

## Contents

1. [Setting up an energy balance model](#section1)
2. [Ice advance and retreat in the EBM](#section2)
3. [Snowball Earth: onset](#section3)
4. [Escape from the Snowball](#section4)

We will use an energy balance model that is very similar to what we set up in Lab 3, with one exception: we account for variations in energy with latitude, and energy transport across latitude bands.

____________

## 1. Setting up an energy balance model
____________

As we have seen in lecture, Earth's energy budget is an important control on the climate we experience.  Radiation energy coming in from the sun can be reflected (in the atmosphere or from the surface) or absorbed.  The portion of incoming solar radiation $Q$ that is reflected, $Q_{reflected}$, is set by the **planetary albedo**, $\alpha$.

As we saw in lecture,

```python
Q = 341.3 # W/m2, the incoming solar radiation
Q_reflected = 101.9 # W/m2, the reflected shortwave radiation
alpha = Q_reflected/Q

print('The planetary albedo is {:.2f}'.format(alpha)) # make a nicely formatted print statement to 2 sig figs
```

    The planetary albedo is 0.30

The portion of incoming solar radiation that is not reflected is the **Absorbed Shortwave Radiation**,
\begin{equation}
ASR=Q-Q_{reflected}=Q (1-\alpha),
\end{equation}
controlled by the albedo $\alpha$.

The heat energy emitted to space at the top of the atmosphere is the **Outgoing Longwave Radiation**, $OLR$.  The total **energy budget** of the Earth system is the balance between energy going out ($OLR$) and coming in ($ASR$):

\begin{align}
\frac{dE}{dt} &= ASR - OLR \\
    &= Q (1-\alpha) - OLR,
\end{align}
where we see $\alpha$ is a key parameter.

This is the basis of the simple **energy balance model** we first set up in Climlab.  Today, we will use it to explore the ice-albedo feedback.

### Our EBM: latitude-dependent, annual mean

Here, we'll use an energy balance model that accounts for radiation fluxes that vary with latitude $\phi$, but averages out seasonal changes in an annual mean.

The equation the model will solve for us is below.
\begin{align*}
C(\phi) \frac{\partial T_s}{\partial t} = & ~(1-\alpha) ~ Q - \left( A + B~T_s \right) +  \\
& \frac{D}{\cos\phi } \frac{\partial }{\partial \phi} \left(\cos\phi  ~ \frac{\partial T_s}{\partial \phi} \right)
\end{align*}

**Questions:**

*1.1. With your lab partner, diagram the physical meaning of each term in the equation.*

*1.2. How do you expect the heat capacity, $C(\phi)$, to vary with latitude?  Why?*

1.1
* $\phi$ = latitude
* $\alpha$ = albedo
* $\frac{\partial T_s}{\partial t}$ = Change in surface temperature over time
* $\frac{\partial T_s}{\partial \phi}$ = Change in latitude
* Q = incoming solar radiation
* A = 210 W m^-2 emission at 0ºC
* B = 2 W m^-2 increase in emission per degree, related to net longwave climate feedback
* D = 0.6 W m^-2 thermal diffusivity of the climate system
1.2
We expect $C($\phi\$) to be higher in areas with a lot of water such as the equator.

____________

## 2. Interactive snow and ice line in the EBM
____________

### Temperature-dependent ice line

Let the surface albedo be larger wherever the temperature is below some threshold $T_f$:

$$ \alpha\left(\phi, T(\phi) \right) = \left\{\begin{array}{ccc}
\alpha_0 + \alpha_2 P_2(\sin\phi) & ~ & T(\phi) > T_f \\
a_i & ~ & T(\phi) \le T_f \\
\end{array} \right. $$

```python
#  for convenience, set up a dictionary with our reference parameters
param = {'A':210, 'B':2, 'a0':0.3, 'a2':0.078, 'ai':0.62, 'Tf':-10.}
model1 = climlab.EBM_annual(name='Annual EBM with ice line',
                            num_lat=180, D=0.55, **param )
print(model1)
```

    climlab Process of type <class 'climlab.model.ebm.EBM_annual'>.
    State variables and domain shapes:
      Ts: (180, 1)
    The subprocess tree:
    Annual EBM with ice line: <class 'climlab.model.ebm.EBM_annual'>
       LW: <class 'climlab.radiation.aplusbt.AplusBT'>
       insolation: <class 'climlab.radiation.insolation.AnnualMeanInsolation'>
       albedo: <class 'climlab.surface.albedo.StepFunctionAlbedo'>
          iceline: <class 'climlab.surface.albedo.Iceline'>
          warm_albedo: <class 'climlab.surface.albedo.P2Albedo'>
          cold_albedo: <class 'climlab.surface.albedo.ConstantAlbedo'>
       SW: <class 'climlab.radiation.absorbed_shorwave.SimpleAbsorbedShortwave'>
       diffusion: <class 'climlab.dynamics.meridional_heat_diffusion.MeridionalHeatDiffusion'>

Because we provided a parameter `ai` for the icy albedo, our model now contains several sub-processes contained within the process called `albedo`. Together these implement the step-function formula above.

The process called `iceline` simply looks for grid cells with temperature below $T_f$ and adjusts their albedo.

```python
print(model1.param)
```

    {'timestep': 350632.51200000005, 'S0': 1365.2, 's2': -0.48, 'A': 210, 'B': 2, 'D': 0.55, 'Tf': -10.0, 'water_depth': 10.0, 'a0': 0.3, 'a2': 0.078, 'ai': 0.62}

```python
model1.integrate_years(5)
```

    Integrating for 450 steps, 1826.2110000000002 days, or 5 years.
    Total elapsed time is 5.000000000000044 years.

```python
f = state_plot(model1)
```

Grey shading on the above plots indicates where there is ice in the model.  We can find the same information by querying the model `icelat` attribute:

```python
model1.icelat
```

    array([-70.,  70.])

### Sudden cooling

What happens if we force the model to be colder? Let's store our current model state, then make a new model like it that will decrease the temperature by 20 $^{\circ}$C everywhere.

**Exercise: figure out how to make a model clone and decrease the temperature by 20 $^{\circ}$C everywhere.**  Call it m2 so that the plotting cell will recognize it.

```python
## Assign current model diagnostics to separate variables
Tequil = np.array(model1.Ts)
ALBequil = np.array(model1.albedo)
OLRequil = np.array(model1.OLR)
ASRequil = np.array(model1.ASR)
```

```python
m2 = ... #your work here
```

```python
m2.compute_diagnostics()
f3 = state_plot(m2)
```

```python
f2 = state_plot(model1)
```

**Question:**

*2.1. Compare and contrast the global patterns of temperature and albedo in this new, colder climate versus the previous simulation.*

Let's look at the radiative effect - how does the absorbed shortwave radiation change with this colder climate?

```python
my_ticks = [-90,-60,-30,0,30,60,90]
lat = model1.lat

fig = plt.figure( figsize=(12,5) )

ax1 = fig.add_subplot(1,2,1)
ax1.plot(lat, Tequil, label='equil')
ax1.plot(lat, m2.state['Ts'], label='pert' )
ax1.grid()
ax1.legend()
ax1.set_xlim(-90,90)
ax1.set_xticks(my_ticks)
ax1.set_xlabel('Latitude')
ax1.set_ylabel('Temperature (degC)')

ax2 = fig.add_subplot(1,2,2)
ax2.plot( lat, ASRequil, label='equil')
ax2.plot( lat, m2.diagnostics['ASR'], label='pert' )
ax2.grid()
ax2.legend()
ax2.set_xlim(-90,90)
ax2.set_xticks(my_ticks)
ax2.set_xlabel('Latitude')
ax2.set_ylabel('ASR (W m$^{-2}$)')
```

This tells us that making the climate colder, and allowing the ice edge to advance, tends to decrease the absorbed shortwave radiation (ASR).  That is, ice advance due to cooling is a ***positive feedback*** that will tend to lead to more cooling and more ice advance.

**Question:**
*2.2. Repeat the comparison with a warmer climate.  What is responsible for the differences?*

____________

## 3. Snowball Earth onset
____________

In section 2 we forced the model to be cooler everywhere.  Now, we'll examine some real conditions of the past that could have forced a cooling.

The radiation coming in from the sun has not always been what it was today.  Millions of years ago, the Sun was less bright, and as a result there was less solar energy entering the Earth system.  Luckily, it is easy to use our ClimLab energy balance model to investigate a past climate with a weaker sun.

The model parameter `S0` is the solar constant, describing the flux of solar radiation.

```python
m3 = climlab.process_like(model1)
```

```python
m3.subprocess.insolation.S0
```

```python
m3.icelat
```

Let's decrease the solar constant to examine past conditions.

```python
m3.subprocess.insolation.S0 = 1300.
m3.integrate_years(100.)
m3.icelat
```

```python
f3 = state_plot(m3)
```

*3.1. Find a value or values of S0 that results in 100% ice cover (entire plot shaded).*

*3.2. Compare the value of S0 that you found with the inferred history of S0 (from the textbook or another source). Was it possible to reach Snowball Earth conditions under realistic historical values of S0?*

Is it possible to produce Snowball Earth with orbital parameters only?  Climlab can help us assess.  Modify the eccentricity, longitude of perihelion, and obliquity below to see if you can produce a Snowball Earth with realistic values.

```python
## alternative: orbital forcing
m3.subprocess.insolation.orb = {'ecc': 0.017236, 'long_peri': 281.37, 'obliquity': 22.9}
m3.integrate_years(100.)
m3.icelat
```

*What do you conclude from this experiment?*

____________

## 4. Escape from the Snowball
____________

We have read that there were other important factors in Earth's past climate.

**Discussion:**
With your neighbor and then in class, propose a geological mechanism that could have ended Snowball Earth.

____________

## Credits

This notebook was developed by [Lizz Ultee](http://ehultee.github.io), based in large part on the `ClimateModeling_courseware` resources of author of [Brian E. J. Rose](http://www.atmos.albany.edu/facstaff/brose/index.html).
____________

```python

```