#### _Speech Processing Labs 2021: Signals 2: IIR Filters (Extensions derivations)_

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

import cmath
from math import floor
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
plt.style.use('ggplot')

from dspMisc import *
```

# IIR Filter Derivations

    This notebook is very much optional extension and is really just here for completeness!

### Learning Outcomes
* Observe how IRR filter coefficients can be derived from poles and zeros

### Need to know
* Topic Videos: Impulse Response, Source Filter Model
* [IIR filters](./signals-2-3-iir-filters.ipynb)

Equation alert: If you're viewing this on github, please note that the equation rendering is not always perfect. You should view the notebooks through a jupyter notebook server for an accurate view.

## Designing a filter using poles and zeros

The following gives a very, very terse overview of how you can design an IIR filter frequency response using the concept of poles and filters. This requires a fair grip on solving polynomials and complex numbers, but once you've got that it's actually quite nice!

If you're interested in learning more about this (or signal processing in general), it's probably time to pick up a Digital Signal Processing textbook. Rick Lyon's [Understanding Digital Signal Processing](https://www.dsprelated.com/books/6.php) is probably the most accessible I've come across. The example filters above were taken from that book. The rest of this notebook basically gives a excessively detailed derivation going from poles and filters to coefficients for those examples for completeness.

Previous SLP students have also recommended [The Scientist and Engineer's Guide to
Digital Signal Processing](http://www.dspguide.com/pdfbook.htm), though the free online version renders a bit weirdly!.

IIR filters have the following form:

$$ y[n] = \sum_{k=0}^N b(k) x[n-k] + \sum_{k=0}^M a(k) y[n-k]$$

Unlike an FIR filter, we can't just put an impulse in and find out what the frequency response of the filter is.  Instead, to understand the properties of an IIR filter, we use the **z-transform**.  The z-transform is the discrete version of the continuous Laplace transform.  We won't get into the details here, but if you do read more about this you'll see that you can think of the Laplace transform as a generalization of the Fourier Transform. So the intuitions you've developed about how the Fourier Transform works (e.g. finding correlations between inputs and basis functions) will help understand the Laplace Transform too.

We use the z-transform to convert filter difference equations to an equivalent form in the z-domain (z being a complex number).  We rewrite our difference equation in terms of the z-transformed input $x[n]$ and output $y[n]$ functions, which we denote as $X(z)$ and $Y(z)$ respectively.  Some algebra then allows us to define a **transfer function** $H(z)$:

$$ Y(z) = H(z)X(z)$$

So, $H(z)$ is what our filter does to our input to get our output in the z-domain. You can think of it as the filter part of the source filter model (the source being our impulse train and/or white noise).

Unsurprisingly, we can actually define $H(z)$ in terms of our the feedforward coefficients $b(k)$ and the feedback coefficients $a(k)$. What's surprising is that we can define it in a relatively simple way:

$$ H(z) = \frac{Y(z)}{X(z)} = \frac{\sum_{k=0}^N b(k)z^{-k}}{1- \sum_{k=1}^M a(k)z^{-k}}\\ $$

You can think of  $z^{-k}$ as referring to a position in the sequence $k$ steps back in time.

We can analyze $H(z)$ to figure out how the filter is going to behave. To do this we need to factorize the numerator and denominator polynomials of $H(z)$.

$$
\begin{align}
H(z) &= \frac{\sum_{k=0}^N b(k)z^{-k}}{1- \sum_{k=1}^M a(k)z^{-k}}\\
&= G\frac{(z-z_0)(z-z_1)...(z-z_N)}{(z-p_0)(z-p_1)...(z-p_M)}
\end{align}
$$

In the equation above, $G$ is an arbitrary gain constant which doesn't affect the frequency response or filter stability. The polynomial roots, $z_n$ and $p_m$ are called *zeros* and *poles* respectively.  These are the (complex) values of $z$ which cause the numerator (zeros) or the denominator (poles) to equal 0.  We can interpret them as follows:

* **Zeros:** when the numerator goes to zero, the transfer function $H(z)$ goes to zero, and we see decrease (attenuation) in the frequency magnitude response.
* **Poles:** when the denominator goes to zero, the transfer function $H(z)$ goes to infinity, leads to an increase (gain) the frequency magnitude response (like vowel formants!)

It may not look like it right now, but it's way easier to analyze IIR filters in the z-domain than in the time domain! Once we know the poles and the zeros of the transfer function, we basically know what the frequency response will look like and whether the filter is stable or not.  The main things to note are:

* If a pole falls inside the unit-circle (of the complex z-plane), the filter output will be **stable**.
    * Given an impulse input, $y[n]$ it will eventually go to zero as $n$ increases.
* If a pole falls outside the unit circle, the output will be **unstable**
    * Given an impulse input, $y[n]$ will keep increasing in magnitude as more time passes (definitely not the behaviour of the vocal tract!).
* If a pole lies exactly on the unit circle, the output will oscillate between finite values determined by the pole locations forever.

An IIR filter is only as stable as it's weakest link - if any poles sit outside the unit circle, the filter as a whole will be unstable.  Zeros don't affect filter stability.

Another property of poles (and zeros) is that, if the denominator polynomial has real-valued coefficients (i.e. the $a(k)$ values), poles (and zeros) either come in pairs of complex conjugates (i.e. if $p=a+ib$ is a pole, so is $p^*=a-ib$), or it's a value on the real axis of the z-plane.  We can use these facts to design filters using the factorized form of the transfer function, which will then allow us to derived the filter coefficients $a(k)$ and $b(k)$. We'll see some examples this below.

The frequency response of the filter is $H(z)$ evaluated at $e^{j\omega}$, where $-\pi \lt \omega \lt \pi$ represents frequencies in terms of radians per sample ($\pi$ radians/sample $\mapsto$ a frequency matching half the sampling frequency (i.e. $f_s/2$).

$$ H(\omega) = H(e^{j\omega}) = \frac{\sum_{k=0}^N b(k)e^{-jk\omega}}{1 -\sum_{k=0}^M a(k)e^{-jk\omega}} $$

These example IIR filters from Richard J Lyon's book
[Understanding Digital Signal Processing](https://www.dsprelated.com/books/6.php), which is a very nice introduction to the topic.

Working through from the factorized form to the coefficients is not hard, but it is very annoying and it's very easy to make mistakes, so I've given all the steps expanding out the factorized forms of the filter examples he gives in his book below.

### Ex 1:

$$
\begin{aligned}
H_1(z) &= \frac{G}{z-0.8}\\
       &= G\frac{1}{1-(0.8z^{-1})}
\end{aligned}
$$

* One pole on the real axis, inside the unit circle, stable peak at $\omega=0$ rad/sample
* no zeros

```python
a = [0.8]
b = [1]
```

### Ex 2:

$$
\begin{aligned}
H_1(z) &= G\frac{(z-1)}{z-(-0.8)}\\
       &= G\frac{1 - z^{-1}}{1 -  (-0.8z^{-1})}
\end{aligned}
$$

* pole on real axis at -0.8, inside the unit circle, stable peak at $\omega=-\pi$ rad/sample, i.e. $f_s/2$ Hz which is 8000/2 = 4000 Hz in our setup above.
* one zero at z=1, this causes attenuation at $\omega = 0$ rads/sample (0 Hz)

```python
a = [-0.8]
b = [1,-1]
```

### Ex 3
$$
\begin{align}
H_3(z) &= G\frac{[z-(-0.707 + j0.707)][z-(-0.707 - j0.707)]}{[z-(0.283+j0.283)][z-(0.283-j0.283)]} \\
&= G\frac{z^2 -z(-0.707 - j0.707) - z(-0.707 + j0.707) + (-0.707 + j0.707)(-0.707 - j0.707) }{z^2 - z(0.283-j0.283) -z(0.283+j0.283) + (0.283+j0.283)0.283-j0.283)} \\
&= G\frac{z^2 +z0.707 + j0.707z + z0.707 - j0.707z + 1 }{z^2 - 0.283z + j0.283z - 0.283z - j0.283z + 0.16017} \\
&= G\frac{z^2 +z[0.707  + 0.707] + 1 }{z^2 - 0.283z - 0.283z  + 0.16017} \\
&= G\frac{z^2 + 1.414z + 1 }{z^2 - 0.566z + 0.16017} \\
&= G\frac{1 + 1.414z^{-1} + 1z^{-2} }{1 - 0.566z^{-1} + 0.16017z^{-2}} \\
\end{align}
$$

* The poles are inside the unit circle, this produces a low pass filter  (pass band is 'bordered' by +/- $\omega = \pi/4$ rad/sample, 1000 Hz)
    * $(0.283+j0.283) \mapsto 0.400e^{j\pi/4}$
* zeros on the unit circle, attenuation at $3\pi/4$ rad/sample = 3000 Hz.
    * $(-0.707+j0.707) \mapsto 1e^{j3\pi/4} $

```python
a = [0.566, -0.16017]
b = [1, 1.414, 1]

```

### Ex 4:
$$
\begin{align}
H_4(z) &= G\frac{[z-1][z-(-0.707 + j0.707)][z-(-0.707 - j0.707)]}{[z-(0.283+j0.283)][z-(0.283-j0.283)]} \\
&= G\frac{[z-1] [z^2 + 1.414z + 1] }{z^2 - 0.566z + 0.16017}\\
&= G\frac{z^3 + 1.414z^2 + 1z -z^2 - 1.414z - 1}{z^2 - 0.566z + 0.16017}\\
&= G\frac{z^3 + 0.414z^2 - 0.414z - 1}{z^2 - 0.566z + 0.16017}\\
&= G\frac{1 + 0.414z^{-1} - 0.414z^{-2} - z^{-3}}{1 - 0.566z^{-1} + 0.16017z^{-2}}\\
\end{align}
$$

* We've added an extra zero at $z=0$ compared to example 3.  This means that we get a valley (attenuation) in the frequency response at $\omega = 0$ (0 Hz).

```python
a = [0.566, -0.16017]
b = [1, 0.414, -0.414, -1]
```

### Ex 5:
$$
\begin{align}
H_5(z) &= G\frac{[z-1][z-(-0.707 + j0.707)][z-(-0.707 - j0.707)]}{[z-(0.636+j0.636)][z-(0.636-j636)]} \\
&= G\frac{[z-1] [z^2 + 1.414z + 1] }{z^2 - 1.272z + 0.808992}\\
&= G\frac{1 + 0.414z^{-1} - 0.414z^{-2} - z^{-3}}{1 - 1.272z^{-1} +  0.808992z^{-2}}\\
\end{align}
$$

* poles are nearer to the unit circle (magnitude 0.80899), so the peaks associated with them (+/- $\omega = \pi/4$ rad/sec (1000 Hz with our sample rate) are sharper.

```python
a = [1.272, -0.80899]
b = [1, 0.414, -0.414, -1]
```

### Ex 6:
$$
\begin{align}
H_6(z) &= G \frac{[z-1][z-(-0.707 + j0.707)][z-(-0.707 - j0.707)]}{[z-(0.777+j0.777)][z-(0.777-777)]} \\
&= G \frac{1 + 0.414z^{-1} - 0.414z^{-2} - z^{-3}}{1 - 1.554z^{-1} +  1.207z^{-2}}\\
\end{align}
$$

* poles are now outside the unit circle (magnitude 1.207), so the filter is now unstable: as n increases, the value of $y[n]$ eventually goes to infinity.
* It still has the same shape frequency response though!

```python
a = [1.554, -1.207]
b = [1, 0.414, -0.414, -1]
```

```python

```

```python

```


# The Variational Quantum Eigensolver Algorithm

```python
import numpy as np

from VariationalForms import VariationalForms

from qiskit import *

from qiskit.visualization import *

from qiskit.providers.aer import QasmSimulator, StatevectorSimulator, UnitarySimulator

```

```python
qasm_sim = QasmSimulator()
state_sim = StatevectorSimulator()
unit_sim = UnitarySimulator()

vf = VariationalForms()
```

# Variational Forms

### Rotation Gates

\begin{align}
    \mathbf{R_z(\theta) = e^{-i\frac{\theta}{2}\sigma_z}}
\end{align}

\begin{align}
    \mathbf{R_y(\phi) = e^{-i\frac{\phi}{2}\sigma_y}}
\end{align}

## $R(\theta, \phi)$

$R(\theta, \phi)$ Rotation gate, where:
$R(\theta, \phi) = $

```python
vf.draw(vf.R(0, 0))
```

```python
vf.decompose(vf.R(0, 0))
```

## $R(\theta, \phi)^\dagger$

$R(\theta, \phi)^\dagger$ Rotation gate

```python
vf.draw(vf.Rd(0, 0))
```

```python
vf.decompose(vf.Rd(0, 0))
```

## $A(\theta, \phi)$

$A(\theta, \phi)$ 2 Qubit A gate. Preserves **Particle-Number Symmetry** and **Time-Reversal Symmetry**

```python
vf.draw(vf.A(0, 0))
```

```python
vf.decompose(vf.A(0, 0))
```

```python

```

    Figure(748.797x204.68)


```python
# plotting libs
import matplotlib as mpl
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
%matplotlib inline
%config InlineBackend.figure_format
mpl.rcParams['figure.figsize'] = (9, 6)

# standard
from collections import defaultdict
from datetime import datetime
from functools import partial
from multiprocessing import Pool, cpu_count
import os
import pickle

# third party
import covidcast
import numpy as np
import pandas as pd
from statsmodels.tsa.tsatools import lagmat

