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 5 -- Soil column under its own weight

The differential equation of the equilibrium conditions for this 1D example reads (soil mechanical sign convention)

$$
\sigma_{zz,z} + \varrho g = 0
$$

Using the constitutive law

$$
\sigma_{zz} = E_\text{s} \epsilon_{zz}
$$

and the kinematical relationship

$$
\epsilon_{zz} = - u_{z,z}
$$

we arrive at a second-order ordinary differential equation with constant coefficients in terms of the vertical displacement:

$$
u_{z,zz} = \frac{\varrho g}{E_\text{s}}
$$

Observe the distinction between a general balance law, a specific constitutive law, and geometric relations (kinematics). The combination yields the problem-specific governing equation.

Integrating twice yields:

$$
u_z = \frac{\varrho g}{2E_\text{s}} z^2 + C_1 z + C_2
$$

The boundary conditions

$$
u_z(z=0) = 0 \qquad \text{and} \qquad \sigma_{zz}(z=H) \propto u_{z,z}(z=H) = 0
$$

allow us to determine the integration constants

$$
C_1 = -\frac{\varrho g H}{E_\text{s}} \qquad \text{and} \qquad C_2 = 0
$$

Thus, the final solution of the problem is given as

$$
u_z(z) = \frac{\varrho gH^2}{2E_\text{s}} \left[ \left(\frac{z}{H}\right)^2 - 2 \frac{z}{H} \right]
$$

The stress distribution follows as

$$
\sigma_{zz}(z) = \varrho g H \left(1 - \frac{z}{H} \right)
$$

```python
import numpy as np #numerical methods
#import sympy as sp #symbolic operations
import matplotlib.pyplot as plt #plotting
#sp.init_printing(use_latex='mathjax') #makes sympy output look nice

#Some plot settings
import plot_functions.plot_settings
```

We insert the dimensionless coordinate $\xi = \displaystyle \frac{z}{H} \in [0,1]$ for simplicity and provide the analytical solutions as functions:

```python
#xi is the dimensionless coordinate z/H
def displacement_analytical(xi):
    return xi**2. - 2. * xi

def stress_analytical(xi):
    return 1. - xi
```

```python
coord = np.linspace(0.,1.,100)
```

```python
plt.plot(-displacement_analytical(coord),coord,label='settlement')
plt.plot(stress_analytical(coord),coord,label='vertical stress')
plt.legend()
plt.xlabel('$-u_z\ /\ \\varrho g H^2 (2 E_\\mathrm{s})^{-1}$, $\\sigma_{zz}\ /\ \\varrho g H$')
plt.ylabel('$z\ /\ H$');
```

We now seek approximate displacement solutions to the second-order ODE

$$
    E_\text{s} u_{z,zz} = \varrho g
$$

and will use different methods.

## Collocation method

Choosing a linear Ansatz $\tilde{u}_z = a z$ does not work since $u_{z,zz} = 0$. A quadratic Ansatz of the following form

$$
    \tilde{u}_z = \frac{a}{2} z^2 - a H z
$$

fulfills the boundary conditions $u_z(z = 0) = 0$ and $u_{z,z}(z = H) = 0$ with

$$
    \tilde{u}_{z,z} = a (z - H)
$$

Substituting this Ansatz into the ODE yields

$$
    a = \frac{\varrho g}{E_\text{s}}
$$

We see that the choice of a collocation point has become arbitrary and the solution is directly determined by the boundary conditions and the ODE itself.

The solution is thus given as

$$
    \tilde{u}_z = \frac{\varrho g H^2}{2E_\text{s}} \left[ \left(\frac{z}{H}\right)^2 - 2 \frac{z}{H} \right]
$$

which turns out to be the exact solution.

To better illustrate the method, let's choose a trigonometric Ansatz:

$$
    \tilde{u}_z = a \sin \left( \frac{\pi z}{2H} \right)
$$

which also fulfills all boundary conditions. Substitution into the ODE yields

$$
    \frac{\pi^2 a}{4 H^2} \sin \left( \frac{\pi}{2} \xi \right) = -\frac{\varrho g}{E_\text{s}}
$$

We have one unknown, so we need one collocation point. We conveniently choose that at $\xi_1 = 1/3$ because $\sin (\pi/6) = 0.5$. Thus, we find

$$
    a = -\frac{8 \varrho g H^2}{E_\text{s} \pi^2}
$$

Thus, the solution reads

$$
    \tilde{u}_z = -\frac{\varrho g H^2}{2 E_\text{s}} \frac{16}{\pi^2}  \sin \left( \frac{\pi }{2} \xi \right)
$$

and

$$
    \sigma_{zz} = \varrho g H \frac{4}{\pi} \cos \left( \frac{\pi }{2} \xi \right)
$$

```python
def displacement_collocation(xi): #in \rho g H^2/(2E_s)
    return - 16/(np.pi**2) * np.sin(np.pi/2 * xi)

def stress_collocation(xi): # in \rho g H
    return 4/np.pi * np.cos(np.pi/2 * xi)
```

```python
plt.plot(-displacement_analytical(coord),coord,label='settlement analytical')
plt.plot(stress_analytical(coord),coord,label='vertical stress analytical')
plt.plot(-displacement_collocation(coord),coord,label='settlement collocation',ls='--')
plt.plot(stress_collocation(coord),coord,label='vertical stress collocation',ls='--')
plt.legend()
plt.xlabel('$-u_z\ /\ \\varrho g H^2 (2 E_\\mathrm{s})^{-1}$, $\\sigma_{zz}\ /\ \\varrho g H$')
plt.ylabel('$z/H$');
```

The approximation is neither qualitatively nor quantitatively very well. We saw above that the solution for $a$ is strongly dependent on the location of the collocation point. Let's try another: $\xi = 2/3$. Then, $\sin (\pi/3) = \sqrt(3)/2$ and

$$
     a = -\frac{8 \varrho g H^2}{\sqrt{3}E_\text{s} \pi^2}
$$

Thus, the solution reads

$$
    \tilde{u}_z = -\frac{\varrho g H^2}{2 E_\text{s}} \frac{16}{\sqrt{3}\pi^2}  \sin \left( \frac{\pi }{2} \xi \right)
$$

and

$$
    \sigma_{zz} = \varrho g H \frac{4}{\sqrt{3}\pi} \cos \left( \frac{\pi }{2} \xi \right)
$$

```python
plt.plot(-displacement_analytical(coord),coord,label='settlement analytical')
plt.plot(stress_analytical(coord),coord,label='vertical stress analytical')
plt.plot(-displacement_collocation(coord)/np.sqrt(3),coord,label='settlement collocation',ls='--')
plt.plot(stress_collocation(coord)/np.sqrt(3),coord,label='vertical stress collocation',ls='--')
plt.legend()
plt.xlabel('$-u_z\ /\ \\varrho g H^2 (2 E_\\mathrm{s})^{-1}$, $\\sigma_{zz}\ /\ \\varrho g H$')
plt.ylabel('$z/H$');
```

## Bubnov-Galerkin scheme

The residual is in our case defined as

$$
r(z) = E_\text{s} \tilde{u}_{z,zz} - \varrho g
$$

We multiply the residual by a test function of which we demand that it vanishes on the Dirichlet boundary (where $u_z = 0$): $\eta(z=0) = 0$. We demand that the weighted residual vanished on average over the domain:

$$
0 = \int \limits_0^H r(z) \eta(z)\,\text{d}z = \int \limits_0^H \left( E_\text{s} \tilde{u}_{z,zz} - \varrho g \right) \eta\,\text{d}z
$$

Now we perform partial integration to shift one of the derivatives from $\tilde{u}_z$ over to the test function $\eta$. The advantage of this is that $C^0$-continuous functions could be used to approximate the displacement (i.e. piecewise linear approximations). This is not exploited in this demonstration example but will be useful for our finite elements.

$$
\begin{aligned}
0 &= \int \limits_0^H \left(\left[ E_\text{s} \tilde{u}_{z,z}\eta\right]_{,z} - E_\text{s} \tilde{u}_{z,z} \eta_{,z} - \varrho g \eta \right) \,\text{d}z
\\
0 &= \left[ E_\text{s} \tilde{u}_{z,z}\eta\right]_0^H - \int \limits_0^H \left(E_\text{s} \tilde{u}_{z,z} \eta_{,z} + \varrho g \eta \right) \,\text{d}z
\end{aligned}
$$

Because $\eta(z=0)=0$, only the top boundary term remains as the naturally appearing boundary condition (natural or Neumann bc). It corresponds to the traction on the top:
$$
\begin{aligned}
0 &= E_\text{s} \tilde{u}_{z,z}(z=H)\eta(z=H) - \int \limits_0^H \left(E_\text{s}\tilde{u}_{z,z} \eta_{,z} + \varrho g \eta \right) \,\text{d}z
\\
&=-\bar{\sigma}_{zz}(z=H)\eta(z=H) - \int \limits_0^H \left(E_\text{s} \tilde{u}_{z,z} \eta_{,z} + \varrho g \eta \right) \,\text{d}z
\end{aligned}
$$

Bringing all known quantities to the right hand side yields:
$$
\int \limits_0^H E_\text{s} \tilde{u}_{z,z} \eta_{,z}\,\text{d}z = -\bar{\sigma}_{zz}|_H\eta|_H - \int \limits_0^H \varrho g \eta \,\text{d}z
$$

Since the traction on top is 0 in our case, the boundary term simply vanishes and we're left with:
$$
\int \limits_0^H E_\text{s} \tilde{u}_{z,z} \eta_{,z}\,\text{d}z = - \int \limits_0^H \varrho g \eta \,\text{d}z
$$

To solve this, we substitute a linear Ansatz $\tilde{u}_z = a z$ which fulfils the essential boundary conditions ($u_z|_0 = 0$) and take the test function to be of the same form: $\eta = b z$.

$$
\int \limits_0^H E_\text{s} a b \,\text{d}z = - \int \limits_0^H \varrho g b z \,\text{d}z
$$

Since $b$ is a constant, we can pull it out of the integral:

$$
\begin{aligned}
 0 &= b \left[ \int \limits_0^H E_\text{s} a \,\text{d}z + \int \limits_0^H \varrho g z \,\text{d}z \right]
\end{aligned}
$$

We furthermore demand that this equation hold for arbitrary $b$ such that the term in the brackets has to vanish, giving us an equation to determine a:

$$
 \int \limits_0^H E_\text{s} \,\text{d}z\ a = - \int \limits_0^H \varrho g z \,\text{d}z
$$

Here we see the stiffness matrix, the unknown $a$, and the right-hand-side terms. We will see this pattern over and over again so that we always end up with a linear system in the form:

$$
    \mathbf{K} \mathbf{u} = \mathbf{f}^\text{ext} \qquad \text{or} \qquad \mathbf{f}^\text{int} = \mathbf{f}^\text{ext}
$$

In this simple example, the stiffness matrix contains only the material's stiffness.

The solution here simply follows as
$$
    a = -\frac{\int \limits_0^H \varrho g z \,\text{d}z}{\int \limits_0^H E_\text{s} \,\text{d}z} = - \frac{\varrho gH}{2E_\text{s}}
$$

Therefore,
$$
    \tilde{u}_z = - \frac{\varrho g H^2}{2E_\text{s}} \frac{z}{H}
$$
and
$$
    \sigma_{zz} = \frac{\varrho g H}{2}
$$

```python
def displacement_galerkin(xi):
    return -xi

def stress_galerkin(xi):
    return 1./2.*np.sign(xi+1.e-16)
```

```python
plt.plot(-displacement_analytical(coord),coord,label='settlement analytical')
plt.plot(stress_analytical(coord),coord,label='vertical stress analytical')
plt.plot(-displacement_galerkin(coord),coord,label='settlement Galerkin',ls='--')
plt.plot(stress_galerkin(coord),coord,label='vertical stress Galerkin',ls='--')
plt.legend()
plt.xlabel('$-u_z\ /\ \\varrho g H^2 (2 E_\\mathrm{s})^{-1}$, $\\sigma_{zz}\ /\ \\varrho g H$')
plt.ylabel('$z/H$');
```

We observe that stress result corresponds to the mean stress and that both displacement boundary values match the exact solution, but the internal distribution is only as good as the approximation.
We furthermore see, that the traction boundary condition cannot be met well due to the constant-stress approximation.

So let's try a trigonometric Ansatz. Due to the properties of the Galerkin scheme, the Ansatz only has to fulfill the essential boundary conditions. Creating a quadratic Ansatz is thus a lot easier than in the collocation example. However, let's simply re-use

$$
    \tilde{u}_z = a \sin \left( \frac{\pi z}{2H} \right) \qquad \text{and} \qquad \eta_z = b \sin \left( \frac{\pi z}{2H} \right)
$$

Substitution into

$$
\int \limits_0^H E_\text{s} \tilde{u}_{z,z} \eta_{,z}\,\text{d}z = - \int \limits_0^H \varrho g \eta \,\text{d}z
$$

yields

