# FINAL PROJECT
## MATH 104C
## Amir Ostrowsky
## June 5th, 2020

# 3.
Acoustic waves. The air pressure $p(x, t)$ in an organ pipe is governed by the wave equation
$$\frac{\partial^2p}{\partial t^2}=c^2\frac{\partial^2p}{\partial x^2}\quad\quad 0<x<l,\quad\quad t>0$$

where $l$ is the length of the pipe and $c$ is a constant. If the pipe is open, the boundary conditions are given by
$$p(0,t)=p_0\quad\quad \textrm{ and } p(l,t)=p_0.$$
If the pupe is closed at the end $x=l$ the boundary conditions are
$$p(0,t)=p_0\quad\quad \textrm{ and } \frac{\partial p}{\partial x}p(l,t)=0.$$
Assume that $c=1, l=1$, and the initial conditions are
$$p(x,0)=p_0\cos(2\pi x)\quad\quad \textrm{ and } \frac{\partial p}{\partial x}p(x,0)=0\quad\quad 0\leq x\leq1.$$

## 3.a.

Write down an explicit finite difference method for $(5)$ and give stability conditions and the order of the method.

Consider the finite difference method to approximate $P_{tt}=c^2P_{xx}$

$$\frac{p_j^{n+1}-2p_j^n+p_j^{n-1}}{k^2} = c^2 \frac{p_{j-1}^n-2p_j^n+p_{j+1}^n}{h^2}\quad\quad (*)$$

with $c=1$ and $\lambda = \frac kh$ the explicit form becomes
$$p_j^{n+1}=2p_j^n-p_j^{n-1}+\frac{k^2}{h^2}(p_{j-1}^n-2p_j^n+p_{j+1}^n)\\
p_j^{n+1}=\lambda^2 p_{j-1}^n+2(p_j^n-\lambda^2 p_j^n)+\lambda^2 p_{j+1}^n-p_j^{n-1}$$

we find the order of the finite difference method by finding the truncation error $\;t_j^n(k,h)$ for $\Delta t = k, \Delta x = h$.

$$\begin{align}
P(t_n,x_{j+1})= P(t_n,x_j+h) &= P(t_n,x_j)+P_x(t_n,x_j)h+P_{xx}(t_n,x_j)\frac{h^2}{2!}+P_{xxx}(t_n,x_j)\frac{h^3}{3!}+\dots&\quad\quad(1)\\
P(t_n,x_{j-1}) = P(t_n,x_j-h) &= P(t_n,x_j)-P_x(t_n,x_j)h+P_{xx}(t_n,x_j)\frac{h^2}{2!}-P_{xxx}(t_n,x_j)\frac{h^3}{3!}+\dots&\quad\quad(2)\\
P(t_{n+1},x_j) = P(t_n+k,x_j) &= P(t_n,x_j)+P_t(t_n,x_j)k+P_{tt}(t_n,x_j)\frac{k^2}{2!}+P_{ttt}(t_n,x_j)\frac{k^3}{3!}+\dots&\quad\quad(3)\\
P(t_{n-1},x_j)=P(t_n-k,x_j) &= P(t_n,x_j)-P_{t}(t_n,x_j)k+P_{tt}(t_n,x_j)\frac{k^2}{2!}-P_{ttt}(t_n,x_j)\frac{k^3}{3!}+\dots&\quad\quad(4)\\
\textrm{then the second central }&\textrm{difference formulas with respect to } h \textrm{ and } k \textrm{ are}& \\
D^2_{h}P(t_n,x_j)=\frac{1}{h^2}((1)-2P(t_n,x_j)+(2)) &= 2P_{xx}(t_n,x_j)\frac{1}{2}+2P_{xxxx}(t_n,x_j)\frac{h^2}{4!}+\dots&\quad\quad(5)\\
D^2_{k}P(t_n,x_j)=\frac{1}{k^2}((1)-2P(t_n,x_j)+(2)) &= 2P_{tt}(t_n,x_j)\frac{1}{2}+2P_{tttt}(t_n,x_j)\frac{k^2}{4!}+\dots&\quad\quad(6)\\
\textrm{then truncation error } t_j^{n}(k,h)&=(6)-(5)\\
&=c_4^*(k^2+h^2)+\mathcal O(k^4)+\mathcal O(h^4) \textrm{ for some real } c_4^*\\
&= \mathcal O(k^2)+\mathcal O(h^2)\\
\end{align}$$
Then the finite difference method $(*)$ is of order $\mathcal O(\Delta t)^2+\mathcal O(\Delta x)^2$.

To find the stability conditions of the method we apply the Von Neumann Analysis
with $P_j^n=g^ne^{ijmh}$ to $(*)$

$$\begin{align}
\frac{1}{k^2}(g^{n+1}e^{ijmh}-2g^ne^{ijmh}+g^{n-1}e^{ijmh})&=\frac{1}{h^2}(g^ne^{i(j-1)mh}-2g^ne^{ijmh}+g^ne^{i(j+1)mh})\\
\textrm{with }\lambda=\frac kh \textrm{ and dividing by for } g^ne^{ijmh}\\
g-2+\frac 1g &=\lambda (e^{-imh}-2+e^{imh})\\
\textrm{ with } e^{imh}+e^{-imh} &= 2\cos(mh)\\
 \textrm{ then with the half angle formula } 2\cos(mh)-2 & = -2(1-\cos(mh))=-2\left(2\left(\sin^2\left({\frac{mh}{2}}\right)\right)\right) \\
\textrm{then }g-2+\frac 1g &=-4\sin^2\left( \frac{mh}{2}\right)\\
\textrm{becomes }\left(\sqrt{g}-\frac {1}{\sqrt{g}}\right)^2 &=\left(\pm2i\lambda \sin\left( \frac{mh}{2}\right)\right)^2\\
\sqrt{g}-\frac {1}{\sqrt{g}} &=\pm2i\lambda \sin\left( \frac{mh}{2}\right) \textrm{ then multiplying by }\sqrt{g} \textrm{ we get}\\
g \pm 2i\lambda\sin\left(\frac{mh}{2}\right)\sqrt{g}-1&=0 \textrm{ which is quadratic in } \sqrt{g}\\
\textrm{then }\sqrt{g}_{\pm} &= \pm i\lambda\sin\left(\frac{mh}{2}\right)\pm\sqrt{1-\lambda^2\sin\left(\frac{mh}{2}\right)}\\
\textrm{then }g_{\pm} &= \left(\pm i\lambda\sin\left(\frac{mh}{2}\right)\pm\sqrt{1-\lambda^2\sin\left(\frac{mh}{2}\right)}\right)^2\\
&=-\lambda^2\sin^2\left(\frac{mh}{2}\right) \pm 2i\lambda\sin\left(\frac{mh}{2}\right)\sqrt{1-\lambda^2\sin\left(\frac{mh}{2}\right)} \\
&+\left(1-\lambda^2\sin\left(\frac{mh}{2}\right)\right)\\
\textrm{then taking the modulus of }&g \textrm{ we get }\\
\left|g_{\pm}\right|&=\lambda^4\sin^2\left(\frac{mh}{2}\right) +1-2\lambda^2\sin^2\left(\frac{mh}{2}\right)+\lambda^4\sin^2\left(\frac{mh}{2}\right)\\
&-2\lambda^2\sin^2\left(\frac{mh}{2}\right)\left( 1-\lambda^2\sin\left(\frac{mh}{2}\right) \right)\\
&+1-2\lambda^2\sin\left(\frac{mh}{2}\right)+\lambda^4\sin^2\left(\frac{mh}{2}\right)\\
&-4\lambda^2\sin^2\left(\frac{mh}{2}\right)\left( 1-\lambda^2\sin^2\left(\frac{mh}{2}\right) \right)\\
&\leq \lambda^4 +1 -2\lambda^2+\lambda^4-2\lambda^2(1-\lambda^2)\\
&+1-2\lambda^2 +\lambda^4-\lambda^4(1-\lambda^2)\\
&=\lambda^6+4\lambda^4-6\lambda^2+2\\
\textrm{then } |g_{\pm}| &\leq 1 \textrm{ if and only if } |\lambda|\leq 1\\
\textrm{for the } g_+=g_- \textrm{ solution, if  }mh&=0,\pi \textrm{ and }|\lambda|=1 \\
\textrm{then } u_j^n \textrm{ has linear growth  } n &\rightarrow\ \infty \textrm{ which grows slower than }g^{n-1} \\
 \textrm{ since the finite difference method}  &\textrm{ is of second order in }t
\end{align}$$

Thus, the the method is stable if and only if $|\lambda| \leq 1$

which is the CFL condition $|a|\frac{\Delta t}{\Delta x}\leq 1$ with $a=c^2$.

## 3.b.
Implement your method given in $a)$ for the open pipe with $p_0 = 0.9$, and with
step sizes $k = h = 0.05$.

Plot your numerical solution at $t = 0.5$ and $t = 1.0$.

We find initial condition for $n=1$. With $\frac{\partial p}{\partial t}(x,0)=0$ from $(9)$ we approximate $p(1,x) \approx P_1(x)$. We use the forward finite difference formula  $P_1(x)=\frac{p_j^1(x_j)-p_j^0(x_j)}{\Delta t}=0$ for $j = 1,\dots, M-1$ with respect to $t$.  Then $p_j^1(x_j)=p_j^0(x_j)$ for $j = 1,\dots, M-1$

```python
'''
WAVE EQUATION OF SECOND ORDER EQUATIONS FINITE DIFFERENCE METHOD waveEq2ndOrdersFD():
INPUT:,pj,pjPlus1,pjnMinus1,lmbda
'pj' is the previous space step's and current x step's approximation node,
'pjMinus1' is the previous time step's and next x step's approximation node,
'pjPlus1' is the next tim
'lmbda' is Courant number
OUTPUT: finite difference scheme forward in time and space

INITIAL VALUE FUNCTION i): u_0ANDu_1()
INPUT:
'x' a value from the domain of function for approximation,
OUTPUT: initial values at initial time step t = 0

GHOST POINT APPROXIMATION OF RIGHT BOUNDARY POINT ghostPointFD():
INPUT:
'pMminus1','pM','lmbda' defined above,
'pMnMinus1' 2 time steps agoth current x step's approximation node,
OUTPUT: an approximation of the right boundary point

OPEN PIPE AIR PRESSURE APPROXIMATION openPipeAirPressure():
'p0' is the initial value of the boundary,
'h' derivative of approximation with respect to space variable 'x',
'k' derivative of approximation with respect to time variable 't',
'l' is the length of space of the pipe,
'T' is the length of time

CONSTANTS, ARRAYS AND VARIABLES:
'N' number of time steps,
'M' number of mesh grid steps,
'delt' cumulative time per iteration
'''

import numpy as np
import math
import matplotlib
import matplotlib.pyplot as plt
import mpl_toolkits
from matplotlib import rc
from IPython.display import display, Math, Latex

# Say, "the default sans-serif font is Palatino Linotype"
matplotlib.rcParams['font.serif'] = "Palatino Linotype"
# Then, "ALWAYS use serif fonts"
matplotlib.rcParams['font.family'] = "DejaVu Sans"
# set the default Latex font for Latex text
plt.rcParams["mathtext.fontset"] = "cm"
# set default font size
matplotlib.rcParams.update({'font.size': 16})
```