# first party
os.chdir("../code")
from delay import get_international_delays, get_delay_distribution
from conv1d import *
from weekday import Weekday, dow_adjust_cases
os.chdir("../notebooks")
```

## Delay distributions

Delay distribution found by fitting a gamma distribution to linelist data.

```python
## Data path globals
data_dir = "../data/"
fl_line_data = data_dir + "FL_line_list.csv"
int_line_data = data_dir + "international_line_list.tar"
us_zip_data_path = "~/Documents/covid-19/geographical_scope/02_20_uszips.csv"
```

```python
international_delays = get_international_delays(data_path=int_line_data)
delay_dist = get_delay_distribution(international_delays)
intl_delay_dist = get_delay_distribution(international_delays)

# https://www.arcgis.com/home/item.html?id=37abda537d17458bae6677b8ab75fcb9
florida_df = pd.read_csv(fl_line_data, parse_dates=["Case_", "EventDate", "ChartDate"])
florida_delays = (florida_df.ChartDate - florida_df.EventDate).dt.days
florida_delays = florida_delays[florida_delays.gt(0) & florida_delays.lt(60)]
fl_delay_dist = get_delay_distribution(florida_delays)

plt.plot(intl_delay_dist, label="international")
plt.plot(fl_delay_dist, label="florida")
plt.title("Symptom-onset to report delay")
plt.legend()
plt.show()
```

Pull case data.

```python
start_date = datetime(2020, 3, 15)
end_date = datetime(2020, 9, 15)

cases_df = covidcast.signal(
        'indicator-combination',
        'confirmed_7dav_incidence_num',
        start_date, end_date,
        geo_type='county')

cumulative_cases_df = covidcast.signal(
        'indicator-combination',
        'confirmed_7dav_cumulative_num',
        end_date, end_date,
        geo_type='county')

thresh_geos = cumulative_cases_df[cumulative_cases_df.value > 500].geo_value
```

```python
# get all florida fips codes
geo_map = pd.read_csv(
            us_zip_data_path,
            usecols=["fips", "state_id", "population"],
            dtype={"state_id": str},
            converters={"fips": lambda x: str(x).zfill(5)},
        )

florida_geo = geo_map[geo_map.state_id.eq("FL")]
florida_population = florida_geo.groupby("fips").population.sum().reset_index()
florida_fips = florida_geo.fips.unique()
```

```python
cases_df = cases_df.set_index(["geo_value"])
geos = cases_df.index.unique()
geos = geos[geos.isin(florida_fips)] # only keep florida geos
geos = geos[geos.isin(thresh_geos)]  # counties with >500 cumulative cases
```

### Models

```python
def regress(X, y):
    """Simple OLS, adds intercept. Returns coefs and fitted vals."""
    X = np.hstack((np.ones((y.size, 1)), X.reshape(-1, 1)))
    beta = np.linalg.inv(X.T @ X + 1e-4 * np.eye(X.shape[1])) @ X.T @ y
    return beta, (X @ beta)

def strawman(x_tilde, delay):
    """
    The prediction x_hat at time t is x_tilde (infections)
    at time t-1:

      x_hat{t} = x_tilde{t-1}

    Returns a tuple of the corresponding time indices to the
    convolution y_hat = delay(*)x_hat.

    """
    n = x_tilde.shape[0]
    x_hat = lagmat(x_tilde, 1).reshape(-1,)
    x_pred = x_tilde[-1] # prediction for next time point
    x_hat = np.append(x_hat, x_pred)

    y_hat = Conv1D.freq_conv(x_hat, delay)[:n]
    return (np.arange(1, n+1), y_hat)

def ar(x_tilde, delay, p, lam=1):
    """
    Fit an AR(p) model on x_tilde.

    Returns a tuple of the corresponding time indices to the
    convolution y_hat = delay(*)x_hat.

    """
    n = x_tilde.shape[0]

    # lag matrix
    Z = np.hstack((np.ones((n, 1)), lagmat(x_tilde, maxlag=p)))[p:]
    Z1 = np.concatenate([[1], np.flip(x_tilde[-p:])])
    beta = np.linalg.solve(Z.T @ Z + lam*np.eye(p+1), Z.T @ x_tilde[p:])
    x_hat = Z @ beta
    x_pred = beta.T @ Z1 # prediction for next time point
    x_hat = np.append(x_hat, x_pred)

    y_hat = Conv1D.freq_conv(x_hat, delay)[:(Z.shape[0]+1)]
    return (np.arange(p, n+1), y_hat)

def sf(x_tilde, train_signals, test_signals, delay, alpha=1):
    """
    Sensor fusion.

    - [TODO] Fit on more than one location. Currently, H=1.
    - [TODO] This does not properly fit sensors sequentially.

    Args:
        x_tilde: infections vector
        train_signals: raw signal matrix to become sensors
        test_signals: raw signal vector at nowcasting time
        delay: reporting delay vector
        alpha: l2 regularization parameter

    """
    n, p = train_signals.shape
    p += 1 # add intercept
    H = np.ones((p, 1))

    # fit sensors
    Z = np.full((n, p), np.nan)
    z1 = np.full((p,), np.nan)
    for j in range(p):
        if j < (p-1):
            beta, fitted_vals = regress(train_signals[:, j], x_tilde)
            Z[:,j] = fitted_vals
            z1[j] = beta.T @ np.array([1, test_signals[j]])
        else:
            Z[:, j] = np.ones(n)
            z1[j] = 1

    X = x_tilde.reshape(-1, 1)
    G = Z - X @ H.T
    #G = G - G.mean(axis=0)
    cov_G = G.T @ G / n
    Ri = np.linalg.inv(alpha * cov_G + (1 - alpha) * np.eye(p))
    beta = np.linalg.inv(H.T @ Ri @ H) @ H.T @ Ri
    x_hat = Z @ beta.T
    x_pred = beta @ z1 # prediction for next time point
    x_hat = np.append(x_hat, x_pred)

    y_hat = Conv1D.freq_conv(x_hat, delay)[:(n+1)]
    return (None, y_hat)
```

## Nowcasting Miami-Dade (12086)

```python
geo = "12086"
cases = cases_df.loc[geo].sort_values(by='time_value')
n = cases.value.shape[0]

fb_df = covidcast.signal(
        'fb-survey',
        'smoothed_hh_cmnty_cli',
        start_date, end_date,
        geo_type='county',
        geo_values=geo)

fb = fb_df[['time_value', 'value']]
plt.plot(fb.time_value, fb.value)
plt.xlabel("Time")
plt.ylabel("Facebook CLI")
plt.title(f"{geo} Facebook CLI")
plt.show()
```

### Deconvolution Approaches

(1) Direct deconvolution in the frequency domain.

(2) Solve

\begin{equation}
\hat{x}=\underset{x}{\operatorname{argmin}} \frac{1}{2}\|y-W x\|_{2}^{2}+\lambda\left\|D^{(k+1)} x\right\|_{1}
\end{equation}

where $D^{(k+1)}$ is the discrete difference operator, and $W$ the convolution matrix.

(3) Solve

\begin{equation}
\hat{x}=\underset{x}{\operatorname{argmin}} \frac{1}{2}\|W^{-1}y-x\|_{2}^{2}+\lambda\left\|D^{(k+1)} x\right\|_{1}.
\end{equation}

where $W^{-1}$ is the deconvolution matrix.

```python
dates = cases.time_value

direct = Conv1D.freq_deconv(cases.value, delay_dist)[:n]
admm_v1 = admm_deconvolution(cases.value, delay_dist,
                             lam=20, rho=20, n_iters=500, k=2)
admm_v2 = admm_deconvolution_v2(cases.value, delay_dist,
                                lam=200, rho=200, n_iters=500, k=2)

plt.plot(dates, cases.value, label="raw", color="black", linewidth=4, alpha=0.2)
plt.plot(dates, direct, label="[direct] symptom-onset")
plt.plot(dates, admm_v1, label="[admm] symptom-onset")
plt.plot(dates, admm_v1, label="[admm_v2] symptom-onset", linestyle="--")
plt.title("Miami-Dade: symptomatic infection curve from raw cases")
plt.ylim(-1000, 1.5 * max(cases.value))
plt.legend()
plt.show()

plt.plot(dates, cases.value, label="raw", color="black", linewidth=4, alpha=0.2)
plt.plot(dates, Conv1D.freq_conv(direct, delay_dist)[:n], label="[direct] cases")
plt.plot(dates, Conv1D.freq_conv(admm_v1, delay_dist)[:n], label="[admm] cases")
plt.plot(dates, Conv1D.freq_conv(admm_v1, delay_dist)[:n],
         label="[admm_v2] cases", linestyle="--")
plt.title("Miami-Dade: reconvolved case curve from symptom-onset curve")
plt.legend()
plt.show()

adj_cases = dow_adjust_cases(cases, lam=10)
direct = Conv1D.freq_deconv(adj_cases, delay_dist)[:n]
admm_v1 = admm_deconvolution(adj_cases, delay_dist, lam=20, rho=20, n_iters=500, k=2)
admm_v2 = admm_deconvolution_v2(adj_cases, delay_dist, lam=200, rho=200, n_iters=500, k=2)
plt.plot(dates, adj_cases, label="dow-adj", color="black", linewidth=4, alpha=0.2)
plt.plot(dates, direct, label="[direct] symptom-onset")
plt.plot(dates, admm_v1, label="[admm] symptom-onset")
plt.plot(dates, admm_v1, label="[admm_v2] symptom-onset", linestyle="--")
plt.title("Miami-Dade: symptomatic infection curve from dow-adj cases")
plt.legend()
plt.show()

plt.plot(dates, cases.value, label="raw", color="black", linewidth=4, alpha=0.2)
plt.plot(dates, Conv1D.freq_conv(direct, delay_dist)[:n], label="[direct] cases")
plt.plot(dates, Conv1D.freq_conv(admm_v1, delay_dist)[:n], label="[admm] cases")
plt.plot(dates, Conv1D.freq_conv(admm_v1, delay_dist)[:n],
         label="[admm_v2] cases", linestyle="--")
plt.title("Miami-Dade: reconvolved case curve from dow-adjusted symptom-onset curve")
plt.legend()
plt.show()
```

### Miami-Dade linelist

Comparison with International, Florida, and Miami-Dade linelists.

```python
miami_dade = florida_df[florida_df.County.eq("Dade")]
miami_dade = miami_dade[miami_dade.ChartDate >= datetime(2020, 6, 1)]
miami_dade_delays = (miami_dade.ChartDate - miami_dade.EventDate).dt.days
miami_dade_delays = miami_dade_delays[miami_dade_delays.gt(0) & miami_dade_delays.lt(60)]
miami_dade_delay_dist = get_delay_distribution(miami_dade_delays)

plt.plot(intl_delay_dist, label="international")
plt.plot(fl_delay_dist, label="florida")
plt.plot(miami_dade_delay_dist, label="miami-dade")
plt.legend()
plt.show()
```

## Nowcasting

We will compare (semi)retrospective performance of the strawman, AR(7), and SF method.

All three methods will recieve identical symptom onset curves. The symptom onset curve is estimated using method (3) of deconvolution (via ADMM), with the three different delay distributions.

```python
intl_preds = defaultdict(lambda: {'time': [],'pred': []})
fl_preds = defaultdict(lambda: {'time': [],'pred': []})
miami_dade_preds = defaultdict(lambda: {'time': [],'pred': []})

t1 = n
t0 = t1-100
for t in range(t0, t1):
    train_time = cases.time_value[:t]
    test_time = cases.time_value.values[t]
    train_cases = cases[cases.time_value.isin(train_time)]
    n_train_obs = train_cases.shape[0]
    train_cases = dow_adjust_cases(train_cases, lam=10)

    for delay_dist, preds in zip([intl_delay_dist, fl_delay_dist, miami_dade_delay_dist],
                                 [intl_preds, fl_preds, miami_dade_preds]):
        # get infection curve
        ## commented out - this does CV and is necessary for good performance but slow
        ## after observing the output, we rerun with a fixed lambda for speed
        # sub_infections = infection_curve::get_infection_curve(train_cases, delay_dist)
        sub_infections = np.clip(admm_deconvolution_v2(
            train_cases, delay_dist, 3000, 3000, n_iters=500, k=2), 0, np.inf)

        # strawman
        strawman0 = strawman(sub_infections, delay_dist)
        preds["strawman"]["pred"].append(strawman0[1][-1])
        preds["strawman"]["time"].append(test_time)

        # ar7 model
        n_lags = 7
        ar0 = ar(sub_infections, delay_dist, n_lags)
        preds["ar7"]["pred"].append(ar0[1][-1])
        preds["ar7"]["time"].append(test_time)

        # sf model (fb_survey, ar7, strawman)
        train_fb = fb[fb.time_value.isin(train_time)]
        test_fb = fb[fb.time_value.eq(test_time)]
        train_ar = pd.DataFrame({'time_value': train_time[(n_lags-1):],
                                 'ar': ar0[1]})
        test_ar = ar0[1][-1]

        train_straw = pd.DataFrame({'time_value': train_time,
                                 'strawman': strawman0[1]})
        test_straw = strawman0[1][-1]
        if test_fb.size > 0:
            sf_df = (pd.DataFrame({'infections': sub_infections,
                                   'time_value': train_time})
                     .merge(train_fb, how='left', on='time_value')
                     .merge(train_ar, how='left', on='time_value')
                     .merge(train_straw, how='left', on='time_value'))
            valid_ind = sf_df.notna().all(axis=1)
            sf_df.drop(columns=['time_value', 'infections'], inplace=True)

            if valid_ind.sum() > 100:
                train_infections = sub_infections[valid_ind]
                train_signals = sf_df.to_numpy()[valid_ind, :]
                test_signals = np.array((
                   test_fb.value.values[0],
                    test_ar,
                    test_straw)).flatten()

                sf0 = sf(train_infections, train_signals,
                         test_signals, delay_dist, 0.5)
                preds["sf"]["pred"].append(sf0[1][-1])
                preds["sf"]["time"].append(test_time)

    if t % 20 == 0: print(f"Finished {t}/{t1}")

