# Multiscale Migration Model

### Scroll down to "Editing Variable Values" after running all cells

This notebook implements our model using `numpy`, `haversine`, and `pandas` (with `xlrd`). It has been tested to run on Python 3.6. To start, import the required libraries.

```python
import os
import sys
import subprocess
working_directory = os.path.abspath('')
sys.path.append(os.path.normpath(os.path.join(working_directory, "..", "..")))
```

```python
# These libraries are used later to supply mathematical calculations.
import numpy as np
import pandas as pd
from math import e
from haversine import haversine
import ipywidgets as widgets
from ipywidgets import *
# Visualizaton
import matplotlib
import matplotlib.pyplot as plt
import gos
from gos.visualization import map_plot
%matplotlib inline
# Reduce some visual clutter by only printing ten rows at a time.
# This can be adjusted to match personal preferences.
pd.set_option("display.max_rows", 500)
```

```python
plt.style.use('ggplot')
```

The datasets used in the model are found in the `/data` subdirectory. Many are formatted as either CSV files or XLSX files.

```python
print(subprocess.run(['ls', 'data'], capture_output=True, check=True).stdout.decode('utf-8'))
```

## Use this tool to change the skill level at which migration is studied

```python
skill = IntSlider(min=0, max=100, value=90)
display(skill)
```

The following shortcut functions helps locate these data files easily.

```python
def file_path(name):
    """
    Shortcut function to get the relative path to the directory
    which contains the data.
    """
    return "./data/%s" % name
```

## Cleaning the data.

In this step, we define some helper functions that will help all of our datasets talk to each other.

```python
column_names = ["Name", "Code"]

def country_codes():
    """
    Build country rows from their names, ISO codes, and Numeric
    Country Codes.
    """
    cc = (
        pd.read_csv(
            file_path(
                "Country_List_ISO_3166_Codes_Latitude_Longitude.csv"),
            usecols=[0, 2, 3],
            index_col=1,
            keep_default_na=False))
    #other_codes = pd.read_csv(file_path("other.csv"), index_col=1)

    #
    cc.columns = column_names
    #cc.index.rename("Name")
    return cc #pd.concat([cc, other_codes])

def other_codes():
    other_codes = pd.read_csv(file_path("other.csv"), index_col=1)
    other_codes.columns = column_names[0:1]
    return other_codes
```

```python
world = gos.World(index=set(country_codes().index) | set(other_codes().index))
```

```python
gos.Neighborhood.update(country_codes().groupby("Alpha-3 code")["Name"].apply(list).to_dict())
gos.Neighborhood.update(other_codes().groupby('ISO')["Name"].apply(list).to_dict())
gos.Neighborhood.update(country_codes().groupby("Alpha-3 code")["Code"].apply(list).to_dict())
```

# Importing Data

### Freedom Index

