##### Run this cell to set your notebook up on Google Colab

```python
!apt-get install -y xvfb python-opengl ffmpeg > /dev/null 2>&1

!git clone https://github.com/yfletberliac/rlss2019-hands-on.git > /dev/null 2>&1
!pip install -q torch==1.1.0 torchvision pyvirtualdisplay piglet > /dev/null 2>&1
```

# Deep Q Networks
------------
You can find the original paper [here](https://www.cs.toronto.edu/~vmnih/docs/dqn.pdf).

## Preliminaries: Q Learning

#### Q-Value

**Q-Value** is a measure of the overall expected reward assuming the agent is in state $s$ and performs action $a$, and then continues playing until the end of the episode following some policy $\pi$. It is defined mathematically as:

\begin{equation}
Q^{\pi}\left(s_{t}, a_{t}\right)=E\left[R_{t+1}+\gamma R_{t+2}+\gamma^{2} R_{t+3}+\ldots | s_{t}, a_{t}\right]
\end{equation}

where $R_{t+1}$ is the immediate reward received after performing action $a_{t}$ in state $s_{t}$ and $\gamma$ is the discount factor and controls the importance of the future rewards versus the immediate ones: the lower the discount factor is, the less important future rewards are.

#### Bellman Optimality Equation

Formally, the Bellman equation defines the relationships between a given state (or, in our case, a **state-action pair**) and its successors. While many forms exist, one of the most common is the **Bellman Optimality Equation** for the optimal **Q-Value**, which is given by:

\begin{equation}
Q^{*}(s, a)=\sum_{s^{\prime}, r} p\left(s^{\prime}, r | s, a\right)\left[r+\gamma \max _{a^{\prime}} Q^{*}\left(s^{\prime}, a^{\prime}\right)\right]
\end{equation}

Of course, when no uncertainty exists (transition probabilities are either 0 or 1), we have:

\begin{equation}
Q^{*}(s, a)=r(s, a)+\gamma \max _{a^{\prime}} Q^{*}\left(s^{\prime}, a^{\prime}\right)
\end{equation}

#### Q-Value Iteration

We define the corresponding Bellman backup operator:
\begin{equation}
[\mathcal{T} Q]\left(s, a\right)=r(s, a)+\gamma \max _{a^{\prime}} Q\left(s^{\prime}, a^{\prime}\right)
\end{equation}

$Q$ is a fixed point of $\mathcal{T}$:
\begin{equation}
\mathcal{T} Q^{*}=Q^{*}
\end{equation}

If we apply the Bellman operator $\mathcal{T}$ repeatedly to any initial $Q$, the series converges to $Q^{*}$:
\begin{equation}
Q, \mathcal{T} Q, \mathcal{T}^{2} Q, \cdots \rightarrow Q^{*}
\end{equation}

# Imports

```python
import sys
sys.path.insert(0, './rlss2019-hands-on/utils')
# If using the Docker image, replace by:
# sys.path.insert(0, '../utils')

import gym, random, os.path, math, glob, csv, base64
from pathlib import Path
from timeit import default_timer as timer
from datetime import timedelta
import numpy as np

import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F

import matplotlib
%matplotlib inline

from qfettes_plot import plot_all_data
from qfettes_wrappers import *
from openai_wrappers import make_atari, wrap_deepmind
from gym.wrappers import Monitor

from pyvirtualdisplay import Display
from IPython import display as ipythondisplay
from IPython.display import clear_output
```

------------

# Deep Q learning

Usually in Deep RL, the **Q-Value** is defined as $Q(s,a;\theta)$ where $\theta$ represents the parameters of the function approximation used.

For *MuJoCo* or *Roboschool* environments, we usually use a simple 2- or 3-layer MLP whereas when using **raw pixels for observations** such as in *Atari 2600* games, we usually use a 1-, 2- or 3-layer CNN.

In our case, since we want to train DQN on *CartPole*, we will use a 3-layer perceptron for our function approximation.

## Network declaration

In this section, we build $Q(s,a;\theta)$ function approximation. Since the input is composed of 4 scalars, namely:
[position of cart, velocity of cart, angle of pole, rotation rate of pole]
we build a FCN -> ReLU -> FCN -> ReLU -> FCN neural network. As an exercice, change the architecture of the network:

1. Change the 1st fully-connected layer from 8 hidden neurons to 16
2. Create `self.fc2` in `__init__` with 16 neurons
3. Create `self.fc3` with `self.num_actions` as the output size
4. Add it to the network in `forward` with no activation function

```python
class DQN(nn.Module):
    def __init__(self, input_shape, num_actions):
        super().__init__()

        self.input_shape = input_shape
        self.num_actions = num_actions

        self.fc1 = nn.Linear(self.input_shape[0], 8)
        self.fc2 = ...
        self.fc3 = ...

    def forward(self, x):
        x = F.relu(self.fc2(F.relu(self.fc1(x))))
        x = ...

        return x
```

## Safety checks

#### Network architecture

As a *safety check*, inspect the resulting network in the next cell. For instance, the total number of trainable parameters should change with the architecture. Check the correctness of `in_features` and `out_features`.

```python
env_id = 'CartPole-v0'
env    = gym.make(env_id)
network = DQN(env.observation_space.shape, env.action_space.n)

print("Observation space:\n", env.observation_space.shape, "\n")
print("Network architecture:\n", network, "\n")

model_parameters = filter(lambda p: p.requires_grad, network.parameters())
print("Total number of trainable parameters:\n", sum([np.prod(p.size()) for p in model_parameters]))
```

    Observation space:
     (4,)

    Network architecture:
     DQN(
      (fc1): Linear(in_features=4, out_features=8, bias=True)
      (fc2): Linear(in_features=8, out_features=8, bias=True)
      (fc3): Linear(in_features=8, out_features=2, bias=True)
    )

    Total number of trainable parameters:
     130

#### Run a Policy with Random Actions

What the working environment looks like? It's always useful to know the details about the environment you train your policy on. For instance, its dynamics, the size of action and observation space, etc. Below we display three different random policies on `CartPole-v0`.

```python
display = Display(visible=0, size=(1400, 900))
display.start()

def show_video():
    html = []
    for mp4 in Path("videos").glob("*.mp4"):
        video_b64 = base64.b64encode(mp4.read_bytes())
        html.append(''''''.format(mp4, video_b64.decode('ascii')))
    ipythondisplay.display(ipythondisplay.HTML(data="<br>".join(html)))

env = Monitor(env, './videos', force=True, video_callable=lambda episode: True)

for episode in range(2):
    done = False
    obs = env.reset()
    while not done:
        action = env.action_space.sample()
        obs, reward, done, info = env.step(action)
env.close()
show_video()
```

    xdpyinfo was not found, X start can not be checked! Please install xdpyinfo!

We can see the episode ending prematurely because the pole drops.

-----
**Question**:

It is also important to identify some of the characteristics of the problem. `CartPole-v0` can be described as a **fully-observable**, **deterministic**, **continuous state space**, with a **discrete action space** and **frequent rewards**. Take some time to understand each of these terms :-) Try to find the opposite term for each of them, e.g. deterministic <> stochastic.