plt.title("Miami-Dade")
plt.plot(cases.time_value, cases.value, label="cases", color="black", linewidth=4, alpha=0.8)
for k in preds.keys():
    plt.plot(preds[k]['time'], preds[k]["pred"], label=k, marker=".", alpha=0.5)
plt.axvline(x=cases.time_value.values[t0], color="gray", linestyle=":")
plt.xlim(datetime(2020, 6, 1), end_date)
plt.legend()
plt.xticks(rotation=90)
plt.show()
```

```python
for method in ["strawman", "ar7", "sf"]:
    plt.plot(cases.time_value, cases.value, label="cases", color="black", linewidth=3, alpha=0.8)
    for name, preds in zip(["international", "florida", "miami-dade"],
                           [intl_preds, fl_preds, miami_dade_preds]):
        plt.plot(preds[method]['time'],
                 preds[method]["pred"],
                 label=name, marker=".", alpha=0.8)
    plt.axvline(x=cases.time_value.values[t0], color="gray", linestyle=":")
    plt.legend()
    plt.title(f"Miami-Dade: {method}")
    plt.xticks(rotation=90)
    plt.show()

cases_adj = dow_adjust_cases(cases, lam=10)
cases["adj_value"] = cases_adj
for name, preds in zip(["international", "florida", "miami-dade"],
                       [intl_preds, fl_preds, miami_dade_preds]):
    preds_df = pd.DataFrame({"time_value": preds[method]['time'], 'preds': preds[method]["pred"]})
    tmp = preds_df.merge(cases, how="left", on="time_value")

    plt.hist(np.abs(tmp.adj_value - tmp.preds),
             label=name, alpha=0.6, density=True, bins=20)

plt.legend()
plt.title(f"Miami-Dade")
plt.xlabel("nowcast one-ahead absolute error")
plt.ylabel("density")
plt.axhline(0)
plt.show()
```

```python
## parallelized version

# def batch_nowcast(geo, y, delay, ts, t0, t1, model):
#     """
#     Make nowcasts over time range [t0, t1).

#     Args:
#         y: ordered array of cases
#         delay: ordered array of delay probabilities
#         ts: ordered array of complete timepoints corresponding to y
#         t0: index of starting timepoint for prediction
#         t1: index of ending timepoint for prediction
#         model: fitting function

#     """
#     #infection_curve = InfectionCurve(delay).get_infection_curve
#     res = {"pred": [], "time": []}
#     for t in range(t0, t1):
#         train_time = ts[:t]
#         test_time = ts[t]
#         train_y = y[:t]

#         # get infection curve (x_tilde)
#         x_tilde = np.clip(admm_deconvolution(train_y, delay, 50, 50, n_iters=300, k=2), 0, np.inf)
# #         plt.plot(train_y)
# #         plt.plot(x_tilde)
# #         plt.show()
#         #x_tilde = infection_curve(train_y)

#         # fit model
#         _, out_y = model(x_tilde, delay)
#         res["pred"].append(out_y[-1])
#         res["time"].append(test_time)

#     print(f"{geo}, {res['pred'][-1]:.3f}")
#     return {geo: res}

# ar7 = partial(ar, p=7)
# for method, method_str in zip([strawman, ar7], ["strawman", "ar7"]):
#     n_cpu = min(5, cpu_count())
#     pool = Pool(n_cpu)
#     pool_results = []

#     for geo in sample_geos:
#         cases = cases_df.loc[geo].sort_values(by='time_value')
#         dates = cases.time_value.values
#         cases = dow_adjust(cases)
#         n = cases.shape[0]
#         pool_results.append(
#             pool.apply_async(batch_nowcast,
#                              args=(geo,
#                                    cases,
#                                    delay_dist,
#                                    dates,
#                                    n-50, n,
#                                    method,)
#                             )
#         )

#     pool_results = [proc.get() for proc in pool_results]
#     pickle.dump(pool_results, open(f"{method_str}_fl.p", "wb"))
#     print(f"finished {method_str}")

# # load data
# def unpack(fn):
#     fp = pickle.load(open(fn, "rb"))
#     out = {}
#     for arr in fp:
#         out[list(arr.keys())[0]] = arr[list(arr.keys())[0]]
#     return out

# strawman_res = unpack("strawman_fl.p")
# ar7_res = unpack("ar7_fl.p")

# out_pdf = f"comp_fl.pdf"
# pdf_pages = PdfPages(out_pdf)
# n_plot = len(pool_results)
# n_row = 4
# n_col = 4
# n_plots_per_page = n_row*n_col
# fig, axs = None, None
# j = 0
# for i, geo in enumerate(sample_geos):
#     cases = cases_df.loc[geo].sort_values(by='time_value')
#     if i % n_plots_per_page == 0:
#         fig, axs = plt.subplots(n_row, n_col, figsize=(12, 12), sharex=True)
#         axs = axs.ravel()
#         j = 0

#     axs[j].plot(cases.time_value.values, cases.value.values,
#                 color="black", label="cases", linewidth=3, alpha=0.8)
#     axs[j].plot(strawman_res[geo]["time"], strawman_res[geo]["pred"], label="strawman", marker=".")
#     axs[j].plot(ar7_res[geo]["time"], ar7_res[geo]["pred"], label="ar7", marker=".")
#     axs[j].legend(fontsize=8)
#     axs[j].set_title(geo)
#     axs[j].tick_params(axis='both', which='major', labelsize=5, labelrotation=90)

#     # close the page if needed
#     if (i + 1) % n_plots_per_page == 0 or (i + 1) == n_plot:
#         plt.tight_layout()
#         pdf_pages.savefig(fig)
#         plt.close()
#     j += 1
# pdf_pages.close()
# print(f"Saved to {out_pdf}")

# plt.plot(cases.time_value, cases.value, label="cases", color="black", linewidth=4, alpha=0.8)
# for k in preds.keys():
#     plt.plot(preds[k]['time'], preds[k]["pred"], label=k, marker=".", alpha=0.5)
# plt.axvline(x=cases.time_value.values[t0], color="gray", linestyle=":")
# plt.legend()
# plt.xticks(rotation=90)
# plt.show()
```

```python

```


# This notebook shows how to extract a model from a Latex document and simulate the mode.

## Why specify a model in Latex?

Sometime the **implementation** of a model in software don't match the **specification** of the model in
the text in which the model is presented. It can be a challenge to make sure that the specification is
updated in order to reflect changes made in the implementation.

By extracting the model from a Latex script which describes and specify the model a one can always be sure that simulations reflect the model as described in the paper.

Also the author is forced to make a complete specification of the model, else it won't run.

## The Economic Credit Loss model
This jupyter notebook is inspired by IMF working paper (WP/20/111) The Expected Credit Loss Modeling from a Top-Down Stress Testing Perspective by Marco Gross, Dimitrios Laliotis, Mindaugas Leika, Pavel Lukyantsau. The working paper and the associated material is located https://www.imf.org/en/Publications/WP/Issues/2020/07/03/Expected-Credit-Loss-Modeling-from-a-Top-Down-Stress-Testing-Perspective-49545

from the abstract of the paper:
> The objective of this paper is to present an integrated tool suite for IFRS 9- and CECL compatible
estimation in top-down solvency stress tests. The tool suite serves as an
illustration for institutions wishing to include accounting-based approaches for credit risk
modeling in top-down stress tests.

This is a jupyter notebook build with the purpose of illustrating the conversion of a model from Latex to ModelFlow. The purpose is testing so take very much care.

## Import libraries

```python
%load_ext autoreload
%autoreload 2
import pandas as pd
from IPython.core.display import HTML,Markdown

from modelclass import model
import modeljupytermagic

# some useful stuf
model.widescreen()
pd.set_option('display.max_rows', None, 'display.max_columns', 10, 'precision', 2)
sortdf = lambda df: df[sorted([c for c in df.columns])]
```

## Write a latex script
The model consists of the equations and the lists

The jupyter magic command **%%latexflow** will extract the model. then it will  transform the equations to **ModelFlow** equations and finaly it will create a modelflow **model** instance.

The **model** instance will be able to solve the model.

```python
%%latexflow ecl

### Loans can be in 3 stages
Loans can be in 3 stages, s1,s2 and s3.
New loans will be generated and loans will mature.

Two lists of stages are defined:

List $stage=\{s1, s2,s3\}$

List $stage\_from=\{s1, s2,s3\}$

### A share of the loans in each stage wil transition to the same or another stage in the next time frame:

\begin{equation}
\label{eq:loanfromto}
loan\_transition\_from\_to^{stage\_from,stage}_{t} = loan^{stage\_from}_{t-1} \times TR^{stage\_from,stage}_{t}
\end{equation}

\begin{equation}
\label{eq:transition}
loan\_transition\_to^{stage}_{t}  = \underbrace{\sum_{stage\_from}(loan\_transition\_from\_to^{stage\_from,stage}_{t})}_{Transition}
\end{equation}

### A share of the loans in each stage will mature another share will have to be written off

\begin{equation}
\label{eq:maturing}
loan\_maturing^{stage}_{t}  = M^{stage}_{t} \times loan^{stage}_{t-1}
\end{equation}

\begin{equation}
\label{eq:writeoff}
loan\_writeoff^{stage}_{t}  = WRO^{stage}_{t} \times loan^{stage}_{t-1}
\end{equation}

### So loans in a stage will reflect the inflow and outflow

\begin{equation}
\label{eq:E}
loan^{stage}_{t}  = \underbrace{loan\_transition\_to^{stage}_{t} }_{Transition}
-\underbrace{loan\_maturing^{stage}_{t}}_{Maturing}
-\underbrace{loan\_writeoff^{stage}_{t}}_{Writeoff}
+\underbrace{loan\_new^{stage}_{t}}_{New loans}
\end{equation}

\begin{equation}
\label{eq:new}
loan\_new^{stage}_{t}  = new^{stage}_{t} \times loan^{stage}_{t-1}
\end{equation}

\begin{equation}
\label{eq:g}
loan\_total_{t}  =  \sum_{stage}(loan^{stage}_{t})
\end{equation}

### New loans are only generated in stage 1.

\begin{equation}
\label{eq:E2}
new\_s1_{t} = \frac{loan\_growth_{t} \times loan\_total_{t-1} +
                \sum_{stage}((M^{stage}_{t}+WRO^{stage}_{t})\times loan^{stage}_{t-1})}{(loan\_s1_{t-1})}
\end{equation}

### Performing Loans

\begin{equation}
\label{eq:Performing}
loan\_performing_{t} = loan\_s1_{t}+loan\_s2_{t}
\end{equation}

### Cure

\begin{equation}
\label{eq:cure}
loan\_cure = loan\_transition\_from\_to\_s3\_s1+loan\_transition\_from\_to\_s3\_s2
\end{equation}

### Probability of default (PD)

The point in time PD is the fraction of loans in stage s1 and s2 going into stage s3
\begin{equation}
\label{eq:PDPIT}
PD\_pit= \frac{loan\_transition\_from\_to\_s1\_s3+loan\_transition\_from\_to\_s2\_s3}{loan\_s1+loan\_s2}
\end{equation}

The Troug The  Cycle PD is a slow mowing average of the Point in time PD. The

\begin{equation}
\label{eq:PDTTC}
PD\_TTC = logit^{-1}(logit(PD\_TTC(-1)) + alfa \times \Delta{PD\_pit})
\end{equation}

### And we can specify the dynamic of the transition matrix, based on Z score

Let $\Phi$ be the normal cumulative distribution $\frac{1}{\sqrt{2\pi}}
\int_{-\infty}^x e^{
-\frac{t^2}{2}}dt$

\begin{equation}
\label{eq:tr}
TR^{stage\_from,stage}_{t} = \Phi{\left(\frac{bound\_upper^{stage\_from,stage}-\sqrt{\rho}\times Z_{t}}{\sqrt{1-\rho}}\right)}-\Phi{\left(\frac{bound\_lower^{stage\_from,stage}-\sqrt{\rho}\times Z_{t}}{\sqrt{1-\rho}}\right)}
\end{equation}

