```python
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath('./introduction.ipynb'))))
```

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

from itertools import product
from sklearn.linear_model import LinearRegression, MultiTaskElasticNet, MultiTaskElasticNetCV
from sklearn.ensemble import RandomForestRegressor,RandomForestClassifier
from sklearn.preprocessing import PolynomialFeatures

from ylearn.exp_dataset.exp_data import single_continuous_treatment, single_binary_treatment, multi_continuous_treatment
from ylearn.estimator_model.double_ml import DML4CATE
from ylearn.estimator_model.utils import nd_kron
```

## Single continuous treatment.
Valid estimation models include
- DoubleML
- DeepIV
- CausalTree

### The dataset is generated by the following process [1]
Below we use the einstein notation to alleviate the headaches of specifiying dimensions of tensors.
\begin{align*}
    x & = \beta_i w^i \\
    y & = x \theta(v^i) + \gamma_j w^j + \epsilon \\
    w & \sim \text{Normal}(0, I_{n_w})\\
    v & \sim \text{Uniform}(0, 1)^{n_v}
\end{align*}

[1] https://github.com/microsoft/EconML/blob/main/notebooks/Double%20Machine%20Learning%20Examples.ipynb

```python
train, val, treatment_effect = single_continuous_treatment()
```

```python
train
```

      w_0
      w_1
      w_2
      w_3
      w_4
      w_5
      w_6
      w_7
      w_8
      w_9
      ...
      w_24
      w_25
      w_26
      w_27
      w_28
      w_29
      c_0
      treatment
      outcome
      TE

      1212
      -1.010090
      -1.793770
      -0.704964
      0.083557
      0.481427
      -0.048769
      -0.194278
      -0.817631
      0.267509
      0.388186
      ...
      0.705857
      -1.235907
      -1.240211
      0.307156
      -0.250887
      0.955369
      0.155103
      -0.049943
      -0.391496
      1.363705

      1690
      0.458838
      0.284358
      -0.456428
      1.166463
      -1.124323
      -0.093801
      3.306066
      1.980896
      0.628333
      -0.282192
      ...
      -2.091498
      0.994543
      2.341314
      0.946668
      0.601337
      -0.028910
      0.335181
      -0.495657
      -0.884389
      1.954943

      1198
      -0.894299
      0.720393
      1.682601
      0.215710
      0.732282
      -1.210671
      0.305657
      -0.427936
      3.042174
      -0.673378
      ...
      0.527429
      1.084672
      0.017845
      -1.970667
      0.183399
      0.292241
      0.547452
      1.082807
      3.535996
      2.988898

      401
      -0.298571
      -1.671269
      0.260021
      -1.836556
      1.249896
      -2.137694
      0.234445
      -1.012757
      -0.203846
      -0.045091
      ...
      -0.147662
      0.010062
      -0.701577
      -0.903222
      -0.303378
      -1.196278
      0.477712
      -0.726763
      -5.814209
      2.599774

      512
      -0.008337
      0.312129
      2.535573
      -1.778272
      2.412057
      -2.550452
      0.696291
      1.054566
      1.074932
      -0.494271
      ...
      1.084067
      -1.606550
      0.782404
      0.928994
      -2.466212
      -1.042824
      0.076972
      3.144719
      5.347627
      1.166426

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

      1323
      0.473488
      -0.246827
      -0.829371
      0.167178
      2.051812
      -1.098893
      -2.022564
      1.115111
      -0.843840
      0.027017
      ...
      0.659410
      -0.066118
      0.836309
      -0.187636
      -1.201602
      -0.812195
      0.887783
      1.707117
      9.797455
      5.903625

      929
      0.062928
      0.663101
      -0.003895
      0.283313
      0.150168
      0.676016
      -0.800642
      -2.142750
      1.956881
      -0.023682
      ...
      0.492842
      -0.446452
      -1.392282
      0.161448
      -0.281472
      2.292087
      0.540622
      0.008218
      1.633430
      2.948346

      652
      -0.474054
      -0.208888
      0.843013
      1.233498
      0.275309
      2.065287
      0.029058
      0.007994
      1.498871
      0.045821
      ...
      2.251417
      -0.007995
      -0.729207
      0.872539
      -1.748384
      -0.784325
      0.720702
      -1.758245
      -10.299303
      4.226625

      553
      0.144784
      -0.631134
      -0.768283
      0.646175
      -2.800647
      0.303113
      1.030039
      0.206670
      0.752679
      0.937146
      ...
      -1.607408
      1.784038
      0.048513
      0.593881
      1.885963
      0.295527
      0.217212
      -0.660572
      0.237139
      1.544072

      382
      -1.012182
      -2.310253
      0.726641
      -0.180945
      -0.162788
      0.318379
      -1.055082
      0.725277
      -0.489057
      0.761651
      ...
      0.369722
      -0.378601
      1.938581
      0.781012
      -0.270042
      0.760944
      0.673528
      0.375009
      1.637244
      3.846083

1500 rows × 34 columns

```python
adjustment = train.columns[:-4]
covariate = 'c_0'
outcome = 'outcome'
treatment = 'treatment'
```

```python
dml = DML4CATE(
    x_model=RandomForestRegressor(),
    y_model=RandomForestRegressor(),
    cf_fold=3,
)
dml.fit(
    train,
    outcome,
    treatment,
    adjustment,
    covariate,
)
```

    05-09 17:20:11 I ylearn.e.double_ml.py 563 - _fit_1st_stage: fitting x_model RandomForestRegressor
    05-09 17:20:15 I ylearn.e.double_ml.py 569 - _fit_1st_stage: fitting y_model RandomForestRegressor
    05-09 17:20:18 I ylearn.e.double_ml.py 593 - _fit_2nd_stage: fitting yx_model LinearRegression

    Double Machine Learning Estimator

```python
dml_poly={}
for dg in [3,5]:
    dml_t = DML4CATE(
        x_model=RandomForestRegressor(),
        y_model=RandomForestRegressor(),
        cf_fold=3,
        covariate_transformer = PolynomialFeatures(degree=dg,include_bias=False)
    )
    dml_t.fit(
        train,
        outcome,
        treatment,
        adjustment,
        covariate,
    )
    dml_poly[dg] = dml_t
```

    05-09 17:20:18 I ylearn.e.double_ml.py 563 - _fit_1st_stage: fitting x_model RandomForestRegressor
    05-09 17:20:22 I ylearn.e.double_ml.py 569 - _fit_1st_stage: fitting y_model RandomForestRegressor
    05-09 17:20:25 I ylearn.e.double_ml.py 593 - _fit_2nd_stage: fitting yx_model LinearRegression
    05-09 17:20:25 I ylearn.e.double_ml.py 563 - _fit_1st_stage: fitting x_model RandomForestRegressor
    05-09 17:20:29 I ylearn.e.double_ml.py 569 - _fit_1st_stage: fitting y_model RandomForestRegressor
    05-09 17:20:33 I ylearn.e.double_ml.py 593 - _fit_2nd_stage: fitting yx_model LinearRegression

```python
def exp_te(x): return np.exp(2*x)
dat = np.array(list(product(np.arange(0, 1, 0.01), repeat=1))).ravel()
data_test = pd.DataFrame({'c_0': dat})
true_te = np.array([exp_te(xi) for xi in data_test[covariate]])
ested_te = dml.estimate(data_test).ravel()

```

```python
plt.figure(figsize=(10, 6))
plt.plot(dat, ested_te, label='dml estimated')
for dg, est in dml_poly.items():
    ested_te_dg = est.estimate(data_test).ravel()
    plt.plot(dat, ested_te_dg, label=f'dml estimated with Polynomial(degree={dg})')
plt.plot(dat, true_te, 'b--', label='true')
plt.legend()
plt.show()
```

## Single binary treatment.
Valid estimation models include
- DoubleML
- DeepIV
- CausalTree
- Metalearners
- DoublyRobust

### The dataset is generated by the following process
Below we use the einstein notation to alleviate the headaches of specifiying dimensions of tensors.
\begin{align*}
    x & \sim \text{Bernoulli}(f(w)), \quad f(w) = \sigma(w_i \beta^i + \eta), \quad \eta\sim \text{Uniform}(-1, 1) \\
    y & = x \theta(v^i) + \gamma_j w^j + \epsilon, \quad \epsilon \sim \text{Uniform}(-1, 1)\\
    w & \sim \text{Normal}(0, I_{n_w})\\
    v & \sim \text{Uniform}(0, 1)^{n_v}
\end{align*}

```python
train1, val1, treatment_effect1 = single_binary_treatment()
def exp_te(x): return np.exp(2*x[0])
n = 1000
n_x = 4
X_test1 = np.random.uniform(0, 1, size=(n, n_x))
X_test1[:, 0] = np.linspace(0, 1, n)
data_test_dict = {
    'c_0': X_test1[:, 0],
    'c_1': X_test1[:, 1],
    'c_2': X_test1[:, 2],
    'c_3': X_test1[:, 3],
}
data_test1 = pd.DataFrame(data_test_dict)
true_te = np.array([exp_te(x_i) for x_i in X_test1])
```

```python
adjustment1 = train1.columns[:-7]
covariate1 = train1.columns[-7:-3]
# t_effect1 = train1['t_effect']
treatment = 'treatment'
outcome = 'outcome'
```

```python
dml1 = DML4CATE(
    x_model=RandomForestClassifier(),
    y_model=RandomForestRegressor(),
    cf_fold=1,
    is_discrete_treatment=True,
)
dml1.fit(
    data=train1,
    outcome=outcome,
    treatment=treatment,
    adjustment=adjustment1,
    covariate=covariate1,
)
predicted = dml1.estimate(data_test1)
```

    05-09 17:20:34 I ylearn.e.double_ml.py 563 - _fit_1st_stage: fitting x_model RandomForestClassifier
    05-09 17:20:34 I ylearn.e.double_ml.py 569 - _fit_1st_stage: fitting y_model RandomForestRegressor
    05-09 17:20:35 I ylearn.e.double_ml.py 593 - _fit_2nd_stage: fitting yx_model LinearRegression

```python
data_test1
```

      c_0
      c_1
      c_2
      c_3

      0
      0.000000
      0.818433
      0.631912
      0.956283

      1
      0.001001
      0.167187
      0.670090
      0.500090

      2
      0.002002
      0.993322
      0.075953
      0.045508

      3
      0.003003
      0.973807
      0.526946
      0.705518

      4
      0.004004
      0.064636
      0.876281
      0.067465

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

      995
      0.995996
      0.941305
      0.259215
      0.236504

      996
      0.996997
      0.909833
      0.555308
      0.247834

      997
      0.997998
      0.362549
      0.289977
      0.999383

      998
      0.998999
      0.625134
      0.648058
      0.401298

      999
      1.000000
      0.438546
      0.981148
      0.558519

1000 rows × 4 columns

```python
plt.figure(figsize=(20,5))
plt.subplot(1, 3, 1)
plt.plot(X_test1[:, 0], predicted, label='dml')
# plt.fill_between(X_test[:, 0], lb, ub, alpha=.4)
plt.plot(X_test1[:, 0], true_te, 'b--', label='true effect')
plt.ylabel('Treatment Effect')
plt.xlabel('x')
plt.legend()
plt.show()
```

## Multiple continuous treatment.
Valid estimation models include
- DoubleML
- DeepIV

### The dataset is generated by the following process
Below we use the einstein notation to alleviate the headaches of specifiying dimensions of tensors.
\begin{align*}
    x & = \beta_i w^i + \eta, \quad \eta\sim \text{Uniform}(-1, 1) \\
    y & = x \theta(v) + (x^2) \phi^j(v) + \gamma_k w^k + \epsilon, \quad \epsilon \sim \text{Uniform}(-1, 1)\\
    w & \sim \text{Normal}(0, I_{n_w})\\
    v & \sim \text{Uniform}(0, 1)^{n_v}
\end{align*}

We let
\begin{equation}
    \theta(v) = e^{2v}, \quad \phi(v) = \ln (3v)
\end{equation}

```python
train2, val2, test = multi_continuous_treatment()
data_test2 = test[0]
true_te1 = test[1]
true_te2 = test[2]
```

```python
adjustment2 = train2.columns[:-10]
covariate2 = train2.columns[-10:-5]
treatment2 = train2.columns[-5: -3]
outcome2 = 'outcome'
```

```python
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.multioutput import MultiOutputRegressor