```python
def waveEq2ndOrdersFD(pjMinus1,pj,pjPlus1,pjnMinus1,lmbda):
  return( lmbda**2*pjMinus1+2*(1-lmbda**2)*pj+lmbda**2*pjPlus1-pjnMinus1 )

def u_0ANDu_1(x,p0):
  return (p0*np.cos(2*np.pi*x))

p0=0.9
k=0.05;h=k;l=1.0;lmbda=k/h
N=20;M=20;T=k*N

x = np.linspace(0.0,l,M+1) # domain of x

P = np.full([N+1, M+1], p0) # initialize with boundary conditions

# fill first inner nodes initial value with u_0 till the boundary condition
P[0,1:-1] = [u_0ANDu_1(x[j],p0) for j in range(1,M)]; P[1,:] = P[0,:]

for n in range(1,N):
  P[n+1,1:-1] = [waveEq2ndOrdersFD(P[n,j-1],P[n,j],P[n,j+1],P[n-1,j],lmbda) for j in range(1,M)]
```

```python
for t in range(10,N+1,10):
    delt=k*t
    fig = plt.figure()
    ax = fig.add_subplot (111)
    ax.plot(x,P[t,:])
    ax.set_title("$M=%i,\\; T=%f,\\; N=%i, \\; \\lambda = %f,t = %f$" %(M,T,N,lmbda,delt))
    plt.ylabel('$p(x,t)$')
    plt.xlabel('$x$')
plt.show()
```

## 3.c.
Implement your method given in $a)$ for the closed pipe at $x = l$ with $p_0 = 0.9$, and with step sizes $k = h = 0.05$. Plot your numerical solution at $t = 0.5$ and $t = 1.0$.

When the pipe is closed with $x=l\;$ we find the approximation for the ghost point $p_{M+1}^n$ in

$$\frac{p_M^{n+1}-2p_M^n+p_M^{n-1}}{(\Delta t)^2}=\frac{p_{M+1}^n-2p_m^n+p_{M-1}^n}{(\Delta x)^2}\quad\quad\quad (**)$$
with ghost point $x_{M+1}=x_M+\Delta x = 1+h$ outside the domain.

With closed pipe boundary condition for $x=l$ we use Centered Difference formula by $(7.10)$
$$\frac{\partial p}{\partial x}(x_M,t_n) = 0 \approx \frac{p_{M+1}^n-p_{M-1}^n}{2\Delta x}$$
then $p_{M+1}^n=p_{M-1}^n$.

This implies that $(**)$ becomes
$$\frac{p_M^{n+1}-2p_M^n+p_M^{n-1}}{(\Delta t)^2}=2\frac{p_{M-1}^n-p_M^n}{(\Delta x)^2}$$
then the explicit form is

$$p_M^{n+1}=2\lambda^2p_{M-1}^n+2(1-\lambda^2)p_M^n-p_M^{n-1}$$

```python
def ghostPointFD(pMminus1,pM,pMnMinus1,lmbda):
  return( 2*lmbda**2*pMminus1+2*(1-lmbda**2)*pM-pMnMinus1 )

P = np.full([N+1, M+1], p0) # initialize with left boundary condition

# fill first inner nodes initial value with u_0 till the boundary condition
P[0,1:-1] = [u_0ANDu_1(x[j],p0) for j in range(1,M)]; P[1,:] = P[0,:]

for n in range(1,N):
  P[n+1,1:-1] = [waveEq2ndOrdersFD(P[n,j-1],P[n,j],P[n,j+1],
                                   P[n-1,j],lmbda) for j in range(1,M)]
  P[n+1,M] = ghostPointFD(P[n,M-1],P[n,M],P[n-1,M],lmbda)
```

```python
for t in range(10,N+1,10):
    delt=k*t
    fig = plt.figure()
    ax = fig.add_subplot (111)
    ax.plot(x,P[t,:])
    ax.set_title("$M=%i,\\; T=%f,\\; N=%i, \\; \\lambda = %f,t = %f$" %(M,T,N,lmbda,delt))
    plt.ylabel('$p(x,t)$')
    plt.xlabel('$x$')
plt.show()
```

## 3.d.

Repeat $b)$ for $k = h = 0.025$. Construct a higher order approximation by extrapolating your numerical solutions corresponding to $k = h = 0.05$ and $k = h = 0.025$.

We construct an approximation by Richardson Extrapolation $(7.3)$ in the lecture notes

$$ p^{ext}(t_n,x_j) = \frac13 \big(4p(t_{2n},x_{2j}) -p(t_n,x_j)\big)$$ for $n=0,\dots. N$ and $j = 0,\dots,M $.

```python
def waveEq2ndOrdersFD(pjMinus1,pj,pjPlus1,pjnMinus1,lmbda):
  return( lmbda**2*pjMinus1+2*(1-lmbda**2)*pj+lmbda**2*pjPlus1-pjnMinus1 )

def u_0ANDu_1(x,p0):
  return (p0*np.cos(2*np.pi*x))

p0=0.9
k1=0.05;h1=k1;l=1.0;T=1.0
Pext = np.full([N+1, M+1], p0)

def openPipeAirPressure(p0,h,k,l,T):
  N=int(T/k);M=int(l/h);lmbda=k/h

  x = np.linspace(0.0,l,M+1) # domain of x

  P = np.full([N+1, M+1], p0) # initialize with boundary conditions

  # fill first inner nodes initial value with u_0 till the boundary condition
  P[0,1:-1] = [u_0ANDu_1(x[j],p0) for j in range(1,M)]; P[1,:] = P[0,:]

  for n in range(1,N):
    P[n+1,1:-1] = [waveEq2ndOrdersFD(P[n,j-1],P[n,j],
                                     P[n,j+1],P[n-1,j],lmbda) for j in range(1,M)]

  return(P, x)

Pn,   xM = openPipeAirPressure(p0,h1, k1,l,T)
P2n, x2M = openPipeAirPressure(p0,h1/2,k1/2,l,T)
P2n = P2n[::2,::2]

Pext[:,:] = 1/3.0*(4.0*P2n[:,:]-Pn[:,:])
```

```python
for t in range(10,N+1,10):
    delt=k*t
    fig = plt.figure()
    ax = fig.add_subplot (111)
    ax.plot(x,Pext[t,:])
    ax.set_title("$M=%i,\\; T=%f,\\; N=%i, \\; \\lambda = %f,t = %f$" %(M,T,N,lmbda,delt))
    plt.ylabel('$p(x,t)$')
    plt.xlabel('$x$')
plt.show()
```

What's the order of the new approximation?

Let $\tilde P(t_{2n},x_{2j})$ for $j = 1,\dots, M$ and $n = 1,\dots, N$ be the approximation with double the space and time

points of the original approximation $P(t_n,x_j)$ with only the evenly indexed terms. Then asymptotically
$$\begin{align}
\tilde P(t_{2n},x_{2(j+1)})=P\left(t_n,x_j+\frac h2\right) &= P(t_n,x_j)+c_1P_x(t_n,x_j)\frac h2+c_2P_{xx}(t_n,x_j)\frac{\left(\frac h2\right)^2}{2!}+c_3P_{xxx}(t_n,x_j)\frac{\left(\frac h2\right)^3}{3!}+\dots&\quad\quad(7)\\
\tilde P(t_{2n},x_{2(j-1)})=P\left(t_n,x_j-\frac h2\right) &= P(t_n,x_j)-c_1P_x(t_n,x_j)\frac h2+c_2P_{xx}(t_n,x_j)\frac{\left(\frac h2\right)^2}{2!}-c_3P_{xxx}(t_n,x_j)\frac{\left(\frac h2\right)^3}{3!}+\dots&\quad\quad(8)\\
\tilde P(t_{2(n+1)},x_{2j})=P\left(t_n+\frac k2,x_j\right) &= P(t_n,x_j)+c_1P_t(t_n,x_j)\frac k2+c_2P_{tt}(t_n,x_j)\frac{\left(\frac k2\right)^2}{2!}+c_3P_{ttt}(t_n,x_j)\frac{\left(\frac k2\right)^3}{3!}+\dots&\quad\quad(9)\\
\tilde P(t_{2(n-1)},x_{2j})=P\left(t_n-\frac k2,x_j\right) &= P(t_n,x_j)-c_1P_{t}(t_n,x_j)\frac k2+c_2P_{tt}(t_n,x_j)\frac{\left(\frac k2\right)^2}{2!}-c_3P_{ttt}(t_n,x_j)\frac{\left(\frac k2\right)^3}{3!}+\dots&\quad\quad(10)\\
\textrm{then the second central }&\textrm{difference formulas with respect to } \frac h2 \textrm{ and } \frac k2 \textrm{ are}& \\
D^2_{h/2}P(t_n,x_j)=\frac{1}{\left(\frac h2\right)^2}((7)-2P(t_n,x_j)+(8)) &= 2c_2P_{xx}(t_n,x_j)\frac{1}{2}+2c_4P_{xxxx}(t_n,x_j)\frac{\left(\frac h2\right)^2}{4!}+\dots&\quad\quad(11)\\
D^2_{k/2}P(t_n,x_j)=\frac{1}{\left(\frac k2\right)^2}((9)-2P(t_n,x_j)+(10)) &= 2c_2P_{tt}(t_n,x_j)\frac{1}{2}+2c_4P_{tttt}(t_n,x_j)\frac{\left(\frac k2\right)^2}{4!}+\dots&\quad\quad(12)\\
\textrm{ with }c_i \textrm{ being real coefficients} \textrm{ for } i= 2,4,6,\dots\\
\textrm{then the extrapolation term with respect to }h\\
\textrm{ is } D^2_hP(t_n,x_j) &= \frac 13 (4D^2_{h/2}P(t_n,x_j)-D^2_{h}P(t_n,x_j)) = P_{xx}(t_n,x_j)+\tilde c_4h^4 +\mathcal O(h^6)\\
\textrm{ and with respect to } k \textrm{ is }  D^2_kP(t_n,x_j) &= \frac 13 (4D^2_{k/2}P(t_n,x_j)-D^2_{k}P(t_n,x_j)) = P_{tt}(t_n,x_j)+\ddot{\breve c}_4k^4 +\mathcal O(k^6)\\
\textrm{then the truncation error of the }&\textrm{Richardson Extrapolation is }\\
\tau^{ext}(t_n,x_j) = D^2_kP(t_n,x_j)-D^2_hP(t_n,x_j) &= P_{tt}(t_n,x_j)-P_{xx}(t_n,x_j)+\tilde{c}_4(k^4+h^4)+\mathcal O(k^6+h^6) \\
&=\tilde{\ddot c}_4(k^4+h^4)+\mathcal O(k^6+h^6)\\
\textrm{ for real coefficients } \tilde{c}_4, \ddot{\breve c}_4, \tilde{\ddot c}\\
\end{align}$$
Then the Richardson Extrapolation of the finite difference method is of order $\mathcal O(\Delta t)^4+\mathcal O(\Delta x)^4$.


```python
%matplotlib inline
```

```python
import numpy as np
import matplotlib.pyplot as plt
from sympy.physics.hydrogen import R_nl
from sympy import integrate, oo, var
from numerov import radial_integral
from numerov.basis import generate_basis
```

```python
basis = list(generate_basis(range(4, 7)))
print(basis)
```

    [(4.0, 0.0), (4.0, 1.0), (4.0, 2.0), (4.0, 3.0), (5.0, 0.0), (5.0, 1.0), (5.0, 2.0), (5.0, 3.0), (5.0, 4.0), (6.0, 0.0), (6.0, 1.0), (6.0, 2.0), (6.0, 3.0), (6.0, 4.0), (6.0, 5.0)]