## Experience Replay Memory

As usual RL tasks have no pre-generated training sets which they can learn from, in off-policy learning, our agent must keep records of all the state-transitions it encountered so it can **learn from them later**. The memory-buffer used to store this is often referred to as the **Experience Replay Memory**. There are several types and architectures of these memory buffers — but some very common ones are:
- the *cyclic memory buffers*: they make sure the agent keeps training over its new behavior rather than things that might no longer be relevant
- the *reservoir-sampling-based memory buffers*: they guarantee each state-transition recorded has an even probability to be inserted to the buffer

We use a combination of both.

In `push`:
1. Append the transition to memory
2. Create the if statement which deletes an old transition from the memory

```python
class ExperienceReplayMemory:
    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []

    def push(self, transition):
        # Append the transition below
        ...

        # Now, we need an `if` statement in order to keep the capacity to its limit. Write it below.
        # Hint: `del something` will delete something if something is an array
        if ...:

        raise NotImplementedError

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)
```

------------

Now we have:
- the **DQN** network,
- the **ExperienceReplayMemory**.

Let's build the **Agent** class !

## Agent declaration

In the cell below:
1. Create `self.target_model` in `declare_networks`
2. Complete the epsilon-greedy algorithm in `get_action`

```python
class Agent(object):
    def __init__(self, config, env, log_dir='/tmp/gym'):
        self.log_dir = log_dir
        self.rewards = []
        self.action_log_frequency = config.ACTION_SELECTION_COUNT_FREQUENCY
        self.action_selections = [0 for _ in range(env.action_space.n)]

    # Define the DQN networks
    def declare_networks(self):
        self.model = DQN(self.num_feats, self.num_actions)
        # Create `self.target_model` with the same network architecture
        self.target_model = ...
        raise NotImplementedError

    # Define the Replay Memory
    def declare_memory(self):
        self.memory = ExperienceReplayMemory(self.experience_replay_size)

    # Append the new transition to the Replay Memory
    def append_to_replay(self, s, a, r, s_):
        self.memory.push((s, a, r, s_))

    # Sample transitions from the Replay Memory
    def sample_minibatch(self):
        transitions = self.memory.sample(self.batch_size)
        batch_state, batch_action, batch_reward, batch_next_state = zip(*transitions)

        shape = (-1,)+self.num_feats

        batch_state = torch.tensor(batch_state, device=self.device, dtype=torch.float).view(shape)
        batch_action = torch.tensor(batch_action, device=self.device, dtype=torch.long).squeeze().view(-1, 1)
        batch_reward = torch.tensor(batch_reward, device=self.device, dtype=torch.float).squeeze().view(-1, 1)

        non_final_mask = torch.tensor(tuple(map(lambda s: s is not None, batch_next_state)), device=self.device, dtype=torch.uint8)
        # Sometimes all next states are false
        try:
            non_final_next_states = torch.tensor([s for s in batch_next_state if s is not None], device=self.device, dtype=torch.float).view(shape)
            empty_next_state_values = False
        except:
            non_final_next_states = None
            empty_next_state_values = True

        return batch_state, batch_action, batch_reward, non_final_next_states, non_final_mask, empty_next_state_values

    # Sample action
    def get_action(self, s, eps=0.1):
        with torch.no_grad():
            # Epsilon-greedy
            if np.random.random() >= eps:
                X = torch.tensor([s], device=self.device, dtype=torch.float)
                a = self.model(X).max(1)[1].view(1, 1)
                return a.item()
            else:
                ...
```

-----
**Question**:

Remember we define the objective function as
\begin{equation}
J=\left(r+\gamma \max _{a^{\prime}} Q\left(s^{\prime}, a^{\prime}, \mathbf{\theta}^{-}\right)-Q(s, a, \mathbf{\theta})\right)^{2},
\end{equation}
where $\theta^{-}$ are the target parameters.

Why do we need a target network in the first place ?

## Learning

In the cell below, and from the above objective fonction:
1. Write the value `expected_q_values`
2. Write `diff`
3. The `update` function needs some work

```python
class Learning(Agent):
    def __init__(self, env=None, config=None, log_dir='/tmp/gym'):
        super().__init__(config=config, env=env, log_dir=log_dir)

    # Compute loss from the Bellman Optimality Equation
    def compute_loss(self, batch_vars):
        batch_state, batch_action, batch_reward, non_final_next_states, non_final_mask, empty_next_state_values = batch_vars

        # Estimate
        current_q_values = self.model(batch_state).gather(1, batch_action)

        # Target
        with torch.no_grad():
            max_next_q_values = torch.zeros(self.batch_size, device=self.device, dtype=torch.float).unsqueeze(dim=1)
            if not empty_next_state_values:
                max_next_action = self.get_max_next_state_action(non_final_next_states)
                max_next_q_values[non_final_mask] = self.target_model(non_final_next_states).gather(1, max_next_action)
        # From the equation above, write the value `expected_q_values`.
            expected_q_values = ...

        # From the equation above, write the value `diff`.
        diff = ...
        loss = self.MSE(diff)
        loss = loss.mean()

        raise NotImplementedError
        return loss

    # Update both networks (the agent and the target)
    def update(self, s, a, r, s_, sample_idx=0):
        self.append_to_replay(s, a, r, s_)

        # When not to update ?
        # There is a concise way to write to skip the update, fill in the 2 blanks in the `if` statement below.
        # Hint: the sample count should be < the learn_start hyperparameter and respect the update_freq.
        if ... or ...:
            raise NotImplementedError
            return None

        batch_vars = self.sample_minibatch()
        loss = self.compute_loss(batch_vars)

        # Optimize the model
        self.optimizer.zero_grad()
        loss.backward()
        for param in self.model.parameters():
            param.grad.data.clamp_(-1, 1)
        self.optimizer.step()

        self.update_target_model()
        self.save_td(loss.item(), sample_idx)

    def update_target_model(self):
        # Copy weights from model to target_model following `target_net_update_freq`.
        self.update_count+=1
        if self.update_count % self.target_net_update_freq == 0:
            self.target_model.load_state_dict(self.model.state_dict())
```

## Model declaration

