# Quantization of Signals

*This jupyter notebook is part of a [collection of notebooks](../index.ipynb) on various topics of Digital Signal Processing. Please direct questions and suggestions to [Sascha.Spors@uni-rostock.de](mailto:Sascha.Spors@uni-rostock.de).*

## Characteristic of a Linear Uniform Quantizer

The characteristics of a quantizer depend on the mapping functions $f(\cdot)$, $g(\cdot)$ and the rounding operation $\lfloor \cdot \rfloor$ introduced in the [previous section](introduction.ipynb). A linear quantizer bases on linear mapping functions $f(\cdot)$ and $g(\cdot)$. A uniform quantizer splits the mapped input signal into quantization steps of equal size. Quantizers can be described by their nonlinear in-/output characteristic $x_Q[k] = \mathcal{Q} \{ x[k] \}$, where $\mathcal{Q} \{ \cdot \}$ denotes the quantization process. For linear uniform quantization it is common to differentiate between two characteristic curves, the so called mid-tread and mid-rise. Both are introduced in the following.

### Mid-Tread Characteristic Curve

The in-/output relation of the mid-tread quantizer is given as

\begin{equation}
x_Q[k] = Q \cdot \underbrace{\left\lfloor \frac{x[k]}{Q} + \frac{1}{2} \right\rfloor}_{index}
\end{equation}

