# Tutorial 1: Learning to Predict
**Week 3, Day 4: Reinforcement Learning**

**By Neuromatch Academy**

__Content creators:__ Marcelo Mattar and Eric DeWitt with help from Matt Krause

__Content reviewers:__ Ella Batty, Byron Galbraith and Michael Waskom

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

---

# Tutorial objectives

*Estimated timing of tutorial: 50 min*

In this tutorial, we will learn how to estimate state-value functions in a classical conditioning paradigm using Temporal Difference (TD) learning and examine TD-errors at the presentation of the conditioned and unconditioned stimulus (CS and US) under different CS-US contingencies. These exercises will provide you with an understanding of both how reward prediction errors (RPEs) behave in classical conditioning and what we should expect to see if Dopamine represents a "canonical" model-free RPE.

At the end of this tutorial:
* You will learn to use the standard tapped delay line conditioning model
* You will understand how RPEs move to CS
* You will understand how variability in reward size effects RPEs
* You will understand how differences in US-CS timing effect RPEs

```python
# @title Tutorial slides

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

---
# Setup

```python
# Imports

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

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

```python
# @title Plotting Functions
from matplotlib import ticker

def plot_value_function(V, ax=None, show=True):
  """Plot V(s), the value function"""
  if not ax:
    fig, ax = plt.subplots()

  ax.stem(V, use_line_collection=True)
  ax.set_ylabel('Value')
  ax.set_xlabel('State')
  ax.set_title("Value function: $V(s)$")

  if show:
    plt.show()

def plot_tde_trace(TDE, ax=None, show=True, skip=400):
  """Plot the TD Error across trials"""
  if not ax:
    fig, ax = plt.subplots()

  indx = np.arange(0, TDE.shape[1], skip)
  im = ax.imshow(TDE[:,indx])
  positions = ax.get_xticks()
  # Avoid warning when setting string tick labels
  ax.xaxis.set_major_locator(ticker.FixedLocator(positions))
  ax.set_xticklabels([f"{int(skip * x)}" for x in positions])
  ax.set_title('TD-error over learning')
  ax.set_ylabel('State')
  ax.set_xlabel('Iterations')
  ax.figure.colorbar(im)
  if show:
    plt.show()

def learning_summary_plot(V, TDE):
  """Summary plot for Ex1"""
  fig, (ax1, ax2) = plt.subplots(nrows = 2, gridspec_kw={'height_ratios': [1, 2]})

  plot_value_function(V, ax=ax1, show=False)
  plot_tde_trace(TDE, ax=ax2, show=False)
  plt.tight_layout()
```

```python
# @title Helper Functions and Classes

def reward_guesser_title_hint(r1, r2):
  """"Provide a mildly obfuscated hint for a demo."""
  if (r1==14 and r2==6) or (r1==6 and r2==14):
    return "Technically correct...(the best kind of correct)"

  if  ~(~(r1+r2) ^ 11) - 1 == (6 | 24): # Don't spoil the fun :-)
    return "Congratulations! You solved it!"

  return "Keep trying...."

class ClassicalConditioning:

    def __init__(self, n_steps, reward_magnitude, reward_time):

        # Task variables
        self.n_steps = n_steps
        self.n_actions = 0
        self.cs_time = int(n_steps/4) - 1

        # Reward variables
        self.reward_state = [0,0]
        self.reward_magnitude = None
        self.reward_probability = None
        self.reward_time = None

        self.set_reward(reward_magnitude, reward_time)

        # Time step at which the conditioned stimulus is presented

        # Create a state dictionary
        self._create_state_dictionary()

    def set_reward(self, reward_magnitude, reward_time):

        """
        Determine reward state and magnitude of reward
        """
        if reward_time >= self.n_steps - self.cs_time:
            self.reward_magnitude = 0

        else:
            self.reward_magnitude = reward_magnitude
            self.reward_state = [1, reward_time]

    def get_outcome(self, current_state):

        """
        Determine next state and reward
        """
        # Update state
        if current_state < self.n_steps - 1:
            next_state = current_state + 1
        else:
            next_state = 0

        # Check for reward
        if self.reward_state == self.state_dict[current_state]:
            reward = self.reward_magnitude
        else:
            reward = 0

        return next_state, reward

    def _create_state_dictionary(self):

        """
        This dictionary maps number of time steps/ state identities
        in each episode to some useful state attributes:

        state      - 0 1 2 3 4 5 (cs) 6 7 8 9 10 11 12 ...
        is_delay   - 0 0 0 0 0 0 (cs) 1 1 1 1  1  1  1 ...
        t_in_delay - 0 0 0 0 0 0 (cs) 1 2 3 4  5  6  7 ...
        """
        d = 0

        self.state_dict = {}
        for s in range(self.n_steps):
            if s <= self.cs_time:
                self.state_dict[s] = [0,0]
            else:
                d += 1 # Time in delay
                self.state_dict[s] = [1,d]

class MultiRewardCC(ClassicalConditioning):
  """Classical conditioning paradigm, except that one randomly selected reward,
    magnitude, from a list, is delivered of a single fixed reward."""
  def __init__(self, n_steps, reward_magnitudes, reward_time=None):
    """"Build a multi-reward classical conditioning environment
      Args:
        - nsteps: Maximum number of steps
        - reward_magnitudes: LIST of possible reward magnitudes.
        - reward_time: Single fixed reward time
      Uses numpy global random state.
      """
    super().__init__(n_steps, 1, reward_time)
    self.reward_magnitudes = reward_magnitudes

  def get_outcome(self, current_state):
    next_state, reward = super().get_outcome(current_state)
    if reward:
      reward=np.random.choice(self.reward_magnitudes)
    return next_state, reward

class ProbabilisticCC(ClassicalConditioning):
  """Classical conditioning paradigm, except that rewards are stochastically omitted."""
  def __init__(self, n_steps, reward_magnitude, reward_time=None, p_reward=0.75):
    """"Build a multi-reward classical conditioning environment
      Args:
        - nsteps: Maximum number of steps
        - reward_magnitudes: Reward magnitudes.
        - reward_time: Single fixed reward time.
        - p_reward: probability that reward is actually delivered in rewarding state
      Uses numpy global random state.
      """
    super().__init__(n_steps, reward_magnitude, reward_time)
    self.p_reward = p_reward

  def get_outcome(self, current_state):
    next_state, reward = super().get_outcome(current_state)
    if reward:
      reward*= int(np.random.uniform(size=1)[0] < self.p_reward)
    return next_state, reward
```

---
# Section 1: Temporal difference learning

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

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

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

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

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

display(out)
```

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

__Environment:__

- The agent experiences the environment in episodes or trials.
- Episodes terminate by transitioning to the inter-trial-interval (ITI) state and they are initiated from the ITI state as well. We clamp the value of the terminal/ITI states to zero.
- The classical conditioning environment is composed of a sequence of states that the agent deterministically transitions through. Starting at State 0, the agent moves to State 1 in the first step, from State 1 to State 2 in the second, and so on.  These states represent time in the tapped delay line representation
- Within each episode, the agent is presented a CS and US (reward).
- The CS is always presented at 1/4 of the total duration of the trial. The US (reward) is then delivered after the CS. The interval between the CS and US is specified by `reward_time`.
- The agent's goal is to learn to predict expected rewards from each state in the trial.

**General concepts**

* Return $G_{t}$: future cumulative reward, which can be written in a recursive form
\begin{align}
G_{t} &= \sum \limits_{k = 0}^{\infty} \gamma^{k} r_{t+k+1} \\
&= r_{t+1} + \gamma G_{t+1}
\end{align}
where $\gamma$ is discount factor that controls the importance of future rewards, and $\gamma \in [0, 1]$. $\gamma$ may also be interpreted as probability of continuing the trajectory.
* Value funtion $V_{\pi}(s_t=s)$: expecation of the return
\begin{align}
V_{\pi}(s_t=s) &= \mathbb{E} [ G_{t}\; | \; s_t=s, a_{t:\infty}\sim\pi] \\
& = \mathbb{E} [ r_{t+1} + \gamma G_{t+1}\; | \; s_t=s, a_{t:\infty}\sim\pi]
\end{align}
With an assumption of **Markov process**, we thus have:
\begin{align}
V_{\pi}(s_t=s) &= \mathbb{E} [ r_{t+1} + \gamma V_{\pi}(s_{t+1})\; | \; s_t=s, a_{t:\infty}\sim\pi] \\
&= \sum_a \pi(a|s) \sum_{r, s'}p(s', r)(r + V_{\pi}(s_{t+1}=s'))
\end{align}

**Temporal difference (TD) learning**

* With a Markovian assumption, we can use $V(s_{t+1})$ as an imperfect proxy for the true value $G_{t+1}$ (Monte Carlo bootstrapping), and thus obtain the generalised equation to calculate TD-error:
\begin{align}
\delta_{t} = r_{t+1} + \gamma V(s_{t+1}) - V(s_{t})
\end{align}

* Value updated by using the learning rate constant $\alpha$:
\begin{align}
V(s_{t}) \leftarrow V(s_{t}) + \alpha \delta_{t}
\end{align}

  (Reference: https://web.stanford.edu/group/pdplab/pdphandbook/handbookch10.html)

__Definitions:__

* TD-error:
\begin{align}
\delta_{t} = r_{t+1} + \gamma V(s_{t+1}) - V(s_{t})
\end{align}

* Value updates:
\begin{align}
V(s_{t}) \leftarrow V(s_{t}) + \alpha \delta_{t}
\end{align}

## Coding Exercise 1: TD-learning with guaranteed rewards

Implement TD-learning to estimate the state-value function in the classical-conditioning world with guaranteed rewards, with a fixed magnitude, at a fixed delay after the conditioned stimulus, CS. Save TD-errors over learning (i.e., over trials) so we can visualize them afterwards.

In order to simulate the effect of the CS, you should only update $V(s_{t})$ during the delay period after CS. This period is indicated by the boolean variable `is_delay`. This can be implemented by multiplying the expression for updating the value function by `is_delay`.

Use the provided code to estimate the value function. We will use helper class `ClassicalConditioning`.

```python
def td_learner(env, n_trials, gamma=0.98, alpha=0.001):
  """ Temporal Difference learning

  Args:
    env (object): the environment to be learned
    n_trials (int): the number of trials to run
    gamma (float): temporal discount factor
    alpha (float): learning rate

  Returns:
    ndarray, ndarray: the value function and temporal difference error arrays
  """
  V = np.zeros(env.n_steps) # Array to store values over states (time)
  TDE = np.zeros((env.n_steps, n_trials)) # Array to store TD errors

  for n in range(n_trials):

    state = 0 # Initial state

    for t in range(env.n_steps):

      # Get next state and next reward
      next_state, reward = env.get_outcome(state)

      # Is the current state in the delay period (after CS)?
      is_delay = env.state_dict[state][0]

      ########################################################################
      ## TODO for students: implement TD error and value function update
      # Fill out function and remove
      #raise NotImplementedError("Student excercise: implement TD error and value function update")
      #################################################################################
      # Write an expression to compute the TD-error
      TDE[state, n] = reward + gamma * V[next_state] - V[state]

      # Write an expression to update the value function
      V[state] += alpha * TDE[state, n]  * is_delay

      # Update state
      state = next_state

  return V, TDE

