# 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?


```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);
```


```python
# This cell is mandatory in all Dymos documentation notebooks.
missing_packages = []
try:
    import openmdao.api as om
except ImportError:
    if 'google.colab' in str(get_ipython()):
        !python -m pip install openmdao[notebooks]
    else:
        missing_packages.append('openmdao')
try:
    import dymos as dm
except ImportError:
    if 'google.colab' in str(get_ipython()):
        !python -m pip install dymos
    else:
        missing_packages.append('dymos')
try:
    import pyoptsparse
except ImportError:
    if 'google.colab' in str(get_ipython()):
        !pip install -q condacolab
        import condacolab
        condacolab.install_miniconda()
        !conda install -c conda-forge pyoptsparse
    else:
        missing_packages.append('pyoptsparse')
if missing_packages:
    raise EnvironmentError('This notebook requires the following packages '
                           'please install them and restart this notebook\'s runtime: {",".join(missing_packages)}')
```

(examples:bryson_denham)=
# The Bryson-Denham Problem

The Bryson-Denham problem is a variation of the double integrator problem {cite}`bryson1975applied`. It can be stated as:

Minimize the control effort required to reverse the direction of motion of a frictionless sliding block such that the reversal happens with some limited amount of displacement.

## State and control variables

This system has two state variables, the position ($x$) and velocity ($v$) of the sliding block.

This system has a single control variable ($u$), the acceleration of the block.

The dynamics of the system are governed by

\begin{align}
  \dot{x} &= v \\
  \dot{v} &= u
\end{align}

## Problem Definition

We seek to minimize the time required to exit the well in the positive direction.

\begin{align}
    \mathrm{Minimize} \, J &= 0.5 \int_0^1 u^2 dt
\end{align}

Subject to the initial conditions

\begin{align}
    x_0 &= 0.0 \\
    v_0 &= 1.0
\end{align}

and the terminal constraints

\begin{align}
    x_f &= 0.0 \\
    v_f &= -1.0
\end{align}

In addition, $x$ is consrained to remain below a displacement of 1/9.

\begin{align}
    x &\le 1/9
\end{align}

## Dealing with integral costs in Dymos

In classic optimal control, the objective is often broken into the terminal component (the Mayer term) and the integral component (the Lagrange term).
Dymos does not distinguish between the two.
In this case, since the objective $J$ is an integrated quantity, we add a term to the ODE

\begin{align}
\dot{J} &= 0.5 u^2
\end{align}

## Defining the ODE

The following code implements the equations of motion for the mountain car problem.
Since the rate of $x$ is given by another state ($v$), and the rate of $v$ is given by a control ($u$), there is no need to compute their rates in the ODE.
Dymos can pull their values from those other states and controls.
The ODE, therefore, only needs to compute the rate of change of $J$.

A few things to note:

1. By providing the tag `dymos.state_rate_source:{name}`, we're letting Dymos know what states need to be integrated, there's no need to specify a rate source when using this ODE in our Phase.
2. Pairing the above tag with `dymos.state_units:{units}` means we don't have to specify units when setting properties for the state in our run script.

```python
import numpy as np
import openmdao.api as om

class BrysonDenhamODE(om.ExplicitComponent):

    def initialize(self):
        self.options.declare('num_nodes', types=int)

    def setup(self):
        nn = self.options['num_nodes']

        self.add_input('x', shape=(nn,), units='m')
        self.add_input('v', shape=(nn,), units='m/s')
        self.add_input('u', shape=(nn,), units='m/s**2')

        self.add_output('J_dot', shape=(nn,), units='m**2/s**4',
                        tags=['dymos.state_rate_source:J',
                              'dymos.state_units:m**2/s**3'])

        ar = np.arange(nn, dtype=int)

        self.declare_partials(of='J_dot', wrt='u', rows=ar, cols=ar)

    def compute(self, inputs, outputs):
        u = inputs['u']
        outputs['J_dot'] = 0.5 * u**2

    def compute_partials(self, inputs, partials):
        partials['J_dot', 'u'] = inputs['u']

```

## Solving the Bryson-Denham problem with Dymos

The following script solves the minimum-time mountain car problem with Dymos.
This problem is pretty trivial and can be solved using the SLSQP optimizer in scipy.

To begin, import the packages we require:

```python
import dymos as dm
import matplotlib.pyplot as plt
```

We then instantiate an OpenMDAO problem and set the optimizer and its options.

The call to `declare_coloring` tells the optimizer to attempt to find a sparsity pattern that minimizes the work required to compute the derivatives across the model.

