# Gini Mean Difference Portfolio Optimization

In this notebook we show how we can solve a hard problem using some reformulations.

## 1. Gini Optimization

### 1.1 Original formulation

The Gini mean difference (GMD) is a measure of dispersion and it was introduced in the context of portfolio optimization by __[Yitzhaki (1982)](https://www.researchgate.net/publication/4900733_Stochastic_Dominance_Mean_Variance_and_Gini%27s_Mean_Difference)__. However, this model is not used by practitioners due to the original formulation having a number of variables that increases proportional to $T(T-1)/2$, where $T$ is the number of observations. The original model is presented as follows:

$$
\begin{equation}
\begin{aligned}
& \underset{x,\, d}{\text{min}} & & \frac{1}{T(T-1)} \sum^{T}_{i=1} \sum^{T}_{j > i} d_{i,j} \\
& \text{s.t.} & & \mu x \geq \bar{\mu} \\
& & &  d_{i,j} \geq (r_{i} -r_{j})x \; ; \; \forall \; i,j =1, \ldots, T \; ; \; i < j \\
& & &  d_{i,j} \geq -(r_{i} -r_{j})x \\
& & &  \sum_{i=1}^{N} x_i = 1 \\
& & &  x_{i} \geq 0 \; ; \; \forall \; i =1, \ldots, N \\
& & &  d_{i,j} \geq 0 \; ; \; \forall \; i,j =1, \ldots, T \\
\end{aligned}
\end{equation}
$$

Where $r_{i}$ is the vector of returns in period $i$ and $d$ is an auxiliary variable.

### 1.2 Murray's reformulation

To increase the efficiency of the problem above, __[Murray (2022)](https://github.com/cvxpy/cvxpy/issues/1585)__ proposed the following reformulation:

$$
\begin{equation}
\begin{aligned}
& \underset{x,\, d}{\text{min}} & & \frac{1}{T(T-1)} \sum^{T}_{i=1} \sum^{T}_{j > i} d_{i,j} \\
& \text{s.t.} & & \mu x \geq \bar{\mu} \\
& & & y = r x \\
& & & d \geq M y \\
& & & d \geq -M y \\
& & &  \sum_{i=1}^{N} x_i = 1 \\
& & &  x_{i} \geq 0 \; ; \; \forall \; i =1, \ldots, N \\
& & &  d_{i,j} \geq 0 \; ; \; \forall \; i,j =1, \ldots, T \\
\end{aligned}
\end{equation}
$$

where
$$
M = \begin{bmatrix}
\left. \begin{matrix}
-1 & 1 & 0 & 0 & \ldots & 0 & 0\\
-1 & 0 & 1 & 0 & \ldots & 0 & 0\\
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\
-1 & 0 & 0 & 0 & \ldots & 0 & 1 \\
\end{matrix} \right \} & T-1\\
\left. \begin{matrix}
0 & -1 & 1 & 0 & \ldots & 0 & 0\\
0 &  -1 & 0 & 1 & \ldots & 0 & 0\\
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots &  \vdots \\
0 & -1 & 0 & 0 & \ldots & 0 & 1 \\
\end{matrix} \right \} & T-2\\
\vdots \\
\underbrace{ \left. \begin{matrix}
0 & 0 & 0 & 0 & \ldots & -1 & 1 \\
\end{matrix} \right \} }_{T} & 1 \\
\end{bmatrix}
$$

This reformulation is more efficient for medium scale problems (T<800).

### 1.3 Cajas's reformulation:

__[Cajas (2021)](https://papers.ssrn.com/sol3/papers.cfm?abstract_id=3988927)__ proposed an alternative reformulation based on the ordered weighted averaging (OWA) operator optimization model for monotonic weights proposed by __[Chassein and Goerigk (2015)](https://kluedo.ub.uni-kl.de/frontdoor/deliver/index/docId/3899/file/paper.pdf)__. This formulation works better for large scale problems (T>=800). This formulation is presented as follows:

$$
\begin{equation}
\begin{aligned}
& \min_{\alpha, \, \beta, \, x, \, y} & & \sum^{T}_{i=1} \alpha_{i} + \beta_{i}  \\
& \text{s.t.} & & \mu x \geq \bar{\mu} \\
& & & r x = y \\
& & & \alpha_{i} + \beta_{j} \geq w_{i} y_{j}  \;\;\;\; \forall \; i,j =1, \ldots, T \\
& & &  \sum_{i=1}^{N} x_i = 1 \\
& & &  x_i \geq 0 \; ; \; \forall \; i =1, \ldots, N \\
\end{aligned}
\end{equation}
$$

where $w_{i} =  2 \left ( \frac{2i - 1 - T}{T(T-1)} \right )$.

```python
import numpy as np
import pandas as pd
import cvxpy as cp
import mosek
import scipy.stats as st
from timeit import default_timer as timer
from datetime import timedelta
import warnings

warnings.filterwarnings("ignore")

def gini(mu, returns, D, assets, lift=0):
    (T, N) = returns.shape

    d = cp.Variable((int(T * (T - 1) / 2),1))
    w = cp.Variable((N,1))
    constraints = []

    if lift in ['Murray', 'Yitzhaki']:  # use Murray's reformulation
        if lift == 'Murray':
            ret_w = cp.Variable((T,1))
            constraints.append(ret_w == returns @ w)
            mat = np.zeros((d.shape[0], T))
            """
            We need to create a vector that has the following entries:
                ret_w[i] - ret_w[j]
            for j in range(T), for i in range(j+1, T).
            We do this by building a numpy array of mostly 0's and 1's.
            (It would be better to use SciPy sparse matrix objects.)
            """
            ell = 0
            for j in range(T):
                for i in range(j+1, T):
                    # write to mat so that (mat @ ret_w)[ell] == var_i - var_j
                    mat[ell, i] = 1
                    mat[ell, j] = -1
                    ell += 1
            all_pairs_ret_diff = mat @ ret_w
        elif lift == 'Yitzhaki':  # use the original formulation
            all_pairs_ret_diff = D @ w

        constraints += [d >= all_pairs_ret_diff,
                        d >= -all_pairs_ret_diff,
                        w >= 0,
                        cp.sum(w) == 1,
                       ]

        risk = cp.sum(d) / ((T - 1) * T)

    elif lift == 'Cajas':
        a = cp.Variable((T,1))
        b = cp.Variable((T,1))
        y = cp.Variable((T,1))

        owa_w = []
        for i in range(1,T+1):
            owa_w.append(2*i - 1 - T)
        owa_w = np.array(owa_w) / (T * (T-1))

        constraints = [returns @ w == y,
                       w >= 0,
                       cp.sum(w) == 1]

        for i in range(T):
            constraints += [a[i] + b >= cp.multiply(owa_w[i], y)]

        risk = cp.sum(a + b)

    objective = cp.Minimize(risk * 1000)

    prob = cp.Problem(objective, constraints)

    try:
        prob.solve(solver=cp.MOSEK,
                   mosek_params={mosek.iparam.num_threads: 2},
                   verbose=False,)
        w = pd.DataFrame(w.value)
        w.index = assets
        w = w / w.sum()
    except:
        w = None
    return w

```

```python
####################################
# Calculating the Gini Portfolio
# using three formulations
####################################

rs = np.random.RandomState(123)

sizes = [100]
data = {}
weights = {}
lifts = ['Yitzhaki', 'Murray', 'Cajas']
k = 0
T = [200, 500, 700, 1000]

for t in T:
    for n in sizes:

        cov = rs.rand(n,n) * 1.5 - 0.5
        cov = cov @ cov.T/1000 + np.diag(rs.rand(n) * 0.7 + 0.3)/1000
        mean = np.zeros(n) + 1/1000

        Y = st.multivariate_normal.rvs(mean=mean, cov=cov, size=t, random_state=rs)
        Y = pd.DataFrame(Y)
        assets = ['Asset '+ str(i) for i in range(1,n+1)]
        mu = Y.mean().to_numpy()
        returns = Y.to_numpy()

        D = np.array([]).reshape(0,len(assets))
        for j in range(0, returns.shape[0]-1):
            D = np.concatenate((D, returns[j+1:] - returns[j,:]), axis=0)

        for lift in lifts:
            name = str(lift) + '-' + str(t) + '-' + str(n)
            data[name] = []
            weights[name] = []
            if t >= 700 and lift == 'Yitzhaki':
                continue
            else:
                start = timer()
                w = gini(mu, returns, D, assets, lift=lift)
                end = timer()
                data[name].append(timedelta(seconds=end-start).total_seconds())
                weights[name].append(w)

            k += 1
            print(name)
```

    Yitzhaki-200-100
    Murray-200-100
    Cajas-200-100
    Yitzhaki-500-100
    Murray-500-100
    Cajas-500-100
    Murray-700-100
    Cajas-700-100
    Murray-1000-100
    Cajas-1000-100

```python
keys = list(data.keys())
for i in keys:
    if len(data[i]) == 0:
        del data[i]

pd.options.display.float_format = '{:.2f}'.format

a = pd.DataFrame(data).T
a.columns = ['Time in Seconds']
display(a)
```

      Time in Seconds

      Yitzhaki-200-100
      9.55

      Murray-200-100
      0.88

      Cajas-200-100
      1.59

      Yitzhaki-500-100
      99.33

      Murray-500-100
      11.49

      Cajas-500-100
      13.67

      Murray-700-100
      27.19

      Cajas-700-100
      29.12

      Murray-1000-100
      96.62

      Cajas-1000-100
      80.50

As we can see, as the number of observations $T$ increases the formulation proposed by Cajas (2021) becomes more efficient than Yitzhaki's and Murray's formulations.

```python
b = pd.DataFrame([])
for i in weights.keys():
    if len(weights[i]) == 0:
        continue
    weights[i][0].columns = [i]
    b = pd.concat([b , mu @ weights[i][0]],axis=0)

pd.options.display.float_format = '{:.4%}'.format

display(b)
```

      0

      Yitzhaki-200-100
      -0.0413%

      Murray-200-100
      -0.0413%

      Cajas-200-100
      -0.0414%

      Yitzhaki-500-100
      -0.2315%

      Murray-500-100
      -0.2315%

      Cajas-500-100
      -0.2313%

      Murray-700-100
      0.1601%

      Cajas-700-100
      0.1602%

      Murray-1000-100
      0.0339%

      Cajas-1000-100
      0.0339%

If we calculate the expected returns for each formulation, we can see that the three models give us the same results, which means that these formulations give us the same solution.

For more portfolio optimization models and applications, you can see the CVXPY based library __[Riskfolio-Lib](https://github.com/dcajasn/Riskfolio-Lib)__.

Finally, I hope you liked this example.


# Solving a mixed Neumann-Dirichlet problem

### Background

With Bempp, it is possible to define operators only on segments of a given domain. This makes it possible to solve mixed Neumann-Dirichlet problems. In this tutorial, we solve the Laplace equation inside the unit cube with unit Dirichlet boundary conditions on two sides and unit Neumann boundary conditions on the other four sides.

Denote by $\Gamma_D$ the part of the boundary that holds the Dirichlet boundary conditions and by $\Gamma_N$ the boundary part that holds the Neumann boundary conditions. We denote by $t\in\Gamma_D$ the unknown Neumann data and by $u\in\Gamma_N$ the unknown Dirichlet data. The given Dirichlet data on $\Gamma_D$ is denoted by $g_D$ and the given Neumann data on $\Gamma_N$ is denoted by $g_N$.

From Green's representation theorem it follows that
$$
\begin{align}
\left[\mathsf{V}t\right] (\mathbf{x}) - \left[\mathsf{K}u\right] (\mathbf{x}) &= \left[\tfrac{1}{2}\mathsf{Id} + \mathsf{K}\right]g_D(\mathbf{x}) - \mathsf{V}g_N(\mathbf{x}),\quad \mathbf{x}\in\Gamma_D\\
\left[\mathsf{W}u\right] (\mathbf{x}) + \left[\mathsf{K}'t\right] (\mathbf{x}) &=\left[\tfrac{1}{2}\mathsf{Id} - \mathsf{K}'\right]g_N(\mathbf{x}) - \mathsf{W}g_D(\mathbf{x}),\quad \mathbf{x}\in\Gamma_N
\end{align}
$$
Here (as usual) $\mathsf{V}$, $\mathsf{K}$, $\mathsf{K}'$, $\mathsf{W}$ are the single layer, double layer, adjoint double layer and hypersingular boundary operators.

The difficulty in the implementation is the definition of the discrete function spaces and the treatment of degrees of freedom (dofs) that lie on the interface between $\Gamma_N$ and $\Gamma_D$. In the following, we will go through the implementation and point out how to correctly define all spaces involved.

### Implementation

We start with the usual imports and set the plotting backend to Jupyter Notebook.

```python
import bempp.api
import numpy as np
```

We now define the domain. We use a standard unit cube. In the corresponding function all sides of the cube are already associated with different domain indices. We associate the indices 1 and 3 with the Dirichlet boundary and the other indices with the neumann boundary.

```python
grid = bempp.api.shapes.cube()
dirichlet_segments = [1, 3]
neumann_segments = [2, 4, 5, 6]
```

We can now define the spaces. For the Neumann data, we use discontinuous polynomial basis functions of order 1. For the Dirichlet data, we use continuous basis functions of local polynomial order 2.

We need global spaces for the Dirichlet and Neumann data and suitable spaces on the segments. The space definitions are as follows:

* The ``neumann_space_dirichlet_segment`` space holds the unknown Neumann data $t$ on $\Gamma_D$.

* The ``neumann_space_neumann_segment`` space is defined on $\Gamma_N$ and holds the Neumann data on this segment.

* The ``dirichlet_space_dirichlet_segment`` space is a space of continuous basis functions that holds the Dirichlet data on $\Gamma_D$. Basis functions are allowed to extend into the elements adjacent to $\Gamma_D$. We hence set the parameter ``include_boundary_dofs`` to true. This ensures that we include dofs on the boundary of the segment, and to ensure that the basis functions extend continuously into the complement of the segment we additionally set ``truncate_at_segment_edge`` to false.

* The ``dirichlet_space_neumann_segment`` is defined similarly to the ``dirichlet_space_dirichlet_segment`` but on the open segment $\Gamma_N$. Hence, we require that ``include_boundary_dofs`` is false, which is the default.

```python
global_neumann_space = bempp.api.function_space(grid, "DP", 0)
global_dirichlet_space = bempp.api.function_space(grid, "P", 1)

neumann_space_dirichlet_segment = bempp.api.function_space(
    grid, "DP", 0, segments=dirichlet_segments)

neumann_space_neumann_segment = bempp.api.function_space(
    grid, "DP", 0, segments=neumann_segments)

dirichlet_space_dirichlet_segment = bempp.api.function_space(
    grid, "P", 1, segments=dirichlet_segments, include_boundary_dofs=True,
    truncate_at_segment_edge=False)

dirichlet_space_neumann_segment = bempp.api.function_space(
    grid, "P", 1, segments=neumann_segments)

dual_dirichlet_space = bempp.api.function_space(
    grid, "P", 1, segments=dirichlet_segments,
    include_boundary_dofs=True)
```

In the following, we define all operators on the corresponding spaces and the overall blocked operator.

```python
slp_DD = bempp.api.operators.boundary.laplace.single_layer(
    neumann_space_dirichlet_segment,
    dirichlet_space_dirichlet_segment,
    neumann_space_dirichlet_segment)

dlp_DN = bempp.api.operators.boundary.laplace.double_layer(
    dirichlet_space_neumann_segment,
    dirichlet_space_dirichlet_segment,
    neumann_space_dirichlet_segment)

adlp_ND = bempp.api.operators.boundary.laplace.adjoint_double_layer(
    neumann_space_dirichlet_segment,
    neumann_space_neumann_segment,
    dirichlet_space_neumann_segment)

hyp_NN = bempp.api.operators.boundary.laplace.hypersingular(
    dirichlet_space_neumann_segment,
    neumann_space_neumann_segment,
    dirichlet_space_neumann_segment)

slp_DN = bempp.api.operators.boundary.laplace.single_layer(
    neumann_space_neumann_segment,
    dirichlet_space_dirichlet_segment,
    neumann_space_dirichlet_segment)

dlp_DD = bempp.api.operators.boundary.laplace.double_layer(
    dirichlet_space_dirichlet_segment,
    dirichlet_space_dirichlet_segment,
    neumann_space_dirichlet_segment)

id_DD = bempp.api.operators.boundary.sparse.identity(
    dirichlet_space_dirichlet_segment,
    dirichlet_space_dirichlet_segment,
    neumann_space_dirichlet_segment)

adlp_NN = bempp.api.operators.boundary.laplace.adjoint_double_layer(
    neumann_space_neumann_segment,
    neumann_space_neumann_segment,
    dirichlet_space_neumann_segment)

id_NN = bempp.api.operators.boundary.sparse.identity(
    neumann_space_neumann_segment,
    neumann_space_neumann_segment,
    dirichlet_space_neumann_segment)

hyp_ND = bempp.api.operators.boundary.laplace.hypersingular(
    dirichlet_space_dirichlet_segment,
    neumann_space_neumann_segment,
    dirichlet_space_neumann_segment)

blocked = bempp.api.BlockedOperator(2, 2)

blocked[0, 0] = slp_DD
blocked[0, 1] = -dlp_DN
blocked[1, 0] = adlp_ND
blocked[1, 1] = hyp_NN
```

Next, we define the functions of the Dirichlet and Neumann data and their discretisations on the corresponding segments.

```python
@bempp.api.real_callable
def dirichlet_data(x, n, domain_index, res):
    res[0] = 1

@bempp.api.real_callable
def neumann_data(x, n, domain_index, res):
    res[0] = 1

dirichlet_grid_fun = bempp.api.GridFunction(
    dirichlet_space_dirichlet_segment,
    fun=dirichlet_data,
    dual_space=dual_dirichlet_space)

neumann_grid_fun = bempp.api.GridFunction(
    neumann_space_neumann_segment,
    fun=neumann_data,
    dual_space=dirichlet_space_neumann_segment)

rhs_fun1 = (.5 * id_DD + dlp_DD) * dirichlet_grid_fun \
           - slp_DN * neumann_grid_fun
rhs_fun2 = - hyp_ND * dirichlet_grid_fun \
           + (.5 * id_NN - adlp_NN) * neumann_grid_fun
```

We can now discretise and solve the blocked operator system. We solve without preconditioner. This would cause problems if we were to further increase the degree of the basis functions. Note that Bempp automatically handles the block structure correctly. We have a $2\times 2$ system of operators, hand over as right-hand side two grid functions, and are returned two grid functions as solution.

```python
(neumann_solution, dirichlet_solution), _ = bempp.api.linalg.gmres(blocked, [rhs_fun1, rhs_fun2])
```

We want to recombine the computed Dirichlet and Neumann data with the corresponding known data in order to get Dirichlet and Neumann grid functions defined on the whole grid. To achieve this we define identity operators from $\Gamma_N$ and $\Gamma_D$ into the global Dirichlet and Neumann spaces.

```python
neumann_imbedding_dirichlet_segment = \
    bempp.api.operators.boundary.sparse.identity(
        neumann_space_dirichlet_segment,
        global_neumann_space,
        global_neumann_space)

neumann_imbedding_neumann_segment = \
    bempp.api.operators.boundary.sparse.identity(
        neumann_space_neumann_segment,
        global_neumann_space,
        global_neumann_space)

dirichlet_imbedding_dirichlet_segment = \
    bempp.api.operators.boundary.sparse.identity(
        dirichlet_space_dirichlet_segment,
        global_dirichlet_space,
        global_dirichlet_space)

dirichlet_imbedding_neumann_segment = \
    bempp.api.operators.boundary.sparse.identity(
        dirichlet_space_neumann_segment,
        global_dirichlet_space,
        global_dirichlet_space)

dirichlet = (dirichlet_imbedding_dirichlet_segment * dirichlet_grid_fun +
             dirichlet_imbedding_neumann_segment * dirichlet_solution)

neumann = (neumann_imbedding_neumann_segment * neumann_grid_fun +
           neumann_imbedding_dirichlet_segment * neumann_solution)
```

Finally, we export the solution.

```python
bempp.api.export('mixed_dirichlet_neumann_laplace_solution.msh', grid_function=dirichlet)
```

It should look as follows.

```python
from IPython.display import Image
Image("mixed_dirichlet_neumann_laplace_solution.png")
```


# Input Driven Observations ("GLM-HMM")

Notebook prepared by Zoe Ashwood: feel free to email me with feedback or questions (zashwood at cs dot princeton dot edu).

This notebook demonstrates the "InputDrivenObservations" class, and illustrates its use in the context of modeling decision-making data as in Ashwood et al. (2020) ([Mice alternate between discrete strategies during perceptual
decision-making](https://www.biorxiv.org/content/10.1101/2020.10.19.346353v1.full.pdf)).

Compared to the model considered in the notebook ["2 Input Driven HMM"](https://github.com/lindermanlab/ssm/blob/master/notebooks/2%20Input%20Driven%20HMM.ipynb), Ashwood et al. (2020) assumes a stationary transition matrix where transition probabilities *do not* depend on external inputs. However, observation probabilities now *do* depend on external covariates according to:

for $c \neq C$:
$$
\begin{align}
\Pr(y_t = c \mid z_{t} = k, u_t, w_{kc}) =
\frac{\exp\{w_{kc}^\mathsf{T} u_t\}}
{1+\sum_{c'=1}^{C-1} \exp\{w_{kc'}^\mathsf{T} u_t\}}
\end{align}
$$

and for $c = C$:
$$
\begin{align}
\Pr(y_t = c \mid z_{t} = k, u_t, w_{kc}) =
\frac{1}
{1+\sum_{c'=1}^{C-1} \exp\{w_{kc'}^\mathsf{T} u_t\}}
\end{align}
$$

where $c \in \{1, ..., C\}$ indicates the categorical class for the observation, $u_{t} \in \mathbb{R}^{M}$ is the set of input covariates, and $w_{kc} \in \mathbb{R}^{M}$ is the set of input weights associated with state $k$ and class $c$. These weights, along with the transition matrix and initial state probabilities, will be learned.

In Ashwood et al. (2020), $C = 2$ as $y_{t}$ represents the binary choice made by an animal during a 2AFC (2-Alternative Forced Choice) task. The above equations then reduce to:

$$
\begin{align}
\Pr(y_t = 0 \mid z_{t} = k, u_t, w_{k}) =
\frac{\exp\{w_{k}^\mathsf{T} u_t\}}
{1 + \exp\{w_{k}^\mathsf{T} u_t\}} = \frac{1}
{1 + \exp\{-w_{k}^\mathsf{T} u_t\}}.
\end{align}
$$

$$
\begin{align}
\Pr(y_t = 1 \mid z_{t} = k, u_t, w_{k}) =
\frac{1}
{1 + \exp\{w_{k}^\mathsf{T} u_t\}}.
\end{align}
$$

and only a single weight vector, $w_{k}$, is associated with each state.

## 1. Setup
The line `import ssm` imports the package for use. Here, we have also imported a few other packages for plotting.

```python
import autograd.numpy.random as npr
```

```python
import numpy as np
import numpy.random as npr
import matplotlib.pyplot as plt
import ssm
from ssm.util import find_permutation

npr.seed(0)
```

## 2. Input Driven Observations
We create a HMM with input-driven observations and 'standard' (stationary) transitions with the following line:
```python
        ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs", observation_kwargs=dict(C=num_categories), transitions="standard")
```

As in Ashwood et al. (2020), we are going to model an animal's binary choice data during a decision-making task, so we will set `num_categories=2` because the animal only has two options available to it. We will also set `obs_dim = 1` because the dimensionality of the observation data is 1 (if we were also modeling, for example, the binned reaction time of the animal, we could set `obs_dim = 2`).  For the sake of simplicity, we will assume that an animal's choice in a particular state is only affected by the external stimulus associated with that particular trial, and its innate choice bias. Thus, we will set `input_dim = 2` and we will simulate input data that resembles sequences of stimuli in what follows.  In Ashwood et al. (2020), they found that many mice used 3 decision-making states when performing 2AFC tasks. We will, thus, set `num_states = 3`.

### 2a. Initialize GLM-HMM

```python
# Set the parameters of the GLM-HMM
num_states = 3        # number of discrete states
obs_dim = 1           # number of observed dimensions
num_categories = 2    # number of categories for output
input_dim = 2         # input dimensions

# Make a GLM-HMM
true_glmhmm = ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs",
                   observation_kwargs=dict(C=num_categories), transitions="standard")
```

### 2b. Specify parameters of generative GLM-HMM

Let's update the weights and transition matrix for the true GLM-HMM so as to bring the GLM-HMM to the parameter regime that real animals use (according to Ashwood et al. (2020)):

```python
gen_weights = np.array([[[6, 1]], [[2, -3]], [[2, 3]]])
gen_log_trans_mat = np.log(np.array([[[0.98, 0.01, 0.01], [0.05, 0.92, 0.03], [0.03, 0.03, 0.94]]]))
true_glmhmm.observations.params = gen_weights
true_glmhmm.transitions.params = gen_log_trans_mat
```

```python
# Plot generative parameters:
fig = plt.figure(figsize=(8, 3), dpi=80, facecolor='w', edgecolor='k')
plt.subplot(1, 2, 1)
cols = ['#ff7f00', '#4daf4a', '#377eb8']
for k in range(num_states):
    plt.plot(range(input_dim), gen_weights[k][0], marker='o',
             color=cols[k], linestyle='-',
             lw=1.5, label="state " + str(k+1))
plt.yticks(fontsize=10)
plt.ylabel("GLM weight", fontsize=15)
plt.xlabel("covariate", fontsize=15)
plt.xticks([0, 1], ['stimulus', 'bias'], fontsize=12, rotation=45)
plt.axhline(y=0, color="k", alpha=0.5, ls="--")
plt.legend()
plt.title("Generative weights", fontsize = 15)

plt.subplot(1, 2, 2)
gen_trans_mat = np.exp(gen_log_trans_mat)[0]
plt.imshow(gen_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(gen_trans_mat.shape[0]):
    for j in range(gen_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(gen_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.ylabel("state t", fontsize = 15)
plt.xlabel("state t+1", fontsize = 15)
plt.title("Generative transition matrix", fontsize = 15)
```

### 2c. Create external input sequences

Simulate an example set of external inputs for each trial in a session. We will create an array of size `(num_sess x num_trials_per_sess x num_covariates)`. As in Ashwood et al. (2020), for each trial in a session we will include the stimulus presented to the animal at that trial, as well as a '1' as the second covariate (so as to capture the animal's innate bias for one of the two options available to it). We will simulate stimuli sequences so as to resemble the sequences of stimuli in the International Brain Laboratory et al. (2020) task.

```python
num_sess = 20 # number of example sessions
num_trials_per_sess = 100 # number of trials in a session
inpts = np.ones((num_sess, num_trials_per_sess, input_dim)) # initialize inpts array
stim_vals = [-1, -0.5, -0.25, -0.125, -0.0625, 0, 0.0625, 0.125, 0.25, 0.5, 1]
inpts[:,:,0] = np.random.choice(stim_vals, (num_sess, num_trials_per_sess)) # generate random sequence of stimuli
inpts = list(inpts) #convert inpts to correct format
```

```python
inpts
```

### 2d. Simulate states and observations with generative model

```python
# Generate a sequence of latents and choices for each session
true_latents, true_choices = [], []
for sess in range(num_sess):
    true_z, true_y = true_glmhmm.sample(num_trials_per_sess, input=inpts[sess])
    true_latents.append(true_z)
    true_choices.append(true_y)
```

```python
true_z
```

```python
# Calculate true loglikelihood
true_ll = true_glmhmm.log_probability(true_choices, inputs=inpts)
print("true ll = " + str(true_ll))
```

## 3. Fit GLM-HMM and perform recovery analysis

### 3a. Maximum Likelihood Estimation

Now we instantiate a new GLM-HMM and check that we can recover the generative parameters in simulated data:

```python
new_glmhmm = ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs",
                   observation_kwargs=dict(C=num_categories), transitions="standard")

N_iters = 200 # maximum number of EM iterations. Fitting with stop earlier if increase in LL is below tolerance specified by tolerance parameter
fit_ll = new_glmhmm.fit(true_choices[0], inputs=inpts[0], method="em", num_iters=N_iters, tolerance=10**-4)
```

```python
new_glmhmm.transitions.transition_matrices(true_choices[0], inpts[0], None, None)
```

```python
new_glmhmm.observations.log_likelihoods(true_choices[0], inpts[0], None, None).shape
```

```python
# Plot the log probabilities of the true and fit models. Fit model final LL should be greater
# than or equal to true LL.
fig = plt.figure(figsize=(4, 3), dpi=80, facecolor='w', edgecolor='k')
plt.plot(fit_ll, label="EM")
plt.plot([0, len(fit_ll)], true_ll * np.ones(2), ':k', label="True")
plt.legend(loc="lower right")
plt.xlabel("EM Iteration")
plt.xlim(0, len(fit_ll))
plt.ylabel("Log Probability")
plt.show()
```

### 3b. Retrieved parameters

Compare retrieved weights and transition matrices to generative parameters. To do this, we may first need to permute the states of the fit GLM-HMM relative to the
generative model.  One way to do this uses the `find_permutation` function from `ssm`:

```python
new_glmhmm.permute(find_permutation(true_latents[0], new_glmhmm.most_likely_states(true_choices[0], input=inpts[0])))
```

Now plot generative and retrieved weights for GLMs (analogous plot to Figure S1c in
Ashwood et al. (2020)):

```python
fig = plt.figure(figsize=(4, 3), dpi=80, facecolor='w', edgecolor='k')
cols = ['#ff7f00', '#4daf4a', '#377eb8']
recovered_weights = new_glmhmm.observations.params
for k in range(num_states):
    if k ==0:
        plt.plot(range(input_dim), gen_weights[k][0], marker='o',
                 color=cols[k], linestyle='-',
                 lw=1.5, label="generative")
        plt.plot(range(input_dim), recovered_weights[k][0], color=cols[k],
                     lw=1.5,  label = "recovered", linestyle = '--')
    else:
        plt.plot(range(input_dim), gen_weights[k][0], marker='o',
                 color=cols[k], linestyle='-',
                 lw=1.5, label="")
        plt.plot(range(input_dim), recovered_weights[k][0], color=cols[k],
                     lw=1.5,  label = '', linestyle = '--')
plt.yticks(fontsize=10)
plt.ylabel("GLM weight", fontsize=15)
plt.xlabel("covariate", fontsize=15)
plt.xticks([0, 1], ['stimulus', 'bias'], fontsize=12, rotation=45)
plt.axhline(y=0, color="k", alpha=0.5, ls="--")
plt.legend()
plt.title("Weight recovery", fontsize=15)
```

Now plot generative and retrieved transition matrices (analogous plot to Figure S1c in
Ashwood et al. (2020)):

```python
fig = plt.figure(figsize=(5, 2.5), dpi=80, facecolor='w', edgecolor='k')
plt.subplot(1, 2, 1)
gen_trans_mat = np.exp(gen_log_trans_mat)[0]
plt.imshow(gen_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(gen_trans_mat.shape[0]):
    for j in range(gen_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(gen_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.ylabel("state t", fontsize = 15)
plt.xlabel("state t+1", fontsize = 15)
plt.title("generative", fontsize = 15)

plt.subplot(1, 2, 2)
recovered_trans_mat = np.exp(new_glmhmm.transitions.log_Ps)
plt.imshow(recovered_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(recovered_trans_mat.shape[0]):
    for j in range(recovered_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(recovered_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.title("recovered", fontsize = 15)
plt.subplots_adjust(0, 0, 1, 1)

```

### 3c. Posterior State Probabilities

Let's now plot $p(z_{t} = k|\mathbf{y}, \{u_{t}\}_{t=1}^{T})$, the posterior state probabilities, which give the probability of the animal being in state k at trial t.

```python
# Get expected states:
posterior_probs = [new_glmhmm.expected_states(data=data, input=inpt)[0]
                for data, inpt
                in zip([true_choices[0]], [inpts[0]])]
```

```python
true_choices
```

```python
fig = plt.figure(figsize=(5, 2.5), dpi=80, facecolor='w', edgecolor='k')
sess_id = 0 #session id; can choose any index between 0 and num_sess-1
for k in range(num_states):
    plt.plot(posterior_probs[sess_id][:, k], label="State " + str(k + 1), lw=2,
             color=cols[k])
plt.ylim((-0.01, 1.01))
plt.yticks([0, 0.5, 1], fontsize = 10)
plt.xlabel("trial #", fontsize = 15)
plt.ylabel("p(state)", fontsize = 15)
```

With these posterior state probabilities, we can assign trials to states and then plot the fractional occupancy of each state:

```python
# concatenate posterior probabilities across sessions
posterior_probs_concat = np.concatenate(posterior_probs)
# get state with maximum posterior probability at particular trial:
state_max_posterior = np.argmax(posterior_probs_concat, axis = 1)
# now obtain state fractional occupancies:
_, state_occupancies = np.unique(state_max_posterior, return_counts=True)
state_occupancies = state_occupancies/np.sum(state_occupancies)
```

```python
fig = plt.figure(figsize=(2, 2.5), dpi=80, facecolor='w', edgecolor='k')
for z, occ in enumerate(state_occupancies):
    plt.bar(z, occ, width = 0.8, color = cols[z])
plt.ylim((0, 1))
plt.xticks([0, 1, 2], ['1', '2', '3'], fontsize = 10)
plt.yticks([0, 0.5, 1], ['0', '0.5', '1'], fontsize=10)
plt.xlabel('state', fontsize = 15)
plt.ylabel('frac. occupancy', fontsize=15)
```

## 4. Fit GLM-HMM and perform recovery analysis: Maximum A Priori Estimation

Above, we performed Maximum Likelihood Estimation to retrieve the generative parameters of the GLM-HMM in simulated data. In the small data regime, where we do not have many trials available to us, we may instead want to perform Maximum A Priori (MAP) Estimation in order to incorporate a prior term and restrict the range for the best fitting parameters. Unfortunately, what is meant by 'small data regime' is problem dependent and will be affected by the number of states in the generative GLM-HMM, and the specific parameters of the generative model, amongst other things. In practice, we may perform both Maximum Likelihood Estimation and MAP estimation and compare the ability of the fit models to make predictions on held-out data (see Section 5 on Cross-Validation below).

The prior we consider for the GLM-HMM is the product of a Gaussian prior on the GLM weights, $W$, and a Dirichlet prior on the transition matrix, $A$:

$$
\begin{align}
\Pr(W, A) &= \mathcal{N}(W|0, \Sigma) \Pr(A|\alpha) \\&= \mathcal{N}(W|0, diag(\sigma^{2}, \cdots, \sigma^{2})) \prod_{j=1}^{K} \dfrac{1}{B(\alpha)} \prod_{k=1}^{K} A_{jk}^{\alpha -1}
\end{align}
$$

There are two hyperparameters controlling the strength of the prior: $\sigma$ and $\alpha$.  The larger the value of $\sigma$ and if $\alpha = 1$, the more similar MAP estimation will become to Maximum Likelihood Estimation, and the prior term will become an additive offset to the objective function of the GLM-HMM that is independent of the values of $W$ and $A$.  In comparison, setting $\sigma = 2$ and $\alpha = 2$ will result in the prior no longer being independent of $W$ and $\alpha$.

In order to perform MAP estimation for the GLM-HMM with `ssm`, the new syntax is:

```python
ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs",
             observation_kwargs=dict(C=num_categories,prior_sigma=prior_sigma),
             transitions="sticky", transition_kwargs=dict(alpha=prior_alpha,kappa=0))
```

where `prior_sigma` is the $\sigma$ parameter from above, and `prior_alpha` is the $\alpha$ parameter.

```python
# Instantiate GLM-HMM and set prior hyperparameters
prior_sigma = 2
prior_alpha = 2
map_glmhmm = ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs",
             observation_kwargs=dict(C=num_categories,prior_sigma=prior_sigma),
             transitions="sticky", transition_kwargs=dict(alpha=prior_alpha,kappa=0))
```

```python
# Fit GLM-HMM with MAP estimation:
_ = map_glmhmm.fit(true_choices, inputs=inpts, method="em", num_iters=N_iters, tolerance=10**-4)
```

Compare final likelihood of data with MAP estimation and MLE to likelihood under generative model (note: we cannot use log_probability that is output of `fit` function as this incorporates prior term, which is not comparable between generative and MAP models). We want to check that MAP and MLE likelihood values are higher than true likelihood; if they are not, this may indicate a poor initialization and that we should refit these models.

```python
true_likelihood = true_glmhmm.log_likelihood(true_choices, inputs=inpts)
mle_final_ll = new_glmhmm.log_likelihood(true_choices, inputs=inpts)
map_final_ll = map_glmhmm.log_likelihood(true_choices, inputs=inpts)
```

```python
# Plot these values
fig = plt.figure(figsize=(2, 2.5), dpi=80, facecolor='w', edgecolor='k')
loglikelihood_vals = [true_likelihood, mle_final_ll, map_final_ll]
colors = ['Red', 'Navy', 'Purple']
for z, occ in enumerate(loglikelihood_vals):
    plt.bar(z, occ, width = 0.8, color = colors[z])
plt.ylim((true_likelihood-5, true_likelihood+15))
plt.xticks([0, 1, 2], ['true', 'mle', 'map'], fontsize = 10)
plt.xlabel('model', fontsize = 15)
plt.ylabel('loglikelihood', fontsize=15)
```

## 5. Cross Validation

To assess which model is better - the model fit via Maximum Likelihood Estimation, or the model fit via MAP estimation - we can investigate the predictive power of these fit models on held-out test data sets.

```python
# Create additional input sequences to be used as held-out test data
num_test_sess = 10
test_inpts = np.ones((num_test_sess, num_trials_per_sess, input_dim))
test_inpts[:,:,0] = np.random.choice(stim_vals, (num_test_sess, num_trials_per_sess))
test_inpts = list(test_inpts)
```

```python
# Create set of test latents and choices to accompany input sequences:
test_latents, test_choices = [], []
for sess in range(num_test_sess):
    test_z, test_y = true_glmhmm.sample(num_trials_per_sess, input=test_inpts[sess])
    test_latents.append(test_z)
    test_choices.append(test_y)
```

```python
# Compare likelihood of test_choices for model fit with MLE and MAP:
mle_test_ll = new_glmhmm.log_likelihood(test_choices, inputs=test_inpts)
map_test_ll = map_glmhmm.log_likelihood(test_choices, inputs=test_inpts)
```

```python
fig = plt.figure(figsize=(2, 2.5), dpi=80, facecolor='w', edgecolor='k')
loglikelihood_vals = [mle_test_ll, map_test_ll]
colors = ['Navy', 'Purple']
for z, occ in enumerate(loglikelihood_vals):
    plt.bar(z, occ, width = 0.8, color = colors[z])
plt.ylim((mle_test_ll-2, mle_test_ll+5))
plt.xticks([0, 1], ['mle', 'map'], fontsize = 10)
plt.xlabel('model', fontsize = 15)
plt.ylabel('loglikelihood', fontsize=15)
```

Here we see that the model fit with MAP estimation achieves higher likelihood on the held-out dataset than the model fit with MLE, so we would choose this model as the best model of animal decision-making behavior (although we'd probably want to perform multiple fold cross-validation to be sure that this is the case in all instantiations of test data).

Let's finish by comparing the retrieved weights and transition matrices from MAP estimation to the generative parameters.

```python
map_glmhmm.permute(find_permutation(true_latents[0], map_glmhmm.most_likely_states(true_choices[0], input=inpts[0])))
```

```python
fig = plt.figure(figsize=(6, 3), dpi=80, facecolor='w', edgecolor='k')
cols = ['#ff7f00', '#4daf4a', '#377eb8']
plt.subplot(1,2,1)
recovered_weights = new_glmhmm.observations.params
for k in range(num_states):
    if k ==0: # show labels only for first state
        plt.plot(range(input_dim), gen_weights[k][0], marker='o',
                 color=cols[k],
                 lw=1.5, label="generative")
        plt.plot(range(input_dim), recovered_weights[k][0], color=cols[k],
                     lw=1.5,  label = 'recovered', linestyle='--')
    else:
        plt.plot(range(input_dim), gen_weights[k][0], marker='o',
                 color=cols[k],
                 lw=1.5, label="")
        plt.plot(range(input_dim), recovered_weights[k][0], color=cols[k],
                     lw=1.5,  label = '', linestyle='--')
plt.yticks(fontsize=10)
plt.ylabel("GLM weight", fontsize=15)
plt.xlabel("covariate", fontsize=15)
plt.xticks([0, 1], ['stimulus', 'bias'], fontsize=12, rotation=45)
plt.axhline(y=0, color="k", alpha=0.5, ls="--")
plt.title("MLE", fontsize = 15)
plt.legend()

plt.subplot(1,2,2)
recovered_weights = map_glmhmm.observations.params
for k in range(num_states):
    plt.plot(range(input_dim), gen_weights[k][0], marker='o',
             color=cols[k],
             lw=1.5, label="", linestyle = '-')
    plt.plot(range(input_dim), recovered_weights[k][0], color=cols[k],
                 lw=1.5,  label = '', linestyle='--')
plt.yticks(fontsize=10)
plt.xticks([0, 1], ['', ''], fontsize=12, rotation=45)
plt.axhline(y=0, color="k", alpha=0.5, ls="--")
plt.title("MAP", fontsize = 15)
```

```python
fig = plt.figure(figsize=(7, 2.5), dpi=80, facecolor='w', edgecolor='k')
plt.subplot(1, 3, 1)
gen_trans_mat = np.exp(gen_log_trans_mat)[0]
plt.imshow(gen_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(gen_trans_mat.shape[0]):
    for j in range(gen_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(gen_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.ylabel("state t", fontsize = 15)
plt.xlabel("state t+1", fontsize = 15)
plt.title("generative", fontsize = 15)

plt.subplot(1, 3, 2)
recovered_trans_mat = np.exp(new_glmhmm.transitions.log_Ps)
plt.imshow(recovered_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(recovered_trans_mat.shape[0]):
    for j in range(recovered_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(recovered_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.title("recovered - MLE", fontsize = 15)
plt.subplots_adjust(0, 0, 1, 1)

plt.subplot(1, 3, 3)
recovered_trans_mat = np.exp(map_glmhmm.transitions.log_Ps)
plt.imshow(recovered_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(recovered_trans_mat.shape[0]):
    for j in range(recovered_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(recovered_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.title("recovered - MAP", fontsize = 15)
plt.subplots_adjust(0, 0, 1, 1)
```

## 6. Multinomial GLM-HMM

Until now, we have only considered the case where there are 2 output classes (the Bernoulli GLM-HMM corresponding to `C=num_categories=2`), yet the `ssm` framework is sufficiently general to allow us to fit the multinomial GLM-HMM described in Equations 1 and 2. Here we demonstrate a recovery analysis for the multinomial GLM-HMM.

```python
# Set the parameters of the GLM-HMM
num_states = 4        # number of discrete states
obs_dim = 1           # number of observed dimensions
num_categories = 3    # number of categories for output
input_dim = 2         # input dimensions

# Make a GLM-HMM
true_glmhmm = ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs",
                   observation_kwargs=dict(C=num_categories), transitions="standard")
```

```python
# Set weights of multinomial GLM-HMM
gen_weights = np.array([[[0.6,3], [2,3]], [[6,1], [6,-2]], [[1,1], [3,1]], [[2,2], [0,5]]])
print(gen_weights.shape)
true_glmhmm.observations.params = gen_weights
```

In the above, notice that the shape of the weights for the multinomial GLM-HMM is `(num_states, num_categories-1, input_dim)`.  Specifically, we only learn `num_categories-1` weight vectors (of size `input_dim`) for a given state, and we set the weights for the other observation class to zero. Constraining the weight vectors for one class is important if that we want to be able to identify generative weights in simulated data. If we didn't do this, it is easy to see that one could generate the same observation probabilities with a set of weight vectors that are offset by a constant vector $w_{k}$ (the index k indicates that a different offset vector could exist per state):
$$
\begin{align}
\Pr(y_t = c \mid z_{t} = k, u_t, w_{kc}) =
\frac{\exp\{w_{kc}^\mathsf{T} u_t\}}
{\sum_{c'=1}^C \exp\{w_{kc'}^\mathsf{T} u_t\}} = \frac{\exp\{(w_{kc}-w_{k})^\mathsf{T} u_t\}}
{\sum_{c'=1}^C \exp\{(w_{kc'}-w_{k})^\mathsf{T} u_t\}}
\end{align}
$$

Equations 1 and 2 at the top of this notebook already take into account the fact that the weights for a particular class for a given state are fixed to zero (this is why $c = C$ is handled differently).

```python
# Set transition matrix of multinomial GLM-HMM
gen_log_trans_mat = np.log(np.array([[[0.90, 0.04, 0.05, 0.01], [0.05, 0.92, 0.01, 0.02], [0.03, 0.02, 0.94, 0.01], [0.09, 0.01, 0.01, 0.89]]]))
true_glmhmm.transitions.params = gen_log_trans_mat
```

```python
# Create external inputs sequence; compared to the example above, we will increase the number of examples
# (through the "num_trials_per_session" paramater) since the number of parameters has increased
num_sess = 20 # number of example sessions
num_trials_per_sess = 1000 # number of trials in a session
inpts = np.ones((num_sess, num_trials_per_sess, input_dim)) # initialize inpts array
stim_vals = [-1, -0.5, -0.25, -0.125, -0.0625, 0, 0.0625, 0.125, 0.25, 0.5, 1]
inpts[:,:,0] = np.random.choice(stim_vals, (num_sess, num_trials_per_sess)) # generate random sequence of stimuli
inpts = list(inpts)
```

```python
# Generate a sequence of latents and choices for each session
true_latents, true_choices = [], []
for sess in range(num_sess):
    true_z, true_y = true_glmhmm.sample(num_trials_per_sess, input=inpts[sess])
    true_latents.append(true_z)
    true_choices.append(true_y)
```

```python
# plot example data:
fig = plt.figure(figsize=(8, 3), dpi=80, facecolor='w', edgecolor='k')
plt.step(range(100),true_choices[0][range(100)], color = "red")
plt.yticks([0, 1, 2])
plt.title("example data (multinomial GLM-HMM)")
plt.xlabel("trial #", fontsize = 15)
plt.ylabel("observation class", fontsize = 15)
```

```python
# Calculate true loglikelihood
true_ll = true_glmhmm.log_probability(true_choices, inputs=inpts)
print("true ll = " + str(true_ll))
```

```python
# fit GLM-HMM
new_glmhmm = ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs",
                   observation_kwargs=dict(C=num_categories), transitions="standard")

N_iters = 500 # maximum number of EM iterations. Fitting with stop earlier if increase in LL is below tolerance specified by tolerance parameter
fit_ll = new_glmhmm.fit(true_choices, inputs=inpts, method="em", num_iters=N_iters, tolerance=10**-4)
```

```python
# Plot the log probabilities of the true and fit models. Fit model final LL should be greater
# than or equal to true LL.
fig = plt.figure(figsize=(4, 3), dpi=80, facecolor='w', edgecolor='k')
plt.plot(fit_ll, label="EM")
plt.plot([0, len(fit_ll)], true_ll * np.ones(2), ':k', label="True")
plt.legend(loc="lower right")
plt.xlabel("EM Iteration")
plt.xlim(0, len(fit_ll))
plt.ylabel("Log Probability")
plt.show()
```

```python
# permute recovered state identities to match state identities of generative model
new_glmhmm.permute(find_permutation(true_latents[0], new_glmhmm.most_likely_states(true_choices[0], input=inpts[0])))
```

```python
# Plot recovered parameters:
recovered_weights = new_glmhmm.observations.params
recovered_transitions = new_glmhmm.transitions.params

fig = plt.figure(figsize=(16, 8), dpi=80, facecolor='w', edgecolor='k')
plt.subplots_adjust(wspace=0.3, hspace=0.6)

plt.subplot(2, 2, 1)
cols = ['#ff7f00', '#4daf4a', '#377eb8', '#f781bf', '#a65628', '#984ea3', '#999999', '#e41a1c', '#dede00']

for c in range(num_categories):
    plt.subplot(2, num_categories+1, c+1)
    if c < num_categories-1:
        for k in range(num_states):
            plt.plot(range(input_dim), gen_weights[k,c], marker='o',
                 color=cols[k], lw=1.5, label="state " + str(k+1) + "; class " + str(c+1))
    else:
        for k in range(num_states):
            plt.plot(range(input_dim), np.zeros(input_dim), marker='o',
                     color=cols[k], lw=1.5, label="state " + str(k+1) + "; class " + str(c+1), alpha = 0.5)

    plt.axhline(y=0, color="k", alpha=0.5, ls="--")
    plt.yticks(fontsize=10)
    plt.xticks([0, 1], ['', ''])
    if c == 0:
        plt.ylabel("GLM weight", fontsize=15)
    plt.legend()
    plt.title("Generative weights; class " + str(c+1), fontsize = 15)
    plt.ylim((-3, 10))

plt.subplot(2, num_categories+1, num_categories+1)
gen_trans_mat = np.exp(gen_log_trans_mat)[0]
plt.imshow(gen_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(gen_trans_mat.shape[0]):
    for j in range(gen_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(gen_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3', '4'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3', '4'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.ylabel("state t", fontsize = 15)
plt.xlabel("state t+1", fontsize = 15)
plt.title("Generative transition matrix", fontsize = 15)

cols = ['#ff7f00', '#4daf4a', '#377eb8', '#f781bf', '#a65628', '#984ea3', '#999999', '#e41a1c', '#dede00']

for c in range(num_categories):
    plt.subplot(2, num_categories+1, num_categories + c + 2)
    if c < num_categories-1:
        for k in range(num_states):
            plt.plot(range(input_dim), recovered_weights[k,c], marker='o', linestyle = '--',
                 color=cols[k], lw=1.5, label="state " + str(k+1) + "; class " + str(c+1))
    else:
        for k in range(num_states):
            plt.plot(range(input_dim), np.zeros(input_dim), marker='o', linestyle = '--',
                     color=cols[k], lw=1.5, label="state " + str(k+1) + "; class " + str(c+1), alpha = 0.5)

    plt.axhline(y=0, color="k", alpha=0.5, ls="--")
    plt.yticks(fontsize=10)
    plt.xlabel("covariate", fontsize=15)
    if c == 0:
        plt.ylabel("GLM weight", fontsize=15)
    plt.xticks([0, 1], ['stimulus', 'bias'], fontsize=12, rotation=45)
    plt.legend()
    plt.title("Recovered weights; class " + str(c+1), fontsize = 15)
    plt.ylim((-3,10))

plt.subplot(2, num_categories+1, 2*num_categories+2)
recovered_trans_mat = np.exp(recovered_transitions)[0]
plt.imshow(recovered_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(recovered_trans_mat.shape[0]):
    for j in range(recovered_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(recovered_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3', '4'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3', '4'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.ylabel("state t", fontsize = 15)
plt.xlabel("state t+1", fontsize = 15)
plt.title("Recovered transition matrix", fontsize = 15)
```

notebooks/2b Input Driven Observations (GLM-HMM).ipynb


```python
import matplotlib.pyplot as plt
import scipy.stats as st
import seaborn as sns
import pandas as pd
from scipy.stats import norm, uniform, expon
from scipy.integrate import quad
from sympy.solvers import solve
from sympy import Symbol
import numpy as np
from pandas import Series, DataFrame
```

```python
methodeA = Series([79.98, 80.04, 80.02, 80.04, 80.03, 80.03, 80.04,
79.97, 80.05, 80.03, 80.02, 80.00, 80.02])
print(methodeA.mean())
print(methodeA.std())
```

    80.02076923076923
    0.023965787580611863

```python
from scipy.stats import norm
np.random.seed(1)
methodeA_sim1 = Series(np.round(norm.rvs(size=6, loc=80, scale=0.02), 2))
methodeA_sim1
print(methodeA_sim1.mean())
print(methodeA_sim1.std())
```

    79.99333333333333
    0.028751811537128993

```python
np.random.seed(17)
for i in range(5):
    methodeA_sim1 = Series(np.round(norm.rvs(size=6, loc=80, scale=0.02), 2))
    print("Mittelwert:", np.round(methodeA_sim1.mean(), 3))
    print("Standardabw.:", np.round(methodeA_sim1.std(), 3))
    print()
```

    Mittelwert: 80.01
    Standardabw.: 0.027

    Mittelwert: 80.007
    Standardabw.: 0.02

    Mittelwert: 79.992
    Standardabw.: 0.028

    Mittelwert: 79.995
    Standardabw.: 0.016

    Mittelwert: 79.992
    Standardabw.: 0.013

```python
np.random.seed(463137)
methodeA_sim2 = Series(np.round(norm.rvs(size=6, loc=80, scale=0.02), 2))
methodeA_sim2
print(methodeA_sim2.mean())
print(methodeA_sim2.std())
```

    80.04
    0.019999999999998862

```python
np.random.seed(647)
methodeA_sim3 = Series(np.round(norm.rvs(size=13, loc=80, scale= 0.02), 2))
methodeA_sim3
print(methodeA_sim3.mean())
print(methodeA_sim3.std())
```

    79.99307692307691
    0.023588350014576355

```python
norm.cdf(x=499.57, loc=500, scale=1/np.sqrt(100))
```

    8.539905470989153e-06

```python
norm.cdf(x=171.54, loc=180, scale=10/np.sqrt(2))*100
```

    11.576543872967077

```python

```


# Covid-19: From model prediction to model predictive control

## A demo of the stochastic modeling framework

*Original code by Ryan S. McGee. Modified by T.W. Alleman in consultation with the BIOMATH research unit headed by prof. Ingmar Nopens.*

Copyright (c) 2020 by T.W. Alleman, BIOMATH, Ghent University. All Rights Reserved.

Our code implements a SEIRS infectious disease dynamics models with extensions to model the effect quarantining detected cases. Using the concept of 'classes' in Python 3, the code was integrated with our previous work and allows to quickly perform Monte Carlo simulations, calibrate model parameters and calculate an *optimal* government policies using a model predictive controller (MPC). A white paper and souce code of our previous work can be found on the Biomath website.

https://biomath.ugent.be/covid-19-outbreak-modelling-and-control

```python
import numpy as np
import matplotlib.pyplot as plt
from IPython.display import Image
from ipywidgets import interact,fixed,FloatSlider,IntSlider,ToggleButtons
import pandas as pd
import datetime
import scipy
from scipy.integrate import odeint
import matplotlib.dates as mdates
import matplotlib
import scipy.stats as st
import networkx
import models
```

## Introduction

### Model dynamics

In this work, we extended the SEIR model to model the disease spread with a higher resolution. To this end, the *infected pool* is split into four types of infectiousness: 1) supermild (SM): people who show little to no symptoms at all, 2) mild (M): people with noticable, mild covid-19 symptoms, 3) heavy (H): people hospitalised but not in need of intensive care and 4) critical (C): people hospitalised and in need of intensive care. It takes several days before a heavy or critical infection becomes so severe that these patients need to be hospitalised. To this end, the pool of heavily and critically infected patients is split in two: not yet hospitalised (H, C) and hospitalised (HH,CH). The *removed pool* from the classical SEIR model is explicitly split into an immune (I) and dead (F) pool. People from the susceptible, exposed, supermild infected, mild infected and immune pool can be quarantined after having tested positive for covid-19. In the flowchart below, quarantine pools are denoted by the *Q*-suffix. Note that for individuals in the susceptible and immune pools, this corresponds to receiving a false positive test.

