```python
import sympy
from sympy import (Symbol, MatrixSymbol, Matrix, symbols, simplify, eye, sin, cos)

sympy.init_printing(use_latex="mathjax")
```

```python
w = symbols('w_0 w_1 w_2')
w = Matrix(3, 1, w)
R = MatrixSymbol('R', 3, 3)
w_skew = Matrix(3, 3, [0, -w[2], w[1], w[2], 0, -w[0], -w[1], w[0], 0])
w_skew
```

$\displaystyle \left[\begin{matrix}0 & - w_{2} & w_{1}\\w_{2} & 0 & - w_{0}\\- w_{1} & w_{0} & 0\end{matrix}\right]$

```python
v = (w_skew * (eye(3) - R)*eye(3)[:, 0])
v = Matrix(3, 1, Matrix(v))
v_skew = Matrix(3, 3, [0, -v[2], v[1], v[2], 0, -v[0], -v[1], v[0], 0])
dRdv0 = ((w[0] * w_skew + v_skew) * R) / w.norm(2)
dRdv0
```

$\displaystyle \frac{1}{\sqrt{\left|{w_{0}}\right|^{2} + \left|{w_{1}}\right|^{2} + \left|{w_{2}}\right|^{2}}} \left[\begin{matrix}0 & - w_{0} w_{2} + w_{0} R_{1, 0} + w_{1} \left(1 - R_{0, 0}\right) & w_{0} w_{1} + w_{0} R_{2, 0} + w_{2} \left(1 - R_{0, 0}\right)\\w_{0} w_{2} - w_{0} R_{1, 0} - w_{1} \left(1 - R_{0, 0}\right) & 0 & - w_{0}^{2} + w_{1} R_{2, 0} - w_{2} R_{1, 0}\\- w_{0} w_{1} - w_{0} R_{2, 0} - w_{2} \left(1 - R_{0, 0}\right) & w_{0}^{2} - w_{1} R_{2, 0} + w_{2} R_{1, 0} & 0\end{matrix}\right] R$

```python
v = (w_skew * (eye(3) - R)*eye(3)[:, 1])
v = Matrix(3, 1, Matrix(v))
v_skew = Matrix(3, 3, [0, -v[2], v[1], v[2], 0, -v[0], -v[1], v[0], 0])
dRdv1 = ((w[1] * w_skew + v_skew) * R) / w.norm(2)
dRdv1
```

$\displaystyle \frac{1}{\sqrt{\left|{w_{0}}\right|^{2} + \left|{w_{1}}\right|^{2} + \left|{w_{2}}\right|^{2}}} \left[\begin{matrix}0 & - w_{0} \left(1 - R_{1, 1}\right) - w_{1} w_{2} - w_{1} R_{0, 1} & w_{0} R_{2, 1} + w_{1}^{2} - w_{2} R_{0, 1}\\w_{0} \left(1 - R_{1, 1}\right) + w_{1} w_{2} + w_{1} R_{0, 1} & 0 & - w_{0} w_{1} + w_{1} R_{2, 1} + w_{2} \left(1 - R_{1, 1}\right)\\- w_{0} R_{2, 1} - w_{1}^{2} + w_{2} R_{0, 1} & w_{0} w_{1} - w_{1} R_{2, 1} - w_{2} \left(1 - R_{1, 1}\right) & 0\end{matrix}\right] R$

```python
v = (w_skew * (eye(3) - R)*eye(3)[:, 2])
v = Matrix(3, 1, Matrix(v))
v_skew = Matrix(3, 3, [0, -v[2], v[1], v[2], 0, -v[0], -v[1], v[0], 0])
dRdv2 = ((w[2] * w_skew + v_skew) * R) / w.norm(2)
dRdv2
```

$\displaystyle \frac{1}{\sqrt{\left|{w_{0}}\right|^{2} + \left|{w_{1}}\right|^{2} + \left|{w_{2}}\right|^{2}}} \left[\begin{matrix}0 & w_{0} R_{1, 2} - w_{1} R_{0, 2} - w_{2}^{2} & - w_{0} \left(1 - R_{2, 2}\right) + w_{1} w_{2} - w_{2} R_{0, 2}\\- w_{0} R_{1, 2} + w_{1} R_{0, 2} + w_{2}^{2} & 0 & - w_{0} w_{2} - w_{1} \left(1 - R_{2, 2}\right) - w_{2} R_{1, 2}\\w_{0} \left(1 - R_{2, 2}\right) - w_{1} w_{2} + w_{2} R_{0, 2} & w_{0} w_{2} + w_{1} \left(1 - R_{2, 2}\right) + w_{2} R_{1, 2} & 0\end{matrix}\right] R$

```python
f_x, f_y, p_x, p_y = symbols('f_x f_y p_x p_y')
x_w, y_w, z_w = symbols('x_w y_w z_w')

X_w = Matrix(3, 1, [x_w, y_w, z_w])
K = Matrix(3, 3, [f_x, 0, p_x, 0, f_y, p_y, 0, 0, 1])
Rv = Matrix(9, 1, Matrix(R))
t = MatrixSymbol('t', 3, 1)
X_i_homo = Matrix(K * (R * X_w + t))
X_i_homo
X_i = Matrix(X_i_homo[0:2, 0] / X_i_homo[2, 0])

dxdt = X_i.jacobian(t).applyfunc(simplify)
dxdt
```

$\displaystyle \left[\begin{matrix}\frac{f_{x}}{x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}} & 0 & - \frac{f_{x} \left(x_{w} R_{0, 0} + y_{w} R_{0, 1} + z_{w} R_{0, 2} + t_{0, 0}\right)}{\left(x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}\right)^{2}}\\0 & \frac{f_{y}}{x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}} & - \frac{f_{y} \left(x_{w} R_{1, 0} + y_{w} R_{1, 1} + z_{w} R_{1, 2} + t_{1, 0}\right)}{\left(x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}\right)^{2}}\end{matrix}\right]$

```python
dxdRv = X_i.jacobian(Rv).applyfunc(simplify)
dxdRv
```

$\displaystyle \left[\begin{matrix}\frac{f_{x} x_{w}}{x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}} & \frac{f_{x} y_{w}}{x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}} & \frac{f_{x} z_{w}}{x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}} & 0 & 0 & 0 & - \frac{f_{x} x_{w} \left(x_{w} R_{0, 0} + y_{w} R_{0, 1} + z_{w} R_{0, 2} + t_{0, 0}\right)}{\left(x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}\right)^{2}} & - \frac{f_{x} y_{w} \left(x_{w} R_{0, 0} + y_{w} R_{0, 1} + z_{w} R_{0, 2} + t_{0, 0}\right)}{\left(x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}\right)^{2}} & - \frac{f_{x} z_{w} \left(x_{w} R_{0, 0} + y_{w} R_{0, 1} + z_{w} R_{0, 2} + t_{0, 0}\right)}{\left(x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}\right)^{2}}\\0 & 0 & 0 & \frac{f_{y} x_{w}}{x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}} & \frac{f_{y} y_{w}}{x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}} & \frac{f_{y} z_{w}}{x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}} & - \frac{f_{y} x_{w} \left(x_{w} R_{1, 0} + y_{w} R_{1, 1} + z_{w} R_{1, 2} + t_{1, 0}\right)}{\left(x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}\right)^{2}} & - \frac{f_{y} y_{w} \left(x_{w} R_{1, 0} + y_{w} R_{1, 1} + z_{w} R_{1, 2} + t_{1, 0}\right)}{\left(x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}\right)^{2}} & - \frac{f_{y} z_{w} \left(x_{w} R_{1, 0} + y_{w} R_{1, 1} + z_{w} R_{1, 2} + t_{1, 0}\right)}{\left(x_{w} R_{2, 0} + y_{w} R_{2, 1} + z_{w} R_{2, 2} + t_{2, 0}\right)^{2}}\end{matrix}\right]$


Exercise material of the MSc-level course **Numerical Methods in Geotechnical Engineering**.
Held at Technische Universität Bergakademie Freiberg.

Comments to:

*Prof. Dr. Thomas Nagel
Chair of Soil Mechanics and Foundation Engineering
Geotechnical Institute
Technische Universität Bergakademie Freiberg.*

https://tu-freiberg.de/en/soilmechanics

# Exercise 8 - Coupled Problems, monolithic and staggered schemes

Consider two coupled ODEs:

\begin{align}
    \dot{y} &= k_1 (y_\text{eq}-y)(z_\text{eq}-c_1z)
    \\
    \dot{z} &= k_2 (z_\text{eq}-z)^3 + c_2\dot{y}
\end{align}

First, we import some libraries for numerical basics and for plotting.

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

#Some plot settings
plt.style.use('seaborn-deep')
plt.rcParams['lines.linewidth']= 2.0
plt.rcParams['lines.color']= 'black'
plt.rcParams['legend.frameon']=True
plt.rcParams['font.family'] = 'serif'
plt.rcParams['legend.fontsize']=14
plt.rcParams['font.size'] = 14
plt.rcParams['axes.spines.right'] = False
plt.rcParams['axes.spines.top'] = False
plt.rcParams['axes.spines.left'] = True
plt.rcParams['axes.spines.bottom'] = True
plt.rcParams['axes.axisbelow'] = True
plt.rcParams['figure.figsize'] = (8, 6)
```

For simplicity, let's consider a Backward Euler scheme, so that

\begin{align}
    \frac{y_{n+1}- y_n}{t_{n+1} - t_n} &= k_1 (y_\text{eq} - y_{n+1})(z_\text{eq} - c_1 z_{n+1})
    \\
    \frac{z_{n+1}- z_n}{t_{n+1} - t_n} &= k_2 (z_\text{eq} - z_{n+1})^3 + c_2\frac{y_{n+1}- y_n}{t_{n+1} - t_n}
\end{align}

This can be re-cast in a residual form, i.e based on an approximation error:

\begin{align}
    r_y &= \frac{y_{n+1}- y_n}{t_{n+1} - t_n} - k_1 (y_\text{eq} - y_{n+1})(z_\text{eq} - c_1 z_{n+1})
    \\
    r_z &= \frac{z_{n+1}- z_n}{t_{n+1} - t_n} - k_2 (z_\text{eq} - z_{n+1})^3 - c_2 \frac{y_{n+1}- y_n}{t_{n+1} - t_n}
\end{align}

where we would like to find the solutions $y_{n+1}$ and $z_{n+1}$ such that $r_y(y_{n+1},z_{n+1}) = 0$ and $r_z(y_{n+1},z_{n+1}) = 0$ *simultaneously*. To achieve that, we have several options.

## Monolithic scheme

The monolithic scheme iteratively solves both equations simultaneously, i.e. in a monolithic entity. Choosing a Newton-Raphson approach to resolve the non-linearities results in

\begin{align}
    r_y(y_{n+1},z_{n+1}) &\approx r_y(y_n,z_n) + \left. \frac{\partial r_y}{\partial y} \right|_n \underbrace{(y_{n+1} - y_n)}_{\displaystyle \Delta y_{n+1}} + \left. \frac{\partial r_y}{\partial z} \right|_n \underbrace{(z_{n+1} - z_n)}_{\displaystyle \Delta z_{n+1}} \overset{!}{=} 0
    \\
    r_z(y_{n+1},z_{n+1}) &\approx r_z(y_n,z_n) + \left. \frac{\partial r_z}{\partial y} \right|_n \underbrace{(y_{n+1} - y_n)}_{\displaystyle \Delta y_{n+1}} + \left. \frac{\partial r_z}{\partial z} \right|_n \underbrace{(z_{n+1} - z_n)}_{\displaystyle \Delta z_{n+1}} \overset{!}{=} 0
\end{align}

To make the monolithic structure more clearly visible, let's re-arrange this to $-r_i = J_{ij} \Delta u_j$ form and cast it into a vector-matrix format (making the time step index $n+1$ implicit and introducing the iteration counter $k$):

$$
-
\left(
\begin{array}{c}
    r_y(y_k,z_k)
    \\
    r_z(y_k,z_k)
\end{array}
\right)
=
\left.
\left(
\begin{array}{cc}
    \displaystyle \frac{\partial r_y}{\partial y} & \displaystyle \frac{\partial r_y}{\partial z}
    \\
    \displaystyle \frac{\partial r_z}{\partial y} & \displaystyle \frac{\partial r_z}{\partial z}
\end{array}
\right)
 \right|_k
 \left(
\begin{array}{c}
    \Delta y_{k+1}
    \\
    \Delta z_{k+1}
\end{array}
\right)
=
\left.
\left(
\begin{array}{cc}
    J_{yy} & J_{yz}
    \\
    J_{zy} & J_{zz}
\end{array}
\right)
 \right|_k
 \left(
\begin{array}{c}
    \Delta y_{k+1}
    \\
    \Delta z_{k+1}
\end{array}
\right)
$$

```python
k_1 = 2.
y_eq = 1.
k_2 = 2
z_eq = 1.
```

```python
def residual_y(y_i,y_n,z_i,dt,c_1):
    return (y_i - y_n)/dt - k_1 * (y_eq - y_i)*(z_eq - c_1 * z_i)
```

```python
def residual_z(y_i,y_n,z_i,z_n,dt,c_2):
    return (z_i - z_n)/dt - k_2 * (z_eq - z_i)**3 - c_2 * (y_i - y_n)/dt
```

```python
def J_yy(z_i,dt, c_1):
    return 1./dt + k_1*(z_eq - c_1 * z_i)
```

```python
def J_yz(y_i, c_1):
    return k_1 * (y_eq - y_i) * c_1
```

```python
def J_zz(z_i,dt):
    return 1./dt + 3*k_2*(z_eq - z_i)**2
```

```python
def J_zy(dt, c_2):
    return -c_2/dt