```python
%%time
step = 0.0001
mat_numerov = np.zeros((len(basis), len(basis)))
for i, state_1 in enumerate(basis):
    for j, state_2 in enumerate(basis):
        mat_numerov[i, j] = radial_integral(*state_1, *state_2, step=step)
```

    CPU times: user 815 ms, sys: 0 ns, total: 815 ms
    Wall time: 815 ms

```python
fig, ax = plt.subplots()
p = ax.imshow(np.abs(mat_numerov))
plt.colorbar(p)
plt.show()
```

```python
%%time
var("r")
mat_sympy = np.zeros((len(basis), len(basis)))
for i, state_1 in enumerate(basis):
    for j, state_2 in enumerate(basis):
        mat_sympy[i, j] = integrate(R_nl(*state_2, r) * r**3 * R_nl(*state_1, r), (r, 0, oo)).evalf()
```

    CPU times: user 20.5 s, sys: 67.2 ms, total: 20.6 s
    Wall time: 20.6 s

```python
fig, ax = plt.subplots()
p = ax.imshow(np.abs(mat_sympy))
plt.colorbar(p)
plt.show()
```

```python
# diff
fig, ax = plt.subplots()
p = ax.imshow(np.abs(mat_sympy) - np.abs(mat_numerov))
plt.colorbar(p)
plt.show()
```

```python
n1, l1 = 12, 5
n2, l2 = 15, 4
```

```python
integral_sympy = integrate(R_nl(n1, l1, r) * r**3 * R_nl(n2, l2, r), (r, 0, oo)).evalf()
print(integral_sympy)
```

    4.57318723103945

```python
integral_numerov = radial_integral(n1, l1, n2, l2, step=step)
print(integral_numerov)
```

    4.573187231242028

```python
# fractional difference
abs(integral_numerov - integral_sympy) / integral_sympy
```

$\displaystyle 4.42975659646626 \cdot 10^{-11}$

```python
# compare step size
step_values = np.linspace(0.0001, 0.001, 100)
integral_values = np.array([radial_integral(n1, l1, n2, l2, step=st) for st in step_values])

# plot
fig, ax = plt.subplots()
ax.plot(step_values, 100.0 * (integral_values - integral_sympy) / integral_sympy, "x")
ax.axhline(0, c='grey')
ax.set_xlabel("step (a.u.)")
ax.set_ylabel("difference (%)")
plt.show()
```

```python
# sympy
%timeit integrate(R_nl(n1, l1, r) * r**3 * R_nl(n2, l2, r), (r, 0, oo)).evalf()
```

    151 ms ± 22.8 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

```python
# numerov
%timeit radial_integral(n1, l1, n2, l2, step=step)
```

    5.88 ms ± 677 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

```python

```


# PETI DOMAĆI ZADATAK iz predmeta "Primenjena teorija igara" (Applied Game Theory)

Rješenje problema Karnoove tržišne utakmice u slučaju duopola numeričkim putem.

# Potrebni import-i

```
"""
SymPy is a Python library for symbolic mathematics. It aims to become a full-featured computer algebra system (CAS) while keeping the code as simple as possible in order to be comprehensible
and easily extensible. SymPy is written entirely in Python.
"""
from sympy import *
```

# Definisanje potrebnih simbola

```
"""
Transform strings into instances of Symbol class.
symbols() function returns a sequence of symbols with names taken from names argument, which can be a comma or whitespace delimited string, or a sequence of strings.
"""
a1, a2 = symbols('a1 a2')
```

# Definisanje Karnoovog duopola

```
def cournot_competition_duopoly(demand, cost_first_producer, cost_second_producer):
    """
    The cost functions may be the same or different among firms. The market price is set at a level such that demand
    equals the total quantity produced by all firms. Each firm takes the quantity set by its competitors as a given,
    evaluates its residual demand, and then behaves as a monopoly.

    :param demand: total demand
    :param cost_first_producer: cost function for first producer
    :param cost_second_producer: cost function for second producer
    :return: None, prints out information about firms
    """
    utility_1 = demand * a1 - cost_first_producer
    utility_2 = demand * a2 - cost_second_producer

    """
    simplify() attempts to apply all functions it has in an intelligent way to arrive at the simplest form of an
    expression
    """
    u_1 = simplify(utility_1)
    u_2 = simplify(utility_2)

    # To take derivatives, use the diff function
    du1_a1 = diff(u_1, a1)
    du2_a2 = diff(u_2, a2)

    # Solve system to get a1 and a2
    solution = solve((du1_a1, du2_a2), (a1, a2))  # should return dictionary {a1: VALUE_1, a2: VALUE_2}

    """
    Substitution replaces all instances of something in an expression with something else. It is done using the subs
    method. We will use it for assigning numerical values to a1 and a2.
    """
    assign_a1 = u_1.subs(a1, solution[a1])
    assign_a2 = u_2.subs(a2, solution[a2])

    utility_firm_1 = assign_a1.subs(a2, solution[a2])
    utility_firm_2 = assign_a2.subs(a1, solution[a1])

    price_firm_1 = demand.subs(a1, solution[a1])
    price_firm_2 = price_firm_1.subs(a2, solution[a2])

    print("A market-clearing quantities for firms 1 and 2 are: a_1 = {}, a_2 = {}".format(round(solution[a1], 2),
                                                                                round(solution[a2], 2)))
    print("Since the firms are competing by quantity, they set the same price, which is: {}".
          format(round(price_firm_2, 2)))
    print("Firm 1 utility/profit is: {}".format(round(utility_firm_1, 2)))
    print("Firm 2 utility/profit is: {}".format(round(utility_firm_2, 2)))
```

# Glavni program

```
cournot_competition_duopoly(100 - a1 - a2, 1/2 * a1**2, 1/2 * a2**2)
print(83 * "=")

cournot_competition_duopoly(120 - a1 - a2, a1**2, a2**2)
print(83 * "=")

cournot_competition_duopoly(50 - 2 * a1 - 2 * a2, a1 * 2 + 10, a2 * 2 + 10)
print(83 * "=")

cournot_competition_duopoly(50 - 2 * a1 - 2 * a2, a1 * 2 + 10, a2 * 8 + 12)
```

    A market-clearing quantities for firms 1 and 2 are: a_1 = 25.0, a_2 = 25.0
    Since the firms are competing by quantity, they set the same price, which is: 50.0
    Firm 1 utility/profit is: 937.5
    Firm 2 utility/profit is: 937.5
    ===================================================================================
    A market-clearing quantities for firms 1 and 2 are: a_1 = 24.0, a_2 = 24.0
    Since the firms are competing by quantity, they set the same price, which is: 72.0
    Firm 1 utility/profit is: 1152.0
    Firm 2 utility/profit is: 1152.0
    ===================================================================================
    A market-clearing quantities for firms 1 and 2 are: a_1 = 8.0, a_2 = 8.0
    Since the firms are competing by quantity, they set the same price, which is: 18.0
    Firm 1 utility/profit is: 118.0
    Firm 2 utility/profit is: 118.0
    ===================================================================================
    A market-clearing quantities for firms 1 and 2 are: a_1 = 9.0, a_2 = 6.0
    Since the firms are competing by quantity, they set the same price, which is: 20.0
    Firm 1 utility/profit is: 152.0
    Firm 2 utility/profit is: 60.0


```python
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm
import itertools

plt.rcParams["figure.figsize"] = (15, 15)

%matplotlib inline
```

# Modeling the random walk

We assume the following:

1) X and Y are independent

2) Each jump of the fire over given timestep t

Our approach is to treat the spread of fire as a random walk on $x$ and $y$.
For each step of the random walk, we'll treat the velocity as a random variable from a distribution $N(\mu, \sigma^2)$

Our input data is a set of trajectories based off of the given tweets.
Each trajectory gives us a set of jumps associated with the given fire of interest in our region.

```python
numtweets = 30
r = 50
variance = 5
binsize = 4

#x, y, and timestamp of each tweet
testx = np.linspace(0, r, numtweets) + np.random.normal(0,variance,numtweets)
testy = np.linspace(0, r, numtweets) + np.random.normal(0,variance,numtweets)
testt = np.arange(testx.shape[0])

bins = np.arange(min(testt), max(testt), binsize) #make some bins based on binsize
inds = np.digitize(testt, bins) #for each time in testt, assign it to a bin;
#the bin of element i in testt will be the bin for element i in testx and testy also

numbins = len(bins)

x_trajectories = []
y_trajectories = []

for i in range(numbins):
    tempx = []
    tempy = []
    for j in range(len(testt)):
        if inds[j] == i + 1: #since bin 0 is indexed as 1
            tempx += [testx[j]] #if the jth element of testt is in bin i,
            tempy += [testy[j]] #add testx[j] and testy[j] to their appropriate bins
    x_trajectories += [tempx]
    y_trajectories += [tempy]

combinatorial_x_trajectories = list(itertools.product(*x_trajectories))
combinatorial_y_trajectories = list(itertools.product(*y_trajectories))

x_distances = [np.diff(i) for i in combinatorial_x_trajectories]
y_distances = [np.diff(i) for i in combinatorial_y_trajectories]

x_distances = np.reshape(x_distances, -1)
y_distances = np.reshape(y_distances, -1)

vx = x_distances/binsize
vy = y_distances/binsize
```

```python
sns.scatterplot(testx, testy);
```

```python
def MLE_mu(data):
    n = len(data)
    return 1/n * np.sum(data)

def MLE_sigma2(data):
    n = len(data)
    mu_hat = MLE_mu(data)
    return 1/n * np.sum((data - mu_hat)**2)

```

Since the behavior of fire spread is uncertain, we assume each trajectory represents an equally likely path of the the fire.
Based off of this, each link $A\to B$ gives us a "representative" sample of the underlying distribution  of velocities, both of x and y.

Therefore, the approach will to be to calculate the $v_x = d_x/t$ for each link of Tweets per trajectory, then to use these to calculate MLEs for the normal distribution governing the velocities of each hop.

Once we have the normal distribution for these hops, we can use this to predict the probability that the fire has reached some point $A$ by some time $T$.
There are two ways we can do this:

1) Set some timestep $t$ within the range of timesteps that we have seen in the data, and subdivide the desired segment into even chunks such that $\sum t_i = T$ and then add the normals such that the add up to the correct value.
In this case, say we have $n$ chunks,
Then, the probability that a fire has spread to at least to point $A$ is
$$
\begin{align}
P\left(\sum_{i=1}^n x_i \geq A \right) &= P\left(\sum_{i=1}^n tv_i \geq A \right) \\
&= P\left(\sum_{i=1}^n v_i \geq \frac{A}{t} \right)\\
&= P\left(N(n\mu, n\sigma^2)\geq \frac{A}{t} \right)\\
&= P\left(N(n\mu, n\sigma^2)\geq \frac{A}{T/n} \right)\\
&= P\left(N(n\mu, n\sigma^2)\geq \frac{nA}{T} \right)\\
&= P\left(N(\mu, \sigma^2)\geq \frac{A}{T} \right)
\end{align}
$$

2) Find the average velocity required to traverse the whole path in one go and find that probability.
$$
\begin{align}
P(X \geq A) &= P\left(vT \geq A\right)\\
&= P\left(v \geq \frac{A}{T}\right)\\
&= P\left(N(\mu, \sigma^2) \geq \frac{A}{T}\right)\\
\end{align}
$$

Let's apply these ideas below.

First, calculate velocity components based off of the given data:

```python
# vx, vy
```

    (array([ 2.52310822, -1.6544361 ,  1.58150142, ...,  3.94943556,
             0.20712346,  2.1478424 ]),
     array([ 3.51277876,  3.83798799, -0.82903521, ...,  3.69953839,
             1.46260731,  0.83939873]))

```python
# timesteps = np.diff(testt)
# vx = np.diff(testx)/timesteps
# vy = np.diff(testy)/timesteps
```

Get MLE estimates for $v_x$, $v_y$

```python
muhat_y, sigmahat_y = MLE_mu(vy), MLE_sigma2(vy)
muhat_x, sigmahat_x = MLE_mu(vx), MLE_sigma2(vx)
print("means, x and y: " , muhat_x, muhat_y)
print("variances, x and y: ", sigmahat_x, sigmahat_y)
```

    means, x and y:  1.8084442148956452 1.670938735439651
    variances, x and y:  3.3483606084057564 2.7596994513903934

```python
vy_dist = norm(loc = muhat_y,  scale = np.sqrt(sigmahat_y))
vx_dist = norm(loc = muhat_x, scale = np.sqrt(sigmahat_x))
```

```python
predictx = [vx_dist.mean() * t for t in testt]
```

```python
simulated_y_trajectories = []
for i in range(40):
    simulated_position = [0]
    last = 0

    for t in testt[1:]:
        last += vy_dist.rvs()
        simulated_position += [last]

    plt.plot(testt, simulated_position)
    simulated_y_trajectories += [simulated_position]
```

```python
simulated_x_trajectories = []
for i in range(40):
    simulated_position = [0]
    last = 0

    for t in testt[1:]:
        last += vx_dist.rvs()
        simulated_position += [last]

    plt.plot(testt, simulated_position)
    simulated_x_trajectories += [simulated_position]
```

```python
import seaborn as sns
simulated_x_trajectories_plot = np.reshape(simulated_x_trajectories, -1)
simulated_y_trajectories_plot = np.reshape(simulated_y_trajectories, -1)
```

```python
plt.figure(figsize = (10, 10))
sns.kdeplot(simulated_x_trajectories_plot, simulated_y_trajectories_plot, shade = True)
for x_traj, y_traj in zip(simulated_x_trajectories, simulated_y_trajectories):
    plt.plot(x_traj, y_traj)
plt.scatter(testx, testy, label = "test data", color = 'orange')
plt.legend();
```

```python
 def rw_simulation(tweetinputlist, numbins = 5, plotstep = .1, minlat = None, maxlat = None, minlong = None , matlong = None):
  """
  Take in a 2D list, where each element has
  Latitude at index 0
  Longitude at index 1
  Time of tweet at index 2

  Also takes in the number of buckets that the time has been discretized into

  Latitude is y, longitude is x

  plotstep determines the mesh size for which we plot

  Returns:
  2D list, where each element is
  Latitude at index 0
  Longitude at index 1
  Weight on (0,1) at index 2
  """
  ###############
  ### imports ###
  ###############

  import numpy as np
  import matplotlib.pyplot as plt
  from scipy.stats import norm
  import itertools
  import seaborn as sns

  ######################
  ### Helper methods ###
  ######################

  def MLE_mu(data):
      n = len(data)
      return 1/n * np.sum(data)

  def MLE_sigma2(data):
      n = len(data)
      mu_hat = MLE_mu(data)
      return 1/n * np.sum((data - mu_hat)**2)

  #################
  ### Main body ###
  #################

  #unpack the input
  testt = [i[2] for i in tweetinputlist]
  testx = [i[0] for i in tweetinputlist]
  testy = [i[1] for i in tweetinputlist]

  numtweets = 30
  r = 50
  variance = 5
  binsize = 4

  #x, y, and timestamp of each tweet
  testx = np.linspace(0, r, numtweets) + np.random.normal(0,variance,numtweets)
  testy = np.linspace(0, r, numtweets) + np.random.normal(0,variance,numtweets)
  testt = np.arange(testx.shape[0])

  bins = np.arange(min(testt), max(testt), binsize) #make some bins based on binsize
  inds = np.digitize(testt, bins) #for each time in testt, assign it to a bin;
  #the bin of element i in testt will be the bin for element i in testx and testy also

  numbins = len(bins)

  x_trajectories = []
  y_trajectories = []

  for i in range(numbins):
      tempx = []
      tempy = []
      for j in range(len(testt)):
          if inds[j] == i + 1: #since bin 0 is indexed as 1
              tempx += [testx[j]] #if the jth element of testt is in bin i,
              tempy += [testy[j]] #add testx[j] and testy[j] to their appropriate bins
      x_trajectories += [tempx]
      y_trajectories += [tempy]

  combinatorial_x_trajectories = list(itertools.product(*x_trajectories))
  combinatorial_y_trajectories = list(itertools.product(*y_trajectories))

  x_distances = [np.diff(i) for i in combinatorial_x_trajectories]
  y_distances = [np.diff(i) for i in combinatorial_y_trajectories]

  x_distances = np.reshape(x_distances, -1)
  y_distances = np.reshape(y_distances, -1)

  vx = x_distances/binsize
  vy = y_distances/binsize

  muhat_y, sigmahat_y = MLE_mu(vy), MLE_sigma2(vy)
  muhat_x, sigmahat_x = MLE_mu(vx), MLE_sigma2(vx)

  vy_dist = norm(loc = muhat_y,  scale = np.sqrt(sigmahat_y))
  vx_dist = norm(loc = muhat_x, scale = np.sqrt(sigmahat_x))

  simulated_y_trajectories = []
  for i in range(40):
      simulated_position = [0]
      last = 0

      for t in testt[1:]:
          last += vy_dist.rvs()
          simulated_position += [last]

      simulated_y_trajectories += [simulated_position]

  simulated_x_trajectories = []
  for i in range(40):
      simulated_position = [0]
      last = 0

      for t in testt[1:]:
          last += vx_dist.rvs()
          simulated_position += [last]

      simulated_x_trajectories += [simulated_position]

  simulated_x_trajectories_plot = np.reshape(simulated_x_trajectories, -1)
  simulated_y_trajectories_plot = np.reshape(simulated_y_trajectories, -1)

  #################
  ### Viz block ###
  #################

  # plt.figure(figsize = (10, 10))
  # sns.kdeplot(simulated_x_trajectories_plot, simulated_y_trajectories_plot, shade = True)
  # for x_traj, y_traj in zip(simulated_x_trajectories, simulated_y_trajectories):
  #     plt.plot(x_traj, y_traj)
  # plt.scatter(testx, testy, label = "test data", color = 'orange')
  # plt.legend();

  import numpy as np
  import scipy.stats as stats
  from matplotlib.pyplot import imshow

  rvs = np.append(np.array([simulated_x_trajectories_plot]),
                  np.array([simulated_y_trajectories_plot]),
                  axis=1).reshape(-1, 2)

  kde = stats.kde.gaussian_kde(rvs.T)

  # Regular grid to evaluate kde upon
  kde = stats.kde.gaussian_kde(rvs.T)

  # Regular grid to evaluate kde upon
  if minlat == None:
    x_flat = np.arange(rvs[:, 0].min(), rvs[:,0].max(), plotstep)
    y_flat = np.arange(rvs[:, 1].min(), rvs[:,1].max(), plotstep)
  else:
    x_flat = np.arange(minlong, maxlong, .1)
    y_flat = np.arange(minlat, maxlat, .1)

  x,y = np.meshgrid(x_flat,y_flat)
  grid_coords = np.append(x.reshape(-1,1),y.reshape(-1,1),axis=1)

  z = np.array(kde(grid_coords.T))

  return np.array([[x, y, z] for x, y, z in zip(grid_coords[:,0], grid_coords[:,1], z)])

```

```python
simtweets = [list(i) for i in zip(testx, testy, testt)]
```

```python
plotpoints = rw_simulation(simtweets, 5, 1)
```

```python
x, y, z = plotpoints[:, 0], plotpoints[:, 1], plotpoints[:, 2]
```

```python
import pandas as pd
import numpy as np
from scipy.interpolate import griddata
import matplotlib.pyplot as plt
%matplotlib inline

def plot_contour(x,y,z,resolution = 50,contour_method='linear'):
    resolution = str(resolution)+'j'
    X,Y = np.mgrid[min(x):max(x):complex(resolution),   min(y):max(y):complex(resolution)]
    points = [[a,b] for a,b in zip(x,y)]
    Z = griddata(points, z, (X, Y), method=contour_method)
    return X,Y,Z

X,Y,Z = plot_contour(x,y,z,resolution = 50,contour_method='linear')

with plt.style.context("seaborn-white"):
    fig, ax = plt.subplots(figsize=(13,8))
    ax.scatter(x,y, color="black", linewidth=1, edgecolor="ivory", s=50)
    ax.contourf(X,Y,Z)
ax.scatter(testx, testy);
```

# What happens with multiple fires?

Inclusion exclusion:
$$
\begin{align}
P(fire @  A @  t) &= P(fire 1 @  A @  t) + P(fire 2 @  A @  t) - P(fire 1 \cap fire 2 @  A @  t)\\
&= P(fire 1 @  A @  t) + P(fire 2 @  A @  t) - P(fire 1 @  A @  t)P(fire 1 @  A @  t)
\end{align}
$$
for time $t$ and point of interest $A$

```python

```


# Model project: The Meltzer-Richard Model

```python
# Import relevant packages
import numpy as np
import pandas as pd
import re
import math
import scipy as sp
from scipy import linalg, optimize, interpolate
from scipy.stats import gamma, chi2, expon
import sympy as sm
from sympy import sympify
import matplotlib.pyplot as plt

sm.init_printing(use_unicode=True)
```

## Setting up the model

Consider a static economy producing a single commodity. There is a continuum of individuals, differing in one dimension only. As economic agents they  work and consume. As voters they evaluate a simple redistributive program that pays a lump sum to each individual, financed by a proportional income tax.

The **preferences** of the i'th individual are quasi-linear and given by

\\[ w_i = c_i + V(x_i), \\]

where c and x denote consumption and leisure, respectively, and V(x) is a concave utility function. **The private budget constraint** is given by

\\[ c_i \leq (1-\tau)l_i + f, \\]

where $\tau$ is the income tax rate, $l_i$ is individual labor supply and $f$ is a lump sum transfer. The real wage is normalized to 1.

We assume that individual productivity differs and that productivity, in turn, is equivalent to having more "effective time" available. For that reason, individuals are also subject to a **time constraint**:

\\[ 1 + e_i \geq x_i + l_i, \\]

where $e_i$ captures individual productivity - more productive individuals have a larger effective time endowment, $e_i$. We further assume that $e_i$ is distributed according to a known distribution function with mean $e$ and median $e_m < e$

First, we define all symbols and equations that we are going to use:

```python
# Define symbols

w = sm.symbols('w_i')
c = sm.symbols('c_i')
x = sm.symbols('x_i')
tau = sm.symbols('tau')
l = sm.symbols('l_i')
f = sm.symbols('f')
ei = sm.symbols('e_i')
e = sm.symbols('e')
V = sm.symbols('V(x_i)')
```

```python
# Define preferences and the utility function, V(x)
V = sm.ln(x)
preferences = sm.Eq(w, c + V)
preferences
```

```python
# We define the preferences as a sympy-function instead of an equation:
w = c + V
w
```

```python
# Define the privat budget constraint
privat_BC = sm.Eq(c, (1-tau)*l + f)
privat_BC
```

```python
# Define the privat time constraint
privat_timecons = sm.Eq(1 + ei, x + l)
privat_timecons
```