# Initialize classical conditioning class
env = ClassicalConditioning(n_steps=40, reward_magnitude=10, reward_time=10)

# Perform temporal difference learning
V, TDE = td_learner(env, n_trials=20000)

# Visualize
learning_summary_plot(V, TDE)
```

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

*Example output:*

## Interactive Demo 1.1: US to CS Transfer

During classical conditioning, the subject's behavioral response (e.g., salivating) transfers from the unconditioned stimulus (US; like the smell of tasty food) to the conditioned stimulus (CS; like Pavlov ringing his bell) that predicts it. Reward prediction errors play an important role in this process by adjusting the value of states according to their expected, discounted return.

Use the widget below to examine how reward prediction errors change over time.

Before training (orange line), only the reward state has high reward prediction error. As training progresses (blue line, slider), the reward prediction errors shift to the conditioned stimulus, where they end up when the trial is complete (green line).

Dopamine neurons, which are thought to carry reward prediction errors _in vivo_, show exactly the same behavior!

```python
#@title

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

n_trials = 20000

@widgets.interact
def plot_tde_by_trial(trial = widgets.IntSlider(value=5000, min=0, max=n_trials-1 , step=1, description="Trial #")):
  if 'TDE' not in globals():
    print("Complete Exercise 1 to enable this interactive demo!")
  else:

    fig, ax = plt.subplots()
    ax.axhline(0, color='k') # Use this + basefmt=' ' to keep the legend clean.
    ax.stem(TDE[:, 0], linefmt='C1-', markerfmt='C1d', basefmt=' ',
            label="Before Learning (Trial 0)",
            use_line_collection=True)
    ax.stem(TDE[:, -1], linefmt='C2-', markerfmt='C2s', basefmt=' ',
            label="After Learning (Trial $\infty$)",
            use_line_collection=True)
    ax.stem(TDE[:, trial], linefmt='C0-', markerfmt='C0o', basefmt=' ',
            label=f"Trial {trial}",
            use_line_collection=True)

    ax.set_xlabel("State in trial")
    ax.set_ylabel("TD Error")
    ax.set_title("Temporal Difference Error by Trial")
    ax.legend()
```

    interactive(children=(IntSlider(value=5000, description='Trial #', max=19999), Output()), _dom_classes=('widge…

## Interactive Demo 1.2: Learning Rates and Discount Factors

Our TD-learning agent has two parameters that control how it learns: $\alpha$, the learning rate, and $\gamma$, the discount factor. In Exercise 1, we set these parameters to $\alpha=0.001$ and $\gamma=0.98$ for you. Here, you'll investigate how changing these parameters alters the model that TD-learning learns.

Before enabling the interactive demo below, take a moment to think about the functions of these two parameters. $\alpha$ controls the size of the Value function updates produced by each TD-error. In our simple, deterministic world, will this affect the final model we learn? Is a larger $\alpha$ necessarily better in more complex, realistic environments?

The discount rate $\gamma$ applies an exponentially-decaying weight to returns occuring in the future, rather than the present timestep. How does this affect the model we learn? What happens when $\gamma=0$ or $\gamma \geq 1$?

Use the widget to test your hypotheses.

```python
#@title

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

@widgets.interact
def plot_summary_alpha_gamma(alpha = widgets.FloatSlider(value=0.0001, min=0.0001, max=0.1, step=0.0001, readout_format='.4f', description="alpha"),
                             gamma = widgets.FloatSlider(value=0.980, min=0, max=1.1, step=0.010, description="gamma")):
  env = ClassicalConditioning(n_steps=40, reward_magnitude=10, reward_time=10)
  try:
    V_params, TDE_params = td_learner(env, n_trials=20000, gamma=gamma, alpha=alpha)
  except NotImplementedError:
    print("Finish Exercise 1 to enable this interactive demo")

  learning_summary_plot(V_params,TDE_params)
```

    interactive(children=(FloatSlider(value=0.0001, description='alpha', max=0.1, min=0.0001, readout_format='.4f'…

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

---
# Section 2: TD-learning with varying reward magnitudes

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

In the previous exercise, the environment was as simple as possible. On every trial, the CS predicted the same reward, at the same time, with 100% certainty. In the next few exercises, we will make the environment more progressively more complicated and examine the TD-learner's behavior.

## Interactive Demo 2: Match the Value Functions

First, will replace the environment with one that dispenses one of several rewards, chosen at random. Shown below is the final value function $V$ for a TD learner that was trained in an enviroment where the CS predicted a reward of 6 or 14 units; both rewards were equally likely).

Can you find another pair of rewards that cause the agent to learn the same value function? Assume each reward will be dispensed 50% of the time.

Hints:
* Carefully consider the definition of the value function $V$. This can be solved analytically.
* There is no need to change $\alpha$ or $\gamma$.
* Due to the randomness, there may be a small amount of variation.

```python
# @markdown Make sure you execute this cell to enable the widget! Please allow some time for the new figure to load

n_trials = 20000
np.random.seed(2020)
rng_state = np.random.get_state()
env = MultiRewardCC(40, [6, 14], reward_time=10)
V_multi, TDE_multi = td_learner(env, n_trials, gamma=0.98, alpha=0.001)

@widgets.interact
def reward_guesser_interaction(r1 = widgets.IntText(value=0, min=0, max=50, description="Reward 1"),
                               r2 = widgets.IntText(value=0, min=0, max=50, description="Reward 2")):
  try:
    env2 = MultiRewardCC(40, [r1, r2], reward_time=10)
    V_guess, _ = td_learner(env2, n_trials, gamma=0.98, alpha=0.001)
    fig, ax = plt.subplots()
    m, l, _ = ax.stem(V_multi, linefmt='y-', markerfmt='yo', basefmt=' ', label="Target",
            use_line_collection=True)
    m.set_markersize(15)
    m.set_markerfacecolor('none')
    l.set_linewidth(4)
    m, _, _ = ax.stem(V_guess, linefmt='r', markerfmt='rx', basefmt=' ', label="Guess",
                      use_line_collection=True)
    m.set_markersize(15)

    ax.set_xlabel("State")
    ax.set_ylabel("Value")
    ax.set_title("Guess V(s)\n" + reward_guesser_title_hint(r1, r2))
    ax.legend()
  except NotImplementedError:
    print("Please finish Exercise 1 first!")
```

## Think! 2: Examining the TD Error

Run the cell below to plot the TD errors from our multi-reward environment. A new feature appears in this plot? What is it? Why does it happen?

```python
plot_tde_trace(TDE_multi)
```

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

---
# Section 3: TD-learning with probabilistic rewards
*Estimated timing to here from start of tutorial: 40 min*

## Think! 3: Probabilistic rewards

In this environment, we'll return to delivering a single reward of ten units. However, it will be delivered intermittently: on 20 percent of trials, the CS will be shown but the agent will not receive the usual reward; the remaining 80% will proceed as usual.

 Run the cell below to simulate. How does this compare with the previous experiment?

Earlier in the notebook, we saw that changing $\alpha$ had little effect on learning in a deterministic environment. What happens if you set it to an large value, like 1, in this noisier scenario? Does it seem like it will _ever_ converge?

```python
# @markdown Execute this cell to see visualization
np.random.set_state(rng_state) # Resynchronize everyone's notebooks
n_trials = 20000
try:
  env = ProbabilisticCC(n_steps=40, reward_magnitude=10, reward_time=10,
                        p_reward=0.8)
  V_stochastic, TDE_stochastic = td_learner(env, n_trials*2, alpha=1)
  learning_summary_plot(V_stochastic, TDE_stochastic)
except NotImplementedError:
  print("Please finish Exercise 1 first")
```

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

---
# Summary

*Estimated timing of tutorial: 50 min*

In this notebook, we have developed a simple TD Learner and examined how its state representations and reward prediction errors evolve during training. By manipualting its environment and parameters ($\alpha$, $\gamma$), you developed an intuition for how it behaves.

This simple model closely resembles the behavior of subjects undergoing classical conditioning tasks and the dopamine neurons that may underlie that behavior. You may have implemented TD-reset or used the model to recreate a common experimental error. The update rule used here has been extensively studied for [more than 70 years](https://www.pnas.org/content/108/Supplement_3/15647) as a possible explanation for artificial and biological learning.

However, you may have noticed that something is missing from this notebook. We carefully calculated the value of each state, but did not use it to actually do anything. Using values to plan _**Actions**_ is coming up next!

---
# Bonus

## Bonus Think! 1: Removing the CS

In Coding Exercise 1, you (should have) included a term that depends on the conditioned stimulus. Remove it and see what happens. Do you understand why?
This phenomena often fools people attempting to train animals--beware!

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


# Notebook

**Roadmap**

* Data Generating Process
* Objects of Interest
* Processing Specification
* Setting up Simulation
* Conducting Estimation
* Inspection of Results

Over the next couple of lectures, we will then constantly refine the basic code and explore elements of software engineering such as Object-Oriented Programming, Unit Testing, Debugging, and Profiling.

Before we get started, let us import some basic libraries.

```python
%%capture
# Notebook metods
from IPython.core.display import HTML, Image

# Unix Pattern Extensions
import glob

# Operating System Interfaces
import os

# Lexical Analysis
import shlex

# Copy operations
import copy

# Encoders and Decoders
import codecs

# Statistical Modeling and Econometrics
import statsmodels.api as sm

# Plotting
import matplotlib.pyplot as plt
%pylab inline --no-import-all

# Scientific Computing
import numpy as np
from scipy.stats    import norm
from scipy.optimize import minimize
```

## Processing of Model Specification

We manage the model specification in an external text file, which is called *init.ini*. This file will turn out to be useful to provide the parameters for a simulation of a synthetic sample or the initialization of starting values for an estimation.

```python
Image(filename='material/images/init.png', width=1000)
```

Now we will develop a set of function that processes the initialization file.

```python
def _process_bene(list_, dict_, keyword):
    """ This function processes the BENE part of the initialization file.
    """
    # Distribute information
    name, val_treated, val_untreated = list_[0], list_[1], list_[2]

    # Initialize dictionary
    if 'TREATED' not in dict_.keys():
        for subgroup in ['TREATED', 'UNTREATED']:
            dict_[subgroup] = {}
            dict_[subgroup]['coeff'] = []
            dict_[subgroup]['int'] = None
            dict_[subgroup]['sd'] = None

    # Type conversion
    val_treated = float(val_treated)
    val_untreated = float(val_untreated)

    # Collect information
    if name in ['coeff']:
        dict_['TREATED'][name] += [val_treated]
        dict_['UNTREATED'][name] += [val_untreated]
    else:
        dict_['TREATED'][name] = val_treated
        dict_['UNTREATED'][name] = val_untreated

    # Finishing
    return dict_

def _process_not_bene(list_, dict_, keyword):
    """ This function processes all of the initialization file, but the
        BENE section.
    """
    # Distribute information
    name, val = list_[0], list_[1]

    # Prepare container.
    if name not in dict_[keyword].keys():
        if name in ['coeff']:
            dict_[keyword][name] = []

    # Type conversion
    if name in ['agents', 'maxiter']:
        val = int(val)
    elif name in ['source', 'algorithm', 'start', 'version']:
        val = str(val)
    else:
        val = float(val)

    # Collect information
    if name in ['coeff']:
        dict_[keyword][name] += [val]
    else:
        dict_[keyword][name] = val

    # Finishing.
    return dict_

