# Fitting experimental data to thermodynamic models

Besides calculating and visualizing phase equilibria using known models, **pycalphad** also enables you to fit experimental data to a parameterized thermodynamic model. While pycalphad handles construction of the energy functions and equilibrium calculation, a separate module handles the numerical details of fitting disparate sources of data to the thermodynamic functions.

Here we consider the case study of the Al-Zn system. First, we use a TDB file as the setup for our model parameterizations. Note that not all of the parameters have not been filled in. These are what we will fit.

```python
%matplotlib inline

TDB_STR = """
$ Al-Zn
 ELEMENT /-   ELECTRON_GAS              0.0000E+00  0.0000E+00  0.0000E+00!
 ELEMENT VA   VACUUM                    0.0000E+00  0.0000E+00  0.0000E+00!
 ELEMENT AL   FCC_A1                    2.6982E+01  4.5773E+03  2.8322E+01!
 ELEMENT ZN   HCP_ZN                    6.5390E+01  5.6568E+03  4.1631E+01!
 FUNCTION  GHSERAL  298.0
    -7976.15+137.0715*T-24.36720*T*LN(T)-1.884662E-3*T**2-0.877664E-6*T**3
    +74092*T**(-1);                                                   700.00 Y
    -11276.24+223.0269*T-38.58443*T*LN(T)+18.531982E-3*T**2-5.764227E-6*T**3
    +74092*T**(-1);                                                    933.6 Y
    -11277.68+188.6620*T-31.74819*T*LN(T)-1234.26E25*T**(-9);        2900.00 N !
 FUNCTION  GALLIQ   298.0
    +3029.403+125.2307*T-24.36720*T*LN(T)-1.884662E-3*T**2-0.877664E-6*T**3
    +74092*T**(-1)+79.401E-21*T**7;                                   700.00 Y
    -270.6860+211.1861*T-38.58443*T*LN(T)+18.53198E-3*T**2-5.764227E-6*T**3
    +74092*T**(-1)+79.401E-21*T**7;                                    933.6 Y
    -795.7090+177.4100*T-31.74819*T*LN(T);                           2900.00 N !
 FUNCTION  GALHCP   298.0  +5481-1.8*T+GHSERAL#;                      6000  N !
 FUNCTION GHSERZN    298.0  -7285.787+118.4693*T-23.70131*T*LN(T)
     -.001712034*T**2-1.264963E-06*T**3;                              692.7 Y
     -11070.60+172.3449*T-31.38*T*LN(T)+4.70657E+26*T**(-9);           1700 N !
 FUNCTION GZNLIQ     298.0  -1.285170+108.1769*T-23.70131*T*LN(T)
     -.001712034*T**2-1.264963E-06*T**3-3.585652E-19*T**7;            692.7 Y
     -11070.60+172.3449*T-31.38*T*LN(T)+4.70657E+26*T**(-9);           1700 N !
 FUNCTION GZNLIQ     298.14  +7157.213-10.29299*T-3.5896E-19*T**7+GHSERZN#;
                                                                      692.7 Y
     +7450.168-10.737066*T-4.7051E+26*T**(-9)+GHSERZN#;                 1700 N !
 FUNCTION GZNFCC     298.15  +2969.82-1.56968*T+GHSERZN#;               1700 N !
$-------------------------------------------------------------------------------
 TYPE_DEFINITION % SEQ *!
 PHASE LIQUID %  1  1.0  !
 CONSTITUENT LIQUID :AL,ZN :  !
   PARAMETER G(LIQUID,AL;0)      298.15  +GALLIQ#;                      2900 N !
   PARAMETER G(LIQUID,ZN;0)      298.15  +GZNLIQ#;                      1700 N !
   PARAMETER G(LIQUID,AL,ZN;0)   298.15  LIQALZN0H-LIQALZN0S*T;         6000 N !
 PHASE FCC_A1  %  1  1.0  !
 CONSTITUENT FCC_A1  :AL,ZN :  !
   PARAMETER G(FCC_A1,AL;0)      298.15  +GHSERAL#;                     2900 N !
   PARAMETER G(FCC_A1,ZN;0)      298.15  +GZNFCC#;                      1700 N !
   PARAMETER G(FCC_A1,AL,ZN;0)   298.15  FCCALZN0H-FCCALZN0S*T;         6000 N !
   PARAMETER G(FCC_A1,AL,ZN;1)   298.15  FCCALZN1H-FCCALZN1S*T;         6000 N !
   PARAMETER G(FCC_A1,AL,ZN;2)   298.15  FCCALZN2H-FCCALZN2S*T;         6000 N !
 PHASE HCP_A3  %  1  1.0  !
 CONSTITUENT HCP_A3  :AL,ZN :  !
  PARAMETER G(HCP_A3,AL;0)       298.15  +GALHCP#;                      2900 N !
  PARAMETER G(HCP_A3,ZN;0)       298.15  +GHSERZN#;                     1700 N !
  PARAMETER G(HCP_A3,AL,ZN;0)    298.15  HCPALZN0H-HCPALZN0S*T;         6000 N !
  PARAMETER G(HCP_A3,AL,ZN;3)    298.15  HCPALZN3H-HCPALZN3S*T;         6000 N !"""
```

Now we will import pycalphad and load the database.

```python
from pycalphad import Database, Model, binplot
import pycalphad.variables as v
import matplotlib.pyplot as plt
from sympy import Symbol

# Initialize database and choose active phases
db_alzn = Database(TDB_STR)
my_phases_alzn = ['LIQUID', 'FCC_A1', 'HCP_A3']
```

## Prior to fitting

First, we'll show what the phase diagram looks like without any binary interaction parameters. In this diagram only the pure component lattice stabilities of each phase relative to their standard states are considered.

```python
blank_parameters = {
    'LIQALZN0H': 0.0,
    'LIQALZN0S': 0.0,
    'FCCALZN0H': 0.0,
    'FCCALZN0S': 0.0,
    'FCCALZN1H': 0.0,
    'FCCALZN1S': 0.0,
    'FCCALZN2H': 0.0,
    'FCCALZN2S': 0.0,
    'HCPALZN0H': 0.0,
    'HCPALZN0S': 0.0,
    'HCPALZN3H': 0.0,
    'HCPALZN3S': 0.0
    }

models = {
    'LIQUID': Model(db_alzn, ['AL', 'ZN'], 'LIQUID', parameters=blank_parameters),
    'FCC_A1': Model(db_alzn, ['AL', 'ZN', 'VA'], 'FCC_A1', parameters=blank_parameters),
    'HCP_A3': Model(db_alzn, ['AL', 'ZN', 'VA'], 'HCP_A3', parameters=blank_parameters)
    }

fig = plt.figure(figsize=(9,6))
%time ax = binplot(db_alzn, ['AL', 'ZN', 'VA'] , my_phases_alzn, 'X(ZN)', 300, 1000, ax=plt.gca(), model=models)
ax = ax.set_title('Al-Zn with endmembers only', fontsize=20)
```

Now load a test Comma-Separated Values (CSV) file with some phase equilibria data. This data could be from experiment, *ab-initio* methods or another source. Let's plot the data on top of our phase diagram to get an idea of what we're dealing with here.

```python
import pandas as pd
from pycalphad.plot.utils import phase_legend

phase_eq = pd.read_csv('alzn_test_phase_eq.csv')

fig = plt.figure(figsize=(9,6))
ax = binplot(db_alzn, ['AL', 'ZN', 'VA'] , my_phases_alzn, 'X(ZN)', 300, 1000, ax=plt.gca(), model=models)
ax.set_title('Al-Zn with experimental data', fontsize=20)
legend_handles, colors = phase_legend(my_phases_alzn)
for phase, phase_df in phase_eq.groupby('Phase'):
    ax = phase_df.plot(kind='scatter', x='X(ZN)', y='T', ax=ax, edgecolor='black',
                       s=30, xlim=(0,1), marker='s', facecolor='none')
ax = ax.legend(handles=legend_handles, loc='center left', bbox_to_anchor=(1, 0.6))
```

## Starting parameter values
Let's imagine we have some way (first-principles calculations, intuition or both) to guess some starting values for our parameters.

```python
perturbed_parameters = {
    'LIQALZN0H': 10000,
    'LIQALZN0S': 3,
    'FCCALZN0H': 7000,
    'FCCALZN0S': -0.2,
    'FCCALZN1H': 5000,
    'FCCALZN1S': 4,
    'FCCALZN2H': -3000,
    'FCCALZN2S': -3,
    'HCPALZN0H': 17000,
    'HCPALZN0S': 6,
    'HCPALZN3H': -700,
    'HCPALZN3S': 1
    }
fit_models = {
    'LIQUID': Model(db_alzn, ['AL', 'ZN'], 'LIQUID', parameters=perturbed_parameters),
    'FCC_A1': Model(db_alzn, ['AL', 'ZN', 'VA'], 'FCC_A1', parameters=perturbed_parameters),
    'HCP_A3': Model(db_alzn, ['AL', 'ZN', 'VA'], 'HCP_A3', parameters=perturbed_parameters)
    }

fig = plt.figure(figsize=(9,6))
ax = binplot(db_alzn, ['AL', 'ZN', 'VA'] , my_phases_alzn, 'X(ZN)', 300, 1000, ax=plt.gca(), model=fit_models)
ax = phase_eq.plot(kind='scatter', x='X(ZN)', y='T', ax=ax, edgecolor='black',
                       s=30, xlim=(0,1), marker='s', facecolor='none')
ax = ax.set_title('Al-Zn with parameter guesses', fontsize=20)
```

## Constructing a residual and performing a fit

We construct a residual function based on distance from the equilibrium hyperplane (driving force).

```python
from pycalphad.residuals import fit_model
import numpy as np
import pandas as pd

perturbed_parameters = {
    'LIQALZN0H': 10000,
    'LIQALZN0S': 3,
    'FCCALZN0H': 7000,
    'FCCALZN0S': -0.2,
    'FCCALZN1H': 5000,
    'FCCALZN1S': 4,
    'FCCALZN2H': -3000,
    'FCCALZN2S': -3,
    'HCPALZN0H': 17000,
    'HCPALZN0S': 6,
    'HCPALZN3H': -700,
    'HCPALZN3S': 1
    }

param_names, param_values = list(zip(*perturbed_parameters.items()))
phase_eq = pd.read_csv('alzn_test_phase_eq.csv')
param_values = np.array(param_values)

%time result, mi = fit_model(perturbed_parameters, phase_eq, db_alzn, ['AL', 'ZN', 'VA'], ['LIQUID', 'HCP_A3', 'FCC_A1'])

fit_models = {
    'LIQUID': Model(db_alzn, ['AL', 'ZN'], 'LIQUID', parameters=result),
    'FCC_A1': Model(db_alzn, ['AL', 'ZN', 'VA'], 'FCC_A1', parameters=result),
    'HCP_A3': Model(db_alzn, ['AL', 'ZN', 'VA'], 'HCP_A3', parameters=result)
    }

fig = plt.figure(figsize=(9,6))
ax = binplot(db_alzn, ['AL', 'ZN', 'VA'] , my_phases_alzn, 'X(ZN)', 300, 1000, ax=plt.gca(), model=fit_models)
ax = phase_eq.plot(kind='scatter', x='X(ZN)', y='T', ax=ax, edgecolor='black',
                       s=30, xlim=(0,1), marker='s', facecolor='none')
ax = ax.set_title('Al-Zn leastsq fit', fontsize=20)
```

## Getting the fitted parameters
The **result** object contains the fitted values for the parameters.