$$
\int \limits_0^H E_\text{s} \frac{ab\pi^2}{4H^2} \cos^2 \left( \frac{\pi z}{2H} \right) \,\text{d}z = - \int \limits_0^H \varrho g b \sin \left( \frac{\pi z}{2H} \right) \,\text{d}z
$$

The constant $b$ again falls out as arbitrary and we end up with $\mathbf{K} \mathbf{u} = \mathbf{f}^\text{ext}$ in the form

$$
 \int \limits_0^H \frac{E_\text{s} \pi^2}{4H^2} \cos^2 \left( \frac{\pi z}{2H} \right) \,\text{d}z\ a = - \int \limits_0^H \varrho g \sin \left( \frac{\pi z}{2H} \right) \,\text{d}z
$$

Let Python do the work ...

```python
import sympy as sp
sp.init_printing(use_latex='mathjax') #makes sympy output look nice
H,z,E,gamma = sp.symbols('H z E gamma',positive=True)

RHS = -gamma*sp.sin(sp.pi*z/(2*H))
RHS
```

$\displaystyle - \gamma \sin{\left(\frac{\pi z}{2 H} \right)}$

```python
sp.integrate(RHS,(z,0,H))
```

$\displaystyle - \frac{2 H \gamma}{\pi}$

```python
LHS = E*sp.pi**2/(4*H**2)*(sp.cos(sp.pi*z/(2*H)))**2
LHS
```

$\displaystyle \frac{\pi^{2} E \cos^{2}{\left(\frac{\pi z}{2 H} \right)}}{4 H^{2}}$

```python
sp.integrate(LHS,(z,0,H))
```

$\displaystyle \frac{\pi^{2} E}{8 H}$

```python
sp.integrate(RHS,(z,0,H))/sp.integrate(LHS,(z,0,H))
```

$\displaystyle - \frac{16 H^{2} \gamma}{\pi^{3} E}$

We thus find:

$$
    a = -\frac{16 \varrho g H^2}{E_\text{s}\pi^3}
$$

```python
def displacement_galerkin2(xi): #in \rho g H^2/(2E_s)
    return - 32/(np.pi**3) * np.sin(np.pi/2 * xi)

def stress_galerkin2(xi): # in \rho g H
    return 8/np.pi**2 * np.cos(np.pi/2 * xi)
```

```python
plt.plot(-displacement_analytical(coord),coord,label='settlement analytical')
plt.plot(stress_analytical(coord),coord,label='vertical stress analytical')
plt.plot(-displacement_galerkin2(coord),coord,label='settlement Galerkin',ls='--')
plt.plot(stress_galerkin2(coord),coord,label='vertical stress Galerkin',ls='--')
plt.legend()
plt.xlabel('$-u_z\ /\ \\varrho g H^2 (2 E_\\mathrm{s})^{-1}$, $\\sigma_{zz}\ /\ \\varrho g H$')
plt.ylabel('$z/H$');
```

No arbitrary collocation point had to be chosen. Instead, an optimal solution was found. Nevertheless, the stress approximation is still hampered by the choice of the approximation function which determines its shape.

# Ritz-Rayleigh

We can also find a statically admissible solution, knowing that elastic systems seek to attain the minimum of their total elastic potential (Minimum des elastischen Gesamtpotentials).

The total elastic potential $\Pi$ consists of the elastic energy stored in the material $\Pi^\text{int}$ and the work done by the external forces $W^\text{ext}$.

To find the former, consider the elastic energy stored in a linear spring. It follows from the work done on the spring:

$$
    W = \int F\,\text{d}u = \int \limits_0^{\Delta l} k u\,\text{d}u = \frac{k}{2} (\Delta l)^2 = \Delta E
$$

In a similar way we find for a uniaxial soil element at constant stress/strain under uniaxial conditions:

$$
    W = \int F\,\text{d}u = \int \limits_0^{\Delta l} E_\text{s} \epsilon A_0 \,\text{d}u = \int \limits_0^{\Delta \epsilon_{zz}} E_\text{s} \epsilon A_0 l_0 \,\text{d}\epsilon = \frac{E_\text{s}}{2} \epsilon_{zz}^2 V_0 = \Delta E
$$

giving rise to the strain energy density (strain energy per unit volume):

$$
    \psi = \frac{\Delta E}{V_0} = \frac{E_\text{s}}{2} \epsilon_{zz}^2
$$

The elastic energy stored in a soil volume under these specific conditions is thus:

$$
    \Pi^\text{int} = \int \limits_V \frac{E_\text{s}}{2} u_{z,z}^2\,\text{d}V
$$

This change in internal stored energy is caused by the work done by the external forces (in our case only gravity acting in the negative coordinate direction):

$$
    W^\text{ext} = - \int \limits_V \varrho g u_z \,\text{d}V
$$

The total elastic potential $\Pi^\text{int} - W^\text{ext}$ is thus given:

$$
    \Pi = \int \limits_V \frac{E_\text{s}}{2} u_{z,z}^2\,\text{d}V + \int \limits_V \varrho g u_z \,\text{d}V
$$

Now we use again $\tilde{u} = a z$ and demand stationarity of the potential (find $a$ such that $\Pi$ becomes minimal):

$$
    \delta \Pi = \frac{\partial \Pi}{\partial a} \delta a \ \forall \delta a
$$

Therefore, from

$$
    \frac{\partial \Pi}{\partial a} = 0 =  \int \limits_0^H E_\text{s} a \,A_0\text{d}z + \int \limits_0^H \varrho g  z \,A_0\text{d}z
$$

we find

$$
    a = - \frac{\int \limits_0^H \varrho g  z \,\text{d}z}{\int \limits_0^H E_\text{s} \,\text{d}z} = -\frac{\varrho g H}{2E_\text{s}}
$$

# Variational approach, principle of virtual displacements

Consider again the total potential for our special case of a load-free soil column (load-free, but subjected to its own weight)

$$
    \Pi = \int \limits_V \frac{E_\text{s}}{2} u_{z,z}^2\,\text{d}V + \int \limits_V \varrho g u_z \,\text{d}V
$$

We can require stationarity with respect to the as yet unspecified solution functional and its derivative:

$$
    \delta \Pi = 0 = \frac{\partial \Pi}{\partial u_z} \delta u_z + \frac{\partial \Pi}{\partial u_{z,z}} \delta u_{z,z}
$$

The only condition we place is that the variations $\delta u$ vanish on Dirichlet boundaries (we don't need to wiggle the solution where we know it), i.e. in our case: $\delta u(z=0)=0$ (note: the use of this is not seen here, but would be seen if we include a boundary traction into the work formulation).

This yields
$$
    0 = \int \limits_0^H E_\text{s} u_{z,z} \delta u_{z,z} \,\text{d}z + \int \limits_0^H \varrho g \delta u_z \,\text{d}z = \delta \Pi A_0^{-1}
$$

This is the well-known principle of virtual work if we interpret our variations as virtual displacements. Furthermore, recall the formulation of the Galerkin problem:
$$
0 = \int \limits_0^H E_\text{s} \tilde{u}_{z,z} \eta_{,z}\,\text{d}z + \int \limits_0^H \varrho g \eta \,\text{d}z
$$

We see that if we interpret the test functions $\eta$ as virtual displacements $\delta u$, Galerkin gave us the principle of virtual work.

Now that we know this, we can interpret the orthogonality condition in the Galerkin scheme as follows: when considering the integral over the entire domain, the residuals (out-of-balance forces) do not perform any virtual work.


# CSE 330 Numerical Analysis Lab

### Lab 4

### Instructions
---

In this week we will learn a little about numpy polynomial module (`numpy.polynomial`). After that we will do two assignment using this Polynomial class.

Todays assignments are:
1. Implement Hermite Interpolation using Python
2. Implement forward, backward and central differentiaton using python and finding error.

But before we start let’s take a look at numpy polynomial module.

### `numpy.polynomial`
---

In mathematics polynomials are of form,

$$f(x)=a_0+a_1 x+a_2 x^2 + a_3 x^3 + ⋯ + a_n x^n \tag{4.1}$$

Numpy provide a class named Ploynmial in it’s numpy.polynomial module. We can initialize a polynomial by passing the coefficients as an array.

```
import numpy as np
from numpy.polynomial import Polynomial

p1 = Polynomial([1,2,3])    # 1 + 2x + 3x^2
p1
```

$x \mapsto \text{1.0} + \text{2.0}\,x + \text{3.0}\,x^{2}$

Here the variable p1 represents the following equation,

$$x \mapsto \text{1.0} + \text{2.0}\,x + \text{3.0}\,x^{2}$$

We can also use a numpy array to initialize the polynomial. For example,

```
coeff = np.array([-5,0,2,1])
p2 = Polynomial(coeff)
p2
```

$x \mapsto \text{-5.0}\color{LightGray}{ + \text{0.0}\,x} + \text{2.0}\,x^{2} + \text{1.0}\,x^{3}$

Here p2 represents the following equation,

$$x \mapsto \text{-5.0}\color{LightGray}{ + \text{0.0}\,x} + \text{2.0}\,x^{2} + \text{1.0}\,x^{3}$$

We can also call the polynomial as a function of x, i.e. $p1(x)$. For example,

```
print("p1(0.0) = ", p1(0.0))
print("p1(1.0) = ", p1(1.0))
print("p1(2.0) = ", p1(2.0))
```

    p1(0.0) =  1.0
    p1(1.0) =  6.0
    p1(2.0) =  17.0

We can also pass variables in the polynomial function,

```
x = 3
y = p1(x)
print("p1(x) = ", y)
```

    p1(x) =  34.0

Althogh, you might sometimes need to calculate the $Y$-values for a lot of different $X$-values at a time. Howver, passing an array with give you an error.

```
x = [1,2,3]
y = p1(x)
print(y)
```

Numpy array can come to your rescue in this case,

```
x = np.array([1,2,3])
y = p1(x)
print(y)
```

You can also run arithmatic operations on the polynomials. for example,

\begin{align}
y = p1(x) &= 9 + 2 x - 2.5 \times x^3 \\
y = p2(x) &= 6 + 3.5 \times x + 1.9 \times x^2 \\
y = p1(x) + p2(x) &= 15 + 5.5 \times x + 1.9 \times x^2 - 2.5 \times x^3
\end{align}

```
p1 = Polynomial([9.0, 2.0, 0.0, -2.5])
p2 = Polynomial([6.0, 3.5, 1.9])

p3 = p1+p2
p3
```

Here $p3$ represents the following equation,
$$x \mapsto \text{15.0} + \text{5.5}\,x + \text{1.9}\,x^{2} - \text{2.5}\,x^{3}$$

Similarly you can multiply polynomials too. For example,
\begin{align}
g(x) &= -1 + x \\
h(x) &= 6 + 5\times x+x^2 \\
g(x)\times h(x) &= -6 + x + 4 x^2 + x^3
\end{align}

```
g = Polynomial([-1, 1])
h = Polynomial([6, 5, 1])

p = g*h
display(p)
```

Here $p$ represents the following polynomial,
$x \mapsto \text{-6.0} + \text{1.0}\,x + \text{4.0}\,x^{2} + \text{1.0}\,x^{3}$

You can also find roots for a polynomial.

```
f = Polynomial([6, -5, 1]) # 6 - 5x + x^2
f.roots()
```

To find the derivative of a function you can use deriv function.

```
f = Polynomial([6, -5, 1]) # 6 - 5x + x^2

df = f.deriv(1) # -5 + 2x
display(df)

ddf = f.deriv(2) # 2
display(ddf)
```

Here, $df$ represensts the polynomial, $x \mapsto \text{-5.0} + \text{2.0}\,x$, and $ddf$ represenst the polynomial, $x \mapsto \text{2.0}$

Let's move on to todays assignment.

## Assignment 1: Hermite Interpolation
---
If you recall your class lecture you might remember *Hermite interpolation*. If you have $n+1$ data points, $(x_0, y_0), (x_1, y_1), x_2, y_2), ... (x_n, y_n)$ and you happen to know the first-order ferivative at all of these points, namely, $(x_0, y_0 ^\prime ), (x_1, y_1 ^\prime ), x_2, y_2 ^\prime ), ... (x_n, y_n ^\prime )$, Accorrding to hermit interpolation you can fit a polynomial of order $2n-1$.

$$p_{2n+1} = \sum_{k=0}^{n}{f(x_K)h_k(x)+f^\prime (x_k)\hat{h_k}(x)}\tag{4.2}$$
where $h_k$ and $\hat{h_k}$ are defined by Lagrange basis functions,
$$h_k(x) = (1-2(x-x_k)l^\prime_k(x_k))l^2_k(x_k) \tag{4.3}$$
and,
$$\hat{h_k}(x) = (x-x_k)l^2_k(x_k) \tag{4.4}$$

For this assignment you will be provided with the following skeleton code.

