# Exercise 4: Short-time Fourier Transform

Doing this exercise you will learn about the concept of the main lobe width of the spectrum of a window and you will better understand the short-time Fourier transform (STFT). You will also use the STFT to extract basic rhythm related information from an audio signal, implementing an onset detection function, which is one of the rhythm descriptors often used in music information retrieval to detect onsets of acoustic events.

There are four parts in this exercise. 1) Extracting the main lobe of the spectrum of a window, 2) Measuring noise in the reconstructed signal using the STFT model, 3) Computing band-wise energy envelopes of a signal, 4) Computing an onset detection function.

### Relevant Concepts

__Main lobe of the spectrum of a window:__
The width of the main lobe of the magnitude spectrum of a window is an important characteristic used in deciding which window type is best for the analysis of an audio excerpt. There exists a tradeoff between the main lobe width and the side lobe attenuation. Typically for windows with a narrower main lobe, the side lobes are less attenuated.
An interesting fact is that changing the length of a window $M$ doesn't affect the main lobe width of the spectrum of the window in samples. Note that if you use zero-padding for computing the spectrum of a window, the main lobe width will be multiplied by the zero-padding factor.

__Fast Fourier Transform (FFT):__
An efficient way to compute the discrete Fourier transform of a signal is the fast Fourier transform, FFT. The FFT algorithm  factorizes the DFT matrix in order to exploit the symmetries in the DFT equation. FFT computation is specially very efficient when the FFT size is a power of 2. Therefore, whenever possible we use an FFT size that is a power of 2.

__Energy of a signal:__ The energy of a signal $x[n]$ of length $N$ can be computed in the discrete time domain as follows:
\begin{equation}
E=\overset{N-1}{\underset{n=0}{\sum}}\left|x\left[n\right]\right|^{2}
\end{equation}

__Energy in a frequency band:__ Given the DFT spectrum of the signal $X[k]$, the energy $E$ in a specific frequency band spanning the bin index $k_1$ to $k_2$ can be computed as:

\begin{equation}
E=\overset{k_{2}}{\underset{k=k_{1}}{\sum}}\left|X\left[k\right]\right|^{2}
\end{equation}

Note that in this computation the $X[k]$ values are not in decibels (dB). The `stftAnal()` function returns magnitude spectra in dB scale, which should be converted to linear scale before the energy computation. Once the energy is computed it can be converted back to the dB scale as:

\begin{equation}
E_{dB}= 10\,\log_{10}(E)
\end{equation}

__Signal to noise ratio (SNR):__
Signal to noise ratio (SNR) is a frequently used measure to quantify the amount of noise present/added in a signal. In the context of this exercise it can be computed in decibels (dB) as:

\begin{equation}
\mathrm{SNR} = 10\,\log_{10}\left(\frac{E_{\mathrm{signal}}}{E_{\mathrm{noise}}}\right)
\end{equation}

where, $E_{\mathrm{signal}}$ and $E_{\mathrm{noise}}$ are the energy of the signal and noise respectively.

For our case of analysis and synthesis using a specific model (for example, STFT) noise can be thought of as the difference between the input signal and the output signal of the model.

__Onset detection function:__
An onset detection function (ODF) refers to a continuous function (one value per audio frame) often used for detecting acoustic events in an audio stream. In music information retrieval (MIR), ODFs are typically used for detecting onsets of musical notes and percussion strokes. An ODF generally has high values at the onsets of acoustic events. A simple ODF can be computed by taking the difference between the energy values of consecutive frames, as shown below:

\begin{equation}
O(l) = E(l) - E(l-1) ,\,\,\,\,\, l \geq 1
\end{equation}

where, $O(l)$ is the ODF computed at frame index $l$ and $E$ is the energy of the signal in a particular frequency band in decibels (dB). Often, multiple ODFs are computed with different frequency bands across the spectrum.

In order to detect only the onsets of the events and not the offsets, it is a common practice to half wave rectify the ODF and obtain $\bar{O}(l)$. Half wave rectification of the ODF is given by:

 \begin{equation}
 \bar{O}(l)=\begin{cases}
 \begin{array}{c}
 O(l), \, \mathrm{if} \,\,O(l)>0\\
 \,\,\,\,0, \,\,\,\,\, \mathrm{if} \,\, O(l)\leq0
 \end{array}\end{cases}
 \end{equation}

## Part 1 - Extracting the main lobe of the spectrum of a window

Complete the function `extract_main_lobe()` to extract the main lobe of the magnitude spectrum of a window given a window type, `window`, and its length, `M`. The function should return the samples corresponding to the main lobe in decibels (dB).

To compute the spectrum, take the FFT size (`N`) to be 8 times the window length (`N = 8*M`) (For this part, `N` does not need to be a power of 2).

The input arguments to the function are the window type (`window`) and the length of the window (`M`). The function should return a numpy array containing the samples corresponding to the main lobe of the window. In the returned numpy array you should include the samples corresponding to both the local minimas across the main lobe.

The possible window types that you can expect as input are rectangular (`'boxcar'`), `'hamming'` or `'blackmanharris'`.

NOTE: You can approach this question by writing code to find the indices of the local minimas across the main lobe.

_Tip:_ `log10(0)` is not well defined, so its a common practice to add a small value such as `eps = 1e-16`
to the magnitude spectrum before computing it in dB. This is optional and will not affect your answers.
If you find it difficult to concatenate the two halves of the main lobe, you can first center the
spectrum using `fftshift()` and then compute the indexes of the minimas around the main lobe.

```python
import os
import sys
import numpy as np

from scipy.signal import get_window
from scipy.fftpack import fft, fftshift

import math
import matplotlib.pyplot as plt
%matplotlib notebook
eps = np.finfo(float).eps

sys.path.append('../software/models/')
import stft
import utilFunctions as UF
```

```python
def extract_main_lobe(window, M):
    """
    Input:
            window (string): Window type to be used (Either rectangular ('boxcar'), 'hamming' or '
                blackmanharris')
            M (integer): length of the window to be used
    Output:
            a numpy array containing the main lobe of the magnitude spectrum of the window in decibels (dB).
    """

    w = get_window(window, M)         # get the window

    ### Your code here

```

Now run and test the `extract_main_lobe()` function. Here you have 3 test cases:

_Test case 1:_ If you run your code using `window = 'blackmanharris'` and `M = 100`, the output numpy
array should contain 65 samples.

_Test case 2:_ If you run your code using `window = 'boxcar'` and `M = 120`, the output numpy array
should contain 17 samples.

_Test case 3:_ If you run your code using `window = 'hamming'` and `M = 256`, the output numpy array
should contain 33 samples.

To understand better the result you can plot the magnitude spectrum of the window and mark the boundaries of the main lobe.

```python
## Your code here

```

## Part 2 - Measuring noise in the reconstructed signal using the STFT model

Complete the function `compute_snr()` to measure the amount of noise introduced during the analysis and synthesis of a signal using the STFT model. Use SNR (signal to noise ratio) in dB to quantify the amount of noise.
Use the `stft()` function in `stft.py` to do an analysis followed by a synthesis of the input signal.

Use the time domain energy definition to compute the SNR. With the input signal and the obtained output, compute two different SNR values for the following cases:

1. `SNR1`: Over the entire length of the input and the output signals.
2. `SNR2`: For the segment of the signals left after discarding `M` samples from both the start and the end, where `M` is the analysis window length.

Note that the computations are done after STFT analysis and synthesis. The input arguments to the function are the wav file name including the path (`input_file`), window  type (`window`), window length (`M`), FFT size (`N`), and hop size (`H`). The function should return a python tuple of both the SNR values in decibels. Both SNR1 and SNR2 are float values.

```python
def compute_snr(input_file, window, M, N, H):
    """
    Input:
            input_file (string): wav file name including the path
            window (string): analysis window type (choice of rectangular, triangular, hanning, hamming,
                    blackman, blackmanharris)
            M (integer): analysis window length (odd positive integer)
            N (integer): fft size (power of two, > M)
            H (integer): hop size for the stft computation
    Output:
            tuple with the signal to noise ratio over the whole sound and of the sound without the begining and end.
    """
    ## your code here

```

Now run and test `compute_snr()`. Use the following test cases to check your code.

_Test case 1:_ If you run your code using `piano.wav` file with `'blackman'` window, `M = 513`, `N = 2048` and
`H = 128`, the output SNR values should be around: `(67.57748352378475, 304.68394693221649)`.