SLSQP does not internally use this sparsity information to reduce memory and improve performance as some other optimizers do, but the performance due to the increased efficiency in computing derivatives still makes it worthwhile.

```python
#
# Initialize the Problem and the optimization driver
#
p = om.Problem()

p.driver = om.ScipyOptimizeDriver()
p.driver.declare_coloring()
```

Next, we add a Dymos Trajectory group to the problem's model and add a phase to it.

In this case we're using the Radau pseudospectral transcription to solve the problem.

```python
#
# Create a trajectory and add a phase to it
#
traj = p.model.add_subsystem('traj', dm.Trajectory())
tx = transcription=dm.Radau(num_segments=24)
phase = traj.add_phase('phase0', dm.Phase(ode_class=BrysonDenhamODE, transcription=tx))
```

At this point, we set the options on the main variables used in a Dymos phase.

In addition to `time`, we have three states (`x`, `v`, and `J`) and a single control (`u`).

Here we use bounds on the states themselves to constrain the initial and final value of `x` and1 `v`, and the initial value of `J`.
From an optimization perspective, this means that we are removing the first and last values in the state histories of $x$ and $v$ from the vector of design variables.
Their initial and final values will remain unchanged throughout the optimization process.

On the other hand, we could specify `fix_initial=False, fix_final=False` for these values, and Dymos would be free to change them.
We would then need to put a boundary constraint in place to enforce their final values.
Feel free to experiment with different ways of enforcing the boundary constraints on this problem and see how it affects performance.

The scaler values (`ref`) are all set to 1 here.

Bounds on time duration are guesses, and the bounds on the states and controls come from the implementation in the references.

Also, we don't need to specify targets for any of the variables here because their names _are_ the targets in the top-level of the model.
The rate source and units for the states are obtained from the tags in the ODE component we previously defined.

```python
#
# Set the variables
#
phase.set_time_options(fix_initial=True, fix_duration=True)

phase.add_state('x', fix_initial=True, fix_final=True, rate_source='v')
phase.add_state('v', fix_initial=True, fix_final=True, rate_source='u')
phase.add_state('J', fix_initial=True, fix_final=False) # Rate source obtained from tags on the ODE outputs
phase.add_control('u', continuity=True, rate_continuity=False)
```

```python
#
# Minimize time at the end of the phase
#
phase.add_objective('J', loc='final', ref=1)
phase.add_path_constraint('x', upper=1/9)

#
# Setup the Problem
#
p.setup()
```

We then set the initial guesses for the variables in the problem and solve it.

We're using the phase `interp` method to provide initial guesses for the states and controls.
In this case, by giving it two values, it is linearly interpolating from the first value to the second value, and then returning the interpolated value at the input nodes for the given variable.

Finally, we use the `dymos.run_problem` method to execute the problem.
This interface allows us to do some things that the standard OpenMDAO `problem.run_driver` interface does not.
It will automatically record the final solution achieved by the optimizer in case named `'final'` in a file called `dymos_solution.db`.
By specifying `simulate=True`, it will automatically follow the solution with an explicit integration using `scipy.solve_ivp`.
The results of the simulation are stored in a case named `final` in the file `dymos_simulation.db`.
This explicit simulation demonstrates how the system evolved with the given controls, and serves as a check that we're using a dense enough grid (enough segments and segments of sufficient order) to accurately represent the solution.

If those two solution didn't agree reasonably well, we could rerun the problem with a more dense grid.
Instead, we're asking Dymos to automatically change the grid if necessary by specifying `refine_method='ph'`.
This will attempt to repeatedly solve the problem and change the number of segments and segment orders until the solution is in reasonable agreement.

```python
#
# Set the initial values
#
p['traj.phase0.t_initial'] = 0.0
p['traj.phase0.t_duration'] = 1.0

p.set_val('traj.phase0.states:x', phase.interp('x', ys=[0, 0]))
p.set_val('traj.phase0.states:v', phase.interp('v', ys=[1, -1]))
p.set_val('traj.phase0.states:J', phase.interp('J', ys=[0, 1]))
p.set_val('traj.phase0.controls:u', np.sin(phase.interp('u', ys=[0, 0])))

#
# Solve for the optimal trajectory
#
dm.run_problem(p, run_driver=True, simulate=True)

```

## Plotting the solution

The recommended practice is to obtain values from the recorded cases.
While the problem object can also be queried for values, building plotting scripts that use the case recorder files as the data source means that the problem doesn't need to be solved just to change a plot.
Here we load values of various variables from the solution and simulation for use in the animation to follow.