```
import numpy as np
import matplotlib.pyplot as plt
from itertools import combinations
from numpy.polynomial import Polynomial

def l(k, x):
    n = len(x)
    assert (k < len(x))

    x_k = x[k]
    x_copy = np.delete(x, k)

    denominator = np.prod(x_copy - x_k)

    coeff = []

    for i in range(n):
        coeff.append(sum([np.prod(x) for x in combinations(x_copy, i)]) * (-1)**(i) / denominator)

    coeff.reverse()

    return Polynomial(coeff)

def h(k, x):
    l_k = l(k, x)
    l_k_sqr = l_k * l_k
    l_k_prime = l_k.deriv(1)
    coeff = np.array([1+2*x[k]*l_k_prime(x[k]), -2*l_k_prime(x[k])])
    p = Polynomial(coeff)

    return p * l_k_sqr

def h_hat(k, x):
    l_k = l(k, x)
    l_k_sqr = l_k * l_k
    coeff = np.array([-x[k], 1])
    p = Polynomial(coeff)

    return p * l_k_sqr

def hermit(x, y, y_prime):
    assert( len(x) == len(y))
    assert( len(y) == len(y_prime))

    f = Polynomial([0.0])
    for k in range(len(x)):
        pol_x = Polynomial(x[k])
        f += y[k] * h(k, x) + pol_x.deriv(1) * h_hat(k, x)
        pass
    return f
```

Here, method `l`, `h` and `h_hat` calculates the corresponding $l_k$, $h_k$ and $\hat{h_k}$.

Note that, we can rewrite equation 4.3 in this way,
\begin{align}
h_k(x) &= (1-2(x-x_k)l^\prime_k(x_k))l^2_k(x_k) \\
&= (1 - 2xl^\prime_k(x_k) + 2x_kl^\prime_k(x_k))l^2_k(x_k)
\end{align}
Replacing, $l^\prime_k(x_k)$ with $m$ we get,
$$h_k(x) = (1 - 2xm + 2x_km)l^2_k(x_k) \tag{4.5}$$

Equation 4.5 is basically becomes a multiplication of two polynomials.

### Task

method `l` has already been defined for you. Your task is to complete the `h`, `h_hat` and `hermit` method.

Later we will draw some plots to check if the code is working.

```
pi = np.pi

x       = np.array([0.0, pi/2.0,  pi, 3.0*pi/2.0])
y       = np.array([0.0,    1.0, 0.0,       -1.0])
y_prime = np.array([1.0,    0.0, 1.0,        0.0])

n = 1
f3     = hermit(x[:(n+1)], y[:(n+1)], y_prime[:(n+1)])
data   = f3.linspace(n=50, domain=[-3, 3])
test_x = np.linspace(-3, 3, 50, endpoint=True)
test_y = np.sin(test_x)

plt.plot(data[0], data[1])
plt.plot(test_x, test_y)
plt.show()

n = 2
f5     = hermit(x[:(n+1)], y[:(n+1)], y_prime[:(n+1)])
data   = f5.linspace(n=50, domain=[-0.7, 3])
test_x = np.linspace(-2*pi, 2*pi, 50, endpoint=True)
test_y = np.sin(test_x)

plt.plot(test_x, test_y)
plt.plot(data[0], data[1])
plt.show()

n = 3
f7     = hermit(x[:(n+1)], y[:(n+1)], y_prime[:(n+1)])
data   = f7.linspace(n=50, domain=[-0.3, 3])
test_x = np.linspace(-2*pi, 2*pi, 50, endpoint=True)
test_y = np.sin(test_x)

plt.plot(data[0], data[1])
plt.plot(test_x, test_y)
plt.show()

x       = np.array([0.0, 1.0,          2.0       ])
y       = np.array([1.0, 2.71828183,  54.59815003])
y_prime = np.array([0.0, 5.43656366, 218.39260013])
f7      = hermit( x, y, y_prime)
data    = f7.linspace(n=50, domain=[-0.5, 2.2])
test_x  = np.linspace(-0.5, 2.2, 50, endpoint=True)
test_y  = np.exp(test_x**2)

plt.plot(data[0], data[1])
plt.plot(test_x, test_y)
plt.show()

x       = np.array([1.0, 3.0, 5.0])
y       = np.array([5.0, 1.0, 5.0])
y_prime = np.array([-4.0, 0.0, 4.0])
f7      = hermit( x, y, y_prime)
data    = f7.linspace(n=50, domain=[-10, 10])
test_x  = np.linspace(-10, 10, 50, endpoint=True)
test_y  = (test_x-3)**2 + 1

plt.plot(data[0], data[1])
plt.plot(test_x, test_y)
plt.show()
```

## Assignment 2: Differentiation

We have already learnt about *forward differentiation*, *backward diferentiation* and *central differentiation*. In this part of the assignment we will write methods to calculate this values and check how they perform.

The equations are as follows,

\begin{align}
\text{forward differentiation}, f^\prime(x) = \frac{f(x+h)-f(x)}{h} \tag{4.6} \\
\text{backward differentiation}, f^\prime(x) = \frac{f(x)-f(x-h)}{h} \tag{4.7} \\
\text{central differentiation}, f^\prime(x) = \frac{f(x+h)-f(x-h)}{2h} \tag{4.8}
\end{align}

you will be provided with a skeleton code.

```
import numpy as np
import matplotlib.pyplot as plt
from numpy.polynomial import Polynomial

def forward_diff(f, h, x):
    fdiff = (f(x+h) - f(x))/ h
    return fdiff

def backward_diff(f, h, x):
    bdiff = (f(x) - f(x-h))/ h
    return bdiff

def central_diff(f, h, x):
    return (f(x+h) - f(x-h))/ (2*h)

def error(f, f_prime, h, x):
    Y_correct = f_prime(x)
    f_error = []
    b_error = []
    c_error = []

    for h_i in h:
        f_error.append(Y_correct - forward_diff(f, h_i, x))
        b_error.append(Y_correct - backward_diff(f, h_i, x))
        c_error.append(Y_correct - central_diff(f, h_i, x))
        pass
    print(f" f_error = {f_error},\n b_error =  = {b_error},\n c_error = {c_error}")

    f_error = np.array(f_error)
    b_error = np.array(b_error)
    c_error = np.array(c_error)

    return f_error, b_error, c_error
```

Here, `forward_diff`, `backward_diff` and `central_diff` calculates the *forward differentiation*, *backward differentiation* and *central differentiation* respectively. finally the `error` method calculates the different values for various $h$ and returns the errors.

Later we will run some code to test out performance,

```
fig, ax = plt.subplots()
ax.axhline(y=0, color='k')

p = Polynomial([2.0, 1.0, -6.0, -2.0, 2.5, 1.0])
data = p.linspace(domain=[-2.4, 1.5])
ax.plot(data[0], data[1], label='Function')

p_prime = p.deriv(1)
data2 = p_prime.linspace(domain=[-2.4, 1.5])
ax.plot(data2[0], data2[1], label='Derivative')

ax.legend()

##################################

h = 1
fig, bx = plt.subplots()
bx.axhline(y=0, color='k')

x = np.linspace(-2.0, 1.3, 50, endpoint=True)
y = forward_diff(p, h, x)
bx.plot(x, y, label='Forward; h=1')
y = backward_diff(p, h, x)
bx.plot(x, y, label='Backward; h=1')
y = central_diff(p, h, x)
bx.plot(x, y, label='Central; h=1')

data2 = p_prime.linspace(domain=[-2.0, 1.3])
bx.plot(data2[0], data2[1], label='actual')

bx.legend()

##################################

h = 0.1
fig, bx = plt.subplots()
bx.axhline(y=0, color='k')

x = np.linspace(-2.2, 1.3, 50, endpoint=True)
y = forward_diff(p, h, x)
bx.plot(x, y, label='Forward; h=0.1')
y = backward_diff(p, h, x)
bx.plot(x, y, label='Backward; h=0.1')
y = central_diff(p, h, x)
bx.plot(x, y, label='Central; h=0.1')

data2 = p_prime.linspace(domain=[-2.2, 1.3])
bx.plot(data2[0], data2[1], label='actual')

bx.legend()

##################################

h = 0.01
fig, bx = plt.subplots()
bx.axhline(y=0, color='k')

x = np.linspace(-2.2, 1.3, 50, endpoint=True)
y = forward_diff(p, h, x)
bx.plot(x, y, label='Forward; h=0.01')
y = backward_diff(p, h, x)
bx.plot(x, y, label='Backward; h=0.01')
y = central_diff(p, h, x)
bx.plot(x, y, label='Central; h=0.01')

data2 = p_prime.linspace(domain=[-2.2, 1.3])
bx.plot(data2[0], data2[1], label='actual')

bx.legend()

##################################

fig, bx = plt.subplots()
bx.axhline(y=0, color='k')

h = np.array([1., 0.55, 0.3, .17, 0.1, 0.055, 0.03, 0.017, 0.01])
err = error(p, p_prime, h, 2.0)

bx.plot(h, err[0], label='Forward')
bx.plot(h, err[1], label='Backward')
bx.plot(h, err[2], label='Central')
bx.legend()
```


# MCMC from scratch

This notebook is a study of Markov Chain MOnte Carlo methods,
and is inspired, in part, by these great papers:
* A basic overview of sampling methods such as AR, MH and Gibbs
  [Chib, Greenberg (1995)](http://web1.sph.emory.edu/users/hwu30/teaching/statcomp/papers/chibGreenbergMH.pdf)
. Includes a hybrid MH-AR method (whrere the proposal is sampled from using AR)
* A self contained introduction and study of Monte Carlo based on Hamiltonian dynamics
  [Betancourt (2017)](https://arxiv.org/abs/1701.02434.pdf)

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

import numpy as np

import torch
import torch.nn.functional as F

import tqdm
```

Get inflated `rect` for better $2$-d plot aesthetics

```python
def get_rect(data, dim=0, r=5e-2, a=1e-3):
    # deal with nans!
    dims, brdc = [*range(data.dim())], data.dim()*[1]
    dims.pop(dim)
    brdc[dim] = -1
    mask = (~torch.isfinite(data)).sum(dims) == 0
    mask = mask.reshape(brdc[:dim+1])

    # get the enclosing rectangle ...
    (uu, _), (ll, _) = data[mask].max(dim), data[mask].min(dim)

    # ... center it, infalte, ...
    cc = (uu + ll) / 2
    uu, ll = uu - cc, ll - cc
    uu = uu + abs(uu) * r + a
    ll = ll - abs(ll) * r - a

    # ... and translate back
    return ll + cc, uu + cc
```

Create a multi-modal density: we can evaluate it, know its structure,
but not the normalizing constant.

$$
p(\theta)
    = \tfrac1Z \mathop{\mathrm{exp}} \bigl(
        \log f(\theta)
    \bigr)
    \,.
$$

Let's make a banana distribution:

$$
p(x)
    \propto p_{\mathcal{N}(0, 1)} \circ \phi(x)
    \,, $$

where $\phi$ is the based on the
[Banana](https://en.wikipedia.org/wiki/Rosenbrock_function)
function and given by $
\phi
\colon \mathbb{R}^2 \to \mathbb{R}^2
\colon (x, y) \mapsto (a x, b (y-x^2))
$

```python
def log_banana_base(x, a=0.75, b=1.05):
#     phi = torch.stack([a - x[..., 0], b * (x[..., 1] - x[..., 0]**2)], dim=-1)  # a=1.75, b=5
    phi = torch.stack([a * (x[..., 0] - x[..., 1]**2), b * (x[..., 1] - x[..., 0]**2)], dim=-1)
    return -0.5 * torch.norm(phi, p=2, keepdim=False, dim=-1)**2
```

```python
def log_density(x):
    mu = map(torch.tensor, [(2., 2.), (-2., -2.)])
    a, b, s = [+0.75, -0.75], [+3.05, -1.05], [+1, -1]

    compo = map(lambda m, a, b, s: log_banana_base(s*(x - m.to(x)), a, b), mu, a, b, s)
    stacked = torch.stack([*compo], dim=0)

    return torch.logsumexp(stacked, dim=0)
```

```python
def log_funnel(x, a=5e-1, b=1.):
#     return -0.5 * (a * x[..., 1]**2 + torch.exp(- b * x[..., 1]) * x[..., 0]**2)
    return -0.5 * (a * (torch.exp(- b * x[..., 0]) - x[..., 1])**2
                   + torch.exp(- b * x[..., 1]) * x[..., 0]**2)
```

```python
# log_density = log_funnel
```

```python
# def log_density(x):
#     return -0.5 * torch.norm(x, p=2, keepdim=False, dim=-1)**2
```

And a plot of it

```python
mesh = torch.meshgrid(2*[torch.linspace(-5, +16, 101)])
marg = torch.stack(mesh, dim=-1).flatten(0, -2)

z = torch.exp(log_density(marg))
```

```python
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, xlabel=r"$\theta_1$", ylabel=r"$\theta_2$",
                    title="2x'blob' density")

ax.contourf(*mesh, z.reshape_as(mesh[0]), levels=51, cmap=plt.cm.terrain)

plt.show()
```

Plot the gradient field

```python
theta = marg.clone().requires_grad_(True)
log_density(theta).mean().backward()

dz = theta.grad.reshape(*mesh[0].shape, -1)
```

```python
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, xlabel=r"$\theta_1$", ylabel=r"$\theta_2$",
                    title="2x'blob' density")