```

Now, for every time step we advance the system by applying the generalized midpoint Newton-Raphson update rule. Each new time step value then serves as the initial condition for the following linearization step.

We will *recursively* apply the Newton update until the residual is below a specified tolerance:

$\left| r_{n+1}^{i+1}\right| < \epsilon_\text{abs}$

with a suitably chosen absolute tolerance set here to $10^{-6}$.

```python
def integrate_monolithic(dt,c1,c2):
    #Startwerte
    t_end = 5.
    absolute_tolerance = 1.e-6
    max_iter = 1000
    iteration_counter = np.array([0])
    u = [np.array([0., 0.])]#initial values for y and z
    times = np.array([0.])
    #
    while times[-1]+dt < t_end: #repeat the loop as long as the final time step is below the end point
        times = np.append(times,times[-1]+dt) #here define the next time point as the previous time point plus the time increment dt
        u_old = u[-1] #Starting value for recursive update
        i = 0
        #
        while True:
            #evaluate residual
            res = np.array([residual_y(u_old[0],u[-1][0],u_old[1],dt,c1),
                            residual_z(u_old[0],u[-1][0],u_old[1],u[-1][1],dt,c2)])
            #if residual is below tolerance, above maximum iterations, stop iterations
            if (np.linalg.norm(res) < absolute_tolerance or i > max_iter):
                break
            #evaluate Jacobian
            Jac_yy = J_yy(u_old[1],dt,c1)
            Jac_yz = J_yz(u_old[0],c1)
            Jac_zy = J_zy(dt,c2)
            Jac_zz = J_zz(u_old[1],dt)
            Jac = np.array([[Jac_yy,Jac_yz],[Jac_zy,Jac_zz]])
            #perform linear step
            u_new = u_old + np.linalg.solve(Jac,-res)
            #update counter
            i += 1
            u_old = u_new #preparation of next recursion

        u.append(u_new) #append the new found solution to the solution vector
        iteration_counter = np.append(iteration_counter,i) #store how much iterations this time step took to converge

    return times, u, iteration_counter

```

```python
from ipywidgets import widgets
from ipywidgets import interact

@interact(dt=widgets.BoundedFloatText(value=0.2,min=1e-3,max=0.8,description='$\Delta t$ / s'))

def plot_monolithic(dt=0.01):
    time_n, u_n, iters11 = integrate_monolithic(dt,1.,1.)
    u11 = np.array(u_n)
    fig, ax = plt.subplots(ncols=2,nrows=2,figsize=(18,12))
    #solution
    ax[0][0].plot(time_n,u11.T[0],marker='d',label='11')
    ax[0][0].legend()
    ax[0][0].set_xlabel('$t$ / s')
    ax[0][0].set_ylabel('$y$')
    #error
    ax[0][1].plot(time_n,u11.T[1],marker='d',label='11')
    ax[0][1].legend()
    ax[0][1].set_xlabel('$t$ / s')
    ax[0][1].set_ylabel('$z$')
    #iterations
    ax[1][0].plot(iters11,marker='d',label='11')
    ax[1][0].set_xlabel('time step')
    ax[1][0].set_ylabel('number of iterations')
    ax[1][0].legend()
    plt.show()
```

# Tasks
* How does the number of iterations change if the system if you a) uncouple the system or b) choose a uni-lateral coupling?
* How does the system evolve for other parameter settings? Try to understand why. Examples: $(c_1,c_2) = (0,2),(1,-5)$

## Staggered scheme

A staggered scheme works by solving one equation for one unknown while keeping the other unknown(s) fixed. With the new solution, one proceeds to the other equation(s) to solve for the next unknown, while keeping the other unknown(s) fixed. In this way, one iterates between the equations until the solution is stationary.

Choosing a Newton-Raphson approach to resolve the non-linearities results in

\begin{align}
    r_y(y_{n+1},z_n) &\approx r_y(y_n,z_n) + \left. \frac{\partial r_y}{\partial y} \right|_n \underbrace{(y_{n+1} - y_n)}_{\displaystyle \Delta y_{n+1}} \overset{!}{=} 0
    \\
    r_z(y_n,z_{n+1}) &\approx r_z(y_n,z_n) + \left. \frac{\partial r_z}{\partial z} \right|_n \underbrace{(z_{n+1} - z_n)}_{\displaystyle \Delta z_{n+1}} \overset{!}{=} 0
\end{align}

To make the staggered scheme more clearly visible, let's re-arrange this to $-r_i = J_{ij} \Delta u_j$ form, cast it into a vector-matrix format, and introduce the Newton iteration counters $k$ and $l$ as well as the coupling iteration counter $o$:

\begin{align}
            - r_y({}^oy^k_{n+1},{}^o z_{n+1}) &= {}^o J_{yy}|_{n+1}^k \Delta {}^oy_{n+1}^{k+1}
            \\
            - r_z({}^oy_{n+1},{}^o z_{n+1}^k) &= {}^o J_{zz}|_{n+1}^k \Delta {}^oz_{n+1}^{k+1}
\end{align}

The coupling iterations $o$ are the outer loop, while the inner loop consists of the Newton iterations. The overall solution is found when the outer solutions of the individual unknowns $z$ and $y$ no longer change between coupling iterations:

$$
    |{}^oy^k_{n+1} - {}^{o-1}y^k_{n+1}| + |{}^oz^k_{n+1} - {}^{o-1}z^k_{n+1}| < \epsilon_\text{abs}
$$

```python
def integrate_staggered(dt,c1,c2):
    #Startwerte
    t_end = 5.
    absolute_tolerance = 1.e-6
    max_iter = 100
    iteration_counter = np.array([0])
    y = np.array([0.])#initial values for y
    z = np.array([0.])#initial values for z
    times = np.array([0.])
    #
    while times[-1]+dt < t_end: #repeat the loop as long as the final time step is below the end point
        times = np.append(times,times[-1]+dt) #here define the next time point as the previous time point plus the time increment dt
        y_old = y[-1] #Starting value for recursive update
        z_old = z[-1] #Starting value for recursive update
        y_coup = y_old
        z_coup = z_old
        i = 0
        #
        while True:
            while True:
                #evaluate residual
                res = residual_y(y_old,y[-1],z_old,dt,c1)
                #if residual is below tolerance, above maximum iterations, stop iterations
                if (np.abs(res) < absolute_tolerance or i > max_iter):
                    break
                #perform linear step
                y_new = y_old - res/J_yy(z_old,dt,c1)
                #update counter
                i += 1
                y_old = y_new #preparation of next recursion

            while True:
                #evaluate residual
                res = residual_z(y_old,y[-1],z_old,z[-1],dt,c2)
                #if residual is below tolerance, above maximum iterations, stop iterations
                if (np.abs(res) < absolute_tolerance or i > max_iter):
                    break
                #perform linear step
                z_new = z_old - res/J_zz(z_old,dt)
                #update counter
                i += 1
                z_old = z_new #preparation of next recursion

            if ((np.abs(y_coup - y_new) + np.abs(z_coup - z_new)) < absolute_tolerance or i > max_iter):
                break
            y_coup = y_new
            z_coup = z_new

        y = np.append(y,y_new) #append the new found solution to the solution vector
        z = np.append(z,z_new) #append the new found solution to the solution vector
        iteration_counter = np.append(iteration_counter,i) #store how much iterations this time step took to converge

    return times, y, z, iteration_counter

```

```python
from ipywidgets import widgets
from ipywidgets import interact

@interact(dt=widgets.BoundedFloatText(value=0.2,min=1e-3,max=0.8,description='$\Delta t$ / s'))

def plot_staggered(dt=0.01):
    time_n, y, z, iters11 = integrate_staggered(dt,1.,1.)
    fig, ax = plt.subplots(ncols=2,nrows=2,figsize=(18,12))
    #solution
    ax[0][0].plot(time_n,y,marker='d',label='11')
    ax[0][0].legend()
    ax[0][0].set_xlabel('$t$ / s')
    ax[0][0].set_ylabel('$y$')
    #error
    ax[0][1].plot(time_n,z,marker='d',label='11')
    ax[0][1].legend()
    ax[0][1].set_xlabel('$t$ / s')
    ax[0][1].set_ylabel('$z$')
    #iterations
    ax[1][0].plot(iters11,marker='d',label='11')
    ax[1][0].set_xlabel('time step')
    ax[1][0].set_ylabel('number of iterations')
    ax[1][0].legend()
    plt.show()
```

```python
from ipywidgets import widgets
from ipywidgets import interact

@interact(dt=widgets.BoundedFloatText(value=0.2,min=1e-3,max=0.8,description='$\Delta t$ / s'))

def plot_comparison(dt=0.01):

    fig, ax = plt.subplots(ncols=2,nrows=2,figsize=(18,12))
    #set values for coupling parameters
    c1 = [0.,1.]
    c2 = [0.,1.,-5.]
    # run all combinations
    for i in range(len(c1)):
        for j in range(len(c2)):
            #run both schemes
            time_n, u_n, iters_m = integrate_monolithic(dt,c1[i],c2[j])
            u_m = np.array(u_n)
            time_n, y_s, z_s, iters_s = integrate_staggered(dt,c1[i],c2[j])
            #plot color
            color = next(ax[0][0]._get_lines.prop_cycler)['color']
            #solution y
            ax[0][0].plot(time_n,u_m.T[0],ls='--',color=color)
            ax[0][0].plot(time_n,y_s,color=color)
            #

            #solution z
            ax[0][1].plot(time_n,u_m.T[1],ls='--',color=color)
            ax[0][1].plot(time_n,z_s,color=color)

            #iterations
            ax[1][0].plot(iters_m,ls='--',color=color)
            ax[1][0].plot(iters_s,color=color)

            #dummy
            ax[1][1].plot(0,0,color=color,label="$c_1 = %.1f$, $c_2 = %.1f$" %(c1[i],c2[j]))

    ax[0][0].set_xlabel('$t$ / s')
    ax[0][0].set_ylabel('$y$')
    ax[0][1].set_xlabel('$t$ / s')
    ax[0][1].set_ylabel('$z$')
    ax[1][0].set_xlabel('time step')
    ax[1][0].set_ylabel('number of iterations')
    ax[1][1].legend()
    ax[1][1].set_title('monolithic --; staggered __')
    plt.show()
```

## Tasks

* What main difference do you observe between both schemes?
* Run the combination $c_1=-3$, $c_2 = -5$. What do you observe? What happens if you a) increase and b) decrease the time steps?


# An pictoral introduction to Gaussian processes and using them to calculate the Greeks

```python
import numpy as np
from bokeh.plotting import figure, show
from bokeh.io import output_notebook
from bokeh.models import LinearColorMapper, BasicTicker, ColorBar
from bokeh.palettes import Category10, Accent
output_notebook()
```

    Loading BokehJS ...

## Gaussian processes
***

Gaussian processes are **distributions over functions** $f:\mathbb{R}^n \to \mathbb{R}$. Since a Gaussian is uniquely determined by its mean and covariance function, we consider its mean function $m:\mathbb{R}^n \to \mathbb{R}$ and *positive definite* covariance function $k:\mathbb{R}^n \times \mathbb{R}^n \to \mathbb{R}$.

We write this mathematically as
$$f \sim \mathcal{GP}(m,k),$$

and require that, for any finite subset $X =\{\mathbf{x}_1 \ldots \mathbf{x}_n \}$ of the domain of $x$, the marginal distribution is a *multivariate Gaussian distribution*:

$$f(X) \sim \mathcal{N}(m(X), k(X, X))$$

with mean vector $\mathbf{\mu} = m(X)$ and covariance matrix $\Sigma = k(X, X)$.

**While the multivariate Gaussian caputures a finite number of jointly distributed Gaussians, the Gaussian process doesn't have this limitation.** Its mean and covariance are defined by a **function**. Each output of this function is a variable correlated with the other variables in the input domain, as defined by the covariance function. Since functions can have an infinite input domain, the Gaussian process can be interpreted as an infinite dimensional Gaussian random variable.

It is shown in Chapter 3.2 of [Dixon et al.](https://www.springer.com/gp/book/9783030410674) that the predictive
distribution of the Gaussian process evaluated at an arbitrary test point $x_*\in X_*$ is:

$$ f_* | X, Y, x_* \sim \mathcal{N}(\mathbb{E}[f_*| X,Y,x_*], \text{var}[f_*|X,Y,x_*]),$$
where
\begin{align}\label{eq:asasa}
	\mathbb{E}[f_*| X,Y,x_*] =&~ \mu(X_*) + K_{X_*,X}[K_{X,X} + \sigma^2 I]^{-1} Y \qquad \qquad \qquad \qquad (1)
	\\
	\text{var}[f_*|X,Y,x_*] =&~ K_{X_*,X_*} - K_{X_*,X}[K_{X,X} + \sigma^2 I]^{-1}K_{X,X_*}. \qquad\qquad\qquad (2)
\end{align}

Here, $K_{X,X_*}$, $K_{X,X}$ etc, are matrices that consist of the kernel, $k$, evaluated at the corresponding points (in the first case, at $X$ and $X_*$).

## White Noise
***
The simplest possible choice of covariance kernel $k$ is the (Kronecker) delta function $k (x_1,x_2) = \delta(||x_1 - x_2||) = \mathbb{1}_{\{x_1=x_2\}}$.

In this case
$$ \begin{pmatrix} y_0 \\ \vdots\\y_n \end{pmatrix} \sim\mathcal{N}\left( \boldsymbol{0}, \boldsymbol{I_n}\right) $$

```python
def plot_unit_gaussian_samples(D):
    p = figure(plot_width=800, plot_height=500, title='Samples from a unit {}D Gaussian'.format(D))

    xs = np.linspace(0, 1, D)
    for color in Accent[3]:
        ys = np.random.multivariate_normal(np.zeros(D), np.eye(D))
        p.line(xs, ys, line_width=1, color=color)
    return p