### Deterministic vs. Stochastic framework

The extended SEIR model is implemented using two frameworks: a deterministic and a stochastic (network based) framework. **This Jupyter Notebooks is a demo of the deterministic model, the stochastic model works but is not yet integrated with our current methods.** A deterministic implementation of the extended SEIRS model captures important features of infectious disease dynamics, but it assumes uniform mixing of the population (i.e. every individual in the population is equally likely to interact with every other individual). The deterministic approach results in a set of N ordinary differential equations, one for every of the N ’population pools’ considered. The main advantage of the deterministic model is that a low amount of computational resources are required while still maintaining an acceptable accuracy. The deterministic framework allows to rapidly explore scenarios and perform optimisations which require thousands of function evaluations.

However, it is often important to consider the structure of contact networks when studying disease transmission and the effect of interventions such as social distancing and contact tracing. The main drawback of the deterministic approach is the inability to simulate contact tracing, which is one of the most promising measures against the spread of sars-cov-2. For this reason, the SEIRS dynamics depicted in on the above flowchart can be simulated on a stochastic dynamical network. This advantages include a more detailed analysis of the relationship between social network structure and effective transmission rates, including the effect of network-based interventions such as social distancing, quarantining, and contact tracing. The first drawback is the increased amount of computational resources required, leading to the inability to perform optimisations on a personal computer. Instead, high performance computing infrastructure is needed. The second drawback is the need for more data and/or assumptions on social interactions and how government measures affect these social interactions.