dml2 = DML4CATE(
    x_model=MultiOutputRegressor(GradientBoostingRegressor()),
    y_model=GradientBoostingRegressor(),
    cf_fold=1,
    is_discrete_treatment=False,
)
dml2.fit(
    data=train2,
    outcome=outcome2,
    treatment=treatment2,
    adjustment=adjustment2,
    covariate=covariate2,
)
```

    05-09 17:20:36 I ylearn.e.double_ml.py 563 - _fit_1st_stage: fitting x_model MultiOutputRegressor
    05-09 17:20:43 I ylearn.e.double_ml.py 569 - _fit_1st_stage: fitting y_model GradientBoostingRegressor
    05-09 17:20:47 I ylearn.e.double_ml.py 593 - _fit_2nd_stage: fitting yx_model LinearRegression

    Double Machine Learning Estimator

```python
effect2 = dml2.estimate(data=data_test2)
```

```python
X_test = np.random.uniform(0, 1, size=(100, 5))
X_test[:, 0] = np.linspace(0, 1, 100)

plt.figure(figsize=(13, 6))
plt.plot(X_test[:, 0], effect2[:, :, 0], label='dml for treatment 1')
plt.plot(X_test[:, 0], effect2[:, :, 1], label='dml for treatment 2')
plt.plot(X_test[:, 0], true_te1, '--', label='True effect1')
plt.plot(X_test[:, 0], true_te2, '--', label='True effect2')
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.show()

```

## Single continuous treatment observational data

valid estimation models include
- DeepIV
- DoubleML
- CausalTree

We perform our test on Dominick’s dataset, a popular historical dataset of store-level orange juice prices and sales provided by University of Chicago Booth School of Business. \textbf{We will reveal that lower income consumers are more price-sensive.}[1]

[1] https://github.com/microsoft/EconML/blob/main/notebooks/Double%20Machine%20Learning%20Examples.ipynb

```python
import os
import urllib.request

from sklearn.preprocessing import StandardScaler
```

Data preprocessing

```python
file_name = "oj_large.csv"

if not os.path.isfile(file_name):
    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['price'] = np.log(oj_data['price'])
```

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

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

      0
      2
      tropicana
      40
      9.018695
      0
      1.353255
      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
      1.353255
      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
      1.353255
      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
      1.353255
      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
      1.353255
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927

```python
oj_data.columns
```

    Index(['store', 'brand', 'week', 'logmove', 'feat', 'price', 'AGE60', 'EDUC',
           'ETHNIC', 'INCOME', 'HHLARGE', 'WORKWOM', 'HVAL150', 'SSTRDIST',
           'SSTRVOL', 'CPDIST5', 'CPWVOL5'],
          dtype='object')

```python
outcome3 = 'logmove'
treatment3 = 'price'
adjustment3 = [c for c in oj_data.columns if c not in ['price', 'logmove', 'brand', 'week', 'store','INCOME']]
covariate3 = 'INCOME'
```

```python
x_ = oj_data[treatment3].values
y_ = oj_data[outcome3].values
scaler = StandardScaler()
w_ = scaler.fit_transform(oj_data[adjustment3].values)
w_prime = pd.get_dummies(oj_data[['brand']]).values
w = np.concatenate((w_, w_prime), axis=1)
v = scaler.fit_transform(oj_data[[covariate3]].values).ravel()
```

```python
adjustment3.append('brand')
transfomred_data_dict = {treatment3: x_,
                         outcome3: y_,
                         covariate3: v,}
for i, wwww in enumerate(adjustment3):
    transfomred_data_dict[wwww] = w[:, i].ravel()
```

```python
transfomred_data = pd.DataFrame(transfomred_data_dict)
```

```python
transfomred_data.head()
```

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

      0
      1.353255
      9.018695
      -0.225035
      -0.557732
      0.965476
      0.215706
      -0.220052
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      0.0

      1
      1.353255
      8.723231
      -0.225035
      -0.557732
      0.965476
      0.215706
      -0.220052
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      0.0

      2
      1.353255
      8.253228
      -0.225035
      -0.557732
      0.965476
      0.215706
      -0.220052
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      0.0

      3
      1.353255
      8.987197
      -0.225035
      -0.557732
      0.965476
      0.215706
      -0.220052
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      0.0

      4
      1.353255
      9.093357
      -0.225035
      -0.557732
      0.965476
      0.215706
      -0.220052
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      0.0

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

```python
dml3 = DML4CATE(
    x_model=RandomForestRegressor(),
    y_model=RandomForestRegressor(),
    cf_fold=4,
)
dml3.fit(
    data=transfomred_data,
    outcome=outcome3,
    treatment=treatment3,
    adjustment=adjustment3,
    covariate=covariate3,
)
```

    05-09 17:21:29 I ylearn.e.double_ml.py 563 - _fit_1st_stage: fitting x_model RandomForestRegressor
    05-09 17:21:37 I ylearn.e.double_ml.py 569 - _fit_1st_stage: fitting y_model RandomForestRegressor
    05-09 17:21:45 I ylearn.e.double_ml.py 593 - _fit_2nd_stage: fitting yx_model LinearRegression

    Double Machine Learning Estimator

```python
mx = RandomForestRegressor()
my = RandomForestRegressor()
wv = np.concatenate((w, v.reshape(-1, 1)), axis=1)
mx.fit(wv, x_)
my.fit(wv, y_)
yprime = y_ - my.predict(wv).reshape(y_.shape)
xhat = x_ - mx.predict(wv).reshape(x_.shape)
xprime = (xhat * v).reshape(-1, 1)
l = LinearRegression()
l.fit(xprime, yprime.reshape(-1, 1))
(l.coef_ - dml3.yx_model.coef_) / dml3.yx_model.coef_
```

    array([[-1.14037098,  1.21219004]])

```python
test_data = transfomred_data.copy(deep=True)[:X_test.shape[0]]
test_data[covariate3] = X_test.ravel()
effect3 = dml3.estimate(data=test_data)
```

```python
test_data.head()
```

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

      0
      1.353255
      9.018695
      -1.00
      -0.557732
      0.965476
      0.215706
      -0.220052
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      0.0

      1
      1.353255
      8.723231
      -0.98
      -0.557732
      0.965476
      0.215706
      -0.220052
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      0.0

      2
      1.353255
      8.253228
      -0.96
      -0.557732
      0.965476
      0.215706
      -0.220052
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      0.0

      3
      1.353255
      8.987197
      -0.94
      -0.557732
      0.965476
      0.215706
      -0.220052
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      0.0

      4
      1.353255
      9.093357
      -0.92
      -0.557732
      0.965476
      0.215706
      -0.220052
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      0.0

```python
plt.figure(figsize=(10,6))
plt.plot(X_test, effect3.squeeze(), label="OJ Elasticity")
plt.xlabel(r'Scale(Income)')
plt.ylabel('Orange Juice Elasticity')
plt.legend()
plt.title("Orange Juice Elasticity vs Income")
plt.show()
```

## Multiple continuous treatment, multiple outcome observational data.

```python
oj_data = pd.read_csv(file_name)
oj_data['price'] = np.log(oj_data["price"])
groupbylist = ["store", "week", "AGE60", "EDUC", "ETHNIC", "INCOME",
               "HHLARGE", "WORKWOM", "HVAL150",
               "SSTRDIST", "SSTRVOL", "CPDIST5", "CPWVOL5"]
oj_data1 = pd.pivot_table(oj_data,index=groupbylist,
                          columns=oj_data.groupby(groupbylist).cumcount(),
                          values=['logmove', 'price'],
                          aggfunc='sum').reset_index()
oj_data1.columns = oj_data1.columns.map('{0[0]}{0[1]}'.format)
oj_data1 = oj_data1.rename(index=str,
                           columns={"logmove0": "logmove_T",
                                    "logmove1": "logmove_M",
                                    "logmove2":"logmove_D",
                                    "price0":"price_T",
                                    "price1":"price_M",
                                    "price2":"price_D"})

```

```python
oj_data1.head()
```

      store
      week
      AGE60
      EDUC
      ETHNIC
      INCOME
      HHLARGE
      WORKWOM
      HVAL150
      SSTRDIST
      SSTRVOL
      CPDIST5
      CPWVOL5
      logmove_T
      logmove_M
      logmove_D
      price_T
      price_M
      price_D

      0
      2
      40
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927
      9.018695
      8.407378
      9.264829
      1.353255
      1.153732
      0.463734

      1
      2
      46
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927
      8.723231
      8.449343
      8.987197
      1.353255
      1.153732
      0.989541

      2
      2
      47
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927
      8.253228
      9.911456
      8.831712
      1.353255
      0.871293
      0.737164

      3
      2
      48
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927
      8.987197
      7.896553
      7.965546
      1.353255
      1.153732
      0.737164

      4
      2
      50
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927
      9.093357
      8.009031
      7.377759
      1.353255
      1.153732
      0.737164

```python
outcome4 = ['logmove_T', 'logmove_M', 'logmove_D']
treatment4 = ['price_T', 'price_M', 'price_D']
adjustment4 = [c for c in groupbylist if c not in ['week', 'store', 'INCOME']]
covariate4 = ['INCOME']
```

```python
scaler = StandardScaler()
w4 = scaler.fit_transform(oj_data1[adjustment4].values)
x4 = scaler.fit_transform(oj_data1[covariate4].values)
```

```python
for i, name in enumerate(adjustment4):
    oj_data1[name] = w4[:, i]
oj_data1[covariate4] = x4
oj_data1.head()
```

      store
      week
      AGE60
      EDUC
      ETHNIC
      INCOME
      HHLARGE
      WORKWOM
      HVAL150
      SSTRDIST
      SSTRVOL
      CPDIST5
      CPWVOL5
      logmove_T
      logmove_M
      logmove_D
      price_T
      price_M
      price_D

      0
      2
      40
      0.965476
      0.215706
      -0.220052
      -0.225035
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      9.018695
      8.407378
      9.264829
      1.353255
      1.153732
      0.463734

      1
      2
      46
      0.965476
      0.215706
      -0.220052
      -0.225035
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      8.723231
      8.449343
      8.987197
      1.353255
      1.153732
      0.989541

      2
      2
      47
      0.965476
      0.215706
      -0.220052
      -0.225035
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      8.253228
      9.911456
      8.831712
      1.353255
      0.871293
      0.737164

      3
      2
      48
      0.965476
      0.215706
      -0.220052
      -0.225035
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      8.987197
      7.896553
      7.965546
      1.353255
      1.153732
      0.737164

      4
      2
      50
      0.965476
      0.215706
      -0.220052
      -0.225035
      -0.386145
      -1.055453
      0.502548
      -0.860274
      -0.122426
      -0.26456
      -0.282733
      9.093357
      8.009031
      7.377759
      1.353255
      1.153732
      0.737164

```python
oj_data1.columns
```

    Index(['store', 'week', 'AGE60', 'EDUC', 'ETHNIC', 'INCOME', 'HHLARGE',
           'WORKWOM', 'HVAL150', 'SSTRDIST', 'SSTRVOL', 'CPDIST5', 'CPWVOL5',
           'logmove_T', 'logmove_M', 'logmove_D', 'price_T', 'price_M', 'price_D'],
          dtype='object')

```python
from sklearn import linear_model
dml4 = DML4CATE(
    x_model=MultiTaskElasticNetCV(),
    y_model=MultiTaskElasticNetCV(),
    yx_model=None,
    cf_fold=6,
)
dml4.fit(
    data=oj_data1,
    outcome=outcome4,
    treatment=treatment4,
    adjustment=adjustment4,
    covariate=covariate4,
)
```

    05-09 17:21:54 I ylearn.e.double_ml.py 563 - _fit_1st_stage: fitting x_model MultiTaskElasticNetCV
    05-09 17:22:01 I ylearn.e.double_ml.py 569 - _fit_1st_stage: fitting y_model MultiTaskElasticNetCV
    05-09 17:22:08 I ylearn.e.double_ml.py 593 - _fit_2nd_stage: fitting yx_model LinearRegression

    Double Machine Learning Estimator

```python
min_income = -1
max_income = 1
delta = (1 - (-1)) / 100
X_test = np.arange(min_income, max_income + delta - 0.001, delta).reshape(-1, 1)
test_data4 = oj_data1.copy(deep=True)[:X_test.shape[0]]
test_data4[covariate4] = X_test

effect4 = dml4.estimate(test_data4)
```

```python
oj_data1.columns
```

    Index(['store', 'week', 'AGE60', 'EDUC', 'ETHNIC', 'INCOME', 'HHLARGE',
           'WORKWOM', 'HVAL150', 'SSTRDIST', 'SSTRVOL', 'CPDIST5', 'CPWVOL5',
           'logmove_T', 'logmove_M', 'logmove_D', 'price_T', 'price_M', 'price_D'],
          dtype='object')

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

```