_Test case 2:_ If you run your code using `sax-phrase-short.wav` file with `'hamming'` window, `M = 512`,
`N = 1024` and `H = 64`, the output SNR values should be around: `(89.510506656299285, 306.18696700251388)`.

_Test case 3:_ If you run your code using `rain.wav` file with `'hann'` window, `M = 1024`, `N = 2048` and
`H = 128`, the output SNR values should be around: `(74.631476225366825, 304.26918192997738)`.

Due to precision differences on different machines/hardware, compared to the expected SNR values, your
output values can differ by +/-10dB for `SNR1` and +/-100dB for `SNR2`.

```python
## Your code here

```

## Part 3 - Computing band-wise energy envelopes of a signal

Complete the function `compute_eng_env()` to compute band-wise energy envelopes of a given audio signal using the STFT.

Consider two frequency bands, low and high. The low frequency band is the set of frequencies between 0 and 3000 Hz and the high frequency band is the set of frequencies between 3000 and 10000 Hz (excluding the boundary frequencies in both the cases). At a given frame, the value of the energy envelope of a band can be computed as the sum of squared values of all the frequency coefficients in that band. Then you should compute the energy envelopes in decibels.

The input arguments are the wav file name including the path (`input_file`), window type (`window`), window length (`M`), FFT size (`N`) and hop size (`H`). The function should return a numpy array with two columns, where the first column is the energy envelope of the low frequency band and the second column is that of the high frequency band.

Use `stftAnal()` function from `stft.py` module to obtain the magnitude spectra for all the audio frames of the input sound. Then compute the energy values for each frequency band specified. While calculating frequency bins for each frequency band, consider only the bins that are within the specified frequency range. For example, for the low frequency band consider only the bins with frequency > 0 Hz and < 3000 Hz (you can use `np.where()` to find those bin indexes). This way you also remove the DC offset of the signal in energy envelope computation. The frequency corresponding to the bin index `k` can be computed as `k*fs/N`, where `fs` is the sampling rate of the signal.

```python
def compute_eng_env(input_file, window, M, N, H):
    """
    Inputs:
            input_file (string): input sound file (monophonic with sampling rate of 44100)
            window (string): analysis window type (choice of rectangular, triangular, hanning,
                hamming, blackman, blackmanharris)
            M (integer): analysis window size (odd positive integer)
            N (integer): FFT size (power of 2, such that N > M)
            H (integer): hop size for the stft computation
    Output:
            magnitude spectra of sound (2D array)
            2D numpy array with energy envelope of band 0 < f < 3000 Hz (in dB) in first column, [:,0] and
            energy envelope of band 3000 < f < 10000 Hz (in dB) in second column [:,1]
    """

    ### your code here

```

Now run and test `compute_eng_env()`. In order to check it you can run and compare these 3 test cases, ploting the results.

_Test case 1:_ Use `piano.wav` file with `window = 'blackman'`, `M = 513`, `N = 1024` and `H = 128` as input.
The bin indexes of the low frequency band span from 1 to 69 (69 samples) and of the high frequency
band span from 70 to 232 (163 samples).

_Test case 2:_ Use `piano.wav` file with `window = 'blackman'`, `M = 2047`, `N = 4096` and `H = 128` as input.
The bin indexes of the low frequency band span from 1 to 278 (278 samples) and of the high frequency
band span from 279 to 928 (650 samples).

_Test case 3:_ Use `sax-phrase-short.wav` file with `window = 'hamming'`, `M = 513`, `N = 2048` and `H = 256` as
input. The bin indexes of the low frequency band span from 1 to 139 (139 samples) and of the high
frequency band span from 140 to 464 (325 samples).

To get a better understanding of the band-wise energy envelope and its characteristics you can plot the envelopes together with the spectrogram of the signal. You can use `matplotlib` plotting library for this purpose. To visualize the spectrogram of a signal, a good option is to use the function `colormesh()` (you can reuse the code in
`lectures/4-STFT/plots-code/spectrogram.py`). Either overlay the envelopes on the spectrogram
or plot them in a different subplot. Make sure you use the same range of the x-axis for both the spectrogram and the energy envelopes.

You can clearly notice the sharp attacks and decay of the piano notes for test case 1. You can compare this with the output from test case 2 that uses a larger window. You can infer the influence of window size on sharpness of the note attacks.

```python
## Your code here

```

## Part 4 - Computing onset detection function

Complete the function `compute_odf()` to compute a simple onset detection function (ODF) using the STFT. Compute two ODFs one for each of the frequency bands, low and high. The low frequency band is the set of all the frequencies between 0 and 3000 Hz and the high frequency band is the set of all the frequencies between 3000 and 10000 Hz (excluding the boundary frequencies in both the cases).

Start with an initial condition of `ODF(0) = 0` in order to make the length of the ODF same as that of the energy envelope. Remember to apply a half wave rectification on the ODF.

The input arguments to the function are the wav file name including the path (`input_file`), window
type (`window`), window length (`M`), FFT size (`N`), and hop size (`H`). The function should return a numpy
array with two columns, where the first column is the ODF computed on the low frequency band and the
second column is the ODF computed on the high frequency band.

Use `stftAnal()` to obtain the magnitude spectra for all the audio frames. Then compute the band-wise energy envelope for each frequency band specified. Finally compute the half wave rectify version of ODF from each energy envelope.

```python
def compute_odf(input_file, window, M, N, H):
    """
    Inputs:
            input_file (string): input sound file (monophonic with sampling rate of 44100)
            window (string): analysis window type (choice of rectangular, triangular, hanning, hamming,
                blackman, blackmanharris)
            M (integer): analysis window size (odd integer value)
            N (integer): fft size (power of two, bigger or equal than than M)
            H (integer): hop size for the STFT computation
    Output:
            magnitude spectra of sound (2D array)
            2D numpy array with ODF computed in band 0 < f < 3000 Hz (in dB) in first column, [:,0] and
            ODF computed of band 3000 < f < 10000 Hz (in dB) in second column [:,1]
    """

    ### your code here

```

Now run and test `compute_odf()`. In order to check it you can run and compare these 3 test cases. Plot the results.

_Test case 1:_ Use `piano.wav` file with `window = 'blackman'`, `M = 513`, `N = 1024` and `H = 128` as input.
The bin indexes of the low frequency band span from 1 to 69 (69 samples) and of the high frequency
band span from 70 to 232 (163 samples).

_Test case 2:_ Use `piano.wav` file with `window = 'blackman'`, `M = 2047`, `N = 4096` and `H = 128` as input.
The bin indexes of the low frequency band span from 1 to 278 (278 samples) and of the high frequency
band span from 279 to 928 (650 samples).

_Test case 3:_ Use `sax-phrase-short.wav` file with `window = 'hamming'`, `M = 513`, `N = 2048` and `H = 256` as
input. The bin indexes of the low frequency band span from 1 to 139 (139 samples) and of the high
frequency band span from 140 to 464 (325 samples).

To get a better understanding of the ODFs and their characteristics you can plot the ODF functions together with the spectrogram of the signal. Use the same plotting than in the previos part.

In order to identify the actual onsets of a signal you would need to find the peaks, local maxima, of the ODF functions using a magnitude threshold. Which of the 2 ODFs would be most useful? What would be the right threshold for all 3 cases? Could you find a single threshold to be used for all 3 cases?

For test case 1, you should clearly see that the ODFs have sharp peaks at the onset of the piano notes. You should notice 5 peaks that are above 10dB value in the ODF computed on the high frequency band.

```python
## Your code here

```

## Your explanation of Part 4 here


# Further information

## How can we plot a function

It is possible to plot a function using `sympy` using the `sympy.plot` function:

````{tip}
```
sympy.plot(expression)
```
````

So for example, here is a plot of $f(x)=x^2 + 3x + 1$:

```python
import sympy as sym

x = sym.Symbol("x")
sym.plot(x ** 2 + 3 * x + 1)
```

It is possible to specify the x limits and combine it with other plots:

```python
sym.plot(x ** 2 + 3 * x + 1, xlim=(-5, 5))
```

**This plotting solution is fine it you want to take a look at a function
quickly but it is not recommended.** The main library for plotting is called
`matplotlib` and {ref}`matplotlib` covers this.