def _check_integrity_process(dict_):
    """ Check integrity of initFile dict.
    """
    # Antibugging
    assert (isinstance(dict_, dict))

    # Check number of agents
    assert (dict_['BASICS']['agents'] > 0)
    assert (isinstance(dict_['BASICS']['agents'], int))

    # Check optimizer
    assert (dict_['ESTIMATION']['algorithm'] in ['bfgs', 'nm'])

    # Check starting values
    assert (dict_['ESTIMATION']['start'] in ['random', 'init'])

    # Maximum iterations
    assert (dict_['ESTIMATION']['maxiter'] >= 0)

    # Finishing
    return True

def _add_auxiliary(dict_):
    """ Add some auxiliary objects.
    """
    # Antibugging
    assert (isinstance(dict_, dict))

    # Initialize container
    dict_['AUX'] = {}

    # Full set of coefficients.
    for key_ in ['TREATED', 'UNTREATED', 'COST']:
        dict_[key_]['all'] = [dict_[key_]['int']]
        dict_[key_]['all'] += dict_[key_]['coeff']
        dict_[key_]['all'] = np.array(dict_[key_]['all'])

    # Number of covariates
    num_covars_out = len(dict_['TREATED']['all'])
    num_covars_cost = len(dict_['COST']['all'])

    dict_['AUX']['num_covars_out'] = num_covars_out
    dict_['AUX']['num_covars_cost'] = num_covars_cost

    # Number of parameters
    dict_['AUX']['num_paras'] = 2 * num_covars_out + num_covars_cost + 2 + 2

    # Starting values
    dict_['AUX']['init_values'] = []

    for key_ in ['TREATED', 'UNTREATED', 'COST']:
        dict_['AUX']['init_values'] += dict_[key_]['all'].tolist()

    dict_['AUX']['init_values'] += [dict_['TREATED']['sd']]
    dict_['AUX']['init_values'] += [dict_['UNTREATED']['sd']]
    dict_['AUX']['init_values'] += [dict_['DIST']['rho1']]
    dict_['AUX']['init_values'] += [dict_['DIST']['rho0']]

    # Finishing
    return dict_

def _process_cases(list_):
    """ Process cases and determine whether keyword or empty
        line.
    """
    # Antibugging
    assert (isinstance(list_, list))

    # Get information
    is_empty = (len(list_) == 0)

    if not is_empty:
        is_keyword = list_[0].isupper()
    else:
        is_keyword = False

    # Antibugging
    assert (is_keyword in [True, False])
    assert (is_empty in [True, False])

    # Finishing
    return is_empty, is_keyword
```

```python
def process(file_):
    """ Process initialization file.
    """
    # Initialization
    dict_ = {}

    for line in open(file_).readlines():

        # Remove UTF-3 marker
        if line.startswith(codecs.BOM_UTF8):
            line = line[3:]

        # Split line
        list_ = shlex.split(line)

        # Determine special cases
        is_empty, is_keyword = _process_cases(list_)

        # Applicability
        if is_empty:
            continue

        if is_keyword:
            keyword = list_[0]
            dict_[keyword] = {}
            continue

        if keyword not in ['BENE']:
            dict_ = _process_not_bene(list_, dict_, keyword)

        else:
            dict_ = _process_bene(list_, dict_, keyword)

    # Remove BENE
    del dict_['BENE']

    # Add auxiliary objects
    dict_ = _add_auxiliary(dict_)

    # Check quality.
    _check_integrity_process(dict_)

    # Finishing.
    return dict_
```

Let us check if it is all working.

```python
init_dict = process('material/msc/init.ini')
```

# Setting up the Simulation

**Distributinal Assumptions**

*Observables*

\begin{align}
X & \sim \mathbb{N}(0, 1) \\
Z & \sim \mathbb{N}(0, 1) \\
\end{align}

*Unobservables*

\begin{eqnarray*}
\begin{pmatrix}U_{1}\\
U_{0}\\
V
\end{pmatrix} & \sim & \mathbb{N}\left[\left(\begin{array}{c}
0\\
0\\
0
\end{array}\right),\left(\begin{array}{ccc}
\sigma_{U_1}^2 & 0 & \sigma_{U_1,V}\\
0 & \sigma_{U_0}^2 & \sigma_{U_0,V}\\
\sigma_{U_1,V} & \sigma_{U_0,V} & \sigma_{V}^2
\end{array}\right)\right]\\
\end{eqnarray*}

```python
init_dict = process('material/msc/init.ini')
```

```python
def _check_integrity_simulate(Y1, Y0, Y, D):
    """ Check quality of simulated sample.
    """
    assert (np.all(np.isfinite(Y1)))
    assert (np.all(np.isfinite(Y0)))

    assert (np.all(np.isfinite(Y)))
    assert (np.all(np.isfinite(D)))

    assert (Y1.dtype == 'float')
    assert (Y0.dtype == 'float')

    assert (Y.dtype == 'float')
    assert (D.dtype == 'float')

    assert (D.all() in [1.0, 0.0])

def _write_out(Y, D, X, Z, source, unobserved=False, Y1=None, Y0=None):
    """ Write out simulated data to file.
    """

    if not unobserved:

        np.savetxt(source, np.column_stack((Y, D, X, Z)), fmt='%8.3f')

    else:

        assert (isinstance(Y1, np.ndarray))
        assert (isinstance(Y0, np.ndarray))

        np.savetxt(source, np.column_stack((Y, D, X, Z, Y1, Y0)),
                   fmt='%8.3f')
```

```python
def simulate(init_dict, unobserved=False):
    """ Simulate a model based on the initialization file.
    """

    # Antibugging
    assert (isinstance(init_dict, dict))
    assert (unobserved in [True, False])

    # Ensure recomputability
    np.random.seed(123)

    # Distribute information
    num_agents = init_dict['BASICS']['agents']
    source = init_dict['BASICS']['source']

    Y1_coeffs = init_dict['TREATED']['all']
    Y0_coeffs = init_dict['UNTREATED']['all']

    C_coeffs = np.array(init_dict['COST']['all'])

    U1_sd = init_dict['TREATED']['sd']
    U0_sd = init_dict['UNTREATED']['sd']

    V_sd = init_dict['COST']['sd']

    U1V_rho = init_dict['DIST']['rho1']
    U0V_rho = init_dict['DIST']['rho0']

    # Auxiliary objects
    U1V_cov = U1V_rho * U1_sd * V_sd
    U0V_cov = U0V_rho * U0_sd * V_sd

    num_covars_out = Y1_coeffs.shape[0]
    num_covars_cost = C_coeffs.shape[0]

    # Simulate observables
    means = np.tile(0.0, num_covars_out)
    covs = np.identity(num_covars_out)

    X = np.random.multivariate_normal(means, covs, num_agents)

    means = np.tile(0.0, num_covars_cost)
    covs = np.identity(num_covars_cost)

    Z = np.random.multivariate_normal(means, covs, num_agents)

    # Add intercepts. The first column of the X and Z matrix always contains
    # the intercept term. This is exploited throughout the code.
    Z[:,0], X[:, 0] = 1.0, 1.0

    # Construct index of observable characteristics
    Y1_level = np.dot(Y1_coeffs, X.T)
    Y0_level = np.dot(Y0_coeffs, X.T)
    C_level = np.dot(C_coeffs, Z.T)

    # Simulate unobservables
    means = np.tile(0.0, 3)
    vars_ = [U1_sd**2, U0_sd**2, V_sd**2]
    covs = np.diag(vars_)

    covs[0, 2] = U1V_cov
    covs[2, 0] = covs[0, 2]

    covs[1, 2] = U0V_cov
    covs[2, 1] = covs[1, 2]

    U = np.random.multivariate_normal(means, covs, num_agents)

    # Simulate endogenous variables
    Y1 = np.tile(np.nan, num_agents)
    Y0 = np.tile(np.nan, num_agents)
    Y = np.tile(np.nan, num_agents)

    D = np.tile(np.nan, num_agents)

    for i in range(num_agents):

        # Select individual unobservables and observables
        u1, u0, v = U[i, 0], U[i, 1], U[i, 2]

        y1_idx, y0_idx, c_idx = Y1_level[i], Y0_level[i], C_level[i]

        # Decision Rule
        expected_benefits = y1_idx - y0_idx
        cost = c_idx + v

        d = np.float((expected_benefits - cost > 0))

        # Potential outcomes
        y1, y0 = y1_idx + u1, y0_idx + u0

        # Observed outcomes
        y = d * y1 + (1.0 - d) * y0

        # Collect data matrices
        Y[i], Y0[i], Y1[i], D[i] = y, y1, y0, d

    # Check integrity of simulated data
    _check_integrity_simulate(Y1, Y0, Y, D)

    # Save to disk
    _write_out(Y, D, X, Z, source, unobserved, Y1, Y0)

    # Return selected features of data
    return Y1, Y0, D
```

Let us check if it is all working.

```python
init_dict = process('material/msc/init.ini')

Y1, Y0, D = simulate(init_dict)

```

Given our parametrization, let us revisit our objects of interest. We start with the individual-specific benefits. Please note the use of the [StatsModels](http://statsmodels.sourceforge.net/devel/index.html) library,
that provides classes and functions for the estimation of many different statistical models, as well as for conducting statistical tests, and statistical data exploration. Think of it as a replacement for using [R](http://www.r-project.org/). In general, [rpy2](http://rpy.sourceforge.net/rpy2.html) provides a low-level interface to R from Python.

```python
# Auxiliary variables
B = Y1 - Y0

# Create histogram and density estimate of benefits.
kde = sm.nonparametric.KDEUnivariate(B)
kde.fit()

# Initialize canvas
ax = plt.figure(figsize=(12,8)).add_subplot(111)

# Plot histogram and density
ax.hist(B, bins=50, normed=True, color='blue')
ax.plot(kde.support, kde.density, lw=2, color='black')

# Set axis labels
ax.set_xlabel('Individual-Specific Benefits', fontsize=18)
ax.set_ylabel('Density and Histogram', fontsize=18)

# Change background color
ax.set_axis_bgcolor('white')

# Remove first element on y-axis
ax.yaxis.get_major_ticks()[0].set_visible(False)

# Calcuate the average treatment effects
ate, tt, tut = np.mean(B), np.mean(B[D==1]), np.mean(B[D==0])

# Pretty formatting of strings and output
fmt = '     {0:<5}{1:10.2f}\n'

print '\nAverage Treatment Effects\n'
print fmt.format('ATE ', ate)
print fmt.format('TT', tt)
print fmt.format('TUT ', tut)

# Let us add them to our plot.
plt.axvline(x=ate, ymin=0, ymax=5, linewidth=2, color='g')
plt.axvline(x=tt, ymin=0, ymax=5, linewidth=2, color='b')
plt.axvline(x=tut, ymin=0, ymax=5, linewidth=2, color='y')