```python
import lmfit
lmfit.report_fit(mi.params)
```

    [[Variables]]
        FCCALZN2S:  -4.05308569 +/- 37.54771 (926.40%) (init=-3)
        FCCALZN0H:   7282.83341 +/- 1.24e+04 (169.87%) (init= 7000)
        FCCALZN1H:   5525.20712 +/- 1.38e+04 (249.16%) (init= 5000)
        LIQALZN0S:   2.80094034 +/- 53.31068 (1903.31%) (init= 3)
        HCPALZN3S:   3.53758565 +/- 3.05e+03 (86078.63%) (init= 1)
        HCPALZN0S:   4.81686801 +/- 2.80e+03 (58200.67%) (init= 6)
        LIQALZN0H:   10213.0904 +/- 3.89e+04 (380.72%) (init= 10000)
        HCPALZN0H:   16196.0727 +/- 1.53e+06 (9433.03%) (init= 17000)
        FCCALZN2H:  -2497.91920 +/- 2.20e+04 (879.98%) (init=-3000)
        FCCALZN1S:   3.09055428 +/- 24.47553 (791.95%) (init= 4)
        HCPALZN3H:  -454.539117 +/- 1.65e+06 (363112.64%) (init=-700)
        FCCALZN0S:  -0.63905647 +/- 22.31785 (3492.31%) (init=-0.2)
    [[Correlations]] (unreported correlations are <  0.100)
        C(HCPALZN0H, HCPALZN3H)      =  0.998
        C(HCPALZN3S, HCPALZN0S)      =  0.998
        C(FCCALZN0H, FCCALZN0S)      =  0.998
        C(LIQALZN0S, LIQALZN0H)      =  0.997
        C(FCCALZN1H, FCCALZN1S)      =  0.995
        C(HCPALZN0S, HCPALZN3H)      =  0.995
        C(HCPALZN0S, HCPALZN0H)      =  0.993
        C(FCCALZN2S, FCCALZN2H)      =  0.993
        C(HCPALZN3S, HCPALZN3H)      =  0.993
        C(HCPALZN3S, HCPALZN0H)      =  0.988
        C(FCCALZN0H, FCCALZN2H)      = -0.767
        C(FCCALZN2H, FCCALZN0S)      = -0.763
        C(LIQALZN0S, FCCALZN1S)      =  0.755
        C(FCCALZN2S, FCCALZN0H)      = -0.753
        C(FCCALZN2S, FCCALZN0S)      = -0.744
        C(LIQALZN0H, FCCALZN1S)      =  0.731
        C(FCCALZN1H, LIQALZN0S)      =  0.706
        C(FCCALZN1H, LIQALZN0H)      =  0.682
        C(FCCALZN1H, FCCALZN2H)      = -0.437
        C(FCCALZN2H, FCCALZN1S)      = -0.388
        C(FCCALZN2S, FCCALZN1H)      = -0.366
        C(FCCALZN2S, FCCALZN1S)      = -0.312
        C(FCCALZN1H, FCCALZN0S)      =  0.278
        C(FCCALZN2S, LIQALZN0H)      =  0.265
        C(FCCALZN0H, FCCALZN1H)      =  0.259
        C(FCCALZN1S, FCCALZN0S)      =  0.249
        C(FCCALZN0H, FCCALZN1S)      =  0.228
        C(FCCALZN2S, LIQALZN0S)      =  0.226
        C(LIQALZN0H, FCCALZN2H)      =  0.185
        C(FCCALZN0H, LIQALZN0H)      = -0.177
        C(LIQALZN0H, FCCALZN0S)      = -0.153
        C(LIQALZN0S, FCCALZN2H)      =  0.143
        C(FCCALZN0H, LIQALZN0S)      = -0.109
        C(FCCALZN1H, HCPALZN3S)      = -0.101

## Comparing with previous work

Next, we show values from Mey et al.'s 1993 assessment of the Al-Zn system. We have not done any fitting here; we are just showing how arbitrary values for the parameters may be specified and easily used to plot the phase diagram. Note the optional *parameters* keyword argument specified for each call to **Model**. This feature is what makes it easy to perform arbitrary variable substitutions in our thermodynamic models.

```python
# Parameters from Mey (1993)
assessed_parameters = {
    'LIQALZN0H': 10465.3,
    'LIQALZN0S': 3.39259,
    'FCCALZN0H': 7297.5,
    'FCCALZN0S': -0.47512,
    'FCCALZN1H': 6612.9,
    'FCCALZN1S': 4.5911,
    'FCCALZN2H': -3097.2,
    'FCCALZN2S': -3.306635,
    'HCPALZN0H': 18821.0,
    'HCPALZN0S': 8.95255,
    'HCPALZN3H': -702.8,
    'HCPALZN3S': 0.0
    }

# Initialize phase models
assessed_models = {
    'LIQUID': Model(db_alzn, ['AL', 'ZN'], 'LIQUID', parameters=assessed_parameters),
    'FCC_A1': Model(db_alzn, ['AL', 'ZN', 'VA'], 'FCC_A1', parameters=assessed_parameters),
    'HCP_A3': Model(db_alzn, ['AL', 'ZN', 'VA'], 'HCP_A3', parameters=assessed_parameters)
    }

fig = plt.figure(figsize=(9,6))
%time ax = binplot(db_alzn, ['AL', 'ZN', 'VA'] , my_phases_alzn, 'X(ZN)', 300, 1000, ax=plt.gca(), model=assessed_models)
ax = phase_eq.plot(kind='scatter', x='X(ZN)', y='T', ax=ax, edgecolor='black',
                       s=30, xlim=(0,1), marker='s', facecolor='none')
ax = ax.set_title('Al-Zn assessed by Mey (1993)', fontsize=20)
```

## Quantitative comparison of assessment values with fitted values

The difference can be partially understood by the fact that Mey's assessment also relies on thermochemical data while here we only fit to points on the phase diagram.

```python
for param_name, param_value in sorted(result.items()):
    percent_change = abs(param_value - assessed_parameters[param_name]) / abs(assessed_parameters[param_name])
    print('{0}: {1:.2%} difference'.format(param_name, percent_change))
```

    FCCALZN0H: 0.20% difference
    FCCALZN0S: 34.50% difference
    FCCALZN1H: 16.45% difference
    FCCALZN1S: 32.68% difference
    FCCALZN2H: 19.35% difference
    FCCALZN2S: 22.57% difference
    HCPALZN0H: 13.95% difference
    HCPALZN0S: 46.20% difference
    HCPALZN3H: 35.32% difference
    HCPALZN3S: inf% difference
    LIQALZN0H: 2.41% difference
    LIQALZN0S: 17.44% difference

```python

```


```python
%matplotlib inline
%load_ext autoreload
%autoreload 2

import time
import itertools
import h5py

import numpy as np
import multiprocessing as mp
from scipy.stats import norm
from scipy.stats import expon

import matplotlib.pyplot as plt
import matplotlib.cm as cm

import seaborn as sns
sns.set(style="ticks", color_codes=True, font_scale=1.5)
sns.set_style({"xtick.direction": "in", "ytick.direction": "in"})
```

```python
import cossio
import kinetics
```

```python
def pmf2d(xq, qk):
    H, x_edges, y_edges = np.histogram2d(data[:,1],data[:,2], \
            bins=[np.linspace(-12,12,100), np.linspace(-12,12,100)])

    fig, ax = plt.subplots(figsize=(6,5))
    pmf = -np.log(H.transpose())
    pmf -= np.min(pmf)
    #pmf = np.max(pmf) - pmf
    cs = ax.contourf(pmf, extent=[x_edges.min(), x_edges.max(), \
                     y_edges.min(), y_edges.max()], \
                     cmap=cm.rainbow, levels=np.arange(0, 8,0.5))
    #pmf = -np.log(counts.transpose())

    #cs = ax.contourf(counts.transpose()) #, \
            #extent=[xbins.min(), xbins.max(), ybins.min(), ybins.max()], \
            #cmap=cm.rainbow, levels=np.arange(0,10 ,1))
    cbar = plt.colorbar(cs)
    ax.set_xlim(-12,12)
    ax.set_ylim(-12,12)
    ax.set_xlabel('$x$', fontsize=20)
    ax.set_ylabel('$q$', fontsize=20)
    plt.tight_layout()
```

```python
def calc_rates(y):
    lifeA, lifeB = kinetics.calc_life([y])
    meanA = 1./np.exp(np.mean(np.log([x for x in lifeA if x>0])))
    meanB = 1./np.exp(np.mean(np.log([x for x in lifeB if x>0])))
    errorA = meanA/np.sqrt(len(lifeA))
    errorB = meanA/np.sqrt(len(lifeB))
    return np.mean([meanA, meanB]), np.mean([errorA, errorB])
```

```python
def cossio_runner(inp):
    np.random.seed()
    kl = inp[0]
    sc = inp[1]
    numsteps = inp[2]

    Dq = sc*Dx
    x, q = [5., 5.]

    tt, xk, qk = cossio.run_brownian(x0=x, barrier=5., kl=kl, \
                            Dx=Dx, Dq=Dq, numsteps=numsteps, \
                                     fwrite=int(1./dt))
    data = np.column_stack((tt,xk,qk))

    h5file = "data/cossio_kl%g_Dx%g_Dq%g.h5"%(kl, Dx, Dq)
    with h5py.File(h5file, "w") as hf:
        hf.create_dataset("data", data=data)

    return h5file
```

```python
# Globals
dt = 5e-4
Dx = 1. # Diffusion coefficient for molecular coordinate
```

```python
x = np.linspace(-10,10,1000)
fig, ax = plt.subplots(2,1, figsize=(7,5), sharex=True)
Gx = [cossio.Gx(y, barrier=5.) for y in x]
dGqxdx = [cossio.dGqxdx(0, y, barrier=5.) for y in x]
ax[0].plot(x, Gx)
ax[0].set_ylabel('$G(x)$', fontsize=20)
ax[0].set_ylim(-1.1*5,0.5*5)

ax[1].plot(x, dGqxdx)
ax[1].set_xlabel('$x$', fontsize=20)
ax[1].set_ylabel('$\partial G(x)/\partial x$', fontsize=20)
ax[1].hlines(0, -10, 10, linestyle='dashed', linewidth=0.5)

plt.tight_layout(h_pad=0)
```

```python
x = np.linspace(-8,8,1000)
fig, ax = plt.subplots(2,1, figsize=(7,5), sharex=True)
Gx = np.array([cossio.Gx(y, barrier=5.) for y in x])
ax[0].plot(x, Gx)
ax[0].plot(x[450:550], Gx[450:550])
ax[0].set_ylabel('$G(x)$', fontsize=20)
ax[0].set_ylim(-1.1*5,1)

Px = np.exp(-np.array(Gx))
ax[1].plot(x, Px)
ax[1].plot(x[450:550], Px[450:550])
ax[1].set_xlabel('$x$', fontsize=20)
ax[1].set_ylabel('$P(x)$', fontsize=20)
ax[1].hlines(0, -10, 10, linestyle='dashed', linewidth=0.5)

plt.tight_layout(h_pad=0)
```

```python
print "Barrier height: %g kBT"%-np.log(np.trapz(Px[430:570],x[430:570])/np.trapz(Px, x))
print "Barrier height: %g kBT"%-np.log(np.trapz(Px[430:570],x[430:570])/np.trapz(Px, x))
```

    Barrier height: 5.01043 kBT

```python
fig, ax = plt.subplots(3,1, figsize=(7,7), sharex=True)

x = np.linspace(-8,8,1000)
Gx = np.array([cossio.Gx(y, barrier=5.) for y in x])
ax[0].plot(x, Gx)
ax[0].set_ylim(-1.1*5,1)

Px = np.exp(-np.array(Gx))
ax12 = ax[0].twinx()
ax12.plot(x, Px, 'g')
i2 = np.argmin(np.abs(x-2.5))
i1 = np.argmin(np.abs(x+2.5))
print "x1 : ", i1, x[i1]
print "x2 :",  i2, x[i2]

pfold = np.array([np.trapz(np.exp(Gx[i:i2]),x[i:i2])/np.trapz(np.exp(Gx[i1:i2]), x[i1:i2]) \
          if x[i]>-2.5 else 1. for i in range(len(x))])

#ax.set_ylim(-1,1)

Px = np.exp(-np.array(Gx))
ax[1].plot(x, pfold, label=r"$\phi_{A}(x)$")
ax[1].plot(x, 1.-pfold, label=r"$\phi_{B}(x)$")
ax[1].set_ylim(-0.1,1.1)

ax22 = ax[1].twinx()
ax22.plot(x, pfold*(1.-pfold), 'r', label=r"$\phi_{A}(x)\phi_{B}(x)$")
ax22.set_ylim(-0.025,0.3)
ax[2].plot(x,Px*pfold*(1-pfold)/np.trapz(Px,x))

ax[2].set_xlabel('$x$', fontsize=20)
for i in range(3):
    ax[i].axvline(-2.5, -6,1, c='gray', ls='dashed', lw=0.5)
    ax[i].axvline(2.5, -6,1, c='gray', ls='dashed', lw=0.5)

ax[1].legend(loc=1)
#ax22.legend(loc=4)
ax[0].set_ylabel('$G(x)$', fontsize=20, color='b')
ax12.set_ylabel('$P(x)$', fontsize=20, color='g')
ax12.set_yticklabels([])

ax[1].set_ylabel(r'$\phi_{A/B}(x)$', fontsize=20)
ax22.set_ylabel(r'$\phi_{A}(x)\phi_{B}(x)$', fontsize=20, color='r')
ax[2].set_ylabel(r'$P(x)\phi_{A}(x)\phi_{B}(x)$', fontsize=20)
ax[2].set_yticklabels([])

plt.tight_layout(h_pad=0)
```

```python
pTP = 2.*np.trapz(Px[i1:i2]*pfold[i1:i2]*(1-pfold[i1:i2]), x[i1:i2])/np.trapz(Px,x)
print "pTP(Hummer) = %f"%pTP
tTP = np.trapz(Px[i1:i2]*pfold[i1:i2]*(1-pfold[i1:i2]), x[i1:i2])/np.trapz(Px,x)
```

    pTP(Hummer) = 0.002692