```python
sol = om.CaseReader('dymos_solution.db').get_case('final')
sim = om.CaseReader('dymos_simulation.db').get_case('final')

t = sol.get_val('traj.phase0.timeseries.time')
x = sol.get_val('traj.phase0.timeseries.states:x')
v = sol.get_val('traj.phase0.timeseries.states:v')
J = sol.get_val('traj.phase0.timeseries.states:J')
u = sol.get_val('traj.phase0.timeseries.controls:u')
h = np.sin(3 * x) / 3

t_sim = sim.get_val('traj.phase0.timeseries.time')
x_sim = sim.get_val('traj.phase0.timeseries.states:x')
v_sim = sim.get_val('traj.phase0.timeseries.states:v')
J_sim = sim.get_val('traj.phase0.timeseries.states:J')
u_sim = sim.get_val('traj.phase0.timeseries.controls:u')
h_sim = np.sin(3 * x_sim) / 3

fig = plt.figure(constrained_layout=True, figsize=(12, 4))
gs = fig.add_gridspec(2, 2)

x_ax = fig.add_subplot(gs[0, 0])
v_ax = fig.add_subplot(gs[1, 0])
J_ax = fig.add_subplot(gs[0, 1])
u_ax = fig.add_subplot(gs[1, 1])

x_ax.set_ylabel('x ($m$)')
v_ax.set_ylabel('v ($m/s$)')
J_ax.set_ylabel('J ($m^2$/s^3$)')
u_ax.set_ylabel('u ($m/s^2$)')
v_ax.set_xlabel('t (s)')
u_ax.set_xlabel('t (s)')

x_sol_handle, = x_ax.plot(t, x, 'o', ms=1)
v_ax.plot(t, v, 'o', ms=1)
J_ax.plot(t, J, 'o', ms=1)
u_ax.plot(t, u, 'o', ms=1)

x_sim_handle, = x_ax.plot(t_sim, x_sim, '-', ms=1)
v_ax.plot(t_sim, v_sim, '-', ms=1)
J_ax.plot(t_sim, J_sim, '-', ms=1)
u_ax.plot(t_sim, u_sim, '-', ms=1)

for ax in [x_ax, v_ax, J_ax, u_ax]:
    ax.grid(True, alpha=0.2)

plt.figlegend([x_sol_handle, x_sim_handle], ['solution', 'simulation'], ncol=2, loc='lower center');

```

```python
from openmdao.utils.assert_utils import assert_near_equal

assert_near_equal(J[-1, 0], 4.0, tolerance=1.0E-3)
```

## Animating the Solution

The collapsed code cell below contains the code used to produce an animation of the mountain car solution using Matplotlib.

The green area represents the hilly terrain the car is traversing.  The black circle is the center of the car, and the orange arrow is the applied control.

The applied control _generally_ has the same sign as the velocity and is 'bang-bang', that is, it wants to be at its maximum possible magnitude.  Interestingly, the sign of the control flips shortly before the sign of the velocity changes.