```python
xmodel = MultiTaskElasticNetCV()
xmodel.fit(oj_data1[['INCOME'] + adjustment4].values, oj_data1[treatment4])
x_hat = xmodel.predict(oj_data1[['INCOME']+adjustment4].values)
ymodel = MultiTaskElasticNetCV()
ymodel.fit(oj_data1[['INCOME']+adjustment4].values, oj_data1[outcome4])
yhat = ymodel.predict(oj_data1[['INCOME']+adjustment4].values)
# x_hat = dml4.x_hat_dict['paras'][0]
# y_hat = dml4.y_hat_dict['paras'][0]
xprime = nd_kron(oj_data1[treatment4].values - x_hat, oj_data1[['INCOME']].values)
yprime = oj_data1[outcome4].values - yhat
lr = linear_model.LinearRegression()
lr.fit(xprime, yprime)
c = lr.coef_
c
```

    array([[ 0.33285236, -0.05867611, -0.06675376],
           [-0.09049298,  0.51610244, -0.1371895 ],
           [ 0.03084033, -0.11316553,  0.26850266]])


# Ikeda for many ships with quadratic contributions
The method developed in: ([01.03_ikeda_many_dev](06_ikeda/01.03_ikeda_many_dev.ipynb)) will now be attempted for many ships.
Instead of just calculating ikeda for one roll amplitude, it will now be calculated for 2 amplitudes to derive the quadratic part $B_2$

```python
# %load ../../imports.py
"""
These is the standard setup for the notebooks.
"""

%matplotlib inline
%load_ext autoreload
%autoreload 2

#from jupyterthemes import jtplot
#jtplot.style(theme='onedork', context='notebook', ticks=True, grid=False)

import pandas as pd
pd.options.display.max_rows = 999
pd.options.display.max_columns = 999
pd.set_option("display.max_columns", None)
import numpy as np
import os
import matplotlib.pyplot as plt
#plt.style.use('paper')

#import data
import copy
from rolldecay.bis_system import BisSystem
from rolldecay import database
from mdldb.tables import Run

from sklearn.pipeline import Pipeline
from rolldecayestimators.transformers import CutTransformer, LowpassFilterDerivatorTransformer, ScaleFactorTransformer, OffsetTransformer
from rolldecayestimators.direct_estimator_cubic import EstimatorQuadraticB, EstimatorCubic
from rolldecayestimators.ikeda_estimator import IkedaQuadraticEstimator
import rolldecayestimators.equations as equations
import rolldecayestimators.lambdas as lambdas
from rolldecayestimators.substitute_dynamic_symbols import lambdify
import rolldecayestimators.symbols as symbols
import sympy as sp

from sklearn.metrics import r2_score
import rolldecay.paper_writing as paper_writing

```

```python
from pyscores2.indata import Indata
from pyscores2.runScores2 import Calculation
from pyscores2.output import OutputFile
from pyscores2 import TDPError
import pyscores2
from rolldecayestimators.ikeda import Ikeda, IkedaR
from rolldecayestimators.simplified_ikeda_class import SimplifiedIkeda
import subprocess
from rolldecayestimators import measure
import rolldecay
```

```python
df_all_sections_id = pd.read_csv('all_sections.csv', sep=';')
df_all_sections_id.head()
```

```python
section_groups=df_all_sections_id.groupby(by='loading_condition_id')
```

```python
loading_condition_ids = df_all_sections_id['loading_condition_id'].unique()
mask=pd.notnull(loading_condition_ids)
loading_condition_ids=loading_condition_ids[mask]
```

```python
df_rolldecay = database.load(rolldecay_table_name='rolldecay_quadratic_b', limit_score=0.99,
                             exclude_table_name='rolldecay_exclude')
```

```python
mask=df_rolldecay['loading_condition_id'].isin(loading_condition_ids)
df=df_rolldecay.loc[mask].copy()
```

```python
df['BKB'].fillna(0, inplace=True)
```

```python
def add_cScores(sections):
    sections=sections.copy()
    sections['cScores']=sections['area']/(sections['b']*sections['t'])
    mask=sections['cScores']>1
    sections.loc[mask,'cScores']=1
    return sections

def cut_sections(sections, draught):
    sections=sections.copy()
    mask = sections['t']>draught
    sections.loc[mask,'t']=draught
    sections.loc[mask,'area']-=draught*sections['b'].max()  # Assuming rectangular shape
    return sections

def remove_duplicate_sections(sections):
    sections=sections.copy()
    mask=~sections['x'].duplicated()
    sections=sections.loc[mask]
    assert sections['x'].is_unique
    return sections

def too_small_sections(sections):
    sections=sections.copy()
    small = 0.1
    mask=sections['b']==0
    sections.loc[mask,'b']=small
    mask=sections['t']==0
    sections.loc[mask,'t']=small
    mask=sections['area']==0
    sections.loc[mask,'area']=small
    return sections
```

```python
from scipy.integrate import simps
def calculate_lcb(x, area, **kwargs):
    """
    Calculate lcb from AP
    """
    return simps(y=area*x,x=x)/np.trapz(y=area,x=x)

def calculate_dispacement(x, area, **kwargs):
    """
    Calculate displacement
    """
    return np.trapz(y=area,x=x)
```

```python
class DraughtError(ValueError): pass

def define_indata(row, sections, rho=1000,  g=9.81):

    indata = Indata()

    draught=(row.TA+row.TF)/2
    indata.draught=draught
    if draught<=sections['t'].max():
        sections = cut_sections(sections, draught)
    else:
        raise DraughtError('Draught is too large for sections')

    sections=add_cScores(sections)

    indata.cScores=np.array(sections['cScores'])
    indata.ts=np.array(sections['t'])
    indata.bs=np.array(sections['b'])
    indata.zbars=np.zeros_like(sections['b'])  # Guessing...

    beam=sections['b'].max()
    indata.lpp=sections['x'].max()-sections['x'].min()
    #indata.displacement=row.Volume
    indata.displacement=calculate_dispacement(**sections)

    indata.g=g
    indata.kxx=row.KXX
    indata.kyy=row.lpp*0.4
    lcb=calculate_lcb(x=sections['x'], area=sections['area'])
    indata.lcb=lcb-row.lpp/2
    indata.lpp=row.lpp
    indata.projectName='loading_condition_id_%i' % row.loading_condition_id

    indata.rho=rho
    indata.zcg=row.kg-draught
    #indata.waveFrequenciesMin=0.2
    #indata.waveFrequenciesMax=0.5
    #indata.waveFrequenciesIncrement=0.006
    w=row.omega0/np.sqrt(row.scale_factor)
    indata.waveFrequenciesMin=w*0.5
    indata.waveFrequenciesMax=w*2.0
    N=40
    indata.waveFrequenciesIncrement=(indata.waveFrequenciesMax-indata.waveFrequenciesMin)/N
    indata.runOptions["IE"].set_value(1)

    return indata,sections
```

```python
def create_ikeda(row, indata, output_file, fi_a):

    w = row.omega0
    scale_factor=row.scale_factor
    V = row.ship_speed*1.852/3.6/np.sqrt(scale_factor)
    R = 0.01*row.beam/scale_factor
    lBK=row.BKL/scale_factor
    bBK=row.BKB/scale_factor
    ikeda = Ikeda.load_scoresII(V=V, w=w, fi_a=fi_a, indata=indata, output_file=output_file,
                                scale_factor=scale_factor, lBK=lBK, bBK=bBK)

    ikeda.R = R
    return ikeda
```

```python
def calculate_ikeda(ikeda):

    output = pd.DataFrame()
    output['B_44_hat']   = ikeda.calculate_B44()
    output['B_W0_hat']   = float(ikeda.calculate_B_W0())
    output['B_W_hat']    = float(ikeda.calculate_B_W())
    output['B_F_hat']    = ikeda.calculate_B_F()
    output['B_E_hat']    = ikeda.calculate_B_E()
    output['B_BK_hat']   = ikeda.calculate_B_BK()
    output['B_L_hat']    = float(ikeda.calculate_B_L())
    output['Bw_div_Bw0'] = float(ikeda.calculate_Bw_div_Bw0())
    return output
```

```python
def get_B_1_B2(s1,s2,fi_as:np.ndarray):

    # Derive linear and quadratic part for all components:

    if isinstance(s1,pd.Series):
        result=pd.Series(name=row.name)
        columns = s1.index
    elif isinstance(s1,pd.DataFrame):
        result=pd.DataFrame()
        columns = s1.columns
    else:
        raise ValueError('s1 must be pd.Series or pd.DataFrame')

    x = fi_as*8*row.omega0/(3 * np.pi)
    B_2 = (s2 - s1) / (x[1] - x[0])
    B_1 = s1 - B_2 * x[0]

    # Save all of the component as one linear term: _1 and a quadratic term: _2

    for key in columns:
        new_name_1 = '%s_1' % key
        result[new_name_1] = s1[key]

        new_name_2 = '%s_2' % key
        result[new_name_2] = s2[key]

    return result
```

```python
results = pd.DataFrame()
fi_as = np.deg2rad([1,10])
for run_name, row in df.iterrows():
    loading_condition_id=row['loading_condition_id']
    sections = section_groups.get_group(loading_condition_id)
    sections=remove_duplicate_sections(sections)
    sections=too_small_sections(sections)

    try:
        indata,sections_ = define_indata(row, sections)
    except DraughtError as e:
        print('Draught is too large for sections, this loading condition is skipped.')
        continue

    save_name='%s.in' % row.loading_condition_id
    save_path=os.path.join('scores2',save_name)
    indata.save(save_path)

    calculation = Calculation(outDataDirectory='scores2/result')

    # Run scoresII:
    try:
        calculation.run(indata=indata, b_div_t_max=None, timeout=1.0)
    except TDPError:
        print('Dissregarding the TDPError')
        continue
    except pyscores2.LcgError as e:
        print('Disregarded')
        print(e)
        continue
    except subprocess.TimeoutExpired:
        print('Disregarded, scoresII got stuck...')
        continue

    output_file = OutputFile(filePath=calculation.outDataPath)

    ikeda = create_ikeda(row=row, indata=indata, output_file=output_file, fi_a=fi_as)
    result_datas = calculate_ikeda(ikeda)  # DataFrame with two roll amplitudes

    # Derive linear and quadratic part for all components:
    s1=result_datas.iloc[0]
    s2=result_datas.iloc[1]
    result = get_B_1_B2(s1=s1, s2=s2, fi_as=fi_as)

    results=results.append(result)

```

```python
results['id']=results.index
```

```python
results
```

```python
results.head()
```

```python
df.head()
```

## Make comparison with model tests

```python
df2 = df.copy()
scale_factor = df2['scale_factor']
df2['Volume']/=(scale_factor**3)
df2['Disp']=df2['Volume']
df2['beam']/=scale_factor
df2['lpp']/=scale_factor
df2['kg']/=scale_factor
df2['BKB']/=scale_factor
df2['BKL']/=scale_factor

```

```python
df['Disp']=df2['Volume']
results['Disp']=df2.loc[results.index,'Disp']
results['omega0']=df2.loc[results.index,'omega0']

#df_compare = measure.linearized_matrix(df_rolldecay=df2, df_ikeda=results, do_hatify=False)
#mask=df_compare['B_e_hat_ikeda'].notnull()
#df_compare=df_compare.loc[mask].copy()
```

```python
def linearize_model_tests(df_rolldecay, phi_as = np.deg2rad(np.linspace(1,10,10)), g=9.81, rho=1000):

    df = pd.DataFrame()

    for phi_a in phi_as:
        df_ = measure.linearize_model_test(phi_a=phi_a, df_rolldecay=df_rolldecay, g=g, rho=rho)
        df_['phi_a']=phi_a
        df =df.append(df_, ignore_index=True)

    return df

def linearize_ikedas(df_ikeda, phi_as = np.deg2rad(np.linspace(1,10,10)), g=9.81, rho=1000,
                     components = ['B_44_hat', 'B_F_hat', 'B_W_hat', 'B_E_hat', 'B_BK_hat', 'B_L_hat']):

    df = pd.DataFrame()

    for phi_a in phi_as:
        df_ = measure.linearize_si(phi_a=phi_a, df_ikeda=df_ikeda, do_hatify=False, components=components)
        df_['phi_a']=phi_a
        df =df.append(df_, ignore_index=True)

    return df
```

```python
df_rolldecays=linearize_model_tests(df_rolldecay=df2)
df_ikedas=linearize_ikedas(df_ikeda=results)
df_ikedas.dropna(subset=['B_e_hat'], inplace=True)
```

```python
df_compare = pd.merge(left=df_rolldecays, right=df_ikedas, how='inner', left_on=('id','phi_a'), right_on=('id','phi_a'), suffixes=('','_ikeda'))
```