```

```python
show(plot_unit_gaussian_samples(50))
```

This is just white noise! It seems unlikely that this is going to help us predict anything though... **We want something smoother.**

## The squared exponential kernel
***

We therefore consider the __squared exponential kernel__ which is defined by:

$$k(x_1, x_2) = \sigma^2\exp{ \left( -\frac{1}{2 l^2} \lVert x_1 - x_2 \rVert^2 \right)}$$

This function is 1 when $x_1=x_2$ and tends to zero as its arguments drift apart.

```python
def squared_exponential(xs, ys, sigma=1, l=1):
    """Squared Exponential kernel as above but designed to return the whole
    covariance matrix - i.e. the pairwise covariance of the vectors xs & ys.
    Also with two parameters which are discussed at the end."""

    # Pairwise difference matrix.
    dx = np.expand_dims(xs, 1) - np.expand_dims(ys, 0)
    return (sigma ** 2) * np.exp(-((dx / l) ** 2) / 2)

def m(x):
    """The mean function. We can let the mean always be zero."""
    return np.zeros_like(x)
```

```python
N = 100
xs = np.linspace(-5, 5, N)
ys = [squared_exponential(np.array([0]),np.array([x])) for x in xs]

p = figure(plot_width=800, plot_height=500, x_axis_label='x', y_axis_label='k(x,0)', title='The squared exponential k(0,x)')
p.line(xs, ys, line_width=1)
show(p)
```

```python
N = 100
x = np.linspace(-2, 2, N)
y = np.linspace(-2, 2, N)
d = squared_exponential(x, y)

color_mapper = LinearColorMapper(palette="Plasma256", low=0, high=1)

p = figure(plot_width=400, plot_height=400, x_range=(-2, 2), y_range=(-2, 2),
           title='Visualisation of k(x, x\')', x_axis_label='x', y_axis_label='x\'', toolbar_location=None)
p.image(image=[d], color_mapper=color_mapper, x=-2, y=-2, dw=4, dh=4)

color_bar = ColorBar(color_mapper=color_mapper, ticker=BasicTicker(),
                     label_standoff=12, border_line_color=None, location=(0,0))

p.add_layout(color_bar, 'right')

show(p)
```

Now, let's try plotting some samples produced by the squared exponential kernel.

```python
n = 100
xs = np.linspace(-5, 5, n)
K = squared_exponential(xs, xs, sigma=1, l=1)
mu = m(xs)

p = figure(plot_width=800, plot_height=500)
for color in Category10[5]:
    ys = np.random.multivariate_normal(mu, K)
    p.line(xs, ys, line_width=2, color=color)
show(p)
```

## The predictive distribution
***
Let's now try and use Gaussian process regression with the squared exponential kernel to fit a "hidden" function.

We take the hidden function to be a $5^{th}$ order polynomial:
$$f(x)=0.03 x^5 + 0.2 x^4 - 0.1 x^3 - 2.4 x^2 - 2.5 x + 6.$$

```python
coefs = [6, -2.5, -2.4, -0.1, 0.2, 0.03]

def f(x):
    total = 0
    for exp, coef in enumerate(coefs):
        total += coef * (x ** exp)
    return total
```

```python
xs = np.linspace(-5.0, 3.5, 100)
ys = f(xs)

p = figure(plot_width=800, plot_height=400, x_axis_label='x', y_axis_label='f(x)',
           title='The hidden function f(x)')
p.title.text_font_size = "15pt"
p.line(xs, ys, line_width=2)
show(p)
```

Let's suppose we observe the true function on some set $X = \{-4, -1.5, 0, 1.5, 2.5, 2.7\}$. Let's see how the Gaussian process fares over the unseen set $[-8,8]$. We make it a little larger to show that the Gaussian process with squared exponential kernel fares fairly poorly outside of seen values.

```python
x_obs = np.array([-4, -1.5, 0, 1.5, 2.5, 2.7])
y_obs = f(x_obs)
x_s = np.linspace(-8,8, 500)
```

The covariance matrices $K_{X,X}$, $K_{X, X_*}$ and $K_{x_*,X_*}$ are given by:

```python
K = squared_exponential(x_obs, x_obs)
K_s = squared_exponential(x_obs, x_s)
K_ss = squared_exponential(x_s, x_s)
```

Using the formulas in [Dixon et al.](https://www.springer.com/gp/book/9783030410674), and assuming the noiseless case, we find that

```python
K_sTKinv = np.matmul(K_s.T, np.linalg.pinv(K))
mu_s = m(x_s) + np.matmul(K_sTKinv, y_obs - m(x_obs))
Sigma_s = K_ss - np.matmul(K_sTKinv, K_s)
```

We now plot the mean function along with 2 standard deviations of the posterior distribution.

```python
def plot_regression_figure(x_obs,y_obs,x_s,y_true,stds,mu_s):
    p = figure(plot_width=800, plot_height=600, y_range=(-7, 8))
    p.line(x_s, y_true, line_width=3, color='black', alpha=0.7, line_dash='dashed', legend_label='True f(x)')
    p.cross(x_obs, y_obs, size=20, legend_label='Training data')
    err_xs = np.concatenate((x_s, np.flip(x_s, 0))) # concatenating x_s with x_s reversed...
    err_ys = np.concatenate((mu_s + 2 * stds, np.flip(mu_s - 2 * stds, 0)))
    p.patch(err_xs, err_ys, alpha=0.2, line_width=0, color='grey', legend_label='Uncertainty')
    p.line(x_s, mu_s, line_width=3, color='blue', alpha=0.4, legend_label='Mean')
    return p
```

```python
np.random.seed(42)
y_true = f(x_s)
stds = np.sqrt(Sigma_s.diagonal())
p = plot_regression_figure(x_obs,y_obs,x_s,y_true,stds,mu_s)
for count, color in enumerate(Category10[3]):
    y_s = np.random.multivariate_normal(mu_s, Sigma_s)
    p.line(x_s, y_s, line_width=1, color=color, legend_label=f"sample {count+1}")
show(p)
```

As we can see the prediction is poor outside of the range of the points seen. It is also worth noting that the variance of points that have not been seen converges to the *prior* $K_{X_*,X_*}$. This is just the parameter $\sigma$ of the squared exponential, which we have set to $1$.

Fitting a Gaussian kernel is done via maximizing the marginal likelihood  $p(\mathbf{y} \mid X, \theta)$ of the Gaussian process distribution based on the observed data $(X, \mathbf{y})$.

$$\hat{\theta}  = \underset{\theta}{\text{argmax}} \left( p(\mathbf{y} \mid X, \theta) \right)$$

The marginal likelihood of the Gaussian process is the likelihood of a Gaussian distribution which is defined as:

$$p(\mathbf{y} \mid \mu, \Sigma) = \frac{1}{\sqrt{(2\pi)^n \lvert\Sigma\rvert}} \exp{ \left( -\frac{1}{2}(\mathbf{y} - \mu)^{\top} \Sigma^{-1} (\mathbf{y} - \mu) \right)}$$

The mean and covariance are calulated from their parameterized functions using the observed data $X$ as input: $\mu_{\theta} = m_{\theta}(X)$ and $\Sigma_{\theta} = k_{\theta}(X, X)$, so we can write the marginal likelihood as:

$$
p(\mathbf{y} \mid X, \theta) =
\frac{1}{\sqrt{(2\pi)^n \lvert \Sigma_{\theta} \rvert}} \exp{ \left( -\frac{1}{2}(\mathbf{y} - \mu_{\theta})^\top \Sigma_{\theta}^{-1} (\mathbf{y} - \mu_{\theta}) \right)}$$

With $n$ the dimensionality of the marginal and $\lvert \Sigma_{\theta} \rvert$ the deteriminant of the kernel matrix. We can get rid of the exponent by taking the log and maximizing the log marginal likilihood:

$$ \log{p(\mathbf{y} \mid X, \theta)} =
-\frac{1}{2}(\mathbf{y} - \mu_{\theta})^\top \Sigma_{\theta}^{-1} (\mathbf{y} - \mu_{\theta}) - \frac{1}{2} \log{\lvert \Sigma_{\theta} \rvert} - \frac{n}{2} \log{2 \pi}$$

The first term ($-\frac{1}{2} (\mathbf{y} - \mu_{\theta})^\top \Sigma_{\theta}^{-1} (\mathbf{y} - \mu_{\theta})$) is the data-fit while the rest ($-\frac{1}{2}(\log{\lvert \Sigma_{\theta} \rvert} + n\log{2 \pi})$) is a complexity penalty, also known as [differential entropy](https://en.wikipedia.org/wiki/Differential_entropy)[⁽¹⁾](#Sidenotes).

The optimimal parameters $\hat{\theta}$ can then be found by minimizing the negative of the log marginal likilihood:

$$
\hat{\theta}  = \underset{\theta}{\text{argmax}} \left( p(\mathbf{y} \mid X, \theta) \right) = \underset{\theta}{\text{argmin}} { \;-\log{ p(\mathbf{y} \mid X, \theta)}}$$

In practice, this is done via numerical methods. We do not code this up, but instead leverage the scikit-learn Gaussian process functionality to fit the model to data.

## Using Scikit-learn
***

In this section we repreoduce the above, using Scikit-learn instead of first principles. We then use the optimisation functionality of sklearn to tune the hyper-parameters of our kernel and get a better fit.

```python
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import Matern, WhiteKernel, ConstantKernel, RBF
```

The `scikit-learn` API is very consistent across learning methods, and as a result, all functions expect a tabular set of input variables, either as a 2-dimensional NumPy `array` or a pandas `DataFrame`. Since we have only a single input variable here, we can add a second dimension using the `reshape` method:

```python
X_obs= x_obs.reshape(-1, 1)
Y_obs = y_obs.reshape(-1,1)
X_s = x_s.reshape(-1,1)
Y_s = y_true.reshape(-1,1)
X_obs.shape
```

    (6, 1)

```python
kernel = RBF(length_scale=1.0, length_scale_bounds="fixed")
gp = GaussianProcessRegressor(kernel=kernel)
gp.fit(X_obs,Y_obs)
Y_pred, Sigma_pred = gp.predict(X_s, return_std=True)
```

```python
y_pred = Y_pred.reshape(1,500)[0]
sigma_pred = Sigma_pred.reshape(1,500)[0]

np.random.seed(42)
p = plot_regression_figure(x_obs,y_obs,x_s,y_true,sigma_pred,mu_s)
p.title.text = f"The fitted Gaussian process regressor with {kernel} kernel"
p.title.align = "center"
p.title.text_font_size = "20px"
for count, color in enumerate(Category10[3]):
    y_s = gp.sample_y(X_s, random_state=count).reshape(1,len(X_s))[0]
    p.line(x_s, y_s, line_width=1, color=color, legend_label=f"sample {count+1}")
show(p)
```

### ... Now we optimise our kernel

```python
N = 1000
xs = np.linspace(-10, 10, N)
p = figure(plot_width=800, plot_height=500, x_axis_label='x', y_axis_label='k(0,x)', title='The squared exponential k(0,x)',
          y_range=(-.1,1.2))

ls = [0.2,1,5]

for count, color in enumerate(Category10[3]):
    ys = [squared_exponential(np.array([0]),np.array([x]), l=ls[count]) for x in xs]
    p.line(xs, ys, line_width=1, color=color, legend_label=f"l = {ls[count]}")

show(p)
```

```python
p = figure(plot_width=800, plot_height=600, y_range=(-7, 8),title='Different values of l')
p.title.text_font_size = "15pt"
for count, color in enumerate(Category10[4]):
    l = 10**(count-1)
    kernel = RBF(length_scale=l, length_scale_bounds="fixed")
    gp = GaussianProcessRegressor(kernel=kernel)
    y_s = gp.sample_y(X_s, random_state=count).reshape(1,len(X_s))[0]
    p.line(x_s, y_s, line_width=1, color=color, legend_label=f"l = {l}")
show(p)
```

```python
kernel = RBF(length_scale=1.0, length_scale_bounds=(1e-05, 1000.0))
gp = GaussianProcessRegressor(kernel=kernel)
gp.fit(X_obs,Y_obs)
Y_pred, Sigma_pred = gp.predict(X_s, return_std=True)
np.random.seed(42)
p = plot_regression_figure(x_obs,y_obs,x_s,y_true,sigma_pred,mu_s)
p.title.text = f"The fitted Gaussian process regressor with {gp.kernel_} kernel"
p.title.align = "center"
p.title.text_font_size = "20px"
for count, color in enumerate(Category10[3]):
    y_s = gp.sample_y(X_s, random_state=count).reshape(1,len(X_s))[0]
    p.line(x_s, y_s, line_width=1, color=color, legend_label=f"sample {count+1}")
show(p)
```

```python
print(f"We see here that our fit is much the same as for when l=1. In fact our optimal value is {gp.kernel_}")
```

    We see here that our fit is much the same as for when l=1. In fact our optimal value is RBF(length_scale=0.759)

## Changing the Kernel
***

Based on [this article.](https://www.cs.toronto.edu/~duvenaud/cookbook/)

```python

```

```python

```

```python

```

# Calculating the Greeks
***

In what follows, European option prices are generated from the Black-Scholes model.  We begin by building a GP call model, where the input is the underlying price. The delta is then derived and compared with the Black-Scholes (BS)
delta. The exercise is repeated, but using the volatility as the input instead of the underlying price. The vega of the GP is then derived and compared with the BS vega.

```python
import sys
if "../" not in sys.path:
    sys.path.append("../")

from BlackScholes import bsformula

import numpy as np
import scipy as sp

