# ADVI from scratch in JAX

Authors: karm-patel@, murphyk@

In this notebook we apply ADVI (automatic differentiation variational inference) to the beta-binomial model, using a Normal Distribution as Variational Posterior. This involves a change of variable from the unconstrained z in R space to the constrained theta in [0,1] space.

```python
try:
    import jax
except ModuleNotFoundError:
    %pip install -qqq jax jaxlib
    import jax

import jax.numpy as jnp
from jax import lax

try:
    from tensorflow_probability.substrates import jax as tfp
except ModuleNotFoundError:
    %pip install -qqq tensorflow_probability
    from tensorflow_probability.substrates import jax as tfp

try:
    import optax
except ModuleNotFoundError:
    %pip install -qqq optax
    import optax

try:
    from rich import print
except ModuleNotFoundError:
    %pip install -qqq rich
    from rich import print

try:
    from tqdm import trange
except:
    %pip install -qqq tqdm
    from tqdm import trange

import seaborn as sns
import numpy as np
import matplotlib.pyplot as plt
import warnings

warnings.filterwarnings("ignore")
dist = tfp.distributions

plt.rc("font", size=10)  # controls default text sizes
plt.rc("axes", labelsize=12)  # fontsize of the x and y labels
plt.rc("legend", fontsize=12)  # legend fontsize
plt.rc("figure", titlesize=15)  # fontsize of the figure title
```

## Functions
Helper functions which will be used later

```python
def prior_dist():
    return dist.Beta(concentration1 = 1.0, concentration0 = 1.0)

def likelihood_dist(theta):
    return dist.Bernoulli(probs=theta)

def transform_fn(x):
    return 1 / (1 + jnp.exp(-x))  # sigmoid

def positivity_fn(x):
    return jnp.log(1 + jnp.exp(x))  # softplus

def variational_distribution_q(params):
    loc = params["loc"]
    scale = positivity_fn(params["scale"]) # apply softplus
    return dist.Normal(loc, scale)

jacobian_fn = jax.jacfwd(transform_fn) # define function to find jacobian for tranform_fun
```

## Dataset
Now, we will create the dataset. we sample `theta_true` (probability of occurring head) random variable from the prior distribution which is Beta in this case. Then we sample `n_samples` coin tosses from likelihood distribution which is Bernouli in this case.

```python
# preparing dataset
# key = jax.random.PRNGKey(128)
# n_samples = 12
# theta_true =  prior_dist().sample((5,),key)[0]

# dataset = likelihood_dist(theta_true).sample(n_samples,key)

# print(f"Dataset: {dataset}")
# n_heads = dataset.sum()
# n_tails = n_samples - n_heads
```

```python
# Use same data as https://github.com/probml/probml-notebooks/blob/main/notebooks/beta_binom_approx_post_pymc.ipynb
key = jax.random.PRNGKey(128)
dataset = np.repeat([0, 1], (10, 1))
n_samples =  len(dataset)
print(f"Dataset: {dataset}")
n_heads = dataset.sum()
n_tails = n_samples - n_heads
```

Dataset: [0 0 0 0 0 0 0 0 0 0 1]

## Prior, Likelihood, and True Posterior
For coin toss problem, since we know the closed form solution of posterior, we compare the distributions of Prior, Likelihood, and True Posterior below.

```python
# closed form of beta posterior
a = prior_dist().concentration1
b = prior_dist().concentration0

exact_posterior = dist.Beta(concentration1 = a + n_heads, concentration0 = b + n_tails)

theta_range = jnp.linspace(0.01, 0.99, 100)

ax = plt.gca()
ax2 = ax.twinx()
plt2, = ax2.plot(theta_range, exact_posterior.prob(theta_range), "g--", label="True Posterior")
plt3, = ax2.plot(theta_range, prior_dist().prob(theta_range), label="Prior")

likelihood = jax.vmap(lambda x: jnp.prod(likelihood_dist(x).prob(dataset)))(theta_range)
plt1, = ax.plot(theta_range, likelihood, "r-.", label="Likelihood")

ax.set_xlabel("theta")
ax.set_ylabel("Likelihood")
ax2.set_ylabel("Prior & Posterior")
ax2.legend(handles=[plt1,plt2,plt3],bbox_to_anchor=(1.6,1));
```

## Optimizing the ELBO
In order to minimize KL divergence between true posterior and variational distribution, we need to minimize the negative ELBO, as we describe below.

We start with the ELBO, which is given by:
\begin{align}
ELBO(\psi) &= E_{z \sim q(z|\psi)} \left[
  p(\mathcal{D}|z) + \log p(z) - \log q(z|\psi) \right]
\end{align}
where
$\psi = (\mu, \sigma)$ are the variational parameters,
$p(\mathcal{D}|z) =  p(\mathcal{D}|\theta=\sigma(z))$
is the likelihood,
and the prior is given by the change of variables formula:
\begin{align}
p(z) &= p(\theta) | \frac{\partial \theta}{\partial z} |
 =  p(\theta) | J |
\end{align}
where $J$ is the Jacobian of the $z \rightarrow \theta$ mapping.
We will use a Monte Carlo approximation of the expectation over $z$.
We also apply the reparameterization trick
to replace $z \sim q(z|\psi)$ with
\begin{align}
\epsilon &\sim \mathcal{N}(0,1 ) \\
z &= \mu + \sigma \epsilon
\end{align}
Putting it altogether our estimate for the negative ELBO (for a single sample of $\epsilon$) is
\begin{align}
-L(\psi; z) &=  -( \log p(\mathcal{D}|\theta )
+\log p( \theta)   + \log|J_\boldsymbol{\sigma}(z)|)
+ \log q(z|\psi)
\end{align}

```python
def log_prior_likelihood_jacobian(normal_sample, dataset):
    theta = transform_fn(normal_sample)  # transform normal sample to beta sample
    likelihood_log_prob = likelihood_dist(theta).log_prob(dataset).sum() # log probability of likelihood
    prior_log_prob =  prior_dist().log_prob(theta) # log probability of prior
    log_det_jacob = jnp.log(jnp.abs(jnp.linalg.det(jacobian_fn(normal_sample).reshape(1,1)))) # log of determinant of jacobian
    return likelihood_log_prob +  prior_log_prob + log_det_jacob
```

```python
# reference: https://code-first-ml.github.io/book2/notebooks/introduction/variational.html
def negative_elbo(params, dataset, n_samples = 10, key=jax.random.PRNGKey(1)):

    q = variational_distribution_q(params) # Normal distribution.
    q_loc, q_scale = q.loc, q.scale

    std_normal = dist.Normal(0, 1)
    sample_set = std_normal.sample(seed=key, sample_shape=[n_samples,])
    sample_set = q_loc + q_scale * sample_set  # reparameterization trick

    #calculate log joint for each sample of z
    p_log_prob = jax.vmap(log_prior_likelihood_jacobian, in_axes=(0, None))(sample_set, dataset)

    return jnp.mean(q.log_prob(sample_set) - p_log_prob)
```

We now apply stochastic gradient descent to minimize negative ELBO and optimize the variational parameters (`loc` and `scale`)

```python
loss_and_grad_fn = jax.value_and_grad(negative_elbo, argnums=(0))
loss_and_grad_fn = jax.jit(loss_and_grad_fn) # jit the loss_and_grad function

params = {"loc": 0.0, "scale": 0.5}
elbo, grads = loss_and_grad_fn(params, dataset)

print(f"loss: {elbo}")
print(f"grads:\n loc: {grads['loc']}\n scale: {grads['scale']} ")

optimizer = optax.adam(learning_rate = 0.01)
opt_state = optimizer.init(params)
```

loss: 9.09147834777832

grads:
 loc: 4.692040920257568
 scale: 0.6960504651069641

```python
#jax scannable function for training
def train_step(carry, data_output):

    #take carry data
    key = carry["key"]
    elbo = carry["elbo"]
    grads =  carry["grads"]
    params = carry["params"]
    opt_state = carry["opt_state"]
    updates = carry["updates"]

    #training
    key, subkey = jax.random.split(key)
    elbo, grads = loss_and_grad_fn(params, dataset, key = subkey)
    updates, opt_state = optimizer.update(grads, opt_state)
    params = optax.apply_updates(params, updates)

    #forward carry to next iteration by storing it
    carry =  { "key": subkey, "elbo": elbo, "grads": grads, "params": params,
              "opt_state": opt_state, "updates": updates }

    output = { "elbo": elbo, "params": params }

    return carry, output

```

```python
%%time

#dummy iteration to pass carry to jax scannale function train()
key, subkey = jax.random.split(key)
elbo, grads = loss_and_grad_fn(params, dataset, key = subkey)
updates, opt_state = optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)

carry = {
    "key": key,
    "elbo": elbo,
    "grads": grads,
    "params": params,
    "opt_state": opt_state,
    "updates": updates
}

num_iter = 1000
elbos = np.empty(num_iter)

#apply scan() to optimize training loop
last_carry, output = lax.scan(train_step, carry, elbos)

elbo = output["elbo"]
params = output["params"]
optimized_params = last_carry["params"]
```

    CPU times: user 2.09 s, sys: 20.1 ms, total: 2.11 s
    Wall time: 2.03 s

```python
print(params['loc'].shape)
print(params['scale'].shape)
```

(1000,)

(1000,)

We now plot the ELBO

```python
plt.plot(elbo)
plt.xlabel("Iterations")
plt.ylabel("Negative ELBO")
sns.despine()
plt.savefig('advi_beta_binom_jax_loss.pdf')
```

We can see that after 200 iterations ELBO is optimized and not changing too much.

## Samples using Optimized parameters

Now, we take 1000 samples from variational distribution (Normal) and transform them into true posterior distribution (Beta) by applying `tranform_fn` (sigmoid) on samples. Then we compare density of samples with exact posterior.

```python
q_learned = variational_distribution_q(optimized_params)
key = jax.random.PRNGKey(128)
q_learned_samples = q_learned.sample(1000, seed=key) # q(z|D)
transformed_samples = transform_fn(q_learned_samples) #transform Normal samples into Beta samples

theta_range = jnp.linspace(0.01,0.99,100)
plt.plot(theta_range, exact_posterior.prob(theta_range),"r" , label="$p(x)$: true posterior")
sns.kdeplot(transformed_samples, color="blue" ,label="$q(x)$: learned", bw_adjust=1.5, clip=(0.0,1.0), linestyle="--")

plt.xlabel("theta")
plt.legend() #bbox_to_anchor=(1.5, 1));
sns.despine();

plt.savefig('advi_beta_binom_jax_posterior.pdf')
```

We can see that the learned `q(x)` is a reasonably good approximation to the true posterior. It seems to have support over negative theta but this is an artefact of KDE.

```python
#print(transformed_samples)
print(len(transformed_samples))
print(jnp.sum(transformed_samples < 0)) # all samples of thetas should be in [0,1]
print(jnp.sum(transformed_samples > 1)) # all samples of thetas should be in [0,1]
```

1000

0

0

```python
print(q_learned)
print(q_learned.mean())
print(jnp.sqrt(q_learned.variance()))
```

tfp.distributions.Normal("Normal", batch_shape=[], event_shape=[], dtype=float32)

-1.9198773

0.8108946

```python
locs, scales = params["loc"], params["scale"]
sigmas = positivity_fn(jnp.array(scales))

plt.plot(locs, label="mu")
plt.xlabel("Iterations")
plt.ylabel("$E_q[z]$")
plt.legend()
sns.despine();
plt.savefig('advi_beta_binom_jax_post_mu_vs_time.pdf')
plt.show()

plt.plot(sigmas, label="sigma")
plt.xlabel("Iterations")
#plt.ylabel(r'$\sqrt{\text{var}(z)}')
plt.ylabel("$std_{q}[z]$")
plt.legend()
sns.despine();
plt.savefig('advi_beta_binom_jax_post_sigma_vs_time.pdf')
plt.show()

```

## Comparison with pymc.ADVI()

Now, we compare our implementation with pymc's ADVI implementation.

**Note**: For pymc implementation, the code is taken from this notebook: https://github.com/probml/probml-notebooks/blob/main/notebooks/beta_binom_approx_post_pymc.ipynb

```python
try:
    import pymc3 as pm
except ModuleNotFoundError:
    %pip install pymc3
    import pymc3 as pm
try:
    import scipy.stats as stats
except ModuleNotFoundError:
    %pip install scipy
    import scipy.stats as stats

import scipy.special as sp

try:
    import arviz as az
except ModuleNotFoundError:
    %pip install arviz
    import arviz as az

import math
```

```python
a = prior_dist().concentration1
b = prior_dist().concentration0

with pm.Model() as mf_model:
    theta = pm.Beta("theta", a, b)
    y = pm.Binomial("y", n=1, p=theta, observed=dataset)  # Bernoulli
    advi = pm.ADVI()
    tracker = pm.callbacks.Tracker(
        mean=advi.approx.mean.eval,  # callable that returns mean
        std=advi.approx.std.eval,  # callable that returns std
    )
    approx = advi.fit(callbacks=[tracker],n=20000)

trace_approx = approx.sample(1000)
thetas = trace_approx["theta"]
```

  <progress value='20000' class='' max='20000' style='width:300px; height:20px; vertical-align: middle;'></progress>
  100.00% [20000/20000 00:01<00:00 Average Loss = 4.9174]