```python
class Model(Learning):
    def __init__(self, env=None, config=None, log_dir='/tmp/gym'):
        super().__init__(config=config, env=env, log_dir=log_dir)
        self.device = config.device

        # Hyperparameters
        self.gamma = config.GAMMA
        self.target_net_update_freq = config.TARGET_NET_UPDATE_FREQ
        self.experience_replay_size = config.EXP_REPLAY_SIZE
        self.batch_size = config.BATCH_SIZE
        self.learn_start = config.LEARN_START
        self.update_freq = config.UPDATE_FREQ

        # Environment specific parameters
        self.num_feats = env.observation_space.shape
        self.num_actions = env.action_space.n
        self.env = env

        self.declare_networks()
        self.declare_memory()
        self.target_model.load_state_dict(self.model.state_dict())
        self.optimizer = optim.Adam(self.model.parameters(), lr=config.LR)

        # Move to correct device
        self.model = self.model.to(self.device)
        self.target_model.to(self.device)

        self.model.train()
        self.target_model.train()

        self.update_count = 0

    def save_td(self, td, tstep):
        with open(os.path.join(self.log_dir, 'td.csv'), 'a') as f:
            writer = csv.writer(f)
            writer.writerow((tstep, td))

    def get_max_next_state_action(self, next_states):
        return self.target_model(next_states).max(dim=1)[1].view(-1, 1)

    def MSE(self, x):
        return 0.5 * x.pow(2)

    def save_reward(self, reward):
        self.rewards.append(reward)

    def save_action(self, action, tstep):
        self.action_selections[int(action)] += 1.0/self.action_log_frequency
        if (tstep+1) % self.action_log_frequency == 0:
            with open(os.path.join(self.log_dir, 'action_log.csv'), 'a') as f:
                writer = csv.writer(f)
                writer.writerow(list([tstep]+self.action_selections))
            self.action_selections = [0 for _ in range(len(self.action_selections))]

    def save_w(self):
        if not os.path.exists("../saved_agents"):
            os.makedirs("../saved_agents")
        torch.save(self.model.state_dict(), '../saved_agents/model.dump')
        torch.save(self.optimizer.state_dict(), '../saved_agents/optim.dump')
```

## Hyperparameters

```python
class Config(object):
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Main agent variables
        self.GAMMA=0.99
        self.LR=1e-3

        # Epsilon variables
        self.epsilon_start    = 1.0
        self.epsilon_final    = 0.01
        self.epsilon_decay    = 10000
        self.epsilon_by_sample = lambda sample_idx: config.epsilon_final + (config.epsilon_start - config.epsilon_final) * math.exp(-1. * sample_idx / config.epsilon_decay)

        # Memory
        self.TARGET_NET_UPDATE_FREQ = 1000
        self.EXP_REPLAY_SIZE = 10000
        self.BATCH_SIZE = 64

        # Learning control variables
        self.LEARN_START = 1000
        self.MAX_SAMPLES = 50000
        self.UPDATE_FREQ = 1

        # Data logging parameters
        self.ACTION_SELECTION_COUNT_FREQUENCY = 1000

config = Config()
```

## Training

```python
import gym
from openai_monitor import Monitor
from IPython import display
import matplotlib
import matplotlib.pyplot as plt
%matplotlib inline

start=timer()

log_dir = "/tmp/gym/"
try:
    os.makedirs(log_dir)
except OSError:
    files = glob.glob(os.path.join(log_dir, '*.monitor.csv')) \
        + glob.glob(os.path.join(log_dir, '*td.csv')) \
        + glob.glob(os.path.join(log_dir, '*action_log.csv'))
    for f in files:
        os.remove(f)

env_id = 'CartPole-v0'
env    = gym.make(env_id)
env    = Monitor(env, os.path.join(log_dir, env_id))

model  = Model(env=env, config=config, log_dir=log_dir)

episode_reward = 0

observation = env.reset()
for sample_idx in range(1, config.MAX_SAMPLES + 1):

    epsilon = config.epsilon_by_sample(sample_idx)

    action = model.get_action(observation, epsilon)
    # Log action selection
    model.save_action(action, sample_idx)

    prev_observation=observation
    observation, reward, done, _ = env.step(action)
    observation = None if done else observation

    model.update(prev_observation, action, reward, observation, sample_idx)
    episode_reward += reward

    if done:
        observation = env.reset()
        model.save_reward(episode_reward)
        episode_reward = 0
    if sample_idx % 1000 == 0:
        try:
            clear_output(True)
            plot_all_data(log_dir, env_id, 'DQN', config.MAX_SAMPLES, bin_size=(10, 100, 100, 1), smooth=1, time=timedelta(seconds=int(timer()-start)), ipynb=True)
        except IOError:
            pass

model.save_w()
env.close()
```

By observing the plots, does the learning appear to be stable?

If your answer is *yes*, then start a second run, and a third, with the same hyperparameters. ;-)

You have just faced reproducibility concerns, which is quite a serious problem in deep RL and which can be dealt with by e.g. running your experiments on a sufficient number of seeds (~ 6-8 min.)

## Visualize the agent

```python
from gym.wrappers import Monitor

# Loading the agent
fname_model = "../saved_agents/model.dump"
fname_optim = "../saved_agents/optim.dump"
log_dir = "/tmp/gym/"

model  = Model(env=env, config=config, log_dir=log_dir)

if os.path.isfile(fname_model):
    model.model.load_state_dict(torch.load(fname_model))
    model.target_model.load_state_dict(model.model.state_dict())

if os.path.isfile(fname_optim):
    model.optimizer.load_state_dict(torch.load(fname_optim))

env_id = 'CartPole-v0'
env    = gym.make(env_id)
env    = Monitor(env, './videos', force=True, video_callable=lambda episode: True)

for episode in range(3):
    done = False
    obs = env.reset()
    while not done:
        action = model.get_action(obs)
        obs, _, done, _ = env.step(action)

env.close()
show_video()
```

You can experiment with modifying the hypermarameters (learning rate, batch size, experience replay size, etc.) to see if you can make its performance improve !

-------------

Reinforcement Learning Summer School 2019 (Lille, France)/practical_drl_2/practical_drl_2.ipynb


# Kamodo-CLI

The kamodo command-line interface should be installed automatically with pip

    pip install kamodo

!!! warning
    The command line app is a work in progress. Expect changes!

!!! note
    You can run these examples in a notebook with `!kamodo`

Run the command line interface with

    kamodo <key>.<subkey>=value

### instantiation
To create a simple function ready for evaluation:

```console
kamodo model.params.f[cm]=x**2
     lhs   rhs symbol units
f(x)   f  x**2   f(x)    cm
\begin{equation}f{\left(x \right)} [cm] = x^{2}\end{equation}
```

Kamodo will print a description of the loaded model when `verbose` output is on (the default).