from sklearn import gaussian_process
import matplotlib.pyplot as plt
%matplotlib inline
```

## Black-Scholes Model
***
First, set the model parameters

```python
KC = 130    # Call strike
KP = 70     # Put strike
r = 0.002   # risk-free rate
sigma = 0.4 # implied volatility
T = 2.0     # Time to maturity
S0 = 100    # Underlying spot
```

```python
lb = 0   # lower bound on domain
ub = 300 # upper bound on domain
training_number = 100 # Number of training samples
testing_number = 500   # Number of testing samples
sigma_n = 1e-8 # additive noise in GP
```

Define the call and put prices using the BS model

```python
bsformula?
```

```python
call = lambda price, vol: bsformula(1, lb+(ub-lb)*price, KC, r, T, vol, 0)[0]
put = lambda price, vol: bsformula(-1, lb+(ub-lb)*x, KP, r, T, vol, 0)[0]
```

## Delta
***

The Delta ($\Delta$) of an option is the first derivative of the option value with respect to the underlying asset price.
If we hold $-\Delta$ units of the asset (to every 1 unit of the option), we are in a delta neutral position. This means that if we are in a Black - Merton - Scholes financial market (fixed volatility) then price changes of the underlying asset do not matter and we have perfectly hedged away our risk. It is therefore a very useful quantity to know. Furthermore, even if we want to calculate more complicated hedging strategies, a delta-gamma hedge for example, then we need to know the delta.

How do we calculate the delta of an option?

We are going to use equation (1). If we take the derivative of equation (1) at a point $X_*$ we get

\begin{align}
	\frac{\partial}{\partial X_*}\mathbb{E}[f_*| X,Y,x_*] =&~ \frac{\partial}{\partial X_*}\mu(X_*) + \left(\frac{\partial}{\partial X_*} K_{X_*,X}\right) [K_{X,X} + \sigma^2 I]^{-1} Y,
    \\
    \text{where } \left(\frac{\partial}{\partial X_*} K_{X_*,X}\right)=&~ \frac{1}{l^2} (X- X_*) K_{X_*,X}.
\end{align}
Note that $\alpha = [K_{X,X} + \sigma^2 I ]$ is already calculated at training time (for pricing) by Cholesky matrix factorization with $O(n^3)$ complexity, so there is no significant computational overhead from Greeking.

1. Generate the training and testing data, where the input is the gridded underlying and the output are the option prices.

```python
x_train = np.array(np.linspace(0.01, 1.2, training_number), dtype='float32').reshape(training_number, 1)
x_test = np.array(np.linspace(0.01, 1.0, testing_number), dtype='float32').reshape(testing_number, 1)

y_train = []
for idx in range(len(x_train)):
    y_train.append(call(x_train[idx], sigma))
y_train = np.array(y_train)
```

2. Fit the GP model to the generated data

```python
sk_kernel = RBF(length_scale=1.0, length_scale_bounds=(0.01, 10000.0))
gp = gaussian_process.GaussianProcessRegressor(kernel=sk_kernel, n_restarts_optimizer=20, normalize_y=True)
gp.fit(x_train, y_train)   # Here, n_restarts > 1 means we keep restarting the algo. This helps avoid local optima.
```

    C:\Users\Josep\anaconda3\envs\computational-finance\lib\site-packages\sklearn\gaussian_process\_gpr.py:509: ConvergenceWarning: lbfgs failed to converge (status=2):
    ABNORMAL_TERMINATION_IN_LNSRCH.

    Increase the number of iterations (max_iter) or scale the data as shown in:
        https://scikit-learn.org/stable/modules/preprocessing.html
      _check_optimize_result("lbfgs", opt_res)

    GaussianProcessRegressor(kernel=RBF(length_scale=1), n_restarts_optimizer=20,
                             normalize_y=True)

3. Get the model's predicted outputs for each of the test inputs

```python
y_pred, sigma_hat = gp.predict(x_test, return_std=True)
```

    C:\Users\Josep\anaconda3\envs\computational-finance\lib\site-packages\sklearn\gaussian_process\_gpr.py:375: UserWarning: Predicted variances smaller than 0. Setting those variances to 0.
      warnings.warn("Predicted variances smaller than 0. "

4. Derive the GP delta

```python
l = gp.kernel_.length_scale
rbf = gaussian_process.kernels.RBF(length_scale=l)
Kernel = rbf(x_train, x_train) # Calculate K_{X,X}
```

```python
K_y = Kernel + np.eye(training_number) * sigma_n
L = sp.linalg.cho_factor(K_y)  # Computes the Cholesky decomposition, the matric L s.t. L L* = A
alpha = sp.linalg.cho_solve(np.transpose(L), y_train) # Solve the linear equations A x = y_train,
                                                        # given the Cholesky factorization of A.
```

    C:\Users\Josep\anaconda3\envs\computational-finance\lib\site-packages\numpy\core\_asarray.py:83: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray
      return array(a, dtype, copy=False, order=order)

```python
k_s = rbf(x_test, x_train)     # Calculate K_{X,X_*}
k_s_prime = (x_train.T - x_test) * k_s / l**2  # This is just the partial derivative of K_{X,X_*} w.r.t X_*
gp_delta_vec = np.dot(k_s_prime, alpha) / (ub - lb)
```

5. Compare the GP delta with the BS delta

```python
delta = lambda x, y: bsformula(1, lb+(ub-lb)*x, KC, r, T, y, 0)[1]
```

```python
plt.figure(figsize = (10,6),facecolor='white', edgecolor='black')
plt.plot(lb+(ub-lb)*x_test, delta(x_test,sigma), color = 'black', label = 'Exact')
plt.plot(lb+(ub-lb)*x_test, gp_delta_vec, color = 'red', label = 'GP')
plt.grid(True)
plt.xlabel('S')
plt.ylabel('$\Delta$')
plt.legend(loc = 'best', prop={'size':10});
```

6. Show the error between the GP delta and the BS delta

```python
plt.figure(figsize = (10,6),facecolor='white', edgecolor='black')
plt.plot(lb+(ub-lb)*x_test, delta(x_test,sigma) - gp_delta_vec, color = 'black', label = 'GP Error')
plt.grid(True)
plt.xlabel('S')
plt.ylabel('Error in $\Delta$')
plt.legend(loc = 'best', prop={'size':10});
```

**This is insanely good. We have 10 data points, and the error is essentially zero. The exception is at the lower end, outside the scope of seen data points.**

## Gamma
***

Gamma, $\Gamma$, measures the rate of change in the delta with respect to changes in the underlying price. Gamma is the second derivative of the value function with respect to the underlying price.

Similarly to before, we take the second derivative of the posterior mean with respect to $X_*$. This gives us
\begin{align}
	\frac{\partial^2}{\partial X_*^2}\mathbb{E}[f_*| X,Y,x_*] = \frac{1}{l^2} \left(\frac{(X- X_*)^2}{l^2} - \boldsymbol{1} \right) K_{X_*,X}.
\end{align}

```python
k_s_double_prime = ((x_train.T - x_test)**2 / l**2 - np.ones_like(x_test))* k_s / l**2
```

```python
gp_gamma_vec = np.dot(k_s_double_prime, alpha) / (ub - lb)
```

```python
# create figure and axis objects with subplots()
fig,ax = plt.subplots(figsize = (10,6),facecolor='white', edgecolor='black')
# make a plot
ax.plot(lb+(ub-lb)*x_test, gp_delta_vec, label = 'Delta')
# set x-axis label
ax.set_xlabel("Stock price",fontsize=14)
# set y-axis label
ax.set_ylabel('$\Delta$',fontsize=14)
plt.legend(loc = 'upper right', prop={'size':10});
ax2=ax.twinx()
# make a plot with different y-axis using second axis object
ax2.plot(lb+(ub-lb)*x_test, gp_gamma_vec, color = 'black', label = 'Gamma')
ax2.set_ylabel("$\Gamma$",color="black",fontsize=14)
plt.legend(loc = 'lower right', prop={'size':10});
plt.show()
```

## Vega
***
Vega measures sensitivity to volatility. Vega is the derivative of the option value with respect to the volatility of the underlying asset.

We now calculate the vega. This is exactly the same formula as for delta, except for the fact that we fit the GP to the data observed at different values of the volatility $\sigma$.

1. Generate the training and testing data, where the input is the gridded underlying and the output are the option prices. The inputs are again scaled to the unit domain.

```python
x_train = np.array(np.linspace(0.01, 1.2, training_number), dtype='float32').reshape(training_number, 1)
x_test = np.array(np.linspace(0.01, 1.0, testing_number), dtype='float32').reshape(testing_number, 1)
y_train = []
for idx in range(len(x_train)):
    y_train.append(call((S0-lb)/(ub-lb), x_train[idx]))
y_train = np.array(y_train)
```

2. Fit the GP model to the generated data

```python
sk_kernel = RBF(length_scale=1.0, length_scale_bounds=(0.01, 10000.0))
gp = gaussian_process.GaussianProcessRegressor(kernel=sk_kernel, n_restarts_optimizer=20, normalize_y=True)
gp.fit(x_train, y_train)
```

    GaussianProcessRegressor(kernel=RBF(length_scale=1), n_restarts_optimizer=20,
                             normalize_y=True)

3. Get the model's predicted outputs for each of the test inputs

```python
y_pred, sigma_hat = gp.predict(x_test, return_std=True)
```

    C:\Users\Josep\anaconda3\envs\computational-finance\lib\site-packages\sklearn\gaussian_process\_gpr.py:375: UserWarning: Predicted variances smaller than 0. Setting those variances to 0.
      warnings.warn("Predicted variances smaller than 0. "

4. Derive the GP vega

```python
l = gp.kernel_.length_scale
rbf = gaussian_process.kernels.RBF(length_scale=l)

Kernel= rbf(x_train, x_train)
K_y = Kernel + np.eye(training_number) * sigma_n
L = sp.linalg.cho_factor(K_y)
alpha = sp.linalg.cho_solve(np.transpose(L), y_train)
k_s = rbf(x_test, x_train)
```

    C:\Users\Josep\anaconda3\envs\computational-finance\lib\site-packages\numpy\core\_asarray.py:83: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray
      return array(a, dtype, copy=False, order=order)

```python
# k_s_prime = np.zeros([len(x_test), len(x_train)])
# for i in range(len(x_test)):
#     for j in range(len(x_train)):
#         k_s_prime[i, j] = (1.0/l**2) * (x_train[j] - x_test[i]) * k_s[i, j]

k_s_prime = (x_train.T - x_test) * k_s / l**2
gp_vega_vec = np.dot(k_s_prime, alpha)
```

5. Calculate the BS vega

```python
vega = lambda x, y: bsformula(1, lb + (ub-lb) * x, KC, r, T, y, 0)[2]
```

6. Compare the GP vega with the BS vega

```python
plt.figure(figsize = (10,6), facecolor='white', edgecolor='black')
plt.plot(x_test, vega((S0-lb)/(ub-lb), x_test), color = 'black', label = 'Exact')
plt.plot(x_test, gp_vega_vec, color = 'red', label = 'GP')
plt.grid(True)
plt.xlabel('$\\sigma$')
plt.ylabel('$\\nu$')
plt.legend(loc = 'best', prop={'size':10});
```

7. Plot the error between the GP vega and the BS vega

```python
plt.figure(figsize = (10,6), facecolor='white', edgecolor='black')
plt.plot(x_test, vega((S0-lb)/(ub-lb), x_test)-gp_vega_vec, color = 'black', label = 'GP Error')
plt.grid(True)
plt.xlabel('$\\sigma$')
plt.ylabel('Error in $\\nu$')
plt.legend(loc = 'best', prop={'size':10});
```

```python

```

## Resources
***
A lovely interactive explanation of Gausian processes: https://distill.pub/2019/visual-exploration-gaussian-processes/

Model comparison: https://ned.ipac.caltech.edu/level5/Sept13/Trotta/Trotta4.html

```python

```

Chapter3-GPs/extras/An intuitive introduction to Gaussian processes and using them to calculate the Greeks.ipynb


```python
import numpy as np
import pwlf
# generate sin wave data
x = np.linspace(0, 10, num=100)
y = np.sin(x * np.pi / 2)
# add noise to the data
y = np.random.normal(0, 0.05, 100) + y
my_pwlf_2 = pwlf.PiecewiseLinFit(x, y, degree=2)
res2 = my_pwlf_2.fitfast(5, pop=50)
```

```python
import matplotlib.pyplot as plt
%matplotlib inline
# first line
xhat = np.linspace(my_pwlf_2.fit_breaks[0], my_pwlf_2.fit_breaks[1], 100)
yhat = my_pwlf_2.beta[0] + my_pwlf_2.beta[1]*(xhat-my_pwlf_2.fit_breaks[0]) + my_pwlf_2.beta[6]*(xhat-my_pwlf_2.fit_breaks[0])**2
plt.figure()
plt.plot(x, y, 'o')
plt.plot(xhat, yhat, '-.', label='First segment')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.show()
```

```python
# second segment
xhat2 = np.linspace(my_pwlf_2.fit_breaks[1], my_pwlf_2.fit_breaks[2], 100)
yhat2 = (my_pwlf_2.beta[0] +
         (my_pwlf_2.beta[1])*(xhat2-my_pwlf_2.fit_breaks[0]) +
         (my_pwlf_2.beta[2])*(xhat2-my_pwlf_2.fit_breaks[1]) +
         (my_pwlf_2.beta[6])*(xhat2-my_pwlf_2.fit_breaks[0])**2 +
         (my_pwlf_2.beta[7])*(xhat2-my_pwlf_2.fit_breaks[1])**2)
