```python
%matplotlib inline

import os
import re
import urllib.request
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
import itertools

from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
```

In this notebook you will work with a deep generative language model that maps words from a continuous latent space. We will use text data (we will work on the character level) in Spanish and pytorch.

The first section concerns data manipulation and data loading classes necessary for our implementation. You do not need to modify anything in this part of the code.

Let's first download the SIGMORPHON dataset that we will be using for this notebook: these are inflected Spanish words together with some morphosyntactic descriptors. For this notebook we will ignore the morphosyntactic descriptors.

```python
url = "https://raw.githubusercontent.com/ryancotterell/sigmorphon2016/master/data/"
train_file = "spanish-task1-train"
val_file = "spanish-task1-dev"
test_file = "spanish-task1-test"

print("Downloading data files...")
if not os.path.isfile(train_file):
    urllib.request.urlretrieve(url + train_file, filename=train_file)
if not os.path.isfile(val_file):
    urllib.request.urlretrieve(url + val_file, filename=val_file)
if not os.path.isfile(test_file):
    urllib.request.urlretrieve(url + test_file, filename=test_file)
print("Download complete.")
```

    Downloading data files...
    Download complete.

# Data

In order to work with text data, we need to transform the text into something that our algorithms can work with. The first step of this process is converting words into word ids. We do this by constructing a vocabulary from the data, assigning a new word id to each new word it encounters.

```python
UNK_TOKEN = "?"
PAD_TOKEN = "_"
SOW_TOKEN = ">"
EOW_TOKEN = "."

def extract_inflected_word(s):
    """
    Extracts the inflected words in the SIGMORPHON dataset.
    """
    return s.split()[-1]

class Vocabulary:

    def __init__(self):
        self.idx_to_char = {0: UNK_TOKEN, 1: PAD_TOKEN, 2: SOW_TOKEN, 3: EOW_TOKEN}
        self.char_to_idx = {UNK_TOKEN: 0, PAD_TOKEN: 1, SOW_TOKEN: 2, EOW_TOKEN: 3}
        self.word_freqs = {}

    def __getitem__(self, key):
        return self.char_to_idx[key] if key in self.char_to_idx else self.char_to_idx[UNK_TOKEN]

    def word(self, idx):
        return self.idx_to_char[idx]

    def size(self):
        return len(self.char_to_idx)

    @staticmethod
    def from_data(filenames):
        """
            Creates a vocabulary from a list of data files. It assumes that the data files have been
            tokenized and pre-processed beforehand.
        """
        vocab = Vocabulary()
        for filename in filenames:
            with open(filename) as f:
                for line in f:

                    # Strip whitespace and the newline symbol.
                    word = extract_inflected_word(line.strip())

                    # Split the words into characters and assign ids to each
                    # new character it encounters.
                    for char in list(word):
                        if char not in vocab.char_to_idx:
                            idx = len(vocab.char_to_idx)
                            vocab.char_to_idx[char] = idx
                            vocab.idx_to_char[idx] = char

        return vocab
```

```python
# Construct a vocabulary from the training and validation data.
print("Constructing vocabulary...")
vocab = Vocabulary.from_data([train_file, val_file])
print("Constructed a vocabulary of %d types" % vocab.size())
```

    Constructing vocabulary...
    Constructed a vocabulary of 37 types

```python
# some examples
print('e', vocab['e'])
print('é', vocab['é'])
print('ș', vocab['ș'])  # something UNKNOWN
```

    e 8
    é 24
    ș 0

We also need to load the data files into memory. We create a simple class `TextDataset` that stores the data as a list of words:

```python
class TextDataset(Dataset):
    """
        A simple class that loads a list of words into memory from a text file,
        split by newlines. This does not do any memory optimisation,
        so if your dataset is very large, you might want to use an alternative
        class.
    """

    def __init__(self, text_file, max_len=30):
        self.data = []
        with open(text_file) as f:
            for line in f:
                word = extract_inflected_word(line.strip())
                if len(list(word)) <= max_len:
                    self.data.append(word)

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

    def __getitem__(self, idx):
        return self.data[idx]
```

```python
# Load the training, validation, and test datasets into memory.
train_dataset = TextDataset(train_file)
val_dataset = TextDataset(val_file)
test_dataset = TextDataset(test_file)

# Print some samples from the data:
print("Sample from training data: \"%s\"" % train_dataset[np.random.choice(len(train_dataset))])
print("Sample from validation data: \"%s\"" % val_dataset[np.random.choice(len(val_dataset))])
print("Sample from test data: \"%s\"" % test_dataset[np.random.choice(len(test_dataset))])
```

    Sample from training data: "fidas"
    Sample from validation data: "gozarían"
    Sample from test data: "figuraríais"

Now it's time to write a function that converts a word into a list of character ids using the vocabulary we created before. This function is `create_batch` in the code cell below. This function creates a batch from a list of words, and makes sure that each word starts with a start-of-word symbol and ends with an end-of-word symbol. Because not all words are of equal length in a certain batch, words are padded with padding symbols so that they match the length of the largest word in the batch. The function returns an input batch, an output batch, a mask of 1s for words and 0s for padding symbols, and the sequence lengths of each word in the batch. The output batch is shifted by one character, to reflect the predictions that the model is expected to make. For example, for a word
\begin{align}
    \text{e s p e s e m o s}
\end{align}
the input sequence is
\begin{align}
    \text{SOW e s p e s e m o s}
\end{align}
and the output sequence is
\begin{align}
    \text{e s p e s e m o s EOW}
\end{align}

You can see the output is shifted wrt the input, that's because we will be computing a distribution for the next character in context of its prefix, and that's why we need to shift the sequence this way.

Lastly, we create an inverse function `batch_to_words` that recovers the list of words from a padded batch of character ids to use during test time.

```python
def create_batch(words, vocab, device, word_dropout=0.):
    """
    Converts a list of words to a padded batch of word ids. Returns
    an input batch, an output batch shifted by one, a sequence mask over
    the input batch, and a tensor containing the sequence length of each
    batch element.
    :param words: a list of words, each a list of token ids
    :param vocab: a Vocabulary object for this dataset
    :param device:
    :param word_dropout: rate at which we omit words from the context (input)
    :returns: a batch of padded inputs, a batch of padded outputs, mask, lengths
    """
    tok = np.array([[SOW_TOKEN] + list(w) + [EOW_TOKEN] for w in words])
    seq_lengths = [len(w)-1 for w in tok]
    max_len = max(seq_lengths)
    pad_id = vocab[PAD_TOKEN]
    pad_id_input = [
        [vocab[w[t]] if t < seq_lengths[idx] else pad_id for t in range(max_len)]
            for idx, w in enumerate(tok)]

    # Replace words of the input with <unk> with p = word_dropout.
    if word_dropout > 0.:
        unk_id = vocab[UNK_TOKEN]
        word_drop =  [
            [unk_id if (np.random.random() < word_dropout and t < seq_lengths[idx]) else word_ids[t] for t in range(max_len)]
                for idx, word_ids in enumerate(pad_id_input)]

    # The output batch is shifted by 1.
    pad_id_output = [
        [vocab[w[t+1]] if t < seq_lengths[idx] else pad_id for t in range(max_len)]
            for idx, w in enumerate(tok)]

    # Convert everything to PyTorch tensors.
    batch_input = torch.tensor(pad_id_input)
    batch_output = torch.tensor(pad_id_output)
    seq_mask = (batch_input != vocab[PAD_TOKEN])
    seq_length = torch.tensor(seq_lengths)

    # Move all tensors to the given device.
    batch_input = batch_input.to(device)
    batch_output = batch_output.to(device)
    seq_mask = seq_mask.to(device)
    seq_length = seq_length.to(device)

    return batch_input, batch_output, seq_mask, seq_length

def batch_to_words(tensors, vocab: Vocabulary):
    """
    Converts a batch of word ids back to words.
    :param tensors: [B, T] word ids
    :param vocab: a Vocabulary object for this dataset
    :returns: an array of strings (each a word).
    """
    words = []
    batch_size = tensors.size(0)
    for idx in range(batch_size):
        word = [vocab.word(t.item()) for t in tensors[idx,:]]

        # Filter out the start-of-word and padding tokens.
        word = list(filter(lambda t: t != PAD_TOKEN and t != SOW_TOKEN, word))

        # Remove the end-of-word token and all tokens following it.
        if EOW_TOKEN in word:
            word = word[:word.index(EOW_TOKEN)]

        words.append("".join(word))
    return np.array(words)
```

In PyTorch the RNN functions expect inputs to be sorted from long words to shorter ones. Therefore we create a simple wrapper class for the DataLoader class that sorts words from long to short:

```python
class SortingTextDataLoader:
    """
    A wrapper for the DataLoader class that sorts a list of words by their
    lengths in descending order.
    """

    def __init__(self, dataloader):
        self.dataloader = dataloader
        self.it = iter(dataloader)

    def __iter__(self):
        return self

    def __next__(self):
        words = None
        for s in self.it:
            words = s
            break

        if words is None:
            self.it = iter(self.dataloader)
            raise StopIteration

        words = np.array(words)
        sort_keys = sorted(range(len(words)),
                           key=lambda idx: len(list(words[idx])),
                           reverse=True)
        sorted_words = words[sort_keys]
        return sorted_words
```

# Model

## Deterministic language model

In language modelling, we model a word $x = \langle x_1, \ldots, x_n \rangle$  of length $n = |x|$ as a sequence of categorical draws:

\begin{align}
X_i|x_{<i} & \sim \text{Cat}(f(x_{<i}; \theta))
& i = 1, \ldots, n \\
\end{align}

where we use $x_{<i}$ to denote a (possibly empty) prefix string, and thus the model makes no Markov assumption. We map from the conditioning context, the prefix $x_{<i}$, to the categorical parameters (a $v$-dimensional probability vector, where $v$ denotes the size of the vocabulary of the language) using a fixed neural network architecture whose parameters we collectively denote by $\theta$.

This assigns the following likelihood to the word
\begin{align}
    P(x|\theta) &= \prod_{i=1}^n P(x_i|x_{<i}, \theta) \\
    &= \prod_{i=1}^n \text{Cat}(x_i|f(x_{<i}; \theta))
\end{align}
where the categorical pmf is $\text{Cat}(k|\pi) = \prod_{j=1}^v \pi_j^{[k=j]} = \pi_k$.

Suppose we have a dataset $\mathcal D = \{x^{(1)}, \ldots, x^{(N)}\}$ containing $N$ i.i.d. observations. Then we can use the log-likelihood function
\begin{align}
\mathcal L(\theta|\mathcal D) &= \sum_{k=1}^{N} \log P(x^{(k)}| \theta) \\
&= \sum_{k=1}^{N} \sum_{i=1}^{|x^{(k)}|} \log \text{Cat}(x^{(k)}_i|f(x^{(k)}_{<i}; \theta))
\end{align}
 to estimate $\theta$ by maximisation:
 \begin{align}
 \theta^\star = \arg\max_{\theta \in \Theta} \mathcal L(\theta|\mathcal D) ~ .
 \end{align}

We can use stochastic gradient-ascent to find a local optimum of $\mathcal L(\theta|\mathcal D)$, which only requires a gradient estimate:

\begin{align}
\nabla_\theta \mathcal L(\theta|\mathcal D) &= \sum_{k=1}^{|\mathcal D|} \nabla_\theta  \log P(x^{(k)}|\theta) \\
&= \sum_{k=1}^{|\mathcal D|} \frac{1}{N} N \nabla_\theta  \log P(x^{(k)}| \theta)  \\
&= \mathbb E_{\mathcal U(1/N)} \left[ N \nabla_\theta  \log P(x^{(K)}| \theta) \right]  \\
&\overset{\text{MC}}{\approx} \frac{N}{M} \sum_{m=1}^M \nabla_\theta  \log P(x^{(k_m)}|\theta) \\
&\text{where }K_m \sim \mathcal U(1/N)
\end{align}

This is a Monte Carlo (MC) estimate of the gradient computed on $M$ data points selected uniformly at random from $\mathcal D$.

For as long as $f$ remains differentiable wrt to its inputs and parameters, we can rely on automatic differentiation to obtain gradient estimates.