```python
fig,ax=plt.subplots()
df_compare.plot(x='B_e_hat', y='B_e_hat_ikeda', ax=ax, style='o', alpha=0.05, ms=20, markeredgewidth=0)
```

```python
r2_score(y_true=df_compare['B_e_hat'], y_pred=df_compare['B_e_hat_ikeda'])
```

## Also run Simplified Ikeda for comparison

```python
def calculate_si(si):

    output = pd.DataFrame()
    output['B_44_hat'] = si.calculate_B44()
    output['B_W0_hat'] =si.calculate_B_W0()
    output['B_W_hat'] =si.calculate_B_W()
    output['B_F_hat'] =si.calculate_B_F()
    output['B_E_hat'] =si.calculate_B_E()
    output['B_BK_hat'] =si.calculate_B_BK()
    output['B_L_hat'] =si.calculate_B_L()
    output['Bw_div_Bw0'] =si.calculate_Bw_div_Bw0()
    return output
```

```python
inputs_si=pd.DataFrame()
inputs_si['w']=df['omega0']  # Already model scale
scale_factor=df['scale_factor']
inputs_si['V']=df['ship_speed']*1.852/3.6/np.sqrt(scale_factor)
inputs_si['beam']=df['beam']/scale_factor
inputs_si['lpp']=df['lpp']/scale_factor
inputs_si['kg']=df['kg']/scale_factor
inputs_si['volume']=df['Volume']/(scale_factor**3)
draught=(df['TA']+df['TF'])/2
inputs_si['draught']=draught/scale_factor
inputs_si['A0']=df['A0']
inputs_si['lBK']=df['BKL']/scale_factor
inputs_si['bBK']=df['BKB']/scale_factor

inputs_si_1=inputs_si.copy()
inputs_si_2=inputs_si.copy()

inputs_si_1['fi_a']=fi_as[0]
si_1 = SimplifiedIkeda(**inputs_si_1)
results_si_1 = calculate_si(si_1)

inputs_si_2['fi_a']=fi_as[1]
si_2 = SimplifiedIkeda(**inputs_si_2)
results_si_2 = calculate_si(si_2)

results_si=get_B_1_B2(s1=results_si_1, s2=results_si_2, fi_as=fi_as)
results_si.index=df.index
```

```python
results_si['Disp']=df2.loc[results_si.index,'Disp']
results_si['omega0']=df2.loc[results_si.index,'omega0']
results_si['id']=results_si.index
#df_compare_si = measure.linearized_matrix(df_rolldecay=df2, df_ikeda=results_si, do_hatify=False, suffixes=('','_si'))
df_sis=linearize_ikedas(df_ikeda=results_si)
df_sis.dropna(subset=['B_e_hat'], inplace=True)

```

```python
df_compare_ikeda_si = pd.merge(left=df_ikedas, right=df_sis, how='inner', left_on=('id','phi_a'), right_on=('id','phi_a'), suffixes=('_ikeda','_si'))
df_compare_all = pd.merge(left=df_rolldecays, right=df_compare_ikeda_si, how='inner', left_on=('id','phi_a'), right_on=('id','phi_a'), suffixes=('','_ikedas'))
```

```python
fig,ax=plt.subplots()
df_compare_all.plot(x='B_e_hat', y=['B_e_hat_ikeda','B_e_hat_si'], ax=ax, style='o', alpha=0.40, ms=5, markeredgewidth=0)

ax.set_xlabel(r'$\hat{B_{44}}$ (Ikeda)')
ax.set_ylabel(r'$\hat{B_{44}}$ (SI)')

xlim = ax.get_xlim()
ylim = ax.get_ylim()
lim = np.max([xlim[1],ylim[1]])
ax.set_xlim(0,lim)
ax.set_ylim(0,lim)
ax.plot([0,lim],[0,lim],'r-')

ax.grid(True)
ax.set_aspect('equal', 'box')
ax.get_legend().remove()
```

```python
size=3.5
with plt.style.context('paper'):
    fig,ax=plt.subplots()
    fig.set_size_inches(size,size)

    df_compare_all.plot(x='B_e_hat', y='B_e_hat_si', ax=ax, style='+',
                    label=r'$\hat{B}_e^{SI}$')

    df_compare_all.plot(x='B_e_hat', y='B_e_hat_ikeda', ax=ax, style='.',
                    label=r'$\hat{B}_e^{Ikeda}$')

    ax.set_xlabel(r'$\hat{B}_e^{Model}$')
    ax.set_ylabel(r'$\hat{B}_e$')

    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    lim = np.max([xlim[1],ylim[1]])
    ax.set_xlim(0,lim)
    ax.set_ylim(0,lim)
    #ax.set_title('Total roll damping for Ikeda, Simplified Ikeda and model tests')
    ax.plot([0,lim],[0,lim],'r-')

    ax.grid(True)
    ax.set_aspect('equal', 'box')
    ax.legend();
    paper_writing.save_fig(fig=fig, name='si_ikeda_model')
```

```python
df_compare_all['B_44_fraction_si_ikeda'] = df_compare_all['B_e_hat_si']/df_compare_all['B_e_hat_ikeda']

fig,ax=plt.subplots()
df_compare_all.plot(x='B_e_hat', y='B_44_fraction_si_ikeda', ax=ax, style='.')
ax.set_xlabel(r'$\hat{B_{44}}$')
ax.set_ylabel(r'$\frac{B_{44}(SI)}{B_{44}(Ikeda)}$')

ax.get_legend().remove()
```

```python
r2_score(y_true=df_compare_all['B_e_hat'], y_pred=df_compare_all['B_e_hat_ikeda'])
```

```python
r2_score(y_true=df_compare_all['B_e_hat'], y_pred=df_compare_all['B_e_hat_si'])
```

## Investigating the residuals

```python
def calculate_residuals(suffix_true='_ikeda', suffix_prediction='_si'):
    prefixes = ['B_e_hat',
                #'B_W0_hat_e',
                'B_W_hat_e',
                'B_F_hat_e',
                'B_E_hat_e',
                'B_BK_hat_e',
                'B_L_hat_e',
                #'Bw_div_Bw0',
               ]

    for prefix in prefixes:
        residual_name = '%s_residual%s%s' % (prefix, suffix_prediction, suffix_true)
        name_true='%s%s' % (prefix, suffix_true)
        name_prediction='%s%s' % (prefix, suffix_prediction)

        df_compare_all[residual_name] = df_compare_all[name_prediction] - df_compare_all[name_true]

```

```python
df_compare_all.head()
```

```python
calculate_residuals()
df_compare_all['B_e_hat_residual_si_model'] = df_compare_all['B_e_hat_si'] - df_compare_all['B_e_hat']
```

```python
import seaborn as sns;
#sns.set_theme()
```

```python
g=9.81
df_compare_all['draught']=((df_compare_all['TA'] + df_compare_all['TF'])/2)/df_compare_all['scale_factor']
df_compare_all['OG']=df_compare_all['draught']-df_compare_all['kg']
df_compare_all['beam/draught']=df_compare_all['beam']/df_compare_all['draught']
df_compare_all['V']=df_compare_all['ship_speed']*1.852
df_compare_all['Fn']=df_compare_all['V']/(np.sqrt(df_compare_all['lpp']*g))
df_compare_all[r'OG/d']=df_compare_all['OG']/df_compare_all['draught']
df_compare_all[r'LBK/Lpp']=df_compare_all['BKL']/df_compare_all['lpp']
df_compare_all[r'BBK/beam']=df_compare_all['BKB']/df_compare_all['beam']
df_compare_all['omega_hat']=lambdas.omega_hat(beam=df_compare_all['beam'], g=g, omega0=df_compare_all['omega0'])

df_compare_all['Cb']=df_compare_all['Volume']/(df_compare_all['lpp']*df_compare_all['beam']*df_compare_all['draught'])
interesting=['Cb','A0','OG/d','LBK/Lpp','BBK/beam','omega_hat',r'beam/draught', 'Fn']
#sns.lmplot(data=df_compare,y='B_44_hat_residual_si_ikeda', x=interesting, aspect=0.6);
sns.pairplot(df_compare_all,y_vars='B_e_hat_residual_si_ikeda', x_vars=interesting, aspect=0.6);
```

```python
corr = df_compare_all[interesting+['B_e_hat_residual_si_ikeda']].corr().abs()

fig,ax=plt.subplots()
corr.plot(y='B_e_hat_residual_si_ikeda', kind='bar', ax=ax);
ax.set_title('Absolut correlation coefficient for the error');
```

```python
with plt.style.context('paper'):

    y='B_e_hat_residual_si_model'

    ylabel=r'$\hat{B}_{e}^{SI}-\hat{B}_{e}^{Model}$'
    fig,ax=plt.subplots()
    fig.set_size_inches(size,size)

    df_compare_all.plot(x=r'beam/draught',y=y, ax=ax, style='.')

    ax.set_ylabel(ylabel)
    ax.set_xlabel(r'$\frac{beam}{T}$ [-]')
    ax.set_ylabel(ylabel)
    ax.grid(True)
    ax.get_legend().remove()
    paper_writing.save_fig(fig, name='beam_T_residual')
```

```python
import seaborn as  sns

g = sns.jointplot(x=r'beam/draught', y='B_e_hat_residual_si_model', data=df_compare_all,
                  kind="reg", truncate=False,
                  color="m", height=7)
```

```python
labels={
'Cb' : r'$C_b$ [-]',
r'beam/draught' : r'$\frac{beam}{T}$ [-]',
r'OG/d' : r'$\frac{\overline{OG}}{T}$ [-]',
'A0' : r'$A_{0}$ [-]',
r'BBK/beam' : r'$\frac{BK_B}{beam}$ [-]',
r'LBK/Lpp' : r'$\frac{BK_L}{L_{pp}}$ [-]',
r'omega_hat' : r'$\hat{\omega}$ [-]',
r'fi_a' : r'$\phi_a$ [rad]',
r'Fn' : r'$F_n$ [-]',

}

with plt.style.context('paper'):
    fig,ax=plt.subplots()
    fig.set_size_inches(size,size)

    df_compare_all['B_e_hat_residual_si_model_abs'] = df_compare_all['B_e_hat_residual_si_model'].abs()

    for y in interesting:

        df_=df_compare_all.sort_values(by=y)
        y_=df_[y].abs()
        y_-=y_.min()

        ax.plot(df_['B_e_hat_residual_si_model_abs'], y_, '.', label=labels[y])

    ax.set_xlabel(r'$|\hat{B}_{e}^{SI}-\hat{B}_{e}^{Model}|$')
    ax.set_ylabel('change')

    ax.grid(True)
    ax.legend(loc='upper center', bbox_to_anchor=(1.30, 0.8),
          ncol=1)

    paper_writing.save_fig(fig, name='parameter_residual')
```

## Comparing damping contributions SI vs. model test error

```python
suffix_true='_ikeda'
suffix_prediction='_si'
prefixes = [
    'B_W_hat_e_si',
    'B_F_hat_e_si',
    'B_E_hat_e_si',
    'B_BK_hat_e_si',
    'B_L_hat_e_si',]

labels=[
    r'$\hat{B_{W}}$',
    r'$\hat{B_{F}}$',
    r'$\hat{B_{E}}$',
    r'$\hat{B_{BK}}$',
    r'$\hat{B_{L}}$',
    ]

lim = np.max([df_compare_all['B_e_hat_ikeda'].max(),
              df_compare_all['B_e_hat_si'].max(),
             ])
with plt.style.context('paper'):

    fig,ax=plt.subplots()
    fig.set_size_inches(size,size)

    df_compare_all.plot(x='B_e_hat_residual_si_model_abs', y=prefixes, style='.', label=labels, ax=ax)
    ax.set_xlabel(r'$|\hat{B}_{e}^{SI}-\hat{B}_{e}^{Model}|$')

    ax.legend()
    ax.grid(True)
    ax.set_ylabel(r'$\hat{B}$ [-]')

    # (depricated) paper_writing.save_fig(fig, name='component_residual')
```

```python
suffix_true='_ikeda'
suffix_prediction='_si'
prefixes = [
    'B_W_hat_e',
    'B_F_hat_e',
    'B_E_hat_e',
    'B_BK_hat_e',
    'B_L_hat_e',]

x_labels=['%s_ikeda' % label for label in prefixes]
y_labels=['%s_si' % label for label in prefixes]

fig,ax=plt.subplots()
for prefix in prefixes:
    x='%s_ikeda' % prefix
    y='%s_si' % prefix

    df_compare_all.plot(x=x, y=y, label=prefix, ax=ax, style='.')

```