```

# Now creating the model **ecl**

## The model

### Loans can be in 3 stages
Loans can be in 3 stages, s1,s2 and s3.
New loans will be generated and loans will mature.

Two lists of stages are defined:

List $stage=\{s1, s2,s3\}$

List $stage\_from=\{s1, s2,s3\}$

### A share of the loans in each stage wil transition to the same or another stage in the next time frame:

\begin{equation}
\label{eq:loanfromto}
loan\_transition\_from\_to^{stage\_from,stage}_{t} = loan^{stage\_from}_{t-1} \times TR^{stage\_from,stage}_{t}
\end{equation}

\begin{equation}
\label{eq:transition}
loan\_transition\_to^{stage}_{t}  = \underbrace{\sum_{stage\_from}(loan\_transition\_from\_to^{stage\_from,stage}_{t})}_{Transition}
\end{equation}

### A share of the loans in each stage will mature another share will have to be written off

\begin{equation}
\label{eq:maturing}
loan\_maturing^{stage}_{t}  = M^{stage}_{t} \times loan^{stage}_{t-1}
\end{equation}

\begin{equation}
\label{eq:writeoff}
loan\_writeoff^{stage}_{t}  = WRO^{stage}_{t} \times loan^{stage}_{t-1}
\end{equation}

### So loans in a stage will reflect the inflow and outflow

\begin{equation}
\label{eq:E}
loan^{stage}_{t}  = \underbrace{loan\_transition\_to^{stage}_{t} }_{Transition}
-\underbrace{loan\_maturing^{stage}_{t}}_{Maturing}
-\underbrace{loan\_writeoff^{stage}_{t}}_{Writeoff}
+\underbrace{loan\_new^{stage}_{t}}_{New loans}
\end{equation}

\begin{equation}
\label{eq:new}
loan\_new^{stage}_{t}  = new^{stage}_{t} \times loan^{stage}_{t-1}
\end{equation}

\begin{equation}
\label{eq:g}
loan\_total_{t}  =  \sum_{stage}(loan^{stage}_{t})
\end{equation}

### New loans are only generated in stage 1.

\begin{equation}
\label{eq:E2}
new\_s1_{t} = \frac{loan\_growth_{t} \times loan\_total_{t-1} +
                \sum_{stage}((M^{stage}_{t}+WRO^{stage}_{t})\times loan^{stage}_{t-1})}{(loan\_s1_{t-1})}
\end{equation}

### Performing Loans

\begin{equation}
\label{eq:Performing}
loan\_performing_{t} = loan\_s1_{t}+loan\_s2_{t}
\end{equation}

### Cure

\begin{equation}
\label{eq:cure}
loan\_cure = loan\_transition\_from\_to\_s3\_s1+loan\_transition\_from\_to\_s3\_s2
\end{equation}

### Probability of default (PD)

The point in time PD is the fraction of loans in stage s1 and s2 going into stage s3
\begin{equation}
\label{eq:PDPIT}
PD\_pit= \frac{loan\_transition\_from\_to\_s1\_s3+loan\_transition\_from\_to\_s2\_s3}{loan\_s1+loan\_s2}
\end{equation}

The Troug The  Cycle PD is a slow mowing average of the Point in time PD. The

\begin{equation}
\label{eq:PDTTC}
PD\_TTC = logit^{-1}(logit(PD\_TTC(-1)) + alfa \times \Delta{PD\_pit})
\end{equation}

### And we can specify the dynamic of the transition matrix, based on Z score

Let $\Phi$ be the normal cumulative distribution $\frac{1}{\sqrt{2\pi}}
\int_{-\infty}^x e^{
-\frac{t^2}{2}}dt$

\begin{equation}
\label{eq:tr}
TR^{stage\_from,stage}_{t} = \Phi{\left(\frac{bound\_upper^{stage\_from,stage}-\sqrt{\rho}\times Z_{t}}{\sqrt{1-\rho}}\right)}-\Phi{\left(\frac{bound\_lower^{stage\_from,stage}-\sqrt{\rho}\times Z_{t}}{\sqrt{1-\rho}}\right)}
\end{equation}

## The equations in Macro business logic language can be inspected

```python
print(ecl.equations_original)
```

    Do stage_from $ Do stage $
     Frml loanfromto loan_transition_from_to_{stage_from}_{stage}  = loan_{stage_from}(-1)*TR_{stage_from}_{stage} $
    enddo $ enddo $
    Do stage $
     Frml transition loan_transition_to_{stage}   = sum(stage_from,loan_transition_from_to_{stage_from}_{stage}) $
    enddo $
    Do stage $
     Frml maturing loan_maturing_{stage}   = M_{stage}*loan_{stage}(-1) $
    enddo $
    Do stage $
     Frml writeoff loan_writeoff_{stage}   = WRO_{stage}*loan_{stage}(-1) $
    enddo $
    Do stage $
     Frml E loan_{stage}   = loan_transition_to_{stage}-loan_maturing_{stage}-loan_writeoff_{stage}+loan_new_{stage} $
    enddo $
    Do stage $
     Frml new loan_new_{stage}   = new_{stage}*loan_{stage}(-1) $
    enddo $
    Frml g loan_total   = sum(stage,loan_{stage}) $
    Frml E2 new_s1  = ((loan_growth*loan_total(-1)+sum(stage,(M_{stage}+WRO_{stage})*loan_{stage}(-1)))/((loan_s1(-1)))) $
    Frml Performing loan_performing  = loan_s1+loan_s2 $
    Frml cure loan_cure  = loan_transition_from_to_s3_s1+loan_transition_from_to_s3_s2 $
    Frml PDPIT PD_pit = ((loan_transition_from_to_s1_s3+loan_transition_from_to_s2_s3)/(loan_s1+loan_s2)) $
    Frml PDTTC PD_TTC  = logit_inverse(logit(PD_TTC(-1))+alfa*diff(PD_pit)) $
    Do stage_from $ Do stage $
     Frml tr TR_{stage_from}_{stage}  = NORM.CDF((((bound_upper_{stage_from}_{stage}-sqrt(rho)*Z)/(sqrt(1-rho)))))-NORM.CDF((((bound_lower_{stage_from}_{stage}-sqrt(rho)*Z)/(sqrt(1-rho))))) $
    enddo $ enddo $
    LIST STAGE =  STAGE : S1  S2 S3$
    LIST STAGE_FROM =  STAGE_FROM : S1  S2 S3$

### The equations in business logic language can be inspected

```python
print(ecl.equations)
```

    FRML LOANFROMTO LOAN_TRANSITION_FROM_TO_S1_S1  = LOAN_S1(-1)*TR_S1_S1 $
    FRML LOANFROMTO LOAN_TRANSITION_FROM_TO_S1_S2  = LOAN_S1(-1)*TR_S1_S2 $
    FRML LOANFROMTO LOAN_TRANSITION_FROM_TO_S1_S3  = LOAN_S1(-1)*TR_S1_S3 $
    FRML LOANFROMTO LOAN_TRANSITION_FROM_TO_S2_S1  = LOAN_S2(-1)*TR_S2_S1 $
    FRML LOANFROMTO LOAN_TRANSITION_FROM_TO_S2_S2  = LOAN_S2(-1)*TR_S2_S2 $
    FRML LOANFROMTO LOAN_TRANSITION_FROM_TO_S2_S3  = LOAN_S2(-1)*TR_S2_S3 $
    FRML LOANFROMTO LOAN_TRANSITION_FROM_TO_S3_S1  = LOAN_S3(-1)*TR_S3_S1 $
    FRML LOANFROMTO LOAN_TRANSITION_FROM_TO_S3_S2  = LOAN_S3(-1)*TR_S3_S2 $
    FRML LOANFROMTO LOAN_TRANSITION_FROM_TO_S3_S3  = LOAN_S3(-1)*TR_S3_S3 $
    FRML TRANSITION LOAN_TRANSITION_TO_S1   = (LOAN_TRANSITION_FROM_TO_S1_S1+LOAN_TRANSITION_FROM_TO_S2_S1+LOAN_TRANSITION_FROM_TO_S3_S1) $
    FRML TRANSITION LOAN_TRANSITION_TO_S2   = (LOAN_TRANSITION_FROM_TO_S1_S2+LOAN_TRANSITION_FROM_TO_S2_S2+LOAN_TRANSITION_FROM_TO_S3_S2) $
    FRML TRANSITION LOAN_TRANSITION_TO_S3   = (LOAN_TRANSITION_FROM_TO_S1_S3+LOAN_TRANSITION_FROM_TO_S2_S3+LOAN_TRANSITION_FROM_TO_S3_S3) $
    FRML MATURING LOAN_MATURING_S1   = M_S1*LOAN_S1(-1) $
    FRML MATURING LOAN_MATURING_S2   = M_S2*LOAN_S2(-1) $
    FRML MATURING LOAN_MATURING_S3   = M_S3*LOAN_S3(-1) $
    FRML WRITEOFF LOAN_WRITEOFF_S1   = WRO_S1*LOAN_S1(-1) $
    FRML WRITEOFF LOAN_WRITEOFF_S2   = WRO_S2*LOAN_S2(-1) $
    FRML WRITEOFF LOAN_WRITEOFF_S3   = WRO_S3*LOAN_S3(-1) $
    FRML E LOAN_S1   = LOAN_TRANSITION_TO_S1-LOAN_MATURING_S1-LOAN_WRITEOFF_S1+LOAN_NEW_S1 $
    FRML E LOAN_S2   = LOAN_TRANSITION_TO_S2-LOAN_MATURING_S2-LOAN_WRITEOFF_S2+LOAN_NEW_S2 $
    FRML E LOAN_S3   = LOAN_TRANSITION_TO_S3-LOAN_MATURING_S3-LOAN_WRITEOFF_S3+LOAN_NEW_S3 $
    FRML NEW LOAN_NEW_S1   = NEW_S1*LOAN_S1(-1) $
    FRML NEW LOAN_NEW_S2   = NEW_S2*LOAN_S2(-1) $
    FRML NEW LOAN_NEW_S3   = NEW_S3*LOAN_S3(-1) $
    FRML G LOAN_TOTAL   = (LOAN_S1+LOAN_S2+LOAN_S3) $
    FRML E2 NEW_S1  = ((LOAN_GROWTH*LOAN_TOTAL(-1)+((M_S1+WRO_S1)*LOAN_S1(-1)+(M_S2+WRO_S2)*LOAN_S2(-1)+(M_S3+WRO_S3)*LOAN_S3(-1)))/((LOAN_S1(-1)))) $
    FRML PERFORMING LOAN_PERFORMING  = LOAN_S1+LOAN_S2 $
    FRML CURE LOAN_CURE  = LOAN_TRANSITION_FROM_TO_S3_S1+LOAN_TRANSITION_FROM_TO_S3_S2 $
    FRML PDPIT PD_PIT = ((LOAN_TRANSITION_FROM_TO_S1_S3+LOAN_TRANSITION_FROM_TO_S2_S3)/(LOAN_S1+LOAN_S2)) $
    FRML PDTTC PD_TTC  = LOGIT_INVERSE(LOGIT(PD_TTC(-1))+ALFA*((PD_PIT)-(PD_PIT(-1)))) $
    FRML TR TR_S1_S1  = NORM.CDF((((BOUND_UPPER_S1_S1-SQRT(RHO)*Z)/(SQRT(1-RHO)))))-NORM.CDF((((BOUND_LOWER_S1_S1-SQRT(RHO)*Z)/(SQRT(1-RHO))))) $
    FRML TR TR_S1_S2  = NORM.CDF((((BOUND_UPPER_S1_S2-SQRT(RHO)*Z)/(SQRT(1-RHO)))))-NORM.CDF((((BOUND_LOWER_S1_S2-SQRT(RHO)*Z)/(SQRT(1-RHO))))) $
    FRML TR TR_S1_S3  = NORM.CDF((((BOUND_UPPER_S1_S3-SQRT(RHO)*Z)/(SQRT(1-RHO)))))-NORM.CDF((((BOUND_LOWER_S1_S3-SQRT(RHO)*Z)/(SQRT(1-RHO))))) $
    FRML TR TR_S2_S1  = NORM.CDF((((BOUND_UPPER_S2_S1-SQRT(RHO)*Z)/(SQRT(1-RHO)))))-NORM.CDF((((BOUND_LOWER_S2_S1-SQRT(RHO)*Z)/(SQRT(1-RHO))))) $
    FRML TR TR_S2_S2  = NORM.CDF((((BOUND_UPPER_S2_S2-SQRT(RHO)*Z)/(SQRT(1-RHO)))))-NORM.CDF((((BOUND_LOWER_S2_S2-SQRT(RHO)*Z)/(SQRT(1-RHO))))) $
    FRML TR TR_S2_S3  = NORM.CDF((((BOUND_UPPER_S2_S3-SQRT(RHO)*Z)/(SQRT(1-RHO)))))-NORM.CDF((((BOUND_LOWER_S2_S3-SQRT(RHO)*Z)/(SQRT(1-RHO))))) $
    FRML TR TR_S3_S1  = NORM.CDF((((BOUND_UPPER_S3_S1-SQRT(RHO)*Z)/(SQRT(1-RHO)))))-NORM.CDF((((BOUND_LOWER_S3_S1-SQRT(RHO)*Z)/(SQRT(1-RHO))))) $
    FRML TR TR_S3_S2  = NORM.CDF((((BOUND_UPPER_S3_S2-SQRT(RHO)*Z)/(SQRT(1-RHO)))))-NORM.CDF((((BOUND_LOWER_S3_S2-SQRT(RHO)*Z)/(SQRT(1-RHO))))) $
    FRML TR TR_S3_S3  = NORM.CDF((((BOUND_UPPER_S3_S3-SQRT(RHO)*Z)/(SQRT(1-RHO)))))-NORM.CDF((((BOUND_LOWER_S3_S3-SQRT(RHO)*Z)/(SQRT(1-RHO))))) $
    LIST STAGE =  STAGE : S1  S2 S3$
    LIST STAGE_FROM =  STAGE_FROM : S1  S2 S3$

### The model structure can be inspected

```python
ecl.drawmodel(sink='LOAN_TOTAL',HR=0,pdf=0,att=0,size=(12,12))
```

## load data
The data is copy-pasted from the excel sheet
### Load data common for baseline and adverse

```python
%%dataframe T0tr noshow  prefix=TR_  periods=1 melt
	S1	S2	S3
S1	89%	9%	2%
S2	15%	79%	6%
S3	2%	24%	74%
```

```python
%%dataframe  startvalues  show  periods=1
     pd_pit lgd_pit
        1.6%   20%
```

      PD_PIT
      LGD_PIT

      index

      2021
      0.02
      0.2

```python
%%dataframe  Loan_t0  show periods=1 melt
S1	S2	S3 total
loan 500	180	30 710