#### Deterministic equations
The dynamics of the deterministic system are mathematically formulated as the rate of change of each population pool shown in the above flowchart. This results in the following system of ordinary differential equations,

\begin{eqnarray}
\dot{S} &=& - \beta \cdot N_c \Big( \frac{E+SM}{N} \Big) \cdot S  - \theta_{\text{S}} \psi_{\text{FP}} \cdot S + SQ/d_{\text{q,FP}} + \zeta \cdot R,\\
\dot{E} &=&  \beta \cdot N_c  \Big( \frac{E+SM}{N} \Big) \cdot S - B / \sigma - \theta_{\text{E}} \psi_{\text{PP}} \cdot E,\\
\dot{SM} &=& \text{sm}/\sigma  \cdot E - SM/d_{\text{sm}} - \theta_{\text{SM}} \psi_{\text{PP}} \cdot SM,\\
\dot{M} &=&  \text{m} / \sigma \cdot E - M/d_{\text{m}} - \theta_{\text{M}} \psi_{\text{PP}} \cdot M,\\
\dot{H} &=&  \text{h} / \sigma \cdot E + h / \sigma \cdot EQ - H/d_{\text{hospital}},\\
\dot{C} &=&  \text{c} / \sigma \cdot E + c / \sigma \cdot EQ - C/d_{\text{hospital}},\\
\dot{HH} &=& H/d_{\text{hospital}} - HH/d_{h}\\
\dot{CH} &=& C/d_{\text{hospital}} - m_c \cdot CH/d_{\text{cf}} - (1-m_c) \cdot CH/d_{\text{cr}}\\
\dot{F} &=& m_c \cdot CH/d_{\text{cf}}\\
\dot{R} &=& SM/d_{\text{sm}} + M/d_{\text{m}} + HH/d_{\text{h}} + (1-m_h) \cdot CH/d_{\text{cr}} + SMQ/d_{\text{sm}} \\ && + MQ/d_{\text{m}} + RQ/d_{\text{q,FP}} - \zeta \cdot R \\
\dot{SQ} &=& \theta_{\text{S}} \psi_{\text{FP}} \cdot S - SQ/d_{\text{q,FP}} \\
\dot{EQ} &=& \theta_{\text{E}} \psi_{\text{PP}} \cdot E - EQ/\sigma\\
\dot{SMQ} &=& \theta_{\text{SM}} \psi_{\text{PP}} \cdot SM + sm/\sigma \cdot EQ - SMQ/d_{\text{sm}} \\
\dot{MQ} &=& \theta_{\text{M}} \psi_{\text{PP}} \cdot M + m/\sigma \cdot EQ - MQ/d_{\text{m}}\\
\dot{RQ} &=& \theta_{\text{R}} \psi_{\text{FP}} \cdot R - RQ/d_{\text{q,FP}}
\end{eqnarray}