# Add title
plt.suptitle('Distribution of Individual-Specific Benefits', fontsize=20)
```

## Estimation

Now, we will perform Maximum Likelihood Estimation using alternative optimization algorithms. Here is the likelihood function:

\begin{align}
\mathcal{L}(\theta; X, Z) =\sum^N_{i=1} D\mathcal{L_{i,1}} + (1 - D)\mathcal{L_{i,0}},
\end{align}

where

\begin{align}
\mathcal{L_1} = & \log\left(\frac{1}{\sigma_{U_1}}\phi\left(\frac{Y_i - X_i\beta_1}{\sigma_{U_1}}\right)\Phi\left(\frac{Z_i\gamma - \sigma_V/\sigma_{U_1}(Y_i - X_i\beta_1)}{\sqrt{(1 - \rho_{U_1,V})*\sigma^2_{V}}}\right)\right) \\
\mathcal{L_0} =  &\log\left(\frac{1}{\sigma_{U_0}}\phi\left(\frac{Y_i - X_i\beta_0}{\sigma_{U_0}}\right)\Phi\left(\frac{Z_i\gamma - \sigma_V/\sigma_{U_0}(Y_i - X_i\beta_0)}{\sqrt{(1 - \rho_{U_0,V})*\sigma^2_{V}}}\right)\right) \\
\end{align}

```python
def _distribute_parameters(x, init_dict, num_covars_out):
    """ Distribute the parameters.
    """
    # Antibugging
    assert (isinstance(x, np.ndarray))
    assert (isinstance(num_covars_out, int))
    assert (num_covars_out > 0)

    # Initialize containers
    rslt = dict()

    rslt['TREATED'] = dict()
    rslt['UNTREATED'] = dict()
    rslt['COST'] = dict()
    rslt['DIST'] = dict()

    # Distribute parameters
    rslt['TREATED']['all'] = x[:num_covars_out]
    rslt['UNTREATED']['all'] = x[num_covars_out:(2 * num_covars_out)]

    rslt['COST']['all'] = x[(2 * num_covars_out):(-4)]
    rslt['COST']['sd'] = init_dict['COST']['sd']

    rslt['TREATED']['sd'] = np.exp(x[(-4)])
    rslt['UNTREATED']['sd'] = np.exp(x[(-3)])

    rslt['DIST']['rho1'] = -1.0 + 2.0 / (1.0 + float(np.exp(-x[-2])))
    rslt['DIST']['rho0'] = -1.0 + 2.0 / (1.0 + float(np.exp(-x[-1])))

    # Update auxiliary versions
    rslt['AUX'] = dict()

    rslt['AUX']['x_internal'] = x.copy()
    rslt['AUX']['x_internal'][-4] = np.exp(x[(-4)])
    rslt['AUX']['x_internal'][-3] = np.exp(x[(-3)])
    rslt['AUX']['x_internal'][-2] = -1.0 + 2.0 / (1.0 + float(np.exp(-x[-2])))
    rslt['AUX']['x_internal'][-1] = -1.0 + 2.0 / (1.0 + float(np.exp(-x[-1])))

    rslt['AUX']['init_values'] = init_dict['AUX']['init_values']

    # Finishing.
    return rslt

def _max_interface(x, Y, D, X, Z, version, init_dict):
    """ Interface to the SciPy maximization routines.
    """
    # Auxiliary objects
    num_covars_out = X.shape[1]

    # Collect maximization arguments
    rslt = _distribute_parameters(x, init_dict, num_covars_out)

    # Calculate likelihood
    likl = _negative_log_likelihood(rslt, Y, D, X, Z)

    # Finishing.
    return likl

def _negative_log_likelihood(args, Y, D, X, Z):
    """ Negative Log-likelihood function of the generalized Roy model.
    """
    # Distribute arguments
    Y1_coeffs, Y0_coeffs, C_coeffs, choice_coeffs, U1_sd, U0_sd, U1V_rho, \
    U0V_rho, V_sd = _distribute_arguments(args)

    # Auxiliary objects
    num_agents = Y.shape[0]

    # Initialize containers
    likl = np.tile(np.nan, num_agents)
    choice_idx = np.tile(np.nan, num_agents)

    # Likelihood construction.
    for i in range(num_agents):

        g = np.concatenate((X[i, :], Z[i,:]))
        choice_idx[i] = np.dot(choice_coeffs, g)

        # Select outcome information
        if D[i] == 1.00:
            coeffs, rho, sd = Y1_coeffs, U1V_rho, U1_sd
        else:
            coeffs, rho, sd = Y0_coeffs, U0V_rho, U0_sd

        arg_one = (Y[i] - np.dot(coeffs, X[i, :])) / sd
        arg_two = (choice_idx[i] - rho * V_sd * arg_one) / \
                  np.sqrt((1.0 - rho ** 2) * V_sd**2)

        pdf_evals, cdf_evals = norm.pdf(arg_one), norm.cdf(arg_two)

        if D[i] == 1.0:
            contrib = (1.0 / float(sd)) * pdf_evals * cdf_evals
        else:
            contrib = (1.0 / float(sd)) * pdf_evals * (1.0 - cdf_evals)

        likl[i] = contrib

    # Transformations.
    likl = -np.mean(np.log(np.clip(likl, 1e-20, np.inf)))

    # Quality checks.
    assert (isinstance(likl, float))
    assert (np.isfinite(likl))

    # Finishing.
    return likl

def _get_start(which, init_dict, Y, D, X, Z):
    """ Get different kind of starting values.
    """
    # Antibugging.
    assert (which in ['random', 'init', 'auto'])

    # Distribute auxiliary objects
    num_paras = init_dict['AUX']['num_paras']
    num_covars_cost = init_dict['AUX']['num_covars_cost']

    # Construct auxiliary objects
    G = np.concatenate((X, Z[:, 1:]), axis=1)

    # Select relevant values.
    if which == 'random':
        x0 = np.random.uniform(size=num_paras)

        # Variances
        x0[(-4)] = max(x0[(-4)], 0.01)
        x0[(-3)] = max(x0[(-3)], 0.01)

        # Correlations
        x0[(-2)] -= 0.5
        x0[(-1)] -= 0.5

    elif which == 'init':
        x0 = np.array(init_dict['AUX']['init_values'][:])

    elif which == 'auto':

        # Subsetting
        Y1, X1 = Y[D == 1], X[(D == 1), :]
        olsRslt = sm.OLS(Y1, X1).fit()

        # Extract results
        coeffs_treated = olsRslt.params
        sd_treated = np.array(np.sqrt(olsRslt.scale))

        # Subsetting
        Y0, X0 = Y[D == 0], X[(D == 0), :]
        olsRslt = sm.OLS(Y0, X0).fit()

        # Extract results
        coeffs_untreated = olsRslt.params
        sd_untreated = np.array(np.sqrt(olsRslt.scale))

        # Estimate choice model
        probitRslt = sm.Probit(D, G).fit()
        sd = init_dict['COST']['sd']
        coeffs = probitRslt.params*sd

        # Special treatment of cost intercept
        cost_int = coeffs_treated[0] - coeffs_untreated[0] - coeffs[0]

        # Collect results
        x0 = np.concatenate((coeffs_treated, coeffs_untreated))
        x0 = np.concatenate((x0, [cost_int], -coeffs[-(num_covars_cost - 1):]))
        x0 = np.concatenate((x0, [sd_treated, sd_untreated]))
        x0 = np.concatenate((x0, [0.00, 0.00]))

    else:
        raise AssertionError

    # Document starting values
    init_dict['AUX']['start_values'] = x0.copy()

    # Transform to real line
    x0 = _transform_start(x0)

    # Type conversion
    x0 = np.array(x0)

    # Quality assurance.
    assert (np.all(np.isfinite(x0)))

    # Finishing.
    return x0

def _transform_start(x):
    """ Transform starting values to cover the whole real line.
    """

    # Coefficients
    x[:(-4)] = x[:(-4)]

    # Variances
    x[(-4)] = np.log(x[(-4)])
    x[(-3)] = np.log(x[(-3)])

    # Correlations
    transform = (x[(-2)] + 1) / 2
    x[(-2)] = np.log(transform / (1.0 - transform))

    transform = (x[(-1)] + 1) / 2
    x[(-1)] = np.log(transform / (1.0 - transform))

    # Finishing
    return x

def _load_data(init_dict):
    """ Load dataset.
    """
    # Auxiliary objects
    num_covars_out = init_dict['AUX']['num_covars_out']
    num_covars_cost = init_dict['AUX']['num_covars_cost']
    num_agents = init_dict['BASICS']['agents']

    # Read dataset
    data = np.genfromtxt(init_dict['BASICS']['source'])

    # Reshaping, this ensure that the program also runs with just one agent
    # as otherwise only an vector is created. This creates problems for the
    # subsetting of the overall data into the components.
    data = np.array(data, ndmin=2)

    # Distribute data
    Y, D = data[:, 0], data[:, 1]

    X, Z = data[:, 2:(num_covars_out + 2)], data[:, -num_covars_cost:]

    # Finishing
    return Y, D, X, Z

def _distribute_arguments(args):
    """ Distribute arguments for evaluation of criterion function and some
        auxiliary parameters.
    """
    Y1_coeffs = np.array(args['TREATED']['all'])
    Y0_coeffs = np.array(args['UNTREATED']['all'])

    C_coeffs = np.array(args['COST']['all'])

    U1_sd = args['TREATED']['sd']
    U0_sd = args['UNTREATED']['sd']

    U1V_rho = args['DIST']['rho1']
    U0V_rho = args['DIST']['rho0']
    V_sd = args['COST']['sd']

    choice_coeffs = np.concatenate((Y1_coeffs - Y0_coeffs, - C_coeffs))

    # Finishing
    return Y1_coeffs, Y0_coeffs, C_coeffs, choice_coeffs, U1_sd, U0_sd, \
           U1V_rho, U0V_rho, V_sd
```

Let us evaluate the criterion function and see if that part is working.

```python
# Load model information
init_dict  = process('material/msc/init.ini')

Y, D, X, Z = _load_data(init_dict)

likl = _negative_log_likelihood(init_dict, Y, D, X, Z)

print 'Evaluation of criterion function', likl
```

    Evaluation of criterion function -2.44668227139

Finally, we create a function for the estimation. It builds on all our previously defined elements.

```python
def estimate(init_dict):
    """ Estimate our version of the generalized Roy model.
    """
    # Antibugging
    assert (isinstance(init_dict, dict))

    # Load dataset
    Y, D, X, Z = _load_data(init_dict)

    # Create auxiliary objects
    start = init_dict['ESTIMATION']['start']
    maxiter = init_dict['ESTIMATION']['maxiter']

    optimizer = init_dict['ESTIMATION']['algorithm']
    num_covars_out = init_dict['AUX']['num_covars_out']

    # Initialize different starting values
    x0 = _get_start(start, init_dict, Y, D, X, Z)

    # Select optimizer
    if optimizer == 'nm':

        optimizer = 'Nelder-Mead'

    elif optimizer == 'bfgs':

        optimizer = 'BFGS'

    # Provide additional arguments to the optimizer
    opts = dict()

    opts['maxiter'] = maxiter

    # Run optimization or just evaluate function at starting values
    if maxiter == 0:

        # Collect maximization arguments.
        rslt = _distribute_parameters(np.array(x0), init_dict, num_covars_out)

        # Calculate likelihood according to user's request
        likl = _negative_log_likelihood(rslt, Y, D, X, Z)

        # Compile results
        x_rslt, fun, success = x0, likl, False

    else:

        # Check out the SciPy documentation for details about the interface
        # to the `minimize' function that provides a convenient interface to
        # a variety of alternative maximization algorithms. You will also
        # find information about the return information.
        opt_rslt = minimize(_max_interface, x0,
                            args=(Y, D, X, Z, init_dict),
                            method=optimizer, options=opts)

        # Compile results
        x_rslt, fun = opt_rslt['x'], opt_rslt['fun']
        success = opt_rslt['success']

    # Tranformation to internal parameters
    rslt = _distribute_parameters(x_rslt, init_dict, num_covars_out)

    rslt['fval'], rslt['success'] = fun, success

    # Finishing
    return rslt
