# Discrete Data and the Multinomial Distribution

### Preliminaries

- Goal
  - Simple Bayesian and maximum likelihood-based density estimation for discretely valued data sets
- Materials
  - Mandatory
    - These lecture notes
  - Optional
    - Bishop pp. 67-70, 74-76, 93-94

### Discrete Data: the 1-of-K Coding Scheme

- Consider a coin-tossing experiment with outcomes $x \in\{0,1\}$ (tail and head) and let $0\leq \mu \leq 1$ represent the probability of heads. This model can written as a **Bernoulli distribution**:
$$
p(x|\mu) = \mu^{x}(1-\mu)^{1-x}
$$
  - Note that the variable $x$ acts as a (binary) **selector** for the tail or head probabilities. Think of this as an 'if'-statement in programming.

- Now consider a $K$-sided coin (e.g., a six-faced _die_ (pl.: dice)). How should we encode outcomes?

- **Option 1**: $x \in \{1,2,\ldots,K\}$.
  - E.g., for $K=6$, if the die lands on the 3rd face $\,\Rightarrow x=3$.

- **Option 2**:  $x=(x_1,\ldots,x_K)^T$ with **binary selection variables**
$$
x_k = \begin{cases} 1 & \text{if die landed on $k$th face}\\
0 & \text{otherwise} \end{cases}
$$
  - E.g., for $K=6$, if the die lands on the 3rd face $\,\Rightarrow x=(0,0,1,0,0,0)^T$.
  - This coding scheme is called a **1-of-K** or **one-hot** coding scheme.

- It turns out that the one-hot coding scheme is mathematically more convenient!

- Consider a $K$-sided die. We use a one-hot coding scheme. Assume the probabilities $p(x_k=1) = \mu_k$ with  $\sum_k \mu_k  = 1$. The data generating distribution is then (note the similarity to the Bernoulli distribution)

$$
p(x|\mu) = \mu_1^{x_1} \mu_2^{x_2} \cdots \mu_k^{x_k}=\prod_k \mu_k^{x_k} \tag{B-2.26}
$$

- This generalized Bernoulli distribution is called the **categorical distribution** (or sometimes the 'multi-noulli' distribution).

### Bayesian Density Estimation for a Loaded Die

- Now let's proceed with Bayesian density estimation $p(x|\theta)$ for an observed data set $D=\{x_1,\ldots,x_N\}$  of $N$ IID rolls of a $K$-sided die, with

$$
x_{nk} = \begin{cases} 1 & \text{if the $n$th throw landed on $k$th face}\\
0 & \text{otherwise} \end{cases}
$$

##### Model specification

- The data generating PDF is
$$
p(D|\mu) = \prod_n \prod_k \mu_k^{x_{nk}} = \prod_k \mu_k^{\sum_n x_{nk}} = \prod_k \mu_k^{m_k} \tag{B-2.29}
$$
where $m_k= \sum_n x_{nk}$ is the total number of occurrences that we 'threw' $k$ eyes. Note that $\sum_k m_k = N$.
  - This distribution depends on the observations **only** through the quantities $\{m_k\}$.