Next, we solve for leisure, $x_i$, and consumption, $c_i$, respectively, and substitute the expression into the utility function, $w_i$.

```python
# Solve for x
x_from_privat_timecons = sm.solve(privat_timecons,x)

# Solve for c
c_from_privat_BC = sm.solve(privat_BC,c)

# Substitute into w
w_subs = w.subs([(x,x_from_privat_timecons[0]), (c,c_from_privat_BC[0])])
w_subs
```

The individual is thus facing the following maximazation problem

\\[ \max_{l_i} w_i = (1-\tau)l_i + f + V(1+e_i-l_i)\\]

By solving the maximazation problem we can find the individual labor supply, $l_i$

```python
# Diffrentiate with respect to labor supply
FOC = sm.diff(w_subs,l)
FOC
```

```python
# Isolating labor supply
laborsupply = sm.solve(FOC,l)
laborsupply
```

Which can also be stated as

\\[ l_i = 1 + e_i - \frac{1}{1-\tau}\\]

This expression shows that labor supply is decreasing in $\tau$. In other words, a higher tax rate distorts the labor-leisure choice and induces the consumer to work less.

We also introduce the average labor supply, which is given by

\\[ L(\tau) = 1 + e - \frac{1}{1-\tau}\\]

This means that we can state individual labor supply as

\\[ l_i = L(\tau) + (e_i - e)\\]

Lastly, we state the government's budget constrain which is given by

\\[ f \leq \tau L(\tau)\\]

We re-write the labor supply and government's budget as a sympy-function

```python
# Reformulating labor supply
L = sm.Function('L')(tau)
laborsupply = L + (ei-e)
laborsupply
```

```python
# Defining the government's budget constraint
government_BC = tau*L
government_BC
```

We now substitute all of this back into $w_i$ to obtain the indirect utility function which is then a function of redistributive policy $\tau$ :

```python
# Substituting laborsupply and the government's budget constraint into w
indirect_utility = w_subs.subs([(l,laborsupply), (f,government_BC)])
indirect_utility = sm.simplify(indirect_utility)
indirect_utility
```

This expression describes individual utility *given* the tax rate. However, as individuals differ they all prefer a different tax rate. By maximazing the indirect utility function with respect to $\tau$, and afterwards isolating $\tau$ we can obtain the preferred tax rate for individual $i$:

```python
# Diffrentiate with respect to labor supply
pref_tax = sm.diff(indirect_utility, tau)
pref_tax
```

Next, we want to insert the expression for the average labor supple, $L(\tau)$, while keeping the diffrentiated average labor as it is. However, we have struggled to do a "proper" substitution and for that reason we have chosen a qiute manual approach.

```python
# Define diffrentiated average labor supply, explicitly
L_diff = sm.diff(L,tau)
```

```python
# Substitute this into the diffrentiated indirect utility function, pref_tax, by defining the equation once again
pref_tax = sm.Eq(0,e-ei+L_diff-L_diff/(e-(1+e-(1/(1-tau)))+1))
pref_tax
```

```python
# Solve for the preferred tax rate for individual i
pref_tax_solve = sm.solve(pref_tax,tau)
pref_tax_solve
```

This shows that our closed form solution can be written as

\\[ \tau_i = \frac{e-e_i}{-\frac{d}{d\tau}L(\tau)} \\]

Recall that $\frac{d}{d\tau}L(\tau) < 0$. This means that an individual with productivity, and hence income, *below* average will prefer a positive tax rate that is larger the poorer that individual is (the larger is $(e-e_i)$. On the other hand, an individual with productivity *above* average will prefer a negative tax rate (income subsidy, $\tau<0$),  financed by a lump sum tax.

To illustrate this relationship we plot the preferred tax rate as a function of individual productivity, $e_i$

```python
# Define the function L
L = (1+e-1/(1-tau))

# Differentiate L and define the closed form solution, pref_tax, as a function
L_diff = sm.diff(L, tau)
pref_tax = (-e+ei)/(L_diff)
pref_tax
```

```python
# Turn the function into a python function
pref_tax_func = sm.lambdify((e,ei,tau), pref_tax)

# create input-value
e = 5
tau = 0.4
ei_list = [0,1,2,3,4,5,6,7,8,9,10]
pref_tax_list = [] # empty list
for x in ei_list:
    pref_tax_list.append(pref_tax_func(e,x,tau))

# plot
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.plot(ei_list,pref_tax_list)
ax.grid()
ax.set_xlabel('Individual productvity, $e_i$')
ax.set_ylabel('$tau$')
ax.set_title('Preferred level of redistribution')
ax.set_xlim([0,10])
ax.set_ylim([-2,2]);
```

This figure clearly illustrates that the preferred tax rate is a negative function of individual productivity and hence income.

As all individuals in this model differs and thus prefers a different tax rate the interesting question is what tax rate will be chosen? To answer this question we assume that two political candidates (who only care about winning the election) are running for office and committing themselves to electoral platforms formulated over the tax rate. Whoever wins the election eneacts the preannounced policy.

Next, we introduce and apply the *median voter theorem*, which in short (and imprecise) terms says that the only relevant person in the electorate is the median voter. The median voter is decisive because (under the assmuption that preferences are monotonic in $e_i$) if that person agrees to a policy then at least half of the electorate will also prefer that policy over an alternative policy which is further away from the median. Hence, if either of the candidates were to announce a policy $\tau'$ that was different from $\tau^m$ (the policy preferred by the median voter) the other candidate could ensure victory by announcing a policy in the interval $(\tau',\tau^m$). Thus, the equilibrium tax rate must coincide with the policy preferred by the median voter:

\\[ \tau_m = \frac{e-e_m}{-\frac{d}{d\tau}L(\tau)} \\]

This result implies that the skewness of the income distribution influences the degree of redistribution in the society.

## Model output

The Meltzer-Richard Model as presented in 1.1 illustrates that the preferred degree of redistribution depends on the income distribution in the society. In specific, the model specifies that the relation between the mean income and the median income defines the preferred degree of redistribution. In the following we will apply the model in a simulation based on the income distribution distributed according to a gamma distribution, a chi2 distribution and an exponential distribution, respectively.

At first, we will illustrate it by dragging a mean and median from a gamma distribution and calculate the preferred redistribution. In the following section we will extend with the chi2 distribution and the exponential distribution.

```python
# Dataframe for draws from distribution
df = pd.DataFrame()

# draws from gamma distribution
size = 200000
df['gamma'] = gamma.rvs(5, size=size)

# print median and mean
tau_ = 0.4
print(' Mean: \n',round(df['gamma'].mean(),3))
print('\n Median: \n',round(df['gamma'].median(),3))

# print results
print('\n \u03C4: \n', round(pref_tax_func(df['gamma'].mean(),df['gamma'].median(),tau_),4))
```

     Mean:
     5.002

     Median:
     4.671

     τ:
     0.1194

Thus, in a simple stylized model society where the mean income is 5 and the median income is 4.67, the preferred tax rate is 0,12.

We will now drag mean and median from the other distributions as well, whereby we can show that the relation between the mean and the median (i.e. the skewness of the income distribution) defines the preferred tax rate according to this model and the median voter theorem.

```python
# draws from chi2 and exponential distribution
df['chi2'] = chi2.rvs(5, size=size)
df['expon'] = expon.rvs(5, size=size)

# print mean and median
print(' Mean: \n Gamma:',round(df['gamma'].mean(),3),
      '\n Chi2:', round(df['chi2'].mean(),3),
      '\n Exponential:', round(df['expon'].mean(),3))

print('\n Median: \n Gamma:',round(df['gamma'].median(),3),
      '\n Chi2:', round(df['chi2'].median(),3),
      '\n Exponential:', round(df['expon'].median(),3))
```

     Mean:
     Gamma: 5.002
     Chi2: 4.998
     Exponential: 5.999

     Median:
     Gamma: 4.671
     Chi2: 4.347
     Exponential: 5.69

The relation between the mean the median differs between the three distributions which we will illustrate in a figure below.

```python
# Add figure
fig = plt.figure(figsize=(8, 6), dpi=95, facecolor='w', edgecolor='k')
ax1 = fig.add_axes([0.1, 0.04, 0.3, 0.4])
ax2 = fig.add_axes([0.5, 0.04, 0.3, 0.4])
ax3 = fig.add_axes([0.9, 0.04, 0.3, 0.4])

# subplot titles
ax1.title.set_text('$\mathrm{Gamma}$')
ax2.title.set_text('$\mathrm{Chi}$' +' '+'$\mathrm{squared}$')
ax3.title.set_text('$\mathrm{Exponential}$')

# plot distriutions
ax1.hist(df['gamma'], color = 'grey', bins = 100)
ax2.hist(df['chi2'], color = 'grey', bins = 100)
ax3.hist(df['expon'], color = 'grey', bins = 100)

# add median and mean lines
ax1.axvline(df['gamma'].mean(), color='blue', linewidth=2)
ax1.axvline(df['gamma'].median(), color='green', linewidth=2)

# Add arrows annotating the means:
x0 = df['gamma'].mean()
ax1.annotate('Mean: {:0.2f}'.format(x0), xy=(x0, 1), xytext=(25, 25),
            xycoords=('data', 'axes fraction'), textcoords='offset points',
            horizontalalignment='left', verticalalignment='center',
            arrowprops=dict(arrowstyle='-|>', fc='black', shrinkA=0, shrinkB=0,
                            connectionstyle='angle,angleA=0,angleB=90,rad=10'),
            )

x0 = df['gamma'].median()
ax1.annotate('Median: {:0.2f}'.format(x0), xy=(x0, 1), xytext=(-25, 25),
            xycoords=('data', 'axes fraction'), textcoords='offset points',
            horizontalalignment='right', verticalalignment='center',
            arrowprops=dict(arrowstyle='-|>', fc='black', shrinkA=0, shrinkB=0,
                            connectionstyle='angle,angleA=0,angleB=90,rad=10'),
            )

# add median and mean lines
ax2.axvline(df['chi2'].mean(), color='blue', linewidth=2)
ax2.axvline(df['chi2'].median(), color='green', linewidth=2)

x0 = df['chi2'].mean()
ax2.annotate('Mean: {:0.2f}'.format(x0), xy=(x0, 1), xytext=(25, 25),
            xycoords=('data', 'axes fraction'), textcoords='offset points',
            horizontalalignment='left', verticalalignment='center',
            arrowprops=dict(arrowstyle='-|>', fc='black', shrinkA=0, shrinkB=0,
                            connectionstyle='angle,angleA=0,angleB=90,rad=10'),
            )

x0 = df['chi2'].median()
ax2.annotate('Median: {:0.2f}'.format(x0), xy=(x0, 1), xytext=(-25, 25),
            xycoords=('data', 'axes fraction'), textcoords='offset points',
            horizontalalignment='right', verticalalignment='center',
            arrowprops=dict(arrowstyle='-|>', fc='black', shrinkA=0, shrinkB=0,
                            connectionstyle='angle,angleA=0,angleB=90,rad=10'),
            )

# add median and mean lines
ax3.axvline(df['expon'].mean(), color='blue', linewidth=2)
ax3.axvline(df['expon'].median(), color='green', linewidth=2)

x0 = df['expon'].mean()
ax3.annotate('Mean: {:0.2f}'.format(x0), xy=(x0, 1), xytext=(25, 25),
            xycoords=('data', 'axes fraction'), textcoords='offset points',
            horizontalalignment='left', verticalalignment='center',
            arrowprops=dict(arrowstyle='-|>', fc='black', shrinkA=0, shrinkB=0,
                            connectionstyle='angle,angleA=0,angleB=90,rad=10'),
            )

x0 = df['expon'].median()
ax3.annotate('Median: {:0.2f}'.format(x0), xy=(x0, 1), xytext=(-25, 25),
            xycoords=('data', 'axes fraction'), textcoords='offset points',
            horizontalalignment='right', verticalalignment='center',
            arrowprops=dict(arrowstyle='-|>', fc='black', shrinkA=0, shrinkB=0,
                            connectionstyle='angle,angleA=0,angleB=90,rad=10'),
            )

plt.show()
```