```

Now, let us actually run some estimations. Our current implementation allow to easily switch between alternative optimzation algorithms and starting values.

```python
# Process model specification
init_dict = process('material/msc/init.ini')

# Simulate a synthetic sample
simulate(init_dict)

# Estimate the generalized Roy model
for algorithm in ['bfgs', 'nm']:

    init_dict['ESTIMATION']['algorithm'] = algorithm

    for start in ['random', 'init']:

        init_dict['ESTIMATION']['start'] = start

        # Monitoring
        print '\n\n Current Request \n'

        print ' Algorithm: ', algorithm

        print ' Start:     ', start

        # Run estimation
        rslt = estimate(init_dict)

        # Inspect subset of results
        print ' Variances: ', rslt['TREATED']['sd'], rslt['UNTREATED']['sd']
```

     Current Request

     Algorithm:  bfgs
     Start:      random
     Variances:  0.0854101015863 0.580645990592

     Current Request

     Algorithm:  bfgs
     Start:      init
     Variances:  0.02 0.02

     Current Request

     Algorithm:  nm
     Start:      random
     Variances:  0.981374779551 0.694354634872

     Current Request

     Algorithm:  nm
     Start:      init
     Variances:  0.02 0.02

## Inspection of Results

After conduction an estimation, the focus shifts back to the economic interpretation of our results. As we often have to study the results from hundreds of different estimation runs, it is convenient to have a function set up that produces the main objects of interest.

```python
def inspect(rslt, init_dict):
    """ This function simulates a sample from the estimates of the model
        and reports the average effects of treatment in a file.
    """
    # Antibugging
    assert (isinstance(rslt, dict))
    assert (isinstance(init_dict, dict))

    # Update results
    modified_init = copy.deepcopy(init_dict)

    for key_ in rslt.keys():

        if key_ in ['fval', 'success']:
            continue

        for subkey in rslt[key_].keys():

            modified_init[key_][subkey] = rslt[key_][subkey]

    # Modified dataset
    modified_init['BASICS']['file'] = 'simulated.grm.txt'

    # Simulate from estimation results
    Y1, Y0, D = simulate(modified_init, True)

    # Calculate the average treatment effectsa
    B = Y1 - Y0

    effects = []
    effects += [np.mean(B)]
    effects += [np.mean(B[D == 1])]
    effects += [np.mean(B[D == 0])]

    # Print selected results to file
    with open('results.grm.txt', 'w') as file_:

        file_.write('\n softEcon: Generalized Roy Model')
        file_.write('\n -------------------------------\n')

        # Average effects of treatment
        fmt = '     {0:<5}{1:10.2f}\n\n'

        file_.write('\n Average Treatment Effects\n\n')

        for i, label in enumerate(['ATE', 'TT', 'TUT']):

            str_ = fmt.format(label, effects[i])

            file_.write(str_)

        file_.write('\n Parameters\n\n')
        file_.write('     Start    Finish\n\n')

        num_paras = init_dict['AUX']['num_paras']

        # Structural parameters
        x0, x = init_dict['AUX']['start_values'], rslt['AUX']['x_internal']

        fmt = '{0:10.2f}{1:10.2f}\n'

        for i in range(num_paras):

            str_ = fmt.format(x0[i], x[i])

            file_.write(str_)

```

Let us try it out.

```python
# Process initialization file
init_dict = process('material/msc/init.ini')

# Simulate synthetic sample
simulate(init_dict)

# Estimate model
rslt = estimate(init_dict)

# Write results
inspect(rslt, init_dict)

# Inspect the results
%cat results.grm.txt
```

     softEcon: Generalized Roy Model
     -------------------------------

     Average Treatment Effects

         ATE       -0.40

         TT        -0.46

         TUT       -0.29

     Parameters

         Start    Finish

          0.50      0.50
          0.20      0.20
          0.52      0.52
          0.52      0.52
          0.10      0.10
          0.12      0.12
          0.30      0.30
          0.52      0.52
          0.20      0.20
          0.12      0.12
          0.25      0.25
          0.35      0.35
          0.02      0.02
          0.02      0.02
          0.10      0.10
          0.12      0.12

## Cleanup

```python
# Create list of all files generated the notebook
files = glob.glob('*.grm.*')

# Remove files
for file_ in files:

    os.remove(file_)
```

**Formatting**

```python
import urllib; from IPython.core.display import HTML
HTML(urllib.urlopen('http://bit.ly/1OKmNHN').read())
```

<link href='http://fonts.googleapis.com/css?family=Alegreya+Sans:100,300,400,500,700,800,900,100italic,300italic,400italic,500italic,700italic,800italic,900italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Arvo:400,700,400italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=PT+Mono' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Shadows+Into+Light' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Philosopher:400,700,400italic,700italic' rel='stylesheet' type='text/css'>

<style>

@font-face {
    font-family: "Computer Modern";
    src: url('http://mirrors.ctan.org/fonts/cm-unicode/fonts/otf/cmunss.otf');
}

#notebook_panel { /* main background */
    background: #888;
    color: #f6f6f6;
}

div.cell { /* set cell width to about 80 chars */
    width: 1000px;
}

div #notebook { /* centre the content */
    background: #fff; /* white background for content */
    width: 1200px;
    margin: auto;
    padding-left: 1em;
}

#notebook li { /* More space between bullet points */
margin-top:0.8em;
}

/* draw border around running cells */
div.cell.border-box-sizing.code_cell.running {
    border: 3px solid #111;
}

/* Put a solid color box around each cell and its output, visually linking them together */
div.cell.code_cell {
    background-color: rgba(171,165,131,0.3);
    border-radius: 10px; /* rounded borders */
    padding: 1em;
    margin-top: 1em;
}

div.text_cell_render{
    font-family: 'Arvo' sans-serif;
    line-height: 130%;
    font-size: 150%;
    width:900px;
    margin-left:auto;
    margin-right:auto;
}

/* Formatting for header cells */
.text_cell_render h1 {
    font-family: 'Philosopher', sans-serif;
    font-weight: 400;
    font-size: 32pt;
    line-height: 100%;
    color: rgb(12,85,97);
    margin-bottom: 0.1em;
    margin-top: 0.1em;
    display: block;
}

.text_cell_render h2 {
    font-family: 'Philosopher', serif;
    font-weight: 700;
    font-size: 24pt;
    line-height: 100%;
    color: rgb(171,165,131);
    margin-bottom: 0.1em;
    margin-top: 0.1em;
    display: block;
}

.text_cell_render h3 {
    font-family: 'Philosopher', serif;
    margin-top:12px;
    margin-bottom: 3px;
    font-style: italic;
    color: rgb(95,92,72);
}

.text_cell_render h4 {
    font-family: 'Philosopher', serif;
}

.text_cell_render h5 {
    font-family: 'Alegreya Sans', sans-serif;
    font-weight: 300;
    font-size: 16pt;
    color: grey;
    font-style: italic;
    margin-bottom: .1em;
    margin-top: 0.1em;
    display: block;
}

.text_cell_render h6 {
    font-family: 'PT Mono', sans-serif;
    font-weight: 300;
    font-size: 10pt;
    color: grey;
    margin-bottom: 1px;
    margin-top: 1px;
}

.CodeMirror{
        font-family: "PT Mono";
        font-size: 120%;
}

</style>


```python
import numpy as np
import matplotlib.pyplot as plt
import xlrd
import sympy as sp
from matplotlib.ticker import MultipleLocator, FormatStrFormatter ,FuncFormatter
import pandas as pd
from scipy import optimize
from scipy.interpolate import spline
import os

strt=["1tally  14        nps =   ","1tally  24        nps =   ",
      "1tally  34        nps =   ","1tally  44        nps =   ",
      "1tally  54        nps =   ","1tally  64        nps =   ",]#14 24
un=["g/kg","g/kg","mm","mm","g/cm$^3$","g/cm$^3$","mm","$^{\circ}$$C$"]
st=["o","h","H","s","d","^","<",">","1","2","p"]
c=["r","b","g","c","m","y","k","blueviolet","lightgreen","darkmagenta","lightgreen"]
la=["地层水矿化度($g/kg$)","井眼矿化度($g/kg$)","厚度($mm$)","厚度($mm$)","泥浆密度($g/cm^3$)","泥浆密度($g/cm^3$)","井径($mm$)",
   "井温($^{\circ}$$C$)"]
col=["孔隙度（pu）","近中子计数器-1","近中子计数器-2","远中子计数器-1","远中子计数器-2","远中子计数器-1","远中子计数器-2",
     "第1组远中子计数和","第2组远中子计数和","第1组近远中子比","第2组近远中子比","总近中子计数和","误差","总远中子计数和","误差",
     "总近远中子比"]
def linersmooth7(a): #七点平滑
    b=[]
    N=len(a)-1
    for i in range(3,N-2): #不要前三个和最后三个点
        b.append((130*a[i]+140*(a[i-1]+a[i+1])+65*(a[i-2]+a[i+2])+40*(a[i-3]+a[i+3]))/610) #公式形式
    return b
def dele(a): #y轴去掉了前三个和后三个点，x轴数据也要去掉前三个和后三个
    b=a.copy()
    for i in range(0,3):
        del b[0]
        b.pop()
    return b
```

```python
 #用flag给指示是否要将标准情况放入其中
#存在用到标准曲线的情况
```

```python
data=pd.read_excel("E:\王馨奕\随钻-19-4\MCNP程序及结果\\02-中子图版\\00-标准曲线\\标准曲线.xlsx",index_col=0)
y0=data["总近远中子比"].values#取得标准刻度井计算结果
x0=data["孔隙度（pu）"].values
```

```python

```

```python
def rf(flag,x0,excelname):
    path=os.getcwd()
    tn=[[] for i in range(7)]#每个中子计数器的计数值,第一个是记孔隙度值
    jx=[]#文件名-环境变量值
    aimf=[]#目标文件
    ind=[]
    err=[[] for i in range(7)]
    files= os.listdir(path)#读文件名
    for file in files: #遍历文件夹
         if os.path.isdir(file):
                if (file!=".ipynb_checkpoints"):#判断是否是文件夹，是文件夹才打开
                    jx.append(file)
    jx=np.array(jx)#将字符转换为数字，去掉引号
    jx=jx.astype(np.int)#环境因素
    jx.sort()
    for i in range(len(jx)):
        path_dir1="%s"%(jx[i])
        path_list=os.listdir(path_dir1)
        for filename in path_list:
            i=0
            if(i<11):
                if os.path.splitext(filename)[1] == '.o':
                    name=os.path.splitext(filename)[0]
                    aimf.append(name)
                    i+=1
        aimf=np.array(aimf)#将字符转换为数字，去掉引号
        aimf=aimf.astype(np.int)#环境因素
        aimf.sort()