```

      S1
      S2
      S3
      TOTAL

      LOAN
      500.0
      180.0
      30.0
      710.0

      var_name
      LOAN_S1
      LOAN_S2
      LOAN_S3
      LOAN_TOTAL

      index

      2021
      500.0
      180.0
      30.0
      710.0

```python
%%dataframe upper_bin prefix=bound_upper_ nshow periods=7 melt
	S1	S2	S3
S1	10000	-1,34	-2,33
S2	10000	1,08	-1,64
S3	10000	1,88	0,58

```

```python
%%dataframe lower_bin prefix=bound_lower_ show periods=7 melt
	S1	S2	S3
S1	-1,34	-2,33	-10000
S2	1,08	-1,64	-10000
S3	1,88	0,58	-10000
```

      S1
      S2
      S3

      S1
      -1.34
      -2.33
      -10000.0

      S2
      1.08
      -1.64
      -10000.0

      S3
      1.88
      0.58
      -10000.0

      var_name
      BOUND_LOWER_S1_S1
      BOUND_LOWER_S2_S1
      BOUND_LOWER_S3_S1
      BOUND_LOWER_S1_S2
      BOUND_LOWER_S2_S2
      BOUND_LOWER_S3_S2
      BOUND_LOWER_S1_S3
      BOUND_LOWER_S2_S3
      BOUND_LOWER_S3_S3

      index

      2021
      -1.34
      1.08
      1.88
      -2.33
      -1.64
      0.58
      -10000.0
      -10000.0
      -10000.0

      2022
      -1.34
      1.08
      1.88
      -2.33
      -1.64
      0.58
      -10000.0
      -10000.0
      -10000.0

      2023
      -1.34
      1.08
      1.88
      -2.33
      -1.64
      0.58
      -10000.0
      -10000.0
      -10000.0

      2024
      -1.34
      1.08
      1.88
      -2.33
      -1.64
      0.58
      -10000.0
      -10000.0
      -10000.0

      2025
      -1.34
      1.08
      1.88
      -2.33
      -1.64
      0.58
      -10000.0
      -10000.0
      -10000.0

      2026
      -1.34
      1.08
      1.88
      -2.33
      -1.64
      0.58
      -10000.0
      -10000.0
      -10000.0

      2027
      -1.34
      1.08
      1.88
      -2.33
      -1.64
      0.58
      -10000.0
      -10000.0
      -10000.0

```python
%%dataframe parameters show periods=7
rho
0.2
```

      RHO

      index

      2021
      0.2

      2022
      0.2

      2023
      0.2

      2024
      0.2

      2025
      0.2

      2026
      0.2

      2027
      0.2

### Create the Static dataframe, common to scenarios

```python
staticdf = pd.concat([T0tr_melted,Loan_t0_melted,startvalues,upper_bin_melted,
                      lower_bin_melted,parameters],axis=1)
HTML(staticdf.T.style.render().replace('nan',''))
```

      index
      2021
      2022
      2023
      2024
      2025
      2026
      2027

      TR_S1_S1
      0.89

      TR_S2_S1
      0.15

      TR_S3_S1
      0.02

      TR_S1_S2
      0.09

      TR_S2_S2
      0.79

      TR_S3_S2
      0.24

      TR_S1_S3
      0.02

      TR_S2_S3
      0.06

      TR_S3_S3
      0.74

      LOAN_S1
      500.00

      LOAN_S2
      180.00

      LOAN_S3
      30.00

      LOAN_TOTAL
      710.00

      PD_PIT
      0.02

      LGD_PIT
      0.20

      BOUND_UPPER_S1_S1
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00

      BOUND_UPPER_S2_S1
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00

      BOUND_UPPER_S3_S1
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00

      BOUND_UPPER_S1_S2
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34

      BOUND_UPPER_S2_S2
      1.08
      1.08
      1.08
      1.08
      1.08
      1.08
      1.08

      BOUND_UPPER_S3_S2
      1.88
      1.88
      1.88
      1.88
      1.88
      1.88
      1.88

      BOUND_UPPER_S1_S3
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33

      BOUND_UPPER_S2_S3
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64

      BOUND_UPPER_S3_S3
      0.58
      0.58
      0.58
      0.58
      0.58
      0.58
      0.58

      BOUND_LOWER_S1_S1
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34

      BOUND_LOWER_S2_S1
      1.08
      1.08
      1.08
      1.08
      1.08
      1.08
      1.08

      BOUND_LOWER_S3_S1
      1.88
      1.88
      1.88
      1.88
      1.88
      1.88
      1.88

      BOUND_LOWER_S1_S2
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33

      BOUND_LOWER_S2_S2
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64

      BOUND_LOWER_S3_S2
      0.58
      0.58
      0.58
      0.58
      0.58
      0.58
      0.58

      BOUND_LOWER_S1_S3
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00

      BOUND_LOWER_S2_S3
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00

      BOUND_LOWER_S3_S3
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00

      RHO
      0.20
      0.20
      0.20
      0.20
      0.20
      0.20
      0.20

### Load data specific for the scenarios

```python
%%dataframe  inf_baseline  nshow t periods=7 melt
    m     wro
s1  5%   0
s2  3.8%   0
s3  0    7.5%
```

```python
%%dataframe  inf_adverse  nshow t periods=7 melt
    m     wro
s1  3.8%   0
s2  2.5%   0
s3  0    6.3%
```

```python
%%dataframe projection_baseline nshow
Z  loan_growth
 0          0.01
-0,47	    0.01
-0,42	    0.01
-0,38	    0.01
-0,36	    0.01
-0,34	    0.01
-0,33	    0.01
```

```python
%%dataframe projection_adverse nshow
 Z        loan_growth
 0          0.01
  -0,65     -0.01
  -0,84     -0.008
   -0,99    -0.006
    -0,69   -0.004
    -0,39   -0.002
    -0,24   -0.0
```

### Create a dataframe for baseline and adverse scenario

```python
baseupdate    = pd.concat([inf_baseline_melted, projection_baseline],axis=1).pipe(sortdf)
adverseupdate =  pd.concat([inf_adverse_melted, projection_adverse],axis=1).pipe(sortdf)
getnotnul = lambda x: x.loc[:,(x!=0.0).any(axis=0)]

display(Markdown('## Baseline scenario'))
display(baseupdate.pipe(getnotnul))
display(Markdown('## Adverse scenario'))
display(adverseupdate.pipe(getnotnul))
```

## Baseline scenario

      LOAN_GROWTH
      M_S1
      M_S2
      WRO_S3
      Z

      index

      2021
      0.01
      0.05
      0.04
      0.07
      0.00

      2022
      0.01
      0.05
      0.04
      0.07
      -0.47

      2023
      0.01
      0.05
      0.04
      0.07
      -0.42

      2024
      0.01
      0.05
      0.04
      0.07
      -0.38

      2025
      0.01
      0.05
      0.04
      0.07
      -0.36

      2026
      0.01
      0.05
      0.04
      0.07
      -0.34

      2027
      0.01
      0.05
      0.04
      0.07
      -0.33

## Adverse scenario

      LOAN_GROWTH
      M_S1
      M_S2
      WRO_S3
      Z

      index

      2021
      1.00e-02
      0.04
      0.03
      0.06
      0.00

      2022
      -1.00e-02
      0.04
      0.03
      0.06
      -0.65

      2023
      -8.00e-03
      0.04
      0.03
      0.06
      -0.84

      2024
      -6.00e-03
      0.04
      0.03
      0.06
      -0.99

      2025
      -4.00e-03
      0.04
      0.03
      0.06
      -0.69

      2026
      -2.00e-03
      0.04
      0.03
      0.06
      -0.39

      2027
      -0.00e+00
      0.04
      0.03
      0.06
      -0.24

```python
baseline = pd.concat([staticdf, inf_baseline_melted, projection_baseline],axis=1)
adverse  = pd.concat([staticdf, inf_adverse_melted, projection_adverse],axis=1)
baseline.index = pd.period_range(start=2021,freq = 'Y',periods=7)
adverse.index  = pd.period_range(start=2021,freq = 'Y',periods=7)

display(baseline.pipe(sortdf).T)
```

      2021
      2022
      2023
      2024
      2025
      2026
      2027

      BOUND_LOWER_S1_S1
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34

      BOUND_LOWER_S1_S2
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33

      BOUND_LOWER_S1_S3
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00

      BOUND_LOWER_S2_S1
      1.08
      1.08
      1.08
      1.08
      1.08
      1.08
      1.08

      BOUND_LOWER_S2_S2
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64

      BOUND_LOWER_S2_S3
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00

      BOUND_LOWER_S3_S1
      1.88
      1.88
      1.88
      1.88
      1.88
      1.88
      1.88

      BOUND_LOWER_S3_S2
      0.58
      0.58
      0.58
      0.58
      0.58
      0.58
      0.58

      BOUND_LOWER_S3_S3
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00
      -10000.00

      BOUND_UPPER_S1_S1
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00

      BOUND_UPPER_S1_S2
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34
      -1.34

      BOUND_UPPER_S1_S3
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33
      -2.33

      BOUND_UPPER_S2_S1
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00

      BOUND_UPPER_S2_S2
      1.08
      1.08
      1.08
      1.08
      1.08
      1.08
      1.08

      BOUND_UPPER_S2_S3
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64
      -1.64

      BOUND_UPPER_S3_S1
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00
      10000.00

      BOUND_UPPER_S3_S2
      1.88
      1.88
      1.88
      1.88
      1.88
      1.88
      1.88

      BOUND_UPPER_S3_S3
      0.58
      0.58
      0.58
      0.58
      0.58
      0.58
      0.58

      LGD_PIT
      0.20
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      LOAN_GROWTH
      0.01
      0.01
      0.01
      0.01
      0.01
      0.01
      0.01

      LOAN_S1
      500.00
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      LOAN_S2
      180.00
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      LOAN_S3
      30.00
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      LOAN_TOTAL
      710.00
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      M_S1
      0.05
      0.05
      0.05
      0.05
      0.05
      0.05
      0.05

      M_S2
      0.04
      0.04
      0.04
      0.04
      0.04
      0.04
      0.04

      M_S3
      0.00
      0.00
      0.00
      0.00
      0.00
      0.00
      0.00

      PD_PIT
      0.02
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      RHO
      0.20
      0.20
      0.20
      0.20
      0.20
      0.20
      0.20

      TR_S1_S1
      0.89
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      TR_S1_S2
      0.09
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      TR_S1_S3
      0.02
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      TR_S2_S1
      0.15
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      TR_S2_S2
      0.79
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      TR_S2_S3
      0.06
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      TR_S3_S1
      0.02
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      TR_S3_S2
      0.24
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      TR_S3_S3
      0.74
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN

      WRO_S1
      0.00
      0.00
      0.00
      0.00
      0.00
      0.00
      0.00

      WRO_S2
      0.00
      0.00
      0.00
      0.00
      0.00
      0.00
      0.00

      WRO_S3
      0.07
      0.07
      0.07
      0.07
      0.07
      0.07
      0.07

      Z
      0.00
      -0.47
      -0.42
      -0.38
      -0.36
      -0.34
      -0.33

## Run baseline

```python
base_result  = ecl(baseline,keep='Baseline')
```

    Will start calculating: testmodel
    2022  solved
    2023  solved
    2024  solved
    2025  solved
    2026  solved
    2027  solved
    testmodel calculated

### Save model and baseline results

```python
ecl.modeldump('ecl.pcim')
```

## Run adverse

```python
adverse_result = ecl(adverse,keep = 'Adverse')
```

    Will start calculating: testmodel
    2022  solved
    2023  solved
    2024  solved
    2025  solved
    2026  solved
    2027  solved
    testmodel calculated

## Inspect Results

```python
with ecl.set_smpl('2021','2027'):
    ecl.keep_plot('loan_total',showtype='growth',legend=False);
```

```python
ecl.keep_plot('tr_*',showtype='level',legend=False);
```

```python
ecl.PD_PIT.tracepre(up=4,HR=1,filter=30,pdf=1)
```

```python

```

Examples/Economic Credit Loss/ECL Z setup.ipynb


# Deep Equilibrium Nets for Guvenen (2009)

#### By  Matias Covarrubias and Min Fang

In this notebook, we use `TensorFlow` to solve [Guvenen (2009)](http://users.econ.umn.edu/~guvenen/HABHET2008.pdf) with _deep equilibrium nets_ method by [Azinovic, Gaegauf, & Scheidegger (2020)](https://papers.ssrn.com/sol3/papers.cfm?abstract_id=3393482).

## The model

GDSGE offers a good [summary](http://www.gdsge.com/example/Guvenen2009/Guvenen2009.html).

## Table of contents

0. [Set up workspace](#workspace)
1. [Model calibration](#modelcal)
2. [_Deep equilibrium net_ hyper-parameters](#deqnparam)
    1. [Neural network](#nn)
3. [Economic model](#econmodel)
    1. [Current period (t)](#currentperiod)
    2. [Next period (t+1)](#nextperiod)
    3. [Cost/Euler function](#cost)
4. [Training](#training)

## 0. Set up workspace

First, we need to set up the workspace. All of the packages are standard python packages. This version of the _deep equilibrium net_ notebook will be computed with `TensorFlow 2`. Make sure you are working in an environment with TF2 installed.

The only special module is `utils` from which we import a mini-batch function `random_mini_batches` and a function that initializes the neural network weights `initialize_nn_weight`.

### Saving and continuing training

You can save and resume training by saving and loading the tensorflow session and data starting point.
* The saved session stores the neural network weights and the optimizer's state. If you have saved a session that you would like to reload, set `sess_path` to the session checkpoint path. For example, to load the 100th episode's session set `sess_path = './output/sess_100.ckpt'`. Otherwise, set `sess_path` to `None` to train from scratch. Currently, this script saves the session at the end of each [episode](#deqnparam).
* The saved data starting point stores the an exogeneous shock and a capital distribution, which can be used to simulate states into the future from. If you have saved a starting point that you would like to reload, set `data_path` to the numpy data path. For example, to load the 100th episode's starting point set `data_path = './output/data_100.npy'`. Otherwise, set `data_path` to `None` to train from scratch.

```python
import tensorflow as tf
print('tf version:', tf.__version__)
```

    tf version: 2.4.1

```python
%matplotlib notebook

