# Session 10: Unsupervised clustering with K-means

------------------------------------------------------
Introduction to Data Science & Machine Learning

*Pablo M. Olmos olmos@tsc.uc3m.es*

------------------------------------------------------

# Unsupervised Learning

In this notebook, we will study a particularly simple method to perform unsupervised learning over a dataset. In this case, we will use [K-means](https://www.datascience.com/blog/k-means-clustering) to divide our population in a set of **clusters**. In unsupervised learning, we expect that the structure behind our data can be interpreted or exploited for future applications.

This post is inspired by the excellent book [Python Data Science Handbook by Jake VanderPlas](https://jakevdp.github.io/PythonDataScienceHandbook/index.html). I mostly added all the theoretical background.

## Clustering with K-means

The K-means algorithm is rather simple. Given a **distance metric** $d(\mathbf{x},\mathbf{x}')$ and a number $K$ of clusters selected in advance, we perform as follows.

Until convergence:

1. Initialize $K$ *centroids* at random among points in the dataset. Denote the centroids as $\mathbf{c}_k$, $k=1,\ldots,K$.
2. For $i=1,\ldots,N$, assign point $\mathbf{x}^{(i)}$ to the cluster with the closest centroid. If $A(\mathbf{x}^{(i)})\in\{1,\ldots,K\}$ is the cluster assigned to point $\mathbf{x}^{(i)}$ then

\begin{align}
A(\mathbf{x}^{(i)}) = \arg \min_{\{1,\ldots,K\}} d(\mathbf{x}^{(i)},\mathbf{c}_k)
\end{align}

3. Recompute centroids as the mean across all points assigned to the same cluster. For $k=1,\ldots,K$,

\begin{align}
\mathbf{c}_k = \frac{1}{N_k} \sum_{i: A(\mathbf{x}^{(i)})= k} \mathbf{x}^{(i)},
\end{align}
where $\frac{1}{N_k}$ is the number of points assigned to cluster $k$.

4. Declare convergence if no changes between two consecutive iterations or maximum number of iterations achieved.

```python
import matplotlib.pyplot as plt
import numpy as np
import scipy.io
from sklearn.datasets.samples_generator import make_blobs
from sklearn.cluster import KMeans,MiniBatchKMeans

%matplotlib inline
```

### Lets create a toy data set

```python
X, y_true = make_blobs(n_samples=500, centers=6,
                       cluster_std=0.6, random_state=22)

plt.scatter(X[:, 0], X[:, 1], s=50);
plt.xlabel('$x_1$')
plt.ylabel('$x_2$')
```

```python
#Lets run Kmeans clustering
K=4

kmeans = KMeans(n_clusters=K)
kmeans.fit(X)
y_kmeans = kmeans.predict(X)
```

Let's visualize the results by plotting the data colored by these labels. We will also plot the cluster centers as determined by the k-means estimator:

```python
plt.scatter(X[:, 0], X[:, 1], c=y_kmeans, s=50, cmap='viridis')

centers = kmeans.cluster_centers_
plt.scatter(centers[:, 0], centers[:, 1], c='black', s=200, alpha=0.5)
```

## Selecting the number of clusters

The number $K$ of clusters represent the model complexity. When we set large number of clusters, we are assuming tha the data structure is very complex, and we might be certainly **overfitting** our data.

How to diagnose overfitting in unsupervised clustering? There are several approaches. A common one is to set a model performance metric and analyze how it depends on the model complexity.

In $K$-means, the natural choice is to analyze the **average distance of points w.r.t. their cluster**. This is in fact the function that we minimize to derive the $K$-means algorithm!
\begin{align}
\mathcal{L}(\mathbf{c}_1,\ldots,\mathbf{c}_K) = \frac{1}{K} \sum_{k=1}^K \sum_{i: A(\mathbf{x}^{(i)})= k} \frac{d(\mathbf{x}^{(i)},\mathbf{c}_k)}{N_k}
\end{align}

If we optimize the above function w.r.t. $\mathbf{c}_k$, $k=1,\ldots,K$, by implementing a [**coordinate gradient descent**](https://en.wikipedia.org/wiki/Coordinate_descent) we derive the $K$-means algorithm.

Lets plot the evolution of $\mathcal{L}(\mathbf{c}_1,\ldots,\mathbf{c}_K)$ for different values of $K$.

```python
kmeans.inertia_
```

    3213.6836424069884

```python
K_list = range(2,15)

L = []

for k in K_list:

    kmeans = KMeans(n_clusters=k)
    kmeans.fit(X)
    L.append((kmeans.inertia_)/(0.0+k*X.shape[0]))

plt.figure()
plt.plot(K_list,L)
plt.xlabel('$K$')
plt.ylabel('$L$')
plt.grid()
```

Clearly, the cost function becomes flat for large $K$ values. This means that we are making our model **more complex** but we are not essentially explaining our data better! A sign of **overfitting**.

Given the above figure, it would make sense to select a value of $K$ around 6 (which is also the true value in this case).

Be aware that this plot is dependent of the number of data points. Also, that the $K$-means solution is **highly dependent** on the intialization. Actually, it is usual to run it several times and select the solution that provides the smallest $\mathcal{L}(\mathbf{c}_1,\ldots,\mathbf{c}_K)$ function. In the [**Sklearn K-means library**](http://scikit-learn.org/stable/modules/generated/sklearn.cluster.KMeans.html), it is controlled by the parameter *n_init*.

```python
N_list =[20,50,100,500]

K_list = range(2,8)

plt.figure()

for n in N_list:

    X, y_true = make_blobs(n_samples=n, centers=6,
                       cluster_std=0.6, random_state=22)

    L = []

    for k in K_list:

        kmeans = KMeans(n_clusters=k,n_init=5)
        kmeans.fit(X)
        L.append((kmeans.inertia_)/(0.0+k*X.shape[0]))

    plt.plot(K_list,L,label='N_points='+str(n))

plt.xlabel('$K$')
plt.ylabel('$L$')
plt.grid()
plt.legend()
```

## Mini-batch $K$-means

Updating centroids at every $K$-means update using all datapoints can be costly for very large databases. An scalable and effective solution uses small **mini-batches** of data at every iteration.  This is known as **mini-batch** $K$-means, and it is also provided in the [**sklearn library**](http://scikit-learn.org/stable/modules/generated/sklearn.cluster.MiniBatchKMeans.html#sklearn.cluster.MiniBatchKMeans).

Compare the running times of the following two examples.

```python
# Without mini-batches. Full K-means

n =int(1e5)

K_list = range(2,8)

X, y_true = make_blobs(n_samples=n, centers=6,
                       cluster_std=0.6, random_state=22)
L = []

batch = 100

for k in K_list:

    kmeans = MiniBatchKMeans(n_clusters=k,batch_size=batch,n_init=5)
    kmeans.fit(X)
    L.append((kmeans.inertia_)/(0.0+k*X.shape[0]))

plt.figure()
plt.plot(K_list,L)
plt.xlabel('$K$')
plt.ylabel('$L$')
plt.grid()
```

**Summary**: $K$-means is a simple and cheap algorithm for unsupervised clustering. However, it has fundamental weaknesses that we will explain when we introduce probabilisitic clustering using **Gaussian Mixture Models** and the **Expectation Maximization Algorithm**.

```python

```


$\newcommand{\pr}{\textrm{Pr}}$
$\newcommand{\l}{\left}$
$\newcommand{\r}{\right}$
$\newcommand\given[1][]{\:#1\vert\:}$
$\newcommand{\var}{\textrm{Var}}$
$\newcommand{\mc}{\mathcal}$
$\newcommand{\lp}{\left(}$
$\newcommand{\rp}{\right)}$
$\newcommand{\lb}{\left\{}$
$\newcommand{\rb}{\right\}}$
$\newcommand{\iid}{\textrm{i.i.d. }}$
$\newcommand{\ev}{\textrm{E}}$
$\newcommand{\odds}{\textrm{odds}}$

So far, we've looked at simple posterior distributions that we can evaluate
exactly. However, sometimes we may be interested in more complex posteriors
that we cannot evaluate numerically. We may also be interested in other quantities
like the $\pr\lp \theta \in A \given y_1, \cdots, y_n\rp$ for arbitrary sets $A$.
We may also want to calculate the posterior of some function of $\theta$
like $\l|\theta_1 - \theta_2\r|$. We can use Monte Carlo sampling to do this.

# 4.1 The Monte Carlo method

**Monte Carlo approximation** is based on random sampling. Let $\theta$ be a parameter
of interest and let $y_1, \cdots, y_n$ be the numerical values of a sample from a distribution
$p\lp y_1, \cdots, y_n \given \theta \rp$. Suppose we take $S$ independent, random samples
from the posterior distribution $p\lp \theta \given y_1, \cdots, y_n \rp$:
$$\theta^{\lp 1\rp}, \theta^{\lp 2\rp}, \cdots, \theta^{\lp S\rp} \sim \iid p\lp \theta \given y_1, \cdots, y_n \rp$$
Since we are sampling from the posterior, these $S$ samples will approximate the posterior,
and the accuracy will increase as $S$ increases. The empirical distribution
$\lb\theta^{\lp 1\rp}, \theta^{\lp 1\rp}, \cdots, \theta^{\lp S\rp}\rb$ is known as the
**Monte Carlo approximation** to $p\lp \theta \given y_1, \cdots, y_n \rp$. As $S\rightarrow\infty$,

* $\bar{\theta} = \sum_{s=1}^S \theta^{\lp s\rp} / S \rightarrow \ev\l[\theta \given y_1, \cdots, y_n\r]$ (the empirical mean of the samples goes to the expected value of the posterior)
* $\sum_{s=1}^S \lp\theta^{\lp s\rp} - \bar{\theta}\rp^2 / \lp S-1\rp \rightarrow \var\l[\theta \given y_1, \cdots, y_n\r]$ (the empirical variance of the sample goes to the variance of the posterior)
* #$\lp \theta^{\lp S\rp} \leq c\rp / S \rightarrow \pr\lp\theta \leq c \given y_1, \cdots, y_n\rp$ (the empirical
probability of observing a sample less than $c$ approximates the cdf of the posterior)
* the empirical distribution of $\lb \theta^{\lp 1\rp}, \cdots, \theta^{\lp S\rp} \rb
\rightarrow p\lp \theta \given y_1, \cdots, y_n \rp$
* the median of $\lb \theta^{\lp 1\rp}, \cdots, \theta^{\lp S\rp} \rb \rightarrow \theta_{1/2}$
* the $\alpha$-percentile of $\lb \theta^{\lp 1\rp}, \cdots, \theta^{\lp S\rp} \rb \rightarrow \theta_\alpha$

Almost any aspect of a posterior distribution can be approximated arbitrarily exactly using Monte Carlo,
and this is true for almost all functions $g\lp\theta\rp$ as well.

You can look at plots of estimated values to see when you have performed enough Monte Carlo
samples. For instance, you could plot the empirical mean of the samples (which would be the
expected value of the posterior distribution). With enough samples, this estimate should
converge to some value.

If you are using the mean of the Monte Carlo samples $\bar{\theta}$ to estimate the mean of the posterior,
the Central Limit Theorem says that $\hat{\theta}$ is approximately normally distributed with mean
$\ev\l[\theta\given y_1,\cdots,y_n\r]$ and standard deviation $\sqrt{\var\l[\theta\given y_1,\cdots,y_n\r] / S}$.
The Monte Carlo standard error is the approximation to this standard deviation $\sqrt{\hat{\sigma}^2 / S}$ where
$\hat{\sigma}$ is the empirical standard deviation. An approximate 95% confidence interval for the posterior
mean of $\theta$ is $\hat{\theta} \pm 2\sqrt{\hat{\sigma}^2/S}$. You can choose $S$ large enough to get the
desired precision for $\theta$.

# 4.2 Posterior inference for arbitrary functions

Suppose we are interested in the posterior distribution of some computable function $g\lp\theta\rp$
of $\theta$. For the binomial model, we may be interested in the log odds
$$\log \odds\lp\theta\rp = \log\frac{\theta}{1-\theta} = \gamma$$
The law of large numbers says that if we sample from the posterior distribution of $\theta$,
then the average value of $\log\frac{\theta^{\lp S\rp}}{1-\theta^{\lp S\rp}}$ converges to
$\ev\l[\log\frac{\theta}{1-\theta}\given y_1, \cdots, y_n\r]$. You can also calculate
$g\lp\theta^{\lp 1\rp}\rp, \cdots, g\lp\theta^{\lp S\rp}\rp$ independently as you sample.
From these samples, as $S\rightarrow\infty$, you get all of the information from the bullet
points above (mean, variance, empirical distribution, median, $\alpha$-percentile).

# 4.3 Sampling from predictive distributions

The predictive distribution of a random variable $\tilde{Y}$ is a probability distribution for
$\tilde{Y}$ such that

* known quantities have been conditioned out
* unknown quantities have been integrated out

A predictive distribution that integrates over unknown parameters but is not conditional
on observed data is called a **prior predictive distribution**. Such a distribution
can be useful for evaluating whether a prior distribution for $\theta$ actually
translates into reasonable prior beliefs for observable data $\tilde{Y}$.

After we observe data, the **posterior predictive distribution** is
$$\pr\lp\tilde{Y} = \tilde{y} \given Y_1=y_1, \cdots, Y_n=y_n\rp =
\int p\lp\tilde{y}\given\theta\rp p\lp\theta\given y_1, \cdots, y_n\rp d\theta$$
In the case of a Poisson model with a gamma prior, the posterior predictive distribution
was a negative binomial.

Often, we can sample from $p\lp\theta\given y_1, \cdots, y_n\rp$ and $p\lp y\given\theta\rp$
but it is too complicated to sample from the posterior predictive distribution. We can use
a Monte Carlo procedure to sample indirectly from the posterior predictive distribution instead.
So we want a set of samples from $\tilde{Y}$ from its posterior  predictive distribution. We can
do this by sampling
\begin{align}
&\textrm{sample } \theta^\lp 1\rp \sim p\lp\theta\given y_1, \cdots, y_n\rp,
\textrm{sample } \tilde{y}^\lp 1\rp \sim p\lp\tilde{y}\given \theta^\lp 1\rp \rp \\
&\vdots \\
&\textrm{sample } \theta^\lp S\rp \sim p\lp\theta\given y_1, \cdots, y_n\rp,
\textrm{sample } \tilde{y}^\lp S\rp \sim p\lp\tilde{y}\given \theta^\lp S\rp \rp
\end{align}
In other words, you sample from the posterior of $\theta$ and use that sample to sample
$\tilde{Y}$. These joint samples constitute $S$ indepdendent samples from the joint posterior
distribution of $\lp \theta, \tilde{Y}\rp$, and the samples of $\tilde{Y}$ constitute
$S$ independent samples from the marginal posterior distribution of $\tilde{Y}$ which
is the posterior predictive distribution.

# 4.4 Posterior predictive model checking

We can plot the posterior predictive distribution (made from sampling) and the
empirical distribution of our data to see whether they match. You would imagine
that the distributions should be similar or else it seems that the predictive
distribution would not be good for predicting future values from the empirical
distribution.

There are a couple resons the empirical (observed) distribution and predictive distribution
may not look similar:

* the empirical distribution doesn't match the exact distribution of the population from which the data
were sampled. This could be due to sampling variability, especially when the number of samples is small
* our model is not the correct model for this data

We can evaluate these two possibilities by performing repeated simulations of the observed
dataset and checking to see how likely the observed data are given repeated samplings of the
posterior. Note that you could check "how likely the observed data are" in different ways.
For instance, you could compare the empirical mean versus the means from the samplings.
In the book, they compare the ratio of 1s versus 2s for a Poisson model. The ratio for the
observed data looks much different than the ratios from the samplings which suggests the model
is not accurately capturing the shape of the distribution. However, just because a model doesn't
accurately capture one part of the data doesn't mean it's useless. The model may still recapitulate
the mean or variance. For instance, the Poisson model may still provide consistente estimates of
the mean and variance for a distribution where the mean and variance are roughly equal.

readings/A First Course in Bayesian Statistical Methods/4. Monte Carlo approximation.ipynb


# CONTROLLABILITY OF BIOLOGICAL SYSTEMS

This notebook explores the construction and interpretation of transfer functions for more complex networks.

# Preliminaries

```python
!pip -q install controlSBML
import controlSBML as ctl
import control
from controlSBML.util import makeSimulationTimes

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import tellurium as te
import sympy
print("controlSBML version: " + ctl.__version__)
```

    controlSBML version: 0.2.10

# Transfer Functions in the ``control`` Package

Suppose our transfer function is $G(s) = \frac{1}{s} \frac{k_1}{k_2 + s}$.

The ``control`` packages provides a way to construct transfer functions that are a ratio of
polynomials in $s$.

## Representing Polynomials

A polynomial in $s$ is represented as a ``python`` list.
For example, $s + k_2$ is the polynomial $(1)s^1 + (k_2) s^0$.
It is represented by the list ``[1, k2]``, where ``k2`` is the floating point value of $k_2$.
The interpretation is that the last element of the list is the coefficient of $s^0$;
elements that preceed the last element represent successive powers of $s$.

**Question: What is the list that represents the polynomial $3 S^3 + 2s$?**

## Creating Transfer Functions

To create a transfer function, you provide the numerator and denominator polynomial.
Consider the construction of the transfer function for
$G(s) = \frac{1}{s} \frac{k_1}{k_2 + s} = \frac{k_1}{s^2 + k_2 s}$.

**Question: What are the poles of $G(s)$?**

```python
k1 = 1
k2 = 2
tf = control.TransferFunction([k1], [1, k2, 0])
tf
```

$$\frac{1}{s^2 + 2 s}$$

## Using ``TransferFunction`` objects

There are several ways to use transfer function objects.
You can find the poles, calculate DC gain, and simulate impulse and step responses.

### Poles of a TransferFunction

```python
# The poles should be 0, -k2 = -2
tf.pole()
```

    array([-2.,  0.])

### DC Gain of a TransferFunction

```python
tf.dcgain()
```

    inf

**Question: Why is DC Gain = $\infty$?**

### Impulse Response

```python
# You can use controlSBML to make a sequence of simulation times. By default it's in increments
# of 0.1 from 0 to 5.
TIMES = ctl.makeSimulationTimes(end_time=10, points_per_time=5)
TIMES
```

    array([ 0. ,  0.2,  0.4,  0.6,  0.8,  1. ,  1.2,  1.4,  1.6,  1.8,  2. ,
            2.2,  2.4,  2.6,  2.8,  3. ,  3.2,  3.4,  3.6,  3.8,  4. ,  4.2,
            4.4,  4.6,  4.8,  5. ,  5.2,  5.4,  5.6,  5.8,  6. ,  6.2,  6.4,
            6.6,  6.8,  7. ,  7.2,  7.4,  7.6,  7.8,  8. ,  8.2,  8.4,  8.6,
            8.8,  9. ,  9.2,  9.4,  9.6,  9.8, 10. ])

```python
result = control.impulse_response(tf, T=TIMES)
```

``result`` is complicated objects with many properties.
* ``result.outputs`` - ``result.outputs[n]`` is a time series for the output ``n``
* ``result.t`` - times simulated (which should be the same as ``TIMES``)

```python
# Since we're doing single input single output (SISO) systems,
# you can simplify matters by just flattening result.t, result.outputs
plt.plot(result.t.flatten(), result.outputs.flatten())
```

**Question: The above impulse response converges to 0.5. How is this predicted by $G(s)$?**

Hint: Recall that for an impulse input, $y(\infty) = lim_{s \rightarrow 0} s G(s)$.

### Step Response

**Question: Does $G(s)$ converge for a step response? If so, to what value?**

Hint: Recall that for a step input, $y(\infty) = lim_{s \rightarrow 0}  G(s)$.

Consider $H(s) = \frac{k_1}{s + k_2}$.

**Question: Construct ``control.TransferFunction`` object for $H(s)$. if $k_2 = 3$.**

**REMOVE**

```python
k2 = 3
tf = control.TransferFunction([k1], [1, k2])
tf
```

$$\frac{1}{s + 3}$$

```python
result = control.step_response(tf, T=TIMES)
plt.plot(result.t.flatten(), result.outputs.flatten())
```

**Question: The above step response converges to 0.33. How is this predicted by $H(s)$?**

# State Space Representation in the ``control`` Package

Consider the reaction network

    S1 -> S2 + S3; k1*S1
    S2 -> S3;      k2*S2
    S3 ->   ;  k3*S3

    k1 = 0.5
    k2 = 1.5
    k3 = 0.75

**Question: What are the system equations for this network?**

**REMOVE**

\begin{eqnarray}
\dot{S}_1 & = & - k_1 S_1 \\
\dot{S}_2 & = & k_1 S_1 - k_2 S_2 \\
\dot{S}_3 & = & k_1 S_1 + k_2 S_2 - k_3 S_3 \\
\end{eqnarray}

A state space representation of a SISO system is:

\begin{eqnarray}
\dot{\bf x} & = & {\bf A} {\bf x} + {\bf b} u \\
y & = & {\bf c} {\bf x}
\end{eqnarray}

where
* ${\bf x}$ is a vector of state
* $u$ is the scalar input
* $y$ is the scalar output
* ${\bf A}$, a matrix of constants, describes the relationships between state variables
* ${\bf b}$ is a column vector that indicates how the input affects each state
* ${\bf c}$ is a row vector that specifies how each state contributes to the output

### Creating State Space Objects for SISO Systems

The ${\bf A}$ matrix is an alternative representation of the state equations.
The rows are the equations.
In our example, row 2 is the equation for $\dot{S}_2$.
(Of course, the python index for this row is 1.)
The columns, represent the state variable.
So, column 3 is for $S_3$.

That is
$\dot{\bf x} =
{\bf A} =
\begin{bmatrix}
-k_1 & 0 & 0\\
k_1 & -k_2 & 0  \\
k_1 & k_2 & -k_3
\end{bmatrix}
{\bf x} + {\bf b}u$.

**Question: If the input is $S_1$, what is ${\bf b}$?**

**REMOVE**

${\bf b} =
\begin{bmatrix}
1 \\
0 \\
0 \\
\end{bmatrix}.
$

### Creating ``control.StateSpace`` objects

Consider a system with $S_1$ as input and $S_3$ as output.

```python
k1 = 0.5
k2 = 1.5
k3 = 0.75

A = np.array( [ [-k1, 0, 0],
    [k1, -k2, 0],
     [k1, k2, -k3]])
B = [1, 0, 0]
C = [0, 0, 1]
state_space = control.StateSpace(A, B, C, 0)
state_space
```

\[
\left(\begin{array}{rllrllrll|rll}
-0.&\hspace{-1em}5&\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.&\hspace{-1em}5&\hspace{-1em}\phantom{\cdot}&-1.&\hspace{-1em}5&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
0.&\hspace{-1em}5&\hspace{-1em}\phantom{\cdot}&1.&\hspace{-1em}5&\hspace{-1em}\phantom{\cdot}&-0.&\hspace{-1em}75&\hspace{-1em}\phantom{\cdot}&0\phantom{.}&\hspace{-1em}&\hspace{-1em}\phantom{\cdot}\\
\hline
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)
\]

### Using State Space Objects

* simulation: ``control.input_output_response``
* convert to transfer function (since this is easier to deriving the transfer function: ``control.ss2tf``

```python
tf = control.ss2tf(state_space)
tf
```

$$\frac{0.5 s + 1.5}{s^3 + 2.75 s^2 + 2.25 s + 0.5625}$$

```python
tf.dcgain()
```

    2.6666666666666683

# mTOR Signaling

For your homework, you create a state space object for an SBML model.
1. Load the model into Tellurium
1. Simulate to the time for the operating point.
1. Construct ${\bf A}$ from the Jacobian (``getFullJacobian``).
1. Construct ${\bf b}$ based on the index of the chemical species being controlled.
1. Construct ${\bf c}$ based on the index of the chemical species for the output.
1. ``state_space = control.StateSpace(A, b, c, 0)``.

## Constructing State Space Objects from Tellurium

```python
rr = te.loadSBMLModel("https://www.ebi.ac.uk/biomodels/model/download/BIOMD0000000823.2?filename=Varusai2018.xml")
#print(rr.getAntimony())
```

There's a some work to get the indices right to construct ${\bf b}$, ${\bf c}$,
and to get the species names.

## Constructing ``control`` objects using ``controlSBML``

```python
ctlsb = ctl.ControlSBML("https://www.ebi.ac.uk/biomodels/model/download/BIOMD0000000823.2?filename=Varusai2018.xml",
                       input_names=["pAkt"], output_names=["mTORC1_DEPTOR"])
mtor_nl = ctlsb.makeNonlinearIOSystem("mtor")
mtor_ss = ctlsb.makeStateSpace()
MTOR_TF = ctlsb.makeTransferFunction(atol=1e-5, time=1)
```

## Controllability Analysis

Suppose that we want to control the complex ``mTORC1_DEPTOR``.
What inputs should we control?
We start by controlling the amount of insullin.

```python
MTOR_TF
```

$$\frac{-0.03472 s^3 - 0.1617 s^2 - 0.09749 s - 0.001885}{s^5 + 11.05 s^4 + 34.96 s^3 + 29.05 s^2 + 7.379 s + 0.1362}$$

```python
MTOR_TF.dcgain()
```

    -0.013840450064546183

```python
MTOR_TF.pole()
```

    array([-5.99984811, -3.96034236, -0.57142857, -0.50165764, -0.02      ])


# Tutorial 4: A two-asset HANK model

In this notebook we solve the two-asset HANK model from Auclert, Bardóczy, Rognlie, Straub (2021): "Using the Sequence-Space Jacobian to Solve and Estimate Heterogeneous-Agent Models" ([link to paper](https://www.bencebardoczy.com/publication/sequence-jacobian/sequence-jacobian.pdf)).

New concepts:
- **Solved block**: an extension of simple blocks that enables much more efficient DAG representations of large macro models.
- **Re-using saved Jacobians**: as the cost of these computations becomes non-trivial, avoiding redundancy becomes key.
- **Fine-tuning options**: how to access and modify various options for each (block, method) pair

For more examples and information on the SSJ toolkit, please visit our [GitHub page](https://github.com/shade-econ/sequence-jacobian).

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

from sequence_jacobian import simple, solved, combine, create_model  # functions
from sequence_jacobian import grids, hetblocks                       # modules
```

## 1 Model description

The household problem is characterized by the Bellman equation
$$
\begin{align} \tag{1}
V_t(e, b_{-}, a_{-}) = \max_{c, b, a} &\left\{\frac{c^{1-\sigma}}{1-\sigma} + \beta \mathbb{E}_t V_{t+1}(e', b, a) \right\}
\\
c + a + b &= z_t(e) + (1 + r_t^a)a_{-} + (1 + r_t^b)b_{-} - \Psi(a, a_{-})
\\
a &\geq \underline{a}, \quad b \geq \underline{b},
\end{align}
$$

where $z_t(e)$ is labor income and the adjustment cost function is specified as

$$
\Psi(a, a_{-}) =  \frac{\chi_1}{\chi_2}\left|\frac{a - (1 + r_t^a) a_{-}}{(1 + r_{t}^a) a_{-} + \chi_0}\right|^{\chi_2} \left[(1 + r_t^a) a_{-} + \chi_0 \right],
$$

with $\chi_0, \chi_1 > 0$ and $\chi_2 > 1.$ For the full description of the model, including the problems of the other agents, please see appendix B.3 of the paper.

We embed this household block in a New Keynesian model with sticky prices, sticky wages, and capital adjustment costs. Thanks to the **solved blocks** (in green), we can write a DAG for this model in just 3 unknowns $(r, w, Y)$ and 3 targets, asset market clearing, Fisher equation, wage Phillips curve.

## 2 Define solved blocks

Solved blocks are miniature models embedded as blocks inside of our larger model. Like simple blocks, solved blocks correspond to aggregate equilibrium conditions: they map sequences of aggregate inputs directly into sequences of aggregate outputs. The difference is that in the case of simple blocks, this mapping has to be analytical, while solved blocks are designed to accommodate implicit relationships that can only be evaluated numerically.

Such implicit mappings between variables become more common as macro complexity increases. Solved blocks are a valuable tool to simplify the DAG of large macro models.

### 2.1 Price setting (NKPC-p)
The Phillips curve characterizes $(\pi)$ conditional on $(Y, mc, r):$

$$
\log(1+\pi_t) = \kappa_p \left(mc_t - \frac{1}{\mu_p} \right) + \frac{1}{1+r_{t+1}} \frac{Y_{t+1}}{Y_t} \log(1+\pi_{t+1})
$$

Inflation shows up with two different time displacements, which means that inflation in any given period depends on the entire sequence of $(Y, mc, r)$. Simple blocks are not meant to represent such relationships. Instead, we write a function that returns the residual of the equation, and use the decorator `@solved` to make it into a `SolvedBlock`.

```python
@solved(unknowns={'pi': (-0.1, 0.1)}, targets=['nkpc'], solver="brentq")
def pricing_solved(pi, mc, r, Y, kappap, mup):
    nkpc = kappap * (mc - 1/mup) + Y(+1) / Y * (1 + pi(+1)).apply(np.log) / \
           (1 + r(+1)) - (1 + pi).apply(np.log)
    return nkpc
```

When our routines encounter a solved block in `blocks`, they compute its Jacobian via the the implicit function theorem, as if it was a model on its own. Given the Jacobian, the rest of the code applies without modification.

### 2.2 Equity price (equity & dividend)
The no arbitrage condition characterizes $(p)$ conditional on $(d, p, r).$

$$
p_t = \frac{d_{t+1} + p_{t+1}}{1 + r_{t+1}}
$$

```python
@solved(unknowns={'p': (5, 15)}, targets=['equity'], solver="brentq")
def arbitrage_solved(div, p, r):
    equity = div(+1) + p(+1) - p * (1 + r(+1))
    return equity
```

### 2.3 Investment with adjustment costs (prod)

Sometimes multiple equilibrium conditions can be combined in a self-contained solved block. Investment subject to capital adjustment costs is such a case. In particular, we can use the following four equations to solve for $(K, Q)$ conditional on $(Y, w, r)$.

 - Production:

 $$
 Y_t = Z_t K_{t-1}^\alpha N_t^{1-\alpha}
 $$

 - Labor demand:

 $$
 w_t = (1-\alpha)\frac{Y_t}{N_t} mc_t
 $$

 - Investment equation:

$$
Q_t = 1 + \frac{1}{\delta \epsilon_I}\left(\frac{K_t-K_{t-1}}{K_{t-1}}\right)
$$

- Valuation equation

$$
(1+r_{t})Q_{t} = \alpha Z_{t+1} \left(\frac{N_{t+1}}{K_t}\right)^{1-\alpha} mc_{t+1} - \left[\frac{K_{t+1}}{K_t} - (1-\delta) + \frac{1}{2\delta \epsilon_I}\left(\frac{K_{t+1} - K_t}{K_t}\right)^2\right] + \frac{K_{t+1}}{K_t}Q_{t+1}
$$

Solved blocks that contain multiple simple blocks have to be initialized with the `CombinedBlock.solved` method instead of the decorator `@solved`.

```python
@simple
def labor(Y, w, K, Z, alpha):
    N = (Y / Z / K(-1) ** alpha) ** (1 / (1 - alpha))
    mc = w * N / (1 - alpha) / Y
    return N, mc

@simple
def investment(Q, K, r, N, mc, Z, delta, epsI, alpha):
    inv = (K / K(-1) - 1) / (delta * epsI) + 1 - Q
    val = alpha * Z(+1) * (N(+1) / K) ** (1 - alpha) * mc(+1) -\
        (K(+1) / K - (1 - delta) + (K(+1) / K - 1) ** 2 / (2 * delta * epsI)) +\
        K(+1) / K * Q(+1) - (1 + r(+1)) * Q
    return inv, val

production = combine([labor, investment])                              # create combined block
production_solved = production.solved(unknowns={'Q': 1., 'K': 10.},    # turn it into solved block
                                      targets=['inv', 'val'],
                                      solver='broyden_custom')
```

## 3 Build DAGs

One for transition dynamics (pictured above) and one for calibrating the steady state.

### Step 1: Adapt HA block

We developed an efficient backward iteration function to solve the Bellman equation in (1). Although we view this as a contribution on its own, discussing the algorithm goes beyond the scope of this notebook. If you are interested in how we solve a two-asset model with convex portfolio-adjustment costs in discrete time, please see appendix E of the paper for a detailed description and `sequence_jacobian/hetblocks/hh_twoasset.py` for the implementation.

Here, we take this generic two-asset model off the shelf, and embed it in our New Keynesian model with the help of two hetinputs.

```python
def make_grids(bmax, amax, kmax, nB, nA, nK, nZ, rho_z, sigma_z):
    b_grid = grids.agrid(amax=bmax, n=nB)
    a_grid = grids.agrid(amax=amax, n=nA)
    k_grid = grids.agrid(amax=kmax, n=nK)[::-1].copy()
    e_grid, _, Pi = grids.markov_rouwenhorst(rho=rho_z, sigma=sigma_z, N=nZ)
    return b_grid, a_grid, k_grid, e_grid, Pi

def income(e_grid, tax, w, N):
    z_grid = (1 - tax) * w * N * e_grid
    return z_grid

hh = hetblocks.hh_twoasset.hh
hh_ext = hh.add_hetinputs([income, make_grids])
```

### Step 2: Complete dynamic DAG with simple blocks

We have set up all the blocks in the `sequence_jacobian/examples/two_asset.py` module. We omit the step-by-step discussion of these blocks since they should be familiar from the other model notebooks.

```python
import sequence_jacobian.examples.two_asset as m

blocks = [hh_ext, production_solved, pricing_solved, arbitrage_solved,
          m.dividend, m.taylor, m.fiscal, m.share_value,
          m.finance, m.wage, m.union, m.mkt_clearing]

hank = create_model(blocks, name='Two-Asset HANK')

print(*hank.blocks, sep='\n')
```

    <SolvedBlock 'labor_to_investment_combined_solved'>
    <SolvedBlock 'pricing_solved'>
    <SimpleBlock 'wage'>
    <SimpleBlock 'taylor'>
    <SimpleBlock 'dividend'>
    <SolvedBlock 'arbitrage_solved'>
    <SimpleBlock 'share_value'>
    <SimpleBlock 'finance'>
    <SimpleBlock 'fiscal'>
    <HetBlock 'hh' with hetinput 'make_grids_marginal_cost_grid' and with hetoutput `adjustment_costs'>
    <SimpleBlock 'union'>
    <SimpleBlock 'mkt_clearing'>

### Step 3: Complete calibration DAG

Analytical:
- find TFP `Z` to hit target for output `Y`
- find markup `mup` to hit target for total wealth `p + Bg`
- find capital share `alpha` to hit target for capital `K`
- find wage `w` to hit Phillips curve given zero inflation
- find disutility of labor `vphi` to hit wage Phillips curve given a target for employment

Numerical:
- find discount factor `beta` to satisfy asset market clearing given an interest rate `r`
- find adjustment cost scale `chi1` to hit target for average liquid wealth `Bh`

```python
blocks_ss = [hh_ext, m.partial_ss, m.union_ss,
             m.dividend, m.taylor, m.fiscal, m.share_value, m.finance, m.mkt_clearing]

hank_ss = create_model(blocks_ss, name='Two-Asset HANK SS')

print(hank_ss)
print(f"Inputs: {hank_ss.inputs}")

```

    <Model 'Two-Asset HANK SS'>
    Inputs: ['beta', 'eis', 'chi0', 'chi1', 'chi2', 'N', 'bmax', 'amax', 'kmax', 'nB', 'nA', 'nK', 'nZ', 'rho_z', 'sigma_z', 'Y', 'K', 'r', 'tot_wealth', 'Bg', 'delta', 'muw', 'frisch', 'pi', 'kappap', 'epsI', 'rstar', 'phi', 'G', 'Bh', 'omega']

## 4 Results

We cover how to pass precomputed Jacobians to the main methods. This is useful when methods that need Jacobians are used repeatedly. These are
- Solve methods: `solve_impulse_linear`, `solve_impulse_nonlinear`
- Jacobian methods: `jacobian`, `solve_jacobian`

### 4.1 Calibrate steady state

Use the calibration DAG to internally calibrate the seven parameters (analytical + numerical). Evaluate the dynamic DAG at the resulting steady state `cali`.

```python
calibration = {'Y': 1., 'N': 1.0, 'K': 10., 'r': 0.0125, 'rstar': 0.0125, 'tot_wealth': 14,
               'delta': 0.02, 'pi': 0., 'kappap': 0.1, 'muw': 1.1, 'Bh': 1.04, 'Bg': 2.8,
               'G': 0.2, 'eis': 0.5, 'frisch': 1., 'chi0': 0.25, 'chi2': 2, 'epsI': 4,
               'omega': 0.005, 'kappaw': 0.1, 'phi': 1.5, 'nZ': 3, 'nB': 50, 'nA': 70, 'nK': 50,
               'bmax': 50, 'amax': 4000, 'kmax': 1, 'rho_z': 0.966, 'sigma_z': 0.92}

unknowns_ss = {'beta': 0.976, 'chi1': 6.5}
targets_ss = {'asset_mkt': 0., 'B': 'Bh'}

cali = hank_ss.solve_steady_state(calibration, unknowns_ss, targets_ss, solver='broyden_custom')
```

Verify solution, generate `ss` from dynamic DAG.

```python
ss =  hank.steady_state(cali)

print(f"Liquid assets: {ss['B']: 0.2f}")
print(f"Asset market clearing: {ss['asset_mkt']: 0.2e}")
print(f"Goods market clearing (untargeted): {ss['goods_mkt']: 0.2e}")
```

    Liquid assets:  1.04
    Asset market clearing:  8.22e-13
    Goods market clearing (untargeted):  3.29e-08

### 4.2 Linearized impulse responses

As before, we can compute the general equilibrium Jacobians $G$ which is sufficient to map any shock into impulse responses.

When the cost of computing a block Jacobian is non-trivial, it's a good idea to precompute it. We can supply block Jacobians for specific blocks via the `Js=` keyword argument. The precomputed Jacobians will only be used if they are **complete** (have all required inputs and outputs) and have the right **size** (truncation horizon).

```python
exogenous = ['rstar', 'Z', 'G']
unknowns = ['r', 'w', 'Y']
targets = ['asset_mkt', 'fisher', 'wnkpc']
T = 300

J_ha = hh_ext.jacobian(ss, inputs=['N', 'r', 'ra', 'rb', 'tax', 'w'], T=T)
G = hank.solve_jacobian(ss, unknowns, targets, exogenous, T=T, Js={'hh': J_ha})
```

The time saving from re-using the Jacobian of the household block is considerable.

```python
%time G = hank.solve_jacobian(ss, unknowns, targets, exogenous, T=T, Js={'hh': J_ha})
```

    Wall time: 433 ms

```python
%time G = hank.solve_jacobian(ss, unknowns, targets, exogenous, T=T)
```

    Wall time: 4.94 s

Note that some block Jacobians may be precomputed even if others are changing. For example, we can re-use `J_ha` while evalutating the model likelihood for 100,000 draws of price and wage adjustment cost.

When we're not planning to change any part of the model, it's even better to store the `H_U` directly. (To be precise, we store the LU-factorized version of the matrix, which facilitates operations with its inverse.) That way, we save time on computing and packing all the block Jacobians.

```python
from sequence_jacobian.classes import FactoredJacobianDict

H_U = hank.jacobian(ss, unknowns, targets, T=T, Js={'hh': J_ha})
H_U_factored = FactoredJacobianDict(H_U, T)
```

```python
%time G = hank.solve_jacobian(ss, unknowns, targets, exogenous, T=T, Js={'hh': J_ha}, H_U_factored=H_U_factored)
```

    Wall time: 343 ms

Let's plot some impulse responses:

```python
rhos = np.array([0.2, 0.4, 0.6, 0.8])
drstar = -0.0025 * rhos ** (np.arange(T)[:, np.newaxis])
dY = 100 * G['Y']['rstar'] @ drstar

plt.plot(dY[:21])
plt.title(r'Output response to 25 bp monetary policy shocks with $\rho=(0.2 ... 0.8)$')
plt.xlabel('quarters')
plt.ylabel('% deviation from ss')
plt.show()
```

### 4.3 Nonlinear impulse responses

Let's compute the nonlinear impulse response for the $\rho=0.6$ shock above.
- Don't forget to use the saved Jacobian.
- Note how to look up and change options specific to (block type, method) pairs.

```python
hank['pricing_solved'].solve_impulse_nonlinear_options
```

    {'tol': 1e-08, 'maxit': 30, 'verbose': True}

By default, `SolvedBlock.solve_impulse_linear` prints the error in each iteration (`verbose=True`). Let's turn this off for the internal solved blocks.

```python
td_nonlin = hank.solve_impulse_nonlinear(ss, unknowns, targets, {"rstar": drstar[:, 2]},
                                         Js={'hh': J_ha}, H_U_factored=H_U_factored,
                                         options={'pricing_solved': {'verbose': False},
                                                  'arbitrage_solved': {'verbose': False},
                                                  'labor_to_investment_combined_solved': {'verbose': False}})
```

    Solving Two-Asset HANK for ['r', 'w', 'Y'] to hit ['asset_mkt', 'fisher', 'wnkpc']
    On iteration 0
       max error for asset_mkt is 3.92E-06
       max error for fisher is 2.50E-03
       max error for wnkpc is 4.72E-08
    On iteration 1
       max error for asset_mkt is 2.66E-04
       max error for fisher is 1.55E-06
       max error for wnkpc is 2.15E-05
    On iteration 2
       max error for asset_mkt is 7.56E-06
       max error for fisher is 9.69E-08
       max error for wnkpc is 6.57E-07
    On iteration 3
       max error for asset_mkt is 4.01E-07
       max error for fisher is 2.24E-09
       max error for wnkpc is 1.64E-08
    On iteration 4
       max error for asset_mkt is 2.20E-08
       max error for fisher is 1.06E-10
       max error for wnkpc is 7.46E-10
    On iteration 5
       max error for asset_mkt is 1.23E-09
       max error for fisher is 5.44E-12
       max error for wnkpc is 3.72E-11

We see rapid convergence and mild nonlinearities in the solution.

```python
dY_nonlin = 100 * td_nonlin['Y']

plt.plot(dY[:21, 2], label='linear', linestyle='-', linewidth=2.5)
plt.plot(dY_nonlin[:21], label='nonlinear', linestyle='--', linewidth=2.5)
plt.title(r'Consumption response to 1% monetary policy shock')
plt.xlabel('quarters')
plt.ylabel('% deviation from ss')
plt.legend()
plt.show()
```

Alternatively, we can compute the impulse response to a version of the shock scaled down to 10% of its original size.

```python
td_nonlin = hank.solve_impulse_nonlinear(ss, unknowns, targets, {"rstar": 0.1 * drstar[:, 2]},
                                         Js={'hh': J_ha},
                                         options={'pricing_solved': {'verbose': False},
                                                  'arbitrage_solved': {'verbose': False},
                                                  'labor_to_investment_combined_solved': {'verbose': False}})
```

    Solving Two-Asset HANK for ['r', 'w', 'Y'] to hit ['asset_mkt', 'fisher', 'wnkpc']
    On iteration 0
       max error for asset_mkt is 3.92E-06
       max error for fisher is 2.50E-04
       max error for wnkpc is 4.72E-08
    On iteration 1
       max error for asset_mkt is 3.20E-06
       max error for fisher is 1.63E-08
       max error for wnkpc is 2.10E-06
    On iteration 2
       max error for asset_mkt is 7.71E-08
       max error for fisher is 2.53E-10
       max error for wnkpc is 5.83E-09
    On iteration 3
       max error for asset_mkt is 8.07E-10
       max error for fisher is 1.50E-12
       max error for wnkpc is 5.69E-11

```python
dY_nonlin = 100 * td_nonlin['Y']

plt.plot(0.1*dY[:21, 2], label='linear', linestyle='-', linewidth=2.5)
plt.plot(dY_nonlin[:21], label='nonlinear', linestyle='--', linewidth=2.5)
plt.title(r'Consumption response to 0.1% monetary policy shock')
plt.xlabel('quarters')
plt.ylabel('% deviation from ss')
plt.legend()
plt.show()
```

When the shock is this small, the linear and nonlinear impulse responses agree exactly. This is a good check, amid a highly complex model, that we didn't make any mistakes.