# Integrating the moment equations - mRNA

(c) 2017 Manuel Razo. This work is licensed under a [Creative Commons Attribution License CC-BY 4.0](https://creativecommons.org/licenses/by/4.0/). All code contained herein is licensed under an [MIT license](https://opensource.org/licenses/MIT)

---

```python
import pickle
import dill # Library that we will use to export lambdify functions
# Our numerical workhorses
import numpy as np
# To compute symbolic expressions
import sympy
sympy.init_printing(use_unicode=True, use_latex=True) # print outputs in LaTeX

# Import matplotlib stuff for plotting
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib as mpl

# Seaborn, useful for graphics
import seaborn as sns

# Import the utils for this project
import chann_cap_utils as chann_cap

# Set PBoC plotting style
chann_cap.set_plotting_style()

# Magic function to make matplotlib inline; other style specs must come AFTER
%matplotlib inline

# This enables SVG graphics inline (only use with static plots (non-Bokeh))
%config InlineBackend.figure_format = 'svg'

tmpdir = '../../tmp/'
```

### $\LaTeX$ macros
$\newcommand{kpon}{k^{(p)}_{\text{on}}}$
$\newcommand{kpoff}{k^{(p)}_{\text{off}}}$
$\newcommand{kron}{k^{(r)}_{\text{on}}}$
$\newcommand{kroff}{k^{(r)}_{\text{off}}}$
$\newcommand{rm}{r _m}$
$\newcommand{gm}{\gamma _m}$
$\newcommand{mm}{\left\langle m \right\rangle}$
$\newcommand{ee}[1]{\left\langle #1 \right\rangle}$
$\newcommand{bb}[1]{\mathbf{#1}}$

# Two-state promoter (mRNA)

In this section we will try to solve the differential equations for the mRNA distribution moments for the two-state unregulated promoter.

Up to this stage we wrote the chemical master equation in matrix notation and did some tricks when it came to compute the $n^{\text{th}}$ moment. Now we will use the equations we obtained for the time derivative of the mRNA distribution moments and try to solve them.

Let's first define the `sympy` variables that we will need for the two- and three-state promoter.

Note: It is important to specify the variables we are defining are real positive nubers. Otherwise `sympy` would be unable to solve the system.

```python
# Define rate constant variables
kp_off, kp_on = sympy.symbols('{k_{off}^{(p)}} {k_{on}^{(p)}}',
                              real=True, positive=True)
kr_off, kr_on = sympy.symbols('{k_{off}^{(r)}} {k_{on}^{(r)}}',
                              real=True, positive=True)

# Define degradation rate and production rate
rm, gm = sympy.symbols('r_m gamma_m', real=True, positive=True)

# Define the time variable
t = sympy.symbols('t', real=True, positive=True)

kp_off, kp_on, kr_off, kr_on, rm, gm, t
```

Let's now define a function that generates a `sympy` variable of the form $\left\langle x^n \right\rangle$ to easily define the variables to save the moments.

```python
def moment_sympy(var, moment, function=None):
    '''
    Shortcut to generate variables of the form <var**n> in sympy to facilitate the
    definition of the moment variables
    Parameters
    ----------
    var : string.
        variable name string that should be put inside the <> parenthesis.
    moment : int.
        exponent of the moment to be defined.
    function : variable
        if given the returned variables are functions rather than symbols
    '''
    if function == None:
        if moment == None:
            return sympy.Symbol('{\\langle ' + var + '\\rangle}',
                               real=True, positive=True)
        else:
            return sympy.Symbol('{\\langle ' + var + '^{:d}'.format(moment) +\
                                '\\rangle}', real=True, positive=True)

    else:
        if moment == None:
            return sympy.Function('\\langle ' + var + '\\rangle',
                                 real=True, positive=True)(function)
        else:
            return sympy.Function('\\langle ' + var + '^{:d}'.format(moment) +\
                                '\\rangle', real=True, positive=True)(function)
```

Now let's define the moment variables for the unregulated promoter.

```python
# Define the string for the states
m_unreg = ['m_{P_{unreg}}', 'm_{E_{unreg}}']

# Define the zeroth moments
m0_unreg = [moment_sympy(x, 0, t) for x in m_unreg]
# Define the first moments
m1_unreg = [moment_sympy(x, 1, t) for x in m_unreg]
# Define the second moments
m2_unreg = [moment_sympy(x, 2, t) for x in m_unreg]
# Define the second moments
m3_unreg = [moment_sympy(x, 3, t) for x in m_unreg]

m0_unreg, m1_unreg, m2_unreg, m3_unreg
```

Let's now define the global moments.

```python
# define "global" moments
mm0_unreg = moment_sympy('\mathbf{m}_u', 0, t)
mm1_unreg = moment_sympy('\mathbf{m}_u', None, t)
mm2_unreg = moment_sympy('\mathbf{m}_u', 2, t)
mm3_unreg = moment_sympy('\mathbf{m}_u', 3, t)

mm0_unreg, mm1_unreg, mm2_unreg, mm3_unreg
```

Now we will define the state transition matrix $\mathbf{K}^u$, the mRNA production matrix $\mathbf{R}^u_m$ and the mRNA degradation matrix $\mathbf{\Gamma}^u_m$.

```python
# Define the rate constant matrix
Kmat_unreg = sympy.Matrix([[-kp_off, kp_on],
                           [kp_off, -kp_on]])
# Define the production matrix
Rmat_unreg = sympy.Matrix([[rm, 0],
                           [0, 0]])
# Define the degradation matrix
Gmat_unreg = sympy.Matrix([[gm, 0],
                           [0, gm]])

Kmat_unreg, Rmat_unreg, Gmat_unreg
```

### Zeroth moment two-state promoter
In order to solve for the two-state promoter we have two equations to use
$$
\bb{1} = \mathbf{u} \cdot \left\langle \mathbf{m}^0 \right\rangle,
$$
where $u = (1, 1)$ and $\left\langle \mathbf{m}^0 \right\rangle$ is the vector containing all the states, and
$$
{d \over dt} \ee{\bb{m}^0} = \mathbf{K}^u_m \left\langle \mathbf{m}^0 \right\rangle.
$$

We will use the linear algebra solver where we write an extended matrix with the last column being the right hand side of these equations.

Let's now define this system.

```python
dm0dt_unreg = sympy.Matrix([m.diff() for m in m0_unreg])
dm0dt_rhs = Kmat_unreg * sympy.Matrix(m0_unreg)

sys = [sympy.Eq(dm0dt_unreg[i], dm0dt_rhs[i]) for i in range(2)]
sys
```

Let's now solve the system.

```python
# Solving the system using dsolve
dm0dt_sol = sympy.dsolve(sys)

# Split left hand side and right hand side
dm0dt_rhs = [x.rhs for x in dm0dt_sol]
dm0dt_lhs = [x.lhs for x in dm0dt_sol]

dm0dt_sol
```

Now let's use the initial conditions to find the value of the constants. For the initial conditions we will simply call $\ee{m_{X_{unreg}}}(0) = \ee{m_{X_{0}}}$

```python
# Define the string for the states
m_unreg = ['m_{P_{0}}', 'm_{E_{0}}']
m0_unreg_init = [moment_sympy(x, 0) for x in m_unreg]
m0_unreg_init
```

Let's substitute these initial conditions and equate them to the right hand side of the solution evaluated at $t=0$

```python
# Building system evaluated at the inital condition
dm0dt_zero = [sympy.Eq(m0_unreg_init[x],
                        dm0dt_rhs[x].subs(t, 0)) for x in range(2)]
dm0dt_zero
```

Finally, let's solve for the constants.

```python
constants = sympy.solve(dm0dt_zero, sympy.symbols('C1 C2'))
constants
```

We can use the fact that the sum of initial conditions should add up to one.

```python
# Substitute the sum of the initial conditions
constants[sympy.Symbol('C1')] = constants[sympy.Symbol('C1')].\
                                subs(sum(m0_unreg_init), 1)
constants
```

And finally substitute these constants into the solution.

```python
# substitute constants
dm0dt_sol = [x.subs(constants) for x in dm0dt_sol]

# Split left hand side and right hand side
dm0dt_rhs = [x.rhs for x in dm0dt_sol]
dm0dt_lhs = [x.lhs for x in dm0dt_sol]
# save solution as dictionary (useful for later one substitutions)
dm0dt_dict = dict(zip(dm0dt_lhs, dm0dt_rhs))

dm0dt_sol
```

To check that this solution makes sense let's take the limit then $t \rightarrow \infty$ to see if we recover the steady state solution we derived before.

```python
# Taking the limit when time goes to infinity
[sympy.limit(x.rhs, t, 'oo') for x in dm0dt_sol]
```

---

### First moment two-state promoter

Now that we have the zeroth moment that indicates the probability of being on either state let's now compute the first moment. For this we have the matrix equation
$$
{d \over dt}\ee{\bb{m}^1} =
\left( \mathbf{K}^u_m - \mathbf{\Gamma}^u_m  \right)\ee{\bb{m}^1}
             + \mathbf{R}^u_m \left\langle \mathbf{m}^0 \right\rangle.
$$

Let's define this system.

```python
dm1dt_unreg = sympy.Matrix([m.diff() for m in m1_unreg])
dm1dt_rhs = (Kmat_unreg - Gmat_unreg)* sympy.Matrix(m1_unreg) +\
            Rmat_unreg * sympy.Matrix(m0_unreg)

sys = [sympy.Eq(dm1dt_unreg[i], dm1dt_rhs[i]) for i in range(2)]
sys
```

Let's now substitute the solution for the zeroth moment.

Note: This could have been done in a single step, but it is good to see every step when troubleshooting.

```python
sys = [x.subs(dm0dt_dict) for x in sys]
sys
```

Now let's try to solve the system

```python
# Solving the system using dsolve
dm1dt_sol = sympy.dsolve(sys)

# Split left hand side and right hand side
dm1dt_rhs = [x.rhs for x in dm1dt_sol]
dm1dt_lhs = [x.lhs for x in dm1dt_sol]

dm1dt_sol
```

Sympy cannot apparently solve this system. But we can try to define a differential equation for the "global" moment $\ee{\bb{m}^1}$ defined as the sum of the first moments for all promoter states.

```python
# Sum both differential equations and ractor terms
dmm1dt_rhs = sum(dm1dt_rhs).collect(m1_unreg[0]).\
                            collect(m1_unreg[1]).\
                            collect(gm).\
                            subs(sum(m1_unreg), mm1_unreg)

# Print differential equation for global moment
sympy.Eq(mm1_unreg.diff(), dmm1dt_rhs)
```

Let's now attempt to solve the system.

```python
# Solve equation for global first moment
dmm1dt_sol = sympy.dsolve(sympy.Eq(mm1_unreg.diff(), dmm1dt_rhs))
dmm1dt_sol
```

Since we know the solution for the zeroth moment, let's try to substitute it.

```python
dmm1dt_sol = dmm1dt_sol.subs(dm0dt_dict)
dmm1dt_sol
```

The symbolic integral must be solvable. Let's try to expand the terms to see each of the individual terms inside the integral.

```python
dmm1dt_sol = dmm1dt_sol.expand()
dmm1dt_sol
```

To perform each of these integrals we will split the left hand side of the equation into the individual components and solve each of them individually.

```python
# Split terms and convert them to list in order to modify them.
dmm1dt_sol_rhs = list(dmm1dt_sol.rhs.args)
dmm1dt_sol_rhs
```

The first integral term is easily solved since the integral is pretty simple.

```python
dmm1dt_sol_rhs[1] = dmm1dt_sol_rhs[1].doit()
dmm1dt_sol_rhs[1]
```

For the second integral we have a case in which `sympy` presents two possible solutions.

```python
integral = dmm1dt_sol_rhs[2].doit()
integral
```

`Sympy` gives one solution for the case when $\gm = \kpon + \kpoff$ and other cases. Since this is a condition that we definitely do not want to impose in our solution we need to choose the general solution.

Given that I do not find a way to tell `sympy` that the equality is not satisfied we will have to
1. Split the terms in the expression.
2. Choose the general solution.
3. Substitute back and compile again.

```python
# split integral into its components.
integral_args = list(integral.args)

# Let's now select the integral part and split it into its components
piecewise_args = integral_args[-2].args

# Now let's select the general solution and substitute it back into the
# split integral terms
integral_args[-2] = piecewise_args[1][0]

# Let's collect the terms and compute the integral
dmm1dt_sol_rhs[2] = np.prod(integral_args).simplify()
dmm1dt_sol_rhs[2]
```

Let's follow the same procedure for the last integral in the solution.

```python
integral = dmm1dt_sol_rhs[3].doit()
integral
```

```python
# split integral into its components.
integral_args = list(integral.args)

# Let's now select the integral part and split it into its components
piecewise_args = integral_args[-2].args

# Now let's select the general solution and substitute it back into the
# split integral terms
integral_args[-2] = piecewise_args[1][0]

# Let's collect the terms and compute the integral
dmm1dt_sol_rhs[3] = np.prod(integral_args).simplify()
dmm1dt_sol_rhs[3]
```

Having done this we can put the solution for ${d \over dt}\ee{\bb{m}^1}$ together.

```python
dmm1dt_sol = sympy.Eq(dmm1dt_sol.lhs, sum(dmm1dt_sol_rhs))
dmm1dt_sol
```

And just as with the zeroth moment solution we will use initial conditions of the form $\ee{\bb{m}}(0) = \ee{\bb{m}_0}$

```python
# Define the string for the states
m1_unreg_init = moment_sympy('\mathbf{m}_0', 1)
m1_unreg_init
```

Let's substitute $t=0$ in the solution.

```python
dmm1dt_zero = sympy.Eq(m1_unreg_init, dmm1dt_sol.rhs.subs(t, 0))
dmm1dt_zero
```

And solve for the constant.

```python
constant = sympy.solve(dmm1dt_zero, sympy.Symbol('C1'), dict=True)[0]
constant
```

Now we can substitute this result into the solution.

```python
# substitute constants
dmm1dt_sol = dmm1dt_sol.subs(constant)

# Split left hand side and right hand side
dmm1dt_rhs = dmm1dt_sol.rhs
dmm1dt_lhs = dmm1dt_sol.lhs

# save solution as dictionary (useful for later one substitutions)
dmm1dt_dict = {dmm1dt_rhs, dmm1dt_lhs}

dmm1dt_sol
```

To check that this solution makes sense let's take the limit then $t \rightarrow \infty$ to see if we recover the steady state solution we derived before.

```python
# Taking the limit when time goes to infinity
sympy.limit(dmm1dt_sol.rhs, t, 'oo').collect(gm)
```

---

### Second moment two-state promoter

For the second moment we have an equation of the form
$$
{d\ee{\bb{m}^2} \over dt} = \left( \mathbf{K}^u_m - 2 \mathbf{\Gamma}^u_n \right)
             \left\langle \mathbf{m}^2 \right\rangle
             + \left( 2 \mathbf{R}^u_m + \mathbf{\Gamma}^u_m \right)
             \left\langle \mathbf{m}^1 \right\rangle
             + \mathbf{R}^u_m
             \left\langle \mathbf{m}^0 \right\rangle
$$

Let's again define this system.

```python
dm2dt_unreg = sympy.Matrix([m.diff() for m in m2_unreg])
dm2dt_rhs = (Kmat_unreg - 2 * Gmat_unreg)* sympy.Matrix(m2_unreg) +\
            (2 * Rmat_unreg + Gmat_unreg) * sympy.Matrix(m1_unreg) +\
            Rmat_unreg * sympy.Matrix(m0_unreg)

sys = [sympy.Eq(dm2dt_unreg[i], dm2dt_rhs[i]) for i in range(2)]
sys
```

As with the first moment we will define a differential equation for the global moment $\ee{\bb{m}^2}$.

```python
# Sum both differential equations and ractor terms
dmm2dt_rhs = sum(dm2dt_rhs).collect(m2_unreg[0]).\
                            collect(m2_unreg[1]).\
                            collect(m1_unreg[0]).\
                            collect(m1_unreg[1]).\
                            collect(gm).\
                            subs(2 * sum(m2_unreg), 2 * mm2_unreg)

# Print differential equation for global moment
sympy.Eq(mm2_unreg.diff(), dmm2dt_rhs)
```

This indicates that without knowing the individual first moments (rather than the global moment as we solved it) we cannot get to the second moment. This is caused because of the asymmetry of the $\bb{R}_m^u$ matrix.


```python
import numpy as np
np.random.seed(100) # Set the seed for reproducibility

import matplotlib.pyplot as plt
from matplotlib import rcParams
rcParams["font.family"] = "Product Sans"
rcParams["font.serif"] = "Times New Roman"
%matplotlib inline
rcParams['text.latex.preamble'] = r'\usepackage{amsmath} \usepackage{bm} \usepackage{physics} \usepackage{type1ec}'
#rcParams['text.usetex'] = True
%config InlineBackend.figure_format = 'retina' # For high quality figures
```

# $\chi^2$ (chi-squared) minimization

## Week 4, Intro-to-Astro 2020
### Written by Joey Murphy, June 2020
### Adapted from lecture notes by Ian Czekala

## Statistics and astronomy
* Astronomers rely heavily on statistical techniques for data analysis
* *But* physics/astronomy curricula usually neglect foundational courses in statistics and probability theory
* Highly recommend taking courses in these areas if available at your school!
* List of useful references at the end

## $\chi^2$ minimization
* In the *Methods* sections of any astronomy/astrophysics paper, there is almost always some model being fit to noisy data (physical or simulated)
* After reading a couple of papers, you'll likely run into the phrase "the best-fitting model..."
* Sometimes authors report something like "the best-fitting model has $\chi^2 = \:$ xyz" or say that the *minimum* $\chi^2$ value is xyz (where $\chi$ is the Greek letter "chi")

## What is $\chi^2$ and why is it being minimized?

What does "best-fitting" actually mean? "Best" as measured how?

* First we need to talk about
    * Measurement errors (uncertainties)
    * Probability density functions
    * Likelihood

## Measurement errors (uncertainties)

* Any measurement (whether it's your height, the voltage across a resistor, or the position of a star) has inherent error
* The size and source of this error is related to how the measurement was made
    * *Question:* What's one source of error you might encounter when measuring your height?
    * *Question:* What's one source of error you might encounter when measuring the poisition of a star on the sky? (like Gaia)

## Measurement errors (uncertainties)
* Let's consider the task of measuring the length (in cm!) of a house key using a ruler
* The ruler has tick marks for milimeters and centimeters

## Measurement errors (uncertainties)
* For this example, let's assume the error on each measurement is about 0.05 cm (i.e. half of the width between the finest tick marks)
* You take a handful of measurements of the key and get a collection of values that are similar, but not exactly the same
* Why different values?
    * Maybe the key isn't always lined up at the edge of the ruler
    * Maybe you're looking at the ruler from a different angle
    * etc...

**What is the best-fitting value for the true length of the key, given our measurements?**

## Measurement errors (uncertainties)
* Let's "take" some measurements and then make a plot to visualize their scatter

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

# Generate the measurement data
mu_key    = 6.0  # cm. The *actual* length of the key (unknown to us)
sigma_key = 0.05 # cm. Uncertainty on each measurement, based on the tick marks on our ruler
n_measure = 100  # Take 100 measurements
X = np.random.normal(mu_key, sigma_key, n_measure) # Generate the "measurements"
```

```python
# Plot the data
plt.plot(np.arange(n_measure), X, '.', label='Measured value')
plt.plot(np.arange(n_measure), mu_key * np.ones(n_measure), '--', alpha=0.5, label='True value')

plt.xlabel('Measurement number', fontsize=14)
plt.ylabel('Length [cm]', fontsize=14)
plt.legend(fontsize=14, framealpha=0.3)

plt.show()
```

```python

```

Let's look back at how the "measurements" were generated in the previous plot.

Here's the line we used to create the data:

```python
X = np.random.normal(mu_key, sigma_key, n_measure) # Generate the "measurements"
```

The call to np.random.normal "draws" n_measure values randomly from a *Gaussian* (i.e. "Normal") probability distribution, centered at value mu_key and of width sigma_key.

In probability theory, this is written as

\begin{align}
X & \sim \mathcal{N}(\mu_\mathrm{key}, \sigma^2_\mathrm{key})
\end{align}

where the tilde ($\sim$) says that the data, $X$, is *drawn from* a Normal distribution ($\mathcal{N}$) with *mean* $\mu_\mathrm{key}$ and *standard deviation* $\sigma_\mathrm{key}$.

(n.b. $\sigma^2$ is referred to as the *variance*, and can sometimes be more convenient to use than $\sigma$ itself.)

Here we are *assuming* our measurements to be *Gaussian random variables*. In many cases, this is usually a reasonable assumption.

## Probability density functions
* How do we describe a probability distribution with a mathematical function?
* A *probability density function* (sometimes abbreviated as pdf) describes the relative likelihood of observing data at different locations in the domain
* For $\mathcal{N}(\mu, \sigma^2)$, the mathematical form of the probability density function is

\begin{align}
p(x | \mu, \sigma^2) & = \frac{1}{\sigma \sqrt{2 \pi}} \mathrm{exp}\Big[-\frac{1}{2} \big(\frac{x-\mu}{\sigma}\big)^2  \Big]
\end{align}

The left-hand side of the equation reads: the likelihood of observing the datum, $x$, *given* (vertical bar) mean, $\mu$, and variance $\sigma^2$. i.e. if we *know* the mean and variance of the Normal distribution that describes the data, we can calculate the relatively probability of observing $x$ using the righ-hand side of the equation.

## Probability density functions

Let's make a plot of the probability density function for the distribution that describes our measurements of the house key. On the same plot, let's show the histogram of the data we actually observered.

```python
def gaussian_pdf(x, mu=0, sigma=1):
    '''
    A gaussian pdf. Units of 1/x.

    Args
    ----------
    x (float): Location to evaluate the gaussian pdf.
    mu (float): Optional, default = 0. Mean of the gaussian.
    sigma (float): Optional, default = 1. Standard deviation.

    Returns
    ----------
    float: N(mu, sigma) evaluated at x.
    '''
    assert sigma > 0, "Standard deviation must be positive."
    return (1 / np.sqrt(2 * np.pi) / sigma) * np.exp(-0.5 * ((x - mu)/sigma)**2)
```

```python
?np.linspace
```

```python
X = np.random.normal(mu_key, sigma_key, n_measure)
x = np.linspace(5.75, 6.25, 100) #(start, stop, number)
pdf = gaussian_pdf(x, mu = mu_key, sigma = sigma_key)

fig, ax = plt.subplots()
pdf_label = '$\mu = \:$' + f'{mu_key} cm' + '\n' + '$\sigma = \:$' + f'{sigma_key} cm'
pdf_line = ax.plot(x, pdf, '--', label=pdf_label)
ax1 = ax.twinx() # Different y-axes but same x-axis for pdf plot and histogram
hist_label = 'Data'
hist = ax1.hist(X, label=hist_label, color='orange', alpha=0.6, bins=15)

ax.set_xlabel('Length [cm]', fontsize=14)
ax.set_ylabel('PDF [cm$^{-1}$]', fontsize=14)
ax1.set_ylabel('Frequency of occurence in data (N)', fontsize=14)
ax.legend(fontsize=14, loc='upper left')
ax1.legend(fontsize=14, loc='upper right')
plt.show()
```

## Likelihood

* So now that we have a probabilistic distribution, $\mathcal{N}(\mu_\mathrm{key}, \sigma^2_\mathrm{key})$, that we think describes the data, $X$, we might ask ourselves:
    * What is the *likelihood* (i.e. probability) of observing our *collection* of measurements, given  $\mu_\mathrm{key}$ and $\sigma^2_\mathrm{key}$?

If the measurements are *independent* (i.e. One measurement has no impact on the outcome of another. Think flipping a coin) then the likelihood of the data is

\begin{align}
\mathcal{L} = \Pi_i^N p(x_i | \mu_\mathrm{key}, \sigma^2_\mathrm{key})
\end{align}

where $\Pi$ means product. i.e. Multiply together all of the individual probabilities for each measurement.

## Likelihood

Often, it's convenient to work with the natural logarithm of the likelihood, since the logarithm of a product is the sum of the logarithms of each factor. The "log-likelihood" is then

\begin{align}
\mathrm{ln} \mathcal{L} & = \Sigma_i^N \mathrm{ln}\Big[ p(x_i | \mu_\mathrm{key}, \sigma^2_\mathrm{key}) \Big]
\end{align}

Plugging in for $p(x_i | \mu_\mathrm{key}, \sigma^2_\mathrm{key})$ and doing some algebra, we have

\begin{align}
\mathrm{ln} \mathcal{L} & = \Sigma_i^N \mathrm{ln}\Big[ \frac{1}{\sigma_\mathrm{key} \sqrt{2 \pi}} \mathrm{exp}\Big[-\frac{1}{2} \big(\frac{x_i-\mu_\mathrm{key}}{\sigma_\mathrm{key}}\big)^2  \Big] \Big] \\
& = \Sigma_i^N \mathrm{ln}\Big[ \frac{1}{\sigma_\mathrm{key} \sqrt{2 \pi}} \Big] + \mathrm{ln}\Big[ \mathrm{exp}\Big[-\frac{1}{2} \big(\frac{x_i-\mu_\mathrm{key}}{\sigma_\mathrm{key}}\big)^2  \Big] \Big] \\
& = \Sigma_i^N \mathrm{ln}\Big[ \frac{1}{\sigma_\mathrm{key} \sqrt{2 \pi}} \Big] + \Big[-\frac{1}{2} \big(\frac{x_i-\mu_\mathrm{key}}{\sigma_\mathrm{key}}\big)^2  \Big]
\end{align}

## Likelihood

\begin{align}
\mathrm{ln} \mathcal{L} & = \Sigma_i^N \mathrm{ln}\Big[ \frac{1}{\sigma_\mathrm{key} \sqrt{2 \pi}} \Big] + \Big[-\frac{1}{2} \big(\frac{x_i-\mu_\mathrm{key}}{\sigma_\mathrm{key}}\big)^2  \Big]
\end{align}

Notice that the first term in the sum doesn't depend on the summation index, $i$, so we can get rid of it and replace the equality with proportionality:

\begin{align}
\mathrm{ln} \mathcal{L} & \propto \Sigma_i^N \Big[-\frac{1}{2} \big(\frac{x_i-\mu_\mathrm{key}}{\sigma_\mathrm{key}}\big)^2  \Big]
\end{align}

Let's define $\chi^2 \equiv \Sigma_i^N \big(\frac{x_i - \mu_\mathrm{key}}{\sigma_\mathrm{key}}\big)^2$. Then we have

\begin{align}
\mathrm{ln} \mathcal{L} & \propto -\frac{1}{2} \chi^2
\end{align}

$\chi^2$ is related to the log-likelihood of our data!

## $\chi^2$ minimization

Finally we've found out what $\chi^2$ is
\begin{align}
\chi^2 & \equiv \Sigma_i^N \big(\frac{x_i - \mu_\mathrm{key}}{\sigma_\mathrm{key}}\big)^2
\end{align}

and how it's related to the (log) likelihood of our data

\begin{align}
\mathrm{ln} \mathcal{L} & \propto -\frac{1}{2} \chi^2
\end{align}

after we made some assumptions about the underlying probability distribution that describes our measurements.

**What is the best-fitting value for the true length of the key, given our data?**

We want to *maximize* the log-likelihood of observing our data with respect to $\mu_\mathrm{key}$ (since, after all, it's what we observed). This corresponds to *minimizing* the $\chi^2$ metric.

## $\chi^2$ minimization

How to find the extremum? Take the derivative and set it equal to zero!

\begin{align}
\frac{d \mathrm{ln}\mathcal{L}}{d \mu_\mathrm{key}} \Big|_{\hat{\mu}_\mathrm{key}} = \Sigma_i^N \frac{x_i - \hat{\mu}_\mathrm{key}}{\sigma_\mathrm{key}^2} = 0,
\end{align}

where $\hat{\mu}_\mathrm{key}$ is the most likely or best-fitting value for our estimate of the true length of the key, given our data. Rearranging, we have

\begin{align}
\hat{\mu}_\mathrm{key} & = \frac{1}{N} \Sigma_i^N x_i.
\end{align}

In other words, in this special case where $\sigma^2_\mathrm{key}$ was the same for every measurement, the best-fitting value of $\mu_\mathrm{key}$ is the average of all of our measurement values!

```python
X = np.random.normal(mu_key, sigma_key, n_measure)
# If we didn't do all of this analysis by hand, and just wanted to minimize chi^2 in code, what would that look like?
from scipy.optimize import minimize

def chi_squared(mu, sigma, data):
    '''
    Args
    ----------
    mu (float): Underlying mean of distribution (unknown to us).
    sigma (float): Measurement error.
    data (ndarray): Vector of observed data

    Returns
    ----------
    float: chi-squared value
    '''
    return np.sum(((data - mu)/sigma)**2)

mu_0 = 4.0 # cm. Our initial guess at the true length of the key, just from eye-balling it.
min_results = minimize(chi_squared, mu_0, args=(sigma_key, X)) # Minimize chi-squared with respect to mu

print(f"Best-fitting value of mu_key from chi-squared minimization = {min_results.x[0]:.3f} cm")
print(f"Average of the data = {np.average(X):.3f} cm")
```

    Best-fitting value of mu_key from chi-squared minimization = 6.005 cm
    Average of the data = 6.005 cm

```python
print(f" Answer = {minimize(chi_squared, mu_0, args=(sigma_key, X)).x[0]:.5f}")
```

     Answer = 6.00546

## $\chi^2$ minimization

* Minimizing $\chi^2$ gave us the result we were expecting from the analysis we did by hand!
* If you remember how we generated the synthetic data, we gave the key a true length of 6.0 cm, so our estimate is right on target

## Conclusion

What did we just do?

* We wanted to estimate the true length of a house key, given some (uncertain) measurements
* Acquired the data (which had **uniform measurement errors**)
* Assumed that the **underlying distribution** that described the data was a **Gaussian**
* Assumed that the **measurements were independent**
* Wrote down an analytic form for the (log) likelihood of the data observed
* Minimized the $\chi^2$ metric (i.e. maximimized the likelihood of the observed data) with respect to the parameter of interest ($\mu_\mathrm{key}$)

## External resources

(List originally compiled by Ian Czekala)

* *Data Analysis: A Bayesian Tutorial*, by Sivia and Skilling [(link)](https://www.amazon.com/Data-Analysis-Bayesian-Devinderjit-Sivia/dp/0198568320/ref=sr_1_1?ie=UTF8&qid=1500081454&sr=8-1&keywords=data+analysis+a+bayesian+tutorial). This is a very useful “handbook” to work through as an introduction to Bayesian data analysis. I highly recommend this as a first text if you can access it through your school's online library.

* Hogg et al. 2010: *Fitting a line*. [(link)](https://arxiv.org/abs/1008.4686) An applied example of Bayesian analysis for the simple example of line fitting. The simple title belies how dense this paper actually is... beware!

* Hogg et al. 2012: *Probability calculus* [(link)](https://arxiv.org/abs/1205.4446) Good notes about what you can and cannot do with probability calculus.

* *Doing Bayesian Data Analysis*, by Kruschke [(link)](https://sites.google.com/site/doingbayesiandataanalysis/). A nice introductory text to Bayesian Analysis (with puppies!)

* *Bayesian Data Analysis*, by Gelman et al. [(link)](http://www.stat.columbia.edu/~gelman/book/) The definitive reference book on Bayesian analysis and hierarchical Bayesian analysis. A dense but worthwhile book.

```python

```


# Fun with Hidden Markov Models
*by Loren Lugosch*

This notebook introduces the Hidden Markov Model (HMM), a simple model for sequential data.

We will see:
- what an HMM is and when you might want to use it;
- the so-called "three problems" of an HMM; and
- how to implement an HMM in PyTorch.

(The code in this notebook can also be found at https://github.com/lorenlugosch/pytorch_HMM.)

EDIT: The first version of this notebook normalized the transition matrix along the wrong axis. I've fixed that now.

A hypothetical scenario
------

To motivate the use of HMMs, imagine that you have a friend who gets to do a lot of travelling. Every day, this jet-setting friend sends you a selfie from the city they’re in, to make you envious.

How would you go about guessing which city the friend is in each day, just by looking at the selfies?

If the selfie contains a really obvious landmark, like the Eiffel Tower, it will be easy to figure out where the photo was taken. If not, it will be a lot harder to infer the city.

But we have a clue to help us: the city the friend is in each day is not totally random. For example, the friend will probably remain in the same city for a few days to sightsee before flying to a new city.

## The HMM setup

The hypothetical scenario of the friend travelling between cities and sending you selfies can be modeled using an HMM.

An HMM models a system that is in a particular state at any given time and produces an output that depends on that state.

At each timestep or clock tick, the system randomly decides on a new state and jumps into that state. The system then randomly generates an observation. The states are "hidden": we can't observe them. (In the cities/selfies analogy, the unknown cities would be the hidden states, and the selfies would be the observations.)

Let's denote the sequence of states as $\mathbf{z} = \{z_1, z_2, \dots, z_T \}$, where each state is one of a finite set of $N$ states, and the sequence of observations as $\mathbf{x} = \{x_1, x_2, \dots, x_T\}$. The observations could be discrete, like letters, or real-valued, like audio frames.

An HMM makes two key assumptions:
- **Assumption 1:** The state at time $t$ depends *only* on the state at the previous time $t-1$.
- **Assumption 2:** The output at time $t$ depends *only* on the state at time $t$.

These two assumptions make it possible to efficiently compute certain quantities that we may be interested in.

## Components of an HMM
An HMM has three sets of trainable parameters.

- The **transition model** is a square matrix $A$, where $A_{s, s'}$ represents $p(z_t = s|z_{t-1} = s')$, the probability of jumping from state $s'$ to state $s$.

- The **emission model** $b_s(x_t)$ tells us $p(x_t|z_t = s)$, the probability of generating $x_t$ when the system is in state $s$. For discrete observations, which we will use in this notebook, the emission model is just a lookup table, with one row for each state, and one column for each observation. For real-valued observations, it is common to use a Gaussian mixture model or neural network to implement the emission model.

- The **state priors** tell us $p(z_1 = s)$, the probability of starting in state $s$. We use $\pi$ to denote the vector of state priors, so $\pi_s$ is the state prior for state $s$.

Let's program an HMM class in PyTorch.

```
import torch
import numpy as np

class HMM(torch.nn.Module):
  """
  Hidden Markov Model with discrete observations.
  """
  def __init__(self, M, N):
    super(HMM, self).__init__()
    self.M = M # number of possible observations
    self.N = N # number of states

    # A
    self.transition_model = TransitionModel(self.N)

    # b(x_t)
    self.emission_model = EmissionModel(self.N,self.M)

    # pi
    self.unnormalized_state_priors = torch.nn.Parameter(torch.randn(self.N))

    # use the GPU
    self.is_cuda = torch.cuda.is_available()
    if self.is_cuda: self.cuda()

class TransitionModel(torch.nn.Module):
  def __init__(self, N):
    super(TransitionModel, self).__init__()
    self.N = N
    self.unnormalized_transition_matrix = torch.nn.Parameter(torch.randn(N,N))

class EmissionModel(torch.nn.Module):
  def __init__(self, N, M):
    super(EmissionModel, self).__init__()
    self.N = N
    self.M = M
    self.unnormalized_emission_matrix = torch.nn.Parameter(torch.randn(N,M))
```

To sample from the HMM, we start by picking a random initial state from the state prior distribution.

Then, we sample an output from the emission distribution, sample a transition from the transition distribution, and repeat.

(Notice that we pass the unnormalized model parameters through a softmax function to make them into probabilities.)

```
def sample(self, T=10):
  state_priors = torch.nn.functional.softmax(self.unnormalized_state_priors, dim=0)
  transition_matrix = torch.nn.functional.softmax(self.transition_model.unnormalized_transition_matrix, dim=0)
  emission_matrix = torch.nn.functional.softmax(self.emission_model.unnormalized_emission_matrix, dim=1)

  # sample initial state
  z_t = torch.distributions.categorical.Categorical(state_priors).sample().item()
  z = []; x = []
  z.append(z_t)
  for t in range(0,T):
    # sample emission
    x_t = torch.distributions.categorical.Categorical(emission_matrix[z_t]).sample().item()
    x.append(x_t)

    # sample transition
    z_t = torch.distributions.categorical.Categorical(transition_matrix[:,z_t]).sample().item()
    if t < T-1: z.append(z_t)

  return x, z

# Add the sampling method to our HMM class
HMM.sample = sample
```

Let's try hard-coding an HMM for generating fake words. (We'll also add some helper functions for encoding and decoding strings.)

We will assume that the system has one state for generating vowels and one state for generating consonants, and the transition matrix has 0s on the diagonal---in other words, the system cannot stay in the vowel state or the consonant state for one than one timestep; it has to switch.

Since we pass the transition matrix through a softmax, to get 0s we set the unnormalized parameter values to $-\infty$.

```
import string
alphabet = string.ascii_lowercase

def encode(s):
  """
  Convert a string into a list of integers
  """
  x = [alphabet.index(ss) for ss in s]
  return x

def decode(x):
  """
  Convert list of ints to string
  """
  s = "".join([alphabet[xx] for xx in x])
  return s

# Initialize the model
model = HMM(M=len(alphabet), N=2)

# Hard-wiring the parameters!
# Let state 0 = consonant, state 1 = vowel
model.unnormalized_state_priors[0] = 0.    # Let's start with a consonant more frequently
model.unnormalized_state_priors[1] = -0.5
print("State priors:", torch.nn.functional.softmax(model.unnormalized_state_priors, dim=0))

# In state 0, only allow consonants; in state 1, only allow vowels
vowel_indices = torch.tensor([alphabet.index(letter) for letter in "aeiou"])
consonant_indices = torch.tensor([alphabet.index(letter) for letter in "bcdfghjklmnpqrstvwxyz"])
model.emission_model.unnormalized_emission_matrix[0, vowel_indices] = -np.inf
model.emission_model.unnormalized_emission_matrix[1, consonant_indices] = -np.inf
print("Emission matrix:", torch.nn.functional.softmax(model.emission_model.unnormalized_emission_matrix, dim=1))

# Only allow vowel -> consonant and consonant -> vowel
model.transition_model.unnormalized_transition_matrix[0,0] = -np.inf  # consonant -> consonant
model.transition_model.unnormalized_transition_matrix[0,1] = 0.       # vowel -> consonant
model.transition_model.unnormalized_transition_matrix[1,0] = 0.       # consonant -> vowel
model.transition_model.unnormalized_transition_matrix[1,1] = -np.inf  # vowel -> vowel
print("Transition matrix:", torch.nn.functional.softmax(model.transition_model.unnormalized_transition_matrix, dim=0))

```

    State priors: tensor([0.6225, 0.3775], device='cuda:0', grad_fn=<SoftmaxBackward>)
    Emission matrix: tensor([[0.0000, 0.0896, 0.1045, 0.0080, 0.0000, 0.0454, 0.0211, 0.0725, 0.0000,
             0.0305, 0.0452, 0.0340, 0.0340, 0.0068, 0.0000, 0.2051, 0.0106, 0.0567,
             0.0259, 0.0527, 0.0000, 0.0185, 0.0388, 0.0331, 0.0641, 0.0027],
            [0.1517, 0.0000, 0.0000, 0.0000, 0.2175, 0.0000, 0.0000, 0.0000, 0.2027,
             0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.2571, 0.0000, 0.0000, 0.0000,
             0.0000, 0.0000, 0.1709, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]],
           device='cuda:0', grad_fn=<SoftmaxBackward>)
    Transition matrix: tensor([[0., 1.],
            [1., 0.]], device='cuda:0', grad_fn=<SoftmaxBackward>)

Try sampling from our hard-coded model:

```
# Sample some outputs
for _ in range(4):
  sampled_x, sampled_z = model.sample(T=5)
  print("x:", decode(sampled_x))
  print("z:", sampled_z)
```

    x: xutek
    z: [0, 1, 0, 1, 0]
    x: bipuk
    z: [0, 1, 0, 1, 0]
    x: pogut
    z: [0, 1, 0, 1, 0]
    x: coraw
    z: [0, 1, 0, 1, 0]

## The Three Problems

In a [classic tutorial](https://www.cs.cmu.edu/~cga/behavior/rabiner1.pdf) on HMMs, Lawrence Rabiner describes "three problems" that need to be solved before you can effectively use an HMM. They are:
- Problem 1: How do we efficiently compute $p(\mathbf{x})$?
- Problem 2: How do we find the most likely state sequence $\mathbf{z}$ that could have generated the data?
- Problem 3: How do we train the model?

In the rest of the notebook, we will see how to solve each problem and implement the solutions in PyTorch.

### Problem 1: How do we compute $p(\mathbf{x})$?

#### *Why?*
Why might we care about computing $p(\mathbf{x})$? Here's two reasons.
* Given two HMMs, $\theta_1$ and $\theta_2$, we can compute the likelihood of some data $\mathbf{x}$ under each model, $p_{\theta_1}(\mathbf{x})$ and $p_{\theta_2}(\mathbf{x})$, to decide which model is a better fit to the data.

  (For example, given an HMM for English speech and an HMM for French speech, we could compute the likelihood given each model, and pick the model with the higher likelihood to infer whether the person is speaking English or French.)
* Being able to compute $p(\mathbf{x})$ gives us a way to train the model, as we will see later.

#### *How?*
Given that we want $p(\mathbf{x})$, how do we compute it?

We've assumed that the data is generated by visiting some sequence of states $\mathbf{z}$ and picking an output $x_t$ for each $z_t$ from the emission distribution $p(x_t|z_t)$. So if we knew $\mathbf{z}$, then the probability of $\mathbf{x}$ could be computed as follows:

$$p(\mathbf{x}|\mathbf{z}) = \prod_{t} p(x_t|z_t) p(z_t|z_{t-1})$$

However, we don't know $\mathbf{z}$; it's hidden. But we do know the probability of any given $\mathbf{z}$, independent of what we observe. So we could get the probability of $\mathbf{x}$ by summing over the different possibilities for $\mathbf{z}$, like this:

$$p(\mathbf{x}) = \sum_{\mathbf{z}} p(\mathbf{x}|\mathbf{z}) p(\mathbf{z}) = \sum_{\mathbf{z}} \prod_{t} p(x_t|z_t) p(z_t|z_{t-1})$$

The problem is: if you try to take that sum directly, you will need to compute $N^T$ terms. This is impossible to do for anything but very short sequences. For example, let's say the sequence is of length $T=100$ and there are $N=2$ possible states. Then we would need to check $N^T = 2^{100} \approx 10^{30}$ different possible state sequences.

We need a way to compute $p(\mathbf{x})$ that doesn't require us to explicitly calculate all $N^T$ terms. For this, we use the forward algorithm.

________

The Forward Algorithm

> for $s=1 \rightarrow N$:\
>       $\alpha_{s,1} := b_s(x_1) \cdot \pi_s$
>
> for $t = 2 \rightarrow T$:\
>       for $s = 1 \rightarrow N$:\
>
> $\alpha_{s,t} := b_s(x_t) \cdot \underset{s'}{\sum} A_{s, s'} \cdot \alpha_{s',t-1} $
>
> $p(\mathbf{x}) := \underset{s}{\sum} \alpha_{s,T}$\
> return $p(\mathbf{x})$
________

The forward algorithm is much faster than enumerating all $N^T$ possible state sequences: it requires only $O(N^2T)$ operations to run, since each step is mostly multiplying the vector of forward variables by the transition matrix. (And very often we can reduce that complexity even further, if the transition matrix is sparse.)

There is one practical problem with the forward algorithm as presented above: it is prone to underflow due to multiplying a long chain of small numbers, since probabilities are always between 0 and 1. Instead, let's do everything in the log domain. In the log domain, a multiplication becomes a sum, and a sum becomes a [logsumexp](https://en.wikipedia.org/wiki/LogSumExp).

________

The Forward Algorithm (Log Domain)

> for $s=1 \rightarrow N$:\
>       $\text{log }\alpha_{s,1} := \text{log }b_s(x_1) + \text{log }\pi_s$
>
> for $t = 2 \rightarrow T$:\
>       for $s = 1 \rightarrow N$:\
>
> $\text{log }\alpha_{s,t} := \text{log }b_s(x_t) +  \underset{s'}{\text{logsumexp}} \left( \text{log }A_{s, s'} + \text{log }\alpha_{s',t-1} \right)$
>
> $\text{log }p(\mathbf{x}) := \underset{s}{\text{logsumexp}} \left( \text{log }\alpha_{s,T} \right)$\
> return $\text{log }p(\mathbf{x})$
________

Now that we have a numerically stable version of the forward algorithm, let's implement it in PyTorch.

```
def HMM_forward(self, x, T):
  """
  x : IntTensor of shape (batch size, T_max)
  T : IntTensor of shape (batch size)

  Compute log p(x) for each example in the batch.
  T = length of each example
  """
  if self.is_cuda:
  	x = x.cuda()
  	T = T.cuda()

  batch_size = x.shape[0]; T_max = x.shape[1]
  log_state_priors = torch.nn.functional.log_softmax(self.unnormalized_state_priors, dim=0)
  log_alpha = torch.zeros(batch_size, T_max, self.N)
  if self.is_cuda: log_alpha = log_alpha.cuda()

  log_alpha[:, 0, :] = self.emission_model(x[:,0]) + log_state_priors
  for t in range(1, T_max):
    log_alpha[:, t, :] = self.emission_model(x[:,t]) + self.transition_model(log_alpha[:, t-1, :])

  # Select the sum for the final timestep (each x may have different length).
  log_sums = log_alpha.logsumexp(dim=2)
  log_probs = torch.gather(log_sums, 1, T.view(-1,1) - 1)
  return log_probs

def emission_model_forward(self, x_t):
  log_emission_matrix = torch.nn.functional.log_softmax(self.unnormalized_emission_matrix, dim=1)
  out = log_emission_matrix[:, x_t].transpose(0,1)
  return out

def transition_model_forward(self, log_alpha):
  """
  log_alpha : Tensor of shape (batch size, N)
  Multiply previous timestep's alphas by transition matrix (in log domain)
  """
  log_transition_matrix = torch.nn.functional.log_softmax(self.unnormalized_transition_matrix, dim=0)

  # Matrix multiplication in the log domain
  out = log_domain_matmul(log_transition_matrix, log_alpha.transpose(0,1)).transpose(0,1)
  return out

def log_domain_matmul(log_A, log_B):
	"""
	log_A : m x n
	log_B : n x p
	output : m x p matrix

	Normally, a matrix multiplication
	computes out_{i,j} = sum_k A_{i,k} x B_{k,j}

	A log domain matrix multiplication
	computes out_{i,j} = logsumexp_k log_A_{i,k} + log_B_{k,j}
	"""
	m = log_A.shape[0]
	n = log_A.shape[1]
	p = log_B.shape[1]

	log_A_expanded = torch.stack([log_A] * p, dim=2)
	log_B_expanded = torch.stack([log_B] * m, dim=0)

	elementwise_sum = log_A_expanded + log_B_expanded
	out = torch.logsumexp(elementwise_sum, dim=1)

	return out

TransitionModel.forward = transition_model_forward
EmissionModel.forward = emission_model_forward
HMM.forward = HMM_forward
```

Try running the forward algorithm on our vowels/consonants model from before:

```
x = torch.stack( [torch.tensor(encode("cat"))] )
T = torch.tensor([3])
print(model.forward(x, T))

x = torch.stack( [torch.tensor(encode("aba")), torch.tensor(encode("abb"))] )
T = torch.tensor([3,3])
print(model.forward(x, T))
```

    tensor([[-7.5605]], device='cuda:0', grad_fn=<GatherBackward>)
    tensor([[-7.1584],
            [   -inf]], device='cuda:0', grad_fn=<GatherBackward>)

When using the vowel <-> consonant HMM from above, notice that the forward algorithm returns $-\infty$ for $\mathbf{x} = \text{"abb"}$. That's because our transition matrix says the probability of vowel -> vowel and consonant -> consonant is 0, so the probability of $\text{"abb"}$ happening is 0, and thus the log probability is $-\infty$.

#### *Side note: deriving the forward algorithm*

If you're interested in understanding how the forward algorithm actually computes $p(\mathbf{x})$, read this section; if not, skip to the next part on "Problem 2" (finding the most likely state sequence).

To derive the forward algorithm, start by deriving the forward variable:

$$\begin{align}
    \alpha_{s,t} &= p(x_1, x_2, \dots, x_t, z_t=s) \\
     &= p(x_t | x_1, x_2, \dots, x_{t-1}, z_t = s) \cdot p(x_1, x_2, \dots, x_{t-1}, z_t = s)  \\
    &= p(x_t | z_t = s) \cdot p(x_1, x_2, \dots, x_{t-1}, z_t = s) \\
    &= p(x_t | z_t = s) \cdot \left( \sum_{s'} p(x_1, x_2, \dots, x_{t-1}, z_{t-1}=s', z_t = s) \right)\\
    &= p(x_t | z_t = s) \cdot \left( \sum_{s'} p(z_t = s | x_1, x_2, \dots, x_{t-1}, z_{t-1}=s') \cdot p(x_1, x_2, \dots, x_{t-1}, z_{t-1}=s') \right)\\
    &= \underbrace{p(x_t | z_t = s)}_{\text{emission model}} \cdot \left( \sum_{s'} \underbrace{p(z_t = s | z_{t-1}=s')}_{\text{transition model}} \cdot \underbrace{p(x_1, x_2, \dots, x_{t-1}, z_{t-1}=s')}_{\text{forward variable for previous timestep}} \right)\\
    &= b_s(x_t) \cdot \left( \sum_{s'} A_{s, s'} \cdot \alpha_{s',t-1} \right)
    \end{align}$$

I'll explain how to get to each line of this equation from the previous line.

Line 1 is the definition of the forward variable $\alpha_{s,t}$.

Line 2 is the chain rule ($p(A,B) = p(A|B) \cdot p(B)$, where $A$ is $x_t$ and $B$ is all the other variables).

In Line 3, we apply Assumption 2: the probability of observation $x_t$ depends only on the current state $z_t$.

In Line 4, we marginalize over all the possible states in the previous timestep $t-1$.

In Line 5, we apply the chain rule again.

In Line 6, we apply Assumption 1: the current state depends only on the previous state.

In Line 7, we substitute in the emission probability, the transition probability, and the forward variable for the previous timestep, to get the complete recursion.

The formula above can be used for $t = 2 \rightarrow T$. At $t=1$, there is no previous state, so instead of the transition matrix $A$, we use the state priors $\pi$, which tell us the probability of starting in each state. Thus for $t=1$, the forward variables are computed as follows:

$$\begin{align}
\alpha_{s,1} &= p(x_1, z_1=s) \\
  &= p(x_1 | z_1 = s) \cdot p(z_1 = s)  \\
&= b_s(x_1) \cdot \pi_s
\end{align}$$

Finally, to compute $p(\mathbf{x}) = p(x_1, x_2, \dots, x_T)$, we marginalize over $\alpha_{s,T}$, the forward variables computed in the last timestep:

$$\begin{align*}
p(\mathbf{x}) &= \sum_{s} p(x_1, x_2, \dots, x_T, z_T = s) \\
&= \sum_{s} \alpha_{s,T}
\end{align*}$$

You can get from this formulation to the log domain formulation by taking the log of the forward variable, and using these identities:
- $\text{log }(a \cdot b) = \text{log }a + \text{log }b$
- $\text{log }(a + b) = \text{log }(e^{\text{log }a} + e^{\text{log }b}) = \text{logsumexp}(\text{log }a, \text{log }b)$

### Problem 2: How do we compute $\underset{\mathbf{z}}{\text{argmax }} p(\mathbf{z}|\mathbf{x})$?

Given an observation sequence $\mathbf{x}$, we may want to find the most likely sequence of states that could have generated $\mathbf{x}$. (Given the sequence of selfies, we want to infer what cities the friend visited.) In other words, we want $\underset{\mathbf{z}}{\text{argmax }} p(\mathbf{z}|\mathbf{x})$.

We can use Bayes' rule to rewrite this expression:
    $$\begin{align*}
    \underset{\mathbf{z}}{\text{argmax }} p(\mathbf{z}|\mathbf{x}) &= \underset{\mathbf{z}}{\text{argmax }} \frac{p(\mathbf{x}|\mathbf{z}) p(\mathbf{z})}{p(\mathbf{x})} \\
    &= \underset{\mathbf{z}}{\text{argmax }} p(\mathbf{x}|\mathbf{z}) p(\mathbf{z})
    \end{align*}$$

Hmm! That last expression, $\underset{\mathbf{z}}{\text{argmax }} p(\mathbf{x}|\mathbf{z}) p(\mathbf{z})$, looks suspiciously similar to the intractable expression we encountered before introducing the forward algorithm, $\underset{\mathbf{z}}{\sum} p(\mathbf{x}|\mathbf{z}) p(\mathbf{z})$.

And indeed, just as the intractable *sum* over all $\mathbf{z}$ can be implemented efficiently using the forward algorithm, so too this intractable *argmax* can be implemented efficiently using a similar divide-and-conquer algorithm: the legendary Viterbi algorithm!

________

The Viterbi Algorithm

> for $s=1 \rightarrow N$:\
>       $\delta_{s,1} := b_s(x_1) \cdot \pi_s$\
>       $\psi_{s,1} := 0$
>
> for $t = 2 \rightarrow T$:\
>       for $s = 1 \rightarrow N$:\
>             $\delta_{s,t} := b_s(x_t) \cdot \left( \underset{s'}{\text{max }} A_{s, s'} \cdot \delta_{s',t-1} \right)$\
            $\psi_{s,t} := \underset{s'}{\text{argmax }} A_{s, s'} \cdot \delta_{s',t-1}$
>
> $z_T^* := \underset{s}{\text{argmax }} \delta_{s,T}$\
> for $t = T-1 \rightarrow 1$:\
      $z_{t}^* := \psi_{z_{t+1}^*,t+1}$
>
> $\mathbf{z}^* := \{z_{1}^*, \dots, z_{T}^* \}$\
return $\mathbf{z}^*$
________

The Viterbi algorithm looks somewhat gnarlier than the forward algorithm, but it is essentially the same algorithm, with two tweaks: 1) instead of taking the sum over previous states, we take the max; and 2) we record the argmax of the previous states in a table, and loop back over this table at the end to get $\mathbf{z}^*$, the most likely state sequence. (And like the forward algorithm, we should run the Viterbi algorithm in the log domain for better numerical stability.)

Let's add the Viterbi algorithm to our PyTorch model:

```
def viterbi(self, x, T):
  """
  x : IntTensor of shape (batch size, T_max)
  T : IntTensor of shape (batch size)
  Find argmax_z log p(x|z) for each (x) in the batch.
  """
  if self.is_cuda:
    x = x.cuda()
    T = T.cuda()

  batch_size = x.shape[0]; T_max = x.shape[1]
  log_state_priors = torch.nn.functional.log_softmax(self.unnormalized_state_priors, dim=0)
  log_delta = torch.zeros(batch_size, T_max, self.N).float()
  psi = torch.zeros(batch_size, T_max, self.N).long()
  if self.is_cuda:
    log_delta = log_delta.cuda()
    psi = psi.cuda()

  log_delta[:, 0, :] = self.emission_model(x[:,0]) + log_state_priors
  for t in range(1, T_max):
    max_val, argmax_val = self.transition_model.maxmul(log_delta[:, t-1, :])
    log_delta[:, t, :] = self.emission_model(x[:,t]) + max_val
    psi[:, t, :] = argmax_val

  # Get the log probability of the best path
  log_max = log_delta.max(dim=2)[0]
  best_path_scores = torch.gather(log_max, 1, T.view(-1,1) - 1)

  # This next part is a bit tricky to parallelize across the batch,
  # so we will do it separately for each example.
  z_star = []
  for i in range(0, batch_size):
    z_star_i = [ log_delta[i, T[i] - 1, :].max(dim=0)[1].item() ]
    for t in range(T[i] - 1, 0, -1):
      z_t = psi[i, t, z_star_i[0]].item()
      z_star_i.insert(0, z_t)

    z_star.append(z_star_i)

  return z_star, best_path_scores # return both the best path and its log probability

def transition_model_maxmul(self, log_alpha):
  log_transition_matrix = torch.nn.functional.log_softmax(self.unnormalized_transition_matrix, dim=0)

  out1, out2 = maxmul(log_transition_matrix, log_alpha.transpose(0,1))
  return out1.transpose(0,1), out2.transpose(0,1)

def maxmul(log_A, log_B):
	"""
	log_A : m x n
	log_B : n x p
	output : m x p matrix

	Similar to the log domain matrix multiplication,
	this computes out_{i,j} = max_k log_A_{i,k} + log_B_{k,j}
	"""
	m = log_A.shape[0]
	n = log_A.shape[1]
	p = log_B.shape[1]

	log_A_expanded = torch.stack([log_A] * p, dim=2)
	log_B_expanded = torch.stack([log_B] * m, dim=0)

	elementwise_sum = log_A_expanded + log_B_expanded
	out1,out2 = torch.max(elementwise_sum, dim=1)

	return out1,out2

TransitionModel.maxmul = transition_model_maxmul
HMM.viterbi = viterbi
```

Try running Viterbi on an input sequence, given the vowel/consonant HMM:

```
x = torch.stack( [torch.tensor(encode("aba")), torch.tensor(encode("abb"))] )
T = torch.tensor([3,3])
print(model.viterbi(x, T))
```

    ([[1, 0, 1], [1, 0, 0]], tensor([[-7.1584],
            [   -inf]], device='cuda:0', grad_fn=<GatherBackward>))

For $\mathbf{x} = \text{"aba"}$, the Viterbi algorithm returns $\mathbf{z}^* = \{1,0,1\}$. This corresponds to "vowel, consonant, vowel" according to the way we defined the states above, which is correct for this input sequence. Yay!

For $\mathbf{x} = \text{"abb"}$, the Viterbi algorithm still returns a $\mathbf{z}^*$, but we know this is gibberish because "vowel, consonant, consonant" is impossible under this HMM, and indeed the log probability of this path is $-\infty$.

Let's compare the "forward score" (the log probability of all possible paths, returned by the forward algorithm) with the "Viterbi score" (the log probability of the maximum likelihood path, returned by the Viterbi algorithm):

```
print(model.forward(x, T))
print(model.viterbi(x, T)[1])
```

    tensor([[-7.1584],
            [   -inf]], device='cuda:0', grad_fn=<GatherBackward>)
    tensor([[-7.1584],
            [   -inf]], device='cuda:0', grad_fn=<GatherBackward>)

The two scores are the same! That's because in this instance there is only one possible path through the HMM, so the probability of the most likely path is the same as the sum of the probabilities of all possible paths.

In general, though, the forward score and Viterbi score will always be somewhat close. This is because of a property of the $\text{logsumexp}$ function: $\text{logsumexp}(\mathbf{x}) \approx \max (\mathbf{x})$. ($\text{logsumexp}$ is sometimes referred to as the "smooth maximum" function.)

```
x = torch.tensor([1., 2., 3.])
print(x.max(dim=0)[0])
print(x.logsumexp(dim=0))
```

    tensor(3.)
    tensor(3.4076)

### Problem 3: How do we train the model?

Earlier, we hard-coded an HMM to have certain behavior. What we would like to do instead is have the HMM learn to model the data on its own. And while it is possible to use supervised learning with an HMM (by hard-coding the emission model or the transition model) so that the states have a particular interpretation, the really cool thing about HMMs is that they are naturally unsupervised learners, so they can learn to use their different states to represent different patterns in the data, without the programmer needing to indicate what each state means.

Like many machine learning models, an HMM can be trained using maximum likelihood estimation, i.e.:

$$\theta^* = \underset{\theta}{\text{argmin }} -\sum_{\mathbf{x}^i}\text{log }p_{\theta}(\mathbf{x}^i)$$

where $\mathbf{x}^1, \mathbf{x}^2, \dots$ are training examples.

The standard method for doing this is the Expectation-Maximization (EM) algorithm, which for HMMs is also called the "Baum-Welch" algorithm. In EM training, we alternate between an "E-step", where we estimate the values of the latent variables, and an "M-step", where the model parameters are updated given the estimated latent variables. (Think $k$-means, where you guess which cluster each data point belongs to, then reestimate where the clusters are, and repeat.) The EM algorithm has some nice properties: it is guaranteed at each step to decrease the loss function, and the E-step and M-step may have an exact closed form solution, in which case no pesky learning rates are required.

But because the HMM forward algorithm is differentiable with respect to all the model parameters, we can also just take advantage of automatic differentiation methods in libraries like PyTorch and try to minimize $-\text{log }p_{\theta}(\mathbf{x})$ directly, by backpropagating through the forward algorithm and running stochastic gradient descent. That means we don't need to write any additional HMM code to implement training: `loss.backward()` is all you need.

Here we will implement SGD training for an HMM in PyTorch. First, some helper classes:

```
import torch.utils.data
from collections import Counter
from sklearn.model_selection import train_test_split

class TextDataset(torch.utils.data.Dataset):
  def __init__(self, lines):
    self.lines = lines # list of strings
    collate = Collate() # function for generating a minibatch from strings
    self.loader = torch.utils.data.DataLoader(self, batch_size=1024, num_workers=1, shuffle=True, collate_fn=collate)

  def __len__(self):
    return len(self.lines)

  def __getitem__(self, idx):
    line = self.lines[idx].lstrip(" ").rstrip("\n").rstrip(" ").rstrip("\n")
    return line

class Collate:
  def __init__(self):
    pass

  def __call__(self, batch):
    """
    Returns a minibatch of strings, padded to have the same length.
    """
    x = []
    batch_size = len(batch)
    for index in range(batch_size):
      x_ = batch[index]

      # convert letters to integers
      x.append(encode(x_))

    # pad all sequences with 0 to have same length
    x_lengths = [len(x_) for x_ in x]
    T = max(x_lengths)
    for index in range(batch_size):
      x[index] += [0] * (T - len(x[index]))
      x[index] = torch.tensor(x[index])

    # stack into single tensor
    x = torch.stack(x)
    x_lengths = torch.tensor(x_lengths)
    return (x,x_lengths)
```

Let's load some training/testing data. By default, this will use the unix "words" file, but you could also use your own text file.

```
!wget https://raw.githubusercontent.com/lorenlugosch/pytorch_HMM/master/data/train/training.txt

filename = "training.txt"

with open(filename, "r") as f:
  lines = f.readlines() # each line of lines will have one word

alphabet = list(Counter(("".join(lines))).keys())
train_lines, valid_lines = train_test_split(lines, test_size=0.1, random_state=42)
train_dataset = TextDataset(train_lines)
valid_dataset = TextDataset(valid_lines)

M = len(alphabet)
```

    --2020-02-17 15:17:50--  https://raw.githubusercontent.com/lorenlugosch/pytorch_HMM/master/data/train/training.txt
    Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 151.101.0.133, 151.101.64.133, 151.101.128.133, ...
    Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|151.101.0.133|:443... connected.
    HTTP request sent, awaiting response... 200 OK
    Length: 2493109 (2.4M) [text/plain]
    Saving to: ‘training.txt’

    training.txt        100%[===================>]   2.38M  --.-KB/s    in 0.1s

    2020-02-17 15:17:50 (24.5 MB/s) - ‘training.txt’ saved [2493109/2493109]

We will use a Trainer class for training and testing the model:

```
from tqdm import tqdm # for displaying progress bar

class Trainer:
  def __init__(self, model, lr):
    self.model = model
    self.lr = lr
    self.optimizer = torch.optim.Adam(model.parameters(), lr=self.lr, weight_decay=0.00001)

  def train(self, dataset):
    train_loss = 0
    num_samples = 0
    self.model.train()
    print_interval = 50
    for idx, batch in enumerate(tqdm(dataset.loader)):
      x,T = batch
      batch_size = len(x)
      num_samples += batch_size
      log_probs = self.model(x,T)
      loss = -log_probs.mean()
      self.optimizer.zero_grad()
      loss.backward()
      self.optimizer.step()
      train_loss += loss.cpu().data.numpy().item() * batch_size
      if idx % print_interval == 0:
        print("loss:", loss.item())
        for _ in range(5):
          sampled_x, sampled_z = self.model.sample()
          print(decode(sampled_x))
          print(sampled_z)
    train_loss /= num_samples
    return train_loss

  def test(self, dataset):
    test_loss = 0
    num_samples = 0
    self.model.eval()
    print_interval = 50
    for idx, batch in enumerate(dataset.loader):
      x,T = batch
      batch_size = len(x)
      num_samples += batch_size
      log_probs = self.model(x,T)
      loss = -log_probs.mean()
      test_loss += loss.cpu().data.numpy().item() * batch_size
      if idx % print_interval == 0:
        print("loss:", loss.item())
        sampled_x, sampled_z = self.model.sample()
        print(decode(sampled_x))
        print(sampled_z)
    test_loss /= num_samples
    return test_loss
```

Finally, initialize the model and run the main training loop. Every 50 batches, the code will produce a few samples from the model. Over time, these samples should look more and more realistic.

```
# Initialize model
model = HMM(N=64, M=M)

# Train the model
num_epochs = 10
trainer = Trainer(model, lr=0.01)

for epoch in range(num_epochs):
        print("========= Epoch %d of %d =========" % (epoch+1, num_epochs))
        train_loss = trainer.train(train_dataset)
        valid_loss = trainer.test(valid_dataset)

        print("========= Results: epoch %d of %d =========" % (epoch+1, num_epochs))
        print("train loss: %.2f| valid loss: %.2f\n" % (train_loss, valid_loss) )
```

      0%|          | 0/208 [00:00<?, ?it/s]

    ========= Epoch 1 of 10 =========

      0%|          | 1/208 [00:00<01:21,  2.53it/s]

    loss: 37.83818817138672
    CnJDVadnIj
    [18, 25, 11, 54, 20, 3, 10, 49, 0, 2]

    RobzPWmmX
    [55, 43, 56, 61, 38, 54, 21, 21, 31, 63]
    wnWnwDCoBM
    [14, 33, 58, 63, 21, 23, 17, 32, 43, 63]
    hLcxLtbumv
    [58, 53, 61, 34, 1, 48, 48, 62, 15, 58]
    wM-fUt-tLr
    [27, 5, 14, 45, 43, 27, 51, 51, 54, 24]

     25%|██▍       | 51/208 [00:13<00:43,  3.57it/s]

    loss: 32.814849853515625
    utrsceValB
    [45, 5, 14, 22, 56, 5, 45, 42, 44, 50]
    gbUnHnDDYM
    [46, 8, 54, 56, 44, 29, 0, 26, 33, 46]
    svgrSPKyHD
    [26, 19, 37, 50, 50, 56, 44, 19, 52, 63]
    PSAIyXlCHu
    [15, 39, 35, 20, 29, 13, 37, 17, 26, 46]
    IjiHtZ-nio
    [7, 37, 19, 6, 54, 12, 44, 27, 3, 46]

     49%|████▊     | 101/208 [00:27<00:29,  3.63it/s]

    loss: 29.635101318359375
    IaryUenEVo
    [18, 22, 27, 3, 27, 46, 14, 6, 35, 46]
    UiePuZneiI
    [16, 42, 21, 3, 7, 58, 3, 29, 4, 27]
    stliozerda
    [22, 16, 5, 4, 23, 53, 63, 44, 48, 2]
    pzpurSsloo
    [26, 34, 12, 29, 15, 22, 6, 21, 3, 12]
    biboakqVhl
    [7, 5, 55, 2, 6, 35, 46, 26, 33, 21]

     73%|███████▎  | 151/208 [00:40<00:14,  3.87it/s]

    loss: 28.170347213745117
    puHpurjlos
    [32, 46, 15, 35, 46, 15, 63, 27, 34, 15]
    pibpaitrei
    [62, 44, 8, 44, 21, 44, 6, 35, 46, 44]
    bensPeaeos
    [48, 46, 15, 22, 35, 38, 53, 35, 46, 15]
    lirrtaqoci
    [62, 34, 46, 15, 22, 53, 22, 2, 14, 63]
    purlayIlXd
    [48, 46, 15, 22, 27, 14, 22, 21, 4, 10]

     97%|█████████▋| 201/208 [00:53<00:01,  3.59it/s]

    loss: 26.471345901489258
    geresodePc
    [31, 63, 40, 63, 15, 3, 31, 63, 27, 14]
    sertVmelpn
    [48, 46, 15, 26, 27, 35, 46, 15, 26, 38]
    qabdellyni
    [48, 34, 24, 10, 21, 44, 21, 4, 27, 42]
    untenomcmy
    [62, 27, 16, 46, 27, 46, 15, 57, 54, 4]
    Pubndytang
    [48, 46, 26, 27, 31, 4, 61, 46, 27, 18]

    100%|██████████| 208/208 [00:55<00:00,  4.68it/s]

    loss: 26.017606735229492
    inzilanuru
    [63, 27, 61, 63, 2, 34, 27, 46, 15, 16]

      0%|          | 0/208 [00:00<?, ?it/s]

    ========= Results: epoch 1 of 10 =========
    train loss: 30.53| valid loss: 26.40

    ========= Epoch 2 of 10 =========

      0%|          | 1/208 [00:00<01:14,  2.77it/s]

    loss: 26.479572296142578
    lplincirmn
    [22, 26, 33, 34, 29, 5, 4, 15, 4, 27]
    unauatorbu
    [49, 18, 42, 14, 42, 29, 3, 15, 54, 53]
    tecu-elibk
    [26, 46, 41, 46, 53, 22, 2, 34, 56, 44]
    thgUrencgo
    [26, 33, 21, 46, 15, 63, 27, 16, 51, 34]
    rtafltoBee
    [15, 14, 34, 45, 44, 29, 46, 44, 21, 4]

     25%|██▍       | 51/208 [00:13<00:41,  3.76it/s]

    loss: 25.50018310546875
    anEitorisZ
    [49, 27, 31, 63, 35, 46, 44, 4, 15, 63]
    ropoberpsi
    [48, 46, 15, 3, 31, 63, 15, 26, 15, 4]
    uionwemidn
    [49, 25, 34, 27, 31, 3, 31, 63, 10, 63]
    alcireytlt
    [46, 15, 35, 46, 44, 21, 4, 22, 5, 14]
    beweleones
    [62, 63, 31, 63, 44, 21, 3, 31, 63, 15]

     49%|████▊     | 101/208 [00:27<00:31,  3.40it/s]

    loss: 24.771900177001953
    SYbrinibby
    [48, 34, 45, 44, 42, 31, 46, 45, 44, 19]
    niniswetus
    [31, 63, 27, 46, 15, 35, 63, 27, 46, 15]
    soborsliig
    [62, 34, 38, 46, 44, 53, 5, 4, 34, 18]
    chiceuspuo
    [48, 33, 14, 61, 46, 46, 15, 48, 46, 34]
    sangmaropc
    [2, 34, 27, 58, 7, 46, 44, 46, 15, 14]

     73%|███████▎  | 151/208 [00:40<00:15,  3.62it/s]

    loss: 24.587112426757812
    unigJzerpa
    [49, 27, 63, 26, 33, 61, 63, 15, 26, 34]
    uasmdencBn
    [62, 46, 15, 50, 10, 46, 15, 2, 34, 27]
    suststyotu
    [48, 53, 22, 35, 22, 5, 4, 5, 26, 46]
    crowhiolap
    [26, 33, 3, 26, 33, 2, 34, 29, 6, 35]
    Qiresononi
    [48, 46, 15, 26, 15, 3, 29, 34, 27, 63]

     97%|█████████▋| 201/208 [00:54<00:01,  3.61it/s]

    loss: 24.65005111694336
    peptismigy
    [26, 63, 15, 35, 14, 22, 7, 42, 29, 4]
    proekicaki
    [26, 33, 46, 63, 15, 4, 61, 6, 44, 42]
    chRngIveif
    [26, 33, 42, 27, 58, 0, 23, 21, 4, 5]
    niusphideg
    [31, 63, 53, 22, 35, 33, 63, 31, 63, 15]
    fapeleauca
    [48, 34, 26, 46, 44, 21, 3, 53, 61, 46]

    100%|██████████| 208/208 [00:55<00:00,  4.52it/s]

    loss: 24.449634552001953
    Costaparra
    [48, 46, 15, 35, 46, 24, 26, 33, 12, 34]

      0%|          | 0/208 [00:00<?, ?it/s]

    ========= Results: epoch 2 of 10 =========
    train loss: 25.07| valid loss: 24.47

    ========= Epoch 3 of 10 =========

      0%|          | 1/208 [00:00<01:13,  2.80it/s]

    loss: 24.374300003051758
    prulesange
    [48, 33, 46, 44, 46, 15, 34, 27, 58, 63]
    jomucrlasi
    [48, 46, 31, 46, 15, 33, 2, 46, 15, 14]
    soryithode
    [48, 46, 15, 4, 22, 35, 33, 3, 31, 63]
    unceyiphiz
    [49, 27, 26, 46, 17, 42, 26, 33, 42, 61]
    tissanarip
    [26, 46, 15, 22, 6, 35, 46, 15, 42, 35]

     25%|██▍       | 51/208 [00:13<00:47,  3.31it/s]

    loss: 24.072898864746094
    wroneneste
    [26, 33, 3, 35, 63, 27, 63, 15, 35, 42]
    caovadiarc
    [48, 46, 3, 31, 42, 31, 14, 34, 15, 26]
    GWlerlgrca
    [48, 46, 44, 63, 15, 31, 63, 15, 61, 6]
    changokers
    [26, 33, 34, 27, 58, 34, 36, 63, 15, 35]
    Lrsophiali
    [48, 15, 26, 34, 26, 33, 2, 6, 44, 42]

     49%|████▊     | 101/208 [00:27<00:29,  3.64it/s]

    loss: 24.010425567626953
    slyeezarvi
    [22, 21, 19, 21, 46, 61, 6, 44, 31, 42]
    lediFlefli
    [48, 46, 29, 14, 6, 44, 63, 10, 21, 4]
    suctenphos
    [48, 46, 15, 35, 63, 27, 26, 33, 46, 15]
    valilrobbu
    [31, 46, 44, 4, 61, 33, 46, 24, 51, 46]
    esadicsest
    [46, 15, 42, 31, 14, 61, 35, 63, 15, 35]

     73%|███████▎  | 151/208 [00:40<00:15,  3.58it/s]

    loss: 24.108936309814453
    mRerphammi
    [7, 32, 63, 15, 26, 33, 46, 7, 7, 4]
    moeteinert
    [48, 46, 63, 15, 14, 34, 27, 63, 15, 35]
    uniztinalQ
    [62, 12, 42, 61, 35, 42, 29, 6, 44, 21]
    Dralhysaso
    [62, 12, 46, 15, 33, 19, 35, 6, 35, 46]
    chissicabs
    [26, 33, 46, 15, 35, 42, 61, 46, 24, 14]

     97%|█████████▋| 201/208 [00:54<00:01,  3.55it/s]

    loss: 23.95977783203125
    hcheddemiY
    [33, 26, 33, 63, 10, 31, 63, 31, 42, 61]
    disyscXomi
    [48, 46, 15, 19, 15, 26, 33, 3, 7, 42]
    cangitesmi
    [48, 34, 27, 35, 42, 35, 63, 22, 7, 42]
    Gueleronin
    [48, 46, 63, 44, 63, 15, 3, 29, 42, 29]
    isulwisaza
    [46, 15, 46, 44, 21, 4, 35, 42, 61, 6]

    100%|██████████| 208/208 [00:56<00:00,  3.89it/s]

    loss: 23.793458938598633
    mestlasDam
    [31, 63, 15, 35, 33, 46, 15, 28, 6, 35]

      0%|          | 0/208 [00:00<?, ?it/s]

    ========= Results: epoch 3 of 10 =========
    train loss: 24.17| valid loss: 23.99

    ========= Epoch 4 of 10 =========

      0%|          | 1/208 [00:00<01:17,  2.66it/s]

    loss: 24.02060317993164
    corypulsyt
    [48, 46, 15, 19, 32, 63, 44, 22, 17, 26]
    varectauyp
    [12, 46, 15, 63, 15, 35, 34, 27, 19, 32]
    loinusiuss
    [48, 46, 49, 27, 46, 15, 22, 53, 22, 35]
    malershiag
    [48, 46, 44, 63, 15, 22, 35, 14, 34, 26]
    haraeiglyl
    [48, 46, 15, 6, 63, 3, 58, 5, 19, 44]

     25%|██▍       | 51/208 [00:13<00:42,  3.69it/s]

    loss: 24.00276756286621
    rerUertivu
    [48, 46, 15, 26, 46, 15, 35, 42, 31, 6]
    tagstertio
    [26, 3, 58, 22, 35, 63, 15, 35, 14, 34]
    lwauedirsl
    [48, 33, 3, 53, 63, 10, 46, 15, 35, 12]
    granilomal
    [62, 12, 46, 27, 42, 44, 3, 7, 6, 44]
    heroresysy
    [33, 63, 15, 46, 15, 63, 15, 4, 22, 17]

     49%|████▊     | 101/208 [00:27<00:30,  3.52it/s]

    loss: 23.493682861328125
    gisegoacia
    [18, 46, 15, 3, 58, 12, 46, 15, 14, 34]
    Petiudnton
    [48, 46, 15, 3, 53, 10, 27, 35, 34, 27]
    amotlylyyh
    [3, 7, 46, 26, 33, 4, 5, 19, 32, 33]
    dettedisti
    [48, 46, 15, 35, 63, 10, 14, 22, 35, 42]
    hlyctumfma
    [33, 21, 4, 61, 35, 46, 7, 39, 38, 6]

     73%|███████▎  | 151/208 [00:40<00:15,  3.73it/s]

    loss: 23.660245895385742
    vigelionho
    [31, 4, 29, 63, 44, 4, 34, 27, 33, 46]
    smptuspoma
    [22, 7, 32, 35, 46, 15, 26, 46, 7, 6]
    chelthitre
    [26, 33, 46, 44, 26, 33, 42, 35, 31, 63]
    derbomping
    [31, 63, 15, 8, 46, 7, 32, 42, 29, 58]
    Tidaticmme
    [48, 42, 61, 6, 35, 46, 15, 7, 7, 63]

     97%|█████████▋| 201/208 [00:54<00:01,  3.72it/s]

    loss: 23.784198760986328
    hemeronede
    [48, 46, 7, 63, 15, 3, 31, 46, 31, 63]
    babBicasal
    [48, 6, 45, 44, 4, 61, 6, 35, 46, 44]
    maibvonoel
    [48, 46, 46, 24, 31, 46, 15, 46, 46, 44]
    scaanurtip
    [22, 26, 46, 49, 27, 46, 15, 26, 46, 26]
    andertrysh
    [34, 27, 31, 63, 15, 26, 33, 46, 22, 33]

    100%|██████████| 208/208 [00:55<00:00,  4.53it/s]

    loss: 23.924264907836914
    plabontowe
    [62, 12, 3, 31, 34, 27, 26, 46, 7, 63]

      0%|          | 0/208 [00:00<?, ?it/s]

    ========= Results: epoch 4 of 10 =========
    train loss: 23.82| valid loss: 23.75

    ========= Epoch 5 of 10 =========

      0%|          | 1/208 [00:00<01:15,  2.75it/s]

    loss: 23.799543380737305
    ronyglydob
    [48, 34, 27, 47, 18, 21, 19, 10, 46, 24]
    pocaristom
    [48, 46, 26, 46, 15, 14, 22, 35, 46, 7]
    geltesylec
    [48, 46, 44, 26, 46, 15, 19, 44, 63, 15]
    epshypedin
    [46, 32, 22, 33, 19, 32, 63, 10, 17, 27]
    otlelisynt
    [3, 26, 12, 46, 44, 4, 22, 17, 27, 26]

     25%|██▍       | 51/208 [00:13<00:40,  3.84it/s]

    loss: 23.391658782958984
    utionerisy
    [49, 15, 14, 34, 29, 63, 15, 4, 22, 17]
    Malymotado
    [48, 34, 44, 19, 7, 46, 26, 46, 31, 3]
    qigbsmeste
    [48, 42, 41, 24, 22, 7, 46, 22, 35, 63]
    Thorsetian
    [48, 33, 46, 27, 22, 46, 35, 14, 34, 29]
    odobointia
    [3, 31, 46, 24, 46, 49, 27, 35, 14, 34]

     49%|████▊     | 101/208 [00:27<00:28,  3.80it/s]

    loss: 23.864551544189453
    gindusalia
    [48, 49, 27, 10, 46, 15, 46, 44, 4, 6]
    yolampewoi
    [8, 46, 44, 46, 7, 32, 63, 15, 3, 42]
    moncrautib
    [48, 34, 27, 26, 12, 46, 46, 35, 14, 45]
    dachontode
    [48, 46, 41, 33, 46, 27, 26, 46, 31, 63]
    nomQotacri
    [48, 46, 7, 11, 3, 26, 34, 26, 33, 42]

     73%|███████▎  | 151/208 [00:40<00:15,  3.70it/s]

    loss: 23.713638305664062
    kraphreesd
    [62, 12, 46, 32, 33, 12, 46, 63, 15, 31]
    dekication
    [31, 63, 31, 42, 61, 6, 35, 14, 34, 27]
    eslytychyr
    [46, 22, 21, 19, 5, 19, 26, 33, 19, 44]
    lymerselew
    [48, 46, 31, 63, 15, 22, 63, 44, 63, 15]
    tedmeredem
    [26, 46, 10, 31, 63, 15, 63, 10, 46, 7]

     97%|█████████▋| 201/208 [00:53<00:01,  3.72it/s]

    loss: 23.801746368408203
    cybyanguud
    [48, 46, 24, 4, 34, 27, 18, 49, 27, 10]
    semicatera
    [48, 46, 7, 42, 61, 6, 35, 63, 15, 3]
    ompilatian
    [3, 7, 32, 42, 44, 6, 35, 14, 34, 29]
    apholticka
    [46, 32, 33, 34, 27, 35, 42, 61, 36, 6]
    mingyntert
    [25, 42, 29, 58, 17, 27, 35, 63, 15, 35]

    100%|██████████| 208/208 [00:55<00:00,  4.46it/s]

    loss: 23.500385284423828
    GensexJuls
    [48, 46, 27, 35, 63, 15, 30, 53, 15, 22]

      0%|          | 0/208 [00:00<?, ?it/s]

    ========= Results: epoch 5 of 10 =========
    train loss: 23.64| valid loss: 23.61

    ========= Epoch 6 of 10 =========

      0%|          | 1/208 [00:00<01:30,  2.28it/s]

    loss: 23.998214721679688
    merypiston
    [48, 46, 15, 19, 32, 63, 15, 8, 34, 29]
    kequitcong
    [48, 46, 30, 53, 42, 35, 16, 34, 27, 58]
    untraryote
    [49, 27, 26, 12, 46, 15, 19, 34, 29, 63]
    dephylumos
    [48, 46, 32, 33, 19, 44, 50, 7, 34, 15]
    abralalicr
    [46, 24, 2, 46, 44, 46, 44, 4, 61, 12]

     25%|██▍       | 51/208 [00:14<00:42,  3.71it/s]

    loss: 23.56419563293457
    strablbieo
    [22, 35, 12, 6, 45, 44, 25, 42, 63, 3]
    sobdicaies
    [48, 46, 24, 10, 14, 61, 46, 42, 63, 15]
    pabgistati
    [48, 46, 24, 58, 42, 22, 35, 6, 35, 42]
    rurchfromo
    [12, 46, 15, 26, 33, 62, 12, 46, 7, 34]
    dellesityp
    [31, 46, 44, 21, 63, 15, 4, 5, 19, 32]

     49%|████▊     | 101/208 [00:27<00:28,  3.70it/s]

    loss: 23.507736206054688
    Iryplalodi
    [62, 12, 19, 32, 44, 46, 44, 3, 31, 42]
    syismembit
    [48, 19, 4, 22, 7, 46, 7, 24, 42, 26]
    mockermmgl
    [48, 46, 15, 57, 63, 15, 38, 25, 58, 12]
    untarislnc
    [49, 27, 35, 46, 12, 42, 22, 17, 27, 26]
    Guchiperri
    [48, 46, 26, 33, 42, 32, 63, 15, 12, 42]

     73%|███████▎  | 151/208 [00:41<00:15,  3.74it/s]

    loss: 23.453380584716797
    gnatonovia
    [41, 38, 6, 35, 34, 27, 3, 31, 42, 6]
    griarcryos
    [62, 12, 42, 6, 15, 26, 12, 19, 3, 22]
    zariogrurs
    [48, 46, 15, 42, 3, 58, 12, 46, 15, 22]
    Fontaterad
    [48, 34, 27, 35, 6, 35, 63, 15, 46, 31]
    Polocksnol
    [48, 46, 44, 46, 61, 57, 22, 29, 6, 44]

     97%|█████████▋| 201/208 [00:54<00:01,  3.72it/s]

    loss: 23.498559951782227
    etecymably
    [13, 26, 46, 15, 19, 7, 6, 45, 21, 19]
    stesemanop
    [22, 35, 63, 15, 63, 1, 34, 29, 3, 32]
    prulyssati
    [62, 12, 46, 44, 19, 22, 35, 6, 35, 14]
    flatsoarin
    [18, 21, 6, 15, 22, 46, 46, 15, 14, 31]
    rommamedam
    [48, 46, 7, 37, 49, 7, 46, 31, 46, 7]

    100%|██████████| 208/208 [00:56<00:00,  4.32it/s]

    loss: 23.944332122802734
    sanginaciu
    [48, 34, 29, 58, 42, 29, 6, 16, 4, 50]

      0%|          | 0/208 [00:00<?, ?it/s]

    ========= Results: epoch 6 of 10 =========
    train loss: 23.52| valid loss: 23.50

    ========= Epoch 7 of 10 =========

      0%|          | 1/208 [00:00<01:15,  2.74it/s]

    loss: 23.66443634033203
    aftelysheb
    [13, 55, 35, 63, 21, 19, 22, 33, 63, 62]
    editepadac
    [49, 10, 42, 35, 46, 32, 46, 10, 3, 26]
    omonsucide
    [3, 11, 34, 27, 22, 46, 61, 14, 31, 63]
    ocKogytabu
    [46, 15, 4, 3, 58, 19, 26, 3, 24, 53]
    orsmpetedy
    [46, 15, 22, 7, 32, 63, 35, 63, 10, 19]

     25%|██▍       | 51/208 [00:14<00:44,  3.49it/s]

    loss: 23.37617301940918
    eerpombavo
    [46, 46, 15, 26, 46, 7, 51, 46, 31, 34]
    bonconatio
    [48, 34, 29, 16, 34, 29, 6, 35, 14, 34]
    gonctisedi
    [48, 34, 27, 16, 35, 14, 22, 63, 10, 4]
    tanroperot
    [26, 46, 27, 12, 3, 32, 63, 15, 3, 26]
    hapterdert
    [33, 46, 32, 35, 63, 44, 31, 63, 15, 35]

     49%|████▊     | 101/208 [00:27<00:28,  3.74it/s]

    loss: 23.183443069458008
    erbbsonine
    [46, 15, 8, 24, 2, 34, 27, 42, 29, 63]
    vacterarir
    [48, 46, 15, 35, 63, 12, 46, 12, 42, 44]
    hocrenanae
    [33, 46, 26, 12, 63, 31, 6, 27, 46, 3]
    unendesizk
    [49, 27, 46, 27, 31, 63, 15, 14, 61, 57]
    untaticens
    [49, 27, 35, 6, 35, 14, 61, 63, 15, 22]

     73%|███████▎  | 151/208 [00:41<00:18,  3.04it/s]

    loss: 23.409622192382812
    lapitaders
    [48, 46, 32, 42, 26, 46, 31, 63, 15, 38]
    arinercras
    [49, 15, 4, 35, 63, 15, 26, 12, 46, 15]
    rulitytent
    [12, 46, 44, 4, 5, 19, 35, 63, 27, 35]
    sobrigaero
    [48, 46, 24, 12, 42, 41, 6, 63, 15, 3]
    fpavelogon
    [13, 26, 46, 31, 63, 44, 3, 58, 34, 29]

     97%|█████████▋| 201/208 [00:54<00:01,  3.72it/s]

    loss: 23.292179107666016
    pruthededi
    [62, 12, 46, 26, 33, 63, 10, 3, 31, 4]
    caedoplycn
    [26, 46, 27, 10, 3, 32, 21, 19, 26, 12]
    cringiatio
    [26, 12, 42, 29, 58, 4, 6, 35, 14, 34]
    increxshta
    [49, 27, 26, 12, 63, 15, 22, 33, 14, 34]
    unegigneri
    [49, 27, 3, 58, 42, 41, 29, 63, 15, 14]

    100%|██████████| 208/208 [00:56<00:00,  4.70it/s]

    loss: 23.704559326171875
    sninbadomi
    [22, 35, 14, 29, 9, 46, 10, 3, 38, 4]

      0%|          | 0/208 [00:00<?, ?it/s]

    ========= Results: epoch 7 of 10 =========
    train loss: 23.41| valid loss: 23.40

    ========= Epoch 8 of 10 =========

      0%|          | 1/208 [00:00<01:24,  2.45it/s]

    loss: 23.13858985900879
    hasstacare
    [48, 46, 15, 22, 26, 46, 26, 46, 15, 63]
    qusionblyl
    [30, 53, 22, 14, 34, 29, 45, 21, 19, 44]
    loglymifis
    [48, 46, 58, 21, 19, 7, 4, 39, 42, 22]
    ungexserme
    [49, 27, 58, 63, 15, 22, 63, 15, 38, 46]
    necsetrecr
    [31, 63, 15, 22, 46, 26, 12, 3, 26, 12]

     25%|██▍       | 51/208 [00:13<00:41,  3.78it/s]

    loss: 23.059940338134766
    cercitylyc
    [48, 46, 44, 61, 4, 5, 19, 44, 19, 26]
    Mocracrala
    [48, 46, 26, 12, 46, 26, 12, 6, 44, 6]
    gasroectiv
    [48, 46, 15, 38, 46, 46, 61, 35, 14, 31]
    trumflerun
    [26, 12, 46, 7, 39, 21, 63, 12, 46, 27]
    panosickex
    [48, 34, 29, 3, 38, 4, 61, 57, 63, 15]

     49%|████▊     | 101/208 [00:27<00:29,  3.59it/s]

    loss: 23.657520294189453
    rHaysminan
    [12, 3, 52, 19, 22, 7, 42, 29, 6, 27]
    trolikenme
    [26, 12, 46, 44, 4, 36, 63, 27, 51, 46]
    dicroperte
    [31, 4, 26, 12, 3, 32, 63, 15, 35, 63]
    welizerrer
    [48, 46, 44, 4, 36, 63, 15, 12, 63, 15]
    matioblesw
    [25, 6, 35, 14, 34, 45, 21, 63, 15, 60]

     73%|███████▎  | 151/208 [00:41<00:14,  3.81it/s]

    loss: 23.064796447753906
    wighmeponz
    [48, 46, 41, 33, 38, 46, 26, 34, 27, 43]
    phipereTmi
    [32, 33, 42, 32, 63, 12, 63, 15, 7, 46]
    umposscers
    [49, 7, 32, 46, 15, 22, 35, 63, 15, 38]
    cronlishla
    [26, 12, 34, 29, 44, 4, 22, 33, 12, 46]
    Aserictich
    [13, 22, 63, 15, 4, 61, 35, 14, 61, 33]

     97%|█████████▋| 201/208 [00:54<00:02,  3.48it/s]

    loss: 23.43466567993164
    epoftienes
    [49, 32, 46, 55, 35, 14, 34, 29, 63, 15]
    untartrivi
    [49, 27, 26, 46, 15, 26, 12, 14, 31, 4]
    stosairido
    [22, 26, 3, 22, 46, 46, 15, 4, 29, 3]
    Palydyluti
    [48, 46, 44, 19, 10, 19, 44, 53, 35, 14]
    unmelyasma
    [49, 27, 38, 46, 44, 0, 6, 22, 28, 6]

    100%|██████████| 208/208 [00:56<00:00,  3.67it/s]

    loss: 23.20675277709961
    corotocity
    [26, 46, 15, 3, 35, 46, 15, 4, 5, 19]

      0%|          | 0/208 [00:00<?, ?it/s]

    ========= Results: epoch 8 of 10 =========
    train loss: 23.34| valid loss: 23.35

    ========= Epoch 9 of 10 =========

      0%|          | 1/208 [00:00<01:16,  2.70it/s]

    loss: 23.150440216064453
    pemacislee
    [48, 46, 7, 46, 61, 4, 22, 21, 63, 3]
    sorenaiabl
    [48, 46, 15, 63, 27, 46, 42, 6, 45, 21]
    bricteecan
    [62, 12, 42, 61, 35, 63, 46, 61, 6, 27]
    vialoughon
    [31, 4, 6, 44, 3, 53, 41, 33, 34, 29]
    traftoveno
    [26, 33, 46, 55, 35, 3, 31, 63, 27, 46]

     25%|██▍       | 51/208 [00:14<00:41,  3.78it/s]

    loss: 23.249794006347656
    trecaliabl
    [62, 12, 46, 61, 6, 44, 4, 6, 45, 21]
    Marvansubs
    [48, 46, 15, 31, 6, 27, 22, 56, 24, 25]
    roncactomo
    [48, 34, 27, 26, 46, 61, 35, 3, 54, 34]
    phaueragol
    [48, 33, 6, 27, 63, 12, 6, 58, 34, 27]
    resdlentet
    [12, 46, 15, 10, 21, 46, 27, 35, 63, 26]

     49%|████▊     | 101/208 [00:28<00:29,  3.61it/s]

    loss: 23.418804168701172
    unninerwos
    [49, 27, 38, 4, 29, 63, 15, 60, 46, 15]
    rossyngisn
    [12, 46, 15, 22, 17, 29, 58, 4, 22, 35]
    morsangelb
    [48, 46, 15, 38, 3, 29, 58, 63, 44, 51]
    wertossict
    [48, 46, 15, 35, 46, 15, 38, 4, 61, 35]
    gerheayatr
    [48, 46, 15, 33, 63, 3, 0, 6, 35, 12]

     73%|███████▎  | 151/208 [00:41<00:15,  3.70it/s]

    loss: 23.153179168701172
    fativedere
    [18, 46, 35, 14, 31, 63, 31, 63, 15, 63]
    vadontubro
    [31, 6, 10, 46, 27, 26, 56, 24, 2, 34]
    warsoiermo
    [48, 46, 15, 38, 3, 42, 63, 15, 38, 46]
    shaunorone
    [22, 33, 46, 20, 31, 63, 15, 34, 29, 63]
    modndoisyl
    [48, 46, 10, 27, 31, 46, 42, 22, 17, 44]

     97%|█████████▋| 201/208 [00:55<00:01,  3.69it/s]

    loss: 23.24637222290039
    fleuredant
    [18, 21, 46, 20, 12, 46, 31, 6, 27, 35]
    erpeceinea
    [46, 15, 26, 3, 26, 46, 42, 29, 63, 6]
    enosyperid
    [49, 27, 3, 38, 19, 32, 63, 12, 42, 10]
    carcrodent
    [48, 46, 15, 26, 12, 3, 31, 63, 27, 35]
    brousesact
    [62, 12, 3, 53, 22, 63, 38, 46, 61, 5]

    100%|██████████| 208/208 [00:56<00:00,  4.54it/s]

    loss: 23.462291717529297
    petasbopae
    [48, 46, 26, 46, 15, 38, 3, 32, 6, 63]

      0%|          | 0/208 [00:00<?, ?it/s]

    ========= Results: epoch 9 of 10 =========
    train loss: 23.30| valid loss: 23.32

    ========= Epoch 10 of 10 =========

      0%|          | 1/208 [00:00<01:13,  2.81it/s]

    loss: 23.180316925048828
    Totickenat
    [48, 46, 26, 42, 61, 57, 63, 29, 6, 35]
    relomelosw
    [12, 46, 44, 3, 11, 63, 21, 3, 22, 26]
    scrygobroi
    [22, 26, 12, 19, 8, 46, 24, 12, 3, 42]
    spterverop
    [22, 32, 35, 63, 15, 31, 63, 12, 3, 32]
    plopensyou
    [62, 12, 3, 32, 63, 27, 35, 19, 3, 53]

     25%|██▍       | 51/208 [00:13<00:42,  3.65it/s]

    loss: 23.340988159179688
    Vopaforedi
    [48, 46, 32, 3, 39, 46, 44, 63, 10, 17]
    thaurnosle
    [26, 33, 46, 46, 15, 38, 46, 22, 21, 63]
    niphreltip
    [48, 42, 32, 33, 12, 46, 44, 35, 42, 32]
    uneangidea
    [49, 27, 63, 6, 29, 58, 14, 31, 63, 6]
    Cytricatel
    [48, 19, 26, 12, 42, 61, 6, 35, 63, 44]

     49%|████▊     | 101/208 [00:27<00:28,  3.75it/s]

    loss: 23.096843719482422
    sphaladica
    [13, 32, 33, 6, 44, 6, 10, 14, 61, 6]
    unsporicta
    [49, 27, 22, 32, 46, 15, 42, 61, 35, 46]
    inizedleca
    [49, 27, 4, 36, 63, 10, 12, 63, 61, 6]
    peldiblegl
    [32, 46, 44, 31, 14, 45, 21, 3, 58, 21]
    tongicubro
    [48, 34, 29, 58, 4, 61, 56, 24, 2, 3]

     73%|███████▎  | 151/208 [00:41<00:16,  3.48it/s]

    loss: 22.76116371154785
    alislydupo
    [13, 44, 4, 22, 21, 19, 10, 56, 32, 46]
    bleneoidos
    [62, 12, 63, 27, 63, 3, 42, 31, 46, 22]
    outatinine
    [3, 53, 35, 6, 35, 14, 29, 4, 29, 63]
    flewrastra
    [18, 21, 63, 26, 12, 46, 15, 26, 12, 46]
    veldebosca
    [48, 46, 44, 31, 63, 51, 46, 22, 26, 46]

     97%|█████████▋| 201/208 [00:54<00:01,  3.68it/s]

    loss: 23.19066047668457
    sthlityssi
    [22, 26, 33, 21, 4, 5, 19, 22, 38, 42]
    ralcubfert
    [12, 6, 44, 26, 20, 55, 39, 63, 15, 35]
    tontermipe
    [26, 46, 27, 35, 63, 15, 25, 42, 32, 63]
    insinially
    [49, 27, 38, 4, 29, 4, 6, 44, 21, 19]
    waarganery
    [48, 46, 46, 15, 8, 34, 29, 63, 15, 17]

    100%|██████████| 208/208 [00:56<00:00,  4.32it/s]

    loss: 23.30339813232422
    Picrertifi
    [48, 46, 61, 12, 63, 15, 35, 14, 39, 42]
    ========= Results: epoch 10 of 10 =========
    train loss: 23.27| valid loss: 23.30

You may wish to try different values of $N$ and see what the impact on sample quality is.

It's also interesting to look at the visited states. For instance, I noticed the following interesting behavior:
- ~~If weight decay is used (`weight_decay=0.00001` in the initialization of the optimizer), the model almost always picks one of two states as the initial state, and uses one of these states to emit vowels and the other to emit consonants (without us programming that behavior or telling the model which letters are vowels!).~~
- ~~But if no weight decay is used, then the model always starts from one particular state. I suspect that this is because weight decay encourages the state priors to be not too big, and if one state prior is very big, the model will always start from that state.~~

EDIT: This behavior disappeared when I fixed the normalization of the transition matrix. But there are other interesting behaviors now:

```
x = torch.tensor(encode("quack")).unsqueeze(0)
T = torch.tensor([5])
print(model.viterbi(x,T))

x = torch.tensor(encode("quick")).unsqueeze(0)
T = torch.tensor([5])
print(model.viterbi(x,T))

x = torch.tensor(encode("qurck")).unsqueeze(0)
T = torch.tensor([5])
print(model.viterbi(x,T)) # should have lower probability---in English only vowels follow "qu"

x = torch.tensor(encode("qiick")).unsqueeze(0)
T = torch.tensor([5])
print(model.viterbi(x,T)) # should have lower probability---in English only "u" follows "q"

```

    ([[30, 53, 6, 16, 57]], tensor([[-15.3841]], device='cuda:0', grad_fn=<GatherBackward>))
    ([[30, 53, 42, 61, 57]], tensor([[-12.1762]], device='cuda:0', grad_fn=<GatherBackward>))
    ([[30, 53, 15, 16, 57]], tensor([[-18.2916]], device='cuda:0', grad_fn=<GatherBackward>))
    ([[30, 53, 42, 61, 57]], tensor([[-20.4187]], device='cuda:0', grad_fn=<GatherBackward>))

## Conclusion

HMMs used to be very popular in natural language processing, but they have largely been overshadowed by neural network models like RNNs and Transformers. Still, it is fun and instructive to study the HMM; some commonly used machine learning techniques like [Connectionist Temporal Classification](https://www.cs.toronto.edu/~graves/icml_2006.pdf) are inspired by HMM methods. HMMs are [still used in conjunction with neural networks in speech recognition](https://arxiv.org/abs/1811.07453), where the assumption of a one-hot state makes sense for modelling phonemes, which are spoken one at a time.

## Acknowledgments

This notebook is based partly on Lawrence Rabiner's excellent article "[A Tutorial on Hidden Markov Models and Selected Applications in Speech Recognition](https://www.cs.cmu.edu/~cga/behavior/rabiner1.pdf)", which you may also like to check out. Thanks also to Dima Serdyuk and Kyle Gorman for their feedback on the draft.