```python
plt.plot(advi.hist, label="ELBO")
plt.xlabel("Iterations")
plt.ylabel("ELBO")
plt.legend()
sns.despine();
plt.savefig('advi_beta_binom_pymc_loss.pdf')
plt.show()
```

```python
print(f"ELBO comparison for last 1% iterations:\nJAX ELBO: {elbo[-10:].mean()}\nPymc ELBO: {advi.hist[-100:].mean()}")
```

ELBO comparison for last 1% iterations:
JAX ELBO: 4.899871826171875
Pymc ELBO: 4.941951796589394

## True posterior, JAX q(x), and pymc q(x)

```python
plt.plot(theta_range, exact_posterior.prob(theta_range), "b--", label="$p(x)$: True Posterior")
sns.kdeplot(transformed_samples, color="red", label="$q(x)$: learnt - jax", clip=(0.0,1.0), bw_adjust=1.5)
sns.kdeplot(thetas, label="$q(x)$: learnt - pymc", clip=(0.0,1.0), bw_adjust=1.5)
plt.xlabel("theta")
plt.legend(bbox_to_anchor=(1.3, 1))
sns.despine()
```

## Plot of loc and scale for variational distribution

```python
fig1, (ax1,ax2) =  plt.subplots(1,2, figsize = (10,4), sharey=True)
locs, scales = params["loc"], params["scale"]

#plot loc

#JAX
ax1.plot(locs, label="JAX: loc")
ax1.set_ylabel("loc")
ax1.legend();

#pymc
ax2.plot(tracker['mean'], label="Pymc: loc")
ax2.legend()
sns.despine()

#plot scale
fig2, (ax3,ax4) =  plt.subplots(1,2, figsize = (10,4), sharey=True)

#JAX
ax3.plot(positivity_fn(jnp.array(scales)), label="JAX: scale"); #apply softplus on scale
ax3.set_xlabel("Iterations")
ax3.set_ylabel("scale")
ax3.legend();

#pymc
ax4.plot(tracker['std'], label="Pymc: scale")
ax4.set_xlabel("Iterations")
ax4.legend();

sns.despine();
```

## References:

1. ADVI paper: https://arxiv.org/abs/1603.00788
2. Blog: https://code-first-ml.github.io/book2/notebooks/introduction/variational.html
3. Blog: https://luiarthur.github.io/statorial/varinf/introvi/
4. Video: https://www.youtube.com/watch?v=HxQ94L8n0vU
5. Github issue: https://github.com/pyro-ppl/pyro/issues/3016#:~:text=loc%3D%27upper%20right%27)-,Bandwidth%20adjustment,-Another%20thing%20to
6. Blog: https://ericmjl.github.io/dl-workshop/02-jax-idioms/02-loopy-carry.html

```python

```


# Computing moments after cell dvision

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

---

```python
import os
import pickle
import glob
import cloudpickle
import git

import itertools # Library we'll use to generate possible pairs of numbers

# Numerical workhorse
import numpy as np
import pandas as pd

# To compute symbolic expressions
import sympy
sympy.init_printing(use_unicode=True, use_latex=True) # print outputs in LaTeX

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

# Seaborn, useful for graphics
import seaborn as sns

# This enables high resolution graphics inline
%config InlineBackend.figure_format = 'retina'

# Import the project utils
import ccutils

# Find home directory for repo
repo = git.Repo("./", search_parent_directories=True)
homedir = repo.working_dir

tmpdir = f'{homedir}/tmp/'
figdir = f'{homedir}/fig/moment_dynamics_numeric/'
datadir = f'{homedir}/data/csv_maxEnt_dist/'
```

```python
# Set PBoC plotting format
ccutils.viz.set_plotting_style()
# Increase dpi
mpl.rcParams['figure.dpi'] = 110
```