plt.plot(x, y, 'o')
plt.plot(xhat, yhat, '-.', label='First segment')
plt.plot(xhat2, yhat2, '-.', label='Second segment')

plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.show()
```

```python
# Third segment
xhat3 = np.linspace(my_pwlf_2.fit_breaks[2], my_pwlf_2.fit_breaks[3], 100)
yhat3 = (my_pwlf_2.beta[0] +
         (my_pwlf_2.beta[1])*(xhat3-my_pwlf_2.fit_breaks[0]) +
         (my_pwlf_2.beta[2])*(xhat3-my_pwlf_2.fit_breaks[1]) +
         (my_pwlf_2.beta[3])*(xhat3-my_pwlf_2.fit_breaks[2]) +
         (my_pwlf_2.beta[6])*(xhat3-my_pwlf_2.fit_breaks[0])**2 +
         (my_pwlf_2.beta[7])*(xhat3-my_pwlf_2.fit_breaks[1])**2 +
         (my_pwlf_2.beta[8])*(xhat3-my_pwlf_2.fit_breaks[2])**2)
plt.plot(x, y, 'o')
plt.plot(xhat, yhat, '-.', label='First segment')
plt.plot(xhat2, yhat2, '-.', label='Second segment')
plt.plot(xhat3, yhat3, '-.', label='Third segment')

plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.show()
```

```python
# Fourth segment
xhat4 = np.linspace(my_pwlf_2.fit_breaks[3], my_pwlf_2.fit_breaks[4], 100)
yhat4 = (my_pwlf_2.beta[0] +
         (my_pwlf_2.beta[1])*(xhat4-my_pwlf_2.fit_breaks[0]) +
         (my_pwlf_2.beta[2])*(xhat4-my_pwlf_2.fit_breaks[1]) +
         (my_pwlf_2.beta[3])*(xhat4-my_pwlf_2.fit_breaks[2]) +
         (my_pwlf_2.beta[4])*(xhat4-my_pwlf_2.fit_breaks[3]) +
         (my_pwlf_2.beta[6])*(xhat4-my_pwlf_2.fit_breaks[0])**2 +
         (my_pwlf_2.beta[7])*(xhat4-my_pwlf_2.fit_breaks[1])**2 +
         (my_pwlf_2.beta[8])*(xhat4-my_pwlf_2.fit_breaks[2])**2 +
         (my_pwlf_2.beta[9])*(xhat4-my_pwlf_2.fit_breaks[3])**2)
plt.plot(x, y, 'o')
plt.plot(xhat, yhat, '-.', label='First')
plt.plot(xhat2, yhat2, '-.', label='Second')
plt.plot(xhat3, yhat3, '-.', label='Third')
plt.plot(xhat4, yhat4, '-.', label='Fourth')

plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.show()
```

```python
# Fifth segment
xhat5 = np.linspace(my_pwlf_2.fit_breaks[4], my_pwlf_2.fit_breaks[5], 100)
yhat5 = (my_pwlf_2.beta[0] +
         (my_pwlf_2.beta[1])*(xhat5-my_pwlf_2.fit_breaks[0]) +
         (my_pwlf_2.beta[2])*(xhat5-my_pwlf_2.fit_breaks[1]) +
         (my_pwlf_2.beta[3])*(xhat5-my_pwlf_2.fit_breaks[2]) +
         (my_pwlf_2.beta[4])*(xhat5-my_pwlf_2.fit_breaks[3]) +
         (my_pwlf_2.beta[5])*(xhat5-my_pwlf_2.fit_breaks[4]) +
         (my_pwlf_2.beta[6])*(xhat5-my_pwlf_2.fit_breaks[0])**2 +
         (my_pwlf_2.beta[7])*(xhat5-my_pwlf_2.fit_breaks[1])**2 +
         (my_pwlf_2.beta[8])*(xhat5-my_pwlf_2.fit_breaks[2])**2 +
         (my_pwlf_2.beta[9])*(xhat5-my_pwlf_2.fit_breaks[3])**2 +
         (my_pwlf_2.beta[10])*(xhat5-my_pwlf_2.fit_breaks[4])**2)
plt.plot(x, y, 'o')
plt.plot(xhat, yhat, '-.', label='First')
plt.plot(xhat2, yhat2, '-.', label='Second')
plt.plot(xhat3, yhat3, '-.', label='Third')
plt.plot(xhat4, yhat4, '-.', label='Fourth')
plt.plot(xhat5, yhat5, '-.', label='Fifth')

plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.show()
```

```python
from sympy import Symbol
from sympy.utilities import lambdify
x = Symbol('x')

def get_symbolic_eqn(pwlf_, segment_number):
    if pwlf_.degree < 1:
        raise ValueError('Degree must be at least 1')
    if segment_number < 1 or segment_number > pwlf_.n_segments:
        raise ValueError('segment_number not possible')
    # assemble degree = 1 first
    for line in range(segment_number):
        if line == 0:
            my_eqn = pwlf_.beta[0] + (pwlf_.beta[1])*(x-pwlf_.fit_breaks[0])
        else:
            my_eqn += (pwlf_.beta[line+1])*(x-pwlf_.fit_breaks[line])
    # assemble all other degrees
    if pwlf_.degree > 1:
        for k in range(2, pwlf_.degree + 1):
            for line in range(segment_number):
                beta_index = pwlf_.n_segments*(k-1) + line + 1
                my_eqn += (pwlf_.beta[beta_index])*(x-pwlf_.fit_breaks[line])**k
    return my_eqn.simplify()

eqn1 = get_symbolic_eqn(my_pwlf_2, 1)

```

```python
eqn_list = []
f_list = []
for i in range(my_pwlf_2.n_segments):
    eqn_list.append(get_symbolic_eqn(my_pwlf_2, i + 1))
    print('Equation number: ', i + 1)
    print(eqn_list[-1])
    f_list.append(lambdify(x, eqn_list[-1]))
```

    Equation number:  1
    -0.953964059782599*x**2 + 1.89945177490653*x + 0.00538634182565454
    Equation number:  2
    0.951561315686298*x**2 - 5.69747505830914*x + 7.5772216545711
    Equation number:  3
    -0.949735350431857*x**2 + 9.48218236957122*x - 22.720785454735
    Equation number:  4
    0.926850298824217*x**2 - 12.9824424358344*x + 44.5102742956827
    Equation number:  5
    -1.03016230425747*x**2 + 18.5306546317065*x - 82.3508513333073

```python
yhat1 = f_list[0](xhat)
yhat2 = f_list[1](xhat2)
yhat3 = f_list[2](xhat3)
yhat4 = f_list[3](xhat4)
yhat5 = f_list[4](xhat5)
xtemp = np.linspace(0,10, 1000)
plt.plot(my_pwlf_2.x_data, my_pwlf_2.y_data, 'o')
plt.plot(xtemp, my_pwlf_2.predict(xtemp), '-k', label='Predict')
plt.plot(xhat, yhat1, '-.', label='First')
plt.plot(xhat2, yhat2, '-.', label='Second')
plt.plot(xhat3, yhat3, '-.', label='Third')
plt.plot(xhat4, yhat4, '-.', label='Fourth')
plt.plot(xhat5, yhat5, '-.', label='Fifth')

plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.show()
```


```python
import torch
import torch.distributions as td
import torch.optim as opt
import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as stats
```

In this notebook I'll try to perform numerical variational bayesian analysis that previously I've done with **Pyro** package.

```python
# In the meantime I'll define thin wrappers around the probability distributions
class Bernoulli:
    def __init__(self, p):
        self.p = p

    def sample(self, size=1):
        return stats.bernoulli.rvs(p=self.p, size=size)

class Uniform:
    def __init__(self, start, end):
        self.start = start
        self.end = end

    def sample(self, size=1):
        return stats.uniform.rvs(loc=self.start, scale=self.end-self.start, size=size)

    def pdf(self, x):
        return stats.uniform.pdf(x, loc=self.start, scale=self.end-self.start)

    def mean(self):
        return stats.uniform.mean(loc=self.start, scale=self.end-self.start)

class Beta:
    def __init__(self, alpha, beta):
        self.alpha = alpha
        self.beta = beta

    def pdf(self, X):
        return stats.beta.pdf(X, a=self.alpha, b=self.beta)

    def mean(self):
        return stats.beta.mean(a=self.alpha, b=self.beta)

class Normal:
    def __init__(self, mu, sigma):
        self.mu = mu
        self.sigma = sigma

    def pdf(self, X):
        return stats.norm.pdf(X, loc=self.mu, scale=self.sigma)

    def sample(self, size=1):
        return stats.norm.rvs(loc=self.mu, scale=self.sigma, size=size)

    def mean(self):
        return self.mu
```

Let's consider a simple example, where:
- Prior $\mathbb{P}(\theta) \sim U(0, 1)$
- Likelihood $\mathbb{P}(X | \theta) \sim B(\theta)$

```python
N = 100
```

```python
Prior = Uniform(0, 1)
```

```python
hidden_theta = Prior.sample()[0]
hidden_theta
```

    0.9510983673224221

```python
Likelihood = Bernoulli(hidden_theta)
```

```python
X = Likelihood.sample(size=N)
```

```python
X
```

    array([1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
           1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1,
           1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
           0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

The whole point of **variational bayesian inference**, is to fit a **variational distribution** $q_\phi$ parametrized by a set of **variational parameters** $\phi$ which is meant to approximate the true posterior $p(\theta | x)$.

Approximation is done by minimizing KL divergence between $q_\phi(\theta)$ and $p(\theta | x)$

The good thing is, we don't need to know the posterior, as long as we can calculate gradient of a KL divergence between $q$ and $p$ with respect to variational parameters $\phi$.

\begin{align}
D_{KL}(q_\phi(\theta) || p(\theta | x))
&= \int_\Theta q(\theta) \log \frac{q(\theta)}{p(\theta|x)} d\theta \\
&= \mathbb{E}_q \log \frac{q(\theta)}{p(\theta|x)} \\
&= \mathbb{E}_q \log q(\theta) - \mathbb{E}_q \log p(\theta|x) \\
&= \mathbb{E}_q \log q(\theta) - \mathbb{E}_q \log \frac{p(\theta, x)}{p(x)} \\
&= \log p(x) - \left( \mathbb{E}_q \log p(\theta, x) - \mathbb{E}_q \log q(\theta) \right)\\
\end{align}

The second bit of that is called **ELBO**, which stands for **Evidence Lower BOund**. **Evidence** usually means the probability of data given the model, which is
$$
p(x) = \int_\Theta p(x | \theta) p(\theta) d\theta
$$

We assume the model a the beginning of our inference so the value of $p(x)$ is given and does not depend on any parameters we consider.

Because KL divergence is always positive, **ELBO** has to be lower or equal to the logarithm of **evidence**. Thus maximizing the **ELBO** is our goal in computational variational inference, which should result in distribution $q_\phi$ that will approximate closely posterior probability distribution $p(\theta|x)$.

```python
# Define our paramters
# We'll parametrize a Beta distribution
a = torch.tensor(2.0, dtype=torch.float32, requires_grad=True)
b = torch.tensor(2.0, dtype=torch.float32, requires_grad=True)
print(a, b)
```

    tensor(2., requires_grad=True) tensor(2., requires_grad=True)

```python
# Prior model defined using PyTorch distributions
TPrior = td.Uniform(0, 1)
```

```python
X_tensor = torch.from_numpy(X)
print(X_tensor)
```

    tensor([1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
            1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1,
            1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1,
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
            1, 1, 1, 1])

```python
def calculate_elbo():
    # Our variational distribution
    q_distribution = td.Beta(a, b)
    # All samples are calculated with respect to q distribution, therefore we sample from q
    theta_sample = q_distribution.rsample()

    # First bit, model log prob assuming theta is the hidden param
    log_prob_theta_model = TPrior.log_prob(theta_sample)
    log_prob_data_model = td.Bernoulli(theta_sample).log_prob(X_tensor.to(torch.float32)).sum()

    # Second bit, variational distribution log prob, assuming theta is the hidden param
    log_prob_q = q_distribution.log_prob(theta_sample)

    return log_prob_theta_model + log_prob_data_model - log_prob_q
```

```python
OPTIMIZATION_STEPS = 30_000
```

```python
optimizer = opt.Adam([a, b], lr=0.1)
```

```python
# metrics
elbo_history = []
a_history = []
b_history = []
```

```python
for i in range(OPTIMIZATION_STEPS):
    optimizer.zero_grad()

    elbo = calculate_elbo()
    loss = -elbo
    loss.backward()

    elbo_history.append(elbo.item())
    a_history.append(a.item())
    b_history.append(b.item())

    optimizer.step()

```

```python
plt.plot(elbo_history)
plt.title("ELBO")
plt.xlabel("step")
plt.ylabel("loss");
```

```python
print("a=", a.item())
print("b=", b.item())
```

    a= 102.12264251708984
    b= 6.05571174621582

```python
successes = X_tensor.sum().item()
failures = (1 - X_tensor).sum().item()
```

```python
print("s =", successes)
print("f =", failures)
```

    s = 95
    f = 5

```python
plt.subplot(1,2,1)
plt.plot([0,OPTIMIZATION_STEPS],[successes + 1, successes + 1], 'k:')
plt.plot(a_history)
plt.ylabel('a')

plt.subplot(1,2,2)
plt.ylabel('b')
plt.plot([0,OPTIMIZATION_STEPS],[failures + 1, failures + 1], 'k:')
plt.plot(b_history)
plt.tight_layout()
```

```python
A = a.item()
B = b.item()
```

```python
MLE = (A - 1) / (A + B - 2)
```

```python
TruePosterior = Beta(X.sum().item() + 1, (1-X).sum().item() + 1)
EstimatedPosterior = Beta(A, B)
```

```python
fig, axs = plt.subplots(1, 1)
axs.set_title("Prior vs Posterior")
support = np.linspace(0.0, 1.0, 1000)