#### Stochastic equations
Consider a network graph G like the one depicted below, representing individuals (nodes) and their interactions (edges). Each individual (node) has a state (S, E, SM, M, H, C, HH, CH, R, F, SQ, SMQ, MQ or RQ). The set of nodes adjacent (connected by an edge) to an individual defines their set of "close contacts" (highlighted in black). At a given time, each individual makes contact with a random individual from their set of close contacts with probability (1-p) or with a random individual from anywhere in the network (highlighted in blue) with probability p. The latter global contacts represent individuals interacting with the population at large. These are individuals outside of ones inner social circle, such as on public transit, at an event or chatting to an old acquaintance etc. When a susceptible individual interacts with an exposed or infectious individual they become exposed themselves. The parameter p defines the locality of the network: for p=0 an individual only interacts with their close contacts, while p=1 represents a uniformly mixed population. **For a sufficiently large network G, and $p = 1$, the stochastic and deterministic framework will yield the same result**. Social distancing interventions ifluences both the locality and connectivity of the network. It is expected that social distancing lowers the average number of interactions an individual has each day (lower connectivity of the network). But as a result, the fraction of interactions in the inner circle increases, lowering p.

Each node $i$ has a state $X_i$ that updates according to a probability transition rate. The dynamics of the stochastic system are mathematically formulated as probabilities of transitioning between the pools shown in the model flowchart. In addition, back tracking of confirmed cases can be included in the model,