### evaluation

To evaluate a model, provide arguments for each function.

```console
kamodo model.params.f[cm]=x**2 model.evaluate.f.x=[-3,-1,1,3] verbose=false
f(x=[-3, -1, 1, 3]) cm =

[9 1 1 9]

```

### visualization

To visualize model output, provide plot parameters similar to evaluation

```console
kamodo model.params.f[cm]=x**2 model.plot.f.x.min=-2 model.plot.f.x.max=2 model.plot.f.x.n=25 verbose=False
```

An interactive plot will open with your figure

{! notebooks/outputs/2019-12-20/13-28-47/temp-plot.html !}

### Kamodofied models

To work with a kamodofied model, specify the `model.class`.

```console
kamodo model.class=kamodo.readers.tiegcm.TIEGCM_Kamodo model.params.filename=$PWD/s001.nc  model.plot.EFLUX.lon=[0]
```

### Configuration

Kamodo can use configuration files so that arguments do not have to be passed manually.
To create your own configuration, create a config.yaml file in your project's directory:
```yaml
{! notebooks/config.yaml !}
```
Running kamodo from a directory containing the above `config.yaml` will produce the same plot as before:

```python
! kamodo
```

         lhs   rhs symbol units
    f(x)   f  x**2   f(x)    cm
    \begin{equation}f{\left(x \right)} [cm] = x^{2}\end{equation}

### Configuration Priority

Kamodo is built on hydra, which prioritizes configuration using the following rules:

1. If there are two configurations that define the same value, the second one would win.
2. If two configurations are contributing to the same dictionary the result would be the combined dictionary.

### Help
```console
kamodo --help
A low-coding command line interface for Kamodo

This application allows users to work with kamodo-compatible models and data
directly from the command line.

Custom models, data, and expressions may be composed by editing config files
without needing to write python.

== Configuration groups ==
Compose your configuration from those groups (group=option)

== Config ==
Override anything in the config (foo.bar=value)
model:
  class: kamodo.Kamodo
  evaluate: {}
  fig_layout: {}
  params: {}
  plot: {}
  plot_conf:
    animation_opts: null
    auto_open: true
    auto_play: true
    config: null
    filename: temp-plot.html
    image: null
    image_filename: plot_image
    include_mathjax: cdn
    include_plotlyjs: true
    link_text: Export to plot.ly
    output_type: file
    show_link: false
    validate: true
verbose: true

Powered by Hydra (https://hydra.cc)
Use --hydra-help to view Hydra specific help
```

### tab completion

Kamodo supports tab completion for bash. To set up bash tab completion, run the following:

    eval "$(kamodo -sc install=bash)"


```python
%load_ext autoreload
%autoreload 2
```

```python
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression, Lasso, LogisticRegression
from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.preprocessing import PolynomialFeatures, StandardScaler
from sklearn.pipeline import Pipeline
import scipy.special
```

# NLSYM DATA

```python
# Preprocess data
df = pd.read_csv("data/card.csv")
data_filter = df['educ'].values >= 6
T = df['educ'].values[data_filter]
Z = df['nearc4'].values[data_filter]
y = df['lwage'].values[data_filter]

# Impute missing values with mean, add dummy columns
# I excluded the columns 'weights' as we don't know what it is
X_df = df[['exper', 'expersq']].copy()
X_df['fatheduc'] = df['fatheduc'].fillna(value=df['fatheduc'].mean())
X_df['fatheduc_nan'] = df['fatheduc'].isnull()*1
X_df['motheduc'] = df['motheduc'].fillna(value=df['motheduc'].mean())
X_df['motheduc_nan'] = df['motheduc'].isnull()*1
X_df[['momdad14', 'sinmom14', 'reg661', 'reg662',
        'reg663', 'reg664', 'reg665', 'reg666', 'reg667', 'reg668', 'reg669', 'south66']] = df[['momdad14', 'sinmom14',
        'reg661', 'reg662','reg663', 'reg664', 'reg665', 'reg666', 'reg667', 'reg668', 'reg669', 'south66']]
X_df[['black', 'smsa', 'south', 'smsa66']] = df[['black', 'smsa', 'south', 'smsa66']]
columns_to_scale = ['fatheduc', 'motheduc', 'exper', 'expersq']
scaler = StandardScaler()
X_raw = X_df.values[data_filter]
X_df[columns_to_scale] = scaler.fit_transform(X_df[columns_to_scale])
X = X_df.values[data_filter]
```

    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\preprocessing\data.py:645: DataConversionWarning: Data with input dtype int64, float64 were all converted to float64 by StandardScaler.
      return self.partial_fit(X, y)
    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\base.py:464: DataConversionWarning: Data with input dtype int64, float64 were all converted to float64 by StandardScaler.
      return self.fit(X, **fit_params).transform(X)

```python
X_df.columns
```

    Index(['exper', 'expersq', 'fatheduc', 'fatheduc_nan', 'motheduc',
           'motheduc_nan', 'momdad14', 'sinmom14', 'reg661', 'reg662', 'reg663',
           'reg664', 'reg665', 'reg666', 'reg667', 'reg668', 'reg669', 'south66',
           'black', 'smsa', 'south', 'smsa66'],
          dtype='object')

# 2. Semi-Synthetic Data with Binary Instrument and Continuous Treatment

Data generating process uning real covariates $X$ and instrument $Z$, but synthetic $y$ and $T$ based on the "intent to treat" instrument setting with non-compliance. The instrument corresponds to a fully randomized recommendation of treatment. Then each sample complies with the recommendation to some degree. This probability also depends on an unobserved confounder that has a direct effect on the outcome. Moreover, compliance also depends on the observed feature $X$.

\begin{align}
X \tag{ real features}\\
Z \tag{real instrument}\\
\nu \sim \; & \text{U}[0, 1] \tag{unobserved confounder}\\
C = \; & c\cdot X[i], \; c \;(const)\sim \text{U}[.2, .3] \tag{compliance degree}\\
T = \; & C\cdot Z + g(X) + \nu  \tag{treatment}\\
y \sim \; & \text{Normal}(\mu=\theta(X) \cdot (T + \nu) + f(X),\; \sigma=.1) \tag{outcome}
\end{align}

Moreover:
\begin{align}
\theta(X) = \; & \alpha + \beta \cdot X[i] \tag{CATE}\\
f(X) = \; & X[i] \tag{Nuissance function}\\
g(X) = \; & X[i] \tag{Nuissance function}
\end{align}