From the figure, we notes that the chi2 distribution implies the biggest difference between the mean income and the median income, whereas the exponential distribution implies the smallest difference. As presented in 1.1, we will expect the preferred level of redistribution to be greater in the society illustrated by the chi2 distribution. This we will now test by calculating the preferred tax level for the three distributions.

```python
# calculating results
gamma = pref_tax_func(df['gamma'].mean(),df['gamma'].median(),tau_)
chi2 = pref_tax_func(df['chi2'].mean(),df['chi2'].median(),tau_)
expon = pref_tax_func(df['expon'].mean(),df['expon'].median(),tau_)

#printing results
print('\u03C4 \n:'
      'gamma:', round(gamma,3), '\n'
     'chi2:', round(chi2,3), '\n'
     'expon:', round(expon,3))
```

    τ
    :gamma: 0.119
    chi2: 0.234
    expon: 0.111

Thus, we can see that the preferred level of redistribution is greater in the model, where the difference between the mean and the median is bigger.


# Covid-19: From model prediction to model predictive control

## A demo of the deterministic and stochastic, national-level, age-stratified modeling framework

*Original code by Ryan S. McGee. Modified by T.W. Alleman in consultation with the BIOMATH research unit headed by prof. Ingmar Nopens.*

Copyright (c) 2020 by T.W. Alleman, BIOMATH, Ghent University. All Rights Reserved.

In `src/covid19model/models.py`, two national-level extended SEIRD compartmental models are coded,
1. COVID19_SEIRD: a deterministic model, using ordinary differential equations (ODEs). The equations are integrated using Scipy's `odeint` with a variable timestep, it is a 'continuous-time' model.
2. COVID19_SEIRD_sto: a stochastic model, using stochastic difference equations (SDEs). The equations are iterated using a fixed timestep of one day and is said to be a 'discrete-time' model.

The deterministic model was used for all results detailed in our first COVID-19 related preprint: *A deterministic, age-stratified, extended SEIRD model for assessing the effect of non-pharmaceutical interventions on SARS-CoV-2 spread in Belgium*, which can be found here: https://www.medrxiv.org/content/10.1101/2020.07.17.20156034v2

The documentation given in this tutorial has been redrafted from the above preprint. References have been omitted for the sake of simplicity.

```python
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import pandas as pd
import datetime
import scipy
import json
import random
```

Load the covid 19 custom development code

```python
from covid19model.optimization import objective_fcns
from covid19model.models import models
from covid19model.data import google, sciensano, polymod, model_parameters
from covid19model.visualization.output import population_status, infected
from covid19model.visualization.optimization import plot_fit, traceplot
```

```python
# OPTIONAL: Load the "autoreload" extension so that package code can change
%load_ext autoreload
# OPTIONAL: always reload modules so that as you change code in src, it gets loaded
%autoreload 2
```

## Introduction

### Model dynamics

The SEIR(D) model was first proposed in 1929 by two Scottish scientists. It is a compartmental model that subdivides the human population into four groups: 1) susceptible individuals, 2) exposed individuals in the latent phase 1, 3) infectious individuals capable of transmitting the disease and 4) individuals removed from the population either through immunization or death. Despite being a simple and idealized reality, the SEIR(D) dynamics are used extensively to predict the outbreak of infectious diseases and this was no different during the SARS-CoV-2 outbreak earlier this year.

In this work, we extend the SEIRD model to incorporate more expert knowledge on SARS-CoV-2. For that purpose, the infectious compartment is split into four parts. The first is a period of presymptomatic infectiousness because several studies have shown that pre-symptomatic transmission is a dominant transmission mechanism of SARS-CoV-2. After the period of pre-symptomatic transmission, three possible infectious outcomes are modeled: (1) Asymptomatic outcome, for individuals who show no symptoms at all, (2) Mild outcome, for individuals with mild symptoms who recover at home, and (3) Hospitalization, when mild symptoms worsen. Children and young adults have a high propensity to experience an asymptomatic or mild outcome, while older individual have a high propensity to be hospitalized.

In general, Belgian hospitals have two wards for COVID-19 patients: 1) Cohort, where patients are not monitored permanently and 2) Intensive care, for patients with the most severe symptoms. Intensive care includes permanent monitoring, the use of ventilators or the use of extracorporeal membrane oxygenation (ECMO). Patients generally spend limited time in the emergency room and/or in a buffer ward before going to Cohort. After spending limited time in Cohort, some patients are transferred to ICU. Patients can perish in both wards, but mortalities are generally lower in Cohort. After a stay in an ICU, patients return to Cohort for recovery in the hospital. During the recovery stay, mortality is limited. The above is a short summary of hospital dynamics based on interviewing Ghent University hospital staff and examining the hospital data.

### Model framework and equations

#### Age-stratification