- Here is the `sympy` documentation for plotting:
  <https://docs.sympy.org/latest/modules/plotting.html>
- Here is the official `matplotlib` documentation: <https://matplotlib.org>


# Part 1 - Scalars and Vectors

For the questions below it is not sufficient to simply provide answer to the questions, but you must solve the problems and show your work using python (the NumPy library will help a lot!) Translate the vectors and matrices into their appropriate python  representations and use numpy or functions that you write yourself to demonstrate the result or property.

```
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
```

## 1.1 Create a two-dimensional vector and plot it on a graph

```
v1 = [300,250]
plt.arrow(0,0, 300, 250, head_width=10, head_length=10, color='r');
plt.title('Red Vector')
plt.axis([0, 350, 0, 300]); #([-x, x, -y, y]) limits of the graph - BRACKETS
```

## 1.2 Create a three-dimensional vecor and plot it on a graph

```
vectors = np.array( [[0,0,0,-1,-0,1]])
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
for vector in vectors:
  v = np.array([vector[3],vector[4],vector[5]])
  vlength=np.linalg.norm(v)
  ax.quiver(vector[0], vector[1], vector[2], vector[3], vector[4], vector[5],
         pivot='tail',length=vlength,arrow_length_ratio=0.3/vlength)
ax.set_xlim([-3,3])
ax.set_ylim([-3,3])
ax.set_zlim([0,3]);
```

## 1.3 Scale the vectors you created in 1.1 by $5$, $\pi$, and $-e$ and plot all four vectors (original + 3 scaled vectors) on a graph. What do you notice about these vectors?

```
from math import e, pi
```

```
v1 = [300, 250]
v2 =[(300*5), (250*5)]
v3 = [(300*pi), (250*pi)]
v4 = [(300*-e), (250*-e)]
```

```
#the vectors and scalers
v1 = [300, 250]
v2 =[(300*5), (250*5)]
v3 = [(300*pi), (250*pi)]
v4 = [(300*-e), (250*-e)]

#plot the vectors
original = plt.arrow(0,0, v1[0], v1[1], head_width=100
                     , head_length=100, color='r');
scaled_by_5 = plt.arrow(0,0, v2[0], v2[1], head_width=100
                        , head_length=100, color='b');
scaled_by_pi = plt.arrow(0,0, v3[0], v3[1], head_width=100
                         , head_length=100, color='y');
scaled_by_neg_e = plt.arrow(0,0, v4[0], v4[1], head_width=100
                            , head_length=100, color='black');

#reference lings and limits
plt.grid()
plt.axhline(y=0, alpha=.5)
plt.axvline(x=0, alpha=.5)
plt.axis([-1600, 1600, -1600, 1600]);
```

These vectors are colinear

## 1.4 Graph vectors $\vec{a}$ and $\vec{b}$ and plot them on a graph

\begin{align}
\vec{a} = \begin{bmatrix} 5 \\ 7 \end{bmatrix}
\qquad
\vec{b} = \begin{bmatrix} 3 \\4 \end{bmatrix}
\end{align}

```
a = np.array([5,7])
b = np.array([3,4])
plt.arrow(0,0, a[0], a[1], head_width=.2, head_length=.2, color='r');
plt.arrow(0,0, b[0], b[1], head_width=.2, head_length=.2, color='b');

plt.axis([0,8,0,8]);
```

## 1.5 find $\vec{a} - \vec{b}$ and plot the result on the same graph as $\vec{a}$ and $\vec{b}$. Is there a relationship between vectors $\vec{a} \thinspace, \vec{b} \thinspace \text{and} \thinspace \vec{a-b}$

```
a = np.array([5,7])
b = np.array([3,4])
difference = a-b
plt.arrow(0,0, a[0], a[1], head_width=.2, head_length=.2, color='r');
plt.arrow(0,0, b[0], b[1], head_width=.2, head_length=.2, color='b');
plt.arrow(0,0, difference[0], difference[1], head_width=.2, head_length=.2, color='k');

print('The result is:',difference)
plt.axis([0,8,0,8]);c = a-b
```

The relationship is taking the values from the first and subtracting the values of the 2nd. The 3rd vector is equal to (2, 3)

## 1.6 Find $c \cdot d$

\begin{align}
\vec{c} = \begin{bmatrix}7 & 22 & 4 & 16\end{bmatrix}
\qquad
\vec{d} = \begin{bmatrix}12 & 6 & 2 & 9\end{bmatrix}
\end{align}

```

```

##  1.7 Find $e \times f$

\begin{align}
\vec{e} = \begin{bmatrix} 5 \\ 7 \\ 2 \end{bmatrix}
\qquad
\vec{f} = \begin{bmatrix} 3 \\4 \\ 6 \end{bmatrix}
\end{align}

```

```

## 1.8 Find $||g||$ and then find $||h||$. Which is longer?

\begin{align}
\vec{e} = \begin{bmatrix} 1 \\ 1 \\ 1 \\ 8 \end{bmatrix}
\qquad
\vec{f} = \begin{bmatrix} 3 \\3 \\ 3 \\ 3 \end{bmatrix}
\end{align}

```

```

## 1.9 Show that the following vectors are orthogonal (perpendicular to each other):

\begin{align}
\vec{g} = \begin{bmatrix} 1 \\ 0 \\ -1  \end{bmatrix}
\qquad
\vec{h} = \begin{bmatrix} 1 \\ \sqrt{2} \\ 1 \end{bmatrix}
\end{align}

```

```

# Part 2 - Matrices

## 2.1 What are the dimensions of the following matrices? Which of the following can be multiplied together? See if you can find all of the different legal combinations.
\begin{align}
A = \begin{bmatrix}
1 & 2 \\
3 & 4 \\
5 & 6
\end{bmatrix}
\qquad
B = \begin{bmatrix}
2 & 4 & 6 \\
\end{bmatrix}
\qquad
C = \begin{bmatrix}
9 & 6 & 3 \\
4 & 7 & 11
\end{bmatrix}
\qquad
D = \begin{bmatrix}
1 & 0 & 0 \\
0 & 1 & 0 \\
0 & 0 & 1
\end{bmatrix}
\qquad
E = \begin{bmatrix}
1 & 3 \\
5 & 7
\end{bmatrix}
\end{align}

```

```

## 2.2 Find the following products: CD, AE, and BA. What are the dimensions of the resulting matrices? How does that relate to the dimensions of their factor matrices?

```

```

## 2.3  Find $F^{T}$. How are the numbers along the main diagonal (top left to bottom right) of the original matrix and its transpose related? What are the dimensions of $F$? What are the dimensions of $F^{T}$?

\begin{align}
F =
\begin{bmatrix}
20 & 19 & 18 & 17 \\
16 & 15 & 14 & 13 \\
12 & 11 & 10 & 9 \\
8 & 7 & 6 & 5 \\
4 & 3 & 2 & 1
\end{bmatrix}
\end{align}

```

```

# Part 3 - Square Matrices

## 3.1 Find $IG$ (be sure to show your work) 😃

\begin{align}
G=
\begin{bmatrix}
12 & 11 \\
7 & 10
\end{bmatrix}
\end{align}

```

```

## 3.2 Find $|H|$ and then find $|J|$.

\begin{align}
H=
\begin{bmatrix}
12 & 11 \\
7 & 10
\end{bmatrix}
\qquad
J=
\begin{bmatrix}
0 & 1 & 2 \\
7 & 10 & 4 \\
3 & 2 & 0
\end{bmatrix}
\end{align}

```

```

## 3.3 Find H^{-1} and then find J^{-1}

```

```

## 3.4 Find $HH^{-1}$ and then find $G^{-1}G$. Is $HH^{-1} == G^{-1}G$? Why or Why not?

# Stretch Goals:

A reminder that these challenges are optional. If you finish your work quickly we welcome you to work on them. If there are other activities that you feel like will help your understanding of the above topics more, feel free to work on that. Topics from the Stretch Goals sections will never end up on Sprint Challenges. You don't have to do these in order, you don't have to do all of them.

- Write a function that can calculate the dot product of any two vectors of equal length that are passed to it.
- Write a function that can calculate the norm of any vector
- Prove to yourself again that the vectors in 1.9 are orthogonal by graphing them.
- Research how to plot a 3d graph with animations so that you can make the graph rotate (this will be easier in a local notebook than in google colab)
- Create and plot a matrix on a 2d graph.
- Create and plot a matrix on a 3d graph.
- Plot two vectors that are not collinear on a 2d graph. Calculate the determinant of the 2x2 matrix that these vectors form. How does this determinant relate to the graphical interpretation of the vectors?