axs.plot(support, Prior.pdf(support), label="Prior")
axs.fill_between(support, 0, Prior.pdf(support), alpha=0.2)
axs.plot(support, TruePosterior.pdf(support), label="TruePosterior")
axs.fill_between(support, 0, TruePosterior.pdf(support), alpha=0.2)
axs.plot(support, EstimatedPosterior.pdf(support), label="EstimatedPosterior")
axs.fill_between(support, 0, EstimatedPosterior.pdf(support), alpha=0.2)
axs.axvline(hidden_theta, color='red', linestyle='--', label='True value of param')
axs.axvline(MLE, color='blue', linestyle='--', label='Maximum likelihood estimator')

axs.legend()
```


## Data Analytics

### Basic Univariate Statistics in Python

#### Michael Pyrcz, Associate Professor, University of Texas at Austin

##### [Twitter](https://twitter.com/geostatsguy) | [GitHub](https://github.com/GeostatsGuy) | [Website](http://michaelpyrcz.com) | [GoogleScholar](https://scholar.google.com/citations?user=QVZ20eQAAAAJ&hl=en&oi=ao) | [Book](https://www.amazon.com/Geostatistical-Reservoir-Modeling-Michael-Pyrcz/dp/0199731446) | [YouTube](https://www.youtube.com/channel/UCLqEr-xV-ceHdXXXrTId5ig)  | [LinkedIn](https://www.linkedin.com/in/michael-pyrcz-61a648a1)

### Data Analytics: Basic Univariate Statistics

Here's a demonstration of calculation of univariate statistics in Python. This demonstration is part of the resources that I include for my courses in Spatial / Subsurface Data Analytics and Geostatistics at the Cockrell School of Engineering and Jackson School of Goesciences at the University of Texas at Austin.

We will cover the following statistics:

#### Measures of Centrality
* Arithmetic Average / Mean
* Median
* Mode (most frequent binned)
* Geometric Mean
* Harmonic Mean
* Power Law Average

#### Measures of Dispersion
* Population Variance
* Sample Variance
* Population Standard Deviation
* Sample Standard Deviation
* Range
* Percentile w. Tail Assumptions
* Interquartile Range

#### Tukey Outlier Test
* Lower Quartile/P25
* Upper Quartile/P75
* Interquartile Range
* Lower Fence
* Upper Fence
* Calculating Outliers

#### Measures of Shape
* Skew
* Excess Kurtosis
* Pearson' Mode Skewness
* Quartile Skew Coefficient

#### Nonparmetric Cumulative Distribution Functions (CDFs)
* plotting a nonparametric CDF
* fitting a parametric distribution and plotting

I have a lecture on these univariate statistics available on [YouTube](https://www.youtube.com/watch?v=wAcbA2cIqec&list=PLG19vXLQHvSB-D4XKYieEku9GQMQyAzjJ&index=11&t=0s).

#### Getting Started

Here's the steps to get setup in Python with the GeostatsPy package:

1. Install Anaconda 3 on your machine (https://www.anaconda.com/download/).
2. From Anaconda Navigator (within Anaconda3 group), go to the environment tab, click on base (root) green arrow and open a terminal.
3. In the terminal type: pip install geostatspy.
4. Open Jupyter and in the top block get started by copy and pasting the code block below from this Jupyter Notebook to start using the geostatspy functionality.

You will need to copy the data file to your working directory.  The dataset is available on my GitHub account in my GeoDataSets repository at:

* Tabular data - [2D_MV_200wells.csv](https://github.com/GeostatsGuy/GeoDataSets/blob/master/2D_MV_200wells.csv)

#### Importing Packages

We will need some standard packages. These should have been installed with Anaconda 3.

```python
import numpy as np                        # ndarrys for gridded data
import pandas as pd                       # DataFrames for tabular data
import os                                 # set working directory, run executables
import matplotlib.pyplot as plt           # plotting
from scipy import stats                   # summary statistics
import scipy                              # statistics
import statistics as stats                # statistics like the mode
from scipy.stats import norm              # fitting a Gaussian distribution
```

#### Set the Working Directory

I always like to do this so I don't lose files and to simplify subsequent read and writes (avoid including the full address each time). Set this to your working directory, with the above mentioned data file.

```python
os.chdir("d:/PGE383")                     # set the working directory
```

#### Loading Data

Let's load the provided multivariate, spatial dataset.  '2D_MV_200wells.csv' is available at https://github.com/GeostatsGuy/GeoDataSets.  It is a comma delimited file with X and Y coordinates,facies 1 and 2 (1 is sandstone and 2 interbedded sand and mudstone), porosity (fraction), permeability (mDarcy) and acoustic impedance (kg/m2s*10^6). We load it with the pandas 'read_csv' function into a data frame we called 'df' and then preview it by printing a slice and by utilizing the 'head' DataFrame member function (with a nice and clean format, see below).

```python
df = pd.read_csv("2D_MV_200wells.csv")                      # read a .csv file in as a DataFrame
#print(df.iloc[0:5,:])                                       # display first 4 samples in the table as a preview
df.head()                                                   # we could also use this command for a table preview
```

      X
      Y
      facies_threshold_0.3
      porosity
      permeability
      acoustic_impedance

      0
      565
      1485
      1
      0.1184
      6.170
      2.009

      1
      2585
      1185
      1
      0.1566
      6.275
      2.864

      2
      2065
      2865
      2
      0.1920
      92.297
      3.524

      3
      3575
      2655
      1
      0.1621
      9.048
      2.157

      4
      1835
      35
      1
      0.1766
      7.123
      3.979

Let's extract one of the features, porosity, into a 1D ndarray and do our statistics on porosity.

* then we can use NumPy's statistics methods

```python
por = df['porosity'].values
```

Now let's go through all the univariate statistics listed above one-by-one.

#### Measures of Central Tendency

Let's start with measures of central tendency.

##### The Arithmetic Average / Mean

\begin{equation}
\overline{x} = \frac{1}{n}\sum^n_{i=1} x_i
\end{equation}

```python
por_average = np.average(por)
print('Porosity average is ' + str(round(por_average,2)) + '.')
```

    Porosity average is 0.15.

##### Median

\begin{equation}
P50_x = F^{-1}_{x}(0.50)
\end{equation}

```python
por_median = np.median(por)
print('Porosity median is ' + str(round(por_median,2)) + '.')
```

    Porosity median is 0.15.

##### Mode

The most common value. To do this we should bin the data, like into histogram bins/bars.  To do this we will round the data to the 2nd decimal place.  We are assume bin boundaries, $0.01, 0.02,\ldots, 0.30$.

```python
por_mode = stats.mode(np.round(por,2))
print('Porosity mode is ' + str(round(por_mode,2)) + '.')
```

    Porosity mode is 0.14.

##### Geometric Mean

\begin{equation}
\overline{x}_G = ( \prod^n_{i=1} x_i )^{\frac{1}{n}}
\end{equation}

```python
por_geometric = scipy.stats.mstats.gmean(por)
print('Porosity geometric mean is ' + str(round(por_geometric,2)) + '.')
```

    Porosity geometric mean is 0.15.

##### Harmonic Mean

\begin{equation}
\overline{x}_H = \frac{n}{\sum^n_{i=1} \frac{1}{x_i}}
\end{equation}

```python
por_hmean = scipy.stats.mstats.hmean(por)
print('Porosity harmonic mean is ' + str(round(por_hmean,2)) + '.')
```

    Porosity harmonic mean is 0.14.

##### Power Law Average

\begin{equation}
\overline{x}_p = (\frac{1}{n}\sum^n_{i=1}{x_i^{p}})^\frac{1}{p}
\end{equation}

```python
power = 1.0
por_power = np.average(np.power(por,power))**(1/power)
print('Porosity law mean for p = ' + str(power) + ' is ' + str(round(por_power,2)) + '.')
```

    Porosity law mean for p = 1.0 is 0.15.

#### Measures of Dispersion

##### Population Variance

\begin{equation}
\sigma^2_{x} = \frac{1}{n}\sum^n_{i=1}(x_i - \mu)
\end{equation}

```python
por_varp = stats.pvariance(por)
print('Porosity population variance is ' + str(round(por_varp,4)) + '.')
```

    Porosity population variance is 0.0011.

##### Sample Variance

\begin{equation}
\sigma^2_{x} = \frac{1}{n-1}\sum^n_{i=1}(x_i - \overline{x})^2
\end{equation}

```python
por_var = stats.variance(por)
print('Porosity sample variance is ' + str(round(por_var,4)) + '.')
```

    Porosity sample variance is 0.0011.

##### Population Standard Deviation

\begin{equation}
\sigma_{x} = \sqrt{ \frac{1}{n}\sum^n_{i=1}(x_i - \mu)^2 }
\end{equation}

```python
por_stdp = stats.pstdev(por)
print('Porosity sample variance is ' + str(round(por_stdp,4)) + '.')
```

    Porosity sample variance is 0.0329.

##### Sample Standard Deviation

\begin{equation}
\sigma_{x} = \sqrt{ \frac{1}{n-1}\sum^n_{i=1}(x_i - \mu)^2 }
\end{equation}

```python
por_std = stats.stdev(por)
print('Porosity sample variance is ' + str(round(por_std,4)) + '.')
```

    Porosity sample variance is 0.0329.

##### Range

\begin{equation}
range_x = P100_x - P00_x
\end{equation}

```python
por_range = np.max(por) - np.min(por)
print('Porosity range is ' + str(round(por_range,2)) + '.')
```

    Porosity range is 0.17.

##### Percentile

\begin{equation}
P(p)_x = F^{-1}_{x}(p)
\end{equation}

```python
p_value = 13
por_percentile = np.percentile(por,p_value)
print('Porosity ' + str(int(p_value)) + 'th percentile is ' + str(round(por_percentile,2)) + '.')
```

    Porosity 13th percentile is 0.11.

##### Inter Quartile Range

\begin{equation}
IQR = P(0.75)_x - P(0.25)_x
\end{equation}

```python
por_iqr = scipy.stats.iqr(por)
print('Porosity interquartile range is ' + str(round(por_iqr,2)) + '.')
```

    Porosity interquartile range is 0.04.

#### Tukey Test for Outliers

Let's demonstrate the Tukey test for outliers based on the lower and upper fences.

\begin{equation}
fence_{lower} = P_x(0.25) - 1.5 \times [P_x(0.75) - P_x(0.25)]
\end{equation}

\begin{equation}
fence_{upper} = P_x(0.75) + 1.5 \times [P_x(0.75) - P_x(0.25)]
\end{equation}

Then we declare samples values above the upper fence or below the lower fence as outliers.

```python
p25, p75 = np.percentile(por, [25, 75])
lower_fence = p25 - por_iqr * 1.5
upper_fence = p75 + por_iqr * 1.5
por_outliers = por[np.where((por > upper_fence) | (por < lower_fence))[0]]
print('Porosity outliers by Tukey test include ' + str(por_outliers) + '.')
por_outliers_indices = np.where((por > upper_fence) | (por < lower_fence))[0]
print('Porosity outlier indices by Tukey test are ' + str(por_outliers_indices) + '.')
```

    Porosity outliers by Tukey test include [0.06726 0.05    0.06092].
    Porosity outlier indices by Tukey test are [110 152 198].

#### Measures of Shape

##### Pearson's Mode Skewness

\begin{equation}
skew = \frac{3 (\overline{x} - P50_x)}{\sigma_x}
\end{equation}

```python
por_skew = (por_average - por_median)/por_std
print('Porosity skew is ' + str(round(por_skew,2)) + '.')
```

    Porosity skew is -0.03.

##### Population Skew, 3rd Central Moment

\begin{equation}
\gamma_{x} = \frac{1}{n}\sum^n_{i=1}(x_i - \mu)^3
\end{equation}

```python
por_cm = scipy.stats.moment(por,moment=3)
print('Porosity 3rd cenral moment is ' + str(round(por_cm,7)) + '.')
```

    Porosity 3rd cenral moment is -1.22e-05.

##### Quartile Skew Coefficient

\begin{equation}
QS = \frac{(P75_x - P50_x) - (P50_x - P25_x)}{(P75_x - P25_x)}
\end{equation}

```python
por_qs = ((np.percentile(por,75)-np.percentile(por,50))
          -(np.percentile(por,50)-np.percentile(por,25))) /((np.percentile(por,75))-np.percentile(por,25))
print('Porosity quartile skew coefficient is ' + str(round(por_qs,2)) + '.')
```

    Porosity quartile skew coefficient is 0.14.

#### Plot the Nonparametric CDF

Let's demonstrate plotting a nonparametric cumulative distribution function (CDF) in Python

```python
# sort the data:
por_sort = np.sort(por)

# calculate the cumulative probabilities assuming known tails
p = np.arange(len(por)) / (len(por) - 1)

# plot the cumulative probabilities vs. the sorted porosity values
plt.subplot(122)
plt.scatter(por_sort, p, c = 'red', edgecolors = 'black', s = 10, alpha = 0.7)
plt.xlabel('Porosity (fraction)'); plt.ylabel('Cumulative Probability'); plt.grid();
plt.title('Nonparametric Porosity CDF')
plt.ylim([0,1]); plt.xlim([0,0.25])
plt.subplots_adjust(left=0.0, bottom=0.0, right=2.0, top=1.2, wspace=0.2, hspace=0.3)
```

#### Fit a Gaussian Distribution

Let's fit a Gaussian distribution

* we get fancy with Maximuum Likelihood Estimation (MLE) for the Gaussian parametric distribution fit mean and standard deviation

```python
por_values = np.linspace(0.0,0.25,100)
fit_mean, fit_stdev = norm.fit(por,loc = por_average, scale = por_std) # fit MLE of the distribution
cumul_p = norm.cdf(por_values, loc = fit_mean, scale = fit_stdev)

