# 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.ortho_forest 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:    8.6s
    [Parallel(n_jobs=-1)]: Done 248 tasks      | elapsed:   10.1s
    [Parallel(n_jobs=-1)]: Done 888 tasks      | elapsed:   13.9s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:   14.6s 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:    2.4s
    [Parallel(n_jobs=-1)]: Done 981 tasks      | elapsed:    5.4s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:    5.4s finished

    <econml.ortho_forest.DMLOrthoForest at 0x299345fc630>

```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:    8.9s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:   11.0s 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   9 tasks      | elapsed:    2.0s
    [Parallel(n_jobs=-1)]: Done  85 out of 100 | elapsed:    4.5s remaining:    0.7s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:    4.5s 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.4s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:    4.8s finished

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

      point_estimate
      stderr
      zstat
      pvalue
      ci_lower
      ci_upper

      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, max_samples=subsample_ratio/2,
                       verbose=0, 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, 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:    7.1s
    [Parallel(n_jobs=-1)]: Done 176 tasks      | elapsed:    8.0s
    [Parallel(n_jobs=-1)]: Done 200 out of 200 | elapsed:    8.1s 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 200 out of 200 | elapsed:    0.7s finished

    <econml.ortho_forest.DROrthoForest at 0x29939f5aac8>

```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:    8.0s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:    8.9s 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.4s 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)
treatment_effects2 = est2.effect(X_test)
te_lower2, te_upper2 = est2.effect_interval(X_test, alpha=0.01)
```

## 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:    8.4s
    [Parallel(n_jobs=-1)]: Done 112 tasks      | elapsed:   10.2s
    [Parallel(n_jobs=-1)]: Done 272 tasks      | elapsed:   13.0s
    [Parallel(n_jobs=-1)]: Done 500 out of 500 | elapsed:   16.9s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.4s
    [Parallel(n_jobs=-1)]: Done 208 tasks      | elapsed:    4.0s
    [Parallel(n_jobs=-1)]: Done 500 out of 500 | elapsed:    8.5s finished

    <econml.ortho_forest.DROrthoForest at 0x299366477f0>

```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:   10.0s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:   11.9s 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  18 tasks      | elapsed:    1.8s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:    4.0s 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:    1.4s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:    3.0s finished

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

      point_estimate
      stderr
      zstat
      pvalue
      ci_lower
      ci_upper

      T0
      T1
      T2
      T0
      T1
      T2
      T0
      T1
      T2
      T0
      T1
      T2
      T0
      T1
      T2
      T0
      T1
      T2

      0
      1.010
      -0.986
      0.036
      0.162
      0.149
      0.225
      6.246
      -6.614
      0.161
      0.0
      0.0
      0.872
      0.744
      -1.231
      -0.335
      1.276
      -0.741
      0.407

      1
      1.012
      -0.991
      0.035
      0.160
      0.150
      0.225
      6.342
      -6.597
      0.157
      0.0
      0.0
      0.876
      0.750
      -1.238
      -0.335
      1.275
      -0.744
      0.406

      2
      1.018
      -0.988
      0.035
      0.161
      0.147
      0.226
      6.323
      -6.707
      0.154
      0.0
      0.0
      0.877
      0.753
      -1.230
      -0.337
      1.283
      -0.746
      0.407

      3
      1.021
      -0.993
      0.035
      0.160
      0.149
      0.225
      6.387
      -6.668
      0.155
      0.0
      0.0
      0.877
      0.758
      -1.238
      -0.335
      1.283
      -0.748
      0.405

      4
      1.025
      -0.994
      0.037
      0.160
      0.148
      0.227
      6.388
      -6.733
      0.161
      0.0
      0.0
      0.872
      0.761
      -1.237
      -0.337
      1.288
      -0.751
      0.410

```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:    7.4s
    [Parallel(n_jobs=-1)]: Done 232 tasks      | elapsed:    8.0s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:    9.3s 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 888 tasks      | elapsed:    1.4s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:    1.7s finished

    <econml.ortho_forest.DMLOrthoForest at 0x29938f1cc88>

```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:   12.1s
    [Parallel(n_jobs=-1)]: Done 101 out of 101 | elapsed:   29.3s 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  16 tasks      | elapsed:    6.4s
    [Parallel(n_jobs=-1)]: Done 101 out of 101 | elapsed:   27.1s 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


```python
%matplotlib widget
```

```python
from bmcs_cross_section.mkappa.mkappa import MKappa
from bmcs_cross_section.api import TShape, ReinfLayer, EC2
from bmcs_cross_section.matmod import PWLConcreteMatMod, EC2ConcreteMatMod

from bmcs_beam.beam_config.boundary_conditions import BoundaryConfig
from bmcs_beam.beam_config.boundary_conditions import BoundaryConditions
from bmcs_beam.bending.deflection_profile import DeflectionProfile
from bmcs_beam.beam_config.beam_design import BeamDesign
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
import bmcs_utils.api as bu
```

```python
def get_Fu_and_Fs(dp, mc, bl1, n_i, rho_range=(0.0025, 0.025), slenderness_range=(3, 35)):
    rho_range = np.linspace(*rho_range, n_i)
    slenderness_range = np.linspace(*slenderness_range, n_i)

    d = mc.cross_section_shape_.B -  bl1.z
    b = mc.cross_section_shape_.B
    area_g = b * d

    rho_grid, sl_grid = np.meshgrid(rho_range, slenderness_range)
    F_u_grid = np.zeros_like(rho_grid)
    F_s_grid = np.zeros_like(rho_grid)

    _ , ax = plt.subplots()
    ax.set_xlabel(r'$w$ [mm]')
    ax.set_ylabel(r'$F$ [KN]')

    for sl_idx in range(0,len(slenderness_range)):
        for rho_idx in range(0,len(rho_range)):
            rho = rho_grid[rho_idx, sl_idx]
            sl = sl_grid[rho_idx, sl_idx]

            print('parameter combination', rho, sl)

            # assigning the grid area (area_g) to the reinforcement area variable
            A_j_g = rho * area_g
            bl1.A = A_j_g

            # assigning the grid length (L_g) to the beam length variable
            L_g = sl * d
            dp.beam_design.L = L_g

            # running the deflection analysis
            F_data, w_data = dp.get_Fw()

            # plotting, post-processing & saving the data
            ax.plot(w_data, F_data/1000, label = "rho={}%-sl={} ".format(rho * 100, sl))

            w_s = dp.beam_design.L / 250
            F_u = max(F_data)
            F_s = np.interp(w_s, w_data, F_data, right=F_u*2)

            F_u_grid[rho_idx,sl_idx] = F_u
            F_s_grid[rho_idx,sl_idx] = F_s

    return F_u_grid, F_s_grid, rho_grid, sl_grid
```

```python
f_ck = 30
mc = MKappa(name='moment-curvature',n_m=100, n_kappa = 100, low_kappa=0.0, high_kappa=0.00009)
mc.matrix = 'piecewise linear'
mc.matrix_.trait_set(
    E_cc = EC2.get_E_cm(f_ck),
    E_ct = EC2.get_E_cm(f_ck),
    eps_cr = 0.000125 * ((f_ck**0.5)/(f_ck/10)**0.3)/((30**0.5)/(30/10)**0.3),
    eps_cy = EC2.get_eps_c2(f_ck),
    eps_cu = EC2.get_eps_cu2(f_ck),
    mu = 0.00
)
mc.cross_section_shape = 'rectangle'
mc.cross_section_shape_.trait_set(B = 1000, H = 300)
csl = mc.cross_section_layout
bl1 = ReinfLayer(name='steel 1', z=100, A=231.04, count=1 ,E=203000, f_sy=560)
bl1.matmod = 'carbon'
csl.add_layer(bl1)
```

    C:\Users\srate\miniconda3\envs\bmcs_env\lib\site-packages\traits\trait_types.py:180: ComplexWarning: Casting complex values to real discards the imaginary part
      return nb_float(value)

```python
dp = DeflectionProfile(mc=mc)
dp.beam_design.beam_conf_name = BoundaryConfig.FOUR_PB
dp.beam_design.beam_conf_name.first_load_distance = 0
```