```python
def dgp_bin_Z_cont_T(X, Z, hetero_col, true_fn, random_seed=None):
    np.random.seed(random_seed)
    n, d = X.shape
    nu = np.random.uniform(-1, 1, size=(n,))
    c = np.random.uniform(0.2, 0.3)
    C = c * X[:, hetero_col] # Compliers when recomended
    T = C * Z + X[:, hetero_col] + nu # Treatment with compliance
    y = true_fn(X) * (T + nu)  + 0.05*X[:, hetero_col] + np.random.normal(0, .1, size=(n,))
    return y, T

hetero_col = 4 # Mother's education
hetero_col_2 = 7
true_fn = lambda X: 0.1 + 0.05*X[:, hetero_col] - 0.1*X[:, hetero_col_2]

np.random.seed(1237)
y, T = dgp_bin_Z_cont_T(X_raw, Z, hetero_col, true_fn)

plt.figure(figsize=(10, 2))
plt.subplot(1, 2, 1)
plt.hist(T[Z==0])
plt.title("T[Z=0]: Total: {}, Mean(T): {:.2f}".format(T[Z==0].shape[0], np.mean(T[Z==0])))
plt.subplot(1, 2, 2)
plt.hist(T[Z==1])
plt.title("T[Z=1]: Total: {}, Mean(T): {:.2f}".format(T[Z==1].shape[0], np.mean(T[Z==1])))
plt.show()
```

# ANALYSIS

### Defining some hyperparameters

```python
random_seed = 12345 # random seed for each experiment
N_SPLITS = 10 # number of splits for cross-fitting
COV_CLIP = 20/X.shape[0] # covariance clipping in driv
print(COV_CLIP)
```

    0.006686726847208292

### Defining some generic non-parametric regressors and classifiers

```python
from utilities import RegWrapper
from sklearn.model_selection import GridSearchCV
from sklearn.linear_model import LassoCV, LogisticRegressionCV
from xgboost import XGBClassifier, XGBRegressor
from xgb_utilities import XGBWrapper

# XGB forest models for Regression and Classification
model = lambda: XGBWrapper(XGBRegressor(gamma=0.001, n_estimators=100, min_child_weight=20, n_jobs=10),
                           early_stopping_rounds=5, eval_metric='rmse', binary=False)

model_clf = lambda: RegWrapper(XGBWrapper(XGBClassifier(gamma=0.001, n_estimators=100, min_child_weight=20, n_jobs=10),
                                          early_stopping_rounds=5, eval_metric='logloss', binary=True))
```

### Some utility functions

```python
def nuisance_diagnostic(cate, nuisance_model, property_name, property_fn,
                        index_names=None, statistic=np.std, threshold=None):
    std = statistic([property_fn(ns) for ns in cate.fitted_nuisances[nuisance_model]], axis=0)
    if hasattr(std, '__len__'):
        if threshold is None:
            coefs = np.argmax(std).flatten()
        else:
            coefs = np.argwhere(std >= threshold).flatten()
        if index_names is None:
            index_names = np.arange(std.shape[0])
        for high_var in coefs:
            plt.title("{}: {}[{}] Across Folds".format(nuisance_model, property_name, index_names[high_var]))
            plt.plot([property_fn(ns)[high_var] for ns in cate.fitted_nuisances[nuisance_model]])
            plt.xlabel('fold')
            plt.ylabel('property')
            plt.show()
    else:
        plt.title("{}: {} Across Folds".format(nuisance_model, property_name))
        plt.plot([property_fn(ns) for ns in cate.fitted_nuisances[nuisance_model]])
        plt.xlabel('fold')
        plt.ylabel('property')
        plt.show()

```

# ATE via DMLATEIV

```python
from dml_ate_iv import DMLATEIV

np.random.seed(random_seed)

# We need to specify models to be used for each of these residualizations
model_Y_X = lambda: model() # model for E[Y | X]
model_T_X = lambda: model() # model for E[T | X]. We use a regressor since T is continuous
model_Z_X = lambda: model_clf() # model for E[Z | X]. We use a classifier since Z is binary

dmlate = DMLATEIV(model_Y_X(), model_T_X(), model_Z_X(),
                  n_splits=N_SPLITS, # n_splits determines the number of splits to be used for cross-fitting.
                  binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
                  binary_treatment=False # a flag whether to stratify cross-fitting by treatment
                 )
```

    C:\ProgramData\Anaconda3\lib\site-packages\numba\errors.py:105: UserWarning: Insufficiently recent colorama version found. Numba requires colorama >= 0.3.9
      warnings.warn(msg)

```python
# We fit DMLATEIV with these models
dmlate.fit(y, T, X, Z)
```

    <dml_ate_iv.DMLATEIV at 0x18d162284a8>

```python
# We call effect() to get the ATE
ta_effect = dmlate.effect()
```

```python
# Comparison with true ATE
print("ATE Estimate: {:.3f}".format(ta_effect))
print("True ATE: {:.3f}".format(np.mean(true_fn(X_raw))))
# CATE MSE
print("CATE MSE: {:.2f}".format(np.mean((true_fn(X_raw) - ta_effect)**2)))
```

    ATE Estimate: 0.648
    True ATE: 0.609
    CATE MSE: 0.03

```python
# We can call normal_effect_interval to get confidence intervals based
# based on the asympotic normal approximation
ta_effect = dmlate.normal_effect_interval(lower=2.5, upper=97.5)
# Comparison with true ATE
print("ATE Estimate Interval: ({:.3f}, {:.3f})".format(ta_effect[0], ta_effect[1]))
print("True ATE: {:.3f}".format(np.mean(true_fn(X_raw))))
```

    ATE Estimate Interval: (0.601, 0.696)
    True ATE: 0.609

```python
def get_dmlateiv_coverage(true_effect, iteration):
    y, T = dgp_bin_Z_cont_T(X_raw, Z, hetero_col, true_fn, random_seed=iteration)
    dmlate = DMLATEIV(model_Y_X(), model_T_X(), model_Z_X(),
                  n_splits=N_SPLITS, # n_splits determines the number of splits to be used for cross-fitting.
                  binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
                  binary_treatment=False # a flag whether to stratify cross-fitting by treatment
                 )
    dmlate.fit(y, T, X, Z)
    left, right = dmlate.normal_effect_interval(lower=2.5, upper=97.5)
    if true_effect >= left and true_effect <= right:
        return 1
    return 0

from joblib import Parallel, delayed
n_experiments=100
true_ate = np.mean(true_fn(X_raw))
if True:
    contains_truth = np.array(Parallel(n_jobs=-1, verbose=3)(
            delayed(get_dmlateiv_coverage)(true_ate, it) for it in range(n_experiments)))
    print("Coverage: {}".format(contains_truth.mean()))
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 12 concurrent workers.
    [Parallel(n_jobs=-1)]: Done   8 tasks      | elapsed:   30.7s

    Coverage: 0.72

    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:  2.0min finished

# ATE and CATE via DMLIV

```python
from dml_iv import DMLIV
from utilities import SelectiveLasso, SeparateModel
from sklearn.linear_model import LassoCV, LogisticRegressionCV
from econml.utilities import hstack