```python
parts=[
'B_W_hat_e_si',
'B_F_hat_e_si',
'B_E_hat_e_si',
'B_BK_hat_e_si',
'B_L_hat_e_si',
]

df_compare_all['B_e_hat_residual_si_model_abs'] = df_compare_all['B_e_hat_residual_si_model'].abs()
corr = df_compare_all[parts+['B_e_hat_residual_si_model_abs']].corr().abs()
corr.drop(index='B_e_hat_residual_si_model_abs', inplace=True)

fig,ax=plt.subplots()
corr.plot(y='B_e_hat_residual_si_model_abs', kind='bar', ax=ax);
ax.set_title('Absolut correlation coefficient for the error');
ax.legend().remove()
```

## Save to table

```python
mean_='$R^2$'
number_of_points = 'Number of points'

validation = pd.DataFrame(columns=[mean_, number_of_points])

s = pd.Series(name='Ikeda')
s[mean_] = r2_score(y_true=df_compare_all['B_e_hat'], y_pred=df_compare_all['B_e_hat_ikeda'])
s[number_of_points] = int(len(df_compare_all))
validation=validation.append(s)

s = pd.Series(name='SI no limits')
s[mean_] = r2_score(y_true=df_compare_all['B_e_hat'], y_pred=df_compare_all['B_e_hat_si'])
s[number_of_points] = int(len(df_compare_all))
validation=validation.append(s)

validation[number_of_points]=validation[number_of_points].astype(int)
```

```python
validation
```

```python
latex = validation.to_latex(float_format='%0.2f', na_rep='')
name='si_ikeda_validation'
file_path = os.path.join(rolldecay.equations_path,name)

# (Uncomment this one if you want to regenerate this table)
paper_writing.save_table(file_path=file_path, tabular_tex=latex, label='tab:si_ikeda_validation',
           caption='Validation of SI and Ikeda')
```

```python
len(df_compare_all['id'].unique())
```

```python
len(df_compare_all['loading_condition_id'].unique())
```

```python
len(df_compare_all['model_number'].unique())
```

```python
#df_compare_all['draught']=(df_compare['TA'] + df_compare['TF'])/2

df_compare_all['CB'] = df_compare_all['Volume']/(df_compare_all['lpp']*df_compare_all['beam']*df_compare_all['draught'])
df_compare_all['g'] = 9.81
df_compare_all[r'B/d']=df_compare_all['beam']/df_compare_all['draught']
df_compare_all['OG']=df_compare_all['draught']-df_compare_all['kg']
df_compare_all[r'OG/d']=df_compare_all['OG']/df_compare_all['draught']
df_compare_all['CMID']=df_compare_all['A0']
df_compare_all[r'bBk/B']=df_compare_all['BKB']/df_compare_all['beam']
df_compare_all[r'lBk/LPP']=df_compare_all['BKL']/df_compare_all['lpp']
df_compare_all['OMEGA_hat']=lambdas.omega_hat(beam=df_compare_all['beam'], g=df_compare_all['g'], omega0=df_compare_all['omega0'])
```

```python
from rolldecayestimators.simplified_ikeda import limits_kawahara
limits_kawahara=pd.Series(limits_kawahara)
```

```python
def cut(df, limits, item):
    df=df.copy()
    mask = ((df[item] >= limits[0]) & (df[item] <= limits[1])) | (df[item]==0)
    df=df.loc[mask].copy()
    return df

def plot_select(df_compare, limits_used):

    df_compare_within = df_compare.copy()

    fig,axes=plt.subplots(nrows=len(limits_used))
    fig.set_size_inches(15,12)

    for (item, limits),ax in zip(limits_used.items(),axes):

        n_bins = 20
        bins=np.linspace(df_compare[item].min(), df_compare[item].max(),n_bins)
        df_compare.hist(item, ax=ax, bins=bins, color='gray')

        df_compare_within = cut(df=df_compare_within, limits=limits, item=item)
        df_compare_within.hist(item, ax=ax, bins=bins, color='blue')
        ax.set_xlabel(item)
        ax.set_title('')

        ylims = ax.get_ylim()
        ax.fill_between(limits, [ylims[1],ylims[1]], y2=0, color='green', alpha=0.5, label='valid')

    plt.tight_layout()
    return df_compare_within
```

```python
df_ = plot_select(df_compare_all, limits_kawahara)
```

```python
len(df_)
```

```python

```


```python
%pylab inline
import sk_dsp_comm.sigsys as ss
#from sk_dsp_comm.pyaudio_helper import pyaudio_helper as pah
import sk_dsp_comm.fir_design_helper as fir_d
import scipy.signal as signal
import scipy.io as io
from ipywidgets import interact, interactive, fixed, interact_manual
import ipywidgets as widgets
from IPython.display import Audio, display
from IPython.display import Image, SVG
```

```python
pylab.rcParams['savefig.dpi'] = 100 # default 72
%config InlineBackend.figure_formats=['svg'] # SVG inline viewing
```

# Introduction
A simplified block diagram of PyAudio *streaming-based* (nonblocking) signal processing when using `pyaudio_helper` and `ipython` widgets.

```python
Image("audio_files/pyaudio_dsp_IO.png", width="90%")
```

## Available Audio I/O Devices
If you add or delete devices by plugging or unplugging USB audio ibterface, this list becomdes invalid. Restart the kernel and run again to get the correct device index list. For two channel apps both the input and output devices must support two channels. For the Sabrent USB audio devices, which has one input and two outputs, Windows for example may improperly list the devices as having two inputs.

```python
pah.available_devices()
```

Index 0 device name = Built-in Microph, inputs = 2, outputs = 0

Index 1 device name = Built-in Output, inputs = 0, outputs = 2

# Real-Time Loop Through
Here we set up a simple `callback` function that passes the input samples directly to the output. The module `pyaudio_helper` provides a class for managing a `pyaudio` stream object, capturing the samples processed by the `callback` function, and collection of performance metrics. Once the `callback` function is written/declared a `DSPIOStream` object can be created and then the `stream(Tsec)` method can be executed to start the input/output processing, e.g.,

```python
import pyaudio_helper as pah

DSP_IO = pah.DSPIOStream(callback,in_idx, out_idx)
DSP_IO.interactive_stream(Tsec = 2, numChan = 1)
```
where `in_idx` is the index of the chosen input device found using `available_devices()` and similarly `out_idx` is the index of the chosen output device.

* The `callback` function must be written first as the function name used by the object to call the *callback*.

### A Minimal Callback
No globals required here as there is no instrumentation configured, externally defined algorithm coefficients, and no widgets being used.

```python
# define a pass through, y = x, callback
def callback(in_data, frame_count, time_info, status):
    # convert byte data to ndarray
    in_data_nda = np.frombuffer(in_data, dtype=np.int16)
    #***********************************************
    # DSP operations here
    # Here we simply pass the input to the output, i.e.
    # y[n] = x[n]
    x = in_data_nda.astype(float32)
    y = x
    # Typically more DSP code here
    #
    #***********************************************
    # Convert from float back to int16
    y = y.astype(int16)
    # Convert ndarray back to bytes
    return y.tobytes(), pah.pyaudio.paContinue
```

### A Basic Callback
This callback makes use of the instrumentation capabilities of the `DSPIOStream` and also has a simple lowpass filter waiting *in-the-wings* if a line of code in commented out and a following line is uncomments, e.g.,
```python
    #y = x
    # Typically more DSP code here
    y, zi = signal.lfilter(b,a,x,zi=zi) # for FIR or simple IIR
```
Notice that `globals` are now used for the `DSP_IO` object, the filter coefficients in arrays, `a` and `b`, and also the filter states held in the array `zi`. In its present form te filtering is commented out, but can be uncommented to allow a simple 1st-order IIR lowpass filter to operate on one channel of audio streaming through the system.

```python
# Add a simple IIR LPF
fs = 48000 # Assummed sampling rate
f3 = 1000 # Hz
a = [1, -exp(-2*pi*f3/fs)]
b = [1 - exp(-2*pi*f3/fs)]
zi = signal.lfiltic(b,a,[0])
```

```python
# define a pass through, y = x, callback
def callback(in_data, frame_length, time_info, status):
    global DSP_IO, b, a, zi
    DSP_IO.DSP_callback_tic()
    # convert byte data to ndarray
    in_data_nda = np.frombuffer(in_data, dtype=np.int16)
    #***********************************************
    # DSP operations here
    # Here we apply a linear filter to the input
    x = in_data_nda.astype(float32)
    y = x
    # Typically more DSP code here
    #y, zi = signal.lfilter(b,a,x,zi=zi) # for FIR or simple IIR
    #***********************************************
    # Save data for later analysis
    # accumulate a new frame of samples
    DSP_IO.DSP_capture_add_samples(y)
    #***********************************************
    # Convert from float back to int16
    y = y.astype(int16)
    DSP_IO.DSP_callback_toc()
    # Convert ndarray back to bytes
    #return (in_data_nda.tobytes(), pyaudio.paContinue)
    return y.tobytes(), pah.pyaudio.paContinue
```

```python
DSP_IO = pah.DSPIOStream(callback,in_idx=0,out_idx=1,fs=48000,Tcapture=0)
```

Index 0 device name = Built-in Microph, inputs = 2, outputs = 0

Index 1 device name = Built-in Output, inputs = 0, outputs = 2

```python
DSP_IO.interactive_stream(Tsec=0,numChan=1)
```

```python
Image("audio_files/start_stop_stream.png", width='55%')
```