ax.contourf(*mesh, z.reshape_as(mesh[0]), levels=51, cmap=plt.cm.terrain)

if True:
    ax.quiver(*mesh, dz[..., 0], dz[..., 1], pivot='mid',
              color="fuchsia", scale_units='xy', angles='uv', scale=.5, alpha=0.5)

plt.show()
```

Now let's create some samplers!

```python
class Proposal:
    def sample(self, n_samples=1, *, at=None):
        raise NotImplementedError()

    def log_prob(self, x, *, at=None):
        raise NotImplementedError()
```

Simulate paths

```python
def sample_chain(sampler, n_steps=501, n_chains=15):
    chain = [torch.randn(n_chains, 2) * 15]
#     chain = [torch.randn(1, 2).repeat(n_chains, 1) * 15]
    for _ in tqdm.trange(n_steps):
        chain.append(sampler.sample(1, at=chain[-1]))

    return torch.stack(chain, dim=-2)
```

A path plotter using quiver.

```python
def plot_path(path, ax=None, **kwargs):
    ax = plt.gca() if ax is None else ax

    uv, xy = path[1:] - path[:-1], path[:-1]

    size = np.linalg.norm(uv, axis=-1, keepdims=False)
    stuck = path[1:][size == 0]

    ax.scatter(stuck[:, 0], stuck[:, 1], c="k", alpha=0.1)

    return ax.quiver(xy[:, 0], xy[:, 1], uv[:, 0], uv[:, 1],
                     scale_units='xy', angles='xy', scale=1.,
                     **kwargs)
```

Check out the Markov Chain paths

```python
def plot_chains(paths, log_density):
    fig = plt.figure(figsize=(16, 16))
    fig.patch.set_alpha(1.0)

    # get the mesh
    ll, ur = get_rect(paths.flatten(0, -2))
    mesh = torch.meshgrid(*map(torch.linspace, ll, ur, [201, 201]))
    marg = torch.stack(mesh, dim=-1).flatten(0, -2)
    z = torch.exp(log_density(marg))

    ax = fig.add_subplot(111, xlabel=r"$\theta_1$", ylabel=r"$\theta_2$")

    ax.contourf(*mesh, z.reshape_as(mesh[0]), levels=21,
               cmap=plt.cm.terrain, alpha=0.05, zorder=10)

    colours = plt.cm.Accent(np.linspace(0, 1, num=len(paths)))
    for i, col in enumerate(colours):
        pts = paths[i].numpy()
        plot_path(pts, color=col, alpha=.5)
    #     ax.scatter(pts[:, 0], pts[:, 1], color=col, s=10, alpha=0.05)

    plt.show()
```

#### Notation

Consider a measurable space $(\Omega, \mathcal{F}, \mu)$.
The Markov Chain sampler needs:
* *(proposal)* the transition `kernel`
$P \colon \Omega \times \mathcal{F} \to [0, 1]$

In MC-MC we typically consider kernels over Lebsegue carrier measures
$\mu = d{x}$ and defined by

$$
P(x, d\omega)
    = q(\omega \vert x) \mu(d\omega) + r(x) \delta_x(d\omega)
    \,, $$

where $Q(\bullet \vert x) = q(\cdot \vert x) \mu(d\omega)$ is a nonnegative measure
with $Q(\Omega\vert x) \leq 1$ and  $q(x \vert x) = 0$.

What does this notation mean?

Well, $\delta_x(\cdot)$ defines a probability measure on the
measurable space $(\Omega, \mathcal{F})$ according to

$$
\delta_x
\colon \mathcal{F} \to [0, 1]
\colon A \mapsto 1_A(x)
    \,. $$

This notation is a shorthand for (like in SDE)
$$
P(x, dy) = q(y \vert x) dy + r(x) \delta_x(dy)
    \Leftrightarrow
    P(x, A)
        = \int_A q(y \vert x) dy + r(x) \delta_x(dy)
        = Q(A \vert x) + r(x) 1_A(x)
    \,, $$

This implies that for $P(x, \cdot)$ to be a probability measure
we need $r(x) = 1 - Q(\Omega \vert x)$.

#### Did i really need this proof? NO!

$\delta_x(\emptyset) = 1_\emptyset(x) = 0$, and
$$
\delta_x\bigl(\biguplus_{n\geq 1} A_n\bigr)
    = 1_{\uplus_{n\geq 1} A_n}(x)
    = \begin{cases}
    1 & \exists{n\geq1}\, x \in A_n\\
    0 &
    \end{cases}
    = \sum_{n\geq 1} 1_{A_n}(x)
    = \sum_{n\geq 1} \delta_x(A_n)
    \,. $$

Via the MCT this implies that $\int f \delta_x(d\omega) = f(x)$

```python
class Target(Proposal):
    def __init__(self, fn):
        self.fn = fn

    def log_prob(self, x, *, at=None):
        return self.fn(x)
```

```python
class IndependentProposal(Proposal):
    def __init__(self, distribution):
        self.distribution = distribution

    def log_prob(self, x, *, at=None):
        return self.distribution.log_prob(x)

    def sample(self, n_samples=1, *, at):
        *head, n_features = at.shape
        return self.distribution.sample((*head, n_samples)).squeeze(-2)
```

```python
class RandomWalkProposal(Proposal):
    def __init__(self, distribution):
        self.distribution = distribution

    def log_prob(self, x, *, at):
        return self.distribution.log_prob(x - at)

    def sample(self, n_samples=1, *, at):
        *head, n_features = at.shape

        step = self.distribution.sample((*head, n_samples))
        return (at.unsqueeze(-2) + step).squeeze(-2)
```

To evaluate gradients of a scalar function in bulk, observe the following:
$$
(\nabla_\theta f(\theta_i))_{i}
    \colon (\theta_i)_{i} \mapsto
        \bigl(
            \tfrac{\partial}{\partial \theta_j}
                \sum_i f(\theta_i)
        \bigr)_{ji} \mathbf{1}
    \,. $$

```python
def batch_grad(x, f, **kwargs):
    theta = x.clone().requires_grad_(True)
    f(theta, **kwargs).sum().backward()
    return theta.grad
```

```python
def f(x, q):
    return - log_density(x) + 0.5 * torch.norm(q, p=2, dim=-1)**2

from torch.autograd import grad

x = torch.randn(10, 2).requires_grad_(True)
q = torch.randn_like(x).requires_grad_(True)

par = x, q
dx, dq = grad(f(*par).sum(), par, create_graph=False)
```

A handy function for clipping a tenspr's norm:
$$
    x \mapsto x \min\bigl\{\tfrac{C}{\| x \|_p}, 1 \bigr\}
    \,. $$

```python
def clip_norm(x, max_norm, p=2, dim=-1, inplace=True):
    mul = torch.Tensor.mul_ if inplace else torch.Tensor.mul

    # an ell-p norm grad clipping step
    norm = torch.norm(x, p=p, dim=dim, keepdim=True)
    x = mul(x, torch.clamp(max_norm / norm, max=1))

    return x
```

### Langevin proposal

The SDE

$$
dX_t
    = \mu_t(X_t) dt
    + \sigma_t(X_t) dW_t
    \,, $$

can be numerically approximated using Euler-Maruyama method: for $\delta > 0$

$$
X_{t+\delta} - X_t
    = \mu_t(X_t) \delta
    + \sigma_t(X_t) \sqrt{\delta} \xi_t
    \,, $$

where $\xi_t \sim \mathcal{N}(0, I)$.

Let's consider Langevin's Îto diffusion:

$$
d\theta_t
    = \nabla_\theta \log \pi(\theta_t) dt
    + \sqrt{2} dW_t
    \,. $$

The Euler-Maruyama step with $\delta > 0$ yields the following
finite-difference approximation

$$
\theta_{t + \delta}
    = \theta_t
    + \delta \nabla_\theta \log \pi(\theta_t)
    + \sqrt{2 \delta} \xi_t
    \,. $$

Thus the proposal is

$$
q(\theta \vert x)
    = \mathcal{N}\bigl(
        \theta \,\big\vert\,
        x + \delta \nabla_\theta \log \pi(\theta) \vert_{\theta=x},
        2 \delta I
    \bigr)
    \,. $$

Recall that, if $x = \mu+ \sigma \xi$ and $\xi \sim p_\xi$, then
$$
p_x(x)
    = \tfrac1\sigma p_\xi(\tfrac{x - \mu}\sigma)
    \,, \text{ or }\,
    p_\xi(\xi)
        = \sigma p_x(\mu + \sigma \xi)
    \,. $$

```python
import math

class LangevinProposal(Proposal):  # , RandomWalkProposal):
    def __init__(self, proposal, target, delta=1e-2, clip_grad=0):
        self.target, self.delta = target, delta
        self.proposal, self.clip_grad = proposal, clip_grad

    def log_prob(self, x, *, at):
        z = (x - self.mu(at=at)) / math.sqrt(2 * self.delta)

        return self.proposal.log_prob(z) - 0.5 * math.log(2 * self.delta)

    def mu(self, *, at):
        grad = batch_grad(at, self.target.log_prob)

        # an ell-2 norm grad clipping step
        if self.clip_grad > 0:
            grad = clip_norm(grad, self.clip_grad, p=2, dim=-1)

        return at + grad * self.delta

    def sample(self, n_samples=1, *, at):
        *head, n_features = at.shape

        step = math.sqrt(2 * self.delta) * self.proposal.sample((*head, n_samples))
        return (self.mu(at=at).unsqueeze(-2) + step).squeeze(-2)
```

$x \in \Omega^{b_1 \times \ldots \times b_p}$ for $\Omega \subseteq \mathbb{R}$
and any transition kernel must

The MH kernel is given by the density:
$$
P(x, dy)
    = p(x, y) dy + r(x) \delta_x(dy)
    \,. $$

### Metropolis-Hastings Proposal

For some transition density (or pmf) $p(y \vert x)$ the transition kernel
$$
P(x, dy)
%     = p(dy \vert x) + r(x) \delta_x(dy)
    = p(y \vert x) dy + r(x) \delta_x(dy)
    \,, $$
with $r(x) = 1 - \int_\Omega p(dy \vert x) \leq 1$ has $\pi$ as the stationary
distribution if $p$ complies with the detailed (micro-) balance condition:
$
\pi(dx) p(dy \vert x) = \pi(dy) p(dx \vert y)
$.

Indeed, we have
$$
\begin{align}
\int_\Omega \pi(dx) P(x, B)
    &= \int_\Omega \pi(dx) \int_B P(x, dy)
    = \int_\Omega \int_B p(dy \vert x) \pi(dx)
      + \int_\Omega \int_B r(x) \delta_x(dy) \pi(dx)
    \\
    &= \int^x_\Omega \int^y_B p(dy \vert x) \pi(dx)
      + \int_\Omega r(x) 1_B(x) \pi(dx)
    = \int^y_B \int^x_\Omega p(dx \vert y) \pi(dy) + \int_B r(x) \pi(dx)
    \\
    &= \int_B (1 - r(y)) \pi(dy) + \int_B r(x) \pi(dx)
    = \pi(B)
    \,,
\end{align}$$

Thus for a chosen proposal density $q(dy \vert x)$ we need to find
a transition density $p(dy \vert x)$, that satisfies the balance.

**(heuristic)** Let's use importance sampling (analogue of): introduce a rv that
controls the transitions and adjusts the resulting `density` so that is has the
needed mass.

Consider $
p(y\vert x)
    = q(y\vert x) \alpha(y \vert x)
$, where $\alpha$ enforces `reversibility`:
$$
\pi(dx) q(dy\vert x) \alpha(y \vert x)
    = \pi(dy) q(dx\vert y) \alpha(x \vert y)
    \,. $$

Assuming $\pi(dx) = \pi(x) dx$ and $q(dy \vert x) = q(y\vert x) dy$ we
may observe the following:
* If $\pi(x) q(y\vert x) > \pi(y) q(x\vert y)$ then $
\alpha(y \vert x)
    = \tfrac{\pi(y) q(x\vert y)}{\pi(x) q(y\vert x)}
$ and $\alpha(x \vert y) = 1$

* If $\pi(x) q(y\vert x) < \pi(y) q(x\vert y)$ then $
\alpha(x \vert y)
    = \tfrac{\pi(x) q(y\vert x)}{\pi(y) q(x\vert y)}
$ and $\alpha(y \vert x) = 1$

Thus the sought $\alpha$ is
$$
\alpha(y\vert x)
    = \min\Bigl\{
        1, \frac{
            \overbrace{
                \pi(y) q(x\vert y)
            }^{
                y \to x
            }
        }{
            \underbrace{\pi(x) q(y\vert x)}_{
                x \to y
            }
        }
    \Bigr\}
    \,. $$

It is called the *probability of move*

Thus the basic step of MH MC sampler is: given $x_t$ do
1. sample $y \sim q(y \vert x_t)$
2. independently draw $u \sim \mathrm{U}[0, 1]$ and
    * set $x_{t+1} = y$ if $u \leq \alpha(y\vert x_t)$
    * put $x_{t+1} = x_t$ otherwise

```python
class MetropolisHastingsProposal(Proposal):
    def __init__(self, proposal, target):
        self.target, self.proposal = target, proposal

    def sample(self, n_samples=1, *, at):
        *head, n_features = at.shape
        head = [1] if not head else head

        p, q = self.target, self.proposal

        assert n_samples == 1
        prop, curr = q.sample(n_samples, at=at), at
        # curr = curr.unsqueeze(-2)  # DON'T : dimension growth! but gives branching paths!

        # \log \pi(prop) q(curr \vert prop)
        log_alpha  = q.log_prob(curr, at=prop) + p.log_prob(prop)

        # - \log \pi(curr) q(prop \vert curr)
        log_alpha -= q.log_prob(prop, at=curr) + p.log_prob(curr)

        alpha = torch.exp(torch.clamp(log_alpha, max=0))

        accept = torch.rand_like(log_alpha) <= alpha

        # print(accept.float().mean(), alpha)
        self.alpha_, self.log_alpha_ = alpha, log_alpha

        return torch.where(accept.unsqueeze(-1), prop, curr)