```python
h5file = "data/cossio_kl0_Dx1_Dq0.h5"

print h5file
f = h5py.File(h5file, 'r')
data = np.array(f['data'])
f.close()

fig, ax = plt.subplots(figsize=(12,2), sharex=True,sharey=True)
ax.plot(data[:,0],data[:,1],'.', markersize=1)
ax.set_ylim(-10,10)
ax.set_xlim(0,25e4)
ax.set_ylabel('x')
ax.set_xlabel('time')
plt.tight_layout(h_pad=0)

fig, ax = plt.subplots(figsize=(6,4))
hist, bin_edges = np.histogram(data[:,1], bins=np.linspace(-10,10,50), normed=True)
bin_centers = [0.5*(bin_edges[i]+bin_edges[i+1]) \
               for i in range(len(bin_edges)-1)]
ax.plot(bin_centers, -np.log(hist))
ax.set_xlim(-7,7)
ax.set_ylim(0,10)
ax.set_xlabel('x')
ax.set_ylabel('F ($k_BT$)')

tau_f, tau_u, data_f, data_u, tp_f, tp_u, recrossings = \
        kinetics.lifetimes(data, f_bound=-2.5, u_bound=2.5)

tpt_u = []
tpt_f = []
for td in tp_u:
    try:
        tpt_u.append(td[:,0][-1] - td[:,0][0])
    except IndexError:
        pass
for td in tp_f:
    try:
        tpt_f.append(td[:,0][-1] - td[:,0][0])
    except IndexError:
        pass

fig, ax = plt.subplots(2,1,figsize=(6,4))
histogram, bin_edges, patches = \
        ax[0].hist(np.log(tau_f + tau_u), \
                   bins=np.linspace(0,12,20), \
                   normed=True, cumulative=False, alpha=0.75)

param = expon.fit([tau_f + tau_u], floc=0)
tau_fold = param[1]
tau_error = param[1]/np.sqrt(len(tau_f + tau_u))

lx = np.linspace(0,12,200)
pdf = 1/param[1]*np.exp(lx)*np.exp(-1/param[1]*(np.exp(lx)))
ax[0].plot(lx, pdf, lw=3, color='b')

histogram, bin_edges, patches = \
        ax[1].hist([x for x in (tpt_f + tpt_u) ], \
                   bins=np.linspace(0,14,15), \
                   normed=True, cumulative=False, alpha=0.75)

tpt = np.mean([x for x in (tpt_f + tpt_u) if x >= 0])
tpt_error = tpt/np.sqrt(len([x/2. for x in (tpt_f + tpt_u) if x >= 0]))

ax[0].set_yticklabels([])
ax[1].set_yticklabels([])
ax[0].set_xlabel(r'ln($\tau$)')
ax[1].set_xlabel(r'$\tau_\mathrm{TP}$')
ax[0].set_ylabel(r'P(ln$\tau$)')
ax[1].set_ylabel(r'P($\tau_\mathrm{TP}$)')

plt.tight_layout(h_pad=0.5, w_pad=0)
```

```python
print "Folding time: %4.1f +/- %4.1f"%(tau_fold, tau_error)
print "Transition path time: %4.1f +/- %4.1f"%(tpt, tpt_error)
print "Barrier Estimate : %4.1f +/-%4.1f kBT"%(-np.log(tpt/tau_fold), -np.log(tpt/(tau_fold+tau_error)) + np.log(tpt/tau_fold) )
```

    Folding time: 776.4 +/- 43.3
    Transition path time:  1.6 +/-  0.1
    Barrier Estimate :  6.2 +/- 0.1 kBT

```python
pTP = (np.sum(tpt_f) + np.sum(tpt_u))/float(len(data))
print "pTP (direct) = %4.2e"%(pTP)
```

    pTP (direct) = 2.07e-03

```python
from scipy.optimize import curve_fit
def parabola(x, a, b, c):
    return a*(x-b)**2 + c

fig, ax = plt.subplots()
x = np.linspace(-10,10,1000)
Gx = [cossio.Gx(y, barrier=5.) for y in x]
ax.plot(x, Gx)
ax.set_ylabel('$G(x)$', fontsize=20)
ax.set_ylim(-1.1*5,0.5*5)
ax.set_xlabel('x')

popt_f, pcov = curve_fit(parabola, x[0:400], Gx[0:400])
print " kappa_f: %g kBT/nm**2"%popt_f[0]

popt_b, pcov = curve_fit(parabola, x[400:600], Gx[400:600])
print " kappa_b: %g kBT/nm**2"%popt_b[0]

popt_u, pcov = curve_fit(parabola, x[600:], Gx[600:])
#print popt_u
print " kappa_u: %g kBT/nm**2"%popt_u[0]

ax.plot(np.linspace(3,5,50),parabola(np.linspace(3,5,50), popt_u[0], popt_u[1], popt_u[2]),\
        'hotpink', lw=3)
ax.plot(np.linspace(-5,-3,50),parabola(np.linspace(-5,-3,50), popt_f[0], popt_f[1], popt_f[2]), \
        'lightgreen', lw=3)
ax.plot(np.linspace(-1,1,50),parabola(np.linspace(-1,1,50), popt_b[0], popt_b[1], popt_b[2]), \
        'cyan', lw=3)

```

```python
from scipy.special import erf
from sympy import *

fig, ax = plt.subplots()
histogram, bin_edges, patches = \
        ax.hist([x for x in (tpt_f + tpt_u) ], \
                   bins=np.linspace(0,14,15), \
                   normed=True, cumulative=False, alpha=0.75)

tpt = np.mean([x for x in (tpt_f + tpt_u) ])
tpt_error = tpt/np.sqrt(len([x/2. for x in (tpt_f + tpt_u) if x >= 0]))

t=np.linspace(0.1,14,100)
ax.plot(t,2*np.sqrt(5)/(1-erf(5))*np.exp(-5*mpmath.coth(popt_f[0]*t)), lw=4, ls='dashed')
print mpmath.coth(popt_f[0]*1)
ax.plot(t,2*5*1/0.625*np.exp(-1/0.625*t), lw=4, ls='dashed')
ax.set_ylim(0,1.5*np.max(histogram))
ax.set_xlim(0,8)
```

```python

```


# SQL 入門

> 基本查詢：隨堂練習參考解答

郭耀仁

```python
import sqlite3
import pandas as pd
from sqlFrameCheck import checkAnsQuery
from test_queries.test_queries_02 import extract_test_queries as etq

conn_twelection = sqlite3.connect('twelection.db')
conn_nba = sqlite3.connect('nba.db')
```

## 隨堂練習：以 `||`  運算符將 `presidential2020` 表格中 `county`、`town` 與 `village` 連結誠一個新的欄位並命名為 `combined_key`，再以 `DISTINCT` 暸解台灣獨一的選舉區有哪些

```python
ans_query = """
SELECT DISTINCT county || town || village AS combined_key
  FROM presidential2020;
"""
# 試跑看看結果
pd.read_sql(ans_query, conn_twelection)
```

      combined_key

      0
      南投縣中寮鄉中寮村

      1
      南投縣中寮鄉內城村

      2
      南投縣中寮鄉八仙村

      3
      南投縣中寮鄉和興村

      4
      南投縣中寮鄉崁頂村

      ...
      ...

      7732
      高雄市鼓山區鼓岩里

      7733
      高雄市鼓山區鼓峰里

      7734
      高雄市鼓山區龍井里

      7735
      高雄市鼓山區龍子里

      7736
      高雄市鼓山區龍水里

7737 rows × 1 columns

## 測資比對

```python
caq = checkAnsQuery(etq('0201'), ans_query, conn_twelection)
caq.run_test()
```

    測資比對正確！

## 隨堂練習：計算 NBA 球員的身體質量指數（BMI）並由大到小排序，選擇 `firstName`、`lastName` 與 `bmi` 這三個變數

```python
ans_query = """
SELECT firstName,
       lastName,
       weightKilograms / (heightMeters*heightMeters) AS bmi
  FROM players
 ORDER BY bmi DESC;
"""
# 試跑看看結果
pd.read_sql(ans_query, conn_nba)
```

      firstName
      lastName
      bmi

      0
      Zion
      Williamson
      32.853790

      1
      Jaylen
      Adams
      30.487623

      2
      Jarrell
      Brantley
      29.518950

      3
      Eric
      Paschall
      29.512295

      4
      Andre
      Drummond
      29.262204

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

      499
      Kostas
      Antetokounmpo
      20.964312

      500
      Justin
      Holiday
      20.814203

      501
      Patrick
      McCaw
      20.321279

      502
      Corey
      Brewer
      20.100858

      503
      Isaac
      Bonga
      19.801500

504 rows × 3 columns

## 測資比對

```python
caq = checkAnsQuery(etq('0202'), ans_query, conn_nba)
caq.run_test()
```

    測資比對正確！

## 隨堂練習：使用 `WHERE` 將 `presidential2020` 中臺北市的資料篩選出來

```python
ans_query = """
SELECT *
  FROM presidential2020
 WHERE county='臺北市';
"""
# 試跑看看結果
pd.read_sql(ans_query, conn_twelection)
```

      county
      town
      village
      office
      number
      candidates
      votes

      0
      臺北市
      中山區
      下埤里
      524
      1
      宋楚瑜/余湘
      48

      1
      臺北市
      中山區
      下埤里
      524
      2
      韓國瑜/張善政
      377

      2
      臺北市
      中山區
      下埤里
      524
      3
      蔡英文/賴清德
      581

      3
      臺北市
      中山區
      下埤里
      525
      1
      宋楚瑜/余湘
      53

      4
      臺北市
      中山區
      下埤里
      525
      2
      韓國瑜/張善政
      342

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

      5179
      臺北市
      萬華區
      騰雲里
      1032
      2
      韓國瑜/張善政
      528

      5180
      臺北市
      萬華區
      騰雲里
      1032
      3
      蔡英文/賴清德
      456

      5181
      臺北市
      萬華區
      騰雲里
      1033
      1
      宋楚瑜/余湘
      43

      5182
      臺北市
      萬華區
      騰雲里
      1033
      2
      韓國瑜/張善政
      502

      5183
      臺北市
      萬華區
      騰雲里
      1033
      3
      蔡英文/賴清德
      512

5184 rows × 7 columns

## 測資比對

```python
caq = checkAnsQuery(etq('0203'), ans_query, conn_twelection)
caq.run_test()
```

    測資比對正確！

## 隨堂練習：使用 `WHERE` 將 `presidential2020` 表格中的六都（臺北市、新北市、桃園市、臺中市、臺南市與高雄市）的資料篩選出來

```python
ans_query = """
SELECT *
  FROM presidential2020
 WHERE county IN ('臺北市', '新北市', '桃園市', '臺中市', '臺南市', '高雄市');
"""
# 試跑看看結果
pd.read_sql(ans_query, conn_twelection)
```

      county
      town
      village
      office
      number
      candidates
      votes

      0
      新北市
      三峽區
      三峽里
      2010
      1
      宋楚瑜/余湘
      11

      1
      新北市
      三峽區
      三峽里
      2010
      2
      韓國瑜/張善政
      83

      2
      新北市
      三峽區
      三峽里
      2010
      3
      蔡英文/賴清德
      124

      3
      新北市
      三峽區
      中埔里
      2014
      1
      宋楚瑜/余湘
      51

      4
      新北市
      三峽區
      中埔里
      2014
      2
      韓國瑜/張善政
      364

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

      32668
      高雄市
      鼓山區
      龍水里
      1268
      2
      韓國瑜/張善政
      477

      32669
      高雄市
      鼓山區
      龍水里
      1268
      3
      蔡英文/賴清德
      647

      32670
      高雄市
      鼓山區
      龍水里
      1269
      1
      宋楚瑜/余湘
      39

      32671
      高雄市
      鼓山區
      龍水里
      1269
      2
      韓國瑜/張善政
      520

      32672
      高雄市
      鼓山區
      龍水里
      1269
      3
      蔡英文/賴清德
      673

32673 rows × 7 columns

## 測資比對

```python
caq = checkAnsQuery(etq('0204'), ans_query, conn_twelection)
caq.run_test()
```

    測資比對正確！

## 隨堂練習：查詢 `careerSummaries` 生涯場均得分 `ppg` 超過 20 分的球員 ID `personId`、場均得分 `ppg`

```python
ans_query = """
SELECT personId,
       ppg
  FROM careerSummaries
 WHERE ppg > 20;
"""
# 試跑看看結果
pd.read_sql(ans_query, conn_nba)
```

      personId
      ppg

      0
      2544
      27.1

      1
      2546
      23.6

      2
      201142
      27.0

      3
      201566
      23.2

      4
      201933
      21.7

      5
      201935
      25.1

      6
      201939
      23.5

      7
      202681
      22.4

      8
      203076
      24.0

      9
      203078
      21.0

      10
      203081
      24.0

      11
      203954
      24.1

      12
      1626157
      22.7

      13
      1626164
      22.3

      14
      1628378
      22.7

      15
      1629027
      23.6

      16
      1629029
      24.4

      17
      1629627
      23.6