#### iMic Noise Capture Gain Flatness
With the iMic plugged in the input/output device indices can be reconfigured to use the iMic index for both the input output streams. The [Analog Discovery (AD2)](https://store.digilentinc.com/analog-discovery-2-100msps-usb-oscilloscope-logic-analyzer-and-variable-power-supply/) is then used to drive a white noise test signal into the ADC and capture the output from the DAC. This allows us to measure the ADC-DAC frequency response using a long-term time average spectral estimate capabilities of the AD2. A second test capture is to use `DSP_IO.DSP_capture_add_samples(y)` to capture the response of the ADC alone, and perform spectral analysis here in the Jupyter notebook. For this capture we set `Tcapture=20`s two cells above and `Tsec=20` one cell above. A comparison of the ADC-alone and ADC-DAC spectrum normalized to look like the frequency response is done in the cell below.

```python
f_AD,Mag_AD = loadtxt('audio_files/Loop_through_noise_SA_iMic.csv',
                       delimiter=',',skiprows=6,unpack=True)
plot(f_AD,Mag_AD-Mag_AD[100])
ylim([-10,5])
xlim([0,20e3])
ylabel(r'ADC Gain Flatness (dB)')
xlabel(r'Frequency (Hz)')
legend((r'ADC-DAC from AD2 SA dB Avg',))
title(r'Loop Through Gain Flatness using iMic at $f_s = 48$ kHz')
grid();
```

The callback stats when capturing data using `DSP_IO.DSP_capture_add_samples(y)` and a plot of the time domain samples.

```python
Nstop = 1000
plot(arange(0,len(DSP_IO.data_capture[:Nstop]))/48000,DSP_IO.data_capture[:Nstop])
DSP_IO.stream_stats()
```

Note for a attributes used in the above examples the `frame_length` is always 1024 samples and the sampling rate $f_s = 48$ ksps. The ideal callback period is this
$$
    T_{cb} = \frac{1024}{480100} = 21.33\ \text{(ms)}
$$

Next consider what the captures `tic` and `toc` data revels about the processing. Calling the method `cb_active_plot()` produces a plot similar to what an electrical engineer would see what using a logic analyzer to show the time spent in an *interrupt service routine* of an *embedded system*. The latency is also evident. You expect to see a minimum latency of two frame lengths (input buffer fill and output buffer fill),e.g.,

$$
    T_\text{latency} >= 2\times \frac{1024}{48000} \times 1000 = 42.6\ \text{(ms)}
$$

The host processor is multitasking, so the latency can be even greater. A true real-time DSP system would give the signal processing high priority and hence much lower is expected, particularly if the `frame_length` can be made small.

## Real-Time Filtering
Here we set up a `callback` function that filters the input samples and then sends them to the output.

```python
import pyaudio_helper as pah

DSP_IO = pah.DSPIOStream(callback,in_idx, out_idx)
DSP_IO.interactive_stream(2,1)
```
where `in_idx` is the index of the chosen input device found using `available_devices()` and similarly `out_idx` is the index of the chosen output device.

* The `callback` function must be written first as the function name is used by the object to call the *callback*
* To demonstrate this we first design some filters that can be used in testing

```python
b = fir_d.fir_remez_bpf(2700,3200,4800,5300,.5,50,48000,18)
a = [1]
fir_d.freqz_resp_list([b],[1],'dB',48000)
ylim([-60,5])
grid();
zi = signal.lfiltic(b,a,[0])
```

```python
f_AD,Mag_AD = loadtxt('audio_files/FIR_BPF_2700_3200_4800_5300_p5dB_50dB_48k.csv',
                       delimiter=',',skiprows=6,unpack=True)
plot(f_AD,Mag_AD-max(Mag_AD)+1)
f = arange(0,20000,10)
w,H_BPF = signal.freqz(b,1,2*pi*f/48000)
plot(f,20*log10(abs(H_BPF)))
ylabel(r'Gain (dB)')
xlabel(r'Frequency (Hz)')
legend((r'AD2 Noise Measured',r'Design Theory'))
title(r'4 kHz 182-Tap FIR Bandpass Design at $f_s = 48$ kHz')
ylim([-60,5])
xlim([2000,8000])
grid();
```

```python
# Design an IIR Notch
b, a = ss.fir_iir_notch(2000,48000,r= 0.9)
fir_d.freqz_resp_list([b],[a],'dB',48000,4096)
ylim([-60,5])
grid();
zi = signal.lfiltic(b,a,[0])
```

Create some global variables for the filter coefficients and the filter state array (recall that a filter has memory).

```python
# define callback (#2)
def callback2(in_data, frame_count, time_info, status):
    global DSP_IO, b, a, zi
    DSP_IO.DSP_callback_tic()
    # convert byte data to ndarray
    in_data_nda = np.frombuffer(in_data, dtype=np.int16)
    #***********************************************
    # DSP operations here
    # Here we apply a linear filter to the input
    x = 5*in_data_nda.astype(float32)
    #y = x
    # The filter state/(memory), zi, must be maintained from frame-to-frame
    # for FIR or simple IIR
    y, zi = signal.lfilter(b,a,x,zi=zi)
    # for IIR use second-order sections
    #y, zi = signal.sosfilt(sos,x,zi=zi)
    #***********************************************
    # Save data for later analysis
    # accumulate a new frame of samples
    DSP_IO.DSP_capture_add_samples(y)
    #***********************************************
    # Convert from float back to int16
    y = y.astype(int16)
    DSP_IO.DSP_callback_toc()
    return y.tobytes(), pah.pyaudio.paContinue
```

```python
DSP_IO = pah.DSPIOStream(callback2,2,2,fs=48000,Tcapture=0)
```

```python
DSP_IO.interactive_stream(Tsec=0,numChan=1)
```

```python
Image("audio_files/start_stop_stream.png", width='55%')
```

## Playback Only Using an Audio Loop
A playback audio loop is created using the `pah.loop_audio` class filled with samples input from a `wav` file. In the example below we take a two-channel (stereo) `wav` file and convert to one channel.

```python
# define callback (3)
# Here we configure the callback to play back a wav file
def callback3(in_data, frame_count, time_info, status):
    global DSP_IO, x
    DSP_IO.DSP_callback_tic()

    # Ignore in_data when generating output only
    #***********************************************
    global x
    # Note wav is scaled to [-1,1] so need to rescale to int16
    y = 32767*x.get_samples(frame_count)
    # Perform real-time DSP here if desired
    #
    #***********************************************
    # Save data for later analysis
    # accumulate a new frame of samples
    DSP_IO.DSP_capture_add_samples(y)
    #***********************************************
    # Convert from float back to int16
    y = y.astype(int16)
    DSP_IO.DSP_callback_toc()
    return y.tobytes(), pah.pyaudio.paContinue
```

```python
fs, x_wav2 = ss.from_wav('Music_Test.wav')
x_wav = (x_wav2[:,0] + x_wav2[:,1])/2
x = pah.loop_audio(x_wav)
DSP_IO = pah.DSPIOStream(callback3,0,1,fs=44100,Tcapture=2)
DSP_IO.interactive_stream(20) # play for 20s but capture only the last 2s
```

```python
Image("audio_files/start_stop_stream.png", width='55%')
```

```python
Npts = 96000
Nstart = 0
plot(arange(len(DSP_IO.data_capture[Nstart:Nstart+Npts]))*1000/44100,
     DSP_IO.data_capture[Nstart:Nstart+Npts]/2**(16-1))
title(r'A Portion of the capture buffer')
ylabel(r'Normalized Amplitude')
xlabel(r'Time (ms)')
grid();
```

```python
Image("audio_files/music_buffer_plot.png", width="75%")
```

Finally, the spectrum of the output signal. To apply custom scaling we use a variation of `psd()` found in the `sigsys` module. If we are plotting the spectrum of *white* noise sent through a filter, the output PSD will be of the form $\sigma_w^2|H(e^{j2\pi f/f_s})|^2$, where $\sigma_w^2$ is the variance of the noise driving the filter. You may choose to overlay a plot of

# Widgets Examples

## Stereo Gain Sliders

```python
L_gain = widgets.FloatSlider(description = 'L Gain',
                continuous_update = True,
                value = 1.0,
                min = 0.0,
                max = 2.0,
                step = 0.01,
                orientation = 'vertical')
R_gain = widgets.FloatSlider(description = 'R Gain',
                continuous_update = True,
                value = 1.0,
                min = 0.0,
                max = 2.0,
                step = 0.01,
                orientation = 'vertical')

#widgets.HBox([L_gain, R_gain])
```

```python
# L and Right Gain Sliders
def callback(in_data, frame_count, time_info, status):
    global DSP_IO, L_gain, R_gain
    DSP_IO.DSP_callback_tic()
    # convert byte data to ndarray
    in_data_nda = np.frombuffer(in_data, dtype=np.int16)
    # separate left and right data
    x_left,x_right = DSP_IO.get_LR(in_data_nda.astype(float32))
    #***********************************************
    # DSP operations here
    y_left = x_left*L_gain.value
    y_right = x_right*R_gain.value

    #***********************************************
    # Pack left and right data together
    y = DSP_IO.pack_LR(y_left,y_right)
    # Typically more DSP code here
    #***********************************************
    # Save data for later analysis
    # accumulate a new frame of samples
    DSP_IO.DSP_capture_add_samples_stereo(y_left,y_right)
    #***********************************************
    # Convert from float back to int16
    y = y.astype(int16)
    DSP_IO.DSP_callback_toc()
    # Convert ndarray back to bytes
    #return (in_data_nda.tobytes(), pyaudio.paContinue)
    return y.tobytes(), pah.pyaudio.paContinue
```

```python
DSP_IO = pah.DSPIOStream(callback,0,1,fs=48000,t_capture=0)
start_stop_buttons = DSP_IO.interactive_stream(0,2)
widgets.VBox([start_stop_buttons, widgets.HBox([L_gain, R_gain])])
```

```python
Image("audio_files/left_right_gain.png", width="65%")
```

## Cross Panning

```python
panning = widgets.FloatSlider(description = 'Panning (%)',
                            continuous_update = True,       # Continuous updates
                            value = 50.0,
                            min = 0.0,
                            max = 100.0,
                            step = 0.1,
                            orientation = 'horizontal')
#display(panning)
```

```python
# Panning
def callback(in_data, frame_count, time_info, status):
    global DSP_IO, panning
    DSP_IO.DSP_callback_tic()
    # convert byte data to ndarray
    in_data_nda = np.frombuffer(in_data, dtype=np.int16)
    # separate left and right data
    x_left,x_right = DSP_IO.get_LR(in_data_nda.astype(float32))
    #***********************************************
    # DSP operations here
    y_left = (100-panning.value)/100*x_left \
              + panning.value/100*x_right
    y_right = panning.value/100*x_left \
              + (100-panning.value)/100*x_right

    #***********************************************
    # Pack left and right data together
    y = DSP_IO.pack_LR(y_left,y_right)
    # Typically more DSP code here
    #***********************************************
    # Save data for later analysis
    # accumulate a new frame of samples
    DSP_IO.DSP_capture_add_samples_stereo(y_left,y_right)
    #***********************************************
    # Convert from float back to int16
    y = y.astype(int16)
    DSP_IO.DSP_callback_toc()
    # Convert ndarray back to bytes
    #return (in_data_nda.tobytes(), pyaudio.paContinue)
    return y.tobytes(), pah.pyaudio.paContinue
```

```python
FRAMES = 512
# Create streaming object
DSP_IO = pah.DSPIOStream(callback,0,1,
                           fs=48000,
                           frame_length = FRAMES,
                           t_capture=0)

# interactive_stream runs in a thread
#so widget can be used
start_stop_buttons = DSP_IO.interactive_stream(0,2)

# display panning widget
widgets.VBox([start_stop_buttons, panning])
```

```python
Image("audio_files/cross_panning.png", width='55%')
```

## Three Band Equalizer
Here we consider a three-band equalizer operating on a music loop. Each peaking filter has system function in the $z$-domain defined by
$$
H_{pk}(z) = C_\text{pk}\frac{1 + b_1 z^{-1} + b_2 z^{-2}}{1 + a_1 z^{-1} + a_2 z^{-2}}
$$

where the filter coefficients are given by
$$\begin{align}
C_\text{pk} &= \frac{1+k_q\mu}{1+k_q}\\
   k_q &= \frac{4}{1+\mu} \tan\left(\frac{2\pi f_c/f_s}{2Q}\right) \\
   b_1 &= \frac{-2\cos(2\pi f_c/f_s)}{1+k_q\mu} \\
   b_2 &= \frac{1-k_q\mu}{1+k_q\mu} \\
   a_1 &= \frac{-2\cos(2\pi f_c/f_s)}{1+k_q} \\
   a_2 &= \frac{1 - k_q}{1+k_q}
\end{align}$$

where
$$
\mu = 10^{G_\text{dB}/20},\ \ Q \in [2, 10]
$$

and and $f_c$ is the center frequency in Hz relative to sampling rate $f_s$ in Hz, and $G_\text{dB}$ is the peaking filter gain in dB. Conveniently, the function `peaking` is available in the module `sk_dsp_comm.sigsys`.

```python
band1 = widgets.FloatSlider(description = '100 Hz',
                            continuous_update = True,       # Continuous updates
                            value = 20.0,
                            min = -20.0,
                            max = 20.0,
                            step = 1,
                            orientation = 'vertical')
band2 = widgets.FloatSlider(description = '1000 Hz',
                            continuous_update = True,       # Continuous updates
                            value = 10.0,
                            min = -20.0,
                            max = 20.0,
                            step = 1,
                            orientation = 'vertical')
band3 = widgets.FloatSlider(description = '8000 Hz',
                            continuous_update = True,       # Continuous updates
                            value = -10.0,
                            min = -20.0,
                            max = 20.0,
                            step = 1,
                            orientation = 'vertical')

Gain = widgets.FloatSlider(description = 'Gain',
                continuous_update = True,
                value = 0.2,
                min = 0.0,
                max = 2.0,
                step = 0.01,
                orientation = 'vertical')

#widgets.HBox([Gain,band1,band2,band3])
```

```python
b_b1,a_b1 = ss.peaking(band1.value,100,Q=3.5,fs=48000)
zi_b1 = signal.lfiltic(b_b1,a_b1,[0])
b_b2,a_b2 = ss.peaking(band2.value,1000,Q=3.5,fs=48000)
zi_b2 = signal.lfiltic(b_b2,a_b2,[0])
b_b3,a_b3 = ss.peaking(band3.value,8000,Q=3.5,fs=48000)
zi_b3 = signal.lfiltic(b_b3,a_b3,[0])
b_12,a_12 = ss.cascade_filters(b_b1,a_b1,b_b2,a_b2)
b_123,a_123 = ss.cascade_filters(b_12,a_12,b_b3,a_b3)
f = logspace(log10(50),log10(10000),100)
w,H_123 = signal.freqz(b_123,a_123,2*pi*f/48000)
semilogx(f,20*log10(abs(H_123)))
ylim([-20,20])
ylabel(r'Gain (dB)')
xlabel(r'Frequency (Hz)')
grid();
```

```python
# define a pass through, y = x, callback
def callback(in_data, frame_count, time_info, status):
    global DSP_IO, zi_b1,zi_b2,zi_b3, x
    global Gain, band1, band2, band3
    DSP_IO.DSP_callback_tic()
    # convert byte data to ndarray
    in_data_nda = np.frombuffer(in_data, dtype=np.int16)
    #***********************************************
    # DSP operations here
    # Here we apply a linear filter to the input
    #x = in_data_nda.astype(float32)
    x = Gain.value*20000*x_loop.get_samples(frame_count)
    # DSP code here
    b_b1,a_b1 = ss.peaking(band1.value,100,Q=3.5,fs=48000)
    z1, zi_b1 = signal.lfilter(b_b1,a_b1,x,zi=zi_b1)
    b_b2,a_b2 = ss.peaking(band2.value,1000,Q=3.5,fs=48000)
    z2, zi_b2 = signal.lfilter(b_b2,a_b2,z1,zi=zi_b2)
    b_b3,a_b3 = ss.peaking(band3.value,8000,Q=3.5,fs=48000)
    y, zi_b3 = signal.lfilter(b_b3,a_b3,z2,zi=zi_b3)
    #***********************************************
    # Save data for later analysis
    # accumulate a new frame of samples
    DSP_IO.DSP_capture_add_samples(y)
    #***********************************************
    # Convert from float back to int16
    y = y.astype(int16)
    DSP_IO.DSP_callback_toc()
    # Convert ndarray back to bytes
    #return (in_data_nda.tobytes(), pyaudio.paContinue)
    return y.tobytes(), pah.pyaudio.paContinue
```

```python
fs, x_wav2 = ss.from_wav('audio_files/Music_Test.wav')
x_wav = (x_wav2[:,0] + x_wav2[:,1])/2
x_loop = pah.loop_audio(x_wav)
DSP_IO = pah.DSPIOStream(callback,0,1,fs=44100,t_capture=0)
start_stop_buttons = DSP_IO.interactive_stream(0,1)
widgets.VBox([start_stop_buttons, widgets.HBox([Gain,band1,band2,band3])])
```

```python
Image("audio_files/three_band_widgets.png", width="55%")
```

```python
f_AD,Mag_AD = loadtxt('audio_files/ThreeBand_Peak_100_p20_1k_p10_8k_m10_fs_48k.csv',
                       delimiter=',',skiprows=6,unpack=True)
semilogx(f_AD,Mag_AD+55)
semilogx(f,20*log10(abs(H_123)))
ylabel(r'Gain (dB)')
xlabel(r'Frequency (Hz)')
legend((r'AD2 Noise Measured',r'Design Theory'))
title(r'Three Band Equalizer: $f_{center} = [100,1000,800]$, $Q = 3.5$')
ylim([-20,20])
xlim([50,10000])
grid();
```


# Zscores

In this notebook, we plot **Figures 4(a) and 4(b)**. To do so, we need to transform the ratings into zscore. Once again, we need to parse the ratings per year.

*No requirement is needed to run this notebook.*

**Benchmark time**: This notebook has been run on a Dell Latitude (ElementaryOS 0.4.1 Loki, i7-7600U, 16GB RAM).

```python
import os
os.chdir('..')
```

```python
# Helpers functions
from python.helpers import parse

# Libraries for preparing data
import json
import gzip
import numpy as np
import pandas as pd
from datetime import datetime

# Libraries for plotting
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter
import matplotlib

# Folders
data_folder = '../data/'
fig_folder = '../figures/'

# For the Python notebook
%matplotlib inline
%reload_ext autoreload
%autoreload 2

# General info for plotting
colors = {'ba': (232/255,164/255,29/255),
          'rb': (0/255,152/255,205/255)}

labels = {'ba': 'BeerAdvocate', 'rb': 'RateBeer'}

# Check that folders exist
if not os.path.exists(data_folder + 'tmp'):
    os.makedirs(data_folder + 'tmp')

if not os.path.exists(data_folder + 'prepared'):
    os.makedirs(data_folder + 'prepared')

if not os.path.exists(fig_folder):
    os.makedirs(fig_folder)
```

## Get the ratings for the matched beers

In the next cell, we compute the ratings for the matched beers only. This is done in order to speed up the process a little bit.

```python
%%time

# Load DF of matched beers
df = pd.read_csv(data_folder + 'matched/beers.csv', header=[0,1])
# Get the ids
beer_ids = {'ba': np.array(df['ba']['beer_id']), 'rb': np.array(df['rb']['beer_id'])}

# Go through BA and RB
for key in ['ba', 'rb']:
    print('Parse {} ratings'.format(key.upper()))

    # Get iterator
    gen = parse(data_folder + key + '/ratings.txt.gz'.format(key))

    # Open the new gzip file
    file_ = gzip.open(data_folder + 'matched/ratings_{}.txt.gz'.format(key), 'wb')

    # Go through iterator
    for item in gen:
        # Check if the beer corresponds to a matched beers
        if int(item['beer_id']) in beer_ids[key]:
            # If yes, add it to the new gzip file
            for key_dct in item.keys():
                file_.write('{}: {}\n'.format(key_dct, item[key_dct]).encode('utf-8'))
            file_.write('\n'.encode('utf-8'))

    file_.close()
```

    Parse BA ratings
    Parse RB ratings
    CPU times: user 19min 26s, sys: 2.61 s, total: 19min 29s
    Wall time: 19min 31s

```python
%%time

# Get the ratings for every year for both BA and RB
ratings_year = {'ba': {}, 'rb': {}}

# Go through RB and BA
for key in ratings_year.keys():
    print('Parsing {} reviews.'.format(key.upper()))

    # Get the iterator
    gen = parse(data_folder + '/matched/ratings_{}.txt.gz'.format(key))

    # Go through the iterator
    for item in gen:

        # Get the data and then the year
        date = int(item['date'])
        year = datetime.fromtimestamp(date).year

        if year not in ratings_year[key].keys():
            ratings_year[key][year] = []

        # Add the rating in the correct year
        ratings_year[key][year].append(float(item['rating']))

with open(data_folder + 'tmp/rating_per_year_matched.json', 'w') as outfile:
    json.dump(ratings_year, outfile)
```

    Parsing BA reviews.
    Parsing RB reviews.
    CPU times: user 1min 33s, sys: 524 ms, total: 1min 33s
    Wall time: 1min 34s

## Compute the zscores

First, we need to get the parameters for the zscores (mean and std). The zscore is defined as:

\begin{equation}
z = \frac{x - \bar{x}_y}{\sigma_y}
\end{equation}
where $x$ is the rating, $\bar{x}_y$ is the average rating for the year $y$, and $\sigma_y$ is the std for the year $y$. The year $y$ corresponds to the year in which the rating $x$ was given.

```python
%%time

# Get the file rating_per_year from notebook 3
with open(data_folder + 'tmp/rating_per_year_matched.json', 'r') as infile:
    ratings_year = json.load(infile)

# Prepare the dict of mean and std per year to compute the zscore
z_score_params = {}

# Go through BA and RB
for key in ratings_year.keys():
    z_score_params[key] = {}
    # Go through each year
    for y in ratings_year[key].keys():
        # Add the average and STD
        z_score_params[key][y] = {'mean': np.mean(ratings_year[key][y]), 'std': np.std(ratings_year[key][y])}
        if z_score_params[key][y]['std'] == 0:
            z_score_params[key][y]['std'] = 1

# Fill some missing years
z_score_params['ba'][1996] = {'mean': 0, 'std': 1}
z_score_params['ba'][1997] = {'mean': 0, 'std': 1}

# And save the file
with open(data_folder + 'tmp/z_score_params_matched_ratings.json', 'w') as file:
    json.dump(z_score_params, file)
```

    CPU times: user 444 ms, sys: 8 ms, total: 452 ms
    Wall time: 450 ms

```python
%%time

with open('../data/tmp/z_score_params_matched_ratings.json') as file:
    z_score_params = json.load(file)

# Prepare the dict with zscores for ba and rb
zscore = {'ba': [], 'rb': []}

# Beers zscore
beers_zscore = {}

# Go through BA and RB
for key in zscore.keys():

    beers_zscore[key] = {}

    print('Parsing {} reviews.'.format(key.upper()))

    # Get the iterator with the ratings fo the matched beers
    gen = parse(data_folder + 'matched/ratings_{}.txt.gz'.format(key))

    # Go through the iterator
    for item in gen:

        # Get the year
        date = int(item['date'])
        year = str(datetime.fromtimestamp(date).year)

        # Get the rating
        rat = float(item['rating'])

        # Compute its zscore based on the year
        zs = (rat-z_score_params[key][year]['mean'])/z_score_params[key][year]['std']
        zscore[key].append(zs)

        # get the beer id
        beer_id = int(item['beer_id'])

        # Add it to the dict if needed
        if beer_id not in beers_zscore[key].keys():
            beers_zscore[key][beer_id] = []

        # Add the zscore to the beer with corresponding id
        beers_zscore[key][beer_id].append(zs)

# Save the file
with open(data_folder + 'prepared/zscore_matched.json', 'w') as outfile:
    json.dump(zscore, outfile)
```

    Parsing BA reviews.
    Parsing RB reviews.
    CPU times: user 1min 38s, sys: 508 ms, total: 1min 38s
    Wall time: 1min 38s

```python
%%time

# Compute avg zscore for each beer
beers_avg_zscore = {}
for key in beers_zscore.keys():
    beers_avg_zscore[key] = {}
    for id_ in beers_zscore[key].keys():
        beers_avg_zscore[key][id_] = np.mean(beers_zscore[key][id_])

# Add the zscores to the matched beers
df = pd.read_csv(data_folder + 'matched/beers.csv', header=[0,1])

# Go through RB and BA
for key in ['ba', 'rb']:

    zscores = []
    nbr_ratings = []
    # Go through all matched beers
    for i in df.index:
        row = df.iloc[i][key]

        # Get the id
        beer_id = int(row['beer_id'])

        # append the avg zscore
        if beer_id not in beers_avg_zscore[key].keys():
            zscores.append(np.nan)
            nbr_ratings.append(0)
        else:
            zscores.append(beers_avg_zscore[key][beer_id])
            nbr_ratings.append(len(beers_zscore[key][beer_id]))

    # Add the column
    df.loc[:, (key, 'zscore')] = zscores
    df.loc[:, (key, 'nbr_ratings')] = nbr_ratings

# Save again
df.to_csv(data_folder + 'matched/beers.csv', index=False)
```

    CPU times: user 41.9 s, sys: 148 ms, total: 42.1 s
    Wall time: 42.1 s

## Plot the zscores for the ratings of the matched beers

The next cell plots Figure 4(a).

```python
with open(data_folder + 'prepared/zscore_matched.json', 'r') as infile:
    zscore = json.load(infile)

years = list(range(2001, 2018))
x = np.array(range(len(years)))

plt.figure(figsize=(5, 3), frameon=False)

sns.set_context("paper")
sns.set(font_scale = 1.1)
sns.set_style("white", {
    "font.family": "sans-serif",
    "font.serif": ['Helvetica'],
    "font.scale": 2
})
sns.set_style("ticks", {"xtick.major.size": 4,
                        "ytick.major.size": 4})

ax = plt.subplot(111)
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)

nbr_bins = 45

bin_edges = {'ba': None, 'rb': None}
hist = {'ba': None, 'rb': None}

for key in bin_edges.keys():
    a, bin_edges[key] = np.histogram(zscore[key], range=(-7, 3), bins=nbr_bins)
    hist[key], a = np.histogram(zscore[key], range=(-7, 3), bins=nbr_bins+1)
    hist[key] = hist[key]/np.sum(hist[key])*100

for key in hist.keys():
    width = bin_edges[key][1] - bin_edges[key][0]
    plt.bar(bin_edges[key], hist[key], width = width,
            facecolor=colors[key], edgecolor='grey', alpha=0.5,
            label=labels[key])

plt.ylabel('Percentage')
plt.xlabel('Rating (standardized)')

plt.xlim([-6.1, 3.5])

leg = plt.legend()
leg.get_frame().set_linewidth(0.0)

plt.savefig(fig_folder + 'zscore_matched_beers.pdf', bbox_inches='tight')
```

# Hexagon plot

The next cell plots Figure 4(b).

```python
# Load files
df = pd.read_csv(data_folder + 'matched/beers.csv', header=[0,1])

# params
min_nbr_ratings = 5
df = df[(df['ba']['nbr_ratings'] >= min_nbr_ratings) & (df['rb']['nbr_ratings'] >= min_nbr_ratings)]
df.index = range(len(df))

idx = df[df['ba']['beer_id'] == 78599].index[0]

ticks = [0, 1, 2, 3, 5, 10, 20, 30, 50, 100, 176]
log_ticks = []
for t in ticks:
    log_ticks.append(np.log10(t+1))

key = 'zscore'

plt.figure(figsize=(4, 3), frameon=False)

sns.set_context("paper")
sns.set(font_scale = 1.1)
sns.set_style("white", {
    "font.family": "sans-serif",
    "font.serif": ['Helvetica'],
    "font.scale": 2
})
sns.set_style("ticks", {"xtick.major.size": 4,
                        "ytick.major.size": 4})

ax = plt.subplot(111)
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)

hb = ax.hexbin(df['rb'][key], df['ba'][key], gridsize=50, bins='log', cmap='Greys')
ax.plot([-5, 2], [-5, 2], 'red', linewidth=0.8)

ax.axis([-5, 2, -5, 2])
plt.xticks(list(range(-5, 3)), list(range(-5, 3)))

plt.xlabel('RateBeer')
plt.ylabel('BeerAdvocate')

cb = plt.colorbar(hb, ax=ax, ticks=log_ticks)
cb.ax.set_yticklabels(ticks)
ax.set_facecolor('white')
row = df.loc[idx]
name = row['rb']['beer_name']
x = row['rb'][key]
y = row['ba'][key]

plt.plot([x], [y], 'x', color='r', mew=2, ms=6)

plt.arrow(-0.5, -4, x+0.35, 2.6, linewidth=1, head_width=0.2, head_length=0.5, fc='k', ec='k')

plt.text(-0.5, -4.7, name, horizontalalignment='center', verticalalignment='bottom')

plt.savefig(fig_folder + 'hexhist_zscores_example.pdf', bbox_inches='tight')
```


```
# This notebook illustrates how to use numpyro
# https://github.com/pyro-ppl/numpyro

# Speed comparison with TFP
# https://rlouf.github.io/post/jax-random-walk-metropolis/
# Speed comparison with pymc3
# https://www.kaggle.com/s903124/numpyro-speed-benchmark

```

# Installation

```
# Standard Python libraries
from __future__ import absolute_import, division, print_function, unicode_literals

import os
import time
#import numpy as np
#np.set_printoptions(precision=3)
import glob
import matplotlib.pyplot as plt
import PIL
import imageio

from IPython import display
%matplotlib inline

import sklearn

import seaborn as sns;
sns.set(style="ticks", color_codes=True)

import pandas as pd
pd.set_option('precision', 2) # 2 decimal places
pd.set_option('display.max_rows', 20)
pd.set_option('display.max_columns', 30)
pd.set_option('display.width', 100) # wide windows
```

```
import jax
import jax.numpy as np
import numpy as onp # original numpy

print("jax version {}".format(jax.__version__))
print("jax backend {}".format(jax.lib.xla_bridge.get_backend().platform))
```

    jax version 0.2.7
    jax backend gpu

```
# https://github.com/pyro-ppl/numpyro
!pip install numpyro

# It seems that numpyro installs jaxlib for CPU
#https://github.com/pyro-ppl/numpyro/issues/531
```

    Collecting numpyro
      Downloading https://files.pythonhosted.org/packages/35/f1/7bada66676245f9e085b870b1051ba183b377af287002e10a2e1bea1b498/numpyro-0.4.1-py3-none-any.whl (176kB)
         |████████████████████████████████| 184kB 8.6MB/s
    Requirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from numpyro) (4.41.1)
    Collecting jax==0.2.3
      Downloading https://files.pythonhosted.org/packages/d7/b2/738298445cb0d9445e84f58f1fdaf73aa7b1d4199e6360620461d6fe3a8b/jax-0.2.3.tar.gz (473kB)
         |████████████████████████████████| 481kB 12.4MB/s
    Collecting jaxlib==0.1.56
      Downloading https://files.pythonhosted.org/packages/aa/44/16d06ee6418ae1b020b0722f7b7465baa08031a85728392e5413dd4e3e04/jaxlib-0.1.56-cp36-none-manylinux2010_x86_64.whl (32.1MB)
         |████████████████████████████████| 32.1MB 111kB/s
    Requirement already satisfied: numpy>=1.12 in /usr/local/lib/python3.6/dist-packages (from jax==0.2.3->numpyro) (1.19.4)
    Requirement already satisfied: absl-py in /usr/local/lib/python3.6/dist-packages (from jax==0.2.3->numpyro) (0.10.0)
    Requirement already satisfied: opt_einsum in /usr/local/lib/python3.6/dist-packages (from jax==0.2.3->numpyro) (3.3.0)
    Requirement already satisfied: scipy in /usr/local/lib/python3.6/dist-packages (from jaxlib==0.1.56->numpyro) (1.4.1)
    Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from absl-py->jax==0.2.3->numpyro) (1.15.0)
    Building wheels for collected packages: jax
      Building wheel for jax (setup.py) ... done
      Created wheel for jax: filename=jax-0.2.3-cp36-none-any.whl size=542178 sha256=f258c0d1f96711cc0b308e64517b4d916ae57c44003c7a217fc8b6cf71fdccd8
      Stored in directory: /root/.cache/pip/wheels/12/30/5d/24b5503a9bbf06bdd0d57bd20a87ef56125581b862731e4a2d
    Successfully built jax
    Installing collected packages: jax, jaxlib, numpyro
      Found existing installation: jax 0.2.7
        Uninstalling jax-0.2.7:
          Successfully uninstalled jax-0.2.7
      Found existing installation: jaxlib 0.1.57+cuda101
        Uninstalling jaxlib-0.1.57+cuda101:
          Successfully uninstalled jaxlib-0.1.57+cuda101
    Successfully installed jax-0.2.3 jaxlib-0.1.56 numpyro-0.4.1