# Modelproject:

**Start by importing packages and modules**

```python
import sympy as sm
from scipy import optimize
import matplotlib.pyplot as plt
import ipywidgets as widgets
```

To print nicely:

```python
sm.init_printing(use_unicode=True)
```

# The basic Solow model

The Solow model is a model that describes the process of capital accumulation by which saving and investment become new capital.
The essential feature of the Solow model is that it incorporates the dynamic link between the flows of savings and investment and the stock of capital.
All in all the model describes how capital evolves as a result of capital accumulation, how the labour force evolves as a result of population growth, and how total production and income evolves as a consequence of the evolutions of the total input of capital and labour.
The Solow model predicts that in the long run the capital intensity and GDP pr. worker converge to particular steady state levels, $k^*$ nad $y^*$, respectively.
The steady states shows us which fundamental characters of an economy that can create a high level of income per worker in the long run according to the Solow model.
If you are sitting down thinking:

The Solow model shows growth in money and we like steady states.

The basic Solow model is as follows:

\begin{equation}
Y_t = BK_t^{\alpha}L_t^{1-\alpha} \tag{1}
\end{equation}
\begin{equation}
r_t = \alpha B \left(\frac{K_t}{L_t}\right)^{\alpha-1} \tag{2}
\end{equation}
\begin{equation}
w_t = (1-\alpha) B \left(\frac{K_t}{L_t}\right)^{\alpha-1} \tag{3}
\end{equation}
\begin{equation}
S_t = sY_t \tag{4}
\end{equation}
\begin{equation}
K_{t+1} - K_t = S_t -\delta K_t \tag{5}
\end{equation}
\begin{equation}
L_{t+1} = (1-n)L_t \tag{6}
\end{equation}

We also define the following equations:
\begin{equation}
y_t = \frac{Y_t}{L_t} = Bk_t^{\alpha}
\end{equation}
\begin{equation}
k_t = \frac{K_t}{L_t}
\end{equation}
The meaning of all the letters can be seen below.

## Defining symbols

First we define the necessary variables so we can use sympy:

```python
Y = sm.symbols("Y_t")         #GDP in the economy
y = sm.symbols("y_t")         #GDP per worker
K = sm.symbols("K_t")         #Capital in the economy
k = sm.symbols("k_t")         #Capital per worker
r = sm.symbols("r_t")         #Real rental rate
w = sm.symbols("w_t")         #Real wage
s = sm.symbols("s")           #Rate of savings
n = sm.symbols("n")           #Population growth
g = sm.symbols("g")           #Growth in technology
B = sm.symbols("B")           #Total factor productivity
L = sm.symbols("L_t")         #Labour
alpha = sm.symbols("alpha")   #Return on capital
delta = sm.symbols("delta")   #Depreciation rate on capital
kstar = sm.symbols("k_t^*")   #Steady state for capital
ystar = sm.symbols("y^*")     #Steady state for GDP per worker
```

It is easy to verify equation (2) and (3) by differentiating equation (1) with respect to respectively $K_t$ and $L_t$:

$r_t$:

```python
sm.simplify(sm.diff(B*K**alpha*L**(1-alpha), K))
```

$w_t$:

```python
sm.simplify(sm.diff(B*K**alpha*L**(1-alpha), L))
```

## Solving for steady state symbolically

We find the steady state by solving the equation $k_{t+1}-k_t = \frac{1}{1+n} (sBk_t^{\alpha} -(n+\delta)k_t)$ for the constant solution where $k_{t+1} = k_t = k^*$. This gives us the equation $sBk^{\alpha} - (n+\delta)k = 0$ to solve.

```python
ksteady = sm.solve(s*B*k**alpha-(n+delta)*k, k)[0]
sm.Eq(kstar , ksteady)
```

This helps us solving for $y^*$:

```python
eq1 = sm.simplify(B*ksteady**alpha)
sm.Eq(ystar, eq1)
```

## Solving for steady state numerically

We can also find the steady state values as real numbers. There is more ways of doing this, as shown below.

### Newton method

First we start of with the Newton method:

```python
def basic_newton(x, alpha=1/3, s=0.4, B=1, n=0.01, delta=0.01):
    """
    Args:
    x      (float)          : a value which is estimated to be around the steady state, if the function returns 0, set this higher.
    alpha  (float, optional): return on physical capital
    s      (float, optional): savings/investments
    B      (int,   optional): totalfactorproductivity
    n      (float, optional): increase in labour
    delta  (float, optional): depreciation rate on capital

    Returns:
    Steady state value for k*
    """
    alpha = alpha
    s = s
    B = B
    n = n
    delta = delta
    return optimize.newton(lambda k: s*B*k**alpha - (n+delta)*k, x)
```

With this method, we can find the steady state value and compare it to the symbolical solution:

```python
print(f'Steady state with Newton method is: {basic_newton(100):.4f}')
print(f'Steady state symbolically is: {((1*0.4)/(0.01+0.01))**(-1/(1/3-1)):.4f}')
```

    Steady state with Newton method is: 89.4427
    Steady state symbolically is: 89.4427

### Bisect method

We can also try with the bisect method:

```python
def basic_bisect(a,b, alpha=1/3, s=0.4, B=1, n=0.01, delta=0.01):
    """
    Args:
    a      (float)          : starting point of the interval [a,b]
    b      (float)          : ending point of the interval [a,b], if the function returns 0, set this higher.
    alpha  (float, optional): return on physical capital
    s      (float, optional): savings/investments
    B      (int,   optional): totalfactorproductivity
    n      (float, optional): increase in labour
    delta  (float, optional): depreciation rate on capital

    Returns:
    Steady state value for k*
    """
    alpha = alpha
    s = s
    B = B
    n = n
    delta = delta
    return optimize.bisect(lambda k: s*B*k**alpha - (n+delta)*k, a,b, full_output=False)
print(f'Steady state with bisect method is: {basic_bisect(1,100):.4f}')
```

    Steady state with bisect method is: 89.4427

We see that we again get the same steady state as both the newton method and the symbolical method.

### Brent-method

The last one we will try is the Brent-method:

```python
def basic_brentq(a,b, alpha=1/3, s=0.4, B=1, n=0.01, delta=0.01):
    """
    Args:
    a      (float)          : starting point of the interval [a,b]
    b      (float)          : ending point of the interval [a,b], if the function returns 0, set this higher.
    alpha  (float, optional): return on physical capital
    s      (float, optional): savings/investments
    B      (int,   optional): totalfactorproductivity
    n      (float, optional): increase in labour
    delta  (float, optional): depreciation rate on capital

    Returns:
    Steady state value for k*
    """
    alpha = alpha
    s = s
    B = B
    n = n
    delta = delta
    return optimize.brentq(lambda k: s*B*k**alpha - (n+delta)*k, a,b, full_output=False)
print(f'Steady state with brentq method is: {basic_brentq(1,100):.4f}')
```

    Steady state with brentq method is: 89.4427

### Comparing

Below we see the 3 numerical solutions compared to the symbolical solution. We notice that even though they look much the same, there is still a minor difference when we reach the last decimals.

```python
print(f'Steady state with Newton method is: {basic_newton(100):0.15f}')
print(f'Steady state with bisect method is: {basic_bisect(1,100):0.15f}')
print(f'Steady state with brentq method is: {basic_brentq(1,100):0.15f}')
print(f'Steady state symbolically is:       {((1*0.4)/(0.01+0.01))**(-1/(1/3-1)):0.15f}')
```

    Steady state with Newton method is: 89.442719099991592
    Steady state with bisect method is: 89.442719099990555
    Steady state with brentq method is: 89.442719099991578
    Steady state symbolically is:       89.442719099991535

## Simulating the basic Solow model