```

From [Betancourt (2017) p.11](https://arxiv.org/abs/1701.02434)
* **(first phase)** MC travels towards the typical set. MC-based
estimators have strong bias (*burn-in*)

* **(second phase)** MC `persists through the first sojourn across the typical set`.
Accuracy of estimators improves as the bias from burn-in dampens

* **(third phase)** MC continues and gradually refines its exploration
of the typical set.

Use a random walk transition kernel

```python
from  torch.distributions import MultivariateNormal

proposal = RandomWalkProposal(MultivariateNormal(torch.zeros(2), torch.eye(2)))
rwmh = MetropolisHastingsProposal(proposal, Target(log_density))

plot_chains(sample_chain(rwmh, n_chains=15), log_density)
```

```python
gauss = MultivariateNormal(torch.zeros(2), torch.eye(2))
# ldp = LangevinProposal(gauss, Target(log_density), delta=1e-1, clip_grad=1e1)
ldp = LangevinProposal(gauss, Target(log_density), delta=1e-2, clip_grad=5e2)
ldmh = MetropolisHastingsProposal(ldp, Target(log_density))

plot_chains(sample_chain(ldmh, n_steps=501, n_chains=15), log_density)
```

From [Betancourt (2017) pp. 16-19](https://arxiv.org/abs/1701.02434)

Although MH with Random Walk proposal (aka RandomWalk Metropolis, RWMH) is simple
and intuitive clear, it drammatically suffers from the curse of dimensionality
and the complexity of the target distribution.

> ... the volume exterior to the typical set overwhelms the interior volume
and almost every RWMH chain gets stuck outside of the typical set towards the
tails, due to low acceptance rate, induced by negligible densities. ... In the
worst case RWMH won't even complete a single sojourn through the typical set.

The idea behind Hamiltonian MC is to use first order information,
$\nabla_\theta \log \pi(\theta)$, about the target distribution $\pi(\theta)$
to make informed moves towards the typical set. However, by itself
the gradient pulls towrads a mode of $\pi(\theta)$ and would make
the chain collapse in it, which is not the typical set (what is it then?)

So in HMC we consider $\pi(\theta, m) = \pi(\theta) p(m \vert \theta)$
where we have introduced auxiliary random varaible, momentum $m$. This
`lifts the target distribution onto a joint probability distribution on
pahse space` If the momentum is marginalized, then the original target
density is recovred, which means that during sampling we can simply discard
$m$ when requesting a sampel of $\theta$.

Consider $\pi(\theta, m) = \exp\bigl\{ - H(\theta, m) \bigr\}$. Then for a
trajectory $t\mapsto (\theta_t, m_t)$ so stay on the level set of $H$
we must have $\tfrac{d}{d t} H(\theta_t, m_t) = 0$, i.e.

$$
dH = \nabla_\theta^\top H(\theta_t, m_t) \dot{\theta}_t dt
    + \nabla_m^\top H(\theta_t, m_t) \dot{m}_t dt
    = 0
    \,, $$

which is satisfied when
$$
\dot{\theta}_t = \nabla_m H(\theta_t, m_t)
    \,,\,
    \dot{m}_t = \nabla_\theta H(\theta_t, m_t)
    \,. $$

Typically the Hamiltonian is decomposed as

$$
H(\theta, m)
    = - \log \pi(\theta) p(m \vert \theta)
    = \underbrace{- \log \pi(\theta)}_{\text{potential}}
    + \bigl(
        \underbrace{- \log p(m \vert \theta)}_{\text{kinetic}}
    \bigr)
    \,. $$

From [Betancourt (2017) pp. 27](https://arxiv.org/abs/1701.02434)

The choice of the kinetic energy term is what determines the interaction if the Chain with the target.
Herein lies the scope of HMC design.

The randomness in HMC scheme comes from the randomness of the `momentun lift`,
whereas the traversal of the level of $H$ is deterministic.

Since the HMC trajectory journies along the level set in the phase-space,
it is natural to factorize the canonical distribution $\pi(\theta, m)$,
byt foliating into concentric level sets:

$$
\pi(\theta, m)
    = \pi(x_E \vert H(x_E) = E)
    \, \pi(H(x_E) = E)
    \,. $$

Lifts determine jumps beteween the energy levels, while each tajectory explores
the corresponding level set $\{H(\theta, m) = E\}$. Thus we get two phases

* deterministic traversal of energy level sets (how long we integrate)
* stochastic exploration between level sets (how quickly jumps `diffuse accross energies typical to the energy marginal distribution`

The following kinetic energy is called Euclidean-Gaussian if $G(\theta) = \Sigma$,
and Riemannian-Gaussian if $G(\theta) = $.

* we can optimize $\Sigma$ using the extendend burn-in pahse
[Betancourt (2017) pp. 31](https://arxiv.org/abs/1701.02434)

* if $G(\theta)$ resembles the Hessian of the target, i.e. its Fisher info-matrix,
then the energy level exploration would be uniform and efficient

Choosing the integration time: if we integrate for a short period of time,
we risk having insufficient diversity on the samples as they would tend to
clump together. On the other hand, long integration times can degrade exploration
in case when the level sets are `topologically` compact (is there any other
notion of compactness?).

Dynamic Ergodicity for orbits $\phi = \{(\theta_t, m_t)\colon t \geq 0\}$ states
that a uniform temporal sample form a trajectory resemples a uniform spatial sample
from $\phi$.

Simplectic integrators [Betancourt (2017) pp. 36](https://arxiv.org/abs/1701.02434)

### Let's try to integrate something

The Cauchy problem for ODE $$
\dot{x}
    = A x\,, x(0) = x_0
\,. $$

```python
def euler(x, f, eps=1e-3):
    return x + f(x) * eps
```

```python
grad = lambda x: np.dot(x, np.array([
    [-1., +2.],
    [-10., -1.],
]))

grad = lambda x: np.dot(x, np.array([
    [0, +1.],
    [-1., 0],
]))
```

```python
paths = [np.random.randn(5, 2)]
for _ in range(100):
    paths.append(euler(paths[-1], grad, eps=1e-1))

paths = np.stack(paths, axis=-2)
```

```python
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111)

colours = plt.cm.Reds(np.linspace(0.5, 1, num=len(paths)))
for i, col in enumerate(colours):
    plot_path(paths[i], color=col, alpha=1)
```

```python
np_mesh = np.meshgrid(
    np.linspace(-5, +5, num=101),
    np.linspace(-5, +5, num=101),
)
np_marg = np.stack(np_mesh, axis=-1).reshape(-1, 2)
```

```python
uv = grad(np_marg).reshape(*np_mesh[0].shape, -1)
```

```python
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, xlabel=r"$\theta_1$", ylabel=r"$\theta_2$",
                    title="2x'blob' density")

ax.quiver(*np_mesh, uv[..., 0], uv[..., 1], pivot='mid',
          color="fuchsia", scale=2000., alpha=0.5)

plt.show()
```

### For Hamiltonian MonteCarlo we need a hamiltonian

Lets introduce an auxiliary variable $q \approx \dot{\theta}$ with
$p(q\vert \theta) = \mathcal{N}(q \vert 0, G(\theta))$. Then the
joint density of $(\theta, q)$

$$
p(\theta, q)
    = p(\theta) p(q\vert \theta)
    \propto \exp{(- H(\theta, q))}
    \,. $$

$$
H(\theta, q)
    = - \ell(\theta)
    + \tfrac12\log\det 2 \pi G(\theta)
    + \tfrac12 q^\top G(\theta)^{-1} q
    \,, $$
with $\ell(\theta) = \log p(\theta)$.

Hamiltonian dynamics in the case of $q$ behaving like a momentum
of $\theta$ is
$$
\begin{align}
    \dot{\theta}
        =& \partial_q H(\theta, q)
        \,, \\
   - \dot{q}
        =& \partial_\theta H(\theta, q)
        \,,
\end{align}
$$

Now the derivative of $\log-\det$ is

$$
\partial_\theta \log \det G(\theta)
    = - \partial_\theta \log \det G(\theta)^{-1}
    = - \mathrm{tr}
        \tfrac1{\det G(\theta)^{-1}} (\det G(\theta)^{-1})
            G(\theta)^\top \partial_\theta G(\theta)^{-1}
    = - \mathrm{tr} G(\theta) \partial_\theta G(\theta)^{-1}
    \,, $$

and the second term is

$$
\partial_\theta \tfrac12 q^\top G(\theta)^{-1} q
    = \tfrac12 \mathrm{tr} q q^\top \partial_\theta G(\theta)^{-1}
    \,. $$

And $
\partial_\theta G(\theta)^{-1}
    = - G(\theta)^{-1}
        \bigl( \partial_\theta G(\theta) \bigr)
    G(\theta)^{-1}
$

The dynamics is thus

$$
\begin{align}
    \dot{\theta}
        =& G(\theta)^{-1} q
        \,, \\
    - \dot{q}
%         =& - \partial_\theta \ell(\theta)
%            + \tfrac12 \mathrm{tr} \bigl(
%                q q^\top - G(\theta)
%            \bigr) \partial_\theta G(\theta)^{-1}
        =& - \partial_\theta \ell(\theta)
           - \tfrac12 \mathrm{tr} G(\theta)^{-1} \bigl(
               q q^\top - G(\theta)
           \bigr) G(\theta)^{-1}
           \partial_\theta G(\theta)
        \,,
\end{align}
$$

For $G(\theta) = \Sigma$ we have

$$
\dot{\theta} = \Sigma^{-1} q
    \,, \dot{q} = \partial_\theta \ell(\theta)
    \,, $$

$$
\ddot{x} = F(x)
    \Leftrightarrow
    \dot{x} = v
    \,, \dot{v} = F(x)
\,. $$

$$
\begin{pmatrix}
    \dot{q} \\ \dot{p}
\end{pmatrix}
    = \begin{pmatrix}
        0 & \partial_p \\
        - \partial_q & 0
    \end{pmatrix}
    H(q, p)
    \,, $$

For Euclidean-Gaussian kinetic energy we have

$$
H(\theta, q)
    = - \log p(\theta)
    + \tfrac12\log\det 2 \pi \Sigma
    + \tfrac12 q^\top \Sigma^{-1} q
    \,. $$

```python
def hamiltonian(x, q):
    return - log_density(x) + 0.5 * torch.norm(q, p=2, dim=-1)**2