\begin{eqnarray}
\text{Pr}(X_i = S \rightarrow E) &=&  \Big[ \underbrace{p \cdot \frac{\beta (E + SM)}{N}}_{\text{infection through random encounter}} \\
&& + \underbrace{(1-p) \cdot \frac{\beta \sum_{j \in C_G(i)} (\delta_{X_j = E}+\delta_{X_j = SM})}{C_G (i)}}_{\text{infection through inner circle}} \Big] \delta_{X_i=S},\\
\text{Pr}(X_i = E \rightarrow SM) &=&  (sm/\sigma) \cdot \delta_{X_i = E},\\
\text{Pr}(X_i = E \rightarrow M) &=&  (m/\sigma) \cdot \delta_{X_i = E},\\
\text{Pr}(X_i = E \rightarrow H) &=&  (h/\sigma) \cdot \delta_{X_i = E},\\
\text{Pr}(X_i = E \rightarrow C) &=&  (c/\sigma) \cdot \delta_{X_i = E},\\
\text{Pr}(X_i = H \rightarrow HH) &=&  (1/d_{\text{hospital}}) \delta_{X_i = H},\\
\text{Pr}(X_i = C \rightarrow CH) &=&  (1/d_{\text{hospital}}) \delta_{X_i = C},\\
\text{Pr}(X_i = SM \rightarrow R) &=& (1/d_{\text{sm}}) \cdot \delta_{X_i = SM},\\
\text{Pr}(X_i = M \rightarrow R) &=& (1/d_{\text{m}}) \cdot \delta_{X_i = M},\\
\text{Pr}(X_i = HH \rightarrow R) &=& (1/d_{\text{h}}) \cdot \delta_{X_i = HH},\\
\text{Pr}(X_i = CH \rightarrow R) &=& (1-m_c)/d_{\text{cr}} \cdot \delta_{X_i = HH},\\
\text{Pr}(X_i = CH \rightarrow F) &=& (m_c/d_{\text{cf}}) \cdot \delta_{X_i = HH},\\
\text{Pr}(X_i = S \rightarrow SQ) &=& \Big[\theta_S + \phi_S \big[ \sum_{j \in C_G(i)} (\delta_{X_j = SQ}+\delta_{X_j = EQ}+\delta_{X_j = SMQ} \\
&&+\delta_{X_j = MQ}+\delta_{X_j = RQ}+\delta_{X_j = HH}+\delta_{X_j = CH}) \big] \Big] \cdot \psi_{\text{FP}} \cdot \delta_{X_i = S}, \\
\text{Pr}(X_i = E \rightarrow EQ) &=& \Big[\theta_E + \phi_E \big[ \sum_{j \in C_G(i)} (\delta_{X_j = SQ}+\delta_{X_j = EQ}+\delta_{X_j = SMQ} \\
&&+\delta_{X_j = MQ}+\delta_{X_j = RQ}+\delta_{X_j = HH}+\delta_{X_j = CH}) \big] \Big] \cdot \psi_{\text{PP}} \cdot \delta_{X_i = E}, \\
\text{Pr}(X_i = SM \rightarrow SMQ) &=& \Big[\theta_{SM} + \phi_{SM} \big[ \sum_{j \in C_G(i)} (\delta_{X_j = SQ}+\delta_{X_j = EQ}+\delta_{X_j = SMQ} \\
&&+\delta_{X_j = MQ}+\delta_{X_j = RQ}+\delta_{X_j = HH}+\delta_{X_j = CH}) \big] \Big] \cdot \psi_{\text{PP}} \cdot \delta_{X_i = SM}, \\
\text{Pr}(X_i = M \rightarrow MQ) &=& \Big[\theta_{M} + \phi_{M} \big[ \sum_{j \in C_G(i)} (\delta_{X_j = SQ}+\delta_{X_j = EQ}+\delta_{X_j = SMQ} \\
&&+\delta_{X_j = MQ}+\delta_{X_j = RQ}+\delta_{X_j = HH}+\delta_{X_j = CH}) \big] \Big] \cdot \psi_{\text{PP}} \cdot \delta_{X_i = M} \\
\text{Pr}(X_i = R \rightarrow RQ) &=& \Big[\theta_{R} + \phi_{R} \big[ \sum_{j \in C_G(i)} (\delta_{X_j = SQ}+\delta_{X_j = EQ}+\delta_{X_j = SMQ} \\
&&+\delta_{X_j = MQ}+\delta_{X_j = RQ}+\delta_{X_j = HH}+\delta_{X_j = CH}) \big] \Big] \cdot \psi_{\text{FP}} \cdot \delta_{X_i = R}\\
\text{Pr}(X_i = SQ \rightarrow S) &=& (1/d_{q,FP}) \cdot \delta_{X_i = SQ},\\
\text{Pr}(X_i = EQ \rightarrow SMQ) &=& sm \cdot \sigma \cdot \delta_{X_i = EQ},\\
\text{Pr}(X_i = EQ \rightarrow MQ) &=& m \cdot \sigma \cdot \delta_{X_i = EQ},\\
\text{Pr}(X_i = EQ \rightarrow H) &=& h \cdot \sigma \cdot \delta_{X_i = EQ},\\
\text{Pr}(X_i = EQ \rightarrow C) &=& h \cdot \sigma \cdot \delta_{X_i = EQ},\\
\text{Pr}(X_i = SMQ \rightarrow R) &=& h (1/d_{sm}) \cdot \delta_{X_i = SMQ},\\
\text{Pr}(X_i = MQ \rightarrow R) &=& h (1/d_{m}) \cdot \delta_{X_i = MQ},\\
\text{Pr}(X_i = RQ \rightarrow R) &=& h (1/d_{\text{hospital}}) \cdot \delta_{X_i = RQ},\\
\text{Pr}(X_i = R \rightarrow S) &=& \zeta \cdot \delta_{X_i = R},\\
\end{eqnarray}

where $\delta_{X_i = A} = 1$ if the state of $X_i$ is A, or 0 if not, and where $C_G(i)$ denotes the set of close contacts of node i (adjacent nodes). For large populations and $p = 1$, this stochastic model approaches the same dynamics as the deterministic SEIRS model.

### Model parameters

In the above equations, S stands for susceptible, E for exposed, SM for supermild, M for mild, H for heavy, C for critical, HH for heavy and hospitalised, CH for critical and hospitalised, F for dead, R for immune. The quarantined states are SQ, which stands for susceptible and quarantined, EQ for exposed and quarantined, SMQ for supermild infected and quarantined, MQ for mildly infected and quarantined and RQ for recovered and quarantined. N stands for the total population. The clinical parameters are: sm, m, z, h: the chance of having an asymptotic, mild, heavy or critical infection. Based on reported cases in China and travel data, Li et al. (2020b) estimated that 86 % of coronavirus infections in the country were "undocumented" in the weeks before officials instituted stringent quarantines. A biotech company in Iceland that has tested more than 9,000 people found that around 50 percent of those who tested positive said they were asymptomatic. We assume that 50 % of all infected cases are supermild (asymptotic) and hence,

$$sm = 0.50$$

Based on previously reported estimates of the distribution between mild, severe and critical cases (Wu and McGoogan, 2020), it is then calculated that the chance of contrapting a mild, severe and critical infection is,

$$m = (1 − 0.50) · 0.81 = 0.405,$$
$$h = (1 − 0.50) · 0.14 = 0.070,$$
$$c = (1 − 0.50) · 0.05 = 0.025.$$

