# Metrics in GPyTorch

In this notebook, we will see how to evaluate GPyTorch models with probabilistic metrics.

**Note:** It is encouraged to check the Simple GP Regression notebook first if not done already. We'll reuse most of the code from there.

We'll be modeling the function

\begin{align}
y &= \sin(2\pi x) + \epsilon \\
  \epsilon &\sim \mathcal{N}(0, 0.04)
\end{align}

```python
import math
import torch
import gpytorch
from matplotlib import pyplot as plt

%matplotlib inline
%load_ext autoreload
%autoreload 2
```

    The autoreload extension is already loaded. To reload it, use:
      %reload_ext autoreload

In the next cell, we set up the train and test data.

```python
# Training data is 100 points in [0,1] inclusive regularly spaced
train_x = torch.linspace(0, 1, 100)
# True function is sin(2*pi*x) with Gaussian noise
train_y = torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * math.sqrt(0.04)

test_x = torch.linspace(0, 1, 51)
test_y = torch.sin(test_x * (2 * math.pi)) + torch.randn(test_x.size()) * math.sqrt(0.04)
```

In the next cell, we define a simple GP regression model.

```python
# We will use the simplest form of GP model, exact inference
class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

# initialize likelihood and model
likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = ExactGPModel(train_x, train_y, likelihood)
```

Our model is ready for hyperparameter learning, but, first let us check how it performs on the test data.

```python
model.eval()
with torch.no_grad():
    untrained_pred_dist = likelihood(model(test_x))
    predictive_mean = untrained_pred_dist.mean
    lower, upper = untrained_pred_dist.confidence_region()

f, ax = plt.subplots(1, 1, figsize=(4, 3))
# Plot training data as black stars
ax.plot(train_x, train_y, 'k*')
# Plot predictive means as blue line
ax.plot(test_x, predictive_mean, 'b')
# Shade between the lower and upper confidence bounds
ax.fill_between(test_x, lower, upper, alpha=0.5)
ax.set_ylim([-3, 3])
ax.legend(['Observed Data', 'Mean', 'Confidence'], bbox_to_anchor=(1.6,1));
```

Visually, this does not look like a good fit. Now, let us train the model hyperparameters.

```python
# this is for running the notebook in our testing framework
import os
smoke_test = ('CI' in os.environ)
training_iter = 2 if smoke_test else 50

# Find optimal model hyperparameters
model.train()

# Use the adam optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)  # Includes GaussianLikelihood parameters

# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

for i in range(training_iter):
    # Zero gradients from previous iteration
    optimizer.zero_grad()
    # Output from model
    output = model(train_x)
    # Calc loss and backprop gradients
    loss = -mll(output, train_y)
    loss.backward()
    print('Iter %d/%d - Loss: %.3f   lengthscale: %.3f   noise: %.3f' % (
        i + 1, training_iter, loss.item(),
        model.covar_module.base_kernel.lengthscale.item(),
        model.likelihood.noise.item()
    ))
    optimizer.step()
```

    Iter 1/50 - Loss: 0.944   lengthscale: 0.693   noise: 0.693
    Iter 2/50 - Loss: 0.913   lengthscale: 0.644   noise: 0.644
    Iter 3/50 - Loss: 0.879   lengthscale: 0.598   noise: 0.598
    Iter 4/50 - Loss: 0.841   lengthscale: 0.555   noise: 0.554
    Iter 5/50 - Loss: 0.798   lengthscale: 0.514   noise: 0.513
    Iter 6/50 - Loss: 0.750   lengthscale: 0.475   noise: 0.474
    Iter 7/50 - Loss: 0.698   lengthscale: 0.439   noise: 0.437
    Iter 8/50 - Loss: 0.645   lengthscale: 0.404   noise: 0.402
    Iter 9/50 - Loss: 0.594   lengthscale: 0.372   noise: 0.369
    Iter 10/50 - Loss: 0.547   lengthscale: 0.342   noise: 0.339
    Iter 11/50 - Loss: 0.505   lengthscale: 0.315   noise: 0.310
    Iter 12/50 - Loss: 0.466   lengthscale: 0.291   noise: 0.284
    Iter 13/50 - Loss: 0.430   lengthscale: 0.271   noise: 0.259
    Iter 14/50 - Loss: 0.395   lengthscale: 0.254   noise: 0.236
    Iter 15/50 - Loss: 0.361   lengthscale: 0.240   noise: 0.215
    Iter 16/50 - Loss: 0.327   lengthscale: 0.229   noise: 0.196
    Iter 17/50 - Loss: 0.294   lengthscale: 0.220   noise: 0.178
    Iter 18/50 - Loss: 0.261   lengthscale: 0.214   noise: 0.162
    Iter 19/50 - Loss: 0.228   lengthscale: 0.209   noise: 0.147
    Iter 20/50 - Loss: 0.196   lengthscale: 0.206   noise: 0.134
    Iter 21/50 - Loss: 0.164   lengthscale: 0.205   noise: 0.122
    Iter 22/50 - Loss: 0.133   lengthscale: 0.204   noise: 0.110
    Iter 23/50 - Loss: 0.104   lengthscale: 0.205   noise: 0.100
    Iter 24/50 - Loss: 0.076   lengthscale: 0.207   noise: 0.091
    Iter 25/50 - Loss: 0.049   lengthscale: 0.210   noise: 0.083
    Iter 26/50 - Loss: 0.025   lengthscale: 0.214   noise: 0.076
    Iter 27/50 - Loss: 0.003   lengthscale: 0.218   noise: 0.069
    Iter 28/50 - Loss: -0.017   lengthscale: 0.223   noise: 0.063
    Iter 29/50 - Loss: -0.033   lengthscale: 0.228   noise: 0.058
    Iter 30/50 - Loss: -0.047   lengthscale: 0.233   noise: 0.053
    Iter 31/50 - Loss: -0.057   lengthscale: 0.239   noise: 0.049
    Iter 32/50 - Loss: -0.065   lengthscale: 0.243   noise: 0.045
    Iter 33/50 - Loss: -0.069   lengthscale: 0.248   noise: 0.042
    Iter 34/50 - Loss: -0.071   lengthscale: 0.252   noise: 0.039
    Iter 35/50 - Loss: -0.071   lengthscale: 0.255   noise: 0.037
    Iter 36/50 - Loss: -0.069   lengthscale: 0.257   noise: 0.035
    Iter 37/50 - Loss: -0.066   lengthscale: 0.257   noise: 0.033
    Iter 38/50 - Loss: -0.062   lengthscale: 0.257   noise: 0.032
    Iter 39/50 - Loss: -0.058   lengthscale: 0.255   noise: 0.031
    Iter 40/50 - Loss: -0.055   lengthscale: 0.252   noise: 0.030
    Iter 41/50 - Loss: -0.054   lengthscale: 0.248   noise: 0.029
    Iter 42/50 - Loss: -0.053   lengthscale: 0.243   noise: 0.029
    Iter 43/50 - Loss: -0.053   lengthscale: 0.237   noise: 0.029
    Iter 44/50 - Loss: -0.055   lengthscale: 0.230   noise: 0.029
    Iter 45/50 - Loss: -0.057   lengthscale: 0.223   noise: 0.029
    Iter 46/50 - Loss: -0.060   lengthscale: 0.216   noise: 0.029
    Iter 47/50 - Loss: -0.063   lengthscale: 0.208   noise: 0.030
    Iter 48/50 - Loss: -0.066   lengthscale: 0.202   noise: 0.030
    Iter 49/50 - Loss: -0.068   lengthscale: 0.196   noise: 0.031
    Iter 50/50 - Loss: -0.070   lengthscale: 0.190   noise: 0.032

In the next cell, we reevaluate the model on the test data.

```python
model.eval()
with torch.no_grad():
    trained_pred_dist = likelihood(model(test_x))
    predictive_mean = trained_pred_dist.mean
    lower, upper = trained_pred_dist.confidence_region()

f, ax = plt.subplots(1, 1, figsize=(4, 3))
# Plot training data as black stars
ax.plot(train_x, train_y, 'k*')
# Plot predictive means as blue line
ax.plot(test_x, predictive_mean, 'b')
# Shade between the lower and upper confidence bounds
ax.fill_between(test_x, lower, upper, alpha=0.5)
ax.set_ylim([-3, 3])
ax.legend(['Observed Data', 'Mean', 'Confidence'], bbox_to_anchor=(1.6,1));
```

Now our model seems to fit well on the data. It is not always possible to visually evaluate the model in high dimensional cases. Thus, now we evaluate the models with help of probabilistic metrics. We have saved predictive distributions from untrained and trained models as `untrained_pred_dist` and `trained_pred_dist` respectively.

## Negative Log Predictive Density (NLPD)

Negative Log Predictive Density (NLPD) is the most standard probabilistic metric for evaluating GP models. In simple terms, it is negative log likelihood of the test data given the predictive distribution. It can be computed as follows:

```python
init_nlpd = gpytorch.metrics.negative_log_predictive_density(untrained_pred_dist, test_y)
final_nlpd = gpytorch.metrics.negative_log_predictive_density(trained_pred_dist, test_y)

print(f'Untrained model NLPD: {init_nlpd:.2f}, \nTrained model NLPD: {final_nlpd:.2f}')
```

    Untrained model NLPD: 0.88,
    Trained model NLPD: -0.21

## Mean Standardized Log Loss (MSLL)

This metric computes average negative log likelihood of all test points w.r.t their univariate predicitve densities. For more details, see "page No. 23, Gaussian Processes for Machine Learning, Carl Edward Rasmussen and Christopher K. I. Williams, The MIT Press, 2006. ISBN 0-262-18253-X"

```python
init_msll = gpytorch.metrics.mean_standardized_log_loss(untrained_pred_dist, test_y)
final_msll = gpytorch.metrics.mean_standardized_log_loss(trained_pred_dist, test_y)

print(f'Untrained model MSLL: {init_msll:.2f}, \nTrained model MSLL: {final_msll:.2f}')
```

    Untrained model MSLL: 0.82,
    Trained model MSLL: -0.49

It is also possible to calculate the quantile coverage error with `gpytorch.metrics.quantile_coverage_error` function.

```python
quantile = 95
qce = gpytorch.metrics.quantile_coverage_error(trained_pred_dist, test_y, quantile=quantile)
print(f'Quantile {quantile}% Coverage Error: {qce:.2f}')
```

    Quantile 95% Coverage Error: 0.01

## Mean Squared Error (MSE)

Mean Squared Error (MSE) is the mean of the squared difference between the test observations and the predictive mean. It is a well-known conventional metric for evaluating regression models. However, it can not take uncertainty into account unlike NLPD, MLSS and ACE.

```python
init_mse = gpytorch.metrics.mean_squared_error(untrained_pred_dist, test_y, squared=True)
final_mse = gpytorch.metrics.mean_squared_error(trained_pred_dist, test_y, squared=True)

print(f'Untrained model MSE: {init_mse:.2f}, \nTrained model MSE: {final_mse:.2f}')
```

    Untrained model MSE: 0.20,
    Trained model MSE: 0.04

## Mean Absolute Error (MAE)

Mean Absolute Error (MAE) is the mean of the absolute difference between the test observations and the predictive mean. It is less sensitive to the outliers than MSE.

```python
init_mae = gpytorch.metrics.mean_absolute_error(untrained_pred_dist, test_y)
final_mae = gpytorch.metrics.mean_absolute_error(trained_pred_dist, test_y)

print(f'Untrained model MAE: {init_mae:.2f}, \nTrained model MSE: {final_mae:.2f}')
```

    Untrained model MAE: 0.38,
    Trained model MSE: 0.16


---
title: "Convolutions of random variables — Part 1: Introduction"
date: 2021-03-14
categories: [statistics, convolutions]
tags: [statistics, convolutions, uniform distribution, sums of random variables]
preview_image: /enrightward.github.io/assets/img/convolutions/convolution-formula.png
---

## 1. Introduction

Suppose that $X$ and $Y$ are two continuous random variables. Given their PDFs PDFs $f_{X}$ and $f_{Y}$, how do you compute the $f_{X + Y}$ of the sum of $X+Y$? In this post, we'll define, compute examples of and prove some facts about the _convolution_ $f_{X} * f_{Y}$ of $f_{X}$ and $f_{Y}$, itself a PDF, and moreover equal to $f_{X + Y}$, provided $X$ and $Y$ are independent. A minor adaptations yield a method for computing $f_{Y - X}$. We will explicitly compute the PDFs $f_{X+Y}$ and $f_{X-Y}$ for $X, Y \sim U(0, 1)$. Although doing all the algebra rigorously is quite fiddly, it turns out there is a nice way of visualising these convolutions is the area of intersection of two colliding rectanlges in the plane, which I [learnt about on Wikipedia](https://en.wikipedia.org/wiki/Convolution). I have also made my own version of this visualisation below, using the [desmos graphical calculator app](https://www.desmos.com/).