```

Hence the dynamics obeys

$$
\dot{\theta} = \Sigma^{-1} q
    \,,\, \dot{q} = - \nabla_\theta H(\theta, q) = \nabla_\theta \log p(\theta)
    \,. $$

For $\tau = s t$, $s = \pm 1$ we have:

$$
\frac{d x}{d \tau}
    = \frac{d x}{d t} \frac{d t}{d \tau}
    % = \frac{d x}{d t} \frac1s
    = s \frac{d x}{d t}
    \,, $$

whence the dynamics becomes

$$
\dot{\theta} = \Sigma^{-1} s q
    \,,\, \dot{q} = s \nabla_\theta \log p(\theta)
    \,. $$

Consider an autonomous ODE

$$
\dot{x} = f(x)
    \,, x(0) = x_0
    \,, $$

and suppose $\tau$ is a differentiable time transformation. Then
$y(t) = x(\tau(t))$ satisfies

$$
\dot{y}
    = \frac{d y(t)}{dt}
    = \frac{d x(\tau(t))}{d\tau} \frac{d \tau(t)}{dt}
    = f(x(\tau(t))) \dot{\tau}
    = f(y) \dot{\tau}
    \,, $$

whence $y$ satisfies the ODE

$$
\dot{y} = \dot{\tau} f(y)
    \,, y(\tau^{-1}(0)) = x_0
    \,. $$

In particular if $\tau(t) = t_0 - t$ then

$$
\dot{y} = - f(y)
    \,, y(t_0) = x_0
    \,. $$

```python
def integrate_(H, x, q, *, epsilon=1e-3, M=None, inplace=False, clip_grad=0):
    # applies the one-step leapforg integrator for \ddot{x} = - M \nabla_x H(x)
    r"""Leap-frog integrator for \dot{x} = q, \dot{q} = - M \nabla_x H(x)"""
    add = torch.Tensor.add_ if inplace else torch.Tensor.add
    sub = torch.Tensor.sub_ if inplace else torch.Tensor.sub
    if not inplace:
        x, q = x.clone(), q.clone()

    grad = batch_grad(x, H, q=q)
    while True:
        # q_{\tfrac12} = q_0 - \nabla_x H(x_0) \tfrac\epsilon2
        q = sub(q, grad * epsilon / 2)

        # x_1 = x_0 + M q_{\tfrac12} \epsilon
        x = add(x, epsilon * (q if M is None else torch.matmul(q, M)))

        # q_1 = q_{\tfrac12} - \nabla_x H(x_1) \tfrac\epsilon2
        grad = batch_grad(x, H, q=q)
        if clip_grad > 0:
            grad = clip_norm(grad, clip_grad, p=2, dim=-1)

        q = sub(q, grad * epsilon / 2)

        # preemptively flip the sign of the momentum to ease time-reversibility
        yield x, -q

    # (todo) add stopped process t \wedge n_j
    # (todo) understand what NUTS does and it does
    # (todo) investigate rare NANs
    # (todo) add time direction