```python
import matplotlib as mpl
from matplotlib import animation

BLOCK_SIZE = 0.02
X_EXTENTS = (-0.025, 0.175)

fig = plt.figure(constrained_layout=True, figsize=(12, 6))
gs = fig.add_gridspec(3, 2)

anim_ax = fig.add_subplot(gs[:, 0])
anim_ax.set_aspect('equal')

x_ax = fig.add_subplot(gs[0, 1:])
v_ax = fig.add_subplot(gs[1, 1:])
u_ax = fig.add_subplot(gs[2, 1:])

x_ax.set_ylabel('x')
v_ax.set_ylabel('v')
u_ax.set_ylabel('u')
u_ax.set_xlabel('t')

# set up the subplots as needed
anim_ax.set_xlim(X_EXTENTS);
y_extents = [-np.diff(X_EXTENTS) / 2, np.diff(X_EXTENTS) / 2]
anim_ax.set_ylim(y_extents);
anim_ax.set_xlabel('x');
anim_ax.set_yticks([])
time_text = anim_ax.text(0.05, 0.9, 't = 0.000 s', fontsize=16, transform=anim_ax.transAxes)

x_sol_line, = x_ax.plot(t, x, 'o', ms=1, label='solution', color='k')
v_ax.plot(t, v, 'o', ms=1, color='k')
u_ax.plot(t, u, 'o', ms=1, color='k')

x_sim_line, = x_ax.plot([], [], '-', linewidth=3, label='simulation', color='tab:blue')
v_sim_line, = v_ax.plot([], [], '-', linewidth=3, color='tab:red')
u_sim_line, = u_ax.plot([], [], '-', linewidth=3, color='tab:orange')

plt.figlegend(ncol=2, handles=[x_sol_line, x_sim_line], loc='upper center',
              bbox_to_anchor=(0.78,0.98))

x_ax.grid(alpha=0.2)
txt_x = x_ax.text(0.8, 0.1, f'x = {x_sim[0, 0]:6.3f}', horizontalalignment='left',
                  verticalalignment='center', transform=x_ax.transAxes)

v_ax.grid(alpha=0.2)
txt_v = v_ax.text(0.8, 0.1, f'v = {v_sim[0, 0]:6.3f}', horizontalalignment='left',
                  verticalalignment='center', transform=v_ax.transAxes)

u_ax.grid(alpha=0.2)
txt_u = u_ax.text(0.8, 0.1, f'u = {u_sim[0, 0]:6.3f}', horizontalalignment='left',
                  verticalalignment='center', transform=u_ax.transAxes)

x_terrain = np.linspace(-0.025, 0.175, 2);
h_terrain = 0 * x_terrain
terrain_line, = anim_ax.plot(x_terrain, h_terrain, '-', color='tab:gray', lw=2);
terrain = anim_ax.fill_between(x_terrain, h_terrain, h_terrain-1, color='tab:green');
wall = mpl.patches.Rectangle(xy=(1/9 + BLOCK_SIZE, 0), width=BLOCK_SIZE, height=3*BLOCK_SIZE, color='tab:gray')
anim_ax.add_patch(wall)
block = mpl.patches.Rectangle(xy=(0, 0), width=BLOCK_SIZE, height=BLOCK_SIZE, alpha=0.5)
anim_ax.add_patch(block)
v_vec = anim_ax.quiver(x_sim[0,0] + BLOCK_SIZE / 2, BLOCK_SIZE * 0.25, v_sim[0,0], [0], scale=10, angles='xy', color='tab:red')
u_vec = anim_ax.quiver(x_sim[0,0] + BLOCK_SIZE / 2, BLOCK_SIZE * 0.75, u_sim[0,0], [0], scale=50, angles='xy', color='tab:orange')

# See https://brushingupscience.com/2019/08/01/elaborate-matplotlib-animations/ for quiver animation

ANIM_DURATION = 3
PRE_PAUSE_DURATION = 1
POST_PAUSE_DURATION = 0.5
ANIM_FPS = 20

num_points = t_sim.size
num_frames = int(ANIM_DURATION * ANIM_FPS)
pre_pause_frames = int(PRE_PAUSE_DURATION * ANIM_FPS)
post_pause_frames = int(POST_PAUSE_DURATION * ANIM_FPS)

idx_from_frame_num = np.linspace(0, num_points-1, num_frames, dtype=int)

def drawframe(n):

    if n < pre_pause_frames:
        idx = 0
    elif n - pre_pause_frames >= idx_from_frame_num.size:
        idx = num_points - 1
    else:
        idx = idx_from_frame_num[n - pre_pause_frames]

    x = x_sim[idx]
    v = v_sim[idx]
    u = u_sim[idx]
    t = t_sim[idx]
    h = 0
    block.set_xy((x[0], h))

    v_vec.set_offsets((x[0] + BLOCK_SIZE/2, BLOCK_SIZE * 0.25))
    v_vec.set_UVC(v, 0)

    u_vec.set_offsets((x[0] + BLOCK_SIZE/2, BLOCK_SIZE * 0.75))
    u_vec.set_UVC(u, 0)

    x_sim_line.set_data(t_sim[:idx], x_sim[:idx])
    v_sim_line.set_data(t_sim[:idx], v_sim[:idx])
    u_sim_line.set_data(t_sim[:idx], u_sim[:idx])

    txt_x.set_text(f'x = {x[0]:6.3f}')
    txt_v.set_text(f'v = {v[0]:6.3f}')
    txt_u.set_text(f'u = {u[0]:6.3f}')

    time_text.set_text(f't = {t[0]:6.3f}')

    return block, u_vec, v_vec, x_sim_line, v_sim_line, u_sim_line, time_text

# # blit=True re-draws only the parts that have changed.
# # repeat_delay has no effect when using to_jshtml, so pad drawframe to show the final frame for PAUSE_FRAMES extra frames.
anim = animation.FuncAnimation(fig, drawframe, frames=pre_pause_frames + num_frames + post_pause_frames, interval=1000/ANIM_FPS, blit=True);
plt.close()  # Don't let jupyter display the un-animated plot

from IPython.display import HTML
with open('anim.html', 'w') as f:
    f.write(anim.to_jshtml())
HTML('anim.html')
```

## References

```{bibliography}
:filter: docname in docnames
```