# Import modules
import os
import re
from datetime import datetime

import numpy as np
from utils import initialize_nn_weight, random_mini_batches

import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 12})
plt.rc('xtick', labelsize='small')
plt.rc('ytick', labelsize='small')
std_figsize = (4, 4)

# Set the seed for replicable results
seed = 0
np.random.seed(seed)
tf.random.set_seed(seed) # was "tf.set_random_seed(seed)" in tf1

# Helper variables
eps = 0.00001  # Small epsilon value

# Make output directory to save network weights and starting point
if not os.path.exists('./output'):
    os.mkdir('./output')

# Path to saved tensorflow session
sess_path = None
# Path to saved data starting point
data_path = None
```

## 1. Model calibration

```python

```

```python

```

## 2. _Deep equilibrium net_ hyper-parameters

Here, we will briefly outline the _deep equilibrium net_ architecture and choice of hyper-parameters. By hyper-parameters, we refer too all of the parameters that can be chosen freely by the modeller. This is in contrast to "parameters", which refers to the neural network weights that are learned during training.
We found it helpful to augment the state we pass to the neural network with redundant information, such as aggregate variables and the distribution of financial wealth.
In this notebook, we do the same.

  * Neural network architecture (see NN diagram below):
    * The neural network takes an extended state that is 32-dimensional:
        * Natural state:
            * 1 dimension for the shock index,
            * 1 dimension for aggregate capital
            * 1 dimension for bond position
        * Redundant extension: (do we need this?)
            * 1 dimension for the current TFP,
            * 1 dimension for the current depreciation,
            * 1 dimension for aggregate capital,
            * 1 dimension for aggregate labor supply,
            * 1 dimension for the return on capital,
            * 1 dimension for the wage,
            * 1 dimension for the aggregate production,
            * 6 dimensions for the distribution of financial wealth,
            * 6 dimensions for the distribution of labor income,
            * 6 dimensions for the distribution of total income.
    * The first and second hidden layers have 100 and 50 nodes, respectively.
    * The output layer is 5-dimensional.
  * Training hyper-parameters:
    * We simulate 5'000 episodes.*
    * Per episode, we compute 20 epochs.
    * We use a minibatch size of 512.
    * Each episode is 10'240 periods long.
    * The learning rate is set to 0.00001.

\* In the _deep equilibrium net_ framework we iterate between a simulation and training phase (see [training](#training)). We first simulate a training dataset based on the current state of the neural network. That is, we simulate a random sequence of exogenous shocks for which we compute the capital holdings of the agents. We call the set of simulated periods an _episode_. 5'000 episodes means that we re-simulate the training data 5'000 times.

Note that this calibration is not optimal for the economic model being solved. This script is intended to provide a simple introduction to _deep equilibrium nets_ and not an indepth discussion on optimizing the performance.

DO CRRA first;
DO use less controls as possible;

```python
### Using subclassing to build layers and model
import tensorflow as tf
from tensorflow import keras