```
import jax
import jax.numpy as np
import numpy as onp # original numpy
from jax import random

print("jax version {}".format(jax.__version__))
print("jax backend {}".format(jax.lib.xla_bridge.get_backend().platform))
```

    jax version 0.2.7
    jax backend gpu

# Distributions

```
import numpyro
import numpyro.distributions as dist
from numpyro.diagnostics import hpdi
from numpyro.distributions.transforms import AffineTransform
from numpyro.infer import MCMC, NUTS, Predictive

rng_key = random.PRNGKey(0)
rng_key, rng_key_ = random.split(rng_key)
```

## 1d Gaussian

```
# 2 independent 1d gaussians (ie 1 diagonal Gaussian)
mu = 1.5
sigma = 2
d = dist.Normal(mu, sigma)
dir(d)
```

    ['__call__',
     '__class__',
     '__delattr__',
     '__dict__',
     '__dir__',
     '__doc__',
     '__eq__',
     '__format__',
     '__ge__',
     '__getattribute__',
     '__gt__',
     '__hash__',
     '__init__',
     '__init_subclass__',
     '__le__',
     '__lt__',
     '__module__',
     '__ne__',
     '__new__',
     '__reduce__',
     '__reduce_ex__',
     '__repr__',
     '__setattr__',
     '__sizeof__',
     '__str__',
     '__subclasshook__',
     '__weakref__',
     '_batch_shape',
     '_event_shape',
     '_validate_args',
     '_validate_sample',
     'arg_constraints',
     'batch_shape',
     'enumerate_support',
     'event_dim',
     'event_shape',
     'expand',
     'expand_by',
     'has_enumerate_support',
     'icdf',
     'is_discrete',
     'loc',
     'log_prob',
     'mask',
     'mean',
     'reparametrized_params',
     'sample',
     'sample_with_intermediates',
     'scale',
     'set_default_validate_args',
     'shape',
     'support',
     'to_event',
     'tree_flatten',
     'tree_unflatten',
     'variance']