- We need a prior for the parameters $\mu = (\mu_1,\mu_2,\ldots,\mu_K)$. In the [binary coin toss example](https://nbviewer.jupyter.org/github/bertdv/BMLIP/blob/master/lessons/notebooks/03-Bayesian-Machine-Learning.ipynb#beta-prior),
we used a [beta distribution](https://en.wikipedia.org/wiki/Beta_distribution) that was conjugate with the binomial and forced us to choose prior pseudo-counts.

- The generalization of the beta prior to the $K$ parameters $\{\mu_k\}$ is the [Dirichlet distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution):
$$
p(\mu|\alpha) = \mathrm{Dir}(\mu|\alpha) = \frac{\Gamma\left(\sum_k \alpha_k\right)}{\Gamma(\alpha_1)\cdots \Gamma(\alpha_K)} \prod_{k=1}^K \mu_k^{\alpha_k-1}
$$

##### Inference for $\{\mu_k\}$

- The posterior for  $\{\mu_k\}$ can be obtained through Bayes rule:

$$\begin{align*}
p(\mu|D,\alpha) &\propto p(D|\mu) \cdot p(\mu|\alpha) \\
  &\propto  \prod_k \mu_k^{m_k} \cdot \prod_k \mu_k^{\alpha_k-1} \\
  &= \prod_k \mu_k^{\alpha_k + m_k -1}\\
  &\propto \mathrm{Dir}\left(\mu\,|\,\alpha + m \right) \tag{B-2.41} \\
  &= \frac{\Gamma\left(\sum_k (\alpha_k + m_k) \right)}{\Gamma(\alpha_1+m_1) \Gamma(\alpha_2+m_2) \cdots \Gamma(\alpha_K + m_K)} \prod_{k=1}^K \mu_k^{\alpha_k + m_k -1}
\end{align*}$$
where $m = (m_1,m_2,\ldots,m_K)^T$.

- Again, we recognize the $\alpha_k$'s as prior pseudo-counts and the Dirichlet distribution shows to be a [conjugate prior](https://en.wikipedia.org/wiki/Conjugate_prior) to the categorical/multinomial:

$$\begin{align*}
\underbrace{\text{Dirichlet}}_{\text{posterior}} &\propto \underbrace{\text{categorical}}_{\text{likelihood}} \cdot \underbrace{\text{Dirichlet}}_{\text{prior}}
\end{align*}$$

- This is actually a generalization of the conjugate relation that we found for the binary coin toss:

$$\begin{align*}
\underbrace{\text{Beta}}_{\text{posterior}} &\propto \underbrace{\text{binomial}}_{\text{likelihood}} \cdot \underbrace{\text{Beta}}_{\text{prior}}
\end{align*}$$

##### Prediction of next toss for the loaded die

- Let's apply what we have learned about the loaded die to compute the probability that we throw the $k$th face at the next toss.

$$\begin{align*}
p(x_{\bullet,k}=1|D)  &= \int p(x_{\bullet,k}=1|\mu)\,p(\mu|D) \,\mathrm{d}\mu \\
  &= \int_0^1 \mu_k \times  \mathcal{Dir}(\mu|\,\alpha+m) \,\mathrm{d}\mu  \\
  &= \mathrm{E}\left[ \mu_k \right] \\
  &= \frac{m_k + \alpha_k }{ N+ \sum_k \alpha_k}
\end{align*}$$

- This result is simply a generalization of [**Laplace's rule of succession**](https://en.wikipedia.org/wiki/Rule_of_succession).

### Categorical, Multinomial and Related Distributions

- In the above derivation, we noticed that the data generating distribution for $N$ die tosses $D=\{x_1,\ldots,x_N\}$ only depends on the **data frequencies** $m_k$:
$$
p(D|\mu) = \prod_n \underbrace{\prod_k \mu_k^{x_{nk}}}_{\text{categorical dist.}} = \prod_k \mu_k^{\sum_n x_{nk}} = \prod_k \mu_k^{m_k} \tag{B-2.29}
$$

- A related distribution is the distribution over data frequency observations $D_m=\{m_1,\ldots,m_K\}$, which is called the **multinomial distribution**,
$$
p(D_m|\mu) =\frac{N!}{m_1! m_2!\ldots m_K!} \,\prod_k \mu_k^{m_k}\,.
$$

- Verify for yourself that ([Exercise](https://nbviewer.jupyter.org/github/bertdv/BMLIP/blob/master/lessons/exercises/Exercises.ipynb)):
  - the categorial distribution is a special case of the multinomial for $N=1$.
  - the Bernoulli is a special case of the categorial distribution for $K=2$.
  - the binomial is a special case of the multinomial for $K=2$.

### Maximum Likelihood Estimation for the Multinomial

##### Maximum likelihood as a special case of Bayesian estimation
- We can get the maximum likelihood estimate $\hat{\mu}_k$for $\mu_k$ based on $N$ throws of a $K$-sided die from the Bayesian framework by using a uniform prior for $\mu$ and taking the mode of the posterior for $\mu$:
$$\begin{align*}
\hat{\mu}_k &= \arg\max_{\mu_k} p(D|\mu) \\
&= \arg\max_{\mu_k} p(D|\mu)\cdot \mathrm{Uniform}(\mu) \\
&= \arg\max_{\mu_k} p(D|\mu) \cdot \left.\mathrm{Dir}(\mu|\alpha)\right|_{\alpha=(1,1,\ldots,1)} \\
&= \arg\max_{\mu_k} \left.p(\mu|D,\alpha)\right|_{\alpha=(1,1,\ldots,1)}  \\
&= \arg\max_{\mu_k} \left.\mathrm{Dir}\left( \mu | m + \alpha \right)\right|_{\alpha=(1,1,\ldots,1)} \\
&= \frac{m_k}{\sum_k m_k} = \frac{m_k}{N}
\end{align*}$$
where we used the fact that the [maximum of the Dirichlet distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution#Mode) $\mathrm{Dir}(\{\alpha_1,\ldots,\alpha_K\})$ is obtained at
$(\alpha_k-1)/(\sum_k\alpha_k - K)$.

##### Maximum likelihood estimation by optimizing a constrained log-likelihood
- Of course, we shouldn't have to go through the full Bayesian framework to get the maximum likelihood estimate. Alternatively, we can find the maximum of the likelihood directly.

- The log-likelihood for the multinomial distribution is given by
$$\begin{align*}
\mathrm{L}(\mu) &\triangleq \log p(D_m|\mu) \propto \log \prod_k \mu_k^{m_k} =  \sum_k m_k \log \mu_k
\end{align*}$$

- When doing ML estimation, we must obey the constraint $\sum_k \mu_k  = 1$, which can be accomplished by a Lagrange multiplier. The **augmented log-likelihood** with Lagrange multiplier is then

$$
\mathrm{L}^\prime(\mu) = \sum_k m_k \log \mu_k  + \lambda \cdot \left(1 - \sum_k \mu_k \right)
$$

- Set derivative to zero yields the **sample proportion** for $\mu_k$
$$\begin{equation*}
\nabla_{\mu_k}   \mathrm{L}^\prime = \frac{m_k }
{\hat\mu_k } - \lambda  \overset{!}{=} 0 \; \Rightarrow \; \hat\mu_k = \frac{m_k }{N}
\end{equation*}$$
where we get $\lambda$ from the constraint
$$\begin{equation*}
\sum_k \hat \mu_k = \sum_k \frac{m_k}
{\lambda} = \frac{N}{\lambda} \overset{!}{=}  1
\end{equation*}$$

### Recap Maximum Likelihood Estimation for Gaussian and Multinomial Distributions

Given $N$ IID observations $D=\{x_1,\dotsc,x_N\}$.

- For a **multivariate Gaussian** model $p(x_n|\theta) = \mathcal{N}(x_n|\mu,\Sigma)$, we obtain ML estimates

$$\begin{align}
\hat \mu &= \frac{1}{N} \sum_n x_n \tag{sample mean} \\
\hat \Sigma &= \frac{1}{N} \sum_n (x_n-\hat\mu)(x_n - \hat \mu)^T \tag{sample variance}
\end{align}$$

- For discrete outcomes modeled by a 1-of-K **categorical distribution** we find

$$\begin{align}
\hat\mu_k  = \frac{1}{N} \sum_n x_{nk} \quad \left(= \frac{m_k}{N} \right) \tag{sample proportion}
\end{align}$$

- Note the similarity for the means between discrete and continuous data.

```julia
open("../../styles/aipstyle.html") do f
    display("text/html", read(f, String))
end
```

<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=Nixie+One' rel='stylesheet' type='text/css'>

```julia

```


# Homework 9.2: MLE of microtubule catastrophe data (40 pts)

[Dataset download](https://s3.amazonaws.com/bebi103.caltech.edu/data/gardner_time_to_catastrophe_dic_tidy.csv)

Refresh yourself about the microtubule catastrophe data we explored in previous homeworks. We will again work with this data set here.

**a)** In their [paper](http://dx.doi.org/10.1016/j.cell.2011.10.037), Gardner, Zanic, and coworkers modeled microtubule catastrophe times as Gamma distributed. Perform a maximum likelihood estimate for the parameters of the Gamma distribution. Because you showed in a previous homework that there is little difference between labeled and unlabeled tubulin, you only need to work this out for the labeled tubulin now and in part (b). Be sure to include confidence intervals or a confidence region for your MLE and discuss the method you used to get the confidence intervals or confidence region.

```python
import numpy as np
import pandas as pd
import os
import scipy
import scipy.optimize
import scipy.stats as st
import tqdm
import math

rg = np.random.default_rng()

data_path = "../data"
fname = os.path.join(data_path, "gardner_time_to_catastrophe_dic_tidy.csv")
df = pd.read_csv(fname)
df
```

      Unnamed: 0
      time to catastrophe (s)
      labeled

      0
      0
      470.0
      True

      1
      1
      1415.0
      True

      2
      2
      130.0
      True

      3
      3
      280.0
      True

      4
      4
      550.0
      True

      ...
      ...
      ...
      ...

      301
      301
      180.0
      False

      302
      302
      145.0
      False

      303
      303
      745.0
      False

      304
      304
      390.0
      False

      305
      305
      1410.0
      False

306 rows × 3 columns

We made a separate dataframe for the data labeled "True", and obtained the time stamps of those labeled "True".

```python
df_true = df[df["labeled"] == True]
true_time = df_true["time to catastrophe (s)"].values
```

Now, we defined the bootstrap to return us alpha and beta.

```python
def bootstrap_param(data, size=1):
    """Parametric bootstrap replicates of parameters of
    Normal distribution."""
    bs_mean = np.empty(size)
    bs_sd = np.empty(size)
    bs_alpha = np.empty(size)
    bs_beta = np.empty(size)

    for i in range(size):
        bs_sample = np.random.choice(data, size=len(data))
        bs_mean[i] = np.mean(bs_sample)
        bs_sd[i] = np.std(bs_sample)
        # Since the mean of the gamma distribution is alpha/beta, and variance is alpha/beta^2, mean/var = beta.
        bs_beta[i] = bs_mean[i]/(bs_sd[i]**2)
        # With the beta value, we solve for alpha by just saying beta * mean = alpha.
        bs_alpha[i] = bs_beta[i] * bs_mean[i]

    return bs_alpha, bs_beta
```

We run the bootstrap for our "True" time stamps.

```python
bs_alpha_true, bs_beta_true = bootstrap_param(
    true_time, size=100000
)

bs_alpha_true
```

    array([2.42338039, 2.50186784, 2.05350071, ..., 2.35151568, 2.06215175,
           2.1495506 ])

We can obtain the alpha and beta values as an expression of the mean and variance.

```python
mean_true = np.mean(true_time)
var_true = np.var(true_time)

# Since the mean of the gamma distribution is alpha/beta, and variance is alpha/beta^2, mean/var = beta.
beta_true = mean_true / var_true

# With the beta value, we solve for alpha by just saying beta * mean = alpha.
alpha_true = beta_true * mean_true
print(alpha_true, beta_true)
```

    2.2239376038081673 0.005046250504393196

```python
print('α: {:.4f} | 95% conf int α: {}'.format(np.mean(bs_alpha_true),str(np.percentile(bs_alpha_true, [2.5, 97.5]))))
print('β2: {:.4f} | 95% conf int β2: {}'.format(np.mean(bs_beta_true),str(np.percentile(bs_beta_true, [2.5, 97.5]))))
```

    α: 2.2670 | 95% conf int α: [1.82320381 2.82034726]
    β2: 0.0052 | 95% conf int β2: [0.00406756 0.00656671]

As the values of the parameters alpha and beta do not overlap, and fall inside the 95% confidence interval for the labeled microtubules, we can reasonably assume that the MLE value we have obtained for the Gamma distribution is reasonable.

**b)** Obtain a maximum likelihood estimate for the parameters $\beta_1$ and $\beta_2$ from the model you derived in homework 5.2. As a reminder, you derived that the PDF for microtubule catastrophe times is

$$\begin{align}
f(t;\beta_1, \beta_2) = \frac{\beta_1\beta_2}{\beta_2 - \beta_1}\left(\mathrm{e}^{-\beta_1 t} - \mathrm{e}^{-\beta_2 t}\right).
\end{align}$$

Again, include confidence intervals. **Be careful**; this is a *very* tricky calculation. It is possible to analytically compute the MLE. If you choose to do it numerically, you need to think about what happens when $\beta_1 \approx \beta_2$. You may also need to think about how you will handle the [log of sums of exponentials](https://en.wikipedia.org/wiki/LogSumExp).

$$\begin{align}\tag{1}
f(t;\beta_1, \beta_2) = \frac{\beta_1 \beta_2}{\beta_2 - \beta_1}\left(\mathrm{e}^{-\beta_1 t} - \mathrm{e}^{-\beta_2 t}\right)
\end{align}$$

$$\begin{align}\tag{2}
L(t;\beta_1, \beta_2)= \ln{f(t;\beta_1, \beta_2)} = \ln{\left(\beta_1\beta_2\right)} - \ln{\left(\beta_2-\beta_1\right)}+\ln{\left(\mathrm{e}^{-\beta_1 t} - \mathrm{e}^{-\beta_2 t}\right)}
\end{align}$$

Our third term is the log of sum of exponentials, and we can change this term to:

$$\begin{align}\tag{3}
y = \ln{\left(\mathrm{e}^{-\beta_1 t} - \mathrm{e}^{-\beta_2 t}\right)}\\
\Leftrightarrow \mathrm{e}^{y} = \sum_{i=1}^{n}\mathrm{e}^{x_{i}}\\
\Leftrightarrow \mathrm{e}^{\beta_1 t}\mathrm{e}^{y} = \mathrm{e}^{-\beta_1 t}\sum_{i=1}^{n}\mathrm{e}^{x_{i}}\\
\Leftrightarrow \mathrm{e}^{y + \beta_1 t} = \sum_{i=1}^{n}\mathrm{e}^{x_{i} + \beta_1 t} \\
\Leftrightarrow y + \beta_1 t = \ln{\sum_{i=1}^{n}\mathrm{e}^{x_{i} + \beta_1 t}}\\
\Leftrightarrow y = -\beta_1 t + \ln{\sum_{i=1}^{n}\mathrm{e}^{x_{i} + \beta_1 t}}\\
\Leftrightarrow y = -\beta_1 t + \ln{\mathrm{e}^{\left(\beta_1-\beta_2\right)t}}
\end{align}$$

As such, our initial log equation will become as follows:
$$\begin{align}\tag{4}
\ln{f(t;\beta_1, \beta_2)} = \ln{\left(\beta_1\beta_2\right)} - \ln{\left(\beta_2-\beta_1\right)}+\ln{\mathrm{e}^{\left(\beta_1-\beta_2\right)t}}\\
= \ln{\beta_1} + \ln{\beta_2} - \ln{(\beta_2-\beta_1)} -\beta_1 t - ln{\mathrm{e}^{\left(\beta_1-\beta_2\right)t}} \\
= \ln{\beta_1} + \ln{\beta_2} - \ln{(\beta_2-\beta_1)} -\beta_1 t - \left(\beta_1-\beta_2\right)t
\end{align}$$

Taking the partial derivatives of each terms gives:
$$\begin{align}\tag{5}
\frac{\partial L}{\partial \beta_1} = \frac{1}{\beta_1} + 0 - \left(\frac{1}{\beta_2-\beta_1}\left(0-1\right)\right) - t - t + 0 \\
= \frac{1}{\beta_1} + \frac{1}{\beta_2-\beta_1} - 2t
\end{align}$$
$$\begin{align}\tag{6}
\frac{\partial L}{\partial \beta_2} = 0 + \frac{1}{\beta_2}-\left(\frac{1}{\beta_2-\beta_1}\left(1-0\right)\right)-0-0+t\\
= \frac{1}{\beta_2} - \frac{1}{\beta_2-\beta_1} + t
\end{align}$$
$$\begin{align}\tag{7}
\frac{\partial L}{\partial t} = 0 + 0 - 0 -\beta_1 -\beta_1 + \beta_2\\
= \beta_2 - 2\beta_1\\
\end{align}$$

We can evaluate the turning point by setting $\frac{\partial L}{\partial t} = 0$, which gives the expression:
$$\begin{align}\tag{8}
\beta_1 = \frac{1}{2}\beta_2 \Leftrightarrow 2\beta_1 = \beta_2
\end{align}$$

Whereas we can rearrange our initial equation to express $t$ in terms of $\beta_1$ and $\beta_2$:
$$\begin{align}\tag{9}
t = \frac{\ln{\beta_1} + \ln{\beta_2} - \ln{\left(\beta_2-\beta_1\right)}}{2\beta_1-\beta_2}
\end{align}$$

As such, the turning point are points where-in:
$$\begin{align}\tag{10}
\beta_1 = \frac{1}{2} \wedge \beta_2 = 1
\end{align}$$
$$\begin{align}\tag{11}
2\beta_1 - \beta_2 \neq 0 \wedge t = \frac{\ln{\beta_1} + \ln{\beta_2} - \ln{\left(\beta_2-\beta_1\right)}}{2\beta_1-\beta_2}
\end{align}$$

Using a numerical approach, we defined the log likelihood function using the gamma distribution, as well as the corresponding MLE.

```python
def log_like_iid_gamma(params, n):

    beta1, beta2 = params
    mlesum = 0

    for t in n:
        mlesum += np.log(beta1)+ np.log(beta2) - np.log((beta2-beta1)) + np.log(np.exp(-beta1*t) - np.exp(-beta2*t))
        if (math.isnan(mlesum)):
            return -np.inf
    return mlesum

def mle_iid_custom(n):
    """Perform maximum likelihood estimates for parameters for i.i.d."""
#     with warnings.catch_warnings():
#         warnings.simplefilter("ignore")

    res = scipy.optimize.minimize(
        fun=lambda params, n: -log_like_iid_gamma(params, n),
        x0=np.array([0.0035, 0.0045]),
        args=(n,),
        method='Powell',
    )
    if res.success:
        return res.x
    else:
        raise RuntimeError('Convergence failed with message', res.message)
```

Now we run our MLE to obtain our beta values.

```python
mle_labeled = mle_iid_custom(df_true["time to catastrophe (s)"].values)
```

    C:\Users\sr_pa\AppData\Local\Temp/ipykernel_30112/519852826.py:7: RuntimeWarning: invalid value encountered in log
      mlesum += np.log(beta1)+ np.log(beta2) - np.log((beta2-beta1)) + np.log(np.exp(-beta1*t) - np.exp(-beta2*t))
    C:\Users\sr_pa\AppData\Local\Temp/ipykernel_30112/519852826.py:7: RuntimeWarning: overflow encountered in exp
      mlesum += np.log(beta1)+ np.log(beta2) - np.log((beta2-beta1)) + np.log(np.exp(-beta1*t) - np.exp(-beta2*t))
    D:\Anaconda3\envs\py38\lib\site-packages\scipy\optimize\optimize.py:2215: RuntimeWarning: invalid value encountered in double_scalars
      tmp2 = (x - v) * (fx - fw)
    D:\Anaconda3\envs\py38\lib\site-packages\scipy\optimize\optimize.py:2216: RuntimeWarning: invalid value encountered in double_scalars
      p = (x - v) * tmp2 - (x - w) * tmp1

```python
print('β1: ' + str(mle_labeled[0]))
print('β2: ' + str(mle_labeled[1]))
```

    β1: 0.004491379510103659
    β2: 0.004582668616663547

We can define a function to draw our bootstrap replicates, as such:

```python
def draw_bs_sample(data):
    """Draw a bootstrap sample from a 1D data set."""
    return rg.choice(data, size=len(data))

def draw_bs_reps_mle(mle_fun, data, args=(), size=1, progress_bar=False):

    if progress_bar:
        iterator = tqdm.tqdm(range(size))
    else:
        iterator = range(size)

    return np.array([mle_fun(draw_bs_sample(data), *args) for _ in iterator])
```

We run our bootstrap to obtain the confidence intervals for our beta values.

```python
bs_reps = draw_bs_reps_mle(
    mle_iid_custom, df_true["time to catastrophe (s)"].values, size=100,
)
```

    C:\Users\sr_pa\AppData\Local\Temp/ipykernel_30112/519852826.py:7: RuntimeWarning: invalid value encountered in log
      mlesum += np.log(beta1)+ np.log(beta2) - np.log((beta2-beta1)) + np.log(np.exp(-beta1*t) - np.exp(-beta2*t))
    C:\Users\sr_pa\AppData\Local\Temp/ipykernel_30112/519852826.py:7: RuntimeWarning: overflow encountered in exp
      mlesum += np.log(beta1)+ np.log(beta2) - np.log((beta2-beta1)) + np.log(np.exp(-beta1*t) - np.exp(-beta2*t))

```python
conf_int = np.percentile(bs_reps, [2.5, 97.5], axis=0)
conf_int = np.transpose(conf_int)

print('β1: {:.4f} | 95% conf int β1: {}'.format(np.mean(bs_reps[:,0]),str(conf_int[0])))
print('β2: {:.4f} | 95% conf int β2: {}'.format(np.mean(bs_reps[:,1]),str(conf_int[1])))
```

    β1: 0.0044 | 95% conf int β1: [0.00396105 0.0045624 ]
    β2: 0.0047 | 95% conf int β2: [0.00440738 0.00547511]

Comparing the parameter estimates from both the analytical and numerical approach gives very different results; in our analytical approach we were not able to find a single point for the MLE, whereas using bootstrapping methods we can obtain precise beta values with a small confidence interval range.

## Attributes:

Done altogether as a group


```python
import holoviews as hv
hv.extension('bokeh')
hv.opts.defaults(hv.opts.Curve(width=500),
                 hv.opts.Scatter(width=500, size=4),
                 hv.opts.Histogram(width=500),
                 hv.opts.Slope(color='k', alpha=0.5, line_dash='dashed'),
                 hv.opts.HLine(color='k', alpha=0.5, line_dash='dashed'))
```

```python
import numpy as np
import scipy.stats
import pandas as pd
import sklearn
```

# Linear models and Basis functions

In previous lessons we learned how to fit lines and hyperplanes to data. But the most general form of a linear model is

$$
g_\theta(x_i) = \sum_{k=0}^K \theta_k \phi_k(x_i),
$$

where $\phi_k: \mathbb{R}^D \to \mathbb{R}$ is a set of basis functions. Note that $K$ and $D$ are not necessarily related

In this lesson we will

- review some examples of basis functions
- learn how to perform linear regression with basis functions using Python
- learn how to avoid overfitting the data using regularization techniques

## Basis functions

### Polynomials

For a unidimensional variable $x_i \in \mathbb{R}, i=1,\ldots,N$, a general polynomial basis is defined

$$
\phi_k(x_i) = x_i^k
$$

which yields a K-degree polynomial model

$$
g_\theta(x_i) = \sum_{k=0}^K \theta_k x_i^k = \theta_0 + \theta_1 x_i + \theta_2 x_i^2 + \ldots + \theta_K x_i^K \quad \forall i
$$

### Trigonometric

If we want to model periodic behavior a trigonometric basis is a suitable choice. For a unidimensional variable $x_i \in \mathbb{R}, i=1,\ldots,N$, a trigonometric basis with period $P=1/f_0$ is

$$
\phi_k(x_i) = \begin{cases} 1 & k=0 \\ \cos(2\pi k f_0 x_i) & k \in [1, K] \\  \sin(2\pi k f_0 x_i) & k \in [K+1, 2K] \end{cases}
$$

which yields a trigonometric model with $2K+1$ parameters

$$
g_\theta(x_i) = \theta_0 + \sum_{k=1}^K \theta_k \cos(2\pi k f_0 x_i) + \sum_{k=1}^K \theta_{k+K} \sin(2\pi k f_0 x_i)
$$

### Interactions between variables

We can create a basis that models linear and non-linear interactions between our independent variables

For example if we have a bidimensional variable $x_i = (w_{i}, v_{i}), i=1,\ldots,N$ a model with interactions up to the second degree would be

$$
g_\theta(x_i) = \theta_0 + \theta_1 w_i + \theta_2 v_i + \theta_3 w_i^2 + \theta_4 v_i w_i + + \theta_5 v_i^2
$$

## Regression with basis functions using `scikit`-learn

We can create a polynomial basis from our variables using

```python
PolynomialFeatures(degree=2, # Degree of the polinomial
                   interaction_only=False, # Return only products between features
                   include_bias=True, # Include the intercept (constant) term
                   ...
                  )
```

The `fit_transform` function of this object receives the data and returns the transformed data. For example if our dataset has two independent variables $x=[w, v]$, then `PolynomialFeatures(degree=2)` would return $[1, w, v, w^2, wv, v^2]$.

The `sklearn` library also has [`LinearRegression`](https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.LinearRegression.html) which is a wrapper for `linalg.lstsq`. The advantage of using `LinearRegression` instead of `linalg.lstsq` is that we can use `sklearn.pipeline` to create the polynomial features and train the model together

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

non_linear_regressor = Pipeline([('features', PolynomialFeatures(2)),
                                 ('regressor', LinearRegression())])
non_linear_regressor
```

Using the data from the previous lesson let's train a polynomial regressor to predict `consumption` as a function of  `temperature`

```python
df = pd.read_csv('data/ice_cream.csv', header=0, index_col=0)
df.columns = ['Consumption', 'Income', 'Price', 'Temperature']
Y = df["Consumption"].values
X = df["Temperature"].values

model = non_linear_regressor.fit(X.reshape(-1, 1), Y)

# The coef_ attribute returns the value of the fitted parameters (theta)
display(model['regressor'].coef_)
```

And we can plot the result of our polynomial regression

```python
hatx = np.linspace(10, 90, num=100)
haty = model.predict(hatx.reshape(-1, 1))

hv.Overlay([hv.Curve((hatx, haty), 'Temperature', 'Consumption', label='model'),
            hv.Scatter((X, Y), label='data').opts(color='k')]).opts(legend_position='top_left')
```

We have move from lines to parabolas!

What happens if we change the `degree` of our polynomial features?

```python
hatx = np.linspace(np.amin(X), np.amax(X), num=1000)
haty = {}
for degree in [1, 2, 3, 4, 5, 10, 20, 30]:
    non_linear_regressor = Pipeline([('features', PolynomialFeatures(degree)),
                                     ('regressor', LinearRegression())])
    model = non_linear_regressor.fit(X.reshape(-1, 1), Y)
    haty[degree] = model.predict(hatx.reshape(-1, 1))
```

```python
hMap = hv.HoloMap(kdims='degree')
for degree, haty_ in haty.items():
    p_model = hv.Curve((hatx, haty_), 'Temperature', 'Consumption', label='model')
    p_data = hv.Scatter((X, Y), label='data').opts(color='k')
    hMap[degree] = hv.Overlay([p_model, p_data]).opts(legend_position='top_left')

hMap
```

:::{danger}

As `degree` grows we start overfitting the data more and more

:::

Let's define what is overfitting and how to combat it using regularization

## Overfitting and regularization

In the previous example `degree` represents the complexity of the model. In general, more complex models give more flexibility to fit the data

But too much complexity causes **overfitting**

- the model fits the noise
- we can't extract the underlying behavior
- the model **does not generalize** to new data

Ways to avoid overfitting

- Using low complexity models
- Set complexity using cross-validation
- Regularization

In what follows we will focus on the latter

### The Bias-Variance trade-off

Let's assume that our data can be modeled as a "true model" plus gaussian noise

$$
y = f(x) + \varepsilon
$$

and that we use a linear model to find $f(x) = \sum_k \theta_k \phi_k(x)$

We can measure the quality of our model with the Mean Square Error (MSE)

$$
\begin{align}
\mathbb{E}[(y - \hat y)^2] &= \mathbb{E}[y^2 -2 y \hat y +\hat y^2] \nonumber \\
&= \mathbb{E}[(f+\varepsilon)^2 -2 (f+\varepsilon) \hat y +\hat y^2] \nonumber \\
&= \mathbb{E}[(f^2 +2 f \varepsilon + \varepsilon^2 -2 (f+\varepsilon) \hat y +\hat y^2] \nonumber \\
&= \mathbb{E}[\varepsilon^2] + f^2  -2 f \mathbb{E}[\hat y]  +\mathbb{E}[\hat y^2]  \pm \mathbb{E}[\hat y]^2  \nonumber \\
&= \mathbb{E}[\varepsilon^2] + (f - \mathbb{E}[\hat y])^2  +\mathbb{E}[(\hat y - \mathbb{E}[\hat y])^2]  \nonumber \\
&= \sigma^2 + (f - \mathbb{E}[\hat y])^2  + \text{Var}[\hat y]  \nonumber
\end{align}
$$

:::{important}

The MSE can be decomposed as the irreducible error (data noise) + the squared bias of the estimator + the variance of the estimator

:::

The MSE can be small if either the bias or the variance are small. More complex models tend to have have lower bias and higher variance

The Gauss-Markov theorem says that OLS has the minimum variance among the unbiased estimator. But zero-bias models are not necessarily good (overfit).

If we are overfitting the data we may want to trade variance for bias. This can be achieved by penalizing the complexity of the model: this is **regularization**.

### Bayesian (MAP) Least Squares and Ridge regression

If we assume a Gaussian likelihood and a Gaussian prior we can write the log joint as

$$
\begin{align}
\log p({x}, \theta) &= \log \prod_{i=1}^N \mathcal{N}(y_i | f_\theta(x_i), \sigma^2) + \log \prod_{j=1}^M \mathcal{N}(\theta_j | 0, \sigma_0^2) \nonumber \\
&= -\frac{N}{2} \log(2\pi \sigma^2) - \frac{1}{2\sigma^2} (Y-X\theta)^T (Y - X\theta) -\frac{M}{2} \log(2\pi \sigma_0^2) - \frac{1}{2\sigma_0^2} \|\theta\|^2 \nonumber
\end{align}
$$

The MAP estimator of $\theta$ is given by

$$
\begin{align}
\hat \theta &= \text{arg}\max_\theta \log p({x}, \theta) \nonumber  \\
&= \text{arg}\max_\theta  - \frac{1}{2\sigma^2} (Y-X\theta)^T (Y - X\theta) - \frac{1}{2\sigma_0^2} \|\theta\|^2 \nonumber \\
&= \text{arg}\min_\theta  (Y-X\theta)^T (Y - X\theta) + \lambda \|\theta\|^2 \nonumber
\end{align}
$$

where $\lambda = \frac{\sigma^2}{\sigma_0^2}$

The solution is obtained by taking the derivative on $\theta$

$$
\frac{d}{d\theta} (Y-X\theta)^T (Y - X\theta) + \lambda \|\theta\|^2  = -X^T (Y - X\theta) + \lambda \theta = 0
$$

and finally

$$
\hat \theta = (X^T X + \lambda I)^{-1} X^T Y
$$

which is known as **Ridge regression** and **Tikhonov regularization**

:::{important}

The gaussian prior is equivalent to a restriction on the $L_2$ norm of $\theta$. Adding this prior forces the solution to be smooth

:::

:::{note}

Adding different priors yield different regularization effects. For example a Laplacian prior yields a $L_1$ norm  on $\theta$ which forces the solution to be sparse

:::

In general, regularizing a model consists of adding a "penalty term" or restriction to the cost/loss function. Typically the additional term will penalize overly complex models. Because of this regularization can help to avoid overfitting with complex models or ill-posed problems, e.g. when we have more parameters than data samples

But there is no free lunch. We now have the additional task of choosing $\lambda$

- Cross-validation: Minimize validation error
- L-curve: Plot $ \log (Y-X\theta)^T (Y - X\theta)$ vs $ \log \|\theta\|^2$ and find the elbow

### Example: Ridge Regression in Python

The `sklearn` library provides

```python
sklearn.linear_modelRidge(alpha=1.0, # The regularization parameter (lambda)
                          fit_intercept=True, # Whether to include the intercept (constant)
                          normalize=False, # Subtract mean and divide std from the data
                          ...
)
```

We will create a pipeline for the polynomial features plus the ridge regression and explore the influence of the regularization parameter

```python
from sklearn.linear_model import Ridge
hatx = np.linspace(np.amin(X), np.amax(X), num=1000)
haty = {}
for degree in [1, 2, 3, 5, 20]:
    for lamb in [0.0, 1e-3, 1e-1,  10]:
        non_linear_regressor = Pipeline([('features', PolynomialFeatures(degree)),
                                         ('regressor', Ridge(normalize=True, alpha=lamb))])
        model = non_linear_regressor.fit(X.reshape(-1, 1), Y)
        haty[degree, lamb] = model.predict(hatx.reshape(-1, 1))
```

```python
hMap = hv.HoloMap(kdims=['degree', 'lambda'])
for (degree, lamb), haty_ in haty.items():
    p_model = hv.Curve((hatx, haty_), 'Temperature', 'Consumption', label='model')
    p_data = hv.Scatter((X, Y), label='data').opts(color='k')
    hMap[degree, lamb] = hv.Overlay([p_model, p_data]).opts(legend_position='top_left')

hMap
```

Observations:

- Regularization decreases overfitting in complex models
- But an excessive penalization may induce a trivial solution, for example a straigh horizontal line

## Extra topics

### A note on independence and correlation

Independence implies uncorrelatedness, but the reverse is not true. Two variables can have zero correlation but still be dependent. Also remember that linear regression (correlation) is only sensitive to linear relationships

If we are interested in testing independence we could use:

$$
p(x,y) = p(x)p(y)
$$

Several methods are based on this, for example Shannon's **Mutual Information**

$$
I(X,Y) = \int \int f_{XY}(x,y) \log \frac{f_{XY}(x,y)}{f_{X}(x) f_Y(y)} dx dy
$$

and the [Correlation distance](https://arxiv.org/pdf/0803.4101.pdf)

$$
R(X,Y) = \int \int |f_{XY}(x,y)  - f_{X}(x) f_Y(y)| dx dy
$$

Although these methods require that we estimate the joint and the marginals (KDE, Histogram, Parametric). For categorical variables we can use the **chi square test**

### Related topics

- (Hastie 3.4 and 3.8) L1 regularization and Least Absolute Shrinkage and Selection Operator (LASSO)
- Robust regression: Least absolute regression and M-estimators for data with outliers (non-Gaussian)
- (Hastie 6 & Bishop 6) Kernel (non-parametric) regression

Some of these topics can be found at [Huijse, Regresión](https://docs.google.com/presentation/d/1UUpK4zSdzRcS79V7_wU9nXe-sR7qYLEWhbmid-Rfp1k/edit#slide=id.g28044c0f85_0_34)

I also suggests the following lecture by Judea Pearl on **causality**

```python
YouTubeVideo('ZaPV1OSEpHw')
```


## Questionário 61 (Q61)

Orientações:

- Registre suas respostas no questionário de mesmo nome no SIGAA.
- O tempo de registro das respostas no questionário será de 10 minutos. Portanto, resolva primeiro as questões e depois registre-as.
- Haverá apenas 1 (uma) tentativa de resposta.
- Submeta seu arquivo-fonte (utilizado para resolver as questões) em formato _.ipynb_ pelo SIGAA anexando-o à Tarefa denominada "Envio de arquivo" correspondente ao questionário.

*Nota:* o arquivo-fonte será utilizado apenas como prova de execução da tarefa. Nenhuma avaliação será feita quanto ao estilo de programação.

```python
import sympy as sym
from sympy import Symbol, pprint
import numpy as np
import matplotlib.pyplot as plt
```

**Questão 1.** Observe a figura abaixo e julgue os itens a seguir.

```{figure} ../figs/q/q61.png
---
width: 300px
name: convex
---
```

i) existe uma função convexa entre as quatro plotadas.

ii) uma entre as funções plotadas possui convexidade parcial.

iii) duas entre as funções plotadas não são convexas.

Assinale a alternativa correta.

A. São corretos i) e ii), apenas.

B. Apenas i) é correto.

C. São corretos i) e iii), apenas.

D. n.d.a.

```python
plt.figure(figsize=(14,4))
plt.subplot(141)
x1 = np.linspace(-10, 10, 100)
plt.plot(np.sin(x1),c='r')
plt.xticks([]); plt.yticks([]);
plt.title('(a)')

plt.subplot(142)
x2 = np.linspace(-2, 2, 100)
plt.plot(x2, np.exp(x2)*10*np.sin(6*x2))
plt.xticks([]); plt.yticks([]);
plt.title('(b)')

plt.subplot(143)
x3 = np.arange(-100, 100, 1)
plt.plot(x3, x3**2, c='orange')
plt.xticks([]); plt.yticks([]);
plt.title('(c)')

plt.subplot(144)
x4 = np.arange(-100, 0, 1)
plt.plot(x4, x4**3,c='m')
plt.xticks([]); plt.yticks([]);
plt.title('(d)')

plt.show()
```

## Gabarito
Alternativa **A**

**Questão 2.** A função a seguir simula a curva do _potencial de ação_ de uma membrana:

$$P(x) = \dfrac{1.0}{(x - 0.5)^2 + 0.01} - \dfrac{1.0}{(x - 0.8)^2 + 0.04} - 70.$$

Use computação simbólica para calcular uma aproximação para $P'(x=0)$ e assinale a alternativa correta.

A. -67.62

B. 0.25

C. 11.33

D. 0.00

Nota: Use `sympy.subs(x,x0)`.

```python
x1 = np.random.normal(0,1,10)
```

```python
#Gráfico da função
plt.plot((1.0/(x1-0.5)**2)-1.0/((x1-0.8)**2 +0.04) - 70 , label='$P(x)$', c='g');
plt.legend()
plt.title('Potencial de ação de uma membrana')
plt.show()
#plt.savefig("../figs/q/q61-2.png'")
```

## Gabarito
Alternativa **C**

```python
x = sym.symbols('x')
p = (1.0/((x-0.5)**2+0.01))-1.0/((x-0.8)**2 +0.04) - 70
```

```python
p
```

$\displaystyle -70 + \frac{1.0}{\left(x - 0.5\right)^{2} + 0.01} - \frac{1.0}{\left(x - 0.8\right)^{2} + 0.04}$

```python
dp = sym.diff(p,x)
dp.subs(x,0)
```

$\displaystyle 11.3326917958273$

**Questão 3.** Considere a função

$$f(x) =  - \dfrac{1}{e^x \text{sen}(6x)},$$

definida no domínio $[-0.5,-0.1]$. Assinale a alternativa correta:

A. $f(x)$ não é convexa e $f'(x) = -\frac{e^{x}}{\text{sen}{\left(6 x \right)}} + \frac{6 e^{- x} \cos{\left(6 x \right)}}{\text{sen}^{2}{\left(6 x \right)}}$

B. $f(x)$ é convexa e $f'(x) = \frac{e^{- x}}{\text{sen}{\left(6 x \right)}} + \frac{6 e^{- x} \cos{\left(6 x \right)}}{\text{sen}^{2}{\left(6 x \right)}}$

C. $f(x)$ não é convexa e $f'(x) = \frac{e^{x}}{\text{sen}{\left(6 x \right)}} + \frac{6 e^{- x} \cos{\left(6 x \right)}}{\text{sen}^{2}{\left(6 x \right)}}$

D. $f(x)$ é convexa e $f'(x) = -\frac{e^{- x}}{\text{sen}{\left(6 x \right)}} + \frac{6 e^{- x} \cos{\left(6 x \right)}}{\text{sen}^{2}{\left(6 x \right)}}$

## Gabarito
Alternativa **B**.

```python
# domínio
a,b = -0.5,-0.1
x = sym.symbols('x')
c = 6

# função e valores
f = -1/(sym.exp(x)*sym.sin(c*x))

df = f.diff(x)

dom = np.linspace(a,b)
plt.plot(dom, -1/(np.exp(dom)*np.sin(c*dom)));
```

```python

```


## Getting Started

In the following we like to introduce *pymoo* by presenting an example optimization scenario. This guide goes through the most important steps to get started with our framework. First, we present an example optimization problem to be solved using *pymoo*. Second, we show how to formulate the optimization problem in our framework and how to instantiate an algorithm object to be used for optimization. Then, a termination criterion for the algorithm is defined and the optimization method is called. Finally, we quickly show a possible post-processing step which analyzes the optimization run performance.

### Multi-Objective Optimization

In general, multi-objective optimization has several objective functions with subject to inequality and equality constraints to optimize <cite data-cite="multi_objective_book"></cite>. The goal is to find a set of solutions that do not have any constraint violation and are as good as possible regarding all its objectives values. The problem definition in its general form is given by:

\begin{align}
\begin{split}
\min \quad& f_{m}(x) \quad \quad \quad \quad m = 1,..,M  \\[4pt]
\text{s.t.}   \quad& g_{j}(x) \leq 0  \quad \; \; \,  \quad j = 1,..,J \\[2pt]
\quad& h_{k}(x) = 0        \quad  \; \; \quad k = 1,..,K \\[4pt]
\quad& x_{i}^{L} \leq x_{i} \leq x_{i}^{U}  \quad i = 1,..,N \\[2pt]
\end{split}
\end{align}

The formulation above defines a multi-objective optimization problem with $N$ variables, $M$ objectives, $J$ inequality and $K$ equality constraints. Moreover, for each variable $x_i$ lower and upper variable boundaries ($x_i^L$ and $x_i^U$) are defined.

### Example Optimization Problem

In the following, we investigate exemplarily a bi-objective optimization with two constraints.
The selection of a suitable optimization problem was made based on having enough complexity for the purpose of demonstration, but not being too difficult to lose track of the overall idea. Its definition is given by:

\begin{align}
\begin{split}
\min \;\; & f_1(x) = (x_1^2 + x_2^2) \\
\max \;\; & f_2(x) = -(x_1-1)^2 - x_2^2 \\[1mm]
\text{s.t.} \;\; & g_1(x) = 2 \, (x_1 - 0.1) \, (x_1 - 0.9) \leq 0\\
& g_2(x) = 20 \, (x_1 - 0.4) \, (x_1 - 0.6) \geq 0\\[1mm]
& -2 \leq x_1 \leq 2 \\
& -2 \leq x_2 \leq 2
\end{split}
\end{align}

 It consists of two objectives ($M=2$) where $f_1(x)$ is minimized and $f_2(x)$ maximized. The optimization is with subject to two inequality constraints ($J=2$) where $g_1(x)$ is formulated as a less than and $g_2(x)$ as a greater than constraint. The problem is defined with respect to two variables ($N=2$), $x_1$ and $x_2$, which both are in the range $[-2,2]$. The problem does not contain any equality constraints ($K=0$).

```python
import numpy as np

X1, X2 = np.meshgrid(np.linspace(-2, 2, 500), np.linspace(-2, 2, 500))

F1 = X1**2 + X2**2
F2 = (X1-1)**2 + X2**2
G = X1**2 - X1 + 3/16

G1 = 2 * (X1[0] - 0.1) * (X1[0] - 0.9)
G2 = 20 * (X1[0] - 0.4) * (X1[0] - 0.6)

import matplotlib.pyplot as plt
plt.rc('font', family='serif')

levels = [0.02, 0.1, 0.25, 0.5, 0.8]
plt.figure(figsize=(7, 5))
CS = plt.contour(X1, X2, F1, levels, colors='black', alpha=0.5)
CS.collections[0].set_label("$f_1(x)$")

CS = plt.contour(X1, X2, F2, levels, linestyles="dashed", colors='black', alpha=0.5)
CS.collections[0].set_label("$f_2(x)$")

plt.plot(X1[0], G1, linewidth=2.0, color="green", linestyle='dotted')
plt.plot(X1[0][G1<0], G1[G1<0], label="$g_1(x)$", linewidth=2.0, color="green")

plt.plot(X1[0], G2, linewidth=2.0, color="blue", linestyle='dotted')
plt.plot(X1[0][X1[0]>0.6], G2[X1[0]>0.6], label="$g_2(x)$",linewidth=2.0, color="blue")
plt.plot(X1[0][X1[0]<0.4], G2[X1[0]<0.4], linewidth=2.0, color="blue")

plt.plot(np.linspace(0.1,0.4,100), np.zeros(100),linewidth=3.0, color="orange")
plt.plot(np.linspace(0.6,0.9,100), np.zeros(100),linewidth=3.0, color="orange")

plt.xlim(-0.5, 1.5)
plt.ylim(-0.5, 1)
plt.xlabel("$x_1$")
plt.ylabel("$x_2$")

plt.legend(loc='upper center', bbox_to_anchor=(0.5, 1.12),
          ncol=4, fancybox=True, shadow=False)

plt.tight_layout()
plt.show()
```

The figure above shows the contours of the problem. The contour lines of the objective function $f_1(x)$ is represented by a solid and $f_2(x)$ by a dashed line. The constraints $g_1(x)$ and $g_2(x)$ are parabolas which intersect the $x_1$-axis at $(0.1, 0.9)$ and $(0.4, 0.6)$. The pareto-optimal set is illustrated by a thick orange line. Through the combination of both constraints the pareto-set is split into two parts.
Analytically, the pareto-optimal set is  given by $PS = \{(x_1, x_2) \,|\, (0.1 \leq x_1 \leq 0.4) \lor (0.6 \leq x_1 \leq 0.9) \, \land \, x_2 = 0\}$ and the Pareto-front by $f_2 = (\sqrt{f_1} - 1)^2$ where $f_1$ is defined in $[0.01,0.16]$ and $[0.36,0.81]$.

### Problem Definition

In *pymoo*, we consider pure minimization problems for optimization in all our modules. However, without loss of generality an objective which is supposed to be maximized, can be multiplied by $-1$ and be minimized. Therefore, we minimize $-f_2(x)$ instead of maximizing $f_2(x)$ in our optimization problem. Furthermore, all constraint functions need to be formulated as a $\leq 0$ constraint.
The feasibility of a solution can, therefore, be expressed by:

$$ \begin{cases}
\text{feasible,} \quad \quad \sum_i^n \langle g_i(x)\rangle = 0\\
\text{infeasbile,} \quad \quad \quad \text{otherwise}\\
\end{cases}
$$

$$
\text{where} \quad \langle g_i(x)\rangle =
\begin{cases}
0, \quad \quad \; \text{if} \; g_i(x) \leq 0\\
g_i(x), \quad  \text{otherwise}\\
\end{cases}
$$

For this reason, $g_2(x)$ needs to be multiplied by $-1$ in order to flip the $\geq$ to a $\leq$ relation. We recommend the normalization of constraints to give equal importance to each of them.
For $g_1(x)$, the coefficient results in $2 \cdot (-0.1) \cdot (-0.9) = 0.18$ and for $g_2(x)$ in $20 \cdot (-0.4) \cdot (-0.6) = 4.8$, respectively. We achieve normalization of constraints by dividing $g_1(x)$ and $g_2(x)$ by its corresponding coefficient.

Finally, the optimization problem to be optimized using *pymoo* is defined by:

\begin{align}
\label{eq:getting_started_pymoo}
\begin{split}
\min \;\; & f_1(x) = (x_1^2 + x_2^2) \\
\min \;\; & f_2(x) = (x_1-1)^2 + x_2^2 \\[1mm]
\text{s.t.} \;\; & g_1(x) = 2 \, (x_1 - 0.1) \, (x_1 - 0.9)  \, /  \,  0.18 \leq 0\\
& g_2(x) = - 20 \, (x_1 - 0.4) \, (x_1 - 0.6) \,  /  \,  4.8 \leq 0\\[1mm]
& -2 \leq x_1 \leq 2 \\
& -2 \leq x_2 \leq 2
\end{split}
\end{align}

Next, the derived problem formulation is implemented in Python. Each optimization problem in *pymoo* has to inherit from the [Problem](problems/index.ipynb) class. First, by calling the `super()` function the problem properties such as the number of variables `n_var`, objectives `n_obj` and constraints `n_constr` are initialized. Furthermore, lower `xl` and upper variables boundaries `xu` are supplied as a NumPy array. Additionally, the evaluation function _evaluate needs to be overwritten from the superclass.
The method takes a two-dimensional NumPy array x with n rows and m columns as an input. Each row represents an individual and each column an optimization variable. After doing the necessary calculations, the objective values have to be added to the dictionary out with the key `F` and the constraints with key `G`.

```python
import numpy as np
from pymoo.util.misc import stack
from pymoo.model.problem import Problem

class MyProblem(Problem):

    def __init__(self):
        super().__init__(n_var=2,
                         n_obj=2,
                         n_constr=2,
                         xl=np.array([-2,-2]),
                         xu=np.array([2,2]))

    def _evaluate(self, x, out, *args, **kwargs):
        f1 = x[:,0]**2 + x[:,1]**2
        f2 = (x[:,0]-1)**2 + x[:,1]**2

        g1 = 2*(x[:, 0]-0.1) * (x[:, 0]-0.9) / 0.18
        g2 = - 20*(x[:, 0]-0.4) * (x[:, 0]-0.6) / 4.8

        out["F"] = np.column_stack([f1, f2])
        out["G"] = np.column_stack([g1, g2])

    # --------------------------------------------------
    # Pareto-front - not necessary but used for plotting
    # --------------------------------------------------
    def _calc_pareto_front(self, flatten=True, **kwargs):
        f1_a = np.linspace(0.1**2, 0.4**2, 100)
        f2_a = (np.sqrt(f1_a) - 1)**2

        f1_b = np.linspace(0.6**2, 0.9**2, 100)
        f2_b = (np.sqrt(f1_b) - 1)**2

        a, b = np.column_stack([f1_a, f2_a]), np.column_stack([f1_b, f2_b])
        return stack(a, b, flatten=flatten)

    # --------------------------------------------------
    # Pareto-set - not necessary but used for plotting
    # --------------------------------------------------
    def _calc_pareto_set(self, flatten=True, **kwargs):
        x1_a = np.linspace(0.1, 0.4, 50)
        x1_b = np.linspace(0.6, 0.9, 50)
        x2 = np.zeros(50)

        a, b = np.column_stack([x1_a, x2]), np.column_stack([x1_b, x2])
        return stack(a,b, flatten=flatten)

problem = MyProblem()

```

Because we consider a test problem where the optimal solutions in design and objective space are known, we have implemented the `_calc_pareto_front` and `_calc_pareto_set` functions to observe the convergence of the algorithm later on. For the optimization run itself the methods need not to be overwritten. **So, no worries if you are investigating benchmark or real-world optimization problems.**

Moreover, we would like to mention that in many test optimization problems implementation already exist. For example, the test problem *ZDT1* can be initiated by:

```python
from pymoo.factory import get_problem

zdt1 = get_problem("zdt1")
```

Our framework has various single- and many-objective optimization test problems already implemented. Furthermore, a more advanced guide for custom problem definitions is available. In case problem functions are computationally expensive, parallelization of the evaluations functions might be an option.

[Optimization Test Problems](problems/index.ipynb) |
[Define a Custom Problem](problems/custom.ipynb) |
[Parallelization](problems/parallelization.ipynb)
[Callback](interface/callback.ipynb)

### Initialize an Algorithm

Moreover, we need to initialize a method to optimize the problem.
In *pymoo* factory methods create an `algorithm` object to be used for optimization. For each of those methods an API documentation is available and through supplying different parameters, algorithms can be customized in a plug-and-play manner.
Depending on the optimization problem different algorithms can be used to optimize the problem. Our framework offers various [Algorithms](algorithms/index.ipynb) which can be used to solve problems with different characteristics.

In general, the choice of a suitable algorithm for optimization problems is a challenge itself. Whenever problem characteristics are known beforehand, we recommended using those through customized operators.
However, in our case the optimization problem is rather simple, but the aspect of having two objectives and two constraints should be considered. For this reason, we decided to use [NSGA-II](algorithms/nsga2.ipynb) with its default configuration with minor modifications. We chose a population size of 40 (`pop_size=40`) and decided instead of generating the same number of offsprings to create only 10 (`n_offsprings=40`). This is a greedier  variant and improves the convergence of rather simple optimization problems without difficulties regarding optimization, such as the existence of local Pareto fronts.
Moreover, we enable a duplicate check (`eliminate_duplicates=True`) which makes sure that the mating produces offsprings which are different with respect to themselves and the existing population regarding their design space values. To illustrate the customization aspect, we listed the other unmodified default operators in the code-snippet below.

```python
from pymoo.algorithms.nsga2 import NSGA2
from pymoo.factory import get_sampling, get_crossover, get_mutation

algorithm = NSGA2(
    pop_size=40,
    n_offsprings=10,
    sampling=get_sampling("real_random"),
    crossover=get_crossover("real_sbx", prob=0.9, eta=15),
    mutation=get_mutation("real_pm", eta=20),
    eliminate_duplicates=True
)

```

The `algorithm` object contains the implementation of NSGA-II  with the custom settings supplied to the factory method.

### Define a Termination Criterion

Furthermore, a termination criterion needs to be defined to finally start the optimization procedure. Different kind of [Termination Criteria](interface/termination.ipynb) are available. Here, since the problem is rather simple, we run the algorithm for some number of generations.

```python
from pymoo.factory import get_termination

termination = get_termination("n_gen", 40)
```

Instead of number of generations (or iterations) other criteria such as the number of function evaluations or the improvement in design or objective space from the last to the current generation can be used.

### Optimize

Finally, we are solving the problem with the algorithm and termination criterion we have defined.

```python
from pymoo.optimize import minimize

res = minimize(problem,
               algorithm,
               termination,
               seed=1,
               pf=problem.pareto_front(use_cache=False),
               save_history=True,
               verbose=True)

```

    ==========================================================================================
    n_gen |  n_eval |   cv (min)   |   cv (avg)   |     igd      |      gd      |      hv
    ==========================================================================================
        1 |      40 |  0.00000E+00 |  2.36399E+01 |  0.355456661 |  0.005362269 |  0.259340233
        2 |      50 |  0.00000E+00 |  1.15486E+01 |  0.355456661 |  0.005362269 |  0.259340233
        3 |      60 |  0.00000E+00 |  5.277918607 |  0.355456661 |  0.005362269 |  0.259340233
        4 |      70 |  0.00000E+00 |  2.406068542 |  0.336177786 |  0.003791374 |  0.267747228
        5 |      80 |  0.00000E+00 |  0.908316880 |  0.146187673 |  0.030825699 |  0.326821728
        6 |      90 |  0.00000E+00 |  0.264746300 |  0.146187673 |  0.030825699 |  0.326821728
        7 |     100 |  0.00000E+00 |  0.054063822 |  0.127166639 |  0.048434536 |  0.326821728
        8 |     110 |  0.00000E+00 |  0.003060876 |  0.116635438 |  0.040756917 |  0.333328216
        9 |     120 |  0.00000E+00 |  0.00000E+00 |  0.085912985 |  0.018668919 |  0.342436730
       10 |     130 |  0.00000E+00 |  0.00000E+00 |  0.077470550 |  0.033493862 |  0.355627493
       11 |     140 |  0.00000E+00 |  0.00000E+00 |  0.077352198 |  0.040282741 |  0.356652764
       12 |     150 |  0.00000E+00 |  0.00000E+00 |  0.053162450 |  0.004785549 |  0.392484461
       13 |     160 |  0.00000E+00 |  0.00000E+00 |  0.044223026 |  0.004743970 |  0.407719675
       14 |     170 |  0.00000E+00 |  0.00000E+00 |  0.038805143 |  0.005120891 |  0.410670860
       15 |     180 |  0.00000E+00 |  0.00000E+00 |  0.037729492 |  0.004271793 |  0.412100954
       16 |     190 |  0.00000E+00 |  0.00000E+00 |  0.036972091 |  0.007588808 |  0.414446121
       17 |     200 |  0.00000E+00 |  0.00000E+00 |  0.034660980 |  0.006036578 |  0.419665136
       18 |     210 |  0.00000E+00 |  0.00000E+00 |  0.033516235 |  0.005483394 |  0.420111173
       19 |     220 |  0.00000E+00 |  0.00000E+00 |  0.024622717 |  0.005292631 |  0.427207379
       20 |     230 |  0.00000E+00 |  0.00000E+00 |  0.023034030 |  0.004382377 |  0.440473619
       21 |     240 |  0.00000E+00 |  0.00000E+00 |  0.016906921 |  0.004430595 |  0.452594419
       22 |     250 |  0.00000E+00 |  0.00000E+00 |  0.016592436 |  0.004155185 |  0.453482517
       23 |     260 |  0.00000E+00 |  0.00000E+00 |  0.016321930 |  0.004527221 |  0.453699792
       24 |     270 |  0.00000E+00 |  0.00000E+00 |  0.013987946 |  0.004004738 |  0.455365890
       25 |     280 |  0.00000E+00 |  0.00000E+00 |  0.013454690 |  0.004313685 |  0.455631510
       26 |     290 |  0.00000E+00 |  0.00000E+00 |  0.012960456 |  0.004264848 |  0.455925267
       27 |     300 |  0.00000E+00 |  0.00000E+00 |  0.012683256 |  0.004123014 |  0.456126518
       28 |     310 |  0.00000E+00 |  0.00000E+00 |  0.012559351 |  0.004359358 |  0.456110146
       29 |     320 |  0.00000E+00 |  0.00000E+00 |  0.012228516 |  0.003453151 |  0.456745602
       30 |     330 |  0.00000E+00 |  0.00000E+00 |  0.011661748 |  0.003452586 |  0.457087950
       31 |     340 |  0.00000E+00 |  0.00000E+00 |  0.011210955 |  0.003806385 |  0.457014270
       32 |     350 |  0.00000E+00 |  0.00000E+00 |  0.010909137 |  0.003449675 |  0.457290422
       33 |     360 |  0.00000E+00 |  0.00000E+00 |  0.010816517 |  0.003320362 |  0.457462656
       34 |     370 |  0.00000E+00 |  0.00000E+00 |  0.010324079 |  0.003311767 |  0.457900371
       35 |     380 |  0.00000E+00 |  0.00000E+00 |  0.010070886 |  0.003101986 |  0.458298915
       36 |     390 |  0.00000E+00 |  0.00000E+00 |  0.009995280 |  0.003347647 |  0.458312542
       37 |     400 |  0.00000E+00 |  0.00000E+00 |  0.009619601 |  0.003028411 |  0.458519325
       38 |     410 |  0.00000E+00 |  0.00000E+00 |  0.009319637 |  0.002690949 |  0.458767148
       39 |     420 |  0.00000E+00 |  0.00000E+00 |  0.009261538 |  0.002697147 |  0.458861977
       40 |     430 |  0.00000E+00 |  0.00000E+00 |  0.009261263 |  0.002742830 |  0.458918287

The [Result](interface/result.ipynb) object provides the corresponding X and F values and some more information.

### Visualize

The optimization results are illustrated below (design and objective space). The solid line represents the analytically derived Pareto set and front in the corresponding space and the circles solutions found by the algorithm. It can be observed that the algorithm was able to converge, and a set of nearly-optimal solutions was obtained.

```python
from pymoo.visualization.scatter import Scatter

# get the pareto-set and pareto-front for plotting
ps = problem.pareto_set(use_cache=False, flatten=False)
pf = problem.pareto_front(use_cache=False, flatten=False)

# Design Space
plot = Scatter(title = "Design Space", axis_labels="x")
plot.add(res.X, s=30, facecolors='none', edgecolors='r')
plot.add(ps, plot_type="line", color="black", alpha=0.7)
plot.do()
plot.apply(lambda ax: ax.set_xlim(-0.5, 1.5))
plot.apply(lambda ax: ax.set_ylim(-2, 2))
plot.show()

# Objective Space
plot = Scatter(title = "Objective Space")
plot.add(res.F)
plot.add(pf, plot_type="line", color="black", alpha=0.7)
plot.show()
```

Visualization is an important post-processing step in multi-objective optimization. Although it seems to be pretty easy for our example optimization problem, it becomes much more difficult in higher dimensions where trade-offs between solutions are not easily observable. For visualizations in higher dimensions, various more advanced [Visualizations](visualization/index.ipynb) are implemented in our framework.

### Performance Tracking

If the optimization scenario is repetitive it makes sense to track the performance of the algorithm. Because we have stored the history of the optimization run, we can now analyze the convergence over time. To measure the performance, we need to decide what metric to be used. Here, we are using Hypervolume. Of course, other [Performance Indicators](misc/performance_indicator.ipynb) are available as well.

```python
import matplotlib.pyplot as plt
from pymoo.performance_indicator.hv import Hypervolume

# create the performance indicator object with reference point (4,4)
metric = Hypervolume(ref_point=np.array([1.0, 1.0]))

# collect the population in each generation
pop_each_gen = [a.pop for a in res.history]

# receive the population in each generation
obj_and_feasible_each_gen = [pop[pop.get("feasible")[:,0]].get("F") for pop in pop_each_gen]

# calculate for each generation the HV metric
hv = [metric.calc(f) for f in obj_and_feasible_each_gen]

# function evaluations at each snapshot
n_evals = np.array([a.evaluator.n_eval for a in res.history])

# visualze the convergence curve
plt.plot(n_evals, hv, '-o')
plt.title("Convergence")
plt.xlabel("Function Evaluations")
plt.ylabel("Hypervolume")
plt.show()
```

We hope you have enjoyed the getting started guide. For more topics we refer to each section covered by on the [landing page](https://pymoo.org). If you have any question or concern do not hesitate to [contact us](contact.rst).


## CCNSS 2018 Module 3: Reinforcement learning and planning
# Tutorial 2: Multi-armed bandits

# **Introduction**

In this tutorial, we will learn about different ways to select actions based on value estimates in an k-armed bandit context.

**Environment:**

Consider the following learning problem. You are faced repeatedly with a choice among $k$ different options, or actions. After each choice you receive a numerical reward chosen from a stationary probability distribution that depends on the action you selected. Your objective is to maximize the expected total reward over some time period, for example, over 1000 action selections, or time steps.
This is the original form of the k-armed bandit problem, so named by analogy to a slot machine.

**Definitions:**

1. Action-value function
\begin{align}
q (a) = \mathbb{E} [r_{t} | a_{t} = a]
\end{align}

2. Action-value updates:
\begin{align}
q_{t}(a) \leftarrow q_{t}(a) + \alpha (r_{t} - q_{t}(a))
\end{align}

3. Greedy policy:
\begin{align}
a_{t} = \text{argmax}_{a} \; q_{t} (a)
\end{align}

4. $\epsilon$-greedy policy:
\begin{align}
P (a_{t} = a) =
        \begin{cases}
        1 - \epsilon    & \quad \text{if } a_{t} = \text{argmax}_{a} \; q_{t} (a) \\
        \epsilon/(N-1)        & \quad \text{else}
        \end{cases}
\end{align}

**Please run the following code for your implementations:**

```
% matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from pylab import *
```

```
class n_armed_bandit():
    """
    World: N-Armed bandit.
    Only one state, multiple actions.
    Each action returns different amount of reward.
    For each action, rewards are randomly sampled from normal distribution,
        with a mean associated with that arm and unit variance.
    """

    def __init__(self, arm_number):
        self.name = "n_armed_bandit"
        self.n_states = 1
        self.n_actions = arm_number
        self.dim_x = 1
        self.dim_y = 1

        self.mu = [np.random.normal(0,1) for a in range(self.n_actions)]

    def get_outcome(self, state, action):

        self.rewards = [np.random.normal(self.mu[i],1) for i in range(self.n_actions)]
        next_state = None

        reward = self.rewards[action]
        return int(next_state) if next_state is not None else None, reward
```

```
class drifting_n_armed_bandit():
    """
    World: N-Armed bandit.
    Only one state, multiple actions.
    Each action returns different amounts of rewards.
    For each action, rewards are randomly sampled from normal distribution,
        with a mean associated with that arm and unit variance.
    In the case of the non-stationary bandit, the mean reward associated with each arm
        follows a Gaussian random.
    """

    def __init__(self, arm_number, drift):
        self.name = "drifting_n_armed_bandit"
        self.n_states = 1
        self.n_actions = arm_number
        self.dim_x = 1
        self.dim_y = 1

        self.mu_min = 0.25
        self.mu_max = 0.75
        self.drift = drift

        self.mu = [np.random.normal(0,0) for a in range(self.n_actions)]

    def update_mu(self):
        self.mu += np.random.normal(0, self.drift, self.n_actions)

    def get_outcome(self, state, action):

        self.update_mu()
        self.rewards = [np.random.normal(self.mu[i],1) for i in range(self.n_actions)]
        next_state = None

        reward = self.rewards[action]
        return int(next_state) if next_state is not None else None, reward
```

```
def e_greedy_policy(q, epsilon):
    if np.random.uniform(0,1) > epsilon:
        return np.argmax(q)
    else:
        # choose randomly from all but the highest q value
        a_list = np.arange(len(q))
        arg_max = np.argmax(q)
        choices = np.delete(a_list, arg_max)
        return np.random.choice(np.delete(a_list, arg_max))
```

# **Exercise 1:**

Estimate the action value function of each of the $k$ actions in the bandit problem using the above update rule while choosing actions using an $\epsilon$-greedy policy to choose actions. Run for 1000 steps (i.e. take 1000 actions for a set of bandits that each have a fixed probability of returning a reward).

Plot reward obtained on each timestep for a randomly chosen bandit problem (i.e. a set of bandit arms with randomly chosen probabilities of providing rewards).

```

```

# **Exercise 2:**

1. Repeat above for 1000 runs (i.e. reset the bandit environment 1000 times, each time the bandit arms will get associated with randomly chosen probabilities of returning a reward), for different values of $\epsilon \in \{0,0.1,0.01\}$. Use $\alpha = 0.01$.
2. Plot performance curves i.e. fraction of runs on which optimal action was chosen at each time step, for the different values of $\epsilon$. At each step check over all 1000 runs to determine the fraction of runs in which the optimal action was chosen.
3. Change the learning rate $\alpha$ and plot performance curves.
4. Plot average rewards and performance curves using decaying learning rates of $\alpha = 1/N(a)$, where $N(a)$ is the number of times bandit arm $a$ was chosen.
5. Can you describe how $\epsilon$ and $\alpha$ affect the agent's performance? Tip: pay attention to the y-axis in your plots.
6. (Optional) For which value of $\epsilon$ will the above perform best in the long run in terms of cumulative reward and probability of selecting the best action? How much better will it be? Express your answer quantitatively.

```

```

# **Exercise 3:**

Estimate action values using an $\epsilon$-greedy policy, using *optimistic starts* i.e. initialise your action value function to be a value greater than zero ($q_{0}(a)> 0$).

Use $q_{0}(a) = 5$ and a fixed learning rate $\alpha = 0.1$.

1. Plot performance curves of $\epsilon$-greedy action selection (for $\epsilon = 0.1$) and using optimistic starts with greedy action selection. Overlay these plots.
2. How does performance change for different values of $\epsilon$?
3. Change $q_{0}(a)$ and plot performance curves for different optimistic initial values.
4. Can you explain why the initialisation is causing a difference in performance?

```

```

# **Exercise 4:**

Estimate action values using an $\epsilon$-greedy policy in the context of a nonstationary bandit problem. Use $\epsilon = 0.1$.

In this case, the expected rewards from each of the bandit arms drifts slowly and independently, based on a Gaussian random walk (with mean 0 and SD 0.025 and reflecting boundaries at 0.25 and 0.75). Number of bandit arms $k = 4$.

1. What differences do you expect when using a fixed vs decaying learning rate in this context?
2. For a single run, plot the expected reward from each bandit arm and the reward obtained by the agent. Run for 10000 steps. Try with fixed $\alpha = 0.1$ and decaying learning rates $\alpha = 1/N(a)$.
3. Simulate 2000 runs and plot performance curves i.e. the fraction of runs on which the optimal arm was chosen at each step.
4. How does the performance change for different learning rates? Plot performance curves using different values of $\alpha = \{0.001, 0.01, 0.1, 0.2\}$ on the same plot.

```

```

# **(Optional) Exercise 5:**

In the case of a decaying learning rate  of $\alpha = 1/N(a)$, at the $n$th timestep of encountering the $a$th bandit arm, each past instance of choosing that arm is equally weighted.

Can you analytically show that the above statement is true?

When $\alpha \neq 1/N(a)$, the resulting estimate is a weighted sum of all previously encountered rewards from that arm.
1. Compute analytically or plot the weighting of all previous rewards at the $n$th step in the case of constant learning rate of $\alpha = 0.01$.
2. What would be the weighting of all previous rewards from a given bandit arm at the $n$th step for an arbitrary sequence of learning rates $\{ \alpha_{1}, ... \alpha_{n}\}$?

```

```


###### Content under Creative Commons Attribution license CC-BY 4.0, code under MIT license (c)2014 L.A. Barba, C.D. Cooper, G.F. Forsyth.

# Spreading out

Welcome to the fifth, and last, notebook of Module 4 "_Spreading out: diffusion problems,"_ of our fabulous course **"Practical Numerical Methods with Python."**

In this course module, we have learned about explicit and implicit methods for parabolic equations in 1 and 2 dimensions. So far, all schemes have been first-order in time and second-order in space. _Can we do any better?_ We certainly can: this notebook presents the Crank-Nicolson scheme, which is a second-order method in both time and space! We will continue to use the heat equation to guide the discussion, as we've done throughout this module.

## Crank-Nicolson scheme

The [Crank Nicolson scheme](http://en.wikipedia.org/wiki/Crank–Nicolson_method) is a popular second-order, implicit method used with parabolic PDEs in particular. It was developed by John Crank and [Phyllis Nicolson](http://en.wikipedia.org/wiki/Phyllis_Nicolson). The main idea is to take the average between the solutions at $t^n$ and $t^{n+1}$ in the evaluation of the spatial derivative. Why bother doing that? Because the time derivative will then be discretized with a centered scheme, giving second-order accuracy!

Remember the 1D heat equation from the [first notebook](https://github.com/numerical-mooc/numerical-mooc/blob/master/lessons/04_spreadout/01_Heat_Equation_1D_Explicit.ipynb)? Just to refresh your memory, here it is:

\begin{equation}
\frac{\partial T}{\partial t} = \alpha \frac{\partial^2T}{\partial x^2}.
\end{equation}

In this case, the Crank-Nicolson scheme leads to the following discretized equation:

\begin{eqnarray}
\frac{T^{n+1}_i - T^n_i}{\Delta t} = & \nonumber \\
\alpha \cdot \frac{1}{2} &\left(
\frac{T^{n+1}_{i+1} - 2T^{n+1}_i + T^{n+1}_{i-1}}{\Delta x^2}
+ \frac{T^n_{i+1}-2T^n_i + T^n_{i-1}}{\Delta x^2}
\right)
\end{eqnarray}

Notice how the both time indices $n$ and $n+1$ appear on the right-hand side. You know we'll have to rearrange this equation, right? Now look at the stencil and notice that we are using more information than before in the update.

#### Figure 2. Stencil of the Crank-Nicolson scheme.

Rearraning terms so that everything that we don't know is on the left side and what we do know on the right side, we get

\begin{eqnarray}
-T^{n+1}_{i-1} + 2\left(\frac{\Delta x^2}{\alpha\Delta t}+1\right)T^{n+1}_i - T^{n+1}_{i+1} = & \nonumber\\
T^{n}_{i-1} + & 2\left(\frac{\Delta x^2}{\alpha\Delta t}-1\right)T^{n}_i + T^{n}_{i+1}
\end{eqnarray}

Again, we are left with a linear system of equations. Check out the left side of that equation: it looks a lot like the matrix from [notebook 2](https://github.com/numerical-mooc/numerical-mooc/blob/master/lessons/04_spreadout/02_Heat_Equation_1D_Implicit.ipynb), doesn't it? Apart from the slight modification in the $T_i^{n+1}$ term, the left side of the equation is pretty much the same. What about the right-hand side? Sure, it looks quite different, but that is not a problem, we know all those terms!

Things don't change much for boundary conditions, either. We've seen all the cases already. Say $T_0^{n+1}$ is a Dirichlet boundary. Then the equation for $i=1$ becomes

\begin{eqnarray}
 2\left(\frac{\Delta x^2}{\alpha\Delta t}+1\right)T^{n+1}_1 - T^{n+1}_{2} = & \nonumber\\
 T^{n}_{0} + & 2\left(\frac{\Delta x^2}{\alpha\Delta t}-1\right)T^{n}_1 + T^{n}_{2} + T^{n+1}_{0}
\end{eqnarray}

And if we have a Neumann boundary $\left(\left.\frac{\partial T}{\partial x}\right|_{x=L} = q\right)$ at $T_{n_x-1}^{n+1}$? We know this stuff, right? For $i=n_x-2$ we get

\begin{eqnarray}
-T^{n+1}_{n_x-3} + \left(2\frac{\Delta x^2}{\alpha\Delta t}+1\right)T^{n+1}_{n_x-2} = & \nonumber\\
T^{n}_{n_x-3} + & 2\left(\frac{\Delta x^2}{\alpha\Delta t}-1\right)T^{n}_{n_x-2} + T^{n}_{n_x-1} + q\Delta x
\end{eqnarray}

The code will look a lot like the implicit method from the [second notebook](https://github.com/numerical-mooc/numerical-mooc/blob/master/lessons/04_spreadout/02_Heat_Equation_1D_Implicit.ipynb). Only some terms of the matrix and right-hand-side vector will be different, which changes some of our custom functions.

### The linear system

Just like in [notebook 2](https://github.com/numerical-mooc/numerical-mooc/blob/master/lessons/04_spreadout/02_Heat_Equation_1D_Implicit.ipynb), we need to solve a linear system on every time step of the form:

$$[A][T^{n+1}_\text{int}] = [b]+[b]_{b.c.}$$

The coefficient matrix is very similar to the previous case, but the right-hand side changes a lot:

The coefficient matrix is very similar to the previous case, but the right-hand side changes a lot:

\begin{align}\left[ \begin{array}{cccccc}
 2\left(\frac{1}{\sigma}+1\right) & -1 & 0 & \cdots & & 0 \\
 -1 & 2\left(\frac{1}{\sigma}+1\right) & -1 & 0 & \cdots & 0 \\
 0 & & \ddots& & & \vdots \\
 \vdots & & & & 2\left(\frac{1}{\sigma}+1\right)& \\
 0 & \cdots & & & -1 & \left(2\frac{1}{\sigma}+1\right) \end{array} \right]
 \cdot
 \left[ \begin{array}{c}
 T_1^{n+1} \\ T_2^{n+1} \\ \vdots \\ \\ T_{N-2}^{n+1} \end{array} \right]
 =
 \left[ \begin{array}{c}
 T_0^n + 2\left(\frac{1}{\sigma}-1\right)T_1^n + T_2^n \\  T_1^n + 2\left(\frac{1}{\sigma}-1\right)T_2^n + T_3^n \\ \vdots \\ \\  T_{n_x-3}^n + 2\left(\frac{1}{\sigma}-1\right)T_{n_x-2}^n + T_{n_x-1}^n \end{array} \right]
 +
 \begin{bmatrix}
  T_0^{n+1}\\
  0\\\\
  \vdots\\\\
  0\\
  q\Delta x
 \end{bmatrix}
 \end{align}

Let's write a function that will create the coefficient matrix and right-hand-side vectors for the heat conduction problem from [notebook 2](https://github.com/numerical-mooc/numerical-mooc/blob/master/lessons/04_spreadout/02_Heat_Equation_1D_Implicit.ipynb): with Dirichlet boundary at $x=0$ and zero-flux boundary $(q=0)$ at $x=L$.

```python
import numpy
from scipy.linalg import solve
```

```python
def generateMatrix(N, sigma):
    """ Computes the matrix for the diffusion equation with Crank-Nicolson
        Dirichlet condition at i=0, Neumann at i=-1

    Parameters:
    ----------
    N: int
        Number of discretization points
    sigma: float
        alpha*dt/dx^2

    Returns:
    -------
    A: 2D numpy array of float
        Matrix for diffusion equation
    """

    # Setup the diagonal
    d = 2*numpy.diag(numpy.ones(N-2)*(1+1./sigma))

    # Consider Neumann BC
    d[-1,-1] = 1+2./sigma

    # Setup upper diagonal
    ud = numpy.diag(numpy.ones(N-3)*-1, 1)

    # Setup lower diagonal
    ld = numpy.diag(numpy.ones(N-3)*-1, -1)

    A = d + ud + ld

    return A
```

```python
def generateRHS(T, sigma):
    """ Computes right-hand side of linear system for diffusion equation
        with backward Euler

    Parameters:
    ----------
    T: array of float
        Temperature at current time step
    sigma: float
        alpha*dt/dx^2

    Returns:
    -------
    b: array of float
        Right-hand side of diffusion equation with backward Euler
    """

    b = T[1:-1]*2*(1./sigma-1) + T[:-2] + T[2:]
    # Consider Dirichlet BC
    b[0] += T[0]

    return b
```

We will solve the linear system at every time step. Let's define a function to step in time:

```python
def CrankNicolson(T, A, nt, sigma):
    """ Advances diffusion equation in time with Crank-Nicolson

    Parameters:
    ----------
    T: array of float
        initial temperature profile
    A: 2D array of float
        Matrix with discretized diffusion equation
    nt: int
        number of time steps
    sigma: float
        alpha*td/dx^2

    Returns:
    -------
    T: array of floats
        temperature profile after nt time steps
    """

    for t in range(nt):
        Tn = T.copy()
        b = generateRHS(Tn, sigma)
        # Use numpy.linalg.solve
        T_interior = solve(A,b)
        T[1:-1] = T_interior
        # Enforce Neumann BC (Dirichlet is enforced automatically)
        T[-1] = T[-2]

    return T
```

And we are good to go! First, let's setup our initial conditions, and the matrix

```python
L = 1
nx = 21
alpha = 1.22e-3

dx = L/(nx-1)

Ti = numpy.zeros(nx)
Ti[0] = 100

sigma = 0.5
dt = sigma * dx*dx/alpha
nt = 10

A = generateMatrix(nx, sigma)
```

Check the matrix...

```python
print(A)
```

    [[ 6. -1.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.
       0.]
     [-1.  6. -1.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.
       0.]
     [ 0. -1.  6. -1.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.
       0.]
     [ 0.  0. -1.  6. -1.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.
       0.]
     [ 0.  0.  0. -1.  6. -1.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.
       0.]
     [ 0.  0.  0.  0. -1.  6. -1.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.
       0.]
     [ 0.  0.  0.  0.  0. -1.  6. -1.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.
       0.]
     [ 0.  0.  0.  0.  0.  0. -1.  6. -1.  0.  0.  0.  0.  0.  0.  0.  0.  0.
       0.]
     [ 0.  0.  0.  0.  0.  0.  0. -1.  6. -1.  0.  0.  0.  0.  0.  0.  0.  0.
       0.]
     [ 0.  0.  0.  0.  0.  0.  0.  0. -1.  6. -1.  0.  0.  0.  0.  0.  0.  0.
       0.]
     [ 0.  0.  0.  0.  0.  0.  0.  0.  0. -1.  6. -1.  0.  0.  0.  0.  0.  0.
       0.]
     [ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0. -1.  6. -1.  0.  0.  0.  0.  0.
       0.]
     [ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0. -1.  6. -1.  0.  0.  0.  0.
       0.]
     [ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0. -1.  6. -1.  0.  0.  0.
       0.]
     [ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0. -1.  6. -1.  0.  0.
       0.]
     [ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0. -1.  6. -1.  0.
       0.]
     [ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0. -1.  6. -1.
       0.]
     [ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0. -1.  6.
      -1.]
     [ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0. -1.
       5.]]

Looks okay! Now, step in time

```python
T = CrankNicolson(Ti.copy(), A, nt, sigma)
```

And plot,

```python
from matplotlib import pyplot
%matplotlib inline
from matplotlib import rcParams
rcParams['font.family'] = 'serif'
rcParams['font.size'] = 16
```

```python
x = numpy.linspace(0,L,nx)

pyplot.plot(x, T, color='#003366', ls='-', lw=3);
```

Works nicely. But wait! This method has elements of explicit and implicit discretizations. Is it *conditionally stable* like forward Euler, or *unconditionally stable* like backward Euler? Try out different values of `sigma`. You'll see Crank-Nicolson is an *unconditionally stable scheme* for the diffusion equation!

## Accuracy & convergence

Using some techniques you might have learned in your PDE class, such as separation of variables, you can get a closed expression for the rod problem. It looks like this:

\begin{eqnarray}
T(x,t) = & \nonumber \\
100 - \sum_{n=1}^{\infty} & \frac{400}{(2n-1)\pi}\sin\left(\frac{(2n-1)\pi}{2L}x\right) \exp\left[-\alpha\left(\frac{(2n-1)\pi}{2L}\right)^2t\right]
\end{eqnarray}

Unfortunately, the analytical solution is a bit messy, but at least it gives a good approximation if we evaluate it for large $n$. Let's define a function that will calculate this for us:

```python
from math import pi
def T_analytical(x, t, n_max, alpha, L):
    """Computes the exact solution for 1D diffusion with T=100 at x=0 and dT/dx=0 at x=L

    Paramters:
    ---------
    x    : array of float
        Spatial position
    t    : float
        Evaluation time
    n_max: int
        Number of terms to evaluate expression
    alpha: float
        diffusion coefficient
    L    : float
        Size of rod

    Returns:
    -------
    T    : array of float
        Temperature at each location x
    """
    T = 100
    for n in range(1,n_max+1):
        k = (2*n-1)*pi/(2*L)

        summation = 400/((2*n-1)*pi) * numpy.sin(k*x) * numpy.exp(-alpha*k*k*t)
        T -= summation

    return T
```

And let's see how that expression looks for the time where we left the numerical solution

```python
T_exact = T_analytical(x, dt*nt, 100, alpha, L)
pyplot.plot(x, T_exact, color='#003366', ls='-', lw=3);
```

```python
T1 = T_analytical(x, .2, 100, alpha, L)
T2 = T_analytical(x, .2, 200, alpha, L)
numpy.sqrt(numpy.sum((T1-T2)**2)/numpy.sum(T2**2))
```

    6.9279171182600926e-13

That looks like it should. We'll now use this result to study the convergence of the Crank-Nicolson scheme.

### Time convergence

We said this method was second-order accurate in time, remember? That's in theory, but we should test that the numerical solution indeed behaves like the theory says.

Leaving $\Delta x$ constant, we'll run the code for different values of $\Delta t$ and compare the result at the same physical time, say $t=n_t\cdot\Delta t=10$, with the analytical expression above.

The initial condition of the rod problem has a very sharp gradient: it suddendly jumps from $0{\rm C}$ to $100{\rm C}$ at the boundary. To resolve that gradient to the point that it doesn't affect time convergence, we would need a very fine mesh, and computations would be very slow. To avoid this issue, we will start from $t=1$ rather than starting from $t=0$.

First, let's define a function that will compute the $L_2$-norm of the error:

```python
def L2_error(T, T_exact):
    """Computes L2 norm of error

    Parameters:
    ----------
    T      : array of float
        array with numerical solution
    T_exact: array of float
        array with exact solution
    Returns:
    -------
    e: L2 norm of error
    """

    e = numpy.sqrt(numpy.sum((T-T_exact)**2)/numpy.sum(T_exact)**2)

    return e
```

For fun, let's compare the Crank-Nicolson schem with the implicit (a.k.a., backward) Euler scheme. We'll borrow some functions from [notebook 2](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/04_spreadout/02_Heat_Equation_1D_Implicit.ipynb) to do this.

```python
def generateMatrix_btcs(N, sigma):
    """ Computes the matrix for the diffusion equation with backward Euler
        Dirichlet condition at i=0, Neumann at i=-1

    Parameters:
    ----------
    T: array of float
        Temperature at current time step
    sigma: float
        alpha*dt/dx^2

    Returns:
    -------
    A: 2D numpy array of float
        Matrix for diffusion equation
    """

    # Setup the diagonal
    d = numpy.diag(numpy.ones(N-2)*(2+1./sigma))

    # Consider Neumann BC
    d[-1,-1] = 1+1./sigma

    # Setup upper diagonal
    ud = numpy.diag(numpy.ones(N-3)*-1, 1)

    # Setup lower diagonal
    ld = numpy.diag(numpy.ones(N-3)*-1, -1)

    A = d + ud + ld

    return A
```

```python
def generateRHS_btcs(T, sigma):
    """ Computes right-hand side of linear system for diffusion equation
        with backward Euler

    Parameters:
    ----------
    T: array of float
        Temperature at current time step
    sigma: float
        alpha*dt/dx^2

    Returns:
    -------
    b: array of float
        Right-hand side of diffusion equation with backward Euler
    """
    b = numpy.zeros_like(T)

    b = T[1:-1]*1./sigma
    # Consider Dirichlet BC
    b[0] += T[0]

    return b
```

```python
def implicit_btcs(T, A, nt, sigma):
    """ Advances diffusion equation in time with implicit central scheme

    Parameters:
    ----------
    T: array of float
        initial temperature profile
    A: 2D array of float
        Matrix with discretized diffusion equation
    nt: int
        number of time steps
    sigma: float
        alpha*td/dx^2

    Returns:
    -------
    T: array of floats
        temperature profile after nt time steps
    """

    for t in range(nt):
        Tn = T.copy()
        b = generateRHS_btcs(Tn, sigma)
        # Use numpy.linalg.solve
        T_interior = solve(A,b)
        T[1:-1] = T_interior
        # Enforce Neumann BC (Dirichlet is enforced automatically)
        T[-1] = T[-2]

    return T
```

Now, let's do the runs!

```python
nx = 1001
dx = L/(nx-1)

dt_values = numpy.asarray([1.0, 0.5, 0.25, 0.125])
error = numpy.zeros(len(dt_values))
error_ftcs = numpy.zeros(len(dt_values))

t_final = 10
t_initial = 1

x = numpy.linspace(0,L,nx)

Ti =  T_analytical(x, t_initial, 100, alpha, L)
T_exact = T_analytical(x, t_final, 100, alpha, L)

for i,dt in enumerate(dt_values):
    sigma = alpha*dt/dx**2

    nt = int((t_final-t_initial)/dt)

    A = generateMatrix(nx, sigma)

    A_btcs = generateMatrix_btcs(nx, sigma)

    T = CrankNicolson(Ti.copy(), A, nt, sigma)

    error[i] = L2_error(T,T_exact)

    T = implicit_btcs(Ti.copy(), A_btcs, nt, sigma)

    error_ftcs[i] = L2_error(T,T_exact)

```

And plot,

```python
pyplot.figure(figsize=(8,8))
pyplot.grid(True)
pyplot.xlabel(r'$\Delta t$', fontsize=18)
pyplot.ylabel(r'$L_2$-norm of the error', fontsize=18)
pyplot.axis('equal')
pyplot.loglog(dt_values, error, color='k', ls='--', lw=2, marker='o')
pyplot.loglog(dt_values, error_ftcs, color='k', ls='--', lw=2, marker='s')
pyplot.legend(['Crank-Nicolson', 'FTCS']);
```

```python
error
```

    array([  3.81125927e-05,   9.41813943e-06,   2.25089054e-06,
             4.63970974e-07])

See how the error drops four times when the time step is halved? This method is second order in time!

Clearly, Crank-Nicolson (circles) converges faster than backward Euler (squares)! Not only that, but also the error curve is shifted down: Crank-Nicolson is more accurate.

If you look closely, you'll realize that the error in Crank-Nicolson decays about twice as fast than backward Euler: it's a second versus first order method!

### Spatial convergence

To study spatial convergence, we will run the code for meshes with 21, 41, 81 and 161 points, and compare them at the same non-dimensional time, say $t=20$.

Let's start by defining a function that will do everything for us

```python
nx_values = numpy.asarray([11, 21, 41, 81, 161])

dt = 0.1
error = numpy.zeros(len(nx_values))

t_final = 20

x = numpy.linspace(0,L,nx)

for i,nx in enumerate(nx_values):

    dx = L/(nx-1)
    x = numpy.linspace(0,L,nx)

    sigma = alpha*dt/dx**2

    nt = int(t_final/dt)

    A = generateMatrix(nx, sigma)

    Ti = numpy.zeros(nx)
    Ti[0] = 100

    T = CrankNicolson(Ti.copy(), A, nt, sigma)

    T_exact = T_analytical(x, t_final, 100, alpha, L)

    error[i] = L2_error(T,T_exact)
```

And plot!

```python
pyplot.figure(figsize=(8,8))
pyplot.grid(True)
pyplot.xlabel(r'$n_x$', fontsize=18)
pyplot.ylabel(r'$L_2$-norm of the error', fontsize=18)
pyplot.axis('equal')
pyplot.loglog(nx_values, error, color='k', ls='--', lw=2, marker='o');
```

That looks good! See how for each quadrant we go right, the error drops two quadrants going down (and even a bit better!).

##### Dig deeper

Let's re-do the spatial convergence, but comparing at a much later time, say $t=1000$.

```python
nx_values = numpy.asarray([11, 21, 41, 81, 161])

dt = 0.1
error = numpy.zeros(len(nx_values))

t_final = 1000

x = numpy.linspace(0,L,nx)

for i,nx in enumerate(nx_values):

    dx = L/(nx-1)
    x = numpy.linspace(0,L,nx)

    sigma = alpha*dt/dx**2

    nt = int(t_final/dt)

    A = generateMatrix(nx, sigma)

    Ti = numpy.zeros(nx)
    Ti[0] = 100

    T = CrankNicolson(Ti.copy(), A, nt, sigma)

    T_exact = T_analytical(x, t_final, 100, alpha, L)

    error[i] = L2_error(T,T_exact)
```

```python
pyplot.figure(figsize=(8,8))
pyplot.grid(True)
pyplot.xlabel(r'$n_x$', fontsize=18)
pyplot.ylabel(r'$L_2$-norm of the error', fontsize=18)
pyplot.xlim(1, 1000)
pyplot.ylim(1e-5, 1e-2)
pyplot.loglog(nx_values, error, color='k', ls='--', lw=2, marker='o');
```

```python
error
```

    array([  3.59569224e-03,   1.34923182e-03,   4.90905474e-04,
             1.76021088e-04,   6.26681531e-05])

Wait, convergence is not that great now! It's not as good as second order, but not as bad as first order. *What is going on?*

Remember our implementation of the boundary conditions? We used

\begin{equation}\frac{T^{n}_{N-1} - T^{n}_{N-2}}{\Delta x} = q.\end{equation}

Well, that is a **first-order** approximation!

But, why doesn't this affect our solution at an earlier time? Initially, temperature on the right side of the rod is zero and the gradient is very small in that region; at that point in time, errors there were negligible. Once temperature starts picking up, we start having problems.

**Boundary conditions can affect the convergence and accuracy of your solution!**

---
###### The cell below loads the style of the notebook

```python
from IPython.core.display import HTML
css_file = '../../styles/numericalmoocstyle.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=Nixie+One' rel='stylesheet' type='text/css'>
<style>

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

#notebook_panel { /* main background */
    background: rgb(245,245,245);
}

div.cell { /* set cell width */
    width: 750px;
}

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

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

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

/* Put a solid color box around each cell and its output, visually linking them*/
div.cell.code_cell {
    background-color: rgb(256,256,256);
    border-radius: 0px;
    padding: 0.5em;
    margin-left:1em;
    margin-top: 1em;
}

div.text_cell_render{
    font-family: 'Alegreya Sans' sans-serif;
    line-height: 140%;
    font-size: 125%;
    font-weight: 400;
    width:600px;
    margin-left:auto;
    margin-right:auto;
}

/* Formatting for header cells */
.text_cell_render h1 {
    font-family: 'Nixie One', serif;
    font-style:regular;
    font-weight: 400;
    font-size: 45pt;
    line-height: 100%;
    color: rgb(0,51,102);
    margin-bottom: 0.5em;
    margin-top: 0.5em;
    display: block;
}

.text_cell_render h2 {
    font-family: 'Nixie One', serif;
    font-weight: 400;
    font-size: 30pt;
    line-height: 100%;
    color: rgb(0,51,102);
    margin-bottom: 0.1em;
    margin-top: 0.3em;
    display: block;
}

.text_cell_render h3 {
    font-family: 'Nixie One', serif;
    margin-top:16px;
    font-size: 22pt;
    font-weight: 600;
    margin-bottom: 3px;
    font-style: regular;
    color: rgb(102,102,0);
}

.text_cell_render h4 {    /*Use this for captions*/
    font-family: 'Nixie One', serif;
    font-size: 14pt;
    text-align: center;
    margin-top: 0em;
    margin-bottom: 2em;
    font-style: regular;
}

.text_cell_render h5 {  /*Use this for small titles*/
    font-family: 'Nixie One', sans-serif;
    font-weight: 400;
    font-size: 16pt;
    color: rgb(163,0,0);
    font-style: italic;
    margin-bottom: .1em;
    margin-top: 0.8em;
    display: block;
}

.text_cell_render h6 { /*use this for copyright note*/
    font-family: 'PT Mono', sans-serif;
    font-weight: 300;
    font-size: 9pt;
    line-height: 100%;
    color: grey;
    margin-bottom: 1px;
    margin-top: 1px;
}

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

</style>


---
title: "Integration test for your Gibbs sampler"
excerpt: "Making your MCMC code ready for production."
---

This summer at [Civis Analytics](https://www.civisanalytics.com/), my project was to derive and implement a Gibbs sampler for one of the company's core algorithms. While I've written several Gibbs samplers, they are one-off projects that pay little attention to robustness and maintainability. Now that my code will be read by other engineers and used to make business decisions worth many (undisclosed number) of dollars, I think a lot more about how to productionize MCMC code.

Traditionally, the way a Gibbs sampler is taught in [lecture notes](https://people.duke.edu/~ccc14/sta-663/MCMC.html#gibbs-sampler) or [tutorials](https://theclevermachine.wordpress.com/2012/11/05/mcmc-the-gibbs-sampler/) is a giant loop iterating through each parameter's full conditional. This no-frill Gibbs sampler is fast to run, relying on `numpy` operations. However, it can only be integration tested, which is slow and non-deterministic.

By contrast, [Grosse and Duvenaud (2014)](https://arxiv.org/abs/1412.5218) proposes writing the Gibbs sampler in a modular way, in which each parameter's full conditional is refactored into its own function. This modular Gibbs sampler facilitates unit testing, which is fast and deterministic. However, it is slow to run due to the overhead of creating classes and objects.

In my work, I combine both approaches to ensure that our production code is both fast to run and to test. In this first post, I write the fast, no-frill Gibbs sampler that can be integration tested. Then, in [the second post]({{ site.baseurl }}{% post_url 2017-08-16-productionizing-Gibbs-sampler-unit-test %}), I write the modular Gibbs sampler that can be unit tested, following [Grosse and Duvenaud (2014)](https://arxiv.org/abs/1412.5218). Finally, we can use the result from the modular, unit-tested Gibbs sampler to check the fast, no-frill Gibbs sampler. Once we're assured that the fast, no-frill Gibbs is correct, we can use it in production.

## Fast, no-frill Gibbs sampler

As an example throughout this series, I'll use the Gibbs sampler for the univariate normal model. It's a simple and common model, which you can find implemented in other tutorials and textbooks to compare with my approach.

Below is the model's setup and its full conditionals, taken from [Hoff (2009), ch. 5](https://www.amazon.com/Bayesian-Statistical-Methods-Springer-Statistics/dp/1441928286). You can skip the math -- I write it here in case you want to check that it matches my code.

**Likelihood**

$$Data = Y_1, \dots, Y_n \sim i.i.d. N(\theta, \sigma^2)$$

**Prior**

$$
\begin{align}
p(\theta) &\sim N(\mu_0, \tau^2_0) \\
p(\sigma^2) &\sim Inverse-Gamma(\nu_0 / 2, \nu_0 \sigma^2_0 / 2)
\end{align}
$$

**Full conditional**

$$
\begin{align}
p(\theta | \tilde \sigma^2, Data) &= N(\mu_n, \tau^2_n) \\
p( \sigma^2 | \theta, Data) &= Inverse-Gamma(\frac{\nu_n}{2}, \frac{\nu_n \sigma_n^2(\theta)}{2})
\end{align}
$$

where

$$
\begin{align}
\tau_n^2 = \frac{1}{\frac{1}{\tau_0^2} + \frac{n}{\sigma^2}} \qquad &\text{and} \qquad \mu_n = \tau_n^2 \left( \frac{\mu_0}{\tau_0^2} + \frac{n\bar y}{\sigma^2} \right) \\
\nu_n = \nu_0 + n \qquad &\text{and} \qquad \sigma^2_n(\theta) = \frac{1}{\nu_n} \left[ \nu_0\sigma_0^2 + \sum (y_i - \theta)^2\right]
\end{align}
$$

We then draw \\( \theta \\) and \\( \sigma^2 \\) iteratively:

$$
\begin{align}
\theta &\sim p(\theta_s | \sigma_{s-1}^2, data) \\
\sigma^2 &\sim p(\sigma_{s}^2 | \theta_s, data)
\end{align}
$$

**We implement this iterative updating in a big loop as follows**

```python
import numpy as np
import scipy.stats
import pandas as pd

def gibbs_simple(S, y, prior, rng=None):
    """
    A no-frill Gibbs sampler

    Parameters
    ----------
    S : int, the number of samples
    y : data vector
    prior : dict, prior parameters
    rng : int, random seed
    """
    n = len(y)
    ybar = y.mean()

    # Initialize storage
    theta_samples = np.empty(S)
    sigma2_samples = np.empty(S)

    # Starting value as the sample variance and mean
    sigma2_samples[0] = y.var()
    theta_samples[0] = ybar

    # Big loop
    for s in range(1, S):
        # Update theta
        tau2_n = 1 / (1 / prior['tau2_0'] + n / sigma2_samples[s - 1])
        mu_n = tau2_n * (prior['mu_0'] / prior['tau2_0'] + n * ybar / sigma2_samples[s - 1])
        theta_samples[s] = scipy.stats.norm(mu_n, np.sqrt(tau2_n)).rvs(random_state=rng)

        # Update sigma2
        nu_n = prior['nu_0'] + n
        nu_sigma2_n = prior['nu_0'] * prior['sigma2_0'] + sum((y - theta_samples[s]) ** 2)
        sigma2_samples[s] = scipy.stats.invgamma(nu_n / 2, scale=nu_sigma2_n / 2).rvs(random_state=rng)

    return {'theta': theta_samples, 'sigma2': sigma2_samples}
```

## Integration test

This no-frill Gibbs sampler is fast, taking full advantage of `numpy` operations. We can test its correctness by generating data from known parameters, then check whether the Gibbs sampler produces posterior estimates that are "close enough" to the known parameter values. From a programming perspective, this can be considered an [integration test](http://softwaretestingfundamentals.com/integration-testing/).

```python
# Generate the data from known parameters
true_theta = 2
true_sigma2 = 3.5
y = np.random.normal(true_theta, np.sqrt(true_sigma2), size = 1000)
prior = {'mu_0': 0, 'tau2_0': 10000, 'nu_0': 1, 'sigma2_0': 1}

# Run the Gibbs sampler
samples_gibbs_simple = gibbs_simple(1000, y, prior)
```

We can inspect the result visually, checking that the posterior distributions (the black density plot) covers the true parameter values (the red line).

```python
from plotnine import *
from plotnine import options

options.figure_size = (3, 3)
ggplot(data=pd.DataFrame(samples_gibbs_simple)) + \
    geom_density(aes(x='theta')) + \
    geom_vline(aes(xintercept=true_theta), color='red')
```

```python
options.figure_size = (3, 3)
ggplot(data=pd.DataFrame(samples_gibbs_simple)) + \
    geom_density(aes(x='sigma2')) + \
    geom_vline(aes(xintercept=true_sigma2), color='red')
```

We should also encode this visual check into a formal test, which can be put in an automated testing framework like `pytest` or `nose`:

```python
# Checking that the posterior mean of a parameter is "close enough" to its true value
assert np.allclose(samples_gibbs_simple['theta'].mean(), true_theta, rtol=0.99)
assert np.allclose(samples_gibbs_simple['sigma2'].mean(), true_sigma2, rtol=0.99)
```

The integration test "works," but suffers from several drawbacks. First, the integration test can be misleading because it relies on checking a "close enough" result. Second, the integration test is slow because it has to generate the posterior sample each time. In [Part 2]({{ site.baseurl }}{% post_url 2017-08-16-productionizing-Gibbs-sampler-unit-test %}), I will discuss how to avoid these problems with a modular Gibbs sampler that can be unit-tested.