# Orthogonal Random Forest and Causal Forest: Use Cases and Examples

Causal Forests and Generalized Random Forests are a flexible method for estimating treatment effect heterogeneity with Random Forests. Orthogonal Random Forest (ORF) combines orthogonalization, a technique that effectively removes the confounding effect in two-stage estimation, with generalized random forests. Due to the orthogonalization aspect of this method, the ORF performs especially well in the presence of high-dimensional confounders. For more details, see [this paper](https://arxiv.org/abs/1806.03467) or the [EconML docummentation](https://econml.azurewebsites.net/).

The EconML SDK implements the following OrthoForest variants:

* DMLOrthoForest: suitable for continuous or discrete treatments

* DROrthoForest: suitable for discrete treatments

* CausalForest: suitable for both discrete and continuous treatments

In this notebook, we show the performance of the ORF on synthetic and observational data.

## Notebook Contents

1. [Example Usage with Continuous Treatment Synthetic Data](#1.-Example-Usage-with-Continuous-Treatment-Synthetic-Data)
2. [Example Usage with Binary Treatment Synthetic Data](#2.-Example-Usage-with-Binary-Treatment-Synthetic-Data)
3. [Example Usage with Multiple Treatment Synthetic Data](#3.-Example-Usage-with-Multiple-Treatment-Synthetic-Data)
4. [Example Usage with Real Continuous Treatment Observational Data](#4.-Example-Usage-with-Real-Continuous-Treatment-Observational-Data)

```python
import econml
```

```python
# Main imports
from econml.orf import DMLOrthoForest, DROrthoForest
from econml.dml import CausalForestDML
from econml.sklearn_extensions.linear_model import WeightedLassoCVWrapper, WeightedLasso, WeightedLassoCV

# Helper imports
import numpy as np
from itertools import product
from sklearn.linear_model import Lasso, LassoCV, LogisticRegression, LogisticRegressionCV
import matplotlib.pyplot as plt

%matplotlib inline
```

# 1. Example Usage with Continuous Treatment Synthetic Data

## 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 = 1000
n_w = 30
support_size = 5
n_x = 1
# 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])
T = np.dot(W[:, support_T], coefs_T) + eta_sample(n)
Y = TE * T + np.dot(W[:, support_Y], coefs_Y) + epsilon_sample(n)

# ORF parameters and test data
subsample_ratio = 0.3
lambda_reg = np.sqrt(np.log(n_w) / (10 * subsample_ratio * n))
X_test = np.array(list(product(np.arange(0, 1, 0.01), repeat=n_x)))
```

## 1.2. Train Estimator

**Note:** The models in the final stage of the estimation (``model_T_final``, ``model_Y_final``) need to support sample weighting.

If the models of choice do not support sample weights (e.g. ``sklearn.linear_model.LassoCV``), the ``econml`` packages provides a convenient wrapper for these models ``WeightedModelWrapper`` in order to allow sample weights.

```python
est = DMLOrthoForest(
    n_trees=1000, min_leaf_size=5,
    max_depth=50, subsample_ratio=subsample_ratio,
    model_T=Lasso(alpha=lambda_reg),
    model_Y=Lasso(alpha=lambda_reg),
    model_T_final=WeightedLasso(alpha=lambda_reg),
    model_Y_final=WeightedLasso(alpha=lambda_reg),
    global_residualization=False,
    random_state=123)
```

To use the built-in confidence intervals constructed via Bootstrap of Little Bags, we can specify `inference="blb"` at `fit` time or leave the default `inference='auto'` which will automatically use the Bootstrap of Little Bags.

```python
est.fit(Y, T, X=X, W=W, inference="blb")
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   17.8s
    [Parallel(n_jobs=-1)]: Done 248 tasks      | elapsed:   19.0s
    [Parallel(n_jobs=-1)]: Done 888 tasks      | elapsed:   21.9s
    [Parallel(n_jobs=-1)]: Done 985 out of 1000 | elapsed:   22.4s remaining:    0.2s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:   22.4s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.0s
    [Parallel(n_jobs=-1)]: Done 368 tasks      | elapsed:    1.7s
    [Parallel(n_jobs=-1)]: Done 984 tasks      | elapsed:    4.6s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:    4.6s finished

    <econml.orf._ortho_forest.DMLOrthoForest at 0x29e25d51438>

```python
# Calculate treatment effects
treatment_effects = est.effect(X_test)
```

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

```python
# Calculate default (90%) confidence intervals for the test data
te_lower, te_upper = est.effect_interval(X_test, alpha=0.01)
```

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

```python
res = est.effect_inference(X_test)
```

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

```python
res.summary_frame().head()
```

      point_estimate
      stderr
      zstat
      pvalue
      ci_lower
      ci_upper

      X

      0
      1.161
      0.183
      6.339
      0.0
      0.860
      1.462

      1
      1.171
      0.177
      6.628
      0.0
      0.881
      1.462

      2
      1.182
      0.171
      6.925
      0.0
      0.901
      1.462

      3
      1.192
      0.165
      7.228
      0.0
      0.921
      1.463

      4
      1.202
      0.160
      7.533
      0.0
      0.940
      1.465

```python
res.population_summary()
```

Uncertainty of Mean Point Estimate

  mean_point	 stderr_mean	 zstat	 pvalue	 ci_mean_lower	 ci_mean_upper

     3.179	      0.287	    11.06	   0.0	      2.706	         3.652

Distribution of Point Estimate

  std_point	 pct_point_lower	 pct_point_upper

    1.715	        1.212	           6.168

Total Variance of Point Estimate

  stderr_point	 ci_point_lower	 ci_point_upper

      1.739	         1.16	           6.247

Note: The stderr_mean is a conservative upper bound.

Similarly we can estimate effects and get confidence intervals and inference results using a `CausalForest`.

```python
est2 = CausalForestDML(model_t=Lasso(alpha=lambda_reg),
                       model_y=Lasso(alpha=lambda_reg),
                       n_estimators=4000, min_samples_leaf=5,
                       max_depth=50,
                       verbose=0, random_state=123)
est2.tune(Y, T, X=X, W=W)
est2.fit(Y, T, X=X, W=W)
treatment_effects2 = est2.effect(X_test)
te_lower2, te_upper2 = est2.effect_interval(X_test, alpha=0.01)
```

## 1.3. Performance Visualization

```python
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.title("ContinuousOrthoForest")
plt.plot(X_test, treatment_effects, label='ORF estimate')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.fill_between(X_test[:, 0], te_lower, te_upper, label="90% BLB CI", alpha=0.3)
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.subplot(1, 2, 2)
plt.title("CausalForest")
plt.plot(X_test, treatment_effects2, label='ORF estimate')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.fill_between(X_test[:, 0], te_lower2, te_upper2, label="90% BLB CI", alpha=0.3)
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.show()
```

# 2. Example Usage with Binary Treatment Synthetic Data

## 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
# DGP constants
np.random.seed(1234)
n = 1000
n_w = 30
support_size = 5
n_x = 1
# 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)

# ORF parameters and test data
subsample_ratio = 0.4
X_test = np.array(list(product(np.arange(0, 1, 0.01), repeat=n_x)))
```

## 2.2. Train Estimator

```python
est = DROrthoForest(
    n_trees=200, min_leaf_size=10,
    max_depth=30, subsample_ratio=subsample_ratio,
    propensity_model = LogisticRegression(C=1/(X.shape[0]*lambda_reg), penalty='l1', solver='saga'),
    model_Y = Lasso(alpha=lambda_reg),
    propensity_model_final=LogisticRegression(C=1/(X.shape[0]*lambda_reg), penalty='l1', solver='saga'),
    model_Y_final=WeightedLasso(alpha=lambda_reg)
)
```

```python
est.fit(Y, T, X=X, W=W)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   18.4s
    [Parallel(n_jobs=-1)]: Done 176 tasks      | elapsed:   19.4s
    [Parallel(n_jobs=-1)]: Done 200 out of 200 | elapsed:   19.5s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.1s
    [Parallel(n_jobs=-1)]: Done 200 out of 200 | elapsed:    0.7s finished

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

```python
# Calculate treatment effects for the default treatment points T0=0 and T1=1
treatment_effects = est.effect(X_test)
```

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

```python
# Calculate default (90%) confidence intervals for the default treatment points T0=0 and T1=1
te_lower, te_upper = est.effect_interval(X_test, alpha=0.01)
```

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

```python
est2 = CausalForestDML(model_y=Lasso(alpha=lambda_reg),
                       model_t=LogisticRegression(C=1/(X.shape[0]*lambda_reg)),
                       n_estimators=200, min_samples_leaf=5,
                       max_depth=50, max_samples=subsample_ratio/2,
                       discrete_treatment=True,
                       random_state=123)
est2.fit(Y, T, X=X, W=W, cache_values=True)
treatment_effects2 = est2.effect(X_test)
te_lower2, te_upper2 = est2.effect_interval(X_test, alpha=0.01)
```

```python
est2.summary()
```

    Population summary of CATE predictions on Training Data

Uncertainty of Mean Point Estimate

  mean_point	 stderr_mean	  zstat	 pvalue	 ci_mean_lower	 ci_mean_upper

     3.088	      0.157	    19.677	   0.0	      2.83	          3.346

Distribution of Point Estimate

  std_point	 pct_point_lower	 pct_point_upper

    1.757	        0.846	           6.851

Total Variance of Point Estimate

  stderr_point	 ci_point_lower	 ci_point_upper

      1.764	         0.89	           6.791

Doubly Robust ATE on Training Data Results

   	   point_estimate	 stderr	  zstat	 pvalue	 ci_lower	 ci_upper

  ATE	      3.158	      0.082	 38.551	   0.0	    3.023	    3.292

Doubly Robust ATT(T=0) on Training Data Results

   	   point_estimate	 stderr	  zstat	 pvalue	 ci_lower	 ci_upper

  ATT	       3.1	       0.096	 32.322	   0.0	    2.942	    3.258

Doubly Robust ATT(T=1) on Training Data Results

   	   point_estimate	 stderr	  zstat	 pvalue	 ci_lower	 ci_upper

  ATT	      3.218	      0.134	 23.965	   0.0	    2.997	    3.439

Note: The stderr_mean is a conservative upper bound.

## 2.3. Performance Visualization

```python
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.title("DiscreteTreatmentOrthoForest")
plt.plot(X_test, treatment_effects, label='ORF estimate')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.fill_between(X_test[:, 0], te_lower, te_upper, label="90% BLB CI", alpha=0.3)
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.subplot(1, 2, 2)
plt.title("CausalForest")
plt.plot(X_test, treatment_effects2, label='ORF estimate')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.fill_between(X_test[:, 0], te_lower2, te_upper2, label="90% BLB CI", alpha=0.3)
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.show()
```

# 3. Example Usage with Multiple Treatment Synthetic Data

## 3.1. DGP
We use the following DGP:

\begin{align}
Y = & \sum_{t=1}^{n_{\text{treatments}}} 1\{T=t\}\cdot \theta_{T}(X) + \langle W, \gamma\rangle + \epsilon, \; \epsilon \sim \text{Unif}(-1, 1), \\
\text{Pr}[T=t \mid W] \propto & \exp\{\langle W, \beta_t \rangle\}, \;\;\;\; \forall t\in \{0, 1, \ldots, n_{\text{treatments}}\}
\end{align}

where $W$ is a matrix of high-dimensional confounders, $\beta_t, \gamma$ are sparse.

For this particular example DGP we used $n_{\text{treatments}}=3$ and
\begin{align}
\theta_1(x) = & \exp( 2 x_1 ),\\
\theta_2(x) = &  3 \cdot \sigma(100\cdot (x_1 - .5)),\\
\theta_3(x) = & -2 \cdot \sigma(100\cdot (x_1 - .25)),
\end{align}
where $\sigma$ is the sigmoid function.

```python
def get_test_train_data(n, n_w, support_size, n_x, te_func, n_treatments):
    # 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, n_treatments))
    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([te_func(x_i, n_treatments) for x_i in X])
    log_odds = np.dot(W[:, support_T], coefs_T)
    T_sigmoid = np.exp(log_odds)
    T_sigmoid = T_sigmoid/np.sum(T_sigmoid, axis=1, keepdims=True)
    T = np.array([np.random.choice(n_treatments, p=p) for p in T_sigmoid])
    TE = np.concatenate((np.zeros((n,1)), TE), axis=1)
    Y = TE[np.arange(n), T] + np.dot(W[:, support_Y], coefs_Y) + epsilon_sample(n)
    X_test = np.array(list(product(np.arange(0, 1, 0.01), repeat=n_x)))

    return (Y, T, X, W), (X_test, np.array([te_func(x, n_treatments) for x in X_test]))
```

```python
import scipy.special
def te_func(x, n_treatments):
    return [np.exp(2*x[0]), 3*scipy.special.expit(100*(x[0] - .5)) - 1, -2*scipy.special.expit(100*(x[0] - .25))]

np.random.seed(123)
(Y, T, X, W), (X_test, te_test) = get_test_train_data(2000, 3, 3, 1, te_func, 4)
```

## 3.2. Train Estimator

```python
est = DROrthoForest(n_trees=500, model_Y = WeightedLasso(alpha=lambda_reg))
```

```python
est.fit(Y, T, X=X, W=W)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   31.1s
    [Parallel(n_jobs=-1)]: Done 112 tasks      | elapsed:   33.8s
    [Parallel(n_jobs=-1)]: Done 272 tasks      | elapsed:   37.9s
    [Parallel(n_jobs=-1)]: Done 500 out of 500 | elapsed:   44.1s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.5s
    [Parallel(n_jobs=-1)]: Done 208 tasks      | elapsed:    6.0s
    [Parallel(n_jobs=-1)]: Done 500 out of 500 | elapsed:   15.9s finished

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

```python
# Calculate marginal treatment effects
treatment_effects = est.const_marginal_effect(X_test)
```

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

```python
# Calculate default (90%) marginal confidence intervals for the test data
te_lower, te_upper = est.const_marginal_effect_interval(X_test, alpha=.01)
```

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

```python
res = est.const_marginal_effect_inference(X_test)
```

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

```python
res.summary_frame()
```

      point_estimate
      stderr
      zstat
      pvalue
      ci_lower
      ci_upper

      X
      T

      0
      T0_1
      1.012
      0.158
      6.420
      0.000
      0.753
      1.272

      T0_2
      -0.984
      0.155
      -6.362
      0.000
      -1.239
      -0.730

      T0_3
      0.036
      0.227
      0.159
      0.874
      -0.337
      0.410

      1
      T0_1
      1.017
      0.159
      6.417
      0.000
      0.756
      1.278

      T0_2
      -0.986
      0.153
      -6.457
      0.000
      -1.238
      -0.735

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      98
      T0_2
      1.966
      0.208
      9.435
      0.000
      1.624
      2.309

      T0_3
      -2.024
      0.163
      -12.447
      0.000
      -2.291
      -1.756

      99
      T0_1
      6.870
      0.239
      28.736
      0.000
      6.477
      7.264

      T0_2
      1.968
      0.212
      9.264
      0.000
      1.618
      2.317

      T0_3
      -2.019
      0.164
      -12.296
      0.000
      -2.289
      -1.749

300 rows × 6 columns

```python
est2 = CausalForestDML(model_y=Lasso(alpha=lambda_reg),
                       model_t=LogisticRegression(C=1/(X.shape[0]*lambda_reg)),
                       n_estimators=4000, min_samples_leaf=5,
                       max_depth=50, max_samples=subsample_ratio/2,
                       discrete_treatment=True,
                       random_state=123)
est2.fit(Y, T, X=X, W=W)
treatment_effects2 = est2.const_marginal_effect(X_test)
te_lower2, te_upper2 = est2.const_marginal_effect_interval(X_test, alpha=.01)
```

## 3.3. Performance Visualization

```python
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.title("DiscreteTreatmentOrthoForest")
y = treatment_effects
colors = ['b', 'r', 'g']
for it in range(y.shape[1]):
    plt.plot(X_test[:, 0], te_test[:, it], '--', label='True effect T={}'.format(it), color=colors[it])
    plt.fill_between(X_test[:, 0], te_lower[:, it], te_upper[:, it], alpha=0.3, color='C{}'.format(it))
    plt.plot(X_test, y[:, it], label='ORF estimate T={}'.format(it), color='C{}'.format(it))
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.subplot(1, 2, 2)
plt.title("CausalForest")
y = treatment_effects2
colors = ['b', 'r', 'g']
for it in range(y.shape[1]):
    plt.plot(X_test[:, 0], te_test[:, it], '--', label='True effect T={}'.format(it), color=colors[it])
    plt.fill_between(X_test[:, 0], te_lower2[:, it], te_upper2[:, it], alpha=0.3, color='C{}'.format(it))
    plt.plot(X_test, y[:, it], label='ORF estimate T={}'.format(it), color='C{}'.format(it))
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.show()
```

# 4. Example Usage with Real 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 `ContinuousTreatmentOrthoForest` 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)
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']]].values)
W2 = pd.get_dummies(oj_data[['brand']]).values
W = np.concatenate([W1, W2], axis=1)
X = oj_data[['INCOME']].values
```

## 4.2. Train Estimator

```python
# Define some parameters
n_trees = 1000
min_leaf_size = 50
max_depth = 20
subsample_ratio = 0.04
```

```python
est = DMLOrthoForest(
        n_trees=n_trees, min_leaf_size=min_leaf_size, max_depth=max_depth,
        subsample_ratio=subsample_ratio,
        model_T=Lasso(alpha=0.1),
        model_Y=Lasso(alpha=0.1),
        model_T_final=WeightedLassoCVWrapper(cv=3),
        model_Y_final=WeightedLassoCVWrapper(cv=3)
       )
```

```python
est.fit(Y, T, X=X, W=W)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   38.4s
    [Parallel(n_jobs=-1)]: Done 232 tasks      | elapsed:   39.4s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:   41.2s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.1s
    [Parallel(n_jobs=-1)]: Done 368 tasks      | elapsed:    1.9s
    [Parallel(n_jobs=-1)]: Done 984 tasks      | elapsed:    5.4s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:    5.5s finished

    <econml.orf._ortho_forest.DMLOrthoForest at 0x29e2c3095f8>

```python
min_income = 10.0
max_income = 11.1
delta = (max_income - min_income) / 100
X_test = np.arange(min_income, max_income + delta - 0.001, delta).reshape(-1, 1)
```

```python
# Calculate marginal treatment effects
treatment_effects = est.const_marginal_effect(X_test)
```

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

```python
# Calculate default (90%) marginal confidence intervals for the test data
te_upper, te_lower = est.const_marginal_effect_interval(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  17 tasks      | elapsed:    6.3s
    [Parallel(n_jobs=-1)]: Done 101 out of 101 | elapsed:   27.3s finished

```python
est2 = CausalForestDML(model_y=WeightedLassoCVWrapper(cv=3),
                       model_t=WeightedLassoCVWrapper(cv=3),
                       n_estimators=n_trees, min_samples_leaf=min_leaf_size, max_depth=max_depth,
                       max_samples=subsample_ratio/2,
                       random_state=123)
est2.fit(Y, T, X=X, W=W)
treatment_effects2 = est2.effect(X_test)
te_lower2, te_upper2 = est2.effect_interval(X_test)
```

## 4.3. Performance Visualization

```python
# Plot Orange Juice elasticity as a function of income
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.plot(X_test.flatten(), treatment_effects, label="OJ Elasticity")
plt.fill_between(X_test.flatten(), te_lower, te_upper, label="90% BLB CI", alpha=0.3)
plt.xlabel(r'$\log$(Income)')
plt.ylabel('Orange Juice Elasticity')
plt.legend()
plt.title("Orange Juice Elasticity vs Income: ContinuousTreatmentOrthoForest")
plt.subplot(1, 2, 2)
plt.plot(X_test.flatten(), treatment_effects2, label="OJ Elasticity")
plt.fill_between(X_test.flatten(), te_lower2, te_upper2, label="90% BLB CI", alpha=0.3)
plt.xlabel(r'$\log$(Income)')
plt.ylabel('Orange Juice Elasticity')
plt.legend()
plt.title("Orange Juice Elasticity vs Income: CausalForest")
plt.show()
```

```python

```

notebooks/Causal Forest and Orthogonal Random Forest Examples.ipynb


# 附录 E：Lagrange 乘子

## 等式约束下的拉格朗日乘子

拉格朗日乘子（`Lagrange multipliers`），又叫不定乘子（`undetermined multipliers`），用于找到多变量函数在有约束条件下的驻点。

考虑最大化函数 $f(x_1, x_2)$，约束条件为

$$
g(x_1, x_2) = 0
$$

一个可能的方法是通过这个约束求出 $x_2 = h(x_1)$ 再带入函数消去 $x_2$，但很多时候函数 $h$ 是很难找到的。

一个更普遍和简单的方法是引入一个拉格朗日乘子 $\lambda$。

现在考虑一个 $D$ 维变量 $\mathbf x$ 的优化问题：

$$
\begin{align}
\max_{\mathbf x}~& f(\mathbf x) \\
s.t.~& g(\mathbf x) = 0
\end{align}
$$

$g(\mathbf x) = 0$ 定义了空间中一个 $D-1$ 维的超平面。

我们注意到，梯度 $\triangledown g(\mathbf x)$ 平行于这个超平面 $g(\mathbf x) = 0$。

考虑超平面上的两点 $\mathbf x, \mathbf{x+\epsilon}$，做泰勒展开有：

$$
g(\mathbf{x+\epsilon}) \approx g(\mathbf{x}) + \mathbf\epsilon^{\text{T}} \triangledown g(\mathbf x)
$$

因为它们都在超平面上，所以有 $\mathbf\epsilon^{\text{T}} \triangledown g(\mathbf x) \approx 0$，令 $\|\mathbf \epsilon\|\to 0$，我们有 $\mathbf\epsilon^{\text{T}} \triangledown g(\mathbf x) = 0$，而 $\epsilon$ 平行于超平面，所以 $\triangledown g(\mathbf x)$ 平行于超平面。

在这个超平面上，最大化 $f(\mathbf x)$ 的点 $\mathbf x^\star$ 也需要满足 $\triangledown f(\mathbf x^\star)$ 垂直于超平面 $g(\mathbf x) = 0$，否则，我们总可以在这个超平面上移动一点距离使得 $f(\mathbf x)$ 增大。

因此 $\triangledown f(\mathbf x^\star)$ 与 $\triangledown g(\mathbf x^\star)$ 应当平行，这意味着存在 $\lambda \neq 0$，使得：

$$
\triangledown f + \lambda \triangledown g = 0
$$

此时我们可以引入一个拉格朗日函数：

$$
L(\mathbf x, \lambda) \equiv f(\mathbf x) + \lambda g(\mathbf x)
$$

其驻点满足 $\triangledown_x L = 0$ 即 $
\triangledown f + \lambda \triangledown g = 0
$，以及偏导 $\frac{\partial L}{\partial \lambda}$ 使得 $g(\mathbf x) = 0$。

因此，问题转化为求这个拉格朗日函数的驻点 $(x^\star, \lambda)$。

由于拉格朗日函数是无约束的，我们很容易通过求偏导得到驻点的值。

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

%matplotlib inline

fig, ax = plt.subplots()

an = np.linspace(0, 2*np.pi, 100)

xx, yy = -20 / (3 - np.cos(an)), -7 / (2 + np.sin(an)) + np.cos(an)

ax.plot(xx, yy, 'r')

ax.set_xlim(-11, -3)
ax.set_ylim(-8, -1)

x, y = xx[30], yy[30]

ax.annotate("",
            xy=(x, y),
            xytext=(x+0.4, y+1),
            arrowprops=dict(arrowstyle="<-",
                            connectionstyle="arc3")
           )

ax.annotate("",
            xy=(x, y),
            xytext=(x-0.4, y-1),
            arrowprops=dict(arrowstyle="<-",
                            connectionstyle="arc3")
           )

ax.text(-5, -7, "$g(\mathbf{x}) = 0$", fontsize="xx-large")

ax.text(x+0.5, y+1, r"$\triangledown f(\mathbf{x})$", fontsize="xx-large")
ax.text(x-0.8, y-1.5, r"$\triangledown g(\mathbf{x})$", fontsize="xx-large")
ax.text(x+0.2, y, '$x_A$', fontsize="xx-large")

ax.set_axis_off()

ax.plot([x], [y], 'k.', linewidth=30)

plt.show()
```

## 例子

考虑这样的问题：

$$
\begin{align}
\max_{x_1, x_2}~& f(x_1, x_2)=1-x_1^2-x_2^2 \\
s.t.~& g(x_1,x_2) = x_1 + x_2 - 1 = 0
\end{align}
$$

对应的拉格朗日函数为

$$
L(\mathbf x, \lambda) = 1-x_1^2-x_2^2 + \lambda(x_1 + x_2 - 1)
$$

驻点满足：

$$
\begin{align}
-2x_1 + \lambda & = 0\\
-2x_2 + \lambda & = 0\\
x_1 + x_2 - 1 & = 0 \\
\end{align}
$$

因此得到 $(x_1^\star, x_2^\star)=(\frac 1 2, \frac 1 2)$，$\lambda = 1$。

```python
an = np.linspace(0, 2*np.pi, 100)

xx = np.linspace(-1, 2)
yy = 1 - xx

fig, ax = plt.subplots()

for i in xrange(1, 4):
    ax.plot(i / np.sqrt(8.0) * np.cos(an), i / np.sqrt(8.0) * np.sin(an), 'b')

ax.plot(xx, yy, 'r')

ax.axis('equal')

ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')

ax.xaxis.set_ticks_position('bottom')
ax.spines['bottom'].set_position(('data',0))

ax.yaxis.set_ticks_position('left')
ax.spines['left'].set_position(('data',0))

ax.set_xticks([])
ax.set_yticks([])

ax.text(2.6, -0.2, "$x_1$", fontsize="xx-large")
ax.text(-0.3, 1.8, "$x_2$", fontsize="xx-large")
ax.text(1, -1.3, "$g(x_1, x_2) = 0$", fontsize="xx-large")

ax.annotate(r"$(x_1^\star, x_2^\star)$",
            fontsize="xx-large",
            xytext=(1, 1),
            xy=(0.5, 0.5),
            arrowprops=dict(arrowstyle="->",
                            connectionstyle="arc3")
           )

plt.show()
```

## 带不等式约束的拉格朗日乘子

现在考虑不等式约束的问题：

$$
\begin{align}
\max_{\mathbf x}~& f(\mathbf x) \\
s.t.~& g(\mathbf x) \geq 0
\end{align}
$$

分成两种情况考虑：

- $g(\mathbf x) < 0$，此时不等式约束不起作用（`inactive`），驻点的条件为 $\triangledown f(\mathbf x)=0$，因此必须有 $\lambda=0$。

- $g(\mathbf x) = 0$，此时不等式约束起作用（`active`），驻点的条件为 $\triangledown f(\mathbf x) + \lambda \triangledown g(\mathbf x)=0$，且 $\lambda \neq 0$。除此之外，$\triangledown f(\mathbf x)$ 应当与 $\triangledown g(\mathbf x)$ 的方向（指向 $g(\mathbf x)>0$ 的区域）相反才能得到最大值（即 $f$ 的梯度方向要远离 $g(\mathbf x)>0$ 的区域），因此我们必然有 $\lambda > 0$。

在这两种情况下，我们都有 $\lambda g(\mathbf x) = 0$，因此问题转换为求解拉格朗日函数

$$
L(\mathbf x, \lambda) \equiv f(\mathbf x) + \lambda g(\mathbf x)
$$

在约束

$$
\begin{align}
g(\mathbf x)&\geq 0\\
\lambda &\geq 0\\
\lambda g(\mathbf x) &= 0
\end{align}
$$

下的驻点，这些条件就是著名的 KKT 条件。

对于最小化问题，我们只需要将拉格朗日函数变为：

$$
L(\mathbf x, \lambda) \equiv f(\mathbf x) - \lambda g(\mathbf x)
$$

```python
fig, ax = plt.subplots()

an = np.linspace(0, 2*np.pi, 100)

xx, yy = -20 / (3 - np.cos(an)), -7 / (2 + np.sin(an)) + np.cos(an)

ax.plot(xx, yy, 'r')

ax.fill(xx, yy, color='gold')

ax.set_xlim(-11, -3)
ax.set_ylim(-8, -1)

x, y = xx[30], yy[30]

ax.annotate("",
            xy=(x, y),
            xytext=(x+0.4, y+1),
            arrowprops=dict(arrowstyle="<-",
                            connectionstyle="arc3")
           )

ax.annotate("",
            xy=(x, y),
            xytext=(x-0.4, y-1),
            arrowprops=dict(arrowstyle="<-",
                            connectionstyle="arc3")
           )

ax.text(-5, -7, "$g(\mathbf{x}) = 0$", fontsize="xx-large")
ax.text(-9, -3.5, "$g(\mathbf{x}) > 0$", fontsize="xx-large")

ax.text(x+0.5, y+1, r"$\triangledown f(\mathbf{x})$", fontsize="xx-large")
ax.text(x-0.8, y-1.5, r"$\triangledown g(\mathbf{x})$", fontsize="xx-large")
ax.text(x+0.2, y, '$x_A$', fontsize="xx-large")

ax.plot([x], [y], 'k.', linewidth=30)

ax.set_axis_off()

plt.show()
```

## 多约束问题

考虑这个问题：

$$
\begin{align}
\max_{\mathbf x}~& f(\mathbf x) &\\
s.t.~& g_j(\mathbf x) \geq 0, &&j = 1,\dots,J \\
& h_k(\mathbf x) \geq 0, &&k = 1,\dots,K \\
\end{align}
$$

很容易推广得到，拉格朗日函数为：

$$
\begin{align}
& L(\mathbf {x, \lambda, \mu}) \equiv f(\mathbf x) + \sum_{j=1}^J \lambda_j g(\mathbf x) + \sum_{k=1}^K \mu_k h(\mathbf x) \\
s.t.~& \mu_k \geq 0, \mu_k h(\mathbf x) = 0, k = 1,\dots,K
\end{align}
$$


# Parametric Model-Based regression

    Notebook version: 1.3 (Sep 20, 2019)

    Author: Jesús Cid-Sueiro (jesus.cid@uc3m.es)
            Jerónimo Arenas García (jarenas@tsc.uc3m.es)

    Changes: v.1.0 - First version, expanding some cells from the Bayesian Regression
                     notebook
             v.1.1 - Python 3 version.
             v.1.2 - Revised presentation.
             v.1.3 - Updated index notation

    Pending changes: * Include regression on the stock data

```python
# Import some libraries that will be necessary for working with data and displaying plots

# To visualize plots in the notebook
%matplotlib inline

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import scipy.io       # To read matlab files
import pylab
```

## A quick note on the mathematical notation

In this notebook we will make extensive use of probability distributions. In general, we will use capital leters
${\bf X}$, $S$, $E$ ..., to denote random variables, and lower-case letters ${\bf x}$, $s$, $\epsilon$ ..., to denote the values they can take.

In general, we will use letter $p$ for probability density functions (pdf). When necessary, we will use, capital subindices to make the random variable explicit. For instance, $p_{{\bf X}, S}({\bf x}, s)$ would be the joint pdf of random variables ${\bf X}$ and $S$ at values ${\bf x}$ and $s$, respectively.

However, to avoid a notation overload, we will omit subindices when they are clear from the context. For instance, we will use $p({\bf x}, s)$ instead of $p_{{\bf X}, S}({\bf x}, s)$.

## 1. Model-based parametric regression

### 1.1. The regression problem

Given an observation vector ${\bf x}$, the goal of the regression problem is to find a function $f({\bf x})$ providing *good* predictions about some unknown variable $s$. To do so, we assume that a set of *labelled* training examples, $\{{\bf x}_k, s_k\}_{k=0}^{K-1}$ is available.

The predictor function should make good predictions for new observations ${\bf x}$ not used during training. In practice, this is tested using a second set (the *test set*) of labelled samples.

### 1.2. The underlying model assumption

Many regression algorithms are grounded on the idea that all samples from the training set have been generated
independently by some common stochastic process.

If $p({\bf x}, s)$ were known, we could apply estimation theory to estimate $s$ for a given ${\bf x}$ using $p$. For instance, we could apply any of the following classical estimates:

   * Maximum A Posterior (MAP): $$\hat{s}_{\text{MAP}} = \arg\max_s p(s| {\bf x})$$
   * Minimum Mean Square Error (MSE): $$\hat{s}_{\text{MSE}} = \mathbb{E}\{S |{\bf x}\} = \int s \, p(s| {\bf x}) \, ds $$

Note that, since these estimators depend on $p(s |{\bf x})$, knowing the posterior distribution of the target variable is enough, and we do not need to know the joint distribution $p({\bf x}, s)$.

More importantly, note that **if we knew the underlying model, we would not need the data** in ${\cal D}$ to make predictions on new data.

#### Exercise 1:

Assume the target variable $s$ is a scaled noisy version of the input variable $x$:
$$
s = 2 x + \epsilon
$$
where $\epsilon$ is Gaussian a noise variable with zero mean and unit variance, which does not depend on $x$.

  1. Compute the target model $p(s| x)$
  2. Compute prediction $\hat{s}_\text{MAP}$ for an arbitrary input $x$
  3. Compute prediction $\hat{s}_\text{MSE}$ for an arbitrary input $x$
  4. Compute prediction $\hat{s}_\text{MSE}$ for input $x=4$

#### Solution:

[comment]: # (<SOL>)
[comment]: # (</SOL>)

### 1.3. Model-based regression

In practice, the underlying model is usually unknown.

Model based-regression methods exploit the idea of using the training data to estimate the posterior distribution $p(s|{\bf x})$ and then apply estimation theory to make predictions.

### 1.4. Parametric model-based regression

In some cases, we may have a partial knowledge about the underlying mode. In this notebook we will assume that $p$ belongs to a parametric family of distributions $p(s|{\bf x},{\bf w})$, where ${\bf w}$ is some unknown parameter.

#### Exercise 2:

Assume the target variable $s$ is a scaled noisy version of the input variable $x$:
$$
s = w x + \epsilon
$$
where $\epsilon$ is Gaussian a noise variable with zero mean and unit variance, which does not depend on $x$. Assume that $w$ is known.
  1. Compute the target model $p(s| x, w)$
  2. Compute prediction $\hat{s}_\text{MAP}$ for an arbitrary input $x$
  3. Compute prediction $\hat{s}_\text{MSE}$ for an arbitrary input $x$

#### Solution:

[comment]: # (<SOL>)
[comment]: # (</SOL>)

We will use the training data to estimate ${\bf w}$

The estimation of ${\bf w}$ from a given dataset $\mathcal{D}$ is the goal of the following sections

## 2. Maximum Likelihood parameter estimation.

The ML (Maximum Likelihood) principle is well-known in statistics and can be stated as follows: take the value of the parameter to be estimated (in our case, ${\bf w}$) that best explains the given observations (in our case, the training dataset $\mathcal{D}$). Mathematically, this can be expressed as follows:
$$
\hat{\bf w}_{\text{ML}} = \arg \max_{\bf w} p(\mathcal{D}|{\bf w})
$$

#### Exercise 3:

All samples in dataset ${\cal D} = \{(x_k, s_k), k=0,\ldots,K-1 \}$
$$
s_k = w \cdot x_k + \epsilon_k
$$
where $\epsilon_k$ are i.i.d. (independent and identically distributed) Gaussian noise random variables with zero mean and unit variance, which do not depend on $x_k$.

Compute the ML estimate, $\hat{w}_{\text{ML}}$, of $w$.

#### Solution:

[comment]: # (<SOL>)
# </SOL>

* **4.2.** Compute the ML estimate

```python
# wML = <FILL IN>

print("The ML estimate is {}".format(wML))
```

* **4.3.** Plot the likelihood as a function of parameter $w$ along the interval $-0.5\le w \le 2$, verifying that the ML estimate takes the maximum value.

```python
sigma_eps = 1
K = len(s)
wGrid = np.arange(-0.5, 2, 0.01)

p = []
for w in wGrid:
    d = s - X*w
    # p.append(<FILL IN>)

# Compute the likelihood for the ML parameter wML
# d = <FILL IN>
# pML = [<FILL IN>]

# Plot the likelihood function and the optimal value
plt.figure()
plt.plot(wGrid, p)
plt.stem([wML], pML)
plt.xlabel('$w$')
plt.ylabel('Likelihood function')
plt.show()
```

* **4.4.** Plot the prediction function on top of the data scatter plot

```python
xgrid = np.arange(0, 1.2, 0.01)
# sML = <FILL IN>

plt.figure()
plt.scatter(X, s)
# plt.plot(<FILL IN>)
plt.xlabel('x')
plt.ylabel('s')
plt.axis('tight')
plt.show()
```

### 2.1. Model assumptions

In order to solve exercise 4 we have taken advantage of the statistical independence of the noise components. Some independence assumptions are required in general to compute the ML estimate in other scenarios.

In order to estimate ${\bf w}$ from the training data in a mathematicaly rigorous and compact form let us group the target variables into a vector
$$
{\bf s} = \left(s_0, \dots, s_{K-1}\right)^\top
$$
and the input vectors into a matrix
$$
{\bf X} = \left({\bf x}_0, \dots, {\bf x}_{K-1}\right)^\top
$$

We will make the following assumptions:

   * A1. All samples in ${\cal D}$ have been generated by the same distribution, $p({\bf x}, s \mid {\bf w})$
   * A2. Input variables ${\bf x}$ do not depend on ${\bf w}$. This implies that
$$
p({\bf X} \mid {\bf w}) = p({\bf X})
$$
   * A3. Targets $s_{0},\ldots, s_{K-1}$ are statistically independent, given ${\bf w}$ and the inputs ${\bf x}_0,\ldots, {\bf x}_{K-1}$, that is:
$$
p({\bf s} \mid {\bf X}, {\bf w}) = \prod_{k=0}^{K-1} p(s_k \mid {\bf x}_k, {\bf w})
$$

Since ${\cal D} = ({\bf X}, {\bf s})$, we can write
$$p(\mathcal{D}|{\bf w})
    = p({\bf s}, {\bf X}|{\bf w})
    = p({\bf s} | {\bf X}, {\bf w}) p({\bf X}|{\bf w})
$$
Using assumption A2,
$$
p(\mathcal{D}|{\bf w})
    = p({\bf s} | {\bf X}, {\bf w}) p({\bf X})
$$

and, finally, using assumption A3, we can express the estimation problem as the computation of

\begin{align}
\hat{\bf w}_{\text{ML}}
                &= \arg \max_{\bf w} p({\bf s}|{\bf X},{\bf w})  \\
   \qquad \quad &= \arg \max_{\bf w} \prod_{k=0}^{K-1} p(s_k \mid {\bf x}_k, {\bf w})  \\
   \qquad \quad &= \arg \max_{\bf w} \sum_{k=0}^{K-1}\log  p(s_k \mid {\bf x}_k, {\bf w})
\end{align}

Any of the last three terms can be used to optimize ${\bf w}$. The sum in the last term is usually called the **log-likelihood** function, $L({\bf w})$, whereas the product in the previous line is simply referred as the **likelihood** function.

### 2.2. Summary.

Let's summarize what we need to do in order to design a regression algorithm based on ML estimation:

1. Assume a parametric data model $p(s| {\bf x},{\bf w})$
2. Using the data model and the i.i.d. assumption, compute $p({\bf s}| {\bf X},{\bf w})$.
3. Find an expression for ${\bf w}_{\text{ML}}$
4. Assuming ${\bf w} = {\bf w}_{\text{ML}}$, compute the MAP or the minimum MSE estimate of $s$ given ${\bf x}$.

## 3. ML estimation for a Gaussian model.

### 3.1. Step 1: The Gaussian generative model

Let us assume that the target variables $s_k$ in dataset $\mathcal{D}$ are given by
$$
s_k = {\bf w}^\top {\bf z}_k + \varepsilon_k
$$

where ${\bf z}_k$ is the result of some transformation of the inputs, ${\bf z}_k = T({\bf x}_k)$, and $\varepsilon_k$ are i.i.d. instances of a Gaussian random variable with mean zero and varianze $\sigma_\varepsilon^2$, i.e.,
$$
p_E(\varepsilon) = \frac{1}{\sqrt{2\pi}\sigma_\varepsilon}
                   \exp\left(-\frac{\varepsilon^2}{2\sigma_\varepsilon^2}\right)
$$

Assuming that the noise variables are independent of ${\bf x}$ and ${\bf w}$, then, for a given ${\bf x}$ and ${\bf w}$, the target variable is gaussian with mean ${\bf w}^\top {\bf z}_k$ and variance $\sigma_\varepsilon^2$
$$
p(s|{\bf x}, {\bf w}) = p_E(s-{\bf w}^\top{\bf z}) =
    \frac{1}{\sqrt{2\pi}\sigma_\varepsilon}
    \exp\left(-\frac{(s-{\bf w}^\top{\bf z})^2}{2\sigma_\varepsilon^2}\right)
$$

### 3.2. Step 2: Likelihood function

Now we need to compute the likelihood function $p({\bf s}, {\bf X} | {\bf w})$. If the samples are i.i.d. we can write
$$
p({\bf s}| {\bf X}, {\bf w})
    = \prod_{k=0}^{K-1} p(s_k| {\bf x}_k, {\bf w})
    = \prod_{k=0}^{K-1} \frac{1}{\sqrt{2\pi}\sigma_\varepsilon}
    \exp\left(-\frac{\left(s_k-{\bf w}^\top{\bf z}_k\right)^2}{2\sigma_\varepsilon^2}\right) \\
    = \left(\frac{1}{\sqrt{2\pi}\sigma_\varepsilon}\right)^K
      \exp\left(-\sum_{k=1}^K \frac{\left(s_k-{\bf w}^\top{\bf z}_k\right)^2}{2\sigma_\varepsilon^2}\right) \\
$$

Finally, grouping variables ${\bf z}_k$ in
$${\bf Z} = \left({\bf z}_0, \dots, {\bf z}_{K-1}\right)^\top$$
we get
$$
p({\bf s}| {\bf X}, {\bf w})
    = \left(\frac{1}{\sqrt{2\pi}\sigma_\varepsilon}\right)^K
      \exp\left(-\frac{1}{2\sigma_\varepsilon^2}\|{\bf s}-{\bf Z}{\bf w}\|^2\right)
$$

### 3.3. Step 3: ML estimation.

The <b>maximum likelihood</b> solution is then given by:
$$
{\bf w}_{ML} = \arg \max_{\bf w} p({\bf s}|{\bf w}) = \arg \min_{\bf w} \|{\bf s} - {\bf Z}{\bf w}\|^2
$$
Note that $\|{\bf s} - {\bf Z}{\bf w}\|^2$ is the sum or the squared prediction errors (Sum of Squared Errors, SSE) for all samples in the dataset. This is also called the * **Least Squares** * (LS) solution.

The LS solution can be easily computed by differentiation,
$$
\nabla_{\bf w} \|{\bf s} - {\bf Z}{\bf w}\|^2\Bigg|_{{\bf w} = {\bf w}_\text{ML}}
= - 2 {\bf Z}^\top{\bf s} + 2 {\bf Z}^\top{\bf Z} {\bf w}_{\text{ML}}
= {\bf 0}
$$
and it is equal to
$$
{\bf w}_\text{ML} = ({\bf Z}^\top{\bf Z})^{-1}{\bf Z}^\top{\bf s}
$$

### 3.4. Step 4: Prediction function.

The last step consists on computing an estimate of $s$ by assuming that the true value of the weight parameters is ${\bf w}_\text{ML}$. In particular, the minimum MSE estimate is
$$
\hat{s}_\text{MSE} = \mathbb{E}\{s|{\bf x},{\bf w}_\text{ML}\}
$$

Knowing that, given ${\bf x}$ and ${\bf w}$, $s$ is normally distributed with mean ${\bf w}^\top {\bf z}$ we can write
$$
\hat{s}_\text{MSE} = {\bf w}_\text{ML}^\top {\bf z}
$$

#### Exercise 5:

Assume that the targets in the one-dimensional dataset given by

```python
X = np.array([0.15, 0.41, 0.53, 0.80, 0.89, 0.92, 0.95])
s = np.array([0.09, 0.16, 0.63, 0.44, 0.55, 0.82, 0.95])
```

have been generated by the polynomial Gaussian model
$$
s = w_0 + w_1 x + w_2 x^2 + \epsilon
$$
(i.e., with ${\bf z} = T(x) = (1, x, x^2)^\intercal$) with noise variance

```python
sigma_eps = 0.3
```

* **5.1.** Compute the ML estimate.

```python
# Compute the extended input matrix Z
nx = len(X)
# Z = <FILL IN>

# Compute the ML estimate using linalg.lstsq from Numpy.
# wML = <FILL IN>
print(wML)

```

* **5.2.** Compute the value of the log-likelihood function for ${\bf w}={\bf w}_\text{ML}$.

```python
K = len(s)

# Compute the likelihood for the ML parameter wML
# d = <FILL IN>

# LwML = [<FILL IN>]

print(LwML)
```

* **5.3.** Plot the prediction function over the data scatter plot

```python
xgrid = np.arange(0, 1.2, 0.01)
nx = len(xgrid)
# Compute the input matrix for the grid data in x
# Z = <FILL IN>

# sML = <FILL IN>

plt.figure()
plt.scatter(X, s)
# plt.plot(<FILL IN>)
plt.xlabel('x')
plt.ylabel('s')
plt.axis('tight')
plt.show()
```

#### Exercise 6:

Assume the dataset $\mathcal{D} = \{(x_k, s_k, k=0,\ldots, K-1\}$ contains i.i.d. samples from a distribution with posterior density given by
$$
p(s \mid x, w) = w x \exp(- w x s), \qquad s\ge0, \,\, x\ge 0, \,\, w\ge 0
$$

* **6.1.** Determine an expression for the likelihood function

**Solution**:
<SOL>
</SOL>

* **6.2.** Draw the likelihood function for the dataset in **Exercise 4** in the range $0\le w\le 6$.

```python
K = len(s)
wGrid = np.arange(0, 6, 0.01)

p = []
Px = np.prod(X)
xs = np.dot(X,s)
for w in wGrid:
    # p.append(<FILL IN>)

plt.figure()
# plt.plot(<FILL IN>)
plt.xlabel('$w$')
plt.ylabel('Likelihood function')
plt.show()
```

* **6.3.** Determine the maximum likelihood coefficient, $w_\text{ML}$.

(*Hint: you can maximize the log-likelihood function instead of the likelihood function in order to simplify the differentiation*)

**Solution**:

<SOL>
</SOL>

* **6.4.** Compute $w_\text{ML}$ for the dataset in **Exercise 4**

```python
# wML = <FILL IN>
print(wML)
```

* **6.5.** Assuming $w = w_\text{ML}$, compute the prediction function based on the estimate $s_{MSE}$

**Solution**:

<SOL>
</SOL>

* **6.6.** Plot the prediction function obtained in ap. 6.5, and compare it with the linear predictor in exercise 4

```python
xgrid = np.arange(0.1, 1.2, 0.01)
# sML = <FILL IN>

plt.figure()
plt.scatter(X, s)
# plt.plot(<FILL IN>)
plt.xlabel('x')
plt.ylabel('s')
plt.axis('tight')
plt.show()
```

Subjectively, we can see that the predictor computed in exercise 6 does not fit the given data very well. This could be a false perception. If the data have been truly generated by the parametric model assumed in exercise 6 (i.e. $p(s \mid x, w) = w x \exp(- w x s)$, the apparent missbehavior of the estimator could be caused by the natural randomness of the data, and a greater amount of data would show a better adjustement.

Alternative, it may be the case the model assumed in sec. 6 is incorrect. Again, more data would be useful to asses that.

This shows that the choice of the data model is important. In many applications, no parametric data model is available, and the data scientist must make a choice based on the nature of the data or any previous knowledge about the statistical behavior of the data.

If no previous information is available, the data scientist can try different models, and compare using validation data and some cross validation technique.


# Setting up Polytropic [TOV](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation) Initial Data

## Authors: Phil Chang, Zachariah B. Etienne

## This module sets up initial data for a [TOV](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation) star in *spherical, isotropic coordinates*

**This module has been validated to exhibit convergence to zero of the Hamiltonian constraint violation at the expected order to the exact solution (see [start-to-finish TOV module](Tutorial-Start_to_Finish-BSSNCurvilinear-Setting_up_TOV_initial_data.ipynb) for full test). Note that convergence at the surface of the star is lower order due to the sharp drop to zero in $T^{\mu\nu}$.**

### NRPy+ Source Code for this module: [TOV/TOV_Solver.py](../edit/TOV/TOV_Solver.py)

$$\label{top}$$

## Introduction

[Tolman-Oppenheimer-Volkoff (TOV)](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation) initial data are constructed from the solution of a system of [ODEs](https://en.wikipedia.org/wiki/Ordinary_differential_equation). The solution is spherically symmetric.

This module first constructs the system of ODEs and uses [NumPy](http://www.numpy.org/) to numerically solve the system and output the result to a file.

Then, a small C code is written to read in the file, interpolate the data to any desired radius, and construct the ADM spacetime quantities and corresponding $T^{\mu\nu}$.

## Table of Contents

1. [Step 1](#step1): Overview of the [Tolman-Oppenheimer-Volkoff (TOV)](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation) solution to Einstein's equations
1. [Step 2](#step2): Python code for solving the TOV equations using NumPy's ODE solver, outputting the data to a file
    1. [Step 2.1](#step2p1): Validation of the TOV output file against a trusted version
1. [Step 3](#step3): C code for reading in data file and interpolating data to any desired radius
    1. [Step 3.1](#step3p1): Reading in the data file, generated in [Step 2](#step2).
    1. [Step 3.2](#step3p2): Interpolating the data to any desired radius
1. [Step 4](#step4): Construct ADM spacetime quantities and $T^{\mu\nu}$ from TOV solution
    1. [Step 4.1](#step4p1): Load NRPy+ modules and set basic parameters
    1. [Step 4.2](#step4p2): Constructing ADM spacetime quantities from the TOV data
    1. [Step 4.3](#step4p3): Constructing $T^{\mu\nu}$ from the TOV data
1. [Step 5](#step5): NRPy+ code validation: Confirm that ADM spacetime quantities and $T^{\mu\nu}$

## Step 1: Overview of the [Tolman-Oppenheimer-Volkoff (TOV)](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation) solution to Einstein's equations \[Back to [top](#top)\]
$$\label{step1}$$

The [TOV line element](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation) in terms of the *Schwarzschild coordinate* $r$ is written (in the $-+++$ form):
$$
ds^2 = - c^2 e^\nu dt^2 + \left(1 - \frac{2Gm}{rc^2}\right)^{-1} dr^2 + r^2 d\Omega^2,
$$
where $m(r)$ is the mass-energy enclosed at a given $r$, and is equal to the total star's mass outside the stellar radius $r=R$.

In terms of the *isotropic coordinate* $\bar{r}$ with $G=c=1$ (i.e., the coordinate system and units we'd prefer to use), the ($-+++$ form) line element is written:
$$
ds^2 = - e^{\nu} dt^2 + e^{4\phi} \left(d\bar{r}^2 + \bar{r}^2 d\Omega^2\right),
$$
where $\phi$ here is the *conformal factor*.

Setting components of the above line element equal to one another, we get (in $G=c=1$ units):

\begin{align}
r^2 &= e^{4\phi} \bar{r}^2 \implies e^{4\phi} = \frac{r^2}{\bar{r}^2} \\
\left(1 - \frac{2m}{r}\right)^{-1} dr^2 &= e^{4\phi} d\bar{r}^2 \\
\implies \frac{d\bar{r}(r)}{dr} &= \left(1 - \frac{2m}{r} \right)^{-1/2} \frac{\bar{r}(r)}{r}.
\end{align}

The TOV equations provide radial ODEs for the pressure and $\nu$ (from [the Wikipedia article on the TOV solution](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation)):

\begin{align}
\frac{dP}{dr} &= - \frac{1}{r} \left( \frac{\rho + P}{2} \right) \left(\frac{2 m}{r} + 8 \pi r^2 P\right) \left(1 - \frac{2 m}{r}\right)^{-1} \\
\frac{d \nu}{d r} &= \frac{1}{r}\left(1 - \frac{2 m}{r}\right)^{-1} \left(\frac{2 m}{r} + 8 \pi r^2 P\right) \\
\end{align}

Assuming a polytropic equation of state, which relates the pressure $P$ to the baryonic rest-mass density $\rho_B$,

$$
P(\rho_B) = K \rho_B^\Gamma,
$$
the specific internal energy will be given by
$$
\epsilon = \frac{P}{\rho_B (\Gamma - 1)},
$$

so the total mass-energy density $\rho$ is given by
$$
\rho = \rho_B (1 + \epsilon).
$$

Given this, the mass-energy $m(r)$ density is the solution to the ODE:
$$
\frac{m(r)}{dr} = 4\pi r^2 \rho(r)
$$

Thus the full set of ODEs that need to be solved is given by

\begin{align}
\frac{dP}{dr} &= - \frac{1}{r} \left( \frac{\rho + P}{2} \right) \left(\frac{2 m}{r} + 8 \pi r^2 P\right) \left(1 - \frac{2 m}{r}\right)^{-1} \\
\frac{d \nu}{d r} &= \frac{1}{r}\left(1 - \frac{2 m}{r}\right)^{-1} \left(\frac{2 m}{r} + 8 \pi r^2 P\right) \\
\frac{m(r)}{dr} &= 4\pi r^2 \rho(r) \\
\frac{d\bar{r}(r)}{dr} &= \left(1 - \frac{2m}{r} \right)^{-1/2} \frac{\bar{r}(r)}{r}.
\end{align}

## Step 2: Python code for solving the TOV equations using NumPy's ODE solver, outputting the data to a file \[Back to [top](#top)\]

$$\label{step2}$$

The following code solves the above equations, and was largely written by **Phil Chang**.

```python
## TOV SOLVER FOR SIMPLE POLYTROPES.
## Authors: Phil Chang, Zachariah B. Etienne

# Full documentation for this module may be found in the NRPy+ tutorial Jupyter notebook:
#  Tutorial-Start_to_Finish-BSSNCurvilinear-Setting_up_TOV_initial_data.ipynb

# Inputs:
# * Output data file name
# * rho_baryon_central, the central density of the TOV star.
# * n, the polytropic equation of state index. n=1 models cold, degenerate neutron star matter.
# * K_Polytrope, the polytropic constant.
# * Verbose output toggle (default = True)

# Output: An initial data file (default file name = "outputTOVpolytrope.txt") that well
#         samples the (spherically symmetric) solution both inside and outside the star.
#         It is up to the initial data module to perform the 1D interpolation to generate
#         the solution at arbitrary radius. The file has the following columns:
# Column 1: Schwarzschild radius
# Column 2: rho(r), *total* mass-energy density (as opposed to baryonic rest-mass density)
# Column 3: P(r), Pressure
# Column 4: m(r), mass enclosed
# Column 5: e^{nu(r)}, g_{tt}(r)
# Column 6: e^{4 phi(r)}, conformal factor g_{rr}(r)
# Column 7: rbar(r), Isotropic radius

# rbar refers to the isotropic radius, and
# R_Schw refers to the Schwarzschild radius

import numpy as np
import scipy.integrate as si
import math

n = 1.
rho_baryon_central = 0.129285

P0 = 1. # ZACH NOTES: CHANGED FROM 100.
gamma = 1. + 1./n
gam1 = gamma - 1.

def TOV_pressure( rho_baryon) :
    return P0*rho_baryon**gamma

def TOV_rhs(r_Schw, y) :
# In \tilde units
#
    P    = y[0]
    m    = y[1]
    nu   = y[2]
    rbar = y[3]

    dPdrSchw    = 0.
    drbardrSchw = 0.

    rho_baryon = (P/P0)**(1./gamma)
    rho = rho_baryon + P/gam1 # rho is the *total* mass-energy density!
    if( r_Schw < 1e-4 or m <= 0.):
        m = 4*math.pi/3. * rho*r_Schw**3
        dPdrSchw = -(rho + P)*(4.*math.pi/3.*r_Schw*rho + 4.*math.pi*r_Schw*P)/(1.-8.*math.pi*rho*r_Schw*r_Schw)
        drbardrSchw = 1./(1. - 8.*math.pi*rho*r_Schw*r_Schw)**0.5
    else:
        dPdrSchw = -(rho + P)*(m + 4.*math.pi*r_Schw**3*P)/(r_Schw*r_Schw*(1.-2.*m/r_Schw))
        drbardrSchw = 1./(1. - 2.*m/r_Schw)**0.5*rbar/r_Schw

    dmdrSchw  =  4.*math.pi*r_Schw*r_Schw*rho
    dnudrSchw = -2./(P + rho)*dPdrSchw
    return [dPdrSchw, dmdrSchw, dnudrSchw, drbardrSchw]

def integrateStar( P, dumpData = False):
    integrator = si.ode(TOV_rhs).set_integrator('dop853')
    y0 = [P, 0., 0., 0.]
    integrator.set_initial_value(y0,0.)
    dr_Schw = 1e-5
    P = y0[0]

    PArr      = []
    r_SchwArr = []
    mArr      = []
    nuArr     = []
    rbarArr   = []

    r_Schw = 0.

    while integrator.successful() and P > 1e-9*y0[0] :
        P, m, nu, rbar = integrator.integrate(r_Schw + dr_Schw)
        r_Schw = integrator.t

        dPdrSchw, dmdrSchw, dnudrSchw, drbardrSchw = TOV_rhs( r_Schw+dr_Schw, [P,m,nu,rbar])
        dr_Schw = 0.1*min(abs(P/dPdrSchw), abs(m/dmdrSchw))
        dr_Schw = min(dr_Schw, 1e-2)
        PArr.append(P)
        r_SchwArr.append(r_Schw)
        mArr.append(m)
        nuArr.append(nu)
        rbarArr.append(rbar)

    M = mArr[-1]
    R_Schw = r_SchwArr[-1]

    # Apply integration constant to ensure rbar is continuous across TOV surface
    for ii in range(len(rbarArr)):
        rbarArr[ii] *= 0.5*(np.sqrt(R_Schw*(R_Schw - 2.0*M)) + R_Schw - M) / rbarArr[-1]

    nuArr_np = np.array(nuArr)
    # Rescale solution to nu so that it satisfies BC: exp(nu(R))=exp(nutilde-nu(r=R)) * (1 - 2m(R)/R)
    #   Thus, nu(R) = (nutilde - nu(r=R)) + log(1 - 2*m(R)/R)
    nuArr_np = nuArr_np - nuArr_np[-1] + math.log(1.-2.*mArr[-1]/r_SchwArr[-1])

    r_SchwArrExtend_np = 10.**(np.arange(0.01,5.0,0.01))*r_SchwArr[-1]

    r_SchwArr.extend(r_SchwArrExtend_np)
    mArr.extend(r_SchwArrExtend_np*0. + M)
    PArr.extend(r_SchwArrExtend_np*0.)
    exp2phiArr_np = np.append( np.exp(nuArr_np), 1. - 2.*M/r_SchwArrExtend_np)
    nuArr.extend(np.log(1. - 2.*M/r_SchwArrExtend_np))
    rbarArr.extend( 0.5*(np.sqrt(r_SchwArrExtend_np**2 - 2.*M*r_SchwArrExtend_np) + r_SchwArrExtend_np - M) )

    # Appending to a Python array does what one would reasonably expect.
    #   Appending to a numpy array allocates space for a new array with size+1,
    #   then copies the data over... over and over... super inefficient.
    r_SchwArr_np          = np.array(r_SchwArr)
    PArr_np               = np.array(PArr)
    rho_baryonArr_np      = (PArr_np/P0)**(1./gamma)
    mArr_np               = np.array(mArr)
    rbarArr_np            = np.array(rbarArr)
    confFactor_exp4phi_np = (r_SchwArr_np/rbarArr_np)**2

    # Compute the *total* mass-energy density (as opposed to the *baryonic* mass density)
    rhoArr_np = []
    for i in range(len(rho_baryonArr_np)):
        rhoArr_np.append(rho_baryonArr_np[i] + PArr_np[i]/(gamma - 1.))

    print(len(r_SchwArr_np),len(rhoArr_np),len(PArr_np),len(mArr_np),len(exp2phiArr_np))
    np.savetxt("outputTOVpolytrope.txt", zip(r_SchwArr_np,rhoArr_np,PArr_np,mArr_np,exp2phiArr_np,confFactor_exp4phi_np,rbarArr_np),
               fmt="%.15e")

    return R_Schw, M

R_Schw_TOV,M_TOV = integrateStar(TOV_pressure(rho_baryon_central), True)
print("Just generated a TOV star with R_Schw = "+str(R_Schw_TOV)+" , M = "+str(M_TOV)+" , M/R_Schw = "+str(M_TOV/R_Schw_TOV)+" .")
```

    (1051, 1051, 1051, 1051, 1051)
    Just generated a TOV star with R_Schw = 0.956568142523 , M = 0.14050303285288188 , M/R_Schw = 0.1468824086931645 .

## Step 2.1: Validation of the TOV output file against a trusted version
$$\label{step2p1}$$
\[Back to [top](#top)\]

Here, as a code validation check, we verify agreement in the SymPy expressions for these TOV initial data between

1. this tutorial and
2. the NRPy+ [TOV.TOV_Solver](../edit/TOV/TOV_Solver.py) module.

```python
import filecmp

import TOV.TOV_Solver as TOV
TOV.TOV_Solver("outputTOVpolytrope-validation.txt",rho_baryon_central=0.129285, n_Polytrope=1.0, K_Polytrope=1.0,
               verbose = False)

if filecmp.cmp('outputTOVpolytrope.txt',
               'outputTOVpolytrope-validation.txt') == False:
    print("ERROR: TOV initial data test FAILED!")
    exit(1)
else:
    print("TOV initial data test PASSED.")
```

    TOV initial data test PASSED.

## Step 3: C code for reading TOV star initial data and interpolating data to numerical grids \[Back to [top](#top)\]

$$\label{step3}$$

### Step 3.1: C code for reading TOV star initial data file \[Back to [top](#top)\]
$$\label{step3p1}$$

```python
%%writefile TOV/read_datafile__set_arrays.h

int read_datafile__set_arrays(FILE *in1Dpolytrope, REAL *r_Schw_arr,REAL *rho_arr,REAL *P_arr,REAL *M_arr,REAL *expnu_arr,REAL *exp4phi_arr,REAL *rbar_arr) {
  char * line = NULL;

  size_t len = 0;
  ssize_t read;

  int which_line = 0;
  while ((read = getline(&line, &len, in1Dpolytrope)) != -1) {
    // Define the line delimiters (i.e., the stuff that goes between the data on a given
    //     line of data.  Here, we define both spaces " " and tabs "\t" as data delimiters.
    const char delimiters[] = " \t";

    //Now we define "token", a pointer to the first column of data
    char *token;

    //Each successive time we call strtok(NULL,blah), we read in a new column of data from
    //     the originally defined character array, as pointed to by token.

    token=strtok(line, delimiters); if(token==NULL) { printf("BADDDD\n"); return 1; }
    r_Schw_arr[which_line]  = strtod(token, NULL); token = strtok( NULL, delimiters );
    rho_arr[which_line]     = strtod(token, NULL); token = strtok( NULL, delimiters );
    P_arr[which_line]       = strtod(token, NULL); token = strtok( NULL, delimiters );
    M_arr[which_line]       = strtod(token, NULL); token = strtok( NULL, delimiters );
    expnu_arr[which_line]   = strtod(token, NULL); token = strtok( NULL, delimiters );
    exp4phi_arr[which_line] = strtod(token, NULL); token = strtok( NULL, delimiters );
    rbar_arr[which_line]    = strtod(token, NULL);

    which_line++;
  }
  free(line);
  return 0;
}
```

    Overwriting TOV/read_datafile__set_arrays.h

### Step 3.2: C code for interpolating TOV data to desired radial coordinate $\bar{r}$ \[Back to [top](#top)\]
$$\label{step3p2}$$

The TOV data file written in [Step 2](#step2) above stored $\left(r,\rho(r),P(r),M(r),e^{\nu(r),\bar{r}}\right)$, where $\rho(r)$ is the total mass-energy density (cf. $\rho_{\text{baryonic}}$) at a fixed number of radii. This C code implements [Lagrange polynomial interpolation](https://en.wikipedia.org/wiki/Polynomial_interpolation) to interpolate $\left(r,\rho(r),P(r),M(r),e^{\nu(r)}\right)$ to any given input radius $\bar{r}$.

```python
%%writefile TOV/TOV_interpolate_1D.h

void TOV_interpolate_1D(REAL rrbar,const REAL Rbar,const int Rbar_idx,const int interp_stencil_size,
                        const int numlines_in_file,const REAL *r_Schw_arr,const REAL *rho_arr,const REAL *P_arr,const REAL *M_arr,const REAL *expnu_arr,const REAL *exp4phi_arr,const REAL *rbar_arr,
                        REAL *rho,REAL *P,REAL *M,REAL *expnu,REAL *exp4phi) {
  // Find interpolation index using Bisection root-finding algorithm:
  int bisection_idx_finder(const REAL rr, const int numlines_in_file, const REAL *rbar_arr) {
    int x1 = 0;
    int x2 = numlines_in_file-1;
    REAL y1 = rrbar-rbar_arr[x1];
    REAL y2 = rrbar-rbar_arr[x2];
    if(y1*y2 >= 0) {
      fprintf(stderr,"INTERPOLATION BRACKETING ERROR %e | %e %e\n",rr,y1,y2);
      exit(1);
    }
    for(int i=0;i<numlines_in_file;i++) {
      int x_midpoint = (x1+x2)/2;
      REAL y_midpoint = rrbar-rbar_arr[x_midpoint];
      if(y_midpoint*y1 < 0) {
        x2 = x_midpoint;
        y2 = y_midpoint;
      } else {
        x1 = x_midpoint;
        y1 = y_midpoint;
      }
      if( abs(x2-x1) == 1 ) {
        // If rbar_arr[x1] is closer to rrbar than rbar_arr[x2] then return x1:
        if(fabs(rrbar-rbar_arr[x1]) < fabs(rrbar-rbar_arr[x2])) return x1;
        // Otherwiser return x2:
        return x2;
      }
    }
    fprintf(stderr,"INTERPOLATION BRACKETING ERROR: DID NOT CONVERGE.\n");
    exit(1);
  }

  // For this case, we know that for all functions, f(r) = f(-r)
  if(rrbar < 0) rrbar = -rrbar;

  // First find the central interpolation stencil index:
  int idx = bisection_idx_finder(rrbar,numlines_in_file,rbar_arr);

#ifdef MAX
#undef MAX
#endif
#define MAX(A, B) ( ((A) > (B)) ? (A) : (B) )

  int idxmin = MAX(0,idx-interp_stencil_size/2-1);

#ifdef MIN
#undef MIN
#endif
#define MIN(A, B) ( ((A) < (B)) ? (A) : (B) )

  // -= Do not allow the interpolation stencil to cross the star's surface =-
  // max index is when idxmin + (interp_stencil_size-1) = Rbar_idx
  //  -> idxmin at most can be Rbar_idx - interp_stencil_size + 1
  if(rrbar < Rbar) {
    idxmin = MIN(idxmin,Rbar_idx - interp_stencil_size + 1);
  } else {
    idxmin = MAX(idxmin,Rbar_idx+1);
    idxmin = MIN(idxmin,numlines_in_file - interp_stencil_size + 1);
  }

  // Now perform Lagrange polynomial interpolation:

  // First set the interpolation coefficients:
  REAL rbar_sample[interp_stencil_size];
  for(int i=idxmin;i<idxmin+interp_stencil_size;i++) {
    rbar_sample[i-idxmin] = rbar_arr[i];
  }
  REAL l_i_of_r[interp_stencil_size];
  for(int i=0;i<interp_stencil_size;i++) {
    REAL numer = 1.0;
    REAL denom = 1.0;
    for(int j=0;j<i;j++) {
      numer *= rrbar - rbar_sample[j];
      denom *= rbar_sample[i] - rbar_sample[j];
    }
    for(int j=i+1;j<interp_stencil_size;j++) {
      numer *= rrbar - rbar_sample[j];
      denom *= rbar_sample[i] - rbar_sample[j];
    }
    l_i_of_r[i] = numer/denom;
  }

  // Then perform the interpolation:
  *rho = 0.0;
  *P = 0.0;
  *M = 0.0;
  *expnu = 0.0;
  *exp4phi = 0.0;

  REAL r_Schw = 0.0;
  for(int i=idxmin;i<idxmin+interp_stencil_size;i++) {
    r_Schw   += l_i_of_r[i-idxmin] * r_Schw_arr[i];
    *rho     += l_i_of_r[i-idxmin] * rho_arr[i];
    *P       += l_i_of_r[i-idxmin] * P_arr[i];
    *M       += l_i_of_r[i-idxmin] * M_arr[i];
    *expnu   += l_i_of_r[i-idxmin] * expnu_arr[i];
    *exp4phi += l_i_of_r[i-idxmin] * exp4phi_arr[i];
  }

  if(rrbar > Rbar) {
    *rho   = 0;
    *P     = 0;
    *M     = M_arr[Rbar_idx+1];
    *expnu = 1. - 2.*(*M) / r_Schw;
    *exp4phi = pow(r_Schw / rrbar,2.0);
  }
}
```

    Overwriting TOV/TOV_interpolate_1D.h

## Step 4: Construct ADM spacetime quantities and $T^{\mu\nu}$  from TOV solution \[Back to [top](#top)\]
$$\label{step4}$$

## Step 4.1: Load NRPy+ modules and set basic parameters \[Back to [top](#top)\]
$$\label{step4p1}$$

```python
# Step 4.1: Load NRPy+ modules and set basic parameters

import sympy as sp
import NRPy_param_funcs as par
import indexedexp as ixp
import grid as gri
from outputC import *
import reference_metric as rfm

# All gridfunctions will be written in terms of spherical coordinates (r, th, ph):
r,th,ph = sp.symbols('r th ph', real=True)

thismodule = "TOV"

DIM = 3
par.set_parval_from_str("grid::DIM",DIM)

# Input parameters read in from the TOV data file:
rbar,expnu,exp4phi,P,rho = par.Cparameters("REAL", thismodule,
                                           ["rbar","expnu","exp4phi","P", "rho"],
                                           [ 1e300,  1e300,    1e300,1e300,1e300]) # Must be read from TOV data
                                                                      # file; set to crazy values to ensure this
```

## Step 4.2: Constructing ADM spacetime quantities from the TOV data \[Back to [top](#top)\]

$$\label{step4p2}$$

The [TOV line element](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation) in *Schwarzschild coordinates* is written (in the $-+++$ form):
$$
ds^2 = - c^2 e^\nu dt^2 + \left(1 - \frac{2GM}{rc^2}\right)^{-1} dr^2 + r^2 d\Omega^2.
$$

In *isotropic coordinates* with $G=c=1$ (i.e., the coordinate system and units we'd prefer to use), the ($-+++$ form) line element is written:
$$
ds^2 = - e^{\nu} dt^2 + e^{4\phi} \left(d\bar{r}^2 + \bar{r}^2 d\Omega^2\right),
$$
where $\phi$ here is the *conformal factor*.

The ADM 3+1 line element for this diagonal metric in isotropic spherical coordinates is given by:
$$
ds^2 = (-\alpha^2 + \beta_k \beta^k) dt^2 + \gamma_{\bar{r}\bar{r}} d\bar{r}^2 + \gamma_{\theta\theta} d\theta^2+ \gamma_{\phi\phi} d\phi^2,
$$

from which we can immediately read off the ADM quantities:
\begin{align}
\gamma_{\bar{r}\bar{r}} &= e^{4\phi}\\
\gamma_{\theta\theta} &= e^{4\phi} \bar{r}^2 \\
\gamma_{\phi\phi} &= e^{4\phi} \bar{r}^2 \sin^2 \theta \\
K_{ij} &= 0 \\
\alpha &= e^{\nu(\bar{r})/2} \\
\beta^i = B^i &= 0 \\
\end{align}

```python
# Step 4.2: Construct ADM quantities:

# *** The physical spatial metric in spherical basis ***
# In isotropic coordinates,
#  gamma_{ij} = e^{4 phi} eta_{ij},
# where eta is the flat-space 3-metric in spherical coordinates
gammaSphDD = ixp.zerorank2()
gammaSphDD[0][0] = exp4phi
gammaSphDD[1][1] = exp4phi * rbar**2
gammaSphDD[2][2] = exp4phi * rbar**2*sp.sin(th)**2

# *** The extrinsic curvature in spherical basis ***
# K_{ij} = 0 for the TOV solution
KSphDD = ixp.zerorank2()

# *** The lapse and shift in spherical basis ***
# alpha = exp^{nu/2} for the TOV solution
# \beta^i = 0 for the TOV solution
alphaSph = sp.sqrt(expnu)
betaSphU = ixp.zerorank1()
BSphSphU = ixp.zerorank1()
```

## Step 4.3: Constructing $T^{\mu\nu}$ from the TOV data \[Back to [top](#top)\]
$$\label{step4p3}$$

We will also need the stress-energy tensor $T^{\mu\nu}$. [As discussed here](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation), the stress-energy tensor is diagonal:

\begin{align}
T^t_t &= -\rho \\
T^i_j &= P \delta^i_j \\
\text{All other components of }T^\mu_\nu &= 0.
\end{align}

Since $\beta^i=0$ the inverse metric expression simplifies to (Eq. 4.49 in [Gourgoulhon](https://arxiv.org/pdf/gr-qc/0703035.pdf)):
$$
g^{\mu\nu} = \begin{pmatrix}
-\frac{1}{\alpha^2} & \frac{\beta^i}{\alpha^2} \\
\frac{\beta^i}{\alpha^2} & \gamma^{ij} - \frac{\beta^i\beta^j}{\alpha^2}
\end{pmatrix} =
\begin{pmatrix}
-\frac{1}{\alpha^2} & 0 \\
0 & \gamma^{ij}
\end{pmatrix},
$$

and since the 4-metric is diagonal we get

\begin{align}
\gamma^{\bar{r}\bar{r}} &= e^{-4\phi}\\
\gamma^{\theta\theta} &= e^{-4\phi}\frac{1}{\bar{r}^2} \\
\gamma^{\phi\phi} &= e^{-4\phi}\frac{1}{\bar{r}^2 \sin^2 \theta}.
\end{align}

Thus raising $T^\mu_\nu$ yields a diagonal $T^{\mu\nu}$

\begin{align}
T^{tt} &= -g^{tt} \rho = \frac{1}{\alpha^2} \rho = e^{-\nu(\bar{r})} \rho \\
T^{\bar{r}\bar{r}} &= g^{\bar{r}\bar{r}} P = \frac{1}{e^{4 \phi}} P \\
T^{\theta\theta} &= g^{\theta\theta} P = \frac{1}{e^{4 \phi}\bar{r}^2} P\\
T^{\phi\phi} &= g^{\phi\phi} P = \frac{1}{e^{4\phi}\bar{r}^2 \sin^2 \theta} P
\end{align}

```python
# Step 4.3: Construct T^{mu nu}:
T4UU = ixp.zerorank2(DIM=4)

# T^tt = e^(-nu) * rho
T4UU[0][0] = rho / expnu
# T^{ii} = P / gamma_{ii}
for i in range(3):
    T4UU[i+1][i+1] = P / gammaSphDD[i][i]
```

## Step 5: NRPy+ Module Code Validation \[Back to [top](#top)\]
$$\label{step5}$$

Here, as a code validation check, we verify agreement in the SymPy expressions for TOV initial data between

1. this tutorial and
2. the NRPy+ [TOV.TOV_ADM_T4UUmunu](../edit/TOV/TOV_ADM_T4UUmunu.py) module.

As all input quantities are functions of $r$, we will simply read the solution from file and interpolate it to the values of $r$ needed by the initial data.

1. First we define functions ID_TOV_ADM_quantities() and ID_TOV_TUPMUNU() that call the [1D TOV interpolator function](../edit/TOV/tov_interp.h) to evaluate the ADM spacetime quantities and $T^{\mu\nu}$, respectively, at any given point $(r,\theta,\phi)$ in the Spherical basis. All quantities are defined as above.
1. Next we will construct the BSSN/ADM source terms $\{S_{ij},S_{i},S,\rho\}$ in the Spherical basis
1. Then we will perform the Jacobian transformation on $\{S_{ij},S_{i},S,\rho\}$ to the desired (xx0,xx1,xx2) basis
1. Next we call the *Numerical* Spherical ADM$\to$Curvilinear BSSN converter function to convert the above ADM quantities to the rescaled BSSN quantities in the desired curvilinear coordinate system: [BSSN/ADM_Numerical_Spherical_or_Cartesian_to_BSSNCurvilinear.py](../edit/BSSN/ADM_Numerical_Spherical_or_Cartesian_to_BSSNCurvilinear.py); [\[**tutorial**\]](Tutorial-ADM_Initial_Data-Converting_Numerical_ADM_Spherical_or_Cartesian_to_BSSNCurvilinear.ipynb).

### Output to $\LaTeX$, then PDF

```python
!jupyter nbconvert --to latex --template latex_nrpy_style.tplx Tutorial-ADM_and_T4UUmunu_Initial_Data-TOV.ipynb
!pdflatex -interaction=batchmode Tutorial-ADM_and_T4UUmunu_Initial_Data-TOV.tex
!pdflatex -interaction=batchmode Tutorial-ADM_and_T4UUmunu_Initial_Data-TOV.tex
!pdflatex -interaction=batchmode Tutorial-ADM_and_T4UUmunu_Initial_Data-TOV.tex
!rm -f Tut*.out Tut*.aux Tut*.log
```

    [NbConvertApp] Converting notebook Tutorial-ADM_and_T4UUmunu_Initial_Data-TOV.ipynb to latex
    [NbConvertApp] Writing 67639 bytes to Tutorial-ADM_and_T4UUmunu_Initial_Data-TOV.tex
    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode
    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode
    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode


```python
%matplotlib notebook
%matplotlib inline

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

# Amortization

A quick review of what I said in lecture, what the book says, and what you might find on the internet

## Learning Objectives

At the end of this lesson, you will be equipped to:

- Recognize different forms of the amortization equation
- Choose the correct form
- Apply the equation with the appropriate variables

## Annuities

Annuities: a series of equal dollar payments coming at the end of a certain time period for a specified number of time periods (n). Examples include mortgages, life insurance benefits, lottery payments, retirement payments.

### Compound annuities

Compound Annuities: depositing an equal sum of money at the end of each time period for a certain number of periods and allowing the money to grow. An example might be having $\$50$ taken out of each paycheck and put in a Christmas account earning a $9\%$ Annual Percentage Rate.

### Future value of an annuity (FVA) equation

This equation is used to determine the future value of a stream of deposits/payments (PMT) invested at a specific interest rate (i), for a specific number of periods (n). For example, the future value of your 401k contributions can be calculated using this equation:

$$
\begin{align}
FVA &= PMT\left(\frac{(1 + i)^n – 1}{i}\right)\\
FVA &= \mbox{the future value, in today’s dollars, of a sum of
money}\\
PMT &= \mbox{the payment made at the end of each time period}\\
i &= \mbox{interest rate}\\
n &= \mbox{number of periods}
\end{align}
$$

```python
def fva(pmt, i, n):
    """the future value, in today’s dollars, of a sum of money
    :param pmt: the payment made at the end of each time period
    :param i: interest rate
    :param n: number of periods
    """
    num = pmt*(pow((1+i),n) - 1)
    return num/i
```

### Present value of an annuity (PVA) equation

This equation is used to determine the present value of a future stream of payments, such as your pension fund or insurance benefits.

$$
\begin{align}
PVA &=PMT\frac{1–(1+i)^{-n}}{i}\\
PVA &= \mbox{the present value, in today’s dollars, of a sum of money}\\
PMT &= \mbox{the payment to be made at the end of each time period}\\
i &= \mbox{interest rate}\\
n &= \mbox{number of periods}
\end{align}
$$

### Solve for PMT

$$
\begin{align}
PMT &= PVA\frac{i}{1–(1+i)^{-n}}
\end{align}
$$

```python
def pva(pmt, i, n):
    """the present value, in today’s dollars, of a sum of money
    :param PMT: the payment to be made at the end of each time period
    :param i: interest rate
    :param n: number of periods
    """
    num = pmt*(1 - pow((1+i),-n))
    return num/i
```

```python
def pmt(pva, i, n):
    """the payment to be made at the end of each time period
    :param pva: the present value in todays dollars, of the annuity
    :param i: interest rate
    :param n: number of periods
    """
    num = pva*i
    denom =1 - pow((1+i),-n)
    return num/denom
```

```python
# To pay off 3 million dollars, over the course of 30 years, at an interest rate of 10%
i = 0.1
n = 30
pva = 3.0e6

print("Each payment: " + '%e' % (pmt(pva, i, n)) + "\n"
+ "Sum of Payments: " + '%e' % (pmt(pva, i, n)*n))

```

    Each payment: 3.182377e+05
    Sum of Payments: 9.547132e+06

## What about the book?

So, in the book, the equation capturing this situation is:

\begin{align}
F &= \left[ \frac{i(1+i)^n}{(1+i)^n - 1} P\right]N
\end{align}

One wants to rearrange this in many ways. First, I really don't like P being inside those brackets. Yikes.

\begin{align}
F &= \left[ \frac{i(1+i)^n}{(1+i)^n - 1} \right]PN
\end{align}

Much better. We also want to solve for the payment. What term or combination of terms would represent the amount of each equal payment?

At first glance, the payment might seem to be P.

**P is not the payment. The payment is F/N.** P is the present value of the annuity (we called it PVA in the last example).

\begin{align}
\frac{F}{N} &= \left[ \frac{i(1+i)^n}{(1+i)^n - 1} \right]P
\end{align}

This doesn't look like the same equation as before, but it is. Let's show that by multiplying the top and bottom of the right hand side by $(1+i)^{-n}$.

\begin{align}
\frac{F}{N} &= \left[ \frac{i(1+i)^{n-n}}{(1+i)^{n-n} - 1(1+i)^{-n}} \right]P\\
\frac{F}{N} &= \left[ \frac{i}{1 - (1+i)^{-n}} \right]P
\end{align}

Now, it's clearly the same equation. But be careful about the variables. If you do the wrong thing, you'll get a very large number for the payment (because you'll be solving for the present value of the annuity whose future value is the amount of the loan. **Very easy mistake to make!**

```python

```