$d_{sm}$ , $d_m$ , $d_h$ : the number of symptomic days in case of a supermild, mild or heavy infection. $d_{hf}$ is the time from hospitalisation until death in case of a critical infection while $d_{hr}$ is the recovery time for critical cases. Zhou et al. (2020) performed a retrospective study on 191 Chinese hospital patients and determined that the time from illness onset to discharge or death was 22.0 days (18.0-25.0, IQR) and 18.5 days (15.0-22.0, IQR) for survivors and victims respectively. Using available preliminary data, the World Health Organisation estimated the median time from onset to clinical recovery for mild cases to be approximately 2 weeks and to be 3-6 weeks for patients with severe or critical disease (WHO, 2020). Based on this report, we assume a recovery time of three weeks for heavy infections. d hospital : the time before heavily or critically infected patients reach the hospital. On average this is 9.1 days (Li et al., 2020a). m c : the mortality in case of a critical infection, which is roughly 50% (Wu and McGoogan, 2020). It is assumed that if the number of critical cases surpasses the total number of ICU beds, the mortality is changed to the average between 0.49 and 1, with the fraction of patients in critical condition receiving care as a weight. This is formulated mathematically in the following way,

\begin{equation}
m_h = \underbrace{\Big( \frac{ICU_{\text{max}}}{H} \Big) \cdot 0.49}_{\text{critical patient receives care}} + \underbrace{\Big( \frac{H - ICU_{\text{max}}}{H} \Big) \cdot 1.00}_{\text{critical patient receives no care}}\ .
\end{equation}

$\sigma$: the incubation period in days. The incubation period is assumed to be Erlang distributed as reported by Li et al. (2020a). The average incubation time is 5.2 days. $\zeta$: can be used to model the effect of re-susceptibility and seasonality of a disease. Throughout this demo, we assume $\zeta = 0$ because data on seasonality is not yet available at the moment. We thus assume permanent immunity after recovering from the infection.

The transmission rate of the disease in the deterministic model depends on the product of three contributions. The first contribution, $(E+SM)/N$, is the probability of encountering a contagious individual. The second contribution, $N_c$, is the average number of human-to-human interactions per day. As previously explained, **in the deterministic framework, this means all $N_c$ contacts within the population are random**. In this work, we explictly split $\beta$ and $N_c$ because this has the following advantages: 1) $\beta$ is now a disease characteric, independent of social interactions. 2) The goal of this work is to demonstrate the concept of social control to contain the outbreak using model predictive control. By splitting $\beta$ and $N_c$, the controlled variable is the number of random contacts $N_c$ which is more comprehendible to the reader of this text. The number of human-human interactions per day is estimated using the Social Contact Rates (SOCRATES) Data Tool (Willem et al., 2020). The dataset on human-human interactions for Belgium is based on a 2008 study by Mossong, which kept track of 750 participants and their interactions (8878 interactions) in Belgium. The dataset includes both physical and non-physical interactions of any duration. The third contribution, $\beta$, is the probability of contracting sars-cov-2 when encoutering a contagious individual. The testing and quarantine parameters are: $\theta_{S}$, $\theta_{E}$, $\theta_{SM}$, $\theta_{M}$, $\theta_{R}$ expressed as the number of susceptibles, exposed, supermild and mild individuals tested each day. It is difficult to model a concrete testing scenario in the deterministic framework because no back tracing can be coupled to the tests. The deterministic model only considers random testing of the population. It uses the total daily testing capacity to calculate the amount of individuals in the S, E, SM, M and R pool that receive a test. $\psi_{PP}$: probability of correctly identifiying and quarantining an exposed or infected person. $\psi_{FP}$: probability of falsly identifying and quarantining a susceptible or recovered individual. Both parameters can be calculated using the reliability of the test and Bayes theorem. For now, we assume no false positives. $d_{\text{q,FP}}$: the duration of the quarantine in the case of a false positive, assumed to be 14 days.

### Social interaction data

#### Social Contact Rates (SOCRATES) Data Tool

https://lwillem.shinyapps.io/socrates_rshiny/

1. What is the average number of daily human-to-human contacts of the Belgian population? Include all ages, all genders and both physical and non-physical interactions of any duration. To include all ages, type: *0,60+* in the *Age Breaks* dialog box.
2. What is the average number of physical human-to-human contacts of the Belgian population? Include all ages, all genders and all durations of physical contact.
3. What is the average number of physical human-to-human contacts of at least 1 hour of the Belgian population?
4. Based on the above results, how would you estimate $N_c$ in the deterministic model?
5. Based on the above results, how would you estimate $p$ in the stochastic model? Recall that $p$ is the fraction of *random contacts* a person has on a daily basis, while $(1-p)$ is the fraction of *inner circle contacts* a person has on a daily basis.

#### Google COVID-19 Community Mobility Reports

https://www.google.com/covid19/mobility/

#### London School of Hygiene

https://www.thelancet.com/journals/lanpub/article/PIIS2468-2667(20)30073-6/fulltext

## Performing simulations

### Without age-structuring

#### The 'SEIRSNetworkModel' class

The *SEIRSNetworkModel* class uses the exact same class structure and function names as the *SEIRSAgeModel* deterministic class, which is detailed below.

As of now (20/04/2020), the SEIRSNetworkModel contains 5 functions which can be grouped in two parts: 1) functions to run and visualise simulations and 2) functions to perform parameter estimations and visualse the results. Implementing the model predictive controller is straightforward and can easily be done. However, the optimisation problem is really difficult and requires thousands of function evaluations. Given the large amount of computational resources required to run just one stochastic simulation, it is highly unlikely that the model predictive controller will ever be used to optimize government policy. The MPC functions will be implemented and their usefullness assessed after a calibration is performed. Also, scenario specific functions will be added over the course of next week.

#### Creating a SEIRSNetworkModel object

Before a stochastic simulation can be performed, the interaction network G, which determines the connectivity of the network model, must be defined. In the example below, an interaction network under normal circumstances and an interaction network under distancing measures is initiated. Switching between networks is possible during a simulation.

```python
# Construct the network G
numNodes = 6000
baseGraph    = networkx.barabasi_albert_graph(n=numNodes, m=7)
# Baseline normal interactions:
G_norm     = models.custom_exponential_graph(baseGraph, scale=200)
models.plot_degree_distn(G_norm, max_degree=40)

# Construct the network G under social distancing
numNodes = 6000
baseGraph    = networkx.barabasi_albert_graph(n=numNodes, m=2)
# Baseline normal interactions:
G_dist     = models.custom_exponential_graph(baseGraph, scale=20000)
models.plot_degree_distn(G_dist, max_degree=40)
```

```python
model = models.SEIRSNetworkModel(
                                 # network connectivty
                                 G = G_norm,
                                 p       = 0.6,
                                 # clinical parameters
                                 beta    = 0.05,
                                 sigma   = 5.2,
                                 zeta    = 0,
                                 sm      = 0.50,
                                 m       = (1-0.50)*0.81,
                                 h       = (1-0.50)*0.15,
                                 c       = (1-0.50)*0.04,
                                 dsm     = 14,
                                 dm      = 14,
                                 dhospital = 1,
                                 dh      = 21,
                                 dcf     = 18.5,
                                 dcr     = 22.0,
                                 mc0     = 0.49,
                                 ICU     = 2000,
                                 # testing
                                 theta_S = 0,
                                 theta_E = 0,
                                 theta_SM= 0,
                                 theta_M = 0,
                                 theta_R = 0,
                                 psi_FP = 0,
                                 psi_PP = 1,
                                 dq     = 14,
                                 # back-tracking
                                 phi_S   = 0,
                                 phi_E   = 0,
                                 phi_SM  = 0,
                                 phi_R   = 0,
                                 # initial condition
                                 initN = 11.43e6, #results are extrapolated to entire population
                                 initE = 1,
                                 initSM = 0,
                                 initM = 0,
                                 initH = 0,
                                 initC = 0,
                                 initHH = 0,
                                 initCH = 0,
                                 initR = 0,
                                 initF = 0,
                                 initSQ = 0,
                                 initEQ = 0,
                                 initSMQ = 0,
                                 initMQ = 0,
                                 initRQ = 0,
                                 # monte-carlo sampling
                                 monteCarlo = False,
                                 n_samples = 100
                            )
```

#### Altering an object variable after intialisation

After initialising our 'model' it is still possible to change variables using the following syntax.

```python
model.beta = 0.30
```

#### Running your first simulation

A simulation is run by using the attribute function *sim*, which uses one argument, the simulation time T, as its input.

```python
y = model.sim(100)
```

    t = 0.59
    t = 20.29
    t = 30.14
    t = 40.01
    t = 50.00
    t = 60.00
    t = 70.01
    t = 80.01
    t = 90.00
    t = 100.00

For advanced users: the numerical results of the simulation can be accessed directly be calling *object.X* or *object.sumX* where X is the name of the desired population pool. Both are numpy arrays. *Ojbect.X* is a 3D array of the following dimensions:
- x-dimension: number of age categories,
- y-dimesion: tN: total number of timesteps taken (one per day),
- z-dimension: n_samples: total number of monte-carlo simulations performed.

Object.sumX is a 2D array containing only the results summed over all age categorie and has the following dimensions,
- x-dimesion: tN: total number of timesteps taken (one per day),
- y-dimension: n_samples: total number of monte-carlo simulations performed.

#### Visualising the results

To quickly visualise simulation results, two attribute functions were created. The first function, *plotPopulationStatus*, visualises the number of susceptible, exposed, infected and recovered individuals in the population. The second function, *plotInfected*, by default visualises the number of heavily and critically infected individuals. Both functions require no user input to work but both have some optional arguments,

> plotPopulationStatus(filename),
> - filename: string with a filename + extension to save the figure. The figure is not saved per default.

> plotInfected(asymptotic, mild, filename),
> - asymptotic: set to *True* to include the supermild pool in the visualisation.
> - mild: set to *True* to include the mild pool in the visualisation.
> - filename: string with a filename + extension to save the figure. The figure is not saved per default.

```python
model.plotPopulationStatus()
model.plotInfected()
```

#### The use of checkpoints to change parameters on the fly

A cool feature of the original SEIRSplus package by Ryan McGee was the use of so-called *checkpoints* dictionary to change simulation parameters on the fly. In our modification, this feature is preserved. Below you can find an example of a *checkpoints* dictionary. The simulation will be started with the previously initialised parameters. After 40 days, social interaction will be limited by lowering the network connectivity to an average of 4 edges per node. The chance of random encounters is lowered to 10%. After 80 days, social restrictions are lifted and beta once more assumes its *business-as-usual* value. *checkpoints* is the only optional argument of the *sim* functions and is set to *None* per default.

```python
# Create checkpoints dictionary
chk = {'t':       [40,80],
       'G':       [G_dist,G_norm],
       'p':       [0.01,0.3],
      }
# Run simulation
y = model.sim(100,checkpoints=chk)
# Visualise
model.plotPopulationStatus()
model.plotInfected()
```

## Calibrating $\beta$  in a *business-as-usual* scenario ($N_c = 11.2$)

### Performing a least-squares fit

The 'SEIRSNetworkModel' class contains a function to fit the model to selected data (*fit*) and one function to visualise the result (*plotFit*). Our code uses the **genetic algorithm** from scipy to perform the optimisation. The *fit* function has the following basic syntax,

> sim(data, parNames, positions, bounds, weights)
> - data: a list containing the dataseries (dtype=np array) to fit the model to. For now, dataseries must be of equal length.
> - parNames: a list containing the names (dtype=string) of the model parameters to be fitted.
> - positions: a list containing the model states (dtype=np.array) used to calculate the sum of least squares. Each dataseries must be matched to a certain (sum of) model state(s). If multiple entries are provided these are added togheter. The order of the states is given according to the following vector, where S has index 0: (S, E, SM, M, H, C, HH, CH, R, F, SQ, EQ, SMQ, MQ, RQ).