```python
def simulate_basic(k,alpha,delta,s,n,B,T):
    """
    Args:
    k      (int)  : initial value for physical capital per worker
    alpha  (float): return on physical capital
    delta  (float): depreciation rate on capital
    s      (float): savings/investments
    n      (float): increase in labour
    B      (int)  : totalfactorproductivity
    T      (int)  : periods to simulate through

    Returns:
    A plot showing how capital per worker accumulates over time"""

    #First we make a list for our future values to be in
    k_path = [k]

    #Then we make which will contain the values for the 45-degree-line
    Degreeline = [0]

    #We generate the 45-degree-line for the basic Solow model
    for t in range(1,T):
        line = (n + delta)*t
        Degreeline.append(line)

    #We generate the Solow movement
    for t in range(1,T):
        k_plus = s*B*t**alpha
        k_path.append(k_plus)

    #Plotting the stuff
    plt.figure(figsize=(5,5))
    plt.plot(k_path[:T], label='$sBk_t^{\u03B1}$', color = 'blue')
    plt.plot(Degreeline[:T], label = '$(n+\delta)k_t$', color = 'black')
    plt.xlim(0,T)
    plt.ylim(0,Degreeline[-1])
    plt.xlabel('$k_t$')
    plt.grid(True)
    plt.legend()

    return plt.show()
```

Example on a simulation with reasonable values:

```python
simulate_basic(0,1/3,0.01,0.4,0.01,1,200)
```

Below, you can try changing the parameters yourself to see how they influence the steady state.