## 測資比對

```python
caq = checkAnsQuery(etq('0205'), ans_query, conn_nba)
caq.run_test()
```

    測資比對正確！

## 隨堂練習：從 `careerSummaries` 中找出生涯助攻失誤比（assists/turnovers）最高的前 10 個球員 ID `personId`、助攻失誤比 `ast_to_ratio`

\begin{equation}
ast\_to\_ratio = \frac{assists}{turnovers}
\end{equation}

```python
ans_query = """
SELECT personId,
       CAST(assists AS REAL) / CAST(turnovers AS REAL) AS ast_to_ratio
  FROM careerSummaries
 ORDER BY ast_to_ratio DESC
 LIMIT 10;
"""
# 試跑看看結果
pd.read_sql(ans_query, conn_nba)
```

      personId
      ast_to_ratio

      0
      1629669
      6.333333

      1
      1628420
      5.448980

      2
      1626145
      4.775862

      3
      1629162
      4.032258

      4
      1629648
      4.000000

      5
      101108
      3.935682

      6
      1627853
      3.539823

      7
      1629723
      3.500000

      8
      1627832
      3.209964

      9
      202397
      3.071053

## 測資比對

```python
caq = checkAnsQuery(etq('0206'), ans_query, conn_nba)
caq.run_test()
```

    測資比對正確！


```python
# ORIGINAL RGB FILE
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import random
import pystan

from IPython.display import display

from matplotlib import rc
rc("font", family="serif", size=11)
```

First we will read in a file that contains the overall list of all the available datasets, with corresponding star IDs, $\nu_{max}$ and $\sigma_{\nu_{max}}$ values. From this we will take a random selection for our analysis.

```python
#random.seed(121)
#output = pd.read_csv('../data/output_1000stars.csv', usecols=range(1,4))
output = pd.read_csv('../data/output_back_filesremoved.csv', usecols=['kic','numax', 'numax_err'])
output = output.loc[(np.abs(output['numax'] - 75) < 30.0)].reset_index(drop=True)
# Select a random number of stars from the file, recording their IDs
IDs = [random.choice(output['kic']) for i in range(6)]
IDs = [10802837,12003742,2283075,4760954,11358669,11521197]
# Create a list of the numax values corresponding to each ID
Numax = [(output.loc[(output['kic'] == IDs[i])]).iloc[0]['numax'] for i in range(len(IDs))]
Numax_err = [(output.loc[(output['kic'] == IDs[i])]).iloc[0]['numax_err'] for i in range(len(IDs))]
```

Using our list of selected stars, we can then pull the respective files containing the induvidual mode frequencies and errors. The different files will contain varying numbers of modes; in order to successfully parse the data through Stan, one must ensure that all the arrays are of the same length. To do this, for every dataset that contains less modes than the most complete dataset, we can produce additional modes based on the average frequency spacing $\left<\Delta\nu\right>$ for each star, ensuring to assign these points with a large uncertainty so that Stan does not use this as restrictive points during the fit (they are essentially ignored).

```python
modesID = [pd.read_csv('../data/rgbmodes/modes_'+str(IDs[i])+'.csv', usecols=['f0', 'f0_err', 'A0'])
           for i in np.arange(0,len(IDs),1)]
lenmodes = [len(modesID[i]) for i in range(len(IDs))]
maxmodes = max(lenmodes)
arr_n = np.zeros([len(IDs),maxmodes])
arr_freq = np.zeros([len(IDs),maxmodes])
arr_freqerr = np.zeros([len(IDs),maxmodes])

dnu_avgID = []
for i in np.arange(0,len(IDs),1):
    modesID[i] = modesID[i].sort_values(by=['f0'])
    modesID[i] = modesID[i].set_index(np.arange(0,len(modesID[i]),1))
    modesID[i]['dnu'] = (modesID[i]['f0'].diff(2).shift(-1))/2
    dnu_avg = (np.mean(modesID[i]['dnu']))
    dnu_avgID.append(dnu_avg)

    n_min = int(modesID[i]['f0'].min() / dnu_avg)
    n = np.arange(n_min, n_min+len(modesID[i]), 1)
    modesID[i].insert(loc=0, column='n', value=n)

    # Loop to ensure all arrays are the same length
    if lenmodes[i] < maxmodes:
        l = lenmodes[i]
        while l < maxmodes:
            newrow = {'n': int(np.max(modesID[i]['n'])+1),
                      'f0': np.max(modesID[i]['f0'])+dnu_avgID[i],
                      'f0_err': 100000}
            modesID[i] = modesID[i].append(newrow, ignore_index=True)
            l += 1

    arr_n[i,:] = modesID[i]['n']
    arr_freq[i,:] = modesID[i]['f0']
    arr_freqerr[i,:] = modesID[i]['f0_err']

    #display(modesID[i])

dnu_avgarr = np.asarray(dnu_avgID)

#Define start parameters
eps = []
for i in range(len(IDs)):
    epsilon1 = np.median((modesID[i].f0 % dnu_avg) / dnu_avg)
    eps.append(epsilon1)
epsilon = np.asarray(eps)

numax_obs = Numax
nmax = numax_obs/dnu_avgarr - epsilon
alpha = 0.015*dnu_avgarr**(-0.32)
A = 0.06*dnu_avgarr**(-0.88)
G = np.ones([len(IDs)])*3.08
tau = np.ones([len(IDs)])*200
phi = np.ones([len(IDs)])*1.71
```

The data for the different stars can be viewed on an Échelle Diagram.

```python
plt.figure(1)
for i in np.arange(0,len(IDs),1):
    plt.scatter(modesID[i]['f0'] % dnu_avgID[i], modesID[i]['f0'], label=r'ID = '+str(IDs[i]))

plt.ylabel('Frequency')
plt.xlabel(r'Mod. Freq. Spacing ('+ str(u"%.3f" % dnu_avg) +') $\mu$Hz')
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.5)
#plt.savefig('seminar/manyrgbs.pdf')

plt.figure(2)
for i in np.arange(0,len(IDs),1):
    plt.scatter(modesID[i]['n'], modesID[i]['f0'], label=r'ID = '+str(IDs[i]))
plt.xlabel('Radial Order ($n$)')
plt.ylabel('Frequency')
plt.show()
```

With all our required data, we can parse the data into a stan model constructed to perform a hierarchical Bayesian model. Hyperpriors have been introduced on a number of parameters based on values in the Vrard (2014) paper.

The model follows the form:
\begin{equation}
    \left(n + \epsilon + \frac{\alpha}{2}\left(n-n_{max}\right)^2 + \frac{\mathcal{AG}}{2\pi}\sin\left(\frac{2\pi(n-n_{max})}{\mathcal{G}} + \phi \right)\right) e^{\frac{n}{\tau}}
\end{equation}
where
\begin{equation}
    n_{max} = \frac{\nu_{max}}{\Delta\nu} - \epsilon
\end{equation}

$$ A = A_0 e^{-n_{max}/\tau} \longrightarrow A_0 = A_{n_{max}} e^{n_{max}/\tau}$$

```python
code = '''
functions {
    real glitch(real n, real dnu, real numax, real epsilon, real alpha, real A, real G, real phi, real tau){
        real nmax = numax/dnu - epsilon;
        return (n + epsilon + alpha/2 * (n-nmax)^2 +
                (A*G)/(2*pi()) * sin((2*pi()*(n-nmax))/G + phi)) * dnu * exp(-(n-nmax)/tau);
    }
}

data {
    int N;                  //number of stars
    int M;                  //number of modes
    real n[N, M];           //radial mode number
    real freq[N, M];        //l=0 frequencies
    real freq_err[N, M];
    real numax_obs[N];
    real numax_err[N];
    real dnu_guess[N];
}

parameters {
    // Normal Parameters
    real dnu[N];
    real numax[N];
    real<lower = -2.0*pi(), upper = 2.0*pi()> phi[N];
    real<lower = 0> A[N];
    real G[N];
    real<lower = 0> tau[N];

    // Hierarchical Parameters
    real eps_std[N];
    real<lower=0> eps_sig;
    real epsA;
    real epsB;

    real al_std[N];
    real<lower=0> al_sig;
    real alA;
    real alB;

    //real A_std[N];
    //real<lower=0> A_sig;
    //real AA;
    //real AB;

    //real G_std[N];
    //real<lower=0> G_sig;
    //real GA;

}

transformed parameters {
    real epsilon[N];
    real alpha[N];
    //real A[N];
    //real G[N];

    for (i in 1:N){
        epsilon[i] = eps_std[i] * eps_sig + (epsA + epsB * log(dnu[i]));
        alpha[i] = al_std[i] * al_sig + (alA * (dnu[i])^(-alB));
        //A[i] = A_std[i] * A_sig + (AA * (dnu[i])^(-AB));
        //G[i] = G_std[i] * G_sig + GA;
    }
}

model {
    real mod[M];
    for (i in 1:N){
        for (j in 1:M){
            mod[j] = glitch(n[i,j], dnu[i], numax[i], epsilon[i], alpha[i], A[i], G[i], phi[i], tau[i]);
            // mod[j] ~ normal(freq[i], freq_err[i]);
        }
        freq[i,:] ~ normal(mod, freq_err[i,:]);
        dnu[i] ~ normal(dnu_guess[i], dnu_guess[i]*0.001);
        numax[i] ~ normal(numax_obs[i], numax_err[i]);
        A[i] ~ lognormal(log(0.06*dnu[i]^(-0.88)), 0.4);
    }

    //epsilon ~ uniform(-1.0, 2.0);
    //nmax ~ normal(10, 4);
    //alpha ~ lognormal(log(0.015*dnu^(-0.32)), 0.3);

    G ~ normal(3.08, 0.65);
    //numax ~ normal(numax_obs, numax_err);
    tau ~ normal(500, 100);

    // Hierarchical Parameters
    eps_std ~ normal(0, 1);
    eps_sig ~ normal(0, 0.5);
    epsA ~ normal(0.601, 0.25);
    epsB ~ normal(0.632, 0.25);

    al_std ~ normal(0, 1);
    al_sig ~ normal(0, 0.5);
    alA ~ normal(0.015, 0.005);
    alB ~ normal(0.32, 0.08);

    //A_std ~ student_t(8, 0, 1);
    //A_sig ~ normal(0, 0.5);
    //AA ~ normal(0.06, 0.02);
    //AB ~ normal(0.88, 0.05);

    //G_std ~ normal(0, 1);
    //G_sig ~ normal(0.65, 0.5);
    //GA ~ normal(3.08, 0.1);
}
'''
sm = pystan.StanModel(model_code=code)
```

    INFO:pystan:COMPILING THE C++ CODE FOR MODEL anon_model_afc862f242f8cfcf619568f553027bec NOW.

```python
no_tau = pd.read_csv('no_tau_models.csv', names=['dnu', 'numax', 'epsilon', 'alpha', 'A', 'G', 'phi', 'tau'])

stan_data = {'N': len(IDs),
         'M': maxmodes,
         'n': arr_n,
         'freq': arr_freq,
         'freq_err': arr_freqerr,
         'numax_obs': Numax,
         'numax_err': Numax_err,
         'dnu_guess': dnu_avgarr
        }
start = {'dnu': dnu_avgarr,
         'numax': Numax,
         'eps_sig': 0.01,
         'al_sig': 0.01,
         'A_sig': 0.01,
         #'G_sig': 0.65,
         'epsilon': no_tau['epsilon'],
         'alpha': no_tau['alpha'],
         'A': no_tau['A'],
         'G': no_tau['G'],
         'phi': no_tau['phi'],
         'epsA': 0.601,
         'epsB': 0.632,
         'alA': 0.015,
         'alB': 0.32,
         'AA': 0.07,
         'AB': 0.87,
         #'GA': 3.08
         'tau': np.ones([len(IDs)])*800
    }
nchains = 4

fit = sm.sampling(data=stan_data, iter=5000, chains=nchains, init=[start for n in range(nchains)])
```

    WARNING:pystan:1 of 10000 iterations ended with a divergence (0.01%).
    WARNING:pystan:Try running with adapt_delta larger than 0.8 to remove the divergences.
    WARNING:pystan:9999 of 10000 iterations saturated the maximum tree depth of 10 (99.99%)
    WARNING:pystan:Run again with max_treedepth larger than 10 to avoid saturation