#         print(aimf)
        for i in range(11):#读到的文件文件名做index将中子计数值存储下来
            ind.append(path_dir1)
            tn[0].append(x0[i])#
            err[0].append(x0[i])
            path_dir=path+"\\"+path_dir1+"\\"+"%s.o"%(aimf[i])#读到了含有输出文件的文件路径
            i=0
            j=0
            f=open(path_dir,"r")#读取该文件
            for n in range(6):
                for line in f:
                    if(line.find(strt[n])!=-1):
                        i=1
                    if(i==1):
                        j+=1
                    if(j==11):
                        m=line.strip()
                        tn[n+1].append(m.split()[0])
                        err[n+1].append(m.split()[1])
                        i=0
                        j=0
                        break
        aimf=[]
    nu=len(tn[0])
#     for i in range(len(tn)):
#         print(len(tn[i]))
    tn=np.array(tn)#将字符转换为数字，去掉引号
    tn=tn.astype(np.float)
    err=np.array(err)#将字符转换为数字，去掉引号
    err=err.astype(np.float)
    t0=tn.T#nu*6维矩阵转换
    t0=t0.tolist()#numpy to list
    print("已读取到路径%s:文件%s"%(path,jx))
    print("读到%s份输出文件"%(nu))
    for i in range(nu):
#         t0[i].append(tn[0][i])#参考组1
#         t0[i].append(tn[1][i])#参考组2
        t0[i].append(tn[3][i]+tn[5][i])#1组远
        t0[i].append(tn[4][i]+tn[6][i])#2组远
        t0[i].append((tn[1][i])/(tn[3][i]+tn[5][i]))#1组比
        t0[i].append((tn[2][i])/(tn[4][i]+tn[6][i]))#2组比
        t0[i].append(tn[2][i]+tn[1][i])#总近和
        t0[i].append(err[2][i]+err[1][i])#总近和误差
        t0[i].append(tn[6][i]+tn[3][i]+tn[4][i]+tn[5][i])#总远和
        t0[i].append(err[6][i]+err[3][i]+err[4][i]+err[5][i])#总远和
        t0[i].append((tn[2][i]+tn[1][i])/(tn[6][i]+tn[3][i]+tn[4][i]+tn[5][i]))#总比
        R=(tn[2][i]+tn[1][i])/(tn[6][i]+tn[3][i]+tn[4][i]+tn[5][i])
    print("已计算三组近远中子比")
    print("数据存入%s×%s矩阵"%(len(t0),len(t0[0])))
    data=pd.DataFrame(data=t0,index=ind,columns=col)
    data.to_excel("%s.xlsx"%(excelname),encoding="utf_8_sig")
    print("存文件:%s.xlsx"%(excelname))
    if (len(tn[0])==len(t0)):
        print("正确录入数据")
    else:
        print("未正确录入数据")
    return (t0,jx,len(jx),err)
```

```python
t0,jx,f,err=rf(0,x0,"间隙300")
```

    已读取到路径E:\王馨奕\随钻-19-4\MCNP程序及结果\02-中子图版\4间隙校正\300:文件[ 0  5 10 15 20 25 30 40 50]
    读到99份输出文件
    已计算三组近远中子比
    数据存入99×16矩阵
    存文件:间隙300.xlsx
    正确录入数据

```python
#画图
```

```python
def f2(x,A,B,C):
    return A*x*x+B*x+C
def pic1(name_pic,flag):
    x=np.arange(0,55,5)
    D=jx
    d=len(D)
    un=["g/kg","g/kg","mm","mm","g/cm$^3$","g/cm$^3$","mm","$^{\circ}$$C$"]
    st=["o","h","H","s","d","^","<",">","1","2","p"]
    c=["r","b","g","c","m","y","k","blueviolet","lightgreen","darkmagenta","lightgreen"]
    plt.figure(figsize=(6.67, 5), dpi=100) #图片大小：6.75cm×5cm
    plt.rcParams['font.sans-serif']=['SimHei'] #中文默认黑体
    plt.rcParams['axes.unicode_minus']=False #显示负号
    for i in range(6):
        A2,B2,C2=optimize.curve_fit(f2,x,y1[i])[0]
        x2=np.arange(0,50,0.5)
        y2=A2*x2*x2+B2*x2+C2
        plt.plot(x2,y2,color="%s"%(c[i]))
        plt.plot(x,y1[i],"%s"%(st[i]),color="%s"%(c[i]),label="%.1f%s:%.4f×x$^2$%+.2f×x%+.2f"%(D[i],un[flag],A2,B2,C2))
    plt.legend()
    plt.xticks(fontproperties='Times New Roman',fontsize=13)
    plt.yticks(fontproperties='Times New Roman',fontsize=13)
    plt.xlabel("孔隙度($pu$)",fontproperties='SimHei',fontsize=13)
    plt.ylabel("近远热中子比值",fontproperties='SimHei',fontsize=13)
    plt.grid(linewidth='1.5' ,linestyle='-',color="black",alpha=0.5)
    plt.savefig('%s趋势1.png'%(name_pic), format='png', dpi=100,pad_inches=0.0,bbox_inches = 'tight')
    plt.savefig('%s趋势1.svg'%(name_pic), format='svg', dpi=100,pad_inches=0.0,bbox_inches = 'tight')
    plt.show()
    plt.figure(figsize=(6.67, 5), dpi=100) #图片大小：6.75cm×5cm
    plt.rcParams['font.sans-serif']=['SimHei'] #中文默认黑体
    plt.rcParams['axes.unicode_minus']=False #显示负号
    x=np.arange(0,55,5)
    for i in range(0,11):
        A,B,C=optimize.curve_fit(f2,D,[x[i] for x in y1])[0]# l=[0---250] 取不同空隙度对应的比值 横行
        x2=np.arange(min(D),max(D)+0.01,0.01)
        y2=f2(x2,A,B,C)
        plt.plot(x2,y2,color="%s"%(c[i])) #先画拟合的直线，再画实际结果
        plt.plot(D,[x[i] for x in y1],"%s"%(st[i]),color="%s"%(c[i]),label="%spu"%(x[i]),markersize=7)
    plt.xticks(fontproperties='Times New Roman',fontsize=17)
    plt.yticks(fontproperties='Times New Roman',fontsize=17)
    plt.xlabel("%s"%(la[flag]),fontproperties='SimHei',fontsize=17)
    plt.ylabel("近远热中子比值",fontproperties='SimHei',fontsize=17)
    plt.legend(bbox_to_anchor=(1,1.001),borderaxespad=0.1,edgecolor="black",handletextpad=-0.3)
    plt.grid(linewidth='1.5' ,color="black",alpha=0.5)
    plt.savefig('%s趋势2.png'%(name_pic), format='png', dpi=100,pad_inches=0.0,bbox_inches = 'tight')
    plt.savefig('%s趋势2.svg'%(name_pic), format='svg', dpi=100,pad_inches=0.0,bbox_inches = 'tight')
    plt.show()
```

```python
len(t0[0])
```

    16

```python
y1=[[] for i in range(f)]
for m in range(f):
    for i in range(len(x0)):
        y1[m].append(t0[i+m*len(x0)][15])#读到的y[0]为标准曲线的值
#     print(len(y1[m]))
pic1("间隙300",3)
```

```python
def dom(name_pic,flag):
    #计算视孔隙度
    ynew=[]
    x=np.arange(0,55,5)
    for i in range(0,len(jx)):
        co = np.polyfit(y1[i],x,2) # 用多项式拟合
        A2=co[0]
        B2=co[1]
        C2=co[2]
        x2=np.arange(1,20,0.2)
        re = np.polyval(co, x2)
        X=sp.Symbol('X')
        temp=[]
        for j in range(len(x)): #解方程 从近远中子比解孔隙度
            y3=sp.solve(A2*X*X+B2*X+C2-x[j],X) #求出来比值 solve函数
            temp.append(y3[1])
        ynew.append(temp)
    for i in range(len(jx)):
        ynew[i][0]=sp.re(ynew[i][0])+sp.im(ynew[i][0])
    ynew=np.array(ynew)
    pnew=[]
    A,B,C=optimize.curve_fit(f2,y0,x)[0]
    for i in range(len(jx)):
        pnew.append(A*ynew[i]*ynew[i]+B*ynew[i]+C)#带入标准曲线得到视孔隙度7个井径
    plt.figure(figsize=(6.67, 2.5), dpi=100) #图片大小：6.75cm×5cm
    plt.rcParams['font.sans-serif']=['SimHei'] #中文默认黑体
    plt.rcParams['axes.unicode_minus']= False #显示负号
    st=["o","h","H","s","d","^","<",">","1","2","p"]
    c=["r","b","g","c","m","y","k","blueviolet","lightgreen","darkmagenta","lightgreen"]
    pnew=np.array(pnew)
    l=jx
    for i in range(0,11):
        k=[x[i] for x in pnew]
        k=np.array(k,dtype='float')
        plt.plot(k,l,"%s"%(st[i]),color="%s"%(c[i]),markersize=5)#,"%s"%(st[i]),color="%s"%(c[i]),markersize=5
        z1 = np.polyfit(l,k,2) # 用3次多项式拟合
        z = np.polyval(z1, l)
        plt.plot(z,l,"-",color="%s"%(c[i]))
    plt.xticks(fontproperties='Times New Roman',fontsize=13)
    plt.yticks(fontproperties='Times New Roman',fontsize=13)
    plt.ylabel("%s"%(la[flag]),fontproperties='SimHei',fontsize=13)
    plt.xlabel("孔隙度($pu$)[视石灰岩]",fontproperties='SimHei',fontsize=13)
    #plt.legend(prop=font1,bbox_to_anchor=(1,1.001),borderaxespad=0.1,edgecolor="black",handletextpad=-0.3)
    plt.xlim(0,50)#x轴范围
    plt.ylim(min(l),max(l))
    plt.grid(linewidth='1.5' ,color="black",alpha=0.5)
    # plt.tight_layout()
    plt.savefig('%s校正.png'%(name_pic), format='png', dpi=100,bbox_inches = 'tight')
    plt.savefig('%s校正.svg'%(name_pic), format='svg', dpi=100,bbox_inches = 'tight')
    plt.show()
    plt.figure(figsize=(6.67, 2.5), dpi=100) #图片大小：6.75cm×5cm
    plt.rcParams['font.sans-serif']=['SimHei'] #中文默认黑体
    plt.rcParams['axes.unicode_minus']= False #显示负号
    st=["o","h","H","s","d","^","<",">","1","2","p"]
    c=["r","b","g","c","m","y","k","blueviolet","lightgreen","darkmagenta","lightgreen"]
    pnew=np.array(pnew)
    l=jx
    for i in range(0,11):
        k=[x[i] for x in pnew]
        k=np.array(k,dtype='float')
        z1 = np.polyfit(l,k,2) # 用3次多项式拟合
        z = np.polyval(z1, l)
        plt.plot(z,l,"-",color="#054E9F")
    plt.xticks(fontproperties='Times New Roman',fontsize=13)
    plt.yticks(fontproperties='Times New Roman',fontsize=13)
    plt.ylabel("%s"%(la[flag]),fontproperties='SimHei',fontsize=13)
    plt.xlabel("孔隙度($pu$)[视石灰岩]",fontproperties='SimHei',fontsize=13)
    #plt.legend(prop=font1,bbox_to_anchor=(1,1.001),borderaxespad=0.1,edgecolor="black",handletextpad=-0.3)
    plt.xlim(0,50)#x轴范围
    plt.ylim(min(l),max(l))
    plt.grid(linewidth='1.5' ,color="black",alpha=0.5)
    # plt.tight_layout()
    plt.savefig('%s校正2.png'%(name_pic), format='png', dpi=100,bbox_inches = 'tight')
    plt.savefig('%s校正2.svg'%(name_pic), format='svg', dpi=100,bbox_inches = 'tight')
    plt.show()