The following arguments are optional,
> - checkpoints: checkpoint dictionary can be used to calibrate under specific scenarios such as policy changes (default: None).
> - setvar: True to replace fitted values in model object after fit is performed (default: False).
> - disp: Show sum-of-least-squares after each optimisation iteration (default: True).
> - polish: True to use a Nelder–Mead simplex to polish the final result (default: True).
> - maxiter: Maximum number of iterations (default: 30).
> - popsize: Population size of genetic algorithm (default: 10).

The genetic algorithm will by default use all cores available for the optimisation. Using the *fit* attribute, it is possible to calibrate any number of model parameters to any sets of data. We do note that fitting the parameters sm,m,h and c requires modification of the source code. In the example below, the transmission parameter $\beta$ is sought after using two dataseries. The first is the number of patients in need of intensive care and the second is the total number of people in the hospital. The number of patients in ICU is matched with the CH population pool while the number of hospitalisations is matched with the sum of the HH and CH population pools.

```python
# hardcoded data vectors
ICUvect= np.array([[5,24,33,53,79,100,130,164,238,290,322,381,474,605,690,789,867,927,1021,1088,1144,1205,1245,1261,1257,1260,1276,1285,1278,1262,1232,1234,1223]])
hospital = np.array([[58,97,163,264,368,496,648,841,1096,1380,1643,1881,2137,2715,3068,3640,4077,4468,4884,4975,5206,5358,5492,5509,5600,5738,5692,5590,5610,5635,5409,5393,5536]])
# vector with dates
index=pd.date_range('2020-03-13', freq='D', periods=ICUvect.size)
# data series used to calibrate model must be given to function 'plotFit' as a list
idx = -26
index = index[0:idx]
data=[np.transpose(ICUvect[:,0:idx]),np.transpose(hospital[:,0:idx])]
# set optimisation settings
parNames = ['beta','p'] # must be a list!
positions = [np.array([7]),np.array([6,7])] # must be a list!
bounds=[(1,100),(0.1,1),(0.1,1)] # must be a list!
weights = np.array([1,0])
# run optimisation
theta = model.fit(data,parNames,positions,bounds,weights,setvar=True,maxiter=10,popsize=5)
```

### Visualising the fit

Visualising the resulting fit is easy and can be done using the plotFit function. The functions uses the following basic syntax,

> plotFit(index,data,positions)
> - index: vector with timestamps corresponding to data.
> - data: list containing the dataseries (dtype=np array) to fit the model to. For now, dataseries must be of equal length.
> - positions: list containing the model states (dtype=np.array) used to calculate the sum of least squares.

The following arguments are optional,
> - dataMkr: list containing the markers (dtype=str) to be used to visualise the data. Default value works up to five dataseries and is equal to: ['o','v','s','*','^'].
>- modelClr: list containing the colors (dtype=str) to be used to visualise the model fit. Default value works up to five dateseries and is equal to: ['green','orange','red','black','blue'].
> - legendText: tuple containing the legend entries. Disabled per default.
> - titleText: string containing the fit title. Disable per default.
> - filename: string with a filename + extension to save the figure. The figure is not saved per default.

```python
# plot result
model.plotFit(index,data,positions,modelClr=['red','orange'],legendText=('ICU (model)','ICU (data)','Hospital (model)','Hospital (data)'),titleText='Belgium')
```

# Code works untill here

To be continued...

## Model Predictive control (MPC)

### Optimising government policy

#### Process control for the layman

As we have the impression that the control part, which we see as our main addition to the problem, is more difficult to grasp for the layman, here is a short intro to process control. Experts in control are welcome to skip this section.

A predictive model consists of a set of equations and aims to predict how the system will behave in the future given a certain input. Process control flips this around and aims at determining what input is needed to achieve a desired system behavior (= goal). It is a tool that helps us in “controlling” how we want a system to behave. It is commonly applied in many industries, but also in our homes (e.g. central heating, washing machine). It's basically everywhere. Here's how it works. An algorithm monitors the deviation between the goal and the true system value and then computes the necessary action to "drive" the system to its goal by means of an actuator (in industry this is typically a pump or a valve). Applying this to Covid-19, the government wants to "control" the spread of the virus in the population by imposing measures (necessary control actions) on the public (which is the actuator here) and achieve the goal that the number of severely sick people does not become larger than can be handled by the health care system. However, the way the population behaves is a lot more complex compared to the heating control in our homes since not only epidemiology (virus spread) but also different aspects of human behavior on both the individual and the societal level (sociology, psychology, economy) are involved. This leads to multiple criteria we want to ideally control simultaneously and we want to use the "smartest" algorithm we can get our hands on.

#### The optimizePolicy function

The 'SEIRSNetworkModel' class contains an implementation of the MPC controller in the function *optimizePolicy*. For now, the controller minimises a weighted squared sum-of-errors between multiple setpoints and model predictions. The algorithm can use any variable to control the virus outbreak, but we recommend sticking with the number of random daily contacts $N_c$ and the total number of random tests ('totalTests') as only these have been tested. We also recommend disabling age-structuring in the model before running the MPC as this feature requires discretisation of the interaction matrix to work which is not yet implemented. Future work will extend the MPC controller to work with age-structuring feature inherent to the model. Future work is also aimed at including an economic cost function to discriminate between control handles. Our MPC uses the **genetic algorithm** from scipy.optimize to perform the optimisation, we recommend using at least a populationsize of 20 and at least 100 iterations to ensure that the trajectory is 'optimal'. The *optimizePolicy* function has the following basic syntax,

> optimizePolicy(parNames, bounds, setpoints, positions, weights)
> - parNames: a list containing the names (dtype=string) of the model parameters to be used as a control handle.
> - bounds: A list containing the lower- and upper boundaries of each parameter to be used as a control handle. Each entry in the list should be a 1D numpy array containing the lower- and upper bound for the respective control handle.
> - setpoints: A list with the numerical values of the desired model output.
> - positions: a list containing the model states (dtype=np.array) used to calculate the sum of least squares. Each modelouput in the given position is matched with a provided setpoint. If multiple position entries are provided, the output in these positions is added togheter. The order of the states is given according to the following vector, where S has index 0: (S, E, SM, M, H, C, HH, CH, R, F, SQ, EQ, SMQ, MQ, RQ).

The following arguments are optional,
> - policy_period: length of one policy interval (default: 7 days).
> - N: number of future policy intervals to be optimised, also called 'control horizon' (default: 6).
> - P: number of policy intervals over which the sum of squared errors is calculated, also called 'prediction horizon' (default:12).
> - disp: Show sum-of-least-squares after each optimisation iteration (default: True).
> - polish: True to use a Nelder–Mead simplex to polish the final result (default: True).
> - maxiter: Maximum number of iterations (default: 100).
> - popsize: Population size of genetic algorithm (default: 20).

The function returns a one-dimensional list containing the optimal values of the control handles. The length of this list is equal to the length of the control horizon (N) times the number of control handles. The list thus lists all control handles and their optimal values in their respective order. **The optimal policy is assigned to the SEIRSNetworkModel object and is only overwritten when a new optimisation is performed. Future work could include the creation of a new object for every optimal policy.** The genetic algorithm will by default use all cores available for the optimisation.

```python
parNames = ['Nc','totalTests']
bounds = [np.array([0,11.2]),np.array([0,1e6])]
setpoints = [1200,5000]
positions = [np.array([7]),np.array([6,7])]
weights = [1,0]
model.optimizePolicy(parNames,bounds,setpoints,positions,weights,policy_period=30,N=6,P=12,polish=False,maxiter=1,popsize=10)
```

### Visualising the effect of government policy

Visualising the resulting optimal policy is easy and can be done using the plotOptimalPolicy function. We note that the functionality of*plotOptimalPolicy** is for now, very basic and will be extended in the future. The function is heavily based on the *plotInfected* visualisation. The function uses the following basic syntax,

> plotOptimalPolicy(parNames,setpoints,policy_period)
> - parNames: a list containing the names (dtype=string) of the model parameters to be used as a control handle.
> - setpoints: A list with the numerical values of the desired model output.
> - policy_period: length of one policy interval (default: 7 days).

The following arguments are optional,
> - asymptotic: set to *True* to include the supermild pool in the visualisation.
> - mild: set to *True* to include the mild pool in the visualisation.
> - filename: string with a filename + extension to save the figure. The figure is not saved per default.

```python
model.plotOptimalPolicy(parNames,setpoints,policy_period=14)
```

## Scenario-specific extensions

### *realTimeScenario*

The 'SEIRSNetworkModel' class contains one function to quickly perform and visualise scenario analysis for a given country. The user is obligated to supply the function with: 1) a set of dataseries, 2) the date at which the data starts, 3) the positions in the model output that correspond with the dataseries and 4) a checkpoints dictionary containing the past governement actions, from hereon referred to as the *pastPolicy* dictionary. If no additional arguments are provided, the data and the corresponding model fit are visualised from the user supplied start date up until the end date of the data plus 14 days. The end date of the visualisation can be altered by defining the optional keyworded argument *T_extra* (default: 14 days). Optionally a dictionary of future policies can be used to simulate scenarios starting on the first day after the end date of the dataseries. The function *realTimeScenario* accomplishes this by merging both the *pastPolicy* and *futurePolicy* dictionaries using the backend function *mergeDict()*. The syntax without optional arguments is as follows,

> realTimeScenario(startDate, data, positions, pastPolicy)
> - startDate: a string with the date corresponding to the first entry of the dataseries (format: 'YYYY-MM-DD').
> - data: a list containing the dataseries (dtype=np array) to fit the model to. For now, dataseries must be of equal length and start on the same day.
> - positions: a list containing the model states (dtype=np.array) used to calculate the sum of least squares. Each dataseries must be matched to a certain (sum of) model state(s). If multiple entries are provided these are added togheter. The order of the states is given according to the following vector, where S has index 0: (S, E, SM, M, H, C, HH, CH, R, F, SQ, EQ, SMQ, MQ, RQ).
> - pastPolicy: a checkpoints dictionary containing past government actions.

The following (simulation) arguments are optional,
> - futurePolicy: a checkpoint dictionary used to simulate scenarios in the future (default: None). By default, time '1' in this dictionary is the date of the first day after the end of the data.
> - T_extra: Extra simulation time after last date of the data if no futurePolicy dictionary is provided. Extra simulation time after last time in futurePolicy dictionary.

The following arguments are for visualisation,
> - dataMkr: list containing the markers (dtype=str) to be used to visualise the data. Default value works up to five dataseries and is equal to: ['o','v','s','*','^'].
>- modelClr: list containing the colors (dtype=str) to be used to visualise the model fit. Default value works up to five dateseries and is equal to: ['green','orange','red','black','blue'].
> - legendText: tuple containing the legend entries. Disabled per default.
> - titleText: string containing the fit title. Disable per default.
> - filename: string with a filename + extension to save the figure. The figure is not saved per default.

```python
# Define data as a list containing data timeseries
data=[np.transpose(ICUvect),np.transpose(hospital)]
# Create a dictionary of past policies
pastPolicy = {'t':       [11],
              'G':       [G_dist],
              'p':       [0.03]
              }
# Create a dictionary of future policies
futurePolicy = {'t':     [1],
              'G':       [G_norm],
              'p':       [0.6]
              }
# Define the data corresponding to the first data entry
startDate='2020-03-13'
# Run realTimeScenario
model.realTimeScenario(startDate,data,positions,pastPolicy,futurePolicy=futurePolicy,T_extra=7,
                       modelClr=['red','orange'],legendText=('ICU (model)','Hospital (model)','ICU (data)','Hospital (data)'),
                       titleText='Belgium')
```