```python
print(fit)
```

    Inference for Stan model: anon_model_afc862f242f8cfcf619568f553027bec.
    4 chains, each with iter=5000; warmup=2500; thin=1;
    post-warmup draws per chain=2500, total post-warmup draws=10000.

                 mean se_mean     sd   2.5%    25%    50%    75%  97.5%  n_eff   Rhat
    dnu[1]       5.49  2.2e-4 5.5e-3   5.48   5.49   5.49    5.5    5.5    655   1.01
    dnu[2]       5.41  1.7e-4 4.8e-3    5.4   5.41   5.41   5.41   5.42    790    1.0
    dnu[3]       6.47  2.2e-4 5.2e-3   6.46   6.47   6.47   6.48   6.48    589   1.01
    dnu[4]       4.76  1.4e-4 4.4e-3   4.75   4.76   4.76   4.76   4.77    962    1.0
    dnu[5]       8.49  2.0e-4 5.4e-3   8.48   8.49   8.49   8.49    8.5    762   1.01
    dnu[6]       6.17  2.2e-4 6.0e-3   6.16   6.17   6.17   6.18   6.19    728    1.0
    numax[1]    55.85    0.01   0.34  55.18  55.62  55.84  56.07  56.53    670   1.01
    numax[2]     58.5    0.02   0.41  57.71  58.22   58.5  58.79  59.32    531   1.01
    numax[3]    65.89    0.02   0.35  65.22  65.66  65.89  66.12  66.56    431   1.01
    numax[4]    44.68    0.01   0.32  44.06  44.46  44.67  44.89  45.31    830    1.0
    numax[5]    97.43  9.9e-3   0.26  96.92  97.25  97.43  97.61  97.96    708    1.0
    numax[6]    67.26    0.01   0.33   66.6  67.04  67.26  67.48  67.92    700    1.0
    phi[1]       0.59  8.3e-3   0.21   0.13   0.46    0.6   0.73   0.97    641   1.01
    phi[2]      -0.65  9.2e-3   0.22  -1.07  -0.79  -0.65   -0.5  -0.22    550   1.01
    phi[3]       0.93  6.9e-3   0.15   0.64   0.83   0.93   1.03   1.22    454    1.0
    phi[4]      -0.43    0.03   0.43   -1.1  -0.73  -0.51  -0.22   0.66    240   1.02
    phi[5]       2.28  6.4e-3   0.18   1.92   2.17   2.29    2.4    2.6    763    1.0
    phi[6]       0.46    0.01   0.23  -0.05   0.31   0.47   0.62   0.87    472   1.01
    A[1]         0.03  2.2e-4 5.0e-3   0.02   0.02   0.03   0.03   0.04    531   1.01
    A[2]         0.09  6.2e-4   0.01   0.07   0.08   0.09    0.1   0.12    517    1.0
    A[3]         0.03  9.2e-5 2.7e-3   0.02   0.03   0.03   0.03   0.03    861    1.0
    A[4]         0.03  3.0e-4 7.9e-3   0.02   0.03   0.03   0.04   0.05    676   1.01
    A[5]         0.02  9.8e-5 2.7e-3   0.02   0.02   0.02   0.02   0.03    737    1.0
    A[6]         0.02  1.6e-4 3.3e-3   0.02   0.02   0.02   0.02   0.03    445   1.01
    G[1]         4.22    0.03   0.56   3.06   3.86   4.23   4.59   5.29    341   1.02
    G[2]         2.39  3.9e-3   0.08   2.25   2.33   2.38   2.44   2.57    449    1.0
    G[3]         2.73  2.9e-3   0.09   2.56   2.67   2.73   2.78   2.91    934    1.0
    G[4]         2.93    0.03   0.37   2.52   2.71   2.84   3.02   4.02    181   1.02
    G[5]         2.74  3.3e-3   0.09   2.58   2.68   2.74    2.8   2.95    790   1.01
    G[6]          6.4    0.02   0.43   5.53   6.12    6.4   6.69   7.25    631    1.0
    tau[1]      791.7    3.29  80.78 639.58 735.34 789.61 846.92 951.94    602   1.01
    tau[2]     915.86    2.43  66.61 789.76 870.04 913.29 960.83 1049.6    749    1.0
    tau[3]     1124.4    2.37  66.06  996.8 1079.7 1122.4 1168.8 1254.3    777    1.0
    tau[4]     749.71    2.24  67.42 622.18 703.47 748.25 794.93 883.92    909   1.01
    tau[5]     1165.1    2.26  60.83 1048.3 1123.2 1164.4 1205.5 1288.4    723    1.0
    tau[6]     848.02    2.71  73.41 705.79 798.69  847.4 897.86 992.09    736    1.0
    eps_std[1]  -0.87    0.03   0.38  -1.59  -1.13  -0.87  -0.61  -0.13    205   1.01
    eps_std[2]   -0.5    0.03   0.37  -1.18  -0.76  -0.51  -0.25   0.23    184   1.01
    eps_std[3]  -0.54    0.03    0.4  -1.27  -0.82  -0.55  -0.28   0.24    181   1.01
    eps_std[4]   1.47    0.05   0.76    0.1   0.93   1.43   1.97   3.02    274    1.0
    eps_std[5]   -0.6    0.03   0.45  -1.45  -0.91   -0.6   -0.3   0.29    188   1.01
    eps_std[6]  -0.67    0.03   0.39  -1.38  -0.94  -0.67  -0.41   0.09    188   1.01
    eps_sig      0.43  7.3e-3   0.16   0.22   0.31   0.39    0.5   0.85    488    1.0
    epsA         0.32  9.2e-3   0.23  -0.11   0.17   0.32   0.48   0.77    603    1.0
    epsB         0.02  7.5e-3   0.15  -0.26  -0.07   0.02   0.12   0.35    412    1.0
    al_std[1]    1.77    0.03   0.65   0.64   1.31   1.72   2.19   3.14    543   1.01
    al_std[2]    0.19    0.03   0.45  -0.76  -0.07   0.22   0.48   1.07    295   1.01
    al_std[3]    0.32    0.02   0.36  -0.48   0.11   0.33   0.55   1.02    242   1.02
    al_std[4]    0.35    0.03   0.44  -0.56    0.1   0.36   0.62   1.21    206   1.01
    al_std[5]   -0.21    0.04    0.5  -1.46  -0.47  -0.11   0.15   0.52    127   1.05
    al_std[6]     0.3    0.03   0.54  -0.88-6.2e-3   0.32   0.62   1.37    346   1.01
    al_sig     8.4e-3  2.8e-4 5.3e-3 2.3e-3 4.8e-3 7.2e-3   0.01   0.02    357   1.03
    alA          0.02  2.6e-4 4.4e-3   0.01   0.02   0.02   0.02   0.03    302   1.01
    alB          0.26  3.7e-3   0.08   0.12   0.21   0.26   0.31   0.42    418   1.02
    epsilon[1] 3.4e-3  3.5e-4   0.01  -0.02-3.2e-3 3.5e-3   0.01   0.02    846    1.0
    epsilon[2]   0.14  3.4e-4   0.01   0.12   0.14   0.14   0.15   0.16    890    1.0
    epsilon[3]   0.13  3.6e-4 8.5e-3   0.11   0.12   0.13   0.14   0.15    561   1.01
    epsilon[4]   0.89  3.4e-4 9.5e-3   0.88   0.89   0.89    0.9   0.91    794    1.0
    epsilon[5]   0.11  2.7e-4 7.4e-3   0.09    0.1   0.11   0.11   0.12    767   1.01
    epsilon[6]   0.08  4.4e-4   0.01   0.06   0.07   0.08   0.09    0.1    712    1.0
    alpha[1]     0.03  2.4e-4 4.3e-3   0.02   0.02   0.03   0.03   0.04    303   1.03
    alpha[2]     0.02  6.9e-5 2.0e-3   0.01   0.01   0.02   0.02   0.02    808    1.0
    alpha[3]     0.02  1.2e-5 6.6e-4   0.01   0.02   0.02   0.02   0.02   3291    1.0
    alpha[4]     0.02  9.3e-5 1.8e-3   0.01   0.02   0.02   0.02   0.02    388   1.01
    alpha[5]     0.01  7.6e-6 5.1e-4   0.01   0.01   0.01   0.01   0.01   4510    1.0
    alpha[6]     0.02  1.4e-4 2.9e-3   0.01   0.01   0.02   0.02   0.02    426    1.0
    lp__       -282.3    0.27    6.2 -295.3 -286.2 -282.0 -277.9 -270.8    516   1.01

    Samples were drawn using NUTS at Thu Feb 21 14:48:06 2019.
    For each parameter, n_eff is a crude measure of effective sample size,
    and Rhat is the potential scale reduction factor on split chains (at
    convergence, Rhat=1).

```python
fit.plot()
plt.show()
```

```python
#%matplotlib inline
def model(n,dnu,numax,eps,alpha,A,G,phi,tau):
    nmax = numax/dnu - eps
    s = (n + eps + ((alpha/2)*(nmax-n)**2) +
         (((A*G)/(2*np.pi)) * np.sin((2*np.pi*(n-nmax))/G + phi)))*dnu*np.exp(-(n-nmax)/(tau))
    return s

no_osc = np.loadtxt('/home/daniel/Documents/y4project/Dan_notebooks/no_osc_models.csv', delimiter=',')

fig = plt.figure(1, figsize=(12,4))
ax  = fig.add_subplot(111)
ax.set_position([0.1,0.1,0.6,0.85])

import itertools
colours = itertools.cycle(('r', 'orange', 'olive', 'teal', 'sienna', 'indigo'))

params = np.zeros([len(IDs), 8])

for i in np.arange(0,len(IDs),1):
    # Iterate point colour
    colour = next(colours)
    pts = modesID[i].loc[(modesID[i]['f0_err'] < 10e3)]
    # Plot all frequencies, and errorbars for the 'true' data
    ax.scatter(modesID[i]['f0'] % dnu_avgID[i], modesID[i]['f0'],
                label='', zorder=-1, facecolors='none', edgecolors=colour)
    ax.scatter(pts['f0'] % dnu_avgID[i], pts['f0'],
                label=r'ID = '+str(IDs[i]), c=colour)
    ax.errorbar(pts['f0'] % dnu_avgID[i], pts['f0'], xerr=pts['f0_err'],
                        zorder=0, fmt="none", label='', c='k',
                        capsize=2, markersize=4, elinewidth=1)

    params[i] = [np.mean(fit['dnu'], axis=0)[i],
           np.mean(fit['numax'], axis=0)[i],
           np.mean(fit['epsilon'], axis=0)[i],
           np.mean(fit['alpha'], axis=0)[i],
           np.mean(fit['A'], axis=0)[i],
           np.mean(fit['G'], axis=0)[i],
           np.mean(fit['phi'], axis=0)[i],
           np.mean(fit['tau'], axis=0)[i]]

    # Plot the modelled fit
    g2 = model(arr_n[i], params[i,0], params[i,1], params[i,2], params[i,3],
               params[i,4], params[i,5], params[i,6], params[i,7])
    ax.plot(g2 % dnu_avgID[i], g2, c=colour)
    ax.plot(no_osc[i] % dnu_avgID[i], no_osc[i], linestyle=':', c=colour)

np.savetxt('tau_models.csv', params, delimiter=',')

plt.legend()
plt.ylabel('Frequency')
plt.xlabel(r'Mod. Freq. Spacing ($\mu$Hz)')
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
#plt.savefig('seminar/somergbseizure.pdf', bbox='tight_layout')
```

```python
%matplotlib inline
plt.figure(2)
for i in np.arange(0,len(IDs),1):
    plt.scatter(modesID[i]['n'], modesID[i]['f0'], label=r'ID = '+str(IDs[i]))
plt.xlabel('Radial Order ($n$)')
plt.ylabel('Frequency')
plt.legend()
plt.show()

fig, ax = plt.subplots(len(IDs))
for i in np.arange(0,len(IDs),1):
    ax[i].errorbar(modesID[i]['n'], modesID[i]['f0'] - g2, yerr=modesID[i]['f0_err']*10)
    #ax[i].set_xlim([6,20])
    #ax[i].set_ylim([-3.5, 3.5])
    ax[i].set_ylim([-15, 15])

plt.show()
```

```python

```

```python
g2 = model(arr_n[5],
           np.mean(fit['dnu'], axis=0)[5],
           np.mean(fit['numax'], axis=0)[5],
           np.mean(fit['epsilon'], axis=0)[5],
           np.mean(fit['alpha'], axis=0)[5],
           np.mean(fit['A'], axis=0)[5],
           np.mean(fit['G'], axis=0)[5],
           np.mean(fit['phi'], axis=0)[5],
           1000
              )

plt.plot(g2 % dnu_avgID[5], g2)

pts = modesID[5].loc[(modesID[5]['f0_err'] < 10e3)]
plt.scatter(pts['f0'] % dnu_avgID[5], pts['f0'],
                label=r'ID = '+str(IDs[5]))
plt.errorbar(pts['f0'] % dnu_avgID[5], pts['f0'], xerr=pts['f0_err'],
                        zorder=0, fmt="none", label='', c='k',
                        capsize=2, markersize=4, elinewidth=1)
```