where $Q$ denotes the constant quantization step size and $\lfloor \cdot \rfloor$ the [floor function](https://en.wikipedia.org/wiki/Floor_and_ceiling_functions) which maps a real number to the largest integer not greater than its argument. Without restricting $x[k]$ in amplitude, the resulting quantization indexes are [countable infinite](https://en.wikipedia.org/wiki/Countable_set). For a finite number of quantization indexes, the input signal has to be restricted to a minimal/maximal amplitude $x_\text{min} < x[k] < x_\text{max}$ before quantization. The resulting quantization characteristic of a linear uniform mid-tread quantizer is shown below

The term mid-tread is due to the fact that small values $|x[k]| < \frac{Q}{2}$ are mapped to zero.

#### Example - Mid-tread quantization of a sine signal

The quantization of one period of a sine signal $x[k] = A \cdot \sin[\Omega_0\,k]$ by a mid-tread quantizer is simulated. $A$ denotes the amplitude of the signal, $x_\text{min} = -1$ and $x_\text{max} = 1$ are the smallest and largest output values of the quantizer, respectively.

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

A = 1.2  # amplitude of signal
Q = 1/10  # quantization stepsize
N = 2000  # number of samples

def uniform_midtread_quantizer(x, Q):
    '''Uniform mid-tread quantizer with limiter.'''
    # limiter
    x = np.copy(x)
    idx = np.where(np.abs(x) >= 1)
    x[idx] = np.sign(x[idx])
    # linear uniform quantization
    xQ = Q * np.floor(x/Q + 1/2)

    return xQ

def plot_signals(x, xQ):
    '''Plot continuous, quantized and error signal.'''
    e = xQ - x
    plt.figure(figsize=(10, 6))
    plt.plot(x, label=r'signal $x[k]$')
    plt.plot(xQ, label=r'quantized signal $x_Q[k]$')
    plt.plot(e, label=r'quantization error $e[k]$')
    plt.xlabel(r'$k$')
    plt.axis([0, N, -1.1*A, 1.1*A])
    plt.legend()
    plt.grid()

# generate signal
x = A * np.sin(2*np.pi/N * np.arange(N))
# quantize signal
xQ = uniform_midtread_quantizer(x, Q)
# plot signals
plot_signals(x, xQ)
```

**Exercise**

* Change the quantization stepsize `Q` and the amplitude `A` of the signal. Which effect does this have on the quantization error?

Solution: The smaller the quantization step size, the smaller the quantization error is for $|x[k]| < 1$. Note, the quantization error is not bounded for $|x[k]| > 1$ due to the clipping of the signal $x[k]$.

### Mid-Rise Characteristic Curve

The in-/output relation of the mid-rise quantizer is given as

\begin{equation}
x_Q[k] = Q \cdot \Big( \underbrace{\left\lfloor\frac{ x[k] }{Q}\right\rfloor}_{index}  + \frac{1}{2} \Big)
\end{equation}

where $\lfloor \cdot \rfloor$ denotes the floor function. The quantization characteristic of a linear uniform mid-rise quantizer is illustrated below

The term mid-rise copes for the fact that $x[k] = 0$ is not mapped to zero. Small positive/negative values around zero are mapped to $\pm \frac{Q}{2}$.

#### Example - Mid-rise quantization of a sine signal

The previous example is now reevaluated using the mid-rise characteristic

```python
A = 1.2  # amplitude of signal
Q = 1/10  # quantization stepsize
N = 2000  # number of samples

def uniform_midrise_quantizer(x, Q):
    '''Uniform mid-rise quantizer with limiter.'''
    # limiter
    x = np.copy(x)
    idx = np.where(np.abs(x) >= 1)
    x[idx] = np.sign(x[idx])
    # linear uniform quantization
    xQ = Q * (np.floor(x/Q) + .5)

    return xQ

# generate signal
x = A * np.sin(2*np.pi/N * np.arange(N))
# quantize signal
xQ = uniform_midrise_quantizer(x, Q)
# plot signals
plot_signals(x, xQ)
```

**Exercise**

* What are the differences between the mid-tread and the mid-rise characteristic curves for the given example?

Solution: The mid-tread and the mid-rise quantization of the sine signal differ for signal values smaller than half of the quantization interval. Mid-tread has a representation of $x[k] = 0$ while this is not the case for the mid-rise quantization.

**Copyright**

This notebook is provided as [Open Educational Resource](https://en.wikipedia.org/wiki/Open_educational_resources). Feel free to use the notebook for your own purposes. The text is licensed under [Creative Commons Attribution 4.0](https://creativecommons.org/licenses/by/4.0/), the code of the IPython examples under the [MIT license](https://opensource.org/licenses/MIT). Please attribute the work as follows: *Sascha Spors, Digital Signal Processing - Lecture notes featuring computational examples*.


# Spectral Estimation of Random Signals

*This jupyter notebook is part of a [collection of notebooks](../index.ipynb) on various topics of Digital Signal Processing.

## Introduction

In the preceding sections various statistical measures have been introduced to characterize random processes and signals. For instance, the probability density function (PDF) $p_x(\theta)$, the mean value $\mu_x$, the auto-correlation function (ACF) $\varphi_{xx}[\kappa]$ and its Fourier transformation, the power spectral density (PSD) $\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$. For many random processes whose internal structure is known, these measures can be derived in closed-form. However, for practical random signals measures of interest have to be estimated from a limited number of samples. These estimated quantities can e.g. be used to fit a parametric model of the random process or as parameters in algorithms.

### Problem Statement

The estimation of the spectral properties of a random signal is of special interest for spectral analysis. The discrete Fourier transform (DFT) of a random signal is also random. It is not very well suited to gain insights into the average spectral structure of a random signal. We aim at estimating the PSD $\hat{\Phi}_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ of a wide-sense stationary (WSS) and ergodic process from a limited number of samples. This is known as [*spectral (density) estimation*](https://en.wikipedia.org/wiki/Spectral_density_estimation). Many techniques have been developed for this purpose. They can be classified into

1. non-parametric and
2. parametric

techniques. Non-parametric techniques estimate the PSD of the random signal without assuming any particular structure for the generating random process. In contrary, parametric techniques assume that the generating random process can be modeled by a few parameters. Their aim is to estimate these parameters in order to characterize the spectral properties of the random signal.

### Evaluation

Various measures have been introduced in order to evaluate the performance of a particular estimation technique. The estimate $\hat{\Phi}_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ can be regarded as a random signal itself. The performance of an estimator is therefore evaluated in a statistical sense. For the PSD, the following metrics are of interest

#### Bias

The [bias of an estimator](https://en.wikipedia.org/wiki/Estimator#Bias)

\begin{equation}
b_{\hat{\Phi}_{xx}}
= E\{ \hat{\Phi}_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) - \Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})\}
= E\{ \hat{\Phi}_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) \} - \Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})
\end{equation}

quantifies the difference between the estimated $\hat{\Phi}_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ and the true $\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$. An estimator is
* biased if $b_{\hat{\Phi}_{xx}} \neq 0$ and
* bias-free if $b_{\hat{\Phi}_{xx}} = 0$.

#### Variance

The [variance of an estimator](https://en.wikipedia.org/wiki/Estimator#Variance)

\begin{equation}
\sigma^2_{\hat{\Phi}_{xx}} = E \left\{ \left(\hat{\Phi}_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) - E\{ \hat{\Phi}_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})\} \right)^2 \right\}
\end{equation}

quantifies its quadratic deviation from its mean value $E\{ \hat{\Phi}_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})\}$.

#### Consistency

A [consistent estimator](https://en.wikipedia.org/wiki/Estimator#Consistency) is an estimator for which the following conditions hold for a large number $N$ of samples:

1. the estimator is unbiased
    $$ \lim_{N \to \infty} b_{\hat{\Phi}_{xx}} = 0 $$

2. its variance converges towards zero
    $$ \lim_{N \to \infty} \sigma^2_{\hat{\Phi}_{xx}} = 0 $$

3. it converges in probability to the true $\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$
    $$ \lim_{N \to \infty} \Pr \left\{ | \hat{\Phi}_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) -  \Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})| > \alpha \right\} = 0$$
    where $\alpha > 0$ denotes a (small) constant.

The last condition ensures that a consistent estimator does not generate outliers. Consistency is a desired property of an estimator. It ensures that if the number of samples $N$ increases towards infinity, the resulting estimates converges towards the true PSD.

#### Example - Discrete Fourier transform of sample functions

The following example computes and plots the magnitude of the discrete Fourier transform (DFT) $|X_n[\mu]|$ of an ensemble of random signals $x_n[k]$. In the plot, each color denotes one sample function and $\Omega[\mu] = \frac{2 \pi}{N} \mu$ where $N$ denotes the length of the sample function. The magnitude spectra are plotted as continuous signals for ease of illustration.

```python
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as sig

N = 256  # number of samples
M = 5  # number of sample functions

# generate random signal
np.random.seed(1)
x = np.random.normal(size=(M, N))
h = sig.firwin2(N, [0, .4, .42, .65, .67, 1], [0, 0, 1, 1, 0, 0])
x = [np.convolve(xi, h, mode='same') for xi in x]

# DFT of signal
X = np.fft.rfft(x, axis=1)
Om = np.linspace(0, np.pi, X.shape[1])

# plot signal and its spectrum
plt.figure(figsize=(10,4))
plt.plot(Om, np.abs(X.T))
plt.title('Magnitude spectrum')
plt.xlabel(r'$\Omega[\mu]$')
plt.ylabel(r'$|X[\mu]|$')
plt.axis([0, np.pi, 0, 30]);
```

**Exercise**

* Increase the number `N` of samples. What changes? What does not change with respect to the evaluation criteria introduced above?
* Is the DFT of a single sample function a consistent estimator for the spectral properties of a random process?

Solution: Increasing the number of samples does only lead to an increase in the number of discrete frequencies $\mu$. The amplitude of the fluctuations (variance) of the spectra within $1.3 < \Omega < 2$ is not decreased when increasing the number of samples. The DFT of a single sample function is hence not a consistent estimator since at least the second condition is violated.


# Multi-fidelity Kriging (MFK) with noisy observations

Andrés F. LOPEZ-LOPERA ONERA/DTIS/M2CI - Juin 2021

To use SMT models, please follow this link: https://github.com/SMTorg/SMT/blob/master/README.md. The documentation is available here: http://smt.readthedocs.io/en/latest/

**Reference paper:** https://www.sciencedirect.com/science/article/pii/S0965997818309360?via%3Dihub

**Preprint:** https://www.researchgate.net/profile/Mohamed_Amine_Bouhlel/publication/331976718_A_Python_surrogate_modeling_framework_with_derivatives/links/5cc3cebd299bf12097829631/A-Python-surrogate-modeling-framework-with-derivatives.pdf

Cite us:

M.-A. Bouhlel, J. T. Hwang, N. Bartoli, R. Lafage, J. Morlier, J .R.R.A Martins (2019), A Python surrogate modeling framework with derivatives, Advances in Engineering Software, 102662

@article{SMT2019,
  title={A Python surrogate modeling framework with derivatives},
  author={Mohamed Amine Bouhlel and John T. Hwang and Nathalie Bartoli and R{\'e}mi Lafage and Joseph Morlier and Joaquim R. R. A. Martins},
  journal={Advances in Engineering Software},
  pages={102662},
  year={2019},
  publisher={Elsevier}
}

```python
import numpy as np
from matplotlib import pyplot as plt
plt.rcParams['figure.figsize'] = [8, 5]

from smt.applications import MFK

%matplotlib inline
```

# 1. Problem statement

In this notebook, we focus on surrogate multi-fidelity Kriging models accouting for a noise term. Here we focus on a multi-fidelity GP framework with two levels of fidelity. The extension to more than two levels of fidelity can be obtained recursively as it is shown by Le Gratiet (2013). We denote $Y_{h}$ and $Y_{l}$ as the processes GPs related to the high and low-fidelity, respectively. To account for additive input-varying noisy observations, we can consider the system given by

\begin{align}
	Y_l^{\text{noise}}(x_i) &= Y_{l}(x_i) + \varepsilon_{l,i},
	\label{eq:MFsurrogateHeteroNoiseLF}	\\
	Y_h^{\text{noise}}(x_j) &= \rho(x_j) Y_{l}(x_j) + \nu(x_j) + \varepsilon_{h,j},
	\label{eq:MFsurrogateHeteroNoise}
\end{align}

where $\rho: \mathcal{D} \to \mathbb{R}$ is a scale factor between $Y_{l}$ and $Y_{h}$, and $\nu: \mathcal{D} \to \mathbb{R}$ the discrepancy function tasked with capturing the differences between both fidelity levels beyond scaling. We assume that $Y_l$ and $\nu$ are independent centred GP with covariance functions $k_l$ and $k_\nu$. For the additive Gaussian noise terms $\varepsilon_{l,1}, \ldots, \varepsilon_{l,n_l}$, $\varepsilon_{h,1}, \ldots, \varepsilon_{h,n_h}$, we assume that they are independent, and independent of $Y_l$ and $\nu$. This means that the covariance matrix $\Omega_l$ and $\Omega_h$ are given by:

$$\Omega_l = \begin{bmatrix} \tau_{l,1}^2 & \ldots & 0 \\ \vdots & \ddots & \vdots \\ 0 & \ldots & \tau_{l,n_l}^2 \end{bmatrix}, \qquad \Omega_h = \begin{bmatrix} \tau_{h,1}^2 & \ldots & 0 \\ \vdots & \ddots & \vdots \\ 0 & \ldots & \tau_{h,n_h}^2 \end{bmatrix},$$

with noise variances $\tau_{l,1}^2, \ldots, \tau_{l, n_l}^2, \tau_{h,1}^2, \ldots, \tau_{h, n_h}^2 \in \mathbb{R}^+$. Then, due to the linearity, $Y_l^{\text{noise}}$ and $Y_h^{\text{noise}}$ are also centred GP-distributed with covariance functions given by

\begin{align}
    \widetilde{k}_{l}(x, x') &= {k}_{l}(x, x') + \tau_{l,x}^2 \delta_{x}(x'),	\label{eq:MFsurrogateHeteroNoiseLFkernel} \\
    \widetilde{k}_{h}(x, x') &= \rho(x) \rho(x') {k}_{l}(x, x') + k_{\nu}(x, x') + \tau_{h,x}^2 \delta_{x}(x'),
\end{align}

where $\delta_{x}(x')$ is the Dirac delta function that is equal to one if $x = x'$, and zero otherwise. Then, since $Y_l^{\text{noise}}$ and $Y_h^{\text{noise}}$ are GP-distributed, one can establish the conditional formulas for computing the Gaussian distribution

$$Y_{h}|\{Y_{l}(x_{l,1}) + \varepsilon_{l,1} = y_{l, 1}, \ldots, Y_{l}(x_{l,n_{l}})  + \varepsilon_{l,n_l} = y_{l, n_{l}}, Y_{h}(x_{h,1})  + \varepsilon_{h,1} = y_{h, 1}, \ldots, Y_{h}(x_{h,n_{h}}) + \varepsilon_{h,n_h} = y_{h, n_{h}} \},$$

with observations $y_{l, 1}, \ldots, y_{l, n_l}$ (low-fidelity) and $y_{h, 1}, \ldots, y_{h, n_h}$ (high-fidelity). Those formulas follow the same structure as the ones for the noise-free case but replacing ${k}_{l}(x, x')$ to $\widetilde{k}_{l}(x, x')$ and ${k}_{h}(x, x')$ to $\widetilde{k}_{h}(x, x')$.

Two cases can be considered:

- **Homoscedastic (default) case:** the noise variances are considered to be equal, i.e.: $\tau^2 = \tau_1^2 = \cdots = \tau_n^2$. The value of $\tau^2$ can be estimated via maximum likelihood.
- **Heterocedastic case:** the noise variances vary across the observations, i.e.: $\tau_1^2 \ne \cdots \ne \tau_n^2$. Those variances can only be estimated if repetitions of observations are given. Developments here are based on pointwise sensible estimates [1] and on the implementations in [2].

In the following, we focus on the heteroscedastic case. We refer to the notebook ``SMT_Noise.ipynb`` for an example of the single-fidelity case.µ

**References**

[1] Loic Le Gratiet (2013). "Multi-Fidelity Gaussian Process Regression for Computer Experiments." Theses, Université Paris-Diderot - Paris VII.

[2] Bruce Ankenman, Barry L. Nelson and Jeremy Staum (2010). "Stochastic Kriging for Simulation Metamodeling." Operations Research. Vol. 58, No. 2 , pp. 371-382.

[3] Olivier Roustant, David Ginsbourger and Yves Deville (2012). "DiceKriging, DiceOptim: Two R Packages for the analysis of computer experiments by kriging-based metamodeling and optimization." Journal of Statistical Software, Vol. 51, No. 1, pp. 1-55.

## Toy data

```python
def LF_function(x): # low-fidelity function
    return 0.5*((x*6-2)**2)*np.sin((x*6-2)*2)+(x-0.5)*10. - 5

def HF_function(x): # high-fidelity function
    return ((x*6-2)**2)*np.sin((x*6-2)*2)

x = np.linspace(0, 1, 101, endpoint=True).reshape(-1, 1)

plt.figure()
plt.plot(x, HF_function(x), '--C0', label='High Fidelity (HF)')
plt.plot(x, LF_function(x), '--C1', label='Low Fidelity (LF)')
plt.ylim(-10, 17); plt.xlim(-0.1, 1.1)
plt.xlabel('x'); plt.ylabel('y')
plt.legend();
```

```python
#Expensive DOE with 4 points
Xt_e = np.array([0.0, 0.3, 0.6, 1.]).reshape(-1, 1)
#Cheap DOE with 12 points
Xt_c = np.linspace(1/9, 8/9, 8, endpoint=True).reshape(-1, 1)
# The two DOEs must be nest: Xt_e must be included within Xt_c
Xt_c = np.concatenate((Xt_c,Xt_e),axis=0)

# Evaluate the HF and LF functions
yt_e = HF_function(Xt_e)
yt_c = LF_function(Xt_c)

plt.figure()
plt.plot(x, HF_function(x), '--C0', label='High Fidelity (HF)')
plt.plot(x, LF_function(x), '--C1', label='Low Fidelity (LF)')
plt.scatter(Xt_e, yt_e, color='C0', label='HF doe')
plt.scatter(Xt_c, yt_c, color='C1', label='LF doe')
plt.ylim(-10, 17); plt.xlim(-0.1, 1.1)
plt.xlabel('x'); plt.ylabel('y')
plt.legend();
```

## MFK with given noise variance parameters

```python
# %% Build the MFK object
yt_joint = np.concatenate((yt_c, yt_e))
noise0 = [np.array([2e-01, 7e-02, 2e-03, 6e-08,
                    5e-02, 5e-02, 4e-02, 3e-02,
                    2e-01, 1e-02, 5e-06, 1e-01])/2,
          np.array([1.5, 0.4, 0.01, 0.1])/5]
theta0 = np.array([[0.5],[0.1]])
np.random.seed(7)
sm = MFK(theta0=theta0, theta_bounds = [1e-1, 20],
         noise0=noise0, use_het_noise = True,
         propagate_uncertainty=False, n_start=1)

# low-fidelity dataset names being integers from 0 to level-1
sm.set_training_values(Xt_c, yt_c, name=0)
# high-fidelity dataset without name
sm.set_training_values(Xt_e, yt_e)
# train the model
sm.train()

# test training
ntest = 101
nlvl = len(sm.X)
x = np.linspace(0, 1, ntest, endpoint=True).reshape(-1, 1)
# HF
y = sm.predict_values(x)
var = sm.predict_variances(x)
# LF
y0 = sm._predict_intermediate_values(x, 1)
var0, _ =  sm.predict_variances_all_levels(x)
var0 = var0[:,0].reshape(-1,1)
```

    ___________________________________________________________________________

                                        MFK
    ___________________________________________________________________________

     Problem size

          # training points.        : 4

    ___________________________________________________________________________

     Training

       Training ...
       Training - done. Time (sec):  0.0506232
    ___________________________________________________________________________

     Evaluation

          # eval points. : 101

       Predicting ...
       Predicting - done. Time (sec):  0.0009887

       Prediction time/pt. (sec) :  0.0000098

```python
# creating a single-fidelity GP model to compare with
theta0 = sm.options["theta0"]
noise0 = sm.options["noise0"]

# %% Kriging on monofidelity HF doe
sm_monoHF =  MFK(theta0=theta0[1].tolist(), theta_bounds = [1e-1, 20],
                 noise0=[noise0[1]], use_het_noise=True, n_start=1)

# %% Kriging on monofidelity LF Doe
sm_monoLF =  MFK(theta0=theta0[0].tolist(), theta_bounds = [1e-1, 20],
                 noise0=[noise0[0]], use_het_noise=True, n_start=1)

# train the model
sm_monoHF.set_training_values(Xt_e, yt_e)
sm_monoHF.train()

sm_monoLF.set_training_values(Xt_c, yt_c)
sm_monoLF.train()

# test training
y_monoHF = sm_monoHF.predict_values(x)
var_monoHF = sm_monoHF.predict_variances(x)

y_monoLF = sm_monoLF.predict_values(x)
var_monoLF = sm_monoLF.predict_variances(x)
```

    ___________________________________________________________________________

                                        MFK
    ___________________________________________________________________________

     Problem size

          # training points.        : 4

    ___________________________________________________________________________

     Training

       Training ...
       Training - done. Time (sec):  0.0120628
    ___________________________________________________________________________

                                        MFK
    ___________________________________________________________________________

     Problem size

          # training points.        : 12

    ___________________________________________________________________________

     Training

       Training ...
       Training - done. Time (sec):  0.0134740
    ___________________________________________________________________________

     Evaluation

          # eval points. : 101

       Predicting ...
       Predicting - done. Time (sec):  0.0000000

       Prediction time/pt. (sec) :  0.0000000

    ___________________________________________________________________________

     Evaluation

          # eval points. : 101

       Predicting ...
       Predicting - done. Time (sec):  0.0000000

       Prediction time/pt. (sec) :  0.0000000

```python
plt.figure(figsize = (15, 5))
plt.subplot(1,2,1)
plt.plot(x, HF_function(x), '--C0', label='High Fidelity (HF)')
plt.plot(x, LF_function(x), '--C1', label='Low Fidelity (LF)')
plt.plot(x, y, 'C0', label='MFGP - HF')
plt.fill_between(np.ravel(x), np.ravel(y-3*np.sqrt(var)),
                  np.ravel(y+3*np.sqrt(var)),
                  color='C0',alpha=0.2, label ='Confidence Interval 99% HF')
plt.plot(x, y0, 'C1', label='MFGP - LF')
plt.fill_between(np.ravel(x), np.ravel(y0-3*np.sqrt(var0)),
                  np.ravel(y0+3*np.sqrt(var0)),
                  color='C1',alpha=0.2, label ='Confidence Interval 99% LF')
plt.errorbar(Xt_e, yt_e, yerr=3*np.sqrt(noise0[-1]),
             fmt="o", color="C0", label='HF doe')
plt.errorbar(Xt_c, yt_c, yerr=3*np.sqrt(noise0[0]),
             fmt="o", color="C1", label='LF doe')
plt.ylim(-10, 17); plt.xlim(-0.1, 1.1)
plt.xlabel('x'); plt.ylabel('y')
plt.legend()
plt.title('Multi-fidelity GPs (MFGP)')

plt.subplot(1,2,2)
plt.plot(x, HF_function(x), '--C0', label='High Fidelity (HF)')
plt.plot(x, LF_function(x), '--C1', label='Low Fidelity (LF)')
plt.plot(x, y_monoHF, 'C2', label='GP - HF')
plt.fill_between(np.ravel(x), np.ravel(y_monoHF-3*np.sqrt(var_monoHF)),
                  np.ravel(y_monoHF+3*np.sqrt(var_monoHF)),
                  color='C2',alpha=0.2, label ='Confidence Interval 99% HF')
plt.plot(x, y_monoLF, 'C3', label='GP - LF')
plt.fill_between(np.ravel(x), np.ravel(y_monoLF-3*np.sqrt(var_monoLF)),
                  np.ravel(y_monoLF+3*np.sqrt(var_monoLF)),
                  color='C3',alpha=0.2, label ='Confidence Interval 99% LF')
plt.errorbar(Xt_e, yt_e, yerr=3*np.sqrt(noise0[-1]),
             fmt="o", color="C0")
plt.errorbar(Xt_c, yt_c, yerr=3*np.sqrt(noise0[0]),
             fmt="o", color="C1")
plt.ylim(-10, 17); plt.xlim(-0.1, 1.1)
plt.xlabel('x'); plt.ylabel('y')
plt.title('Single-fidelity GPs')
plt.legend();
```

## MFK with repetitions of observations

```python
# adding noisy repetitions
Xt_c_reps = Xt_c.copy()
yt_c_reps = yt_c.copy()
Xt_e_reps = Xt_e.copy()
yt_e_reps = yt_e.copy()

for i in range(4):
    Xt_c_reps = np.concatenate((Xt_c_reps, Xt_c))
    Xt_e_reps = np.concatenate((Xt_e_reps, Xt_e))
    np.random.seed(i)
    yt_c_reps = np.concatenate((yt_c_reps,
                                yt_c + 0.1*np.std(yt_c)*np.random.normal(size=yt_c.shape)))
    yt_e_reps = np.concatenate((yt_e_reps,
                                yt_e + 0.1*np.std(yt_e)*np.random.normal(size=yt_e.shape)))

plt.figure()
plt.plot(x, HF_function(x), '--C0', label='High Fidelity (HF)')
plt.plot(x, LF_function(x), '--C1', label='Low Fidelity (LF)')
plt.scatter(Xt_e_reps, yt_e_reps, color='C0', label='HF doe')
plt.scatter(Xt_c_reps, yt_c_reps, color='C1', label='LF doe')
plt.ylim(-10, 17); plt.xlim(-0.1, 1.1)
plt.xlabel('x'); plt.ylabel('y')
plt.legend();
```

```python
# %% Build the MFK object
yt_joint = np.concatenate((yt_c_reps, yt_e_reps))
theta0 = np.array([[0.5],[0.1]])
np.random.seed(7)
sm = MFK(theta0=theta0,
         noise0=noise0, eval_noise=True,
         use_het_noise = True,
         propagate_uncertainty = False, n_start=1)

# low-fidelity dataset names being integers from 0 to level-1
sm.set_training_values(Xt_c_reps, yt_c_reps, name=0)
# high-fidelity dataset without name
sm.set_training_values(Xt_e_reps, yt_e_reps)
# train the model
sm.train()

# test training
# HF
y = sm.predict_values(x)
var = sm.predict_variances(x)
# LF
y0 = sm._predict_intermediate_values(x, 1)
var0, _ =  sm.predict_variances_all_levels(x)
var0 = var0[:,0].reshape(-1,1)
```

    ___________________________________________________________________________

                                        MFK
    ___________________________________________________________________________

     Problem size

          # training points.        : 20

    ___________________________________________________________________________

     Training

       Training ...
       Training - done. Time (sec):  0.0481076
    ___________________________________________________________________________

     Evaluation

          # eval points. : 101

       Predicting ...
       Predicting - done. Time (sec):  0.0019996

       Prediction time/pt. (sec) :  0.0000198

```python
plt.figure(figsize=(7,5))
plt.plot(x, HF_function(x), '--C0', label='High Fidelity (HF)')
plt.plot(x, LF_function(x), '--C1', label='Low Fidelity (LF)')
plt.plot(x, y, 'C0', label='MFGP - HF')
plt.fill_between(np.ravel(x), np.ravel(y-3*np.sqrt(var)),
                 np.ravel(y+3*np.sqrt(var)),
                 color='C0',alpha=0.2 ,label ='Confidence Interval 99% HF')
plt.plot(x, y0, 'C1', label='MFGP - LF')
plt.fill_between(np.ravel(x), np.ravel(y0-3*np.sqrt(var0)),
                  np.ravel(y0+3*np.sqrt(var0)),
                  color='C1',alpha=0.2,label ='Confidence Interval 99% LF')
plt.scatter(Xt_e_reps, yt_e_reps, color='C0', label='HF doe with repetitions')
plt.scatter(Xt_c_reps, yt_c_reps, color='C1', label='LF doe with repetitions')
plt.xlabel('x'); plt.ylabel('y')
plt.ylim(-10, 17); plt.xlim(-0.1, 1.1)
plt.legend();
```

```python

```


# Two component Universe: Choice of equations

The equations for perturbations are not independent, so one can chose any of the Einstein equations to get the solutions. Here we will try to understand the relation in more detail.

**Warning** this is clearly a very hard problem. You can safely skip it, as many seasoned researchers working in the area do, but if you manage to solve it, you'll learn a lot about differential equations, Einstein equations, numerical mysteries.

In these notebook you will find the numerical solutions for the evolution of perturbations in a Universe filled with radiation and non-relativistic matter, solved using the energy momentum conservation equations and two different Einsein equations - either *'00'* or *'0i'* component. In both cases the initial conditions used were the lower order in $\eta$ adiabatic initial conditions. The mode was choen to enter the horion at matter domination (MD). You can see in the plots in the notebook, that the solutions are very different: the Newton potential $\Phi$ decreases to $9/10\Phi_{(i)}$ in one case (correct result), while _increases_ if another set of equations is used. At the same time it is expected that any of the Einstein equations should lead to the sae correct answer.

Can you spot what is the problem? *Hint:* look very carefully at the initial conditions using *all* the equations.

```python
from numpy import pi, sqrt
import numpy as np

from scipy.integrate import solve_ivp
from scipy.interpolate import interp1d

import matplotlib.pyplot as plt
%matplotlib inline
%config InlineBackend.figure_formats={'retina','png'}

from numba import jit
```

Define the equations.  Convention is that the vector of unknows `y[0:5]` contains $(a, \Phi, \delta_m, v_m, \delta_r, v_r)$.

By convention $a_0=1$ in the code.

Verison 1 of the equations: the Poisson equation (00) for gravity perturbations, and conservationequations:
$$\begin{align}
  -k^2{\Phi}-3\frac{a'}{a}{\Phi}'-3\frac{{a'}^2}{a^2}{\Phi}
    &= 4\pi Ga^2\left[ \rho_{{\mathrm{m}}}{\delta}_{{\mathrm{m}}}+\rho_{{\mathrm{r}}}{\delta}_{{\mathrm{r}}} \right],  \label{eq:pert2c1} \\
{\delta}_{{\mathrm{m}}}' - k^2v_{{\mathrm{m}}}                     &= 3{\Phi}',  \label{eq:pert2c2} \\
  v_{{\mathrm{m}}}' + \frac{a'}{a}v_{{\mathrm{m}}}                 &= -{\Phi}.   \label{eq:pert2c3} \\
  {\delta}_{{\mathrm{r}}}'-\frac{4}{3}k^2v_{{\mathrm{r}}}            &= 4{\Phi}',  \label{eq:pert2c4} \\
  v_{{\mathrm{r}}}' + \frac{1}{4}{\delta}_{{\mathrm{r}}}             &= -{\Phi}.   \label{eq:pert2c5}
\end{align}
$$
with $4\pi G a^2\rho_m=\frac{3}{2}H_0^2\Omega_m\frac{a_0^3}{a}$, $4\pi G a^2\rho_r=\frac{3}{2}H_0^2\Omega_r\frac{a_0^4}{a^2}$

Also, the Friedman equation for scale factor is used to get the exact background expansion
$$
a'/a = a H \equiv a H_0 \sqrt{\Omega_m\left(\frac{a_0}{a}\right)^3+\Omega_r\left(\frac{a_0}{a}\right)^4}
$$

```python
@jit
def eqns1(tau, Y, H0, OM, Or, k):
    a, Phi, dm, vm, dr, vr = Y
    Ha = a*H0*sqrt(OM/a**3+Or/a**4)
    k2 = k*k
    da = Ha*a
    dPhi = -(3/2*H0*H0*(dm*OM/a+dr*Or/(a*a))+(k2+3*Ha*Ha)*Phi)/3/Ha
    ddm = 3*dPhi+k2*vm
    dvm = -Phi-Ha*vm
    ddr = 4*dPhi+4/3*k2*vr
    dvr = -Phi-dr/4
    return np.array((da, dPhi, ddm, dvm, ddr, dvr))
```

Verison 2 of the equations: the (_0i_) equation for gravity perturbations, and conservationequations:
$$\begin{align}
  {\Phi}' + \frac{a'}{a}{\Phi}
    &= -4\pi Ga^2 (\rho_mv_m + \frac{4}{3}\rho_rv_r),
  \notag\\
{\delta}_{{\mathrm{m}}}' - k^2v_{{\mathrm{m}}}                     &= 3{\Phi}',  \label{eq:pert2c2} \\
  v_{{\mathrm{m}}}' + \frac{a'}{a}v_{{\mathrm{m}}}                 &= -{\Phi}.   \label{eq:pert2c3} \\
  {\delta}_{{\mathrm{r}}}'-\frac{4}{3}k^2v_{{\mathrm{r}}}            &= 4{\Phi}',  \label{eq:pert2c4} \\
  v_{{\mathrm{r}}}' + \frac{1}{4}{\delta}_{{\mathrm{r}}}             &= -{\Phi}.   \label{eq:pert2c5}
\end{align}
$$

```python
@jit
def eqns2(tau, Y, H0, OM, Or, k):
    a, Phi, dm, vm, dr, vr = Y
    Ha = a*H0*sqrt(OM/a**3+Or/a**4)
    k2 = k*k
    da = Ha*a
    dPhi = -Ha*Phi-(3/2*H0*H0*a*a*(OM/a**3*vm+4/3*Or/a**4*vr))
    ddm = 3*dPhi+k2*vm
    dvm = -Phi-Ha*vm
    ddr = 4*dPhi+4/3*k2*vr
    dvr = -Phi-dr/4
    return np.array((da, dPhi, ddm, dvm, ddr, dvr))
```

These are adiabatic initial condition to the lowest order in $\eta$
$$
\begin{align}
{\delta}_{m}=\frac{3}{4}{\delta}_r=
    -\frac{3}{2}{\Phi}&=1,\\
    v_{m}=v_r=\frac{1}{3}{\eta}
\end{align}
$$
And early time scale factor
$$
a = H_0\sqrt{\Omega_r}\eta
$$
(we neglect the change of the number of degrees of freedom for simplicity)

```python
def adiabatic_ic(eta,k):
    a = H0*eta*sqrt(Or)
    Phi = -2/3
    dm = -3/2*Phi
    dr = -2*Phi
    vm = -Phi/2*eta
    vr = -Phi/2*eta
    return np.array((a, Phi, dm, vm, dr, vr))
```

These are the service functions that solve the differential equations from eta1 to eta2

```python
def solve(eta1, eta2, ic):
    return solve_ivp(eqns1, (eta1, eta2), ic, args=(H0, OM, Or, k), rtol=1e-8, atol=1e-10)
```

```python
def solve2(eta1, eta2, ic):
    return solve_ivp(eqns2, (eta1, eta2), ic, args=(H0, OM, Or, k), rtol=1e-8, atol=1e-10)
```

## Mode entering at MD

Let us define the Universe with 'reasonably realistic' parameters — a small amount of radiation today, and select the mode _k_ that enters horizon at MD

```python
k=0.0008

H0=0.7*100*3.33564e-6 # 1/Mpc
Or=1e-4
OM=1-Or
```

Let us solve the two versions of the equations

```python
eta1 = 1
eta2 = 18000
res = solve(eta1, eta2, adiabatic_ic(eta1,k))
res2 = solve2(eta1, eta2, adiabatic_ic(eta1,k))
```

You may safely ignore this -- numerically find the RD/MD equality and crossing moments of time, used only for illustration later

```python
eta_at_rho_m_over_r = interp1d(OM/(Or/res.y[0,:]),res.t)
eta_at_crossing = interp1d(H0*np.sqrt(OM/res.y[0,:]**3+Or/res.y[0,:]**4)/(k/res.y[0,:]), res.t)
eta_eq = eta_at_rho_m_over_r(1.)
eta_x = eta_at_crossing(1.)
```

So, why the two solutions do not coincide?!

```python
plt.semilogx(res.t,-res.y[1,:], label=r"$-\Phi$ equation set 1")
plt.semilogx(res2.t,-res2.y[1,:], '--', label=r"$-\Phi$ equation set 2")

plt.axhline(0, color='black', lw=1)
plt.axhline(2/3, color='black', lw=1)
plt.axhline(2/3*9/10, color='black', lw=1)
plt.axvline(eta_eq, color='red', label="RDMD equality")
plt.axvline(eta_x, color="green", label="horizon crossing")
plt.xlabel(r"$a_0\eta$")
plt.legend()
```

This one is a log-log plot just to reassure you that teh matter perturbations grow after the horizon crossing (while the Newton potanital form the last plot does not really change after the horizon crossing).

```python
plt.loglog(res.t,res.y[2,:], label=r"$\delta_m$ equation set 1")
plt.loglog(res2.t,res2.y[2,:], '--', label=r"$\delta_m$ equation set 2")

plt.axhline(0, color='black', lw=1)
plt.axvline(eta_eq, color='red', label="RDMD equality")
plt.axvline(eta_x, color="green", label="horizon crossing")
plt.xlabel(r"$a_0\eta$")
plt.legend()
```

```python

```


# Unit 4 - Causal Inference

We are going to combine our knowledge of regression and classification methods to estimate causal effects in observational data.

We will play around with two datasets of two different kinds:
- Average effect (sodium)
- Heterogeneous effect (toy dataset)

Steps we are going to take:
1. Importing packages.
2. Average effect:
  1. Problem description.
  2. Evaluation metrics.
  3. Loading and exploring the data.
  4. Data pre-processing.
  5. Training the models.
  6. Performance evaluation.
3. Heterogeneous effect:
  1. Repeat the same steps as for the average effect.

## Step 1 - packages

Google Colab has loads of pre-installed packages already. What if we need more? We can install additional packages by knowing that:
- The exclamation mark ('!') gives us access to the Linux command line that runs behind our notebook.
- We can access the default python's package manager with the command 'pip'.

In our case, we are interested in using the [EconML](https://econml.azurewebsites.net/index.html) package, which is not part of the default Colab environment. We can install it via the following command:

```python
!pip install econml
```

    Collecting econml
      Downloading econml-0.13.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.3 MB)
         |████████████████████████████████| 3.3 MB 25.7 MB/s
    Collecting sparse
      Downloading sparse-0.13.0-py2.py3-none-any.whl (77 kB)
         |████████████████████████████████| 77 kB 6.3 MB/s
    Requirement already satisfied: scipy>1.4.0 in /usr/local/lib/python3.7/dist-packages (from econml) (1.4.1)
    Collecting shap<0.40.0,>=0.38.1
      Downloading shap-0.39.0.tar.gz (356 kB)
         |████████████████████████████████| 356 kB 41.5 MB/s
    Requirement already satisfied: joblib>=0.13.0 in /usr/local/lib/python3.7/dist-packages (from econml) (1.1.0)
    Requirement already satisfied: scikit-learn>0.22.0 in /usr/local/lib/python3.7/dist-packages (from econml) (1.0.2)
    Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from econml) (1.21.5)
    Requirement already satisfied: statsmodels>=0.10 in /usr/local/lib/python3.7/dist-packages (from econml) (0.10.2)
    Requirement already satisfied: lightgbm in /usr/local/lib/python3.7/dist-packages (from econml) (2.2.3)
    Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from econml) (1.3.5)
    Collecting dowhy<0.7
      Downloading dowhy-0.6-py3-none-any.whl (123 kB)
         |████████████████████████████████| 123 kB 48.3 MB/s
    Collecting pydot>=1.4
      Downloading pydot-1.4.2-py2.py3-none-any.whl (21 kB)
    Requirement already satisfied: sympy>=1.4 in /usr/local/lib/python3.7/dist-packages (from dowhy<0.7->econml) (1.7.1)
    Requirement already satisfied: networkx>=2.0 in /usr/local/lib/python3.7/dist-packages (from dowhy<0.7->econml) (2.6.3)
    Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas->econml) (2.8.2)
    Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas->econml) (2018.9)
    Requirement already satisfied: pyparsing>=2.1.4 in /usr/local/lib/python3.7/dist-packages (from pydot>=1.4->dowhy<0.7->econml) (3.0.7)
    Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.7.3->pandas->econml) (1.15.0)
    Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn>0.22.0->econml) (3.1.0)
    Requirement already satisfied: tqdm>4.25.0 in /usr/local/lib/python3.7/dist-packages (from shap<0.40.0,>=0.38.1->econml) (4.63.0)
    Collecting slicer==0.0.7
      Downloading slicer-0.0.7-py3-none-any.whl (14 kB)
    Requirement already satisfied: numba in /usr/local/lib/python3.7/dist-packages (from shap<0.40.0,>=0.38.1->econml) (0.51.2)
    Requirement already satisfied: cloudpickle in /usr/local/lib/python3.7/dist-packages (from shap<0.40.0,>=0.38.1->econml) (1.3.0)
    Requirement already satisfied: patsy>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from statsmodels>=0.10->econml) (0.5.2)
    Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.7/dist-packages (from sympy>=1.4->dowhy<0.7->econml) (1.2.1)
    Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from numba->shap<0.40.0,>=0.38.1->econml) (57.4.0)
    Requirement already satisfied: llvmlite<0.35,>=0.34.0.dev0 in /usr/local/lib/python3.7/dist-packages (from numba->shap<0.40.0,>=0.38.1->econml) (0.34.0)
    Building wheels for collected packages: shap
      Building wheel for shap (setup.py) ... done
      Created wheel for shap: filename=shap-0.39.0-cp37-cp37m-linux_x86_64.whl size=491660 sha256=27ca490e8eb2bc1a76aee810c37643e1b95513224e296551e2f4765379ab7ee5
      Stored in directory: /root/.cache/pip/wheels/ca/25/8f/6ae5df62c32651cd719e972e738a8aaa4a87414c4d2b14c9c0
    Successfully built shap
    Installing collected packages: slicer, pydot, sparse, shap, dowhy, econml
      Attempting uninstall: pydot
        Found existing installation: pydot 1.3.0
        Uninstalling pydot-1.3.0:
          Successfully uninstalled pydot-1.3.0
    Successfully installed dowhy-0.6 econml-0.13.0 pydot-1.4.2 shap-0.39.0 slicer-0.0.7 sparse-0.13.0

Note this command has to be executed everytime you restart the notebook.

Note to people using local Jupyter notebooks - you will have to execute the pip command in your local environment. It may not work from the notebook itself.

Now, we can import anything we need from EconML, together with the rest of the packages we need.

```python
from econml.metalearners import XLearner
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import numpy as np
import pandas as pd
import scipy.stats as st
import matplotlib.pyplot as plt
```

    /usr/local/lib/python3.7/dist-packages/statsmodels/tools/_testing.py:19: FutureWarning: pandas.util.testing is deprecated. Use the functions in the public API at pandas.testing instead.
      import pandas.util.testing as tm

## Step 2 - Average effect

### Step 2.1 - problem setting

This is a simplified version of a model that simulates the effect of sodium intake on blood pressure. Officially proposed in ([Luque-Fernandez et al. 2019](https://academic.oup.com/ije/article/48/2/640/5248195)).

We are given three covariates: age (A), sodium (S), and blood pressure (B). We are interested in the effect of sodium intake on blood pressure. In addition, we know that age affects both variables - confounder. Here is the assumed causal graph:

The code that implements the simulation can be found in [sodium_generate.py](https://github.com/dmachlanski/CE888_2022/blob/main/lab4/data/sodium_generate.py) script. This data generation process can be also described as follows:

$$A = \mathcal{N}(65, 5)$$ \\
$$
S=
    \begin{cases}
        1 & \text{ if } A/18 + \mathcal{N}(0, 1) > 3.5 \\
        0 & \text{ otherwise }
    \end{cases}
$$ \\
$$B = ATE \times S + 2 A + \mathcal{N}(0, 1)$$ \\
$$ATE = 1.05$$

Note we assume the same effect across the entire population and fix it to an arbitrary scalar.

With this setting, we generate 10,000 samples and obtain file [sodium_10k.npz](https://github.com/dmachlanski/CE888_2022/blob/main/lab4/data/sodium_10k.npz).

Because we know the true ATE used in the data generation process, we can evaluate our models' performance by measuring how close their predicted ATE ($\widehat{ATE}$) is to the true one. More formally, we can write ATE as:

$$ATE = \mathbb{E}[\mathcal{Y}_1 - \mathcal{Y}_0]$$

And predicted ATE as:

$$\widehat{ATE} = \frac{1}{n}\sum \limits_{i=1}^{n}(\hat{y}_1^{(i)} - \hat{y}_0^{(i)})$$

Where $\hat{y}_t^{(i)}$ denotes predicted potential outcome for treatment $t$ and individual $(i)$. We can thus see that ATE is essentially the average of individual treatment effects, that is:

$$\widehat{ATE} = \frac{1}{n}\sum \limits_{i=1}^{n}\widehat{ITE}^{(i)}$$

One commonly used metric that we will use here measures the absolute difference between the predicted and true ATE, that is:

$$\epsilon_{ATE} = \left| \widehat{ATE} - ATE \right|$$

This metric works properly when the same average effect applies to all units in the dataset. In some cases, effects can vary across groups of units. In other words, we can say that the effect is heterogeneous. For those cases, it is useful to also include another metric that captures the errors made per each individual predicted ITE. This metric is PEHE:

$$\epsilon_{PEHE}=\sqrt{\frac{1}{n}\sum_{i=1}^{n}( \hat{y}_1^{(i)}  - \hat{y}_0^{(i)} - (\mathcal{Y}_1^{(i)} - \mathcal{Y}_0^{(i)}))^2}$$

Which can be simplified to:

$$\epsilon_{PEHE}=\sqrt{\frac{1}{n}\sum_{i=1}^{n}( \widehat{ITE}^{(i)} - ITE^{(i)})^2}$$

### Step 2.2 - metrics

Here, we are going to implement the two metrics discussed in the previous step 2.1: $\epsilon_{ATE}$ and $\epsilon_{PEHE}$.

You should have already implemented both metrics as part of the Moodle quiz. Feel free to copy the code from your quiz answers or write the code from scratch.

```python
def abs_ate(effect_true, effect_pred):
    """
    Absolute error for the Average Treatment Effect (ATE)
    :param effect_true: true treatment effect value
    :param effect_pred: predicted treatment effect value
    :return: absolute error on ATE
    """
    # YOUR CODE HERE
    return np.abs(np.mean(effect_pred) - np.mean(effect_true))

def pehe(effect_true, effect_pred):
    """
    Precision in Estimating the Heterogeneous Treatment Effect (PEHE)
    :param effect_true: true treatment effect value
    :param effect_pred: predicted treatment effect value
    :return: PEHE
    """
    # YOUR CODE HERE
    return np.sqrt(np.mean((effect_true - effect_pred)**2))

```

### Step 2.3 - data

The sodium dataset is in the form of numpy's '.npz' file. Thus, we cannot use the usual 'read_csv' from pandas. This file must be downloaded first and loaded locally.

We can do this by either manually uploading the file to Colab, or via the following command that downloads the file to Colab automatically:

```python
!wget https://github.com/dmachlanski/CE888_2022/raw/main/lab4/data/sodium_10k.npz
```

    --2022-03-31 19:51:55--  https://github.com/dmachlanski/CE888_2022/raw/main/lab4/data/sodium_10k.npz
    Resolving github.com (github.com)... 140.82.114.4
    Connecting to github.com (github.com)|140.82.114.4|:443... connected.
    HTTP request sent, awaiting response... 302 Found
    Location: https://raw.githubusercontent.com/dmachlanski/CE888_2022/main/lab4/data/sodium_10k.npz [following]
    --2022-03-31 19:51:55--  https://raw.githubusercontent.com/dmachlanski/CE888_2022/main/lab4/data/sodium_10k.npz
    Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...
    Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.
    HTTP request sent, awaiting response... 200 OK
    Length: 320962 (313K) [application/octet-stream]
    Saving to: ‘sodium_10k.npz’

    sodium_10k.npz      100%[===================>] 313.44K  --.-KB/s    in 0.006s

    2022-03-31 19:51:55 (47.7 MB/s) - ‘sodium_10k.npz’ saved [320962/320962]

You can check in the panel on the left that the file is there.

Now, let's load the file with numpy's 'load' function

```python
data = np.load('sodium_10k.npz')
```

We can see what variable names are available by examining the 'files' property.

```python
data.files
```

    ['x', 't', 'y', 'ite']

Variable names:
- X -> A (age; background variable)
- T -> S (sodium intake; treatment variable)
- Y -> B (blood pressure; outcome variable)
- ITE -> individual treatment effect

```python
for f in data.files:
  print(f'{f}: {data[f].shape}')
```

    x: (10000,)
    t: (10000,)
    y: (10000,)
    ite: (10000,)

As we can see above, we have four single-dimensional vectors, each of which consisting of 10,000 samples. We are going to reshape them to 2D vectors as some methods prefer this representation of data, though we will have to switch back and forth between 1D and 2D representation often as this requirement varies across methdos.

```python
X = data['x'].reshape(-1, 1)
T = data['t'].reshape(-1, 1)
Y = data['y'].reshape(-1, 1)
ite = data['ite']
```

Let's plot the distributions of all three.

```python
bins=20
fig, axs = plt.subplots(1, 4, figsize=(16, 4))
axs[0].hist(X, bins=bins)
axs[1].hist(T, bins=bins)
axs[2].hist(Y, bins=bins)
axs[3].hist(ite, bins=bins)
plt.show()
```

X and Y clearly follow the shape of the normal distribution, as expected. T is binary, i.e., consists of 0 and 1 values. All ITE values are slightly higher than 1.0, which is correct as we know the ATE is 1.05.

```python
np.mean(ite)
```

Let's do some scatter plots as well to explore the data. We will limit ourselves to just a fraction of the data for visualisation purposes, but feel free to play around with the 'limit' number.

```python
limit = 100
plt.scatter(X[T == 1][:limit], Y[T == 1][:limit], c = "r", label = "Treated")
plt.scatter(X[T == 0][:limit], Y[T == 0][:limit], c = "g", label = "Control")
plt.legend(ncol=2)
```

The effect is difficult to notice as it is quite small (1.05), nevertheless, in some regions the treated units (red) are a bit higher than the control ones (gree), which demonstrates the effect (treated units have higher outcomes).

### Step 2.4 - data pre-processing

Split the data into training and testing for evaluation purposes with 80/20 ratio.

```python
x_train, x_test, t_train, t_test, y_train, y_test, ite_train, ite_test = train_test_split(X, T, Y, ite, test_size=0.2)

scaler_x = StandardScaler()
x_train = scaler_x.fit_transform(x_train)
x_test = scaler_x.transform(x_test)
```

### Step 2.5 - training

We are going to use three estimators here:
- Random forest (RF; regressor).
- Random forest with Inverse Propensity Weighting (IPW).
- X-learner with RF as base learners.

#### Random Forest

##### Description

With random forests we simply model $P(Y|X, T)$, that is, we regress Y on X and T. Once trained, the model predicts potential outcomes $\hat{y}_t^{(i)}$. Thus, to obtain ITE predictions for each individual, we need to predict both $\hat{y}_0$ and $\hat{y}_1$, that is, predict the outcomes for both control and treated settings for each individual.

By calculating the difference between the treated and control outcome, we can obtain the effect per given individual. In other words:

$$\widehat{ITE}^{(i)} = \hat{y}_1^{(i)} - \hat{y}_0^{(i)}$$

##### Usage

```python
# Random Forest with default parameters.
rf = RandomForestRegressor()

# Train on the training data.
# In the supervised setting, our usual X here consists of X and T.
# The target is Y - the outcome.
# Input: [X, T], output: Y.
# Hint: use np.concatenate or np.hstack functions to merge X and T.

xt_train = np.concatenate([x_train,t_train], axis = 1) # YOUR CODE HERE
rf.fit(xt_train, y_train.flatten())

# Predictions
# Note we set T to a specific value for ALL individuals.
# These are interventional distributions - P(Y|X, do(T=t)).

# Predict Y_0 given [X, 0]
# Hint: use np.zeros_like or np.zeros functions to get the zeros for all units.
# Hint 2: you still have to merge X and zeros (see training above).
xt0_test = np.concatenate([x_test,np.zeros_like(t_test)], axis = 1) # YOUR CODE HERE
rf_y0_test = rf.predict(xt0_test)

# Predict Y_1 given [X, 1]
# Hint: use np.ones_like or np.ones functions to get the ones for all units.
# Hint 2: you still have to merge X and ones (see training above).
xt1_test = np.concatenate([x_test,np.ones_like(t_test)], axis = 1) # YOUR CODE HERE
rf_y1_test = rf.predict(xt1_test)

# Compute ITEs
# ITE = Y_1 - Y_0
rf_te_test = rf_y1_test - rf_y0_test # YOUR CODE HERE
```

#### Random Forest with IPW

##### Description

This is a similar approach to the previous one, but extends it further with the Inverse Propensity Weighting (IPW). We proceed as follows:

1. Use random forest classifier to model unit's probability of receiving the treatment, that is, $P(t_i|x_i)$. Input: X, target: T. Note this is a classic binary classification problem. $P(t_i|x_i)$ is called a **propensity score**.
2. Use trained classifier to predict propensity scores for each individual. Compute weights per each sample as:

$$w_i = \frac{t_i}{P(t_i|x_i)} + \frac{1-t_i}{1-P(t_i|x_i)}$$

Intuition: Dominant group gets smaller weights. In practice, gives more importance to treated samples.

3. Use sample weights in random forest regressor training.
4. Proceed with predictions as in the previous section.

We can compute the sample weights with the following function:

```python
def get_ps_weights(clf, x, t):
  ti = np.squeeze(t)
  clf.fit(x, ti)
  ptx = clf.predict_proba(x).T[1].T + 0.0001 # add a small value to avoid dividing by 0
  # Given ti and ptx values, compute the weights wi (see formula above):
  wi = (ti/ptx) + (1 - ti)/(1 - ptx)# YOUR CODE HERE
  return wi
```

##### Usage

```python
# Get the sample weights
prop_clf = RandomForestClassifier()
weights = get_ps_weights(prop_clf, x_train, t_train)

# Train the regressor
rf_ipsw = RandomForestRegressor()

rf_ipsw.fit(xt_train, y_train.flatten(), sample_weight=weights) #pass the weights

# Make predictions
rf_ipsw_y0_test = rf_ipsw.predict(xt0_test) # YOUR CODE HERE
rf_ipsw_y1_test = rf_ipsw.predict(xt1_test) # YOUR CODE HERE

# Compute ITEs
rf_ipsw_te_test = rf_ipsw_y1_test - rf_ipsw_y0_test # YOUR CODE HERE
```

#### X-learner

##### Formal description

A meta-learner implemented via EconML. Uses provided regressors and classifiers to solve smaller sub-problems. Models the effect directly instead of the outcomes, predicting ITEs as a consequence.

Originally introduced in [(Künzel et al. 2019)](http://arxiv.org/abs/1706.03461). The modelling process of the X-learner can be divided into three stages.

**Stage 1**

Use provided regressors to model $\mathcal{Y}_0$ and $\mathcal{Y}_1$ separately. More formally, the response functions are:

$$\mu_0(x) = \mathbb{E}[\mathcal{Y}_0|X=x]$$
$$\mu_1(x) = \mathbb{E}[\mathcal{Y}_1|X=x]$$

We denote estimated functions as $\hat{\mu}_0$ and $\hat{\mu}_1$.

**Stage 2**

Define imputed treatment effects as:

$$\mathcal{D}_0^{(i)} = \hat{\mu}_1(X_0^{(i)}) - \mathcal{Y}_0^{(i)}$$
$$\mathcal{D}_1^{(i)} = \mathcal{Y}_1^{(i)} - \hat{\mu}_0(X_1^{(i)})$$

Use provided regressors to model $\mathcal{D}_0$ and $\mathcal{D}_1$ separately. The response functions are formally defined as:

$$\tau_0(x) = \mathbb{E}[\mathcal{D}_0|X=x]$$
$$\tau_1(x) = \mathbb{E}[\mathcal{D}_1|X=x]$$

We denote estimated functions as $\hat{\tau}_0$ and $\hat{\tau}_1$.

**Stage 3**

The final treatment effect estimate is a weighted average of the two estimates from Stage 2:

$$\hat{\tau}(x) = g(x)\hat{\tau}_0(x) + (1 - g(x))\hat{\tau}_1(x)$$

Where $g \in [0, 1]$ is a weight function. In practice, $g$ can be modelled as a propensity score function $e$, formally written as:

$$e(x) = \mathbb{E}[T|X=x]$$

Using a provided classifier, we can obtain an estimate $\hat{e}$ that can be used in place of $g$. That is:

$$\hat{\tau}(x) = \hat{e}(x)\hat{\tau}_0(x) + (1 - \hat{e}(x))\hat{\tau}_1(x)$$

In summary, we perform the following steps:
1. Training (inputs: X, T, Y; outputs: $\hat{\tau}_0$, $\hat{\tau}_1$, $\hat{e}$):
  1. Model the outcomes with a provided regressor class and obtain estimators $\hat{\mu}_0$ and $\hat{\mu}_1$.
  2. Make predictions $\hat{\mu}_1(X_0^{(i)})$ and $\hat{\mu}_0(X_1^{(i)})$.
  3. Compute imputed treatment effects $\mathcal{D}_0$ and $\mathcal{D}_1$.
  4. Model the imputed treatment effects with a provided regressor class and obtain estimators $\hat{\tau}_0$ and $\hat{\tau}_1$.
  5. Model the propensity scores with a provided classifier class and obtain an estimator $\hat{e}$.
2. Prediction (inputs: X; outputs: $\widehat{ITE}$):
  1. Use estimator $\hat{\tau}$ to obtain individual treatment effect estimates.

##### Usage

We continue with random forests here by providing the X-learner with RF regressor and RF classifier as base learners. Note it distinguishes X from T in its input - it is necessary to learn the propensity score estimator. Also, the final prediction does not depend on T! Though it does indirectly by predicting the propensity score internally.

```python
# X-Learner
xl = XLearner(models=RandomForestRegressor(), propensity_model=RandomForestClassifier())
xl.fit(y_train, t_train.flatten(), X=x_train)

xl_te_test = xl.effect(x_test)
```

    Function get_feature_names is deprecated; get_feature_names is deprecated in 1.0 and will be removed in 1.2. Please use get_feature_names_out instead.

Just 3 lines of code for such a complex estimator!

### Step 2.6 - evaluation

We will examine our estimators' performance from different perspectives:
- Metrics: $ϵ_{ATE}$ and $ϵ_{PEHE}$
- Confidence intervals of predicted ATEs
- Visualisations of predicted effects

#### Metrics

```python
# Error on ATE
# Use our 'abs_ate' function, true ITEs and predicted ITEs to get the measurements.
rf_ate_test = abs_ate(ite_test,rf_te_test) # YOUR CODE HERE
rf_ipsw_ate_test = abs_ate(ite_test,rf_ipsw_te_test) # YOUR CODE HERE
xl_ate_test = abs_ate(ite_test,xl_te_test) # YOUR CODE HERE

# PEHE
# Use our 'pehe' function, true ITEs and predicted ITEs to get the measurements.
rf_pehe_test = pehe(ite_test,rf_te_test) # YOUR CODE HERE
rf_ipsw_pehe_test = pehe(ite_test,rf_ipsw_te_test) # YOUR CODE HERE
xl_pehe_test = pehe(ite_test,xl_te_test) # YOUR CODE HERE
```

```python
results = []
results.append(['RF', rf_ate_test, rf_pehe_test])
results.append(['RF (IPW)', rf_ipsw_ate_test, rf_ipsw_pehe_test])
results.append(['XL', xl_ate_test, xl_pehe_test])

cols = ['Method', 'ATE test', 'PEHE test']

df = pd.DataFrame(results, columns=cols)
df
```

      Method
      ATE test
      PEHE test

      0
      RF
      0.013131
      0.953157

      1
      RF (IPW)
      0.010231
      0.944660

      2
      XL
      0.038464
      0.848420

      <button class="colab-df-convert" onclick="convertToInteractive('df-a2f71e5a-ee41-45d5-9db4-6a8b73fc7939')"
              title="Convert this dataframe to an interactive table."
              style="display:none;">

  <svg xmlns="http://www.w3.org/2000/svg" height="24px"viewBox="0 0 24 24"
       width="24px">
    <path d="M0 0h24v24H0V0z" fill="none"/>
    <path d="M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z"/><path d="M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z"/>
  </svg>
      </button>

#### Confidence intervals

```python
def mean_ci(data, ci=0.95):
  l_mean = np.mean(data)
  lower, upper = st.t.interval(ci, len(data)-1, loc=l_mean, scale=st.sem(data))
  return l_mean, lower, upper
```

```python
rf_ate_bounds = mean_ci(rf_te_test)
rf_ipsw_ate_bounds = mean_ci(rf_ipsw_te_test)
xl_ate_bounds = mean_ci(xl_te_test)

results = []
results.append(['RF', rf_ate_bounds[0], rf_ate_bounds[1], rf_ate_bounds[2]])
results.append(['RF (IPW)', rf_ipsw_ate_bounds[0], rf_ipsw_ate_bounds[1], rf_ipsw_ate_bounds[2]])
results.append(['XL', xl_ate_bounds[0], xl_ate_bounds[1], xl_ate_bounds[2]])

cols = ['Method', 'ATE mean', 'CI lower', 'CI upper']

df = pd.DataFrame(results, columns=cols)
df
```

      Method
      ATE mean
      CI lower
      CI upper

      0
      RF
      1.036869
      0.995064
      1.078674

      1
      RF (IPW)
      1.039769
      0.998335
      1.081202

      2
      XL
      1.088464
      [1.0512874657096427]
      [1.1256405485575294]

      <button class="colab-df-convert" onclick="convertToInteractive('df-b1ba055b-2dbd-4299-a43a-60b1009a8395')"
              title="Convert this dataframe to an interactive table."
              style="display:none;">

  <svg xmlns="http://www.w3.org/2000/svg" height="24px"viewBox="0 0 24 24"
       width="24px">
    <path d="M0 0h24v24H0V0z" fill="none"/>
    <path d="M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z"/><path d="M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z"/>
  </svg>
      </button>

#### Visualisations

```python
plt.figure()
plt.boxplot([rf_te_test, rf_ipsw_te_test, xl_te_test.flatten()], labels=['RF', 'RF (IPW)', 'X-learner'])
plt.ylabel('Treatment Effect')

plt.show()
```

```python
plt.figure(figsize=(12, 10))
m_size = 10
plt.scatter(x_test, rf_te_test, label="RF", s=m_size)
plt.scatter(x_test, rf_ipsw_te_test, label="RF (IPW)", s=m_size)
plt.scatter(x_test, xl_te_test, label="X-learner", s=m_size)
plt.xlabel('X')
plt.ylabel('Treatment Effect')
plt.legend()
plt.show()
```

## Step 3 - Heterogeneous effect

In the previous dataset, we observed the same effect value across all units. Such homogeneity does not always apply. In the second part of this exercise, we are going to explore a toy data where the effect is cleary heterogeneous.

### Step 3.1 - Problem setting

This dataset was generated via the [hte_generate.py](https://github.com/dmachlanski/CE888_2022/blob/main/lab4/data/hte_generate.py) script. Similarly to the previous dataset, the usual 'triangle' causal graph applies here as well.

Here, we are interested in estimating the effect of T on Y, while also knowing that X affects both covariates. In addition, we know there is more than one effect group, which will become apparent once we start exploring the data visually.

### Step 3.2 - Metrics

We can use the same metrics as in the previous dataset ($ϵ_{ATE}$ and $ϵ_{PEHE}$), so there is no need to implement more - just reuse the existing code.

### Step 3.3 - Load the data

The data can be found [here](https://github.com/dmachlanski/CE888_2022/tree/main/lab4/data). It has been already split into training and testing sets (hte_train.csv and hte_test.csv).

```python
# Load the dataset
# 1. Load hte_train.csv from: https://raw.githubusercontent.com/dmachlanski/CE888_2022/main/lab4/data/hte_train.csv
# 2. Load hte_test.csv from: https://raw.githubusercontent.com/dmachlanski/CE888_2022/main/lab4/data/hte_test.csv

# YOUR CODE HERE
df_train = pd.read_csv("https://raw.githubusercontent.com/dmachlanski/CE888_2022/main/lab4/data/hte_train.csv")
df_test = pd.read_csv("https://raw.githubusercontent.com/dmachlanski/CE888_2022/main/lab4/data/hte_test.csv")
```

```python
df_train.head()
```

      x
      t
      y
      ite

      0
      0.023976
      1.0
      2.429799
      2.3

      1
      -0.094244
      1.0
      2.363082
      2.3

      2
      -0.020031
      1.0
      2.572256
      2.3

      3
      0.006930
      1.0
      2.605929
      2.3

      4
      1.156410
      1.0
      0.346002
      -3.7

      <button class="colab-df-convert" onclick="convertToInteractive('df-52dd8b3f-a006-41b8-9e2e-e7fa371853fc')"
              title="Convert this dataframe to an interactive table."
              style="display:none;">

  <svg xmlns="http://www.w3.org/2000/svg" height="24px"viewBox="0 0 24 24"
       width="24px">
    <path d="M0 0h24v24H0V0z" fill="none"/>
    <path d="M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z"/><path d="M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z"/>
  </svg>
      </button>

```python
df_test.head()
```

      x
      t
      y
      ite

      0
      0.012826
      1.0
      2.492201
      2.3

      1
      0.041816
      1.0
      2.388398
      2.3

      2
      -0.096298
      1.0
      2.568154
      2.3

      3
      -0.027440
      1.0
      2.649401
      2.3

      4
      0.063349
      1.0
      2.514691
      2.3

      <button class="colab-df-convert" onclick="convertToInteractive('df-cff2e035-f09e-4b9b-bf04-7921d2d45c29')"
              title="Convert this dataframe to an interactive table."
              style="display:none;">

  <svg xmlns="http://www.w3.org/2000/svg" height="24px"viewBox="0 0 24 24"
       width="24px">
    <path d="M0 0h24v24H0V0z" fill="none"/>
    <path d="M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z"/><path d="M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z"/>
  </svg>
      </button>

Now explore the data a bit. You can follow exactly the same process as with the sodium dataset - make histograms and scatter plots. Feel free to reuse the code from the previous parts of the exercise.

```python
# Scatter plots - feel free to reuse the code from the sodium example.
# Hint: use DataFrame.loc to access treated and control units, or store the columns as numpy arrays.

# YOUR CODE HERE
x_train = df_train['x'].values.reshape(-1, 1)
t_train = df_train['t'].values.reshape(-1, 1)
y_train = df_train['y'].values.reshape(-1, 1)
ite_train = df_train['ite']

```

```python
# Histograms - again, feel free to reuse the pre-existing code.

# YOUR CODE HERE
limit = 100
plt.scatter(x_train[t_train == 1][:limit], y_train[t_train == 1][:limit], c = "r", label = "Treated")
plt.scatter(x_train[t_train == 0][:limit], y_train[t_train == 0][:limit], c = "g", label = "Control")
plt.legend(ncol=2)
```

```python
bins=20
fig, axs = plt.subplots(1, 4, figsize=(16, 4))
axs[2].hist(y_train, bins=bins)
axs[3].hist(ite_train, bins=bins)
axs[1].hist(t_train, bins=bins)
axs[0].hist(x_train, bins=bins)
plt.show()
```

Once you do the plots, it should be quite obvious by now that we have two heterogeneous effect groups: ITE is either -3.7 or 2.3.

### Step 3.4 Data pre-processing

There is little (if any) data pre-processing necessary here. The dataset is already divided into training and testing. The only part left is scaling, depending on the type of models you plan to use.

```python
# YOUR CODE HERE
x_test = df_test['x'].values.reshape(-1, 1)
t_test = df_test['t'].values.reshape(-1, 1)
y_test = df_test['y'].values.reshape(-1, 1)
ite_test = df_test['ite']
scaler_x = StandardScaler()
x_train = scaler_x.fit_transform(x_train)
x_test = scaler_x.transform(x_test)
```

### Step 3.5 Training

This is the modelling part. You should have the following types of models:
- Regressor (scikit-learn)
- **Weighted** regressor (IPW, scikit-learn)
- [CATE](https://econml.azurewebsites.net/reference.html#cate-estimators) estimator from EconML

For each of the above types, choose the models you like - it is up to you!

Steps:
1. Train the models on training data.
2. Predict outcomes and calculate effects, or predict effects directly, depending on the method used, on test data.

```python
# YOUR CODE HERE (add new code cells as necessary)
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
```

```python
random = RandomForestRegressor()

xt_train = np.concatenate([x_train,t_train], axis = 1)
random.fit(xt_train, y_train.flatten())

xt0_test = np.concatenate([x_test,np.zeros_like(t_test)], axis = 1)
random_y0_test = random.predict(xt0_test)
xt1_test = np.concatenate([x_test,np.ones_like(t_test)], axis = 1)
random_y1_test = random.predict(xt1_test)
random_te_test = random_y1_test - random_y0_test

prop_random = RandomForestClassifier()
weights = get_ps_weights(prop_random, x_train, t_train)

random_ipsw = RandomForestRegressor()
random_ipsw.fit(xt_train, y_train.flatten(), sample_weight=weights)
random_ipsw_y0_test = random_ipsw.predict(xt0_test)
random_ipsw_y1_test = random_ipsw.predict(xt1_test)
random_ipsw_te_test = random_ipsw_y1_test - random_ipsw_y0_test
```

X-**Learner**

```python
xl = XLearner(models=RandomForestRegressor(), propensity_model=RandomForestClassifier())
xl.fit(y_train, t_train.flatten(), X=x_train)
xl_te_test = xl.effect(x_test)
```

    Function get_feature_names is deprecated; get_feature_names is deprecated in 1.0 and will be removed in 1.2. Please use get_feature_names_out instead.

### Step 3.6 Evaluation

Here, use both metrics $ϵ_{ATE}$ and $ϵ_{PEHE}$ to evaluate the performance of your models. Feel free to perform any other form of evaluation as we did before (confidence intervals, box plots, etc.).

```python
# Evaluate your predictions with 'abs_ate' and 'pehe' functions.

# Error on ATE
rf_ate_test = abs_ate(ite_test,random_te_test)
rf_ipsw_ate_test = abs_ate(ite_test,random_ipsw_te_test)
xl_ate_test = abs_ate(ite_test,xl_te_test)

# Error on PEHE
rf_pehe_test = pehe(ite_test,random_te_test)
rf_ipsw_pehe_test = pehe(ite_test,random_ipsw_te_test)
xl_pehe_test = pehe(ite_test,xl_te_test)
```

## Final comment

Make sure to upload this notebook to your GitHub repo once you finished!


# Tutoriel MID Générique
Objectif : Illustrer les propriétés de la MID.

Concrètement, nous considérons un système qui a un controle par une loi de commande retardée (expression mathématique) et l'on prend compte et tirons avantage de ce retard dans le calcul des gains de la loi de commande retardée pour stabiliser le système.

 Informations importantes  :

- Stabilité : Lorsque l'on a l'équation décrivant un système dans l'espace de Laplace (transformée de Laplace), la stabilité est traduite par l'ensemble des racines (complexes) de l'équation ayant leurs parties réelles négatives. On dit que les racines sont à gauche (de l'axe des ordonnées).
- Transformée de Laplace (juste pour info) : Dans notre cas, la partie la plus importante dans la transformée de Laplace est la dérivation, c'est ce qui nous permet de passer d'une équation différentielle à une équation polynomiale. L'idée très grossière est qu'une dérivée d'ordre $n$ dans l'espace "classique"/"temporel" correspond à une multiplication par un terme $s^n$ dans l'espace de Laplace (fréquentiel). Donc, par exemple, $4\times k \times y''(t)\rightarrow 4ks^2y(s)$

On traite une équation différentielle du type :
$$
\frac{d^n y(t)}{dt^n} + \sum\limits_{i=0}^{n-1}a_i\frac{d^i y(t)}{dt^i} + \sum\limits_{j=0}^m \alpha_j\frac{d^j y(t-\tau)}{dt^j}=0
$$

- $n, m\in\mathbb{N}, n>m$
- $a_i, \alpha_j \in \mathbb{R}$
- $s\in\mathbb{C}$,
- $\tau \in \mathbb{R}$

Par une transformée de Laplace :
$$
Q(s) = s^n+a_{n-1}s^{n-1}+...+a_0 + \left[\alpha_ms^m+...+\alpha_0\right]e^{-s\tau}
$$

Cette équation est polynomiale et non différentielle ce qui la rend beaucoup plus simple à traiter.

On souhaite réaliser ce que l'on appelle un placement de pôle. Cela signifie dans ce cas que l'on veut trouver les coefficients ($a_i$ et $\alpha_j$ parfois notés $b_j$) soient de telle sorte que $Q(s_0)=0$.

On connaît donc une valeur de racine à "forcer" pour le quasipolynôme en jouant sur les coefficients.

Le code suivant a été développé sur la base du pseudo-code de la MID Générique, décrit par Franck, que vous trouverez à la fin du document.

## Packages nécessaires

On utilise numpy pour les opérations mathématiques numériques telles que les opérations matricielles (produit matriciel par exemple).

Sympy permet de réaliser des calculs en symbolique sur Python et donc de travailler avec des expressions mathématiques et en numérique.

```python
# pip install sympy si nécessaire (module not found), pareil avec cxroots
%matplotlib notebook
import numpy as np
import sympy as sp
import cxroots as cx
import matplotlib.pyplot as plt
```

    D:\ProgramData\Anaconda3\lib\site-packages\numpy\_distributor_init.py:32: UserWarning: loaded more than 1 DLL from .libs:
    D:\ProgramData\Anaconda3\lib\site-packages\numpy\.libs\libopenblas.NOIJJG62EMASZI6NYURL6JBKM4EVBGM7.gfortran-win_amd64.dll
    D:\ProgramData\Anaconda3\lib\site-packages\numpy\.libs\libopenblas.PYQHXLVVQ7VESDPUVUADXEVJOBGHJPAY.gfortran-win_amd64.dll
      stacklevel=1)

## Paramètres
$n$ est le degré du polynôme non retardé, $m$ est celui du polynôme retardé (celui multiplié par l'exponentielle).

$s_0$ est le taux de décroissance exponentiel. Il s'agit de la racine dominante que l'on souhaite assigner (celle la plus à droite lorsqu'on plot les racines). Celui-ci doit être négatif pour que le système soit stable. Un $s_0$ positif signifie donc que le système est instable.

$\tau$ est le retard considéré.

```python
n = 2
m = 1
s0 = -5
value_tau = 1
```

## Création de symboles

Les variables s, tau, a et alpha vont désormais correspondre à des symboles et non des variables numériques.

On considère dans nos cas concrets que les $a_i$ correspondent aux paramètres du système (de son équation) et les $\alpha_j$ correspondent aux gains de la loi de commande retardée.

```python
s = sp.symbols('s')  # define variable s for our problem to be solved
tau = sp.symbols('tau')  # define variable tau : delay

a = sp.symbols(["a{:d}".format(i) for i in range(n)], real=True)
alpha = sp.symbols(["alpha{:d}".format(i) for i in range(m + 1)], real=True)

```

## Création du quasipolynome

```python
Polynomial = s**n + np.array(a).dot([s**i for i in range(n)]) # Revient à faire s^n + a_{n-1}^{n-1}...
Delayed = np.array(alpha).dot([s**i for i in range(m+1)])*sp.exp(-s*tau) # Revient à faire
#b^m*s^m + b_{m-1}^{m-1}...
Q = Polynomial + Delayed
Q
```

$\displaystyle a_{0} + a_{1} s + s^{2} + \left(\alpha_{0} + \alpha_{1} s\right) e^{- s \tau}$

## Créer le système à résoudre

Dans ce cadre de MID générique, on a $n+m+1$ inconnues qui sont les $a_i$ ($n$ variables de $0$ à $n-1$ inclus) et $\alpha_j$ ($m+1$ variables de $0$ à $m$ inclus).

On parle donc de multiplicité de $n+m+1$ et on souhaite donc avoir $n+m+1$ équations. Pour l'instant nous avons une première équation qui correspond à notre quasipolynôme (=0). Étant en multiplicité $n+m+1$ en ***Multiplicity Induced Dominancy*** (MID), on va chercher à ce que la racine (en $s$) que l'on souhaite placer en $s_0$ soit de multiplicité $n+m+1$. Cela revient à ce que les dérivées successives de $Q$ par rapport à $s$ soient nulles lorsque $s=s_0$.

Il faut donc calculer les $n+m$ dérivées suivantes ce qui se fait très simplement avec Sympy en utilisant la méthode diff en spécifiant en argument la variable de différenciation.

```python
SysDerivatif = [Q]
print(Q)
for i in range(n+m+1):
    DerniereDerivee = SysDerivatif[-1]
    SysDerivatif.append(DerniereDerivee.diff(s)) # Dérivée par rapport à s
    print(SysDerivatif[-1])
```

    a0 + a1*s + s**2 + (alpha0 + alpha1*s)*exp(-s*tau)
    a1 + alpha1*exp(-s*tau) + 2*s - tau*(alpha0 + alpha1*s)*exp(-s*tau)
    -2*alpha1*tau*exp(-s*tau) + tau**2*(alpha0 + alpha1*s)*exp(-s*tau) + 2
    3*alpha1*tau**2*exp(-s*tau) - tau**3*(alpha0 + alpha1*s)*exp(-s*tau)
    -4*alpha1*tau**3*exp(-s*tau) + tau**4*(alpha0 + alpha1*s)*exp(-s*tau)

## Résolution du système

Après avoir créé le système dérivatif, il suffit de le résoudre avec Sympy en spécifiant bien que l'on souhaite obtenir les expressions des $a_i$ et des $\alpha_j$.

Ici, on a bien $n+m+1 = 4$ ce qui nous retourne $a_0, a_1, \alpha_0$ et $\alpha_1$ en gardant bien à l'esprit que $a_2=1$.

Le ".args[0]" sert à extraire les solutions à partir du solveur qui retourne un set.

```python
sol = sp.linsolve(SysDerivatif[:-1], alpha + a).args[0] # Solveur selon les alpha et les a
sol
```

$\displaystyle \left( \frac{2 s e^{s \tau}}{\tau} - \frac{6 e^{s \tau}}{\tau^{2}}, \  - \frac{2 e^{s \tau}}{\tau}, \  s^{2} + \frac{4 s}{\tau} + \frac{6}{\tau^{2}}, \  - 2 s - \frac{4}{\tau}\right)$

## Vérification de l'ordre des inconnues

On vérifie que les variables retournées par le solveur sont dans ce sens : $\alpha_0,\alpha_1,a_0,a_1$.

```python
alpha + a
```

    [alpha0, alpha1, a0, a1]

```python
solNum = sol.subs({s : s0})
solNum = solNum.subs({tau : value_tau})
solNum
```

$\displaystyle \left( - \frac{16}{e^{5}}, \  - \frac{2}{e^{5}}, \  11, \  6\right)$

```python
a_num = list(solNum[m + 1:])
alpha_num = list(solNum[:m + 1])
print(a_num, alpha_num)
```

    [11, 6] [-16*exp(-5), -2*exp(-5)]

## Rootfinding

Après avoir calculé ces coefficients, on souhaite vérifier de manière visuelle que les résultats trouvés fonctionnent. Pour cela, on reconstitue le quasipolynôme mais en remplaçant cette fois ci les coefficients par les valeurs trouvées. On remplace également $\tau$ par sa valeur et l'on dérive ensuite ce quasipolynôme par rapport à s. La dérivée est utile pour le rootfinding qui est l'algorithme qui va nous trouver les racines du quasipolynôme dans un contour donné.

La définition d'un contour est indispensable pour du rootfinding complexe. On a ici choisi un rectangle entre $-100$ et $10$ en abscisse et $-100$ et $-100$ en ordonnée.

```python
QNumerique = s**n + np.array(a_num).dot([s**i for i in range(n)])+\
            np.array(alpha_num).dot([s**i for i in range(m+1)])*sp.exp(-s*tau)

# Possibilité de faire ça de manière plus propre avec un "zip"

QNumerique = QNumerique.subs(tau, value_tau)
QNumerique
```

$\displaystyle s^{2} + 6 s + \left(- \frac{2 s}{e^{5}} - \frac{16}{e^{5}}\right) e^{- s} + 11$

```python
sysRootFinding = [QNumerique, QNumerique.diff(s)]
sysFunc = [sp.lambdify(s, i) for i in sysRootFinding]
rect = cx.Rectangle([-100, 10], [-100, 100])
roots = rect.roots(sysFunc[0], sysFunc[1], rootErrTol=1e-5, absTol=1e-5, M = n + m + 1)
print(roots)
xroot = np.real(roots[0])
yroot = np.imag(roots[0])
```

     Multiplicity |               Root
    ------------------------------------------------
          1       | -8.903772521426 -98.849979505853i
          1       | -8.903772521426 +98.849979505853i
          1       | -8.838409259546 -92.560024714680i
          1       | -8.838409259546 +92.560024714680i
          1       | -8.768480972175 -86.269136627702i
          1       | -8.768480972175 +86.269136627702i
          1       | -8.693302366816 -79.977103585485i
          1       | -8.693302366816 +79.977103585485i
          1       | -8.612021236139 -73.683643809653i
          1       | -8.612021236139 +73.683643809653i
          1       | -8.523559382989 -67.388373443179i
          1       | -8.523559382989 +67.388373443179i
          1       | -8.426524806933 -61.090755200800i
          1       | -8.426524806933 +61.090755200800i
          1       | -8.319076500982 -54.790012244816i
          1       | -8.319076500982 +54.790012244816i
          1       | -8.198707288311 -48.484976045378i
          1       | -8.198707288311 +48.484976045378i
          1       | -8.061876389753 -42.173799847042i
          1       | -8.061876389753 +42.173799847042i
          1       | -7.903345461542 -35.853373281832i
          1       | -7.903345461542 +35.853373281832i
          1       | -7.714870396864 -29.517990804516i
          1       | -7.714870396864 +29.517990804516i
          1       | -7.482294141462 -23.155832835555i
          1       | -7.482294141462 +23.155832835555i
          1       | -7.177781931717 -16.737274851129i
          1       | -7.177781931717 +16.737274851129i
          1       | -6.730697330726 -10.155954800599i
          1       | -6.730697330726 +10.155954800599i
          4       | -5.000356963113 -0.000000000000i

```python
plt.figure(figsize=(8,6))
plt.axhline(linewidth=2, color='black')
plt.axvline(linewidth=2, color='black')
plt.axvspan(0, 1, alpha=0.5, color='red')
plt.scatter(xroot, yroot)
plt.grid()
plt.title("Root spectrum")
plt.xlabel(r"$Re(s)$")
plt.ylabel(r"$Im(s)$")
```

    <IPython.core.display.Javascript object>

    Text(0, 0.5, '$Im(s)$')

En résumé :

```python

```


# Exactly simulation of Kitaev Honeycomb model on IBMQ

# Jordan-Wigner transformation

We can use JW transformation to go from spin basis to fermion basis. The JW transformation for 2D Kitaev Honeycomb model is:

$$a_{i,j}=(\Pi_{j'<j}\Pi_{i'}\sigma_{i',j'}^z)(\Pi_{i'<i}\sigma^z_{i',j})\frac{\sigma^x+i\sigma_y}{2}$$
$$a_{i,j}^\dagger=(\prod_{j'<j}\prod_{i'}\sigma_{i',j'}^z)(\prod_{i'<i}\sigma^z_{i',j})\frac{\sigma^x-i\sigma_y}{2}$$

$$\{a_{i,j},a_{i',j'}\}=0,\quad \{a_{i,j},a_{i',j'}^\dagger\}=\delta_{ii'}\delta_{jj'}$$
The first product represent every site above $j$-th column and the second product represent all the sites after site $(i,j)$ following the blue string.

In this case we can restore the anticommutation relation for fermions. Very similar to XY model, we have :$$\sigma^z_{i,j} = 2a^\dagger_{i,j}a_{i,j}-1$$
$$\sigma^z_{i,j}a^\dagger_{i,j}=(2a^\dagger_{i,j}a_{i,j}-1)a^\dagger_{i,j}=2a^\dagger_{i,j}(1-a^\dagger_{i,j}a_{i,j})-a^\dagger_{i,j}=a^\dagger_{i,j}(1-2a_{i,j}^\dagger a_{i,j})=-a^\dagger_{i,j}\sigma^z_{i,j}$$
The interaction terms in Hamiltonian transform into:

$$
\begin{align}
\sigma^x_{i,j}\sigma^x_{i+1,j}&=(\sigma^+_{i,j}+\sigma^-_{i,j})_w(\sigma^+_{i+1,j}+\sigma^-_{i+1,j})_d\\
&=(\prod_{j'<j}\prod_{i'}\sigma_{i',j'}^z)(\prod_{i'<i}\sigma^z_{i',j})(a_{i,j}^\dagger+a_{i,j})(\prod_{j'<j}\prod_{i'}\sigma_{i',j'}^z)(\prod_{i'<i+1}\sigma^z_{i',j})(a_{i+1,j}^\dagger+a_{i+1,j})\\
&=(\prod_{i'<i}\sigma^z_{i',j})(a_{i,j}^\dagger+a_{i,j})(\prod_{i'<i+1}\sigma^z_{i',j})(a_{i+1,j}^\dagger+a_{i+1,j})\\
&=(a_{i,j}^\dagger+a_{i,j})_w\sigma^z_{i,j}(a_{i+1,j}^\dagger+a_{i+1,j})_d\\
&=-(a_{i,j}^\dagger-a_{i,j})_w(a_{i+1,j}^\dagger+a_{i+1,j})_d
\end{align}
$$

$$\sigma^y_{i,j}\sigma^y_{i+1,j}=(a^\dagger_{i-1,j}+a_{i-1,j})_d(a_{i,j}^\dagger-a_{i,j})_w$$

$$\sigma^z_{i,j}\sigma^z_{i+1,j}=(2a^\dagger_{i,j}a_{i,j}-1)(2a^\dagger_{i+1,j}a_{i+1,j}-1)$$

The Hamiltonian becomes $$\hat{H}=J_x\sum_{x\text{ links}}(a_{i,j}^\dagger-a_{i,j})(a_{i+1,j}^\dagger+a_{i+1,j})-J_y\sum_{y\text{ links}}(a_{i-1,j}^\dagger+a_{i-1,j})(a_{i,j}^\dagger-a_{i,j})-J_z\sum_{z\text{ links}}(2a^\dagger_{i,j}a_{i,j}-1)(2a^\dagger_{i+1,j}a_{i+1,j}-1)$$

# Transform into Majorana fermions

$$c_w=i(a_{i,j}^\dagger-a_{i,j}),\quad d_w=a_{i,j}^\dagger+a_{i,j},\quad \text{where $i+j$ even=white}$$

$$c_d=a_{i,j}^\dagger+a_{i,j},\quad d_d=i(a_{i,j}^\dagger-a_{i,j}),\quad \text{where $i+j$ odd=dark}$$

The Hamiltonian becomes:

$$\hat{H}=-iJ_x\sum_{x\text{ links}}c_wc_d+iJ_y\sum_{y\text{ links}}c_dc_w-iJ_z\sum_{z \text{ links}}(id_dd_w)c_dc_w$$

By putting everything in a square lattice with the unit cell containing one while site and one dark site

we can rewrite the Hamiltonian interms of the position $\vec{r}$ of unitcell:

$$\hat{H}=i\sum_{x\text{ links}}(J_xc_{w,r}c_{d,r+r_1}+iJ_yc_{d,r}c_{w,r+r_2}-iJ_z\alpha_rc_{d,r}c_{w,r})$$

we define a new operator $\alpha_r=(id_{d,r}d_{w,r})$. It commutes with the $\hat{H}$ so they can be diagonalized simutanuously.

```python

```