**Note:** This post is quite technical and calculation-heavy. To get a high-level feel for convolutions, together with a discussion of convolutional neural networks, see for instance [this lovely post from Christopher Olah](https://colah.github.io/posts/2014-07-Understanding-Convolutions/).

## 2. Convolutions

As motivation, suppose first $X$ and $Y$ are independent discrete random variables, taking values on the integers, and let $Z := X + Y$. Then the PDF of $Z$ is given by:

\begin{align}
\mathbb{Pr}(Z = z) = \sum_{k \in \mathbb{Z}} \mathbb{Pr}(X = k) \, \mathbb{Pr}(Y = z-k).
\end{align}

Now assume $X$ and $Y$ are continuous, independent random variables with PDFs $f_{X}$ and $f_{Y}$, taking values on the whole real line. What is the analogue of the above summation formula? We have for any $z \in \mathbb{R}$ that:

\begin{align}
F_{Z}(z) = \mathbb{Pr}(Z \le z) = \int \int_{S(z)} f_{X}(x) \, f_{Y}(y) \, dy \, dx,
\end{align}

where $S(z) := \{ (x, y) \in \mathbb{R}^{2} \mid x + y \le z \}$. We can integrate over $S(z)$ by allowing $x$ to vary freely, and $y$ to vary over the half-open interval $(-\infty, z-x]$:

\begin{align}
F_{Z}(z) = \int_{-\infty}^{\infty} f_{X}(x) \left\{ \int_{-\infty}^{z-x} f_{Y}(y) \, dy \right\} dx
= \int_{-\infty}^{\infty} f_{X}(x) \, F_{Y}(z-x) \, dx.
\end{align}

To pass to the PDF $f_{Z}(z)$, we must take the derivative of $F_{Z}(z)$ with respect to $z$. Using the [Leibniz integral rule](https://en.wikipedia.org/wiki/Leibniz_integral_rule) to interchange the order of differentiation and integration, we have

\begin{align}
f_{Z}(z) &= F_{Z}^{\prime}(z) =
\frac{d}{dz} \left( \int_{-\infty}^{\infty} f_{X}(x) \, F_{Y}(z-x) \, dx \right) \newline
&= \int_{-\infty}^{\infty} \frac{d}{dz} \left( f_{X}(x) \, F_{Y}(z-x) \, dx \right) \newline
&= \int_{-\infty}^{\infty} f_{X}(x) \, f_{Y}(z-x) \, dx,
\end{align}

where the equality uses the fact that $f_{X}(x)$ is independent of $z$, and also that the derivative of $F_{Y}(z-x)$ with respect to $z$ is $f_{Y}(z-x)$, by the chain rule. We refer to the integral as the _convolution_ of the functions $f_{X}$ and $f_{Y}$. In this language, we have shown that the PDF of the sum of two independent random variables is given by the convolution of their PDFs.

## 3. PDFs of sums and differences of two random variables

We showed in the previous section that

\begin{align}
f_{X+Y}(z) = \int_{-\infty}^{\infty} f_{X}(x) \, f_{Y}(z-x) \, dx, \newline
F_{X+Y}(z) = \int_{-\infty}^{\infty} f_{X}(x) \, F_{Y}(z-x) \, dx.
\end{align}

A similar argument, starting with the variable $Z = Y - X$, shows that

\begin{align}
f_{Y-X}(z) &= \int_{-\infty}^{\infty} f_{X}(x) \, f_{Y}(z+x) \, dx, \newline
F_{X-Y}(z) &= \int_{-\infty}^{\infty} f_{X}(x) \, F_{Y}(z+x) \, dx.
\end{align}

In the next section, we compute examples these PDFs in the case where $X, Y \sim U(0, 1)$.

## 4. The sum of two uniform random variables

Suppose now that $X, Y \sim U(0, 1)$. To calculate the sum and difference of their PDFs, we introduce "indicator function" notation: For any element $x \in \mathbb{R}$ and subset $S \subseteq \mathbb{R}$, define

\begin{align}
{\mathbb{1}}_{S}(x) := \left\{ \begin{array}{ll}
1, & \textrm{if } x \in S,  \newline
0, & \textrm{otherwise.}
\end{array} \right.
\end{align}

Then for the uniform random variables $X$ and $Y$,

\begin{align}
f_{X}(x) &= f_{Y}(x) = {\mathbb{1}}_{\lbrack 0, 1 \rbrack}(x), \newline
F_{X}(x) &= F_{Y}(x) = x \cdot {\mathbb{1}}_{\lbrack 0, 1 \rbrack}(x).
\end{align}

We now compute $f_{X+Y}(z)$. By the formula above, $f_{X}(x) = \mathbb{1}_{[0, 1]}(x)$, but what is $f_{Y}(z -x)$? To answer this, note that $0 \le z - x \le 1$ if and only if $z - 1 \le x \le z$, and so:

\begin{align}
f_{Y}(z - x) = {\mathbb{1}}_{[0, 1]}(z - x) = {\mathbb{1}}_{[z-1, z]}(x).
\end{align}

Using the convolution formula, we now have:

\begin{align}
f_{X+Y}(z) &= \int_{-\infty}^{\infty} f_{X}(x) \, f_{Y}(z-x) \, dx \newline
&= \int_{-\infty}^{\infty} {\mathbb{1}}_{[0, 1]}(x) \, {\mathbb{1}}_{[z-1, z]}(x) \, dx \newline
&= \int_{-\infty}^{\infty} {\mathbb{1}}_{[0, 1] \cap [z-1, z]}(x) \, dx.
\end{align}

The last equality holds because the product $\mathbb{1}_{S}(x) \cdot \mathbb{1}_{T}(x)$ equals $1$ if and only if $x$ belongs to both $S$ and $T$, and is zero otherwise, showing the general identity $\mathbb{1}_{S} \cdot \mathbb{1}_{T} = \mathbb{1}_{S \cap T}$. Set now $I: = [0, 1]$ and $J := [z-1, z]$ to ease notation. Clearly, $I \cap J$ is an interval whose endpoints, hence length, depend on $z$. Concretely:

\begin{align}
I \cap J = \left\{ \begin{array}{ll}
\emptyset, & \textrm{if } z \le 0, \; \textrm{or} \; z \ge 2, \newline
[z-1, 1], & \textrm{if } 1 \le z \le 2, \newline
[0, z], & \textrm{if } 0 \le z \le 1.
\end{array} \right.
\end{align}

Geometrically, the inequality conditions correspond to the fact that $I$ and $J$ either (i) Don't overlap at all, or else they overlap partially, and (ii) Either and $I$ is on the left, or else (iii) $J$ is on the left. For an interval $S = (a, b)$, the integral over the whole real line of $\mathbb{1}_{S}(x)$ is equal to the length $b-a$ of $S$. It follows that

\begin{align}
f_{X+Y}(z) &= \int_{-\infty}^{\infty} {\mathbb{1}}_{I \cap J}(x) \, dx =
\left\{ \begin{array}{ll}
0, & \textrm{if } z \le 0, \; \textrm{or} \; z \ge 2, \newline
2-z, & \textrm{if } 1 \le z \le 2, \newline
z, & \textrm{if } 0 \le z \le 1.
\end{array} \right.
\end{align}

We plot this, using the `matplotlib` code in the cell below. The function `graph_piecemeal` is (only slightly) adapted from code posted by stackoverflow user [rputikar](https://stackoverflow.com/users/1888184/rputikar), in response to [this question](https://stackoverflow.com/questions/14000595/graphing-an-equation-with-matplotlib).

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

def _set_ax_colour(ax, colour):
    """Set the spine, label and ticks of the
    axis colour to `colour`"""
    ax.spines['bottom'].set_color(colour)
    ax.spines['top'].set_color(colour)
    ax.spines['left'].set_color(colour)
    ax.spines['right'].set_color(colour)
    ax.xaxis.label.set_color(colour)
    ax.yaxis.label.set_color(colour)
    ax.title.set_color(colour)
    ax.tick_params(axis='x', colors=colour)
    ax.tick_params(axis='y', colors=colour)

def graph_piecemeal(boundaries, formulae, resolution, figsize, title):
    """Display a piecemeal graph defined by the list of
    strings `formulae`, separated by the list of numbers
    `boundaries`. Each piece of the x_axis is broken into
    `resolution` parts for graphing. There must be one less
    formula than boundary."""
    assert len(formulae) == len(boundaries) - 1

    # setup axes for dark background
    fig, ax = plt.subplots(figsize=figsize)
    _set_ax_colour(ax, 'white')
    ax.set_title(title, fontsize=12)

    # define `intervals` as consecutive boundary points.
    intervals = [np.linspace(a, b, resolution) for a, b in zip(boundaries, boundaries[1:])]
    assert len(intervals) == len(formulae)

    # graph pieces (= (interval, formula) pairs) in a loop
    for interval, formula in zip(intervals, formulae):
        x = np.array(interval)
        y = eval(formula)
        plt.plot(x, y, color='blue')

    plt.show()
```

```python
# Define the boundaries between the formulae. They correspond to the four regions
# z < 0, 0 < z < 1, 1 < z < 2, z > 2 in the definition of f_{X+Y}
boundaries = [-1, 0, 1, 2, 3]

# Define the different formulae for each region
formulae = ['x*0', 'x', '2 - x', 'x*0']

resolution = 5
figsize = (10, 6)
title = 'PDF of X+Y, where X and Y uniform'
graph_piecemeal(boundaries, formulae, resolution, figsize, title)
```

##  5. The difference of two uniform random variables

We now compute the PDF of $Z = Y - X$. The method parallels that of the sum $X+Y$, so we skip some steps. As noted above, the general formula is:

\begin{align}
f_{Y-X}(z) &= \int_{-\infty}^{\infty} f_{X}(x) \, f_{Y}(z+x) \, dx.
\end{align}

This time, to find $f_{Y}(z+x)$ we note that $0 \le z + x \le 1$ if and only if $-z \le x \le 1 - z$, and so:

\begin{align}
f_{Y}(z + x) = \mathbb{1}_{[0, 1]}(z + x) = \mathbb{1}_{[-z, 1-z]}(x).
\end{align}

The convolution now reads:

\begin{align}
f_{X+Y}(z) = \int_{-\infty}^{\infty} \mathbb{1}_{[0, 1] \cap [-z, 1-z]}(x) \, dx.
\end{align}

If now $I: = [0, 1]$ and $J := [-z, 1-z]$, then  $I \cap J$ depends on $z$ as follows:

\begin{align}
I \cap J = \left\{ \begin{array}{ll}
\emptyset, & \textrm{if } z \le -1, \; \textrm{or} \; z \ge 1, \newline
[-z, 1], & \textrm{if } -1 \le z \le 0, \newline
[0, 1-z], & \textrm{if } 0 \le z \le 1.
\end{array} \right.
\end{align}

It follows that:

\begin{align}
f_{Y-X}(z) &= \int_{-\infty}^{\infty} \mathbb{1}_{I \cap J}(x) \, dx =
\left\{ \begin{array}{ll}
0, & \textrm{if } z \le -1, \; \textrm{or} \; z \ge 1, \newline
1 + z, & \textrm{if } -1 \le z \le 0, \newline
1 - z, & \textrm{if }  0 \le z \le 1.
\end{array} \right.
\end{align}

As above, let's display the graph of $f_{Y-X}(z)$.

```python
# Define the boundaries between the formulae. They correspond to the four regions
# z < 0, 0 < z < 1, 1 < z < 2, z > 2 in the definition of f_{X+Y}
boundaries = [-2, -1, 0, 1, 2]

# Define the different formulae for each region
formulae = ['x*0', '1+x', '1-x', 'x*0']

resolution = 5
figsize = (10, 6)
title = 'PDF of Y-X, where X and Y uniform'
graph_piecemeal(boundaries, formulae, resolution, figsize, title)
```

## 6. Intuition about convolutions of uniform PDFs

So $f_{Y-X}(x)$ is the same shape as $f_{X+Y}(x)$, but shifted one unit to the left. Could we have seen this in advance, you ask? Why, yes! It follows from the fact that if $X \sim U(0, 1)$, then $X-1 = -X$ as random variables, since:

\begin{align}
f_{-X}(x) = \mathbb{1}_{[-1, 0]}(x) = f_{X-1}(x),
\end{align}

implying also $Y-X=X+Y-1$. It follows:

\begin{align}
f_{Y-X}(x) = f_{X+Y-1}(x) = f_{X+Y}(x+1),
\end{align}

where the last equality holds because $f_{Z-a}(z) = f_{Z}(z+a)$ for any random variable $Z$ and constant $a$. Note that the equivalence of $-X$ and $X-1$ is also quite intuitive, for it says that if you sample a real number $x$ uniformly from $[0, 1]$, then either subtracting $1$ to get $x-1$, or negating to get $-x$ are probabilistically the same: You land (uniformly) in the range $[-1, 0]$.

You might also ask: Could have predicted the "pyramid" shape of $f_{X+Y}(z)$ in the first place? The answer is also yes, I think, at least once you know the formula:

\begin{align}
f_{X+Y}(z) = \int_{-\infty}^{\infty} \mathbb{1}_{[0, 1] \cap [-z, 1-z]}(x) \, dx.
\end{align}

This integral invites us to consider the indicator function of the set $I \cap J$, where $I := [0, 1]$ and $J := [-z, 1-z]$. This function is constant and equal to one on $I \cap J$, and is equal to zero every where else, so the value of this integral at any given $z$ is the length of $I \cap J$, or equivalenty, the area under the curve $y=1$ on $I \cap J$. The first interval, $I$, is fixed, and both have constant length $1$, but $J$ slides around as we vary $z$. So, computing the length of $I \cap J$ for any particular $z$ is the same as understanding how much $I$ and $J$ overlap for $z$. The rate of change of $f_{X+Y}(z)$, as $z$ changes, is simply the rate of change of this overlap area. Since the overlapping areas are squares, overlap increases linearly until it is complete (corresponding to the apex of the pyramid), then decreases linearly at the same rate. This is shown below in the following animation, made in [desmos](https://www.desmos.com/calculator/):

## 7. Roundup

We defined, for two continuous random variables $X$ and $Y$, the _convolution_ $f_{X} * f_{Y}$ of $f_{X}$ and $f_{Y}$. The convolution is itself a PDF. In the case where $X$ and $Y$ are independent, $f_{X} * f_{Y}$ is equal to the PDF $f_{X + Y}$ of the sum of $X+Y$. We then computed the PDFs for $f_{X+Y}$ and $f_{X-Y}$ explicitly. Algebraiclly, the answer in each case was an ugly-looking four-part piecemeal function. Visually, however, these PDFs are two, translated copies of the same "pyramid" function. Using some basic reasoning and a visualisation of the changing area of colliding rectangles in the plane, we were saw where the pyramid came from. Starting from the observation that if $X \sim U(0,1)$, then $1-X=-X$, we were able to explain the translation. In the [next post](https://enrightward.github.io/enrightward.github.io/posts/convolutions-part-2/), we'll use convolutions to explicitly compute and visualise the PDF $f_{aX+bY}$ of a linear combination of $X \sim U(0,1)$.


# Computational Astrophysics
## Elliptic PDEs. Examples

---
## Eduard Larrañaga

Observatorio Astronómico Nacional\
Facultad de Ciencias\
Universidad Nacional de Colombia

---

### About this notebook

In this notebook we present some of the techniques used to solve the Poisson equation.

`A. Garcia. Numerical Methods for Physics. (1999). Chapter 6 - 7 `

---

##  Poisson Equation as an ODE Example. Pre-Supernova Model

A supernova is a process in which a dying star liberates a huge quantity of energy and mass.
 In this section we will solve Poisson equation to obtain the gravitational potential inside the mass distribution of a dying star.

In the data file `presupernovaStar.dat` we find the information of a pre-supernova stellar model. The columns in this file are:

1. First column: grid point index
2. Second column: enclosed mass of the star ($\text{gr}$)
3. Third column: radius of the distribution ($\text{cm}$)
4. Fourth column: temperature ($\text{K}$)
5. Fifth column: mass density ($\text{gr/cm}^3$)

**The other columns in the data file are not important for this exercise**

First, we will read the data and make a plot of some of the variables in order to visualize the information.

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

# Load data of rho vs. radius
(index, mass, dataradius, temperature, density, var1, var2, var3) = np.loadtxt('presupernovaStar.dat', unpack = True)

fig, ax = plt.subplots(3,1, figsize=(7,10))
ax[0].plot(np.log10(dataradius), density)
ax[0].set_ylabel(r'$\rho (r)$')

ax[1].plot(np.log10(dataradius), mass)
ax[1].set_ylabel(r'$M (r)$')

ax[2].plot(np.log10(dataradius), temperature)
ax[2].set_ylabel(r'$T (r)$')
ax[2].set_xlabel(r'$\log_{10} [r]$')

plt.show()
```

The first radius in the data file is

```python
dataradius[0]
```

    4296013.288986084

---

We will restrict the data from this value up to a radius of $r < 10^9 \,\mathrm{cm}$. Hence we will define an equidistant grid and we will interpolate the values of the density in this grid using the function `numpy.interp()`.

Information of this function can be found at
https://numpy.org/doc/stable/reference/generated/numpy.interp.html

```python
# Equidistant grid definition
n = 1000
r_0 = dataradius[0]
r_f = 1e9
radius = np.linspace(r_0, r_f, n)

# Interpolation of the density in the equidistant grid
rho = np.interp(radius, dataradius, density)

# Plot of the data points together with the interpolated density
plt.figure(figsize=(7,5))
plt.scatter(np.log10(dataradius), density, color='black')
plt.plot(np.log10(radius), rho, 'r')

plt.xlabel(r'$\log_{10} [r]$')
plt.ylabel(r'$\rho(r)$')
plt.show()
```

Now that we have the interpolated density, we implement the solution of the Poisson equation with the ODE method.

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

# Global constants (cgs units)
G = 6.67e-8

# ODE definition
def ODE(r, q0, rho0):
    '''
    ------------------------------------------
    ODE(r,q0)
    ------------------------------------------
    ODEs system for the Poisson Equation
    Arguments:
    r: radius
    q0: numpy array with the initial condition
        data:
        q0[0] = Phi
        q0[1] = dPhi/dr
    ------------------------------------------
    '''
    Phi = q0[0]
    z = q0[1]
    f = np.zeros(2)
    f[0] = z
    f[1] = 4*np.pi*G*rho0 - 2*z/r
    return f

def FEuler(h, r0, q0, rho0):
    '''
    ------------------------------------------
    FEuler(h, r0, q0, rho0)
    ------------------------------------------
    Forward Euler's method for solving a ODEs
    system.
    Arguments:
    h: stepsize for the iteration
    r0: independent parameter initial value
    q0: numpy array with the initial values of
        the functions in the ODEs system
    ------------------------------------------
    '''
    f = ODE(r0, q0, rho0)
    q1 = q0 + h*f
    return q1

# Constant stepsize defined by the number of steps in the grid
h = radius[1] - radius[0]

# Arrays to store the solution
Q = np.zeros([2,n]) # Euler's Method information

# Initial Conditions
Q[0,0] = 0. # Initial guess for the value of the potential at the center
Q[1,0] = 0. # Derivative of the potential at the center

# Main loops for solving the problem
for i in range(1,n):
    q0 = Q[:,i-1]
    qf = FEuler(h, radius[i-1], q0, rho[i-1])
    Q[:,i] = qf[:]

plt.figure(figsize=(7,5))
plt.plot(radius,Q[0])
plt.xlabel(r'$r$')
plt.ylabel(r'$\Phi(r)$')
plt.show()
```

---

## Poisson Equation solved by the Matrix Method. Homogenoeous Sphere

In the matrix method, the Poisson Equation is discretized with centered differences. Imposing the boundary condition $\frac{\partial \Phi}{ \partial r} = 0$ gives

\begin{equation}
\Phi_{-1} = \Phi_{0}\,\,.
\end{equation}

Then, the Poisson equation is written as linear system

\begin{equation}
J \mathbf{\Phi} = \mathbf{b}\,\,,
\label{eq:pde_poisson2}
\end{equation}

where $\Phi$ = $(\Phi_0, \cdots , \Phi_{n-1})^T$ (for a grid with $n$
points labeled $0$ to $n-1$) and  $\mathbf{b} = 4\pi G
(\rho_0, \cdots , \rho_{n-1})^T$.

The matrix $J$ has tri-diagonal form and can be explicitely given as

1. $i=j=0$:
\begin{equation}
J_{00} = - \frac{1}{(\Delta r)^2} - \frac{1}{r_0 \Delta r}\,\,,
\end{equation}

2. $i=j$:
\begin{equation}
J_{ij} = \frac{-2}{(\Delta r)^2}\,\,,
\end{equation}

3. $i+1=j$:
\begin{equation}
J_{ij} = \frac{1}{(\Delta r)^2} + \frac{1}{r_i \Delta r}\,\,,
\end{equation}

4. $i-1=j$:
\begin{equation}
J_{ij} = \frac{1}{(\Delta r)^2} - \frac{1}{r_i \Delta r}\,\,.
\end{equation}

Now we solve the pre-spuernova stellar model problem using this description.

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

# Global constants (cgs units)
G = 6.67e-8

# Constant stepsize defined by the number of steps in the grid
dr = radius[1] - radius[0]

b = 4*np.pi*G*rho
J = np.zeros([n,n])

# Definition of the matrix J
J[0,0] = -1/dr**2 - 1/(dr*radius[0])

for i in range(1,n):
    J[i,i] = -2/dr**2
    J[i,i-1] = 1/dr**2 - 1/(dr*radius[i])

for i in range(0,n-1):
    J[i,i+1] = 1/dr**2 + 1/(dr*radius[i])

# Forward Elimination
d = np.zeros(n)
w = np.zeros(n)

d[0] = J[0,0]
w[0] = b[0]/d[0]

for i in range(1,n):
    d[i] = J[i,i] - J[i,i-1]*J[i-1,i]/d[i-1]
    w[i] = (b[i] - w[i-1]*J[i,i-1])/d[i]

# Back-Substitution
y = np.zeros(n)

y[-1] = w[-1]
for i in range(n-2,-1,-1):
    y[i] = w[i] - y[i+1]*J[i,i+1]/d[i]

plt.figure(figsize=(7,5))
plt.plot(radius, y)
plt.xlabel(r'$r$')
plt.ylabel(r'$\Phi(r)$')
plt.show()

```

```python

```


# Staff Scheduling For Stochastic Demands
In this article, I'll introduce a simple and efficient model for solving a staff scheduling problem.

## Motivation and Problem Definition
Imagine that Alice owns a restaurant. Her awesome fashion restaurant aims to provide great dishes all day long, and she plans to hire some people to run the restaurant. Since this is an awesome fashion restaurant, Alice requires that all chefs have to be professionally trained, and so do the servers. Some employees are more properly trained so that they can be either a chef or a server, but they can't do both at the same time since that's unsanitary and keep changing jobs make the employees angry. In fact, since employees hate changing jobs all the time, we want to add some penalties to compensate for the loss in employees' emotions. The whole day is divided into 8 shifts, each shift is 6-hour-long, and the nearby two shifts have three-hour-long overlaps.

While this is a hypothesis scenario for staff scheduling, the model I'll represent can use for more scenarios. For example, consider, again, a restaurant that provides two special ingredients that require some certain equipment to cook, but a specific type of equipment can cook both, such as the rice cooker (?). Moreover, those pieces of equipment can't be used continuously for too long, or they may break early.

## Basic Model
Here I'll present a simple model to solve the problem above. Before diving into the model itself, we need to define some parameters:

| Set  | Definition                                  |
| ---- | ------------------------------------------- |
| $I$  | Schedules of when to work and when to leave |
| $J$  | Set of jobs                                 |
| $T$  | Time period in a day, $T = {0,...,23}$      |

| Parameter | Definition                                                   |
| --------- | ------------------------------------------------------------ |
| $A_{it}$  | Schedule $i$ include period $t$ or not (binary)              |
| $D_{jt}$  | Demand of worker needed on job $j$ at period $t$\\           |
| $E_{j}$   | Number of workers who have the skill for job $j$             |
| $B$       | Number of workers who have both skills, $B < E_j$            |
| $C_{t}$   | Cost of hiring a worker to work at period $t$                |
| $F$       | Cost of letting a worker change jobs at the middle of a day  |

| Decision Variable | Definition                                                   |
| ----------------- | ------------------------------------------------------------ |
| $x_{i}$           | Number of worker work on schedule $i$ (integer)              |
| $y_{jt}$          | Number of worker work on job $j$ on period $t$ (integer)     |
| $z_{jt}$          | Difference of number of workers on job $j$ at the start of period $t$ (integer) |

The definition above is quite straightforward, I believe. Then, we'll formulate the model step by step.
First of all, Alice's main goal is to minimize the personnel costs of running her restaurant, so the objective function would be the cost of hiring her employees and the cost of making her employees change their jobs.

$$
\min \sum_{j \in J} \sum_{t \in T} y_{jt} C_{t} + \sum_{t \in T}\sum_{j \in J} z_{jt} F
$$

Then, at each period, the number of staff on duty has to equal the number of staff assigned to the schedules which contain that period.
$$
\sum_{j \in J} y_{jt} = \sum_{i \in I} x_{i} A_{it} \quad \forall t \in T
$$
Okay, I know this might sound redundant or useless, but I'll show how it can be useful in the extended problem.

Most importantly, the number of staff on duty must be larger or equal to the demands, otherwise, the staff would be too busy and resign.
$$
y_{jt} \geq D_{jt} \quad \forall j \in J,\ t \in T
$$

At each period, the number of staff on duty must be smaller than the total staff the restaurant hired.
$$
\sum_{i \in I} x_{i} \leq \sum_{j \in J} E_j  - B
$$

Similarly, the number of staff on duty with a certain skill should not exceed the number of employees with that skill. For example, Alice can't assign 20 chefs to work if she only hires 10 chefs.
$$
y_{jt} \leq E_j \quad \forall j \in J,\ t \in T
$$

Finally, we want to define $z$ as the number of job changes between periods.
$$
z_{jt} = \mid y_{jt} - y_{j,t-1} \mid \quad \forall j \in J,\ t \in \{1,...23\}
$$

Then, the range of all variables must be non-negative.
$$
x_i, y_{jt}, z_{jt}, r_{jt} \geq 0 \quad \forall i,j,t
$$

```python
from IPython.display import display, Math, Latex
```

The complete model would be:

\begin{align}
    \min \quad
    &  \sum_{j \in J} \sum_{t \in T} y_{sjt} C_{t}  % hiring cost
        + \sum_{t \in T}\sum_{j \in J} z_{sjt} F \\ % change job cost
    % constraints
    \mbox{s.t.} \quad
    % 分配到工作的人數要符合上班人數
    & \sum_{j \in J} y_{jt} = \sum_{i \in I} x_{i} A_{it}
        & \forall t \in T\\
    % 值班人數要滿足需求
    & y_{jt} \geq D_{jt}
        & \forall j \in J,\ t \in T\\
    % 總上班人數 <= 總員工數
    & \sum_{i \in I} x_{i} \leq \sum_{j \in J} E_j  - B\\
    % 任一時段，每個技能的值班人數 <= 總持有技能人數
    & y_{jt} \leq E_j
        & \forall j \in J,\ t \in T\\
    % 任一時段，外包人數 <= 外包上限
    & l_{jt} \leq L_{jt}
        & \forall j \in J,\ t \in T\\
    % 定義 z: 轉換次數
    & z_{jt} = \mid y_{jt} - y_{j,t-1} \mid
        & \forall j \in J,\ t \in \{1,...23\}\\  % 中途轉換次數(取絕對值)
    & x_i, y_{jt}, z_{jt}, r_{jt} \geq 0 & \forall i,j,t
\end{align}

## Weighted Sum Method
So far so good, until one day Alice realized that the restaurant's demand can't be accurately predicted. In particular, she found that the demand is quite low on rainy days because people would rather stay home, and the demand is high on pleasant days. Fortunately, the weather forecast is quite well so that she knows the probability distribution of the weather tomorrow. She then wants to consider all scenarios, instead of betting on which scenario would realize, and decide how many employees have to work tomorrow.

We'll start with adding a set $S$ indicating all scenarios and change the definition of the parameters and decision variables.

| Set  | Definition                                  |
| ---- | ------------------------------------------- |
| $I$  | Schedules of when to work and when to leave |
| $J$  | Set of jobs, $|J| = 2$                      |
| $T$  | Time period in a day, $T = {0,...,23}$      |

| Parameter | Definition                                                   |
| --------- | ------------------------------------------------------------ |
| $A_{it}$  | Schedule $i$ include period $t$ or not (binary)              |
| $D_{jt}$  | Demand of worker needed on job $j$ at period $t$\\           |
| $E_{j}$   | Number of workers who have the skill for job $j$             |
| $B$       | Number of workers who have both skills, $B < E_j$            |
| $C_{t}$   | Cost of hiring a worker to work at period $t$                |
| $F$       | Cost of letting a worker change jobs at the middle of a day  |

| Decision Variable | Definition                                                   |
| ----------------- | ------------------------------------------------------------ |
| $x_{i}$           | Number of worker work on schedule $i$ (integer)              |
| $y_{sjt}$         | Number of worker work on job $j$ on period $t$ in scenario $s$ (integer) |
| $z_{sjt}$         | Number of worker changing job at the start of period $t$ in scenario $s$ (integer) |

In the objective functions, we want to take all scenarios into consideration. More specifically, we want to optimize the expected value of the cost in all scenarios.
$$
\min \sum_{s \in S} P_s  (\sum_{j \in J} \sum_{t \in T} y_{jt} C_{t} + \sum_{t \in T}\sum_{j \in J} z_{jt} F)
$$
The remaining is quite the same, except we now need to consider more scenarios. The full model is shown as below:

\begin{align}
    % objective: min cost
    \min \quad
    & \sum_{s \in S} P_s
    \left( \sum_{j \in J} \sum_{t \in T} y_{sjt} C_{t}  % hiring cost
        + \sum_{t \in T}\sum_{j \in J} z_{sjt} F \right) \\ % change job cost
    % constraints
    \mbox{s.t.} \quad
    % 分配到工作的人數要符合上班人數
    & \sum_{j \in J} y_{sjt} = \sum_{i \in I} x_{i} A_{it}
        & \forall t \in T, \ s \in S\\
    % 值班人數要滿足需求
    & y_{sjt} \geq D_{sjt}
        & \forall j \in J,\ t \in T, \ s \in S\\
    % 總上班人數 <= 總員工數
    & \sum_{i \in I} x_{i} \leq \sum_{j \in J} E_j  - B\\
    % 任一時段，每個技能的值班人數 <= 總持有技能人數
    & y_{sjt} \leq E_j
        & \forall j \in J,\ t \in T, \ s \in S\\
    % 定義 z: 轉換次數
    & z_{sjt} = \mid y_{sjt} - y_{s,j,t-1} \mid
        & \forall j \in J,\ t \in \{1,...23\}, \ s \in S\\  % 中途轉換次數(取絕對值)
    & x_i, y_{sjt}, z_{st}, r_{sjt} \geq 0 & \forall i,j,t,s
\end{align}

## Outsourcing
However, you might notice that the model above is quite useless because it always tries to satisfy the highest demand in each period. This is as a matter of course, for Alice does not allow lack of labor, so the model always tries to find a distribution of labor force that satisfies all scenarios. While in some extreme cases, overestimating the required labor forces may help since the highest possible demand might realize, in most cases, this is just a waste of money. To solve this issue, Alice decided that she can accept outsourcing, that is, if in the morning of a day, Alice found that the demand is higher than expected, she can call for help, and a group of poor students will come to her restaurants and help. Of course, the cost of hiring the students temporarily is higher than the cost of existing employees, otherwise, she can always call the students for help, instead of hiring the employees herself.

As a side note, the idea of "outsourcing" does not need to be actually hiring people, you can also interpret it as the penalty for lack of labor.

We add two new parameters $L_{jt}$, indicating the upper bound of outsourcing workers with the skill for job $j$ on period $t$, and its cost $G_{jt}$. In addition, we add a new decision variable $l_{sjt}$ to decide how many outsourcing workers on job $j$, period $t$ in scenario $s$.

| Parameter | Definition                                                   |
| --------- | ------------------------------------------------------------ |
| $A_{it}$  | Schedule $i$ include period $t$ or not (binary)              |
| $D_{jt}$  | Demand of worker needed on job $j$ at period $t$\\           |
| $E_{j}$   | Number of workers who have the skill for job $j$             |
| $L_{jt}$  | Number of outsourcing workers who have the skill for job $j$ on period $t$ |
| $B$       | Number of workers who have both skills, $B < E_j$            |
| $C_{t}$   | Cost of hiring a worker to work at period $t$                |
| $F$       | Cost of letting a worker change jobs at the middle of a day  |
| $G_{jt}$  | Cost of satisfied an unit of demand by outsourcing           |

| Decision Variable | Definition                                                   |
| ----------------- | ------------------------------------------------------------ |
| $x_{i}$           | Number of worker work on schedule $i$ (integer)              |
| $y_{sjt}$         | Number of worker work on job $j$ on period $t$ in scenario $s$ (integer) |
| $z_{sjt}$         | Number of worker changing job at the start of period $t$ in scenario $s$ (integer) |
| $l_{sjt}$         | Number of outsourcing on job $j$, period $t$ in scenario $s$ |

We'll first change the definition of the objective functions, taking the cost of outsourcing into account.

$$
\sum_{s \in S} P_s ( \sum_{j \in J} \sum_{t \in T} y_{sjt} C_{t}+ \sum_{t \in T}\sum_{j \in J} l_{sjt}G_{jt} + \sum_{t \in T}\sum_{j \in J} z_{sjt} F)
$$

The constraints also have to be slightly modified. To satisfy the demand, the summation of employees on duty and the outsourcing workers on duty have to be larger or equals to

$$
y_{sjt} + l_{sjt} \geq D_{sjt} \quad \forall j \in J,\ t \in T, \ s \in S\\
$$

Remind that we need to first decide how many employees have to be on duty before learning the actual scenario, but we can postpone the decision of hiring the outsourcing workers until we learn the actual scenario.

Then, the range of $l_{sjt}$.
$$
l_{sjt} \leq L_{jt} \quad \forall j \in J,\ t \in T, \ s \in S
$$

The full model is shown as follows.

\begin{align}
    % objective: min cost
    \min \quad
    & \sum_{s \in S} P_s
    \left( \sum_{j \in J} \sum_{t \in T} y_{sjt} C_{t}  % hiring cost
        + \sum_{t \in T}\sum_{j \in J} l_{sjt}G_{jt}  % outsourcing
        + \sum_{t \in T}\sum_{j \in J} z_{sjt} F \right) \\ % change job cost
    % objective: max redundant
    \max \quad
        & \sum_{s \in S} P_s \left( \sum_{j \in J} \sum_{t \in T} r_{sjt} \right)\\
    % objective: min variance of redundant
    %\min \quad
    %    & \sum_{s \in S} P_s \left( \sum_{j \in J} \sum_{t \in T} (r_{sjt} - \overline{r})^2 \right)\\
    % constraints
    \mbox{s.t.} \quad
    % 分配到工作的人數要符合上班人數
    & \sum_{j \in J} y_{sjt} = \sum_{i \in I} x_{i} A_{it}
        & \forall t \in T, \ s \in S\\
    % 值班人數要滿足需求
    & y_{sjt} + l_{sjt} \geq D_{sjt}
        & \forall j \in J,\ t \in T, \ s \in S\\
    % 定義 r: redundant (恆正)
    & r_{sjt} = y_{sjt} + l_{sjt} - D_{sjt}
        & \forall j \in J,\ t \in T, \ s \in S\\
    % 總上班人數 <= 總員工數
    & \sum_{i \in I} x_{i} \leq \sum_{j \in J} E_j  - B\\
    % 任一時段，每個技能的值班人數 <= 總持有技能人數
    & y_{sjt} \leq E_j
        & \forall j \in J,\ t \in T, \ s \in S\\
    % 任一時段，外包人數 <= 外包上限
    & l_{sjt} \leq L_{jt}
        & \forall j \in J,\ t \in T, \ s \in S\\
    % 定義 r 的平均值
    & \overline{r} = \sum_{s \in S}\sum_{j \in J}\sum_{t \in T} r_{sjt}\ /\ (S \times J \times T)\\
    % 定義 z: 轉換次數
    & z_{sjt} = \mid y_{sjt} - y_{s,j,t-1} \mid
        & \forall j \in J,\ t \in \{1,...23\}, \ s \in S\\  % 中途轉換次數(取絕對值)
    & x_i, y_{sjt}, z_{st}, r_{sjt} \geq 0 & \forall i,j,t,s
\end{align}

## 2-Stage Recourse
While the weighted-sum method mentioned above is good, we can further improve it by postponing the decision. More specifically, we can separate the model into two stages. In the first stage, we only know the probability distribution of the scenario, and we want to decide how many workers have to be on duty. In the second stage, we already learn the actual scenario and have to decide who has to work at which periods, and whether we need outsourcing. This is called the "2-stage recourse" method.

In the first stage, we only decide $x$, the expected number of workers on duty.

\begin{align}
    \min \quad & \sum_{i \in I} \sum_{t \in T} x_{i} A_{it} C_{t}  % hiring cost
    + \mathbb{E}(\pi_s \mid s \in S)\\
    \mbox{s.t.} \quad
    % % enough manpower
    % & \sum_{i \in I} x_{i} A_{it} \geq \sum_{j \in J} D_{sjt}
    %     & \forall t \in T\\ % 各時段上班人數 >= 所需人數
    & x_i \geq 0 & \forall i \in I
\end{align}

The $\pi$ is the output of the second stage model given the scenario $s$. We need to calculate the expected value of cost given by the second stage models, and optimize $x$ according to it.

In the second stage, we already know the actual scenario, it is at this point that we assign the employees to jobs and hire the outsourcing workers if required.

\begin{align}
    \min \quad & \pi_s = \sum_{t \in T}\sum_{j \in J} z_{sjt} F % change job cost
        + \sum_{t \in T} \sum_{j \in J} l_{sjt} G_{jt}\\
    \mbox{s.t.} \quad
    % constraint
    & \sum_{j \in J} y_{sjt} = \sum_{i \in I} x_{i} A_{it}
        & \forall t \in T\\  % 各時段分配到工作的人數要符合上班人數
    & y_{sjt} + l_{sjt} \geq D_{sjt}
        & \forall j \in J,\ t \in T\\  % 值班加外包要能滿足需求
    & y_{sjt} \leq E_j
        & \forall j \in J,\ t \in T\\  % 任一時段，每個技能的值班人數 <= 總持有技能人數
    & l_{sjt} \leq L_{jt}
        & \forall j \in J,\ t \in T\\  % 任一時段，外包人數 <= 上限
    & z_{sjt} = \mid y_{sjt} - y_{s,j,t-1} \mid
        & \forall j \in J,\ t \in \{1,...23\}
        \label{s.t.absolute}\\  % 中途轉換次數(取絕對值)
    & x_i, y_{sjt}, z_{st}, r_{sjt}, l_{sjt} \geq 0 & \forall i,j,t
\end{align}

Notice that in the second stage model, the $s$ is actually fixed, since we've learned the scenario.

## Simulation
Alice is now satisfied with the model. However, before applying to her awesome fashion restaurant, she wants to prove that this model is efficient. Therefore, we need to run several simulations.

In the most simple case, assuming that the number of customers is a bell shape. We first need to generate a demand that is bell-shaped.

```python
from scipy.stats import norm
import random
import numpy as np
import matplotlib.pyplot as plt
def demand_generation_normal(length):
    a = np.array(range(0, length)) - length // 2
    a = a / 6
    return (norm().pdf(a) * 70).astype(int)
demand_generation_normal(24), plt.plot(demand_generation_normal(24))
```

Looks good. But as mentioned, there are several scenarios, and we need to create several possible ones. In each scenario, the height, the center, and the variance of the bell shapes can be different. We'll try to randomly generate severals.

```python
def demand_generation_normal(length, offset=0, do_random=False):
    a = np.roll(np.array(range(0, length)) - length // 2, offset if not do_random else int(random.uniform(-2, 2)))
    a = a / ( random.uniform(3, 9) if do_random else 6 )
    return (norm().pdf(a) * ( random.uniform(80, 120) if do_random else 100 )).astype(int)
plt.plot(demand_generation_normal(24, do_random=True))
plt.plot(demand_generation_normal(24, do_random=True))
plt.plot(demand_generation_normal(24, do_random=True))
```

We can now see several different demands in different scenarios. After having this demand generator, we can run the simulation.
For example, we now have three scenarios of demand with equal probability:
$$
s_{low}, j_1=[1, 1, 2, 3, 4, 5, 6, 7, 7, 8, 9, 9, 9, 9, 9, 8, 7, 7, 6, 5, 4, 3, 2, 1] \\
s_{low}, j_2=[1, 1, 1, 2, 3, 4, 4, 5, 6, 7, 7, 7, 7, 7, 7, 7, 6, 5, 4, 4, 3, 2, 1, 1] \\
s_{mid}, j_1=[ 8, 10, 12, 14, 15, 17, 18, 19, 19, 19, 18, 17, 15, 14, 12, 10,  8, 6,  4,  3,  2,  3,  4,  6] \\
s_{mid}, j_2=[ 6,  4,  3,  2,  2,  2,  3,  4,  6,  8,  9, 11, 12, 14, 14, 15, 15, 15, 14, 14, 12, 11,  9,  8] \\
s_{high}, j_1=[23, 26, 27, 29, 29, 29, 27, 26, 23, 21, 18, 15, 12,  9,  6,  5,  3, 5,  6,  9, 12, 15, 18, 21] \\
s_{high}, j_2=[18, 16, 14, 12,  9,  7,  5,  4,  3,  4,  5,  7,  9, 12, 14, 16, 18, 21, 22, 23, 23, 23, 22, 21] \\
$$

The simulation of the basic model gives $x=[9.0, 32.0, 0, 28, 0, 26, 9, 32]$, with the cost 8,476.

The simulation of two-stage recourse method gives $x = [9, 32, 0, 28, 0, 26, 9, 32]$.

| Scenario | Cost  | Outsourcing |
| -------- | ----- | ----------- |
| Low      | 8,310 | 0           |
| Mid      | 8,370 | 0           |
| High     | 8,750 | 2           |

We can see that the model actually decided to assign fewer workers and use outsourcing in high demands. In addition, in the low and mid scenarios, the cost is lower.

---

But one example doesn't speak much, we then run the simulation 100 times. The basic model gives the average cost of 8,946.28, and the average outsourcing of 22.5.

And the two-stage recourse model gives:

| Scenario | Average Cost | Average Outsourcing |
| -------- | ------------ | ------------------- |
| Low      | 8117.35      | 0                   |
| Mid      | 8409.35      | 2.62                |
| High     | 11191.65     | 29.42               |

---

The above simulation assumes that the possibilities of all scenarios are the same. If the probability is also randomized, the basic model gives the expected cost of 8651.13, and the outsourcing amount of 43.35. The two-stage recourse model gives:

| Scenario | Average Cost | Average Outsourcing |
| -------- | ------------ | ------------------- |
| Low      | 7655.4       | 0.02                |
| Mid      | 8230.0       | 5.45                |
| High     | 11228.7      | 34.34               |

Notice that the basic model still often overestimate the labor force needed, and the two-stage recourse model can postpone the decision so that it usually performs better. Hope that Alice likes this model.

## Conclusion
This simple model shows how integer programming can be used for scheduling problems under stochastic demands. We also demonstrated that the model is efficient in terms of optimizing the cost. However, optimizing the cost may not be the only objective. In fact, my teammate wrote another tutorial, showing that how one can further enhance this model by the epsilon constraints. The link will be provided later.

Most of our model is based on the implementation of [1], but to simplify the problem, I removed the multi-criteria parts (which, as I mentioned, is covered by my teammate) and use the traditional integer programming method instead of the genetic algorithm. I hope this article helps readers understand how elegant the solutions provided by Operations Research are.

This article is the final project of [Operations Research Applications and Implementation](https://nol.ntu.edu.tw/nol/coursesearch/print_table.php?course_id=725%20U3690&class=&dpt_code=7050&ser_no=64554&semester=109-2&lang=CH) at Dept. of Information Management, National Taiwan University, in partnership with Ting Lin, under the supervision of Prof. Lee.

All the code used in this tutorial is open-sourced and published at [the GitHub repo](https://github.com/s3131212/Staff-Scheduling-For-Stochastic-Demands).

## Reference
[1] Cai, X. & Li, K. N., 2000. "A genetic algorithm for scheduling staff of mixed skills under multi-criteria," European Journal of Operational Research, Elsevier, vol. 125(2), pages 359-369, September.


```python
%matplotlib inline
import matplotlib.pyplot as plt

x = [1, 2, 3, 4, 5, 6, 7]
y = [64.3, 63.8, 63.6, 64.0, 63.5, 63.2, 63.1]

plt.plot(x, y, marker='o', color='red')
plt.grid(color='0.8')
plt.show()
```

```python
%matplotlib inline
import matplotlib.pyplot as plt

x = list(range(1, 11))
y = []
for i in range(len(x)):
    y.append(3 * x[i] - 24)

plt.plot(x, y)
plt.grid(color='0.8')
plt.show()
```

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

x = np.arange(-1.0, 1.01, 0.01)
y = 3 * x * x * x

plt.plot(x, y)
plt.grid(color='0.8')
plt.show()
```

```python
from sympy import Symbol, solve

a = Symbol('a')
b = Symbol('b')
ex1 = a + b - 1
ex2 = 5 * a + b - 3

print(solve((ex1, ex2)))
```

    {a: 1/2, b: 1/2}

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

x = np.arange(-1, 6)
y = 1 / 2 * x + 1 / 2
y2 = -2 * x + 7

plt.plot(x, y)
plt.plot(x, y2)
plt.axis('equal')
plt.grid(color='0.8')
plt.show()
```

```python
from sympy import Symbol, solve

x = Symbol('x')
y = Symbol('y')

ex1 = -3 / 2 * x + 6 - y
ex2 = 1 / 2 * x + 2 - y
print(solve((ex1, ex2)))
```

    {x: 2.00000000000000, y: 3.00000000000000}

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

a1 = (5 - 1) / (6 - 0)
b1 = 1

cx = (0 + 6) / 2
cy = (1 + 5) / 2

a2 = -1 / a1

b2 = cy - a2 * cx

x = np.arange(0, 7)
y1 = a1 * x + b1
y2 = a2 * x + b2

plt.plot(x, y1)
plt.plot(x, y2)
plt.axis('equal')
plt.grid(color='0.8')
plt.show()
```

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

th = np.arange(0, 360)

r = 5
x = r * np.cos(np.radians(th)) + 3
y = r * np.sin(np.radians(th)) - 1

plt.plot(x, y)
plt.axis('equal')
plt.grid(color='0.8')
plt.show()
```

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

rad = np.arctan2(3, 4) # 높이 = 3, 밑변 = 4 인 삼각형의 각도 구하기
print(rad) # 호도법
th = np.degrees(rad)
print(th) # 각도법
```

    0.6435011087932844
    36.86989764584402

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

r = 3000
x = np.arange(-r, r + 1)
y = np.sqrt(r**2 - x**2)

plt.plot(x, y)
plt.plot(x, -y)
plt.axis('equal')
plt.grid(color='0.8')
plt.show()
```

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

a = 200
b = 300

r = 300
x = np.arange(a - r, a + r + 1)
y = np.sqrt(r**2 - (x - a)**2) + b
y2 = -y + 2 * b

plt.plot(x, y)
plt.plot(x, y2)
plt.axis('equal')
plt.grid(color='0.8')
plt.show()
```

```python
import math

x = 1
y = 5
a = 1
b = 1
c = -4

math.fabs(a * x + b * y + c) / math.sqrt(a**2 + b**2)
```

    1.414213562373095

```python
import math

x = [1, 3, 6]
y = [5, 1, 4]

a = math.sqrt((x[1] - x[0]) ** 2 + (y[1] - y[0]) ** 2)
b = math.sqrt((x[2] - x[1]) ** 2 + (y[2] - y[1]) ** 2)
c = math.sqrt((x[2] - x[0]) ** 2 + (y[2] - y[0]) ** 2)

s = (a + b + c) / 2
math.sqrt(s * (s - a) * (s - b) * (s - c))
```

    8.999999999999996

```python

```


# Statistics

## 1.1 Sales for the past week were the following amounts: [3505, 2400, 3027, 2798, 3700, 3250, 2689]. Without using library functions, what is the mean, variance, and standard deviation of of sales from last week? (for extra bonus points, write your own function that can calculate these two values for any sized list)

```python
# not suing bessel's correction, since this is not a sample

sales = [3505, 2400, 3027, 2798, 3700, 3250, 2689]

def mean(x):
    return sum(x) / len(x)

def variance(x):
    mean_x = mean(x)
    return sum(map(lambda _: (_ - mean_x) ** 2, x)) / len(x)

def std_dev(x):
    return variance(x) ** (1/2)

print('mean', mean(sales), 'variance', variance(sales), 'std_dev', std_dev(sales))
```

    mean 3052.714285714286 variance 183761.06122448976 std_dev 428.67360686714756

```python
sum(map(lambda x:x**2, sales))
```

    66519779

## 1.2 Find the covariance between last week's sales numbers and the number of customers that entered the store last week: [127, 80, 105, 92, 120, 115, 93] (you may use library functions for calculating the covariance since we didn't specifically talk about its formula)

```python
import pandas as pd

customers = [127, 80, 105, 92, 120, 115, 93]

df = pd.DataFrame({'sales':sales,'customers':customers})
```

```python
df.sales.var(ddof=0)
```

    183761.06122448976

```python
# This method only computes sample covariance (See the docs)
df.cov()
```

      sales
      customers

      sales
      214387.904762
      7604.357143

      customers
      7604.357143
      290.952381

## 1.3 Find the standard deviation of customers who entered the store last week. Then, use the standard deviations of both sales and customers to standardize the covariance to find the correlation coefficient that summarizes the relationship between sales and customers. (You may use library functions to check your work.)

```python
# incorrect, because our std are POPULATION parameters, not sample statistics
df.cov().loc['sales', 'customers'] / (std_dev(sales) * std_dev(customers))
```

    1.123306307450706

```python
# this computes the POPULATION correlation, even though ddof=1.
# Bessels' correction cancels out with a factor in the numerator.
df.cov().loc['sales', 'customers'] / (df.sales.std(ddof=1) * df.customers.std(ddof=1))
```

    0.9628339778148909

```python
df.corr()
```

      sales
      customers

      sales
      1.000000
      0.962834

      customers
      0.962834
      1.000000

## 1.4 Use pandas to import a cleaned version of the titanic dataset from the following link: [Titanic Dataset](https://raw.githubusercontent.com/Geoyi/Cleaning-Titanic-Data/master/titanic_clean.csv)

## Calculate the variance-covariance matrix and correlation matrix for the titanic dataset's numeric columns. (you can encode some of the categorical variables and include them as a stretch goal if you finish early)

```python
df.cov?
```

```python
df = pd.read_csv('https://raw.githubusercontent.com/Geoyi/'
                 'Cleaning-Titanic-Data/master/titanic_clean.csv')
```

```python
df.cov()
```

      Unnamed: 0
      pclass
      survived
      age
      sibsp
      parch
      fare
      body
      has_cabin_number

      Unnamed: 0
      143117.500000
      284.357034
      -53.967125
      -1442.939812
      25.828746
      1.172783
      -9410.735123
      591.579132
      -95.438885

      pclass
      284.357034
      0.701969
      -0.127248
      -3.954605
      0.053090
      0.013287
      -24.227788
      -2.876653
      -0.249992

      survived
      -53.967125
      -0.127248
      0.236250
      -0.314343
      -0.014088
      0.034776
      6.146023
      0.000000
      0.061406

      age
      -1442.939812
      -3.954605
      -0.314343
      165.850021
      -2.559806
      -1.459378
      114.416613
      81.622922
      1.463138

      sibsp
      25.828746
      0.053090
      -0.014088
      -2.559806
      1.085052
      0.336833
      8.641768
      -8.708471
      -0.003946

      parch
      1.172783
      0.013287
      0.034776
      -1.459378
      0.336833
      0.749195
      9.928031
      4.237190
      0.013316

      fare
      -9410.735123
      -24.227788
      6.146023
      114.416613
      8.641768
      9.928031
      2678.959738
      -179.164684
      10.976961

      body
      591.579132
      -2.876653
      0.000000
      81.622922
      -8.708471
      4.237190
      -179.164684
      9544.688567
      3.625689

      has_cabin_number
      -95.438885
      -0.249992
      0.061406
      1.463138
      -0.003946
      0.013316
      10.976961
      3.625689
      0.174613

```python
df.corr()
```

      Unnamed: 0
      pclass
      survived
      age
      sibsp
      parch
      fare
      body
      has_cabin_number

      Unnamed: 0
      1.000000
      0.897822
      -0.293717
      -0.296172
      0.065594
      0.003584
      -0.481215
      0.015558
      -0.603727

      pclass
      0.897822
      1.000000
      -0.312469
      -0.366370
      0.060832
      0.018322
      -0.558629
      -0.034642
      -0.713857

      survived
      -0.293717
      -0.312469
      1.000000
      -0.050199
      -0.027825
      0.082660
      0.244265
      NaN
      0.302250

      age
      -0.296172
      -0.366370
      -0.050199
      1.000000
      -0.190747
      -0.130872
      0.171892
      0.059059
      0.271887

      sibsp
      0.065594
      0.060832
      -0.027825
      -0.190747
      1.000000
      0.373587
      0.160238
      -0.099961
      -0.009064

      parch
      0.003584
      0.018322
      0.082660
      -0.130872
      0.373587
      1.000000
      0.221539
      0.051099
      0.036806

      fare
      -0.481215
      -0.558629
      0.244265
      0.171892
      0.160238
      0.221539
      1.000000
      -0.043110
      0.507253

      body
      0.015558
      -0.034642
      NaN
      0.059059
      -0.099961
      0.051099
      -0.043110
      1.000000
      0.083796

      has_cabin_number
      -0.603727
      -0.713857
      0.302250
      0.271887
      -0.009064
      0.036806
      0.507253
      0.083796
      1.000000

# Orthogonality

## 2.1 Plot two vectors that are orthogonal to each other. What is a synonym for orthogonal?

```python
import matplotlib.pyplot as plt

plt.axes().set_aspect('equal')
plt.arrow(0,0,2,1)
plt.arrow(0,0,1,-2);
plt.xlim(0,4)
plt.ylim(2,-2);
```

## 2.2 Are the following vectors orthogonal? Why or why not?

\begin{align}
a = \begin{bmatrix} -5 \\ 3 \\ 7 \end{bmatrix}
\qquad
b = \begin{bmatrix} 6 \\ -8 \\ 2 \end{bmatrix}
\end{align}

```python
import numpy as np

a = np.array([-5, 3, 7])
b = np.array([6, -8, 2])

np.dot(a,b)
```

    -40

No, the dot product is not 0.

## 2.3 Compute the following values: What do these quantities have in common?

## What is $||c||^2$?

## What is $c \cdot c$?

## What is $c^{T}c$?

\begin{align}
c = \begin{bmatrix} 2 & -15 & 6 & 20 \end{bmatrix}
\end{align}

```python
c = np.array([2, -15, 6, 20])

print(np.linalg.norm(c) ** 2)

print(np.dot(c, c))

print(np.matmul(np.transpose(c), c))
```

    665.0
    665
    665

# 🤯

# Unit Vectors

## 3.1 Using Latex, write the following vectors as a linear combination of scalars and unit vectors:

\begin{align}
d = \begin{bmatrix} 7 \\ 12 \end{bmatrix}
\qquad
e = \begin{bmatrix} 2 \\ 11 \\ -8  \end{bmatrix}
\end{align}

\begin{align}
d = 7\hat{i} + 12\hat{j}
\qquad
e = 2\hat{i} + 11\hat{j} - 8\hat{k}
\end{align}

But also . . .

```python
v = np.array([7, 12])
np.sqrt((np.sum(v**2)))
```

    13.892443989449804

```python
v = np.array([2, 11, -8])
np.sqrt((np.sum(v**2)))
```

    13.74772708486752

```python
7/(7**2 + 12**2)**(1/2)
```

    0.5038710255240862

```python
np.linalg.norm(np.array([0.14, 0.79, -0.58]))
```

    0.99

\begin{align}
d = 13.9 \begin{bmatrix} 0.50 \\ 0.86 \end{bmatrix}
\qquad
e = 13.7 \begin{bmatrix} 0.14 \\ 0.79 \\ -0.58  \end{bmatrix}
\end{align}

## 3.2 Turn vector $f$ into a unit vector:

\begin{align}
f = \begin{bmatrix} 4 & 12 & 11 & 9 & 2 \end{bmatrix}
\end{align}

```python
v = np.array([4, 12, 11, 9, 2])
v / np.sqrt((np.sum(v**2)))
```

    array([0.20908335, 0.62725005, 0.57497921, 0.47043754, 0.10454167])

# Linear Independence / Dependence

## 4.1 Plot two vectors that are linearly dependent and two vectors that are linearly independent (bonus points if done in $\mathbb{R}^3$).

```python
plt.axes().set_aspect('equal')
plt.arrow(0,0,2,1,head_width=.1, head_length=0.1)
plt.arrow(0,0,1,-2,head_width=.1, head_length=0.1)
plt.arrow(0,0,1,2,head_width=.1, head_length=0.1)
plt.arrow(0,0,0.5,1,head_width=.1, head_length=0.1)
plt.xlim(0,4)
plt.ylim(-2,2);
```

# Span

## 5.1 What is the span of the following vectors?

\begin{align}
g = \begin{bmatrix} 1 & 2 \end{bmatrix}
\qquad
h = \begin{bmatrix} 4 & 8 \end{bmatrix}
\end{align}

R, they are on the same line

## 5.2 What is the span of $\{l, m, n\}$?

\begin{align}
l = \begin{bmatrix} 1 & 2 & 3 \end{bmatrix}
\qquad
m = \begin{bmatrix} -1 & 0 & 7 \end{bmatrix}
\qquad
n = \begin{bmatrix} 4 & 8  & 2\end{bmatrix}
\end{align}

R^3

# Basis

## 6.1 Graph two vectors that form a basis for $\mathbb{R}^2$

```python
plt.axes().set_aspect('equal')
plt.arrow(0,0,2,1,head_width=.1, head_length=0.1)
plt.arrow(0,0,1,-2,head_width=.1, head_length=0.1)
plt.xlim(0,4)
plt.ylim(-2,2);
```

## 6.2 What does it mean to form a basis?

The basis is formed by linear independent vectors that can be scaled and combined to fill a vector space $V$

# Rank

## 7.1 What is the Rank of P?

\begin{align}
P = \begin{bmatrix}
1 & 2 & 3 \\
 -1 & 0 & 7 \\
4 & 8  & 2
\end{bmatrix}
\end{align}

```
[ 1, 2, 3]
[-1, 0, 7]
[ 4, 8, 2] / 4

[ 1, 2, 3  ] -v3
[-1, 0, 7  ]
[ 0, 2,  .5]

[ 1, 0, 2.5]
[-1, 0, 7  ] +v1
[ 0, 2,  .5] / 2

[ 1, 0, 2.5 ]
[ 0, 0, 9.5 ] / 9.5
[ 0, 1,  .25]

[ 1, 0, 2.5 ] -2.5  * v3
[ 0, 1,  .25] - .25 * v3
[ 0, 0, 1   ]

[ 1, 0, 0]
[ 0, 1, 0]
[ 0, 0, 1]
```

R^3

## 7.2 What does the rank of a matrix tell us?

The number of effective dimensions the matrix reduces to.

# Linear Projections

## 8.1 Line $L$ is formed by all of the vectors that can be created by scaling vector $v$
\begin{align}
v = \begin{bmatrix} 1 & 3 \end{bmatrix}
\end{align}

\begin{align}
w = \begin{bmatrix} -1 & 2 \end{bmatrix}
\end{align}

## find $proj_{L}(w)$

## graph your projected vector to check your work (make sure your axis are square/even)

\begin{align}
proj_{L}(w) =  \frac{w \cdot v}{v \cdot v}v
\end{align}

```python
w = np.array([-1, 2])
v = np.array([ 1, 3])
proj = np.multiply(np.multiply(w, v) / np.multiply(v, v), v)

plt.axes().set_aspect('equal')

plt.arrow(0, 0, *v, linewidth=3, head_width=.05, head_length=0.05, color ='green')
plt.arrow(0, 0, *w, linewidth=3, head_width=.05, head_length=0.05, color ='red')
plt.arrow(0, 0, *proj, linewidth=3, head_width=.05, head_length=0.05, color ='gray')

plt.xlim(-1,2)
plt.ylim(-1,4)
plt.show();
```

# Stretch Goal

## For vectors that begin at the origin, the coordinates of where the vector ends can be interpreted as regular data points. (See 3Blue1Brown videos about Spans, Basis, etc.)

## Write a function that can calculate the linear projection of each point (x,y) (vector) onto the line y=x. run the function and plot the original points in blue and the new projected points on the line y=x in red.

## For extra points plot the orthogonal vectors as a dashed line from the original blue points to the projected red points.

```python
import pandas as pd
import matplotlib.pyplot as plt

# Creating a dataframe for you to work with -Feel free to not use the dataframe if you don't want to.
x_values = [1, 4, 7, 3, 9, 4, 5 ]
y_values = [4, 2, 5, 0, 8, 2, 8]

data = {"x": x_values, "y": y_values}

df = pd.DataFrame(data)

df.head()

plt.scatter(df.x, df.y)
plt.show()
```


```python
import numpy
import sympy
from matplotlib import pyplot, cm
%matplotlib inline
pyplot.rcParams['font.family'] = 'serif'
pyplot.rcParams['font.size'] = 16
sympy.init_printing()
```

```python
import urllib.request

# Download and read the data file.
url = ('https://github.com/numerical-mooc/numerical-mooc/blob/master/'
       'lessons/04_spreadout/data/uvinitial.npz?raw=true')
filepath = 'uvinitial.npz'
urllib.request.urlretrieve(url, filepath);
```

```python
# Read the initial fields from the file.
# Load the initial data from file.
filepath = 'uvinitial.npz'
with open(filepath, 'rb') as infile:
    data = numpy.load(infile)
    u0, v0 = data['U'], data['V']
```

```python
# Plot the initial fields.
fig, (ax1, ax2) = pyplot.subplots(ncols=2, figsize=(9.0, 4.0))
ax1.set_title('$u$')
ax1.imshow(u0, cmap=cm.Greys_r)
ax1.axis('off')
ax2.set_title('$v$')
ax2.imshow(v0, cmap=cm.Greys_r)
ax2.axis('off');
```

It was the scheme to achieve second-order accuracy in both space and time. Numerical dispersion occurs when a higher order discretisation scheme is used to improve accuracy of the result. This is due to the truncation error of the discretisation, a second order upwind method, the leading truncation error is odd, and odd order derivatives contribute to numerical dispersion.

```python
def ftcs(u0, v0, nt, dt, dx, dy, Du, Dv, F, k):
    # Define some constants.
    sigma_xu = Du * dt / dx**2
    sigma_yu = Du * dt / dy**2
    sigma_xv = Dv * dt / dx**2
    sigma_yv = Dv * dt / dy**2
    # Integrate in time
    u = u0.copy()
    v = v0.copy()
    u_hist = []
    v_hist = []
#     ny, nx = u.shape
    for n in range(nt):
        u[1:-1, 1:-1] = u[1:-1, 1:-1] - dt*u[1:-1, 1:-1] * v[1:-1, 1:-1]**2 + dt * F * (1 - u[1:-1, 1:-1]) + sigma_xu * (u[1:-1, 2:] - 2.0 * u[1:-1, 1:-1] + u[1:-1, :-2]) + sigma_yu * (u[2:, 1:-1] - 2.0 * u[1:-1, 1:-1] + u[:-2, 1:-1])
        v[1:-1, 1:-1] = v[1:-1, 1:-1] + dt*u[1:-1, 1:-1] * v[1:-1, 1:-1]**2 - dt* (F + k) * v[1:-1, 1:-1] + sigma_xv * (v[1:-1, 2:] - 2.0 * v[1:-1, 1:-1] + v[1:-1, :-2]) + sigma_yv * (v[2:, 1:-1] - 2.0 * v[1:-1, 1:-1] + v[:-2, 1:-1])
        # Apply Neumann conditions
        u[-1, :] = u[-2, :]
        u[:, -1] = u[:, -2]
        u[0, :] = u[1, :]
        u[:, 0] = u[:, 1]
        v[-1, :] = v[-2, :]
        v[:, -1] = v[:, -2]
        v[0, :] = v[1, :]
        v[:, 0] = v[:, 1]

        u_hist.append(u.copy())
        v_hist.append(v.copy())

    return u_hist, v_hist

```

```python
# Set spatial parameters.
Lx, Ly = 5.0, 5.0  # domain dimensions
nx, ny = 192, 192  # number of points in each direction
dx, dy = Lx / (nx - 1), Ly / (ny - 1)  # grid spacings

# Create the gridline locations.
x = numpy.linspace(0.0, Lx, num=nx)
y = numpy.linspace(0.0, Ly, num=ny)

# Set parameters of the pattern.
Du, Dv = 0.00016, 0.00008  # rates of diffusion
F, k = 0.035, 0.060  # parameters to feed and kill

# Set temporal parameters.
t0 = 0.0
t = 8000
dt = 9.0 * dx**2 / (40.0 * max(Du, Dv))  # time-step size
nt = 1000  # number of time steps
```

```python
u_hist, v_hist= ftcs(u0, v0, nt, dt, dx, dy, Du, Dv, F, k)

```

```python
import ipywidgets

def plot_species_concentration(u_hist, state):
    fig, ax = pyplot.subplots(figsize=(6.0, 6.0))
    ax.imshow(u_hist[state], vmin=0.5, vmax=1.0,
              interpolation='bicubic', cmap=cm.Greys_r)
    ax.axis('off')

def interactive_figure(u_hist):
    nt = len(u_hist) - 1
    state_slider = ipywidgets.IntSlider(value=0, min=0, max=nt,
                                        step=1,
                                        description='state')
    w = ipywidgets.interactive(plot_species_concentration,
                               u_hist=ipywidgets.fixed(u_hist),
                               state=state_slider)
    return w

# Display the history of the first species concentration.
display(interactive_figure(u_hist))
```

    interactive(children=(IntSlider(value=0, description='state', max=999), Output()), _dom_classes=('widget-inter…

```python
a = u_hist[-1].copy()
u_sample = a[100, ::40]
print(u_sample)
```

    [0.9985286  0.59597234 0.39688384 0.74579087 0.57737006]

```python
help(numpy.stack)
```

    Help on function stack in module numpy:

    stack(arrays, axis=0, out=None)
        Join a sequence of arrays along a new axis.

        The ``axis`` parameter specifies the index of the new axis in the
        dimensions of the result. For example, if ``axis=0`` it will be the first
        dimension and if ``axis=-1`` it will be the last dimension.

        .. versionadded:: 1.10.0

        Parameters
        ----------
        arrays : sequence of array_like
            Each array must have the same shape.

        axis : int, optional
            The axis in the result array along which the input arrays are stacked.

        out : ndarray, optional
            If provided, the destination to place the result. The shape must be
            correct, matching that of what stack would have returned if no
            out argument were specified.

        Returns
        -------
        stacked : ndarray
            The stacked array has one more dimension than the input arrays.

        See Also
        --------
        concatenate : Join a sequence of arrays along an existing axis.
        block : Assemble an nd-array from nested lists of blocks.
        split : Split array into a list of multiple sub-arrays of equal size.

        Examples
        --------
        >>> arrays = [np.random.randn(3, 4) for _ in range(10)]
        >>> np.stack(arrays, axis=0).shape
        (10, 3, 4)

        >>> np.stack(arrays, axis=1).shape
        (3, 10, 4)

        >>> np.stack(arrays, axis=2).shape
        (3, 4, 10)

        >>> a = np.array([1, 2, 3])
        >>> b = np.array([2, 3, 4])
        >>> np.stack((a, b))
        array([[1, 2, 3],
               [2, 3, 4]])

        >>> np.stack((a, b), axis=-1)
        array([[1, 2],
               [2, 3],
               [3, 4]])

```python

```


```python
from devito import *
from examples.seismic.source import DGaussSource, TimeAxis
from examples.seismic import plot_image
import numpy as np

from sympy import init_printing, latex
init_printing(use_latex='mathjax')
```

```python
# Initial grid: 1km x 1km, with spacing 100m
extent = (2000., 2000.)
shape = (81, 81)
x = SpaceDimension(name='x', spacing=Constant(name='h_x', value=extent[0]/(shape[0]-1)))
z = SpaceDimension(name='z', spacing=Constant(name='h_z', value=extent[1]/(shape[1]-1)))
grid = Grid(extent=extent, shape=shape, dimensions=(x, z))
```

```python
# Timestep size from Eq. 7 with V_p=6000. and dx=100
t0, tn = 0., 200.
dt = 1e2*(1. / np.sqrt(2.)) / 60.
time_range = TimeAxis(start=t0, stop=tn, step=dt)

src = DGaussSource(name='src', grid=grid, f0=0.01, time_range=time_range, a=0.004)
src.coordinates.data[:] = [1000., 1000.]
```

```python
#NBVAL_SKIP

src.show()
```

```python
# Now we create the velocity and pressure fields
p = TimeFunction(name='p', grid=grid, staggered=NODE, space_order=2, time_order=1)
v = VectorTimeFunction(name='v', grid=grid, space_order=2, time_order=1)
```

```python
from devito.finite_differences.operators import div, grad
t = grid.stepping_dim
time = grid.time_dim

# We need some initial conditions
V_p = 4.0
density = 1.

ro = 1/density
l2m = V_p*V_p*density

# The source injection term
src_p = src.inject(field=p.forward, expr=src)

# 2nd order acoustic according to fdelmoc
u_v_2 = Eq(v.forward, solve(v.dt - ro * grad(p), v.forward))
u_p_2 = Eq(p.forward, solve(p.dt - l2m * div(v.forward), p.forward))
```

```python
u_v_2
```

$\displaystyle \left[\begin{matrix}\operatorname{v_{x}}{\left(t + dt,x + \frac{h_{x}}{2},z \right)}\\\operatorname{v_{z}}{\left(t + dt,x,z + \frac{h_{z}}{2} \right)}\end{matrix}\right] = \left[\begin{matrix}dt \left(\frac{\partial}{\partial x} p{\left(t,x,z \right)} + \frac{\operatorname{v_{x}}{\left(t,x + \frac{h_{x}}{2},z \right)}}{dt}\right)\\dt \left(\frac{\partial}{\partial z} p{\left(t,x,z \right)} + \frac{\operatorname{v_{z}}{\left(t,x,z + \frac{h_{z}}{2} \right)}}{dt}\right)\end{matrix}\right]$

```python
u_p_2
```

$\displaystyle p{\left(t + dt,x,z \right)} = dt \left(16.0 \frac{\partial}{\partial x} \operatorname{v_{x}}{\left(t + dt,x + \frac{h_{x}}{2},z \right)} + 16.0 \frac{\partial}{\partial z} \operatorname{v_{z}}{\left(t + dt,x,z + \frac{h_{z}}{2} \right)} + \frac{p{\left(t,x,z \right)}}{dt}\right)$

```python
op_2 = Operator([u_v_2, u_p_2] + src_p)
```

```python
#NBVAL_IGNORE_OUTPUT

# Propagate the source
op_2(time=src.time_range.num-1, dt=dt)
```

    Operator `Kernel` ran in 0.02 s

    PerformanceSummary([(PerfKey(name='section0', rank=None),
      PerfEntry(time=0.012256000000000008, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
     (PerfKey(name='section1', rank=None),
      PerfEntry(time=0.0027869999999999982, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])

```python
#NBVAL_SKIP

# Let's see what we got....
plot_image(v[0].data[0])
plot_image(v[1].data[0])
plot_image(p.data[0])
```

```python
norm_p = norm(p)
assert np.isclose(norm_p, .35098, atol=1e-4, rtol=0)
```

```python
# # 4th order acoustic according to fdelmoc
p4 = TimeFunction(name='p', grid=grid, staggered=NODE, space_order=4, time_order=1)
v4 = VectorTimeFunction(name='v', grid=grid, space_order=4, time_order=1)
src_p = src.inject(field=p4.forward, expr=src)
u_v_4 = Eq(v4.forward, solve(v4.dt - ro * grad(p4), v4.forward))
u_p_4 = Eq(p4.forward, solve(p4.dt - l2m * div(v4.forward), p4.forward))
```

```python
#NBVAL_IGNORE_OUTPUT
op_4 = Operator([u_v_4, u_p_4] + src_p)
# Propagate the source
op_4(time=src.time_range.num-1, dt=dt)
```

    Operator `Kernel` ran in 0.03 s

    PerformanceSummary([(PerfKey(name='section0', rank=None),
      PerfEntry(time=0.019839999999999976, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
     (PerfKey(name='section1', rank=None),
      PerfEntry(time=0.0033089999999999986, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])

```python
#NBVAL_SKIP

# Let's see what we got....
plot_image(v4[0].data[-1])
plot_image(v4[1].data[-1])
plot_image(p4.data[-1])
```

```python
norm_p = norm(p)
assert np.isclose(norm_p, .35098, atol=1e-4, rtol=0)
```


# Tutorial 3: A one-asset HANK model

In this notebook we solve the one-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:
- **Hetinputs and hetoutputs**: adapt off-the-shelf HA blocks to new macro models.
- **Calibration DAG**: exploit analytical part of internal calibration

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, create_model  # functions
from sequence_jacobian import hetblocks, grids      # modules
```

## 1 Model description

The model is a HA version of the textbook New Keynesian model. Technology is linear in labor, prices are sticky á la Rotemberg, and monetary policy follows a Taylor rule. The model can be summarized in sequence form as

$$
\textbf{F}_t(\textbf{X}, Z) \equiv
\begin{pmatrix}
Y_t - Z_t L_t
\\
Y_t \left[1 - \frac{\mu}{\mu-1}\frac{1}{2\kappa} \log(1 + \pi_t)^2\right] - w_t L_t - d_t
\\
r_t B - \tau_t
\\
r^*_t + \phi \pi_t - i_t
\\
1 + r_t - \frac{1+i_{t-1}}{1+\pi_t}
\\
\kappa \left(\frac{w_t}{Z_t} - \frac{1}{\mu} \right)  + \frac{1}{1+r_{t+1}} \frac{Y_{t+1}}{Y_t} \log(1+\pi_{t+1}) - \log(1+\pi_t)
\\
\mathcal{A}_t(\{r_s, w_s, \tau_s, d_s\})  - B
\\
\mathcal{N}_t(\{r_s, w_s, \tau_s, d_s\}) - L_t
\end{pmatrix}
= \begin{pmatrix} 0 \\ 0 \\ 0 \\ 0 \\ 0 \\ 0 \\ 0 \\ 0\end{pmatrix},
\qquad t = 0, 1, \dots
$$

where the endogenous variables are $\textbf{X} = (Y, L, r, w, d, \pi, \tau, i)$ and the exogenous variables are $\textbf{Z}=(r^*, Z)$.

The asset demand and labor supply functions $\{\mathcal{A}, \mathcal{L}\}$ follow from the household block with Bellman equation

$$
\begin{align} \tag{HH}
V_t(e, a_{-}) = \max_{c, n, a} &\left\{\frac{c^{1-\sigma}}{1-\sigma} - \varphi \frac{n^{1+\nu}}{1+\nu} + \beta \mathbb{E}_t\left[V_{t+1}(e', a)|e \right] \right\}
\\
c + a &= (1 + r_t)a_{-} + w_t e n - \tau_t \bar{\tau}(e) + d_t \bar{d}(e)
\\
a &\geq 0
\end{align}
$$

where $\bar\tau(e)$ and $\bar d(e)$ are skill-specific incidence rules for taxes and dividends.

We can think of the model as a directed acyclical graph (DAG) with **3 endogenous inputs** and write it as an implicit function

$$
H(\pi, Y, w; \epsilon, Z) = 0.
$$

## 2 Embed HA block

As we have seen in the Krusell-Smith notebook, the main task in setting up HA blocks is to write a backward iteration function that represents the Bellman equation. This has to be a single step of an iterative solution method such as value function iteration. For the standard income fluctuation problem with endogenous labor supply we're dealing with here, the endogenous gridpoint method of [Carroll (2006)](https://www.sciencedirect.com/science/article/pii/S0165176505003368) is the best practice.

Solving the endogenous-labor problem via EGM is standard but somewhat tedious and so the details are left to ``sequence_jacobian/hetblocks/hh_labor.py``. Instead we will focus on how to adapt this off-the-shelf HetBlock to our specific macro enviroment.

```python
hh = hetblocks.hh_labor.hh

print(hh)
print(f'Inputs: {hh.inputs}')
print(f'Macro outputs: {hh.outputs}')
print(f'Micro outputs: {hh.internals}')
```

    <HetBlock 'hh'>
    Inputs: ['a_grid', 'we', 'T', 'r', 'beta', 'eis', 'frisch', 'vphi', 'Pi']
    Macro outputs: ['A', 'C', 'N']
    Micro outputs: ['D', 'Dbeg', 'Pi', 'Va', 'a', 'c', 'n']

The HetBlock `hh` corresponds to the general problem:

$$
\begin{align} \tag{HH-general}
V_t(e, a_{-}) = \max_{c, n, a} &\left\{\frac{c^{1-\sigma}}{1-\sigma} - \varphi \frac{n^{1+\nu}}{1+\nu} + \beta \mathbb{E}_t\left[V_{t+1}(e', a)|e\right] \right\}
\\
c + a &= (1 + r_t)a_{-} + w_t(e) n + T_t(e)
\\
a &\geq 0
\end{align}
$$

That is, households take as given the sequence of interest rates $r_t$, and skill-specific wages $w_t(e)$ and  transfers $T_t(e).$ In the context of this particular HANK model, transfers equal dividends minus taxes. But it's easy to imagine many other cases. Rather than writing a specific backward iteration function for each of them, we can just supply a function that specifies how the $\{w_t(e), T_t(e)\}$ are determined in this particular case. We refer such functions as **hetinput**.

In addition, we need to report effective labor supply $ne = n\cdot e$ to resolve labor market clearing. We can do so by attaching a **hetoutput** function to the core HetBlock.

### 2.1 Hetinputs

Let's start with the hetinputs. These functions will be evaluated before the core HetBlock (i.e. the backward iteration).

```python
def make_grid(rho_e, sd_e, nE, amin, amax, nA):
    e_grid, pi_e, Pi = grids.markov_rouwenhorst(rho=rho_e, sigma=sd_e, N=nE)
    a_grid = grids.agrid(amin=amin, amax=amax, n=nA)
    return e_grid, pi_e, Pi, a_grid

def transfers(pi_e, Div, Tax, e_grid):
    # hardwired incidence rules are proportional to skill; scale does not matter
    tax_rule, div_rule = e_grid, e_grid
    div = Div / np.sum(pi_e * div_rule) * div_rule
    tax = Tax / np.sum(pi_e * tax_rule) * tax_rule
    T = div - tax
    return T

def wages(w, e_grid):
    we = w * e_grid
    return we
```

- The second hetinput `transfers` takes inputs (`e_grid` and `pi_e`) that are produced by the first hetinput `make_grid`. Such *acylic* dependence between hetinputs is allowed. The block processes the inputs and outputs of the hetinput functions and puts them in a correct order of evaluation.
- Scalar-valued inputs of hetinputs may be time-varying. For example, aggregate dividends and taxes (`Div` and `Tax`) are determined endogenously in the HANK model and passed on to Bellman equation through the hetinput `transfers`. Thus, we can compute Jacobians with respect to them.

Let's attach these hetinputs to the household block using (the aptly-named) ``HetBlock.add_hetinput`` method.

```python
hh1 = hh.add_hetinputs([make_grid, transfers, wages])

print(hh1)
print(f'Inputs: {hh1.inputs}')
```

    <HetBlock 'hh' with hetinput 'make_grid_transfers'>
    Inputs: ['r', 'beta', 'eis', 'frisch', 'vphi', 'rho_e', 'sd_e', 'nE', 'amin', 'amax', 'nA', 'Div', 'Tax', 'w']

Note that `hh1` has only scalar-valued inputs.

### 2.2 Hetoutputs

Hetoutput functions are analogous to hetinputs. They are called after the backward iteration has converged. Thus, they may take multidimensional outputs of the backward iteration function as well as of the hetinputs as inputs.

```python
def labor_supply(n, e_grid):
    ne = e_grid[:, np.newaxis] * n
    return ne
```

Let's attach this hetoutput to the household block using ``HetBlock.add_hetoutput`` method.

```python
hh_ext = hh1.add_hetoutputs([labor_supply])

print(hh_ext)
print(f'Outputs: {hh_ext.outputs}')
```

    <HetBlock 'hh' with hetinput 'make_grid_transfers' and with hetoutput `labor_supply'>
    Outputs: ['A', 'C', 'N', 'NE']

- The signature of HetBlock `hh_ext` now references both hetinputs and hetoutputs.
- Aggregate outputs now include effective labor supply `NE`.

### 2.3 Taking stock

The SSJ toolkit comes with 3 generic HetBlocks, located in `sequence_jacobian/hetblocks`.
- `hh_sim`: standard incomplete markets model
- `hh_labor`: standard incomplete markets model with frictionless labor supply
- `hh_twoasset`: two-asset model with convex portfolio adjustment cost

By using hetinputs and hetoutputs, these core blocks may be embedded in different macro environments. This is the simplest way of using the SSJ toolkit, which may suffice for many applications.

If you wish to solve a model that's not just a variation on these off-the-shelf HetBlocks, there's two cases to consider.
1. The model fits into the HetBlock paradigm. E.g., standard incomplete markets models with additional choices such as search intensity. All you need to do is write a new backward iteration function. Use it to instantiate a new HetBlock and get all the HetBlock methods for free.
2. The model does not fit the HetBlock paradigm. E.g., models in which discrete endogenous states. In this case, we recommend that you "bring your own Jacobian". That is, solve the Jacobian of your block outside the SSJ toolkit. Once you turn them into an instance of `JacobianDict` (like we did in section 4 of the Krusell-Smith notebook), you can include them in models in lieu of an actual block. This is sufficient for using linear solution methods (`impulse_linear`, `jacobian` and their `solved_` versions) at the macro model level.

## 3 Calibrating the steady state
Similarly to the RBC example, we calibrate the discount factor $\beta$ and disutility of labor $\varphi$ to hit a target for the interest rate and effective labor $L=1.$ Additionally we calibrate the wage $w$ such that the Phillips curve relation is satisfied in steady state for zero inflation $\pi=0$.

Note that the mapping from $\beta$ to asset market clearing given $r$ and that from $\varphi$ to average effective labor supply involve the household block. As such, we must rely on numerical root-finding. In contrast, the $w$ is easy to characterize analytically from the ss-version of NKPC:
$$
w = \frac{Z}{\mu} \tag{ss wage}
$$
This situation is very common in HA-DSGE models. Although using a numerical root-finder on every unknown jointly is likely to succeed in simple models, exploiting analytical solutions becomes crucial in more complicated models.

The simplest way of doing so is to work with two DAGs: one for calibration and one for transition dynamics. Typically, the two DAGs share some, but not all of their blocks. Furthermore, they may have different unknowns and targets.

### 3.1 Blocks for SS-DAG

We already implemented the household block. Let's define the rest as simple blocks.

```python
@simple
def firm(Y, w, Z, pi, mu, kappa):
    L = Y / Z
    Div = Y - w * L - mu/(mu-1)/(2*kappa) * (1+pi).apply(np.log)**2 * Y
    return L, Div

@simple
def monetary(pi, rstar, phi):
    r = (1 + rstar(-1) + phi * pi(-1)) / (1 + pi) - 1
    return r

@simple
def fiscal(r, B):
    Tax = r * B
    return Tax

@simple
def mkt_clearing(A, NE, C, L, Y, B, pi, mu, kappa):
    asset_mkt = A - B
    labor_mkt = NE - L
    goods_mkt = Y - C - mu/(mu-1)/(2*kappa) * (1+pi).apply(np.log)**2 * Y
    return asset_mkt, labor_mkt, goods_mkt

@simple
def nkpc_ss(Z, mu):
    w = Z / mu
    return w
```

- These blocks---with the exception of `nkpc_ss`--- are valid in transition as well as in steady state.
- Notice that we wrap `np.log()` in an `.apply()` inside simple blocks. This is necessary for functions whose name includes a dot. The reason is that it would interfere with some internal processing that simple blocks do to handle leads, lags, and references to steady state.

```python
blocks_ss = [hh_ext, firm, monetary, fiscal, mkt_clearing, nkpc_ss]

hank_ss = create_model(blocks_ss, name="One-Asset HANK SS")

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

    <Model 'One-Asset HANK SS'>
    Inputs: ['beta', 'eis', 'frisch', 'vphi', 'rho_e', 'sd_e', 'nE', 'amin', 'amax', 'nA', 'Y', 'Z', 'pi', 'mu', 'kappa', 'rstar', 'phi', 'B']

Print inputs so we don't have to remember everything.

```python
calibration = {'eis': 0.5, 'frisch': 0.5, 'rho_e': 0.966, 'sd_e': 0.5, 'nE': 7,
               'amin': 0.0, 'amax': 150, 'nA': 500, 'Y': 1.0, 'Z': 1.0, 'pi': 0.0,
               'mu': 1.2, 'kappa': 0.1, 'rstar': 0.005, 'phi': 1.5, 'B': 5.6}

unknowns_ss = {'beta': 0.986, 'vphi': 0.8}
targets_ss = {'asset_mkt': 0, 'labor_mkt': 0}

ss0 = hank_ss.solve_steady_state(calibration, unknowns_ss, targets_ss, solver="hybr")
```

Let's see the targets and Walras's law.

```python
print(f"Asset market clearing: {ss0['asset_mkt']: 0.2e}")
print(f"Labor market clearing: {ss0['labor_mkt']: 0.2e}")
print(f"Goods market clearing (untargeted): {ss0['goods_mkt']: 0.2e}")
```

    Asset market clearing:  1.42e-14
    Labor market clearing:  4.44e-16
    Goods market clearing (untargeted): -5.26e-09

Looks good. Let's also plot the labor supply policy as function of assets for each skill type. We see that poorer and more productive households choose to work more.

```python
plt.plot(ss0.internals['hh']['a_grid'], ss0.internals['hh']['n'].T)
plt.xlabel('Assets'), plt.ylabel('Labor supply')
plt.show()
```

## 4 Linearized dynamics using Jacobians
Recall that we reduced the model to an implicit equation
$$
H(\pi, Y, w; r^*, Z) =
\begin{pmatrix}
\kappa \left(\frac{w_t}{Z_t} - \frac{1}{\mu} \right)  + \frac{1}{1+r_{t+1}} \frac{Y_{t+1}}{Y_t} \log(1+\pi_{t+1}) - \log(1+\pi_t)
\\
\mathcal{A}_t(\{r_s, w_s, \tau_s, d_s\})  - B
\\
\mathcal{L}_t(\{r_s, w_s, \tau_s, d_s\}) - L_t
\end{pmatrix}
= \begin{pmatrix} 0 \\ 0 \\ 0\end{pmatrix},
$$

to be solved for $U=(\pi, w, Y)$ given any $Z=(Z, r^*)$. The rest of the endogenous variables are be obtained as explicit functions of $(\pi, w, Y; Z, r^*)$ along the DAG, but it would be tedious to write them out.

Keep in mind that the implicit function theorem implies that the response of unknowns is

$$
dU = \underbrace{-H_U^{-1}H_Z}_{G_U} dZ \tag{1}
$$

Recall that we already solved for a steady state and stored is as `ss`. Furthermore, we will use a 300-period truncation horizon.

### 4.1 Set up the DAG

Let's set up the second (main) DAG. We just have to replace `nkpc_ss` with a full Phillips curve.

```python
@simple
def nkpc(pi, w, Z, Y, r, mu, kappa):
    nkpc_res = kappa * (w / Z - 1 / mu) + Y(+1) / Y * (1 + pi(+1)).apply(np.log) / (1 + r(+1))\
               - (1 + pi).apply(np.log)
    return nkpc_res

blocks = [hh_ext, firm, monetary, fiscal, mkt_clearing, nkpc]
hank = create_model(blocks, name="One-Asset HANK")

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

    <SimpleBlock 'monetary'>
    <SimpleBlock 'nkpc'>
    <SimpleBlock 'fiscal'>
    <SimpleBlock 'firm'>
    <HetBlock 'hh' with hetinput 'make_grid_transfers' and with hetoutput `labor_supply'>
    <SimpleBlock 'mkt_clearing'>

- Recall that `create_model` performs a topological sort to put the blocks in a (not necessarily unique) correct order of evaluation. Here, `monetary` must be first, then `nkpc`, `fiscal`, `firm` in any order, then `hh`, and finally `mkt_clearing`. Note that `hank.blocks` is a list in such an order.

Before we proceed, let's make sure that `ss0` is consistent with the second DAG. Just evaluate the `hank` at `ss0` and verify that the equilibrium conditions hold as expected.

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

for k in ss0.keys():
    assert np.all(np.isclose(ss[k], ss0[k]))
```

### 4.2 Obtain GE Jacobian

With the model object `hank` in hand, we can get the general equilibrium Jacobians by using its `solve_jacobian` method.

```python
# setup
T = 300
exogenous = ['rstar', 'Z']
unknowns = ['pi', 'w', 'Y']
targets = ['nkpc_res', 'asset_mkt', 'labor_mkt']

# general equilibrium jacobians
G = hank.solve_jacobian(ss, unknowns, targets, exogenous, T=T)

print(G)
```

    <JacobianDict outputs=['pi', 'w', 'Y', 'A', 'C', 'N', 'NE', 'L', 'Div', 'r', 'Tax', 'asset_mkt', 'labor_mkt', 'goods_mkt', 'nkpc_res'], inputs=['rstar', 'Z']>

Under the hood, the `solve_jacobian` method performs the following steps:
 - computes the partial Jacobians $\mathcal{J}^{o,i}$ for all blocks (if their Jacobian is not supplied already), only with respect to the inputs that actually change: unknowns, exogenous shocks, outputs of earlier blocks
 - forward accumulates partial Jacobians $\mathcal{J}^{o,i}$ to form total Jacobians $\mathbf{J}^{o,i}$
 - packs $\mathbf{J}^{o,i}$ to form $\mathbf{H_U}$ and $\mathbf{H_Z}$
 - solves for the GE Jacobians for unknowns $\mathbf{G_U} = \mathbf{H_U}^{-1}\mathbf{H_Z}$
 - forward accumulates GE Jacobians to obtain $\mathbf{G}$ for other endogenous variables

### 4.3 Results

Now let's consider 25 basis point monetary policy shocks with different persistences and plot the response of inflation.

```python
rhos = np.array([0.2, 0.4, 0.6, 0.8, 0.9])

drstar = -0.0025 * rhos ** (np.arange(T)[:, np.newaxis])
dpi = G['pi']['rstar'] @ drstar

plt.plot(10000 * dpi[:21])
plt.title(r'Inflation responses monetary policy shocks')
plt.xlabel('quarters')
plt.ylabel('bp deviation from ss')
plt.show()
```

## 5. Nonlinear dynamics
Conceptually there's nothing new in this section compared to the Krusell-Smith notebook. We're going to implement a quasi-Newton algorithm to solve for the unknown sequences $U=(\pi, w, Y)$ given some sequences of shocks $(r^*, Z).$ We initialize the algorithm by the naive guess that the variables in $U$ stay constant at their steady-state level. Then we evaluate the DAG and update the guess using the inverse Jacobian $H_U^{-1}.$ The algorithm converges in a few steps, despite the presence of substantial nonlinearities.

### 5.1 A typical monetary policy shock
Note that the linearized solution ignores price adjustment costs. For a monetary policy shock of typical size and persistence, this does not really matter.

```python
rho_r, sig_r = 0.61, -0.01/4
rstar_shock_path = {"rstar": sig_r * rho_r ** (np.arange(T))}

td_nonlin = hank.solve_impulse_nonlinear(ss, unknowns, targets, rstar_shock_path)
td_lin = hank.solve_impulse_linear(ss, unknowns, targets, rstar_shock_path)

dC_nonlin = 100 * td_nonlin['C']
dC_lin = 100 * td_lin['C']

plt.plot(dC_lin[:21], label='linear', linestyle='-', linewidth=2.5)
plt.plot(dC_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()
```

### 5.2 A very large monetary policy shock
However, the nonlinearities may become substantial for very large or persistent monetary policy shocks. Reassuringly, the Jacobian still works well as an updating rule. The quasi-Newton method for a 10% monetary policy shock (extremely large!) still converges below in just 9 iterations, despite nonlinearities evident in the results.

```python
rho_r, sig_r = 0.61, -0.10/4
rstar_shock_path = {"rstar": sig_r * rho_r ** (np.arange(T))}

td_nonlin = hank.solve_impulse_nonlinear(ss, unknowns, targets, rstar_shock_path)
td_lin = hank.solve_impulse_linear(ss, unknowns, targets, rstar_shock_path)

dC_nonlin = 100 * td_nonlin['C']
dC_lin = 100 * td_lin['C']

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