```python

```


## Example 1: `Hello world`

```python
# Print the character string `Hello world!`
print('Hello world!')
```

    Hello world!

## Example 2: Make a plot of the US unemployment rate

Download unemployment rate data from FRED (https://fred.stlouisfed.org/series/UNRATE/) and make a well-labeled plot.

```python
# Import the pandas library as pd
import pandas as pd

# Import the plotting library matplotlib.pyplot as plt
import matplotlib.pyplot as plt

# Download the unemployment rate data (PROVIDED)
unemployment = pd.read_csv('https://fred.stlouisfed.org/data/UNRATE.txt',skiprows=24,sep='\s+',index_col=0,parse_dates = True)

# Print first 5 rows of unemployment variable
unemployment.head()
```

      VALUE

      DATE

      1948-01-01
      3.4

      1948-02-01
      3.8

      1948-03-01
      4.0

      1948-04-01
      3.9

      1948-05-01
      3.5

```python
# Create a well-labeled plot of the US unemployment rate
unemployment.plot(legend=False,lw=3,alpha=0.75)
plt.ylabel('Percent')
plt.xlabel('Date')
plt.title('US Unemployment Rate')

# Save the figure to the current working directory at 120 dots per inch resolution
plt.savefig('unemployment.png',dpi=120)
```

## Example 2: Date and time

Use the `datetime` module to print the current date and time

```python
# Import the datetime module
import datetime

# Create a variable called `today` that stores the current date and time
today = datetime.datetime.today()

# Print the value of `today`
print(today)
```

    2020-01-07 11:23:25.333697

```python
# Print today's date formatted as: month day, year
print(today.strftime('%B %d, %Y'))
```

    January 07, 2020

```python
# Print the time component of `today` as: HH:MM:SS
print(today.strftime('%I:%M:%S'))
```

    11:23:25

## Example 3: Approximate a geometric series

Consider the infinite geometric series:

\begin{align}
\sum_{k=0}^{\infty} r^k & = r + r + r^2 + r^3 + \cdots
\end{align}

If $|r|<1$, then the series converges to:

\begin{align}
\frac{1}{1+r}
\end{align}

We can verify this fact numerically by computing the *truncated* series, i.e., compute the series to a finite number of terms. We will do this for $r=0.2$. The idea is to create a variable $s$ that is initially equal to zero. Then we let $k$ increment from 0, 1, up to $N$ and each time add $r^k$ to $s$.

```python
# Set the value of r
r = 0.2

# Create a variable `N` that stores the desired number of terms to compute
N = 25

# Initialize a variable that will store the value of the summation. Set equal to tzero
s = 0

# Iterate over values of k fron 0 to N. Add the new term to the sum and print the current value of the sum
for k in range(N+1):
    s+=r**k
    print(s)
```

    1.0
    1.2
    1.24
    1.248
    1.2496
    1.2499200000000001
    1.2499840000000002
    1.2499968000000001
    1.2499993600000001
    1.249999872
    1.2499999744
    1.24999999488
    1.249999998976
    1.2499999997952
    1.24999999995904
    1.249999999991808
    1.2499999999983618
    1.2499999999996725
    1.2499999999999347
    1.2499999999999871
    1.2499999999999976
    1.2499999999999996
    1.25
    1.25
    1.25
    1.25

## Example 4: Draw random values from a normal probability distribution

The `numpy` library has a bunch of powerful tools for working with numbers. Here we'll use `numpy` to draw a random sample from the normal probability distribution. Specifically, 100 values from $\mathcal{N}(2,0.2^2)$.

```python
# Import the numpy numerical tools library as np
import numpy as np

# Set the seed or state of the numpy random number generator
np.random.seed(126)

# Create a variable called 'x' that contains 100 draws from the normal(2,0.2^2) didistribution
x = np.random.normal(loc=2,scale=0.2,size=100)

plt.plot(x)
```

```python
# Print the mean of the values in the variable 'x'
print(np.mean(x))
```

    1.983109892630882

```python
# Print the standard deviation of the values in the variable 'x'
print(np.std(x))
```

    0.21843109555638324


```python
from IPython.display import Image
Image('../../../python_for_probability_statistics_and_machine_learning.jpg')
```

[Python for Probability, Statistics, and Machine Learning](https://www.springer.com/fr/book/9783319307152)

```python
import numpy as np
np.random.seed(12345)
```

# Conditional Expectation and Mean Square Error

In this section, we work through a detailed example using conditional
expectation and optimization methods.  Suppose we have two fair six-sided die
($X$ and $Y$) and we want to measure the sum of the two variables as $Z=X+Y$.
Further, let's suppose that given $Z$, we want the best estimate of $X$ in the
mean-squared-sense. Thus, we want to minimize the following:

$$
J(\alpha) = \sum ( x - \alpha z )^2 \mathbb{P}(x,z)
$$

 where $\mathbb{P}$ is the probability mass function for this problem.
The idea is that when we have solved this problem, we will have a function of
$Z$ that is going to be the minimum MSE  estimate of $X$.  We can substitute in
for $Z$ in $J$ and get:

$$
J(\alpha) = \sum ( x - \alpha (x+y) )^2 \mathbb{P}(x,y)
$$

  Let's work out the steps in Sympy in the following:

```python
import sympy as S
from sympy.stats import density, E, Die

x=Die('D1',6)     # 1st six sided die
y=Die('D2',6)     # 2nd six sides die
a=S.symbols('a')
z = x+y           # sum of 1st and 2nd die
J = E((x-a*(x+y))**2) # expectation
print S.simplify(J)
```

 With all that setup we can now use basic calculus to minimize the
objective function $J$,

```python
sol,=S.solve(S.diff(J,a),a) # using calculus to minimize
print sol # solution is 1/2
```

**Programming Tip.**

Sympy has a `stats` module that can do some basic work with expressions
involving probability densities and expectations. The above code uses its `E`
function to compute the expectation.

 This says that $z/2$ is the MSE estimate of $X$ given $Z$ which means
geometrically (interpreting the MSE as a squared distance weighted by the
probability mass function) that $z/2$ is as *close* to $x$ as we are going to
get for a given $z$.

Let's look at the same problem using the conditional expectation operator $
\mathbb{E}(\cdot|z) $ and apply it to our definition of $Z$. Then

$$
\mathbb{E}(z|z)=\mathbb{E}(x+y|z)=\mathbb{E}(x|z)+\mathbb{E}(y|z)=z
$$

 using the linearity of the expectation. Now, since by the
symmetry of the problem (i.e., two identical die), we have

$$
\mathbb{E}(x|z)=\mathbb{E}(y|z)
$$

 we can plug this in and solve

$$
2 \mathbb{E}(x|z)=z
$$

 which once again gives,

$$
\mathbb{E}(x|z)  =\frac{z}{2}
$$

```python
%matplotlib inline

from numpy import arange,array
from matplotlib.pylab import subplots, cm
from sympy import Integer
fig,ax = subplots()
v = arange(1,7) + arange(1,7)[:,None]
foo=lambda i: density(z)[Integer(i)].evalf() # some tweaks to get a float out
Zmass=array(map(foo,v.flat),dtype=float).reshape(6,6)

pc=ax.pcolor(arange(1,8),arange(1,8),Zmass,cmap=cm.gray)
_=ax.set_xticks([(i+0.5) for i in range(1,7)])
_=ax.set_xticklabels([str(i) for i in range(1,7)])
_=ax.set_yticks([(i+0.5) for i in range(1,7)])
_=ax.set_yticklabels([str(i) for i in range(1,7)])
for i in range(1,7):
    for j in range(1,7):
        _=ax.text(i+.5,j+.5,str(i+j),fontsize=18,fontweight='bold',color='goldenrod')

_=ax.set_title(r'Probability Mass for $Z$',fontsize=18)
_=ax.set_xlabel('$X$ values',fontsize=18)
_=ax.set_ylabel('$Y$ values',fontsize=18);
cb=fig.colorbar(pc)
_=cb.ax.set_title(r'Probability',fontsize=12)
#fig.savefig('fig-probability/Conditional_expectation_MSE_001.png')
```

The values of $Z$ are in yellow with the corresponding values for $X$ and $Y$ on the axes. The gray scale colors indicate the underlying joint probability density.

 which is equal to the estimate we just found by minimizing the MSE.
Let's explore this further with [Figure](#fig:Conditional_expectation_MSE_001).  [Figure](#fig:Conditional_expectation_MSE_001) shows the values of $Z$ in yellow with
the corresponding values for $X$ and $Y$ on the axes.  Suppose $z=2$, then the
closest $X$ to this is $X=1$, which is what $\mathbb{E}(x|z)=z/2=1$ gives. What
happens when $Z=7$? In this case, this value is spread out diagonally along the
$X$ axis so if $X=1$, then $Z$ is 6 units away, if $X=2$, then $Z$ is 5 units
away and so on.

Now, back to the original question, if we had $Z=7$ and we wanted
to get as close as we could to this using $X$, then why not choose
$X=6$ which is only one unit away from $Z$? The problem with doing
that is $X=6$ only occurs 1/6 of the time, so we are not likely to
get it right the other 5/6 of the time. So, 1/6 of the time we are
one unit away but 5/6 of the time we are much more than one unit
away. This means that the MSE score is going to be worse. Since
each value of $X$ from 1 to 6 is equally likely, to play it safe,
we choose $7/2$ as the estimate, which is what the conditional
expectation suggests.

We can check this claim with samples using Sympy below:

```python
import numpy as np
from sympy import stats
# Eq constrains Z
samples_z7 = lambda : stats.sample(x, S.Eq(z,7))
#using 6 as an estimate
mn= np.mean([(6-samples_z7())**2 for i in range(100)])
#7/2 is the MSE estimate
mn0= np.mean([(7/2.-samples_z7())**2 for i in range(100)])
print 'MSE=%3.2f using 6 vs MSE=%3.2f using 7/2 ' % (mn,mn0)
```

**Programming Tip.**

The `stats.sample(x, S.Eq(z,7))` function call samples the `x` variable subject
to a condition on the `z` variable. In other words, it generates random samples
of `x` die, given that the sum of the outcomes of that die and the `y` die add
up to `z==7`.

 Please run the above code repeatedly in the Jupyter/IPython
Notebook corresponding to this section until you have convinced
yourself  that the $\mathbb{E}(x|z)$ gives the lower MSE every
time.  To push this reasoning, let's consider the case where the
die is so biased so that the outcome of `6` is ten times more
probable than any of the other outcomes. That is,

$$
\mathbb{P}(6) = 2/3
$$

 whereas $\mathbb{P}(1)=\mathbb{P}(2)=\ldots=\mathbb{P}(5)=1/15$.
We can explore this using Sympy as in the following:

```python
# here 6 is ten times more probable than any other outcome
x=stats.FiniteRV('D3',{1:1/15., 2:1/15.,
                       3:1/15., 4:1/15.,
                       5:1/15., 6:2/3.})
```

 As before, we construct the sum of the two dice, and plot the
corresponding probability mass function in [Figure](#fig:Conditional_expectation_MSE_002).  As compared with [Figure](#fig:Conditional_expectation_MSE_001), the probability mass has been shifted
away from the smaller numbers.

```python
z = x + y
foo=lambda i: density(z)[S.Integer(i)].evalf() # some tweaks to get a float out
v = np.arange(1,7) + np.arange(1,7)[:,None]
Zmass=np.array(map(foo,v.flat),dtype=float).reshape(6,6)

from matplotlib.pylab import subplots, cm
fig,ax=subplots()
pc=ax.pcolor(np.arange(1,8),np.arange(1,8),Zmass,cmap=cm.gray)
_=ax.set_xticks([(i+0.5) for i in range(1,7)])
_=ax.set_xticklabels([str(i) for i in range(1,7)])
_=ax.set_yticks([(i+0.5) for i in range(1,7)])
_=ax.set_yticklabels([str(i) for i in range(1,7)])
for i in range(1,7):
    for j in range(1,7):
        _=ax.text(i+.5,j+.5,str(i+j),fontsize=18,fontweight='bold',color='goldenrod')

_=ax.set_title(r'Probability Mass for $Z$; Nonuniform case',fontsize=16)
_=ax.set_xlabel('$X$ values',fontsize=18)
_=ax.set_ylabel('$Y$ values',fontsize=18);
cb=fig.colorbar(pc)
_=cb.ax.set_title(r'Probability',fontsize=12)
#fig.savefig('fig-probability/Conditional_expectation_MSE_002.png')
```

The values of $Z$ are in yellow with the corresponding values for $X$ and $Y$ on the axes.

Let's see what the conditional expectation says about how we can estimate $X$
from $Z$.

```python
E(x, S.Eq(z,7)) # conditional expectation E(x|z=7)
```

    5.00000000000000

  Now that we have $\mathbb{E}(x|z=7) = 5$, we can generate
samples as before and see if this gives the minimum MSE.

```python
samples_z7 = lambda : stats.sample(x, S.Eq(z,7))
#using 6 as an estimate
mn= np.mean([(6-samples_z7())**2 for i in range(100)])
#5 is the MSE estimate
mn0= np.mean([(5-samples_z7())**2 for i in range(100)])
print 'MSE=%3.2f using 6 vs MSE=%3.2f using 5 ' % (mn,mn0)
```

    MSE=4.08 using 6 vs MSE=2.20 using 5

Using a simple example, we have emphasized the connection between minimum mean
squared error problems and conditional expectation. Hopefully, the last two
figures helped expose the role of the probability density.  Next, we'll
continue revealing  the true power of the conditional expectation as we
continue to develop corresponding geometric intuition.

```python

```


### Installation

```
!pip install qiskit
```

    Collecting qiskit
      Downloading https://files.pythonhosted.org/packages/96/6b/1c1a5ed225ac2e42ebc07df49e9c1ce0d3221c6d0744e7d6d5983ff0e69e/qiskit-0.23.6.tar.gz
    Collecting qiskit-terra==0.16.4
      Downloading https://files.pythonhosted.org/packages/56/e2/c256863ec8fab2162bf17d7e9c8564186a4636b4f24b5ea5d38a014b809e/qiskit_terra-0.16.4-cp36-cp36m-manylinux2010_x86_64.whl (8.5MB)
         |████████████████████████████████| 8.5MB 8.2MB/s
    Collecting qiskit-aer==0.7.5
      Downloading https://files.pythonhosted.org/packages/0e/4e/cf4493e8d8a7785b945b483429aa4b7960e30576eb8f3b24444991781460/qiskit_aer-0.7.5-cp36-cp36m-manylinux2010_x86_64.whl (17.6MB)
         |████████████████████████████████| 17.6MB 248kB/s
    Collecting qiskit-ibmq-provider==0.11.1
      Downloading https://files.pythonhosted.org/packages/32/b9/f99bec4fdc4dec234d8a85a8da378750441a203c614be35353f5e8738316/qiskit_ibmq_provider-0.11.1-py3-none-any.whl (195kB)
         |████████████████████████████████| 204kB 29.7MB/s
    Collecting qiskit-ignis==0.5.2
      Downloading https://files.pythonhosted.org/packages/b2/1d/d5df8cf48cf20a26deb6e0535f020cd2ad625cb27457abb0ab39c2104e38/qiskit_ignis-0.5.2-py3-none-any.whl (203kB)
         |████████████████████████████████| 204kB 36.4MB/s
    Collecting qiskit-aqua==0.8.2
      Downloading https://files.pythonhosted.org/packages/0e/82/528741676960a7fd1c75151639601826118f2729ad53034fa3b39290c056/qiskit_aqua-0.8.2-py3-none-any.whl (2.1MB)
         |████████████████████████████████| 2.1MB 34.4MB/s
    Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.6/dist-packages (from qiskit-terra==0.16.4->qiskit) (1.19.5)
    Collecting sympy>=1.3
      Downloading https://files.pythonhosted.org/packages/ff/69/b16fc81b939d3efdd0b552f2e3e54f7fa1423d0c320cced2e69e675dde26/sympy-1.7.1-py3-none-any.whl (5.9MB)
         |████████████████████████████████| 5.9MB 39.5MB/s
    Collecting contextvars>=2.4; python_version < "3.7"
      Downloading https://files.pythonhosted.org/packages/83/96/55b82d9f13763be9d672622e1b8106c85acb83edd7cc2fa5bc67cd9877e9/contextvars-2.4.tar.gz
    Requirement already satisfied: scipy>=1.4 in /usr/local/lib/python3.6/dist-packages (from qiskit-terra==0.16.4->qiskit) (1.4.1)
    Requirement already satisfied: dill>=0.3 in /usr/local/lib/python3.6/dist-packages (from qiskit-terra==0.16.4->qiskit) (0.3.3)
    Collecting fastjsonschema>=2.10
      Downloading https://files.pythonhosted.org/packages/89/1c/8be51fa42aadc1c1611a52b866e1a5a1032a504f24789cf140b4e6d7c940/fastjsonschema-2.15.0-py3-none-any.whl
    Requirement already satisfied: psutil>=5 in /usr/local/lib/python3.6/dist-packages (from qiskit-terra==0.16.4->qiskit) (5.4.8)
    Requirement already satisfied: python-dateutil>=2.8.0 in /usr/local/lib/python3.6/dist-packages (from qiskit-terra==0.16.4->qiskit) (2.8.1)
    Collecting retworkx>=0.5.0
      Downloading https://files.pythonhosted.org/packages/b7/21/5d4498cedb23927a31173176975e2162e0d6a2b80846ef254dd5500e10da/retworkx-0.7.2-cp36-cp36m-manylinux2010_x86_64.whl (1.3MB)
         |████████████████████████████████| 1.3MB 44.1MB/s
    Requirement already satisfied: jsonschema>=2.6 in /usr/local/lib/python3.6/dist-packages (from qiskit-terra==0.16.4->qiskit) (2.6.0)
    Collecting ply>=3.10
      Downloading https://files.pythonhosted.org/packages/a3/58/35da89ee790598a0700ea49b2a66594140f44dec458c07e8e3d4979137fc/ply-3.11-py2.py3-none-any.whl (49kB)
         |████████████████████████████████| 51kB 3.8MB/s
    Collecting python-constraint>=1.4
      Downloading https://files.pythonhosted.org/packages/37/8b/5f1bc2734ca611943e1d6733ee244238679f6410a10cd45ede55a61a8402/python-constraint-1.4.0.tar.bz2
    Requirement already satisfied: networkx>=2.2 in /usr/local/lib/python3.6/dist-packages (from qiskit-terra==0.16.4->qiskit) (2.5)
    Requirement already satisfied: cython>=0.27.1 in /usr/local/lib/python3.6/dist-packages (from qiskit-aer==0.7.5->qiskit) (0.29.21)
    Collecting pybind11>=2.4
      Downloading https://files.pythonhosted.org/packages/8d/43/7339dbabbc2793718d59703aace4166f53c29ee1c202f6ff5bf8a26c4d91/pybind11-2.6.2-py2.py3-none-any.whl (191kB)
         |████████████████████████████████| 194kB 59.2MB/s
    Requirement already satisfied: requests>=2.19 in /usr/local/lib/python3.6/dist-packages (from qiskit-ibmq-provider==0.11.1->qiskit) (2.23.0)
    Requirement already satisfied: nest-asyncio!=1.1.0,>=1.0.0 in /usr/local/lib/python3.6/dist-packages (from qiskit-ibmq-provider==0.11.1->qiskit) (1.5.1)
    Collecting websockets>=8
      Downloading https://files.pythonhosted.org/packages/bb/d9/856af84843912e2853b1b6e898ac8b802989fcf9ecf8e8445a1da263bf3b/websockets-8.1-cp36-cp36m-manylinux2010_x86_64.whl (78kB)
         |████████████████████████████████| 81kB 5.5MB/s
    Collecting requests-ntlm>=1.1.0
      Downloading https://files.pythonhosted.org/packages/03/4b/8b9a1afde8072c4d5710d9fa91433d504325821b038e00237dc8d6d833dc/requests_ntlm-1.1.0-py2.py3-none-any.whl
    Requirement already satisfied: urllib3>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from qiskit-ibmq-provider==0.11.1->qiskit) (1.24.3)
    Requirement already satisfied: scikit-learn>=0.17 in /usr/local/lib/python3.6/dist-packages (from qiskit-ignis==0.5.2->qiskit) (0.22.2.post1)
    Requirement already satisfied: setuptools>=40.1.0 in /usr/local/lib/python3.6/dist-packages (from qiskit-ignis==0.5.2->qiskit) (53.0.0)
    Requirement already satisfied: fastdtw in /usr/local/lib/python3.6/dist-packages (from qiskit-aqua==0.8.2->qiskit) (0.3.4)
    Requirement already satisfied: pandas in /usr/local/lib/python3.6/dist-packages (from qiskit-aqua==0.8.2->qiskit) (1.1.5)
    Requirement already satisfied: h5py in /usr/local/lib/python3.6/dist-packages (from qiskit-aqua==0.8.2->qiskit) (2.10.0)
    Collecting dlx
      Downloading https://files.pythonhosted.org/packages/54/c0/b8fb5bb727e983b6f5251433ef941b48f38c65bb0bd6ec509e9185bcd406/dlx-1.0.4.tar.gz
    Collecting quandl
      Downloading https://files.pythonhosted.org/packages/c2/58/9f0e69d836045e3865d263e9ed49f42b23a58526fdabb30f74c430baee3f/Quandl-3.6.0-py2.py3-none-any.whl
    Requirement already satisfied: dataclasses; python_version < "3.7" in /usr/local/lib/python3.6/dist-packages (from qiskit-aqua==0.8.2->qiskit) (0.8)
    Collecting docplex==2.15.194
      Downloading https://files.pythonhosted.org/packages/9b/76/ebf5b75a25c4357c9a4758cae7fea778dd7f282b83c2f22e150dc7e3c852/docplex-2.15.194.tar.gz (582kB)
         |████████████████████████████████| 583kB 24.8MB/s
    Collecting yfinance
      Downloading https://files.pythonhosted.org/packages/7a/e8/b9d7104d3a4bf39924799067592d9e59119fcfc900a425a12e80a3123ec8/yfinance-0.1.55.tar.gz
    Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.6/dist-packages (from sympy>=1.3->qiskit-terra==0.16.4->qiskit) (1.1.0)
    Collecting immutables>=0.9
      Downloading https://files.pythonhosted.org/packages/4a/52/e64a14a99c509cbdfe0405e9f076aef0331cb9548a3efa1d5bacd524978a/immutables-0.15-cp36-cp36m-manylinux1_x86_64.whl (100kB)
         |████████████████████████████████| 102kB 7.8MB/s
    Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.6/dist-packages (from python-dateutil>=2.8.0->qiskit-terra==0.16.4->qiskit) (1.15.0)
    Requirement already satisfied: decorator>=4.3.0 in /usr/local/lib/python3.6/dist-packages (from networkx>=2.2->qiskit-terra==0.16.4->qiskit) (4.4.2)
    Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19->qiskit-ibmq-provider==0.11.1->qiskit) (2.10)
    Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19->qiskit-ibmq-provider==0.11.1->qiskit) (3.0.4)
    Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19->qiskit-ibmq-provider==0.11.1->qiskit) (2020.12.5)
    Collecting cryptography>=1.3
      Downloading https://files.pythonhosted.org/packages/f8/1f/acde6ff69864c5e78b56488e3afd93c1ccc8c2651186e2a5f93d93f64859/cryptography-3.4.6-cp36-abi3-manylinux2014_x86_64.whl (3.2MB)
         |████████████████████████████████| 3.2MB 30.8MB/s
    Collecting ntlm-auth>=1.0.2
      Downloading https://files.pythonhosted.org/packages/ff/84/97c550164b54942b0e908c31ef09d9469f3ba4cd7332a671e2125732f63b/ntlm_auth-1.5.0-py2.py3-none-any.whl
    Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.6/dist-packages (from scikit-learn>=0.17->qiskit-ignis==0.5.2->qiskit) (1.0.0)
    Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.6/dist-packages (from pandas->qiskit-aqua==0.8.2->qiskit) (2018.9)
    Collecting inflection>=0.3.1
      Downloading https://files.pythonhosted.org/packages/59/91/aa6bde563e0085a02a435aa99b49ef75b0a4b062635e606dab23ce18d720/inflection-0.5.1-py2.py3-none-any.whl
    Requirement already satisfied: more-itertools in /usr/local/lib/python3.6/dist-packages (from quandl->qiskit-aqua==0.8.2->qiskit) (8.7.0)
    Requirement already satisfied: multitasking>=0.0.7 in /usr/local/lib/python3.6/dist-packages (from yfinance->qiskit-aqua==0.8.2->qiskit) (0.0.9)
    Collecting lxml>=4.5.1
      Downloading https://files.pythonhosted.org/packages/bd/78/56a7c88a57d0d14945472535d0df9fb4bbad7d34ede658ec7961635c790e/lxml-4.6.2-cp36-cp36m-manylinux1_x86_64.whl (5.5MB)
         |████████████████████████████████| 5.5MB 45.9MB/s
    Requirement already satisfied: cffi>=1.12 in /usr/local/lib/python3.6/dist-packages (from cryptography>=1.3->requests-ntlm>=1.1.0->qiskit-ibmq-provider==0.11.1->qiskit) (1.14.4)
    Requirement already satisfied: pycparser in /usr/local/lib/python3.6/dist-packages (from cffi>=1.12->cryptography>=1.3->requests-ntlm>=1.1.0->qiskit-ibmq-provider==0.11.1->qiskit) (2.20)
    Building wheels for collected packages: qiskit, contextvars, python-constraint, dlx, docplex, yfinance
      Building wheel for qiskit (setup.py) ... done
      Created wheel for qiskit: filename=qiskit-0.23.6-cp36-none-any.whl size=2883 sha256=448fe804f4a1d8dd5b5d2765a10ad01dd0a05e3b9249ec7f645ac6036864be98
      Stored in directory: /root/.cache/pip/wheels/b1/5e/2a/a9afb84cb4e8a08ea994982aaf040b2c668ef5fee81430f1f5
      Building wheel for contextvars (setup.py) ... done
      Created wheel for contextvars: filename=contextvars-2.4-cp36-none-any.whl size=7667 sha256=a97389bed2439405119384255d0e97f6aa93618a48faa1d9e3ce29b8479dd5f5
      Stored in directory: /root/.cache/pip/wheels/a5/7d/68/1ebae2668bda2228686e3c1cf16f2c2384cea6e9334ad5f6de
      Building wheel for python-constraint (setup.py) ... done
      Created wheel for python-constraint: filename=python_constraint-1.4.0-py2.py3-none-any.whl size=24079 sha256=668507f30864f03c81d8f46d59340e421eae5278fb09f8ac4a3181145eb02416
      Stored in directory: /root/.cache/pip/wheels/34/31/15/7b070b25d0a549d20ce2e9fe6d727471c2c61ef904720fd40c
      Building wheel for dlx (setup.py) ... done
      Created wheel for dlx: filename=dlx-1.0.4-cp36-none-any.whl size=5711 sha256=284359d03e3718772785b0dbadd4d875597a379ede34bc035c6de1881b908aaa
      Stored in directory: /root/.cache/pip/wheels/bb/ba/15/fdd0deb104df3254912998150ba9245668db06b00af5912d1a
      Building wheel for docplex (setup.py) ... done
      Created wheel for docplex: filename=docplex-2.15.194-cp36-none-any.whl size=645116 sha256=71e8bec82b7555ce4dce6921a7e66ce71fe0c5c8a742f15ca5bdf0cedb8dd814
      Stored in directory: /root/.cache/pip/wheels/ce/fa/5d/8f4d5fed1b8edb72e453cb2ac6fb75c6776b49d7174eb70457
      Building wheel for yfinance (setup.py) ... done
      Created wheel for yfinance: filename=yfinance-0.1.55-py2.py3-none-any.whl size=22616 sha256=beab70c102162821a68879f43e634c3d77db5ef64eb3a6fa0f53c3cde6fa1462
      Stored in directory: /root/.cache/pip/wheels/04/98/cc/2702a4242d60bdc14f48b4557c427ded1fe92aedf257d4565c
    Successfully built qiskit contextvars python-constraint dlx docplex yfinance
    Installing collected packages: sympy, immutables, contextvars, fastjsonschema, retworkx, ply, python-constraint, qiskit-terra, pybind11, qiskit-aer, websockets, cryptography, ntlm-auth, requests-ntlm, qiskit-ibmq-provider, qiskit-ignis, dlx, inflection, quandl, docplex, lxml, yfinance, qiskit-aqua, qiskit
      Found existing installation: sympy 1.1.1
        Uninstalling sympy-1.1.1:
          Successfully uninstalled sympy-1.1.1
      Found existing installation: lxml 4.2.6
        Uninstalling lxml-4.2.6:
          Successfully uninstalled lxml-4.2.6
    Successfully installed contextvars-2.4 cryptography-3.4.6 dlx-1.0.4 docplex-2.15.194 fastjsonschema-2.15.0 immutables-0.15 inflection-0.5.1 lxml-4.6.2 ntlm-auth-1.5.0 ply-3.11 pybind11-2.6.2 python-constraint-1.4.0 qiskit-0.23.6 qiskit-aer-0.7.5 qiskit-aqua-0.8.2 qiskit-ibmq-provider-0.11.1 qiskit-ignis-0.5.2 qiskit-terra-0.16.4 quandl-3.6.0 requests-ntlm-1.1.0 retworkx-0.7.2 sympy-1.7.1 websockets-8.1 yfinance-0.1.55