An example design for $f$ is:
\begin{align}
\mathbf x_i &= \text{emb}(x_i; \theta_{\text{emb}}) \\
\mathbf h_0 &\in \theta_{\text{rnn}} \\
\mathbf h_i &= \text{rnn}(\mathbf h_{i-1}, \mathbf x_{i-1}; \theta_{\text{rnn}}) \\
f(x_{<i}; \theta) &= \text{softmax}(\text{dense}_v(\mathbf h_{i};  \theta_{\text{out}}))
\end{align}
where
* $\text{emb}$ is a fixed embedding layer with parameters $\theta_{\text{emb}}$;
* $\text{rnn}$ is a recurrent architecture with parameters $\theta_{\text{rnn}}$, e.g. an LSTM or GRU, and $\mathbf h_0$ is part of the architecture's parameters;
* $\text{dense}_v$ is a dense layer with $v$ outputs (vocabulary size) and parameters $\theta_{\text{out}}$.

In what follows we show how to extend this model with a continuous latent word embedding.

## Deep generative language model

We want to model a word $x$ as a draw from the marginal of deep generative model $p(z, x|\theta) = p(z)P(x|z, \theta)$. Note that we use $p(\cdot)$ for pdfs and $P(\cdot)$ for pmfs.

### Generative model

The generative story is:
\begin{align}
    Z & \sim \mathcal N(0, I) \\
    X_i | z, x_{<i} &\sim \text{Cat}(f(z, x_{<i}; \theta)) & i=1, \ldots, n
\end{align}
where $z \in \mathbb R^D$ and  we impose a standard Gaussian prior on latent space. Other choices of prior can induce interesting properties in latent space, however, in this notebook, we use the Gaussian for its simplicity.

It is easy to design $f$ by a simple modification of the deterministic design shown before:
\begin{align}
\mathbf x_i &= \text{emb}(x_i; \theta_{\text{emb}}) \\
\mathbf h_0 &= \tanh(\text{dense}(z; \theta_{\text{init}})) \\
\mathbf h_i &= \text{rnn}(\mathbf h_{i-1}, \mathbf x_{i-1}; \theta_{\text{rnn}}) \\
f(x_{<i}; \theta) &= \text{softmax}(\text{dense}_v(\mathbf h_{i};  \theta_{\text{out}}))
\end{align}
where we just initialise the recurrent cell using $z$. Note we could also use $z$ in other places, for example, as additional input to every update of the recurrent cell $\mathbf h_i = \text{rnn}(\mathbf h_{i-1}, [\mathbf x_{i-1}, z])$. This is an architecture choice which like many others can only be judged empirically or on the basis of practical convenience.

### Parameter estimation