# plot the cumulative probabilities vs. the sorted porosity values
plt.subplot(122)
plt.scatter(por_sort, p, c = 'red', edgecolors = 'black', s = 10, alpha = 0.7)
plt.plot(por_values,cumul_p, c = 'black')
plt.xlabel('Porosity (fraction)'); plt.ylabel('Cumulative Probability'); plt.grid();
plt.title('Nonparametric Porosity CDF')
plt.ylim([0,1]); plt.xlim([0,0.25])
plt.subplots_adjust(left=0.0, bottom=0.0, right=2.0, top=1.2, wspace=0.2, hspace=0.3)
```

#### Comments

This was a basic demonstration of univariate statistics in Python.

I have other demonstrations on the basics of working with DataFrames, ndarrays, univariate statistics, plotting data, declustering, data transformations, trend modeling and many other workflows available at [Python Demos](https://github.com/GeostatsGuy/PythonNumericalDemos) and a Python package for data analytics and geostatistics at [GeostatsPy](https://github.com/GeostatsGuy/GeostatsPy).

I hope this was helpful,

*Michael*

#### The Author:

### Michael Pyrcz, Associate Professor, University of Texas at Austin
*Novel Data Analytics, Geostatistics and Machine Learning Subsurface Solutions*

With over 17 years of experience in subsurface consulting, research and development, Michael has returned to academia driven by his passion for teaching and enthusiasm for enhancing engineers' and geoscientists' impact in subsurface resource development.

For more about Michael check out these links:

#### [Twitter](https://twitter.com/geostatsguy) | [GitHub](https://github.com/GeostatsGuy) | [Website](http://michaelpyrcz.com) | [GoogleScholar](https://scholar.google.com/citations?user=QVZ20eQAAAAJ&hl=en&oi=ao) | [Book](https://www.amazon.com/Geostatistical-Reservoir-Modeling-Michael-Pyrcz/dp/0199731446) | [YouTube](https://www.youtube.com/channel/UCLqEr-xV-ceHdXXXrTId5ig)  | [LinkedIn](https://www.linkedin.com/in/michael-pyrcz-61a648a1)

#### Want to Work Together?

I hope this content is helpful to those that want to learn more about subsurface modeling, data analytics and machine learning. Students and working professionals are welcome to participate.

* Want to invite me to visit your company for training, mentoring, project review, workflow design and / or consulting? I'd be happy to drop by and work with you!

* Interested in partnering, supporting my graduate student research or my Subsurface Data Analytics and Machine Learning consortium (co-PIs including Profs. Foster, Torres-Verdin and van Oort)? My research combines data analytics, stochastic modeling and machine learning theory with practice to develop novel methods and workflows to add value. We are solving challenging subsurface problems!

* I can be reached at mpyrcz@austin.utexas.edu.

I'm always happy to discuss,

*Michael*

Michael Pyrcz, Ph.D., P.Eng. Associate Professor The Hildebrand Department of Petroleum and Geosystems Engineering, Bureau of Economic Geology, The Jackson School of Geosciences, The University of Texas at Austin

#### More Resources Available at: [Twitter](https://twitter.com/geostatsguy) | [GitHub](https://github.com/GeostatsGuy) | [Website](http://michaelpyrcz.com) | [GoogleScholar](https://scholar.google.com/citations?user=QVZ20eQAAAAJ&hl=en&oi=ao) | [Book](https://www.amazon.com/Geostatistical-Reservoir-Modeling-Michael-Pyrcz/dp/0199731446) | [YouTube](https://www.youtube.com/channel/UCLqEr-xV-ceHdXXXrTId5ig)  | [LinkedIn](https://www.linkedin.com/in/michael-pyrcz-61a648a1)

```python

```


# Tutorial Brief

**Disclaimer:** Notebook taken from https://www.youtube.com/channel/UC98x5I1LVPhtnUHDyujq7zg

SymPy is symbolic mathematics library written completely in Python and doesn't require any dependencies.

Finding Help:

- http://docs.sympy.org/latest/index.html
- http://sympy.org/en/features.html
- http://nbviewer.ipython.org/github/jrjohansson/scientific-python-lectures/blob/master/Lecture-5-Sympy.ipynb

NumPy
 Base N-dimensional array package

SciPy
 Fundamental library for scientific computing

Matplotlib
 Comprehensive 2D Plotting

IPython
 Enhanced Interactive Console

SymPy
 Symbolic mathematics

Pandas
 Data structures & analysis

> ###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 and does not require any external libraries.
> **http://sympy.org/**

# Import SymPy

```python
from sympy import *
import math
```

```python
3 + math.sqrt(3)
```

    4.732050807568877

```python
expr = 3 * sqrt(3)
expr
```

    3*sqrt(3)

```python
init_printing(use_latex='mathjax')
```

```python
expr
```

    3*sqrt(3)

```python
expr = sqrt(8)
expr
```

    2*sqrt(2)

# `symbols()` & `Symbol()`

```python
x, y = symbols("x y")
```

```python
expr = x**2 + y**2
expr
```

```python
expr = (x+y)**3
expr
```

## Assumptions for symbols

```python
a = Symbol("a")
```

```python
a.is_imaginary
```

```python
b = Symbol("b", integer=True)
```

```python
b.is_imaginary
```

```python
c = Symbol("c", positive=True)
```

```python
c.is_positive
```

```python
c.is_imaginary
```

# Imaginary Numbers

```python
I
```

```python
I ** 2
```

# `Rational()`

```python
Rational(1,3)
```

```python
Rational(1,3) + Rational(1,2)
```

# Numerical evaluation

```python
expr = Rational(1,3) + Rational(1,2)
N(expr)
```

```python
N(pi, 100)
```

```python
pi.evalf(100)
```

# `subs()`

```python
expr = x**2 + 2*x + 1
expr
```

```python
expr.subs(x, 1)
```

```python
expr = pi * x**2
expr
```

```python
expr.subs(x, 3)
```

```python
N(_)
```

# `factor()` and `expand()`

```python
expr = (x + y) ** 2
expr
```

```python
expand(expr)
```

```python
factor(_)
```

# `simplify()`

```python
expr = (2*x + Rational(1,3)*x + 4) / x
expr
```

```python
simplify(expr)
```

```python
expr = "(2*x + 1/3*x + 4)/x"
simplify(expr)
```

```python
expr = sin(x)/cos(x)
expr
```

```python
simplify(expr)
```

# `apart()` and `together()`

```python
expr = 1/(x**2 + 2*x)
expr
```

```python
apart(expr)
```

```python
together(_)
```

# Calculus

```python
diff(sin(x), x)
```

```python
diff(log(x**2 + 1) + 2*x, x)
```

```python
integrate(cos(x), x)
```

```python
Integral(sin(x), (x,0,pi))
```

```python
N(_)
```

# `Sum()`

```python
expr = Sum(1/(x**2 + 2*x), (x, 1, 10))
expr
```

```python
expr.doit()
```

# `Product()`

```python
expr = Product(1/(x**2 + 2*x), (x, 1, 10))
expr
```

```python
expr.doit()
```

# `Solve()`

```python
expr = 2*x + 1
solve(expr)
```

```python
expr = x**2 - 1
solve(expr)
```

```python
expr_1 = 2*x + y + 3
expr_2 = 2*y - x
solve([expr_1, expr_2],(x,y))
```

# Units

```python
from sympy.physics import units as u
```

```python
5. * u.milligram
```

```python
1./2 * u.inch
```

```python
1. * u.nano
```

```python
u.watt
```

```python
u.ohm
```

## Converting from Kilometers/hours to Miles/hours

```python
kmph = u.km / u.hour
mph = u.mile / u.hour

N(mph / kmph)
```

```python
80 * N(mph / kmph)
```

# Working with NumPy / Pandas and Matplotlib

```python
def sympy_expr(x_val):
    expr = x**2 + sqrt(3)*x - Rational(1,3)
    return expr.subs(x, x_val)
```

```python
sympy_expr(3)
```

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

```python
list1 = np.arange(1,1000)
list2 = pd.Series(list1)
```

```python
%timeit [sympy_expr(item) for item in list1]
%timeit [sympy_expr(item) for item in list2]
```

```python
%timeit np.vectorize(sympy_expr)(list1)
%timeit list2.apply(sympy_expr)
```

```python
expr = x**2 + sqrt(3)*x - Rational(1,3)

lf = lambdify(x, expr)
```

```python
%timeit lf(list1)
%timeit lf(list2)
```

```python
fig = plt.figure()
axes = fig.add_subplot(111)

x_vals = np.linspace(-5.,5.)
y_vals = lf(x_vals)

axes.grid()
axes.plot(x_vals, y_vals)

plt.show();
```

```python

```


# Estimating MI by noise contrastive method (InfoNCE)

## Setup

The following cell imports the necessary packages:

```python
import torch
import torch.nn as nn
import os
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
import argparse
from torch.autograd import Variable
import itertools
from tqdm import tqdm
from data.mix_gaussian import MixedGaussian
from data.gaussian import Gaussian
from model.utils import *
from torch.utils.tensorboard import SummaryWriter
import datetime
```

Use GPU if GPU is available.

```python
cuda = True if torch.cuda.is_available() else False
FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor
torch.set_default_tensor_type(FloatTensor)

name = './results/InfoNCE'   # filename
chkpt_name = name+'.pt'      # checkpoint
writer = SummaryWriter('./results/log/InfoNCE')
```

Set random seed for reproducibility.

```python
SEED = 0
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
np.random.seed(SEED)
```

Specify the hyper-parameters:

```python
# arguments
parser = argparse.ArgumentParser()
parser.add_argument("--rho", type=float, default=0.9, help="coefficient of Gaussian")
parser.add_argument("--d", type=int, default=6, help="dimension of X & Y")
parser.add_argument("--sample_size", type=int, default=400, help="sample size")
parser.add_argument("--gamma", type=float, default=0, help="clipping parameter")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--n_iters", type=int, default=60000, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=100, help="size of the batches")
parser.add_argument("--lr", type=float, default=1e-4, help="adam: learning rate")
parser.add_argument("--hidden_dim", type=int, default=100, help="Hidden dimension")
parser.add_argument("--c_0_1_ratio", type=float, default=1, help="Ratio of samples with label 0 and samples with label 1 ")
parser.add_argument("--ma_rate", type=float, default=0.1, help="move average rate")
parser.add_argument("--ma_ef", type=float, default=1, help="move average ef")

opt, unknown = parser.parse_known_args()
```

Specify the distribution to generate the data samples:

```python
# Two choices here: 'Gaussian' and 'MixedGaussian'
density = 'Gaussian'
```

## Model

Define the function `generate_data` for generating Gaussian or Mixed Gaussian distributions.

- When generating Gaussian distrition, ${X}$ and ${Y}$ are distributed as

\begin{align}
p_{XY} = \prod_{i=1}^d \mathcal{N} \left( \mathbf{0}, \begin{bmatrix}1 & \rho \\ \rho & 1\\ \end{bmatrix} \right)
\end{align}

- When generating Mixed Gaussian distributions, ${X}$ and ${Y}$ are distributed as

\begin{align}
p_{XY} = \frac{1}{2} \prod_{i=1}^d \mathcal{N} \left( \mathbf{0}, \begin{bmatrix}1 & \rho \\ \rho & 1\\ \end{bmatrix} \right)
+ \frac{1}{2} \prod_{i=1}^d \mathcal{N} \left( \mathbf{0}, \begin{bmatrix}1 & {-} \rho \\ {-} \rho & 1\\ \end{bmatrix} \right)
\end{align}

```python
def generate_data(distribution='Gaussian', rho=0.9):
    # rho is the covariance for generating distributions
    # mu1 and mu2 are means for generating Mixed Gaussian distribution
    mu1 = 0
    mu2 = 0

    # mg is an object of Class Gaussian or MixedGaussian
    if distribution=='Gaussian':
        mg = Gaussian(sample_size=opt.sample_size,rho=rho)
    else:
        mg = MixedGaussian(sample_size=opt.sample_size, mean1=mu1, mean2=mu2,rho1=rho, rho2=-rho)

    # Calculate the ground truth MI between X and Y for (X, Y) from mg
    mi = mg.ground_truth * opt.d

    # Create X, Y for storing generated samples
    X = np.zeros((opt.sample_size,opt.d))
    Y = np.zeros((opt.sample_size,opt.d))

    # Generate samples of random variable X,Y and XY
    for j in range(opt.d):
        # In each iteration, mg.data will generate samples of two dimensions where one dimension for X and another for Y respectively
        data = mg.data
        X[:,j] = data[:,0]
        Y[:,j] = data[:,1]
    X = torch.Tensor(X)
    Y = torch.Tensor(Y)
    XY = torch.cat((X, Y), dim=1)

    return XY, X, Y, mi
```

Define the neural network.

```python
class Net(nn.Module):
    # Inner class that defines the neural network architecture
    def __init__(self, input_size=2, hidden_size=100, sigma=0.02):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, 1)
        nn.init.normal_(self.fc1.weight, std=sigma)
        nn.init.constant_(self.fc1.bias, 0)
        nn.init.normal_(self.fc2.weight, std=sigma)
        nn.init.constant_(self.fc2.bias, 0)
        nn.init.normal_(self.fc3.weight, std=sigma)
        nn.init.constant_(self.fc3.bias, 0)

    def forward(self, input):
        output = F.elu(self.fc1(input))
        output = F.elu(self.fc2(output))
        output = self.fc3(output)

        return output