```

The Hamiltinian Metropolis Hastings proposal is the same as MH
but with varaites in phase-space and

$$
\alpha(y\vert x)
    = \min\Bigl\{
        1, \frac{
            p(\theta_y, -q_y)
        }{
            p(\theta_x, q_x)
        } \frac{
            h_T((\theta_x, q_x) \vert (\theta_y, -q_y))
        }{
            h_T((\theta_y, -q_y) \vert (\theta_x, q_x))
        }
    \Bigr\}
    \,, $$

where $
h_T((\theta, q) \vert (\theta_0, q_0))
    = \delta_{\theta - \theta_T} \delta_{q - (-q_T)}
$
and $(\theta_T, q_T)$ is the integral of the Hamiltonian ODE at $T$

$$
\begin{align}
    \dot{\theta}
        =& \partial_q H(\theta, q)
        \,, \\
   \dot{q}
        =& - \partial_\theta H(\theta, q)
        \,,
\end{align}
$$

with initial conditions $(\theta_0, q_0)$. Thus

$$
\alpha((\theta_T, -q_T)\vert (\theta_0, q_0))
\alpha_T(\theta_0, q_0)
    = \min\Bigl\{
        1, \frac{
            p(\theta_T, -q_T)
        }{
            p(\theta_0, q_0)
        } \frac{
            h_T((\theta_0, q_0) \vert (\theta_T, -q_T))
        }{
            h_T((\theta_T, -q_T) \vert (\theta_0, q_0))
        }
    \Bigr\}
%     = \min\Bigl\{
%         1, \frac{
%             \exp{-H(\theta_T, -q_T)}
%         }{
%             \exp{-H(\theta_0, q_0)}
%         }
    = \min\Bigl\{
        1, \exp{\{H(\theta_0, q_0) - H(\theta_T, -q_T)\}}
    \Bigr\}
    \,. $$

The need for mnegation is to make sure that momentum is flipped and,
since hamiltonian dynamics is deterministic and time-reversible, that
when integrating back in time from $(\theta_T, -q_T)$ we recover $(\theta_0, q_0)$.

```python
class HamiltonianMHProposal(Proposal):
    def __init__(self, proposal, target, time=2e-2, delta=1e-3, clip_grad=0):
        self.target, self.proposal = target, proposal
        self.delta, self.time = delta, time
        self.clip_grad = clip_grad

    def log_prob(self, x, q, *, at=None):
        return self.target.log_prob(x) + self.proposal.log_prob(q, at=x)

    def hamiltonian(self, x, q, *, at=None):
        return - self.log_prob(x, q, at=at)

    def sample(self, n_samples=1, *, at):
        *head, n_features = at.shape
        head = [1] if not head else head

        assert n_samples == 1
        # sample the initial momentum q (from kinetic)
        curr = at, self.proposal.sample(n_samples, at=at)

        # integrate \dot{x} = q, \dot{q} = -\nabla_x H(x, q)
        #  from x_0, q_0 until T flipping the sign
        integrator = integrate_(self.hamiltonian, *curr,
                                epsilon=self.delta, inplace=False,
                                clip_grad=self.clip_grad)
        for _, prop in zip(range(-int(-self.time // self.delta)), integrator):
            # consume the integrator until the required number of steps is done
            pass
        del integrator

        # \log q(curr \vert prop) \pi(prop) - \log q(prop \vert curr) \pi(curr)
        log_alpha = self.log_prob(*prop) - self.log_prob(*curr)
        alpha = torch.exp(torch.clamp(log_alpha, max=0))
        accept = torch.rand_like(log_alpha) < alpha

        self.alpha_, self.log_alpha_ = alpha, log_alpha

        return torch.where(accept.unsqueeze(-1), prop[0], curr[0])
```

Let's integrate some!

```python
x = torch.randn(255, 2)
q = torch.randn_like(x)
until, epsilon = 0.25, 5e-3

int_ = integrate_(hamiltonian, x, q, epsilon=epsilon, inplace=False)
path = [(x, q)]
path.extend((x, q) for _, (x, q) in zip(tqdm.trange(-int(-until // epsilon)), int_))
```

```python
int_ = integrate_(hamiltonian, *path[-1], epsilon=epsilon, inplace=False)

path.extend((x, q) for _, (x, q) in zip(tqdm.trange(-int(-until // epsilon)), int_))
```

```python
x, q = zip(*path)
paths = torch.stack(x, dim=-2)
```

```python
torch.norm(paths[:, 0] - paths[:, -1], dim=-1)
```

```python
ll, ur = get_rect(paths.flatten(0, -2))

mesh = torch.meshgrid(*map(torch.linspace, ll, ur, [201, 201]))
marg = torch.stack(mesh, dim=-1).flatten(0, -2)

z = torch.exp(log_density(marg))
```

```python
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111)

ax.contourf(*mesh, z.reshape_as(mesh[0]), levels=51, cmap=plt.cm.terrain)

colours = plt.cm.Reds(np.linspace(0.5, 1, num=len(paths)))
for i, col in enumerate(colours):
    plot_path(paths[i][-100:], color=col, alpha=0.5)
```

Simulate paths

```python
from  torch.distributions import MultivariateNormal

idp = IndependentProposal(MultivariateNormal(torch.zeros(2), torch.eye(2)))
hmh = HamiltonianMHProposal(idp, Target(log_density), time=2e-1, delta=1e-2, clip_grad=100)

plot_chains(sample_chain(hmh, n_chains=15), log_density)
```

```python
assert False
```

```python
# run the sampler
from pyhmc import hmc
```

```python
# !pip install git+https://github.com/rmcgibbo/pyhmc
```

```python
def logprob(x):
    x = torch.from_numpy(x)

    logp = log_density(x)
    grad = batch_grad(x, log_density)

    return logp.numpy(), grad.numpy()
```

```python
from pyhmc import hmc
```

```python
samples = hmc(logprob, x0=np.random.randn(2), args=(), n_samples=10000)
```

```python
plot_chains(torch.from_numpy(samples[np.newaxis]), log_density)
```

```python
paths[0][torch.isfinite(paths[0])]
```

```python
mask = (~torch.isfinite(paths)).sum(dim=-1).sum(dim=-1)
paths[mask > 0]
```

```python
def leap_frog(nabla, q, p, grad=None, eps=0.01):
    r"""Leap-frog integrator for \dot{p} = - \nabla_q V(q), \dot{q} = p"""
    # p_{\tfrac12} = p_0 - \tfrac\epsilon2 \nabla_q V(q_0)
    grad = nabla(q) if grad is None else grad  # .detach()
    p.sub_(- eps * grad / 2)

    # q_1 = q_0 + \epsilon p_{\tfrac12}
    q.add_(eps * p)

    # p_1 = p_{\tfrac12} - \tfrac\epsilon2 \nabla_q V(q_1)
    grad = nabla(q)
    p.sub_(- eps * grad / 2)
    return p, q, grad

```

```python
mu = torch.tensor([
    [-3., -3.],
    [+3., -3.],
    [+3., +3.],
])
```

```python
cov = torch.tensor[]
```

```python
torch.triangular_solve()
```

```python
def log_prob(value, mu, loc):
    pass
```

```python
a = Normal(mu, 1.)
```

```python
target.rsample((10,)).mean(0)
```

```python
target.log_prob??
```

```python
target.scale
```

```python
q = mu[0].clone().requires_grad_(True)
target.log_prob(q).mean(0)
```

```python

```

```python
x = torch.randn(15, 2)
q = torch.randn_like(x)
eps = 5e-2

# applies the one-step leapforg integrator for \ddot{x} = \nabla_x \ell(x)
r"""Leap-frog integrator for \dot{x} = q, \dot{q} = \nabla_x \ell(x)"""

paths = [x]
grad = batch_grad(x, log_density)
for _ in range(2500):
    # (todo) add stopped process t \wedge n_j
    # (todo) understand what NUTS does and it does
    # (todo) investigate rare NANs

    # leap-frog integrator of \ddot{x} = f(x)
    # p_{\tfrac12} = p_0 + \nabla_q V(q_0) \tfrac\epsilon2
    q_half = q + grad * eps / 2

    # q_1 = q_0 + p_{\tfrac12} \epsilon
    x = x + q_half * eps

    # p_1 = p_{\tfrac12} + \nabla_q V(q_1) \tfrac\epsilon2
    grad = batch_grad(x, log_density)
    q = q_half + grad * eps / 2

    paths.append(x)

paths = torch.stack(paths, dim=-2)
```

```python

```

```python
until, epsilon = 1., 1e-3
```

```python
x = torch.randn(15, 2)
q = torch.randn_like(x)
```

```python
x, q
```

```python
integrator = integrate_(hamiltonian, x, q, epsilon=epsilon, inplace=True)
for _ in zip(range(-int(-until // epsilon)), integrator):
    pass
```

```python

```

```python
paths = sample_chain(ldmh, n_chains=15)
```

```python
paths[0]
```

```python
clip_grad = 10.
grad = batch_grad(paths[0], ldmh.target.log_prob)
```


# Activations
> Activations used between layers

```
# default_exp activations
```

```
# export
import torch
import torch.nn as nn
import torch.nn.functional as F
```

```
# hide
import matplotlib.pyplot as plt

torch.Tensor.ndim = property(lambda x: x.dim())
%matplotlib inline
```

```
# export
def relu(inplace=True):
    return nn.ReLU(inplace=inplace)

ReLU = Relu = relu
```

```
a = torch.randn(3, 3)
a, relu()(a)
```

    (tensor([[0.0000, 0.0000, 0.1396],
             [0.7610, 0.0000, 0.1118],
             [0.0000, 0.0000, 0.0000]]),
     tensor([[0.0000, 0.0000, 0.1396],
             [0.7610, 0.0000, 0.1118],
             [0.0000, 0.0000, 0.0000]]))

```
x = torch.arange(-5, 3, 0.05)
y = Relu(inplace=False)(x)
plt.plot(x, y);
```

```
# export
def LeakyReLU(inplace=True, negative_slope=0.01):
    return nn.LeakyReLU(negative_slope, inplace=inplace)

leaky_relu = leakyrelu = LeakyReLU
```

```
a = torch.randn(3, 3)
a, leaky_relu()(a)
```

    (tensor([[ 0.4458,  0.0084,  1.1934],
             [-0.0066, -0.0087, -0.0084],
             [ 0.1442,  0.6694,  0.2855]]),
     tensor([[ 0.4458,  0.0084,  1.1934],
             [-0.0066, -0.0087, -0.0084],
             [ 0.1442,  0.6694,  0.2855]]))

```
x = torch.arange(-5, 3, 0.05)
y = LeakyReLU(inplace=False, negative_slope=0.5)(x)
plt.plot(x, y);
```

## Mish
$$
\begin{align}
\text{mish}(x) &= \text{tanh}\left(\text{softplus}(x)\right)\\
\text{tanh}(x) &= \frac{\exp(2x) - 1}{\exp(2x) + 1}\\
\text{softplus}(x) &= \ln(1 + \exp(x))
\end{align}
$$

```
# export
class Mish(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x * (torch.tanh(F.softplus(x)))

mish = Mish
```

```
x = torch.arange(-5, 3, 0.05)
y = Mish()(x)
plt.plot(x, y);
```

## Sigmoid Activation
Usually used to convert a number into a probability:
$$
\sigma(x) = \frac{1}{1+\exp(-x)}
$$

```
# export
sigmoid = nn.Sigmoid
```

```
sigmoid()(torch.randn((3,1)))
```

    tensor([[0.3305],
            [0.4534],
            [0.6104]])

## Softmax Activation
Converts multicategorical numerical numbers to a probability distribution. Suppose $y_ij$ is the output for the i-th instance for the j-th category. Then the probability is:
$$
p(y_{ij}) = \frac{\exp(y_{ij})}{\sum_j \exp(y_{ij})}
$$

```
# export
softmax = nn.Softmax(dim=-1)
```

```
y = torch.randn((2,3))
p = softmax(y)
print(p)
print(p.sum(dim=-1, keepdim=True))
```

    tensor([[0.2092, 0.1568, 0.6340],
            [0.2197, 0.3284, 0.4519]])
    tensor([[1.],
            [1.]])

```
# export
__activations__ = {
    'relu': relu(),
    'leaky_relu': leaky_relu(),
    'mish': mish(),
    'sigmoid': sigmoid(),
    'softmax': softmax
}

def get_activation(activation):
    return __activations__[activation.lower()]
```

```
# hide
from nbdev.export import *
notebook2script()
```

    Converted Activations.ipynb.
    Converted Layers.ipynb.
    Converted Model.ipynb.
    Converted data.ipynb.
    Converted index.ipynb.
    Converted losses.ipynb.

```

```


```python
%config IPCompleter.greedy = True
%config InlineBackend.figure_format = 'retina'
%matplotlib inline
%load_ext tensorboard

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sn
import tensorflow as tf
from datetime import datetime

pd.set_option('mode.chained_assignment', None)
sn.set(rc={'figure.figsize':(9,9)})
sn.set(font_scale=1.4)

# make results reproducible
seed = 0
np.random.seed(seed)
tf.random.set_seed(seed)

!pip install pydot
!rm -rf ./logs/func/
```

    Requirement already satisfied: pydot in /Users/samholt/anaconda3/lib/python3.7/site-packages (1.4.1)
    Requirement already satisfied: pyparsing>=2.1.4 in /Users/samholt/anaconda3/lib/python3.7/site-packages (from pydot) (2.4.5)

# Working with Tensorflow 2.0

[**TensorFlow**](https://www.tensorflow.org/) is an end-to-end open source platform for machine learning. It has a comprehensive, flexible ecosystem of tools, libraries, and community resources that lets researchers push the state-of-the-art in ML and developers can easily build and deploy ML-powered applications.

TensorFlow library is a free and open-source software library for *dataflow programming* (models a program as a directed graph of the data flowing between operations, thus implementing dataflow principles and architecture) and *differentiable programming* (programs can be differentiated throughout, usually via automatic differentiation. This allows for gradient based optimization of parameters in the program, often via gradient descent) across a range of tasks. It is a symbolic math library, and is also used for machine learning applications such as neural networks. It is used for both research and production at Google.

TensorFlow was originally developed by researchers and engineers working on the Google Brain team within Google's Machine Intelligence Research organization to conduct machine learning and deep neural networks research. The system is general enough to be applicable in a wide variety of other domains, as well. It is one of the most popular libaries for machine learning with neural networks.

TensorFlow provides stable Python and C++ APIs, as well as non-guaranteed backward compatible API for other languages.

### Tensorflow tensors

A [tensor](https://www.tensorflow.org/guide/tensor) is a generalization of vectors and matrices to potentially higher dimensions. Internally, TensorFlow represents tensors as n-dimensional arrays of base datatypes (similar to `numpy.ndarray`).

When we write a TensorFlow program, we mainly operate and pass around `tf.Tensor` objects, which represents a partially defined data object that will eventually produce a value. In a Tensorflow program, we specify how these `tf.Tensor`'s link into a computation graph. We can then run the computation graph to get the output (based on an input), and use differential properties of the graph to optomise for parameter (`tf.Tensor` nodes) in the program. `tf.Tensor`'s can reside in GPU or CPU memory. Tensorflow supports a rich library of operations (`tf.add`, `tf.matmul`, `tf.linalg.inv` etc) that consume and produce `tf.Tensors`'s. These operations also automatically convert to native Python types too.

```python
import tensorflow as tf
print("TensorFlow version : {}\n\n".format(tf.__version__))

print(tf.add(1, 3))
print(tf.add([2, 3], [6, 6]))
print(tf.square(2))
print(tf.reduce_sum([4, 3, 2]))

# Operator overloading is also supported
print(tf.square(3) + tf.square(2))
```

    TensorFlow version : 2.1.0

    tf.Tensor(4, shape=(), dtype=int32)
    tf.Tensor([8 9], shape=(2,), dtype=int32)
    tf.Tensor(4, shape=(), dtype=int32)
    tf.Tensor(9, shape=(), dtype=int32)
    tf.Tensor(13, shape=(), dtype=int32)

Each `tf.Tensor` has a datatype and a shape

```python
x = tf.constant([[4, 3]])
print(x)
print(x.shape)
print(x.dtype)
```

    tf.Tensor([[4 3]], shape=(1, 2), dtype=int32)
    (1, 2)
    <dtype: 'int32'>

Tensors differ from numpy arrays as they are *immutable*, can be computed and stored in the GPU or CPU.

However it easy to convert from a `tf.Tensor` to a `np.ndarray`. Tensorflow automatically converts numpy arrays to `tf.Tensors`.

```python
ndarray = np.array([[1, 2], [3, 4]])

# Tensorflow converts ndarray to a tensor after an operation
tensor = tf.multiply(ndarray, 2)
print('Tensor :\n', tensor)

# Tensorflow converts ndarray to a tensor after operator overloading operation
tensor = tf.constant(2) * ndarray
print('\nTensor :\n', tensor)

# Numpy can convert a tensor to a numpy ndarray
numpy_array = np.multiply(tensor, 2)
print('\nNumpy ndarray :\n', numpy_array)

# It does this by calling the .numpy() method explicity
print('\nNumpy ndarray : \n', tensor.numpy())
```

    Tensor :
     tf.Tensor(
    [[2 4]
     [6 8]], shape=(2, 2), dtype=int64)

    Tensor :
     tf.Tensor(
    [[2 4]
     [6 8]], shape=(2, 2), dtype=int32)

    Numpy ndarray :
     [[ 4  8]
     [12 16]]

    Numpy ndarray :
     [[2 4]
     [6 8]]

We can also index tensors with the same convention as numpy.

```python
print(tensor[0, 0])
print('\n', tensor[:, 1])
```

    tf.Tensor(2, shape=(), dtype=int32)

     tf.Tensor([4 8], shape=(2,), dtype=int32)

Tensorflow operations can be accelerated by parralelised to be computed on a GPU. By default, Tensorflow automatically decides whether to use the GPU or CPU for an operation (copying the memory between them if necessary).

```python
x = tf.random.uniform([4, 4])

print("List if there is a GPU available : ",
      tf.config.experimental.list_physical_devices("GPU"))

print("\nIs the Tensor on GPU #0: ", x.device.endswith('GPU:0'))
```

    List if there is a GPU available :  []

    Is the Tensor on GPU #0:  False

The `tf.Tensor` has the same data type `tf.Dtype` (e.g. {`int32`, `float32`, `string`}) for each element. There are some special types of tensors, such as `tf.Variable`, `tf.constant`, `tf.SparseTensor`, `tf.placeholder`. The data of a tensor is immutable, meaning that in a single run of the graph the data doesn't change.

Depending on the dimension of the tensor, we get the following

|Dimension|Classical mathematical form|
|---|---|
| 0       | Scalar                    |
| 1       | Vector                    |
| 2       | Matrix                    |
| 3       | 3-Tensor                  |
| n       | n-Tensor                  |

A variable, i.e. `tf.Variable` is the best way to represent shared state, and the value can be changed by running operations on it. Higher level libraries like `keras` use `tf.Variable`'s to store model parameters.

```python
zeros = tf.Variable(tf.zeros([2, 3]))
print(zeros)

scalar = tf.Variable(5)
print('\n', scalar)

# To assign a value to an already set variable we use the .assign() method
scalar.assign(4)
print('\n', scalar)

# We can also explicitly read a value using the .read_value() method
print('\n', scalar.read_value())
```

    <tf.Variable 'Variable:0' shape=(2, 3) dtype=float32, numpy=
    array([[0., 0., 0.],
           [0., 0., 0.]], dtype=float32)>

     <tf.Variable 'Variable:0' shape=() dtype=int32, numpy=5>

     <tf.Variable 'Variable:0' shape=() dtype=int32, numpy=4>

     tf.Tensor(4, shape=(), dtype=int32)

### Eager execution

A very useful feature in tensorflow 2.0, enabled by default is *eager execution*, which is an imperative programming enviroment that evaluates operations immediately, without building graphs: operations return definitive values rather than constructing a computational graph to run later. This makes it easy to start working with Tensorflow and debug models. This allows use of Pythons data structures, flexibility how we can structure our code and using Python control flow as well. This supports GPU operations as well.

```python
print('Tensorflow {}, executing early by default\nAre we evaluating operations immediately ? : {}'.format(
    tf.__version__, tf.executing_eagerly()))
```

    Tensorflow 2.1.0, executing early by default
    Are we evaluating operations immediately ? : True

```python
# When we run Tensorflow operations the result is immediately returned
x = tf.constant(2, name='x')
y = tf.constant(3, name='y')
z = x * y

print('z = {}\n As'.format(z))
print('z = {} = x ({}) * y ({})'.format(z, x, y))
```

    z = 6
     As
    z = 6 = x (2) * y (3)

```python
# Visualise the computation graph

# The function to be traced.
@tf.function
def my_func(x, y):
    # A simple hand-rolled layer.
    return x*y

# Set up logging.
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
logdir = 'logs/func/%s' % stamp
writer = tf.summary.create_file_writer(logdir)

# Sample data for your function.
x = tf.constant(2, name='x')
y = tf.constant(3, name='y')

# Bracket the function call with
# tf.summary.trace_on() and tf.summary.trace_export().
tf.summary.trace_on(graph=True, profiler=True)
# Call only one tf.function when tracing.
z = my_func(x, y)
with writer.as_default():
    tf.summary.trace_export(
        name="my_func_trace",
        step=0,
        profiler_outdir=logdir)

print('z = {} = x ({}) * y ({})'.format(z, x, y))

%tensorboard --logdir logs/func
```

    z = 6 = x (2) * y (3)

## Automatic differentiation

[Automatic differentiation](https://www.wikiwand.com/en/Automatic_differentiation) (AD) is a set of techniques to numerically evaluate the derivative of a function specified by a computer program. AD exploits the fact that every computer program, no matter how complicated, executes a sequence of elementary arithmetic operations (addition, subtraction, multiplication, division, etc.) and elementary functions (exp, log, sin, cos, etc.). By applying the chain rule repeatedly to these operations, derivatives of arbitrary order can be computed automatically, accurately to working precision, and using at most a small constant factor more arithmetic operations than the original program.

#### The chain rule, forward and reverse accumulation

Fundamental to AD is the decomposition of differentials provided by the chain rule. For the simple composition

$$\begin{align}
y &= f(g(h(x))) = f(g(h(w_0))) = f(g(w_1)) = f(w_2) = w_3 \\
w_0 &= x \\
w_1 &= h(w_0) \\
w_2 &= g(w_1) \\
w_3 &= f(w_2) = y
\end{align}$$

the chain rule gives

$$\frac{dy}{dx} = \frac{dy}{dw_2} \frac{dw_2}{dw_1} \frac{dw_1}{dx}$$

Usually, two distinct modes of AD are presented, *forward accumulation* (or *forward mode*) and *reverse accumulation* (or *reverse mode*). Forward accumulation specifies that one traverses the chain rule from inside to outside (that is, first compute $dw_1/dx$ and then $dw_2/dw_1$ and at last $dy/dw_2$), while reverse accumulation has the traversal from outside to inside (first compute $dy/dw_2$ and then $dw_2/dw_1$ and at last $dw_1/dx$). More succinctly,
1. *forward accumulation* computes the recursive relation: $\frac{dw_i}{dx} = \frac{dw_i}{dw_{i-1}} \frac{dw_{i-1}}{dx}$ with $w_3 = y$, and,
2. *reverse accumulation* computes the recursive relation: $\frac{dy}{dw_i} = \frac{dy}{dw_{i+1}} \frac{dw_{i+1}}{dw_{i}}$ with $w_0 = x$.
Generally, both forward and reverse accumulation are specific manifestations of applying the operator of program composition, fixing the appropriate one of the two mappings $(w,y)$. Tensorflow uses *reverse accumulation*.

#### Reverse accumulation

In reverse accumulation AD, the *dependent variable* to be differentiated is fixed and the derivative is computed *with respect to* each sub-expression recursively. In a pen-and-paper calculation, the derivative of the *outer* functions is repeatedly substituted in the chain rule:

$$\frac{\partial y}{\partial x}
= \frac{\partial y}{\partial w_1} \frac{\partial w_1}{\partial x}
= \left(\frac{\partial y}{\partial w_2} \frac{\partial w_2}{\partial w_1}\right) \frac{\partial w_1}{\partial x}
= \left(\left(\frac{\partial y}{\partial w_3} \frac{\partial w_3}{\partial w_2}\right) \frac{\partial w_2}{\partial w_1}\right) \frac{\partial w_1}{\partial x}
= \cdots$$

In reverse accumulation, the quantity of interest is the *adjoint*, denoted with a bar ($w̄$); it is a derivative of a chosen dependent variable with respect to a subexpression $w$:

$\bar w = \frac{\partial y}{\partial w}$

[[2](https://en.wikipedia.org/wiki/File:ForwardAccumulationAutomaticDifferentiation.png)]

Reverse accumulation traverses the chain rule from outside to inside, or in the case of the computational graph in the diagram, from top to bottom. The example function is scalar-valued, and thus there is only one seed for the derivative computation, and only one sweep of the computational graph is needed to calculate the (two-component) gradient.
This is only half the work when compared to forward accumulation, but reverse accumulation requires the storage of the intermediate variables $w-i$ as well as the instructions that produced them in a data structure known as a Wengert list (or "tape") which may consume significant memory if the computational graph is large.
This can be mitigated to some extent by storing only a subset of the intermediate variables and then reconstructing the necessary work variables by repeating the evaluations, a technique known as rematerialization. Checkpointing is also used to save intermediary states (Tensorflow supports checkpointing).

The operations to compute the derivative using reverse accumulation are shown in the table below (note the reversed order):
$$\begin{array}{l}
\text{Operations to compute derivative}
\\ \hline
\bar w_5 = 1 \text{ (seed)}
\\
\bar w_4 = \bar w_5
\\
\bar w_3 = \bar w_5
\\
\bar w_2 = \bar w_3 \cdot w_1
\\
\bar w_1 = \bar w_3 \cdot w_2 + \bar w_4 \cdot \cos w_1
\end{array}
$$

The data flow graph of a computation can be manipulated to calculate the gradient of its original calculation. This is done by adding an adjoint node for each primal node, connected by adjoint edges which parallel the primal edges but flow in the opposite direction. The nodes in the adjoint graph represent multiplication by the derivatives of the functions calculated by the nodes in the primal. For instance, addition in the primal causes fanout in the adjoint; fanout in the primal causes addition in the adjoint; In terms of weight matrices, the adjoint is the transpose. Addition is the covector $[1 \cdots 1]$, since $[1 \cdots 1]
\left[\begin{smallmatrix}x_1 \\ \vdots \\ x_n \end{smallmatrix}\right] = x_1 + \cdots + x_n,$ and fanout is the vector $\left[\begin{smallmatrix}1 \\ \vdots \\ 1 \end{smallmatrix}\right],$ since $\left[\begin{smallmatrix}1 \\ \vdots \\ 1 \end{smallmatrix}\right][x] = \left[\begin{smallmatrix}x \\ \vdots \\ x \end{smallmatrix}\right].$ a unary function $y=f(x)$ in the primal causes $x̄=ȳf′(x)$ in the adjoint; etc.

Reverse accumulation is more efficient than forward accumulation for functions $f : ℝ^{n} → ℝ^{m}$ with $m ≪ n$ as only $m$ sweeps are necessary, compared to $n$ sweeps for forward accumulation.

Reverse mode AD was first published in 1976 by Seppo Linnainmaa.

Backpropagation of errors in multilayer perceptrons, a technique used in machine learning, is a special case of reverse mode AD.

### AD and Gradient Tape in TensorFlow

The `tf.GradientTape` records operations for automatic differentiation, computing the gradient of a computation with respect to its inputs. Operations are recorded if they are executed within this context manager and at least one of their inputs is being "watched". Tensorflow then uses that tape and the gradients associated with each recorded operation to compute the gradients of a "recorded" computation using *reverse accumulation* (explained above).

E.g.

Suppose we have the network for

$$ y = x^2 $$

We can compute analytically $\frac{dy}{dx}$ $$\frac{dy}{dx}=2x$$

```python
x = tf.constant(5.0)

with tf.GradientTape() as t:
    t.watch(x)
    y = x * x

dy_dx = t.gradient(y, x)

print('dy_dx : ', dy_dx)
```

    dy_dx :  tf.Tensor(10.0, shape=(), dtype=float32)

```python
x = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)

with tf.GradientTape() as t:
    t.watch(x)
    k = tf.matmul(x, x)
    y = tf.add(k, tf.ones((2, 2)))

dy_dk = t.gradient(y, k)

print('dy_dk : ', dy_dx)
```

    dy_dk :  tf.Tensor(10.0, shape=(), dtype=float32)

```python
x = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)

with tf.GradientTape(persistent=True) as t:
    t.watch(x)
    k = tf.matmul(x, x)
    y = tf.add(k, tf.ones((2, 2)))

dy_dk = t.gradient(y, k)
dy_dx = t.gradient(y, x)

del t  # Drop the rerefence to the tape when finished with it

print('dy_dk : ', dy_dx)
print('dy_dx : ', dy_dx)
```

    dy_dk :  tf.Tensor(
    [[ 7. 11.]
     [ 9. 13.]], shape=(2, 2), dtype=float32)
    dy_dx :  tf.Tensor(
    [[ 7. 11.]
     [ 9. 13.]], shape=(2, 2), dtype=float32)

Tapes record operations as they execute therefore they naturally handle Python control flow, such as `if` and `while`

```python
def f(x, y):
    output = 6.0
    for i in range(y):
        if i > 2 and i < 6:
            output = tf.multiply(output, x)
    return output

def grad(x, y):
    with tf.GradientTape() as t:
        t.watch(x)
        out = f(x, y)
    return t.gradient(out, x)

x = tf.Variable(2.0)

print(grad(x, 6))
print(grad(x, 5))
print(grad(x, 4))
```

    tf.Tensor(72.0, shape=(), dtype=float32)
    tf.Tensor(24.0, shape=(), dtype=float32)
    tf.Tensor(6.0, shape=(), dtype=float32)

```python
x = tf.constant(3.0)
with tf.GradientTape() as g:
    g.watch(x)
    with tf.GradientTape() as gg:
        gg.watch(x)
        y = x * x

dy_dx = gg.gradient(y, x)
d2y_dx2 = g.gradient(dy_dx, x)
```

## Ridge Regression with TensorFlow

Recall that we looked at ridge regression in section 2, to recap:

https://en.wikipedia.org/wiki/Tikhonov_regularization
$y=wx_i$ where $w\in{\rm I\!R}^{n}$

i.e. if dimensions $n=2$;    $y_{i} = f(x_{i}) = \begin{bmatrix} w_{0}, w_{1} \end{bmatrix} \begin{bmatrix} x_{0} \\ x_{1} \end{bmatrix} = w_{0}x_{0} + w_{1}x_{1}$

Model : $\hat{y}=\mathbf{X}w$

Cost function : $$ J(w)= \sum_i (w^T x_i - y_i)^2 = {\bf ||} \mathbf{X} w-y {\bf ||}^2_{2} $$

$$\min_{w} \, (\mathbf{y} - \mathbf{X} w)^\mathsf{T}(\mathbf{y} - \mathbf{X} w) $$

$$\frac{\partial (\mathbf{y} - \mathbf{X} w)^\mathsf{T}(\mathbf{y} - \mathbf{X} w)}{\partial w}=-2X^T(y-\mathbf{X} w)$$

$$\frac{\partial (\mathbf{y} - \mathbf{X} w)^\mathsf{T}(\mathbf{y} - \mathbf{X} w)}{\partial w}=-2X^T(y-\mathbf{X} w) = 0$$

$$y=\mathbf{X} w$$

### Solution with the Analytic (Exact algebraic solution)
i.e. the analytical solution for the model parameters is to compute :$$w = \mathbf{X}^{-1} y$$

```python
# Recall the Ridge Regression example from section 2-2
from sklearn import datasets
digits = datasets.load_digits()

(X, y) = datasets.load_digits(return_X_y=True)

# Split the labelled dataset into a training set and test set
index_to_split = -1
X_train = X[:index_to_split]
y_train = np.array([y[:index_to_split]])
X_test = X[index_to_split:]
y_test = np.array([y[index_to_split:]])

# Calculate the analytical solution
X_train_inv = tf.linalg.pinv(X_train)
y_train = tf.transpose(tf.Variable(y_train, dtype=tf.float64))

w = tf.matmul(X_train_inv, y_train)
print('w: {}'.format(w.numpy().flatten()))
```

    w: [ 1.62642066e-14  9.40249350e-02 -3.67939617e-03 -7.72167429e-03
      7.49514697e-02  1.15243563e-02 -2.72158636e-02 -7.31904304e-03
      1.00102690e+00 -2.84132215e-02  1.18887868e-01  6.59994113e-02
     -5.58787175e-02 -7.04137172e-02  9.65429164e-02  2.56033005e-01
     -7.29405068e-01  2.40605548e-02  7.76124220e-02 -2.29357716e-02
     -5.66039721e-02  5.81134072e-02 -4.91673806e-02 -2.62755693e-01
     -9.08254598e-01 -1.49425933e-01  5.59359249e-02  8.97465435e-02
      8.39772547e-02  9.83482587e-02  1.86692450e-03 -2.97298764e+00
      2.33680455e-14 -1.54851687e-01 -8.97570476e-03  1.39494684e-01
     -3.67398683e-02  5.46737441e-02 -9.53756017e-03  9.30683775e-16
      1.00930061e-01  1.23968401e-01 -1.37042140e-02  5.34285543e-03
      1.30949507e-01  5.51389202e-02  2.24650594e-02  7.76787978e-03
      6.14807150e-01  2.67348801e-02  1.20233616e-03 -6.19644555e-02
     -2.06960828e-01 -3.39245068e-02  1.05783054e-01 -1.40682284e-01
     -1.01420937e+00 -1.13249749e-01  2.06527116e-02 -4.39417720e-02
      1.85650152e-02 -6.65628180e-02  1.14557345e-02 -5.23335705e-02]

```python
# Test the model
y_hat = X_test @ w
print(y_hat)
print(w.shape)
```

    tf.Tensor([[9.30939946]], shape=(1, 1), dtype=float64)
    (64, 1)

```python
print('y_hat: {}, rounds to {}'.format(y_hat, np.round(y_hat).item(0)))
print('Actual y label: {}'.format(y_test))
plt.figure()
plt.imshow(16 - digits.images[-1], cmap='gray')
plt.show()
```

### Solution with Batch Gradient Descent (Using TensorFlow)

Lets train the ridge regression classifier this time using gradient descent, recall the gradient of the cost function is :

$$ J(w)= \sum_i (w^T x_i - y_i)^2 = {\bf ||} \mathbf{X} w-y {\bf ||}^2_{2} $$

Our update equation for $w$ is

$$ w \leftarrow w - \eta \frac{\partial J(\bf{w}, y_i)}{\partial \bf{w}} $$

We will use `tf.Variable`'s to represent the parameters of the model, and we can use `tf` methods to read and update the stored value such as `tf.assign_sub` etc. These are mutable in the tensorflow computation graph. We will do the following:

1. Define the model
2. Define a loss function
3. Train using the training data and use an *optomizer* to adjust the variables to fit the data

To recap our model is : $\hat{y}=\mathbf{X}w$

```python
# 1. Defining our model
class RidgeRegression(object):
    def __init__(self):
        # Intialise parameter values
        self.w = tf.Variable(tf.random.normal(
            [64, 1], stddev=0.1, dtype=tf.float64, seed=0))

    # Actual model
    def __call__(self, X):
        return X @ self.w

# 2. Define the loss function

def loss(target_y, predicted_y):
    return tf.reduce_mean(tf.square(predicted_y - target_y))

# 3. Using gradient descent define a training function

def train(ridge_regression, train_X, train_y, learning_rate):
    with tf.GradientTape() as t:
        t.watch(ridge_regression.w)
        current_loss = loss(train_y, ridge_regression(train_X))
    dw = t.gradient(current_loss, ridge_regression.w)
    ridge_regression.w.assign_add(- learning_rate * dw)

# # Putting it all together on the digits dataset
ridge_regression = RidgeRegression()

print('Starting loss : ', loss(y_train, ridge_regression(X_train)))
print('Lowest possible training loss : {}\n'.format(loss(y_train, X_train @ w)))

# Collect history of loss function
losses = []
epochs = range(200)
for epoch in epochs:
    current_loss = loss(y_train, ridge_regression(X_train))
    losses.append(current_loss)
    train(ridge_regression, X_train, y_train, learning_rate=0.0001)
    if epoch < 10 or epoch > 190:
        print('Epoch {}: loss={}'.format(epoch, current_loss))

# Let's plot it all
plt.figure(figsize=(13, 13))
plt.plot(epochs, losses, 'r')
plt.plot([loss(y_train, X_train @ w)] * len(epochs), 'b--')
plt.legend(['Loss function', 'Lowest possible training loss'])
plt.xlabel('Epochs')
plt.ylabel('$J(w)$ Cost Function')
plt.show()
```

```python

```