# Double Machine Learning: Use Cases and Examples

Double Machine Learning (DML) is an algorithm that applies arbitrary machine learning methods
to fit the treatment and response, then uses a linear model to predict the response residuals
from the treatment residuals.

The EconML SDK implements the following DML classes:
* LinearDMLCateEstimator: suitable for estimating heterogeneous treatment effects.
* SparseLinearDMLCateEstimator: suitable for the case when $W$ is high dimensional vector and both the first stage and second stage estimate are linear.

In ths notebook, we show the performance of the DML on both synthetic data and observational data.

**Notebook contents:**

1. Example usage with single continuous treatment synthetic data
2. Example usage with single binary treatment synthetic data
3. Example usage with multiple continuous treatment synthetic data
4. Example usage with single continuous treatment observational data
5. Example usage with multiple continuous treatment, multiple outcome observational data

```python
import econml
```

```python
## Ignore warnings
import warnings
warnings.filterwarnings('ignore')
```

```python
# Main imports
from econml.dml import DMLCateEstimator, LinearDMLCateEstimator,SparseLinearDMLCateEstimator

# Helper imports
import numpy as np
from itertools import product
from sklearn.linear_model import Lasso, LassoCV, LogisticRegression, LogisticRegressionCV,LinearRegression,MultiTaskElasticNet,MultiTaskElasticNetCV
from sklearn.ensemble import RandomForestRegressor,RandomForestClassifier
from sklearn.preprocessing import PolynomialFeatures
import matplotlib.pyplot as plt
import matplotlib
from sklearn.model_selection import train_test_split

%matplotlib inline
```

## 1. Example Usage with Single Continuous Treatment Synthetic Data and Model Selection