The [Freedom Index](https://freedomhouse.org/report/freedom-world/freedom-world-2017) comes from Freedom House.

A high freedom index score indicates a politically free country

```python
def freedom_index():
    """
    Read data from the Freedom Index.
    """
    # TODO: Add xlrd to requirements.
    xl = pd.ExcelFile(file_path("Freedom_index.xlsx"))
    xl = xl.parse(1)
    xl.set_index("Country")
    return xl

fi = freedom_index().set_index("Country")
fi.columns = ["Freedom Index"]
fi.plot.hist(bins=10)
#print(len(fi))
```

### A/B Values

These values are used in the return to skill function. These values are based on each country's income distribution.

```python
def ab_values():
    """
    Read generated A/B values for each country.
    """
    return pd.read_excel(file_path("A&B values for RTS.xlsx"))

ab = ab_values()
ab = ab.set_index("Country")
#ab.info()
#print(ab)
ab.plot.hist(subplots=True, sharex=False)
```

### Passport Index

The [Passport Index](https://www.cato.org/human-freedom-index) comes from Arton Capital.

A low passport index indicates a welcoming country.

```python
def passport_index():
    """
    Read data from the Passport Index.
    """
    pi = pd.read_excel(file_path("PassportIndex.xlsx"))
    pi = pi.set_index("Country")
    pi.columns = ["Passport Index"]
    return pi

pi = passport_index()
pi.plot.hist()
```

## Unemployment

```python
unemployment_data = pd.read_csv(file_path("CIA_Unemployment.csv"), index_col=0, usecols=[1, 2])
unemployment_data["Unemployment"] /= 100
unemployment_data.plot.hist()
```

## Population

```python
# Population
population = pd.read_csv(file_path("newPOP.csv"))
population = population.set_index("Country")
population
```

```python
world.update_neighborhoods(ab)
world.update_neighborhoods(pi)
world.update_neighborhoods(unemployment_data)
world.update_neighborhoods(population)
world.update_neighborhoods(fi)
```

## Shared Language

Agents are assigned proficiency in languages spoken in their origin country. Moving to a country with entirely new languages presents a higher migration cost.

\begin{equation}
    L_{O\leftrightarrow D}\ =
    \begin{cases}
        0 & \text{if origin and destination share a spoken language} \\
        1 & \text{otherwise} \\
    \end{cases}
\end{equation}

```python
lang_csv = pd.read_csv(file_path("languages.csv"), index_col=0)
lang_sets = [set([str(y).strip() for y in x[1] if y != ' ']) for x in lang_csv.iterrows()]
overlap = []
for s in lang_sets:
    o = []
    for i in range(len(lang_sets)):
        o.append(len(lang_sets[i].intersection(s)) >= 1)
    overlap.append(o)
lang_data = pd.DataFrame(overlap, index=lang_csv.index, columns=lang_csv.index)
print(len(lang_data))
world.add_matrix("language", 1-lang_data)
```

## UN Migration History

```python
un_pd = pd.read_excel(
        file_path(
            "UN_MigrantStockByOriginAndDestination_2015.xlsx"
        ),
        skiprows=15
    )
un_pd = un_pd.set_index('Unnamed: 1')
un_pd = un_pd.iloc[0:275,7:250]

# TODO: Should we be using the UN numbers for this?
un_pd = un_pd.sort_index().fillna(1)
world.add_matrix("un", un_pd)
```

# Calculations

## Distance

The great circle distance between the average latitude and longitude of each country is used to determine distance between each pair of countries. A greater distance between countries corresponds to a greater cost of migration.

```python
distance_frame = pd.read_csv(
    file_path("Country_List_ISO_3166_Codes_Latitude_Longitude.csv"),
    usecols=[2,4,5],
    index_col=0,
    keep_default_na=False)
locations = [(x[1][0], x[1][1]) for x in distance_frame.iterrows()]
rows = []
for i in range(len(locations)):
    row = []
    for loc in locations:
        row.append(haversine(loc, locations[i]))
    rows.append(row)
distance = pd.DataFrame(rows, distance_frame.index, distance_frame.index)
world.add_matrix("distance", distance / distance.max().max())
```

# Editing Variable Values
Remove the "#" and change the values to best simulate your scenario

To reset values, run all cells above

```python
pd.options.mode.chained_assignment = None  # default='warn'

world.data["Freedom Index"]["AUS"]=50
#world.data["Passport Index"]["AUS"]=80
#world.data["Unemployment"]["AUS"]=.058
#world.data["Population"]["AUS"]=8000000
#world.matrices["un"]["AUS"]["USA"]=77845
```

## Calculate "RTS".
This function measures income potential in a given country based on an agent's skill level (x).

It takes the form:
\begin{equation}
Ae^{Bx}
\end{equation}
Where A and B are constants unique to each country

```python
world.update_neighborhoods(pd.Series(world.data["A"] * e ** (world.data["B"] * skill.value)), "rts")
```

## Country Beta Values

```python
world.update_neighborhoods(pd.Series(world.data["A"] * e ** (world.data["B"] * 30)), "beta")
```

## Freedom Index difference
\begin{equation}
\frac{FI_D-FI_O}{100}
\end{equation}

This value weighs the origin and destination FI values against eachother.

```python
rows=[]
for i in range(len(world.data["Freedom Index"])):
    row = []
    for freedom_index in world.data["Freedom Index"]:
        diff=(freedom_index-(world.data["Freedom Index"][i]))/100.0
        row.append(diff)
    rows.append(row)
fi_diff = (pd.DataFrame(rows, world.data["Freedom Index"].index, world.data["Freedom Index"].index))
print(fi_diff)
```

## Political Barriers
$$  PB = \delta_1 \frac{PIR_{D}}{100} +  \delta_2 (1-\frac{(FI_D-FI_O)}{100}) $$

Political barriers combines passport index rank (PIR) and freedom index scores (FI) for each country. The greater the value, the higher the cost of migration.

```python
delta1=.5
delta2=.25
political_barriers=delta1*world.data["Passport Index"]/100.0 + delta2*(1 - fi_diff)
print(political_barriers)

```

```python
world.add_matrix("Political Barriers", political_barriers)
```

## Migration History
Out Migration:
\begin{equation} OM_{O\rightarrow D}\ = \frac {\text {migrants from origin in destination}} {\text {population of origin}} \end{equation}

Ethnic Enclave:
\begin{equation} EE_{O\leftrightarrow D}\ = \frac {\text{migrants from origin in destination}} {\text {population of destination.}} \end{equation}

Migration History:
\begin{equation}
MH_{O\rightarrow D}\ = \gamma_1 (1-OM_{O\rightarrow D})\ + \gamma_2 (1-EE_{O\rightarrow D})
\end{equation}

```python
gamma1=.5
gamma2=.5
OM=world.matrices['un'].sort_index(axis=1).sort_index(axis=0)/world.data['Population']
#transpose UN matrix for this calculation so that we are dividing by population of destination
EE=world.matrices['un'].T.sort_index(axis=1).sort_index(axis=0)/world.data['Population']
EE=EE.T
MH=gamma1*(OM)+gamma2*(EE)
#print(MH)
max_MH = MH.max().nlargest(10).mean()
#print(max_MH)
MH = 1 - (MH/max_MH)
MH[MH<0] = 0
print(MH)
```

```python
world.add_matrix("Migration History", MH)
```

# Cost

The cost of migration between an origin and destination is the weighted average of the distance, migration history, shared language and political barriers between the two countries.
\begin{equation}
C = { \alpha_1 \frac{D_{O\leftrightarrow D}} {D_{Max}} + \alpha_2 MH_{O\rightarrow D}  + \alpha_3 L_{O\leftrightarrow D} + \alpha_4 PB}
\end{equation}

```python
# Cost
alpha1=.35
alpha2=.35
alpha3=.15
alpha4=.15
c = (alpha1*world.matrices["distance"] +
     alpha2*world.matrices["Migration History"] +
     alpha3*world.matrices["language"] +
     alpha4*world.matrices["Political Barriers"])
world.add_matrix("cost", c * world.data['beta'])
```

```python
map_plot(world.matrices["cost"]["SYR"], title="Costs (SYR)")
```

```python
map_plot(world.matrices["cost"]["FRA"], title="Costs (FRA)")
```

```python
world.update_neighborhoods((1 - world.data["Unemployment"]) * world.data["rts"], "wages")
```

```python
map_plot(world.data["wages"], title="Wages")
```

```python
#beta = world.data.rts.mean()
#beta = RTS_list.median()

world.add_matrix("migration", (pd.DataFrame(
    np.array(
        [[x] * len(world.data) for x in world.data["wages"].values]
    ) - np.array(
        [list(world.data["wages"].values)] * len(world.data)
    ),
    world.data.index,
    world.data.index
    )  - world.matrices["cost"]).clip(lower=0))
world.matrices["migration"] = world.matrices["migration"] / (world.matrices["migration"].sum() + 1)
world.matrices["migration"] = world.matrices["migration"] / world.matrices["migration"].sum(axis=1).max()

# TODO: Why does this require being transposed?
world.matrices["migration"] = (0.15 * world.matrices["migration"].transpose() * world.data["Population"]).transpose()

world.matrices["migration"]
```

```python
map_plot(
   world.matrices["migration"].sum(axis=1)+1,
   title="Immigration Estimations (x={})".format(skill.value),
   normc=matplotlib.colors.LogNorm
)
```

```python
map_plot(
   world.matrices["migration"].sum(),
   title="Estimated Number of Emigrants (x={})".format(skill.value),
   normc=matplotlib.colors.Normalize
)
```

```python
map_plot(
   world.matrices["migration"].sum(axis=1) - world.matrices["migration"].sum(),
   title="Net Migration (x={})".format(skill.value),
   normc=gos.visualization.MidPointNorm
)
```

# Exporting Data

Defining Vectors

```python
immigration = world.matrices["migration"].sum(axis=1)+1
emigration = world.matrices["migration"].sum()
net_migration = world.matrices["migration"].sum(axis=1) - world.matrices["migration"].sum()
```

- To export the following data, remove the "#" and run the cell.
- The excel files will appear in the data folder under multiscale-migration.
- Change the file names (in red) everytime you want to create a spreadsheet with new results. Otherwise new changes will override the original spreadsheets.

```python
immigration.to_excel("data/immigration.xlsx")
emigration.to_excel("data/emigration.xlsx")
net_migration.to_excel("data/net_migration.xlsx")
```

```python

```


# Optimization of a State-to-State Transfer in a Two-Level-System

```
# NBVAL_IGNORE_OUTPUT
%load_ext watermark
import qutip
import numpy as np
import scipy
import matplotlib
import matplotlib.pylab as plt
import krotov
%watermark -v --iversions
```

    qutip       4.3.1
    numpy       1.15.4
    scipy       1.1.0
    matplotlib  3.0.2
    matplotlib.pylab  1.15.4
    krotov      0.0.1
    CPython 3.6.7
    IPython 7.2.0

$\newcommand{tr}[0]{\operatorname{tr}}
\newcommand{diag}[0]{\operatorname{diag}}
\newcommand{abs}[0]{\operatorname{abs}}
\newcommand{pop}[0]{\operatorname{pop}}
\newcommand{aux}[0]{\text{aux}}
\newcommand{opt}[0]{\text{opt}}
\newcommand{tgt}[0]{\text{tgt}}
\newcommand{init}[0]{\text{init}}
\newcommand{lab}[0]{\text{lab}}
\newcommand{rwa}[0]{\text{rwa}}
\newcommand{bra}[1]{\langle#1\vert}
\newcommand{ket}[1]{\vert#1\rangle}
\newcommand{Bra}[1]{\left\langle#1\right\vert}
\newcommand{Ket}[1]{\left\vert#1\right\rangle}
\newcommand{Braket}[2]{\left\langle #1\vphantom{#2} \mid
#2\vphantom{#1}\right\rangle}
\newcommand{op}[1]{\hat{#1}}
\newcommand{Op}[1]{\hat{#1}}
\newcommand{dd}[0]{\,\text{d}}
\newcommand{Liouville}[0]{\mathcal{L}}
\newcommand{DynMap}[0]{\mathcal{E}}
\newcommand{identity}[0]{\mathbf{1}}
\newcommand{Norm}[1]{\lVert#1\rVert}
\newcommand{Abs}[1]{\left\vert#1\right\vert}
\newcommand{avg}[1]{\langle#1\rangle}
\newcommand{Avg}[1]{\left\langle#1\right\rangle}
\newcommand{AbsSq}[1]{\left\vert#1\right\vert^2}
\newcommand{Re}[0]{\operatorname{Re}}
\newcommand{Im}[0]{\operatorname{Im}}$
The purpose of this example is not to solve an especially interesting physical
problem but to give a rather simple example of how the package can be used in
order to solve an optimization problem.

## Define the Hamiltonian

In the
following the Hamiltonian, guess field and
states are defined.

The Hamiltonian
$\op{H}_{0} = - \omega \op{\sigma}_{z}$
represents a
simple qubit with energy
level splitting $\omega$ in the basis
$\{\ket{0},\ket{1}\}$. The control
field
$\epsilon(t)$ is assumed to couple via
the
Hamiltonian $\op{H}_{1}(t) =
\epsilon(t) \op{\sigma}_{x}$ to the qubit,
i.e., the control
field effectively
drives
transitions between both qubit
states. For now, we initialize the control
field as constant.

```
def ham_and_states(omega=1.0, ampl0=0.2):
    """Two-level-system Hamiltonian

    Args:
        omega (float): energy separation of the qubit levels
        ampl0 (float): constant amplitude of the driving field
    """
    H0 = - 0.5 * omega * qutip.operators.sigmaz()
    H1 = qutip.operators.sigmax()

    psi0 = qutip.Qobj(np.array([1,0]))
    psi1 = qutip.Qobj(np.array([0,1]))

    eps0 = lambda t, args: ampl0
    return ([H0, [H1, eps0]], psi0, psi1)

H, psi0, psi1 = ham_and_states()
```

The projectors $\op{P}_0 = \ket{0}\bra{0}$ and $\op{P}_1 = \ket{1}\bra{1}$ are
introduced since they allow for calculating the
population in the respective
states later on.

```
proj0 = psi0 * psi0.dag()
proj1 = psi1 * psi1.dag()
```

## Define the optimization target

First we define the time grid of the
dynamics, i.e., by taking the following
values as an example, we define the
initial state to be at time $t=0$ and
consider a total propagation time of
$T=5$. The entire time grid is divided into
$n_{t}=500$ equidistant time steps.

```
tlist = np.linspace(0, 5, 500)
```

Next, we define the optimization targets, which is technically a list of
objectives, but here it has just one entry defining a simple state-to-state
transfer
from initial state $\ket{\Psi_{\init}} = \ket{0}$ to the target state
$\ket{\Psi_{\tgt}} = \ket{1}$, which we want to reach at final time $T$. Note
that we also have to pass the Hamiltonian $\op{H}(t)$ that determines the
dynamics of
the system to the optimization objective.

```
objectives = [
    krotov.Objective(initial_state=psi0, target=psi1, H=H)
]
```

In addition, we have to define and assign a shape function $S(t)$ for the update
in each control iteration to each
control field that will be updated. This shape
usually takes care of
experimental limits such as the necessity of finite ramps
at the beginning and
end of the control field or other conceivable limitations
for field shapes: wherever $S(t)$ is zero, the optimization will not change the
value of the control from the original guess.

```
def S(t):
    """Shape function for the field update"""
    return krotov.shapes.flattop(t, t_start=0, t_stop=5, t_rise=0.3, t_fall=0.3, func='sinsq')
```

At this point, we also change the initial control field $\epsilon_{0}(t)$ from a
constant to a shaped pulse that switches on smoothly from zero and again
switches off at the final time $T$. We re-use the shape function $S(t)$ that we
defined for the updates for this purpose (although generally, $S(t)$ for the
updates has nothing to with the shape of the control field).

```
def shape_field(eps0):
    """Applies the shape function S(t) to the guess field"""
    eps0_shaped = lambda t, args: eps0(t, args)*S(t)
    return eps0_shaped

H[1][1] = shape_field(H[1][1])
```

Having defined the shape function $S(t)$ and having shaped the guess field, we
now tell the optimization to also use $S(t)$ as the update-shape for
$\epsilon_0(t)$. In addition, we have to choose `lambda_a` for each control
field. It controls the update magnitude of the respective field in each
iteration.

```
pulse_options = {
    H[1][1]: krotov.PulseOptions(lambda_a=5, shape=S)
}
```

It is convenient to introduce the function `print_fidelity`, which can be passed
to the optimization procedure and will be called after each iteration and thus
provides additional feedback about the optimization progress.

```
def print_fidelity(**args):
    F_re = np.average(np.array(args['tau_vals']).real)
    print("    F = %f" % F_re)
    return F_re
```

## Simulate dynamics of the guess field

Before heading towards the optimization
procedure, we first simulate the
dynamics under the guess field
$\epsilon_{0}(t)$.

```
def plot_pulse(pulse, tlist):
    fig, ax = plt.subplots()
    if callable(pulse):
        pulse = np.array([pulse(t, args=None) for t in tlist])
    ax.plot(tlist, pulse)
    ax.set_xlabel('time')
    ax.set_ylabel('pulse amplitude')
    plt.show(fig)
```

The following plot shows the guess field $\epsilon_{0}(t)$, which is, as chosen
above, just a constant field (with a smooth switch-on and switch-off)

```
plot_pulse(H[1][1], tlist)
```

The next line solves the equation of motion for the defined objective, which
contains the initial state $\ket{\Psi_{\init}}$ and the Hamiltonian $\op{H}(t)$
defining its evolution.

```
guess_dynamics = objectives[0].mesolve(tlist, e_ops=[proj0, proj1])
```

The plot of the population dynamics shows that the guess field does not transfer
the initial state $\ket{\Psi_{\init}} = \ket{0}$ to the desired target state
$\ket{\Psi_{\tgt}} = \ket{1}$.

```
def plot_population(result):
    fig, ax = plt.subplots()
    ax.plot(result.times, result.expect[0], label='0')
    ax.plot(result.times, result.expect[1], label='1')
    ax.legend()
    ax.set_xlabel('time')
    ax.set_ylabel('population')
    plt.show(fig)
```

```
plot_population(guess_dynamics)
```

## Optimize

In the following we optimize the guess field $\epsilon_{0}(t)$ such
that the intended state-to-state transfer $\ket{\Psi_{\init}} \rightarrow
\ket{\Psi_{\tgt}}$ is solved.

The cell below carries out the optimization. It
requires, besides the
previously
defined optimization `objectives`, information
about the
optimization functional
$F$ (via `chi_constructor`) and the
propagation method that should be used. In
addition, the number of total
iterations is required and, as an option, we pass
an info-hook that after each
iteration combines a complete printout of the state
of the optimization with the
`print_fidelity` function defined above.

Here, we
choose $F = F_{re}$ with
\begin{equation}
F_{re}
=
\Re\Braket{\Psi(T)}{\Psi_{\tgt}}
\end{equation}

with
$\ket{\Psi(T)}$ the
forward propagated state of $\ket{\Psi_{\init}}$.

```
oct_result = krotov.optimize_pulses(
    objectives,
    pulse_options=pulse_options,
    tlist=tlist,
    propagator=krotov.propagators.expm,
    chi_constructor=krotov.functionals.chis_re,
    info_hook=krotov.info_hooks.chain(
        krotov.info_hooks.print_debug_information, print_fidelity
    ),
    check_convergence=krotov.convergence.check_monotonic_fidelity,
    iter_stop=10,
)
```

    Iteration 0
        objectives:
            1:|(2)⟩ - {[Herm[2,2], [Herm[2,2], u1(t)]]} - |(2)⟩
        adjoint objectives:
            1:⟨(2)| - {[Herm[2,2], [Herm[2,2], u1(t)]]} - ⟨(2)|
        λₐ: 5.00e+00
        S(t) (ranges): [0.000000, 1.000155]
        duration: 0.6 secs (started at 2018-12-20 08:17:54)
        optimized pulses (ranges): [0.00, 0.20]
        backward states: None
        forward states: [1 * ndarray(500)]
        fw_states_T norm: 1.000000
        τ: (2.12e-01:0.50π)
        F = 0.000000
    Iteration 1
        duration: 1.3 secs (started at 2018-12-20 08:17:55)
        optimized pulses (ranges): [0.00, 0.29]
        backward states: [1 * ndarray(500)]
        forward states: [1 * ndarray(500)]
        fw_states_T norm: 1.000000
        τ: (3.16e-01:0.23π)
        F = 0.235157
    Iteration 2
        duration: 1.3 secs (started at 2018-12-20 08:17:56)
        optimized pulses (ranges): [0.00, 0.37]
        backward states: [1 * ndarray(500)]
        forward states: [1 * ndarray(500)]
        fw_states_T norm: 1.000000
        τ: (4.91e-01:0.14π)
        F = 0.444066
    Iteration 3
        duration: 1.3 secs (started at 2018-12-20 08:17:57)
        optimized pulses (ranges): [-0.07, 0.44]
        backward states: [1 * ndarray(500)]
        forward states: [1 * ndarray(500)]
        fw_states_T norm: 1.000000
        τ: (6.44e-01:0.10π)
        F = 0.611354
    Iteration 4
        duration: 1.3 secs (started at 2018-12-20 08:17:59)
        optimized pulses (ranges): [-0.14, 0.50]
        backward states: [1 * ndarray(500)]
        forward states: [1 * ndarray(500)]
        fw_states_T norm: 1.000000
        τ: (7.60e-01:0.08π)
        F = 0.735071
    Iteration 5
        duration: 1.3 secs (started at 2018-12-20 08:18:00)
        optimized pulses (ranges): [-0.20, 0.55]
        backward states: [1 * ndarray(500)]
        forward states: [1 * ndarray(500)]
        fw_states_T norm: 1.000000
        τ: (8.41e-01:0.07π)
        F = 0.821734
    Iteration 6
        duration: 1.4 secs (started at 2018-12-20 08:18:01)
        optimized pulses (ranges): [-0.26, 0.58]
        backward states: [1 * ndarray(500)]
        forward states: [1 * ndarray(500)]
        fw_states_T norm: 1.000000
        τ: (8.96e-01:0.06π)
        F = 0.880461
    Iteration 7
        duration: 1.3 secs (started at 2018-12-20 08:18:03)
        optimized pulses (ranges): [-0.30, 0.61]
        backward states: [1 * ndarray(500)]
        forward states: [1 * ndarray(500)]
        fw_states_T norm: 1.000000
        τ: (9.31e-01:0.05π)
        F = 0.919555
    Iteration 8
        duration: 1.3 secs (started at 2018-12-20 08:18:04)
        optimized pulses (ranges): [-0.34, 0.63]
        backward states: [1 * ndarray(500)]
        forward states: [1 * ndarray(500)]
        fw_states_T norm: 1.000000
        τ: (9.54e-01:0.04π)
        F = 0.945388
    Iteration 9
        duration: 1.6 secs (started at 2018-12-20 08:18:05)
        optimized pulses (ranges): [-0.36, 0.65]
        backward states: [1 * ndarray(500)]
        forward states: [1 * ndarray(500)]
        fw_states_T norm: 1.000000
        τ: (9.69e-01:0.04π)
        F = 0.962447
    Iteration 10
        duration: 1.3 secs (started at 2018-12-20 08:18:07)
        optimized pulses (ranges): [-0.39, 0.66]
        backward states: [1 * ndarray(500)]
        forward states: [1 * ndarray(500)]
        fw_states_T norm: 1.000000
        τ: (9.78e-01:0.03π)
        F = 0.973756

```
oct_result
```

    Krotov Optimization Result
    --------------------------
    - Started at 2018-12-20 08:17:54
    - Number of objectives: 1
    - Number of iterations: 10
    - Reason for termination: Reached 10 iterations
    - Ended at 2018-12-20 08:18:08

## Simulate dynamics of the optimized field

Having obtained the optimized
control field, we can now
plot it and calculate the
population dynamics under
this field.

```
plot_pulse(oct_result.optimized_controls[0], tlist)
```

In contrast to the dynamics under the guess field, the optimized field indeed
drives the initial state $\ket{\Psi_{\init}} = \ket{0}$ to the desired target
state $\ket{\Psi_{\tgt}} = \ket{1}$.

```
opt_dynamics = oct_result.optimized_objectives[0].mesolve(
    tlist, e_ops=[proj0, proj1])
```

```
plot_population(opt_dynamics)
```


```python
from google.colab import drive
drive.mount('/content/drive')
```

    Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3aietf%3awg%3aoauth%3a2.0%3aoob&response_type=code&scope=email%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdocs.test%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive.photos.readonly%20https%3a%2f%2fwww.googleapis.com%2fauth%2fpeopleapi.readonly

    Enter your authorization code:
    ··········
    Mounted at /content/drive

# Neuromatch Academy: Week 2, Day 5, Tutorial 3
# Learning to Act: Q-Learning

__Content creators:__ Marcelo Mattar and Eric DeWitt with help from Byron Galbraith

__Content reviewers:__ Matt Krause and Michael Waskom

---

# Tutorial Objectives

In this tutorial you will learn how to act in the more realistic setting of sequential decisions, formalized by Markov Decision Processes (MDPs). In a sequential decision problem, the actions executed in one state not only may lead to immediate rewards (as in a bandit problem), but may also affect the states experienced next (unlike a bandit problem). Each individual action may therefore affect affect all future rewards. Thus, making decisions in this setting requires considering each action in terms of their expected **cumulative** future reward.

We will consider here the example of spatial navigation, where actions (movements) in one state (location) affect the states experienced next, and an agent might need to execute a whole sequence of actions before a reward is obtained.

By the end of this tutorial, you will learn
* what grid worlds are and how they help in evaluating simple reinforcement learning agents
* the basics of the Q-learning algorithm for estimating action values
* how the concept of exploration and exploitation, reviewed in the bandit case, also applies to the sequential decision setting

---
# Setup

```python
# Imports
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import convolve as conv
```

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

```python
#@title Helper functions
def epsilon_greedy(q, epsilon):
  """Epsilon-greedy policy: selects the maximum value action with probabilty
  (1-epsilon) and selects randomly with epsilon probability.

  Args:
    q (ndarray): an array of action values
    epsilon (float): probability of selecting an action randomly

  Returns:
    int: the chosen action
  """
  if np.random.random() > epsilon:
    action = np.argmax(q)
  else:
    action = np.random.choice(len(q))

  return action

class CliffWorld:
  """
  World: Cliff world.
  40 states (4-by-10 grid world).
  The mapping from state to the grids are as follows:
  30 31 32 ... 39
  20 21 22 ... 29
  10 11 12 ... 19
  0  1  2  ...  9
  0 is the starting state (S) and 9 is the goal state (G).
  Actions 0, 1, 2, 3 correspond to right, up, left, down.
  Moving anywhere from state 9 (goal state) will end the session.
  Taking action down at state 11-18 will go back to state 0 and incur a
      reward of -100.
  Landing in any states other than the goal state will incur a reward of -1.
  Going towards the border when already at the border will stay in the same
      place.
  """
  def __init__(self):
    self.name = "cliff_world"
    self.n_states = 40
    self.n_actions = 4
    self.dim_x = 10
    self.dim_y = 4
    self.init_state = 0

  def get_outcome(self, state, action):
    if state == 9:  # goal state
      reward = 0
      next_state = None
      return next_state, reward
    reward = -1  # default reward value
    if action == 0:  # move right
      next_state = state + 1
      if state % 10 == 9:  # right border
        next_state = state
      elif state == 0:  # start state (next state is cliff)
        next_state = None
        reward = -100
    elif action == 1:  # move up
      next_state = state + 10
      if state >= 30:  # top border
        next_state = state
    elif action == 2:  # move left
      next_state = state - 1
      if state % 10 == 0:  # left border
        next_state = state
    elif action == 3:  # move down
      next_state = state - 10
      if state >= 11 and state <= 18:  # next is cliff
        next_state = None
        reward = -100
      elif state <= 9:  # bottom border
        next_state = state
    else:
      print("Action must be between 0 and 3.")
      next_state = None
      reward = None
    return int(next_state) if next_state is not None else None, reward

  def get_all_outcomes(self):
    outcomes = {}
    for state in range(self.n_states):
      for action in range(self.n_actions):
        next_state, reward = self.get_outcome(state, action)
        outcomes[state, action] = [(1, next_state, reward)]
    return outcomes

def learn_environment(env, learning_rule, params, max_steps, n_episodes):
  # Start with a uniform value function
  value = np.ones((env.n_states, env.n_actions))

  # Run learning
  reward_sums = np.zeros(n_episodes)

  # Loop over episodes
  for episode in range(n_episodes):
    state = env.init_state  # initialize state
    reward_sum = 0

    for t in range(max_steps):
      # choose next action
      action = epsilon_greedy(value[state], params['epsilon'])

      # observe outcome of action on environment
      next_state, reward = env.get_outcome(state, action)

      # update value function
      value = learning_rule(state, action, reward, next_state, value, params)

      # sum rewards obtained
      reward_sum += reward

      if next_state is None:
          break  # episode ends
      state = next_state

      reward_sums[episode] = reward_sum

  return value, reward_sums

def plot_state_action_values(env, value, ax=None):
  """
  Generate plot showing value of each action at each state.
  """
  if ax is None:
    fig, ax = plt.subplots()

  for a in range(env.n_actions):
    ax.plot(range(env.n_states), value[:, a], marker='o', linestyle='--')
  ax.set(xlabel='States', ylabel='Values')
  ax.legend(['R','U','L','D'], loc='lower right')

def plot_quiver_max_action(env, value, ax=None):
  """
  Generate plot showing action of maximum value or maximum probability at
    each state (not for n-armed bandit or cheese_world).
  """
  if ax is None:
    fig, ax = plt.subplots()

  X = np.tile(np.arange(env.dim_x), [env.dim_y,1]) + 0.5
  Y = np.tile(np.arange(env.dim_y)[::-1][:,np.newaxis], [1,env.dim_x]) + 0.5
  which_max = np.reshape(value.argmax(axis=1), (env.dim_y,env.dim_x))
  which_max = which_max[::-1,:]
  U = np.zeros(X.shape)
  V = np.zeros(X.shape)
  U[which_max == 0] = 1
  V[which_max == 1] = 1
  U[which_max == 2] = -1
  V[which_max == 3] = -1

  ax.quiver(X, Y, U, V)
  ax.set(
      title='Maximum value/probability actions',
      xlim=[-0.5, env.dim_x+0.5],
      ylim=[-0.5, env.dim_y+0.5],
  )
  ax.set_xticks(np.linspace(0.5, env.dim_x-0.5, num=env.dim_x))
  ax.set_xticklabels(["%d" % x for x in np.arange(env.dim_x)])
  ax.set_xticks(np.arange(env.dim_x+1), minor=True)
  ax.set_yticks(np.linspace(0.5, env.dim_y-0.5, num=env.dim_y))
  ax.set_yticklabels(["%d" % y for y in np.arange(0, env.dim_y*env.dim_x,
                                                  env.dim_x)])
  ax.set_yticks(np.arange(env.dim_y+1), minor=True)
  ax.grid(which='minor',linestyle='-')

def plot_heatmap_max_val(env, value, ax=None):
  """
  Generate heatmap showing maximum value at each state
  """
  if ax is None:
    fig, ax = plt.subplots()

  if value.ndim == 1:
      value_max = np.reshape(value, (env.dim_y,env.dim_x))
  else:
      value_max = np.reshape(value.max(axis=1), (env.dim_y,env.dim_x))
  value_max = value_max[::-1,:]

  im = ax.imshow(value_max, aspect='auto', interpolation='none', cmap='afmhot')
  ax.set(title='Maximum value per state')
  ax.set_xticks(np.linspace(0, env.dim_x-1, num=env.dim_x))
  ax.set_xticklabels(["%d" % x for x in np.arange(env.dim_x)])
  ax.set_yticks(np.linspace(0, env.dim_y-1, num=env.dim_y))
  if env.name != 'windy_cliff_grid':
      ax.set_yticklabels(
          ["%d" % y for y in np.arange(
              0, env.dim_y*env.dim_x, env.dim_x)][::-1])
  return im

def plot_rewards(n_episodes, rewards, average_range=10, ax=None):
  """
  Generate plot showing total reward accumulated in each episode.
  """
  if ax is None:
    fig, ax = plt.subplots()

  smoothed_rewards = (conv(rewards, np.ones(average_range), mode='same')
                      / average_range)

  ax.plot(range(0, n_episodes, average_range),
          smoothed_rewards[0:n_episodes:average_range],
          marker='o', linestyle='--')
  ax.set(xlabel='Episodes', ylabel='Total reward')

def plot_performance(env, value, reward_sums):
  fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(16, 12))
  plot_state_action_values(env, value, ax=axes[0,0])
  plot_quiver_max_action(env, value, ax=axes[0,1])
  plot_rewards(n_episodes, reward_sums, ax=axes[1,0])
  im = plot_heatmap_max_val(env, value, ax=axes[1,1])
  fig.colorbar(im)
```

---
# Section 1: Markov Decision Processes

```python
#@title Video 1: MDPs and Q-learning
# Insert the ID of the corresponding youtube video
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="8yvwMrUQJOU", width=854, height=480, fs=1)
print("Video available at https://youtu.be/" + video.id)
video
```

    Video available at https://youtu.be/8yvwMrUQJOU

## Section 1.1: Grid Worlds

As pointed out, bandits only have a single state and immediate rewards for our actions. Many problems we are interested in have multiple states and delayed rewards, i.e. we won't know if the choices we made will pay off over time, or which actions we took contributed to the outcomes we observed.

In order to explore these ideas, we turn the a common problem setting: the grid world. Grid worlds are simple environments where each state corresponds to a tile on a 2D grid, and the only actions the agent can take are to move up, down, left, or right across the grid tiles. The agent's job is almost always to find a way to a goal tile in the most direct way possible while overcoming some maze or other obstacles, either static or dynamic.

For our discussion we will be looking at the classic Cliff World, or Cliff Walker, environment. This is a 4x10 grid with a starting position in the lower-left and the goal position in the lower-right. Every tile between these two is the "cliff", and should the agent enter the cliff, they will receive a -100 reward and be sent back to the starting position. Every tile other than the cliff produces a -1 reward when entered. The goal tile ends the episode after taking any action from it.

Given these conditions, the maximum achievable reward is -11 (1 up, 9 right, 1 down). Using negative rewards is a common technique to encourage the agent to move and seek out the goal state as fast as possible.

---
# Section 2: Q-Learning

Now that we have our environment, how can we solve it?

One of the most famous algorithms for estimating action values (aka Q-values) is the Temporal Differences (TD) **control** algorithm known as *Q-learning* (Watkins, 1989).

\begin{align}
Q(s_t,a_t) \leftarrow Q(s_t,a_t) + \alpha \big(r_t + \gamma\max_\limits{a} Q(s_{t+1},a_{t+1}) - Q(s_t,a_t)\big)
\end{align}

where $Q(s,a)$ is the value function for action $a$ at state $s$, $\alpha$ is the learning rate, $r$ is the reward, and $\gamma$ is the temporal discount rate.

The expression $r_t + \gamma\max_\limits{a} Q(s_{t+1},a_{t+1})$ is referred to as the TD target while the full expression
\begin{align}
r_t + \gamma\max_\limits{a} Q(s_{t+1},a_{t+1}) - Q(s_t,a_t),
\end{align}
i.e. the difference between the TD target and the current Q-value, is referred to as the TD error, or reward prediction error.

Because of the max operator used to select the optimal Q-value in the TD target, Q-learning directly estimates the optimal action value, i.e. the cumulative future reward that would be obtained if the agent behaved optimally, regardless of the policy currently followed by the agent. For this reason, Q-learning is referred to as an **off-policy** method.

## Exercise 1: Implement the Q-learning algorithm

In this exercise you will implement the Q-learning update rule described above. It takes in as arguments the previous state $s_t$, the action $a_t$ taken, the reward received $r_t$, the current state $s_{t+1}$, the Q-value table, and a dictionary of parameters that contain the learning rate $\alpha$ and discount factor $\gamma$. The method returns the updated Q-value table. For the parameter dictionary, $\alpha$: `params['alpha']` and $\gamma$: `params['gamma']`.

```python
def q_learning(state, action, reward, next_state, value, params):
  """Q-learning: updates the value function and returns it.

  Args:
    state (int): the current state identifier
    action (int): the action taken
    reward (float): the reward received
    next_state (int): the transitioned to state identifier
    value (ndarray): current value function of shape (n_states, n_actions)
    params (dict): a dictionary containing the default parameters

  Returns:
    ndarray: the updated value function of shape (n_states, n_actions)
  """
  # Q-value of current state-action pair
  q = value[state, action]

  ##########################################################
  ## TODO for students: implement the Q-learning update rule
  # Fill out function and remove
  #raise NotImplementedError("Student excercise: implement the Q-learning update rule")
  ##########################################################

  # write an expression for finding the maximum Q-value at the current state
  if next_state is None:
    max_next_q = 0
  else:
    max_next_q = np.max(value[next_state])

  # write the expression to compute the TD error
  td_error = reward + params['gamma']*max_next_q - q
  # write the expression that updates the Q-value for the state-action pair
  value[state, action] = q +params['alpha']*td_error

  return value
```

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

Now that we have our Q-learning algorithm, let's see how it handles learning to solve the Cliff World environment.

You will recall from the previous tutorial that a major part of reinforcement learning algorithms are their ability to balance exploitation and exploration. For our Q-learning agent, we again turn to the epsilon-greedy strategy. At each step, the agent will decide with probability $1 - \epsilon$ to use the best action for the state it is currently in by looking at the value function, otherwise just make a random choice.

The process by which our the agent will interact with and learn about the environment is handled for you in the helper function `learn_environment`. This implements the entire learning episode lifecycle of stepping through the state observation, action selection (epsilon-greedy) and execution, reward, and state transition. Feel free to review that code later to see how it all fits together, but for now let's test out our agent.

```python
# set for reproducibility, comment out / change seed value for different results
np.random.seed(1)

# parameters needed by our policy and learning rule
params = {
  'epsilon': 0.1,  # epsilon-greedy policy
  'alpha': 0.1,  # learning rate
  'gamma': 1.0,  # discount factor
}

# episodes/trials
n_episodes = 500
max_steps = 1000

# environment initialization
env = CliffWorld()

# solve Cliff World using Q-learning
results = learn_environment(env, q_learning, params, max_steps, n_episodes)
value_qlearning, reward_sums_qlearning = results

# Plot results
plot_performance(env, value_qlearning, reward_sums_qlearning)
```

If all went well, we should see four plots that show different aspects on our agent's learning and progress.

* The top left is a representation of the Q-table itself, showing the values for different actions in different states. Notably, going right from the starting state or down when above the cliff is clearly very bad.
* The top right figure shows the greedy policy based on the Q-table, i.e. what action would the agent take if it only took its best guess in that state.
* The bottom right is the same as the top, only instead of showing the action, it's showing a representation of the maximum Q-value at a particular state.
* The bottom left is the actual proof of learning, as we see the total reward steadily increasing after each episode until asymptoting at the maximum possible reward of -11.

Feel free to try changing the parameters or random seed and see how the agent's behavior changes.

---
# Summary

In this tutorial you implemented a reinforcement learning agent based on Q-learning to solve the Cliff World environment. Q-learning combined the epsilon-greedy approach to exploration-expoitation with a table-based value function to learn the expected future rewards for each state.

---
# Bonus

## SARSA

An alternative to Q-learning, the SARSA algorithm also estimates action values. However, rather than estimating the optimal (off-policy) values, SARSA estimates the **on-policy** action value, i.e. the cumulative future reward that would be obtained if the agent behaved according to its current beliefs.

\begin{align}
Q(s_t,a_t) \leftarrow Q(s_t,a_t) + \alpha \big(r_t + \gamma Q(s_{t+1},a_{t+1}) - Q(s_t,a_t)\big)
\end{align}

where, once again, $Q(s,a)$ is the value function for action $a$ at state $s$, $\alpha$ is the learning rate, $r$ is the reward, and $\gamma$ is the temporal discount rate.

In fact, you will notices that the *only* difference between Q-learning and SARSA is the TD target calculation uses the policy to select the next action (in our case epsilon-greedy) rather than using the action that maximizes the Q-value.

### Exercise 2: Implement the SARSA algorithm

In this exercise you will implement the SARSA update rule described above. Just like Q-learning, it takes in as arguments the previous state $s_t$, the action $a_t$ taken, the reward received $r_t$, the current state $s_{t+1}$, the Q-value table, and a dictionary of parameters that contain the learning rate $\alpha$ and discount factor $\gamma$. The method returns the updated Q-value table. You may use the `epsilon_greedy` function to acquire the next action. For the parameter dictionary, $\alpha$: `params['alpha']`, $\gamma$: `params['gamma']`, and $\epsilon$: `params['epsilon']`.

```python
def sarsa(state, action, reward, next_state, value, params):
  """SARSA: updates the value function and returns it.

  Args:
    state (int): the current state identifier
    action (int): the action taken
    reward (float): the reward received
    next_state (int): the transitioned to state identifier
    value (ndarray): current value function of shape (n_states, n_actions)
    params (dict): a dictionary containing the default parameters

  Returns:
    ndarray: the updated value function of shape (n_states, n_actions)
  """
  # value of previous state-action pair
  q = value[state, action]

  ##########################################################
  ## TODO for students: implement the SARSA update rule
  # Fill out function and remove
  raise NotImplementedError("Student excercise: implement the SARSA update rule")
  ##########################################################

  # select the expected value at current state based on our policy by sampling
  # from it
  if next_state is None:
    policy_next_q = 0
  else:
    # write an expression for selecting an action using epsilon-greedy
    policy_action = ...
    # write an expression for obtaining the value of the policy action at the
    # current state
    policy_next_q = ...

  # write the expression to compute the TD error
  td_error = ...
  # write the expression that updates the Q-value for the state-action pair
  value[state, action] = ...

  return value
```

```python
def sarsa(state, action, reward, next_state, value, params):
  """SARSA: updates the value function and returns it.

  Args:
    state (int): the current state identifier
    action (int): the action taken
    reward (float): the reward received
    next_state (int): the transitioned to state identifier
    value (ndarray): current value function of shape (n_states, n_actions)
    params (dict): a dictionary containing the default parameters

  Returns:
    ndarray: the updated value function of shape (n_states, n_actions)
  """
  # value of previous state-action pair
  q = value[state, action]

  # select the expected value at current state based on our policy by sampling
  # from it
  if next_state is None:
    policy_next_q = 0
  else:
    # write an expression for selecting an action using epsilon-greedy
    policy_action = epsilon_greedy(value[next_state], params['epsilon'])
    # write an expression for obtaining the value of the policy action at the
    # current state
    policy_next_q = value[next_state, policy_action]

  # write the expression to compute the TD error
  td_error = reward + params['gamma'] * policy_next_q - q
  # write the expression that updates the Q-value for the state-action pair
  value[state, action] = q + params['alpha'] * td_error

  return value
```

Now that we have an implementation for SARSA, let's see how it tackles Cliff World. We will again use the same setup we tried with Q-learning.

```python
# set for reproducibility, comment out / change seed value for different results
np.random.seed(1)

# parameters needed by our policy and learning rule
params = {
  'epsilon': 0.1,  # epsilon-greedy policy
  'alpha': 0.1,  # learning rate
  'gamma': 1.0,  # discount factor
}

# episodes/trials
n_episodes = 500
max_steps = 1000

# environment initialization
env = CliffWorld()

# learn Cliff World using Sarsa
results = learn_environment(env, sarsa, params, max_steps, n_episodes)
value_sarsa, reward_sums_sarsa = results

# Plot results
plot_performance(env, value_sarsa, reward_sums_sarsa)
```

We should see that SARSA also solves the task with similar looking outcomes to Q-learning. One notable difference is that SARSA seems to be skittsh around the cliff edge and often goes further away before coming back down to the goal.

Again, feel free to try changing the parameters or random seed and see how the agent's behavior changes.

## On-Policy vs Off-Policy

We have now seen an example of both on- and off-policy learning algorithms. Let's compare both Q-learning and SARSA reward results again, side-by-side, to see how they stack up.

```python
# parameters needed by our policy and learning rule
params = {
  'epsilon': 0.1,  # epsilon-greedy policy
  'alpha': 0.1,  # learning rate
  'gamma': 1.0,  # discount factor
}

# episodes/trials
n_episodes = 500
max_steps = 1000

# environment initialization
env = CliffWorld()

# learn Cliff World using Sarsa
np.random.seed(1)
results = learn_environment(env, q_learning, params, max_steps, n_episodes)
value_qlearning, reward_sums_qlearning = results
np.random.seed(1)
results = learn_environment(env, sarsa, params, max_steps, n_episodes)
value_sarsa, reward_sums_sarsa = results
```

```python
fig, ax = plt.subplots()
ax.plot(reward_sums_qlearning, label='Q-learning')
ax.plot(reward_sums_sarsa, label='SARSA')
ax.set(xlabel='Episodes', ylabel='Total reward')
plt.legend(loc='lower right');
```

On this simple Cliff World task, Q-learning and SARSA are almost indistinguisable from a performance standpoint, but we can see that Q-learning has a slight-edge within the 500 episode time horizon. Let's look at the illustrated "greedy policy" plots again.

```python
fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(16, 6))
plot_quiver_max_action(env, value_qlearning, ax=ax1)
ax1.set(title='Q-learning maximum value/probability actions')
plot_quiver_max_action(env, value_sarsa, ax=ax2)
ax2.set(title='SARSA maximum value/probability actions');
```

What should immediately jump out is that Q-learning learned to go up, then immediately go to the right, skirting the cliff edge, until it hits the wall and goes down to the goal. The policy further away from the cliff is less certain.

SARSA, on the other hand, appears to avoid the cliff edge, going up one more tile before starting over to the goal side. This also clearly solves the challenge of getting to the goal, but does so at an additional -2 cost over the truly optimal route.

Why do you think these behaviors emerged the way they did?


# An Introduction to SageMaker Neural Topic Model

***Unsupervised representation learning and topic extraction using Neural Topic Model***

1. [Introduction](#Introduction)
1. [Data Preparation](#Data-Preparation)
1. [Model Training](#Model-Training)
1. [Model Hosting and Inference](#Model-Hosting-and-Inference)
1. [Model Exploration](#Model-Exploration)

---
## Introduction

Amazon SageMaker Neural Topic Model (NTM) is an unsupervised learning algorithm that attempts to describe a set of observations as a mixture of distinct categories. NTM is most commonly used to discover a user-specified number of topics shared by documents within a text corpus. Here each observation is a document, the features are the presence (or occurrence count) of each word, and the categories are the topics. Since the method is unsupervised, the topics are not specified upfront and are not guaranteed to align with how a human may naturally categorize documents. The topics are learned as a probability distribution over the words that occur in each document. Each document, in turn, is described as a mixture of topics.

In this notebook, we will use the Amazon SageMaker NTM algorithm to train a model on the [20NewsGroups](https://archive.ics.uci.edu/ml/datasets/Twenty+Newsgroups) data set. This data set has been widely used as a topic modeling benchmark.

The main goals of this notebook are as follows:

1. learn how to obtain and store data for use in Amazon SageMaker,
2. create an AWS SageMaker training job on a data set to produce an NTM model,
3. use the model to perform inference with an Amazon SageMaker endpoint.
4. explore trained model and visualized learned topics

If you would like to know more please check out the [SageMaker Neural Topic Model Documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/ntm.html).

### A Brief Overview of SageMaker NTM

Topic models are a classical example of probablistic graphical models that involve challenging posterior inference problems. We implement topic modeling under a neural-network based variational inference framework. The difficult inference problem is framed as an optimization problem solved by scalable methods such as stochastic gradient descent. Compared to conventional inference schemes, the neural-network implementation allows for scalable model training as well as low-latency inference. Furthermore, the flexibility of the neural inference framework allows us to more quickly add new functionalities and serve a wider range of customer use cases.

The high-level diagram of SageMaker NTM is shown below:

$$
\begin{equation}
\begin{split}
\textrm{Encoder} \ q(z\vert x): & \\
& \pi = f_X^{MLP}(x), \quad \mu(x) = l_1(\pi), \quad \log \sigma(x) = l_2(\pi)\\
& h(x, \epsilon) = \mu + \sigma \epsilon, \ \textrm{where} \ \epsilon \sim \mathcal{N}(0,I),\quad z=g(h), \ \textrm{where} \ h \sim \mathcal{N}(\mu, \sigma^2 I) \\
\textrm{Decoder} \ p(x\vert z): & \\
& y(z) = \textrm{softmax}(Wz+b), \quad \log p(x\vert z) = \sum x \odot \log(y(z))
\end{split}
\end{equation}
$$

where $l_1$ and $l_2$ are linear transformations with bias.

### Beyond Text Data

In principle, topic models can be applied to types of data beyond text documents. For example, topic modeling has been applied on network traffice data to discover [peer-to-peer applications usage patterns](http://isis.poly.edu/~baris/papers/GangsOfInternet-CNS13.pdf). We would be glad to hear about your novel use cases and are happy to help provide additional information. Please feel free to post questions or feedback at our [GitHub repository](https://github.com/awslabs/amazon-sagemaker-examples) or in the [Amazon SageMaker](https://forums.aws.amazon.com/forum.jspa?forumID=285) section of AWS Developer Forum.

## Install Python packages

```python
import sys

!{sys.executable} -m pip install "scikit_learn==0.20.0" "nltk==3.4.4"
```

---
## Data Preparation

The 20Newsgroups data set is a collection of approximately 20,000 newsgroup documents, partitioned (nearly) evenly across 20 different newsgroups. This collection has become a popular data set for experiments in text applications of machine learning techniques, such as text classification and text clustering. Here, we will see what topics we can learn from this set of documents with NTM. The data setis available at the UCI Machine Learning Repository at this [location](https://archive.ics.uci.edu/ml/datasets/Twenty+Newsgroups). Please aware of the following requirements about ackonwledge, copyright and availability, cited from the [data set description page](https://archive.ics.uci.edu/ml/machine-learning-databases/20newsgroups-mld/20newsgroups.data.html).

> **Acknowledgements, Copyright Information, and Availability**

>You may use this material free of charge for any educational purpose, provided attribution is given in any lectures or publications that make use of this material.

### Fetching Data Set

First let's define the folder to hold the data and clean the content in it which might be from previous experiments.

```python
import os
import shutil

data_dir = "20_newsgroups_bulk"
if os.path.exists(data_dir):  # cleanup existing data folder
    shutil.rmtree(data_dir)
```

Now we can download the data. We download the [`20 newsgroups dataset`](http://qwone.com/~jason/20Newsgroups/). The `20 newsgroups dataset` consists of 20000 messages taken from 20 Usenet newsgroups.

```python
!aws s3 cp s3://sagemaker-sample-files/datasets/text/20_newsgroups/20_newsgroups_bulk.tar.gz .
```

```python
!tar xzf 20_newsgroups_bulk.tar.gz
!ls 20_newsgroups_bulk
```

```python
file_list = [os.path.join(data_dir, f) for f in os.listdir(data_dir)]
print("Number of files:", len(file_list))
```

```python
import pandas as pd

documents_count = 0
for file in file_list:
    df = pd.read_csv(file, header=None, names=["text"])
    documents_count = documents_count + df.shape[0]
print("Number of documents:", documents_count)
```

This following function will remove the header, footer and quotes (of earlier messages in each text).

```python
def strip_newsgroup_item(item):
    item = strip_newsgroup_header(item)
    item = strip_newsgroup_quoting(item)
    item = strip_newsgroup_footer(item)
    return item
```

```python
from sklearn.datasets.twenty_newsgroups import (
    strip_newsgroup_header,
    strip_newsgroup_quoting,
    strip_newsgroup_footer,
)

data = []

for file in file_list:
    print(f"Processing {file}")
    label = file.split("/")[1]
    df = pd.read_csv(file, header=None, names=["text"])
    df["text"] = df["text"].apply(strip_newsgroup_item)
    data.extend(df["text"].tolist())
```

As we can see below, the entries in the data set are just plain text paragraphs. We will need to process them into a suitable data format.

```python
data[10:13]
```

---
### From Plain Text to Bag-of-Words (BOW)

The input documents to the algorithm, both in training and inference, need to be vectors of integers representing word counts. This is so-called bag-of-words (BOW) representation. To convert plain text to BOW, we need to first "tokenize" our documents, i.e identify words and assign an integer id to each of them.

$$
\text{"cat"} \mapsto 0, \; \text{"dog"} \mapsto 1 \; \text{"bird"} \mapsto 2, \ldots
$$

Then, we count the occcurence of each of the tokens in each document and form BOW vectors as illustrated in the following example:

$$
w = \text{"cat bird bird bird cat"} \quad \longmapsto \quad w = [2, 0, 3, 0, \ldots, 0]
$$

Also, note that many real-world applications have large vocabulary sizes. It may be necessary to represent the input documents in sparse format. Finally, the use of stemming and lemmatization in data preprocessing provides several benefits. Doing so can improve training and inference compute time since it reduces the effective vocabulary size. More importantly, though, it can improve the quality of learned topic-word probability matrices and inferred topic mixtures. For example, the words *"parliament"*, *"parliaments"*, *"parliamentary"*, *"parliament's"*, and *"parliamentarians"* are all essentially the same word, *"parliament"*, but with different conjugations. For the purposes of detecting topics, such as a *"politics"* or *governments"* topic, the inclusion of all five does not add much additional value as they all essentially describe the same feature.

In this example, we will use a simple lemmatizer from [`nltk`](https://www.nltk.org/) package and use `CountVectorizer` in `scikit-learn` to perform the token counting. For more details please refer to their documentation respectively. Alternatively, [`spaCy`](https://spacy.io/) also offers easy-to-use tokenization and lemmatization functions.

---
In the following cell, we use a tokenizer and a lemmatizer from `nltk`. In the list comprehension, we implement a simple rule: only consider words that are longer than 2 characters, start with a letter and match the `token_pattern`.

```python
import nltk

nltk.download("punkt")
nltk.download("wordnet")
from nltk import word_tokenize
from nltk.stem import WordNetLemmatizer
import re

token_pattern = re.compile(r"(?u)\b\w\w+\b")

class LemmaTokenizer(object):
    def __init__(self):
        self.wnl = WordNetLemmatizer()

    def __call__(self, doc):
        return [
            self.wnl.lemmatize(t)
            for t in word_tokenize(doc)
            if len(t) >= 2 and re.match("[a-z].*", t) and re.match(token_pattern, t)
        ]
```

With the tokenizer defined we perform token counting next while limiting the vocabulary size to `vocab_size`. Use a maximum document frequency of 95% of documents (`max_df=0.95`) and a minimum document frequency of 2 documents (`min_df=2`).

```python
import time
import numpy as np
from sklearn.feature_extraction.text import CountVectorizer

vocab_size = 2000
print("Tokenizing and counting, this may take a few minutes...")
start_time = time.time()
vectorizer = CountVectorizer(
    input="content",
    analyzer="word",
    stop_words="english",
    tokenizer=LemmaTokenizer(),
    max_features=vocab_size,
    max_df=0.95,
    min_df=2,
)
vectors = vectorizer.fit_transform(data)
vocab_list = vectorizer.get_feature_names()
print("vocab size:", len(vocab_list))

# random shuffle
idx = np.arange(vectors.shape[0])
np.random.shuffle(idx)
vectors = vectors[idx]

print("Done. Time elapsed: {:.2f}s".format(time.time() - start_time))
```

Optionally, we may consider removing very short documents, the following cell removes documents shorter than 25 words. This certainly depends on the application, but there are also some general justifications. It is hard to imagine very short documents express more than one topic. Topic modeling tries to model each document as a mixture of multiple topics, thus it may not be the best choice for modeling short documents.

```python
threshold = 25
vectors = vectors[
    np.array(vectors.sum(axis=1) > threshold).reshape(
        -1,
    )
]
print("removed short docs (<{} words)".format(threshold))
print(vectors.shape)
```

The output from `CountVectorizer` are sparse matrices with their elements being integers.

```python
print(type(vectors), vectors.dtype)
print(vectors[0])
```

Because all the parameters (weights and biases) in the NTM model are `np.float32` type we'd need the input data to also be in `np.float32`. It is better to do this type-casting upfront rather than repeatedly casting during mini-batch training.

```python
import scipy.sparse as sparse

vectors = sparse.csr_matrix(vectors, dtype=np.float32)
print(type(vectors), vectors.dtype)
```

As a common practice in modeling training, we should have a training set, a validation set, and a test set. The training set is the set of data the model is actually being trained on. But what we really care about is not the model's performance on training set but its performance on future, unseen data. Therefore, during training, we periodically calculate scores (or losses) on the validation set to validate the performance of the model on unseen data. By assessing the model's ability to generalize we can stop the training at the optimal point via early stopping to avoid over-training.

Note that when we only have a training set and no validation set, the NTM model will rely on scores on the training set to perform early stopping, which could result in over-training. Therefore, we recommend always supply a validation set to the model.

Here we use 80% of the data set as the training set and the rest for validation set and test set. We will use the validation set in training and use the test set for demonstrating model inference.

```python
n_train = int(0.8 * vectors.shape[0])

# split train and test
train_vectors = vectors[:n_train, :]
test_vectors = vectors[n_train:, :]

# further split test set into validation set (val_vectors) and test  set (test_vectors)
n_test = test_vectors.shape[0]
val_vectors = test_vectors[: n_test // 2, :]
test_vectors = test_vectors[n_test // 2 :, :]
```

```python
print(train_vectors.shape, test_vectors.shape, val_vectors.shape)
```

---
### Store Data on S3

The NTM algorithm, as well as other first-party SageMaker algorithms, accepts data in [RecordIO](https://mxnet.apache.org/api/python/io/io.html#module-mxnet.recordio) [Protobuf](https://developers.google.com/protocol-buffers/) format. The SageMaker Python API provides helper functions for easily converting your data into this format. Below we convert the from numpy/scipy data and upload it to an Amazon S3 destination for the model to access it during training.

### Setup AWS Credentials

We first need to specify data locations and access roles. In particular, we need the following data:

- The S3 `bucket` and `prefix` that you want to use for training and model data.  This should be within the same region as the Notebook Instance, training, and hosting.
- The IAM `role` is used to give training and hosting access to your data. See the documentation for how to create these.  Note, if more than one role is required for notebook instances, training, and/or hosting, please replace the boto regexp with a the appropriate full IAM role arn string(s).

```python
import os
import sagemaker
from sagemaker import get_execution_role

role = get_execution_role()

sess = sagemaker.Session()
bucket = sess.default_bucket()
```

```python
prefix = "20newsgroups"

train_prefix = os.path.join(prefix, "train")
val_prefix = os.path.join(prefix, "val")
output_prefix = os.path.join(prefix, "output")

s3_train_data = os.path.join("s3://", bucket, train_prefix)
s3_val_data = os.path.join("s3://", bucket, val_prefix)
output_path = os.path.join("s3://", bucket, output_prefix)
print("Training set location", s3_train_data)
print("Validation set location", s3_val_data)
print("Trained model will be saved at", output_path)
```

Here we define a helper function to convert the data to RecordIO Protobuf format and upload it to S3. In addition, we will have the option to split the data into several parts specified by `n_parts`.

The algorithm inherently supports multiple files in the training folder ("channel"), which could be very helpful for large data set. In addition, when we use distributed training with multiple workers (compute instances), having multiple files allows us to distribute different portions of the training data to different workers conveniently.

Inside this helper function we use `write_spmatrix_to_sparse_tensor` function provided by [SageMaker Python SDK](https://github.com/aws/sagemaker-python-sdk) to convert scipy sparse matrix into RecordIO Protobuf format.

```python
# update sagemake package, in order to use write_spmatrix_to_sparse_tensor in the next cell
# !pip install -U sagemaker
```

```python
def split_convert_upload(sparray, bucket, prefix, fname_template="data_part{}.pbr", n_parts=2):
    import io
    import boto3
    import sagemaker.amazon.common as smac

    chunk_size = sparray.shape[0] // n_parts
    for i in range(n_parts):

        # Calculate start and end indices
        start = i * chunk_size
        end = (i + 1) * chunk_size
        if i + 1 == n_parts:
            end = sparray.shape[0]

        # Convert to record protobuf
        buf = io.BytesIO()
        smac.write_spmatrix_to_sparse_tensor(array=sparray[start:end], file=buf, labels=None)
        buf.seek(0)

        # Upload to s3 location specified by bucket and prefix
        fname = os.path.join(prefix, fname_template.format(i))
        boto3.resource("s3").Bucket(bucket).Object(fname).upload_fileobj(buf)
        print("Uploaded data to s3://{}".format(os.path.join(bucket, fname)))
```

```python
split_convert_upload(
    train_vectors, bucket=bucket, prefix=train_prefix, fname_template="train_part{}.pbr", n_parts=8
)
split_convert_upload(
    val_vectors, bucket=bucket, prefix=val_prefix, fname_template="val_part{}.pbr", n_parts=1
)
```

---
## Model Training

We have created the training and validation data sets and uploaded them to S3. Next, we configure a SageMaker training job to use the NTM algorithm on the data we prepared

SageMaker uses Amazon Elastic Container Registry (ECR) docker container to host the NTM training image. The following ECR containers are currently available for SageMaker NTM training in different regions. For the latest Docker container registry please refer to [Amazon SageMaker: Common Parameters](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-algo-docker-registry-paths.html).

```python
import boto3
from sagemaker.image_uris import retrieve

container = retrieve("ntm", boto3.Session().region_name)
```

The code in the cell below automatically chooses an algorithm container based on the current region. In the API call to `sagemaker.estimator.Estimator` we also specify the type and count of instances for the training job. Because the 20NewsGroups data set is relatively small, we have chosen a CPU only instance (`ml.c4.xlarge`), but do feel free to change to [other instance types](https://aws.amazon.com/sagemaker/pricing/instance-types/). NTM fully takes advantage of GPU hardware and in general trains roughly an order of magnitude faster on a GPU than on a CPU. Multi-GPU or multi-instance training further improves training speed roughly linearly if communication overhead is low compared to compute time.

```python
import sagemaker

sess = sagemaker.Session()
ntm = sagemaker.estimator.Estimator(
    container,
    role,
    instance_count=2,
    instance_type="ml.c4.xlarge",
    output_path=output_path,
    sagemaker_session=sess,
)
```

## Hyperparameters

Here we highlight a few hyperparameters. For information about the full list of available hyperparameters, please refer to [NTM Hyperparameters](https://docs.aws.amazon.com/sagemaker/latest/dg/ntm_hyperparameters.html).

- **feature_dim** - the "feature dimension", it should be set to the vocabulary size
- **num_topics** - the number of topics to extract
- **mini_batch_size** - this is the batch size for each worker instance. Note that in multi-GPU instances, this number will be further divided by the number of GPUs. Therefore, for example, if we plan to train on an 8-GPU machine (such as `ml.p2.8xlarge`) and wish each GPU to have 1024 training examples per batch, `mini_batch_size` should be set to 8196.
- **epochs** - the maximal number of epochs to train for, training may stop early
- **num_patience_epochs** and **tolerance** controls the early stopping behavior. Roughly speaking, the algorithm will stop training if within the last `num_patience_epochs` epochs there have not been improvements on validation loss. Improvements smaller than `tolerance` will be considered non-improvement.
- **optimizer** and **learning_rate** - by default we use `adadelta` optimizer and `learning_rate` does not need to be set. For other optimizers, the choice of an appropriate learning rate may require experimentation.

```python
num_topics = 20
ntm.set_hyperparameters(
    num_topics=num_topics,
    feature_dim=vocab_size,
    mini_batch_size=128,
    epochs=100,
    num_patience_epochs=5,
    tolerance=0.001,
)
```

Next, we need to specify how the training data and validation data will be distributed to the workers during training. There are two modes for data channels:

- `FullyReplicated`: all data files will be copied to all workers
- `ShardedByS3Key`: data files will be sharded to different workers, i.e. each worker will receive a different portion of the full data set.

At the time of writing, by default, the Python SDK will use `FullyReplicated` mode for all data channels. This is desirable for validation (test) channel but not suitable for training channel. The reason is that when we use multiple workers we would like to go through the full data set by each of them going through a different portion of the data set, so as to provide different gradients within epochs. Using `FullyReplicated` mode on training data not only results in slower training time per epoch (nearly 1.5X in this example), but also defeats the purpose of distributed training. To set the training data channel correctly we specify `distribution` to be `ShardedByS3Key` for the training data channel as follows.

```python
from sagemaker.inputs import TrainingInput

s3_train = TrainingInput(s3_train_data, distribution="ShardedByS3Key")
```

Now we are ready to train. The following cell takes a few minutes to run. The command below will first provision the required hardware. You will see a series of dots indicating the progress of the hardware provisioning process. Once the resources are allocated, training logs will be displayed. With multiple workers, the log color and the ID following `INFO` identifies logs emitted by different workers.

```python
ntm.fit({"train": s3_train, "validation": s3_val_data})
```

If you see the message

> `===== Job Complete =====`

at the bottom of the output logs then that means training successfully completed and the output NTM model was stored in the specified output path. You can also view information about and the status of a training job using the AWS SageMaker console. Just click on the "Jobs" tab and select training job matching the training job name, below:

```python
print("Training job name: {}".format(ntm.latest_training_job.job_name))
```

## Model Hosting and Inference

A trained NTM model does nothing on its own. We now want to use the model we computed to perform inference on data. For this example, that means predicting the topic mixture representing a given document.

We create an inference endpoint using the SageMaker Python SDK `deploy()` function from the job we defined above. We specify the instance type where inference is computed as well as an initial number of instances to spin up.

```python
from sagemaker.serializers import CSVSerializer
from sagemaker.deserializers import JSONDeserializer

ntm_predictor = ntm.deploy(
    initial_instance_count=1,
    instance_type="ml.m4.xlarge",
    serializer=CSVSerializer(),
    deserializer=JSONDeserializer(),
)
```

Congratulations! You now have a functioning SageMaker NTM inference endpoint. You can confirm the endpoint configuration and status by navigating to the "Endpoints" tab in the AWS SageMaker console and selecting the endpoint matching the endpoint name, below:

```python
print("Endpoint name: {}".format(ntm_predictor.endpoint_name))
```

Let's pass 5 examples from the test set to the inference endpoint

```python
test_data = np.array(test_vectors.todense())
results = ntm_predictor.predict(test_data[:5])
print(results)
```

We can see the output format of SageMaker NTM inference endpoint is a Python dictionary with the following format.

```
{
  'predictions': [
    {'topic_weights': [ ... ] },
    {'topic_weights': [ ... ] },
    {'topic_weights': [ ... ] },
    ...
  ]
}
```

We extract the topic weights, themselves, corresponding to each of the input documents.

```python
predictions = np.array([prediction["topic_weights"] for prediction in results["predictions"]])
print(predictions)
```

---
### Inference with RecordIO Protobuf
The inference endpoint also supports JSON-formatted and RecordIO Protobuf, see [Common Data Formats—Inference](https://docs.aws.amazon.com/sagemaker/latest/dg/cdf-inference.html) for more information.

At the time of writing SageMaker Python SDK does not yet have a RecordIO Protobuf serializer, but it is fairly straightforward to create one as follows.

```python
def recordio_protobuf_serializer(spmatrix):
    import io
    import sagemaker.amazon.common as smac

    buf = io.BytesIO()
    smac.write_spmatrix_to_sparse_tensor(array=spmatrix, file=buf, labels=None)
    buf.seek(0)
    return buf
```

If you decide to compare these results to the known topic weights generated above keep in mind that SageMaker NTM discovers topics in no particular order. That is, the approximate topic mixtures computed above may be (approximate) permutations of the known topic mixtures corresponding to the same documents.

---
Now we can take a look at how the 20 topics are assigned to the 5 test documents with a bar plot.

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

%matplotlib inline

fs = 12
df = pd.DataFrame(predictions.T)
df.plot(kind="bar", figsize=(16, 4), fontsize=fs)
plt.ylabel("Topic assignment", fontsize=fs + 2)
plt.xlabel("Topic ID", fontsize=fs + 2)
```

### Stop / Close the Endpoint

Finally, we should delete the endpoint before we close the notebook.

To restart the endpoint you can follow the code above using the same `endpoint_name` we created or you can navigate to the "Endpoints" tab in the SageMaker console, select the endpoint with the name stored in the variable `endpoint_name`, and select "Delete" from the "Actions" dropdown menu.

```python
sagemaker.Session().delete_endpoint(ntm_predictor.endpoint_name)
```

## Model Exploration

***Note: The following section is meant as a deeper dive into exploring the trained models. The demonstrated functionalities may not be fully supported or guaranteed. For example, the parameter names may change without notice.***

The trained model artifact is a compressed package of MXNet models from the two workers. To explore the model, we first need to install mxnet.

```python
# If you use conda_mxnet_p36 kernel, mxnet is already installed, otherwise, uncomment the following line to install.
# !pip install mxnet
import mxnet as mx
```

Here we download unpack the artifact

```python
model_path = os.path.join(output_prefix, ntm._current_job_name, "output/model.tar.gz")
model_path
```

```python
boto3.resource("s3").Bucket(bucket).download_file(model_path, "downloaded_model.tar.gz")
```

```python
!pwd
```

```python
import tarfile

tarfile.open("downloaded_model.tar.gz").extractall()
```

```python
import zipfile

with zipfile.ZipFile("model_algo-1", "r") as zip_ref:
    zip_ref.extractall("./")
```

We can load the model parameters and extract the weight matrix $W$ in the decoder as follows

```python
model = mx.ndarray.load("params")
W = model["arg:projection_weight"]
```

Matrix $W$ corresponds to the $W$ in the NTM digram at the beginning of this notebook. Each column of $W$ corresponds to a learned topic. The elements in the columns of $W$ corresponds to the pseudo-probability of a word within a topic. We can visualize each topic as a word cloud with the size of each word be proportional to the pseudo-probability of the words appearing under each topic.

```python
import sys

!{sys.executable} -m pip install wordcloud

import wordcloud as wc
```

```python
import matplotlib.pyplot as plt

word_to_id = dict()
for i, v in enumerate(vocab_list):
    word_to_id[v] = i

limit = 24
n_col = 4
counter = 0

plt.figure(figsize=(20, 16))
for ind in range(num_topics):

    if counter >= limit:
        break

    title_str = "Topic{}".format(ind)

    # pvals = mx.nd.softmax(W[:, ind]).asnumpy()
    pvals = mx.nd.softmax(mx.nd.array(W[:, ind])).asnumpy()

    word_freq = dict()
    for k in word_to_id.keys():
        i = word_to_id[k]
        word_freq[k] = pvals[i]

    wordcloud = wc.WordCloud(background_color="white").fit_words(word_freq)

    plt.subplot(limit // n_col, n_col, counter + 1)
    plt.imshow(wordcloud, interpolation="bilinear")
    plt.axis("off")
    plt.title(title_str)
    # plt.close()

    counter += 1
```

```python

```


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

# Support Vector Machines

Support Vector Machines (SVM) originated from the statistical learning theory
developed by Vapnik-Chervonenkis. As such, it represents a deep application of
statistical theory that incorporates the VC dimension concepts we
discussed in the first section. Let's start by looking at some pictures.
Consider the two-dimensional classification problem shown in
[Figure](#fig:svm_001).  [Figure](#fig:svm_001) shows two classes (gray and
white
circles) that can be separated by any of the lines shown. Specifically, any
such separating line can be written as the locus of points ($\mathbf{x}$) in
the two-dimensional plane that satisfy the following,

In the two-dimensional plane, the two classes (gray and white circles) are
easily separated by any one of the lines shown.

$$
\beta_0 + \boldsymbol{\beta}^T \mathbf{x} = 0
$$

 To classify  an arbitrary $\mathbf{x}$ using this line, we just
compute the sign of $\beta_0+\boldsymbol{\beta}^T \mathbf{x}$ and assign one
class to the positive sign and the other class to the negative sign.  To
uniquely specify such a separating line (or, hyperplane in a higher-dimensional
space) we need additional criteria.

[Figure](#fig:svm_002) shows the data with two bordering parallel lines that
form a margin around the central separating line.  The *maximal margin
algorithm* finds the widest margin and the unique separating line.  As a
consequence, the algorithm uncovers the elements in the data that touch the
margins. These are the *support* elements. The other elements
away from the border are not relevent to the solution. This reduces
model variance because the solution is insensitive to the removal of
elements other than these supporting elements (usually a small minority).

The maximal margin algorithm finds the separating line that maximizes the
margin shown. The elements that touch the margins are the support elements. The
dotted elements are not relevent to the solution.

To see how this works for linearly separable classes, consider a
training set consisting of $\lbrace (\mathbf{x},y) \rbrace$ where
$y\in \lbrace -1,1 \rbrace$. For any point $\mathbf{x}_i$, we
compute the functional margin as $\hat{ \gamma_i }=y_i (\beta_0 +
\boldsymbol{\beta}^T \mathbf{x}_i)$. Thus, $\hat{\gamma}_i >0$ when
$\mathbf{x}_i$ is correctly classified. The geometrical margin is
$\gamma = \hat{\gamma}/\lVert\boldsymbol{\beta}\rVert$. When
$\mathbf{x}_i$ is correctly classified, the geometrical margin is
equal to the perpendicular distance from $\mathbf{x}_i$ to the line.
Let's look see how the maximal margin algorithm works.

Let $M$ be the width of the margin.  The maximal margin algorithm is can be
formulated as a quadratic programming problem. We want to simultaneously
maximize the margin $M$ while ensuring that all of the data points are
correctly classified.

$$
\begin{aligned}
& \underset{\beta_0,\boldsymbol{\beta},\lVert\boldsymbol{\beta}\rVert=1}{\text{m
aximize}}
& & M \\\
& \text{subject to:}
& & y_i(\beta_0+\boldsymbol{\beta}^T \mathbf{x}_i) \geq M, \; i = 1, \ldots, N.
\end{aligned}
$$

 The first line says we want to generate a maximum value for $M$ by
adjusting $\beta_0$ and $\boldsymbol{\beta}$ while keeping
$\lVert\boldsymbol{\beta}\rVert=1$. The functional margins for each $i^{th}$
data element are the constraints to the problem and must be satisfied for every
proposed solution. In words, the constraints enforce that the elements have to
be correctly classified and outside of the margin around the separating line.
With some reformulation, it turns out that
$M=1/\lVert\boldsymbol{\beta}\rVert$ and this can be put into the following
standard format,

$$
\begin{aligned}
& \underset{\beta_0,\boldsymbol{\beta}}{\text{minimize}}
& & \lVert\boldsymbol{\beta}\rVert \\\
& \text{subject to:}
& & y_i(\beta_0+\boldsymbol{\beta}^T \mathbf{x}_i) \geq 1, \; i = 1, \ldots, N.
\end{aligned}
$$

 This is a convex optimization problem  and can be solved using
powerful
methods in that area.

The situation becomes more complex when the two classes are not separable and
we have to allow some unavoidable mixing between the two classes in the
solution. This means that the contraints have to modified as in the following,

$$
y_i(\beta_0+\boldsymbol{\beta}^T \mathbf{x}_i) \geq M(1-\xi_i)
$$

 where the $\xi_i$ are the slack variables and represent the
proportional amount tha the prediction is on the wrong side of the margin. Thus,
elements are misclassified when $\xi_i>1$. With these additional variables,
we have a more general formulation of the convex optimization problem,

$$
\begin{aligned}
& \underset{\beta_0,\boldsymbol{\beta}}{\text{minimize}}
& & \lVert\boldsymbol{\beta}\rVert \\\
& \text{subject to:}
& & y_i(\beta_0+\boldsymbol{\beta}^T \mathbf{x}_i) \geq 1-\xi_i, \\\
& & & \xi_i \geq 0, \sum \xi_i \leq \texttt{constant}, \; i = 1, \ldots, N.
\end{aligned}
$$

 which can be rewritten in the following equivalent form,

$$
\begin{equation}
\begin{aligned}
& \underset{\beta_0,\boldsymbol{\beta}}{\text{minimize}}
& & \frac{1}{2}\lVert\boldsymbol{\beta}\rVert + C \sum \xi_i \\\
& \text{subject to:}
& & y_i(\beta_0+\boldsymbol{\beta}^T \mathbf{x}_i) \geq 1-\xi_i, \xi_i \geq 0 \;
i = 1, \ldots, N.
\end{aligned}
\end{equation}
\label{eq:svm} \tag{1}
$$

 Because the $\xi_i$ terms are all positive, the objective
is to maximize the margin (i.e., minimize $\lVert\boldsymbol{\beta}\rVert$)
while minimizing the proportional drift of the predictions to the wrong side
of the margin (i.e., $C \sum \xi_i$). Thus, large values of $C$ shunt
algorithmic focus towards the correctly classified points near the
decision boundary and small values focus on further data. The value $C$ is
a hyperparameter for the SVM.

The good news is that all of these complicated pieces are handled neatly inside
of Scikit-learn. The following sets up the linear *kernel* for the SVM (more on
kernels soon),

```python
from sklearn.datasets import make_blobs
from sklearn.svm import SVC
sv = SVC(kernel='linear')
```

We can create some synthetic data using `make_blobs` and then
fit it to the SVM,

```python
X,y=make_blobs(n_samples=200, centers=2, n_features=2,
               random_state=0,cluster_std=.5)
sv.fit(X,y)
```

    SVC(C=1.0, cache_size=200, class_weight=None, coef0=0.0,
      decision_function_shape='ovr', degree=3, gamma='auto', kernel='linear',
      max_iter=-1, probability=False, random_state=None, shrinking=True,
      tol=0.001, verbose=False)

After fitting, the SVM now has the estimated support vectors and the
coefficients of the $\boldsymbol{\beta}$  in the `sv.support_vectors_` and
`sv.coef_` attributes, respectively. [Figure](#fig:svm_003) shows the two
sample classes (white and gray circles) and the line separating them that was
found by the maximal margin algorithm. The two parallel dotted lines show  the
margin. The large circles enclose the support vectors, which are the data
elements that are relevent to the solution. Notice that only these elements
can touch the edges of the margins.

```python
%matplotlib inline

from matplotlib.pylab import subplots
import numpy as np
xi = np.linspace(X[:,0].min(),X[:,0].max(),100)

fig,ax=subplots()
_=ax.scatter(X[:,0],X[:,1],c=y,s=50,cmap='gray',marker='o',alpha=.3)
_=ax.plot(sv.support_vectors_[:,0],sv.support_vectors_[:,1],'ko',markersize=20,alpha=.2)
_=ax.plot(xi,-sv.coef_[0,0]/sv.coef_[0,1]*xi- sv.intercept_/sv.coef_[0,1],'k',lw=3.)
margin = np.linalg.norm(sv.coef_)
_=ax.plot(xi,-sv.coef_[0,0]/sv.coef_[0,1]*xi-(sv.intercept_+margin/2.)/sv.coef_[0,1],'--k',lw=3.)
_=ax.plot(xi,-sv.coef_[0,0]/sv.coef_[0,1]*xi-(sv.intercept_-margin/2.)/sv.coef_[0,1],'--k',lw=3.)
```

The two class shown (white and gray circles) are linearly separable. The
maximal margin solution is shown by the dark black line in the middle. The
dotted lines show the extent of the margin.  The large circles indicate the
support vectors for the maximal margin solution.

```python

def draw_margins(sv,X,y,ax=None):
    sv.fit(X,y)
    xi = np.linspace(X[:,0].min(),X[:,0].max(),100)
    if ax is None: fig,ax=subplots()
    _=ax.scatter(X[:,0],X[:,1],c=y,s=50,cmap='gray',marker='o',alpha=.3)
    _=ax.plot(sv.support_vectors_[:,0],sv.support_vectors_[:,1],'ko',markersize=20,alpha=.2)
    _=ax.plot(xi,-sv.coef_[0,0]/sv.coef_[0,1]*xi- sv.intercept_/sv.coef_[0,1],'k',lw=3.)
    margin = np.linalg.norm(sv.coef_)
    _=ax.plot(xi,-sv.coef_[0,0]/sv.coef_[0,1]*xi- (sv.intercept_+margin/2.)/sv.coef_[0,1],'--k',lw=3.)
    _=ax.plot(xi,-sv.coef_[0,0]/sv.coef_[0,1]*xi- (sv.intercept_-margin/2.)/sv.coef_[0,1],'--k',lw=3.)
```

```python

X, y = make_blobs(n_samples=50, centers=2, n_features=2,
                  cluster_std=1,random_state=0)

fig,axs = subplots(2,2,sharex=True,sharey=True)
#fig.set_size_inches((12,6))
sv = SVC(kernel='linear',C=.0100)
draw_margins(sv,X,y,ax=axs[0,0])
_=axs[0,0].set_title('C=0.01')
sv = SVC(kernel='linear',C=1)
draw_margins(sv,X,y,ax=axs[0,1])
_=axs[0,1].set_title('C=1')
sv = SVC(kernel='linear',C=100)
draw_margins(sv,X,y,ax=axs[1,0])
_=axs[1,0].set_title('C=100')
sv = SVC(kernel='linear',C=10000)
draw_margins(sv,X,y,ax=axs[1,1])
_=axs[1,1].set_title('C=10000')
```

[Figure](#fig:svm_004) shows what happens when the value of $C$ changes.
Increasing this value emphasizes the $\xi$ part of the objective function in
Equation [eq:svm](#eq:svm). As shown in the top left panel, a small value for
$C$ means that
the algorithm is willing to accept many support vectors at the expense of
maximizing the margin. That is, the proportional amount that predictions are on
the wrong side of the margin is more acceptable with smaller $C$.  As the value
of $C$ increases, there are fewer support vectors because the optimization
process prefers to eliminate support vectors that are far away from the margins
and accept fewer of these that encroach into the margin. Note that as the value
of $C$ progresses through this figure, the separating line tilts slightly.

The maximal margin algorithm finds the separating line that maximizes the
margin shown. The elements that touch the margins are the support elements. The
dotted elements are not relevent to the solution.

## Kernel Tricks

Support Vector Machines provide a powerful method to deal with linear
separations, but they can also apply to non-linear boundaries by
exploiting the so-called *kernel trick*.  The convex optimization
formulation of the SVM includes a *dual* formulation that leads to a
solution that requires only the inner-products of the features. The
kernel trick is to substitute inner-products by nonlinear kernel
functions.  This can be thought of as mapping the original features
onto a possibly infinite dimensional space of new features.  That is,
if the data are not linearly separable in two-dimensional space (for
example) maybe they are separable in three-dimensional space (or
higher)?

To make this concrete, suppose the original input space is
$\mathbb{R}^n$ and we want to use a non-linear mapping
$\psi:\mathbf{x} \mapsto \mathcal{F}$ where $\mathcal{F}$ is an
inner-product space of higher dimension.  The kernel trick is to
calculate the inner-product in $\mathcal{F}$ using a kernel
function, $K(\mathbf{x}_i,\mathbf{x}_j) = \langle
\psi(\mathbf{x}_i),\psi(\mathbf{x}_j)\rangle$. The long way to
compute this is to first compute $\psi(\mathbf{x})$ and then do the
inner-product. The kernel-trick way to do it is to use the kernel
function and avoid computing $\psi$. In other words, the kernel
function returns what the inner-product in $\mathcal{F}$ would have
returned if $\psi$ had been applied. For example, to achieve an
$n^{th}$ polynomial mapping of the input space, we can use
$\kappa(\mathbf{x}_i,\mathbf{x}_j)=(\mathbf{x}_i^T\mathbf{x}_j+\theta)^n$.
For example, suppose the input space is $\mathbb{R}^2$ and
$\mathcal{F}=\mathbb{R}^4$ and we have the following mapping,

$$
\psi(\mathbf{x}) : (x_0,x_1) \mapsto (x_0^2,x_1^2,x_0 x_1, x_1 x_0)
$$

 The inner product in $\mathcal{F}$ is then,

$$
\langle \psi(\mathbf{x}),\psi(\mathbf{y})  \rangle = \langle
\mathbf{x},\mathbf{y}  \rangle^2
$$

 In other words, the kernel is the square of the inner
product in input space. The advantage of using the kernel instead of
simply enlarging the feature space is computational because you only
need to compute the kernel on all distinct pairs of the input space.
The following example should help make this concrete. First we create
some Sympy variables,

```python
import sympy as S
x0,x1=S.symbols('x:2',real=True)
y0,y1=S.symbols('y:2',real=True)
```

Next, we create the $\psi$ function that maps into $\mathbb{R}^4$
and the corresponding kernel function,

```python
psi = lambda x,y: (x**2,y**2,x*y,x*y)
kern = lambda x,y: S.Matrix(x).dot(y)**2
```

Notice that the inner product in $\mathbb{R}^4$  is
equal to the kernel function, which only uses wthe $\mathbb{R}^2$
variables.

```python
print(S.Matrix(psi(x0,x1)).dot(psi(y0,y1)))
print(S.expand(kern((x0,x1),(y0,y1)))) # same as above
```

    x0**2*y0**2 + 2*x0*x1*y0*y1 + x1**2*y1**2
    x0**2*y0**2 + 2*x0*x1*y0*y1 + x1**2*y1**2

**Polynomial Regression Using Kernels.**  Recall our favorite
linear regression problem from the regularization chapter,

$$
\min_{\boldsymbol{\beta}}  \Vert y - \mathbf{X}\boldsymbol{\beta}\Vert^2
$$

 where $\mathbf{X}$ is a $n\times m$ matrix with $m>n$. As
we discussed, there are multiple solutions to this problem. The
least-squares solution is the following:

$$
\boldsymbol{\beta}_{LS}=\mathbf{X}^T(\mathbf{X}\mathbf{X}^T)^{\text{-1}}\mathbf{
y}
$$

 Given a new feature vector $\mathbf{x}$, the corresponding estimator
for $\mathbf{y}$ is the following,

$$
\hat{\mathbf{y}} = \mathbf{x}^T\boldsymbol{\beta}_{LS}=\mathbf{x}^T\mathbf{X}^T(
\mathbf{X}\mathbf{X}^T)^{\text{-1}}\mathbf{y}
$$

 Using the kernel trick, the solution can be written more generally as
the following,

$$
\hat{\mathbf{y}}=\mathbf{k}(\mathbf{x})^T\mathbf{K}^{\text{-1}}\mathbf{y}
$$

 where the $n\times n$ kernel matrix $\mathbf{K}$ replaces
$\mathbf{X}\mathbf{X}^T$ and where $\mathbf{k}(\mathbf{x})$ is a $n$-vector of
components $\mathbf{k}(\mathbf{x})=[\kappa(\mathbf{x}_i,\mathbf{x})]$ and where
$\mathbf{K}_{i,j}=\kappa(\mathbf{x}_i,\mathbf{x}_j)$ for the kernel function
$\kappa$.  With this more general setup, we can substitute
$\kappa(\mathbf{x}_i,\mathbf{x}_j)=(\mathbf{x}_i^T\mathbf{x}_j+\theta)^n$ for
$n^{th}$-order polynomial regression [[bauckhagenumpy]](#bauckhagenumpy). Note
that ridge
regression can also be incorporated by inverting $(\mathbf{K}+\alpha
\mathbf{I})$, which can help stabilize poorly conditioned $\mathbf{K}$ matrices
with a tunable $\alpha$ hyper-parameter [[bauckhagenumpy]](#bauckhagenumpy).

For some kernels, the enlarged $\mathcal{F}$ space is infinite-dimensional.
Mercer's conditions provide technical restrictions on the kernel functions.
Powerful, well-studied kernels have been implemented in Scikit-learn. The
advantage of kernel functions may evaporate for when $n\rightarrow m$ in which
case using the $\psi$ functions instead can be more practicable.

```python
from matplotlib.pylab import cm
xi = np.linspace(X[:,0].min(),X[:,0].max(),100)
yi = np.linspace(X[:,1].min(),X[:,1].max(),100)

fig,ax=subplots()
_=ax.scatter(X[:,0],X[:,1],c=y,s=50,cmap='gray',marker='o',alpha=.3)
Xi,Yi = np.meshgrid(xi,yi)
Zi=sv.predict(np.c_[Xi.ravel(),Yi.ravel()]).reshape(Xi.shape)

_=ax.contourf(Xi,Yi,Zi,cmap=cm.Paired,alpha=0.2);
```