### 1.SwapGate

```
from qiskit import AncillaRegister,QuantumRegister, ClassicalRegister, QuantumCircuit, execute, Aer
```

```
q= QuantumRegister(2)
c= ClassicalRegister(2)
qc= QuantumCircuit(q,c)

qc.x(q[0])
qc.h(q[1])
# qc.cnot(q[0],q[1])
# qc.cnot(q[1],q[0])
# qc.cnot(q[0],q[1])
job = execute(qc,Aer.get_backend('qasm_simulator'),shots=4096)

counts = job.result().get_counts(qc)
print(counts)
qc.measure(q,c)

qc.draw()
```

    {'00': 4096}

      ┌───┐┌─┐
q9_0: ┤ X ├┤M├───
      ├───┤└╥┘┌─┐
q9_1: ┤ H ├─╫─┤M├
      └───┘ ║ └╥┘
c7: 2/══════╩══╩═
            0  1

```
from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit
from numpy import pi

qreg_q = QuantumRegister(2, 'q')
creg_c = ClassicalRegister(2, 'c')
circuit = QuantumCircuit(qreg_q, creg_c)

circuit.x(qreg_q[0])
circuit.swap(qreg_q[0], qreg_q[1])
circuit.measure(qreg_q[0], creg_c[0])
circuit.measure(qreg_q[1], creg_c[1])
circuit.draw()
```

     ┌───┐   ┌─┐