```

The function `_resample` below is for resampling the given data samples for training the neural network.

```python
def _resample(data, batch_size, replace=False):
    # Resample the given data sample.
    index = np.random.choice(
        range(data.shape[0]), size=batch_size, replace=replace)
    batch = data[index]
    return batch
```

The function `infonce_loss` below returns the empirical estimate of the following infoNCE lower bound on the mutual information $I(X\wedge Y)$:

\begin{align}
I_{\text{infoNCE}}(f) := E\left[\log \frac{e^{f(X, Y)}}{E \left[\left. e^{f( X, Y')} \right| X \right]} \right] \leq I(X\wedge Y)
\end{align}

where $f: \mathcal{X}\times \mathcal{Y}\to \mathbb{R}$ is a function to be implemented by a neural network, and $p_{Y'|X,Y}(y'|x,y)=p_{Y}(y)$, i.e., $Y'$ is identically distributed as $Y$ and independent of $(X,Y)$.

```python
def infonce_loss(net, x_samples, y_samples):
    sample_size = x_samples.shape[0]
    x_tile = x_samples.unsqueeze(0).repeat((sample_size, 1, 1))
    y_tile = y_samples.unsqueeze(1).repeat((1, sample_size, 1))

    T0 = net(torch.cat([x_samples,y_samples], dim = -1))
    T1 = net(torch.cat([x_tile, y_tile], dim = -1))

    lower_bound = T0.mean() - (T1.logsumexp(dim = 1).mean() - np.log(sample_size))
    return lower_bound
```

The neural network `discriminator` below implements $f$:

```python
discriminator = Net(input_size=opt.d*2, hidden_size=100)

# move NN model to GPU if GPU is available
if cuda:
    discriminator.cuda()

# Adam optimizer
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
```

Load previous results.

```python
load_available = True # set to False to prevent loading previous results
if load_available and os.path.exists(chkpt_name):
    checkpoint = torch.load(
        chkpt_name, map_location='cuda' if torch.cuda.is_available() else 'cpu')
    mi_list = checkpoint['mi_list']
    model_state = checkpoint['model_state']
    discriminator.load_state_dict(model_state)
    print('Previous results loaded.')
else:
    mi_list = [] # storing the mi estimation of each iteration
```

## Training

Since maximizing the infoNCE bound over $f$ attains the mutual information

\begin{align}
I(X\wedge Y) = \sup_{f} I_{\text{infoNCE}}(f),
\end{align}

we train the neural network to minimize the loss $-I_{\text{infoNCE}}(f)$. The optimal solution $f$ satisfies

\begin{align}
f(x,y)=\log \frac{p_{Y|X}(y|x)}{p_{Y}(y)} + c \quad \forall x\in \mathcal{X}, y \in \mathcal{Y}
\end{align}

for any constant $c$.

The following cell train the neural network using the data samples.

```python
XY, X, Y, Ground_truth = generate_data(distribution=density, rho=opt.rho)

continue_train = False  # set to True to continue to train
if continue_train:
    for i in range(opt.n_iters):
        data_joint = _resample(XY, batch_size=opt.batch_size)
        train_x = data_joint[:,0:opt.d]
        train_y = data_joint[:,opt.d:]

        optimizer_D.zero_grad()
        loss = - infonce_loss(discriminator, train_x, train_y) # negative infonce_bound as the loss
        loss.backward()

        optimizer_D.step()
        with torch.no_grad():
            mi_est = infonce_loss(discriminator, X, Y)
        mi_list.append(mi_est.item())
        writer.add_scalar('mi_list', mi_est.item(), i)
        writer.add_scalar('loss', loss, i)
        if i%500==0:
            print("Iternation: %d, loss: %f, mi_est: %f"%(i, loss.item(), mi_est))

writer.add_graph(discriminator, (XY,))
```

    Iternation: 0, loss: 0.000010, mi_est: -0.000011
    Iternation: 500, loss: -1.986413, mi_est: 2.064315
    Iternation: 1000, loss: -3.324629, mi_est: 3.781751
    Iternation: 1500, loss: -3.816203, mi_est: 4.240057
    Iternation: 2000, loss: -3.713367, mi_est: 4.403550
    Iternation: 2500, loss: -3.796084, mi_est: 4.479726
    Iternation: 3000, loss: -3.887900, mi_est: 4.518673
    Iternation: 3500, loss: -4.057630, mi_est: 4.549497
    Iternation: 4000, loss: -3.867954, mi_est: 4.571889
    Iternation: 4500, loss: -3.878994, mi_est: 4.592578
    Iternation: 5000, loss: -3.913843, mi_est: 4.617288
    Iternation: 5500, loss: -4.089233, mi_est: 4.643414
    Iternation: 6000, loss: -3.991745, mi_est: 4.663235
    Iternation: 6500, loss: -3.934570, mi_est: 4.686573
    Iternation: 7000, loss: -4.032845, mi_est: 4.712495
    Iternation: 7500, loss: -4.063352, mi_est: 4.731988
    Iternation: 8000, loss: -4.008489, mi_est: 4.754710
    Iternation: 8500, loss: -4.079526, mi_est: 4.782144
    Iternation: 9000, loss: -4.210361, mi_est: 4.795047
    Iternation: 9500, loss: -4.119889, mi_est: 4.821354
    Iternation: 10000, loss: -4.167799, mi_est: 4.841077
    Iternation: 10500, loss: -4.207998, mi_est: 4.863544
    Iternation: 11000, loss: -4.259162, mi_est: 4.880834
    Iternation: 11500, loss: -4.085646, mi_est: 4.909222
    Iternation: 12000, loss: -4.237001, mi_est: 4.927190
    Iternation: 12500, loss: -4.238161, mi_est: 4.941082
    Iternation: 13000, loss: -4.208549, mi_est: 4.969678
    Iternation: 13500, loss: -4.337718, mi_est: 4.987714
    Iternation: 14000, loss: -4.229445, mi_est: 5.016221
    Iternation: 14500, loss: -4.249481, mi_est: 5.035546
    Iternation: 15000, loss: -4.330183, mi_est: 5.052544
    Iternation: 15500, loss: -4.326031, mi_est: 5.071255
    Iternation: 16000, loss: -4.261763, mi_est: 5.099937
    Iternation: 16500, loss: -4.250395, mi_est: 5.124137
    Iternation: 17000, loss: -4.325122, mi_est: 5.137771
    Iternation: 17500, loss: -4.382256, mi_est: 5.164788
    Iternation: 18000, loss: -4.259338, mi_est: 5.176157
    Iternation: 18500, loss: -4.387579, mi_est: 5.207092
    Iternation: 19000, loss: -4.284403, mi_est: 5.226067
    Iternation: 19500, loss: -4.365234, mi_est: 5.249267
    Iternation: 20000, loss: -4.372822, mi_est: 5.271327
    Iternation: 20500, loss: -4.368681, mi_est: 5.277494
    Iternation: 21000, loss: -4.467575, mi_est: 5.301701
    Iternation: 21500, loss: -4.394046, mi_est: 5.332717
    Iternation: 22000, loss: -4.323278, mi_est: 5.349476
    Iternation: 22500, loss: -4.371675, mi_est: 5.367954
    Iternation: 23000, loss: -4.457518, mi_est: 5.389325
    Iternation: 23500, loss: -4.489443, mi_est: 5.410488
    Iternation: 24000, loss: -4.405626, mi_est: 5.424552
    Iternation: 24500, loss: -4.381571, mi_est: 5.449566
    Iternation: 25000, loss: -4.434167, mi_est: 5.469605
    Iternation: 25500, loss: -4.520988, mi_est: 5.486578
    Iternation: 26000, loss: -4.437832, mi_est: 5.503848
    Iternation: 26500, loss: -4.452613, mi_est: 5.523114
    Iternation: 27000, loss: -4.480072, mi_est: 5.542090
    Iternation: 27500, loss: -4.533477, mi_est: 5.559884
    Iternation: 28000, loss: -4.506382, mi_est: 5.581320
    Iternation: 28500, loss: -4.476612, mi_est: 5.595648
    Iternation: 29000, loss: -4.548304, mi_est: 5.610945
    Iternation: 29500, loss: -4.542811, mi_est: 5.632243
    Iternation: 30000, loss: -4.514511, mi_est: 5.645619
    Iternation: 30500, loss: -4.521682, mi_est: 5.657323
    Iternation: 31000, loss: -4.531125, mi_est: 5.681959
    Iternation: 31500, loss: -4.500861, mi_est: 5.696266
    Iternation: 32000, loss: -4.536700, mi_est: 5.711816
    Iternation: 32500, loss: -4.556186, mi_est: 5.723273
    Iternation: 33000, loss: -4.522396, mi_est: 5.737668
    Iternation: 33500, loss: -4.533624, mi_est: 5.754815
    Iternation: 34000, loss: -4.552873, mi_est: 5.767499
    Iternation: 34500, loss: -4.548962, mi_est: 5.778892
    Iternation: 35000, loss: -4.574831, mi_est: 5.795615
    Iternation: 35500, loss: -4.564256, mi_est: 5.801952
    Iternation: 36000, loss: -4.561363, mi_est: 5.816156
    Iternation: 36500, loss: -4.543841, mi_est: 5.828671
    Iternation: 37000, loss: -4.563581, mi_est: 5.835651
    Iternation: 37500, loss: -4.556992, mi_est: 5.846558
    Iternation: 38000, loss: -4.559249, mi_est: 5.858206
    Iternation: 38500, loss: -4.577271, mi_est: 5.867131
    Iternation: 39000, loss: -4.577236, mi_est: 5.874681
    Iternation: 39500, loss: -4.589894, mi_est: 5.884233
    Iternation: 40000, loss: -4.569760, mi_est: 5.892098
    Iternation: 40500, loss: -4.589166, mi_est: 5.898573
    Iternation: 41000, loss: -4.597407, mi_est: 5.903266
    Iternation: 41500, loss: -4.582931, mi_est: 5.909054
    Iternation: 42000, loss: -4.581110, mi_est: 5.918843
    Iternation: 42500, loss: -4.590831, mi_est: 5.923814
    Iternation: 43000, loss: -4.581953, mi_est: 5.928774
    Iternation: 43500, loss: -4.586583, mi_est: 5.933340
    Iternation: 44000, loss: -4.589789, mi_est: 5.939531
    Iternation: 44500, loss: -4.592797, mi_est: 5.944098
    Iternation: 45000, loss: -4.594497, mi_est: 5.945866
    Iternation: 45500, loss: -4.598111, mi_est: 5.952806
    Iternation: 46000, loss: -4.598542, mi_est: 5.955683
    Iternation: 46500, loss: -4.592979, mi_est: 5.958319
    Iternation: 47000, loss: -4.597659, mi_est: 5.962470
    Iternation: 47500, loss: -4.594491, mi_est: 5.964867
    Iternation: 48000, loss: -4.598873, mi_est: 5.967290
    Iternation: 48500, loss: -4.601477, mi_est: 5.969827
    Iternation: 49000, loss: -4.601346, mi_est: 5.972457
    Iternation: 49500, loss: -4.601625, mi_est: 5.974484
    Iternation: 50000, loss: -4.600704, mi_est: 5.975776
    Iternation: 50500, loss: -4.601782, mi_est: 5.977545
    Iternation: 51000, loss: -4.599659, mi_est: 5.979087
    Iternation: 51500, loss: -4.602352, mi_est: 5.980538
    Iternation: 52000, loss: -4.602472, mi_est: 5.981858
    Iternation: 52500, loss: -4.604167, mi_est: 5.982779
    Iternation: 53000, loss: -4.604033, mi_est: 5.983732
    Iternation: 53500, loss: -4.603842, mi_est: 5.984695
    Iternation: 54000, loss: -4.603607, mi_est: 5.985449
    Iternation: 54500, loss: -4.603514, mi_est: 5.986165
    Iternation: 55000, loss: -4.604076, mi_est: 5.986738
    Iternation: 55500, loss: -4.604426, mi_est: 5.986921
    Iternation: 56000, loss: -4.603978, mi_est: 5.987811
    Iternation: 56500, loss: -4.604378, mi_est: 5.988275
    Iternation: 57000, loss: -4.604061, mi_est: 5.988607
    Iternation: 57500, loss: -4.604918, mi_est: 5.989032
    Iternation: 58000, loss: -4.604554, mi_est: 5.989225
    Iternation: 58500, loss: -4.604691, mi_est: 5.989665
    Iternation: 59000, loss: -4.604871, mi_est: 5.989891
    Iternation: 59500, loss: -4.604836, mi_est: 5.990000

## MI estimates

Use moving average to smooth the estimated MI in each training iteration.

```python
ma_rate = 0.01  # moving average rate
mi_copy = mi_list.copy()
for k in range(1,len(mi_list)):
    mi_copy[k] = (1-ma_rate) * mi_copy[k-1] + ma_rate * mi_copy[k]
```

Plot the MI estimation curve against the training iteration, together with the ground truth.

```python
plt.plot(mi_copy, label='MI estimate')
plt.axhline(Ground_truth,label='ground truth',linestyle='--',color='red')
for t in range(len(mi_copy)):
    if (mi_copy[t]>.8*Ground_truth):
        plt.axvline(t,label='80% reached',linestyle=':',color='green')
        break
plt.xlabel('number of iterations')
plt.ylabel('MI estimation')
plt.legend()
```

Save the model.

```python
overwrite = True  # set to True to overwrite previously stored results
if overwrite or not os.path.exists(chkpt_name):
    model_state = discriminator.state_dict()
    torch.save({
        'mi_list': mi_list,
        'mi_copy': mi_copy,
        'model_state': model_state
    }, chkpt_name)
    writer.close()
    print('Current results saved.')
```

    Current results saved.