np.random.seed(random_seed)

# For DMLIV we also need a model for E[T | X, Z]. To allow for heterogeneity in the compliance, i.e.
# T = beta(X)*Z + gamma(X)
# we train a separate model for Z=1 and Z=0. The model for Z=1 learns the
# quantity beta(X) + gamma(X) and the model for Z=0 learns gamma(X).
model_T_XZ = lambda: SeparateModel(model(), model())

# We now specify the features to be used for heterogeneity. We will fit a CATE model of the form
#      theta(X) = <theta, phi(X)>
# for some set of features phi(X). The featurizer needs to support fit_transform, that takes
# X and returns phi(X). We need to include a bias if we also want a constant term.
dmliv_featurizer = lambda: PolynomialFeatures(degree=1, include_bias=True)

# Then we need to specify a model to be used for fitting the parameters theta in the linear form.
# This model will minimize the square loss:
#        (Y - E[Y|X] - <theta, phi(X)> * (E[T|X,Z] - E[T|X]))**2
#dmliv_model_effect = lambda: LinearRegression(fit_intercept=False)

# Potentially with some regularization on theta. Here we use an ell_1 penalty on theta
# If we also have a prior that there is no effect heterogeneity we can use a selective lasso
# that does not penalize the constant term in the CATE model
dmliv_model_effect = lambda: SelectiveLasso(np.arange(1, X.shape[1]+1), LassoCV(cv=5, fit_intercept=False))

# We initialize DMLIV with all these models and call fit
cate = DMLIV(model_Y_X(), model_T_X(), model_T_XZ(),
             dmliv_model_effect(), dmliv_featurizer(),
             n_splits=N_SPLITS, # number of splits to use for cross-fitting
             binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
             binary_treatment=False # a flag whether to stratify cross-fitting by treatment
            )
```

```python
cate.fit(y, T, X, Z)
```

    <dml_iv.DMLIV at 0x18d17d0e1d0>

```python
# To get the CATE at every X we call effect(X)
dml_effect = cate.effect(X)
plt.hist(dml_effect, label='est')
plt.hist(true_fn(X_raw), alpha=.2, label='true')
plt.legend()
plt.show()
```

```python
# To get the parameter theta we call coef_. The first entry is the intercept of the CATE model
print(cate.coef_)
```

    [ 0.60084499  0.          0.          0.00162356 -0.          0.11428089
      0.         -0.         -0.          0.         -0.          0.
      0.          0.         -0.         -0.          0.         -0.
     -0.         -0.         -0.00789543 -0.         -0.        ]

```python
# We can average the CATE to get an ATE
print("ATE Estimate: {:.3f}".format(np.mean(dml_effect)))
print("True ATE: {:.3f}".format(np.mean(true_fn(X_raw))))
```

    ATE Estimate: 0.596
    True ATE: 0.609

```python
# We can also see how it compares to the true CATE at each target point and calculate MSE
plt.title("DMLIV CATE as Function of {}. MSE={:.3f}".format(X_df.columns[4], np.mean((dml_effect-true_fn(X_raw))**2)))
plt.scatter(X[:, 4], dml_effect, label='est')
plt.scatter(X[:, 4], true_fn(X_raw), label='true', alpha=.2)
plt.legend()
plt.show()
```

# ATE and Projected CATE via DRIV

```python
from dml_iv import DMLIV
from dr_iv import DRIV, ProjectedDRIV
from utilities import SubsetWrapper, StatsModelLinearRegression, ConstantModel
from sklearn.dummy import DummyRegressor

np.random.seed(random_seed)

# For DRIV we need a model for predicting E[T*Z | X]. We use a classifier
model_TZ_X = lambda: model()

# We also need a model for the final regression that will fit the function theta(X)
# If we want to fit an ATE, we simply fit a constant functin theta(X) = theta
# We can do this with a pipeline where the preprocessing step only creates a bias column
# and the regression step fits a linear regression with no intercept.
# To get normal confidence intervals easily we can use a statsmodels linear regression
# wrapped in an sklearn interface
const_driv_model_effect = lambda: ConstantModel()

# As in OrthoDMLIV we need a perliminary estimator of the CATE.
# We use a DMLIV estimator with no cross-fitting (n_splits=1)
dmliv_prel_model_effect = DMLIV(model_Y_X(), model_T_X(), model_T_XZ(),
                                dmliv_model_effect(), dmliv_featurizer(),
                                n_splits=1, binary_instrument=True, binary_treatment=False)

const_dr_cate = DRIV(model_Y_X(), model_T_X(), model_Z_X(), # same as in DMLATEIV
                        dmliv_prel_model_effect, # preliminary model for CATE, must support fit(y, T, X, Z) and effect(X)
                        model_TZ_X(), # model for E[T * Z | X]
                        const_driv_model_effect(), # model for final stage of fitting theta(X)
                        cov_clip=COV_CLIP, # covariance clipping to avoid large values in final regression from weak instruments
                        n_splits=N_SPLITS, # number of splits to use for cross-fitting
                        binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
                        binary_treatment=False, # a flag whether to stratify cross-fitting by treatment
                        opt_reweighted=False # whether to optimally re-weight samples. Valid only for flexible final model
                       )