The marginal likelihood, necessary for parameter estimation, is now no longer tractable:
\begin{align}
P(x|\theta) &= \int p(z)P(x|z, \theta) \text{d}z \\
&= \int \mathcal N(z|0, I)\prod_{i=1}^n \text{Cat}(x_i|f(z,x_{<i}; \theta) )\text{d}z
\end{align}
the intractability is clear as even if $f$ would be linear (recall it isn't) the Gaussian and the Categorical are not conjugate and the integration has no closed-form solution.

We turn to variational inference and derive a lowerbound $\mathcal E(\theta, \lambda|\mathcal D)$ on the log-likelihood function

\begin{align}
    \mathcal E(\theta, \lambda|\mathcal D) &= \sum_{k=1}^{|\mathcal D|} \mathcal E_k(\theta, \lambda|x^{(k)})
\end{align}

which for a single datapoint $x$ is
\begin{align}
    \mathcal E(\theta, \lambda|x) &= \mathbb{E}_{q(z|x)}\left[\log P(x|z, \theta)\right] - \text{KL}\left(q(z|x, \lambda)||p(z)\right)\\
\end{align}
where we have introduce an independently parameterised auxiliary distribution $q(z|x, \lambda)$. The distribution $q$ which maximises this *evidence lowerbound* (ELBO) is also the distribution that minimises
\begin{align}
\text{KL}(q(z|x, \lambda)||p(z|x, \theta)) = \mathbb E_{q(z|x, \lambda)}\left[ \frac{q(z|x, \lambda)}{p(z|x, \theta)}\right]
\end{align}
 where $p(z|x, \theta) = \frac{p(x, z|\theta)}{p(x|\theta)}$ is our intractable true posterior. For that reason, we think of $q(z|x, \lambda)$ as an *approximate posterior*.

 The approximate posterior is an independent model of the latent variable given the data, for that reason we also call it an *inference model*.
 In this notebook, our inference model will be a diagonal Gaussian, to make sure that we cover the sample space of our latent variable. A diagonal Gaussian models each component of the latent representation independently, and is an example of mean field (MF) inference:

 \begin{align}
    q(z|x, \lambda) &\overset{\text{MF}}{=} \prod_{d=1}^D \mathcal N(z_d|\mu_d(x; \lambda); \sigma_d(x; \lambda)^2) \\
    &= \mathcal N(z|\mu(x; \lambda), \text{diag}(\sigma(x; \lambda)^2))
 \end{align}

 where we employ neural network architectures to map from the word $x$ to a vector of locations (in $\mathbb R^D$) and a vector of scales (in $\mathbb R^D_{>0}$).

For this choice, the KL term in the ELBO is tractable:

\begin{align}
\text{KL}\left(q(z|x, \lambda)||p(z)\right) &= \sum_{d=1}^D \text{KL}\left(q(z_d|x, \lambda)||p(z_d)\right) \\
&= \sum_{d=1}^D \text{KL}\left(\mathcal N(u_d, s^2)|| \mathcal N(0,1)\right) \\
&= - \frac{1}{2} \sum_{d=1}^D \left(1 + \log s_d^2 - u_d^2 - s_d^2 \right)
\end{align}
where $u_d = \mu_d(x; \lambda)$ and $s_d = \sigma_d(x; \lambda)$.

Here's an example design for our inference model:

\begin{align}
\mathbf x_i &= \text{emb}(x_i; \lambda_{\text{emb}}) \\
\mathbf f_i &= \text{rnn}(\mathbf f_{i-1}, \mathbf x_{i}; \lambda_{\text{fwd}}) \\
\mathbf b_i &= \text{rnn}(\mathbf b_{i+1}, \mathbf x_{i}; \lambda_{\text{bwd}}) \\
\mathbf h &= \text{dense}([\mathbf f_{n}, \mathbf b_1]; \lambda_{\text{hid}}) \\
\mu(x; \lambda) &= \text{dense}(\mathbf h; \lambda_{\text{loc}})\\
\sigma(x; \lambda) &= \text{softplus}(\text{dense}(\mathbf h; \lambda_{\text{scale}}))
\end{align}

where we use the $\text{softplus}$ activation to make sure our scales are strictly positive. Note that $\exp$ would also do that job, though $\text{softplus}$ is assymptotically linear with its input, which gives us better gradient dynamics.

Because we have neural networks compute the diagonal Gaussian parameters for us, we call this *amortised* mean field inference.

### Gradient estimation

We have to obtain gradients of the ELBO

\begin{align}
\nabla_\theta  \mathcal E(\theta, \lambda|x) &=  \mathbb{E}_{q(z|x)}\left[\nabla_\theta \log P(x|z, \theta)\right] - \underbrace{\nabla_\theta \text{KL}\left(q(z|x, \lambda)||p(z)\right)}_{=0}
\end{align}

and

\begin{align}
\nabla_\lambda  \mathcal E(\theta, \lambda|x) &= \nabla_\lambda\mathbb{E}_{q(z|x)}\left[\log P(x|z, \theta)\right] - \nabla_\lambda \text{KL}\left(q(z|x, \lambda)||p(z)\right)
\end{align}

Clearly, the gradient for the generative network is easy to estimate using MC, but the gradient for the inference network is more complicated because we don't have an expected gradient, rather the gradient of an expected value.

But recall, that we chose a Gaussian for approximate posterior, and Gaussians are what we call a location-scale family. Every Gaussian variable can be *re-expressed* or  **reparameterised** in terms of the standard Gaussian (a distribution with fixed parameters), that is:

\begin{align}
\epsilon = \frac{z - \mu(x; \sigma)}{\sigma(x; \lambda)} &\sim \mathcal N(0, I) \\
z = \mu(x; \lambda) + \sigma(x;\lambda) \odot \epsilon &\sim \mathcal N(\mu(x; \lambda), \text{diag}(\sigma(x; \lambda)^2))
\end{align}

where $\epsilon$ is distributed by a $D$-dimensional standard Gaussian and $\odot$ denotes elementwise multiplication.

This means we can rewrite the ELBO as

\begin{align}
\mathcal E(\theta, \lambda|x) &= \mathbb{E}_{\epsilon \sim \mathcal N(0, I)}\left[\log P(x|z=\mu(x; \lambda) + \sigma(x; \lambda) \odot \epsilon, \theta)\right] - \text{KL}\left(q(z|x, \lambda)||p(z)\right)
\end{align}

and though we could also rewrite the KL term, we will leave as is because it is tractable to compute.

Now our gradients are both very simple:

\begin{align}
\nabla_\theta \mathcal E(\theta, \lambda|x) &= \mathbb{E}_{\epsilon \sim \mathcal N(0, I)}\left[\nabla_\theta \log P(x|z=\mu(x; \lambda) + \sigma(x; \lambda) \odot \epsilon, \theta)\right] - \underbrace{\nabla_\theta \text{KL}\left(q(z|x, \lambda)||p(z)\right)}_{=0} \\
&\overset{\text{MC}}{\approx} \frac{1}{S} \sum_{s=1}^S \nabla_\theta \log P(x|z^{(s)}, \theta) \\
&\text{where }z^{(s)} = \mu(x; \lambda) + \sigma(x; \lambda) \odot \epsilon^{(s)}\\
&\text{and }\epsilon^{(s)} \sim \mathcal N(0, I)
\end{align}

and

\begin{align}
\nabla_\lambda \mathcal E(\theta, \lambda|x) &=\nabla_\lambda \mathbb{E}_{\epsilon \sim \mathcal N(0, I)}\left[\nabla_\lambda \log P(x|z=\mu(x; \lambda) + \sigma(x; \lambda) \odot \epsilon, \theta)\right] - \nabla_\lambda \underbrace{\text{KL}\left(q(z|x, \lambda)||p(z)\right)}_{\text{tractable}} \\
&\overset{\text{MC}}{\approx} \left(\frac{1}{S} \sum_{s=1}^S \nabla_\lambda \log P(x|z^{(s)}, \theta)\right) -  \nabla_\lambda \underbrace{\text{KL}\left(q(z|x, \lambda)||p(z)\right)}_{\text{tractable}}\\
&\text{where }z^{(s)} = \mu(x; \lambda) + \sigma(x; \lambda) \odot \epsilon^{(s)}\\
&\text{and }\epsilon^{(s)} \sim \mathcal N(0, I)
\end{align}

and note how both, but especially the second, use this notion of *reparameterised sample* to make all sources of stochasticity independent of the parameters of the network.

Gradient estimates of this sort are known in the literature as *reparameterised gradients* and this makes our model an instance of what is known in the literature as a *variational auto-encoder* (VAE).

## Implementation

We start by implementing our generative model, which requires implementing a Gaussian/Normal distribution:

Check the [wikipedia page about Normal distribution](https://en.wikipedia.org/wiki/Normal_distribution) for information such as the functional form of the pdf.

You will need the KL divergence for univariate Normal distributions:

\begin{align}
\text{KL}\left(\mathcal N(\mu_1, \sigma_1^2) || \mathcal N(\mu_2, \sigma_2^2)\right)
&= \frac{1}{2 \sigma_2^2} \left((\mu_1 - \mu_2)^2 + \sigma_1^2 - \sigma_2^2\right) + \log \frac{\sigma_2}{\sigma_1}
\end{align}

```python
class Normal:
    """
    This is a normal distribution
        N(u, s^2)
    thus specified by a location u and a strictly positive scale.

    This class can hold a collection of D independent Gaussian variables
        by having a D-dimensional vector of locations and
        a D-dimensinal vector of scales.
    """

    def __init__(self, loc, scale):
        """
        :param loc: a tensor of locations (real numbers)
        :param scale: a tensor of scales (strictly positive real numbers)
        """
        pass

    def mean(self):
        """For Gaussians this is the location"""
        pass

    def std(self):
        """For Gaussians this is the scale"""
        pass

    def sample(self):
        """
        Returns a reparameterised sample with the shape of the location parameter.
        """
        pass

    def log_pdf(self, x):
        """
        Assess the log probability density of x.

        :param x: a tensor of Gaussian samples (same shape as the location parameter)
        :returns:  tensor of log probabilitie densities
        """
        pass

    def kl(self, other: 'Normal'):
        """
        The KL divergence between two Gaussians
        :returns: a tensor of KL values with the same shape as the parameters of self.
        """
        pass
```

```python
# SOLUTION
import numpy as np

class Normal:
    """
    This is a normal distribution
        N(u, s^2)
    thus specified by a location u and a strictly positive scale.

    This class can hold a collection of D independent Gaussian variables
        by having a D-dimensional vector of locations and
        a D-dimensinal vector of scales.
    """

    def __init__(self, loc, scale):
        """
        :param loc: a tensor of locations (real numbers)
        :param scale: a tensor of scales (strictly positive real numbers)
        """
        self.loc = loc
        self.scale = scale

    def mean(self):
        """For Gaussians this is the location"""
        return self.loc

    def std(self):
        """For Gaussians this is the scale"""
        return self.scale

    def sample(self):
        """
        Returns a reparameterised sample with the shape of the location parameter.
        """
        epsilon = torch.randn_like(self.scale)
        return self.loc + epsilon * self.scale

    def log_pdf(self, x):
        """
        Assess the log probability density of x.

        :param x: a tensor of Gaussian samples (same shape as the location parameter)
        :returns:  tensor of log probabilitie densities
        """
        return -(((x - self.loc) ** 2) / (2 * (self.scale ** 2))) -0.5 * np.log(2 * np.pi) - torch.log(self.scale)

    def kl(self, other: 'Normal'):
        """
        The KL divergence between two Gaussians
        :returns: a tensor of KL values with the same shape as the parameters of self.
        """
        return (1. / (2 * (other.scale**2))) * ((self.loc - other.loc)**2 + self.scale**2 - other.scale**2) + \
                torch.log(other.scale) - torch.log(self.scale)

```

```python
# tests for Normal
torch.manual_seed(0xBADF00D)
dummy_loc = torch.randn(10, requires_grad=True)
dummy_scale = torch.linspace(0.1, 5.0, 10, requires_grad=True)
dummy_samples = Normal(dummy_loc, dummy_scale).sample()
sample_pdf = Normal(dummy_loc, dummy_scale).log_pdf(dummy_samples)

dummy_grads = torch.autograd.grad(dummy_samples.mean(), [dummy_loc, dummy_scale], allow_unused=True)
assert all(grad is not None for grad in dummy_grads), "samples are not differentiable w.r.t. loc and/or scale. "\
                                                      " Make sure you use torch throughout the implementation "
assert sample_pdf.shape == dummy_loc.shape

# kl divergence
dummy_loc_2 = torch.randn(10, requires_grad=True)

normal1 = Normal(dummy_loc, dummy_scale)
normal2 = Normal(dummy_loc_2, 5.1 - dummy_scale)
dummy_kl = normal1.kl(normal2)

dummy_kl_ref = torch.tensor([3.423, 1.468, 0.8952, 0.4297, 0.1506, 0.2868, 0.8887, 4.5, 21.904, 1355.1639])
assert dummy_kl.shape == dummy_kl_ref.shape, "please return a batch of KL divergence with the same shape as self.loc"
assert torch.allclose(dummy_kl, dummy_kl_ref, rtol=1e-3, atol=1e-3), "your KL implementation is off"
```

Then we should implement the inference model $q(z | x, \lambda)$:

```python
class InferenceModel(nn.Module):

    def __init__(self, vocab_size, embedder, hidden_size,
                 latent_size, pad_idx, bidirectional=False):
        """
        Implement the layers in the inference model.

        :param vocab_size: size of the vocabulary of the language
        :param embedder: embedding layer
        :param hidden_size: size of recurrent cell
        :param latent_size: size D of the latent variable
        :param pad_idx: id of the -PAD- token
        :param bidirectional: whether we condition on x via a bidirectional or
          unidirectional encoder
        """
        pass

    def forward(self, x, seq_mask, seq_len) -> Normal:
        """
        Return an inference Gaussian per instance in the mini-batch
        :param x: words [B, T] as token ids
        :param seq_mask: indicates valid positions vs padding positions [B, T]
        :param seq_len: the length of the sequences [B]
        :return: Gaussian approximate posterior
        """
        pass
```

```python
# SOLUTION
class InferenceModel(nn.Module):

    def __init__(self, vocab_size, embedder, hidden_size,
                 latent_size, pad_idx, bidirectional=False):
        """
        :param vocab_size: size of the vocabulary of the language
        :param embedder: embedding layer
        :param hidden_size: size of recurrent cell
        :param latent_size: size D of the latent variable
        :param pad_idx: id of the -PAD- token
        :param bidirectional: whether we condition on x via a bidirectional or
          unidirectional encoder
        """
        super().__init__()
        self.bidirectional = bidirectional

        # We borrow the embedder from the generative model, but we don't
        # want tobackpropagate through it for the inference model. So we
        # need to make sure to call detach on the embeddings later.
        self.embedder = embedder
        emb_size = embedder.embedding_dim

        # Create a (bidirectional) LSTM to encode x.
        self.lstm = nn.LSTM(emb_size, hidden_size, batch_first=True,
                            bidirectional=bidirectional)

        # The output of the LSTM doubles if we use a bidirectional encoder.
        encoding_size = hidden_size * 2 if bidirectional else hidden_size

        # We can let features interact once more
        self.combination_layer = nn.Linear(encoding_size, 2 * latent_size)

        # Create two affine layers to project the encoder final state to
        # the mean and standard deviation of the diagonal Gaussian that
        # we are predicting.
        self.mu_layer = nn.Linear(2 * latent_size, latent_size)
        self.sigma_layer = nn.Linear(2 * latent_size, latent_size)

    def forward(self, x, seq_mask, seq_len) -> Normal:

        # Compute word embeddings and detach them so that no gradients
        # from the infererence model flow through. That's done because
        # this embedding layer was borrowed from the generative model
        # thus its parameters a part of the set \theta
        x_embed = self.embedder(x).detach()
        # Alternatively, we could have construct an independent embedding layer
        # for the inference net, then its parameters would be part of the set
        # \lambda and we would allow updates

        # Encode the sentence using the LSTM.
        hidden = None
        packed_seq = pack_padded_sequence(x_embed, seq_len, batch_first=True)
        _, final = self.lstm(packed_seq, hidden)

        # Take the final output h_T from the LSTM, concatenate the forward
        # and backward directions for the bidirectional case.
        h_T = final[0]
        if self.bidirectional:
            h_T_fwd = h_T[0]
            h_T_bwd = h_T[1]
            h_T = torch.cat([h_T_fwd, h_T_bwd], dim=-1)

        # We make one more transformation
        #  this allows a few more interactions between features
        #  and if we have bidirectional features then the two
        #  directions also interact
        h_T = torch.tanh(self.combination_layer(h_T))

        # Compute the mean and sigma of the diagonal Gaussian distribution.
        # Use a softplus activation for the standard deviation to ensure it's
        # positive.
        mu = self.mu_layer(h_T)
        sigma = F.softplus(self.sigma_layer(h_T))

        # Return the inferred Gaussian distribution q(z|x).
        qz = Normal(mu, sigma)
        return qz
```

```python
# tests for inference model
pad_idx = vocab.char_to_idx[PAD_TOKEN]

dummy_inference_model = InferenceModel(
    vocab_size=vocab.size(),
    embedder=nn.Embedding(vocab.size(), 64, padding_idx=pad_idx),
    hidden_size=128, latent_size=16, pad_idx=pad_idx, bidirectional=True
).to(device=device)
dummy_batch_size = 32
dummy_dataloader = SortingTextDataLoader(DataLoader(train_dataset, batch_size=dummy_batch_size))
dummy_words = next(dummy_dataloader)

x_in, _, seq_mask, seq_len = create_batch(dummy_words, vocab, device)

q_z_given_x = dummy_inference_model.forward(x_in, seq_mask, seq_len)
assert isinstance(q_z_given_x, Normal), "inference model should return a Normal distribution"
assert q_z_given_x.loc.shape == q_z_given_x.scale.shape, "loc and scale must be of the same size"
assert q_z_given_x.loc.shape == (dummy_batch_size, 16), "model must produce [batch_size x latent_size] units."
assert torch.all(q_z_given_x.scale >= 0), "scale can't be negative"
```

Then we should implement the generative model, we call it BowmanLM after one of the [authors of the model](https://arxiv.org/abs/1511.06349).

```python
class BowmanLM(nn.Module):

    def __init__(self, vocab_size, emb_size, hidden_size, latent_size,
                 pad_idx, dropout=0.):
        """
        :param vocab_size: size of the vocabulary of the language
        :param emb_size: dimensionality of embeddings
        :param hidden_size: dimensionality of recurrent cell
        :param latent_size: this is D the dimensionality of the latent variable z
        :param pad_idx: the id reserved to the -PAD- token
        :param dropout: a dropout rate (you can ignore this for now)
        """
        pass

    def init_hidden(self, z):
        """
        Returns the hidden state of the LSTM initialized with a projection of a given z.
        :param z: [B, D]
        :returns: [B, D] hidden state, [B, D] cell state

        """
        pass

    def step(self, prev_x, z, hidden):
        """
        Performs a single LSTM step for a given previous word and hidden state.
        Returns the unnormalized log probabilities (logits) over the vocabulary for this time step.
        :param prev_x: [B] id of the previous token
        :param z: [B, D] latent variable
        :param hidden:  hidden ([B, H] state, [B, H] cell)
        """
        pass

    def forward(self, x, z):
        """
        Performs an entire forward pass given a sequence of words x and a z.
        :param x: [B, T] token ids
        :param z: [B, D] a latent sample
        """
        hidden = self.init_hidden(z)
        outputs = []
        for t in range(x.size(1)):
            prev_x = x[:, t].unsqueeze(-1)
            scores, hidden = self.step(prev_x, z, hidden)
            outputs.append(scores)
        return torch.cat(outputs, dim=1)

    def loss(self, scores, targets, pz, qz, free_nats=0., evaluation=False):
        """
        Computes the terms in the loss (negative ELBO) given the
            scores (unnormalized log-probabilities), targets,
        the prior distribution p(z), and the approximate posterior distribution q(z|x).

        If free_nats is nonzero it will clamp the KL divergence between the posterior
        and prior to that value, preventing gradient propagation via the KL if it's
        below that value.

        If evaluation is set to true, the loss will be summed instead
        of averaged over the batch.

        Returns the reconstruction loss and the KL.

        The loss
        can be computed from those as loss = rec_loss - KL.

        :returns:
            negative log likelihood (scalar), KL (scalar)
            (use mean for training mode and sum for evaluation mode)
        """
        pass
```

```python
# SOLUTION
class BowmanLM(nn.Module):

    def __init__(self, vocab_size, emb_size, hidden_size, latent_size,
                 pad_idx, dropout=0.):
        """
        :param vocab_size: size of the vocabulary of the language
        :param emb_size: dimensionality of embeddings
        :param hidden_size: dimensionality of recurrent cell
        :param latent_size: this is D the dimensionality of the latent variable z
        :param pad_idx: the id reserved to the -PAD- token
        :param dropout: a dropout rate
        """
        super().__init__()
        self.pad_idx = pad_idx
        self.embedder = nn.Embedding(vocab_size, emb_size,
                                     padding_idx=pad_idx)
        self.lstm = nn.LSTM(emb_size, hidden_size, batch_first=True)
        self.bridge = nn.Linear(latent_size, hidden_size)
        self.projection = nn.Linear(hidden_size, vocab_size, bias=False)
        self.dropout_layer = nn.Dropout(p=dropout)

    def init_hidden(self, z):
        """
        Returns the hidden state of the LSTM initialized with a projection of a given z.
        :param z: [B, D]
        :returns: [B, D] hidden state, [B, D] cell state

        """
        h = self.bridge(z).unsqueeze(0)
        c = self.bridge(z).unsqueeze(0)
        return (h, c)

    def step(self, prev_x, z, hidden):
        """
        Performs a single LSTM step for a given previous word and hidden state.
        Returns the unnormalized probabilities over the vocabulary for this time step.
        :param prev_x: [B] id of the previous token
        :param z: [B, D] latent variable
        :param hidden:  hidden ([B, H] state, [B, H] cell)
        """
        x_embed = self.dropout_layer(self.embedder(prev_x))
        output, hidden = self.lstm(x_embed, hidden)
        scores = self.projection(self.dropout_layer(output))
        return scores, hidden

    def forward(self, x, z):
        """
        Performs an entire forward pass given a sequence of words x and a z.
        :param x: [B, T] token ids
        :param z: [B, D] a latent sample
        """
        hidden = self.init_hidden(z)
        outputs = []
        for t in range(x.size(1)):
            prev_x = x[:, t].unsqueeze(-1)
            scores, hidden = self.step(prev_x, z, hidden)
            outputs.append(scores)
        return torch.cat(outputs, dim=1)

    def loss(self, scores, targets, pz, qz, free_nats=0., evaluation=False):
        """
        Computes the terms in the loss (negative ELBO) given the
            scores (unnormalized log-probabilities), targets,
        the prior distribution p(z), and the approximate posterior distribution q(z|x).

        If free_nats is nonzero it will clamp the KL divergence between the posterior
        and prior to that value, preventing gradient propagation via the KL if it's
        below that value.

        If evaluation is set to true, the loss will be summed instead
        of averaged over the batch.

        Returns the reconstruction loss and the KL.

        The loss
        can be computed from those as loss = rec_loss - KL.

        :returns:
            negative log likelihood (scalar), KL (scalar)
            (use mean for training mode and sum for evaluation mode)
        """

        # Approximate E[log P(x|z)].
        scores = scores.permute(0, 2, 1)
        reconstruction_loss = F.cross_entropy(scores, targets,
                                              ignore_index=self.pad_idx,
                                              reduction="none")
        reconstruction_loss = reconstruction_loss.sum(dim=1)

        # Compute the KL divergence and clamp to at least the given amount of free nats.
        KL = qz.kl(pz).sum(dim=1)
        KL = torch.clamp(KL, min=free_nats)

        # For evaluation return the sum of individual components, for
        # training return the mean of those components.
        if evaluation:
            return (reconstruction_loss.sum(), KL.sum())
        else:
            return (reconstruction_loss.mean(), KL.mean())
```

The code below is used to assess the model and also investigate what it learned. We implemented it for you, so that you can focus on the VAE part. It's useful however to learn from this example: we do interesting things like computing perplexity and sampling novel words!

# Evaluation metrics

During training we'd like to keep track of some evaluation metrics on the validation data in order to keep track of how our model is doing and to perform early stopping. One simple metric we can compute is the ELBO on all the validation or test data using a single sample from the approximate posterior $q(z|x)$:

```python
def eval_elbo(model, inference_model, eval_dataset, vocab, device, batch_size=128):
    """
    Computes a single sample estimate of the ELBO on a given dataset.
    """
    dl = DataLoader(eval_dataset, batch_size=batch_size)
    sorted_dl = SortingTextDataLoader(dl)

    # Make sure the model is in evaluation mode (i.e. disable dropout).
    model.eval()

    total_rec_loss = 0.
    total_KL = 0.
    num_words = 0

    # We don't need to compute gradients for this.
    with torch.no_grad():
        for words in sorted_dl:
            x_in, x_out, seq_mask, seq_len = create_batch(words, vocab, device)

            # Infer the approximate posterior and construct the prior.
            qz = inference_model(x_in, seq_mask, seq_len)
            pz = Normal(torch.zeros_like(qz.mean()),
                        torch.ones_like(qz.std()))

            # Compute the unnormalized probabilities using a single sample from the
            # approximate posterior.
            z = qz.sample()
            scores = model(x_in, z)

            # Compute the reconstruction loss and KL divergence.
            reconstruction_loss, KL = model.loss(scores, x_out, pz, qz,
                                                 free_nats=0.,
                                                 evaluation=True)
            total_rec_loss += reconstruction_loss
            total_KL += KL
            num_words += x_in.size(0)

    # Return the average reconstruction loss and KL.
    avg_rec_loss = total_rec_loss / num_words
    avg_KL = total_KL / num_words
    return avg_rec_loss, avg_KL
```

```python
# test for your BowmanLM implementation with elbo. This may take a few seconds
dummy_lm = BowmanLM(vocab.size(), emb_size=64, hidden_size=128,
                        latent_size=16, pad_idx=pad_idx).to(device=device)

!head -n 128 {val_file} > ./dummy_dataset
dummy_data = TextDataset('./dummy_dataset')
dummy_rec_loss, dummy_kl = eval_elbo(dummy_lm, dummy_inference_model,
                                     dummy_data, vocab, device)
print(dummy_rec_loss, dummy_kl)
assert dummy_rec_loss.item() > 0 and dummy_kl.item() > 0
```

    tensor(36.9271) tensor(1.6146)

A common metric to evaluate language models is the perplexity per word. The perplexity per word for a dataset is defined as:

\begin{align}
    \text{ppl}(\mathcal{D}) = \exp\left(-\frac{1}{\sum_{k=1}^{|\mathcal D|} n^{(k)}} \sum_{k=1}^{|\mathcal{D}|} \log P(x^{(k)})\right)
\end{align}

where $n^{(k)} = |x^{(k)}|$ is the number of tokens in a word and $P(x^{(k)})$ is the probability that our model assigns to the datapoint $x^{(k)}$. In order to compute $\log P(x)$ for our model we need to evaluate the integral:

\begin{align}
    P(x) = \int P(x|z) p(z) dz
\end{align}

As this is an integral  cannot be compute in closed-form, we have two options: we can use the earlier derived lower-bound on the log-likelihood, which will give us an upper-bound on the perplexity, or we can make an importance sampling estimate using our approximate posterior distribution. The importance sampling (IS) estimate can be done as:

\begin{align}
\hat P(x) &\overset{\text{IS}}{\approx} \frac{1}{S} \sum_{s=1}^{S} \frac{p(z^{(s)})p(x|z^{(s)})}{q(z^{(s)}|x)} & \text{where }z^{(s)} \sim q(z|x)
\end{align}

where $S$ is the number of samples.

Then our perplexity becomes:
\begin{align}
    &\frac{1}{\sum_{k=1}^{|\mathcal D|} n^{(k)}}  \sum_{k=1}^{|\mathcal D|} \log p(x^{(k)}) \\
    &\approx \frac{1}{\sum_{k=1}^{|\mathcal D|} n^{(k)}}  \sum_{k=1}^{|\mathcal D|} \log \frac{1}{S} \sum_{s=1}^{S} \frac{p(z^{(s)})p(x^{(k)}|z^{(s)})}{q(z^{(s)}|x^{(k)})} \\
\end{align}

We define the function `eval_perplexity` below that implements this importance sampling estimate:

```python
def eval_perplexity(model, inference_model, eval_dataset, vocab, device,
                    n_samples, batch_size=128):
    """
    Estimates the per-word perplexity using importance sampling with the
    given number of samples.
    """

    dl = DataLoader(eval_dataset, batch_size=batch_size)
    sorted_dl = SortingTextDataLoader(dl)

    # Make sure the model is in evaluation mode (i.e. disable dropout).
    model.eval()

    log_px = 0.
    num_predictions = 0
    num_words = 0

    # We don't need to compute gradients for this.
    with torch.no_grad():
        for words in sorted_dl:
            x_in, x_out, seq_mask, seq_len = create_batch(words, vocab, device)

            # Infer the approximate posterior and construct the prior.
            qz = inference_model(x_in, seq_mask, seq_len)
            pz = Normal(torch.zeros_like(qz.mean()),
                        torch.ones_like(qz.std()))

            # Create an array to hold all samples for this batch.
            batch_size = x_in.size(0)
            log_px_samples = torch.zeros(n_samples, batch_size)

            # Sample log P(x) n_samples times.
            for s in range(n_samples):

                # Sample a z^s from the posterior.
                z = qz.sample()

                # Compute log P(x^k|z^s)
                scores = model(x_in, z)
                cond_log_prob = F.log_softmax(scores, dim=-1)
                cond_log_prob = torch.gather(cond_log_prob, 2, x_out.unsqueeze(-1)).squeeze() # B x T
                cond_log_prob = (cond_log_prob * seq_mask.type_as(cond_log_prob)).sum(dim=1) # B

                # Compute log p(z^s) and log q(z^s|x^k)
                prior_log_prob = pz.log_pdf(z).sum(dim=1) # B
                posterior_log_prob = qz.log_pdf(z).sum(dim=1) # B

                # Store the sample for log P(x^k) importance weighted with p(z^s)/q(z^s|x^k).
                log_px_sample = cond_log_prob + prior_log_prob - posterior_log_prob
                log_px_samples[s] = log_px_sample

            # Average over the number of samples and count the number of predictions made this batch.
            log_px_batch = torch.logsumexp(log_px_samples, dim=0) - \
                    torch.log(torch.Tensor([n_samples]))
            log_px += log_px_batch.sum()
            num_predictions += seq_len.sum()
            num_words += seq_len.size(0)

    # Compute and return the perplexity per word.
    perplexity = torch.exp(-log_px / num_predictions)
    NLL = -log_px / num_words
    return perplexity, NLL
```

Lastly, we want to occasionally qualitatively see the performance of the model during training, by letting it reconstruct a given word from the latent space. This gives us an idea of whether the model is using the latent space to encode some semantics about the data. For this we use a deterministic greedy decoding algorithm, that chooses the word with maximum probability at every time step, and feeds that word into the next time step.

```python
def greedy_decode(model, z, vocab, max_len=50):
    """
    Greedily decodes a word from a given z, by picking the word with
    maximum probability at each time step.
    """

    # Disable dropout.
    model.eval()

    # Don't compute gradients.
    with torch.no_grad():
        batch_size = z.size(0)

        # We feed the model the start-of-word symbol at the first time step.
        prev_x = torch.ones(batch_size, 1, dtype=torch.long).fill_(vocab[SOW_TOKEN]).to(z.device)

        # Initialize the hidden state from z.
        hidden = model.init_hidden(z)

        predictions = []
        for t in range(max_len):
            scores, hidden = model.step(prev_x, z, hidden)

            # Choose the argmax of the unnnormalized probabilities as the
            # prediction for this time step.
            prediction = torch.argmax(scores, dim=-1)
            predictions.append(prediction)

            prev_x = prediction.view(batch_size, 1)

        return torch.cat(predictions, dim=1)
```

# Training

Now it's time to train the model. We use early stopping on the validation perplexity for model selection.

```python
# Define the model hyperparameters.
emb_size = 256
hidden_size = 256
latent_size = 16
bidirectional_encoder = True
free_nats = 5.
annealing_steps = 11400
dropout = 0.6
word_dropout = 0.75
batch_size = 64
learning_rate = 0.001
num_epochs = 20
n_importance_samples = 3 # 50

# Create the training data loader.
dl = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
sorted_dl = SortingTextDataLoader(dl)

# Create the generative model.
model = BowmanLM(vocab_size=vocab.size(),
                 emb_size=emb_size,
                 hidden_size=hidden_size,
                 latent_size=latent_size,
                 pad_idx=vocab[PAD_TOKEN],
                 dropout=dropout)
model = model.to(device)

# Create the inference model.
inference_model = InferenceModel(vocab_size=vocab.size(),
                                 embedder=model.embedder,
                                 hidden_size=hidden_size,
                                 latent_size=latent_size,
                                 pad_idx=vocab[PAD_TOKEN],
                                 bidirectional=bidirectional_encoder)
inference_model = inference_model.to(device)

# Create the optimizer.
optimizer = optim.Adam(itertools.chain(model.parameters(),
                                       inference_model.parameters()),
                       lr=learning_rate)

# Save the best model (early stopping).
best_model = "./best_model.pt"
best_val_ppl = float("inf")
best_epoch = 0

# Keep track of some statistics to plot later.
train_ELBOs = []
train_KLs = []
val_ELBOs = []
val_KLs = []
val_perplexities = []
val_NLLs = []

step = 0
training_ELBO = 0.
training_KL = 0.
num_batches = 0
for epoch_num in range(1, num_epochs+1):
    for words in sorted_dl:

        # Make sure the model is in training mode (for dropout).
        model.train()

        # Transform the words to input, output, seq_len, seq_mask batches.
        x_in, x_out, seq_mask, seq_len = create_batch(words, vocab, device,
                                                      word_dropout=word_dropout)

        # Compute the multiplier for the KL term if we do annealing.
        if annealing_steps > 0:
            KL_weight = min(1., (1.0 / annealing_steps) * step)
        else:
            KL_weight = 1.

        # Do a forward pass through the model and compute the training loss. We use
        # a reparameterized sample from the approximate posterior during training.
        qz = inference_model(x_in, seq_mask, seq_len)
        pz = Normal(torch.zeros_like(qz.mean()),
                    torch.ones_like(qz.std()))
        z = qz.sample()
        scores = model(x_in, z)
        rec_loss, KL = model.loss(scores, x_out, pz, qz, free_nats=free_nats)
        loss = rec_loss + KL_weight * KL

        # Backpropagate and update the model weights.
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        # Update some statistics to track for the training loss.
        training_ELBO += -(rec_loss - KL)
        training_KL += KL
        num_batches += 1

        # Every 100 steps we evaluate the model and report progress.
        if step % 100 == 0:
            val_rec_loss, val_KL = eval_elbo(model, inference_model, val_dataset, vocab, device)
            val_ELBO = -(val_rec_loss - val_KL)
            print("(%d) step %d: training ELBO (KL) = %.2f (%.2f) --"
                  " KL weight = %.2f --"
                  " validation ELBO (KL) = %.2f (%.2f)" %
                  (epoch_num, step, training_ELBO/num_batches,
                   training_KL/num_batches, KL_weight, val_ELBO, val_KL))

            # Update some statistics for plotting later.
            train_ELBOs.append((step, (training_ELBO/num_batches).item()))
            train_KLs.append((step, (training_KL/num_batches).item()))
            val_ELBOs.append((step, val_ELBO.item()))
            val_KLs.append((step, val_KL.item()))

            # Reset the training statistics.
            training_ELBO = 0.
            training_KL = 0.
            num_batches = 0

        step += 1

    # After an epoch we'll compute validation perplexity and save the model
    # for early stopping if it's better than previous models.
    print("Finished epoch %d" % (epoch_num))
    val_perplexity, val_NLL = eval_perplexity(model, inference_model, val_dataset, vocab, device,
                                              n_importance_samples)
    val_rec_loss, val_KL = eval_elbo(model, inference_model, val_dataset, vocab, device)
    val_ELBO = -(val_rec_loss - val_KL)

    # Keep track of the validation perplexities / NLL.
    val_perplexities.append((epoch_num, val_perplexity.item()))
    val_NLLs.append((epoch_num, val_NLL.item()))

    # If validation perplexity is better, store this model for early stopping.
    if val_perplexity < best_val_ppl:
        best_val_ppl = val_perplexity
        best_epoch = epoch_num
        torch.save(model.state_dict(), best_model)

    # Print epoch statistics.
    print("Evaluation epoch %d:\n"
          " - validation perplexity: %.2f\n"
          " - validation NLL: %.2f\n"
          " - validation ELBO (KL) = %.2f (%.2f)"
          % (epoch_num, val_perplexity, val_NLL, val_ELBO, val_KL))

    # Also show some qualitative results by reconstructing a word from the
    # validation data. Use the mean of the approximate posterior and greedy
    # decoding.
    random_word = val_dataset[np.random.choice(len(val_dataset))]
    x_in, _, seq_mask, seq_len = create_batch([random_word], vocab, device)
    qz = inference_model(x_in, seq_mask, seq_len)
    z = qz.mean()
    reconstruction = greedy_decode(model, z, vocab)
    reconstruction = batch_to_words(reconstruction, vocab)[0]
    print("-- Original word: \"%s\"" % random_word)
    print("-- Model reconstruction: \"%s\"" % reconstruction)
```

# Let's plot the training and validation statistics:

```python
steps, training_ELBO = list(zip(*train_ELBOs))
_, training_KL = list(zip(*train_KLs))
_, val_ELBO = list(zip(*val_ELBOs))
_, val_KL = list(zip(*val_KLs))
epochs, val_ppl = list(zip(*val_perplexities))
_, val_NLL = list(zip(*val_NLLs))

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 5))

# Plot training ELBO and KL
ax1.set_title("Training ELBO")
ax1.plot(steps, training_ELBO, "-o")
ax2.set_title("Training KL")
ax2.plot(steps, training_KL, "-o")
plt.show()

# Plot validation ELBO and KL
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 5))
ax1.set_title("Validation ELBO")
ax1.plot(steps, val_ELBO, "-o", color="orange")
ax2.set_title("Validation KL")
ax2.plot(steps, val_KL, "-o",  color="orange")
plt.show()

# Plot validation perplexities.
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 5))
ax1.set_title("Validation perplexity")
ax1.plot(epochs, val_ppl, "-o", color="orange")
ax2.set_title("Validation NLL")
ax2.plot(epochs, val_NLL, "-o",  color="orange")
plt.show()
print()
```

Let's load the best model according to validation perplexity and compute its perplexity on the test data:

```python
# Load the best model from disk.
model = BowmanLM(vocab_size=vocab.size(),
                 emb_size=emb_size,
                 hidden_size=hidden_size,
                 latent_size=latent_size,
                 pad_idx=vocab[PAD_TOKEN],
                 dropout=dropout)
model.load_state_dict(torch.load(best_model))
model = model.to(device)

# Compute test perplexity and ELBO.
test_perplexity, test_NLL = eval_perplexity(model, inference_model, test_dataset, vocab,
                                            device, n_importance_samples)
test_rec_loss, test_KL = eval_elbo(model, inference_model, test_dataset, vocab, device)
test_ELBO = test_rec_loss - test_KL
print("test ELBO (KL) = %.2f (%.2f) -- test perplexity = %.2f -- test NLL = %.2f" %
      (test_ELBO, test_KL, test_perplexity, test_NLL))
```

# Qualitative analysis

Let's have a look at what how our trained model interacts with the learned latent space. First let's greedily decode some samples from the prior to assess the diversity of the model:

```python
# Generate 10 samples from the standard normal prior.
num_prior_samples = 10
pz = Normal(torch.zeros(num_prior_samples, latent_size),
            torch.ones(num_prior_samples, latent_size))
z = pz.sample()
z = z.to(device)

# Use the greedy decoding algorithm to generate words.
predictions = greedy_decode(model, z, vocab)
predictions = batch_to_words(predictions, vocab)
for num, prediction in enumerate(predictions):
    print("%d: %s" % (num+1, prediction))
```

Let's now have a look how good the model is at reconstructing words from the test dataset using the approximate posterior mean and a couple of samples:

```python
# Pick a random test word.
test_word = test_dataset[np.random.choice(len(test_dataset))]

# Infer q(z|x).
x_in, _, seq_mask, seq_len = create_batch([test_word], vocab, device)
qz = inference_model(x_in, seq_mask, seq_len)

# Decode using the mean.
z_mean = qz.mean()
mean_reconstruction = greedy_decode(model, z_mean, vocab)
mean_reconstruction = batch_to_words(mean_reconstruction, vocab)[0]

print("Original: \"%s\"" % test_word)
print("Posterior mean reconstruction: \"%s\"" % mean_reconstruction)

# Decode a couple of samples from the approximate posterior.
for s in range(3):
    z = qz.sample()
    sample_reconstruction = greedy_decode(model, z, vocab)
    sample_reconstruction = batch_to_words(sample_reconstruction, vocab)[0]
    print("Posterior sample reconstruction (%d): \"%s\"" % (s+1, sample_reconstruction))
```

We can also qualitatively assess the smoothness of the learned latent space by interpolating between two words in the test set:

```python
# Pick a random test word.
test_word_1 = test_dataset[np.random.choice(len(test_dataset))]

# Infer q(z|x).
x_in, _, seq_mask, seq_len = create_batch([test_word_1], vocab, device)
qz = inference_model(x_in, seq_mask, seq_len)
qz_1 = qz.mean()

# Pick a random second test word.
test_word_2 = test_dataset[np.random.choice(len(test_dataset))]

# Infer q(z|x) again.
x_in, _, seq_mask, seq_len = create_batch([test_word_2], vocab, device)
qz = inference_model(x_in, seq_mask, seq_len)
qz_2 = qz.mean()

# Now interpolate between the two means and generate words between those.
num_words = 5
print("Word 1: \"%s\"" % test_word_1)
for alpha in np.linspace(start=0., stop=1., num=num_words):
    z = (1-alpha) * qz_1 + alpha * qz_2
    reconstruction = greedy_decode(model, z, vocab)
    reconstruction = batch_to_words(reconstruction, vocab)[0]
    print("(1-%.2f) * qz1.mean + %.2f qz2.mean: \"%s\"" % (alpha, alpha, reconstruction))
print("Word 2: \"%s\"" % test_word_2)
```

```python

```


```python
from IPython.core.display import HTML, Image
css_file = 'style.css'
HTML(open(css_file, 'r').read())
```

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

```python
from sympy import init_printing, Matrix, symbols, eye, Rational
from warnings import filterwarnings
```

```python
init_printing(use_latex = 'mathjax')
filterwarnings('ignore')
```

# Diagonalizing a matrix
# Powers of a matrix A

## Definition

* If A is a *n*×*n*, then a non-zero vector **x** in ℝn is called an *eigenvector* of the matrix A if A**x** is a scalar multiple of **x**
* What this suggests is that if you consider the column vector **x** and multiply it by a scalar (here called λ) (which is then parallel to **x**, just of different length) it results in the same solution as multiplying the matrix A by **x**
* Let's try another explanation: if a matrix A, multiplied with a (column) vector (**x**) results in a scalar multiple of that same (column) vector (and is thus parallel to that (column) vector) then this (column) vector is an eigenvector of the matrix A
    * In essence this multiplication of a matrix with a (column) vector produces another vector on the same line as the original vector
    * Depending on the value of this scalar the resulting vector might point in the opposite direction and be shorter or longer than the original
* This scalar multiple is called the eigenvalue
* Matrices can have more than one eigenvalue and eigenvector

## Derivations

* We need to insert an identity matrix of size *n* into the equation that describes the explanation above
$$ {A}\underline{x}={\lambda}\underline{x} \\ {A}\underline{x}={\lambda}{I}\underline{x} \\ {A}\underline{x}-{\lambda}{I}\underline{x}=\underline{0} \\ \left({A}-{\lambda}{I}\right)\underline{x}=\underline{0} $$

* Look at this carefully and you'll notice that we are suggesting the nullspace (eigenspace) of the matrix (A-λI)
* This matrix has to be singular, i.e. have a determinant of 0
$$ \left|{A}-{\lambda}{I}\right|=0 $$
* Solving this equation (called the characteristic equation) will give us the eigenvalues (λ's)
* It will always be a polynomial in λ (called the characteristic polynomial of A), with a leading coefficient of 1 and a degree of *n* corresponding to the size of A
$$ {p}\left({\lambda}\right)={\lambda}^{n}+{c}_{1}{\lambda}^{n-1}+\dots+{c}_{n} $$
* Substituting them back into...
$$ \left({A}\underline{x}-{\lambda}{I}\right)\underline{x}=\underline{0} $$
* ... allows us to calculate the eigenvector(s) **x**

* Let's look at the following matrix A
$$ A=\begin{bmatrix} 0 & 0 & -2 \\ 1 & 2 & 1 \\ 1 & 0 & 3 \end{bmatrix}\\ A-\lambda I=\begin{bmatrix} 0 & 0 & -2 \\ 1 & 2 & 1 \\ 1 & 0 & 3 \end{bmatrix}-\begin{bmatrix} \lambda  & 0 & 0 \\ 0 & \lambda  & 0 \\ 0 & 0 & \lambda  \end{bmatrix}=\begin{bmatrix} -\lambda  & 0 & -2 \\ 1 & 2-\lambda  & 1 \\ 1 & 0 & 3-\lambda  \end{bmatrix}\\ \begin{vmatrix} -\lambda  & 0 & -2 \\ 1 & 2-\lambda  & 1 \\ 1 & 0 & 3-\lambda  \end{vmatrix}=0\\ { \lambda  }^{ 3 }-5{ \lambda  }^{ 2 }+8\lambda -4=0\\ { \lambda  }_{ 1 }=1,\quad { \lambda  }_{ 2 }={ \lambda  }_{ 3 }=2 $$

* Let's start with the first eigenvalue, which is equal to 1 and replace it in A-λI

```python
A = Matrix([[-1, 0 ,-2], [1, 1, 1], [1, 0, 2]])
A
```

$$\left[\begin{matrix}-1 & 0 & -2\\1 & 1 & 1\\1 & 0 & 2\end{matrix}\right]$$

* We now need the nullspace of this matrix

```python
A.nullspace()
```

$$\begin{bmatrix}\left[\begin{matrix}-2\\1\\1\end{matrix}\right]\end{bmatrix}$$

* We knew that this would be 1-dimensional after looking at the row-reduced form

```python
A.rref()
```

$$\begin{pmatrix}\left[\begin{matrix}1 & 0 & 2\\0 & 1 & -1\\0 & 0 & 0\end{matrix}\right], & \begin{bmatrix}0, & 1\end{bmatrix}\end{pmatrix}$$

* It has rank 2 (two pivot column and 1 free variable

* Now for the other 2 eigenvalues, both equaling 2

```python
A = Matrix([[-2, 0, -2], [1, 0, 1], [1, 0, 1]])
A
```

$$\left[\begin{matrix}-2 & 0 & -2\\1 & 0 & 1\\1 & 0 & 1\end{matrix}\right]$$

```python
A.nullspace()
```

$$\begin{bmatrix}\left[\begin{matrix}0\\1\\0\end{matrix}\right], & \left[\begin{matrix}-1\\0\\1\end{matrix}\right]\end{bmatrix}$$

```python
A.rref()
```

$$\begin{pmatrix}\left[\begin{matrix}1 & 0 & 1\\0 & 0 & 0\\0 & 0 & 0\end{matrix}\right], & \begin{bmatrix}0\end{bmatrix}\end{pmatrix}$$

* Only a single pivot column, therefor rank of 1 and two independent (free) variables

* Corresponding to the first eigenvalue we have a single eigenvector that is the basis for a 1-dimensional (line) eigenspace in ℝ3
* Corresponding to the second (and third) eigenvalues we have two basis vectors for a 2-dimensional plane in ℝ3
* Since we are talking about subspaces, we must note that the zero vector must be in both eigenspaces (type of nullspace), but isn't an eigenvector

## The eigenvalues of triangular (upper and lower) and diagonal matrices

* The eigenvalue of these type of matrices are exactly the entries along the main diagonal

## Real and complex eigenvalues

* There will be characteristic polynomials resulting in complex roots
* The consequences of real-valued eigenvalues for a square matrix A of size *n* are the following
    * The system (A-λI)**x**=**0** has non-trivial solutions
    * There is a non-zero vector **x** in ℝn such that A**x**=λ**x**

## The eigenvector matrix S and eigenvalue matrix Λ

* We need to create S from the (column) eigenvectors such that the following holds
$$ {S}^{-1}{A}{S}=\Lambda $$

* As such, S should be square of size *n*×*n* and invertible, so we need *n* independent eigenvectors

* Suppose we have *n* linearly independent eigenvectors of A
* Put them in the columns of S and calculate AS
$$ AS=A\begin{bmatrix} \vdots  & \vdots  & \vdots  & \vdots  \\ \vdots  & \vdots  & \vdots  & \vdots  \\ { x }_{ 1 } & { x }_{ 2 } & \dots  & { x }_{ n } \\ \vdots  & \vdots  & \vdots  & \vdots  \end{bmatrix}=\begin{bmatrix} \vdots  & \vdots  & \vdots  & \vdots  \\ \vdots  & \vdots  & \vdots  & \vdots  \\ { { \lambda  }_{ 1 }x }_{ 1 } & { \lambda  }_{ 2 }{ x }_{ 2 } & \dots  & { \lambda  }_{ n }{ x }_{ n } \\ \vdots  & \vdots  & \vdots  & \vdots  \end{bmatrix}=\begin{bmatrix} \vdots  & \vdots  & \vdots  & \vdots  \\ \vdots  & \vdots  & \vdots  & \vdots  \\ { x }_{ 1 } & { x }_{ 2 } & \dots  & { x }_{ n } \\ \vdots  & \vdots  & \vdots  & \vdots  \end{bmatrix}\begin{bmatrix} { \lambda  }_{ 1 } & 0 & 0 & 0 \\ 0 & { \lambda  }_{ 2 } & 0 & 0 \\ \vdots  & \vdots  & { \dots  } & \vdots  \\ 0 & 0 & 0 & { \lambda  }_{ n } \end{bmatrix}\\ AS=S\Lambda  $$

* From this we have the following
$$ AS=S\Lambda \\ { S }^{ -1 }AS=\Lambda \\ A=S\Lambda { S }^{ -1 } $$
* Later I will use the computer variable D for this diagonal matrix Λ

## The power of a matrix (only for *n* independent eigenvectors)

* We saw in the example section of the last lecture that the following holds
$$ {A}^{2}\underline{x}={\lambda}^{2}{x} $$
* The eigenvectors are the same for A and A2
* We can also see the following
$$ { A }^{ 2 }=S\Lambda { S }^{ -1 }S\Lambda { S }^{ -1 }=S{ \Lambda  }^{ 2 }{ S }^{ -1 } $$

* The power need not be 2, but any *k* which will have S-1 appearing *k*-1 times

* We thus have the following theorems
$$ { A }^{ k }\rightarrow 0\quad \because \quad k\rightarrow \infty ;\quad \left| { \lambda  }_{ i } \right| $$
* ...and...
* If *k* is a positive integer, λ is an eigenvalue of the matrix A, and **x** is a corresponding eigenvector, then λk is an eigenvalue of Ak and **x** is a corresponding eigenvector

## What makes a matrix diagonalizable

* In discussing diagonalization we are concerned with finding a basis for ℝn that consists of eigenvectors of a given square matrix of size *n*
* These bases can tell us about geometric properties of A and it can simplify numerical computations involving A

* We need to answer two question (which are actually the same)
    * Given a square matrix of size *n*, is there a basis for ℝn consisting of eigenvectors?
    * Given a square matrix of size *n*, it there and invertible matrix S, such that S-1AS is a diagonal matrix? (It is the same matrix S referred to above)
* If such a matrix S exists, it is said to diagonalize A (and we will call the resultant diagonal matrix D)

* In short the answer to the above question(s) is yes if A has *n* independent eigenvectors
    * This happens if all λ's are different (none are repeated) (not totally excluded if they are repeated though)

* If they are repeated, we still might have independent eigenvectors, i.e. any size identity matrix (because it is already diagonal)

```python
A = eye(5)
A
```

$$\left[\begin{matrix}1 & 0 & 0 & 0 & 0\\0 & 1 & 0 & 0 & 0\\0 & 0 & 1 & 0 & 0\\0 & 0 & 0 & 1 & 0\\0 & 0 & 0 & 0 & 1\end{matrix}\right]$$

```python
A.eigenvals()
```

$$\begin{Bmatrix}1 : 5\end{Bmatrix}$$

```python
A.eigenvects()
```

$$\begin{bmatrix}\begin{pmatrix}1, & 5, & \begin{bmatrix}\left[\begin{matrix}1\\0\\0\\0\\0\end{matrix}\right], & \left[\begin{matrix}0\\1\\0\\0\\0\end{matrix}\right], & \left[\begin{matrix}0\\0\\1\\0\\0\end{matrix}\right], & \left[\begin{matrix}0\\0\\0\\1\\0\end{matrix}\right], & \left[\begin{matrix}0\\0\\0\\0\\1\end{matrix}\right]\end{bmatrix}\end{pmatrix}\end{bmatrix}$$

* Here we look at a triangular matrix, though

```python
A = Matrix([[2, 1], [0, 2]])
A.eigenvals()
```

$$\begin{Bmatrix}2 : 2\end{Bmatrix}$$

```python
A.eigenvects()
```

$$\begin{bmatrix}\begin{pmatrix}2, & 2, & \begin{bmatrix}\left[\begin{matrix}1\\0\end{matrix}\right]\end{bmatrix}\end{pmatrix}\end{bmatrix}$$

* We can use python™ code to calculate the diagonalized matrix

```python
A = Matrix([[3, -2,  4, -2], [5,  3, -3, -2], [5, -2,  2, -2], [5, -2, -3, 3]])
A
```

$$\left[\begin{matrix}3 & -2 & 4 & -2\\5 & 3 & -3 & -2\\5 & -2 & 2 & -2\\5 & -2 & -3 & 3\end{matrix}\right]$$

```python
S, D = A.diagonalize()
```

```python
S # S, such that A = S times D times the inverse of S
```

$$\left[\begin{matrix}0 & 1 & 1 & 0\\1 & 1 & 1 & -1\\1 & 1 & 1 & 0\\1 & 1 & 0 & 1\end{matrix}\right]$$

```python
D # The diagonal
```

$$\left[\begin{matrix}-2 & 0 & 0 & 0\\0 & 3 & 0 & 0\\0 & 0 & 5 & 0\\0 & 0 & 0 & 5\end{matrix}\right]$$

```python
S * D * S.inv() == A # Checking to see if our statement above is correct
```

    True

```python
S.inv() * A * S == D # Checking to see if our statement above is correct
```

    True

```python
A.eigenvals()
```

$$\begin{Bmatrix}-2 : 1, & 3 : 1, & 5 : 2\end{Bmatrix}$$

* Remember Λ from above?
    * The eigenvalues are precisely the entries along the main diagonal of the diagonal matrix

* To produce the required diagonal matrix manually then will require computing *n* linearly independent eigenvectors for matrix A of size *n* (assuming that it is diagonalizable), creating a matrix with its columns equal to these eigenvectors (called matrix S) and performing the equation S-1AS to calculate the diagonal matrix D ( or Λ)

* Back to the topic of what makes a matrix diagonalizable

* Suppose we have an equation that starts with some vector and every subsequent vector is a matrix A time the previous vector
$$ \underline{u}_{k+1}={A}\underline{u}_{k} $$

* From this arises the following
$$ { \underline { u }  }_{ 1 }=A{ \underline { u }  }_{ 0 }\\ { \underline { u }  }_{ 2 }=A{ A\underline { u }  }_{ 0 }={ A }^{ 2 }{ \underline { u }  }_{ 0 }\\ { \underline { u }  }_{ k }={ A }^{ k }{ \underline { u }  }_{ 0 } $$

* To really solve this problem, rewrite **u**0 as follows (a certain scalar times an eigenvector)
$$ { \underline { u }  }_{ 0 }={ c }_{ 1 }\underline{ x }_{ 1 }+{ c }_{ 2 }\underline{ x }_{ 2 }+\dots +{ c }_{ n }\underline{ x }_{ n } = {S}\underline{c} $$
* Where the S**c** is a linear combination of the individual eigenvectors

* Now multiply both sides by A
$$ A{ \underline { u }  }_{ 0 }={ c }_{ 1 }{A}\underline{ x }_{ 1 }+{ c }_{ 2 }{A}\underline{ x }_{ 2 }+\dots +{ c }_{ n }{A}\underline{ x }_{ n } \\ A{ \underline { u }  }_{ 0 }={ c }_{ 1 }{ \lambda  }_{ 1 }\underline{ x }_{ 1 }+{ c }_{ 2 }{ \lambda  }_{ 2 }\underline{ x }_{ 2 }+\dots +{ c }_{ n }{ \lambda  }_{ n }\underline{ x }_{ n } $$
* Taking a power of A now (i.e. *k*) would be akin to taking each eigenvalue to that power
$$ {A}^{k}{ \underline { u }  }_{ 0 }={ c }_{ 1 }{ \lambda  }_{ 1 }^{k}\underline{ x }_{ 1 }+{ c }_{ 2 }{ \lambda  }_{ 2 }^{k}\underline{ x }_{ 2 }+\dots +{ c }_{ n }{ \lambda  }_{ n }^{k}\underline{ x }_{ n } $$
* This can be written as
$$ \underline{u}_{k} = {A}^{k}\underline{u}_{0}={\Lambda}^{k}{S}\underline{c} $$

* As an example consider the Fibonacci numbers: 0, 1, 1, 2, 3, 5, 8, 13, ...
* What would the 100th number be?
* Consider the following
$$ {F}_{k+2}={F}_{k+1}+{F}_{k} $$
* This is a (second-order) difference equation; think of this example as similar to a second-order differential equation (without derivatives)
* By adding a second equation Fk+1=Fk+1, consider **u**k to be the following vector
$$ \underline{u}_{k}=\begin{bmatrix} {F}_{k+1} \\ {F}_{k} \end{bmatrix} $$
* This means the following
$$ \underline{u}_{k+1}=\begin{bmatrix} 1 & 1 \\ 1 & 0 \end{bmatrix} \begin{bmatrix} {F}_{k+1} \\ {F}_{k} \end{bmatrix}=\begin{bmatrix} {F}_{k+1}+{F}_{k} \\ {F}_{k+1} \end{bmatrix} \\ \underline{u}_{k+1}=\begin{bmatrix} 1 & 1 \\ 1 & 0 \end{bmatrix}\underline{u}_{k} $$

```python
A = Matrix([[1, 1], [1, 0]])
A
```

$$\left[\begin{matrix}1 & 1\\1 & 0\end{matrix}\right]$$

```python
A.eigenvals()
```

$$\begin{Bmatrix}\frac{1}{2} + \frac{\sqrt{5}}{2} : 1, & - \frac{\sqrt{5}}{2} + \frac{1}{2} : 1\end{Bmatrix}$$

```python
A.eigenvects()
```

$$\begin{bmatrix}\begin{pmatrix}\frac{1}{2} + \frac{\sqrt{5}}{2}, & 1, & \begin{bmatrix}\left[\begin{matrix}- \frac{1}{- \frac{\sqrt{5}}{2} + \frac{1}{2}}\\1\end{matrix}\right]\end{bmatrix}\end{pmatrix}, & \begin{pmatrix}- \frac{\sqrt{5}}{2} + \frac{1}{2}, & 1, & \begin{bmatrix}\left[\begin{matrix}- \frac{1}{\frac{1}{2} + \frac{\sqrt{5}}{2}}\\1\end{matrix}\right]\end{bmatrix}\end{pmatrix}\end{bmatrix}$$

```python
S, D = A.diagonalize()
```

```python
D
```

$$\left[\begin{matrix}\frac{1}{2} + \frac{\sqrt{5}}{2} & 0\\0 & - \frac{\sqrt{5}}{2} + \frac{1}{2}\end{matrix}\right]$$

* From above we remember the following
$$ \underline{u}_{k} = {A}^{k}\underline{u}_{0}={\Lambda}^{k}{S}\underline{c} $$
* We have **u**0 contains the first two values

```python
u_zero =  Matrix([1, 0])
u_100 = A ** 100 * u_zero
u_100 # The top value is the 100th Fibonacci number
```

$$\left[\begin{matrix}573147844013817084101\\354224848179261915075\end{matrix}\right]$$

```python
u_four = A ** 4 * u_zero
u_four # If the first number is 0 the the fourth number would be the top value
```

$$\left[\begin{matrix}5\\3\end{matrix}\right]$$

## Example problems

### Example problem 1

* Find an equation for Ck where C is given by the following matrix
$$  $$
* Calculate C100 when *a*=*b*=-1

#### Solution

```python
a, b, k = symbols('a b k')
```

```python
C = Matrix([[2 * b - a, a - b], [2 * b - 2 * a, 2 * a - b]])
C
```

$$\left[\begin{matrix}- a + 2 b & a - b\\- 2 a + 2 b & 2 a - b\end{matrix}\right]$$

* We remember the following
$$ {A}^{k}={S}{\Lambda}^{k}{S}^{-1} $$
* Where Λ is denoted by the computer variable D

```python
S, D = C.diagonalize()
```

```python
S
```

$$\left[\begin{matrix}- \frac{2 a - 2 b}{- 3 a + 3 b + \sqrt{\left(a - b\right)^{2}}} & \frac{2 a - 2 b}{3 a - 3 b + \sqrt{\left(a - b\right)^{2}}}\\1 & 1\end{matrix}\right]$$

* Python™ is not always good at simplifying these
* If you look at it carefully you will note the following

```python
S = Matrix([[1, Rational(1, 2)], [1, 1]])
S
```

$$\left[\begin{matrix}1 & \frac{1}{2}\\1 & 1\end{matrix}\right]$$

```python
D
```

$$\left[\begin{matrix}\frac{a}{2} + \frac{b}{2} - \frac{1}{2} \sqrt{\left(a - b\right)^{2}} & 0\\0 & \frac{a}{2} + \frac{b}{2} + \frac{1}{2} \sqrt{\left(a - b\right)^{2}}\end{matrix}\right]$$

```python
D = Matrix([[b, 0], [0, a]])
D
```

$$\left[\begin{matrix}b & 0\\0 & a\end{matrix}\right]$$

* For the values given, we have the following

```python
C = Matrix([[-1, 0],  [0, -1]])
C
```

$$\left[\begin{matrix}-1 & 0\\0 & -1\end{matrix}\right]$$

```python
S, D = C.diagonalize()
```

```python
D
```

$$\left[\begin{matrix}-1 & 0\\0 & -1\end{matrix}\right]$$

```python
D ** 100
```

$$\left[\begin{matrix}1 & 0\\0 & 1\end{matrix}\right]$$

```python
S * (D ** 100) * S.inv()
```

$$\left[\begin{matrix}1 & 0\\0 & 1\end{matrix}\right]$$

* Doing the same, but with eigenvalues and eigenvectors

```python
C = Matrix([[2 * b - a, a - b], [2 * b - 2 * a, 2 * a - b]])
C
```

$$\left[\begin{matrix}- a + 2 b & a - b\\- 2 a + 2 b & 2 a - b\end{matrix}\right]$$

```python
C.eigenvals()
```

$$\begin{Bmatrix}\frac{a}{2} + \frac{b}{2} - \frac{1}{2} \sqrt{\left(a - b\right)^{2}} : 1, & \frac{a}{2} + \frac{b}{2} + \frac{1}{2} \sqrt{\left(a - b\right)^{2}} : 1\end{Bmatrix}$$

* This simplifies the λ1 = *b* and λ2 = *a*
* That makes Λ (or D) the following

```python
D = Matrix([[b, 0], [0, a]])
D
```

$$\left[\begin{matrix}b & 0\\0 & a\end{matrix}\right]$$

```python
C.eigenvects() # The solution is two tuples, with each being eigenvalue, eigenvector
```

$$\begin{bmatrix}\begin{pmatrix}\frac{a}{2} + \frac{b}{2} - \frac{1}{2} \sqrt{\left(a - b\right)^{2}}, & 1, & \begin{bmatrix}\left[\begin{matrix}- \frac{a - b}{- \frac{3 a}{2} + \frac{3 b}{2} + \frac{1}{2} \sqrt{\left(a - b\right)^{2}}}\\1\end{matrix}\right]\end{bmatrix}\end{pmatrix}, & \begin{pmatrix}\frac{a}{2} + \frac{b}{2} + \frac{1}{2} \sqrt{\left(a - b\right)^{2}}, & 1, & \begin{bmatrix}\left[\begin{matrix}- \frac{a - b}{- \frac{3 a}{2} + \frac{3 b}{2} - \frac{1}{2} \sqrt{\left(a - b\right)^{2}}}\\1\end{matrix}\right]\end{bmatrix}\end{pmatrix}\end{bmatrix}$$

* This simplifies to the following eigenvalue matrix S

```python
S = Matrix([[1, Rational(1, 2)], [1, 1]])
S
```

$$\left[\begin{matrix}1 & \frac{1}{2}\\1 & 1\end{matrix}\right]$$

* We can see if we can get back to C

```python
S * D * S.inv()
```

$$\left[\begin{matrix}- a + 2 b & a - b\\- 2 a + 2 b & 2 a - b\end{matrix}\right]$$

```python
S * D * S.inv() == C
```

    True

* Python™ won't to Dk for you, but it's easy to do yourself

```python
D = Matrix([[b ** k, 0], [0, a ** k]])
D
```

$$\left[\begin{matrix}b^{k} & 0\\0 & a^{k}\end{matrix}\right]$$

* Now we can compute SΛS-1

```python
S * D * S.inv()
```

$$\left[\begin{matrix}- a^{k} + 2 b^{k} & a^{k} - b^{k}\\- 2 a^{k} + 2 b^{k} & 2 a^{k} - b^{k}\end{matrix}\right]$$

* Placing the given values into this equation will give you the same solution for C100 as above

```python

```

Mathematics/Linear Algebra/0.0 MIT-18.06 - Jupyter/Lecture_22_Diagonalization_and_Powers.ipynb


# Code stuff - not slides!

```python
%run ../ML_plots.ipynb
```

    ERROR:root:File `'../ML_plots.ipynb.py'` not found.

# Session 12:
## Supervised learning, part 1

*Andreas Bjerre-Nielsen*

## Agenda
1. [Modelling data](#Modelling-data)
1. [A familiar regression model](#A-familiar-regression-model)
1. [The curse of overfitting](#The-curse-of-overfitting)
1. [Important details](#Implementation-details)

## Vaaaamos

```python
import warnings
from sklearn.exceptions import ConvergenceWarning
warnings.filterwarnings(action='ignore', category=ConvergenceWarning)

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

plt.style.use('default') # set style (colors, background, size, gridlines etc.)
plt.rcParams['figure.figsize'] = 10, 4 # set default size of plots
plt.rcParams.update({'font.size': 18})
```

## Supervised problems (1)
*How do we distinguish between problems?*

```python
f_identify_question
```

## Supervised problems (2)
*The two canonical problems*

```python
f_identify_answer
```

## Supervised problems (3)
*Which models have we seen for classification?*

- .

- .

- .

# Modelling data

## Model complexity (1)
*What does a model of low complexity look like?*

```python
f_complexity[0]
```

## Model complexity (2)
*What does medium model complexity look like?*

```python
f_complexity[1]
```

## Model complexity (3)
*What does high model complexity look like?*

```python
f_complexity[2]
```

## Model fitting (1)
*Quiz (1 min.): Which model fitted the data best?*

```python
f_bias_var['regression'][2]
```

## Model fitting (2)
*What does underfitting and overfitting look like for classification?*

```python
f_bias_var['classification'][2]
```

## Two agendas (1)

What are the objectives of empirical research?

1. *causation*: what is the effect of a particular variable on an outcome?
2. *prediction*: find some function that provides a good prediction of $y$ as a function of $x$

## Two agendas (2)

How might we express the agendas in a model?

$$ y = \alpha + \beta x + \varepsilon $$

- *causation*: interested in $\hat{\beta}$

- *prediction*: interested in $\hat{y}$

## Two agendas (3)

Might these two agendas be related at a deeper level?

Can prediction quality inform us about how to make causal models?

# A familiar regression model

## Estimation (1)
*Do we know already some ways to estimate regression models?*

- Social scientists know all about the Ordinary Least Squares (OLS).
    - OLS estimate both parameters and their standard deviation.
    - Is best linear unbiased estimator under regularity conditions.

*How is OLS estimated?*

- $\beta=(\textbf{X}^T\textbf{X})^{-1}\textbf{X}^T\textbf{y}$
- computation requires non perfect multicollinarity.

## Estimation (2)
*How might we estimate a linear regression model?*

- first order method (e.g. gradient descent)
- second order method (e.g. Newton-Raphson)

*So what the hell was gradient descent?*

- compute errors, multiply with features and update

## Estimation (3)
*Can you explain that in details?*

- Yes, like with Adaline, we minimize the sum of squared errors (SSE):
\begin{align}SSE&=\boldsymbol{e}^{T}\boldsymbol{e}\\\boldsymbol{e}&=\textbf{y}-\textbf{X}\textbf{w}\end{align}

```python
X = np.random.normal(size=(3,2))
y = np.random.normal(size=(3))
w = np.random.normal(size=(3))

e = y-(w[0]+X.dot(w[1:]))
SSE = e.T.dot(e)
```

## Estimation (4)
*And what about the updating..? What is it something about the first order deritative?*

\begin{align}
\frac{\partial SSE}{\partial\hat{w}}=&\textbf{X}^T\textbf{e},\\
 \Delta\hat{w}=&\eta\cdot\textbf{X}^T\textbf{e}=\eta\cdot\textbf{X}^T(\textbf{y}-\hat{\textbf{y}})
\end{align}

```python
eta = 0.001 # learning rate
fod = X.T.dot(e)
update_vars = eta*fod
update_bias = eta*e.sum()
```

## Estimation (5)
*What might some advantages be relative to OLS?*

- Works despite high multicollinarity
- Speed
    - OLS has $\mathcal{O}(K^2N)$ computation time ([read more](https://math.stackexchange.com/questions/84495/computational-complexity-of-least-square-regression-operation))
        - Quadratic scaling in number of variables ($K$).
    - Stochastic gradient descent
        - Likely to converge faster with many observations ($N$)

## Fitting a polynomial (1)
Polyonomial: $f(x) = 2+8*x^4$

Try models of increasing order polynomials.

- Split data into train and test (50/50)

- For polynomial order 0 to 9:
    - Iteration n: $y = \sum_{k=0}^{n}(\beta_k\cdot x^k)+\varepsilon$.
    - Estimate order n model on training data
    - Evaluate with on test data with RMSE:
        - $log RMSE = \log (\sqrt{MSE})$

## Fitting a polynomial (2)
We generate samples of data from true model.

```python
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression

def true_fct(X):
    return 2+X**4

n_samples = 25
n_degrees = 15

np.random.seed(0)

X_train = np.random.normal(size=(n_samples,1))
y_train = true_fct(X_train).reshape(-1) + np.random.randn(n_samples)

X_test = np.random.normal(size=(n_samples,1))
y_test = true_fct(X_test).reshape(-1) + np.random.randn(n_samples)
```

## Fitting a polynomial (3)
We estimate the polynomials

```python
from sklearn.metrics import mean_squared_error as mse

test_mse = []
train_mse = []
parameters = []
degrees = range(n_degrees+1)

for p in degrees:
    X_train_p = PolynomialFeatures(degree=p).fit_transform(X_train)
    X_test_p = PolynomialFeatures(degree=p).fit_transform(X_train)
    reg = LinearRegression().fit(X_train_p, y_train)
    train_mse += [mse(reg.predict(X_train_p),y_train)]
    test_mse += [mse(reg.predict(X_test_p),y_test)]
    parameters.append(reg.coef_)
```

## Fitting a polynomial (4)
*So what happens to the model performance in- and out-of-sample?*

```python
degree_index = pd.Index(degrees,name='Polynomial degree ~ model complexity')
ax = pd.DataFrame({'Train set':train_mse, 'Test set':test_mse})\
    .set_index(degree_index)\
    .plot(figsize=(10,4))
ax.set_ylabel('Mean squared error')
```

## Fitting a polynomial (4)
*Why does it go wrong?*
- more spurious parameters
- the coefficient size increases

## Fitting a polynomial (5)
*What do you mean coefficient size increase?*

```python
order_idx = pd.Index(range(n_degrees+1),name='Polynomial order')
ax = pd.DataFrame(parameters,index=order_idx)\
.abs().mean(1)\
.plot(logy=True)
ax.set_ylabel('Mean parameter size')
```

## Fitting a polynomial (6)
*How else could we visualize this problem?*

```python
f_bias_var['regression'][2]
```

# The curse of overfitting

## Looking for a remedy
*How might we solve the overfitting problem?*

By reducing
- the number of variables
- the coefficient size of variables

## Regularization (1)

*Why do we regularize?*

- To mitigate overfitting > better model predictions

*How do we regularize?*

- We make models which are less complex:
  - reducing the **number** of coefficient;
  - reducing the **size** of the coefficients.

## Regularization (2)

*What does regularization look like?*

We add a penalty term our optimization procedure:

$$ \text{arg min}_\beta \, \underset{\text{MSE}}{\underbrace{E[(y_0 - \hat{f}(x_0))^2]}} + \underset{\text{penalty}}{\underbrace{\lambda \cdot R(\beta)}}$$

Introduction of penalties implies that increased model complexity has to be met with high increases precision of estimates.

## Regularization (3)

*What are some used penalty functions?*

The two most common penalty functions are L1 and L2 regularization.

- L1 regularization (***Lasso***): $R(\beta)=\sum_{j=1}^{p}|\beta_j|$
    - Makes coefficients sparse, i.e. selects variables by removing some (if $\lambda$ is high)

- L2 regularization (***Ridge***): $R(\beta)=\sum_{j=1}^{p}\beta_j^2$
    - Reduce coefficient size
    - Fast due to analytical solution

*To note:* The *Elastic Net* uses a combination of L1 and L2 regularization.

## Regularization (4)

*How the Lasso (L1 reg.) deviates from OLS*

## Regularization (5)

*How the Ridge regression (L2 reg.) deviates from OLS*

## Regularization (6)

*How might we describe the $\lambda$ of Lasso and Ridge?*

These are hyperparameters that we can optimize over.

- More about this tomorrow.

# Implementation details

## The devils in the details (1)

*So we just run regularization?*

# NO

We need to rescale our features:
- convert to zero mean:
- standardize to unit std:

Compute in Python:
- option 1: `StandardScaler` in `sklearn`
- option 2: `(X - np.mean(X)) / np.std(X)`

## The devils in the details (2)
*So we just scale our test and train?*

# NO

Fit to the distribution in the training data first, then rescale train and test! See more [here](https://stats.stackexchange.com/questions/174823/how-to-apply-standardization-normalization-to-train-and-testset-if-prediction-i).

## The devils in the details (3)
*So we just rescale before using polynomial features?*

# NO

Otherwise the interacted varaibles are not gaussian distributed.

## The devils in the details (4)
*Does sklearn's `PolynomialFeatures` work for more than variable?*

# YES!

# The end
[Return to agenda](#Agenda)


# Summary 2020-12-15

```python
# %load imports.py
"""
These is the standard setup for the notebooks.
"""

%matplotlib inline
%load_ext autoreload
%autoreload 2

from jupyterthemes import jtplot
jtplot.style(theme='onedork', context='notebook', ticks=True, grid=False)

import pandas as pd
pd.options.display.max_rows = 999
pd.options.display.max_columns = 999
pd.set_option("display.max_columns", None)
import numpy as np
import os
import matplotlib.pyplot as plt
from collections import OrderedDict
#plt.style.use('paper')

#import data
import copy
from mdldb.run import Run

from sklearn.pipeline import Pipeline
from rolldecayestimators.transformers import CutTransformer, LowpassFilterDerivatorTransformer, ScaleFactorTransformer, OffsetTransformer
from rolldecayestimators.direct_estimator_cubic import EstimatorQuadraticB, EstimatorCubic
from rolldecayestimators.ikeda_estimator import IkedaQuadraticEstimator
import rolldecayestimators.equations as equations
import rolldecayestimators.lambdas as lambdas
from rolldecayestimators.substitute_dynamic_symbols import lambdify
import rolldecayestimators.symbols as symbols
import sympy as sp

from sympy.physics.vector.printing import vpprint, vlatex
from IPython.display import display, Math, Latex

from sklearn.metrics import r2_score
from src.data import database
from mdldb import tables

```

    Duplicate key in file WindowsPath('C:/Users/maa/.matplotlib/stylelib/paper.mplstyle'), line 461 ('figure.figsize   : 5, 3   ## figure size in inches')
    Duplicate key in file WindowsPath('C:/Users/maa/.matplotlib/stylelib/paper.mplstyle'), line 462 ('figure.dpi       : 100        ## figure dots per inch')

## Nomenclature
| Variable | Explain |
|---|---|
|$\pi$| example |

Here is a cell link: [Logbook](#logbook)

# Abstract

Many cost-efficient computation methods have been developed over the years to analyze various aspects of ship hydrodynamics such as:  resistance,  propulsion  and  seakeeping. Getting the  best  possible  accuracy with the lowest possible computational cost is an important factor in a ship’s early design stage.  Potential flow-based analysis partly presents such a solution for seakeeping, with good accuracy for heave and pitch, but not for roll where the roll damping contains both inviscid and viscouseffects.  Roll motion is, however, often a critical degree of freedom that needs to be analyzed since large roll motions can result in cargo shifting or even capsizing. The viscous part of roll damping can be assessed with high accuracy by means of experimental model tests or URANS calculations, but these are generally too expensive in the early design stage of ships. Many semi-empirical formulas to determine viscous damping were therefore developed during the 1970s, where Ikeda’s method is one of the most widely used.  The viscous damping from this method is normally combined with inviscid roll damping from strip theory.

With today’s computational power, more advanced potential flow methods can be used in the seakeeping analysis to enhance the accuracy in the predictions, but still at relatively low computational cost.  This paper investigates the feasibility of combining 3D unsteady fully nonlinearpotential flow (FNPF) theory solved by means of a Boundary ElementMethod  (BEM)  together  with  the  viscous  contributions  from  Ikeda’smethod.

The approach of substituting the inviscid part from Ikeda’s method using strip theory with FNPF is investigated by conducting roll decay simulations. The results estimated by the proposed approach are compared with both the classic strip theory approach and roll decay model tests.
**It is found that potential improvements to the modelling of roll damping can be achieved by introducing FNPF analysis in the Ikeda’s method.**

# Abstract

* **Important**: Good accuracy at low computational cost.

(URANS or model test too expensive)

* → seakeeping: Potential flow

* heave/pitch good!

* **NOT** roll!

* **Solution**: Potential flow + semi empirical viscous roll damping

* Potential flow:
    * Strip theory (milli seconds)
    * Nonlinear 3D methods (hours)

## Background
The roll damping can be divided into various components :
$$B_{44} = B_F + B_E + B_L + B_W + B_{BK}$$

Viscous: $ B_{visc} = B_F + B_E + B_L + B_{BK} $ (Ikeda's method, Simplified Ikeda)

Inviscid: $ B_{invisc} = B_W $ (Potential flow)

$$B^{Ikeda} = B_{invisc}^{1D} + B_{visc}$$

$$B^{Motions} = B_{invisc}^{3D} + B_{visc}$$

### Problems with $B_W$ in Simplified Ikeda

```python
data = {
    'KVLCC2' : {
        'type':'tanker',
        'test data':True,
        'B_W': 'small',
        'bilge keel':False,
        'publish geom':True,
        'publish test':True,
    },
    'DTC' : {
        'type':'container',
        'test data':True,
        'B_W': '?',
        'bilge keel':True,
        'publish geom':True,
        'publish test':True,
    },
    'Wallenius' : {
        'type':'PCTC',
        'test data':True,
        'B_W': 'medium',
        'bilge keel':True,
        'publish geom':False,
        'publish test':True,
    },

}
test_cases = pd.DataFrame(data=data).transpose()
```

```python
def background_colorer(val):
    return 'background-color: %s' % get_color(val)

def text_colorer(val):
    return 'color: %s' % get_color(val)

def get_color(val):

    color = 'none'
    if isinstance(val, bool):
        if val:
            color = 'green'
        else:
            color = 'red'

    return color

```

## Possible test cases:

```python
test_cases.style.applymap(background_colorer).applymap(text_colorer)
```

        	        type	        test data	        B_W	        bilge keel	        publish geom	        publish test

                        KVLCC2
                        tanker
                        True
                        small
                        False
                        True
                        True

                        DTC
                        container
                        True
                        ?
                        ?
                        True
                        True

                        Wallenius
                        PCTC
                        True
                        medium
                        True
                        False
                        True

## Test case: KVLCC2
[04.3_KVLCC2_Ikedas_model_tests](../../notebooks/04.3_KVLCC2_Ikedas_model_tests.ipynb)

* B_W very small!
* B_E is dominating

* B_E decreased, B_L is now dominating
* B_W is a bit larger but still minor

# Abstract conclusion:
#### "It is found that potential improvements to the modelling of roll damping can be achieved by introducing FNPF analysis in the Ikeda’s method"

# Is this what we are aiming for?

## References