```
rng_key, rng_key_ = random.split(rng_key)
nsamples = 1000
ys = d.sample(rng_key_, (nsamples,))
print(ys.shape)
mu_hat = np.mean(ys,0)
print(mu_hat)
sigma_hat = np.std(ys, 0)
print(sigma_hat)
```

    (1000,)
    1.5070927
    2.0493808

## Multivariate Gaussian

```
mu = np.array([-1, 1])
sigma = np.array([1, 2])
Sigma = np.diag(sigma)
d2 = dist.MultivariateNormal(mu, Sigma)
```

```
#rng_key, rng_key_ = random.split(rng_key)
nsamples = 1000
ys = d2.sample(rng_key_, (nsamples,))
print(ys.shape)
mu_hat = np.mean(ys,0)
print(mu_hat)
Sigma_hat = np.cov(ys, rowvar=False) #jax.np.cov not implemented
print(Sigma_hat)
```

    (1000, 2)
    [-1.0127413  1.0091063]
    [[ 0.9770031  -0.00533966]
     [-0.00533966  1.9718108 ]]

## Shape semantics

Numpyro, [Pyro](https://pyro.ai/examples/tensor_shapes.html) and [TFP](https://www.tensorflow.org/probability/examples/Understanding_TensorFlow_Distributions_Shapes) all distinguish between 'event shape' and 'batch shape'.
For a D-dimensional Gaussian, the event shape is (D,), and the batch shape
will be (), meaning we have a single instance of this distribution.
If the covariance is diagonal, we can view this as D independent
1d Gaussians, stored along the batch dimension; this will have event shape () but batch shape (2,).

When we sample from a distribution, we also specify the sample_shape.
Suppose we draw N samples  from a single D-dim diagonal Gaussian,
and N samples from D 1d Gaussians. These samples will have the same shape.
However, the semantics of logprob differs.
We illustrate this below.

```
d2 = dist.MultivariateNormal(mu, Sigma)
print(f'event shape {d2.event_shape}, batch shape {d2.batch_shape}')
nsamples = 3
ys2 = d2.sample(rng_key_, (nsamples,))
print('samples, shape {}'.format(ys2.shape))
print(ys2)

# 2 independent 1d gaussians (same as one 2d diagonal Gaussian)
d3 = dist.Normal(mu, np.diag(Sigma))
print(f'event shape {d3.event_shape}, batch shape {d3.batch_shape}')
ys3 = d3.sample(rng_key_, (nsamples,))
print('samples, shape {}'.format(ys3.shape))
print(ys3)

print(np.allclose(ys2, ys3))
```

    event shape (2,), batch shape ()
    samples, shape (3, 2)
    [[-0.06819373  0.9942934 ]
     [-1.740325   -1.0183868 ]
     [ 0.05969942  2.314332  ]]
    event shape (), batch shape (2,)
    samples, shape (3, 2)
    [[-0.06819373  0.99192965]
     [-1.740325   -1.85443   ]
     [ 0.05969942  2.8587465 ]]
    False

```
y = ys2[0,:] # 2 numbers
print(d2.log_prob(y)) # log prob of a single 2d distribution on 2d input
print(d3.log_prob(y)) # log prob of two 1d distributions on 2d input

```

    -2.6185904
    [-1.35307   -1.6120898]

We can turn a set of independent distributions into a single product
distribution using the [Independent class](http://num.pyro.ai/en/stable/distributions.html#independent)

```
d4 = dist.Independent(d3, 1) # treat the first batch dimension as an event dimensions
print(d4.event_shape)
print(d4.batch_shape)
print(d4.log_prob(y))
```

    (2,)
    ()
    -2.96516

# Posterior inference with MCMC

## Example: 1d Gaussian with unknown mean.

We use the simple example from the [Pyro intro](https://pyro.ai/examples/intro_part_ii.html#A-Simple-Example). The goal is to infer the weight $\theta$ of an object, given noisy measurements $y$. We assume the following model:
$$
\begin{align}
\theta &\sim N(\mu=8.5, \tau^2=1.0)\\
y \sim &N(\theta, \sigma^2=0.75^2)
\end{align}
$$

Where $\mu=8.5$ is the initial guess.

By Bayes rule for Gaussians, we know that the exact posterior,
given a single observation $y=9.5$, is given by

$$
\begin{align}
\theta|y &\sim N(m, s^s) \\
m &=\frac{\sigma^2 \mu + \tau^2 y}{\sigma^2 + \tau^2}
  = \frac{0.75^2 \times 8.5 + 1 \times 9.5}{0.75^2 + 1^2}
  = 9.14 \\
s^2 &= \frac{\sigma^2 \tau^2}{\sigma^2  + \tau^2}
= \frac{0.75^2 \times 1^2}{0.75^2 + 1^2}= 0.6^2
\end{align}
$$

```
mu = 8.5; tau = 1.0; sigma = 0.75; y = 9.5
m = (sigma**2 * mu + tau**2 * y)/(sigma**2 + tau**2)
s2 = (sigma**2 * tau**2)/(sigma**2 + tau**2)
s = np.sqrt(s2)
print(m)
print(s)
```

    9.14
    0.6

```
def model(prior_mean, prior_sd, obs_sd, measurement=None):
    theta = numpyro.sample("theta", dist.Normal(prior_mean, prior_sd))
    return numpyro.sample("y", dist.Normal(theta, obs_sd), obs=measurement)

```

```
nuts_kernel = NUTS(model)
mcmc = MCMC(nuts_kernel, num_warmup=100, num_samples=1000)
mcmc.run(rng_key_, mu, tau, sigma, y)

mcmc.print_summary()
samples  = mcmc.get_samples()

```

    sample: 100%|██████████| 1100/1100 [00:03<00:00, 286.64it/s, 3 steps of size 9.41e-01. acc. prob=0.91]

                    mean       std    median      5.0%     95.0%     n_eff     r_hat
         theta      9.17      0.60      9.13      8.27     10.15    365.16      1.00

    Number of divergences: 0

```

```