```

```python
dom("间隙250",1)
```

```python

```


Natural Language Processing From Scratch

Bruno Gonçalves

bgoncalves@gmail.com

@bgoncalves

# Lesson I - Text representation

In this lesson we will see in some details how we can best represent text in our application. Let's start by importing the modules we will be using:

```python
import string
from collections import Counter
from pprint import pprint
import gzip
import matplotlib.pyplot as plt
import numpy as np

%matplotlib inline
```

We choose a well known nursery rhyme, that has the added distinction of having been the first audio ever recorded, to be the short snippet of text that we will use in our examples:

```python
text = """Mary had a little lamb, little lamb,
    little lamb. Mary had a little lamb
    whose fleece was white as snow.
    And everywhere that Mary went
    Mary went, Mary went. Everywhere
    that Mary went,
    The lamb was sure to go"""
```

## Tokenization

The first step in any analysis is to tokenize the text. What this means is that we will extract all the individual words in the text. For the sake of simplicity, we will assume that our text is well formed and that our words are delimited either by white space or punctuation characters.

```python
def extract_words(text):
    temp = text.split() # Split the text on whitespace
    text_words = []

    for word in temp:
        # Remove any punctuation characters present in the beginning of the word
        while word[0] in string.punctuation:
            word = word[1:]

        # Remove any punctuation characters present in the end of the word
        while word[-1] in string.punctuation:
            word = word[:-1]

        # Append this word into our list of words.
        text_words.append(word.lower())

    return text_words
```

After this step we now have our text represented as an array of individual, lowercase, words:

```python
text_words = extract_words(text)
print(text_words)
```

    ['mary', 'had', 'a', 'little', 'lamb', 'little', 'lamb', 'little', 'lamb', 'mary', 'had', 'a', 'little', 'lamb', 'whose', 'fleece', 'was', 'white', 'as', 'snow', 'and', 'everywhere', 'that', 'mary', 'went', 'mary', 'went', 'mary', 'went', 'everywhere', 'that', 'mary', 'went', 'the', 'lamb', 'was', 'sure', 'to', 'go']

As we saw during the video, this is a wasteful way to represent text. We can be much more efficient by representing each word by a number

```python
word_dict = {}
word_list = []
vocabulary_size = 0
text_tokens = []

for word in text_words:
    # If we are seeing this word for the first time, create an id for it and added it to our word dictionary
    if word not in word_dict:
        word_dict[word] = vocabulary_size
        word_list.append(word)
        vocabulary_size += 1

    # add the token corresponding to the current word to the tokenized text.
    text_tokens.append(word_dict[word])
```

When we were tokenizing our text, we also generated a dictionary **word_dict** that maps words to integers and a **word_list** that maps each integer to the corresponding word.

```python
print("Word list:", word_list, "\n\n Word dictionary:")
pprint(word_dict)
```

    Word list: ['mary', 'had', 'a', 'little', 'lamb', 'whose', 'fleece', 'was', 'white', 'as', 'snow', 'and', 'everywhere', 'that', 'went', 'the', 'sure', 'to', 'go']

     Word dictionary:
    {'a': 2,
     'and': 11,
     'as': 9,
     'everywhere': 12,
     'fleece': 6,
     'go': 18,
     'had': 1,
     'lamb': 4,
     'little': 3,
     'mary': 0,
     'snow': 10,
     'sure': 16,
     'that': 13,
     'the': 15,
     'to': 17,
     'was': 7,
     'went': 14,
     'white': 8,
     'whose': 5}

These two datastructures already proved their usefulness when we converted our text to a list of tokens.

```python
print(text_tokens)
```

    [0, 1, 2, 3, 4, 3, 4, 3, 4, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 0, 14, 0, 14, 0, 14, 12, 13, 0, 14, 15, 4, 7, 16, 17, 18]

Unfortunately, while this representation is convenient for memory reasons it has some severe limitations. Perhaps the most important of which is the fact that computers naturally assume that numbers can be operated on mathematically (by addition, subtraction, etc) in a way that doesn't match our understanding of words.

## One-hot encoding

One typical way of overcoming this difficulty is to represent each word by a one-hot encoded vector where every element is zero except the one corresponding to a specific word.

```python
def one_hot(word, word_dict):
    """
        Generate a one-hot encoded vector corresponding to *word*
    """

    vector = np.zeros(len(word_dict))
    vector[word_dict[word]] = 1

    return vector
```

So, for example, the word "fleece" would be represented by:

```python
fleece_hot = one_hot("fleece", word_dict)
print(fleece_hot)
```

    [0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]

This vector has every element set to zero, except element 6, since:

```python
print(word_dict["fleece"])
fleece_hot[6] == 1
```

    6

    True

## Bag of words

We can now use the one-hot encoded vector for each word to produce a vector representation of our original text, by simply adding up all the one-hot encoded vectors:

```python
text_vector1 = np.zeros(vocabulary_size)

for word in text_words:
    hot_word = one_hot(word, word_dict)
    text_vector1 += hot_word

print(text_vector1)
```

    [6. 2. 2. 4. 5. 1. 1. 2. 1. 1. 1. 1. 2. 2. 4. 1. 1. 1. 1.]

In practice, we can also easily skip the encoding step at the word level by using the *word_dict* defined above:

```python
text_vector = np.zeros(vocabulary_size)

for word in text_words:
    text_vector[word_dict[word]] += 1

print(text_vector)
```

    [6. 2. 2. 4. 5. 1. 1. 2. 1. 1. 1. 1. 2. 2. 4. 1. 1. 1. 1.]

Naturally, this approach is completely equivalent to the previous one and has the added advantage of being more efficient in terms of both speed and memory requirements.

This is known as the __bag of words__ representation of the text. It should be noted that these vectors simply contains the number of times each word appears in our document, so we can easily tell that the word *mary* appears exactly 6 times in our little nursery rhyme.

```python
text_vector[word_dict["mary"]]
```

    6.0

A more pythonic (and efficient) way of producing the same result is to use the standard __Counter__ module:

```python
word_counts = Counter(text_words)
pprint(word_counts)
```

    Counter({'mary': 6,
             'lamb': 5,
             'little': 4,
             'went': 4,
             'that': 2,
             'had': 2,
             'a': 2,
             'everywhere': 2,
             'was': 2,
             'white': 1,
             'snow': 1,
             'sure': 1,
             'to': 1,
             'the': 1,
             'whose': 1,
             'go': 1,
             'as': 1,
             'fleece': 1,
             'and': 1})

From which we can easily generate the __text_vector__ and __word_dict__ data structures:

```python
items = list(word_counts.items())

# Extract word dictionary and vector representation
word_dict2 = dict([[items[i][0], i] for i in range(len(items))])
text_vector2 = [items[i][1] for i in range(len(items))]
```

And let's take a look at them:

```python
print("Text vector:", text_vector2, "\n\nWord dictionary:")
pprint(word_dict2)
```

    Text vector: [1, 2, 4, 5, 1, 2, 1, 4, 6, 1, 1, 1, 1, 1, 1, 2, 1, 2, 2]

    Word dictionary:
    {'a': 15,
     'and': 16,
     'as': 13,
     'everywhere': 17,
     'fleece': 14,
     'go': 12,
     'had': 5,
     'lamb': 3,
     'little': 2,
     'mary': 8,
     'snow': 4,
     'sure': 6,
     'that': 1,
     'the': 10,
     'to': 9,
     'was': 18,
     'went': 7,
     'white': 0,
     'whose': 11}

The results using this approach are slightly different than the previous ones, because the words are mapped to different integer ids but the corresponding values are the same:

```python
for word in word_dict.keys():
    if text_vector[word_dict[word]] != text_vector2[word_dict2[word]]:
        print("Error!")
```

As expected, there are no differences!

## Term Frequency

The bag of words vector representation introduced above relies simply on the frequency of occurence of each word. Following a long tradition of giving fancy names to simple ideas, this is known as __Term Frequency__.

Intuitively, we expect the the frequency with which a given word is mentioned should correspond to the relevance of that word for the piece of text we are considering. For example, **Mary** is a pretty important word in our little nursery rhyme and indeed it is the one that occurs the most often:

```python
sorted(items, key=lambda x:x[1], reverse=True)
```

    [('mary', 6),
     ('lamb', 5),
     ('little', 4),
     ('went', 4),
     ('that', 2),
     ('had', 2),
     ('a', 2),
     ('everywhere', 2),
     ('was', 2),
     ('white', 1),
     ('snow', 1),
     ('sure', 1),
     ('to', 1),
     ('the', 1),
     ('whose', 1),
     ('go', 1),
     ('as', 1),
     ('fleece', 1),
     ('and', 1)]

However, it's hard to draw conclusions from such a small piece of text. Let us consider a significantly larger piece of text, the first 100 MB of the english Wikipedia from: http://mattmahoney.net/dc/textdata. For the sake of convenience, text8.gz has been included in this repository in the **data/** directory. We start by loading it's contents into memory as an array of words:

```python
data = []

for line in gzip.open("data/text8.gz", 'rt'):
    data.extend(line.strip().split())
```

Now let's take a look at the most common words in this large corpus:

```python
counts = Counter(data)

sorted_counts = sorted(list(counts.items()), key=lambda x:x[1], reverse=True)

for word, count in sorted_counts[:10]:
    print(word, count)
```

    the 1061396
    of 593677
    and 416629
    one 411764
    in 372201
    a 325873
    to 316376
    zero 264975
    nine 250430
    two 192644

Surprisingly, we find that the most common words are not particularly meaningful. Indeed, this is a common occurence in Natural Language Processing. The most frequent words are typically auxiliaries required due to gramatical rules.

On the other hand, there is also a large number of words that occur very infrequently as can be easily seen by glancing at the word freqency distribution.

```python
dist = Counter(counts.values())
dist = list(dist.items())
dist.sort(key=lambda x:x[0])
dist = np.array(dist)

norm = np.dot(dist.T[0], dist.T[1])

plt.loglog(dist.T[0], dist.T[1]/norm)
plt.xlabel("count")
plt.ylabel("P(count)")
plt.title("Word frequency distribution")
```

## Stopwords

One common technique to simplify NLP tasks is to remove what are known as Stopwords, words that are very frequent but not meaningful. If we simply remove the most common 100 words, we significantly reduce the amount of data we have to consider while losing little information.

```python
stopwords = set([word for word, count in sorted_counts[:100]])

clean_data = []