```

```python
const_dr_cate.fit(y, T, X, Z, store_final=True)
```

    <dr_iv.DRIV at 0x18d1da9b6d8>

```python
# To get the statsmodel summary we look at the effect_model, which is the pipeline, we then look
# at the reg step of the pipeline which is the statsmodel wrapper and then we look
# at the model attribute of the statsmodel wrapper and print the summary()
const_dr_cate.effect_model.summary()
```

    C:\ProgramData\Anaconda3\lib\site-packages\statsmodels\regression\linear_model.py:1554: RuntimeWarning: invalid value encountered in double_scalars
      return self.ess/self.df_model
    C:\ProgramData\Anaconda3\lib\site-packages\scipy\stats\_distn_infrastructure.py:877: RuntimeWarning: invalid value encountered in greater
      return (self.a < x) & (x < self.b)
    C:\ProgramData\Anaconda3\lib\site-packages\scipy\stats\_distn_infrastructure.py:877: RuntimeWarning: invalid value encountered in less
      return (self.a < x) & (x < self.b)
    C:\ProgramData\Anaconda3\lib\site-packages\scipy\stats\_distn_infrastructure.py:1831: RuntimeWarning: invalid value encountered in less_equal
      cond2 = cond0 & (x <= self.a)

OLS Regression Results

  Dep. Variable:	            y	          R-squared:         	    0.000

  Model:	                   OLS	         Adj. R-squared:    	    0.000

  Method:	             Least Squares	    F-statistic:       	      nan

  Date:	             Sat, 01 Jun 2019	   Prob (F-statistic):	     nan

  Time:	                 16:39:54	       Log-Likelihood:    	  -6509.5

  No. Observations:	        2991	        AIC:               	 1.302e+04

  Df Residuals:	            2990	        BIC:               	 1.303e+04

  Df Model:	                   0

  Covariance Type:	      nonrobust

    	       coef	     std err	      t	      P>|t|	  [0.025	    0.975]

  const	     0.6501	     0.039	    16.668	  0.000	     0.574	     0.727

  Omnibus:	       2934.155	   Durbin-Watson:     	     1.968

  Prob(Omnibus):	   0.000	    Jarque-Bera (JB):  	 1860178.825

  Skew:	            3.818	    Prob(JB):          	      0.00

  Kurtosis:	       124.934	   Cond. No.          	      1.00

Warnings:
[1] Standard Errors assume that the covariance matrix of the errors is correctly specified.

```python
def get_driv_coverage(true_effect, iteration):
    y, T = dgp_bin_Z_cont_T(X_raw, Z, hetero_col, true_fn, random_seed=iteration)
    dmliv_prel_model_effect = DMLIV(model_Y_X(), model_T_X(), model_T_XZ(),
                                dmliv_model_effect(), dmliv_featurizer(),
                                n_splits=1, binary_instrument=True, binary_treatment=True)
    const_dr_cate = DRIV(model_Y_X(), model_T_X(), model_Z_X(), # same as in DMLATEIV
                            dmliv_prel_model_effect, # preliminary model for CATE, must support fit(y, T, X, Z) and effect(X)
                            model_TZ_X(), # model for E[T * Z | X]
                            const_driv_model_effect(), # model for final stage of fitting theta(X)
                            cov_clip=COV_CLIP, # covariance clipping to avoid large values in final regression from weak instruments
                            n_splits=N_SPLITS, # number of splits to use for cross-fitting
                            binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
                            binary_treatment=False # a flag whether to stratify cross-fitting by treatment
                           )
    const_dr_cate.fit(y, T, X, Z, store_final=True)
    left, right = const_dr_cate.effect_model.est.model.conf_int(alpha=0.05)[0]
    if true_effect >= left and true_effect <= right:
        return 1
    return 0
```

```python
from joblib import Parallel, delayed
n_experiments=100
true_ate = np.mean(true_fn(X_raw))
if True:
    contains_truth = np.array(Parallel(n_jobs=-1, verbose=3)(
            delayed(get_driv_coverage)(true_ate, it) for it in range(n_experiments)))
    print("Coverage: {}".format(contains_truth.mean()))
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 12 concurrent workers.
    [Parallel(n_jobs=-1)]: Done   8 tasks      | elapsed:   31.2s

    Coverage: 0.93

    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:  4.3min finished

## Projecting CATE to a pre-chosen subset of variables in final model

```python
from dml_iv import DMLIV
from dr_iv import DRIV, ProjectedDRIV
from utilities import SubsetWrapper

np.random.seed(random_seed)

# We could also fit a projection on a subset of the features by using the
# subset wrapper from our utilities.

# Example: including everything for expository purposes, but any array-like of indices would work
subset_names = set(['motheduc', 'sinmom14'])
# list of indices of features X to use in the final model
feature_inds = np.argwhere([(x in subset_names) for x in X_df.columns.values]).flatten() #[0] #np.arange(X.shape[1])
print(feature_inds)
# Because we are projecting to a low dimensional model space, we can
# do valid inference and we can use statsmodel linear regression to get all
# the hypothesis testing capability
proj_driv_model_effect = lambda: SubsetWrapper(StatsModelLinearRegression(),
                                          feature_inds # list of indices of features X to use in the final model
                                         )
```

    [4 7]

```python
X_df.columns
```

    Index(['exper', 'expersq', 'fatheduc', 'fatheduc_nan', 'motheduc',
           'motheduc_nan', 'momdad14', 'sinmom14', 'reg661', 'reg662', 'reg663',
           'reg664', 'reg665', 'reg666', 'reg667', 'reg668', 'reg669', 'south66',
           'black', 'smsa', 'south', 'smsa66'],
          dtype='object')

```python
proj_dr_cate = const_dr_cate.refit_final(proj_driv_model_effect())
```

```python
# To get the CATE at every X we call effect(X[:, feature_inds])
proj_dr_effect = proj_dr_cate.effect(X[:, feature_inds])
```

```python
# To get the statsmodel summary we look at the effect_model, which is
# an instance of SubsetWrapper, we look at the model of the SubsetWrapper which is
# and instance of the pipeline, we then look at the reg step of the pipeline which is the statsmodel wrapper and
# call summary() of the wrapper (most prob there is a better API for this, but we can go with this for now :)
proj_dr_cate.effect_model.summary(alpha=.05, xname=['const']+list(X_df.columns[feature_inds]))
```

OLS Regression Results

  Dep. Variable:	            y	          R-squared:         	    0.002

  Model:	                   OLS	         Adj. R-squared:    	    0.001

  Method:	             Least Squares	    F-statistic:       	    2.934

  Date:	             Sat, 01 Jun 2019	   Prob (F-statistic):	  0.0533

  Time:	                 16:44:16	       Log-Likelihood:    	  -6506.5

  No. Observations:	        2991	        AIC:               	 1.302e+04

  Df Residuals:	            2988	        BIC:               	 1.304e+04

  Df Model:	                   2

  Covariance Type:	      nonrobust

      	        coef	     std err	      t	      P>|t|	  [0.025	    0.975]

  const	        0.6352	     0.041	    15.433	  0.000	     0.554	     0.716

  motheduc	     0.0915	     0.040	     2.306	  0.021	     0.014	     0.169

  sinmom14	     0.1403	     0.131	     1.070	  0.285	    -0.117	     0.397

  Omnibus:	       2951.753	   Durbin-Watson:     	     1.972

  Prob(Omnibus):	   0.000	    Jarque-Bera (JB):  	 1886680.784

  Skew:	            3.859	    Prob(JB):          	      0.00

  Kurtosis:	       125.798	   Cond. No.          	      3.41

Warnings:
[1] Standard Errors assume that the covariance matrix of the errors is correctly specified.