# Neural network architecture parameters
num_input_nodes = 3  # Dimension of extended state space
num_hidden_nodes1 = 100 # Dimension of hidden layers 1
num_hidden_nodes2 = 50 # Dimension of hidden layers 2
num_output_nodes = 17  # Output dimension
```

```python
### Build Layers
class hidden_layer1(keras.layers.Layer):
    def __init__(self, units=None, input_dim=3):  #should I start with None?
        super(hidden_layer1, self).__init__()
        w_init = tf.random_normal_initializer()
        self.w = tf.Variable(
            initial_value=w_init(shape=(input_dim, units), dtype="float32"),
            trainable=True,
        )
        b_init = tf.zeros_initializer()
        self.b = tf.Variable(
            initial_value=b_init(shape=(units,), dtype="float32"), trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

class hidden_layer2(keras.layers.Layer):
    def __init__(self, units=None, input_dim=50):  #should I start with None?
        super(hidden_layer1, self).__init__()
        w_init = tf.random_normal_initializer()
        self.w = tf.Variable(
            initial_value=w_init(shape=(input_dim, units), dtype="float32"),
            trainable=True,
        )
        b_init = tf.zeros_initializer()
        self.b = tf.Variable(
            initial_value=b_init(shape=(units,), dtype="float32"), trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

class outer_layer(keras.layers.Layer):
    def __init__(self, units=None, input_dim=100):  #should I start with None?
        super(hidden_layer1, self).__init__()
        w_init = tf.random_normal_initializer()
        self.w = tf.Variable(
            initial_value=w_init(shape=(input_dim, units), dtype="float32"),
            trainable=True,
        )
        b_init = tf.zeros_initializer()
        self.b = tf.Variable(
            initial_value=b_init(shape=(units,), dtype="float32"), trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

```

```python
class Guvenen(Model):
  def __init__(self):
    super(Guvenen, self).__init__()
    self.conv1 = Conv2D(32, 3, activation='relu')
    self.flatten = Flatten()
    self.d1 = Dense(128, activation='relu')
    self.d2 = Dense(10, activation='softmax')

  def call(self, x):
    x = self.conv1(x)
    x = self.flatten(x)
    x = self.d1(x)
    return self.d2(x)

model = Guvenen()
```

```python
num_episodes = 5000
len_episodes = 10240
epochs_per_episode = 20
minibatch_size = 512
num_minibatches = int(len_episodes / minibatch_size)
lr = 0.00001

```

### 2.A. Neural network

Since we are using a neural network with 2 hidden layers, the network maps:
$$X \rightarrow \mathcal{N}(X) = \sigma(\sigma(XW_1 + b_1)W_2 + b_2)W_3 + b_3$$
where $\sigma$ is the rectified linear unit (ReLU) activation function and the output layer is activated with the linear function (which is the identity function and hence omitted from the equation). Therefore, we need 3 weight matrices $\{W_1, W_2, W_3\}$ and 3 bias vectors $\{b_1, b_2, b_3\}$ (compare with the neural network diagram above). In total, we train $300+5000+850+167$ parameters:

| $W_1$ | $3 \times 100$  | $= 300$ |
| --- | --- | --- |
| $W_2$ | $100 \times 50$ | $= 5000$ |
| $W_3$ | $50 \times 17$ | $= 850$ |
| $b_1 + b_2 + b_3$ | $100 + 50 + 17$ | $= 167$|

We initialize the neural network parameters with the  `initialize_nn_weight` helper function from `utils`.

Then, we compute the neural network prediction using the parameters in `nn_predict`.

```python
# We create a placeholder for X, the input data for the neural network, which corresponds
# to the state.
X = tf.placeholder(tf.float32, shape=(None, num_input_nodes))
# Get number samples
m = tf.shape(X)[0]

# We create all of the neural network weights and biases. The weights are matrices that
# connect the layers of the neural network. For example, W1 connects the input layer to
# the first hidden layer
W1 = initialize_nn_weight([num_input_nodes, num_hidden_nodes[0]])
W2 = initialize_nn_weight([num_hidden_nodes[0], num_hidden_nodes[1]])
W3 = initialize_nn_weight([num_hidden_nodes[1], num_output_nodes])

# The biases are extra (shift) terms that are added to each node in the neural network.
b1 = initialize_nn_weight([num_hidden_nodes[0]])
b2 = initialize_nn_weight([num_hidden_nodes[1]])
b3 = initialize_nn_weight([num_output_nodes])

# Then, we create a function, to which we pass X, that generates a prediction based on
# the current neural network weights. Note that the hidden layers are ReLU activated.
# The output layer is not activated (i.e., it is activated with the linear function).
def nn_predict(X):
    hidden_layer1 = tf.nn.relu(tf.add(tf.matmul(X, W1), b1))
    hidden_layer2 = tf.nn.relu(tf.add(tf.matmul(hidden_layer1, W2), b2))
    output_layer = tf.add(tf.matmul(hidden_layer2, W3), b3)
    return output_layer

```

## 3. Economic model

In this section, we implement the economics outlined in [the model](#model).

Each period, based on the current distribution of capital and the exogenous state, agents decide whether and how much to save in risky capital and to consume. Their savings together with the labor supplied implies the rest of the economic state (e.g., capital return, wages, incomes, ...). We use the neural network to generate the savings based on the agents' capital holdings at the beginning of the period. The remaining economic state is computed using the equations outlined [above](#model). The economic mechanisms are encoded in helper functions. We create one for the `firm`, the `shocks`, and `wealth` (see cell below).

Then, in the face of future uncertainty, the agents again decide whether and how much to save in risky capital for the next period. We use the neural network to generate new savings for each of the future states (one for each shock). The input state for these network predictions is the next period's capital holding which are current periods savings.

First, we define the helper functions for the `firm`, the `shocks`, and `wealth`.

```python
"""        TO DO: replace by these
          Y = Z*(K^theta);              % output
  W = (1-theta)*Z*(K^theta);    % Wage = F_l
  Div = Y - W - Inv - (1-Pf)*chi*Kss;            % dividends
   b_h = (1-bn_shr)*chi*Kss/mu;
  b_n = bn_shr*chi*Kss/(1-mu);"""

def firm(K, eta, alpha, delta):
    """Calculate return, wage and aggregate production.

    r = eta * K^(alpha-1) * L^(1-alpha) + (1-delta)
    w = eta * K^(alpha) * L^(-alpha)
    Y = eta * K^(alpha) * L^(1-alpha) + (1-delta) * K

    Args:
        K: aggregate capital,
        eta: TFP value,
        alpha: output elasticity,
        delta: depreciation value.

    Returns:
        return: return (marginal product of capital),
        wage: wage (marginal product of labor).
        Y: aggregate production.
    """
    L = tf.ones_like(K)

    r = alpha * eta * K**(alpha - 1) * L**(1 - alpha) + (1 - delta)
    w = (1 - alpha) * eta * K**alpha * L**(-alpha)
    Y = eta * K**alpha * L**(1 - alpha) + (1 - delta) * K

    return r, w, Y

def shocks(z, eta, delta):
    """Calculates tfp and depreciation based on current exogenous shock.

    Args:
        z: current exogenous shock (in {1, 2, 3, 4}),
        eta: tensor of TFP values to sample from,
        delta: tensor of depreciation values to sample from.

    Returns:
        tfp: TFP value of exogenous shock,
        depreciation: depreciation values of exogenous shock.

    """
    tfp = tf.gather(eta, tf.cast(z, tf.int32))
    depreciation = tf.gather(delta, tf.cast(z, tf.int32))
    return tfp, depreciation

def wealth(k, R, l, W):
    """Calculates the wealth of the agents.

    Args:
        k: capital distribution,
        R: matrix of return,
        l: labor distribution,
        W: matrix of wages.

    Returns:
        fin_wealth: financial wealth distribution,
        lab_wealth: labor income distribution,
        tot_income: total income distribution.
    """
    fin_wealth = k * R
    lab_wealth = l * W
    tot_income = tf.add(fin_wealth, lab_wealth)
    return fin_wealth, lab_wealth, tot_income

```

### 3.A. Current period (t)
Using the current state `X` we can calculate the economy. The state is composed of today's shock ($z_t$) and today's capital ($k_t^h$). Note that this constitutes the minimal state. Often, including redundant variables is a simple way to increase the speed of convergence (see the [working paper](https://papers.ssrn.com/sol3/papers.cfm?abstract_id=3393482) for more information on this). However, this notebook attempts to present the most basic example for simplicity.

Using the current state `X`, we predict how much the agents save ($a_t^h$) based on their current capital holdings ($k_t^h$). We do this by passing the state `X` to the neural network.

Then we can calculate the agents' consumptions. To do this, we need to calculate the aggregate capital ($K_t$), the return on capital payed by the firm ($r_t$), and the wages payed by the firm ($w_t$).
$$K_t = \sum_{h=0}^{A-1} k_t^h,$$
$$r_t = \alpha \eta_t K_t^{\alpha - 1} L_t^{1 - \alpha} + (1 - \delta_t),$$
$$w_t = (1 - \alpha) \eta_t K_t^{\alpha} L_t^{-\alpha}.$$

Aggregate labor $L_t$ is always 1 by construction. Then, we can calculate each agent's total income. Finally, we calculate the resulting consumption ($c_t^h$):
$$c_{t}^{h} = r_t k_t^h + w_t l^h_t - a_t^{h}.$$

```python
"""To do: input states Z K bn_shr """

# Today's extended state:
z = X[:, 0]  # exogenous shock
tfp = X[:, 1]  # total factor productivity
depr = X[:, 2]  # depreciation
K = X[:, 3]  # aggregate capital
L = X[:, 4]  # aggregate labor
r = X[:, 5]  # return on capital
w = X[:, 6]  # wage
Y = X[:, 7]  # aggregate production
k = X[:, 8 : 8 + A]  # distribution of capital
fw = X[:, 8 + A : 8 + 2 * A]  # distribution of financial wealth
linc = X[:, 8 + 2 * A : 8 + 3 * A]  # distribution of labor income
inc = X[:, 8 + 3 * A : 8 + 4 * A]   # distribution of total income

""" TO DO:
policy functions c_h c_n Ps Pf Inv bn_shr_next lambdah lambdan;"""
# Today's assets: How much the agents save
# Get today's assets by executing the neural network
a = nn_predict(X)
# The last agent consumes everything they own
a_all = tf.concat([a, tf.zeros([m, 1])], axis=1)

# c_orig: the original consumption predicted by the neural network However, the
#     network can predict negative values before it learns not to. We ensure that
#     the network learns itself out of a bad region by penalizing negative
#     consumption. We ensure that consumption is not negative by including a penalty
#     term on c_orig_prime_1
# c: is the corrected version c_all_orig_prime_1, in which all negative consumption
#     values are set to ~0. If none of the consumption values are negative then
#     c_orig_prime_1 == c_prime_1.
c_orig = inc - a_all
c = tf.maximum(c_orig, tf.ones_like(c_orig) * eps)
```

### 3.B. Next period (t+1)

Now that we have calculated the current economic variables, we simulate one period forward. Since we have 4 possible futures---1 for each of the possible shocks that can realize---we simulate forward for each of the 4 states. Hence, we repeat each calculation 4 times. Below, we first calculate the aggregate variables for the beginning of the period. Then, we calculate the economy for shock $z=1$. The calculations for shock 2, 3, and 4 are analogous and will not be commented on.

From the current period's savings $a_t$ we can calculate the next period's capital holdings $k_{t+1}^h$ and, consequently, aggregate capital $K_{t+1}$. Again, aggregate labor $L_{t+1}$ is 1 by construction.

Then, like above, we calculate economic variables for each shock $z_{t+1} \in \{1, 2, 3, 4\}$. That is, the new state `X'`$=[z_{t+1}, k_{t+1}^h]$ (together with the redundant extensions) is passed to the neural network to generate the agents' savings $a_{t+1}^h$. Then, we use `firm`, the `shocks`, and `wealth` to calculate the return on capital ($r_{t+1}$), the wages ($w_{t+1}$), and the agents' total income. Finally, consumption $c_{t+1}^h$ is calculated.

```python
""" Next period"
  Knext = (1-delta)*K + (a1*((Inv/K)^((xsi-1)/xsi))+a2)*K;
  bh_share_nezt is already defined
"""

# Today's savings become tomorrow's capital holding, but the first agent
# is born without a capital endowment.
k_prime = tf.concat([tf.zeros([m, 1]), a], axis=1)

# Tomorrow's aggregate capital
K_prime_orig = tf.reduce_sum(k_prime, axis=1, keepdims=True)
K_prime = tf.maximum(K_prime_orig, tf.ones_like(K_prime_orig) * eps)

# Tomorrow's labor
l_prime = tf.tile(labor_endow, [m, 1])
L_prime = tf.ones_like(K_prime)

# Shock 1 ---------------------------------------------------------------------
# 1) Get remaining parts of tomorrow's extended state
# Exogenous shock
"""TO DO: construct next period exogenous state Z_prime, the auxiliary variables, the full X next period, and policy functions
Also: figure out if we can use this to construct recursive preferences
construct loop for i in range(n_shocks):"""
z_prime_1 = 0 * tf.ones_like(z)

# TFP and depreciation
tfp_prime_1, depr_prime_1 = shocks(z_prime_1, eta, delta)

# Return on capital, wage and aggregate production
r_prime_1, w_prime_1, Y_prime_1 = firm(K_prime, tfp_prime_1, alpha, depr_prime_1)
R_prime_1 = r_prime_1 * tf.ones([1, A])
W_prime_1 = w_prime_1 * tf.ones([1, A])

# Distribution of financial wealth, labor income, and total income
fw_prime_1, linc_prime_1, inc_prime_1 = wealth(k_prime, R_prime_1, l_prime, W_prime_1)

# Tomorrow's state: Concatenate the parts together
x_prime_1 = tf.concat([tf.expand_dims(z_prime_1, -1),
                       tfp_prime_1,
                       depr_prime_1,
                       K_prime,
                       L_prime,
                       r_prime_1,
                       w_prime_1,
                       Y_prime_1,
                       k_prime,
                       fw_prime_1,
                       linc_prime_1,
                       inc_prime_1], axis=1)

# 2) Get tomorrow's policy
# Tomorrow's capital: capital holding at beginning of period and how much they save
a_prime_1 = nn_predict(x_prime_1)
a_prime_all_1 = tf.concat([a_prime_1, tf.zeros([m, 1])], axis=1)

# 3) Tomorrow's consumption
c_orig_prime_1 = inc_prime_1 - a_prime_all_1
c_prime_1 = tf.maximum(c_orig_prime_1, tf.ones_like(c_orig_prime_1) * eps)
```

#### Repeat for the remaining shocks
Then, we do the same for the remaining shocks ($z = 2,3,4$). Nothing changes in terms of the math for the remaining shocks.

```python
# Shock 2 ---------------------------------------------------------------------
# 1) Get remaining parts of tomorrow's extended state
# Exogenous shock
z_prime_2 = 1 * tf.ones_like(z)

# TFP and depreciation
tfp_prime_2, depr_prime_2 = shocks(z_prime_2, eta, delta)

# return on capital, wage and aggregate production
r_prime_2, w_prime_2, Y_prime_2 = firm(K_prime, tfp_prime_2, alpha, depr_prime_2)
R_prime_2 = r_prime_2 * tf.ones([1, A])
W_prime_2 = w_prime_2 * tf.ones([1, A])

# distribution of financial wealth, labor income, and total income
fw_prime_2, linc_prime_2, inc_prime_2 = wealth(k_prime, R_prime_2, l_prime, W_prime_2)

# Tomorrow's state: Concatenate the parts together
x_prime_2 = tf.concat([tf.expand_dims(z_prime_2, -1),
                       tfp_prime_2,
                       depr_prime_2,
                       K_prime,
                       L_prime,
                       r_prime_2,
                       w_prime_2,
                       Y_prime_2,
                       k_prime,
                       fw_prime_2,
                       linc_prime_2,
                       inc_prime_2], axis=1)

# 2) Get tomorrow's policy
a_prime_2 = nn_predict(x_prime_2)
a_prime_all_2 = tf.concat([a_prime_2, tf.zeros([m, 1])], axis=1)

# 3) Tomorrow's consumption
c_orig_prime_2 = inc_prime_2 - a_prime_all_2
c_prime_2= tf.maximum(c_orig_prime_2, tf.ones_like(c_orig_prime_2) * eps)

# Shock 3 ---------------------------------------------------------------------
# 1) Get remaining parts of tomorrow's extended state
# Exogenous shock
z_prime_3 = 2 * tf.ones_like(z)

# TFP and depreciation
tfp_prime_3, depr_prime_3 = shocks(z_prime_3, eta, delta)

# return on capital, wage and aggregate production
r_prime_3, w_prime_3, Y_prime_3 = firm(K_prime, tfp_prime_3, alpha, depr_prime_3)
R_prime_3 = r_prime_3 * tf.ones([1, A])
W_prime_3 = w_prime_3 * tf.ones([1, A])

# distribution of financial wealth, labor income, and total income
fw_prime_3, linc_prime_3, inc_prime_3 = wealth(k_prime, R_prime_3, l_prime, W_prime_3)

# Tomorrow's state: Concatenate the parts together
x_prime_3 = tf.concat([tf.expand_dims(z_prime_3, -1),
                       tfp_prime_3,
                       depr_prime_3,
                       K_prime,
                       L_prime,
                       r_prime_3,
                       w_prime_3,
                       Y_prime_3,
                       k_prime,
                       fw_prime_3,
                       linc_prime_3,
                       inc_prime_3], axis=1)

# 2) Get tomorrow's policy
# Tomorrow's capital: capital holding at beginning of period and how much they save
a_prime_3 = nn_predict(x_prime_3)
a_prime_all_3 = tf.concat([a_prime_3, tf.zeros([m, 1])], axis=1)

# 3) Tomorrow's consumption
c_orig_prime_3 = inc_prime_3 - a_prime_all_3
c_prime_3 = tf.maximum(c_orig_prime_3, tf.ones_like(c_orig_prime_3) * eps)

# Shock 4 ---------------------------------------------------------------------
# 1) Get remaining parts of tomorrow's extended state
# Exogenous shock
z_prime_4 = 3 * tf.ones_like(z)

# TFP and depreciation
tfp_prime_4, depr_prime_4 = shocks(z_prime_4, eta, delta)

# return on capital, wage and aggregate production
r_prime_4, w_prime_4, Y_prime_4 = firm(K_prime, tfp_prime_4, alpha, depr_prime_4)
R_prime_4 = r_prime_4 * tf.ones([1, A])
W_prime_4 = w_prime_4 * tf.ones([1, A])

# distribution of financial wealth, labor income, and total income
fw_prime_4, linc_prime_4, inc_prime_4 = wealth(k_prime, R_prime_4, l_prime, W_prime_4)

# Tomorrow's state: Concatenate the parts together
x_prime_4 = tf.concat([tf.expand_dims(z_prime_4, -1),
                       tfp_prime_4,
                       depr_prime_4,
                       K_prime,
                       L_prime,
                       r_prime_4,
                       w_prime_4,
                       Y_prime_4,
                       k_prime,
                       fw_prime_4,
                       linc_prime_4,
                       inc_prime_4], axis=1)

# 2) Get tomorrow's policy
# Tomorrow's capital: capital holding at beginning of period and how much they save
a_prime_4 = nn_predict(x_prime_4)
a_prime_all_4 = tf.concat([a_prime_4, tf.zeros([m, 1])], axis=1)

# 3) Tomorrow's consumption
c_orig_prime_4 = inc_prime_4 - a_prime_all_4
c_prime_4 = tf.maximum(c_orig_prime_4, tf.ones_like(c_orig_prime_4) * eps)

```

### 3.C. Cost / Euler function

#### tldr

We train the neural network in an unsupervised fashion by approximating the equilibrium functions. The cost function is composed of the Euler equation, and the punishments for negative consumption and negative aggregate capital.

***

The final key ingredient is the cost function, which encodes the equilibrium functions.

A loss function is needed to train a neural network. In supervised learning, the neural network's prediction is compared to the true value. The weights are updated in the direction that decreases the discripancy between the prediction and the truth. In _deep equilibrium nets_, we approximate all equilibrium functions directly. That is, we update the weights in the direction that minimizes these equilibrium functions. Since the equilibrium functions are approximately 0 in equilibrium, we do not require labeled data and can train in an unsupervised fashion.

Our loss function has 3 components:

  * the Euler equation,
  * the punishments for negative consumption, and
  * the punishments for negative aggregate capital.

The relative errors in the Euler equation is given by:
$$e_{\text{REE}}^i(\mathbf{x}_j) := \frac{u^{\prime -1}\left(\beta \mathbf{E}_{z_j}{r(\hat{\mathbf{x}}_{j,+})u^{\prime}(\hat{c}^{i+1}(\hat{\mathbf{x}}_{j,+}))}\right)}{\hat{c}^i(\mathbf{x}_j)}-1$$

```python
"""Put our equilibrium conditions
  err_bdgt_h = 1 - (W + (Div/mu) + b_h - Pf*(chi*Kss*(1-bn_shr_next)/mu))/c_h; % these are individual consumptions
  err_bdgt_n = 1 - (W + b_n - Pf*(bn_shr_next*chi*Kss/(1-mu)))/c_n;
  foc_stock = 1 - (beta*EEulerstock_future*(Evh_future^((alpha-rhoh)/(1-alpha))))/((c_h^(-rhoh))*Ps);
  foc_bondh = 1 - (beta*EEulerbondh_future*(Evh_future^((alpha-rhoh)/(1-alpha))) + lambdah)/((c_h^(-rhoh))*Pf);
  foc_bondn = 1 - (beta*EEulerbondn_future*(Evn_future^((alpha-rhon)/(1-alpha))) + lambdan)/((c_n^-rhon)*Pf);
  foc_f = 1 - (beta*EEulerf_future*(Evh_future^((alpha-rhoh)/(1-alpha))))/((c_h^(-rhoh))*dIdKp);

  slack_bn = lambdan*(bn_shr_next - bn_shr_lb);    %mun_lw*bn_shr_next;
  slack_bh = lambdah*(bn_shr_ub - bn_shr_next);    %mun_up*(1-bn_shr_next);

  ALSO: try weights for the functions"""

# Prepare transitions to the next periods states. In this setting, there is a 25% chance
# of ending up in any of the 4 states in Z. This has been hardcoded and need to be changed
# to accomodate a different transition matrix.
pi_trans_to1 = p_transition * tf.ones((m, A-1))
pi_trans_to2 = p_transition * tf.ones((m, A-1))
pi_trans_to3 = p_transition * tf.ones((m, A-1))
pi_trans_to4 = p_transition * tf.ones((m, A-1))

# Euler equation
opt_euler = - 1 + (
    (
        (
            beta * (
                pi_trans_to1 * R_prime_1[:, 0:A-1] * c_prime_1[:, 1:A]**(-gamma)
                + pi_trans_to2 * R_prime_2[:, 0:A-1] * c_prime_2[:, 1:A]**(-gamma)
                + pi_trans_to3 * R_prime_3[:, 0:A-1] * c_prime_3[:, 1:A]**(-gamma)
                + pi_trans_to4 * R_prime_4[:, 0:A-1] * c_prime_4[:, 1:A]**(-gamma)
            )
        ) ** (-1. / gamma)
    ) / c[:, 0:A-1]
)

# Punishment for negative consumption (c)
orig_cons = tf.concat([c_orig, c_orig_prime_1, c_orig_prime_2, c_orig_prime_3, c_orig_prime_4], axis=1)
opt_punish_cons = (1./eps) * tf.maximum(-1 * orig_cons, tf.zeros_like(orig_cons))

# Punishment for negative aggregate capital (K)
opt_punish_ktot_prime = (1./eps) * tf.maximum(-K_prime_orig, tf.zeros_like(K_prime_orig))

# Concatenate the 3 equilibrium functions
combined_opt = [opt_euler, opt_punish_cons, opt_punish_ktot_prime]
opt_predict = tf.concat(combined_opt, axis=1)

# Define the "correct" outputs. For all equilibrium functions, the correct outputs is zero.
opt_correct = tf.zeros_like(opt_predict)

```

#### Optimizer

Next, we chose an optimizer; i.e., the algorithm we use to perform gradient descent. We use [Adam](https://arxiv.org/abs/1412.6980), a favorite in deep learning research. Adam uses a parameter specific learning rate and momentum, which encourages gradient descent steps that occur in a consistent direction.

```python
"""all the same"""

# Define the cost function
cost = tf.losses.mean_squared_error(opt_correct, opt_predict)

# Adam optimizer
optimizer = tf.train.AdamOptimizer(learning_rate=lr)

# Clip the gradients to limit the extent of exploding gradients
gvs = optimizer.compute_gradients(cost)
capped_gvs = [(tf.clip_by_value(grad, -1., 1.), var) for grad, var in gvs]

# Define a training step
train_step = optimizer.apply_gradients(capped_gvs)

```

## 4. Training

### tldr

We iterate between simulating a dataset using the current neural network and training on this dataset.

***

In the final stage, we put everything together and train the neural network.

In this section, we iterate between a simulation phase and a training phase. That is, we first simulate a dataset. We simulate a sequence of states ($[z, k]$) based on a random sequence of shocks. The states are computed using the current state of the neural network. We created the `simulate_episodes` helper function to do this. Then, in the training phase, we use the dataset to update our network parameters through multiple epochs. After completion of the training phase, we resimulate a dataset using the new network parameters and repeat.

By computing the error directly after simulating a new dataset, we are able to evaluate our algorithms out-of-sample performance.

First, we define the helper function `simulate_episodes` that simulates the training data used in an episode.

```python
def simulate_episodes(sess, x_start, episode_length, print_flag=True):
    """Simulate an episode for a given starting point using the current
       neural network state.

    Args:
        sess: Current tensorflow session,
        x_start: Starting state to simulate forward from,
        episode_length: Number of steps to simulate forward,
        print_flag: Boolean that determines whether to print simulation stats.

    Returns:
        X_episodes: Tensor of states [z, k] to train on (training set).
    """
    time_start = datetime.now()
    if print_flag:
        print('Start simulating {} periods.'.format(episode_length))
    dim_state = np.shape(x_start)[1]

    X_episodes = np.zeros([episode_length, dim_state])
    X_episodes[0, :] = x_start
    X_old = x_start

    # Generate a sequence of random shocks
    rand_num = np.random.rand(episode_length, 1)

    for t in range(1, episode_length):
        z = int(X_old[0, 0])  # Current period's shock
"""TO DO: change simulation part"""
        # Determine which state we will be in in the next period based on
        # the shock and generate the corresponding state (x_prime)
        if rand_num[t - 1] <= pi_np[z, 0]:
            X_new = sess.run(x_prime_1, feed_dict={X: X_old})
        elif rand_num[t - 1] <= pi_np[z, 0] + pi_np[z, 1]:
            X_new = sess.run(x_prime_2, feed_dict={X: X_old})
        elif rand_num[t - 1] <= pi_np[z, 0] + pi_np[z, 1] + pi_np[z, 2]:
            X_new = sess.run(x_prime_3, feed_dict={X: X_old})
        else:
            X_new = sess.run(x_prime_4, feed_dict={X: X_old})

        # Append it to the dataset
        X_episodes[t, :] = X_new
        X_old = X_new

    time_end = datetime.now()
    time_diff = time_end - time_start
    if print_flag:
        print('Finished simulation. Time for simulation: {}.'.format(time_diff))

    return X_episodes
```

### The true analytic solution
This model can be solved analytically.
Therefore, additionally to the relative errors in the Euler equations, we can compare the solution learned by the neural network directly to the true solution.
The true policy is given by
\begin{align}
\mathbf{a}^{\text{analytic}}_t=
\beta
\begin{bmatrix}
\frac{1-\beta^{A-1}}{1-\beta^{A}} w_t \\
\frac{1-\beta^{A-2}}{1-\beta^{A-1}} r_t k^{1}_t \\
\frac{1-\beta^{A-3}}{1-\beta^{A-2}} r_t k^{2}_t \\
\dots \\
\frac{1-\beta^{1}}{1-\beta^2} r_t k^{A-2}_t \\
\end{bmatrix}.
\end{align}

```python
beta_vec = beta_np * (1 - beta_np ** (A - 1 - np.arange(A-1))) / (1 - beta_np ** (A - np.arange(A-1)))
beta_vec = tf.constant(np.expand_dims(beta_vec, 0), dtype=tf.float32)
a_analytic = inc[:, : -1] * beta_vec
```

### Training the _deep equilibrium net_

Now we can begin training.

```python
"""TO DO: change plots and number of shocks"""

# Helper variables for plotting
all_ages = np.arange(1, A+1)
ages = np.arange(1, A)

# Initialize tensorflow session
sess = tf.Session()

# Initialize interactive plotting
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8, ax9)) = plt.subplots(3, 3, figsize=(18, 18))
plt.ion()
fig.show()
fig.canvas.draw()

# Generate a random starting point
if data_path:
    X_data_train = np.load(data_path)
    print('Loaded initial data from ' + data_path)
    start_episode = int(re.search('_(.*).npy', data_path).group(1))
else:
    X_data_train = np.random.rand(1, num_input_nodes)
    X_data_train[:, 0] = (X_data_train[:, 0] > 0.5)
    X_data_train[:, 1:] = X_data_train[:, 1:] + 0.1
    assert np.min(np.sum(X_data_train[:, 1:], axis=1, keepdims=True) > 0) == True, 'Starting point has negative aggregate capital (K)!'
    print('Calculated a valid starting point')
    start_episode = 0

train_seed = 0

cost_store, mov_ave_cost_store = [], []

time_start = datetime.now()
print('start time: {}'.format(time_start))

# Initialize the random variables (neural network weights)
init = tf.global_variables_initializer()

# Initialize saver to save and load previous sessions
saver = tf.train.Saver()

# Run the initializer
sess.run(init)

if sess_path is not None:
    saver.restore(sess, sess_path)

for episode in range(start_episode, num_episodes):
    # Simulate data: every episode uses a new training dataset generated on the current
    # iteration's neural network parameters.
    X_episodes = simulate_episodes(sess, X_data_train, len_episodes, print_flag=(episode==0))
    X_data_train = X_episodes[-1, :].reshape([1, -1])
    k_dist_mean = np.mean(X_episodes[:, 8 : 8 + A], axis=0)
    k_dist_min = np.min(X_episodes[:, 8 : 8 + A], axis=0)
    k_dist_max = np.max(X_episodes[:, 8 : 8 + A], axis=0)

    ee_error = np.zeros((1, num_agents-1))
    max_ee = np.zeros((1, num_agents-1))

    for epoch in range(epochs_per_episode):
        # Every epoch is one full pass through the dataset. We train multiple passes on
        # one training set before we resimulate a new dataset.
        train_seed += 1
        minibatch_cost = 0

        # Mini-batch the simulated data
        minibatches = random_mini_batches(X_episodes, minibatch_size, train_seed)

        for minibatch_X in minibatches:
            # Run optimization; i.e., determine the cost of each mini-batch.
            minibatch_cost += sess.run(cost, feed_dict={X: minibatch_X}) / num_minibatches
            if epoch == 0:
                # For the first epoch, save the mean and max euler errors for plotting
                # This way, the errors are calculated out-of-sample.
                opt_euler_ = np.abs(sess.run(opt_euler, feed_dict={X: minibatch_X}))
                ee_error += np.mean(opt_euler_, axis=0) / num_minibatches
                mb_max_ee = np.max(opt_euler_, axis=0, keepdims=True)
                max_ee = np.maximum(max_ee, mb_max_ee)

        if epoch == 0:
            # Record the cost and moving average of the cost at the beginning of each
            # episode to track learning progress.
            cost_store.append(minibatch_cost)
            mov_ave_cost_store.append(np.mean(cost_store[-100:]))

        for minibatch_X in minibatches:
            # Take a mini-batch gradient descent training step. That is, update the
            # weights for one mini-batch.
            sess.run(train_step, feed_dict={X: minibatch_X})

    if episode % 20 == 0:
        # Plot
        # Plot the loss function
        ax1.clear()
        line_cost = ax1.plot(np.log10(cost_store), label='Cost')
        line_mov_ave = ax1.plot(np.log10(mov_ave_cost_store), label='Moving average')
        ax1.set_xlabel('Episodes')
        ax1.set_ylabel('Cost [log10]')
        ax1.legend(loc='upper right')

        # Plot the relative errors in the Euler equation
        ax2.clear()
        ee_mean_cost = ax2.plot(ages, np.log10(ee_error).ravel(), 'k-', label='mean')
        ee_max_cost = ax2.plot(ages, np.log10(max_ee).ravel(), '--', label='max')
        ax2.set_xlabel('Age')
        ax2.set_ylabel('Rel EE [log10]')
        ax2.set_xticks(ages)
        ax2.legend()

        # Plot the capital distribution
        ax3.clear()
        k_mean_cost = ax3.plot(all_ages, k_dist_mean, 'k-')
        k_min_cost = ax3.plot(all_ages, k_dist_min, 'k--')
        k_max_cost = ax3.plot(all_ages, k_dist_max, 'k--')
        ax3.set_xlabel('Age')
        ax3.set_ylabel('capital (k)')
        ax3.set_xticks(all_ages)

        # =======================================================================================
        # Sample 50 states and compare the neural network's prediction to the analytical solution
        pick = np.random.randint(len_episodes, size=50)
        random_states = X_episodes[pick, :]

        # Sort the states by the exogenous shock
        random_states_1 = random_states[random_states[:, 0] == 0]
        random_states_2 = random_states[random_states[:, 0] == 1]
        random_states_3 = random_states[random_states[:, 0] == 2]
        random_states_4 = random_states[random_states[:, 0] == 3]

        # Get corresponding capital distribution for plots
        random_k_1 = random_states_1[:, 8 : 8 + A]
        random_k_2 = random_states_2[:, 8 : 8 + A]
        random_k_3 = random_states_3[:, 8 : 8 + A]
        random_k_4 = random_states_4[:, 8 : 8 + A]

        # Generate a prediction using the neural network
        nn_pred_1 = sess.run(a, feed_dict={X: random_states_1})
        nn_pred_2 = sess.run(a, feed_dict={X: random_states_2})
        nn_pred_3 = sess.run(a, feed_dict={X: random_states_3})
        nn_pred_4 = sess.run(a, feed_dict={X: random_states_4})

        # Calculate the analytical solution
        true_pol_1 = sess.run(a_analytic, feed_dict={X: random_states_1})
        true_pol_2 = sess.run(a_analytic, feed_dict={X: random_states_2})
        true_pol_3 = sess.run(a_analytic, feed_dict={X: random_states_3})
        true_pol_4 = sess.run(a_analytic, feed_dict={X: random_states_4})

        ax_list = [ax4, ax5, ax6, ax7, ax8]
        # Plot both
        for i in range(A - 1):
            ax = ax_list[i]

            ax.clear()
            # Plot the true solution with a circle
            ax.plot(random_k_1[:, i], true_pol_1[:, i], 'ro', mfc='none', alpha=0.5, markersize=6, label='analytic')
            ax.plot(random_k_2[:, i], true_pol_2[:, i], 'bo', mfc='none', alpha=0.5, markersize=6)
            ax.plot(random_k_3[:, i], true_pol_3[:, i], 'go', mfc='none', alpha=0.5, markersize=6)
            ax.plot(random_k_4[:, i], true_pol_4[:, i], 'yo', mfc='none', alpha=0.5, markersize=6)
            # Plot the prediction of the neural net
            ax.plot(random_k_1[:, i], nn_pred_1[:, i], 'r*', markersize=2, label='DEQN')
            ax.plot(random_k_2[:, i], nn_pred_2[:, i], 'b*', markersize=2)
            ax.plot(random_k_3[:, i], nn_pred_3[:, i], 'g*', markersize=2)
            ax.plot(random_k_4[:, i], nn_pred_4[:, i], 'y*', markersize=2)
            ax.set_title('Agent {}'.format(i+1))
            ax.set_xlabel(r'$k_t$')
            ax.set_ylabel(r'$a_t$')
            ax.legend()

        ax9.axis('off')
        fig.canvas.draw()
        #========================================================================================

    # Print cost and time log
    print('Episode {}: log10(Cost): {:.4f}; time: {}; time since start: {}'.format(episode,
                                                                                np.log10(cost_store[-1]),
                                                                                datetime.now().time(),
                                                                                datetime.now() - time_start))

    if episode % 100 == 0:
        # Save the tensorflow session
        saver.save(sess, './output/sess_{}.ckpt'.format(episode))
        # Save the starting point
        np.save('./output/data_{}.npy'.format(episode), X_data_train)

```

```python

```