### 1.1. DGP
We use the data generating process (DGP) from [here](https://arxiv.org/abs/1806.03467). The DGP is described by the following equations:

\begin{align}
T =& \langle W, \beta\rangle + \eta, & \;\eta \sim \text{Uniform}(-1, 1)\\
Y =& T\cdot \theta(X) + \langle W, \gamma\rangle + \epsilon, &\; \epsilon \sim \text{Uniform}(-1, 1)\\
W \sim& \text{Normal}(0,\, I_{n_w})\\
X \sim& \text{Uniform}(0,1)^{n_x}
\end{align}

where $W$ is a matrix of high-dimensional confounders and $\beta, \gamma$ have high sparsity.

For this DGP,
\begin{align}
\theta(x) = \exp(2\cdot x_1).
\end{align}

```python
# Treatment effect function
def exp_te(x):
    return np.exp(2*x[0])
```

```python
# DGP constants
np.random.seed(123)
n = 2000
n_w = 30
support_size = 5
n_x = 1
# Outcome support
support_Y = np.random.choice(np.arange(n_w), size=support_size, replace=False)
coefs_Y = np.random.uniform(0, 1, size=support_size)
epsilon_sample = lambda n: np.random.uniform(-1, 1, size=n)
# Treatment support
support_T = support_Y
coefs_T = np.random.uniform(0, 1, size=support_size)
eta_sample = lambda n: np.random.uniform(-1, 1, size=n)

# Generate controls, covariates, treatments and outcomes
W = np.random.normal(0, 1, size=(n, n_w))
X = np.random.uniform(0, 1, size=(n, n_x))
# Heterogeneous treatment effects
TE = np.array([exp_te(x_i) for x_i in X])
T = np.dot(W[:, support_T], coefs_T) + eta_sample(n)
Y = TE * T + np.dot(W[:, support_Y], coefs_Y) + epsilon_sample(n)

Y_train, Y_val, T_train, T_val, X_train, X_val, W_train, W_val = train_test_split(Y, T, X, W, test_size=.2)
# Generate test data
X_test = np.array(list(product(np.arange(0, 1, 0.01), repeat=n_x)))
```

### 1.2. Train Estimator
We train models in three different ways, and compare their performance.

#### 1.2.1. Default Setting

```python
est = LinearDMLCateEstimator(model_y=RandomForestRegressor(),
                             model_t=RandomForestRegressor(),
                             random_state=123)
est.fit(Y_train, T_train, X_train, W_train)
te_pred = est.effect(X_test)
```

#### 1.2.2. Polynomial Features for Heterogeneity

```python
est1 = SparseLinearDMLCateEstimator(model_y=RandomForestRegressor(),
                              model_t=RandomForestRegressor(),
                              featurizer=PolynomialFeatures(degree=3),
                              random_state=123)
est1.fit(Y_train, T_train, X_train, W_train)
te_pred1=est1.effect(X_test)
```

#### 1.2.3. Polynomial Features with regularization

```python
est2 = DMLCateEstimator(model_y=RandomForestRegressor(),
                        model_t=RandomForestRegressor(),
                        model_final=Lasso(alpha=0.1, fit_intercept=False),
                        featurizer=PolynomialFeatures(degree=10),
                        random_state=123)
est2.fit(Y_train, T_train, X_train, W_train)
te_pred2=est2.effect(X_test)
```

#### 1.2.4 Random Forest Final Stage

```python
from econml.dml import ForestDMLCateEstimator

# One can replace model_y and model_t with any scikit-learn regressor and classifier correspondingly
# as long as it accepts the sample_weight keyword argument at fit time.
est3 = ForestDMLCateEstimator(model_y=RandomForestRegressor(),
                              model_t=RandomForestRegressor(),
                              discrete_treatment=False,
                              n_estimators=1000,
                              subsample_fr=.8,
                              min_samples_leaf=10,
                              min_impurity_decrease=0.001,
                              verbose=0, min_weight_fraction_leaf=.01)
est3.fit(Y_train, T_train, X_train, W_train)
te_pred3 = est3.effect(X_test)
```

### 1.3. Performance Visualization

```python
plt.figure(figsize=(10,6))
plt.plot(X_test, te_pred, label='DML default')
plt.plot(X_test, te_pred1, label='DML polynomial degree=3')
plt.plot(X_test, te_pred2, label='DML polynomial degree=10 with Lasso')
plt.plot(X_test, te_pred3, label='ForestDML')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test, expected_te, 'b--', label='True effect')
plt.ylabel('Treatment Effect')
plt.xlabel('x')
plt.legend()
plt.show()
```

### 1.4. Model selection

For the three different models above, we can use score function to estimate the final model performance. The score is the MSE of the final stage Y residual, which can be seen as a proxy of the MSE of treatment effect.

```python
score={}
score["DML default"] = est.score(Y_val, T_val, X_val, W_val)
score["DML polynomial degree=2"] = est1.score(Y_val, T_val, X_val, W_val)
score["DML polynomial degree=10 with Lasso"] = est2.score(Y_val, T_val, X_val, W_val)
score["ForestDML"] = est3.score(Y_val, T_val, X_val, W_val)
score
```

    {'DML default': 1.8232573949078739,
     'DML polynomial degree=2': 1.7052571853668774,
     'DML polynomial degree=10 with Lasso': 2.1530239340849437,
     'ForestDML': 1.7657955605139215}

```python
print("best model selected by score: ",min(score,key=lambda x: score.get(x)))
```

    best model selected by score:  DML polynomial degree=2

```python
mse_te={}
mse_te["DML default"] = ((expected_te - te_pred)**2).mean()
mse_te["DML polynomial degree=2"] = ((expected_te - te_pred1)**2).mean()
mse_te["DML polynomial degree=10 with Lasso"] = ((expected_te - te_pred2)**2).mean()
mse_te["ForestDML"] = ((expected_te - te_pred3)**2).mean()
mse_te
```

    {'DML default': 0.35920075038929977,
     'DML polynomial degree=2': 0.21518984313847153,
     'DML polynomial degree=10 with Lasso': 0.195600010017266,
     'ForestDML': 0.13690181421352374}

```python
print("best model selected by MSE of TE: ", min(mse_te, key=lambda x: mse_te.get(x)))
```

    best model selected by MSE of TE:  ForestDML

## 2. Example Usage with Single Binary Treatment Synthetic Data and Confidence Intervals

### 2.1. DGP
We use the following DGP:

\begin{align}
T \sim & \text{Bernoulli}\left(f(W)\right), &\; f(W)=\sigma(\langle W, \beta\rangle + \eta), \;\eta \sim \text{Uniform}(-1, 1)\\
Y = & T\cdot \theta(X) + \langle W, \gamma\rangle + \epsilon, & \; \epsilon \sim \text{Uniform}(-1, 1)\\
W \sim & \text{Normal}(0,\, I_{n_w}) & \\
X \sim & \text{Uniform}(0,\, 1)^{n_x}
\end{align}

where $W$ is a matrix of high-dimensional confounders, $\beta, \gamma$ have high sparsity and $\sigma$ is the sigmoid function.

For this DGP,
\begin{align}
\theta(x) = \exp( 2\cdot x_1 ).
\end{align}

```python
# Treatment effect function
def exp_te(x):
    return np.exp(2 * x[0])# DGP constants

np.random.seed(123)
n = 1000
n_w = 30
support_size = 5
n_x = 4
# Outcome support
support_Y = np.random.choice(range(n_w), size=support_size, replace=False)
coefs_Y = np.random.uniform(0, 1, size=support_size)
epsilon_sample = lambda n:np.random.uniform(-1, 1, size=n)
# Treatment support
support_T = support_Y
coefs_T = np.random.uniform(0, 1, size=support_size)
eta_sample = lambda n: np.random.uniform(-1, 1, size=n)

# Generate controls, covariates, treatments and outcomes
W = np.random.normal(0, 1, size=(n, n_w))
X = np.random.uniform(0, 1, size=(n, n_x))
# Heterogeneous treatment effects
TE = np.array([exp_te(x_i) for x_i in X])
# Define treatment
log_odds = np.dot(W[:, support_T], coefs_T) + eta_sample(n)
T_sigmoid = 1/(1 + np.exp(-log_odds))
T = np.array([np.random.binomial(1, p) for p in T_sigmoid])
# Define the outcome
Y = TE * T + np.dot(W[:, support_Y], coefs_Y) + epsilon_sample(n)

# get testing data
X_test = np.random.uniform(0, 1, size=(n, n_x))
X_test[:, 0] = np.linspace(0, 1, n)
```

### 2.2. Train Estimator

```python
est = LinearDMLCateEstimator(model_y=RandomForestRegressor(),
                             model_t=RandomForestClassifier(min_samples_leaf=10),
                             discrete_treatment=True,
                             linear_first_stages=False,
                             n_splits=6)
est.fit(Y, T, X, W, inference='statsmodels')
te_pred = est.effect(X_test)
lb, ub = est.effect_interval(X_test, alpha=0.01)
```

```python
est2 = SparseLinearDMLCateEstimator(model_y=RandomForestRegressor(),
                                    model_t=RandomForestClassifier(min_samples_leaf=10),
                                    discrete_treatment=True,
                                    featurizer=PolynomialFeatures(degree=2),
                                    linear_first_stages=False,
                                    n_splits=6)
est2.fit(Y, T, X, W, inference='debiasedlasso')
te_pred2 = est2.effect(X_test)
lb2, ub2 = est2.effect_interval(X_test, alpha=0.01)
```

```python
est3 = ForestDMLCateEstimator(model_y=RandomForestRegressor(),
                              model_t=RandomForestClassifier(min_samples_leaf=10),
                              discrete_treatment=True,
                              n_estimators=1000,
                              subsample_fr=.8,
                              min_samples_leaf=10,
                              min_impurity_decrease=0.001,
                              verbose=0, min_weight_fraction_leaf=.01,
                              n_crossfit_splits=6)
est3.fit(Y, T, X, W, inference='blb')
te_pred3 = est3.effect(X_test)
lb3, ub3 = est3.effect_interval(X_test, alpha=0.01)
```

### 2.3. Performance Visualization

```python
expected_te=np.array([exp_te(x_i) for x_i in X_test])
plt.figure(figsize=(16,6))
plt.subplot(1, 3, 1)
plt.plot(X_test[:, 0], te_pred, label='LinearDML', alpha=.6)
plt.fill_between(X_test[:, 0], lb, ub, alpha=.4)
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.ylabel('Treatment Effect')
plt.xlabel('x')
plt.legend()
plt.subplot(1, 3, 2)
plt.plot(X_test[:, 0], te_pred2, label='SparseLinearDML', alpha=.6)
plt.fill_between(X_test[:, 0], lb2, ub2, alpha=.4)
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.ylabel('Treatment Effect')
plt.xlabel('x')
plt.legend()
plt.subplot(1, 3, 3)
plt.plot(X_test[:, 0], te_pred3, label='ForestDML', alpha=.6)
plt.fill_between(X_test[:, 0], lb3, ub3, alpha=.4)
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.ylabel('Treatment Effect')
plt.xlabel('x')
plt.legend()
plt.show()
```

### 2.4. Other Inferences
#### 2.4.1 Effect Inferences
Other than confidence interval, we could also output other statistical inferences of the effect include standard error, z-test score and p value given each sample $X[i]$.

```python
est.effect_inference(X_test[:10,]).summary_frame(alpha=0.1, value=0, decimals=3)
```

      point_estimate
      stderr
      zstat
      pvalue
      ci_lower
      ci_upper

      0
      0.381
      0.146
      2.605
      0.009
      0.140
      0.621

      1
      0.409
      0.163
      2.512
      0.012
      0.141
      0.677

      2
      0.308
      0.170
      1.811
      0.070
      0.028
      0.588

      3
      0.416
      0.150
      2.780
      0.005
      0.170
      0.662

      4
      0.562
      0.152
      3.700
      0.000
      0.312
      0.811

      5
      0.538
      0.139
      3.879
      0.000
      0.310
      0.766

      6
      0.436
      0.124
      3.531
      0.000
      0.233
      0.639

      7
      0.589
      0.154
      3.825
      0.000
      0.336
      0.842

      8
      0.446
      0.126
      3.530
      0.000
      0.238
      0.653

      9
      0.487
      0.162
      3.011
      0.003
      0.221
      0.753

We could also get the population inferences given sample $X$.

```python
est.effect_inference(X_test).population_summary(alpha=0.1, value=0, decimals=3, tol=0.001)
```

Uncertainty of Mean Point Estimate

  mean_point	 stderr_mean	  zstat	 pvalue	 ci_mean_lower	 ci_mean_upper

     3.369	      0.14	     24.056	   0.0	      3.138	         3.599

Distribution of Point Estimate

  std_point	 pct_point_lower	 pct_point_upper

    1.724	        0.691	           6.067

Total Variance of Point Estimate

  stderr_point	 ci_point_lower	 ci_point_upper

      1.73	          0.682	          6.065

Note: The stderr_mean is a conservative upper bound.

#### 2.4.2 Coefficient and Intercept Inferences
We could also get the coefficient and intercept inference for the final model when it's linear.

```python
est.coef__inference().summary_frame()
```

      point_estimate
      stderr
      zstat
      pvalue
      ci_lower
      ci_upper

      X0
      5.958
      0.227
      26.211
      0.000
      5.584
      6.332

      X1
      -0.058
      0.216
      -0.267
      0.789
      -0.413
      0.297

      X2
      -0.326
      0.218
      -1.491
      0.136
      -0.685
      0.034

      X3
      0.217
      0.211
      1.029
      0.304
      -0.130
      0.564

```python
est.intercept__inference().summary_frame()
```

      point_estimate
      stderr
      zstat
      pvalue
      ci_lower
      ci_upper

      intercept
      0.47
      0.218
      2.154
      0.031
      0.111
      0.829

```python
est.summary()
```

Coefficient Results

   	  point_estimate	 stderr	  zstat	 pvalue	 ci_lower	 ci_upper

  X0	      5.958	      0.227	 26.211	   0.0	    5.584	    6.332

  X1	     -0.058	      0.216	 -0.267	  0.789	  -0.413	    0.297

  X2	     -0.326	      0.218	 -1.491	  0.136	  -0.685	    0.034

  X3	      0.217	      0.211	  1.029	  0.304	   -0.13	    0.564

Intercept Results

      	      point_estimate	 stderr	 zstat	 pvalue	 ci_lower	 ci_upper

  intercept	      0.47	       0.218	 2.154	  0.031	   0.111	    0.829

## 3. Example Usage with Multiple Continuous Treatment Synthetic Data

### 3.1. DGP
We use the data generating process (DGP) from [here](https://arxiv.org/abs/1806.03467), and modify the treatment to generate multiple treatments. The DGP is described by the following equations:

\begin{align}
T =& \langle W, \beta\rangle + \eta, & \;\eta \sim \text{Uniform}(-1, 1)\\
Y =& T\cdot \theta_{1}(X) + T^{2}\cdot \theta_{2}(X) + \langle W, \gamma\rangle + \epsilon, &\; \epsilon \sim \text{Uniform}(-1, 1)\\
W \sim& \text{Normal}(0,\, I_{n_w})\\
X \sim& \text{Uniform}(0,1)^{n_x}
\end{align}

where $W$ is a matrix of high-dimensional confounders and $\beta, \gamma$ have high sparsity.

For this DGP,
\begin{align}
\theta_{1}(x) = \exp(2\cdot x_1)\\
\theta_{2}(x) = x_1^{2}\\
\end{align}

```python
# DGP constants
np.random.seed(123)
n = 6000
n_w = 30
support_size = 5
n_x = 5
# Outcome support
support_Y = np.random.choice(np.arange(n_w), size=support_size, replace=False)
coefs_Y = np.random.uniform(0, 1, size=support_size)
epsilon_sample = lambda n: np.random.uniform(-1, 1, size=n)
# Treatment support
support_T = support_Y
coefs_T = np.random.uniform(0, 1, size=support_size)
eta_sample = lambda n: np.random.uniform(-1, 1, size=n)

# Generate controls, covariates, treatments and outcomes
W = np.random.normal(0, 1, size=(n, n_w))
X = np.random.uniform(0, 1, size=(n, n_x))
# Heterogeneous treatment effects
TE1 = np.array([x_i[0] for x_i in X])
TE2 = np.array([x_i[0]**2 for x_i in X]).flatten()
T = np.dot(W[:, support_T], coefs_T) + eta_sample(n)
Y = TE1 * T + TE2 * T**2 + np.dot(W[:, support_Y], coefs_Y) + epsilon_sample(n)
# Generate test data
X_test = np.random.uniform(0, 1, size=(100, n_x))
X_test[:, 0] = np.linspace(0, 1, 100)
```

### 3.2. Train Estimator

```python
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.multioutput import MultiOutputRegressor
from sklearn.linear_model import ElasticNetCV
est = LinearDMLCateEstimator(model_y=GradientBoostingRegressor(n_estimators=100, max_depth=3, min_samples_leaf=20),
                             model_t=MultiOutputRegressor(GradientBoostingRegressor(n_estimators=100,
                                                                                max_depth=3,
                                                                                min_samples_leaf=20)),
                             featurizer=PolynomialFeatures(degree=2, include_bias=False),
                             linear_first_stages=False,
                             n_splits=5)
```

```python
T = T.reshape(-1,1)
est.fit(Y, np.concatenate((T, T**2), axis=1), X, W, inference='statsmodels')
```

    <econml.dml.LinearDMLCateEstimator at 0x1da9d64bf08>

```python
te_pred = est.const_marginal_effect(X_test)
```

```python
lb, ub = est.const_marginal_effect_interval(X_test, alpha=0.01)
```

### 3.3. Performance Visualization

```python
plt.figure(figsize=(10,6))
plt.plot(X_test[:, 0], te_pred[:, 0], label='DML estimate1')
plt.fill_between(X_test[:, 0], lb[:, 0], ub[:, 0], alpha=.4)
plt.plot(X_test[:, 0], te_pred[:, 1], label='DML estimate2')
plt.fill_between(X_test[:, 0], lb[:, 1], ub[:, 1], alpha=.4)
expected_te1 = np.array([x_i[0] for x_i in X_test])
expected_te2=np.array([x_i[0]**2 for x_i in X_test]).flatten()
plt.plot(X_test[:, 0], expected_te1, '--', label='True effect1')
plt.plot(X_test[:, 0], expected_te2, '--', label='True effect2')
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.show()
```

## 4. Example Usage with Single Continuous Treatment Observational Data

We applied our technique to Dominick’s dataset, a popular historical dataset of store-level orange juice prices and sales provided by University of Chicago Booth School of Business.

The dataset is comprised of a large number of covariates $W$, but researchers might only be interested in learning the elasticity of demand as a function of a few variables $x$ such
as income or education.

We applied the `LinearDMLCateEstimator` to estimate orange juice price elasticity
as a function of income, and our results, unveil the natural phenomenon that lower income consumers are more price-sensitive.

### 4.1. Data

```python
# A few more imports
import os
import pandas as pd
import urllib.request
from sklearn.preprocessing import StandardScaler
```

```python
# Import the data
file_name = "oj_large.csv"

if not os.path.isfile(file_name):
    print("Downloading file (this might take a few seconds)...")
    urllib.request.urlretrieve("https://msalicedatapublic.blob.core.windows.net/datasets/OrangeJuice/oj_large.csv", file_name)
oj_data = pd.read_csv(file_name)
```

```python
oj_data.head()
```

      store
      brand
      week
      logmove
      feat
      price
      AGE60
      EDUC
      ETHNIC
      INCOME
      HHLARGE
      WORKWOM
      HVAL150
      SSTRDIST
      SSTRVOL
      CPDIST5
      CPWVOL5

      0
      2
      tropicana
      40
      9.018695
      0
      3.87
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927

      1
      2
      tropicana
      46
      8.723231
      0
      3.87
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927

      2
      2
      tropicana
      47
      8.253228
      0
      3.87
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927

      3
      2
      tropicana
      48
      8.987197
      0
      3.87
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927

      4
      2
      tropicana
      50
      9.093357
      0
      3.87
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927

```python
# Prepare data
Y = oj_data['logmove'].values
T = np.log(oj_data["price"]).values
scaler = StandardScaler()
W1 = scaler.fit_transform(oj_data[[c for c in oj_data.columns if c not in ['price', 'logmove', 'brand', 'week', 'store','INCOME']]].values)
W2 = pd.get_dummies(oj_data[['brand']]).values
W = np.concatenate([W1, W2], axis=1)
X=scaler.fit_transform(oj_data[['INCOME']].values)
```

```python
## Generate test data
min_income = -1
max_income = 1
delta = (1 - (-1)) / 100
X_test = np.arange(min_income, max_income + delta - 0.001, delta).reshape(-1,1)
```

### 4.2. Train Estimator

```python
est = LinearDMLCateEstimator(model_y=RandomForestRegressor(),model_t=RandomForestRegressor())
est.fit(Y, T, X, W)
te_pred=est.effect(X_test)
```

### 4.3. Performance Visualization

```python
# Plot Oranje Juice elasticity as a function of income
plt.figure(figsize=(10,6))
plt.plot(X_test, te_pred, label="OJ Elasticity")
plt.xlabel(r'Scale(Income)')
plt.ylabel('Orange Juice Elasticity')
plt.legend()
plt.title("Orange Juice Elasticity vs Income")
plt.show()
```

### 4.4. Confidence Intervals

We can also get confidence intervals around our predictions by passing an additional `inference` argument to `fit`.  All estimators support bootstrap intervals, which involves refitting the same estimator repeatedly on subsamples of the original data, but `LinearDMLCateEstimator` also supports a more efficient approach which can be achieved by passing `'statsmodels'` as the inference argument.

```python
est.fit(Y, T, X, W, inference='statsmodels')
te_pred=est.effect(X_test)
te_pred_interval = est.const_marginal_effect_interval(X_test, alpha=0.02)
```

```python
# Plot Oranje Juice elasticity as a function of income
plt.figure(figsize=(10,6))
plt.plot(X_test.flatten(), te_pred, label="OJ Elasticity")
plt.fill_between(X_test.flatten(), te_pred_interval[0], te_pred_interval[1], alpha=.5, label="1-99% CI")
plt.xlabel(r'Scale(Income)')
plt.ylabel('Orange Juice Elasticity')
plt.title("Orange Juice Elasticity vs Income")
plt.legend()
plt.show()
```

## 5. Example Usage with Multiple Continuous Treatment, Multiple Outcome Observational Data

We use the same data, but in this case, we want to fit the demand of multiple brand as a function of the price of each one of them, i.e. fit the matrix of cross price elasticities. It can be done, by simply setting as $Y$ to be the vector of demands and $T$ to be the vector of prices. Then we can obtain the matrix of cross price elasticities.

\begin{align}
Y=[Logmove_{tropicana},Logmove_{minute.maid},Logmove_{dominicks}] \\
T=[Logprice_{tropicana},Logprice_{minute.maid},Logprice_{dominicks}] \\
\end{align}

### 5.1. Data

```python
# Import the data
oj_data = pd.read_csv(file_name)
```

```python
# Prepare data
oj_data['price'] = np.log(oj_data["price"])
# Transform dataset.
# For each store in each week, get a vector of logmove and a vector of logprice for each brand.
# Other features are store specific, will be the same for all brands.
groupbylist = ["store", "week", "AGE60", "EDUC", "ETHNIC", "INCOME",
               "HHLARGE", "WORKWOM", "HVAL150",
               "SSTRDIST", "SSTRVOL", "CPDIST5", "CPWVOL5"]
oj_data1 = pd.pivot_table(oj_data,index=groupbylist,
                          columns=oj_data.groupby(groupbylist).cumcount(),
                          values=['logmove', 'price'],
                          aggfunc='sum').reset_index()
oj_data1.columns = oj_data1.columns.map('{0[0]}{0[1]}'.format)
oj_data1 = oj_data1.rename(index=str,
                           columns={"logmove0": "logmove_T",
                                    "logmove1": "logmove_M",
                                    "logmove2":"logmove_D",
                                    "price0":"price_T",
                                    "price1":"price_M",
                                    "price2":"price_D"})

# Define Y,T,X,W
Y = oj_data1[['logmove_T', "logmove_M", "logmove_D"]].values
T = oj_data1[['price_T', "price_M", "price_D"]].values
scaler = StandardScaler()
W=scaler.fit_transform(oj_data1[[c for c in groupbylist if c not in ['week', 'store', 'INCOME']]].values)
X=scaler.fit_transform(oj_data1[['INCOME']].values)
```

```python
## Generate test data
min_income = -1
max_income = 1
delta = (1 - (-1)) / 100
X_test = np.arange(min_income, max_income + delta - 0.001, delta).reshape(-1, 1)
```

### 5.2. Train Estimator

```python
est = LinearDMLCateEstimator(model_y=MultiTaskElasticNetCV(cv=3, tol=1, selection='random'),
                             model_t=MultiTaskElasticNetCV(cv=3),
                             featurizer=PolynomialFeatures(1),
                             linear_first_stages=True)
est.fit(Y, T, X, W)
te_pred = est.const_marginal_effect(X_test)
```

### 5.3. Performance Visualization

```python
# Plot Oranje Juice elasticity as a function of income
plt.figure(figsize=(18, 10))
dic={0:"Tropicana", 1:"Minute.maid", 2:"Dominicks"}
for i in range(3):
    for j in range(3):
        plt.subplot(3, 3, 3 * i + j + 1)
        plt.plot(X_test, te_pred[:, i, j],
                 color="C{}".format(str(3 * i + j)),
                 label="OJ Elasticity {} to {}".format(dic[j], dic[i]))
        plt.xlabel(r'Scale(Income)')
        plt.ylabel('Orange Juice Elasticity')
        plt.legend()
plt.suptitle("Orange Juice Elasticity vs Income", fontsize=16)
plt.show()
```

**Findings**: Look at the diagonal of the matrix, the TE of OJ prices are always negative to the sales across all the brand, but people with higher income are less price-sensitive. By contrast, for the non-diagonal of the matrix, the TE of prices for other brands are always positive to the sales for that brand, the TE is affected by income in different ways for different competitors. In addition, compare to previous plot, the negative TE of OJ prices for each brand are all larger than the TE considering all brand together, which means we would have underestimated the effect of price changes on demand.

### 5.4. Confidence Intervals

```python
est.fit(Y, T, X, W, inference='statsmodels')
te_pred = est.const_marginal_effect(X_test)
te_pred_interval = est.const_marginal_effect_interval(X_test, alpha=0.02)
```

```python
# Plot Oranje Juice elasticity as a function of income
plt.figure(figsize=(18, 10))
dic={0:"Tropicana", 1:"Minute.maid", 2:"Dominicks"}
for i in range(3):
    for j in range(3):
        plt.subplot(3, 3, 3 * i + j + 1)
        plt.plot(X_test, te_pred[:, i, j],
                 color="C{}".format(str(3 * i + j)),
                 label="OJ Elasticity {} to {}".format(dic[j], dic[i]))
        plt.fill_between(X_test.flatten(), te_pred_interval[0][:, i, j],te_pred_interval[1][:, i,j], color="C{}".format(str(3*i+j)),alpha=.5, label="1-99% CI")
        plt.xlabel(r'Scale(Income)')
        plt.ylabel('Orange Juice Elasticity')
        plt.legend()
plt.suptitle("Orange Juice Elasticity vs Income",fontsize=16)
plt.show()
```

notebooks/Double Machine Learning Examples.ipynb


# ROC / AUC

Table of Contents

```
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from collections import defaultdict
from sklearn.model_selection import train_test_split
from sklearn.base import BaseEstimator
from sklearn.base import TransformerMixin
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import cross_val_score
from sklearn.pipeline import make_pipeline
from sklearn.compose import make_column_transformer

print(f'numpy: {np.__version__}')
print(f'pandas: {pd.__version__}')
```

    numpy: 1.17.5
    pandas: 0.25.3

## ROC/AUC for Binary Calssification

We will analyze a employee retention. Our goal is to find the employees that are likely to leave in the future.

```
hr_retention_url = 'https://raw.githubusercontent.com/martin-fabbri/colab-notebooks/master/data/model-selection/hr_retention.csv'
df = pd.read_csv(hr_retention_url)
df.info()
```

    <class 'pandas.core.frame.DataFrame'>
    RangeIndex: 12000 entries, 0 to 11999
    Data columns (total 7 columns):
    s          12000 non-null float64
    lpr        12000 non-null float64
    np         12000 non-null int64
    anh        12000 non-null int64
    tic        12000 non-null int64
    newborn    12000 non-null int64
    left       12000 non-null int64
    dtypes: float64(2), int64(5)
    memory usage: 656.4 KB

Variables definition:

- **s**: The satisfaction level on a scale of 0 to 1
- **lpe**: Last project evaluation by a client on a scale of 0 to 1
- **np**: Represents the number of projects worked on by employee in the last 12 month
- **anh**: Average number of hours worked in the last 12 month for that employee
- **tic**: Amount of time the employee spent in the company, measured in years
- **newborn**: This variable will take the value 1 if the employee had a newborn within the last 12 month and 0 otherwise
- **left**: 1 if the employee left the company, 0 if they're still working here. This is our response variable

```
df.head()
```

      s
      lpr
      np
      anh
      tic
      newborn
      left

      0
      0.38
      0.53
      2
      157
      3
      0
      1

      1
      0.80
      0.86
      5
      262
      6
      0
      1

      2
      0.11
      0.88
      7
      272
      4
      0
      1

      3
      0.72
      0.87
      5
      223
      5
      0
      1

      4
      0.37
      0.52
      2
      159
      3
      0
      1

```
X = df.drop(['left'], axis=1)
y = df['left']
print('labels distribution:', np.bincount(y) / y.size)

test_size = 0.2
random_state = 43
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size,
                                                    random_state=random_state,
                                                    stratify=y)
```

    labels distribution: [0.83333333 0.16666667]

### Preprocessing

We will perform generic preprocessing task including:

- Normalize numeric columns
- One-hot-encode categorical columns. The ***newborn*** variable is treated as a categorical variable.

```
column_trans = make_column_transformer(
    (StandardScaler(), ['s', 'lpr', 'np', 'anh', 'tic']),
    (OneHotEncoder(), ['newborn']),
    remainder='passthrough'
)
```

## Model Building

```
rf = RandomForestClassifier(max_depth = 4, random_state=43)
pipe = make_pipeline(column_trans, rf)

cross_val_score(pipe, X, y, cv=3, scoring='accuracy').mean()
```

    0.9473333333333334

After training our model, we need to evaluate whether its any good or not and the most straightforward and intuitive metric for a supervised classifier's performance is accuracy. Unfortunately, there are circumstances where simple accuracy does not work well. For example, with a disease that only affects 1 in a million people, a completely bogus screening test that always reports "negative" will be 99.9999% accurate. Unlike accuracy, ROC curves are less sensitive to class imbalance; the bogus screening test would have an AUC of 0.5, which is like not having a test at all.

>First column name | Second column name
>--- | ---
>Row 1, Col 1 | Row 1, Col 2
>Row 2, Col 1 | Row 2, Col 2

---

**ROC curve (Receiver Operating Characteristic)** is a commonly used way to visualize the performance of a binary classifier and AUC (Area Under the ROC Curve) is used to summarize its performance in a single number. Most machine learning algorithms have the ability to produce probability scores that tells us the strength in which it thinks a given observation is positive. Turning these probability scores into yes or no predictions requires setting a threshold; cases with scores above the threshold are classified as positive, and vice versa. Different threshold values can lead to different result:

- A higher threshold is more conservative about labeling a case as positive; this makes it less likely to produce false positive (an observation that has a negative label but gets classified as positive by the model) results but more likely to miss cases that are in fact positive (lower true positive rate)
- A lower threshold produces positive labels more liberally, so it creates more false positives but also generate more true positives

A quick refresher on terminology:

\begin{align}
[\text{true positive rate}]
&= \frac{[\text{# positive data points with positive predictions}]}{\text{[# all positive data points]}} \\
&= \frac{[\text{# true positives}]}{[\text{# true positives}] + [\text{# false negatives}]}
\end{align}

true positive rate is also known as **recall** or **sensitivity**

\begin{align}
[\text{false positive rate}]
&= \frac{[\text{# positive data points with positive predictions}]}{\text{[# all negative data points]}} \\
&= \frac{[\text{# false positives}]}{[\text{# false positives}] + [\text{# true negatives}]}
\end{align}

The ROC curve is created by plotting the true positive rate (when it's actually a yes, how often does it predict yes?) on the y axis against the false positive rate (when it's actually a no, how often does it predict yes?) on the x axis at various cutoff settings, giving us a picture of the whole spectrum of the trade-off we're making between the two measures.

If all these true/false positive terminology is confusing to you, consider reading the material at the following link. [Blog: Simple guide to confusion matrix terminology](http://www.dataschool.io/simple-guide-to-confusion-matrix-terminology/)

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```


Sascha Spors,
Professorship Signal Theory and Digital Signal Processing,
Institute of Communications Engineering (INT),
Faculty of Computer Science and Electrical Engineering (IEF),
University of Rostock,
Germany

# Tutorial Digital Signal Processing

**Random Signals, Ensemble & Temporal Averages**,
Winter Semester 2021/22 (Course #24505)

- lecture: https://github.com/spatialaudio/digital-signal-processing-lecture
- tutorial: https://github.com/spatialaudio/digital-signal-processing-exercises

Feel free to contact lecturer frank.schultz@uni-rostock.de

WIP...

```python
import numpy as np
import matplotlib as mpl
from matplotlib import pyplot as plt
from numpy.random import Generator, PCG64
from scipy import signal
from scipy import stats

def my_xcorr(x, y):
    N, M = len(x), len(y)
    kappa = np.arange(N+M-1) - (M-1)
    ccf = signal.correlate(x, y, mode='full', method='auto')
    return kappa, ccf

if True:  # test my_xcorr with simple example
    x = np.array([0, 1, 0, 0, 0])
    y = np.array([1, 0, 0])
    # plot my_xcorr(x, y) vs. my_xcorr(y, x)
    plt.figure(figsize=(9, 2))
    plt.subplot(1, 2, 1)
    kappa_xy, ccf_xy = my_xcorr(x, y)
    plt.stem(kappa_xy, ccf_xy, basefmt='C0:', use_line_collection=True)
    plt.xlabel(r'$\kappa$')
    plt.ylabel(r'$\varphi_{xy}[\kappa]$')
    plt.title('cross correlation between x and y')
    plt.grid(True)
    plt.subplot(1, 2, 2)
    kappa_yx, ccf_yx = my_xcorr(y, x)
    plt.stem(kappa_yx, ccf_yx, basefmt='C0:', use_line_collection=True)
    plt.xlabel(r'$\kappa$')
    plt.ylabel(r'$\varphi_{yx}[\kappa]$')
    plt.title('cross correlation between y and x')
    plt.grid(True)
```

# First-Order Ensemble Averages

For a probability density function (PDF) $p_x(\theta, k)$ that describes a random process of 'drawing' signal amplitudes $\theta$ for $n$-th sample function $x_n[k]$ over time $k$, we can define the following **expectation**

\begin{equation}
\begin{split}
E\{ f(x[k]) \} = \int\limits_{-\infty}^{\infty} f(\theta) \, p_x(\theta, k) \, \mathrm{d}\theta\\
E\{ f(x[k]) \} = \lim_{N \to \infty} \frac{1}{N} \sum_{n=0}^{N-1} f(x_n[k])
\end{split}
\end{equation}

using the operator or **mapping function** $f(\cdot)$.

Most important are the following **first-order** ensemble averages, also called **univariate** moments, named since **one** random process is involved.

#### Linear mean / 1st raw moment

for $f(\theta)=\theta^1$

\begin{equation}
\begin{split}
\mu_x[k] = E\{ x[k] \} = \int\limits_{-\infty}^{\infty} \theta \, p_x(\theta, k) \, \mathrm{d}\theta\\
\mu_x[k] = E\{ x[k] \} = \lim_{N \to \infty} \frac{1}{N} \sum_{n=0}^{N-1} x_n[k]
\end{split}
\end{equation}

#### Quadratic mean / 2nd raw moment

for $f(\theta)=\theta^2$

\begin{equation}
\begin{split}
E\{ x^2[k] \} = \int\limits_{-\infty}^{\infty} \theta^2 \, p_x(\theta, k) \, \mathrm{d}\theta\\
E\{ x^2[k] \} = \lim_{N \to \infty} \frac{1}{N} \sum_{n=0}^{N-1} x_n^2[k]
\end{split}
\end{equation}

#### Variance / 2nd centralized moment

for $f(\theta) = (\theta - \mu_x[k])^2$

\begin{equation}
\begin{split}
\sigma_x^2[k] = E\{ (x[k] - \mu_x[k])^2 \} = \int\limits_{-\infty}^{\infty} (\theta - \mu_x[k])^2 \, p_x(\theta, k) \, \mathrm{d}\theta\\
\sigma_x^2[k] = E\{ (x[k] - \mu_x[k])^2 \} = \lim_{N \to \infty} \frac{1}{N} \sum_{n=0}^{N-1} (x_n[k] - \mu_x[k])^2
\end{split}
\end{equation}

These three moments are generally linked as
\begin{equation}
E\{ x^2[k] \} = \mu_x^2[k] + \sigma_x^2[k].
\end{equation}

For **stationary processes** these ensemble averages are not longer time-dependent, but rather $\mu_x[k] = \mu_x = \mathrm{const}$, etc. holds.
This implies that the PDF is not changing over time.

# Second-Order Ensemble Averages

The **second-order** ensemble averages, also called **bivariate** moments (**two** random processes are involved) can be derived from

\begin{equation}
E\{ f(x[k_x], y[k_y]) \} = \iint\limits_{-\infty}^{\infty} f(\theta_x, \theta_y) \, p_{xy}(\theta_x, \theta_y, k_x, k_y) \, \mathrm{d}\theta_x\, \mathrm{d}\theta_y
\end{equation}

\begin{equation}
E\{ f(x[k_x], y[k_y]) \} = \lim_{N \to \infty} \frac{1}{N} \sum_{n=0}^{N-1} f(x_n[k_x], y_n[k_y])
\end{equation}

using appropriate **mapping functions** $f(\cdot)$.

For **stationary processes** only the difference $\kappa = k_x - k_y$ is relevant as bivariate PDF
\begin{equation}
p_{xy}(\theta_x, \theta_y, k_x, k_y) = p_{xy}(\theta_x, \theta_y, \kappa).
\end{equation}

For **stationary processes** two important cases lead to fundamental tools for random signal processing:

1. $\kappa = 0$, i.e. $k = k_x = k_y$
2. $\kappa \neq 0$

## Case 1

The general linear mappings

\begin{equation}
\begin{split}
\text{for raw (1,1)-bivariate moment:} \qquad & f(\theta_x, \theta_y)=\theta_x^1 \cdot \theta_y^1,\\
\text{for centralized (1,1)-bivariate moment:} \qquad & f(\theta_x, \theta_y) = (\theta_x - \mu_x[k_x])^1 \cdot (\theta_y - \mu_y[k_y])^1\\
\text{for standardized (1,1)-bivariate moment:} \qquad & f(\theta_x, \theta_y) = \left(\frac{\theta_x - \mu_x[k_x]}{\sigma_x[k_x]}\right)^1 \cdot \left(\frac{\theta_y - \mu_y[k_y]}{\sigma_y[k_y]}\right)^1.
\end{split}
\end{equation}

simplify under the assumption of stationary processes and considering $\kappa=0$, i.e. $k = k_x = k_y$.
The resulting expectations $E\{\cdot\}$ are

- the raw (1,1)-bivariate moment, known as cross-power $\mathrm{P}_{xy}$
- the centralized (1,1)-bivariate moment, known as co-variance $\sigma_{xy}$
- the standardized (1,1)-bivariate moment, known as correlation coefficient $\rho_{xy}$

## Case 2

For $\kappa = k_x - k_y \neq 0$ the raw and centralized moments are of special importance:

\begin{equation}
\begin{split}
\mathrm{raw:}\qquad&\varphi_{xy}[k_x, k_y] = \varphi_{xy}[\kappa] = E\{ x[k] \cdot y[k-\kappa]\} = E\{ x[k+\kappa] \cdot y[k]\}\\
\mathrm{centralized:}\qquad&\psi_{xy}[\kappa] = \varphi_{xy}[\kappa] - \mu_x \mu_y
\end{split}
\end{equation}

The raw moment is known as **cross-correlation** function $\varphi_{xy}[\kappa]$, the centralized moment is known as **cross-covariance** function $\psi_{xy}[\kappa]$.

If for the second process $y$ the process $x$ is considered as

\begin{equation}
\begin{split}
\mathrm{raw:}\qquad&\varphi_{xx}[\kappa] = E\{ x[k] \cdot x[k-\kappa]\} = E\{ x[k+\kappa] \cdot x[k]\}\\
\mathrm{centralized:}\qquad&\psi_{xx}[\kappa] = \varphi_{xx}[\kappa] - \mu^2_x
\end{split}
\end{equation}
the so called **auto-correlation** function $\varphi_{xx}[\kappa]$ and **auto-covariance** function $\psi_{xx}[\kappa]$ are obtained.

The auto- and cross-correlation functions are of fundamental importance for random signal processing, as these are linked to LTI systems.

# Ergodic Processes

Averaging over time is equal to ensemble averages:

\begin{equation}
\overline{ f(x_n[k], x_n[k-\kappa_1], x_n[k-\kappa_2], \dots) } = E\{  f(x[k], x[k-\kappa_1], x[k-\kappa_2], \dots)  \} \;\; \forall n.
\end{equation}

## Wide-Sense Ergodic

\begin{equation}
\overline{ x_n[k] \cdot x_n[k-\kappa] } = E\{  x[k] \cdot x[k-\kappa]  \} \;\; \forall n
\end{equation}

\begin{equation}
\overline{ x_n[k] } = E\{  x[k] \} \;\; \forall n.
\end{equation}

## Important Temporal Averages

The linear mean as temporal average of the $n$-th sample function $x_n[k]$ is for instance given by

\begin{equation}
\overline{x_n[k]} = \lim_{K \to \infty} \frac{1}{2K + 1} \sum_{k = -K}^{K} x_n[k].
\end{equation}

Furthermore, the quadratic mean from simple quadratic mapping is given as
\begin{equation}
\lim_{K \to \infty} \frac{1}{2K + 1} \sum_{k = -K}^{K} x^2_n[k],
\end{equation}

the variance is given as
\begin{equation}
\lim_{K \to \infty} \frac{1}{2K + 1} \sum_{k = -K}^{K} (x_n[k]-\overline{x_n[k]})^2,
\end{equation}

the cross-correlation as
\begin{equation}
\lim_{K \to \infty} \frac{1}{2K + 1} \sum_{k=-K}^{K} x[k] \cdot y[k-\kappa],
\end{equation}

and the auto- correlation as
\begin{equation}
\lim_{K \to \infty} \frac{1}{2K + 1} \sum_{k=-K}^{K} x[k] \cdot x[k-\kappa].
\end{equation}

These equations hold for power signals. Virtually all random signals are power signals rather than energy signals.

# Example: Histogram as PDF Estimate, First-Order Ensemble Averages

of Normal distribution process

```python
# set seed for reproducible results
seed = 1234
stats.norm.random_state = Generator(PCG64(seed))

# create random process based on normal distribution
Ns = 2**10  # number of sample functions for e.g. time instance k=0
loc, scale = 5, 3  # mu, sigma

theta = np.arange(-15, 25, 0.01)  # amplitudes for plotting PDF
# random process object with normal PDF
rv = stats.norm(loc=loc, scale=scale)
# get random data from sample functions
x = stats.norm.rvs(loc=loc, scale=scale, size=Ns)

# plot
fig, ax = plt.subplots(1, 1)
hist_estimate = ax.hist(x, bins='auto', density=True, histtype='bar',
                        color='C0', alpha=0.5, label='histogram')
ax.plot(theta, rv.pdf(theta), 'C0-', lw=2, label='pdf')
ax.set_xlabel(r'$\theta$')
ax.set_ylabel(r'$\hat{p}_x(\theta,k=0)$')
ax.set_title('normalized histogram = PDF estimate')
ax.set_xlim(-15, 25)
ax.legend()
ax.grid(True)

# get histogram data from ax.hist()
edges = hist_estimate[1]
freq = hist_estimate[0]

# simple ensemble averages by numeric integration
# over histogram data as a simple estimate of the pdf
theta_num = edges[:-1]
dtheta = np.diff(edges)
mu = np.sum(theta_num * freq * dtheta)  # mu estimate
qm = np.sum(theta_num**2 * freq * dtheta)  # quadratic mean estimate
sig2 = np.sum((theta_num-mu)**2 * freq * dtheta)  # sigma^2 estimate
print('ideal   ensemble average: mu = %5.2f, mu^2 = %5.2f, sigma^2 = %5.2f, mu^2 + sigma^2 = %5.2f' %
      (loc, loc**2, scale**2, loc**2+scale**2))
print('numeric ensemble average: mu = %5.2f, mu^2 = %5.2f, sigma^2 = %5.2f, mu^2 + sigma^2 = %5.2f' %
      (mu, mu**2, sig2, qm))
print('ideal sigma = %5.2f, numeric sigma = %5.2f' % (scale, np.sqrt(sig2)))

# We should think about:

# play around with Ns: what happens if you increase / decrease Ns in terms of
# the histogram plot and the estimated first-order ensemble averages

# play around with loc==mean, scale==sigma: how is the histogram and pdf
# changed, what tells us the standard deviation in terms of the area under the
# pdf

# ax.hist(x,...) is a handy tool for plotting and getting histogram data
# we have chosen bins='auto', density=True. Calculating these data is not
# trivial, at least if the histogram should represent the data in pdf-like
# form as here. So we should make sure that we are aware of the concepts for
# so called kernel density estimation
# Nice programming task would be manual histogram calc and plot for
# bins=100 and density=False, i.e. a classical manual histogram
```

# Example: Histogram of Gaussian Noise, Cosine and Rectangular Signal

here we use the numpy histogram with fixed number of bins and histogram mode rather than density mode

we don't strictly deal here with random sample functions, but with a amplitude values over time, but it nice to get an idea what a histogram looks like for known signals

```python
Ns = 10000  # number of sample function
Nt = 1  # number of time steps per sample function

# normal pdf
x = np.random.normal(loc=0, scale=1, size=[Ns, 1])
pdf, edges = np.histogram(x[:, 0], bins=200, density=False)
plt.plot(edges[:-1], pdf, 'o-', ms=5, label=r'gaussian PDF, $\sigma=1$')

# cosine signal with peak amplitude 1
x = np.cos(1 * 2*np.pi/Ns*np.arange(0, Ns))
pdf, edges = np.histogram(x, bins=200, density=False)
plt.plot(edges[:-1], pdf, 'o-', ms=5, label='cos')

# rect signal with amplitude 1.5
x = np.cos(1 * 2*np.pi/Ns*np.arange(0, Ns))
x[x >= 0] = +1.5
x[x < 0] = -1.5
pdf, edges = np.histogram(x, bins=200, density=False)
plt.plot(edges[:-1], pdf, 'o-', ms=5, label='rect')

plt.ylim(0, 500)
plt.xlabel(r'$\theta$')
plt.ylabel(r'histogram($\theta$)')
plt.title('frequency (Häufigkeit) of the signal amplitudes within 200 bins')
plt.legend()
plt.grid(True)

# We should think about:

# what happens if we apply a DC component to the cos and rect signal,
# e.g. x += 1

# we should able to predict the green histogram values exactly, how should
# plt.ylim(0, ???) altered to plot the histogram for the rect more nicely
```

#  Example: Higher-Order Ensemble Averages

```python
# create two random processes based on normal distribution
Ns = 2**10  # number of sample functions at certain time instant k
Nt = 1  # number of time steps per sample function
np.random.seed(1)

# 1st process:
locx, scalex = 1, 3
x = np.random.normal(loc=locx, scale=scalex, size=[Ns, Nt])

# 2nd process:
locy, scaley = 2, 4
y = np.random.normal(loc=locy, scale=scaley, size=[Ns, Nt])
```

```python
crosspower = np.mean(x * y)
covariance = np.mean((x-np.mean(x)) * (y-np.mean(y)))
rho = np.mean((x-np.mean(x))/np.std(x) * (y-np.mean(y))/np.std(y))
print('crosspower = %4.3f, covariance = %4.3f,  correlation coefficient rho = %4.3f' %
      (crosspower, covariance, rho))
```

# Ensemble Average vs. Temporal Average

```python
# create random process based on normal distribution
Ns = 4000  # number of sample functions
Nt = 16000  # number of time steps per sample function, 1s for 16kHz sampling
np.random.seed(1)

loc, scale = 5, 3  # mu, sigma
x = np.random.normal(loc=loc, scale=scale, size=[Ns, Nt])

plt.figure(figsize=(10, 3))
plt.subplot(1, 2, 1)
plt.plot(x[:, 0])
plt.xlabel('process index / sample function n')
plt.ylabel(r'x$_n$[k=0]')
plt.title('many sample functions for fixed time')

# check these three cases, the plots nicely show the concepts of ensemble
# averages and temporal average, make sure that you understand this with the
# cos-like patterns either over sample functions or over time
case_str = 'x'  # simulate an ergodic process, i.e. ensemble = temporal average
#case_str = 'cosn'  # simulate an non-stationary process by changing the mean
#case_str = 'cosk'

plt.subplot(1, 2, 2)
if case_str == 'x':  # use x directly
    plt.plot(x[0, :])
    plt.xlabel('time index k')
    plt.ylabel(r'x$_{n=0}$[k]')
    plt.title('one sample function over time')
elif case_str == 'cosn':  # add cosine over sample functions, fixed time
    tmp = 2*np.cos(2 * 2*np.pi/Ns * np.arange(0, Ns)) + 5
    x = x + np.transpose(np.tile(tmp, (Nt, 1)))
    plt.plot(x[:, 0])
    plt.xlabel('process index / sample function n')
    plt.ylabel(r'x$_n$[k=0]')
    plt.title('add cos() over sample function entries')
elif case_str == 'cosk':  # add cosine over time, fixed sample function
    tmp = 2*np.cos(2 * 2*np.pi/Nt * np.arange(0, Nt)) + 5
    x = x + np.tile(tmp, (Ns, 1))
    plt.plot(x[0, :])
    plt.xlabel('time index k')
    plt.ylabel(r'x$_{n=0}$[k]')
    plt.title('add cos() over time')

plt.figure(figsize=(10, 10))
plt.subplot(3, 1, 1)
plt.plot(np.mean(x, axis=0), label='ensemble average for fixed time instance')
plt.plot(np.mean(x, axis=1), label='temporal average for fixed sample function')
plt.plot([0, Nt], [loc, loc])
plt.legend()
plt.title(r'linear mean $E\{x\} = \mu$ of random process')
plt.grid(True)

plt.subplot(3, 1, 2)
plt.plot(np.var(x, axis=0))
plt.plot(np.var(x, axis=1))
plt.plot([0, Nt], [scale**2, scale**2])
plt.title('variance  $E\{(x -E\{x\} )^2\} = \sigma^2$ of random process')
plt.grid(True)

plt.subplot(3, 1, 3)
plt.plot(np.mean(x**2, axis=0))
plt.plot(np.mean(x**2, axis=1))
plt.plot([0, Nt], [loc**2+scale**2, loc**2+scale**2])
plt.title('quadratic mean $E\{x^2\} = \mu^2+\sigma^2$  of random process')
plt.grid(True)
```

# Higher-Order Temporal Averages

```python
# create two random processes based on normal distribution
Ns = 1  # number of sample functions at certain time instant k
Nt = 2**7  # number of time steps per sample function
np.random.seed(1)

# 1st process:
locx, scalex = 1, 3
x = np.random.normal(loc=locx, scale=scalex, size=[Ns, Nt])

# 2nd process:
locy, scaley = 2, 4
y = np.random.normal(loc=locy, scale=scaley, size=[Ns, Nt])
```

## Auto Correlation Function (ACF)

```python
plt.figure(figsize=(10, 3))
plt.subplot(1, 2, 1)
kappa, ccf = my_xcorr(x[0, :], x[0, :])
plt.stem(kappa, ccf, basefmt='C0:', use_line_collection=True)
plt.xlabel(r'$\kappa$')
plt.title(r'Auto Correlation Function $\varphi_{xx}[\kappa]$')
plt.grid(True)
plt.subplot(1, 2, 2)
kappa, ccf = my_xcorr(y[0, :], y[0, :])
plt.stem(kappa, ccf, basefmt='C0:', use_line_collection=True)
plt.xlabel(r'$\kappa$')
plt.title(r'Auto Correlation Function $\varphi_{yy}[\kappa]$')
plt.grid(True)

# check the axial symmetry, why is the peak always at kappa=0
```

## Cross Correlation Function (CCF)

```python
plt.figure(figsize=(10, 3))
plt.subplot(1, 2, 1)
kappa, ccf = my_xcorr(x[0, :], y[0, :])
plt.stem(kappa, ccf, basefmt='C0:', use_line_collection=True)
plt.xlabel(r'$\kappa$')
plt.title(r'Cross Correlation Function $\varphi_{xy}[\kappa]=\phi_{yx}[-\kappa]$')
plt.grid(True)
plt.subplot(1, 2, 2)
kappa, ccf = my_xcorr(y[0, :], x[0, :])
plt.stem(kappa, ccf, basefmt='C0:', use_line_collection=True)
plt.xlabel(r'$\kappa$')
plt.title(r'Cross Correlation Function $\varphi_{yx}[\kappa]=\phi_{xy}[-\kappa]$')
plt.grid(True)

# check the mirrored versions and how they related wrt kappa, x,y-sequence
```

##  Auto Covariance Function

typically no symbol in equations and no abbreviation

```python
plt.figure(figsize=(10, 3))
plt.subplot(1, 2, 1)
kappa, ccf = my_xcorr(x[0, :]-np.mean(x[0, :]), x[0, :]-np.mean(x[0, :]))
plt.stem(kappa, ccf, basefmt='C0:', use_line_collection=True)
plt.xlabel(r'$\kappa$')
plt.title(r'AutoCov$_{xx}[\kappa]$')
plt.grid(True)
plt.subplot(1, 2, 2)
kappa, ccf = my_xcorr(y[0, :]-np.mean(y[0, :]), y[0, :]-np.mean(y[0, :]))
plt.stem(kappa, ccf, basefmt='C0:', use_line_collection=True)
plt.xlabel(r'$\kappa$')
plt.title(r'AutoCov$_{yy}[\kappa]$')
plt.grid(True)

# check the axial symmetry, why is the peak always at kappa=0
```

## Cross Covariance Function

typically no symbol in equations and no abbreviation

```python
plt.figure(figsize=(10, 3))
plt.subplot(1, 2, 1)
kappa, ccf = my_xcorr(x[0, :]-np.mean(x[0, :]), y[0, :]-np.mean(y[0, :]))
plt.stem(kappa, ccf, basefmt='C0:', use_line_collection=True)
plt.xlabel(r'$\kappa$')
plt.title(r'CrossCov$_{xy}[\kappa]$=CrossCov$_{yx}[-\kappa]$')
plt.grid(True)
plt.subplot(1, 2, 2)
kappa, ccf = my_xcorr(y[0, :]-np.mean(y[0, :]), x[0, :]-np.mean(x[0, :]))
plt.stem(kappa, ccf, basefmt='C0:', use_line_collection=True)
plt.xlabel(r'$\kappa$')
plt.title(r'CrossCov$_{yx}[\kappa]$=CrossCov$_{xy}[-\kappa]$')
plt.grid(True)

# check the mirrored versions and how they related wrt kappa, x,y-sequence
```

# **Copyright**

The notebooks are provided as [Open Educational Resources](https://en.wikipedia.org/wiki/Open_educational_resources). Feel free to use the notebooks for your own purposes. The text is licensed under [Creative Commons Attribution 4.0](https://creativecommons.org/licenses/by/4.0/), the code of the IPython examples under the [MIT license](https://opensource.org/licenses/MIT). Please attribute the work as follows: *Frank Schultz, Digital Signal Processing - A Tutorial Featuring Computational Examples* with the URL https://github.com/spatialaudio/digital-signal-processing-exercises


# 第十八讲：线性二次调节

首先回顾一下前面关于马尔可夫决策过程的内容。

* 我们用五元组$\left(S,A,\{P_{sa}\},\gamma,R\right)$定义MDP，我介绍了$0\leq\gamma\lt1$、$R:S\to\mathbb R$等等……
* 定义了价值函数，并介绍了使用值迭代算法求价值函数$\displaystyle V(s):=R(s)+\max_a\gamma\sum_{s'\in S}P_{sa}\left(s'\right)V\left(s'\right)$；（经过有限次的迭代会靠近最优解。）
* 在值迭代后$V$收敛于$V^*$附近，再通过计算$\displaystyle\pi^*=\arg\max_a\sum_{s'\in S}P_{sa}\left(s'\right)V^*\left(s'\right)$得到最优策略$\pi^*$。

在[上一讲](chapter17.ipynb)中，我们了解了对无限状态MDP的求解，本讲我们将介绍一些做了轻微改动的MDP的变型形式——首先是修改奖励函数$R$，然后是对折扣因子$\gamma$的调整。记得在上一讲中，我们提到了无限状态MDP无法直接求解（我们使用了对价值函数的近似来求解）；而在本将中，我们将介绍一种特殊的MDP——即使状态空间是连续的（无限状态）依然可以直接对价值函数进行求解。

## 5. 状态-动作奖励（State-action reward）

我们要介绍的第一个MDP的变型形式就是状态动作奖励，我们将奖励函数改为$R:S,A\to \mathbb R$，也就是从一对“状态-动作”到实数的映射。此时我们就得到了一个状态-动作序列：

$$
\require{AMScd}
\begin{CD}
s_0,a_0 @>>> s_1,a_1 @>>> s_2,a_2 \cdots
\end{CD}\\
$$

根据这个序列，我们的总收益就应该是$R(s_0,a_0)+\gamma R(s_1,a_1)+\gamma^2R(s_2,a_2)+\cdots$，而我们的目标依然是找到一个能够最大化这个总收益的策略。可以看出这里的奖励函数变为一个关于状态和动作的函数（其实我们也可以减少一种变量，使其变回标准的MDP中的只关于状态的函数），使用状态-动作奖励可以使我们更直接的对“在不同动作中有不同成本”的问题进行建模（比如说对一个机器人来说，让机器人动起来的成本会高于让机器人静止的成本，此时给机器人一个静止在原地的指令就会对应更低的成本；或者比如对一个越野车来说，在泥泞道路上行驶比在柏油路上行驶成本更高，此时“在泥泞道路上行驶”的指令就更“贵”）。

这个变动对于MDP来说并不大，我们依然使用贝尔曼方程求解$V^*$：

$$
V^*(s)=\max_a\left(R(s,a)+\gamma\sum_{s'}P_{sa}\left(s'\right)V^*\left(s'\right)\right)
$$

在以前的价值函数中，$\displaystyle\max_a$的目标是未来收益，而现在因为“状态”奖励函数变为了“状态-动作”奖励函数，所以$\displaystyle\max_a$运算符移到了最外面。现在从状态$s$开始求最优价值函数时，最大化的目标分为两个部分：第一部分是新的即时奖励——在状态$s$下执行动作$a$的回报；第二部分是折扣因子$\gamma$乘以未来的预期总收益。这两个部分合起来就是$V^*$——最优价值函数——也就是取等号右侧的所有动作的最优值。

我们再来看值迭代，基本还是原来的算法$\displaystyle V(s):=\max_a\left(R(s,a)+\gamma\sum_{s'}P_{sa}\left(s'\right)V\left(s'\right)\right)$（仍然是用贝尔曼方程更新$V(s)$），迭代直到$V(s)$（近似）收敛于$V^*(s)$。

最后一步就是通过$V^*(s)$求最佳策略$\pi^*(s)$了，基本和以前一样$\displaystyle\pi^*(s)=\arg\max_a\left(R(s,a)+\gamma\sum_{s'}P_{sa}\left(s'\right)V^*\left(s'\right)\right)$

## 6. 有限时域马尔可夫决策过程（Finite horizon MDP）

我们要介绍的第二个MDP的变型形式就是有限时域MDP，这个MDP的五元组为$(S,A,\{P_{sa}\},T,R)$，其中$T$称为时域，它限定了MDP的最大步数，也就是在$T$步之后采取的动作将不会得到奖励。

$$
\require{AMScd}
\begin{CD}
s_0,a_0 @>>> s_1,a_1 @>>> \cdots @>>> s_T,a_T
\end{CD}\\
$$

那么在这种情形下，预期总收益就变成了$R(s_0,a_0)+R(s_1,a_1)+\cdots+R(s_T,a_T)$。与以往的目标相同，我们仍需要找到一个策略来最大化这个预期总收益，唯一的不同就是在$T$步之后MDP就结束了。这也说明策略有可能是**非平稳的（non-stationary）**（“平稳的”是指该项不依赖时间相关的变量），也就是说我们对最优动作的选择会根据所处时间的不同而不同。举个例子：
$
\begin{array}
{|c|c|c|c|c|}
\hline
+1&\cdot&s&\cdots&+10\\
\hline
\end{array}
$，假设机器人现在位于$s$的位置，距离左边的$+1$奖励比较近，距离右边的$+10$奖励比较远，而现在的MDP有个限制条件——不能超过$T$步。则这个例子中的机器人在做关于下一步动作的决策时会考虑：如果MDP刚开始，离结束$T$还比较远，那么它就可能决定向右走去取的$+10$的奖励；而如果已经进行到MDP末尾，剩下的步数并不一定能够到达$+10$奖励，那么它就可能决定向左走去取$+1$的奖励。

再来看非平稳状态转换概率$P_{sa}^{(t)}$，以前的MDP中有$s_{t+1}\sim P_{sa}$，即通过前一个状态的$s,a$就可以确定这个状态的分布，$P_{sa}$不会随时间变化（不依赖于时间相关的变量）。而在有限时域MDP中，不同的时间下状态转换概率是不同的，也就是说决定状态$s_{t+1}$的不只是$s,a$，同时也取决于它处在什么时间$t$。

再举几个例子：

* 在对于飞行器的建模中，飞行器会随着飞行消耗燃料，于是飞行器的总质量会随着时间的推移而变轻。于是对于飞行器的模型，我们对下一步所做的决定不只是取决于当前状态或动作，同时也取决于烧了多少燃料——即时间；
* 在飞机航线规划时，天气会随着时间的推移而不断变化，于是在好天气下某条航线的飞行成本与在恶劣天气下飞行成本是相差很大的，所以此模型同样收到时间因素的影响；
* 在汽车导航规划路线时，同一个区域在不同的时间车流量可能会有很大变化，因此某些区域的形式成本与时间有着直接关系。

最后，奖励函数也可能会随着时间改变，所以预期总收益会变为$R^{(0)}(s_0,a_0)+R^{(1)}(s_1,a_1)+\cdots+R^{(T)}(s_T,a_T)$。所以现在我们需要一个非平稳的随时间调整的策略。

最优价值函数$V_t^*(s)=\mathrm E\left[R^{(t)}(s_t,a_t)+\cdots+R^{(T)}(s_T,a_T)\mid\pi^*,s_t=s\right]$表示在状态$s$及时刻$t$下启动模型一直到结束步骤$T$期间的预期总收益，所以这个最优价值函数也取决于当前的时间以及所剩的时间。我们现在用贝尔曼方程写出最优价值函数的值迭代算法：

$$
\begin{align}
V_T^*(s)&=\max_aR^{(T)}(s,a)\tag{1}\\
V_t^*(s)&=\max_a\left(R^{(t)}(s,a)+\sum_{s'}P_{sa}^{(t)}\left(s'\right)V_{t+1}^*\left(s'\right)\right)\tag{2}
\end{align}
$$

$(2)$式表示，如果从时刻$t$状态$s$开始计算预期总收益，则应为“在$s,a$的即时奖励”加上“未来奖励”（当执行动作$a$后，MDP按照状态转换概率$P_{sa}^{(t)}$进入下一个状态$s'$，那么当MDP进入下一个状态$s'$时间自然应该累加，这时的预期总收益就应该是$V_{t+1}^*\left(s'\right)$，可以参考[第十六讲](chapter16.ipynb)由价值函数的递归定义得出贝尔曼方程的过程）的最大值，这个式子将$V_t$表示为$V_{t+1}$的递推式。而$(1)$式就是这个递推算法的启动项，它表示在时刻$T$只能得到该时间点的即时奖励，之后MDP立刻结束，没有未来奖励。计算的时候我们从$(1)$式开始，得到MDP的最后一项，然后使用$(2)$式向后递推依次得到$V_{T-1},V_{T-2},\cdots,V_0$。

最后再计算最优策略即可：

$$
\pi_t^*(s)=\arg\max_a\left(R^{(t)}(s,a)+\sum_{s'}P_{sa}\left(s'\right)V_{t+1}^*\left(s'\right)\right)
$$

最优策略和上面的最优价值函数一样，使用向后递推依次计算每个时间点对应的策略。

（再提一下折扣因子$\gamma$，在有限时域MDP模型中并没有出现这个参数，因为这个参数的作用从计算角度来说与参数$T$的作用非常相似。在标准MDP模型中，$\gamma$保证了MDP在一定步骤后的“未来”中，执行动作的奖励趋近于令；而在有限时域MDP模型中，$T$直接限定了“未来”步骤的总数。所以，我们通常在模型中只挑选这两个参数中的一个。）

接下来我们将结合状态-动作奖励MDP和有限时域MDP，再加上某些合理的假设，得到一个新的MDP模型，它能够优雅、高效的对大型MDP进行求解。

## 7. 线性二次调节（LQR: Linear Quadratic Regulation）

我们想要把状态-动作奖励MDP和有限时域MDP的思想（动态规划算法）应用到连续状态、甚至是连续动作空间中。仍使用五元组$(S,A,\{P_{sa}\},T,R)$定义MDP，注意这里我们使用有限时序中的$T$而不是折扣因子$\gamma$，其中：
* 状态为$S\in\mathbb R^n$；
* 动作为$A\in\mathbb R^d$；
* 将下一步作为关于当前状态及动作的线性函数$s_{t+1}=A_ts_t+B_ta_t+w_t$，此处的$w_t\sim\mathcal N(0,\varSigma_w)$是一个期望为零、协方差为$\varSigma_w$的高斯噪音项。$A_t\in\mathbb R^{n\times n}, B_t\in\mathbb R^{n\times d}$。需要留意的是我们在这里复用了符号$A$，这里的$A_t$是一个系数矩阵，在后面出现的$A$基本上都表示这个矩阵，而不是上面的动作集合；（在后面的介绍中，我们会假设$A,B,w$都是已知项，而我们的任务是给MDP找到一个最优策略。另外，在后面的证明中，我们会看到噪音项$w_t$并不十分重要。）
* $T$就是时域的限制；
* 奖励函数定义为：
    $$
    \begin{align}
    R^{(t)}(s_t,a_t)&=-\left(s_t^TU_ts_t+a^TV_ta_t\right)\\
    U_t&\in\mathbb R^{n\times n},U_t\geq0\\
    V_t&\in\mathbb R^{d\times d},V_t\geq0\\
    \therefore&\ s_t^TU_ts_t\geq0,v_t^TV_ts_t\geq0\\
    \implies&R^{(t)}(s_t,a_t)\leq0
    \end{align}
    $$
    在上面的定义中，$U_t,V_t$都是半正定矩阵，根据半正定矩阵的性质（可以参考[线性代数第二十八讲](http://nbviewer.jupyter.org/github/zlotus/notes-linear-algebra/blob/master/chapter28.ipynb)）就有奖励函数始终是一个非正数，它只会有成本而并没有奖励。举个例子，比如我们想让直升机模型中$s_t\approx0$，那么可以取$U_t=I,V_t=I$，则$R(s_t,a_t)=-s_t^Ts_t-a_t^Ta_t=-\lVert s_t\rVert^2-\lVert a_t\rVert^2$，这意味着如果直升机“远离原状态”（保持稳定）或“被来回猛拉”（防止高频率的反复操作），则奖励函数将以二次增长来惩罚模型，当然我们也可以修改矩阵$U,V$对角线上的元素以调整不同状态分量的权重。

（在本节有$A=A_1=A_2=\cdots,B=B_1=B_2=\cdots$，我们将假设MDP是静止的。）

### 7.1 非线性函数的线性近似

我们先来讨论线性动态规划的情形，$s_{t+1}=As_t+Ba_t$，其中一种建模方法我们在前面提过，就是做多次试验：

$$
\require{AMScd}
\begin{CD}
s_0^{(1)} @>{a_0^{(1)}}>> s_1^{(1)} @>{a_1^{(1)}}>> s_2^{(1)} @>{a_2^{(1)}}>> \cdots @>{a_{T-1}^{(1)}}>> s_T^{(1)}
\end{CD}\\
\begin{CD}
s_0^{(2)} @>{a_0^{(2)}}>> s_1^{(2)} @>{a_1^{(2)}}>> s_2^{(2)} @>{a_2^{(2)}}>> \cdots @>{a_{T-1}^{(2)}}>> s_T^{(2)}
\end{CD}\\
\vdots\\
\begin{CD}
s_0^{(m)} @>{a_0^{(m)}}>> s_1^{(m)} @>{a_1^{(m)}}>> s_2^{(m)} @>{a_2^{(m)}}>> \cdots @>{a_{T-1}^{(m)}}>> s_T^{(m)}
\end{CD}
$$

我们尝试不同的动作，记录测试得出的的各种结果，并使用线性回归拟合参数：

$$
\arg\min_{A,B}\sum_{i=1}^m\sum_{t=0}^{T-1}\left\lVert s_{t+1}^{(i)}-\left(As_t^{(i)}+Ba_t^{(i)}\right)\right\rVert^2
$$

也就是监督学习中的最小化“A,B”对数据的预测的误差。这是一种对线性动态系统的合理建模。我们再来介绍一种线性模型的建模方法——求非线性模型再将其线性化（linearize a non-linear model），仍然用倒置钟摆的例子：

假设有某个非线性模型$s_{t+1}=f(s_t,a_t)\begin{pmatrix}x_{t+1}\\\dot x_{t+1}\\\theta_{t+1}\\\dot\theta_{t+1}\end{pmatrix}=f\left(\begin{matrix}x_{t}\\\dot x_{t}\\\theta_{t}\\\dot\theta_{t}\end{matrix},a\right)$（$f$是一个非线性函数），也就是说，下一个时间点的状态将是一个关于某些当前状态和当前动作的函数，在本例中的动作就是一个实数，表示小车向左或向右的加速度的大小。 如果将这个式子看做$s_{t+1}$关于$s_t$的函数，那么我们可以做一个横轴为$s_t$，纵轴为$s_{t+1}$的图像，并绘出这个非线性函数。于是，当要对$s_{t+1}$进行预测时，就在横轴上找一个点$\bar s_t$，并在$\bar s_t$求函数的导数/梯度，于是就可以做出在点$\left(\bar s_t,\bar s_{t+1}\right)$处与非线性函数曲线相切的直线，而这条直线就可以近似为非线性函数在该点附近的线性表示。用数学符号表示：

* 假设$s_{t+1}=f(s_t)$（这里为了简化计算先不考虑动作$a_t$），于是有$s_{t+1}\approx f'\left(\bar s_t\right)\left(s_t-\bar s_t\right)+f\left(\bar s_t\right)$，这就把原来的非线性函数写成了一个关于$s_t$的线性函数，而$\bar s_t$是一个常数。最后，就是在式子里对应出矩阵$A,B$了。

现在我们再考虑另一个问题——这条直线能够在多大程度上近似原来的非线性函数？很明显，在$\bar s_t$点附近的近似都可以做到足够好，但是离的越远就差距就可能越大，就会变成一个糟糕的近似。所以，当我们使用LQR线性化非线性函数时，其中一个重要的参数就是用来确定线性化后函数在什么区间上有足够好的近似效果，放在倒置钟摆的例子中就是，我们希望这个系统的绝大部分时间都落在在这个“足够好的近似区间”内，这就是一次“合理的近似”。因此，从经验上讲，我们通常根据“系统大部分时间所在的区间”来确定在哪个状态附近对这个非线性函数做线性化处理。在倒置钟摆系统下，我们希望系统总是处于“零状态”，也就是大部分时间位于轨道的中央附近、小车的速度大约为零、同时杆能够的与轨道近似垂直——这就是一个我们希望做线性化的区间。

写出通用的线性化方程：

$$
s_{t+1}\approx f(\bar s_t,\bar a_t)+(\nabla_sf(\bar s_t,\bar a_t))^T(s_t-\bar s_t)+(\nabla_af(\bar s_t,\bar a_t))^T(a_t-\bar a_t)
$$

选择合适的$\bar s_t,\bar a_t$（在式子里，这两个量是常数）附近的区间做线性化，就可以将$s_{t+1}$近似的表示为关于$s_t,a_t$的线性函数。最后，我们可以将其化为$s_{t+1}=As_t+Ba_t$的形式。

### 7.2 LQR的求解

有了有线性近似非线性函数的思想，回到LQR的目标——在知道矩阵$A,B,U,V$组成的二次型奖励函数后后，找到能够最大化预期收益的策略。

同以前一样，我们要最大化奖励函数$R^{(0)}(s_0,a_0)+R^{(1)}(s_1,a_1)+\cdots+R^{(T)}(s_T,a_T)$，利用前面介绍过的“有限时域MDP”中动态规划算法的思路：
* 先计算出MDP的最后一步$V_T^*(s)$：
    $$
    \begin{align}
    V_T^*(s)&=\max_{a_T}R^{(T)}(s,a)\\
    &=\max_{a_T}\left(-s_T^TU_Ts_T-a_T^TV_Ta_T\right)\\
    &=-s_T^TU_Ts_T
    \end{align}
    $$
    因为我们定义$V_T$是半正定矩阵（$a_T^TV_Ta_T\geq0$），所以要保证上式取最大值，那么最好的方法就是不做任何动作（使$a_T^TV_Ta_T=0$）
* 对应的最后一步的策略为$\pi_T^*(s_T)=\arg\max_{a_T}R^{(T)}(s_t,a_t)=0$；

接下来进入动态规划的步骤，先来考虑如何从已知$V_{t+1}^*$向后递推出$V_{t}^*$：
* 在有限时域MDP的$(2)$式中，我们只需把求和换成积分（因为现在又是求连续状态的情形了）即可：
    $$
    \begin{align}
    V_t^*(s)&=\max_{a_t}\left(R^{(t)}(s_t,a_t)+\sum_{s_{t+1}}P_{s_ta_t}^{(t)}\left(s_{t+1}\right)V_{t+1}^*\left(s_{t+1}\right)\right)\\
    &=\max_{a_t}\left(R^{(t)}(s_t,a_t)+\mathrm E_{s_{t+1}\ \sim P_{s_ta_t}}\left[V_{t+1}^*(s_{t+1})\right]\right)
    \end{align}
    $$
    我们跳过了换成积分的步骤，直接将后面的项写为期望形式。实际上LQR具有性质：每一个价值函数$V_{t+1},V_{t},\cdots$都可以表为二次函数。

    假设$\displaystyle V_{t+1}^*(s_{t+1})=s_{t+1}^T\varPhi_{t+1}s_{t+1}+\varPsi_{t+1},\varPhi\in\mathbb R^{n\times n},\varPsi\in\mathbb R$，即$V_{t+1}^*$是一个关于$s_{t+1}$的二次函数，则选择适当的$\varPhi_t,\varPsi_t$就可以得到$\displaystyle V_t^*(s_t)=s_t^T\varPhi_ts_t+\varPsi_t$。

    那么在上面的$V_T^*(s)$中，$\varPhi=-U_T,\varPsi=0\to V_T^*(s_T)=s_T^T\varPhi_Ts_T+\varPsi_T$。
* 然后就是向后递推过程：
    $$
    V_t^*(s_t)=\max_{a_t}\left(\underbrace{-s_t^TU_ts_t-a_t^TV_ta_t}_{\textrm{(a)}}+\underbrace{\mathrm E_{s_{t+1}\ \sim\mathcal N(A_ts_t+B_ta_t,\varSigma_w)}\left[s_{t+1}^T\varPhi_{t+1}s_{t+1}+\varPsi_{t+1}\right]}_{\textrm{(b)}}\right)
    $$
    式中的$N(A_ts_t+B_ta_t,\varSigma_w)$就是$P_{s_ta_t}$，而期望中的$s_{t+1}^T\varPhi_{t+1}s_{t+1}+\varPsi_{t+1}$就是$V_{t+1}^*(s_{t+1})$。这个式子可以理解为找到动作$a_t$使$\textrm{(a)}+\textrm{(b)}$能够取到最大值，其中$\textrm{(a)}$表示处于状态$s_t$的即时奖励，而$\textrm{(b)}$表示在$s_ta_t$下$V^*$的期望，而$V^*$是下一步的最优价值函数。这个式子就是将奖励函数、状态转换概率、最优价值函数等带入后的最终结果。

    整个等号右边可以看做是对一个“大型”二次函数极大值的求解过程，我们和以前一样，求这个二次函数关于$a_t$的导数并将导数置为零，最终解出$a_t$。现在省略推导步骤，直接写出结果就有：
    $$
    a_t=\underbrace{\left(B_t^T\varPhi_{t+1}B_t-V_t\right)^{-1}B_t^T\varPhi_{t+1}A_t}_{L_t}\cdot s_t
    $$
    要留意的是，可以把前面的部分看做是一个“大型”矩阵$L_t$，于是，$a_t$就是$L_t$与$s_t$的乘积，即$a_t$是关于$s_t$的线性函数。
* 最后是求解最优策略，记得我摸在有限时域MDP中提到，$\pi^*$计算中$\arg\max$的部分就是$V^*$计算中$\max$的部分，而得到的也就是上面的$a_t$：

    $$
    \begin{align}
    \displaystyle \pi_t^*(s_t)&=\arg\max_{a_t}\left(-s_t^TU_ts_t-a_t^TV_ta_t+\mathrm E_{s_{t+1}\ \sim\mathcal N(A_ts_t+B_ta_t,\varSigma_w)}\left[s_{t+1}^T\varPhi_{t+1}s_{t+1}+\varPsi_{t+1}\right]\right)\\
    &=L_ts_t
    \end{align}
    $$
    也就是说，要得到最优策略，我摸只需要计算一个矩阵$L_t$，然后再乘以当前状态$s_t$即可。换句话说，最优策略是当前状态的线性函数。需要注意的是，这不是一个近似的函数，我们不会说“找到一个最优线性策略”、“找到最优策略然后用直线拟合该策略”，这不是用直线去近似最优策略，而是说最佳策略就是一条直线、最优动作就是当前状态的一个线性函数。

    而当我们求出$a_t$之后，把它代回$\displaystyle V_t^*(s_t)=s_t^T\varPhi_ts_t+\varPsi_t$，就会发现在这个二次函数中有：
    $$
    \begin{align}
    \varPhi_t&=A_t^T\left(\varPhi_{t+1}-\varPhi_{t+1}B_t\left(B_t^T\varPhi_{t+1}B_t-V_t\right)^{-1}B_t\varPhi_{t+1}\right)A_t-U_t\tag{3}\\
    \varPsi_t&=-\mathrm{tr}\varSigma_w\varPhi_{t+1}+\varPsi_{t+1}
    \end{align}
    $$
    即$\varPhi_t,\varPsi_t$是关于$\varPhi_{t+1},\varPsi_{t+1}$的线性函数，现在就可以计算$V_t^*(s_t)$了。其中$(3)$式也被称作**离散时间Riccati方程（discrete time Riccati equation）**。

对LQR做一个总结，算法如下：
* 初始化$\varPhi_t=-U_t,\varPsi_t=0$；
* 向后递推计算$\varPhi_t,\varPsi_t$，使用离散时间Riccati方程，将其当做关于$\varPhi_{t+1},\varPsi_{t+1}$的线性函数；（依次求出$t=T-1,T-2,\cdots,0$。）
* 计算$L_t$，这是一个关于$\varPhi_{t+1},\varPsi_{t+1}$的函数；
* 计算$\pi^*(s_t)$，这是一个关于$s_t$的线性函数，系数是$L_t$。


# EpiEstim (python)

## Configure the data

```python
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

plt.style.use('seaborn-poster')
from datetime import datetime as dt
from datetime import timedelta
import glob

from Reff_functions import *
from Reff_constants import *
from scipy.stats import gamma
```

```python
#Code taken from read_in_cases from Reff_functions. Preprocessing was not helpful for this situation.
def read_cases_lambda(case_file_date):
    path = "../data/COVID-19 UoM "+case_file_date+"*.xlsx"

    for file in glob.glob(path):
        df_NNDSS = pd.read_excel(file,
                           parse_dates=['SPECIMEN_DATE','NOTIFICATION_DATE','NOTIFICATION_RECEIVE_DATE','TRUE_ONSET_DATE'],
                           dtype= {'PLACE_OF_ACQUISITION':str})
        df_NNDSS.PLACE_OF_ACQUISITION.fillna('00038888',inplace=True) #Fill blanks with simply unknown

       # df_NNDSS['date_inferred'] = df_NNDSS.TRUE_ONSET_DATE
      #  df_NNDSS.loc[df_NNDSS.TRUE_ONSET_DATE.isna(),'date_inferred'] = df_NNDSS.loc[df_NNDSS.TRUE_ONSET_DATE.isna()].NOTIFICATION_DATE - timedelta(days=5)
      #  df_NNDSS.loc[df_NNDSS.date_inferred.isna(),'date_inferred'] = df_NNDSS.loc[df_NNDSS.date_inferred.isna()].NOTIFICATION_RECEIVE_DATE - timedelta(days=6)
    df_NNDSS['imported'] = df_NNDSS.PLACE_OF_ACQUISITION.apply(lambda x: 1 if x[-4:]=='8888' and x != '00038888' else 0)
    df_NNDSS['local'] = 1 - df_NNDSS.imported

    df_interim = df_NNDSS[['NOTIFICATION_RECEIVE_DATE','STATE','imported','local']]
    #df_interim = df_interim[~np.isnat(df_interim.NOTIFICATION_DATE)] #Get rid of non-existent dates.
    # Importantly, imported and local are indicator variables in df_interim.

    #df_state = df_NNDSS[['NOTIFICATION_DATE','STATE','imported','local']].groupby(['STATE','NOTIFICATION_DATE']).sum()
    return(df_interim)

def tidy_cases_lambda(interim_data, remove_territories=True):
    #Remove non-existent notification dates
    interim_data = interim_data[~np.isnat(interim_data.NOTIFICATION_RECEIVE_DATE)]

    #Filter out territories
    if(remove_territories):
        df_linel = interim_data[(interim_data['STATE']!='NT') & (interim_data['STATE']!='ACT')]

    #Melt down so that imported and local are no longer columns. Allows multiple draws for infection date.
    #i.e. create linelist data
    df_linel = df_linel.melt(id_vars = ['NOTIFICATION_RECEIVE_DATE','STATE'], var_name = 'SOURCE',value_name='n_cases')

    #Reset index or the joining doesn't work
    df_linel = df_linel[df_linel.n_cases!=0]
    df_linel = df_linel.reset_index(drop=True)
    return(df_linel)

date = '10Aug'
df_interim = read_cases_lambda(date)

df_linel = tidy_cases_lambda(df_interim)
```

## Part 1: Inferring infection dates
$\Lambda$ depends on the infection date (ID), while the data contains the notification date (ND). We obtain ID through the following relationship:
$$
ID = ND - reporting\_delay - incubation\_period.
$$

A gamma distribution was fitted to case data using the MLE algorithm to produce distributions for reporting delay and incubation period.

```python
##uncomment for debugging
# notification_dates = df_linel['NOTIFICATION_DATE']
# mean_rd = 5.47
# sd_rd = 4.04
# mean_inc = 2.0
# sd_inc = 1.41
# nreplicates = 3

##gamma draws take arguments (shape, scale)
def draw_inf_dates(df_linelist, shape_rd=2.77, scale_rd=3.17, offset_rd=0,
                    shape_inc=2.0/1.5, scale_inc=1.5, offset_inc=1,nreplicates=1):

    notification_dates = df_linelist['NOTIFICATION_RECEIVE_DATE']
    nsamples = notification_dates.shape[0]

    #    DEFINE DELAY DISTRIBUTION
    #     mean_rd = 5.47
    #     sd_rd = 4.04
    #scale_rd = shape_rd/(scale_rd)**2
    #shape_rd = shape_rd/scale_rd

    # DEFINE INCUBATION PERIOD DISTRIBUTION
    #     mean_inc = 2.0
    #     sd_inc = 1.41
    #scale_inc = (scale_inc)**2/shape_inc #scale**2 = var / shape
    #shape_inc =(scale_inc)**2/scale_inc**2

    #Draw from distributions - these are long vectors
    inc_period = offset_inc+np.random.gamma(shape_inc, scale_inc, size = (nsamples*nreplicates))
    rep_delay = offset_rd+np.random.gamma(shape_rd, scale_rd, size = (nsamples*nreplicates))

    #infection date is id_nd_diff days before notification date. This is also a long vector.
    id_nd_diff = inc_period + rep_delay

    #Minutes aren't included in df. Take the ceiling because the day runs from 0000 to 2359. This can still be a long vector.
    whole_day_diff = np.ceil(id_nd_diff)
    time_day_diffmat = whole_day_diff.astype('timedelta64[D]').reshape((nsamples, nreplicates))

    #Vector must be coerced into a nsamples by nreplicates array. Then each column must be subtracted from notification_dates.
    #Subtract days off of notification dates.

    notification_mat = np.tile(notification_dates, (nreplicates,1)).T #notification_dates is repeated as a column nreplicates times.

    infection_dates = notification_mat - time_day_diffmat

    #Make infection dates into a dataframe
    datecolnames = [*map(str,range(nreplicates))]
    infdates_df = pd.DataFrame(infection_dates,columns = datecolnames)

    #Uncomment this if theres errors
    #print([df_linelist.shape, infdates_df.shape])

    #Combine infection dates and original dataframe
    df_inf = pd.concat([df_linelist, infdates_df], axis=1, verify_integrity=True)

    return(df_inf)
```

```python
df_inf = draw_inf_dates(df_linel, nreplicates=1000)
df_inf.head()
```

      NOTIFICATION_RECEIVE_DATE
      STATE
      SOURCE
      n_cases
      0
      1
      2
      3
      4
      5
      ...
      990
      991
      992
      993
      994
      995
      996
      997
      998
      999

      0
      2020-08-03
      NSW
      imported
      1
      2020-07-16
      2020-07-21
      2020-07-15
      2020-07-22
      2020-07-22
      2020-07-25
      ...
      2020-07-26
      2020-07-13
      2020-07-22
      2020-07-29
      2020-07-14
      2020-07-28
      2020-07-18
      2020-07-21
      2020-07-30
      2020-07-11

      1
      2020-07-29
      NSW
      imported
      1
      2020-07-25
      2020-07-15
      2020-07-12
      2020-07-14
      2020-07-11
      2020-07-14
      ...
      2020-07-17
      2020-07-20
      2020-07-23
      2020-07-16
      2020-07-12
      2020-07-16
      2020-07-24
      2020-07-13
      2020-07-24
      2020-07-15

      2
      2020-08-03
      NSW
      imported
      1
      2020-07-19
      2020-07-09
      2020-07-29
      2020-07-15
      2020-07-28
      2020-07-24
      ...
      2020-07-24
      2020-07-09
      2020-07-26
      2020-07-20
      2020-07-24
      2020-07-24
      2020-07-17
      2020-07-24
      2020-07-23
      2020-07-21

      3
      2020-07-25
      NSW
      imported
      1
      2020-07-17
      2020-07-16
      2020-06-30
      2020-07-09
      2020-07-22
      2020-07-06
      ...
      2020-07-21
      2020-07-19
      2020-06-30
      2020-07-08
      2020-07-14
      2020-07-11
      2020-07-08
      2020-07-19
      2020-07-18
      2020-07-13

      4
      2020-07-31
      NSW
      imported
      1
      2020-07-18
      2020-07-22
      2020-07-16
      2020-07-22
      2020-07-25
      2020-07-12
      ...
      2020-07-21
      2020-07-20
      2020-07-26
      2020-07-22
      2020-07-15
      2020-07-15
      2020-07-24
      2020-07-15
      2020-07-11
      2020-07-20

5 rows × 1004 columns

```python
def index_by_infection_date(infections_wide):
    datecolnames = [*infections_wide.columns[4:]]
    df_combined = infections_wide[['STATE','SOURCE',datecolnames[0],'n_cases']].groupby(['STATE', datecolnames[0],'SOURCE']).sum()

    #For each column (cn=column number): concatenate each sample as a column.
    for cn in range(1,len(datecolnames)):
        df_addin = infections_wide[['STATE','SOURCE',datecolnames[cn],'n_cases']].groupby(['STATE', datecolnames[cn],'SOURCE']).sum()
        df_combined = pd.concat([df_combined,df_addin], axis=1, ignore_index = True)

    #NaNs are inserted for missing values when concatenating. If it's missing, there were zero infections
    df_combined[np.isnan(df_combined)]=0
    #Rename the index.
    df_combined.index.set_names(["STATE","INFECTION_DATE","SOURCE"], inplace=True)

    #return(df_combined)

    ##INCLUDE ALL DAYS WITH ZERO INFECTIONS IN THE INDEX AS WELL.

    # Reindex to include days with zero total infections.
    local_infs = df_combined.xs('local',level='SOURCE')
    imported_infs = df_combined.xs('imported',level='SOURCE')
    statelist = [*df_combined.index.get_level_values('STATE').unique()]

    #Should all states have the same start date? Current code starts from the first case in each state.
    #For the same start date:
    local_statedict = dict(zip(statelist, np.repeat(None, len(statelist))))
    imported_statedict = dict(zip(statelist, np.repeat(None, len(statelist))))

    #Determine start date as the first infection date for all.
    #start_date = np.datetime64("2020-02-01")
    start_date = df_combined.index.get_level_values('INFECTION_DATE').min()

    #Determine end dates as the last infected date by state.
    index_only = df_combined.index.to_frame()
    index_only = index_only.reset_index(drop=True)
    maxdates = index_only.groupby(['STATE'])['INFECTION_DATE'].max()

    for aus_state in statelist:
        state_data = local_infs.xs(aus_state, level='STATE')
        #start_date = state_data.index.min()

        #dftest.index=dftest.reindex(alldates, fill_value=0)

        alldates = pd.date_range(start_date, maxdates[aus_state]) #All days from start_date to the last infection day.
        local_statedict[aus_state] = state_data.reindex(alldates, fill_value=0)

    for aus_state in statelist:
        state_data = imported_infs.xs(aus_state, level='STATE')
        alldates = pd.date_range(start_date, maxdates[aus_state])
        imported_statedict[aus_state] = state_data.reindex(alldates, fill_value=0)

    #Convert dictionaries to data frames
    df_local_inc_zeros = pd.concat(local_statedict)
    df_local_inc_zeros['SOURCE']='local'
    df_imp_inc_zeros = pd.concat(imported_statedict)
    df_imp_inc_zeros['SOURCE']='imported'

    #Merge dataframes and reindex.
    df_inc_zeros = pd.concat([df_local_inc_zeros, df_imp_inc_zeros])

    df_inc_zeros = df_inc_zeros.reset_index()
    df_inc_zeros= df_inc_zeros.groupby(['level_0',"level_1","SOURCE"]).sum()
    df_inc_zeros.index = df_inc_zeros.index.rename(['STATE','INFECTION_DATE',"SOURCE"])

    return(df_inc_zeros)

df_inc_zeros = index_by_infection_date(df_inf)
df_inc_zeros.head()
```

      0
      1
      2
      3
      4
      5
      6
      7
      8
      9
      ...
      990
      991
      992
      993
      994
      995
      996
      997
      998
      999

      STATE
      INFECTION_DATE
      SOURCE

      NSW
      2019-12-14
      imported
      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
      0.0
      0.0
      0.0
      0.0
      0.0

      local
      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
      0.0
      0.0
      0.0
      0.0
      0.0

      2019-12-15
      imported
      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
      0.0
      0.0
      0.0
      0.0
      0.0

      local
      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
      0.0
      0.0
      0.0
      0.0
      0.0

      2019-12-16
      imported
      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
      0.0
      0.0
      0.0
      0.0
      0.0

5 rows × 1000 columns

```python
# unit test to ensure min and max dates captures all imputations of
# infection dates
summary = np.sum(df_inc_zeros, axis=0).describe() #Differences in numbers: start date?
assert summary.loc['max'] == summary.loc['min'], "Min number of cases does not match max number, dates have truncated cases"

```

## Part 2: Calculating Lambda

$$
\Lambda_t(w_s) = \sum_{s=1}^t (I_{t-s}^{local} + I_{t-s}^{imported})w_s = \sum_{s=1}^t I_{t-s}w_s,
$$
where $w_s$ is the probability that the generation interval is $s$ and $I_t$ is the number of infected individuals at time $t$.

### Part 2a: Discretizing the gamma generation interval distribution

In the formula for $\Lambda_t$, we sum over $w$. We should consider generation interval as a discrete random variable here.

```python
#Define gamma distribution for generation interval
#mean_gen = 2
#sd_gen = 1.74
scale_gen = 1#mean_gen/(sd_gen)**2
shape_gen = 2#mean_gen/scale_gen

trunc_days = 21
shift=0
xmids = [x+shift for x in range(trunc_days+1)] #Find midpoints for discretisation

#scipy uses scale in the compsci sense
gamma_vals = gamma.pdf(xmids,
                       a=shape_gen,
                       scale=scale_gen)
disc_gamma = gamma_vals/sum(gamma_vals)
```

```python
#Discretisation error check (should sum to 1)

print("Sum of gamma values is " + str(sum(gamma_vals))+"; \n Sum of discretised gamma values is " + str(sum(disc_gamma)))
```

    Sum of gamma values is 0.9206735842426517;
     Sum of discretised gamma values is 1.0

```python
xrange = np.linspace(0,trunc_days,150)
fig,ax = plt.subplots(figsize=(12,9))
w = ax.bar(xmids,height=disc_gamma, width=1)
ax.set_title("Generation time distribution")
ax.plot(xrange, gamma.pdf(xrange, a=shape_gen, scale=scale_gen), linewidth=4,alpha=0.8, color="orange")
ax.set_xlabel('Days')
plt.show()
```

### Part 2b: Calculating $\Lambda$

```python
def generate_lambda(infection_dates, shape_gen=2, scale_gen=1,
                    trunc_day=21,shift=0, offset=1):
    """
    Given array of infection_dates (N_dates by N_samples), where values are possible
    number of cases infected on this day, generate the force of infection Lambda_t,
    a N_dates-tau by N_samples array.
    """
    from scipy.stats import gamma

    #scale_gen = mean_gen/(sd_gen)**2
    #shape_gen = mean_gen/scale_gen

    xmids = [x+shift for x in range(trunc_days+1)] #Find midpoints for discretisation
    gamma_vals = gamma.pdf(xmids, a=shape_gen, scale=scale_gen) #double check parameterisation of scipy
    #renormalise the pdf
    disc_gamma = gamma_vals/sum(gamma_vals)

    ws = disc_gamma[:trunc_day]
    #offset
    ws[offset:] = disc_gamma[:trunc_day-offset]
    ws[:offset] = 0
    lambda_t = np.zeros(shape=(infection_dates.shape[0]-trunc_day+1, infection_dates.shape[1]))
    for n in range(infection_dates.shape[1]):
        lambda_t[:,n] = np.convolve(infection_dates[:,n], ws, mode='valid')
    return lambda_t

def lambda_all_states(df_infection, **kwargs):
    """
    Use geenrate lambda on every state
    """
    statelist = [*df_infection.index.get_level_values('STATE').unique()]

    lambda_dict ={}
    for state in statelist:
        df_total_infections = df_infection.groupby(['STATE','INFECTION_DATE']).agg(sum)
        lambda_dict[state] = generate_lambda(
             df_total_infections.loc[state].values,
            **kwargs
            )

    return lambda_dict
```

```python
trunc_day = 21
#get all lambdas
lambda_dict = lambda_all_states(df_inc_zeros,
                                shape_gen=2,scale_gen=1,offset=1,
                                trunc_day=trunc_day)

lambda_dict['VIC']
```

    array([[ 0.39957641,  0.        ,  0.        , ...,  0.        ,
             0.        ,  0.        ],
           [ 0.29399189,  0.        ,  0.        , ...,  0.        ,
             0.        ,  0.        ],
           [ 0.16223036,  0.        ,  0.        , ...,  0.        ,
             0.        ,  0.        ],
           ...,
           [73.85978124, 73.30521254, 71.73239971, ..., 70.42277095,
            68.14326479, 75.01023409],
           [47.88760062, 50.945106  , 45.22574027, ..., 44.38950894,
            43.13436371, 48.09216342],
           [28.03491746, 29.56040224, 27.16349104, ..., 24.72758341,
            26.51005578, 28.42930394]])

## 3. Sample from the posterior

Using Cori et al. 2013, the posterior distribution of $R_{t,\tau}$ is a Gamma distribution with parameters shape and scale

\begin{equation}
\left( a + \sum^t_{s = t - \tau +1} I_s , \frac{1}{\frac{1}{b} + \sum^t_{ s = t-\tau + 1} \Lambda_s } \right)
\end{equation}

```python
def Reff_from_case(cases_by_infection, lamb, prior_a=1, prior_b=5, tau=7, samples=1000):
    """
    Using Cori at al. 2013, given case incidence by date of infection, and the force
    of infection \Lambda_t on day t, estimate the effective reproduction number at time
    t with smoothing parameter \tau.

    cases_by_infection: A T by N array, for T days and N samples
    lamb : A T by N array, for T days and N samples
    """
    csum_incidence = np.cumsum(cases_by_infection, axis = 0)
    #remove first few incidences to align with size of lambda
    # Generation interval length 20
    csum_incidence = csum_incidence[20:,:]
    csum_lambda = np.cumsum(lamb, axis =0)

    roll_sum_incidence = csum_incidence[tau:, :] - csum_incidence[:-tau, :]
    roll_sum_lambda = csum_lambda[tau:,:] - csum_lambda[:-tau,:]
    a = prior_a + roll_sum_incidence
    b = 1/(1/prior_b + roll_sum_lambda)

    R = np.random.gamma(a,b) #shape, scale

    #Need to empty R when there is too few cases...

    #Use array inputs to output to same size
    #inputs are T-tau by N, output will be T-tau by N
    #
    return a,b, R

def generate_summary(samples, dates_by='rows'):
    """
    Given an array of samples (T by N) where rows index the dates,
    generate summary statistics and quantiles
    """

    if dates_by=='rows':
        #quantiles of the columns
        ax = 1
    else:
        #quantiles of the rows
        ax = 0
    mean = np.mean(samples, axis = ax)
    bottom, lower, median, upper, top = np.quantile(samples,
                                                     (0.05, 0.25, 0.5, 0.75, 0.95),
                                                     axis =ax)
    std = np.std(samples, axis = ax)
    output = {
        'mean':mean,
        'std':std,
        'bottom':bottom,
        'lower':lower,
        'median':median,
        'upper':upper,
        'top': top,

    }
    return output

def plot_Reff(Reff:dict, dates=None, ax_arg=None, **kwargs):
    """
    Given summary statistics of Reff as a dictionary, plot the distribution over time
    """
    import matplotlib.pyplot as plt

    plt.style.use('seaborn-poster')
    from datetime import datetime as dt

    if ax_arg is None:
        fig, ax = plt.subplots(figsize=(12,9))
    else:
        fig, ax = ax_arg

    color_cycle = ax._get_lines.prop_cycler
    curr_color = next(color_cycle)['color']
    if dates is None:
        dates = range(len(Reff['mean']))

    ax.plot(dates, Reff['mean'], color= curr_color, **kwargs)

    ax.fill_between(dates, Reff['lower'],Reff['upper'], alpha=0.4, color = curr_color)
    ax.fill_between(dates, Reff['bottom'],Reff['top'], alpha=0.4, color= curr_color)

       #grid line at R_eff =1
    ax.set_yticks([1],minor=True,)
    ax.set_yticks([0,2,3],minor=False)
    ax.set_yticklabels([0,2,3],minor=False)
    ax.yaxis.grid(which='minor',linestyle='--',color='black',linewidth=2)

    ax.tick_params(axis='x', rotation = 45)

    return fig, ax

```

## 4. Plot the estimates

```python

tau = 14
prior_a=1
prior_b=5

#get all lambdas
lambda_DL = lambda_all_states(df_inc_zeros)

states = [initial[1] for initial in sorted(list(states_initials.items()))]
states.remove('NT')
states.remove('ACT')

#read in old LSHTM estimates
df_L_R = read_in_LSHTM()
date_filter = pd.date_range(start='2020-03-01',end='2020-08-01')

#prepare NNDSS cases
df_cases = df_interim.groupby(['NOTIFICATION_RECEIVE_DATE','STATE']).agg(sum)
df_cases = df_cases.reset_index()

fig, ax = plt.subplots(nrows=2, ncols=3,
                       sharex=True, sharey=True,
                      figsize=(15,12)
                      )

for i,state in enumerate(states):
    row = i//3
    col = i%3

    lambda_state = lambda_DL[state]
    df_state_I = df_inc_zeros.xs((state,'local'),level=('STATE','SOURCE'))

    a,b,R = Reff_from_case(df_state_I.values,lambda_state,prior_a=1, prior_b=2, tau=tau)

    R_summary = generate_summary(R)

    fig, ax[row,col] = plot_Reff(R_summary,
                                 dates=df_state_I.index.values[trunc_days-1+tau:],
                                 ax_arg=(fig, ax[row,col]),
                                 label='Our Model')

    #plot formatting
    ax[row,col].set_title(state)
    ax[row,col].set_ylim((0,4))
    ax[row,col].set_xlim((pd.to_datetime('2020-03-01'),pd.to_datetime('2020-08-10')))

    #plot cases behind
    ax2 = ax[row,col].twinx()

    ax2.bar(df_cases.loc[df_cases.STATE==state,'NOTIFICATION_RECEIVE_DATE'],
            df_cases.loc[df_cases.STATE==state,'local']+df_cases.loc[df_cases.STATE==state,'imported'],
           color='grey',
            alpha=0.3
           )
    ax2.bar(df_cases.loc[df_cases.STATE==state,'NOTIFICATION_RECEIVE_DATE'],
            df_cases.loc[df_cases.STATE==state,'local'],
           color='grey',
            alpha=0.8
           )

    #plot old LSHTM estimates
    df_june = df_L_R.loc[(df_L_R.date_of_analysis=='2020-07-27')&(df_L_R.state==state)]
    df = df_june.loc[(df_june.date.isin(date_filter))]

    ax[row,col].plot(df.date, df['median'], label='Old LSHTM',color='C1')
    ax[row,col].fill_between(df.date, df['bottom'], df['top'],color='C1', alpha=0.3)
    ax[row,col].fill_between(df.date, df['lower'], df['upper'],color='C1', alpha=0.3)

plt.legend()
plt.savefig("../figs/EpyEstim_tau_"+str(tau)+"_"+date+".png",dpi=300)
plt.show()
```

### Sophie's implementation
Use this as a unit test for changes to the estimator

```python
#ws is the discretised gamma distribution; reversed due to the formula for lambda t.
ws = [*reversed(disc_gamma[:(trunc_days+1)])] #was taken from 1 before.

#Calculate lambda t for a given t in one state.
def calculate_lambda_t(state_df, t, trunc_days = 21, ws=ws):
    #t = predict_date_range[30]
    #state_df = input_state

    tstart= t-np.timedelta64(trunc_days,'D')
    relevant_dates = pd.date_range(tstart, t-np.timedelta64(1,'D'))
    reldates_df = state_df.loc[relevant_dates]

    #Dates don't matter, since we're calculating lambda t for t = t.
    reldates_df = reldates_df.reset_index(drop=True)
    ws_mat = pd.DataFrame(np.tile(ws, (reldates_df.shape[1],1)).T)

    #lambda_t=sum(reldates*ws)
    lambda_t = np.sum(reldates_df.mul(ws_mat), axis=0)
    return(lambda_t)

#Loop over states and times to calculate all lambda t's
#Input: imported/local counts of infections by date and state. Each column should be a different sample.
#Output: Lambda t by date and state. Each column corresponds to a different original sample.
def calculate_all_lambdas(infection_df):

    #Depending on data format, flatten if necessary
    if type(infection_df.index)!=pd.RangeIndex:
        infection_df = infection_df.reset_index()

    #Create data frame with the total number of infections.
    I_total = infection_df.groupby(['STATE',"INFECTION_DATE"]).sum()

    #Find states and preallocate to dict
    statelist = [*I_total.index.get_level_values('STATE').unique()]
    state_dict = dict(zip(statelist, np.repeat(None, len(statelist))))

    predict_reff_from = np.datetime64('2020-02-01')

    #Calculate Reff for each state.
    for state in statelist:
        #print(state)
        input_state_df = I_total.xs(state, level='STATE')
        tmax = input_state_df.index.get_level_values('INFECTION_DATE').max()
        predict_date_range = pd.date_range(predict_reff_from, tmax)

        date_dict = dict(zip(predict_date_range, np.repeat(None, len(predict_date_range))))

        #Find lambda t for every day.
        for t in predict_date_range:
            #print(t)
            date_dict[t] = calculate_lambda_t(input_state_df, t).to_numpy()

        state_dict[state]=date_dict

    #Convert dict to a dataframe
    lambda_df = pd.DataFrame.from_dict({(i,j): state_dict[i][j]
                                for i in state_dict.keys()
                                for j in state_dict[i].keys()},
                                orient='index')
    lambda_df.index = pd.MultiIndex.from_tuples(lambda_df.index,names = ['STATE','INFECTION_DATE'])

    return(lambda_df)

lambdas = calculate_all_lambdas(df_inc_zeros)
```

```python
#test run on a state
state='SA'
tau = 14
#df_VIC = df_inc_zeros.xs(('VIC','local'),level=('STATE','SOURCE'))
#lambda_VIC = generate_lambda(df_VIC.values )

#get all lambdas
lambda_DL = lambda_all_states(df_inc_zeros)

#select lambda for the right state
lambda_VIC = lambda_DL[state]
df_VIC = df_inc_zeros.xs((state,'local'),level=('STATE','SOURCE'))

a,b,R = Reff_from_case(df_VIC.values,lambda_VIC,prior_a=1, prior_b=2, tau = tau)

R_summary = generate_summary(R)

fig, ax = plot_Reff(R_summary, dates=df_VIC.index.values[20+tau:])

a,b,R = Reff_from_case(df_VIC.values[7:],lambdas.loc[state].values,prior_a=1, prior_b=2, tau = tau)

R_summary = generate_summary(R)

fig, ax = plot_Reff(R_summary, dates=df_VIC.index.values[27+tau:],ax_arg =(fig,ax))
   #grid line at R_eff =1

ax2 = ax.twinx()
df_cases = df_interim.groupby(['NOTIFICATION_RECEIVE_DATE','STATE']).agg(sum)
df_cases = df_cases.reset_index()
ax2.bar(df_cases.loc[df_cases.STATE==state,'NOTIFICATION_RECEIVE_DATE'],
        df_cases.loc[df_cases.STATE==state,'local']+df_cases.loc[df_cases.STATE==state,'imported'],
       color='grey',
        alpha=0.3
       )
ax2.bar(df_cases.loc[df_cases.STATE==state,'NOTIFICATION_RECEIVE_DATE'],
        df_cases.loc[df_cases.STATE==state,'local'],
       color='grey',
        alpha=0.8
       )

plt.show()
```


```javascript
%%javascript
MathJax.Hub.Config({
    TeX: { equationNumbers: { autoNumber: "AMS" } }
});

```

    <IPython.core.display.Javascript object>

Optimal transport studies turning one measure to another while some cost of transportation is minimized. It started in 17 century with practical problems. The theory behind optimal transport developed extensively in nineteenth and twenteenth centruries.

In this post, I review the main theoretical ideas of optimal transport and their interpretations.

There are two major historical points in optimal transport marked with two names _Monge (1781)_ and _Kantorovich (1940)_. In order to compress in one sentence, Monge tried to find a deterministic transport map while Kantorovich is fine with a probabilistic map as well.

We start with a definition of _image measure_ that will be recurring over and over in this post.

**Definition 1.** _Given $X, Y$ spaces, take $T:X\to Y$ and $\mu\in\mathcal{P}(X)$. The image measure represented
    $T_{\#}\mu\in\mathcal{P}(Y)$ as_
    \begin{equation}
    \forall A\in \mathcal{B}(Y): \quad (T_\#\mu)(A):= \mu(T^{-1}(A))
    \end{equation}

It can be shown that $T_\#\mu$ is a probability measure on $Y$. Notice that this defintion is one-way, i.e., the pullback operator $S^\#$ does not necessarily give a probability measure.

One practical way to check $\nu\stackrel{?}{=}T_\#\mu$ is to integrate all Borel and Bounded maps $\phi:Y\to\mathbb{R}$ against the measures of both sides
\begin{equation}
\int_{Y} \varphi(y) d \nu(y)=\int_{X} \varphi(T(x)) d \mu(x)
\end{equation}

As a corollary, the _push-forward_ condition is known as
\begin{equation}\label{eq:push_forward_condition}
\nu=T_{\#} \mu \Rightarrow \int_{Y} \varphi d\left(T_{\#} \mu\right)=\int_{X} \varphi \circ T d \mu
\end{equation}
for any $\phi:X\to\mathbb{R}$ Borel and bounded.

Measures can be pushed forward via the composition of functions
\begin{equation}
(S \circ T)_{\#} \mu=S_{\#}\left(T_{\#} \mu\right)
\end{equation}

**Definition 2.** _Given $\mu\in\mathcal{P}(X)$ and $\nu\in\mathcal{P}(Y)$, a map $T:X\to Y$
    is called a **transport map** from $\mu$ to $\nu$ if $T_\#\mu=\nu$. _

With this defintion, given a $\mu$ and $\nu$, there might be no map that transports $\mu$ to $\nu$.
To generalize a little bit, **coupling** is defined as a probabilistic map

**Definition 3.** $\gamma\in\mathcal{P}(X\times Y)$ is called a coupling of $\mu$ and $\nu$ if_
\begin{equation}
\left(\pi_{X}\right)_{\#} \gamma=\mu \text { and }\left(\pi_{Y}\right)_{\#} \gamma=\nu
\end{equation}
where
\begin{equation}
\pi_{X}(x, y)=x, \quad \pi_{Y}(x, y)=y \quad \forall(x, y) \in X \times Y
\end{equation}
According to the push-forward condition \eqref{eq:push_forward_condition}, Borel bounded functions $\phi$ and $\psi$ can be used verify a coupling as
\begin{equation}
\forall \varphi: X \rightarrow \mathbb{R} \quad \int_{X \times Y} \varphi(x) d \gamma(x, y)=\int_{X \times Y} \varphi \circ \pi_{X}(x, y) d \gamma(x, y)=\int_{X} \varphi(x) d \mu(x)
\end{equation}
and
\begin{equation}
\forall \psi: Y \rightarrow \mathbb{R} \quad \int_{X \times Y} \psi(y) d \gamma(x, y)=\int_{X \times Y} \psi \circ \pi_{Y}(x, y) d \gamma(x, y)=\int_{Y} \psi(y) d \nu(y)
\end{equation}

Let denote the set of all couplings of the measures $\mu$ and $\nu$ by $\Gamma(\mu, \nu)$. Despite the transport map that can be non-existing for a given $\mu$ and $\nu$, $\Gamma(\mu, \nu)$ is always non-empty. Because at least $\gamma=\mu \otimes \nu$.

Any transport map $T$ induces a couple denoted by $\gamma_T$ defined as

```javascript
%%javascript
MathJax.Hub.Queue(
  ["resetEquationNumbers", MathJax.InputJax.TeX],
  ["PreProcess", MathJax.Hub],
  ["Reprocess", MathJax.Hub]
);
```

    <IPython.core.display.Javascript object>

```python

```