```python
# We can access the coefficient by looking at the coefficient attribute of the final step of the pipeline
print("Estimated Params: {}, {}".format(proj_dr_cate.intercept_, proj_dr_cate.coef_))
# True coefficients of projection
print("True Params: {}".format(
        LinearRegression(fit_intercept=False).fit(PolynomialFeatures(degree=1,
                                                                     include_bias=True).fit_transform(X[:, feature_inds]),
                                                  true_fn(X_raw)).coef_))
```

    Estimated Params: 0.6351902202433224, [0.091538   0.14032943]
    True Params: [ 0.61740685  0.14934234 -0.1       ]

```python
# We can also evaluate coverage and create prediction intervals using statsmodels attributes
from statsmodels.sandbox.regression.predstd import wls_prediction_std
res = proj_dr_cate.effect_model.model
predictions = res.get_prediction(PolynomialFeatures(degree=1, include_bias=True).fit_transform(X[:, feature_inds]))
frame = predictions.summary_frame(alpha=0.05)
pred = frame['mean']
iv_l = frame['mean_ci_lower']
iv_u = frame['mean_ci_upper']

# This is the true CATE functions
theta_true = true_fn(X_raw)
# This is the true projection of the CATE function on the subspace of linear functions of the
# subset of the features used in the projection
true_proj = LinearRegression().fit(X[:, feature_inds], theta_true).predict(X[:, feature_inds])

# Are we covering the true projection
covered = (true_proj <= iv_u) & (true_proj >= iv_l)
print("Coverage of True Projection: {:.2f}".format(np.mean(covered)))

fig, ax = plt.subplots(figsize=(8,6))

order = np.argsort(X[:, feature_inds[0]])
ax.plot(X[order, feature_inds[0]], iv_u[order], 'r--')
ax.plot(X[order, feature_inds[0]], iv_l[order], 'r--')
ax.plot(X[order, feature_inds[0]], pred[order], 'g--.', label="pred")
ax.plot(X[order, feature_inds[0]], theta_true[order], 'b-', label="True", alpha=.3)
ax.plot(X[order, feature_inds[0]], true_proj[order], 'b-', label="TrueProj", alpha=.3)
ax.legend(loc='best')
plt.show()
```

# CATE via Re-Weighted DRIV: DRIV-RW

## Lasso CATE

```python
from dml_iv import DMLIV
from dr_iv import DRIV, ProjectedDRIV
from utilities import SubsetWrapper, StatsModelLinearRegression, ConstantModel, WeightWrapper
from sklearn.dummy import DummyRegressor

np.random.seed(random_seed)

# For DRIV we need a model for predicting E[T*Z | X]. We use a classifier
model_TZ_X = lambda: model()

# We also need a model for the final regression that will fit the function theta(X)
# This model now needs to accept sample weights at fit time.
const_driv_model_effect = lambda: WeightWrapper(Pipeline([('bias', PolynomialFeatures(degree=1, include_bias=True)),
                                                          ('reg', SelectiveLasso(np.arange(1, X.shape[1]+1),
                                                                                LassoCV(cv=5, fit_intercept=False)))]))

# As in OrthoDMLIV we need a perliminary estimator of the CATE.
# We use a DMLIV estimator with no cross-fitting (n_splits=1)
dmliv_prel_model_effect = DMLIV(model_Y_X(), model_T_X(), model_T_XZ(),
                                dmliv_model_effect(), dmliv_featurizer(),
                                n_splits=1, binary_instrument=True, binary_treatment=False)

const_dr_cate = DRIV(model_Y_X(), model_T_X(), model_Z_X(), # same as in DMLATEIV
                        dmliv_prel_model_effect, # preliminary model for CATE, must support fit(y, T, X, Z) and effect(X)
                        model_TZ_X(), # model for E[T * Z | X]
                        const_driv_model_effect(), # model for final stage of fitting theta(X)
                        cov_clip=COV_CLIP, # covariance clipping to avoid large values in final regression from weak instruments
                        n_splits=N_SPLITS, # number of splits to use for cross-fitting
                        binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
                        binary_treatment=False, # a flag whether to stratify cross-fitting by treatment
                        opt_reweighted=True # whether to optimally re-weight samples. Valid only for flexible final model
                       )
```

```python
const_dr_cate.fit(y, T, X, Z, store_final=True)
```

    <dr_iv.DRIV at 0x18d22b12da0>

```python
# We can average the CATE to get an ATE
dr_effect = const_dr_cate.effect(X)
print("ATE Estimate: {:.3f}".format(np.mean(dr_effect)))
print("True ATE: {:.3f}".format(np.mean(true_fn(X_raw))))
```

    ATE Estimate: 0.599
    True ATE: 0.609

```python
# We can also see how it compares to the true CATE at each target point and calculate MSE
plt.title("DMLIV CATE as Function of {}: MSE={:.3f}".format(X_df.columns[4], np.mean((dr_effect-true_fn(X_raw))**2)))
plt.scatter(X[:, 4], dr_effect, label='est')
plt.scatter(X[:, 4], true_fn(X_raw), label='true', alpha=.2)
plt.legend()
plt.show()
```

## Random Forest CATE

```python
from dml_iv import DMLIV
from dr_iv import DRIV, ProjectedDRIV
from utilities import SubsetWrapper
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor

np.random.seed(random_seed)

rf_driv_model_effect = lambda: RandomForestRegressor(n_estimators=5000, max_depth=3, min_impurity_decrease=0.00001,
                                                     min_samples_leaf=100, bootstrap=True)
```

```python
rf_dr_cate = const_dr_cate.refit_final(rf_driv_model_effect())
```

```python
rf_dr_effect = rf_dr_cate.effect(X)
```

```python
print("ATE Estimate: {:.2f}".format(np.mean(rf_dr_effect)))
print("True ATE: {:.2f}".format(np.mean(true_fn(X_raw))))
```

    ATE Estimate: 0.60
    True ATE: 0.61

```python
plt.title("DRIV CATE: MSE {:.2}".format(np.mean((true_fn(X_raw) - rf_dr_effect)**2)))
plt.scatter(X[:, 4], rf_dr_effect, label='est')
plt.scatter(X[:, 4], true_fn(X_raw), label='true', alpha=.2)
plt.legend()
plt.show()
```

```python
import shap
import pandas as pd

Xdf = pd.DataFrame(X, columns=X_df.columns)
# explain the model's predictions using SHAP values
explainer = shap.TreeExplainer(rf_dr_cate.effect_model)
shap_values = explainer.shap_values(Xdf)

# visualize the first prediction's explanation (use matplotlib=True to avoid Javascript)
shap.force_plot(explainer.expected_value, shap_values[0,:], Xdf.iloc[0,:], matplotlib=True)
```

```python
shap.summary_plot(shap_values, Xdf)
```

```python
shap.summary_plot(shap_values, Xdf, plot_type='bar')
```

```python

```