q_0: ┤ X ├─X─┤M├───
     └───┘ │ └╥┘┌─┐
q_1: ──────X──╫─┤M├
              ║ └╥┘
c: 2/═════════╩══╩═
              0  1

### 2. Bit-flip code and the sign-flip code

```
from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit
from numpy import pi

qreg_q = QuantumRegister(2, 'q')
creg_c = ClassicalRegister(2, 'c')
circuit = QuantumCircuit(qreg_q, creg_c)

circuit.reset(qreg_q[0])
circuit.reset(qreg_q[1])
circuit.h(qreg_q[0])
circuit.cx(qreg_q[0], qreg_q[1])
circuit.z(qreg_q[0])
circuit.measure(qreg_q[0], creg_c[0])
circuit.measure(qreg_q[1], creg_c[1])
circuit.draw()
```

          ┌───┐     ┌───┐┌─┐
q_0: ─|0>─┤ H ├──■──┤ Z ├┤M├
          └───┘┌─┴─┐└┬─┬┘└╥┘
q_1: ─|0>──────┤ X ├─┤M├──╫─
               └───┘ └╥┘  ║
c: 2/═════════════════╩═══╩═
                      1   0


# Orthogonal Random Forest: Use Cases and Examples

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

The EconML SDK implements the following OrthoForest variants:

* ContinuousTreatmentOrthoForest: suitable for continuous treatments

* DiscreteTreatmentOrthoForest: suitable for discrete treatments

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

## Notebook Contents

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

```python
import econml
import warnings
warnings.filterwarnings('ignore')
```

```python
# Main imports
from econml.ortho_forest import ContinuousTreatmentOrthoForest, DiscreteTreatmentOrthoForest
from econml.sklearn_extensions.linear_model import WeightedLassoCVWrapper, WeightedLasso, WeightedLassoCV

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

%matplotlib inline
```

# 1. Example Usage with Continuous Treatment Synthetic Data

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

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

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

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

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

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

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

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

## 1.2. Train Estimator

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

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

```python
est = ContinuousTreatmentOrthoForest(
    n_trees=200, min_leaf_size=5,
    max_depth=50, subsample_ratio=subsample_ratio,
    model_T=Lasso(alpha=lambda_reg),
    model_Y=Lasso(alpha=lambda_reg),
    model_T_final=WeightedLasso(alpha=lambda_reg),
    model_Y_final=WeightedLasso(alpha=lambda_reg),
    random_state=123)
```

To use the built-in confidence intervals constructed via Bootstrap of Little Bags, we need to specify `inference="blb"` at `fit` time.

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

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    6.0s
    [Parallel(n_jobs=-1)]: Done 200 out of 200 | elapsed:    7.7s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.2s
    [Parallel(n_jobs=-1)]: Done 118 out of 200 | elapsed:    1.0s remaining:    0.7s
    [Parallel(n_jobs=-1)]: Done 200 out of 200 | elapsed:    1.5s finished

    <econml.ortho_forest.ContinuousTreatmentOrthoForest at 0x225925d5c88>

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

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

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

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

## 1.3. Performance Visualization

```python
plt.plot(X_test, treatment_effects, label='ORF estimate')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.fill_between(X_test[:, 0], te_lower, te_upper, label="90% BLB CI", alpha=0.3)
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.show()
```

# 2. Example Usage with Binary Treatment Synthetic Data

## 2.1. DGP
We use the following DGP:

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

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

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

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

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

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

## 2.2. Train Estimator

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

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

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.2s
    [Parallel(n_jobs=-1)]: Done 200 out of 200 | elapsed:    1.4s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.2s
    [Parallel(n_jobs=-1)]: Done 200 out of 200 | elapsed:    1.5s finished

    <econml.ortho_forest.DiscreteTreatmentOrthoForest at 0x22591c28e10>

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

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

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

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

## 2.3. Performance Visualization

```python
# Smaller slices
y = treatment_effects
plt.plot(X_test, y, label='ORF estimate')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.fill_between(X_test[:, 0], te_lower, te_upper, label="90% BLB CI", alpha=0.3)
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.show()
```

# 3. Example Usage with Multiple Treatment Synthetic Data

## 3.1. DGP
We use the following DGP:

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

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

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

```python
def get_test_train_data(n, n_w, support_size, n_x, te_func, n_treatments):
    # Outcome support
    support_Y = np.random.choice(range(n_w), size=support_size, replace=False)
    coefs_Y = np.random.uniform(0, 1, size=support_size)
    epsilon_sample = lambda n: np.random.uniform(-1, 1, size=n)
    # Treatment support
    support_T = support_Y
    coefs_T = np.random.uniform(0, 1, size=(support_size, n_treatments))
    eta_sample = lambda n: np.random.uniform(-1, 1, size=n)
    # Generate controls, covariates, treatments and outcomes
    W = np.random.normal(0, 1, size=(n, n_w))
    X = np.random.uniform(0, 1, size=(n, n_x))
    # Heterogeneous treatment effects
    TE = np.array([te_func(x_i, n_treatments) for x_i in X])
    log_odds = np.dot(W[:, support_T], coefs_T)
    T_sigmoid = np.exp(log_odds)
    T_sigmoid = T_sigmoid/np.sum(T_sigmoid, axis=1, keepdims=True)
    T = np.array([np.random.choice(n_treatments, p=p) for p in T_sigmoid])
    TE = np.concatenate((np.zeros((n,1)), TE), axis=1)
    Y = TE[np.arange(n), T] + np.dot(W[:, support_Y], coefs_Y) + epsilon_sample(n)
    X_test = np.array(list(product(np.arange(0, 1, 0.01), repeat=n_x)))

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

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

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

## 3.2. Train Estimator

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

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

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.5s
    [Parallel(n_jobs=-1)]: Done 112 tasks      | elapsed:    4.1s
    [Parallel(n_jobs=-1)]: Done 272 tasks      | elapsed:   10.6s
    [Parallel(n_jobs=-1)]: Done 500 out of 500 | elapsed:   19.5s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.8s
    [Parallel(n_jobs=-1)]: Done 112 tasks      | elapsed:    4.0s
    [Parallel(n_jobs=-1)]: Done 272 tasks      | elapsed:    9.3s
    [Parallel(n_jobs=-1)]: Done 500 out of 500 | elapsed:   20.2s finished

    <econml.ortho_forest.DiscreteTreatmentOrthoForest at 0x22596b66978>

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

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

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

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

## 3.3. Performance Visualization

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

# 4. Example Usage with Real Continuous Treatment Observational Data

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

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

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

## 4.1. Data

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

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

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

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

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

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

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

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

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

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

## 4.2. Train Estimator

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

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

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

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done 232 tasks      | elapsed:    0.5s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:    2.1s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done 232 tasks      | elapsed:    0.5s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:    1.9s finished

    <econml.ortho_forest.ContinuousTreatmentOrthoForest at 0x22596e560b8>

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

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

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

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

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

## 4.3. Performance Visualization

```python
# Plot Orange Juice elasticity as a function of income
plt.plot(X_test.flatten(), treatment_effects, label="OJ Elasticity")
plt.fill_between(X_test.flatten(), te_lower, te_upper, label="90% BLB CI", alpha=0.3)
plt.xlabel(r'$\log$(Income)')
plt.ylabel('Orange Juice Elasticity')
plt.legend()
plt.title("Orange Juice Elasticity vs Income")
plt.show()
```

notebooks/Orthogonal Random Forest Examples.ipynb