for word in data:
    if word not in stopwords:
        clean_data.append(word)

print("Original size:", len(data))
print("Clean size:", len(clean_data))
print("Reduction:", 1-len(clean_data)/len(data))
```

    Original size: 17005207
    Clean size: 9006229
    Reduction: 0.470384041782026

Wow, our dataset size was reduced almost in half!

In practice, we don't simply remove the most common words in our corpus but rather a manually curate list of stopwords. Lists for dozens of languages and applications can easily be found online.

## Term Frequency/Inverse Document Frequency

One way of determining of the relative importance of a word is to see how often it appears across multiple documents. Words that are relevant to a specific topic are more likely to appear in documents about that topic and much less in documents about other topics. On the other hand, less meaningful words (like **the**) will be common across documents about any subject.

To measure the document frequency of a word we will need to have multiple documents. For the sake of simplicity, we will treat each sentence of our nursery rhyme as an individual document:

```python
corpus_text = text.split('.')
corpus_words = []

for document in corpus_text:
    doc_words = extract_words(document)
    corpus_words.append(doc_words)
```

Now our corpus is represented as a list of word lists, where each list is just the word representation of the corresponding sentence:

```python
pprint(corpus_words)
```

    [['mary', 'had', 'a', 'little', 'lamb', 'little', 'lamb', 'little', 'lamb'],
     ['mary',
      'had',
      'a',
      'little',
      'lamb',
      'whose',
      'fleece',
      'was',
      'white',
      'as',
      'snow'],
     ['and', 'everywhere', 'that', 'mary', 'went', 'mary', 'went', 'mary', 'went'],
     ['everywhere',
      'that',
      'mary',
      'went',
      'the',
      'lamb',
      'was',
      'sure',
      'to',
      'go']]

Let us now calculate the number of documents in which each word appears:

```python
document_count = {}

for document in corpus_words:
    word_set = set(document)

    for word in word_set:
        document_count[word] = document_count.get(word, 0) + 1

pprint(document_count)
```

    {'a': 2,
     'and': 1,
     'as': 1,
     'everywhere': 2,
     'fleece': 1,
     'go': 1,
     'had': 2,
     'lamb': 3,
     'little': 2,
     'mary': 4,
     'snow': 1,
     'sure': 1,
     'that': 2,
     'the': 1,
     'to': 1,
     'was': 2,
     'went': 2,
     'white': 1,
     'whose': 1}

As we can see, the word __Mary__ appears in all 4 of our documents, making it useless when it comes to distinguish between the different sentences. On the other hand, words like __white__ which appear in only one document are very discriminative. Using this approach we can define a new quantity, the ___Inverse Document Frequency__ that tells us how frequent a word is across the documents in a specific corpus:

```python
def inv_doc_freq(corpus_words):
    number_docs = len(corpus_words)

    document_count = {}

    for document in corpus_words:
        word_set = set(document)

        for word in word_set:
            document_count[word] = document_count.get(word, 0) + 1

    IDF = {}

    for word in document_count:
        IDF[word] = np.log(number_docs/document_count[word])

    return IDF
```

Where we followed the convention of using the logarithm of the inverse document frequency. This has the numerical advantage of avoiding to have to handle small fractional numbers.

We can easily see that the IDF gives a smaller weight to the most common words and a higher weight to the less frequent:

```python
IDF = inv_doc_freq(corpus_words)

pprint(IDF)
```

    {'a': 0.6931471805599453,
     'and': 1.3862943611198906,
     'as': 1.3862943611198906,
     'everywhere': 0.6931471805599453,
     'fleece': 1.3862943611198906,
     'go': 1.3862943611198906,
     'had': 0.6931471805599453,
     'lamb': 0.28768207245178085,
     'little': 0.6931471805599453,
     'mary': 0.0,
     'snow': 1.3862943611198906,
     'sure': 1.3862943611198906,
     'that': 0.6931471805599453,
     'the': 1.3862943611198906,
     'to': 1.3862943611198906,
     'was': 0.6931471805599453,
     'went': 0.6931471805599453,
     'white': 1.3862943611198906,
     'whose': 1.3862943611198906}

As expected **Mary** has the smallest weight of all words 0, meaning that it is effectively removed from the dataset. You can consider this as a way of implicitly identify and remove stopwords. In case you do want to keep even the words that appear in every document, you can just add a 1. to the argument of the logarithm above:

\begin{equation}
\log\left[1+\frac{N_d}{N_d\left(w\right)}\right]
\end{equation}

When we multiply the term frequency of each word by it's inverse document frequency, we have a good way of quantifying how relevant a word is to understand the meaning of a specific document.

```python
def tf_idf(corpus_words):
    IDF = inv_doc_freq(corpus_words)

    TFIDF = []

    for document in corpus_words:
        TFIDF.append(Counter(document))

    for document in TFIDF:
        for word in document:
            document[word] = document[word]*IDF[word]

    return TFIDF
```

```python
tf_idf(corpus_words)
```

    [Counter({'a': 0.6931471805599453,
              'had': 0.6931471805599453,
              'lamb': 0.8630462173553426,
              'little': 2.0794415416798357,
              'mary': 0.0}),
     Counter({'a': 0.6931471805599453,
              'as': 1.3862943611198906,
              'fleece': 1.3862943611198906,
              'had': 0.6931471805599453,
              'lamb': 0.28768207245178085,
              'little': 0.6931471805599453,
              'mary': 0.0,
              'snow': 1.3862943611198906,
              'was': 0.6931471805599453,
              'white': 1.3862943611198906,
              'whose': 1.3862943611198906}),
     Counter({'and': 1.3862943611198906,
              'everywhere': 0.6931471805599453,
              'mary': 0.0,
              'that': 0.6931471805599453,
              'went': 2.0794415416798357}),
     Counter({'everywhere': 0.6931471805599453,
              'go': 1.3862943611198906,
              'lamb': 0.28768207245178085,
              'mary': 0.0,
              'sure': 1.3862943611198906,
              'that': 0.6931471805599453,
              'the': 1.3862943611198906,
              'to': 1.3862943611198906,
              'was': 0.6931471805599453,
              'went': 0.6931471805599453})]

Now we finally have a vector representation of each of our documents that takes the informational contributions of each word into account. Each of these vectors provides us with a unique representation of each document, in the context (corpus) in which it occurs, making it posssible to define the similarity of two documents, etc.

## Porter Stemmer

There is still, however, one issue with our approach to representing text. Since we treat each word as a unique token and completely independently from all others, for large documents we will end up with many variations of the same word such as verb conjugations, the corresponding adverbs and nouns, etc.

One way around this difficulty is to use stemming algorithm to reduce words to their root (or stem) version. The most famous Stemming algorithm is known as the **Porter Stemmer** and was introduced by Martin Porter in 1980 [Program 14, 130 (1980)](https://dl.acm.org/citation.cfm?id=275705)

The algorithm starts by defining consonants (C) and vowels (V):

```python
V = set('aeiouy')
C = set('bcdfghjklmnpqrstvwxz')
```

The stem of a word is what is left of that word after a speficic ending has been removed. A function to do this is easy to implement:

```python
def get_stem(suffix, word):
    """
        Extract the stem of a word
    """

    if word.lower().endswith(suffix.lower()): # Case insensitive comparison
        return word[:-len(suffix)]

    return None
```

It also defines words (or stems) to be sequences of vowels and consonants of the form:

\begin{equation}
[C](VC)^m[V]
\end{equation}

where $m$ is called the **measure** of the word and [] represent optional sections.

```python
def measure(orig_word):
    """
        Calculate the "measure" m of a word or stem, according to the Porter Stemmer algorthim
    """

    word = orig_word.lower()

    optV = False
    optC = False
    VC = False
    m = 0

    pos = 0

    # We can think of this implementation as a simple finite state machine
    # looks for sequences of vowels or consonants depending of the state
    # in which it's in, while keeping track of how many VC sequences it
    # has encountered.
    # The presence of the optional V and C portions is recorded in the
    # optV and optC booleans.

    # We're at the initial state.
    # gobble up all the optional consonants at the beginning of the word
    while pos < len(word) and word[pos] in C:
        pos += 1
        optC = True

    while pos < len(word):
        # Now we know that the next state must be a vowel
        while pos < len(word) and word[pos] in V:
            pos += 1
            optV = True

        # Followd by a consonant
        while pos < len(word) and word[pos] in C:
            pos += 1
            optV = False

        # If a consonant was found, then we matched VC
        # so we should increment m by one. Otherwise,
        # optV remained true and we simply had a dangling
        # V sequence.
        if not optV:
            m += 1

    return m
```

Let's consider a simple example. The word __crepusculars__ should have measure 4:

[cr] (ep) (usc) (ul) (ars)

and indeed it does.

```python
word = "crepusculars"
print(measure(word))
```

    4

The Porter algorithm sequentially applies a series of transformation rules over a series of 5 steps (step 1 is divided in 3 substeps and step 5 in 2). The rules are only applied if a certain condition is true.

In addition to possibily specifying a requirement on the measure of a word, conditions can make use of different boolean functions as well:

```python
def ends_with(char, stem):
    """
        Checks the ending of the word
    """
    return stem[-1] == char

def double_consonant(stem):
    """
        Checks the ending of a word for a double consonant
    """
    if len(stem) < 2:
        return False

    if stem[-1] in C and stem[-2] == stem[-1]:
        return True

    return False

def contains_vowel(stem):
    """
        Checks if a word contains a vowel or not
    """
    return len(set(stem) & V) > 0
```

Finally, we define a function to apply a specific rule to a word or stem:

```python
def apply_rule(condition, suffix, replacement, word):
    """
        Apply Porter Stemmer rule.
        if "condition" is True replace "suffix" by "replacement" in "word"
    """

    stem = get_stem(suffix, word)

    if stem is not None and condition is True:
        # Remove the suffix
        word = stem

        # Add the replacement suffix, if any
        if replacement is not None:
            word += replacement

    return word
```

Now we can see how rules can be applied. For example, this rule, from step 1b is successfully applied to __pastered__:

```python
word = "plastered"
suffix = "ed"
stem = get_stem(suffix, word)
apply_rule(contains_vowel(stem), suffix, None, word)
```

    'plaster'

While try applying the same rule to **bled** will fail to pass the condition resulting in no change.

```python
word = "bled"
suffix = "ed"
stem = get_stem(suffix, word)
apply_rule(contains_vowel(stem), suffix, None, word)
```

    'bled'

For a more complex example, we have, in Step 4:

```python
word = "adoption"
suffix = "ion"
stem = get_stem(suffix, word)
apply_rule(measure(stem) > 1 and (ends_with("s", stem) or ends_with("t", stem)), suffix, None, word)
```

    'adopt'

In total, the Porter Stemmer algorithm (for the English language) applies several dozen rules (see https://tartarus.org/martin/PorterStemmer/def.txt for a complete list). Implementing all of them is both tedious and error prone, so we abstain from providing a full implementation of the algorithm here. High quality implementations can be found in all major NLP libraries such as [NLTK](http://www.nltk.org/howto/stem.html).

The dificulties of defining matching rules to arbitrary text cannot be fully resolved without the use of Regular Expressions (typically implemented as Finite State Machines like our __measure__ implementation above), a more advanced topic that is beyond the scope of this course.