```python
widgets.interact(simulate_basic,
                k     = widgets.fixed(0),
                alpha = widgets.FloatSlider(description = '$\u03B1$' , min = 0 ,    max = 0.9 , step = 0.05 , value = 0.3),
                delta = widgets.FloatSlider(description = '$\u03B4$' , min = 0 ,    max = 0.1 , step = 0.01 , value = 0.01),
                s     = widgets.FloatSlider(description = 's' ,        min = 0.01 , max = 0.8 , step = 0.05 , value = 0.4),
                n     = widgets.FloatSlider(description ='$n$',        min = 0.01 , max = 0.1 , step = 0.005, value = 0.01),
                B     = widgets.fixed(1),
                T     = widgets.IntSlider(description='$T$' ,          min = 1,     max = 1000, step = 10,    value = 100))
```

    interactive(children=(FloatSlider(value=0.3, description='$α$', max=0.9, step=0.05), FloatSlider(value=0.01, d…

    <function __main__.simulate_basic(k, alpha, delta, s, n, B, T)>

# Extension: The solow model with human capital

We will now introduce human capital which covers things like education and knowledge etc.
The return on human capital is given by $\varphi$.
The savings are now divided into savings in physical capital $s_K$ and human capital $s_H$.
The Solow model with human capital is as follows:

\begin{equation}
Y_t = K_t^{\alpha}H_t^{\varphi}(A_tL_t)^{1-\alpha -\varphi} \tag{1} \\
\end{equation}
\begin{equation}
r_t = \alpha \left(\frac{K_t}{A_t L_t}\right)^{\alpha -1} \left(\frac{H_t}{A_t L_t}\right)^{\varphi} \tag{2}
\end{equation}
\begin{equation}
w_t = (1-\alpha) \left(\frac{K_t}{A_tL_t}\right)^{\alpha} \left(\frac{H_t}{A_tL_t}\right)^{\varphi}A_t \tag{3}
\end{equation}
\begin{equation}
K_{t+1}-K_t = s_K Y_t - \delta K_t \quad , \quad K_0 \text{ given} \tag{4}
\end{equation}
\begin{equation}
H_{t+1}-H_t = s_H Y_t -\delta H_t \quad , \quad H_0 \text{ given} \tag{5}
\end{equation}
\begin{equation}
L_{t+1} = (1+n)L_t \quad , \quad L_0 \text{ given} \tag{6}
\end{equation}
\begin{equation}
A_{t+1} = (1+g)A_t \quad , \quad A_0 \text{ given} \tag{7}
\end{equation}

We also define the following equations:
\begin{equation}
\text{Physical capital per effective worker}: \quad \tilde{k_t} = \frac{K_t}{A_tL_t}
\end{equation}
\begin{equation}
\text{Human capital per effective worker}: \quad \tilde{h_t} = \frac{H_t}{A_tL_t}
\end{equation}
\begin{equation}
\text{Output per effective worker}: \quad \tilde{y_t} = \frac{Y_t}{A_tL_t}
\end{equation}
\begin{equation}
\tilde{y_t} = \tilde{k_t}^{\alpha} \tilde{h_t}^{\varphi}
\end{equation}

## Defining symbols

Defining all the new things as symbols so we can use sympy:

```python
A = sm.symbols("A_t")                 #Technology level
H = sm.symbols("H_t")                 #Human capital
h = sm.symbols("h_t")                 #Human capital per effective worker
sK = sm.symbols("s_K")                #Savings in physical capital
sH = sm.symbols("s_H")                #Savings in human capital
phi = sm.symbols("varphi")            #Return on human capital
ktilde = sm.symbols("ktilde_t")       #Physical capital per effective worker
htilde = sm.symbols("htilde_t")       #Human capital per effective worker
ytilde = sm.symbols("ytilde_t")       #BNP per effective worker
htilde1 = sm.symbols("htilde_t+1")    #Human capital per effective worker in period t+1
ktilde1 = sm.symbols("ktilde_t+1")    #Physical capital per effective worker in period t+1
ktildestar = sm.symbols("ktilde^*")   #Steady state for physical capital per effective worker
htildestar = sm.symbols("htilde^*")   #Steady state for human capital per effective worker
ytildestar = sm.symbols("ytilde^*")   #Steady state for BNP per effective worker
```

We can find the real rental rate and wage just as in the basic Solow model:

$r_t$:

```python
sm.simplify(sm.diff(K**alpha*H**phi*(A*L)**(1-alpha-phi), K))
```

$w_t$ (we use that $H_t = h_tL_t$):

```python
sm.simplify(sm.diff(K**alpha*h**phi*L**(1-alpha)*A**(1-alpha-phi), L))
```

## Solving for steady state symbolically

We will now find the steady state values for variables $\tilde{k^*}, \tilde{h^*}$ and $\tilde{y^*}$ by solving the following equations:

\begin{equation}
\tilde{k}_{t+1} - \tilde{k}_t = \frac{1}{(1+n)(1+g)}(s_K \tilde{k}_t^{\alpha} \tilde{h}_t^{\varphi} - (n+g+\delta +ng)\tilde{k}_t)
\end{equation}
\begin{equation}
\tilde{h}_{t+1} - \tilde{h}_t = \frac{1}{(1+n)(1+g)}(s_H \tilde{k}_t^{\alpha} \tilde{h}_t^{\varphi} - (n+g+\delta +ng)\tilde{h}_t)
\end{equation}

First we make variables which transform $\tilde{h}_t$ into a function og $\tilde{k}_t$ and vice verca:

```python
ktildealpha = ktilde**alpha
htildephi = htilde**phi
htok = (sH/sK)*ktilde                 #Turning h into a function of k
htokphi = ((sH/sK)*ktilde)**phi
ktoh = (sK/sH)*htilde                 #Turning k into a function of h
ktohalpha = ((sK/sH)*htilde)**alpha
```

We set up the equation to solve for $\tilde{h}_{t+1} = \tilde{h}_t = \tilde{h}^*$:

```python
eq = sm.Eq(sH*ktildealpha*htildephi - (n+g+delta+n*g)*htilde)
eq
```

Then we solve it for $\tilde{h}_t$:

```python
hisolated = sm.solve(eq,htilde)[0]
hisolated
```

Now we insert the transformed $\tilde{h}_t$:

```python
htildefinal = sm.Eq(htilde , ((ktoh**(-alpha)*(delta+n*g+g+n)/sH)**(1/(phi-1))))
htildefinal
```

And solve for the steady state:

```python
htildesteady = sm.solve(htildefinal,htilde)[0]
print('We find that the steady state value for htilde is:')
sm.Eq(htildestar, htildesteady)
```

We now do the same for $\tilde{k}_t$:

```python
eq2 = sm.Eq(sK*ktildealpha*htildephi - (n+g+delta+n*g)*ktilde)
eq2
```

```python
kisolated = sm.solve(eq2,ktilde)[0]
kisolated
```

```python
ktildefinal = sm.Eq(ktilde, ((htok**(-phi)*(delta+n*g+g+n)/sK)**(1/(alpha-1))))
ktildefinal
```

```python
ktildesteady = sm.solve(ktildefinal,ktilde)[0]
print('We find that the steady state value for ktilde is:')
sm.Eq(ktildestar, ktildesteady)
```

Now we can plot this into $\tilde{y}^*$:

```python
yfinal = sm.Eq(ytildestar, ktildesteady**alpha * htildesteady**phi)
ytildesteady = sm.solve(yfinal, ytildestar)[0]
sm.Eq(ytildestar, sm.solve(yfinal, ytildestar)[0])
```

We now get that the BNP per capita in steady state must be:

```python
sm.Eq(ystar, A*ytildesteady)
```

We can now create a function which gives us the steady state values for chosen parameter:

```python
ss_htilde = sm.lambdify((sK, sH, n, g, delta, alpha, phi), htildesteady)
ss_ktilde = sm.lambdify((sK, sH, n, g, delta, alpha, phi), ktildesteady)
ss_ytilde = sm.lambdify((sK, sH, n, g, delta, alpha, phi), ytildesteady)
def steady_states(sK, sH, n, g, delta, alpha, phi):

    """
    Args:
    sK    (float): investments in physical income
    sH    (float): investments in human capital
    n     (float): increase in labour
    g     (float): increase in technology
    delta (float): the depreciation rate
    alpha (float): return on physical capital
    phi   (float): return on human capital

    Returns:
    Steady state values for h, k and y"""

    return print(f'Steady state for h is {ss_htilde(sK, sH, n, g, delta, alpha, phi):0.2f}'), print(f'Steady state for k is {ss_ktilde(sK, sH, n, g, delta, alpha, phi):0.2f}'), print(f'Steady state for BNP per worker is {ss_ytilde(sK, sH, n, g, delta, alpha, phi):0.2f}')
```

An example on finding steady states could be:

```python
steady_states(0.2,0.2,0.01,0.02,0.01,1/3,1/3)
```

    Steady state for h is 123.14
    Steady state for k is 123.14
    Steady state for BNP per worker is 24.75

    (None, None, None)

## Solving for steady state numerically

We do the same as in the basic model.

### Newton method

```python
def human_newton(x=1000, alpha=0.3, phi=0.3, delta=0.01, sK=0.2, sH=0.2, n=0.01, g=0.02):
    """
    Args:
    x      (float, optional): a value which is estimated to be around the steady state.
    alpha  (float, optional): return on physical capital
    phi    (float, optional): return on human capital
    delta  (float, optional): depreciation rate
    sK     (float, optional): investments in physical income
    sH     (float, optional): investments in human capital
    n      (float, optional): increase in labour
    g      (float, optional): increase in technology

    Returns:
    Steady state value for k and h*
    """
    alpha = alpha
    sK = sK
    sH = sH
    phi = phi
    n = n
    delta = delta
    g = g
    return print(f'Steady state for k with Newton method is: {optimize.newton(lambda k: sK*k**alpha*((sH/sK)*k)**phi-(n+g+delta+n*g)*k, x):.15f} \nSteady state for h with Newton method is: {optimize.newton(lambda h: sH*((sK/sH)*h)**alpha*h**phi-(n+g+delta+n*g)*h, x):.15f}')
human_newton()
```

    Steady state for k with Newton method is: 55.208996899268655
    Steady state for h with Newton method is: 55.208996899268655

### Bisect method

```python
def human_bisect(a=1, b=100, alpha=0.3, phi=0.3, delta=0.01, sK=0.2, sH=0.2, n=0.01, g=0.02):
    """
    Args:
    a      (float)          : starting point of the interval [a,b]
    b      (float)          : ending point of the interval [a,b], if the function returns 0, set this higher.
    alpha  (float, optional): return on physical capital
    phi    (float, optional): return on human capital
    delta  (float, optional): depreciation rate
    sK     (float, optional): investments in physical income
    sH     (float, optional): investments in human capital
    n      (float, optional): increase in labour
    g      (float, optional): increase in technology

    Returns:
    Steady state value for k and h*
    """
    alpha = alpha
    sK = sK
    sH = sH
    phi = phi
    n = n
    delta = delta
    g = g
    return print(f'Steady state for k with bisect method is: {optimize.bisect(lambda k: sK*k**alpha*((sH/sK)*k)**phi-(n+g+delta+n*g)*k, a, b):.15f} \nSteady state for h with bisect method is: {optimize.bisect(lambda h: sH*((sK/sH)*h)**alpha*h**phi-(n+g+delta+n*g)*h, a, b):.15f}')
human_bisect()
```

    Steady state for k with bisect method is: 55.208996899267603
    Steady state for h with bisect method is: 55.208996899267603

### Brent method

```python
def human_brentq(a=1, b=100, alpha=0.3, phi=0.3, delta=0.01, sK=0.2, sH=0.2, n=0.01, g=0.02):
    """
    Args:
    a      (float)          : starting point of the interval [a,b]
    b      (float)          : ending point of the interval [a,b], if the function returns 0, set this higher.
    alpha  (float, optional): return on physical capital
    phi    (float, optional): return on human capital
    delta  (float, optional): depreciation rate
    sK     (float, optional): investments in physical income
    sH     (float, optional): investments in human capital
    n      (float, optional): increase in labour
    g      (float, optional): increase in technology

    Returns:
    Steady state value for k and h*
    """
    alpha = alpha
    sK = sK
    sH = sH
    phi = phi
    n = n
    delta = delta
    g = g
    return print(f'Steady state for k with brentq method is: {optimize.brentq(lambda k: sK*k**alpha*((sH/sK)*k)**phi-(n+g+delta+n*g)*k, a, b):.15f} \nSteady state for h with brentq method is: {optimize.brentq(lambda h: sH*((sK/sH)*h)**alpha*h**phi-(n+g+delta+n*g)*h, a, b):.15f}')
human_brentq()
```

    Steady state for k with brentq method is: 55.208996899268605
    Steady state for h with brentq method is: 55.208996899268605

We again compare all the numerical steady states below:

```python
print(f'{human_newton()} \n {human_bisect()} \n {human_brentq()}')
```

    Steady state for k with Newton method is: 55.208996899268655
    Steady state for h with Newton method is: 55.208996899268655
    Steady state for k with bisect method is: 55.208996899267603
    Steady state for h with bisect method is: 55.208996899267603
    Steady state for k with brentq method is: 55.208996899268605
    Steady state for h with brentq method is: 55.208996899268605
    None
     None
     None

We see again that the methods is quite similar, but we notice again that there is still very small differences between them when we reach the last decimals.

## Simulating the Solow model with human capital

We will now simulate the Solow model with human capital, to graphically show how to find the steady state.

```python
def simulate_human_capital(htilde,ktilde,alpha,phi,delta,sK,sH,n,g,T):
    """
    Args:
    htilde (int)  : initial value for human capital per effective worker, normally = 0
    ktilde (int)  : initial value for physical capital per effective worker, normally = 0
    alpha  (float): return on physical capital
    phi    (float): return on human capital
    delta  (float): depreciation rate
    sK     (float): investments in physical income
    sH     (float): investments in human capital
    n      (float): increase in labour
    g      (float): increase in technology
    T      (int)  : periods to simulate through

    Returns:
    A plot showing how the capital accumulates over time, as well as steady state values
    """
    #Making a list for all our future values
    htilde_path = [htilde]
    ktilde_path = [ktilde]

    #Adding numbers to the list using equation (5) where htilde is isolated
    for t in range(1,T):
        ktilde = t
        htilde_next = ((ktilde**(-alpha)*(delta+n*g+n+g))/sH)**(1/(phi-1))
        htilde_path.append(htilde_next)

    #Adding numbers to the list using equation (4) where htilde is isolated
    for t in range(1,T):
        ktilde = t
        ktilde_next = ((ktilde**(-alpha+1)*(delta+n*g+n+g))/sK)**(1/(phi))
        ktilde_path.append(ktilde_next)

    #Plotting it all
    plt.figure(figsize=(5,5))
    plt.plot(htilde_path[:T], linewidth = 2, label='$\Delta \~{h} = 0$', color = 'navy')
    plt.plot(ktilde_path[:T], linewidth = 2, label='$\Delta \~{k} = 0$', color = 'lime')
    plt.xlim(0,T)
    plt.ylim(0,T)
    plt.xlabel('$\~{k}$')
    plt.ylabel('$\~{h}$')
    plt.grid(True)
    plt.legend()

    return plt.show(), print(steady_states(sK, sH, n, g, delta, alpha, phi))
```

Example on a simulation with reasonable values:

```python
simulate_human_capital(0,0,1/3,1/3,0.01,0.2,0.2,0.01,0.02,200)
```

Try changing the parametres yourself to see how they influence the steady state.

```python
widgets.interact(simulate_human_capital,
                 htilde = widgets.fixed(0),
                 ktilde = widgets.fixed(0),
                 alpha  = widgets.FloatSlider(description='$\u03B1$', min=0,    max=0.5,  step=0.005,  value=0.3),
                 phi    = widgets.FloatSlider(description='$\u03C6$', min=0,    max=0.5,  step=0.005,  value=0.3),
                 delta  = widgets.FloatSlider(description='$\u03B4$', min=0.01, max=0.1,  step=0.001,  value=0.01),
                 sK     = widgets.FloatSlider(description='$s_K$',    min=0.1,  max=0.4,  step=0.01,   value=0.2),
                 sH     = widgets.FloatSlider(description='$s_H$',    min=0.1,  max=0.4,  step=0.01,   value=0.2),
                 n      = widgets.FloatSlider(description='$n$',      min=0.01, max=0.1,  step=0.0005, value=0.01),
                 g      = widgets.FloatSlider(description='$g$',      min=0.01, max=0.1,  step=0.0005, value=0.02),
                 T      = widgets.IntSlider(description='$T$',        min=1,    max=1000, step=10,     value=100))
```

    interactive(children=(FloatSlider(value=0.3, description='$α$', max=0.5, step=0.005), FloatSlider(value=0.3, d…

    <function __main__.simulate_human_capital(htilde, ktilde, alpha, phi, delta, sK, sH, n, g, T)>

# Conclusion

We find that adding human capital as well as technology level has an impact on how we find the real rental rate as well as the real wages. It doesn't change much in how we find the steady state, it just adds some more parameters. The transition-diagram is changing a lot though, we go from having a diagram with $k_t$ on the x-axis and $k_{t+1}$ on the y-axis, to having $\tilde{k}_t$ on the x-axis and $\tilde{h}_t$ on the y-axis. Furthermore the way you graphically find the steady state, changes a lot. In the basic model you just need a 45-degree-line and the capital movement to find the steady state, while you in the human capital model have to find where the two movements cut each other. In our opinion it would be more correct to use the model with human capital to analyze growth in the real world, as it has more parameters that impact how the economy is moving forward.


# Inference with GPs

The dataset needed for this worksheet [can be downloaded](https://northwestern.box.com/s/el0s1imhdxq5qwvzb4hgap90mxpjcfdq). Once you have downloaded [s9_gp_dat.tar.gz](https://northwestern.box.com/s/el0s1imhdxq5qwvzb4hgap90mxpjcfdq), and moved it to this folder, execute the following cell:

```python
!tar -zxvf s9_gp_dat.tar.gz
!mv *.txt data/
```

    x ./._sample_data.txt
    x ./sample_data.txt
    x ./._sample_data_line.txt
    x ./sample_data_line.txt
    x ./._sample_data_line_truths.txt
    x ./sample_data_line_truths.txt

Here are the functions we wrote in the previous tutorial to compute and draw from a GP:

```python
import numpy as np
from scipy.linalg import cho_factor

def ExpSquaredKernel(t1, t2=None, A=1.0, l=1.0):
    """
    Return the ``N x M`` exponential squared
    covariance matrix between time vectors `t1`
    and `t2`. The kernel has amplitude `A` and
    lengthscale `l`.

    """
    if t2 is None:
        t2 = t1
    T2, T1 = np.meshgrid(t2, t1)
    return A ** 2 * np.exp(-0.5 * (T1 - T2) ** 2 / l ** 2)

def draw_from_gaussian(mu, S, ndraws=1, eps=1e-12):
    """
    Generate samples from a multivariate gaussian
    specified by covariance ``S`` and mean ``mu``.

    (We derived these equations in Day 1, Notebook 01, Exercise 7.)
    """
    npts = S.shape[0]
    L, _ = cho_factor(S + eps * np.eye(npts), lower=True)
    L = np.tril(L)
    u = np.random.randn(npts, ndraws)
    x = np.dot(L, u) + mu[:, None]
    return x.T

def compute_gp(t_train, y_train, t_test, sigma=0, A=1.0, l=1.0):
    """
    Compute the mean vector and covariance matrix of a GP
    at times `t_test` given training points `y_train(t_train)`.
    The training points have uncertainty `sigma` and the
    kernel is assumed to be an Exponential Squared Kernel
    with amplitude `A` and lengthscale `l`.

    """
    # Compute the required matrices
    kernel = ExpSquaredKernel
    Stt = kernel(t_train, A=1.0, l=1.0)
    Stt += sigma ** 2 * np.eye(Stt.shape[0])
    Spp = kernel(t_test, A=1.0, l=1.0)
    Spt = kernel(t_test, t_train, A=1.0, l=1.0)

    # Compute the mean and covariance of the GP
    mu = np.dot(Spt, np.linalg.solve(Stt, y_train))
    S = Spp - np.dot(Spt, np.linalg.solve(Stt, Spt.T))

    return mu, S
```

## The Marginal Likelihood

In the previous notebook, we learned how to construct and sample from a simple GP. This is useful for making predictions, i.e., interpolating or extrapolating based on the data you measured. But the true power of GPs comes from their application to *regression* and *inference*: given a dataset $D$ and a model $M(\theta)$, what are the values of the model parameters $\theta$ that are consistent with $D$? The parameters $\theta$ can be the hyperparameters of the GP (the amplitude and time scale), the parameters of some parametric model, or all of the above.

A very common use of GPs is to model things you don't have an explicit physical model for, so quite often they are used to model "nuisances" in the dataset. But just because you don't care about these nuisances doesn't mean they don't affect your inference: in fact, unmodelled correlated noise can often lead to strong biases in the parameter values you infer. In this notebook, we'll learn how to compute likelihoods of Gaussian Processes so that we can *marginalize* over the nuisance parameters (given suitable priors) and obtain unbiased estimates for the physical parameters we care about.

Given a set of measurements $y$ distributed according to
$$
\begin{align}
    y \sim \mathcal{N}(\mathbf{\mu}(\theta), \mathbf{\Sigma}(\alpha))
\end{align}
$$
where $\theta$ are the parameters of the mean model $\mu$ and $\alpha$ are the hyperparameters of the covariance model $\mathbf{\Sigma}$, the *marginal likelihood* of $y$ is
$$
\begin{align}
    \ln P(y | \theta, \alpha) = -\frac{1}{2}(y-\mu)^\top \mathbf{\Sigma}^{-1} (y-\mu) - \frac{1}{2}\ln |\mathbf{\Sigma}| - \frac{N}{2} \ln 2\pi
\end{align}
$$

where $||$ denotes the determinant and $N$ is the number of measurements. The term *marginal* refers to the fact that this expression implicitly integrates over all possible values of the Gaussian Process; this is not the likelihood of the data given one particular draw from the GP, but given the ensemble of all possible draws from $\mathbf{\Sigma}$.

Exercise 1

Define a function ``ln_gp_likelihood(t, y, sigma, A=1, l=1)`` that returns the log-likelihood defined above for a vector of measurements ``y`` at a set of times ``t`` with uncertainty ``sigma``. As before, ``A`` and ``l`` should get passed direcetly to the kernel function. Note that you're going to want to use [np.linalg.slogdet](https://docs.scipy.org/doc/numpy-1.14.0/reference/generated/numpy.linalg.slogdet.html) to compute the log-determinant of the covariance instead of ``np.log(np.linalg.det)``. (Why?)

```python
def ln_gp_likelihood(t, y, sigma=0, A=1.0, l=1.0):
    """
    Docstring
    """
    C = ExpSquaredKernel(t, A=A, l=l) + np.eye(len(y)) * sigma**2
    term1 = -0.5 * np.dot(y.T, np.linalg.solve(C, y))
    term2 = -0.5 * np.linalg.slogdet(C)[1]
    term3 = -1.0 * len(y) / 2 * np.log(2 * np.pi)

    loglike = term1 + term2 + term3
    return loglike

```

Exercise 2

The following dataset was generated from a zero-mean Gaussian Process with a Squared Exponential Kernel of unity amplitude and unknown timescale. Compute the marginal log likelihood of the data over a range of reasonable values of $l$ and find the maximum. Plot the **likelihood** (not log likelihood) versus $l$; it should be pretty Gaussian. How well are you able to constrain the timescale of the GP?

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

```python
t, y, sigma = np.loadtxt("data/sample_data.txt", unpack=True)
plt.plot(t, y, "k.", alpha=0.5, ms=3)
plt.xlabel("time")
plt.ylabel("data");
```

```python
l = np.linspace(0.05, 15, 1000)
log_likelihoods = np.array([ln_gp_likelihood(t, y, sigma=sigma, A=1.0, l=l_val) for l_val in l])
likelihoods = np.exp(log_likelihoods - np.max(log_likelihoods))

```

```python
plt.plot(l, likelihoods)
plt.xlim(0,1)
plt.show()
print(l[np.argmax(likelihoods)])
```

Exercise 3a

The timeseries below was generated by a linear function of time, $y(t)= mt + b$. In addition to observational uncertainty $\sigma$ (white noise), there is a fair bit of correlated (red) noise, which we will assume is well described
by the squared exponential covariance with a certain (unknown) amplitude $A$ and timescale $l$.

Your task is to estimate the values of $m$ and $b$, the slope and intercept of the line, respectively. In this part of the exercise, **assume there is no correlated noise.** Your model for the $n^\mathrm{th}$ datapoint is thus

$$
\begin{align}
    y_n \sim \mathcal{N}(m t_n + b, \sigma_n\mathbf{I})
\end{align}
$$

and the probability of the data given the model can be computed by calling your GP likelihood function:

```python
def lnprob(params):
    m, b = params
    model = m * t + b
    return ln_gp_likelihood(t, y - model, sigma, A=0, l=1)
```

Note, importantly, that we are passing the **residual vector**, $y - (mt + b)$, to the GP, since above we coded up a zero-mean Gaussian process. We are therefore using the GP to model the **residuals** of the data after applying our physical model (the equation of the line).

To estimate the values of $m$ and $b$ we could generate a fine grid in those two parameters and compute the likelihood at every point. But since we'll soon be fitting for four parameters (in the next part), we might as well upgrade our inference scheme and use the ``emcee`` package to do Markov Chain Monte Carlo (MCMC). If you haven't used ``emcee`` before, check out the first few tutorials on the [documentation page](https://emcee.readthedocs.io/en/latest/). The basic setup for the problem is this:

```python
import emcee
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob)

initial = [4.0, 15.0]
p0 = initial + 1e-3 * np.random.randn(nwalkers, ndim)

print("Running burn-in...")
p0, _, _ = sampler.run_mcmc(p0, nburn)   # nburn = 500 should do
sampler.reset()

print("Running production...")
sampler.run_mcmc(p0, nsteps);            # nsteps = 1000 should do
```

where ``nwalkers`` is the number of walkers (something like 20 or 30 is fine), ``ndim`` is the number of dimensions (2 in this case), and ``lnprob`` is the log-probability function for the data given the model. Finally, ``p0`` is a list of starting positions for each of the walkers. Above we picked some fiducial/eyeballed value for $m$ and $b$, then added a small random number to each to generate different initial positions for each walker. This will initialize all walkers in a ball centered on some point, and as the chain progresses they'll diffuse out and begin to explore the posterior.

Once you have sampled the posterior, plot several draws from it on top of the data. You can access a random draw from the posterior by doing

```python
m, b = sampler.flatchain[np.random.randint(len(sampler.flatchain))]
```

Also plot the **true** line that generated the dataset (given by the variables ``m_true`` and ``b_true`` below). Do they agree, or is there bias in your inferred values? Use the ``corner`` package to plot the joint posterior. How many standard deviations away from the truth are your inferred values?

```python
t, y, sigma = np.loadtxt("data/sample_data_line.txt", unpack=True)
m_true, b_true, A_true, l_true = np.loadtxt("data/sample_data_line_truths.txt", unpack=True)
plt.errorbar(t, y, yerr=sigma, fmt="k.", label="observed")
plt.plot(t, m_true * t + b_true, color="C0", label="truth")
plt.legend(fontsize=12)
plt.xlabel("time")
plt.ylabel("data")
plt.show()
```

```python
def lnprob(params):
    m, b = params
    model = m * t + b
    return ln_gp_likelihood(t, y - model, sigma, A=0, l=1)
```

```python
#!pip install emcee
import emcee
```

```python
nwalkers = 24
ndim = 2
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob)
```

```python
nburn = 500
nsteps = 1000

initial = [4.0, 15.0]
p0 = initial + 1e-3 * np.random.randn(nwalkers, ndim)

print("Running burn-in...")
p0, _, _ = sampler.run_mcmc(p0, nburn)   # nburn = 500 should do
sampler.reset()

print("Running production...")
sampler.run_mcmc(p0, nsteps);            # nsteps = 1000 should do
```

    Running burn-in...
    Running production...

```python
draws = []

for i in range(10):
    m, b = sampler.flatchain[np.random.randint(len(sampler.flatchain))]
    draws.append(m * t + b)
```

```python
plt.errorbar(t, y, yerr=sigma, fmt="k.", label="observed")
plt.plot(t, m_true * t + b_true, color="C0", label="truth", lw=3, ls='--')

for i in range(9):
    plt.plot(t, draws[i], label=None)
plt.plot(t, draws[9], label='draws')

plt.legend(fontsize=12)
plt.xlabel("time")
plt.ylabel("data")
plt.show()
```

```python
from corner import corner
```

```python
fig = corner(sampler.flatchain, truths=(m_true, b_true), range=((0,6), (10, 20)));
```

Exercise 3b

This time, let's actually model the correlated noise. Re-define your ``lnprob`` function to accept four parameters (slope, intercept, amplitude, and timescale). If you didn't before, it's a good idea to enforce some priors to keep the parameters within reasonable (and physical) ranges. If any parameter falls outside this range, have ``lnprob`` return negative infinity (i.e., zero probability).

You'll probably want to run your chains for a bit longer this time, too. As before, plot some posterior samples for the line, as well as the corner plot. How did you do this time? Is there any bias in your inferred values? How does the variance compare to the previous estimate?

```python
def lnprob(params):
    m, b, a, l = params

    if m > 6 or m < 2 or b < 15 or b > 20 or a <=0.0 or l <= 0.0:
        return -np.inf

    model = m * t + b
    return ln_gp_likelihood(t, y - model, sigma, A=a, l=l)
```

```python
nwalkers = 50
ndim = 4
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob)
```

```python
nburn = 1000
nsteps = 500

initial = [4.0, 15.0, 1.0, 1.0]
p0 = initial + 1e-3 * np.random.randn(nwalkers, ndim)

print("Running burn-in...")
p0, _, _ = sampler.run_mcmc(p0, nburn)   # nburn = 500 should do
sampler.reset()

print("Running production...")
sampler.run_mcmc(p0, nsteps);            # nsteps = 1000 should do
```

    Running burn-in...

    /Users/rmorgan/anaconda3/lib/python3.7/site-packages/emcee/ensemble.py:335: RuntimeWarning: invalid value encountered in subtract
      lnpdiff = (self.dim - 1.) * np.log(zz) + newlnprob - lnprob0
    /Users/rmorgan/anaconda3/lib/python3.7/site-packages/emcee/ensemble.py:336: RuntimeWarning: invalid value encountered in greater
      accept = (lnpdiff > np.log(self._random.rand(len(lnpdiff))))

    Running production...

```python
draws = []

for i in range(10):
    m, b, a, l = sampler.flatchain[np.random.randint(len(sampler.flatchain))]
    draws.append(m * t + b)
```

```python
plt.errorbar(t, y, yerr=sigma, fmt="k.", label="observed")
plt.plot(t, m_true * t + b_true, color="C0", label="truth", lw=3, ls='--')

for i in range(9):
    plt.plot(t, draws[i], label=None)
plt.plot(t, draws[9], label='draws')

plt.legend(fontsize=12)
plt.xlabel("time")
plt.ylabel("data")
plt.show()
```

```python
fig = corner(sampler.flatchain, truths=(m_true, b_true, 1.0, 1.0));
```

```python

```

Exercise 3c

If you didn't do this already, re-plot the posterior samples on top of the data, but this time draw them from the GP, *conditioned on the data*. How good is the fit?