```python
dp.app()
```

    M - k recalculated
    No solution The iteration is not making good progress, as measured by the
      improvement from the last ten iterations.

    VBox(children=(HBox(children=(VBox(children=(Tree(layout=Layout(align_items='stretch', border='solid 1px black…

```python
F_u_grid, F_s_grid, rho_grid, sl_grid = get_Fu_and_Fs(dp, mc, bl1, 4, rho_range=(0.01, 0.025), slenderness_range=(3, 35))
```

```python
_, ax = plt.subplots(1,1)

z = F_u_grid / F_s_grid - 1./0.59 # - 1./0.51

CS = ax.contour(rho_grid, sl_grid, z)#, levels=[0])
ax.clabel(CS, inline=1, fontsize=10)
ax.set_ylabel('L/d')
ax.set_xlabel(r'$\rho$ %')
ax.set_ylim(0, 35)
ax.set_xlim(0.0, 0.025)

# Draw C30 curve
# rho_range = np.arange(0.0025, 0.025, 0.0001)
# f_ck = 30
# slenderness = get_slenderness(rho_range, 0,  1, f_ck)
# ax.plot(rho_range, slenderness, c=np.random.rand(3,))
```

    Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …

    (0.0, 0.025)

```python
np.save('F_u_grid_carbon_study' + str(f_ck) + '.npy', F_u_grid)
np.save('F_s_grid_carbon_study' + str(f_ck) + '.npy', F_s_grid)
np.save('rho_grid_carbon_study' + str(f_ck) + '.npy', rho_grid)
np.save('sl_grid_carbon_study' + str(f_ck) + '.npy', sl_grid)
```

    array([[207086.67179165, 207086.67179165],
           [ 17750.28615357,  17750.28615357]])

```python
f_cks = [20, 30, 35, 40, 50, 70]
F_u_grids = []
F_s_grids = []
rho_grids = []
sl_grids = []
reinforcement = 'carbon'
for f_ck in f_cks:
    f_ck = str(f_ck)
    F_u_grids.append(np.load('F_u_grid_' + reinforcement + '_EC2_eq2_tension_c_Homam' + f_ck + '.npy'))
    F_s_grids.append(np.load('F_s_grid_' + reinforcement + '_EC2_eq2_tension_c_Homam' + f_ck + '.npy'))
    rho_grids.append(np.load('rho_grid_' + reinforcement + '_EC2_eq2_tension_c_Homam' + f_ck + '.npy'))
    sl_grids.append(np.load('sl_grid_' + reinforcement + '_EC2_eq2_tension_c_Homam' + f_ck + '.npy'))
```

```python
_, ax = plt.subplots(1,1)

ax.set_ylabel('L/d')
ax.set_xlabel(r'$\rho$ %')
ax.set_ylim(0, 35)
ax.set_xlim(0.0, 0.025)

for f_ck, F_u_grid, F_s_grid, rho_grid, sl_grid in zip(f_cks, F_u_grids, F_s_grids, rho_grids, sl_grids):
    z = F_u_grid / F_s_grid - 1./0.59
    CS = ax.contour(rho_grid, sl_grid, z, levels=[0], colors=[np.random.rand(3,)])
    CS.collections[0].set_label('C' + str(f_ck))
#     ax.clabel(CS, inline=1, fontsize=10)
rho_range = np.arange(0.0025, 0.025, 0.0001)
# f_ck = 30
slenderness = get_slenderness(rho_range, 0,  1, f_ck)
ax.plot(rho_range, slenderness, c=np.random.rand(3,), label = 'EC2 - Steel')

ax.legend()

```

```python

```

```python

```

```python

```

```python

```

```python
dp.beam_design.beam.applied_loads
```

    [(R1, 0, -1, None), (R2, 5000, -1, None), (1000, 2500.00000000000, -1, None)]

```python
bl1.matmod_.trait_set(E=203000, f_sy=560)
```

    <bmcs_cross_section.matmod.reinforcement.SteelReinfMatMod at 0x2bacb161360>

```python
# mc.app()
```

```python
dp.beam_design.get_Q_x()
```

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

```python
dp.beam_design.beam_conf_name = BoundaryConfig.FOUR_PB
dp.beam_design.beam_conf_name.first_load_distance = 0

```

```python
dp.beam_design.beam_conf_name
```

    <BoundaryConfig.FOUR_PB: 1>

```python
M_I, kappa_I = dp.mc.inv_M_kappa
4 * M_I[-1] / dp.beam_design.L
```

    M - k recalculated

    48173.499398734995

```python
dp.F_max
```

    36130.124549051245

```python
mc_params = dict(
            reinforcement_type= ReinforcementType.STEEL,
            n_m=100,
            n_kappa = 100,
            kappa_slider = 0.000001,
            low_kappa = 0,
            high_kappa =  0.00002,
            eps_tu = 0.002508,
            E_cc = 34000,
            E_ct = 34000,
            eps_cr = 0.000132,
            eps_cy = 0.0015312,
            mu = 0.63,
            eps_cu = 0.003498)
mc = MKappa(**mc_params)
b = 285
h = 29
mc.cross_section_shape.B = b
mc.cross_section_shape.H = h

dp = DeflectionProfile(mc=mc)
dp.beam_design.beam_conf_name = BoundaryConfig.FOUR_PB
# beam.apply_load(F, 0, 0)
# bc = BoundaryConditions()
# bc.get_configured_beam()
dp.beam_design.beam_conf_name.first_load_distance = 0
dp.beam_design.L = 3000
# dp.beam_design.beam.apply_load(879898787, 0, 0)
%matplotlib widget
bu.InteractiveWindow([dp,
                      dp.beam_design,
                      dp.mc,
                      dp.mc.cs_design.cross_section_layout,
                      dp.mc.cs_design.cross_section_shape,
                      dp.mc.cross_section_layout.matrix,
                      dp.mc.cross_section_layout],
                      tight_layout=True).interact()

# M = dp.beam_design.get_M_x()
# plt.plot(mc.get_kappa_M(M+1))
# plt.show()
# mc.get_kappa_M(M) + 1

# dp.beam_design.beam.solve_for_reaction_loads(R1, R2)
# print(dp.beam_design.beam.reaction_loads)
# # print(dp.beam_design.beam.plot_bending_moment())
# dp.beam_design.beam.load
# # dp.beam_design.beam.apply_load(879898787, 0, 0)
# dp.beam_design.beam

```

```python
mc.matrix = 'piecewise linear'
css = mc.cross_section_shape = 'T-shape'
csl = mc.cross_section_layout
bl1 = BarLayer(name='steel 1', z=10, ds=10, count=2 )
bl2 = BarLayer(name='steel 2', z=30, ds=10, count=2 )
csl.add_layer(bl1)
csl.add_layer(bl2)

```

```python
mc.cross_section_layout.reinforcement = [Reinforcement(eps_sy=560/203000,\
E=203000, z=26.82, A=231.04)]
```


# Tutorial 2: Principal Component Analysis
**Week 1, Day 5: Dimensionality Reduction**

**By Neuromatch Academy**

__Content creators:__ Alex Cayco Gajic, John Murray

__Content reviewers:__ Roozbeh Farhoudi, Matt Krause, Spiros Chavlis, Richard Gao, Michael Waskom, Siddharth Suresh, Natalie Schaworonkow, Ella Batty

**Our 2021 Sponsors, including Presenting Sponsor Facebook Reality Labs**

---
# Tutorial Objectives

*Estimated timing of tutorial: 45 minutes*

In this notebook we'll learn how to perform PCA by projecting the data onto the eigenvectors of its covariance matrix.

Overview:
- Calculate the eigenvectors of the sample covariance matrix.
- Perform PCA by projecting data onto the eigenvectors of the covariance matrix.
- Plot and explore the eigenvalues.

To quickly refresh your knowledge of eigenvalues and eigenvectors, you can watch this [short video](https://www.youtube.com/watch?v=kwA3qM0rm7c) (4 minutes) for a geometrical explanation. For a deeper understanding, this [in-depth video](https://www.youtube.com/watch?v=PFDu9oVAE-g&list=PLZHQObOWTQDPD3MizzM2xVFitgF8hE_ab&index=14) (17 minutes) provides an excellent basis and is beautifully illustrated.

```python
# @title Tutorial slides

# @markdown These are the slides for the videos in all tutorials today
from IPython.display import IFrame
IFrame(src=f"https://mfr.ca-1.osf.io/render?url=https://osf.io/kaq2x/?direct%26mode=render%26action=download%26mode=render", width=854, height=480)
```

```python
# @title Video 1: PCA
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  from IPython.display import IFrame
  class BiliVideo(IFrame):
      def __init__(self, id, page=1, width=400, height=300, **kwargs):
          self.id=id
          src = 'https://player.bilibili.com/player.html?bvid={0}&page={1}'.format(id, page)
          super(BiliVideo, self).__init__(src, width, height, **kwargs)

  video = BiliVideo(id="BV1hK411H7Zi", width=854, height=480, fs=1)
  print('Video available at https://www.bilibili.com/video/{0}'.format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id="-f6T9--oM0E", width=854, height=480, fs=1, rel=0)
  print('Video available at https://youtube.com/watch?v=' + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

display(out)
```

    Tab(children=(Output(), Output()), _titles={'0': 'Youtube', '1': 'Bilibili'})

---
# Setup

```python
# Imports

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

```python
# @title Figure Settings
import ipywidgets as widgets  # interactive display
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
```

```python
# @title Plotting Functions

def plot_eigenvalues(evals):
  """
  Plots eigenvalues.

  Args:
      (numpy array of floats) : Vector of eigenvalues

  Returns:
    Nothing.

  """
  plt.figure(figsize=(4, 4))
  plt.plot(np.arange(1, len(evals) + 1), evals, 'o-k')
  plt.xlabel('Component')
  plt.ylabel('Eigenvalue')
  plt.title('Scree plot')
  plt.xticks(np.arange(1, len(evals) + 1))
  plt.ylim([0, 2.5])

def plot_data(X):
  """
  Plots bivariate data. Includes a plot of each random variable, and a scatter
  scatter plot of their joint activity. The title indicates the sample
  correlation calculated from the data.

  Args:
    X (numpy array of floats) : Data matrix each column corresponds to a
                                different random variable

  Returns:
    Nothing.
  """

  fig = plt.figure(figsize=[8, 4])
  gs = fig.add_gridspec(2, 2)
  ax1 = fig.add_subplot(gs[0, 0])
  ax1.plot(X[:, 0], color='k')
  plt.ylabel('Neuron 1')
  ax2 = fig.add_subplot(gs[1, 0])
  ax2.plot(X[:, 1], color='k')
  plt.xlabel('Sample Number (sorted)')
  plt.ylabel('Neuron 2')
  ax3 = fig.add_subplot(gs[:, 1])
  ax3.plot(X[:, 0], X[:, 1], '.', markerfacecolor=[.5, .5, .5],
           markeredgewidth=0)
  ax3.axis('equal')
  plt.xlabel('Neuron 1 activity')
  plt.ylabel('Neuron 2 activity')
  plt.title('Sample corr: {:.1f}'.format(np.corrcoef(X[:, 0], X[:, 1])[0, 1]))
  plt.show()

def plot_data_new_basis(Y):
  """
  Plots bivariate data after transformation to new bases. Similar to plot_data
  but with colors corresponding to projections onto basis 1 (red) and
  basis 2 (blue).
  The title indicates the sample correlation calculated from the data.

  Note that samples are re-sorted in ascending order for the first random
  variable.

  Args:
    Y (numpy array of floats) : Data matrix in new basis each column
                                corresponds to a different random variable

  Returns:
    Nothing.
  """

  fig = plt.figure(figsize=[8, 4])
  gs = fig.add_gridspec(2, 2)
  ax1 = fig.add_subplot(gs[0, 0])
  ax1.plot(Y[:, 0], 'r')
  plt.ylabel('Projection \n basis vector 1')
  ax2 = fig.add_subplot(gs[1, 0])
  ax2.plot(Y[:, 1], 'b')
  plt.xlabel('Sample number')
  plt.ylabel('Projection \n basis vector 2')
  ax3 = fig.add_subplot(gs[:, 1])
  ax3.plot(Y[:, 0], Y[:, 1], '.', color=[.5, .5, .5])
  ax3.axis('equal')
  plt.xlabel('Projection basis vector 1')
  plt.ylabel('Projection basis vector 2')
  plt.title('Sample corr: {:.1f}'.format(np.corrcoef(Y[:, 0], Y[:, 1])[0, 1]))
  plt.show()

def plot_basis_vectors(X, W):
  """
  Plots bivariate data as well as new basis vectors.

  Args:
    X (numpy array of floats) : Data matrix each column corresponds to a
                                different random variable
    W (numpy array of floats) : Square matrix representing new orthonormal
                                basis each column represents a basis vector

  Returns:
    Nothing.
  """

  plt.figure(figsize=[4, 4])
  plt.plot(X[:, 0], X[:, 1], '.', color=[.5, .5, .5], label='Data')
  plt.axis('equal')
  plt.xlabel('Neuron 1 activity')
  plt.ylabel('Neuron 2 activity')
  plt.plot([0, W[0, 0]], [0, W[1, 0]], color='r', linewidth=3,
           label='Basis vector 1')
  plt.plot([0, W[0, 1]], [0, W[1, 1]], color='b', linewidth=3,
           label='Basis vector 2')
  plt.legend()
  plt.show()
```

```python
# @title Helper functions

def sort_evals_descending(evals, evectors):
  """
  Sorts eigenvalues and eigenvectors in decreasing order. Also aligns first two
  eigenvectors to be in first two quadrants (if 2D).

  Args:
    evals (numpy array of floats)    : Vector of eigenvalues
    evectors (numpy array of floats) : Corresponding matrix of eigenvectors
                                        each column corresponds to a different
                                        eigenvalue

  Returns:
    (numpy array of floats)          : Vector of eigenvalues after sorting
    (numpy array of floats)          : Matrix of eigenvectors after sorting
  """

  index = np.flip(np.argsort(evals))
  evals = evals[index]
  evectors = evectors[:, index]
  if evals.shape[0] == 2:
    if np.arccos(np.matmul(evectors[:, 0],
                           1 / np.sqrt(2) * np.array([1, 1]))) > np.pi / 2:
      evectors[:, 0] = -evectors[:, 0]
    if np.arccos(np.matmul(evectors[:, 1],
                           1 / np.sqrt(2) * np.array([-1, 1]))) > np.pi / 2:
      evectors[:, 1] = -evectors[:, 1]
  return evals, evectors

def get_data(cov_matrix):
  """
  Returns a matrix of 1000 samples from a bivariate, zero-mean Gaussian

  Note that samples are sorted in ascending order for the first random
  variable.

  Args:
    var_1 (scalar)                     : variance of the first random variable
    var_2 (scalar)                     : variance of the second random variable
    cov_matrix (numpy array of floats) : desired covariance matrix

  Returns:
    (numpy array of floats)            : samples from the bivariate Gaussian,
                                          with each column corresponding to a
                                          different random variable
  """

  mean = np.array([0, 0])
  X = np.random.multivariate_normal(mean, cov_matrix, size=1000)
  indices_for_sorting = np.argsort(X[:, 0])
  X = X[indices_for_sorting, :]
  return X

def calculate_cov_matrix(var_1, var_2, corr_coef):
  """
  Calculates the covariance matrix based on the variances and
  correlation coefficient.

  Args:
    var_1 (scalar)         :  variance of the first random variable
    var_2 (scalar)         :  variance of the second random variable
    corr_coef (scalar)     :  correlation coefficient

  Returns:
    (numpy array of floats) : covariance matrix
  """
  cov = corr_coef * np.sqrt(var_1 * var_2)
  cov_matrix = np.array([[var_1, cov], [cov, var_2]])
  return cov_matrix

def define_orthonormal_basis(u):
  """
  Calculates an orthonormal basis given an arbitrary vector u.

  Args:
    u (numpy array of floats) : arbitrary 2D vector used for new basis

  Returns:
    (numpy array of floats)   : new orthonormal basis columns correspond to
                                basis vectors
  """

  u = u / np.sqrt(u[0] ** 2 + u[1] ** 2)
  w = np.array([-u[1], u[0]])
  W = np.column_stack((u, w))
  return W

def change_of_basis(X, W):
  """
  Projects data onto a new basis.

  Args:
    X (numpy array of floats) : Data matrix each column corresponding to a
                                different random variable
    W (numpy array of floats) : new orthonormal basis columns correspond to
                                basis vectors

  Returns:
    (numpy array of floats)   : Data matrix expressed in new basis
  """

  Y = np.matmul(X, W)
  return Y
```

---
# Section 1: Calculate the eigenvectors of the the sample covariance matrix

As we saw in the lecture, PCA represents data in a new orthonormal basis defined by the eigenvectors of the covariance matrix. Remember that in the previous tutorial, we generated bivariate normal data with a specified covariance matrix $\bf \Sigma$, whose $(i,j)$th element is:
\begin{align}
\Sigma_{ij} = \mathbb{E}[ x_i x_j ] - \mathbb{E}[ x_i] \mathbb{E}[ x_j ] .
\end{align}
However, in real life we don't have access to this ground-truth covariance matrix. To get around this, we can use the sample covariance matrix, $\bf\hat\Sigma$, which is calculated directly from the data. The $(i,j)$th element of the sample covariance matrix is:
\begin{align}
 \hat \Sigma_{ij} =  \frac{1}{N_\text{samples}}{\bf x}_i^T {\bf x}_j - \bar {\bf x}_i \bar{\bf x}_j ,
\end{align}
where ${\bf x}_i = [ x_i(1), x_i(2), \dots,x_i(N_\text{samples})]^T$ is a column vector representing all measurements of neuron $i$, and  $\bar {\bf x}_i$ is the mean of neuron $i$ across samples:
\begin{align}
\bar {\bf x}_i = \frac{1}{N_\text{samples}} \sum_{k=1}^{N_\text{samples}} x_i(k).
\end{align}
If we assume that the data has already been mean-subtracted, then we can write the sample covariance matrix in a much simpler matrix form:
\begin{align}
{\bf \hat \Sigma}
&= \frac{1}{N_\text{samples}} {\bf X}^T {\bf X}.
\end{align}

where $\bf X$ is the full data matrix (each column of $\bf X$ is a different $\bf x_i$).

## Coding Exercise 1.1: Calculate the covariance matrix

Before calculating the eigenvectors, you must first calculate the sample covariance matrix.

**Steps**
* Complete the function `get_sample_cov_matrix` by first subtracting the sample mean of the data, then calculate $\bf \hat \Sigma$ using the equation above.
* Use `get_data` to generate bivariate normal data, and calculate the sample covariance matrix with your finished `get_sample_cov_matrix`. Compare this estimate to the true covariate matrix using `calculate_cov_matrix`. You've seen both `get_data` and `calculate_cov_matrix` in Tutorial 1.

```python
help(get_data)
help(calculate_cov_matrix)
```

    Help on function get_data in module __main__:

    get_data(cov_matrix)
        Returns a matrix of 1000 samples from a bivariate, zero-mean Gaussian

        Note that samples are sorted in ascending order for the first random
        variable.

        Args:
          var_1 (scalar)                     : variance of the first random variable
          var_2 (scalar)                     : variance of the second random variable
          cov_matrix (numpy array of floats) : desired covariance matrix

        Returns:
          (numpy array of floats)            : samples from the bivariate Gaussian,
                                                with each column corresponding to a
                                                different random variable

    Help on function calculate_cov_matrix in module __main__:

    calculate_cov_matrix(var_1, var_2, corr_coef)
        Calculates the covariance matrix based on the variances and
        correlation coefficient.

        Args:
          var_1 (scalar)         :  variance of the first random variable
          var_2 (scalar)         :  variance of the second random variable
          corr_coef (scalar)     :  correlation coefficient

        Returns:
          (numpy array of floats) : covariance matrix

```python
def get_sample_cov_matrix(X):
  """
  Returns the sample covariance matrix of data X

  Args:
    X (numpy array of floats) : Data matrix each column corresponds to a
                                different random variable

  Returns:
    (numpy array of floats)   : Covariance matrix
  """

  #################################################
  ## TODO for students: calculate the covariance matrix
  # Fill out function and remove
  # raise NotImplementedError("Student excercise: calculate the covariance matrix!")
  #################################################

  # Subtract the mean of X
  X = X - np.mean(X,0)
  print(X.shape)

  # Calculate the covariance matrix (hint: use np.matmul)
  cov_matrix = np.matmul(X.T, X)/X.shape[0]
  #cov_matrix = (X.T @ X)/X.shape[0]

  return cov_matrix

# Set parameters
np.random.seed(2020)  # set random seed
variance_1 = 1
variance_2 = 1
corr_coef = 0.8

# Calculate covariance matrix
cov_matrix = calculate_cov_matrix(variance_1, variance_2, corr_coef)
print(cov_matrix)

# Generate data with that covariance matrix
X = get_data(cov_matrix)

# Get sample covariance matrix
sample_cov_matrix = get_sample_cov_matrix(X)
print(sample_cov_matrix)
```

    [[1.  0.8]
     [0.8 1. ]]
    (1000, 2)
    [[0.99315313 0.82347589]
     [0.82347589 1.01281397]]

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

SAMPLE OUTPUT
```
[[1.  0.8]
 [0.8 1. ]]
[[0.99315313 0.82347589]
 [0.82347589 1.01281397]]
```

We get a sample covariance matrix that doesn't look too far off from the true covariance matrix, so that's a good sign!

## Coding Exercise 1.2: Eigenvectors of the covariance matrix

Next you will calculate the eigenvectors of the covariance matrix. Plot them along with the data to check that they align with the geometry of the data.

**Steps:**
* Calculate the eigenvalues and eigenvectors of the sample covariance matrix. (**Hint:** use `np.linalg.eigh`, which finds the eigenvalues of a symmetric matrix).
* Use the provided code to sort the eigenvalues in descending order.
* Plot the eigenvectors on a scatter plot of the data, using the function `plot_basis_vectors`.

```python
help(sort_evals_descending)
```

    Help on function sort_evals_descending in module __main__:

    sort_evals_descending(evals, evectors)
        Sorts eigenvalues and eigenvectors in decreasing order. Also aligns first two
        eigenvectors to be in first two quadrants (if 2D).

        Args:
          evals (numpy array of floats)    : Vector of eigenvalues
          evectors (numpy array of floats) : Corresponding matrix of eigenvectors
                                              each column corresponds to a different
                                              eigenvalue

        Returns:
          (numpy array of floats)          : Vector of eigenvalues after sorting
          (numpy array of floats)          : Matrix of eigenvectors after sorting

```python
#################################################
## TODO for students
# Fill out function and remove
#raise NotImplementedError("Student excercise: calculate and sort eigenvalues")
#################################################

# Calculate the eigenvalues and eigenvectors
evals, evectors = np.linalg.eigh(cov_matrix)

# Sort the eigenvalues in descending order
evals, evectors = sort_evals_descending(evals, evectors)

# Visualize
plot_basis_vectors(X, evectors)
```

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

*Example output:*

---
# Section 2: Perform PCA by projecting data onto the eigenvectors

*Estimated timing to here from start of tutorial: 25 min*

To perform PCA, we will project the data onto the eigenvectors of the covariance matrix, i.e.:
\begin{align}
\bf S = X W
\end{align}
where $\bf S$ is an $N_\text{samples} \times N$ matrix representing the projected data (also called *scores*), and $\bf W$ is an $N\times N$ orthogonal matrix, each of whose columns represents the eigenvectors of the covariance matrix (also called *weights* or *loadings*).

## Coding Exercise 2: PCA implementation

You will now perform PCA on the data using the intuition and functions you have developed so far. Fill in the function below to carry out the steps to perform PCA by projecting the data onto the eigenvectors of its covariance matrix.

**Steps:**
* First subtract the mean and calculate the sample covariance matrix.
* Then find the eigenvalues and eigenvectors and sort them in descending order.
* Finally project the mean-centered data onto the eigenvectors.

```python
help(change_of_basis)
```

    Help on function change_of_basis in module __main__:

    change_of_basis(X, W)
        Projects data onto a new basis.

        Args:
          X (numpy array of floats) : Data matrix each column corresponding to a
                                      different random variable
          W (numpy array of floats) : new orthonormal basis columns correspond to
                                      basis vectors

        Returns:
          (numpy array of floats)   : Data matrix expressed in new basis

```python
def pca(X):
  """
  Sorts eigenvalues and eigenvectors in decreasing order.

  Args:
    X (numpy array of floats): Data matrix each column corresponds to a
                               different random variable

  Returns:
    (numpy array of floats)  : Data projected onto the new basis
    (numpy array of floats)  : Vector of eigenvalues
    (numpy array of floats)  : Corresponding matrix of eigenvectors

  """

  #################################################
  ## TODO for students: calculate the covariance matrix
  # Fill out function and remove
  #raise NotImplementedError("Student excercise: sort eigenvalues/eigenvectors!")
  #################################################

  # Calculate the sample covariance matrix
  cov_matrix = get_sample_cov_matrix(X)

  # Calculate the eigenvalues and eigenvectors
  evals, evectors = np.linalg.eigh(cov_matrix)

  #plt.figure()
  #plt.plot(evals)

  # Sort the eigenvalues in descending order
  evals, evectors =  sort_evals_descending(evals, evectors)

  #plt.figure()
  #plt.plot(evals)

  # Project the data onto the new eigenvector basis
  score = change_of_basis(X, evectors)

  return score, evectors, evals

# Perform PCA on the data matrix X
score, evectors, evals = pca(X)

# Plot the data projected into the new basis
plot_data_new_basis(score)

print(evals)
```

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

*Example output:*

   Finally, we will examine the eigenvalues of the covariance matrix. Remember that each eigenvalue describes the variance of the data projected onto its corresponding eigenvector. This is an important concept because it allows us to rank the PCA basis vectors based on how much variance each one can capture. First run the code below to plot the eigenvalues (sometimes called the "scree plot"). Which eigenvalue is larger?

```python
plot_eigenvalues(evals)
```

## Interactive Demo 2: Exploration of the correlation coefficient

Run the following cell and use the slider to change the correlation coefficient in the data. You should see the scree plot and the plot of basis vectors update.

1.  What happens to the eigenvalues as you change the correlation coefficient?
2.  Can you find a value for which both eigenvalues are equal?
3.  Can you find a value for which only one eigenvalue is nonzero?

```python
# @markdown Make sure you execute this cell to enable the widget!

def refresh(corr_coef=.8):
  cov_matrix = calculate_cov_matrix(variance_1, variance_2, corr_coef)
  X = get_data(cov_matrix)
  score, evectors, evals = pca(X)
  plot_eigenvalues(evals)
  plot_basis_vectors(X, evectors)

_ = widgets.interact(refresh, corr_coef=(-1, 1, .1))
```

    interactive(children=(FloatSlider(value=0.8, description='corr_coef', max=1.0, min=-1.0), Output()), _dom_clas…

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

---
# Summary

*Estimated timing of tutorial: 45 minutes*

- In this tutorial, we learned that goal of PCA is to find an orthonormal basis capturing the directions of maximum variance of the data. More precisely, the $i$th basis vector is the direction that maximizes the projected variance, while being orthogonal to all previous basis vectors. Mathematically, these basis vectors are the eigenvectors of the covariance matrix (also called *loadings*).
- PCA also has the useful property that the projected data (*scores*) are uncorrelated.
- The projected variance along each basis vector is given by its corresponding eigenvalue. This is important because it allows us rank the "importance" of each basis vector in terms of how much of the data variability it explains. An eigenvalue of zero means there is no variation along that direction so it can be dropped without losing any information about the original data.
- In the next tutorial, we will use this property to reduce the dimensionality of high-dimensional data.

---
# Notation

\begin{align}
\mathbf{x_i}  &\quad \text{all measurements of neuron } i\\
\bar{\bf x_i}  &\quad \text{mean across samples for neuron } i \\
\bf \Sigma &\quad \text{covariance matrix}\\
\bf \hat \Sigma &\quad \text{sample covariance matrix}\\
\bf W &\quad \text{weights, loadings matrix}\\
{\bf X} &\quad \text{original data matrix}\\
\bf S &\quad \text{projected matrix, scores}\\
N &\quad \text{data dimensionality}\\
N_\text{samples} &\quad \text{number of samples}\\
\end{align}

---
# Bonus

---
## Bonus Section 1: Mathematical basis of PCA properties

```python
# @title Video 2: Properties of PCA
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  from IPython.display import IFrame
  class BiliVideo(IFrame):
      def __init__(self, id, page=1, width=400, height=300, **kwargs):
          self.id=id
          src = 'https://player.bilibili.com/player.html?bvid={0}&page={1}'.format(id, page)
          super(BiliVideo, self).__init__(src, width, height, **kwargs)

  video = BiliVideo(id="BV1xa4y1a77c", width=854, height=480, fs=1)
  print('Video available at https://www.bilibili.com/video/{0}'.format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id="p56UrMRt6-U", width=854, height=480, fs=1, rel=0)
  print('Video available at https://youtube.com/watch?v=' + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

display(out)
```

    Tab(children=(Output(), Output()), _titles={'0': 'Youtube', '1': 'Bilibili'})


# Finite difference methods for 2D and 3D wave equations

A natural next step is to consider extensions of the methods for
various
variants of the one-dimensional wave equation to two-dimensional (2D) and
three-dimensional (3D) versions of the wave equation.

## Multi-dimensional wave equations

The general wave equation in $d$ space dimensions, with constant
wave velocity $c$,
can be written in the compact form

$$
\begin{equation}
\frac{\partial^2 u}{\partial t^2} = c^2\nabla^2 u\hbox{ for }\xpoint\in\Omega\subset\mathbb{R}^d,\ t\in (0,T] ,
\label{wave:2D3D:model1} \tag{1}
\end{equation}
$$

where

$$
\nabla^2 u = \frac{\partial^2 u}{\partial x^2} +
\frac{\partial^2 u}{\partial y^2} ,
$$

in a 2D problem ($d=2$) and

$$
\nabla^2 u = \frac{\partial^2 u}{\partial x^2} +
\frac{\partial^2 u}{\partial y^2} + \frac{\partial^2 u}{\partial z^2},
$$

in three space dimensions ($d=3$).

Many applications involve variable coefficients, and the general
wave equation in $d$ dimensions is in this case written as

$$
\begin{equation}
\varrho\frac{\partial^2 u}{\partial t^2} = \nabla\cdot (q\nabla u) + f\hbox{ for }\xpoint\in\Omega\subset\mathbb{R}^d,\ t\in (0,T],
\label{wave:2D3D:model2} \tag{2}
\end{equation}
$$

which in, e.g.,  2D becomes

$$
\begin{equation}
\varrho(x,y)
\frac{\partial^2 u}{\partial t^2} =
\frac{\partial}{\partial x}\left( q(x,y)
\frac{\partial u}{\partial x}\right)
+
\frac{\partial}{\partial y}\left( q(x,y)
\frac{\partial u}{\partial y}\right)
+ f(x,y,t)
\thinspace .
\label{_auto1} \tag{3}
\end{equation}
$$

To save some writing and space we may use the index notation, where
subscript $t$, $x$, or $y$ means differentiation with respect
to that coordinate. For example,

$$
\begin{align*}
\frac{\partial^2 u}{\partial t^2} &= u_{tt},\\
\frac{\partial}{\partial y}\left( q(x,y)
\frac{\partial u}{\partial y}\right) &= (q u_y)_y
\thinspace .
\end{align*}
$$

These comments extend straightforwardly to 3D, which means that
the 3D versions of the
two wave PDEs, with and without variable coefficients,
can be stated as

$$
\begin{equation}
u_{tt} = c^2(u_{xx} + u_{yy} + u_{zz}) + f,
\label{wave:2D3D:model1:v2} \tag{4}
\end{equation}
$$

$$
\begin{equation}
\varrho u_{tt} = (q u_x)_x + (q u_y)_y + (q u_z)_z + f\thinspace .
\label{wave:2D3D:model2:v2} \tag{5}
\end{equation}
$$

At *each point* of the boundary $\partial\Omega$ (of $\Omega$) we need
*one* boundary condition involving the unknown $u$.
The boundary conditions are of three principal types:

1. $u$ is prescribed ($u=0$ or a known time variation of $u$ at
    the boundary points, e.g.,
    modeling an incoming wave),

2. $\partial u/\partial n = \boldsymbol{n}\cdot\nabla u$ is prescribed
    (zero for reflecting boundaries),

3. an open boundary condition (also called radiation condition)
    is specified to let waves travel undisturbed out of the domain,
    see [wave:app:exer:radiationBC](#wave:app:exer:radiationBC) for details.

All the listed wave equations with *second-order* derivatives in
time need *two* initial conditions:

1. $u=I$,

2. $u_t = V$.

## Mesh

We introduce a mesh in time and in space. The mesh in time consists
of time points

$$
t_0=0 < t_1 < \cdots < t_{N_t},
$$

normally, for wave equation problems, with a constant spacing $\Delta
t= t_{n+1}-t_{n}$, $n\in\setl{\mathcal{I}_t}$.

Finite difference methods are easy to implement on simple rectangle-
or box-shaped *spatial domains*. More complicated shapes of the
spatial domain require substantially more advanced techniques and
implementational efforts (and a finite element method is usually a more
convenient approach). On a rectangle- or box-shaped domain, mesh
points are introduced separately in the various space directions:

$$
\begin{align*}
&x_0 < x_1 < \cdots < x_{N_x}  \hbox{ in the }x \hbox{ direction},\\
&y_0 < y_1 < \cdots < y_{N_y}  \hbox{ in the }y \hbox{ direction},\\
&z_0 < z_1 < \cdots < z_{N_z}  \hbox{ in the }z \hbox{ direction}\thinspace .
\end{align*}
$$

We can write a general mesh point as $(x_i,y_j,z_k,t_n)$, with
$i\in\mathcal{I}_x$, $j\in\Iy$, $k\in\Iz$, and $n\in\mathcal{I}_t$.

It is a very common choice to use constant mesh spacings:
$\Delta x = x_{i+1}-x_{i}$, $i\in\setl{\mathcal{I}_x}$,
$\Delta y = y_{j+1}-y_{j}$, $j\in\setl{\Iy}$, and
$\Delta z = z_{k+1}-z_{k}$, $k\in\setl{\Iz}$.
With equal mesh spacings one often introduces
$h = \Delta x = \Delta y =\Delta z$.

The unknown $u$ at mesh point $(x_i,y_j,z_k,t_n)$ is denoted by
$u^{n}_{i,j,k}$. In 2D problems we just skip the $z$ coordinate
(by assuming no variation in that direction: $\partial/\partial z=0$)
and write $u^n_{i,j}$.

## Discretization

Two- and three-dimensional wave equations are easily discretized by
assembling building blocks for discretization of
1D wave equations, because the multi-dimensional versions just contain
terms of the same type as those in 1D.

### Discretizing the PDEs

Equation ([4](#wave:2D3D:model1:v2)) can be discretized as

$$
\begin{equation}
[D_tD_t u = c^2(D_xD_x u + D_yD_yu + D_zD_z u) + f]^n_{i,j,k}
\thinspace .
\label{_auto2} \tag{6}
\end{equation}
$$

A 2D version might be instructive to write out in detail:

$$
[D_tD_t u = c^2(D_xD_x u + D_yD_yu) + f]^n_{i,j},
$$

which becomes

$$
\frac{u^{n+1}_{i,j} - 2u^{n}_{i,j} + u^{n-1}_{i,j}}{\Delta t^2}
= c^2
\frac{u^{n}_{i+1,j} - 2u^{n}_{i,j} + u^{n}_{i-1,j}}{\Delta x^2}
+ c^2
\frac{u^{n}_{i,j+1} - 2u^{n}_{i,j} + u^{n}_{i,j-1}}{\Delta y^2}
+ f^n_{i,j},
$$

Assuming, as usual, that all values at time levels $n$ and $n-1$
are known, we can solve for the only unknown $u^{n+1}_{i,j}$. The
result can be compactly written as

$$
\begin{equation}
u^{n+1}_{i,j} = 2u^n_{i,j} + u^{n-1}_{i,j} + c^2\Delta t^2[D_xD_x u + D_yD_y u]^n_{i,j}\thinspace .
\label{wave:2D3D:models:unp1} \tag{7}
\end{equation}
$$

As in the 1D case, we need to develop a special formula for $u^1_{i,j}$
where we combine the general scheme for $u^{n+1}_{i,j}$, when $n=0$,
with the discretization of the initial condition:

$$
[D_{2t}u = V]^0_{i,j}\quad\Rightarrow\quad u^{-1}_{i,j} = u^1_{i,j} - 2\Delta t V_{i,j}
\thinspace .
$$

The result becomes, in compact form,

$$
\begin{equation}
u^{1}_{i,j} = u^0_{i,j} -2\Delta V_{i,j} + {\frac{1}{2}}
c^2\Delta t^2[D_xD_x u + D_yD_y u]^0_{i,j}\thinspace .
\label{wave:2D3D:models:u1} \tag{8}
\end{equation}
$$

The PDE ([5](#wave:2D3D:model2:v2))
with variable coefficients is discretized term by term using
the corresponding elements from the 1D case:

$$
\begin{equation}
[\varrho D_tD_t u = (D_x\overline{q}^x D_x u +
D_y\overline{q}^y D_yu + D_z\overline{q}^z D_z u) + f]^n_{i,j,k}
\thinspace .
\label{_auto3} \tag{9}
\end{equation}
$$

When written out and solved for the unknown $u^{n+1}_{i,j,k}$, one gets the
scheme

$$
\begin{align*}
u^{n+1}_{i,j,k} &= - u^{n-1}_{i,j,k}  + 2u^{n}_{i,j,k} + \\
&\quad \frac{1}{\varrho_{i,j,k}}\frac{1}{\Delta x^2} ( \frac{1}{2}(q_{i,j,k} + q_{i+1,j,k})(u^{n}_{i+1,j,k} - u^{n}_{i,j,k}) - \\
&\qquad\qquad \frac{1}{2}(q_{i-1,j,k} + q_{i,j,k})(u^{n}_{i,j,k} - u^{n}_{i-1,j,k})) + \\
&\quad \frac{1}{\varrho_{i,j,k}}\frac{1}{\Delta y^2} ( \frac{1}{2}(q_{i,j,k} + q_{i,j+1,k})(u^{n}_{i,j+1,k} - u^{n}_{i,j,k}) - \\
&\qquad\qquad\frac{1}{2}(q_{i,j-1,k} + q_{i,j,k})(u^{n}_{i,j,k} - u^{n}_{i,j-1,k})) + \\
&\quad \frac{1}{\varrho_{i,j,k}}\frac{1}{\Delta z^2} ( \frac{1}{2}(q_{i,j,k} + q_{i,j,k+1})(u^{n}_{i,j,k+1} - u^{n}_{i,j,k}) -\\
&\qquad\qquad \frac{1}{2}(q_{i,j,k-1} + q_{i,j,k})(u^{n}_{i,j,k} - u^{n}_{i,j,k-1})) + \\
&\quad \Delta t^2 f^n_{i,j,k}
\thinspace .
\end{align*}
$$

Also here we need to develop a special formula for $u^1_{i,j,k}$
by combining the scheme for $n=0$ with the discrete initial condition,
which is just a matter of inserting
$u^{-1}_{i,j,k}=u^1_{i,j,k} - 2\Delta tV_{i,j,k}$ in the scheme
and solving for $u^1_{i,j,k}$.

### Handling boundary conditions where $u$ is known

The schemes listed above are valid for the internal points in the mesh.
After updating these, we need to visit all the mesh points at the
boundaries and set the prescribed $u$ value.

### Discretizing the Neumann condition

The condition $\partial u/\partial n = 0$ was implemented in 1D by
discretizing it with a $D_{2x}u$ centered difference, followed by
eliminating the fictitious $u$ point outside the mesh by using the
general scheme at the boundary point. Alternatively, one can introduce
ghost cells and update a ghost value for use in the Neumann
condition. Exactly the same ideas are reused in multiple dimensions.

Consider the condition  $\partial u/\partial n = 0$
at a boundary $y=0$ of a rectangular domain $[0, L_x]\times [0,L_y]$ in 2D.
The normal direction is then in $-y$ direction,
so

$$
\frac{\partial u}{\partial n} = -\frac{\partial u}{\partial y},
$$

and we set

$$
[-D_{2y} u = 0]^n_{i,0}\quad\Rightarrow\quad \frac{u^n_{i,1}-u^n_{i,-1}}{2\Delta y} = 0
\thinspace .
$$

From this it follows that $u^n_{i,-1}=u^n_{i,1}$.
The discretized PDE at the boundary point $(i,0)$ reads

$$
\frac{u^{n+1}_{i,0} - 2u^{n}_{i,0} + u^{n-1}_{i,0}}{\Delta t^2}
= c^2
\frac{u^{n}_{i+1,0} - 2u^{n}_{i,0} + u^{n}_{i-1,0}}{\Delta x^2}
+ c^2
\frac{u^{n}_{i,1} - 2u^{n}_{i,0} + u^{n}_{i,-1}}{\Delta y^2}
+ f^n_{i,j},
$$

We can then just insert $u^n_{i,1}$ for $u^n_{i,-1}$ in this equation
and solve for the boundary value $u^{n+1}_{i,0}$, just as was done in 1D.

From these calculations, we see a pattern:
the general scheme applies at the boundary $j=0$ too if we just
replace $j-1$ by $j+1$. Such a pattern is particularly useful for
implementations. The details follow from the explained 1D case
in the section [wave:pde2:Neumann:impl](#wave:pde2:Neumann:impl).

The alternative approach to eliminating fictitious values outside the
mesh is to have $u^n_{i,-1}$ available as a ghost value.  The mesh is
extended with one extra line (2D) or plane (3D) of ghost cells at a
Neumann boundary. In the present example it means that we need a line
with ghost cells below the $y$ axis.  The ghost values must be updated
according to $u^{n+1}_{i,-1}=u^{n+1}_{i,1}$.


```python
import matplotlib.pyplot as plt
%matplotlib widget
import numpy as np
import scipy as sp
import matplotlib as mpl
import matplotlib.pyplot as plt
import chemiscope
from widget_code_input import WidgetCodeInput
from ipywidgets import Textarea
from iam_utils import *
import ase
from ase.io import read, write
```

```python
#### AVOID folding of output cell
```

```python
%%html

<style>
.output_wrapper, .output {
    height:auto !important;
    max-height:4000px;  /* your desired max-height here */
}
.output_scroll {
    box-shadow:none !important;
    webkit-box-shadow:none !important;
}
</style>
```

```python
data_dump = WidgetDataDumper(prefix="module_03")
display(data_dump)
```

```python
module_summary = Textarea("general comments on this module", layout=Layout(width="100%"))
data_dump.register_field("module-summary", module_summary, "value")
display(module_summary)
```

_Reference textbook / figure credits: Giuseppe Grosso, Giuseppe Parravicini, "Solid-state physics", Chapter IX_

# Lattice dynamics of a one-dimensional crystal

Consider a one-dimensional crystal with lattice parameter $a$, so that the positions of atoms in the ideal lattice are $r_n=na$. Each atom may be displaced by a finite amount $u_n$ from the ideal position, so that the actual atomic positions are $r_n=na+u_n$.

_Figure credit: Grosso,Parravicini, "Solid state physics"_

The energy can be computed in terms of a Taylor expansion: if $na$ are equilibrium lattice positions, $\partial E/\partial u_n = 0$ and so the lowest-order term involves the matrix of second derivatives,
$$
E(\mathbf{u}) \approx E_0 + \frac{1}{2} \sum_{ij} \left.\frac{\partial^2 E}{\partial u_i \partial u_j}\right|_{\mathbf{u}=0} u_i u_j
$$

The matrix of second derivatives $D_{ij} \equiv \left.{\partial^2 E}/{\partial u_i \partial u_j}\right|_{\mathbf{u}=0}$ is usually called the *matrix of force constants*, and defines the response of the crystal to perturbations of the atomic positions.

The force acting on the atoms when they are displaced from their equilibrium position is given by $\mathbf{f}=-\partial E/\partial \mathbf{u}$.

 **01** Write down the expression in terms of the elements of $\mathbf{D}$ and $\mathbf{u}$, for a general system.
Consider the case of a 1D bond between two atoms, that keeps them at a distance $d$, with a spring constant $k$. The energy can be written as $E=k(x_1-x_2-d)^2/2$. How would you write this in terms of displacements and a $\mathbf{D}$ matrix?

```python
ex01_txt = Textarea("enter your answer", layout=Layout(width="100%"))
data_dump.register_field("ex01-answer", ex01_txt, "value")
display(ex01_txt)
```

 **02** Write a function that computes the energy and the forces for a lattice containing two atoms, with displacements $u_0$ and $u_1$. Check that the function gives the correct results by validating your input.

```python
# set upt the code widget window
ex01_wci = WidgetCodeInput(
        function_name="lattice_energy_force",
        function_parameters="D_00, D_01, D_10, D_11, u_0, u_1",
        docstring="""
Computes energy and force associated with a given displacement of two atoms in a harmonic lattice.

:param D_ij: elements of the matrix of second derivatives
:param u_i: atomic displacements

:return: A tuple containing the lattice energy and the force, (E, f_0, f_1)
""",
        function_body="""

# Write your solution, and test it by moving the sliders

e = 0.0       # write the potential energy function here
f = [0, 0]    # write the force function here
return e, f[0], f[1]
"""
        )

data_dump.register_field("ex02-function", ex01_wci, "function_body")
```

```python
def draw_forces(ax, D_00, D_01, D_10, D_11, u_0, u_1 ):
    ax.add_artist(plt.Circle((0+u_0, 0), 0.1, color='red'))
    ax.add_artist(plt.Circle((1+u_1, 0), 0.1, color='red'))

    ax.add_artist(plt.Line2D([-0.05,0.05], [-0.05, 0.05], color='k'))
    ax.add_artist(plt.Line2D([-0.05,0.05], [0.05, -0.05], color='k'))
    ax.add_artist(plt.Line2D([1-0.05,1+0.05], [-0.05, 0.05], color='k'))
    ax.add_artist(plt.Line2D([1-0.05,1+0.05], [0.05, -0.05], color='k'))

    U, f0, f1 = ex01_wci.get_function_object()(D_00, D_01, D_10, D_11, u_0, u_1)

    if f0**2>1e-6:
        arr0 = mpl.patches.FancyArrow(0+u_0,0,f0,0,width=0.02)
        ax.add_artist(arr0)
    if f1**2>1e-6:
        arr1 = mpl.patches.FancyArrow(1+u_1,0,f1,0,width=0.02)
        ax.add_artist(arr1)

    ax.set_xlim(-0.5,1.5)
    ax.set_xlabel("x/a")
    ax.set_ylim(-0.5,0.5)
    ax.yaxis.set_ticks([])
    ax.set_aspect('equal')
```

```python
p = WidgetPlot(draw_forces, WidgetParbox(D_00 = (1.0, -2, 2, 0.1, r'$D_{00}$'),
                                         D_01 = (-1.0, -2, 2, 0.1, r'$D_{01}$'),
                                         D_10 = (-1.0, -2, 2, 0.1, r'$D_{10}$'),
                                         D_11 = (1.0, -2, 2, 0.1, r'$D_{11}$'),
                                         u_0 = (-0.1, -0.2, 0.2, 0.01, r'$u_0$'),
                                         u_1 = (0.05, -0.2, 0.2, 0.01, r'$u_0$'),
                                         ));
```

```python
display(p)
```

```python
ex01_wcc = WidgetCodeCheck(ex01_wci, ref_values =
                           {
                               (0,0,0,0,0,0) : (0,0,0),
                               (1,2,3,4,0,0) : (0,0,0),
                               (0,0,0,0,1,2) : (0,0,0),
                               (1,2,3,4,1,2) : (13.5,-6,-10.5),
                               (1,-1,-1,1,1,1) : (0,0,0),
                               (1,1,1,1,1,1) : (2,-2,-2),
                               (1,0,0,1,1,-1) : (1, -1, 1)

                           }, demo = p)
```

```python
display(ex01_wcc)
```

 **03** Set the values of the matrix of force constants to large random values, and then adjust the displacement values so that both atoms have the same displacement $u_0=u_1>0$.
What do you observe? Is this a physical behavior if these two atoms represent a cell in a periodic system?

```python
ex03_txt = Textarea("enter your answer", layout=Layout(width="100%"))
data_dump.register_field("ex03-answer", ex03_txt, "value")
display(ex03_txt)
```

# Properties of the matrix of force constants

When the matrix of force constants describes the interactions in a periodic crystal, it must fulfill several physical conditions:
* The matrix is symmetric $D_{ij}=D_{ji}$
* The elements only depend on the separation between the atoms, $D_{ij} \equiv -K_{|i-j|}$. $\mathbf{K}$ is called the matrix of force constants.
* The matrix must satisfy a condition called *acoustic sum rule* (ASR), $\sum_i D_{ij} = 0$

 **04** Write the expression for the energy of a structure with two atoms, with a $\mathbf{D}$ matrix that fulfills the acoustic sum rule. What happens if the two atoms are displaced by the same amount?
You can go back to the previous exercise, and enter a set of values consistent with the ASR and visualize the effect in practice.

```python
ex04_txt = Textarea("enter your answer", layout=Layout(width="100%"))
data_dump.register_field("ex04-answer", ex04_txt, "value")
display(ex04_txt)
```

# The dispersion relation

We look for a way to express the time dependence of the lattice vibrations, $u_n(t)$. The lattice vibrations must satisfy Newton's equation, which implies that $m\ddot{u}_i=f_i = -\sum_j D_{ij} u_j$.
To solve this in the most general way possible, we express $u_n(t)$ on a plane wave basis,
$$
u_n(t) = \int \,d\omega dq \hat{u}(q,\omega) e^{\mathrm{i} (q n a - \omega t)}.
$$
Given the periodicity of the lattice, one only needs to consider $-\pi/a < q < \pi/a$.
By substituting into the equations of motion, and noting that both the time derivative and the sum over the matrix of force constants are linear operations that commute with the integration, one sees that the only way to ensure a consistent solution for any $\hat{u}(q,\omega)$ is to have
$$
m \omega^2 e^{\mathrm{i} (q i a - \omega t)} = \sum_j D_{ij} e^{\mathrm{i} (q j a - \omega t)},
$$
that is, there is a *dispersion relation* that links $q$ and $\omega$, $m \omega(q)^2 = -\sum_j K_{j} e^{-\mathrm{i} q j a}$.

 **05** Write the function that computes $m \omega(q)^2$ given the force constants $K_{1\ldots 4}$, and plot the frequencies of the lattice vibrations as a function of the wavevector. You can check your solution by comparing the value of the function you computed (red line) with a reference implementation (dashed gray line)

Hints:
* When you change the function, move the sliders to update the plot
* You will need [functions from numpy](https://numpy.org/doc/stable/reference/routines.math.html), that you can access with `np.XXXX` from the function body
* The sum extends to negative values of $j$. Ask yourself what should be the value of $K_{-j}$ given the properties of the matrix of force constants.
* What happens if you combine the terms $j$ and $-j$? Is the dispersion relation real?
* Don't forget there is also a $K_0$! Use the acoustic sum rule to determine its value

```python
# set upt the code widget window
ex02_wci = WidgetCodeInput(
        function_name="m_omega_square",
        function_parameters="K_1, K_2, K_3, K_4, qa",
        docstring="""
Computes the dispersion relation for a one-dimensional lattice given the first four force constants

:param K_i: force constants
:param qa: wavevector to compute the dispersion relation (multiplied by the lattice parameter, so it is adimensional)

:return: the value of m * omega**2
""",
        function_body="""
# Write your solution, and test it by moving the sliders
import numpy as np
return 0.0
"""
        )

data_dump.register_field("ex05-function", ex02_wci, "function_body")
```

```python
def reference_func_02(K_1, K_2, K_3, K_4, q):
    K_0 = -2*(K_1+K_2)-K_3-K_4-K_4-K_3; return -(K_0 + 2*K_1 * np.cos(q*1)+2*K_2 * (2*np.cos(q)**2 - 1)
                                         +2*K_3 * np.cos(q*3)+2*K_4 * (np.cos(q)**4+np.sin(q)**4-6*np.cos(q)**2*np.sin(q)**2 ) )

```

```python
def plot_freq(ax, K_1, K_2, K_3, K_4):
    q = np.linspace(-np.pi, np.pi, 200)
    w2ref = reference_func_02( K_1, K_2, K_3, K_4, q)

    user_function = ex02_wci.get_function_object()
    w2 = [ user_function(K_1, K_2, K_3, K_4, qa) for qa in q]

    ax.plot(q, np.sqrt(np.abs(w2))*np.sign(w2ref), 'r', linewidth=2)
    ax.plot(q, np.sqrt(np.abs(w2ref))*np.sign(w2ref), 'k--')

    ax.set_xlim(-np.pi, np.pi)
    ax.set_xlabel("q a")
    ax.set_ylim(-1,4)
    ax.set_ylabel(r"$\omega\sqrt{m}$")

```

```python
a = WidgetPlot(plot_freq, WidgetParbox(K_1 = (1.0, -2, 2, 0.1, r'$K_{1}$'),
                                       K_2 = (0.0, -2, 2, 0.1, r'$K_{2}$'),
                                       K_3 = (0.0, -2, 2, 0.1, r'$K_{3}$'),
                                       K_4 = (0.0, -2, 2, 0.1, r'$K_{4}$'),
                                       ));
```

```python
ex02_wcc = WidgetCodeCheck(ex02_wci, ref_values = {
    x: reference_func_02(*x) for x in [ (1,2,3,4,0), (1,2,3,4,1), (1,1,1,1,0.5), (-1,1,-2,2,0.25) ]
    }, demo = a)
```

```python
display(ex02_wcc)
```

Make sure the reference implementation matches your function for all values of the force constants.
Observe what happens as you change the parameters.

 **06** Why does it make sense to truncate the expansion at $K_4$? Do you expect the force constants to increase or decrease as $j$ gets larger? Hint: go back to look at the expression of the force acting on atom $0$. How does it depend on the displacement of atom $j$?

```python
ex06_txt = Textarea("enter your answer", layout=Layout(width="100%"))
data_dump.register_field("ex06-answer", ex06_txt, "value")
display(ex06_txt)
```

Note that the plot visualizes $\operatorname{sign}(\omega^2) \sqrt{|m\omega^2|}$, so negative values on the plot correspond to so-called *imaginary frequencies*, i.e. values of $q$ for which $\omega^2$ is negative. Move around the sliders to get some negative lobes

 **07** How will the energy change if you introduce a distortion with a periodicity corresponding to a $q$ associated with a negative $\omega^2$? What does this imply in terms of the stability of the crystal?

```python
ex07_txt = Textarea("enter your answer", layout=Layout(width="100%"))
data_dump.register_field("ex07-answer", ex07_txt, "value")
display(ex07_txt)
```

# Lattice with a basis - The Diatomic spring chain

Consider a one-dimensional lattice with two non-equivalent atoms of masses $M_1$ and $M_2$ in a unit cell of lattice parameter $a_0$. The atoms of mass $M_1$ occupy the sublattice positions $R_{n}^{(1)}$ = $na_0$ and can be displaced by an amount $u_n$, while the atoms of mass $M_2$ occupy the sublattice positions $R_{n}^{(2)}$ = $(n + 1/2)a_0$ and can be displaced by an amount $v_n$.

_Figure credit: Grosso,Parravicini, "Solid state physics"_

The classical equations of motion for the two types of particles are then given by:
$$
M_1\ddot{u}_n = -K(2u_n - v_{n-1}-v_n) \\
M_2\ddot{v}_n = -K(2v_n - u_n - u_{n+1})
$$
To solve this we express $u_n(t)$ and $v_n(t)$ as:
$$
u_n(t) = A_1 e^{i(qna_0 −ωt)}\\
v_n(t) = A_2 e^{i(qna_0 + qa_0/2 −ωt)}
$$
Which gives
$$
-M_1\omega^2A_1 = -K(2A_1 - A_2e^{-iqa_0/2}-A_2e^{iqa_0/2}) \\
-M_2\omega^2A_2 = -K(2A_2 - A_1e^{-iqa_0/2}-A_1e^{iqa_0/2})
$$
This can be written in the matrix form as:
$$
\begin{equation}
   \begin{bmatrix}
      2K-M_1\omega^2 & -2K \cos(qa_0/2) \\
      -2K \cos(qa_0/2) & 2K-M_2\omega^2
   \end{bmatrix}
   \begin{bmatrix}
      A_1 \\
      A_2
   \end{bmatrix}
   = 0
\end{equation}
$$
For this equation, a nontrivial solution exists only if the determinant of the matrix is zero. This leads us to the equation:
$$
(2K-M_1\omega^2)(2K-M_2\omega^2) - 4C^2 \cos^2(qa_0/2)=0
$$
which is a quadratic equation in $\omega^2$, hence,
$$
\omega^2 = K \left( \frac{1}{M_1}+\frac{1}{M_2} \right) \pm K\sqrt{\left( \frac{1}{M_1}+\frac{1}{M_2} \right)^2 - \frac{4\sin^2(qa_0/2)}{M_1 M_2}}
$$

These two values of $\omega$ gives us the dispersion relations, that are illustrated in the figure given below.

```python
#for dispersion relation in diatomic lattice
def diatomic_dispersion(K, M_1, M_2, q):
    w1 = K*(1/M_1 + 1/M_2) + K*(np.sqrt((1/M_1 + 1/M_2)**2 - (4*(np.sin(q/2))**2)/(M_1*M_2)))
    w2 = K*(1/M_1 + 1/M_2) - K*(np.sqrt((1/M_1 + 1/M_2)**2 - (4*(np.sin(q/2))**2)/(M_1*M_2)))
    w = [w1, w2]
    return w
```

```python
def plot_omega(ax, K, M_1, M_2):
    q = np.linspace(-1, 1, 200)
    w_plus = diatomic_dispersion(K, M_1, M_2, q*np.pi)[0]
    w_minus = diatomic_dispersion(K, M_1, M_2, q*np.pi)[1]

    ax.plot(q, np.sqrt(np.abs(w_plus))*np.sign(w_plus), 'r', linewidth=2)
    ax.plot(q, np.sqrt(np.abs(w_minus))*np.sign(w_minus), 'b', linewidth=2)

    ax.set_xlim(-1,1)
    ax.set_xlabel("$q a_0/\pi$")
    ax.set_ylim(-0.1,1.8)
    ax.set_ylabel("$\omega$")

```

```python
A = WidgetPlot(plot_omega, WidgetParbox(M_1 = (9.0, 1, 20, 1, r'$M_{1}$'),
                                       M_2 = (3.0, 1, 20, 1, r'$M_{2}$'),
                                       K = (1.0, 0.1, 2, 0.1, r'$K$'),
                                       ));
```

```python
display(A)
```

The two solutions are associated with two _branches_. The lower one is called _acoustic branch_ and the upper one is called _optical branch_, with a frequency gap between them.

 **08** Play around with the parameters. How does the gap change with mass? What happens when $M_1$ = $M_2$? Explain your observations by relating these results to the monoatomic case.

```python
ex08_txt = Textarea("enter your answer", layout=Layout(width="100%"))
data_dump.register_field("ex08-answer", ex08_txt, "value")
display(ex08_txt)
```