We introduced heterogeneity in the deterministic implementation by means of age-stratification. Every population compartment is split into a number of age classes, the age-groups have different contact rates with other age-groups and the disease progresses differently for each age-group, making the model behaviour more realistic. Our age-stratified model consists of 9 age classes, i.e., [0-10[, [10-20[, [20-30[, [30-40[, [40-50[, [50-60[, [60-70[, [70-80[, [80- $\infty$[. The age-stratified implementation provides a good balance between added complexity and computational resources.

#### Deterministic framework

Our extended SEIRD model is implemented using two frameworks: a deterministic and a stochastic framework. The deterministic equations are obtained by writing down the following equation,

`rate of change = in - out`,

for every of the 11 population compartments. This results in the following system of coupled ordinary differential equations,
\begin{eqnarray}
\dot{S_i} &=& - \beta s_i \sum_{j=1}^{N} N_{c,ij} S_j \Big( \frac{I_j+A_j}{T_j} \Big) + \zeta R_i, \label{eq:system1} \\
\dot{E_i} &=& \beta s_i \sum_{j=1}^{N} N_{c,ij} S_j \Big( \frac{I_j+A_j}{T_j} \Big) - (1/\sigma) \cdot E_i, \label{eq:system2} \\
\dot{I_i} &=& (1/\sigma) E_i - (1/\omega) I_i, \label{eq:system3}\\
\dot{A_i} &=& (\text{a}_i/\omega) I_i - (1/d_{\text{a}}) A_i, \label{eq:system4}\\
\dot{M_i} &=&  ((1-\text{a}_i) / \omega ) I_i - ( (1-h_i)/d_m + h_i/d_{\text{hospital}} ) M_i,\label{eq:system5} \\
\dot{ER_i} &=& (h_i/d_{\text{hospital}}) M_i - (1/d_{\text{ER}}) ER_i, \label{eq:system6} \\
\dot{C_i} &=& c_i (1/d_{\text{ER}}) ER_i  - (m_{C, i}/d_{c,D}) C_i - ((1 - m_{C, i})/d_{c,R}) C_i, \label{eq:system7} \\
\dot{ICU_i} &=& (1-c_i) (1/d_{\text{ER}}) ER_i - (m_{ICU,i}/d_{\text{ICU},D}) ICU_i  \label{eq:system8} \\
&& - ((1-m_{ICU,i})/d_{\text{ICU},R}) ICU_i, \label{eq:system9}\\
\dot{C}_{\text{ICU,rec,i}} &=& ((1-m_{ICU,i})/d_{\text{ICU},R}) ICU_i - (1/d_{\text{ICU,rec}}) C_{\text{ICU,rec,i}}, \label{eq:system10}\\
\dot{D_i} &=&  (m_{ICU,i}/d_{\text{ICU},D}) ICU_i +  (m_{C,i}/d_{\text{c},D}) C_i , \label{eq:system11}\\
\dot{R_i} &=&  (1/d_a) A_i + ((1-h_i)/d_m) M_i + ((1-m_{C,i})/d_{c,R}) C_i \\
&& + (1/d_{\text{ICU,rec}}) C_{\text{ICU,rec,i}} - \zeta R_i \label{eq:system13},
\end{eqnarray}

for $i = 1,2,...,9$. Here, $T_i$ stands for total population, $S_i$ stands for susceptible, $E_i$ for exposed, $I_i$ for pre-symptomatic and infectious, $A_i$ for asymptomatic and infectious, $M_i$ for mildly symptomatic and infectious, $ER_i$ for emergency room and/or buffer ward, $C_i$ for cohort, $C_{\text{ICU,rec,i}}$ for a recovery stay in Cohort coming from Intensive Care, $ICU_i$ for Intensive Care Unit, $D_i$ for dead and $R_i$ for recovered. Using the above notation, all model states are 9x1 vectors,
$$ \mathbf{S} = [S_1(t)\ S_2(t)\ ...\ S_i(t)]^T, $$
where $S_i(t)$ denotes the number of susceptibles in age-class i at time t after the introduction of the virus in the population.

These equations are implemented in the function `COVID19_SEIRD` located in `src/covid19model/models.py`. The integration is performed in `_sim_single` located in `src/covid19model/base.py` by using Scipy's `solve_ivp`. The integrator timestep depends on the rate of change of the system and the solver method is thus referred to as a 'continuous-time' solver. The implementation uses non-integer individuals.

#### Stochastic framework

By defining the probabilities of transitioning (propensities) from one state to another, a system of coupled stochastic difference equations (SDEs) can be obtained. The probability to transition from one state to another is assumed to be exponentially distributed. As an example, consider the average time a patient spends in an ICU when recovering, which is $d_{\text{ICU,R}} = 9.9$ days. The chances of survival in ICU are $(1-m_{\text{ICU,i}})$, where $m_{\text{ICU,i}}$ is the mortality in ICU for an individual in age group $i$. The probability of transitioning from state ICU to state $C_{\text{ICU,rec}}$ on any given day and for an individual in age group $i$ is,

\begin{equation}
P(ICU_i \rightarrow C_{\text{ICU,rec,i}}) = 1 - \text{exp} \Bigg[ - \frac{1-m_{\text{ICU},i}}{d_{\text{ICU,R}}}\Bigg].
\end{equation}

If a transitioning between states is defined as "succes", we can regard the number of individuals transitioning from ICU to a Cohort recovery ward as a binomial experiment. On a given day, the number of individuals transitioning is,

\begin{equation}
(\text{ICU}_i \rightarrow C_{\text{ICU,rec,i}})(k) \sim \text{Binomial}\Bigg(\text{ICU}_i(k),\ 1 - \text{exp}\Bigg[- \frac{1-m_{\text{ICU,i}}}{d_{\text{ICU,R}}}\Bigg]\Bigg).
\end{equation}

For a discrete stepsize $l$, there are 15 possible transitions,
\begin{eqnarray}
(S_i \rightarrow E_i) (k) &\sim& \text{Binomial}\Bigg(S_i(k), 1 - \text{exp}\Bigg[- l \beta s_i \sum_{j=1}^{N} N_{c,ij} S_j \Big( \frac{I_j+A_j}{T_j} \Big) \Bigg]\Bigg)\\
(E_i \rightarrow I_i) (k) &\sim& \text{Binomial}\Bigg(E_i(k), 1 - \text{exp}\Bigg[- l\ \frac{1}{\sigma}\Bigg]\Bigg)\\
(I_i \rightarrow A_i) (k) &\sim& \text{Binomial}\Bigg(I_i(k), 1 - \text{exp}\Bigg[- l\ \frac{a_i}{\omega}\Bigg]\Bigg)\\
(I_i \rightarrow M_i) (k) &\sim& \text{Binomial}\Bigg(I_i(k), 1 - \text{exp}\Bigg[- l\ \frac{1-a_i}{\omega}\Bigg]\Bigg)\\
(A_i \rightarrow R_i) (k) &\sim& \text{Binomial}\Bigg(A_i(k), 1 - \text{exp}\Bigg[- l\ \frac{1}{d_a}\Bigg]\Bigg)\\
(M_i \rightarrow R_i) (k) &\sim& \text{Binomial}\Bigg(M_i(k), 1 - \text{exp}\Bigg[- l\ \frac{1-h_i}{d_m}\Bigg]\Bigg)\\
(M_i \rightarrow ER_i) (k) &\sim& \text{Binomial}\Bigg(M_i(k), 1 - \text{exp}\Bigg[- l\ \frac{h_i}{d_{\text{hospital}}}\Bigg]\Bigg)\\
(ER_i \rightarrow C_i) (k) &\sim& \text{Binomial}\Bigg(ER_i(k), 1 - \text{exp}\Bigg[- l\ \frac{c_i}{d_{\text{ER}}}\Bigg]\Bigg)\\
(ER_i \rightarrow ICU_i) (k) &\sim& \text{Binomial}\Bigg(ER_i(k), 1 - \text{exp}\Bigg[- l\ \frac{1-c_i}{d_{\text{ER}}}\Bigg]\Bigg)\\
(C_i \rightarrow R_i) (k) &\sim& \text{Binomial}\Bigg(C_i(k), 1 - \text{exp}\Bigg[- l\ \frac{1-m_{C,i}}{d_{C,R}}\Bigg]\Bigg)\\
(ICU_i \rightarrow C_{\text{ICU,rec,i}}) (k) &\sim& \text{Binomial}\Bigg(ICU_i(k), 1 - \text{exp}\Bigg[- l\ \frac{1-m_{\text{ICU,i}}}{d_{ICU,R}}\Bigg]\Bigg)\\
(C_{\text{ICU,rec,i}} \rightarrow R_i) (k) &\sim& \text{Binomial}\Bigg(C_{\text{ICU,rec,i}}(k), 1 - \text{exp}\Bigg[- l\ \frac{1}{d_{\text{ICU,rec}}}\Bigg]\Bigg)\\
(C_i \rightarrow D_i) (k) &\sim& \text{Binomial}\Bigg(C_i(k), 1 - \text{exp}\Bigg[- l\ \frac{m_{C,i}}{d_{C,D}}\Bigg]\Bigg)\\
(ICU_i \rightarrow D_i) (k) &\sim& \text{Binomial}\Bigg(ICU_i(k), 1 - \text{exp}\Bigg[- l\ \frac{m_{\text{ICU,i}}}{d_{\text{ICU,D}}}\Bigg]\Bigg)\\
(R_i \rightarrow S_i) (k) &\sim& \text{Binomial}\Bigg(R_i(k), 1 - \text{exp}\Bigg[- l\ \zeta \Bigg]\Bigg)\\
\end{eqnarray}

And the system of equations becomes,

\begin{eqnarray}
S_i(k+1) &=& S_i(k) + (R_i \rightarrow S_i) (k) - (S_i \rightarrow E_i) (k) \\
E_i(k+1) &=& E_i(k) + (S_i \rightarrow E_i) (k) - (E_i \rightarrow I_i) (k) \\
I_i(k+1) &=& I_i(k) + (E_i \rightarrow I_i) (k) - (I_i \rightarrow A_i) - (I_i \rightarrow M_i) (k) \\
A_i(k+1) &=& A_i(k) + (I_i \rightarrow A_i) (k) - (A_i \rightarrow R_i) (k) \\
M_i(k+1) &=& M_i(k) + (I_i \rightarrow M_i) (k) - (M_i \rightarrow R_i) (k) - (M_i \rightarrow ER_i) (k) \\
ER_i(k+1) &=& ER_i(k) + (M_i \rightarrow ER_i) (k) - (ER_i \rightarrow C_i) (k) - (ER_i \rightarrow ICU_i) (k) \\
C_i(k+1) &=& C_i(k) + (ER_i \rightarrow C_i) (k) - (C_i \rightarrow R_i) (k) - (C_i \rightarrow D_i) (k) \\
C_{\text{ICU,rec,i}}(k+1) &=& C_{\text{ICU,rec,i}}(k)  + (ICU_i \rightarrow C_{\text{ICU,rec,i}}) (k) - (C_{\text{ICU,rec,i}} \rightarrow R_i) (k) \\
R_i(k+1) &=& R_i(k) + (A_i \rightarrow R_i) (k)  + (M_i \rightarrow R_i) (k) + (C_i \rightarrow R_i) (k) + (C_{\text{ICU,rec,i}} \rightarrow R_i) (k)  - (R_i \rightarrow S_i) (k) \\
D_i(k+1) &=& D_i(k) + (ICU_i \rightarrow D_i) (k) + (C_i \rightarrow D_i) (k) \\
\end{eqnarray}

These equations are implemented in the function `COVID19_SEIRD_sto` located in `src/covid19model/models.py`. The computation itself is performed in the function `solve_discrete` located in `src/covid19model/base.py`. Please note that the deterministic model uses **differentials** in the model defenition and must be integrated, while the stochastic model uses **differences** and must be iterated. The discrete timestep is fixed at one day. The stochastic implementation only uses integer individuals, which is considered an advantage over the deterministic implementation.

### Transmission rate and social contact data

In our model, the transmission rate of the disease depends on the product of four contributions. The first contribution, $(I+A)/T$, is the fraction of contagious individuals in the population. The second contribution, $\mathbf{N}_c$, is the average number of human-to-human interactions per day. The third contribution, $s_i$, is the relative susceptiblity to SARS-CoV-2 infection in age group $i$, and the fourth contribution, $\beta$, is the probability of contracting COVID-19 when encountering a contagious individual under the assumption of 100 \% susceptibility to SARS-CoV-2 infection. We assume that the per contact transmission probability $\beta$ is independent of age and we will infer its distribution by calibrating the model to national Belgian hospitalization data. The number of human-human interactions, $\mathbf{N}_c$, are both place and age-dependent. These matrices assume the form of a 9x9 *interaction matrix* where an entry X, Y denotes the number of social contacts age group X has with age group Y per day. These matrices are available for homes, schools, workplaces, in public transport, and leisure activities, from a survey study by Lander Willem (2012). The total number of social interactions is given by the sum of the contributions in different places,

\begin{equation}\label{eq:interaction_matrices}
\mathbf{N_{\text{c}}} = \mathbf{N_{\text{c, home}}} + \mathbf{N_{\text{c, schools}}} + \mathbf{N_{\text{c, work}}} + \mathbf{N_{\text{c, transport}}} + \mathbf{N_{\text{c, leisure}}} + \mathbf{N_{\text{c, others}}}
\end{equation}

https://lwillem.shinyapps.io/socrates_rshiny/

```python
# Load the interaction matrices (size: 9x9)
initN, Nc_home, Nc_work, Nc_schools, Nc_transport, Nc_leisure, Nc_others, Nc_total = polymod.get_interaction_matrices()
# Define the number of age categories
levels = initN.size
```

The coefficients of the contact matrices in different places can be altered to model a goverment policy. For instance, to model the Belgian lockdown, the mobility reductions deduced from the Google community mobility reports were used as coefficients for the different interaction matrices. We assumed workplace interactions were down to only 40 % of their prepandemic values before the lockdown.

https://www.google.com/covid19/mobility/

Our code has built-in functionality to download an visualize these mobility reports. To do so, simply run the following code:

```python
df_google = google.get_google_mobility_data(update=False, plot=True)
```

### Social interatia

The model takes into account the effect of *social inertia* when measures are taken. In reality, social restrictions or relaxations represent a change in behaviour which is gradual and cannot be modeled using a step-wise change of the social interaction matrix $\mathbf{N_c}$. This can be seen when closely inspecting the *Google community mobility report* above. Multiple functions can be used to model the effects of social compliance, e.g. a delayed or non-delayed ramp, or a logistic function. In our model, we use a delayed ramp to model compliance,

\begin{equation}
\mathbf{N_{c}}^{k} = \mathbf{N_{\text{c, old}}} + f^{k} (\mathbf{N_{\text{c, new}}} - \mathbf{N_{\text{c, old}}})
\end{equation}

where,
\begin{equation}
    f^k=
\begin{cases}
	0.0,& \text{if } k\leq \tau\\
    \frac{k}{l} - \frac{\tau}{l},& \text{if } \tau < k\leq \tau + l\\
    1.0,              & \text{otherwise}
\end{cases}
\end{equation}

where $\tau$ is the number of days before measures start having an effect and $l$ is the number of additional days after the time delay until full compliance is reached. Both parameters were calibrated to the daily number of hospitalizations in Belgium (notebooks `notebooks/0.1-twallema-calibration-deterministic.ipynb` and `notebooks/0.1-twallema-calibration-stochastic.ipynb`). $k$ denotes the number of days since a change in social policy.

```python
# Load the compliance model
from covid19model.models.compliance import ramp_2
```

### Basic reproduction number

The basic reproduction number $R_0$, defined as the expected number of secondary cases directly generated by one case in a population where all individuals are susceptible to infection, is computed using the next generation matrix (NGM) approach introducted by Diekmann. For our model, the basic reproduction number of age group $i$ is,
\begin{equation}\label{eq:reproduction_number}
R_{0,i} = (a_i d_a + \omega) \beta s_i \sum_{j=1}^{N} N_{c,ij}
\end{equation}
and the population basic reproduction number is calculated as the weighted average over all age groups using the demographics of Belgium. The detailed algebra underlying the computation equation of the basic reproduction number is presented in the supplementary materials of our manuscript (https://www.medrxiv.org/content/10.1101/2020.07.17.20156034v2).

### Model parameters

## Simulations with the deterministic model

### Model initialization

```python
# Load the model parameters using `get_COVID19_SEIRD_parameters()`.
params = model_parameters.get_COVID19_SEIRD_parameters()
# Add the delayed ramp parameters to the parameter dictionary to use compliance.
params.update({'l': 1,
              'tau': 5})
# Define the initial condition: one exposed inidividual in every age category
initial_states = {'S': initN, 'E': np.ones(levels)}
# Initialize the model
model = models.COVID19_SEIRD(initial_states, params, compliance=ramp_2)
```

### Load posterior parameter distributions of calibrated parameters

Both the deterministic and stochastic model were calibrated to the national hospitalization data provided by the Scientific Institute of Public Health (Sciensano).

The parameters $\beta$, $l$, $\tau$ and $\Omega$ were calibrated.

The aim of such a calibration is to obtain a parameter set which leads to a good agreement between the model predictions and the experimental data. Five timeseries starting on Sunday, March 15th, 2020 have been made publically available by the Belgian Scientific Institute of Public Health (Sciensano),

- the total number of patients in the hospitals ($H_{\text{tot}}$)
- the total number of patients in ICU ($ICU_{\text{tot}}$)
- the daily hospitalizations ($H_{\text{in}}$)
- the daily hospital discharges ($H_{\text{out}}$)
- the daily number of deaths ($D_{\text{new}}$)

We assume the timeseries are independent and identically distributed (i.d.d.) sequences of normal random variables having a mean $\hat{y}(k)$ and an unknown variance $\hat{\sigma}^2$. The resulting log-likelihood function is,

\begin{equation}
\ln \big( P( y \mid \mathbf{\hat{\theta}}) \big) = - \frac{1}{2} \sum_{k=1}^{N} \Bigg[ \frac{\big(y(k) - \hat{y}(k, \mathbf{\hat{\theta}})\big)^2}{\hat{\sigma}^2} + \ln\big(2 \pi \hat{\sigma}^2\big) \Bigg]\ ,
\end{equation}

where $\mathbf{\hat{\theta}}$ and $\hat{\sigma}^2$ are computed to maximize the log-likelihood function. The fitting procedure is performed in two steps. First, the transmission parameter $\beta$ is calibrated using the daily hospitalizations from March 15th, 2020 until March 22nd, 2020. These data are used because they do not depend on the durations of the hospitalizations, $d_c$ and $d_{\text{ICU}}$, which are computed from the limited and imperfect hospital dataset. In addition to the calibration of the transmission parameter $\beta$, the number of days between the first infectious individuals entering the country and the onset of data collection (March 15th, 2020) must be estimated. The latter is denoted as excess time, $t_e$. Second, the daily hospitalizations ($H_{\text{in}}$) from March 15th, 2020 until April 20th, 2020 were used to infer the slope of the compliance ramp ($l$), as well as the effect of prevention measures and *network* effects in the workplace and on public transport ($\Omega$).

We performed the above optimizations in two steps. First, we used a particle swarm optimization (PSO) algorithm to scan the parameter space and find the global maximum likelihood. We then used the obtained maximum-likelihood estimates as initial values for the ensemble sampler for Markov Chain Monte Carlo (MCMC) proposed by Goodman 2010. For all parameters, uniform prior distributions were used. The sampler was always run on $2n$ parallel chains, were $n$ was the number of parameters to be sampled from the posterior distributions. For each chain, we drew at least 50-fold more samples than the autocorrelation time. Next, we visualised the samples using the scatterplot matrix proposed by Foreman (2016), further referred to as *cornerplot*. In these visualizations, each one- and two-dimensional projection of the sample is plotted to reveal correlations. During each step of the optimization procedure, previously sampled parameters were propagated.

The calibration can be found in the notebook `notebooks/0.1-twallema-calibration-deterministic.ipynb`. The posterior samples of the calibrated parameters can be found in `data/interim/model_parameters/deterministic_22072020.json`.

```python
# Scrape the publically available data from the Scientific institute of Public Health (Sciensano)
df_sciensano = sciensano.get_sciensano_COVID19_data(update=False)
df_sciensano.tail()
```

<div>
<style scoped>
    .dataframe tbody tr th:only-of-type {
        vertical-align: middle;
    }

    .dataframe tbody tr th {
        vertical-align: top;
    }

    .dataframe thead th {
        text-align: right;
    }
</style>
<table border="1" class="dataframe">
  <thead>
    <tr style="text-align: right;">
      <th></th>
      <th>H_tot</th>
      <th>ICU_tot</th>
      <th>H_in</th>
      <th>H_out</th>
      <th>H_tot_cumsum</th>
      <th>D_tot</th>
      <th>D_25_44</th>
      <th>D_45_64</th>
      <th>D_65_74</th>
      <th>D_75_84</th>
      <th>D_85+</th>
    </tr>
    <tr>
      <th>DATE</th>
      <th></th>
      <th></th>
      <th></th>
      <th></th>
      <th></th>
      <th></th>
      <th></th>
      <th></th>
      <th></th>
      <th></th>
      <th></th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <th>2020-08-06</th>
      <td>278</td>
      <td>65</td>
      <td>28</td>
      <td>39</td>
      <td>687</td>
      <td>2</td>
      <td>0.0</td>
      <td>0.0</td>
      <td>0</td>
      <td>1.0</td>
      <td>1</td>
    </tr>
    <tr>
      <th>2020-08-07</th>
      <td>285</td>
      <td>69</td>
      <td>27</td>
      <td>28</td>
      <td>686</td>
      <td>4</td>
      <td>0.0</td>
      <td>0.0</td>
      <td>0</td>
      <td>1.0</td>
      <td>3</td>
    </tr>
    <tr>
      <th>2020-08-08</th>
      <td>281</td>
      <td>67</td>
      <td>31</td>
      <td>39</td>
      <td>678</td>
      <td>2</td>
      <td>0.0</td>
      <td>0.0</td>
      <td>0</td>
      <td>2.0</td>
      <td>0</td>
    </tr>
    <tr>
      <th>2020-08-09</th>
      <td>285</td>
      <td>69</td>
      <td>20</td>
      <td>14</td>
      <td>684</td>
      <td>4</td>
      <td>1.0</td>
      <td>0.0</td>
      <td>0</td>
      <td>1.0</td>
      <td>2</td>
    </tr>
    <tr>
      <th>2020-08-10</th>
      <td>312</td>
      <td>73</td>
      <td>33</td>
      <td>12</td>
      <td>705</td>
      <td>3</td>
      <td>0.0</td>
      <td>1.0</td>
      <td>1</td>
      <td>0.0</td>
      <td>1</td>
    </tr>
  </tbody>
</table>
</div>

```python
# Load the dictionary containing the posterior parameter distributions obtained from calibrating the model to Belgian hospitalization data
with open('../../data/interim/model_parameters/deterministic_22072020.json', 'r') as fp:
    samples_dict = json.load(fp)

lag_time = samples_dict['lag_time']
model.extraTime = samples_dict['lag_time']
model.parameters['beta'] = np.mean(samples_dict['beta'])
model.parameters['l'] = np.mean(samples_dict['l'])
model.parameters['tau'] = np.mean(samples_dict['tau'])
prevention = np.mean(samples_dict['prevention'])
```

### Simulating

A simulation is run by using the attribute function *sim*, which uses one argument, the simulation time, as its input.

```python
out=model.sim(100)
```

    ['Nc', 'time']

#### Extracting model output

The output format of the deterministic and stochastic model are an `xarray` dataset. This flexible python package introduces labels in the form of dimensions, coordinates and attributes on top of raw NumPy-like arrays. Relevant to our work are the coordinates,
1. 'stratification': the nine different age groups used in the model.
2. 'time': simulation time

and the data variables, which are the 11 model states (S,E,I,...,R,D). Future work will include a spatial extension, which will simply use a new coordinate 'NIS'.

```python
out
```

As an example of output handling, consider the following simple code snippet. Here, the total number of patients in Belgian ICUs over time is given. Summing over all age groups is simple.

```python
plt.plot(out['time'],out['ICU'].sum(dim='stratification'))
```

#### Visualizing the results using builtin methods

Two custom visualization routines are included in `src/covid19model/visualization/output.py` for rapid prototyping, `infected` and `population_status`. For more details, we refer to the docstring in the aforementioned python script.

```python
infected(out)
population_status(out)
```

#### The use of checkpoints to change parameters on the fly

A cool feature of the original SEIRSplus package by Ryan McGee was the use of so-called *checkpoints* dictionary to change simulation parameters on the fly. In our modification, this feature is preserved. Below you can find an example of a *checkpoints* dictionary. The simulation will be started with the previously initialised parameters. After 55 days, social interaction will be limited by changing the interaction matrx $\mathbf{N_c}$ to home contacts only. On day 100, social restrictions are lifted and beta once more assumes its *business-as-usual* value.

```python
# Create checkpoints dictionary
chk = {'time':    [55,100],
       'Nc':      [Nc_home,Nc_total]
      }
# Run simulation
out=model.sim(130,checkpoints=chk)
# Visualise the results
population_status(out)
infected(out)
```

#### Using posterior samples from previously calibrated parameters to quantify uncertainty

```python
fig,ax=plt.subplots()
for i in range(200):
    # Draw a paramter value from the dictionary of posterior samples for beta, l and tau (fitted parameters).
    model.parameters['beta'] = np.random.choice(samples_dict['beta'],1,replace=False)
    # Correlation exists between l, tau and prevention. Thus we must use the same index to ensure we are drawing trios from the posterior distribution!
    idx,model.parameters['l'] = random.choice(list(enumerate(samples_dict['l'])))
    model.parameters['tau'] = samples_dict['tau'][idx]
    prevention = samples_dict['prevention'][idx]
    # Perform simulation
    out=model.sim(130,checkpoints=chk)
    # Visualise the results
    infected(out,ax=ax,alpha=0.02)
```

## Simulations with the stochastic model

### Model initialization

When initializing the stochastic model, a flag `discrete` must be used to switch from the integrator to the iterator.

```python
# Load the parameters using `get_COVID19_SEIRD_parameters()`.
params = parameters.get_COVID19_SEIRD_parameters()
# Add the delayed ramp parameters to the parameter dictionary.
params.update({'l': 1,
              'tau': 5,
              'd': 100}) # overdispersion coefficient (adds variability --> larger value means less influential)
# Delete the following parameters, they are not included in the stochastic implementation
del params['totalTests']
del params['dq']
del params['psi_FP']
del params['psi_PP']
# Define the initial condition: ten exposed inidividual in every age category
# Why 10? When only using one exposed, there is a big chance the chain will go extinct!
initial_states = {'S': initN, 'E': 10*np.ones(levels)}
# Initialize the model, the only difference with the stochastic model is the `discrete` flag
model = models.COVID19_SEIRD_sto(initial_states, params, compliance=ramp_2, discrete=True)
```

### Load posterior parameter distributions of calibrated parameters

The calibration procedure was exactly the same as for the deterministic model. For more details we refer to the calibration notebook, located in `notebooks/0.1-twallema-calibration-stochastic` . The posterior distributions are located in `data/interim/model_parameters/stochastic_15-08-2020.json`.

```python
with open('../../data/interim/model_parameters/stochastic_15-08-2020.json', 'r') as fp:
    samples_dict = json.load(fp)

lag_time = samples_dict['lag_time']
model.extraTime = samples_dict['lag_time']
model.parameters['beta'] = np.mean(samples_dict['beta'])
model.parameters['l'] = np.mean(samples_dict['l'])
model.parameters['tau'] = np.mean(samples_dict['tau'])
prevention = np.mean(samples_dict['prevention'])
```

### Simulating

After the model initialization, the code works in the exact same way as the deterministic model.

```python
fig,ax=plt.subplots()
for i in range(200):
    # Use posterior samples of fitted parameters
    model.parameters['beta'] = np.random.choice(samples_dict['beta'],1,replace=False)
    idx,model.parameters['l'] = random.choice(list(enumerate(samples_dict['l'])))
    model.parameters['tau'] = samples_dict['tau'][idx]
    prevention = samples_dict['prevention'][idx]
    # Perform simulation
    out=model.sim(130,checkpoints=chk)
    # Visualise the results
    infected(out,ax=ax,alpha=0.02)
```


```python
from sympy import *
```

```python
init_printing()
```

```python
from sympy.abc import a, b, c, d, f, g, h, x, y, z
```

```python
expr = a * x**2 + b * x + c
```

```python
expr
```

```python
solve(expr, x)
```

```python
m = Matrix([[1, 2], [3, 4]])
```

```python
m
```

```python
m.inv()
```

```python
m.det()
```

```python
%matplotlib inline
```

```python
plot(x**2)
```