### $\LaTeX$ macros
$\newcommand{kpon}{k^{(p)}_{\text{on}}}$
$\newcommand{kpoff}{k^{(p)}_{\text{off}}}$
$\newcommand{kron}{k^{(r)}_{\text{on}}}$
$\newcommand{kroff}{k^{(r)}_{\text{off}}}$
$\newcommand{rm}{r _m}$
$\newcommand{gm}{\gamma _m}$
$\newcommand{ee}[1]{\left\langle #1 \right\rangle}$
$\newcommand{bb}[1]{\mathbf{#1}}$
$\newcommand{th}[1]{{#1}^{\text{th}}}$
$\newcommand{dt}[1]{{\partial{#1} \over \partial t}}$
$\newcommand{Km}{\bb{K}}$
$\newcommand{Rm}{\bb{R}_m}$
$\newcommand{Re}{\bb{R}_m'}$
$\newcommand{Gm}{\bb{\Gamma}_m}$
$\newcommand{Rp}{\bb{R}_p}$
$\newcommand{Gp}{\bb{\Gamma}_p}$

## Computing the moments of the distribution after cell division.

We have already solved a general form for the dynamics of the moments of the distribution, i.e. we wrote differential equations for the moments ${d\ee{m^x p^y}\over dt}$. Given that we know all parameters for our model we can simply integrate these equations numerically to compute how the moments of the distribution evolve as cells progress through their cell cycle. Once the cell reaches a time $t_d$ when is going to divide the mRNA and proteins that we are interested in undergo a binomial partitioning between the two daughter cells. In other words each molecule flips a coin and decides whether to go to either daughter. The question then becomes given that we have a value for the moment $\ee{m^x p^y}_{t_d}$ at a time before the cell division, what would the value of this moment be after the cell division takes place $\ee{m^x p^y}_{t_o}$?

We know that the probability distribution of mRNA and protein after the cell division $P_{t_o}(m, p)$ is of the form
$$
P_{t_o}(m, p) = \sum_{m', p'} P(m, p \mid m', p') P_{t_d}(m', p'),
$$
where we are basically summing over all the possibilities of having $m'$ mRNA and $p'$ proteins before cell division. The conditional probability $P(m, p \mid m', p')$ is simply given by a product of two binomial distributions, one for the mRNA and one for the protein; this is
$$
P(m, p \mid m', p') = {m' \choose m}\left( {1 \over 2} \right)^{m'}
{p' \choose p}\left( {1 \over 2} \right)^{p'}.
$$
This is because the partitionining of mRNA and protein between cells is independent of each other.

Given this functional form we can then compute the moments after the cell division $\ee{m^x p^y}_{t_o}$ as
$$
\ee{m^x p^y}_{t_o} = \sum_m \sum_p m^x p^y P_{t_o}(m, p).
$$
Substituting the definition of the probability distribution results in
$$
\ee{m^x p^y}_{t_o} = \sum_m \sum_p m^x p^y
\sum_{m', p'} P(m, p \mid m', p') P_{t_d}(m', p').
$$
We can rearange the sums to be
$$
\ee{m^x p^y}_{t_o} = \sum_{m', p'} P_{t_d}(m', p')
\sum_m \sum_p m^x p^y
P(m, p \mid m', p').
$$
We mentioned that the partitioning of mRNA and protein between daughter cells are independent events. This can be expressed as
$$
P(m, p \mid m', p') = P(m \mid m') \cdot P(p \mid p').
$$
With this we can then write the moment $\ee{m^x p^y}_{t_o}$ as
$$
\ee{m^x p^y}_{t_o} = \sum_{m', p'} P_{t_d}(m', p')
\sum_m  m^x  P(m \mid m')
\sum_p p^y P(p \mid p').
$$
Notice that the terms with the conditional distributions are of the form
$$
\sum_z  z^x  P(z \mid z') \equiv \ee{z^x \mid z'}, \;
{\text{ for } z\in \{m, p \}},
$$
i.e. conditional expected values. These conditional expected values are the expected values of a binomial random variable $z \sim \text{Bin}(z', 1/2)$. We can then rewrite the expected values after the cell division as
$$
\ee{m^x p^y}_{t_o} = \sum_{m', p'} \ee{m^x \mid m'} \ee{p^y \mid p'}
P_{t_d}(m', p').
$$

To see how this general formula for the moments after the cell division works let's compute the mean protein per cell after the cell division $\ee{p}_{t_o}$. That is setting $x = 0$, and $y = 1$. This results in
$$
\ee{p}_{t_o} = \sum_{m', p'} \ee{m^0 \mid m'} \ee{p \mid p'} P_{t_d}(m', p').
$$
The zeroth moment $\ee{m^0 \mid m'}$ by definition must be one since we have
$$
\ee{m^0 \mid m'} = \sum_m m^0 P(m \mid m') = \sum_m P(m \mid m') = 1,
$$
since the probability distribution must be normalized. This leaves us then with
$$
\ee{p}_{t_o} = \sum_{m', p'} P_{t_d}(m', p') \ee{p \mid p'}.
$$
If we take the sum over $m'$ we simply compute the marginal probability distribuiton $\sum_{m'} P_{t_d}(m', p') = P_{t_d}(p')$, then we have
$$
\ee{p}_{t_o} = \sum_{p'} \ee{p \mid p'} P_{t_d}(p').
$$
For the particular case of the first moment of the binomial distribution with parameters $p'$ and $1/2$ we know that
$$
\ee{p \mid p'} = {p' \over 2}.
$$
Therefore the moment after division is equal to
$$
\ee{p}_{t_o} = \sum_{p'} {p' \over 2} P_{t_d}(p') = {\ee{p'}_{t_d} \over 2}.
$$
This result makes perfect sense. What this is saying is that the mean protein copy number right after the cell divides is half of the mean protein copy number just before the cell division. That is exactly we would expect.

## Moment Generating Function for the binomial distribution

The result from last section was dependent on us knowing the functional form of the firt moment of the binomial distribution. For higher moments we need some systematic way to compute such moments. Luckily for us we can do so by using the so-called moment generating function (MGF). The MGF of a random variable $X$ is defined as
$$
M_X(t) = \ee{e^{tX}},
$$
where $t$ is a dummy variable. Once we know the MGF we can obtain any moment of the distribution by simply computing
$$
\ee{X^m} = \left. {d^m \over dt^m} M_X(t) \right\vert_{t=0},
$$
i.e. taking the n-th derivative. For the particular case of the binomial distribution $X \sim \text{Bin}(N, q)$ it can be shown that the MGF is of the form
$$
M_X(t) = \left[ (1 - q) + qe^{t} \right]^N.
$$

The objective of this notebook will be to use `sympy` to compute such derivatives and therefore generate explicit formulas for the moments of the distribution right after the cell division.

We begin by initializing the symbols that we will need.

```python
# Define mRNA and protein variables
m, p, = sympy.symbols('m p')
# Define the exponents
x, y = sympy.symbols('x y')
# Define dummy variable for MGF
t, N, q = sympy.symbols('t N q')

m**x, p**y, t, N, q
```

Now let's define the moment generating function.

```python
# Define MGF for binomial N, p
mgf = ((1 - q) + q * sympy.exp(t))**N

mgf
```

As an example let's make sure we can reproduce the known first moment by taking the first derivative.

```python
# Compute first deriative and substitute t = 0.
first_mom = mgf.diff((t,1)).subs({t: 0})

# Simplify and print expression.
first_mom.simplify()
```

This is indeed the expected value of the binomial distribution.
Let's follow up the analysis, compute the second moment and then the variance

```python
# Compute the second moment and substitute t = 0.
second_mom = mgf.diff((t, 2)).subs({t: 0})

# Compute the variance
var = second_mom - first_mom**2

# Print variance
var.simplify()
```

Again, the expected variance of the binomial distribution. Now we are in position of computing the rest of the moments systematically.

```python
# Define moments to compute
mom = np.arange(1, 7)

# Initialize list to store moments
mom_list = list()

# Loop through moments and compute derivative
for i in mom:
    # Compute and store derivative
    mom_list.append(mgf.diff((t, i)).subs({t: 0,
                    q: sympy.Rational(1,2)}).expand())
```

Let's take a look at all the derivatives together.

```python
mom_list
```

What we need to create is a function that for a given pair of exponents $x$ and $y$ computes the product of the two binomials

```python
def binom_prod(x, y):
    '''
    Function that computes the the product of the binomials for a given
    pair of exponents x, y.
    Parameters
    ----------
    x, y : int.
        Exponents of the moment <m**x * p**y>

    Returns
    -------
    binom_pro : Sympy eq.
        Product of the two binomial moments
    '''
    # Compute mRNA related moment
    mRNA_mom = mgf.diff((t, x)).subs({t: 0, N: m,
                                      q: sympy.Rational(1,2)})

    # Compute protein related moment
    protein_mom = mgf.diff((t, y)).subs({t: 0, N: p,
                                         q: sympy.Rational(1,2)})

    # Compute product
    binom_p = mRNA_mom * protein_mom

    return binom_p.expand()
```

Let's look at a couple of examples to gain intuition on these results. For example let's first take a look at the moment $\ee{mp}$.

```python
binom_prod(1, 1)
```

This result indicates is that when computing the moment after cell division $\ee{mp}_{t_o}$ which is of the form
$$
\ee{mp}_{t_o} = \sum_{m', p'} \ee{m \mid m'} \ee{p \mid p'} P_{t_d}(m', p'),
$$
the product $\ee{m \mid m'} \ee{p \mid p'}$ is then
$$
\ee{m \mid m'} \ee{p \mid p'} = {m' p' \over 4}.
$$
Substituting this result into the moment would then give
$$
\ee{mp}_{t_o} = \sum_{m', p'} {m' p' \over 4} P_{t_d}(m', p') =
{\ee{m' p'}_{t_d} \over 4}.
$$
Therefore to compute the moment after cell division $\ee{mp}_{t_o}$ we simply have to divide by 4 the same moment before the cell division.

This was an easy example. Let's take a look at another moment, for example $\ee{p^3}$.

```python
binom_prod(0, 3)
```

Following the same logic as before this result indicates that in order to compute the moment before cell division $\ee{p^3}_{t_o}$ we would compute
$$
\ee{p^3}_{t_o} = {\ee{p^3}_{t_d} \over 8} + {3 \ee{p^2}_{t_d} \over 8}.
$$
So for this particular case the third moment of the protein distribution depends on the third moment and the second moment before the cell division.

Let us then define a vector to collect all moments before the cell division up  the $\ee{m^x p^y}$, i.e.
$$
\bb{\ee{m^x p^y}}_{t_d} = \left(
\ee{m^0 p^0}_{t_d}, \ee{m^1}_{t_d}, \ldots , \ee{m^x p^y}_{t_d}
\right).
$$
Then the any moment after cell division $\ee{m^{x'}, p^{y'}}_{t_o}$ for $x' \leq x$ and $y' \leq y$ can be computed as
$$
\ee{m^{x'} p^{y'}}_{t_o} = \bb{z}_{x'y'} \cdot \bb{\ee{m^x p^y}}_{t_d},
$$
where we define the vector $\bb{z}_{x'y'}$ as the vector containing all the coefficients that we obtain with the product of the two binomial distributions.
For example for the case of the third protein moment $\ee{p^3}_{t_o}$ the vector $\bb{z}_{x'y'}$ would have zeros for all entries except for the corresponding entry for $\ee{p^2}_{t_d}$ and for $\ee{p^3}_{t_d}$, where it would have $3/8$ and $1/8$ accordingly.

If we want then to compute all the moments after the cell division up to $\ee{m^x p^y}_{t_o}$ let us define an equivalent vector
$$
\bb{\ee{m^x p^y}}_{t_o} = \left(
\ee{m^0 p^0}_{t_o}, \ee{m^1}_{t_o}, \ldots , \ee{m^x p^y}_{t_o}
\right).
$$
Then we need to build a square matrix $\bb{Z}$ such that each row of the matrix contains the corresponding vector $\bb{z}_{x' y'}$ for each of the moments. Having this matrix we would simply compute the moments after the cell division as
$$
\bb{\ee{m^x p^x}}_{t_o} = \bb{Z} \cdot \bb{\ee{m^x p^x}}_{t_d}.
$$
In other words, matrix $\bb{Z}$ will contain all the coefficients that we need to multiply by the moments before the cell division in order to obtain the moments after cell division.

Let's systematically compute this matrix for all the moments up to the 6th protein moment.

### Extract coefficients of polynomial equation.

The first function we need to define is one that takes a result from our `binom_prod` function and extracts all the coefficients into a dictionary. For example if we were to feed the result for the third protein moment $\ee{p^3}$ this function must return a dictionary of the form
$$
\{
p^2: {3 \over 8}, p^3: {1 \over 8}
\}.
$$

```python
def coeff_dictionary(eq):
    '''
    Returns a dictionary with each of the coefficients for a given
    equation.

    Parameter
    ---------
    eq : sympy expression.
        Sympy expression for the equation

    Returns
    -------
    coeff_dict : dictionary.
        Dictionary containing all the coefficients of each of the elements
        in the polynomial equation
    '''
    # Find the degree of the eqent for each of the variables
    if eq.has(m):
        m_degree = sympy.Poly(eq).degree(m)
    else:
        m_degree = 0
    if eq.has(p):
        p_degree = sympy.Poly(eq).degree(p)
    else:
        p_degree = 0

    return {m**x * p**y: eq.coeff(m**x * p**y).\
            subs([[m, 0], [p, 0]])
            for x in range(m_degree + 1)
            for y in range(p_degree + 1)}
```

Let's now test the function

```python
coeff_dictionary(binom_prod(0, 3))
```

This looks exactly like we expected. Let's look at a slightly more complicated moment such as $\ee{m^3 p^3}$. First we look at the product from the two binomial distributions.

```python
binom_prod(3, 3)
```

Now we know what entries our dictionary must contain. Let's see if the function actually returns what we expect.

```python
coeff_dictionary(binom_prod(3, 3))
```

This is indeed the expected dictionary, so the function is working.

### Systematically finding the moments necessary to solve for $\bb{\ee{m^x p^y}}$

We established that a moment of the form $\bb{\ee{m^x p^y}}$ depends on lower moments that satisfy two conditions. We need to define a function that given the largest moment $\bb{\ee{m^x p^y}}$ to be included in the matrix it finds all of the moments $\bb{\ee{m^{x'} p^{y'}}}$ that satisfy such conditions.

```python
def expo_pairs(m_expo, p_expo):
    '''
    Finds all of the pairs of exponents x', y' that are necessary to compute the
    moment <m**m_expo * p**p_expo> that satisfy the necessary conditions:
    1) y' <= p_expo
    2) x' + y' <= m_expo + p_expo
    Parameters
    ----------
    m_expo, p_expo: int.
        Exponents of the highest moment to be included in the system.
        m_expo corresponds to the mRNA exponent
        p_expo corresponds to the protein exponent

    Returns
    -------
    pairs : list.
        List of sorted exponent pairs necessary for the computation
    NOTE: The sorting (which is not a necessary feature) given the general
    form of the equation only works up to the moment <m**0 * p**p_expo>.
    Any moment with x'>0 and y'=p_expo will be out of order appended at
    the end of the list.
    '''
    # Find all possible pair of exponents that satisfy
    # x' + y' <= m_expo + p_expo
    expo_pairs = list(itertools.permutations(range(m_expo + p_expo + 1), 2))

    # Add the (num, num) pair that are not being included
    expo_pairs = expo_pairs + [tuple([s, s]) for s in
                               range(max([m_expo, p_expo]) + 1)]

    # Remove pairs that do not satisfy the condition
    # y' <= p_expo
    expo_pairs = [s for s in expo_pairs if s[1] <= p_expo]

    # Remove pairs that do not satisfy the condition
    # x' <= m_expo + 1
#     expo_pairs = [x for x in expo_pairs if x[0] <= m_expo + 1]

    # # Remove pairs that do not satisfy the condition
    # x' + y' <= m_expo + p_expo
    expo_pairs = [s for s in expo_pairs if sum(s) <= m_expo + p_expo]

    ##  Moment sorting ##
    # Initialize list to append sorted moments
    expo_sorted = list()

    # Append mRNA moments
    mRNA_mom = sorted([s for s in expo_pairs if s[1] == 0])
    expo_sorted.append(mRNA_mom)

    # Find each protein moment
    protein_mom = sorted([s for s in expo_pairs if (s[0] == 0) & (s[1] != 0)])

    # Loop through each protein moment and find the cross correlations
    # associated with it
    for pr in protein_mom:
        cross_corr = sorted([s for s in expo_pairs
                             if (s[0] > 0) & (s[1] > 0) & (sum(s) == pr[1])],
                            reverse=True)
        # append it to the list
        expo_sorted.append(cross_corr)
        expo_sorted.append([pr])

    expo_sorted = list(itertools.chain.from_iterable(expo_sorted))

    # Append the other terms that are missing
    missing = [s for s in expo_pairs if s not in expo_sorted]

    return expo_sorted + missing
```

Let's test the function by finding the moments that would be needed to obtain the second moment of the protein distribution.

```python
# Define exponents for highest moment to be inferred
m_expo, p_expo = 0, 2

# Find the list of moments that need to be computed
expo_pairs(m_expo, p_expo)
```

These are indeed the moments needed to compute the second protein moment. Now we are ready to generate the matrix $\bb{Z}$. What we need to do is indicate the exponents $x$ and $y$ of the largest moment to be included, obtain all of the moments that are needed in order to compute the this moment, and loop through each of them extracting the coefficients at each step. This is better explained in the implementation of the function.

```python
def Z_matrix_gen(m_expo, p_expo):
    '''
    Generates the matrix A to compute the moment dynamics of the form
    dµ/dt = A * µ
    It is basically a collection of the coefficients that go along
    each moment.

    Parameters
    ----------
    m_expo, p_expo: int.
        Exponents of the highest moment to be included in the system.
        m_expo corresponds to the mRNA exponent
        p_expo corresponds to the protein exponent

    Returns
    -------
    Z_matrix : 2D-array.
        Sympy matrix containing all of the coefficients to compute the
        moment after cell division.
    exponents : list.
        List of exponents as they are included in the matrix
    '''
    # Find exponents of moments needed to be computed
    exponents = expo_pairs(m_expo, p_expo)

    # Initalize matrix Z
    Z_matrix = sympy.Matrix(np.zeros([len(exponents),
                                      len(exponents)]))

    # Generate dictionary that saves position of each moment
    # on the matrix
    idx_dict = {e: i for i, e in enumerate(exponents)}

    # Loop through moments
    for i, (mexp, pexp) in enumerate(exponents):
        # Obtain coefficients of equation
        mom_coeff = coeff_dictionary(binom_prod(mexp, pexp))

        # Find row index on matrix Z for substitutions
        row_idx = idx_dict[(mexp, pexp)]

        # Loop through coefficients to make substitutions
        for key, value in mom_coeff.items():
            # Find exponents of the moment
            if key.has(m):
                m_degree = sympy.Poly(key).degree(m)
            else:
                m_degree = 0
            if key.has(p):
                p_degree = sympy.Poly(key).degree(p)
            else:
                p_degree = 0

            # Check if moment included in mom_coeff.
            # Sometimes extra moments with no values are added
            if (m_degree, p_degree) in exponents:
                # Find index for columns
                col_idx = idx_dict[(m_degree, p_degree)]
                # Substitute values in corresponding entries
                Z_matrix[row_idx, col_idx] = value

    return Z_matrix, exponents
```

Let's test the function for a simple case. First let's see which moments are needed in order to compute the first protein moment.

```python
# Define exponents for moment
m_exp, p_exp = 0, 1

# Find which moments are needed
pairs = expo_pairs(m_exp, p_exp)
pairs
```

So we need the moments $\ee{m^0 p^0}$, $\ee{m}$, and $\ee{p}$ in order to compute the first protein moment. Let's then list the resulting products of the binomial products.

```python
prods = [binom_prod(*x) for x in pairs]
prods
```

Given these results our function should return a matrix $\bb{Z}$ of the form
$$
\begin{bmatrix}
1 & 0 & 0\\
0 & 1/2 & 0\\
0 & 0 & 1/2
\end{bmatrix}
$$
Let's see if this is what our function actually returns

```python
# Generate Z matrix
Z_mat, expo = Z_matrix_gen(0, 1)
Z_mat
```

That's exactly the matrix we were expecting!

Let's now generate the much larger matrix needed to compute up to the 6th protein moment.

```python
# Define exponents for moment
m_exp, p_exp = 0, 6

# Generate Z matrix
Z_mat, expo = Z_matrix_gen(m_exp, p_exp)
Z_mat
```

The matrix is what it is, now we just need to save it into a `pickle` object to use it in order to calculate the moments after cell division.

```python
# Export matrix and exponents to file
with open('./pkl_files/binom_coeff_matrix.pkl', 'wb') as file:
    pickle.dump(np.array(Z_mat).astype(np.float64), file)
    pickle.dump(expo, file)
```


# Creating a Simple Model for SEI Growth
Before adding a new model, please read the [contribution guidelines](https://github.com/pybamm-team/PyBaMM/blob/master/CONTRIBUTING.md)

In this notebook, we will run through the steps involved in creating a new model within pybamm. We will then solve and plot the outputs of the model. We have chosen to implement a very simple model of SEI growth. We first give a brief derivation of the model and discuss how to nondimensionalise the model so that we can show the full process of model conception to solution within a single notebook.

Note: if you run the entire notebook and then try to evaluate the earlier cells, you will likely receive an error. This is because the state of objects is mutated as it is passed through various stages of processing. In this case, we recommend that you restart the Kernel and then evaluate cells in turn through the notebook.

## A Simple Model of Solid Electrolyte Interphase (SEI) Growth

The SEI is a porous layer that forms on the surfaces of negative electrode particles from the products of electrochemical reactions which consume lithium and electrolyte solvents. In the first few cycles of use, a lithium-ion battery loses a large amount of capacity; this is generally attributed to lithium being consumed to produce SEI. However, after a few cycles, the rate of capacity loss slows at a rate often (but not always) reported to scale with the square root of time. SEI growth is therefore often considered to be limited in some way by a diffusion process.

### Dimensional Model

We shall first state our model in dimensional form, but to enter the model in pybamm, we strongly recommend converting models into dimensionless form. The main reason for this is that dimensionless models are typically better conditioned than dimensional models and so several digits of accuracy can be gained. To distinguish between the dimensional and dimensionless models, we shall always employ a superscript $*$ on dimensional variables.

In our simple SEI model, we consider a one-dimensional SEI which extends from the surface of a planar negative electrode at $x^*=0$ until $x^*=L^*$, where $L^*$ is the thickness of the SEI. Since the SEI is porous, there is some electrolyte within the region $x^*\in[0, L^*]$ and therefore some concentration of solvent, $c^*$. Within the porous SEI, the solvent is transported via a diffusion process according to:
$$
\frac{\partial c^*}{\partial t^*} = - \nabla^* \cdot N^*, \quad N^* = - D^*(c^*) \nabla^* c^* \label{dim:eqn:solvent-diffusion}\tag{1}\\
$$
where $t^*$ is the time, $N^*$ is the solvent flux, and $D^*(c^*)$ is the effective solvent diffusivity (a function of the solvent concentration).

On the electrode-SEI surface ($x^*=0$) the solvent is consumed by the SEI growth reaction, $R^*$. We assume that diffusion of solvent in the bulk electrolyte ($x^*>L^*$) is fast so that on the SEI-electrolyte surface ($x^*=L^*$) the concentration of solvent is fixed at the value $c^*_{\infty}$. Therefore, the boundary conditions are
$$
    N^*|_{x^*=0} = - R^*, \quad c^*|_{x^*=L^*} = c^*_{\infty},
$$
We also assume that the concentration of solvent within the SEI is initially uniform and equal to the bulk electrolyte solvent concentration, so that the initial condition is
$$
   c^*|_{t^*=0} = c^*_{\infty}
$$

Since the SEI is growing, we require an additional equation for the SEI thickness. The thickness of the SEI grows at a rate proportional to the SEI growth reaction $R^*$, where the constant of proportionality is the partial molar volume of the reaction products, $\hat{V}^*$. We also assume that the SEI is initially of thickness $L^*_0$. Therefore, we have
$$
    \frac{d L^*}{d t^*} = \hat{V}^* R^*, \quad L^*|_{t^*=0} = L^*_0
$$

Finally, we assume for the sake of simplicity that the SEI growth reaction is irreversible and that the potential difference across the SEI is constant. The reaction is also assumed to be proportional to the concentration of solvent at the electrode-SEI surface ($x^*=0$). Therefore, the reaction flux is given by
$$
 R^* = k^* c^*|_{x^*=0}
$$
where $k^*$ is the reaction rate constant (which is in general dependent upon the potential difference across the SEI).

### Non-dimensionalisation

To convert the model into dimensionless form, we scale the dimensional variables and dimensional functions. For this model, we choose to scale $x^*$ by the current SEI thickness, the current SEI thickness by the initial SEI thickness, solvent concentration with the bulk electrolyte solvent concentration, and the solvent diffusion with the solvent diffusion in the electrolyte. We then use these scalings to infer the scaling for the solvent flux. Therefore, we have
$$
x^* = L^* x, \quad L^*= L^*_0 L \quad c^* = c^*_{\infty} c, \quad D^*(c^*) = D^*(c^*_{\infty}) D(c), \quad
N^* = \frac{D^*(c^*_{\infty}) c^*_{\infty}}{L^*_0}N.
$$
We also choose to scale time by the solvent diffusion timescale so that
$$
t^* = \frac{(L^*_0)^2}{D^*(c^*_{\infty})}t.
$$
Finally, we choose to scale the reaction flux in the same way as the solvent flux so that we have
$$
    R^* = \frac{D^*(c^*_{\infty}) c^*_{\infty}}{L^*_0} R.
$$

We note that there are multiple possible choices of scalings. Whilst they will all give the ultimately give the same answer, some choices are better than others depending on the situation under study.

### Dimensionless Model

After substituting in the scalings from the previous section, we obtain the dimensionless form of the model given by:

Solvent diffusion through SEI:
\begin{align}
\frac{\partial c}{\partial t} = \frac{\hat{V} R}{L} x \cdot \nabla c - \frac{1}{L}\nabla \cdot N, \quad N = - \frac{1}{L}D(c) \nabla c, \label{eqn:solvent-diffusion}\tag{1}\\
N|_{x=0} = - R, \quad c|_{x=1} = 1  \label{bc:solvent-diffusion}\tag{2} \quad
c|_{t=0} = 1;
\end{align}

Growth reaction:
$$
R = k c|_{x=0}; \label{eqn:reaction}\tag{3}
$$

SEI thickness:
$$
\frac{d L}{d t} = \hat{V} R, \quad L|_{t=0} = 1; \label{eqn:SEI-thickness}\tag{4}
$$
where the dimensionless parameters are given by
$$
 k = \frac{k^* L^*_0}{D^*(c^*_{\infty})}, \quad \hat{V} = \hat{V}^* c^*_{\infty}, \quad
 D(c) = \frac{D^*(c^*)}{D^*(c^*_{\infty})}.  \label{parameters}\tag{5}
$$
In the above, the additional advective term in the diffusion equation arises due to our choice to scale the spatial coordinate $x^*$ with the time-dependent SEI layer thickness $L^*$.

## Entering the Model into PyBaMM

As always, we begin by importing pybamm and changing our working directory to the root of the pybamm folder.

```python
import pybamm
import numpy as np
import os
os.chdir(pybamm.__path__[0]+'/..')
```

A model is defined in six steps:
1. Initialise model
2. Define parameters and variables
3. State governing equations
4. State boundary conditions
5. State initial conditions
6. State output variables

We shall proceed through each step to enter our simple SEI growth model.

#### 1. Initialise model

We first initialise the model using the `BaseModel` class. This sets up the required structure for our model.

```python
model = pybamm.BaseModel()
```

#### 2. Define parameters and variables

In our SEI model, we have two dimensionless parameters, $k$ and $\hat{V}$, and one dimensionless function $D(c)$, which are all given in terms of the dimensional parameters, see (5). In pybamm, inputs are dimensional, so we first state all the dimensional parameters. We then define the dimensionless parameters, which are expressed an non-dimensional groupings of dimensional parameters. To define the dimensional parameters, we use the `Parameter` object to create parameter symbols. Parameters which are functions are defined using  `FunctionParameter` object and should be defined within a python function as shown.

```python
# dimensional parameters
k_dim = pybamm.Parameter("Reaction rate constant")
L_0_dim = pybamm.Parameter("Initial thickness")
V_hat_dim = pybamm.Parameter("Partial molar volume")
c_inf_dim = pybamm.Parameter("Bulk electrolyte solvent concentration")

def D_dim(cc):
    return pybamm.FunctionParameter("Diffusivity", {"Solvent concentration [mol.m-3]": cc})

# dimensionless parameters
k = k_dim * L_0_dim / D_dim(c_inf_dim)
V_hat = V_hat_dim * c_inf_dim

def D(cc):
    c_dim = c_inf_dim * cc
    return D_dim(c_dim) / D_dim(c_inf_dim)
```

We now define the dimensionless variables in our model. Since these are the variables we solve for directly, we do not need to write them in terms of the dimensional variables. We simply use `SpatialVariable` and `Variable` to create the required symbols:

```python
x = pybamm.SpatialVariable("x", domain="SEI layer", coord_sys="cartesian")
c = pybamm.Variable("Solvent concentration", domain="SEI layer")
L = pybamm.Variable("SEI thickness")
```

#### 3. State governing equations

We can now use the symbols we have created for our parameters and variables to write out our governing equations. Note that before we use the reaction flux and solvent flux, we must derive new symbols for them from the defined parameter and variable symbols. Each governing equation must also be stated in the explicit form `d/dt = rhs` since pybamm only stores the right hand side (rhs) and assumes that the left hand side is the time derivative. The governing equations are then simply

```python
# SEI reaction flux
R = k * pybamm.BoundaryValue(c, "left")

# solvent concentration equation
N = - (1 / L) * D(c) * pybamm.grad(c)
dcdt = (V_hat * R) * pybamm.inner(x / L, pybamm.grad(c)) - (1 / L) * pybamm.div(N)

# SEI thickness equation
dLdt = V_hat * R
```

Once we have stated the equations, we can add them to the `model.rhs` dictionary. This is a dictionary whose keys are the variables being solved for, and whose values correspond right hand sides of the governing equations for each variable.

```python
model.rhs = {c: dcdt, L: dLdt}
```

#### 4. State boundary conditions

We only have boundary conditions on the solvent concentration equation. We must state where a condition is Neumann (on the gradient) or Dirichlet (on the variable itself).

The boundary condition on the electrode-SEI (x=0) boundary is:
$$
    N|_{x=0} = - R, \quad  N|_{x=0} = - \frac{1}{L} D(c|_{x=0} )\nabla c|_{x=0}
$$
which is a Neumann condition. To implement this boundary condition in pybamm, we must first rearrange the equation so that the gradient of the concentration, $\nabla c|_{x=0}$, is the subject. Therefore we have
$$
    \nabla c|_{x=0} = \frac{L R}{D(c|_{x=0} )}
$$
which we enter into pybamm as

```python
# electrode-SEI boundary condition (x=0) (lbc = left boundary condition)
D_left = pybamm.BoundaryValue(D(c), "left") # pybamm requires BoundaryValue(D(c)) and not D(BoundaryValue(c))
grad_c_left = L * R / D_left
```

On the SEI-electrolyte boundary (x=1), we have the boundary condition
$$
    c|_{x=1} = 1
$$

which is a Dirichlet condition and is just entered as

```python
c_right = pybamm.Scalar(1)
```

We now load these boundary conditions into the `model.boundary_conditions` dictionary in the following way, being careful to state the type of boundary condition:

```python
model.boundary_conditions = {c: {"left": (grad_c_left, "Neumann"), "right": (c_right, "Dirichlet")}}
```

#### 5. State initial conditions

There are two initial conditions in our model:
$$
    c|_{t=0} = 1, \quad L|_{t=0} = 1
$$

which are simply written in pybamm as

```python
c_init = pybamm.Scalar(1)
L_init = pybamm.Scalar(1)
```

and then included into the `model.initial_conditions` dictionary:

```python
model.initial_conditions = {c: c_init, L: L_init}
```

#### 6. State output variables

We already have everything required in model for the model to be used and solved, but we have not yet stated what we actually want to output from the model. PyBaMM allows users to output any combination of symbols as an output variable therefore allowing the user the flexibility to output important quantities without further tedious postprocessing steps.

Some useful outputs for this simple model are:
- the SEI thickness
- the SEI growth rate
- the solvent concentration

These are added to the model by adding entries to the `model.variables` dictionary

```python
model.variables = {"SEI thickness": L, "SEI growth rate": dLdt, "Solvent concentration": c}
```

We can also output the dimensional versions of these variables by multiplying by the scalings used to non-dimensionalise. By convention, we recommend including the units in the output variables name so that they do not overwrite the dimensionless output variables. To add new entries to the dictionary we used the method `.update()`.

```python
L_dim = L_0_dim * L
dLdt_dim = (D_dim(c_inf_dim) / L_0_dim ) * dLdt
c_dim = c_inf_dim * c

model.variables.update({
                        "SEI thickness [m]": L_dim,
                        "SEI growth rate [m/s]": dLdt_dim,
                        "Solvent concentration [mols/m^3]": c_dim
                        }
                      )

```

The model is now fully defined and ready to be used. If you plan on reusing the model several times, you can additionally set model defaults which may include: a default geometry to run the model on, a default set of parameter values, a default solver, etc.

## Using the Model

The model will now behave in the same way as any of the inbuilt PyBaMM models. However, to demonstrate that the model works we display the steps involved in solving the model but we will not go into details within this notebook.

```python
# define geometry
geometry = pybamm.Geometry(
    {"SEI layer": {x: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)}}}
)

def Diffusivity(cc):
    return cc * 10**(-5)

# parameter values (not physically based, for example only!)
param = pybamm.ParameterValues(
    {
        "Reaction rate constant": 20,
        "Initial thickness": 1e-6,
        "Partial molar volume": 10,
        "Bulk electrolyte solvent concentration": 1,
        "Diffusivity": Diffusivity,
    }
)

# process model and geometry
param.process_model(model)
param.process_geometry(geometry)

# mesh and discretise
submesh_types = {"SEI layer": pybamm.Uniform1DSubMesh}
var_pts = {x: 100}
mesh = pybamm.Mesh(geometry, submesh_types, var_pts)

spatial_methods = {"SEI layer": pybamm.FiniteVolume()}
disc = pybamm.Discretisation(mesh, spatial_methods)
disc.process_model(model)

# solve
solver = pybamm.ScipySolver()
t = np.linspace(0, 100, 100)
solution = solver.solve(model, t)

# Extract output variables
L_out = solution["SEI thickness"]
c_out = solution["Solvent concentration"]
x = np.linspace(0, 1, 100)
```

    2020-05-30 11:30:30,931 - [WARNING] processed_variable.get_spatial_scale(497): No scale set for spatial variable x. Using default of 1 [m].

Using these outputs, we can now plot the SEI thickness as a function of time and also the solvent concentration profile within the SEI. We use a slider to plot the concentration profile at different times.

```python
import matplotlib.pyplot as plt
def plot(t):
    f, (ax1, ax2) = plt.subplots(1, 2 ,figsize=(10,5))
    ax1.plot(solution.t, L_out(solution.t))
    ax1.plot([t], [L_out(t)], 'r.')
    plot_c, = ax2.plot(x * L_out(t), c_out(t, x))
    ax1.set_ylabel('SEI thickness')
    ax1.set_xlabel('t')
    ax2.set_ylabel('Solvent concentration')
    ax2.set_xlabel('x')
    ax2.set_ylim(0, 1.1)
    ax2.set_xlim(0, x[-1]*L_out(solution.t[-1]))
    plt.show()

import ipywidgets as widgets
widgets.interact(plot, t=widgets.FloatSlider(min=0,max=solution.t[-1],step=0.1,value=0));
```

    interactive(children=(FloatSlider(value=0.0, description='t'), Output()), _dom_classes=('widget-interact',))

## Formally adding your model

The purpose of this notebook has been to go through the steps involved in getting a simple model working within PyBaMM. However, if you plan on reusing your model and want greater flexibility then we recommend that you create a new class for your model. We have set out instructions on how to do this in the "Adding a Model" tutorial in the documentation.

```python

```


# Aim

To provide a clear understanding of some of the key measures or reliability and their implementation in `python`.

# How

To do that I will follow references as close as possible while providing the python code of the implementation of the ideas and examples.

Revelle has greatly created and write many papers in the matter and the library `psych`in `R` with the corresponding implementations. We would follow Revelle 2021 and 2017, in particlar Revelle 2021 very closely.

We would expand a bit when calculating the omegas and in particular when using the `Schmid-Leiman` solution, following Wolff and Presing 2005, as it constains a clear explanations of the theory and the estimation, as in this paper authors describe the implementation in `SPSS` and `SAS`, which something similar we are trying to do here in `python`.

## References
* Revelle, Willian. Manuscrip. 2021. An introduction to psychometric theory with applications in R. https://personality-project.org/r/book/Chapter7.pdf
* Revelle, William R. "psych: Procedures for personality and psychological research." (2017).
* Wolff, Hans-Georg, and Katja Preising. "Exploring item and higher order factor structure with the Schmid-Leiman solution: Syntax codes for SPSS and SAS." Behavior Research Methods 37.1 (2005): 48-58.

Examples in R.

This work:
* Rafael Valero Fernández. (2022). reliabiliPy: measures of survey domain reliability in Python with explanations and examples. Cronbach´s Alpha and Omegas. (v0.0.0). Zenodo.[](https://zenodo.org/badge/latestdoi/445846537)

```python
name_notebook = 'revelle_ch7_example' # this is to save
```

```python
%load_ext autoreload
%autoreload 2
import pandas as pd
import numpy as np
import sys
sys.path.append("../")
from reliabilipy import reliability_analysis
```

Dataset as examample in Table 7.5  https://personality-project.org/r/book/Chapter7.pdf

```python
V = pd.DataFrame(np.matrix([[1.0, 0.5, 0.5, 0.5, 0.5, 0.5, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2],
[0.5, 1.0, 0.5, 0.5, 0.5, 0.5, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2],
[0.5, 0.5, 1.0, 0.5, 0.5, 0.5, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2],
[0.5, 0.5, 0.5, 1.0, 0.5, 0.5, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2],
[0.5, 0.5, 0.5, 0.5, 1.0, 0.5, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2],
[0.5, 0.5, 0.5, 0.5, 0.5, 1.0, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2],
[0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 1.0, 0.6, 0.6, 0.6, 0.2, 0.2],
[0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.6, 1.0, 0.6, 0.6, 0.2, 0.2],
[0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.6, 0.6, 1.0, 0.6, 0.2, 0.2],
[0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.6, 0.6, 0.6, 1.0, 0.2, 0.2],
[0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 1.0, 0.7],
[0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.7, 1.0]]))
```

Lets try to stimate different reliability measures

```python
Vx=V.sum().sum()
print(Vx)
```

    53.20000000000001

$$
\lambda_{1}=1-\frac{\operatorname{tr}\left(\mathbf{V}_{\mathbf{x}}\right)}{V_{x}}=\frac{V_{x}-\operatorname{tr}\left(\mathbf{V}_{x}\right)}{V_{x}} .
$$
In the example:
$$
\lambda_{1}=1-\frac{12}{53.2}=\frac{41.2}{53.2}=.774
$$

```python
lambda1 = 1 - np.diag(V).sum() / Vx
print(lambda1)
```

    0.7744360902255639

Let $C_{2}=\mathbf{1}(\mathbf{V}-\operatorname{diag}(\mathbf{V}))^{2} \mathbf{1}^{\prime}$, then
$$
\lambda_{2}=\lambda_{1}+\frac{\sqrt{\frac{n}{n-1} C_{2}}}{V_{x}}=\frac{V_{x}-\operatorname{tr}\left(\mathbf{V}_{x}\right)+\sqrt{\frac{n}{n-1} C_{2}}}{V_{x}}
$$

```python
n=V.shape[0]
C2=((V - np.eye(n)*np.diag(V))**2).sum().sum()
print(C2)
```

    16.320000000000004

```python
lambda2=lambda1+(n/(n-1)*C2)**0.5/Vx
print(lambda2)
```

    0.8537487885405894

$$\lambda_{3}=\lambda_{1}+\frac{\frac{V_{X}-\operatorname{tr}\left(\mathbf{V}_{X}\right)}{n(n-1)}}{V_{X}}=\frac{n \lambda_{1}}{n-1}=\frac{n}{n-1}\left(1-\frac{\operatorname{tr}(\mathbf{V})_{x}}{V_{x}}\right)=\frac{n}{n-1} \frac{V_{x}-\operatorname{tr}\left(\mathbf{V}_{x}\right)}{V_{x}}=\alpha$$

```python
alpha = n/(n-1)*(Vx-np.diag(V).sum())/Vx
print(alpha)
```

    0.8448393711551606

```python
reliability_report = reliability_analysis(correlations_matrix=V)
reliability_report.fit()
print(reliability_report.lambda1)
print(reliability_report.lambda2)
print(reliability_report.alpha_cronbach)
```

    0.7744360902255639
    0.8537487885405894
    0.8448393711551606

# Omegas reliability

Omegas are based on a decomposition of the variance of a test score, $V_x$ into four parts: that due to a general factor, $\vec{g}$, that due to a set of group factors, $\vec{f}$, (factors common to some but not all of the items), specific factors, $\vec{s}$ unique to each item, and $\vec{e}$, random error. (Because specific variance can not be distinguished from random error unless the test is given at least twice, some combine these both into error).

$$
\begin{equation} \label{eq:1}
x = cg + Af + Ds + e
\end{equation}
$$

then the communality of item $j$, based upon general as well as group factors, $h_j^2 = c_j^2 + \sum(f_{ij}^2)$ and the unique variance for the item $u_j^2 = σ_j^2 (1-h_j^2)$ may be used to estimate the test reliability. That is, if $h_j^2$ is the communality of item $j$, based upon general as well as group factors, then for standardized items, $e_j^2 = 1 - h_j^2$ and

$$ω_t =  \frac{1 cc' 1 + 1 AA' 1'}{V_x}$$

$ω_t$ is based upon the sum of squared loadings on all the factors.
$ω_h$ is based upon the sum of the squared loadings on the general factor.

$$ω_h =  \frac{1 cc' 1'}{Vx}$$

Another estimate reported is the omega for an infinite length test with a structure similar to the observed test (omega hierarchical asymptotic). This is found by

$$ω_{limit} = \frac{1 cc' 1'}{1 cc' 1' + 1 AA' 1'}$$

The input to omega may be a correlation matrix or a raw data matrix, or a factor pattern matrix with the factor intercorrelations (Phi) matrix.

## Results in R using Psych
There is an R notebook with the full example, in the same folder that this notebook, with this result:
```bash
Omega
Call: omega(m = correlation_matrix)
Alpha:                 0.84
G.6:                   0.88
Omega Hierarchical:    0.54
Omega H asymptotic:    0.6
Omega Total            0.9
````

## Results using 'reliabiliPy'

```python
reliability_report = reliability_analysis(correlations_matrix=V)
reliability_report.fit()
```

```python
print(reliability_report.omega_total)
print(reliability_report.omega_hierarchical)
print(reliability_report.omega_hierarchical_asymptotic)
```

    0.9022556800089196
    0.5413535677830178
    0.6000001770869052

## Further R output
```bash
    g        F1   F2   F3     h2      u2      p2
v0	0.45	0.55			0.5	0.5	0.40
v1	0.45	0.55			0.5	0.5	0.40
v2	0.45	0.55			0.5	0.5	0.40
v3	0.45	0.55			0.5	0.5	0.40
v4	0.45	0.55			0.5	0.5	0.40
v5	0.45	0.55			0.5	0.5	0.40
v6	0.45		 0.63	   0.6	0.4	0.33
v7	0.45		 0.63	   0.6	0.4	0.33
v8	0.45		 0.63	   0.6	0.4	0.33
v9	0.45		 0.63	   0.6	0.4	0.33
```

With eigenvalues of:
```bash
  g F1* F2* F3*
2.4 1.8 1.6 1.0
```

```python
reliability_report.report_loadings['g'][0]
```

    0.45

```python
reliability_report.report_eigenvalues
```

      g
      F0
      F1
      F2

      eigenvalues
      2.400001
      1.800007
      1.599981
      1.000007

# Estimation of Omegas reliabilities

This section is in case you want to go deeper. This will allow you to understand the calculation and to be able to modify the package.

But how can we estimate that?. Using `Schmid-Leiman` solution. For an explanation on that   Wolff and Preising (2005). There the authors try to do the same that in this work but for SPSS and SAS.

"matrix $\mathbf{F}_{1}$ (i.e., rotated loadings), a first-order factor correlation matrix termed $\mathbf{R}_{1}$, and unique variances $\mathbf{U}_{1}^{2}$ of variables:
$$
\mathbf{R}_{v}=\mathbf{F}_{1} \mathbf{R}_{1} \mathbf{F}_{1}^{\prime}+\mathbf{U}_{1}^{2}
$$

Matrix $\mathbf{F}_{1}$ is of the order $v \times f_{1}$, where $f_{1}$ is the number of first-order factors. $\mathbf{R}_{1}$ is of the order $f_{1} \times f_{1} . \mathbf{U}_{1}^{2}$ is a matrix of the order $v \times v$ and consists of measurement error and specific components of variables.

Second-order factors can be extracted from the factor correlation matrix $\mathbf{R}_{1}$, yielding a matrix of second-order factor loadings $\mathbf{F}_{2}$, a second-order factor correlation matrix $\mathbf{R}_{2}$, and the unique variance of this level, $\mathbf{U}_{2}^{2}$ :

$$
\mathbf{R}_{1}=\mathbf{F}_{2} \mathbf{R}_{2} \mathbf{F}_{2}^{\prime}+\mathbf{U}_{2}^{2}
$$

$\mathbf{F}_{2}$ is of the order $f_{1} \times f_{2}$, where $f_{2}$ is the number of second-order factors. $\mathbf{R}_{2}$ is of the order $f_{2} \times f_{2} . \mathbf{U}_{2}^{2}$ is a matrix of the order $f_{1} \times f_{1}$.

These two analyses comprise a minimal higher order FA consisting of two levels. $\mathbf{F}_{1}$ from Equation 1 provides first-order factor loadings (e.g., paths from $\mathrm{F}_{1}$ to variables in Figure 1A), and $\mathbf{F}_{2}$ yields second-order factor loadings (e.g., the path from Factor 1 to $\mathrm{F}_{1}$ in Figure $1 \mathrm{~A}$ ). If second-order factors are uncorrelated, matrix $\mathbf{R}_{2}$ is an identity matrix and can be ignored. If second-order factors are correlated, analyses may continue to extract third-order factors.

To obtain the SLS, two further calculations are necessary. In a first step, direct effects of higher order factors on variables are determined by multiplying all factor-loading matrices from first to highest order. In the present case, the direct loadings of second-order factors on variables, termed $\mathbf{F}_{2}^{S L S}$, are obtained by multiplying $\mathbf{F}_{1} \times \mathbf{F}_{2}$
$$
\mathbf{F}_{2}^{\mathrm{SLS}}=\mathbf{F}_{1} * \mathbf{F}_{2},
$$
where $\mathbf{F}_{2}^{\text {SLS }}$ is of the order $v \times f_{2}$. To obtain residualized first-order factor loadings $\mathbf{F}_{1}^{\mathrm{SLS}}$, factor loadings $\mathbf{F}_{1}$ are multiplied by the uniqueness $\mathbf{U}_{2}$ of higher order factors:
$$
\mathbf{F}_{1}^{\mathrm{SLS}}=\mathbf{F}_{1} * \mathbf{U}_{2},
$$
where $\mathbf{F}_{1}^{\mathrm{SLS}}$ is of the order $v \times f_{1}$. The uniqueness $\mathbf{U}_{2}$ is equivalent to the square root of the variances of $\mathbf{R}_{1}$ not explained by higher order factors. It is obtained by subtracting the explained variance of a factor from 1:
$$
\mathbf{U}_{2}=\left[\mathbf{I}-\operatorname{diag}\left(\mathbf{F}_{2} \mathbf{R}_{2} \mathbf{F}_{2}^{\prime}\right)\right]^{0.5},
$$
where $\mathbf{I}$ is an identity matrix and diag indicates that only"

# Tranlations in our case and to Python
$$
\mathbf{F}_{2}^{\mathrm{SLS}}=\mathbf{F}_{1} * \mathbf{F}_{2},
$$
implementation in `python`:

```python
general_component = np.dot(reliability_report.fa_f.loadings_,
                           reliability_report.fa_g.loadings_)
general_component=np.abs(general_component)
general_component
```

    array([[0.44721233],
           [0.44721233],
           [0.44721233],
           [0.44721233],
           [0.44721233],
           [0.44721233],
           [0.4472171 ],
           [0.4472171 ],
           [0.4472171 ],
           [0.4472171 ],
           [0.44721084],
           [0.44721084]])

Eigenvalue g

```python
np.dot(general_component.T,general_component)
```

    array([[2.40000082]])

# $u^2$ in our example:

```python
reliability_report.fa_f.get_uniquenesses()
```

    array([0.50000014, 0.50000014, 0.50000014, 0.50000014, 0.50000014,
           0.50000014, 0.39999742, 0.39999742, 0.39999742, 0.39999742,
           0.30000366, 0.30000366])

# $h^2$ in our example:

```python
reliability_report.fa_f.get_communalities()
```

    array([0.49999986, 0.49999986, 0.49999986, 0.49999986, 0.49999986,
           0.49999986, 0.60000258, 0.60000258, 0.60000258, 0.60000258,
           0.69999634, 0.69999634])

To update the group factors

$$
\mathbf{F}_{1}^{\mathrm{SLS}}=\mathbf{F}_{1} * \mathbf{U}_{2},
$$

```python
f_loadings_final = np.zeros(reliability_report.fa_f.loadings_.shape)
f_loadings_final

```

    array([[0., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.]])

```python
reliability_report.fa_g.loadings_
```

    array([[-0.63245363],
           [-0.5773573 ],
           [-0.53451571]])

```python
reliability_report.fa_f.get_uniquenesses()
```

    array([0.50000014, 0.50000014, 0.50000014, 0.50000014, 0.50000014,
           0.50000014, 0.39999742, 0.39999742, 0.39999742, 0.39999742,
           0.30000366, 0.30000366])

```python
for i in range(0,reliability_report.fa_g.get_uniquenesses().__len__()):
    f_loadings_final[:,i] = reliability_report.fa_f.loadings_[:,i]*reliability_report.fa_g.get_uniquenesses()[i]**0.5

f_loadings_final = np.abs(f_loadings_final)
```

The eigenvectors in our example:

```python
np.abs(f_loadings_final).round(3)
```

    array([[0.548, 0.   , 0.   ],
           [0.548, 0.   , 0.   ],
           [0.548, 0.   , 0.   ],
           [0.548, 0.   , 0.   ],
           [0.548, 0.   , 0.   ],
           [0.548, 0.   , 0.   ],
           [0.   , 0.632, 0.   ],
           [0.   , 0.632, 0.   ],
           [0.   , 0.632, 0.   ],
           [0.   , 0.632, 0.   ],
           [0.   , 0.   , 0.707],
           [0.   , 0.   , 0.707]])

The eigenvalues:

```python
np.dot(f_loadings_final.T,f_loadings_final).sum(axis=1)
```

    array([1.80001644, 1.60000108, 1.00001845])

```python
"""save code """
```

    'save code '

```python
import os
order = f'jupytext --to py {name_notebook}.ipynb'
print(order)
os.system(order)
```

    jupytext --to py revelle_ch7_example.ipynb
    [jupytext] Reading revelle_ch7_example.ipynb in format ipynb
    [jupytext] Updating the timestamp of revelle_ch7_example.py

    0


$\newcommand{\Normal}{\mathcal{N}}
\newcommand{\lp}{\left(}
\newcommand{\rp}{\right)}
\newcommand{\lf}{\left\{}
\newcommand{\rf}{\right\}}
\newcommand{\ls}{\left[}
\newcommand{\rs}{\right]}
\newcommand{\lv}{\left|}
\newcommand{\rv}{\right|}
\newcommand{\state}{x}
\newcommand{\State}{\boldx}
\newcommand{\StateR}{\boldX}
\newcommand{\Covariance}{\Sigma}
\newcommand{\CovX}{\Covariance_{\boldX}}
\newcommand{\CovY}{\Covariance_{\boldY}}
\newcommand{\CovZ}{\Covariance_{\boldZ}}
\newcommand{\CovXY}{\Covariance_{\boldX\boldY}}
\newcommand{\hatState}{\hat{\State}}
\newcommand{\StateNum}{N}
\newcommand{\StateDim}{K}
\newcommand{\NumStates}{N}
\newcommand{\StateToState}{A}
\newcommand{\stateToState}{A}
\newcommand{\StateCov}{\Sigma}
\newcommand{\StateJac}{A}
\newcommand{\hatStateCov}{\hat{\StateCov}}
\newcommand{\StateMean}{\boldmu}
\newcommand{\hatStateMean}{\hat{\StateMean}}
\newcommand{\StateToStateHistory}{\boldA}
\newcommand{\StateNoise}{\boldr}
\newcommand{\StateNoiseCov}{R}
\newcommand{\StateHistory}{\boldX}
\newcommand{\StatesHistory}{\StateHistory}
\newcommand{\StateToObserv}{C}
\newcommand{\StateToobserv}{\boldc}
\newcommand{\StateToObservHistory}{\boldC}
\newcommand{\DState}{\bolddelta}
\newcommand{\hatDState}{\hat{\DState}}
\newcommand{\DStateMean}{\boldlambda}
\newcommand{\hatDStateMean}{\hat{\DStateMean}}
\newcommand{\DStateCov}{\Lambda}
\newcommand{\hatDStateCov}{\hat{\DStateCov}}
\newcommand{\DObserv}{\boldgamma}
\newcommand{\hatDObserv}{\hat{\DObserv}}
\newcommand{\observ}{z}
\newcommand{\Observ}{\boldsymbol{\observ}}
\newcommand{\ObservCov}{\Lambda}
\newcommand{\observMean}{\lambda}
\newcommand{\ObservMean}{\boldlambda}
\newcommand{\hatobserv}{\hat{\observ}}
\newcommand{\hatObserv}{\hat{\Observ}}
\newcommand{\hatObservCov}{\hat{\ObservCov}}
\newcommand{\hatobservMean}{\hat{\observMean}}
\newcommand{\hatObservMean}{\hat{\ObservMean}}
\newcommand{\ObservSet}{\ZZ}
\newcommand{\ObservNum}{N}
\newcommand{\ObservDim}{D}
\newcommand{\ObservSourceNum}{M}
\newcommand{\ObservHistory}{\boldZ}
\newcommand{\ObservsHistory}{\ObservHistory}
\newcommand{\Timestamps}{\boldT}
\newcommand{\ObservJac}{H}
\newcommand{\observNoise}{q}
\newcommand{\ObservNoise}{\boldq}
\newcommand{\ObservNoiseCov}{Q}
\newcommand{\ObservNoiseCovHistory}{\boldQ}
\newcommand{\Jacobian}{\boldJ}
\newcommand{\Kalman}{K}
\newcommand{\kalman}{\boldk}
\newcommand{\CC}{\mathbb{C}}
\newcommand{\NN}{\mathbb{N}}
\newcommand{\RR}{\mathbb{R}}
\newcommand{\XX}{\mathbb{X}}
\newcommand{\ZZ}{\mathbb{Z}}
\renewcommand{\AA}{\mathbb{A}}
\newcommand{\boldzero}{\boldsymbol{0}}
\newcommand{\boldone}{\boldsymbol{1}}
\newcommand{\bolda}{\boldsymbol{a}}
\newcommand{\boldb}{\boldsymbol{b}}
\newcommand{\boldc}{\boldsymbol{c}}
\newcommand{\boldd}{\boldsymbol{d}}
\newcommand{\bolde}{\boldsymbol{e}}
\newcommand{\boldf}{\boldsymbol{f}}
\newcommand{\boldg}{\boldsymbol{g}}
\newcommand{\boldh}{\boldsymbol{h}}
\newcommand{\boldi}{\boldsymbol{i}}
\newcommand{\boldj}{\boldsymbol{j}}
\newcommand{\boldk}{\boldsymbol{k}}
\newcommand{\boldl}{\boldsymbol{l}}
\newcommand{\boldm}{\boldsymbol{m}}
\newcommand{\boldn}{\boldsymbol{n}}
\newcommand{\boldo}{\boldsymbol{o}}
\newcommand{\boldp}{\boldsymbol{p}}
\newcommand{\boldq}{\boldsymbol{q}}
\newcommand{\boldr}{\boldsymbol{r}}
\newcommand{\bolds}{\boldsymbol{s}}
\newcommand{\boldt}{\boldsymbol{t}}
\newcommand{\boldu}{\boldsymbol{u}}
\newcommand{\boldv}{\boldsymbol{v}}
\newcommand{\boldw}{\boldsymbol{w}}
\newcommand{\boldx}{\boldsymbol{x}}
\newcommand{\boldy}{\boldsymbol{y}}
\newcommand{\boldz}{\boldsymbol{z}}
\newcommand{\boldA}{\boldsymbol{A}}
\newcommand{\boldB}{\boldsymbol{B}}
\newcommand{\boldC}{\boldsymbol{C}}
\newcommand{\boldD}{\boldsymbol{D}}
\newcommand{\boldE}{\boldsymbol{E}}
\newcommand{\boldF}{\boldsymbol{F}}
\newcommand{\boldH}{\boldsymbol{H}}
\newcommand{\boldJ}{\boldsymbol{J}}
\newcommand{\boldK}{\boldsymbol{K}}
\newcommand{\boldL}{\boldsymbol{L}}
\newcommand{\boldM}{\boldsymbol{M}}
\newcommand{\boldI}{\boldsymbol{I}}
\newcommand{\boldP}{\boldsymbol{P}}
\newcommand{\boldQ}{\boldsymbol{Q}}
\newcommand{\boldR}{\boldsymbol{R}}
\newcommand{\boldS}{\boldsymbol{S}}
\newcommand{\boldT}{\boldsymbol{T}}
\newcommand{\boldO}{\boldsymbol{O}}
\newcommand{\boldU}{\boldsymbol{U}}
\newcommand{\boldV}{\boldsymbol{V}}
\newcommand{\boldW}{\boldsymbol{W}}
\newcommand{\boldX}{\boldsymbol{X}}
\newcommand{\boldY}{\boldsymbol{Y}}
\newcommand{\boldZ}{\boldsymbol{Z}}
\newcommand{\boldXY}{\boldsymbol{XY}}
\newcommand{\boldmu}{\boldsymbol{\mu}}$

### Состояние системы
Теперь рассмотрим чуть более близкую к реальности модель движения материальной точки. В данной модели состояние движущейся описывается следующим вектором:
$$
\State = \begin{pmatrix}
x\\
y\\
\gamma\\
v\\
\omega\\
\end{pmatrix},
$$
где
* $x$ — координата точки по оси $X$;
* $y$ — координата точки по оси $Y$;
* $\gamma$ — угол рыскания; далее просто yaw;
* $v$ — проекция скорости на ось $X$;
* $\omega$ — угловая скорость материальной точки.

### Модель эволюции системы
#### Переход из момента времени $t$ к моменту $t + \Delta t$[toc](#toc)
$$
\State(t + \Delta t) = \boldf(\State(t)) \Rightarrow
\begin{pmatrix}
x(t + \Delta t)\\
y(t + \Delta t)\\
\gamma(t + \Delta t)\\
v(t + \Delta t)\\
\omega(t + \Delta t)
\end{pmatrix}
\approx
\begin{pmatrix}
x(t) + v(t)\cos(\gamma(t))\Delta t\\
y(t) + v(t)\sin(\gamma(t))\Delta t\\
\gamma(t) + \omega(t) \Delta t\\
v(t)\\
\omega(t)\\
\end{pmatrix}.
$$
Заметим, что зависимость между состояниям нелинейная. Т.е. уже нельзя представить переход из момента $t$ в момент $t + \Delta t$ в виде:
$$
\State(t + \Delta t) \approx A(\Delta t) \State(t)
$$
Поэтому в данном случае мы будем использовать **EKF-фильтр**.

#### Якобиан функции перехода $f(x(t))$[toc](#toc)
Итак, мы знаем, что $\State(t) \sim \Normal(\mu(t), \Sigma(t))$. Также мы знаем, что состояние в момент $\State(t + \Delta t)$ выражается через состояние в момент $\State(t)$ через некоторую нелинейную функцию $f(\cdot)$:
$$
\State(t + \Delta t) \approx f(\State(t), \Delta t) = \begin{pmatrix}
x(t) + v(t)\cos(\gamma(t))\Delta t\\
y(t) + v(t)\sin(\gamma(t))\Delta t\\
\gamma(t) + \omega(t) \Delta t\\
v(t)\\
\omega(t)\\
\end{pmatrix}
$$
Поэтому
$$
\State(t + \Delta t) \approx f(\mu(t)) + J_f(\mu(t)) (\State(t) - \mu(t)) + \text{H.O.D.},
$$
где
* $f(\mu(t), \Delta t)$ — состояние в момент времени $t + \Delta t$, при условии, что в момент $t$ мы детерминированно находились в состоянии $\State(t) = \mu(t)$.
* $J_f(\mu(t))$ — значение Якобиана $f(\State)$ по $\State$ в точке $\mu(t)$.
* $\State(t) - \mu(t)$ — отклонение истинного значения состояния $\State(t)$ от ожидаемого $\mu(t)$; имеет распределение $\Normal(\boldzero, \Sigma(t))$, чем и обуславливает нормальное распределение для $\State(t + \Delta t)$.
* H.O.D. — компоненты разложения высшего порядка, которыми пренебрегаем (Higher Order Degress).

Якобиан имеет вид:
$$
J_f(\State(t)) =
\begin{pmatrix}
1 &0 &-v(t)\sin(\gamma(t))\Delta t &\cos(\gamma(t))\Delta t &0\\
0 &1 &v(t)\cos(\gamma(t))\Delta t &\sin(\gamma(t))\Delta t &0\\
0 &0 &1 &0 &\Delta t\\
0 &0 &0 &1 &0\\
0 &0 &0 &0 &1\\
\end{pmatrix}
$$

#### Модель эволюции системы[toc](#toc)
В результате получаем
\begin{align}
&\mu(t + \Delta t) = f(\mu(t), \Delta t),\\
&\Sigma(t + \Delta t) = J_f(\mu(t)) \Sigma(t)J_f(\mu(t))^T + R(t) \Delta t,
\end{align}
где
* $\mu(t)$ — текущее среднее состояния;
* $\Sigma(t)$ — текущая дисперсия состояния;
* $\Delta t$ — шаг по времени;
* $R(t)$ — матрица плотности шума (показывает, как быстро дисперсия состояния наростает со временем).

### Модель наблюдений системы
По сути достаточно лишь указать матрицы наблюдений. Так как в них нет нелинейностей, то тут можно применять обычную обработку показаний (без появления Якобианов и т.п.).
#### GPS
$$
C_{gps} =
\begin{pmatrix}
1 &0 &0 &0 &0\\
0 &1 &0 &0 &0\\
\end{pmatrix}
$$

#### CAN
$$
C_{can} =
\begin{pmatrix}
0 &0 &0 &1 &0\\
\end{pmatrix}
$$

#### IMU
$$
C_{imu} =
\begin{pmatrix}
0 &0 &0 &0 &1\\
\end{pmatrix}
$$

```python
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.patches import Ellipse, Rectangle
from IPython import display
import time
from matplotlib.lines import Line2D

%matplotlib inline

from sdc.timestamp import Timestamp
from sdc.car import Car
from sdc.linear_movement_model import LinearMovementModel
from sdc.cycloid_movement_model import CycloidMovementModel
from sdc.can_sensor import CanSensor
from sdc.gps_sensor import GpsSensor
from sdc.imu_sensor import ImuSensor
```

### Создание автомобиля, настрока движения и установка необходимых датчиков[toc](#toc)
Функция для удобного создания полностью настроенной модели машины:

```python
def create_car(
        initial_position=[5, 5],
        initial_velocity=5,
        initial_omega=0.0,
        initial_yaw=np.pi / 4,
        can_noise_variances=[0.25],   # Стандартное отклонение - 0.5м
        gps_noise_variances=[1, 1],   # Стандартное отклонение - 1м
        imu_noise_variances=None,     # По умолчанию IMU не используется
        random_state=0,
):
    # Начальное состояние автомобиля
    car = Car(
        initial_position=initial_position,
        initial_velocity=initial_velocity,
        initial_yaw=initial_yaw,
        initial_omega=initial_omega)
    # Создание сенсоров
    if can_noise_variances is not None:
        car.add_sensor(CanSensor(noise_variances=can_noise_variances, random_state=random_state))
        random_state += 1
    if gps_noise_variances is not None:
        car.add_sensor(GpsSensor(noise_variances=gps_noise_variances, random_state=random_state))
        random_state += 1
    if imu_noise_variances is not None:
        car.add_sensor(ImuSensor(noise_variances=imu_noise_variances, random_state=random_state))
        random_state += 1
    # Последовательное движение
    movement_model = LinearMovementModel()
    car.set_movement_model(movement_model)
    return car
```

```python
# Калмановская локализация
from sdc.car_plotter import CarPlotter
from sdc.kalman_car import KalmanCar
from sdc.kalman_can_sensor import KalmanCanSensor
from sdc.kalman_gps_sensor import KalmanGpsSensor
from sdc.kalman_imu_sensor import KalmanImuSensor
from sdc.kalman_movement_model import KalmanMovementModel
from sdc.kalman_filter import (
    kalman_transit_covariance,
    kalman_process_observation,
)
```

### Создание калмановской модели машины из обычной модели[toc](#toc)

```python
def create_kalman_car(car, gps_variances=None, can_variances=None, imu_variances=None):
    """Создает калмановскую модель движения автомобиля на основе уже настроенной модели самого автомобиля"""
    # Скорость нарастания дисперсии в секунду
    noise_covariance_density = np.diag([
        0.1,
        0.1,
        0.1,   # Дисперсия yaw
        0.1,   # Дисперсия скорости
        0.1    # Дисперсия угловой скорости
    ])
    # Формирование состояние калмановской локализации
    kalman_car = KalmanCar(
        initial_position=car.initial_position,
        initial_velocity=car.initial_velocity,
        initial_yaw=car.initial_yaw,
        initial_omega=car.initial_omega)
    # Начальная матрица ковариации
    kalman_car.covariance_matrix = noise_covariance_density

    # Модель движения
    kalman_movement_model = KalmanMovementModel(noise_covariance_density=noise_covariance_density)
    kalman_car.set_movement_model(kalman_movement_model)

    for sensor in car.sensors:
        noise_variances = sensor._noise_variances
        if isinstance(sensor, GpsSensor):
            noise_variances = noise_variances if gps_variances is None else gps_variances
            kalman_sensor = KalmanGpsSensor(noise_variances=noise_variances)
        elif isinstance(sensor, CanSensor):
            noise_variances = noise_variances if can_variances is None else can_variances
            kalman_sensor = KalmanCanSensor(noise_variances=noise_variances)
        elif isinstance(sensor, ImuSensor):
            noise_variances = noise_variances if imu_variances is None else imu_variances
            kalman_sensor = KalmanImuSensor(noise_variances=noise_variances)
        else:
            assert False
        kalman_car.add_sensor(kalman_sensor)
    return kalman_car
```

### Визуализация траектории автомобиля[toc](#toc)

```python
initial_position = [20, 20]
initial_velocity = 10
initial_yaw = 0.5
initial_omega = 0.02

car = Car(initial_position=initial_position, initial_velocity=initial_velocity,
          initial_yaw=initial_yaw, initial_omega=initial_omega)
print(car)

# Тестирование последовательного движения
movement_model = LinearMovementModel()
car.set_movement_model(movement_model)

dt = Timestamp.seconds(0.1)
duration = Timestamp.seconds(40)
final_time = car.time + duration
while car.time < final_time:
    car.move(dt)
print(car)

# Отрисовка траектории
fig = plt.figure(figsize=(15, 15))
ax = plt.subplot(111, aspect='equal')
ax.grid(which='both', linestyle='--', alpha=0.5)

car_plotter = CarPlotter(car_width=3, car_height=1.5)
car_plotter.plot_car(ax, car)
car_plotter.plot_trajectory(ax, car, traj_color='k')

# Установка корректных пределов
x_limits, y_limits = car_plotter.get_limits(car)
ax.set_xlim(x_limits)
ax.set_ylim(y_limits);
```

## Тестирование калмановской фильтрации[toc](#toc)

```python
car = create_car(initial_omega=0.05)
kalman_car = create_kalman_car(car)

# Создаем полотно
fig = plt.figure(figsize=(15, 15))
ax = plt.subplot(111, aspect='equal')
# ax.grid(which='both', linestyle='--', alpha=0.5)
real_color = 'green'
obs_color = 'blue'
est_color = 'red'

legend_lines = {'Kalman estimate': Line2D([0], [0], color=est_color, lw=4),
                'Measurement': Line2D([0], [0], color=obs_color, lw=4),
                'Real position': Line2D([0], [0], color=real_color, lw=4)}

# Шаг интегрирования
dt = Timestamp.seconds(0.1)
# Длительность проезда
duration = Timestamp.seconds(40)
final_time = car.time + duration

# Отрисовщик автомобиля и траектории
car_plotter = CarPlotter(
    car_width=3, car_height=1.5,
    real_color=real_color, obs_color=obs_color, pred_color=est_color)
car_plotter.plot_car(ax, car)

# Отрисовка начальной позиции
car_plotter.plot_car(ax, car)
display.clear_output(wait=True)
time.sleep(0.25)
ax.set_xlim(-10, 100)
ax.set_ylim(-10, 100)

while car.time < final_time:
    ax.clear()
    ax.legend([legend_lines['Kalman estimate'],
               legend_lines['Measurement'],
               legend_lines['Real position']], ['Kalman estimate', 'Measurement', 'Real position'])
    # Делаем реальный переход к моменту времени t + dt
    car.move(dt)
    # Делаем предсказание на момент времени t + dt
    kalman_car.move(dt)
    # Теперь обработаем наблюдения в момент t + dt
    for sensor, kalman_sensor in zip(car.sensors, kalman_car.sensors):
        observation = sensor.observe()
        kalman_sensor.process_observation(observation)

    car_plotter.plot_car(ax, car)
    car_plotter.plot_kalman_car(ax, kalman_car)
    car_plotter.plot_trajectory(ax, car, traj_color=real_color)
    car_plotter.plot_trajectory(ax, kalman_car, traj_color=est_color)
    car_plotter.plot_observations(ax, car.gps_sensor.history[:, 0], car.gps_sensor.history[:, 1], color=obs_color)
    ax.set_xlim(-5, 100)
    ax.set_ylim(-5, 100)
    display.clear_output(wait=True)
    display.display(plt.gcf())
    time.sleep(0.05)

display.clear_output(wait=True)
```

road situation analysis/research/road/state estimation/kalman_filter_demo.ipynb


# TESTBED FOR FULL STATE FEEDBACK

This notebook describes the proportional, integral, and differential controllers.

# Preliminaries

```python
!pip install -q control
!pip install -q tellurium
!pip install -q controlSBML

import control
import controlSBML as ctl
from IPython.display import HTML, Math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import sympy
import tellurium as te
```

```python
TIMES = ctl.makeSimulationTimes(0, 5, 500)
```

```python
MODEL = """
-> S0; 5
S0 -> S1; k0*S0
S1 -> S2; k1*S1
S2 -> S1; k2*S2
S2 -> ; k3*S2

k0 = 0.5
k1 = 1
k2 = 2
k3 = 3
S0 = 5
S1 = 0
S2 = 0
"""
RR = te.loada(MODEL)
RR.simulate(0, 10)
RR.plot()
```

# Calculating symbolic eigenvalues

```python
k0, k1, k2, k3, S1, S2, S0 = sympy.symbols("k0, k1, k2, k3, S1, S2, S0")
```

```python
A = sympy.Matrix([[-k1, k2], [k1, -k2-k3]])
A
```

$\displaystyle \left[\begin{matrix}- k_{1} & k_{2}\\k_{1} & - k_{2} - k_{3}\end{matrix}\right]$

```python
Ap = A.subs({k1:1, k2:2, k3:3})
Ap.as_real_imag()[0]
```

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

```python
mat = np.array([[-1, 2], [1, -5]])
np.linalg.eig(mat)
```

    (array([-0.55051026, -5.44948974]),
     array([[ 0.97566304, -0.40997761],
            [ 0.21927526,  0.91209559]]))

```python
3 - np.sqrt(6)
```

    0.5505102572168221

```python
eigenvalues = list(A.eigenvals().keys())
eigenvalues[1].simplify()
```

$\displaystyle - \frac{k_{1}}{2} - \frac{k_{2}}{2} - \frac{k_{3}}{2} + \frac{\sqrt{k_{1}^{2} + 2 k_{1} k_{2} - 2 k_{1} k_{3} + k_{2}^{2} + 2 k_{2} k_{3} + k_{3}^{2}}}{2}$

```python
sympy.expand((k1 + k2 + k3)**2)
```

$\displaystyle k_{1}^{2} + 2 k_{1} k_{2} + 2 k_{1} k_{3} + k_{2}^{2} + 2 k_{2} k_{3} + k_{3}^{2}$

```python
expr1 = -(k1 +k2 + k3) + sympy.sqrt( (k1 + k2 +k3)**2 - 4*(k1*(k3-k2)))
```

```python
sympy.N(expr1.subs({k1: 1, k2: 2, k3: 3}))
```

$\displaystyle -0.34314575050762$

```python
expr2 = -(k1 +k2 + k3) - sympy.sqrt( (k1 + k2 +k3)**2 - 4*(k1*(k3-k2)))
sympy.N(expr2.subs({k1: 1, k2: 2, k3: 3}))
```

$\displaystyle -11.6568542494924$

# DC Gain

```python
Ap = np.array([ [15, 25], [-1, 5]])
Ap_inv = np.linalg.inv(Ap)
print(Ap_inv)
App = np.matmul(Ap_inv, np.array([ [0.5], [0]]))
print(App)
print(np.matmul(np.array([0, 1]), App))
```

    [[ 0.05 -0.25]
     [ 0.01  0.15]]
    [[0.025]
     [0.005]]
    [0.005]

# Full State Feedback

```python
np.matmul(Ap, Ap_inv)
```

    array([[ 1.00000000e+00, -1.38777878e-16],
           [-1.73472348e-18,  1.00000000e+00]])

## Full State Controller

```python

```

```python
def makeController(name, kp_vec):
    """
    Creates a full state feedback controller.

    Parameters
    ----------
    name: str
        Name of the system
    kp_vec: list-float

    Returns
    -------
    control.NonlinearIOSystem
    """
    kp_vec = np.array(kp_vec)
    def outfcn(_, __, u_vec, ___):
        # u_vec: list-float - vector of internal state deviations
        ref = u_vec[0]
        arr = np.array(u_vec[1:])
        output = ref - kp_vec.dot(arr)
        return output
    #
    return control.NonlinearIOSystem(
        None, outfcn, inputs=['ref', 'in1', 'in2'], outputs=['out'],
        name=name)

# Tests
kp_vec = [1, 2]
controller = makeController("controller", kp_vec)
times = [0, 1, 2, 3, 4]
refs = np.repeat(1, len(times))
U = np.array(list(zip(refs, times, times)))
U = U.transpose()
result = control.input_output_response(controller, T=times, U=U)
assert(len(result.y[0]) == len(times))
assert(result.y[0][1] == -2)
print("OK!")
```

    OK!

```python
def makeFullStateFeedbackController(name, ctlsb, factor=1.0, poles=-2, time=0):
    """
    Creates a full state feedback controller for an SBML model where the system is linearized
    at the specified time.

    Parameters
    ----------
    name: str
        Name of the system
    ctlsb: ControlSBML
        SISO system
    factor: float
        factor for adjusting the reference input
    poles: list-float/float
        Desired poles
    time: float
        Time where system is lienarized

    Returns
    -------
    control.NonlinearIOSystem
       inputs:
           <state_variable> (excluding the input)
           ref: reference input
       outputs:
           out
    """
    REF = "ref"
    # Validity Checks
    if len(ctlsb.input_names) != 1:
        raise ValueError("SBML model must have a single input. Has: %s" % str(ctlsb.input_names))
    # Initializations
    state_space = ctlsb.makeStateSpace(time=time)
    controller_input_names = [n for n in ctlsb.state_names if not n in ctlsb.input_names]
    controller_input_names.insert(0, REF)  # first input
    num_state_input = len(controller_input_names) - 1
    try:
        _ = len(poles)
        poles = np.array(poles)
    except:
        poles = np.array([poles + 0.1*n for n in range(num_state_input)])
    # Calculate the gain matrix
    kp_vec = control.place(state_space.A, state_space.B, poles)
    def outfcn(time, _, u_vec, __):
        # u_vec: list-float - reference, state variables
        ref = factor*u_vec[0]
        arr = np.array(u_vec[1:])
        output = ref - kp_vec.dot(arr)
        return output
    #
    return control.NonlinearIOSystem(
        None, outfcn, inputs=controller_input_names, outputs=['out'],
        name=name)

# Tests
ctlsb = ctl.ControlSBML(MODEL, input_names=["S0"], output_names=["S2"])
controller = makeFullStateFeedbackController("controller", ctlsb, factor=1.0, poles=-10, time=1)
times = [0, 1, 2, 3, 4]
U = np.array([(1, 1, 1,) for _ in range(len(times))])
U = U.transpose()
result = control.input_output_response(controller, T=times, U=U)
assert(len(result.y[0]) == len(times))
print("OK!")
```

    OK!

```python
result.y
```

    array([[-79.8, -79.8, -79.8, -79.8, -79.8]])

```python
def calcKp(lambda1, lambda2, k0=0.5, k1=1, k2=2, k3=3):
    """
    Calculates the proportional constants for MODEL.
    """
    kp1 = (-lambda1 - lambda2 - k1 - k2 - k3)/k0
    kp2 = (lambda1*lambda2 - k1*k3 - kp1*k0*(k2 + k3))/(k0*k1)
    return np.array([kp1, kp2])

# Tests
poles = calcKp(-10, -11)
assert(poles[1] == 64)
print("OK!")
```

    OK!

## Network

```python
# Create elements
CTLSB = ctl.ControlSBML(MODEL, input_names=["S0"], output_names=["S1", "S2"])
system = CTLSB.makeNonlinearIOSystem("system")
controller = makeController("controller", calcKp(-10, -10))
```

```python
state_space = ctlsb.makeStateSpace()
state_space
```

\[
\left(\begin{array}{rllrll|rll}
-1\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&2\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0.&\hspace{-1em}5&\hspace{-1em}\phantom{\cdot}\\
1\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&-5\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
\hline
1\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&1\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
\end{array}\right)
\]

```python
control.place(state_space.A, state_space.B, [-10, -11])
```

    array([[30., 64.]])

```python
# Create the closed loop system with a disturbance input.
closed_loop = control.interconnect(
  [system, controller ],
  connections=[
    ['controller.in1', 'system.S1'],
    ['controller.in2', 'system.S2'],
    ["system.S0", "controller.out"],
  ],
  inplist=["controller.ref"],
  outlist = ["controller.out", "system.S1", "system.S2"],
)
```

```python
controller = makeFullStateFeedbackController("controller", CTLSB, factor=1, poles=-0.1)
# Create the closed loop system with a disturbance input.
closed_loop = control.interconnect(
  [system, controller ],
  connections=[
    ['controller.S1', 'system.S1'],
    ['controller.S2', 'system.S2'],
    ["system.S0", "controller.out"],
  ],
  inplist=["controller.ref"],
  outlist = ["controller.out", "system.S1", "system.S2"],
)
```

```python
X0 = ctl.makeStateVector(closed_loop)
times = ctl.makeSimulationTimes(end_time=10, points_per_time=10)
result = control.input_output_response(closed_loop, T=times, U=40, X0=X0)
plt.plot(result.t.flatten(), result.outputs[1].flatten())
plt.plot(result.t.flatten(), result.outputs[2].flatten())
_ = plt.legend(["S1", "S2"])
```

```python
closed_loop.linearize(X0, 5).pole()
```

    array([ 0.00000000e+00, -4.44087789e-09, -9.99999956e-02])


## 2. Introduction to Submodular Optimization

So far we've focused on looking at submodular functions in action, primarily in the context of identifying a good selection of samples for the purpose of training machine learning models. While we have covered at a high level how submodular selection works, in this tutorial we will focus on the math and mechanics behind submodular selection.

As mentioned before, submodular selection is a process of greedily selecting objects from a large set of objects in such a manner that a submodular function is maximized. This general definition allows submodular selection to be applied in a wide variety of areas. In our context we'll be focusing on the selection of samples for data analysis or machine learning purpose and using terminology common to that field.

The equation for a feature based function is below.

\begin{equation}
f(X) = \sum\limits_{u \in U} w_{u} \phi_{u} \left( \sum\limits_{x \in X} m_{u}(x) \right)
\end{equation}

In this equation, $U$ refers to all features in a sample, and $u$ refers to a specific feature. $X$ refers to the original data set that we are selecting from and $x$ refers to a single sample from that data set. $w$ is a vector of weights that indicate how important each feature is, with $w_{u}$ being a scalar referring to how important feature $u$ is. Frequently these weights are uniform. $\phi$ refers to a set of saturating functions, such as $sqrt(X)$ or $log(X + 1)$ that incude a property of "diminishing returns" on the feature values. These diminishing returns will become very important later.

When we maximize a submodular function, our goal is to, for each iteration, select the sample that yields the largest gain when added to the growing subset. This gain is dependant on the items that are already in the subset due to the saturating function. Let's talk through an example of selecting the best subset.

```python
%pylab inline
import seaborn; seaborn.set_style('whitegrid')
```

    Populating the interactive namespace from numpy and matplotlib

Let's generate a two dimensional data set for the purposes of vissualization, where some samples are high in one of the features and some samples are high in the other feature. We do this purposefully so as to illustrate the effect that selection has on the marginal gain of each sample during the selection process.

```python
numpy.random.seed(3)
X = numpy.concatenate([numpy.random.normal([0.5, 0.1], [0.3, 0.05], size=(50, 2)),
                       numpy.random.normal([0.1, 0.5], [0.05, 0.4], size=(50, 2))])
X = numpy.abs(X)

plt.figure(figsize=(8, 6))
plt.title("Randomly generated data", fontsize=16)
plt.scatter(X[:,0], X[:,1], s=20)
plt.xlim(0, 1.4)
plt.ylim(0, 1.4)
plt.show()
```

Now let's implement a function that calculates the gain of each sample with respect to some growing subset. Since we are using a feature based function, the gain requires summing column-wise down each sample in the growing set and then applying the saturating function in order to squash it. We can speed this up significantly by pre-computing the column-wise sums of the current selected set, so that we only need to add in the new sample to be considered.

```python
def gains(X, z=None):
    concave_sums = numpy.sum(z, axis=0) if z is not None else numpy.zeros(X.shape[1])
    concave_func = numpy.log(concave_sums + 1)
    gains = []

    for x in X:
        gain = numpy.sum(numpy.log(concave_sums + x + 1) - concave_func).sum()
        gains.append(gain)

    return gains
```

We can now use this function in order to calculate the gain of each of the samples in our set if we were to use it as the first sample in our subset.

```python
gain1 = gains(X)

plt.figure(figsize=(8, 6))
plt.title("Gain if adding this item", fontsize=16)
plt.scatter(X[:,0], X[:,1], c=gain1, cmap='Purples')
plt.colorbar()
plt.show()
```

We see a clear trend that the samples with high values in either of the two dimensions have high gains, whereas those samples with small values in each dimension have small gains. Given the simplicity of the feature-based selection algorithm this makes sense---high feature values correspond to higher gains.

What happens if we select the sample with the highest gain and then recalculate gains?

```python
idx = numpy.argmax(gain1)
z = [X[idx]]

gain2 = gains(X, z)

plt.figure(figsize=(8, 6))
plt.title("Gain if adding this item", fontsize=16)
plt.scatter(X[:,0], X[:,1], c=gain2, cmap='Purples')
plt.colorbar()

plt.scatter(X[idx, 0], X[idx, 1], c='c', s=75, label="First Selected Sample")
plt.legend(fontsize=14)
plt.show()
```

It looks like we select a sample that has the highest magnitude single feature value. While the saturating function will flatten values as the set gets larger, with small feature values and no samples in the subset yet, the raw feature value is roughly the gain.

Interestingly, it looks like all of the samples that have higher values for the y-axis now have a diminished marginal gain. The highest gain samples now look like they come from those with high x-axis values.

```python
idx = numpy.argmax(gain2)
z += [X[idx]]

gain3 = gains(X, z)

plt.figure(figsize=(8, 6))
plt.title("Gain if adding this item", fontsize=16)
plt.scatter(X[:,0], X[:,1], c=gain3, cmap='Purples')
plt.colorbar()

plt.scatter(X[idx, 0], X[idx, 1], c='c', s=75, label="Highest Gain")
plt.legend(fontsize=14)
plt.show()
```

To complete the selection process, we would greedily select samples from the full set until the desired number of samples have been reached. This requires scanning over the full set (minus the samples that have been selected) one full time for each samle that we would like to select, coting $nm$ time where $n$ is the number of samples in the full set and $m$ is the number of samples that onne would like to select. Because sometimes our goal is to induce a ranking over the full set, i.e., dtermine the order of selection for each sample, this becomes quadratic time.

Is it possible to do better than than this? The short answer is yes, and the reason lies in non-negativity constraint on the input data and on the saturating function.

```python
import pandas

gain_table = pandas.DataFrame({"x" : X[:,0], "y" : X[:,1], "gain 1": gain1, "gain 2": gain2, "gain 3": gain3})
gain_table[['x', 'y', 'gain 1', 'gain 2', 'gain 3']].head()
```

      x
      y
      gain 1
      gain 2
      gain 3

      0
      1.036589
      0.121825
      0.826233
      0.720047
      0.446619

      1
      0.528949
      0.006825
      0.431383
      0.398133
      0.224836

      2
      0.416784
      0.082262
      0.427443
      0.359090
      0.212751

      3
      0.475178
      0.068650
      0.455174
      0.391253
      0.229906

      4
      0.486855
      0.076139
      0.470043
      0.401884
      0.237428

We can exploit the property of diminishing returns to devise a better algorithm for selection known as the "lazy greedy" algorithm. This algorithm relies on the property that each time a sample is added to the growing set, the gains of each other sample diminish.

The algorithm works faily simply. First, a priority queue is built consisting of the samples prioritized by their gain, with the sample with the highest gain at the front. That sample is then selected for addition to the growing subset. After this selection, the gain values for each sample will be different. However, due to this monotonic decrease in gain, should the sample that was second in line at the time of selection remain at the front of the queue after that first element was selected and the sample was r-evaluated, it would be at the front of the queue even if all items were evaluated. It can be selected without needing to re-evaluate all other samples.

tutorials/2. Introduction to Submodular Optimization.ipynb