```python
from sympy.physics.mechanics import *
from sympy import symbols, atan, cos, Matrix
```

```python
q = dynamicsymbols('q:2')
qd = dynamicsymbols('q:2', level=1)
l = symbols('l:2')
m = symbols('m:2')
g, t = symbols('g, t')
```

```python
# Compose World Frame
N = ReferenceFrame('N')
A = N.orientnew('A', 'axis', [q[0], N.z])
B = N.orientnew('B', 'axis', [q[1], N.z])

A.set_ang_vel(N, qd[0] * N.z)
B.set_ang_vel(N, qd[1] * N.z)
```

```python
O = Point('O')
P = O.locatenew('P', l[0] * A.x)
R = P.locatenew('R', l[1] * B.x)
```

```python
O.set_vel(N, 0)
P.v2pt_theory(O, N, A)
R.v2pt_theory(P, N, B)
```

$\displaystyle l_{0} \dot{q}_{0}\mathbf{\hat{a}_y} + l_{1} \dot{q}_{1}\mathbf{\hat{b}_y}$

```python
ParP = Particle('ParP', P, m[0])
ParR = Particle('ParR', R, m[1])
```

```python
FL = [(P, m[0] * g * N.x), (R, m[1] * g * N.x)]
```

```python
# Calculate the lagrangian, and form the equations of motion
Lag = Lagrangian(N, ParP, ParR)
LM = LagrangesMethod(Lag, q, forcelist=FL, frame=N)
lag_eqs = LM.form_lagranges_equations()
```

```python
lag_eqs
```

$\displaystyle \left[\begin{matrix}g l_{0} m_{0} \sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} + g l_{0} m_{1} \sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} + l_{0}^{2} m_{0} \frac{d^{2}}{d t^{2}} \operatorname{q_{0}}{\left(t \right)} + l_{0}^{2} m_{1} \frac{d^{2}}{d t^{2}} \operatorname{q_{0}}{\left(t \right)} + l_{0} l_{1} m_{1} \left(\sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \sin{\left(\operatorname{q_{1}}{\left(t \right)} \right)} + \cos{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{1}}{\left(t \right)} \right)}\right) \frac{d^{2}}{d t^{2}} \operatorname{q_{1}}{\left(t \right)} - l_{0} l_{1} m_{1} \left(- \sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{1}}{\left(t \right)} \right)} + \sin{\left(\operatorname{q_{1}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{0}}{\left(t \right)} \right)}\right) \frac{d}{d t} \operatorname{q_{0}}{\left(t \right)} \frac{d}{d t} \operatorname{q_{1}}{\left(t \right)} + l_{0} l_{1} m_{1} \left(- \sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{1}}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{q_{0}}{\left(t \right)} + \sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{1}}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{q_{1}}{\left(t \right)} + \sin{\left(\operatorname{q_{1}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{q_{0}}{\left(t \right)} - \sin{\left(\operatorname{q_{1}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{q_{1}}{\left(t \right)}\right) \frac{d}{d t} \operatorname{q_{1}}{\left(t \right)}\\g l_{1} m_{1} \sin{\left(\operatorname{q_{1}}{\left(t \right)} \right)} + l_{0} l_{1} m_{1} \left(\sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \sin{\left(\operatorname{q_{1}}{\left(t \right)} \right)} + \cos{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{1}}{\left(t \right)} \right)}\right) \frac{d^{2}}{d t^{2}} \operatorname{q_{0}}{\left(t \right)} - l_{0} l_{1} m_{1} \left(\sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{1}}{\left(t \right)} \right)} - \sin{\left(\operatorname{q_{1}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{0}}{\left(t \right)} \right)}\right) \frac{d}{d t} \operatorname{q_{0}}{\left(t \right)} \frac{d}{d t} \operatorname{q_{1}}{\left(t \right)} + l_{0} l_{1} m_{1} \left(- \sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{1}}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{q_{0}}{\left(t \right)} + \sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{1}}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{q_{1}}{\left(t \right)} + \sin{\left(\operatorname{q_{1}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{q_{0}}{\left(t \right)} - \sin{\left(\operatorname{q_{1}}{\left(t \right)} \right)} \cos{\left(\operatorname{q_{0}}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{q_{1}}{\left(t \right)}\right) \frac{d}{d t} \operatorname{q_{0}}{\left(t \right)} + l_{1}^{2} m_{1} \frac{d^{2}}{d t^{2}} \operatorname{q_{1}}{\left(t \right)}\end{matrix}\right]$

```python
lag_eqs.simplify()
```

```python
lag_eqs
```

$\displaystyle \left[\begin{matrix}l_{0} \left(g m_{0} \sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} + g m_{1} \sin{\left(\operatorname{q_{0}}{\left(t \right)} \right)} + l_{0} m_{0} \frac{d^{2}}{d t^{2}} \operatorname{q_{0}}{\left(t \right)} + l_{0} m_{1} \frac{d^{2}}{d t^{2}} \operatorname{q_{0}}{\left(t \right)} + l_{1} m_{1} \sin{\left(\operatorname{q_{0}}{\left(t \right)} - \operatorname{q_{1}}{\left(t \right)} \right)} \left(\frac{d}{d t} \operatorname{q_{1}}{\left(t \right)}\right)^{2} + l_{1} m_{1} \cos{\left(\operatorname{q_{0}}{\left(t \right)} - \operatorname{q_{1}}{\left(t \right)} \right)} \frac{d^{2}}{d t^{2}} \operatorname{q_{1}}{\left(t \right)}\right)\\l_{1} m_{1} \left(g \sin{\left(\operatorname{q_{1}}{\left(t \right)} \right)} - l_{0} \sin{\left(\operatorname{q_{0}}{\left(t \right)} - \operatorname{q_{1}}{\left(t \right)} \right)} \left(\frac{d}{d t} \operatorname{q_{0}}{\left(t \right)}\right)^{2} + l_{0} \cos{\left(\operatorname{q_{0}}{\left(t \right)} - \operatorname{q_{1}}{\left(t \right)} \right)} \frac{d^{2}}{d t^{2}} \operatorname{q_{0}}{\left(t \right)} + l_{1} \frac{d^{2}}{d t^{2}} \operatorname{q_{1}}{\left(t \right)}\right)\end{matrix}\right]$

```python

```


```python
%reset -f
```

```python
from sympy import *
```

```python
init_printing()
```

## Define variables

```python
x,y,z  = symbols('x y z')
```

```python
f = sin(x)
```

## Differentiate

```python
diff(f,x)
```

## Integrate

```python
integrate(f, x)
```

```python
integrate(f, [x,0,pi])
```

```python

```


# Two Dimensional Fractal Integration

To generalize the concepts given in the previous section for FIFs to the FISs, consider that $\Omega = \{\tilde{P}_j = (\tilde{x}_j, \tilde{y}_j), \; j = 1, 2, 3\}$ is a triangular domain in the plane. Let $P = \{P_i = (x_i, y_i), \; i = 1, \ldots, N\}$ be given points in the plane scattered over $\Omega$ containing the points $\tilde{P}_j, \; j = 1, 2, 3$. A triangulation $T(\Omega)$ of $\Omega$ over the points $P$ is given with the set,
$$
\begin{equation}
T(\Omega) = \{ \Omega^i = \{P_j^i\} : P_j^i \in P, \; j = 1, 2, 3, \;  i  = 1, \ldots, K \}
\end{equation}
$$
Note that $T(\Omega)$ consists of non-degenerate triangles such that $\Omega^i \cap \Omega^j = \emptyset, \; i \neq j$ and $\Omega = \cup_{i  = 1}^K \Omega^i$.

Consider that each $P_i = (x_i, y_i) \in P$ is associated with $z_i \in \mathbb{R}$. An interpolation function--also called interpolant, corresponding to the data set $\{(x_i, y_i, z_i), \; i = 1, \ldots, N\}$  is a continuous function $f: \Omega \mapsto \mathbb{R}$ such that $f(x_i, y_i) = z_i$. The answer to the problem of constructing the interpolation function $f$ is in two folds:

1. Construction of an IFS $\mathcal{I} = \{\Omega \times \mathbb{R}; w_i; \; i = 1, \ldots, K \}$ whose attractor is the graph$G$ of the function $f$ satisfying the interpolation points, i.e.$f(x_i, y_i) = z_i, \; i = 1, \ldots, N$.

2. Construction of a contraction mapping $M: \mathcal{F} \mapsto\mathcal{F}$ where $\mathcal{F} =\{ \tilde{f} \in \mathcal{C}(\Omega) : \tilde{f}(\tilde{x}_j, \tilde{y}_j) = \tilde{z}_j \}$ such that the fixed point $f$ of the mapping $M$ satisfies the interpolation points, i.e. $f(x_i, y_i) = z_i, \; i = 1, \ldots, N$

### Construction of the IFS $\mathcal{I}$

Addressing to the first step, we put forward the following IFS $I = \{ \Omega \times \mathbb{R}; w_i, \; i = 1, \ldots, K \}$ with affine transformations $w_i$,
$$
\begin{aligned}
     w_i(x, y, z) &= (L_i(x, y), F_i(x, y, z))\\
    &=\begin{bmatrix}
    \alpha_1^i & \alpha_2^i & 0 \\
    \alpha_3^i & \alpha_4^i & 0 \\
    \alpha_5^i & \alpha_6^i & \alpha_7^i
    \end{bmatrix}
    \begin{bmatrix}
    x \\ y \\ z
    \end{bmatrix} + \begin{bmatrix}
    \beta_1^i \\ \beta_2^i \\ \beta_3^i
    \end{bmatrix}
     \; i = 1, \ldots, K,
    \end{aligned}
$$
where $L_i : \Omega \mapsto \Omega^i$
$$
L_i(x, y) = \begin{bmatrix}
    \alpha_1^i & \alpha_2^i \\ \alpha_3^i & \alpha_4^i
    \end{bmatrix} \begin{bmatrix}
    x \\ y
    \end{bmatrix} + \begin{bmatrix}
    \beta_1^i \\ \beta_2^i
    \end{bmatrix}
$$
are contraction mappings for the $z$ axis satisfying the boundary conditions,
$$
F_i(\tilde{x}_j, \tilde{y}_j, \tilde{z}_j) = z_j^i, \; j = 1, 2, 3.
$$
where $\alpha_7^i$ are arbitrary contractivity factors satisfying $|\alpha_7^i| < 1, \; i = 1, \ldots, K$, also called as vertical scaling factors.

Given the vertical scaling factors $\alpha_7^i, \; i = 1, \ldots, K$, the coefficients $\alpha_k^i, \; k = 1 , \ldots, 6, \; i = 1, \ldots, K$ can be found using the boundary conditions in (\ref{eq: plane boundary conditions}) and (\ref{eq: z values boundary conditons}) which results in following system of equations,
$$
\begin{aligned}
\alpha_1^i \tilde{x}_j + \alpha_2^i \tilde{y}_j + \beta_1^i &= x_j^i \\
\alpha_3^i \tilde{x}_j + \alpha_4^i \tilde{y}_j + \beta_2^i &= y_j^i \\
\alpha_5^i \tilde{x}_j + \alpha_6^i \tilde{y}_j + \beta_3^i &= z_j^i - \alpha_7^i \tilde{z}_j
\end{aligned}
$$
for $i = 1, \ldots, K, \; j = 1, 2, 3$. This system can be rewritten in block diagonal matrix equation system,
$$
\begin{bmatrix}
P & 0 & \ldots & 0 \\
0 & P & \ldots & 0 \\
\vdots & \vdots & \ddots & \vdots \\
0 & 0 & \ldots & P \\
\end{bmatrix}
\begin{bmatrix}
r_1 \\ r_2 \\ r_3
\end{bmatrix} = \begin{bmatrix}
c_1 \\ c_2 \\ c_3
\end{bmatrix}
$$
where
$$
\begin{aligned}
P &=
\begin{bmatrix}
\tilde{x}_1 & \tilde{y}_1 & 1 \\
\tilde{x}_2 & \tilde{y}_2 & 1 \\
\tilde{x}_3 & \tilde{y}_3 & 1 \\
\end{bmatrix}\\
r_1 &= [r_1^1, \ldots, r_1^i, \ldots, r_1^K], \; r_1^i = [\alpha_1^i, \alpha_2^i, \beta_1^i] \\
r_2 &= [r_2^1, \ldots, r_2^i, \ldots, r_2^K], \; r_2^i = [\alpha_3^i, \alpha_4^i, \beta_2^i]  \\
r_3 &= [r_3^1, \ldots, r_3^i, \ldots, r_3^K], \; r_3^i = [\alpha_5^i, \alpha_6^i, \beta_3^i] \\
c_1 &= [c_1^1, \ldots, c_1^i, \ldots, c_1^K], \; c_1^i = [x_1^i, x_2^i, x_3^i] \\
c_2 &= [c_2^1, \ldots, c_2^i, \ldots, c_2^K], \; c_2^i = [y_1^i, y_2^i, y_3^i] \\
c_3 &= [c_3^1, \ldots, c_3^i, \ldots, c_3^K], \; c_3^i = [z_1^i, z_2^i, z_3^i] - \alpha_7^i [\tilde{z}_1^i, \tilde{z}_2^i, \tilde{z}_3^i]1
\end{aligned}
$$
We  have uncoupled system of equations,
$$
P r_j^i = c_j^i, \; j = 1, 2, 3, \; i = 1, \ldots, N
$$
Since the points $\{\tilde{P}_j, \; j = 1, 2, 3\}$ forms a non-degenerate triangular region $\Omega$, $P^{-1}$ exists and can be solved for the coefficients of the IFS $\mathcal{I}$ which gives,
$$
r_j^i = P^{-1} c_j^i, \; j = 1, 2, 3, \; i = 1, \ldots, K
$$

Inspired by the reasoning given in Theorem 1, we propose the following iterative map to find the fractal interpolant.

Given the triangulation $T(\Omega)$ and the IFS $\mathcal{I}$ defined above, let $(\mathcal{F}, d)$ be a complete metric space such that $\mathcal{F} = \{ \tilde{f} \in \mathcal{C}(\Omega) : \tilde{f}(\tilde{x}_j, \tilde{y}_j) = \tilde{z}_j, \; j = 1, 2, 3 \}$ with the metric $d(\tilde{f}_1, \tilde{f}_2) = max \{ |\tilde{f}_1(x, y) - \tilde{f}_2(x, y)|\}, \; x, y \in \Omega$. Let $M : \mathcal{F} \mapsto \mathcal{F}$ be defined by
$$
(M\tilde{f})(x, y) = F_i(L_i^{-1}(x, y), \tilde{f}(L_i^{-1}(x, y)), \; x, y \in \Omega^i, \; i = 1, \ldots, K
$$
Then, $M$ has a unique fixed point $f$ such that $f = \lim_{n \mapsto \infty} M^n(\tilde{f})$ for any  $\tilde{f} \in \mathcal{F}$ which interpolates the given data set, i.e. $f(x_i, y_i) = z_i, \; i = 1, \ldots, N$.

```julia
# This dimensinal fractal integration.

using FractalTools
using Triangulation
using PyPlot
using Statistics

# Read mesh
datapath = joinpath(dirname(dirname(pathof(FractalTools))),  "data/gmsh_data")
filename = "t3.msh"
domain = read_mesh(joinpath(datapath, filename))
scale_factor = 1   # Rescale the mesh.
for node in domain.nodes
    node.x /= scale_factor
    node.y /= scale_factor
end
domains = triangular_partition(domain)
regions = gmsh2matplotlibtri.(domains)

# Define function
func(x, y) = x.^2 .+ y.^2
number_of_regions = length(regions)
z = Vector{Vector}(undef, number_of_regions)
for k = 1 : number_of_regions
    z[k] = func(regions[k].x, regions[k].y)
end

# Refine region to evaluate interpolated function
subdiv = 2
x, y = refine(gmsh2matplotlibtri(domain), subdiv, true)

# Compute initial and real values.
func0(xi, yi) = 0.

α = 0.001
interpolant, coeffs = fis(regions, z, α=α, func0=func0, num_iter=10, get_coeffs=true)

# Interpolation and error
real_values = func.(x, y)
interpolated_values = interpolant.(x, y)
absolute_error = abs.(interpolated_values - real_values)
relative_error = 100 * map(i -> abs(real_values[i]) <= 0.1 ? absolute_error[i] : absolute_error[i] / abs(real_values[i]),
    1 : length(absolute_error));
```

```julia
nothing
```

```julia
# # Plot the results.
xd = xcoords(domain)
yd = ycoords(domain)
zd = func.(xd, yd)
minx, meanx, maxx = minimum(x), mean(x), maximum(x)
miny, meany, maxy = minimum(y), mean(y), maximum(y)
plot_trisurf(x, y, real_values)

coluors = ["black", "red", "blue", "green", "purple", "magenta"]
for (color, region) in zip(coluors, regions)
    triplot(region, color=color)
    plot(region.x, region.y, ".", color="orange")
end
xlabel("x", labelpad=10, fontsize=12)
ylabel("y", labelpad=10, fontsize=12)
zlabel("z", labelpad=10, fontsize=12)
xticks(range(minx, maxx, length=5))
yticks(range(miny, maxy, length=5))
zticks(range(minimum(real_values), maximum(real_values), length=5));

```

```julia
plot_trisurf(x, y, interpolated_values)
for (color, region) in zip(coluors, regions)
    triplot(region, color=color)
    plot(region.x, region.y, ".", color="orange")
end
xlabel("x", labelpad=10, fontsize=12)
ylabel("y", labelpad=10, fontsize=12)
zlabel("z", labelpad=10, fontsize=12)
xticks(range(minx, maxx, length=5))
yticks(range(miny, maxy, length=5))
zticks(range(minimum(interpolated_values), maximum(interpolated_values), length=5))
tight_layout()
```

```julia
plot_trisurf(x, y, absolute_error)
for (color, region) in zip(coluors, regions)
    triplot(region, color=color)
    plot(region.x, region.y, ".", color="orange")
end
xlabel("x", labelpad=10, fontsize=12)
ylabel("y", labelpad=10, fontsize=12)
zlabel("z", labelpad=10, fontsize=12)
xticks(range(minx, maxx, length=5))
yticks(range(miny, maxy, length=5))
zticks(range(minimum(absolute_error), maximum(absolute_error), length=5))
tight_layout()
```

```julia
plot_trisurf(x, y, relative_error)
for (color, region) in zip(coluors, regions)
    triplot(region, color=color)
    plot(region.x, region.y, ".", color="orange")
end
xlabel("x", labelpad=10, fontsize=12)
ylabel("y", labelpad=10, fontsize=12)
zlabel("z", labelpad=10, fontsize=12)
xticks(range(minx, maxx, length=5))
yticks(range(miny, maxy, length=5))
zticks(range(minimum(relative_error), maximum(relative_error), length=5))
tight_layout()
```

```julia
plot_trisurf(x, y, relative_error_fltr)
for (color, region) in zip(coluors, regions)
    triplot(region, color=color)
    plot(region.x, region.y, ".", color="orange")
end
xlabel("x", labelpad=10, fontsize=12)
ylabel("y", labelpad=10, fontsize=12)
zlabel("z", labelpad=10, fontsize=12)
xticks(range(minx, maxx, length=5))
yticks(range(miny, maxy, length=5))
zticks(range(minimum(relative_error_fltr), maximum(relative_error_fltr), length=5))
tight_layout()
```

```julia
fig, ax = subplots(1)
for region in regions
    ax.triplot(region)
    ax.plot(region.x, region.y, ".", color="orange")
end
ax.set_xlabel("x", fontsize=12)
ax.set_ylabel("y", fontsize=12)
tight_layout()
```


```python
import sympy as sm
import sympy.physics.vector
B = sm.physics.vector.ReferenceFrame("B")
e0,e1,e2,e3,b0,b1,b2,b3 = sm.symbols('e_0:4, e\'_0:4')
((e0+e1+e2+e3)*(b0+b1+b2+b3)).expand()
# e0b0
M = sm.Matrix([
    [e0,e1,e2,e3],
    [e1,e0, 0, 0],
    [e2, 0,e0, 0],
    [e3, 0, 0,e0]])
N = sm.Matrix([
    [b0,b1,b2,b3],
    [b1,b0, 0, 0],
    [b2, 0,b0, 0],
    [b3, 0, 0,b0],])
M*N
S = sm.Matrix([
    [1,0,0,0],
    [0,1,0,0],
    [0,0,1,0],
    [0,0,0,1],])
E1 = sm.Matrix([
    [0,1,0,0],
    [1,0,0,0],
    [0,0,0,0],
    [0,0,0,0],])
E2 = sm.Matrix([
    [0,0,1,0],
    [0,0,0,0],
    [1,0,0,0],
    [0,0,0,0],])
E3 = sm.Matrix([
    [0,0,0,1],
    [0,0,0,0],
    [0,0,0,0],
    [1,0,0,0],])

I = sm.Matrix([
    [1,0,0,0],
    [0,1,0,0],
    [0,0,1,0],
    [0,0,0,1]])
a1 = e1*E1 + e2*E2+ e3*E3
b1 = b1*E1 + b2*E2+ b3*E3
I*(a1*b1)
```

$\displaystyle \left[\begin{matrix}e'_{1} e_{1} + e'_{2} e_{2} + e'_{3} e_{3} & 0 & 0 & 0\\0 & e'_{1} e_{1} & e'_{2} e_{1} & e'_{3} e_{1}\\0 & e'_{1} e_{2} & e'_{2} e_{2} & e'_{3} e_{2}\\0 & e'_{1} e_{3} & e'_{2} e_{3} & e'_{3} e_{3}\end{matrix}\right]$

```python
a = e1*B.x + e2*B.y + e3*B.z
b = b1*B.x + b2*B.y + b3*B.z
(a|b).to_matrix(B)
```

```python
(B.x^B.y)^B.z
```

```python
w|u
```

```python
vi
```

```python
u
```

```python
vi&u
```

```python
# 1/5*B.x*B.y - 2/5*B.y*B.x
# 3/5 + 3/5*B.x*B.y
v
```

```python
sm.Eijk(e1,e2,e3)
```

```python
S = sm.Matrix([
    [1,0,0,0],
    [0,1,0,0],
    [0,0,1,0],
    [0,0,0,1],])
E1 = sm.Matrix([
    [0,e1,0,0],
    [e1,0,0,0],
    [0,0,0,e1],
    [0,0,e1,0],])
E2 = sm.Matrix([
    [0,0,e2,0],
    [0,0,0,e2],
    [e2,0,0,0],
    [0,e2,0,0],])
E3 = sm.Matrix([
    [0,0,0,e3],
    [0,0,e3,0],
    [0,e3,0,0],
    [e3,0,0,0],])

(S+E1+E2+E3)
```

$\displaystyle \left[\begin{matrix}1 & e_{1} & e_{2} & e_{3}\\e_{1} & 1 & e_{3} & e_{2}\\e_{2} & e_{3} & 1 & e_{1}\\e_{3} & e_{2} & e_{1} & 1\end{matrix}\right]$

```python

```


## Mutual information minimization on Toy Gaussians

This is a sanity check involving binary classification of bivariate Gaussians. Data is sampled from bivariate normal distributions as follows:
\begin{align}
x_0 &\sim \mathcal{N}\left(\begin{pmatrix} z \\ 0 \end{pmatrix}, \begin{pmatrix} 1 & 0.5 \\ 0.5 & 1 \end{pmatrix}\right) \\
x_1 &\sim \mathcal{N}\left(\begin{pmatrix} 0 \\ 1.5 \end{pmatrix}, \begin{pmatrix} 1 & 0 \\ 0 & 1 \end{pmatrix}\right)
\end{align}
Where $z \sim \mathcal{N}(0,1)$ simulates uncertainty in the $x$-coordinate of samples drawn from the first distribution. The data is then passed through a shallow neural network, which we refer to as the classification network whose objective is to distinguish between samples drawn from each distribution.

To reduce the dependency of the network on the potentially noisy $x$-variable, we penalize an estimate of the mutual information between the logits of the neural network and the $x$-coordinate of the input data. We optimize a lower bound on the mutual information through a shallow neural network of similar capacity to the classification network, which remains unchanged. The degree to which the output encoding of the classification network should be indepedent of the $x$-coordinate can be controlled through the hyperparameter $\lambda$ in the objective function.

```python
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import scipy
from scipy.interpolate import griddata
from sklearn.metrics import roc_curve, auc, roc_auc_score
plt.style.use('seaborn-darkgrid')
plt.style.use('seaborn-talk')
plt.style.use('seaborn-pastel')

import tensorflow as tf
import tensorflow_probability as tfp
import time, os, sys, glob
from lnc import MI
```

```python
def MINE(x, y, y_prime, training, batch_size, name='MINE', actv=tf.nn.relu, dimension=2, l2_reg=False):
    """
    Mutual Information Neural Estimator
    (x,y):      Drawn from joint
    y_prime:    Drawn from marginal

    returns
    MI:         Lower bound on mutual information between x,y
    """

    init = tf.contrib.layers.xavier_initializer()
    drop_rate = 0.0
    shape = [32,32,32,32]
    kwargs = {'center': True, 'scale': True, 'training': training, 'fused': True, 'renorm': False}

    lambda_l2 = 0.001
    regularization = tf.contrib.layers.l2_regularizer(scale=lambda_l2) if l2_reg else None
    y_prime = tf.random_shuffle(y)

    def statistic_network(t, name='stat_net', reuse=False):
        with tf.variable_scope(name, initializer=init, reuse=reuse) as scope:

            h0 = tf.layers.dense(t, units=shape[0], activation=actv, kernel_regularizer=regularization)
            h1 = tf.layers.dense(h0, units=shape[1], activation=actv, kernel_regularizer=regularization)
            h2 = tf.layers.dense(h1, units=shape[2], activation=actv, kernel_regularizer=regularization)
            out = tf.layers.dense(h2, units=1, kernel_initializer=init, kernel_regularizer=regularization)

        return out

    def log_sum_exp_trick(x, batch_size, axis=1):
        # Compute along batch dimension
        x = tf.squeeze(x)
        x_max = tf.reduce_max(x)
        lse = x_max + tf.log(tf.reduce_sum(tf.exp(x-x_max))) - tf.log(batch_size)
        return lse

    x, y, y_prime = tf.expand_dims(tf.squeeze(x), axis=1), tf.expand_dims(tf.squeeze(y), axis=1), tf.expand_dims(tf.squeeze(y_prime), axis=1)
    z = tf.concat([x,y], axis=1)
    z_prime = tf.concat([x,y_prime], axis=1)
    joint_f = statistic_network(z)
    marginal_f = statistic_network(z_prime, reuse=True)

    print('joint shape', joint_f.shape)
    print('marginal shape', marginal_f.shape)

    MI_lower_bound = tf.squeeze(tf.reduce_mean(joint_f)) - tf.squeeze(log_sum_exp_trick(marginal_f, 1.*batch_size))

    joint_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=joint_f,
        labels=tf.ones_like(joint_f)))
    marginal_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=marginal_f,
        labels=tf.zeros_like(marginal_f)))

    # Approximate the trained discriminator as the optimal discriminator
    JSD_lower_bound = -(joint_loss + marginal_loss) + tf.log(4.0)

    if args.KL_loss:
        lower_bound = MI_lower_bound
    else:
        lower_bound = JSD_lower_bound

    return lower_bound, joint_f, marginal_f

def dense_network(x, training, name='fully_connected', actv=tf.nn.relu, **kwargs):
    # Toy dense network for binary classification
    init = tf.contrib.layers.xavier_initializer()
    shape = [32, 32, 32, 32, 32]
    kwargs = {'center': True, 'scale': True, 'training': training, 'fused': True, 'renorm': True}

    print('Input X shape', x.get_shape())

    with tf.variable_scope(name, initializer=init, reuse=tf.AUTO_REUSE) as scope:
        h0 = tf.layers.dense(x, units=shape[0], activation=actv)
        # h0 = tf.layers.batch_normalization(h0, **kwargs)
        h1 = tf.layers.dense(h0, units=shape[1], activation=actv)
        # h1 = tf.layers.batch_normalization(h1, **kwargs)
        h2 = tf.layers.dense(h1, units=shape[2], activation=actv)
        # h2 = tf.layers.batch_normalization(h2, **kwargs)

    out = tf.layers.dense(h2, units=1, kernel_initializer=init)

    return out

def mutual_information_1D_kraskov(x, y):
    # k-NN based estimate of mutual information
    from lnc import MI
    mi = MI.mi_LNC([x,y],k=5,base=np.exp(1),alpha=0.5)

    return mi

def scope_variables(name):
    with tf.variable_scope(name):
        return tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=tf.get_variable_scope().name)

def plot_decision_surface(sess, model, title='Decision Surface'):
    data, out = [], []

    for _ in range(42):
        data_, out_ = sess.run([model.data, model.softmax])
        data.append(data_)
        out.append(np.squeeze(out_))

    data = np.vstack(data)
    out = np.hstack(out)
    x_grid, y_grid = np.mgrid[-2:3:100j, -2:4:100j]

    z = griddata(points=data[:,:2], values=out, xi=(x_grid, y_grid), method='cubic')
    plt.contourf(x_grid, y_grid, z, 25, vmin=0., vmax=1., cmap=plt.cm.viridis)
    plt.colorbar(label='Output')
    plt.xlabel('$x$')
    plt.ylabel('$y$')
    plt.title(title)
    plt.scatter([0], [0], c="green", linewidths=0, label=r"$\mu_0$")
    plt.scatter([0], [1.5], c="blue", linewidths=0, label=r"$\mu_1$")
    plt.legend(loc="upper left")
    # plt.savefig(os.path.join('results', 'decision_no_pen.pdf'), format='pdf', dpi=256, bbox_inches='tight')

def collect_metrics(sess, model):
    clf_loss, MI_est, joint_loss, y_true, y_prob, MI_kNN, acc, logits, logits_x_shift = sess.run([
        model.clf_cost, model.MI_estimate, model.cost, model.labels, model.softmax,
        model.MI_estimate_kNN, model.accuracy, model.logits, model.logits_x_shift])
    auc = roc_auc_score(y_true=y_true, y_score=y_prob)
    ks = scipy.stats.ks_2samp(np.squeeze(logits), np.squeeze(logits_x_shift))
    return {'clf-loss': clf_loss, 'MI-est': MI_est, 'joint-loss': joint_loss, 'auc': auc,
            'MI-kNN': MI_kNN, 'accuracy': acc, 'ks-test': ks[0], 'ks-p-value': ks[1]}
```

```python
class Model():
    def __init__(self, args):

        tfd = tfp.distributions
        self.global_step = tf.Variable(0, trainable=False)
        self.MI_global_step = tf.Variable(0, trainable=False)

        if args.dims == 2:
            dim = 2
            z = tf.random.normal([args.sample_size])
            z_prime = tf.random.shuffle(z)
            mu_x = tf.stack([z, tf.zeros(args.sample_size)], axis=-1)  # [z,0] broadcasted
            mu_y = [0,1.5]
            cov_x = [[1., 0.5], [0.5, 1.]]
            cov_y = [[1., 0.], [0., 1.]]

            standard_mvn_normal = tfd.MultivariateNormalDiag(loc=np.zeros(dim).astype(np.float32))
            scale_x = tf.cholesky(cov_x)
            scale_y = tf.cholesky(cov_y)

            # u ~ N(0,1), v = scale @ u + loc -> v ~ N(loc, scale)
            self.sample_x = tf.squeeze(tf.matmul(tf.broadcast_to(scale_x, [args.sample_size, dim, dim]),
                tf.expand_dims(standard_mvn_normal.sample(args.sample_size), axis=-1))) + mu_x
            self.sample_y = tf.squeeze(tf.matmul(tf.broadcast_to(scale_y, [args.sample_size, dim, dim]),
                tf.expand_dims(standard_mvn_normal.sample(args.sample_size), axis=-1))) + mu_y

            self.x_data = tf.concat([self.sample_x, tf.zeros([args.sample_size,1])], axis=1)
            self.y_data = tf.concat([self.sample_y, tf.ones([args.sample_size,1])], axis=1)
            x_shift = tf.stack((tf.ones(args.sample_size)*2, tf.zeros(args.sample_size)), axis=1)
            self.x_data_shift = tf.concat([self.sample_x + x_shift, tf.zeros([args.sample_size,1])], axis=1)

            data = tf.random.shuffle(tf.concat([self.x_data, self.y_data], axis=0))
            data_x_shift = tf.random.shuffle(tf.concat([self.x_data_shift, self.y_data], axis=0))
            idx = tf.range(tf.shape(data)[0])
            ridx = tf.random.shuffle(idx)[:args.sample_size]
            self.data = tf.gather(data, ridx)
            self.data_x_shift = tf.gather(data_x_shift, ridx)
            self.labels = tf.squeeze(self.data[:,2])

        with tf.variable_scope('classifier', reuse=tf.AUTO_REUSE) as scope:
            self.logits = dense_network(x=self.data[:,:2], training=True)

        with tf.variable_scope('classifier', reuse=True) as scope:
            self.logits_x_shift = dense_network(x=self.data_x_shift[:,:2], training=True)

        with tf.variable_scope('miner', reuse=tf.AUTO_REUSE) as scope:
            self.MI_estimate, joint_f, marginal_f = MINE(x=self.logits, y=self.data[:,0], y_prime=tf.random_shuffle(self.data[:,0]),
                batch_size=args.sample_size, dimension=dim, training=True, actv=tf.nn.elu)

        self.softmax = tf.nn.sigmoid(self.logits)
        self.pred = tf.squeeze(tf.cast(tf.greater(self.softmax, 0.5), tf.float32))
        self.accuracy = tf.reduce_mean(tf.cast(tf.equal(self.pred, self.labels), tf.float32))

        self.cross_entropy = tf.nn.sigmoid_cross_entropy_with_logits(logits=self.logits, labels=tf.expand_dims(self.labels,axis=-1))
        self.clf_cost = tf.reduce_mean(self.cross_entropy)
        self.MI_estimate_kNN = tf.py_func(mutual_information_1D_kraskov, inp=[tf.squeeze(self.logits),
            tf.squeeze(self.data[:,0])], Tout=tf.float64)

        # Alternate non-saturating updated rule for generator
        # max_E log(D(E,Z)) = min_E H(D(E,Z),1)
        self.NS_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=marginal_f, labels=tf.ones_like(marginal_f)))
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

        if args.penalty:
            if args.alt_update:
                self.cost = self.clf_cost + args.lambda_MI * tf.nn.relu(self.MI_estimate + self.NS_loss)
            else:
                self.cost = self.clf_cost + args.lambda_MI * tf.nn.relu(self.MI_estimate)
        else:
            self.cost = self.clf_cost

        lr = args.lr_E
        theta_f = scope_variables('classifier')
        theta_MINE = scope_variables('miner')
        print('Classifier Parameters', theta_f)
        print('MI-est Parameters', theta_MINE)
        args.optimizer = args.optimizer.lower()
        if args.optimizer == 'adam':
            self.opt = tf.train.AdamOptimizer(lr)
        elif args.optimizer == 'momentum':
            self.opt = tf.train.MomentumOptimizer(lr, momentum=0.9, use_nesterov=True)
        elif args.optimizer == 'rmsprop':
            self.opt = tf.train.RMSPropOptimizer(lr)
        elif args.optimizer == 'sgd':
            self.opt = tf.train.GradientDescentOptimizer(lr)

        with tf.variable_scope('opt', reuse=tf.AUTO_REUSE) as scope:
            self.MI_opt = tf.train.AdamOptimizer(args.lr_T)
            self.grads_clf = self.opt.compute_gradients(self.clf_cost, var_list=theta_f)
            self.grads_mi_clf =  self.opt.compute_gradients(self.MI_estimate, var_list=theta_f)
            self.grads_mi_est = self.MI_opt.compute_gradients(-self.MI_estimate, var_list=theta_MINE)
            self.MINE_opt_op = self.MI_opt.minimize(-self.MI_estimate, global_step=self.MI_global_step,
                    var_list=theta_MINE)

            with tf.control_dependencies(update_ops):
                # Ensures that we execute the update_ops before performing the train_step
                self.opt_op = self.opt.minimize(self.cost, global_step=self.global_step, var_list=theta_f)
```

```python
def train(args):

    start_time = time.time()
    ckpt = tf.train.get_checkpoint_state('checkpoints')

    # Build graph
    model = Model(args)
    saver = tf.train.Saver()
    all_metrics = list()

    sess = tf.InteractiveSession(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=False))
    sess.run(tf.global_variables_initializer())
    sess.run(tf.local_variables_initializer())

    for step in range(int(args.steps)):
        sess.run(model.opt_op)

        # if args.penalty:
        for MI_step in range(args.MI_iters):
            sess.run(model.MINE_opt_op)

        if step % 100 == 0:
            metrics = collect_metrics(sess, model)
            all_metrics.append(metrics)
            print('MI Estimate: {:.3f} | kNN: {:.3f} | Accuracy: {:.3f} | ({:.2f} s)'.format(metrics['MI-est'],
                metrics['MI-kNN'], metrics['accuracy'], time.time()-start_time))

    print("Training Complete. Time elapsed: {:.3f} s".format(time.time()-start_time))

    metrics = pd.concat([pd.DataFrame.from_dict(d, orient='index').transpose() for d in all_metrics], axis=0, ignore_index=True)
    return sess, model, metrics
```

```python
class args(object):
    rho = 0.5
    sample_size = 512
    steps = 8e3
    dims = 2
    lr_E = 1e-4
    lr_T = 1e-3
    optimizer = 'adam'
    name = 'toy-MI'
    lambda_MI = 10.0
    penalty = False
    MI_iters = 4
    KL_loss = False
    alt_update = False
```

```python
sess, model, metrics = train(args)
```

    WARNING:tensorflow:From /usr/local/easybuild/software/Tensorflow/1.13.1-intel-2017.u2-GCC-6.2.0-CUDA10-Python-3.6.4-GPU/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.
    Instructions for updating:
    Colocations handled automatically by placer.
    Input X shape (512, 2)
    WARNING:tensorflow:From <ipython-input-3-5c5c669f09fb>:72: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.
    Instructions for updating:
    Use keras.layers.dense instead.
    Input X shape (512, 2)
    joint shape (512, 1)
    marginal shape (512, 1)
    WARNING:tensorflow:From <ipython-input-4-9775be383715>:57: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.
    Instructions for updating:
    tf.py_func is deprecated in TF V2. Instead, use
        tf.py_function, which takes a python function which manipulates tf eager
        tensors instead of numpy arrays. It's easy to convert a tf eager tensor to
        an ndarray (just call tensor.numpy()) but having access to eager tensors
        means `tf.py_function`s can use accelerators such as GPUs as well as
        being differentiable using a gradient tape.

    Classifier Parameters [<tf.Variable 'classifier/fully_connected/dense/kernel:0' shape=(2, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_1/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_1/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_2/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_2/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/dense/kernel:0' shape=(32, 1) dtype=float32_ref>, <tf.Variable 'classifier/dense/bias:0' shape=(1,) dtype=float32_ref>]
    MI-est Parameters [<tf.Variable 'miner/stat_net/dense/kernel:0' shape=(2, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_1/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_1/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_2/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_2/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_3/kernel:0' shape=(32, 1) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_3/bias:0' shape=(1,) dtype=float32_ref>]
    MI Estimate: -0.000 | kNN: 0.298 | Accuracy: 0.373 | (2.34 s)
    MI Estimate: 0.080 | kNN: 0.207 | Accuracy: 0.795 | (3.56 s)
    MI Estimate: 0.090 | kNN: 0.370 | Accuracy: 0.740 | (4.70 s)
    MI Estimate: 0.070 | kNN: 0.203 | Accuracy: 0.803 | (5.85 s)
    MI Estimate: 0.075 | kNN: 0.211 | Accuracy: 0.756 | (6.99 s)
    MI Estimate: 0.047 | kNN: 0.132 | Accuracy: 0.828 | (8.13 s)
    MI Estimate: 0.073 | kNN: 0.160 | Accuracy: 0.775 | (9.27 s)
    MI Estimate: 0.042 | kNN: 0.166 | Accuracy: 0.779 | (10.42 s)
    MI Estimate: 0.080 | kNN: 0.208 | Accuracy: 0.807 | (11.58 s)
    MI Estimate: 0.028 | kNN: 0.157 | Accuracy: 0.781 | (12.72 s)
    MI Estimate: 0.054 | kNN: 0.207 | Accuracy: 0.787 | (13.86 s)
    MI Estimate: 0.055 | kNN: 0.148 | Accuracy: 0.783 | (15.00 s)
    MI Estimate: 0.025 | kNN: 0.145 | Accuracy: 0.781 | (16.14 s)
    MI Estimate: 0.061 | kNN: 0.153 | Accuracy: 0.787 | (17.28 s)
    MI Estimate: 0.054 | kNN: 0.168 | Accuracy: 0.789 | (18.42 s)
    MI Estimate: 0.049 | kNN: 0.161 | Accuracy: 0.797 | (19.56 s)
    MI Estimate: 0.038 | kNN: 0.157 | Accuracy: 0.773 | (20.70 s)
    MI Estimate: 0.049 | kNN: 0.123 | Accuracy: 0.807 | (21.84 s)
    MI Estimate: 0.082 | kNN: 0.159 | Accuracy: 0.824 | (22.99 s)
    MI Estimate: 0.044 | kNN: 0.175 | Accuracy: 0.793 | (24.13 s)
    MI Estimate: 0.069 | kNN: 0.216 | Accuracy: 0.775 | (25.27 s)
    MI Estimate: 0.062 | kNN: 0.161 | Accuracy: 0.809 | (26.41 s)
    MI Estimate: 0.055 | kNN: 0.173 | Accuracy: 0.797 | (27.55 s)
    MI Estimate: 0.050 | kNN: 0.148 | Accuracy: 0.799 | (28.69 s)
    MI Estimate: 0.080 | kNN: 0.180 | Accuracy: 0.807 | (29.83 s)
    MI Estimate: 0.062 | kNN: 0.220 | Accuracy: 0.789 | (30.97 s)
    MI Estimate: 0.049 | kNN: 0.177 | Accuracy: 0.807 | (32.11 s)
    MI Estimate: 0.065 | kNN: 0.140 | Accuracy: 0.820 | (33.26 s)
    MI Estimate: 0.044 | kNN: 0.167 | Accuracy: 0.768 | (34.40 s)
    MI Estimate: 0.042 | kNN: 0.072 | Accuracy: 0.795 | (35.54 s)
    MI Estimate: 0.082 | kNN: 0.250 | Accuracy: 0.770 | (36.68 s)
    MI Estimate: 0.034 | kNN: 0.141 | Accuracy: 0.805 | (37.80 s)
    MI Estimate: 0.043 | kNN: 0.112 | Accuracy: 0.811 | (38.91 s)
    MI Estimate: 0.069 | kNN: 0.198 | Accuracy: 0.799 | (40.03 s)
    MI Estimate: 0.075 | kNN: 0.182 | Accuracy: 0.768 | (41.15 s)
    MI Estimate: 0.054 | kNN: 0.198 | Accuracy: 0.773 | (42.27 s)
    MI Estimate: 0.047 | kNN: 0.108 | Accuracy: 0.785 | (43.38 s)
    MI Estimate: 0.061 | kNN: 0.157 | Accuracy: 0.770 | (44.49 s)
    MI Estimate: 0.076 | kNN: 0.150 | Accuracy: 0.777 | (45.60 s)
    MI Estimate: 0.064 | kNN: 0.219 | Accuracy: 0.797 | (46.71 s)
    MI Estimate: 0.050 | kNN: 0.183 | Accuracy: 0.816 | (47.82 s)
    MI Estimate: 0.065 | kNN: 0.149 | Accuracy: 0.795 | (48.94 s)
    MI Estimate: 0.066 | kNN: 0.187 | Accuracy: 0.799 | (50.06 s)
    MI Estimate: 0.058 | kNN: 0.159 | Accuracy: 0.795 | (51.19 s)
    MI Estimate: 0.054 | kNN: 0.124 | Accuracy: 0.779 | (52.31 s)
    MI Estimate: 0.036 | kNN: 0.103 | Accuracy: 0.793 | (53.43 s)
    MI Estimate: 0.059 | kNN: 0.210 | Accuracy: 0.826 | (54.55 s)
    MI Estimate: 0.059 | kNN: 0.195 | Accuracy: 0.787 | (55.66 s)
    MI Estimate: 0.060 | kNN: 0.144 | Accuracy: 0.775 | (56.78 s)
    MI Estimate: 0.056 | kNN: 0.173 | Accuracy: 0.797 | (57.90 s)
    MI Estimate: 0.088 | kNN: 0.249 | Accuracy: 0.797 | (59.01 s)
    MI Estimate: 0.061 | kNN: 0.153 | Accuracy: 0.764 | (60.14 s)
    MI Estimate: 0.067 | kNN: 0.184 | Accuracy: 0.797 | (61.25 s)
    MI Estimate: 0.079 | kNN: 0.158 | Accuracy: 0.793 | (62.37 s)
    MI Estimate: 0.080 | kNN: 0.172 | Accuracy: 0.789 | (63.50 s)
    MI Estimate: 0.084 | kNN: 0.286 | Accuracy: 0.820 | (64.62 s)
    MI Estimate: 0.062 | kNN: 0.211 | Accuracy: 0.766 | (65.73 s)
    MI Estimate: 0.068 | kNN: 0.169 | Accuracy: 0.754 | (66.84 s)
    MI Estimate: 0.064 | kNN: 0.133 | Accuracy: 0.799 | (67.96 s)
    MI Estimate: 0.051 | kNN: 0.120 | Accuracy: 0.770 | (69.08 s)
    MI Estimate: 0.095 | kNN: 0.183 | Accuracy: 0.812 | (70.19 s)
    MI Estimate: 0.030 | kNN: 0.099 | Accuracy: 0.781 | (71.31 s)
    MI Estimate: 0.022 | kNN: 0.122 | Accuracy: 0.803 | (72.43 s)
    MI Estimate: 0.078 | kNN: 0.167 | Accuracy: 0.805 | (73.56 s)
    MI Estimate: 0.063 | kNN: 0.238 | Accuracy: 0.779 | (74.68 s)
    MI Estimate: 0.080 | kNN: 0.194 | Accuracy: 0.771 | (75.81 s)
    MI Estimate: 0.044 | kNN: 0.178 | Accuracy: 0.760 | (76.92 s)
    MI Estimate: 0.064 | kNN: 0.162 | Accuracy: 0.832 | (78.04 s)
    MI Estimate: 0.088 | kNN: 0.202 | Accuracy: 0.801 | (79.17 s)
    MI Estimate: 0.026 | kNN: 0.192 | Accuracy: 0.799 | (80.31 s)
    MI Estimate: 0.054 | kNN: 0.129 | Accuracy: 0.793 | (81.44 s)
    MI Estimate: 0.087 | kNN: 0.269 | Accuracy: 0.795 | (82.55 s)
    MI Estimate: 0.078 | kNN: 0.214 | Accuracy: 0.807 | (83.67 s)
    MI Estimate: 0.110 | kNN: 0.205 | Accuracy: 0.801 | (84.79 s)
    MI Estimate: 0.057 | kNN: 0.194 | Accuracy: 0.793 | (85.91 s)
    MI Estimate: 0.041 | kNN: 0.103 | Accuracy: 0.771 | (87.03 s)
    MI Estimate: 0.098 | kNN: 0.234 | Accuracy: 0.820 | (88.14 s)
    MI Estimate: 0.070 | kNN: 0.187 | Accuracy: 0.781 | (89.26 s)
    MI Estimate: 0.067 | kNN: 0.151 | Accuracy: 0.779 | (90.39 s)
    MI Estimate: 0.072 | kNN: 0.212 | Accuracy: 0.799 | (91.52 s)
    Training Complete. Time elapsed: 92.517 s

```python
plot_decision_surface(sess, model, title='Unpenalized Decision Surface')
```

### Switch on MI penalty

```python
sess.close()
args.penalty = True
args.lambda_MI = 1.0
sess_p, model_p, metrics_p = train(args)
```

    Input X shape (512, 2)
    Input X shape (512, 2)
    joint shape (512, 1)
    marginal shape (512, 1)
    Classifier Parameters [<tf.Variable 'classifier/fully_connected/dense/kernel:0' shape=(2, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_1/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_1/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_2/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_2/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/dense/kernel:0' shape=(32, 1) dtype=float32_ref>, <tf.Variable 'classifier/dense/bias:0' shape=(1,) dtype=float32_ref>]
    MI-est Parameters [<tf.Variable 'miner/stat_net/dense/kernel:0' shape=(2, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_1/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_1/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_2/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_2/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_3/kernel:0' shape=(32, 1) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_3/bias:0' shape=(1,) dtype=float32_ref>]
    MI Estimate: -0.002 | kNN: 0.316 | Accuracy: 0.453 | (1.26 s)
    MI Estimate: 0.010 | kNN: 0.202 | Accuracy: 0.299 | (2.60 s)
    MI Estimate: 0.008 | kNN: 0.205 | Accuracy: 0.424 | (3.84 s)
    MI Estimate: 0.042 | kNN: 0.134 | Accuracy: 0.422 | (5.05 s)
    MI Estimate: 0.005 | kNN: 0.137 | Accuracy: 0.449 | (6.27 s)
    MI Estimate: 0.003 | kNN: 0.039 | Accuracy: 0.457 | (7.48 s)
    MI Estimate: 0.004 | kNN: 0.096 | Accuracy: 0.418 | (8.70 s)
    MI Estimate: 0.005 | kNN: 0.043 | Accuracy: 0.418 | (9.92 s)
    MI Estimate: 0.008 | kNN: 0.063 | Accuracy: 0.414 | (11.13 s)
    MI Estimate: 0.005 | kNN: 0.023 | Accuracy: 0.471 | (12.35 s)
    MI Estimate: -0.001 | kNN: 0.056 | Accuracy: 0.498 | (13.56 s)
    MI Estimate: 0.021 | kNN: 0.119 | Accuracy: 0.561 | (14.78 s)
    MI Estimate: 0.003 | kNN: 0.075 | Accuracy: 0.551 | (15.99 s)
    MI Estimate: 0.009 | kNN: 0.083 | Accuracy: 0.584 | (17.21 s)
    MI Estimate: 0.007 | kNN: 0.145 | Accuracy: 0.666 | (18.42 s)
    MI Estimate: 0.010 | kNN: 0.169 | Accuracy: 0.688 | (19.63 s)
    MI Estimate: 0.017 | kNN: 0.131 | Accuracy: 0.723 | (20.84 s)
    MI Estimate: 0.006 | kNN: 0.066 | Accuracy: 0.660 | (22.04 s)
    MI Estimate: 0.008 | kNN: 0.028 | Accuracy: 0.668 | (23.25 s)
    MI Estimate: 0.004 | kNN: 0.072 | Accuracy: 0.682 | (24.48 s)
    MI Estimate: 0.010 | kNN: 0.039 | Accuracy: 0.697 | (25.69 s)
    MI Estimate: 0.012 | kNN: 0.044 | Accuracy: 0.707 | (26.90 s)
    MI Estimate: 0.009 | kNN: 0.122 | Accuracy: 0.744 | (28.10 s)
    MI Estimate: 0.008 | kNN: 0.051 | Accuracy: 0.727 | (29.30 s)
    MI Estimate: 0.016 | kNN: 0.086 | Accuracy: 0.719 | (30.51 s)
    MI Estimate: -0.001 | kNN: 0.136 | Accuracy: 0.744 | (31.71 s)
    MI Estimate: 0.011 | kNN: 0.079 | Accuracy: 0.746 | (32.92 s)
    MI Estimate: 0.016 | kNN: 0.033 | Accuracy: 0.748 | (34.14 s)
    MI Estimate: 0.020 | kNN: 0.100 | Accuracy: 0.734 | (35.35 s)
    MI Estimate: 0.019 | kNN: 0.197 | Accuracy: 0.785 | (36.57 s)
    MI Estimate: 0.027 | kNN: 0.088 | Accuracy: 0.758 | (37.78 s)
    MI Estimate: 0.020 | kNN: 0.097 | Accuracy: 0.830 | (39.00 s)
    MI Estimate: 0.012 | kNN: 0.037 | Accuracy: 0.801 | (40.20 s)
    MI Estimate: 0.021 | kNN: 0.044 | Accuracy: 0.809 | (41.42 s)
    MI Estimate: 0.020 | kNN: 0.051 | Accuracy: 0.770 | (42.63 s)
    MI Estimate: 0.014 | kNN: 0.127 | Accuracy: 0.807 | (43.84 s)
    MI Estimate: 0.013 | kNN: 0.069 | Accuracy: 0.803 | (45.06 s)
    MI Estimate: 0.006 | kNN: 0.018 | Accuracy: 0.777 | (46.26 s)
    MI Estimate: 0.002 | kNN: 0.082 | Accuracy: 0.775 | (47.47 s)
    MI Estimate: 0.015 | kNN: 0.092 | Accuracy: 0.805 | (48.68 s)
    MI Estimate: 0.018 | kNN: 0.071 | Accuracy: 0.799 | (49.88 s)
    MI Estimate: 0.005 | kNN: 0.041 | Accuracy: 0.785 | (51.09 s)
    MI Estimate: 0.008 | kNN: 0.008 | Accuracy: 0.783 | (52.30 s)
    MI Estimate: 0.017 | kNN: 0.062 | Accuracy: 0.771 | (53.52 s)
    MI Estimate: 0.013 | kNN: 0.062 | Accuracy: 0.809 | (54.73 s)
    MI Estimate: -0.001 | kNN: 0.045 | Accuracy: 0.789 | (55.95 s)
    MI Estimate: 0.010 | kNN: 0.048 | Accuracy: 0.793 | (57.16 s)
    MI Estimate: 0.014 | kNN: 0.030 | Accuracy: 0.781 | (58.37 s)
    MI Estimate: 0.030 | kNN: 0.028 | Accuracy: 0.764 | (59.58 s)
    MI Estimate: 0.020 | kNN: 0.055 | Accuracy: 0.764 | (60.79 s)
    MI Estimate: 0.011 | kNN: 0.072 | Accuracy: 0.775 | (62.01 s)
    MI Estimate: 0.022 | kNN: 0.019 | Accuracy: 0.832 | (63.23 s)
    MI Estimate: 0.011 | kNN: 0.049 | Accuracy: 0.766 | (64.44 s)
    MI Estimate: 0.001 | kNN: 0.051 | Accuracy: 0.746 | (65.65 s)
    MI Estimate: 0.017 | kNN: 0.056 | Accuracy: 0.791 | (66.86 s)
    MI Estimate: 0.013 | kNN: 0.085 | Accuracy: 0.807 | (68.07 s)
    MI Estimate: 0.010 | kNN: 0.068 | Accuracy: 0.779 | (69.29 s)
    MI Estimate: 0.022 | kNN: 0.091 | Accuracy: 0.760 | (70.51 s)
    MI Estimate: 0.020 | kNN: 0.095 | Accuracy: 0.787 | (71.73 s)
    MI Estimate: 0.005 | kNN: 0.024 | Accuracy: 0.750 | (72.94 s)
    MI Estimate: 0.018 | kNN: 0.059 | Accuracy: 0.785 | (74.16 s)
    MI Estimate: 0.024 | kNN: 0.081 | Accuracy: 0.809 | (75.38 s)
    MI Estimate: 0.012 | kNN: -0.009 | Accuracy: 0.748 | (76.58 s)
    MI Estimate: 0.012 | kNN: 0.045 | Accuracy: 0.801 | (77.79 s)
    MI Estimate: 0.005 | kNN: 0.058 | Accuracy: 0.785 | (79.00 s)
    MI Estimate: 0.002 | kNN: 0.067 | Accuracy: 0.754 | (80.21 s)
    MI Estimate: 0.007 | kNN: 0.024 | Accuracy: 0.805 | (81.42 s)
    MI Estimate: 0.016 | kNN: 0.050 | Accuracy: 0.779 | (82.64 s)
    MI Estimate: 0.012 | kNN: 0.048 | Accuracy: 0.809 | (83.85 s)
    MI Estimate: 0.006 | kNN: 0.052 | Accuracy: 0.816 | (85.07 s)
    MI Estimate: 0.004 | kNN: 0.070 | Accuracy: 0.801 | (86.27 s)
    MI Estimate: 0.008 | kNN: 0.041 | Accuracy: 0.818 | (87.50 s)
    MI Estimate: 0.006 | kNN: 0.049 | Accuracy: 0.783 | (88.70 s)
    MI Estimate: 0.009 | kNN: 0.040 | Accuracy: 0.803 | (89.93 s)
    MI Estimate: 0.003 | kNN: 0.068 | Accuracy: 0.785 | (91.13 s)
    MI Estimate: 0.016 | kNN: 0.095 | Accuracy: 0.750 | (92.35 s)
    MI Estimate: 0.008 | kNN: 0.087 | Accuracy: 0.793 | (93.56 s)
    MI Estimate: 0.019 | kNN: 0.097 | Accuracy: 0.777 | (94.78 s)
    MI Estimate: 0.015 | kNN: 0.035 | Accuracy: 0.781 | (96.00 s)
    MI Estimate: 0.026 | kNN: 0.054 | Accuracy: 0.820 | (97.20 s)
    Training Complete. Time elapsed: 98.297 s

```python
plot_decision_surface(sess_p, model_p, title='Penalized Decision Surface')
```

### Turn $\lambda$ up a bit

```python
sess_p.close()
args.penalty = True
args.lambda_MI = 10.0
sess_p, model_p, metrics_p = train(args)
```

    Input X shape (512, 2)
    Input X shape (512, 2)
    joint shape (512, 1)
    marginal shape (512, 1)
    Classifier Parameters [<tf.Variable 'classifier/fully_connected/dense/kernel:0' shape=(2, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_1/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_1/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_2/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_2/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/dense/kernel:0' shape=(32, 1) dtype=float32_ref>, <tf.Variable 'classifier/dense/bias:0' shape=(1,) dtype=float32_ref>]
    MI-est Parameters [<tf.Variable 'miner/stat_net/dense/kernel:0' shape=(2, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_1/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_1/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_2/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_2/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_3/kernel:0' shape=(32, 1) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_3/bias:0' shape=(1,) dtype=float32_ref>]
    MI Estimate: -0.008 | kNN: 1.587 | Accuracy: 0.607 | (1.34 s)
    MI Estimate: 0.098 | kNN: 0.335 | Accuracy: 0.789 | (2.66 s)
    MI Estimate: 0.001 | kNN: 0.298 | Accuracy: 0.695 | (3.90 s)
    MI Estimate: 0.004 | kNN: 0.179 | Accuracy: 0.680 | (5.13 s)
    MI Estimate: 0.005 | kNN: 0.173 | Accuracy: 0.666 | (6.35 s)
    MI Estimate: 0.007 | kNN: 0.137 | Accuracy: 0.662 | (7.58 s)
    MI Estimate: 0.002 | kNN: 0.146 | Accuracy: 0.699 | (8.81 s)
    MI Estimate: 0.003 | kNN: 0.010 | Accuracy: 0.701 | (10.04 s)
    MI Estimate: 0.004 | kNN: 0.064 | Accuracy: 0.705 | (11.26 s)
    MI Estimate: 0.003 | kNN: 0.065 | Accuracy: 0.689 | (12.49 s)
    MI Estimate: -0.001 | kNN: 0.016 | Accuracy: 0.680 | (13.71 s)
    MI Estimate: 0.002 | kNN: 0.089 | Accuracy: 0.654 | (14.92 s)
    MI Estimate: -0.001 | kNN: 0.030 | Accuracy: 0.711 | (16.14 s)
    MI Estimate: 0.003 | kNN: 0.021 | Accuracy: 0.709 | (17.35 s)
    MI Estimate: -0.001 | kNN: 0.055 | Accuracy: 0.732 | (18.56 s)
    MI Estimate: 0.001 | kNN: -0.004 | Accuracy: 0.697 | (19.79 s)
    MI Estimate: 0.002 | kNN: 0.068 | Accuracy: 0.723 | (21.01 s)
    MI Estimate: 0.001 | kNN: 0.021 | Accuracy: 0.705 | (22.23 s)
    MI Estimate: 0.003 | kNN: 0.047 | Accuracy: 0.723 | (23.45 s)
    MI Estimate: -0.001 | kNN: 0.021 | Accuracy: 0.711 | (24.66 s)
    MI Estimate: 0.002 | kNN: 0.033 | Accuracy: 0.717 | (25.88 s)
    MI Estimate: -0.003 | kNN: 0.051 | Accuracy: 0.699 | (27.11 s)
    MI Estimate: 0.002 | kNN: 0.040 | Accuracy: 0.680 | (28.33 s)
    MI Estimate: 0.000 | kNN: -0.006 | Accuracy: 0.723 | (29.55 s)
    MI Estimate: 0.002 | kNN: 0.039 | Accuracy: 0.727 | (30.76 s)
    MI Estimate: -0.001 | kNN: 0.069 | Accuracy: 0.705 | (31.98 s)
    MI Estimate: -0.002 | kNN: 0.035 | Accuracy: 0.730 | (33.20 s)
    MI Estimate: -0.000 | kNN: 0.073 | Accuracy: 0.727 | (34.42 s)
    MI Estimate: -0.006 | kNN: 0.060 | Accuracy: 0.734 | (35.64 s)
    MI Estimate: -0.004 | kNN: 0.054 | Accuracy: 0.725 | (36.88 s)
    MI Estimate: -0.004 | kNN: 0.044 | Accuracy: 0.715 | (38.10 s)
    MI Estimate: 0.003 | kNN: 0.058 | Accuracy: 0.736 | (39.31 s)
    MI Estimate: 0.004 | kNN: 0.025 | Accuracy: 0.727 | (40.53 s)
    MI Estimate: 0.000 | kNN: 0.030 | Accuracy: 0.719 | (41.75 s)
    MI Estimate: -0.001 | kNN: -0.006 | Accuracy: 0.742 | (42.98 s)
    MI Estimate: 0.003 | kNN: 0.023 | Accuracy: 0.758 | (44.20 s)
    MI Estimate: 0.002 | kNN: -0.015 | Accuracy: 0.719 | (45.42 s)
    MI Estimate: 0.004 | kNN: 0.052 | Accuracy: 0.738 | (46.63 s)
    MI Estimate: 0.001 | kNN: 0.029 | Accuracy: 0.754 | (47.87 s)
    MI Estimate: -0.001 | kNN: 0.060 | Accuracy: 0.736 | (49.08 s)
    MI Estimate: 0.004 | kNN: -0.031 | Accuracy: 0.734 | (50.29 s)
    MI Estimate: 0.002 | kNN: 0.042 | Accuracy: 0.785 | (51.52 s)
    MI Estimate: -0.001 | kNN: 0.044 | Accuracy: 0.725 | (52.73 s)
    MI Estimate: -0.002 | kNN: 0.027 | Accuracy: 0.766 | (53.96 s)
    MI Estimate: 0.002 | kNN: 0.017 | Accuracy: 0.713 | (55.17 s)
    MI Estimate: -0.001 | kNN: 0.030 | Accuracy: 0.752 | (56.38 s)
    MI Estimate: -0.001 | kNN: 0.037 | Accuracy: 0.715 | (57.60 s)
    MI Estimate: 0.002 | kNN: 0.035 | Accuracy: 0.756 | (58.82 s)
    MI Estimate: 0.001 | kNN: 0.010 | Accuracy: 0.785 | (60.04 s)
    MI Estimate: 0.001 | kNN: 0.090 | Accuracy: 0.775 | (61.25 s)
    MI Estimate: -0.000 | kNN: 0.057 | Accuracy: 0.732 | (62.47 s)
    MI Estimate: -0.000 | kNN: 0.077 | Accuracy: 0.740 | (63.67 s)
    MI Estimate: 0.002 | kNN: 0.012 | Accuracy: 0.771 | (64.87 s)
    MI Estimate: -0.002 | kNN: 0.019 | Accuracy: 0.725 | (66.09 s)
    MI Estimate: -0.000 | kNN: 0.032 | Accuracy: 0.744 | (67.30 s)
    MI Estimate: 0.003 | kNN: 0.074 | Accuracy: 0.793 | (68.51 s)
    MI Estimate: 0.003 | kNN: 0.073 | Accuracy: 0.762 | (69.72 s)
    MI Estimate: 0.003 | kNN: 0.032 | Accuracy: 0.750 | (70.93 s)
    MI Estimate: 0.000 | kNN: 0.006 | Accuracy: 0.773 | (72.13 s)
    MI Estimate: -0.002 | kNN: 0.037 | Accuracy: 0.775 | (73.33 s)
    MI Estimate: -0.001 | kNN: 0.034 | Accuracy: 0.777 | (74.54 s)
    MI Estimate: 0.001 | kNN: 0.034 | Accuracy: 0.756 | (75.75 s)

```python
plot_decision_surface(sess_p, model_p, title='Penalized Decision Surface, $\lambda=10$')
```

### Use KL Loss

```python
sess_p.close()
args.penalty = True
args.lambda_MI = 1.0
args.KL_loss = True
sess_p, model_p, metrics_p_kl = train(args)
```

    Input X shape (512, 2)
    Input X shape (512, 2)
    joint shape (512, 1)
    marginal shape (512, 1)
    Classifier Parameters [<tf.Variable 'classifier/fully_connected/dense/kernel:0' shape=(2, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_1/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_1/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_2/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'classifier/fully_connected/dense_2/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'classifier/dense/kernel:0' shape=(32, 1) dtype=float32_ref>, <tf.Variable 'classifier/dense/bias:0' shape=(1,) dtype=float32_ref>]
    MI-est Parameters [<tf.Variable 'miner/stat_net/dense/kernel:0' shape=(2, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_1/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_1/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_2/kernel:0' shape=(32, 32) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_2/bias:0' shape=(32,) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_3/kernel:0' shape=(32, 1) dtype=float32_ref>, <tf.Variable 'miner/stat_net/dense_3/bias:0' shape=(1,) dtype=float32_ref>]
    WARNING:tensorflow:From /usr/local/easybuild/software/Tensorflow/1.13.1-intel-2017.u2-GCC-6.2.0-CUDA10-Python-3.6.4-GPU/lib/python3.6/site-packages/tensorflow/python/ops/math_ops.py:3066: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.
    Instructions for updating:
    Use tf.cast instead.
    WARNING:tensorflow:From /usr/local/easybuild/software/Tensorflow/1.13.1-intel-2017.u2-GCC-6.2.0-CUDA10-Python-3.6.4-GPU/lib/python3.6/site-packages/tensorflow/python/ops/math_grad.py:102: div (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.
    Instructions for updating:
    Deprecated in favor of operator or tf.math.divide.
    MI Estimate: -0.003 | kNN: 0.224 | Accuracy: 0.516 | (1.40 s)
    MI Estimate: 0.010 | kNN: 0.274 | Accuracy: 0.678 | (2.72 s)
    MI Estimate: 0.022 | kNN: 0.065 | Accuracy: 0.734 | (3.95 s)
    MI Estimate: -0.015 | kNN: 0.057 | Accuracy: 0.740 | (5.17 s)
    MI Estimate: 0.017 | kNN: 0.080 | Accuracy: 0.750 | (6.39 s)
    MI Estimate: 0.015 | kNN: 0.056 | Accuracy: 0.762 | (7.59 s)
    MI Estimate: 0.011 | kNN: 0.075 | Accuracy: 0.777 | (8.80 s)
    MI Estimate: 0.019 | kNN: 0.045 | Accuracy: 0.781 | (10.00 s)
    MI Estimate: 0.024 | kNN: -0.004 | Accuracy: 0.744 | (11.21 s)
    MI Estimate: 0.022 | kNN: 0.049 | Accuracy: 0.791 | (12.43 s)
    MI Estimate: 0.024 | kNN: 0.085 | Accuracy: 0.783 | (13.65 s)
    MI Estimate: 0.017 | kNN: 0.037 | Accuracy: 0.814 | (14.86 s)
    MI Estimate: 0.004 | kNN: 0.031 | Accuracy: 0.762 | (16.07 s)
    MI Estimate: 0.013 | kNN: 0.075 | Accuracy: 0.795 | (17.28 s)
    MI Estimate: 0.021 | kNN: 0.025 | Accuracy: 0.799 | (18.49 s)
    MI Estimate: 0.028 | kNN: 0.085 | Accuracy: 0.779 | (19.70 s)
    MI Estimate: 0.006 | kNN: 0.064 | Accuracy: 0.750 | (20.91 s)
    MI Estimate: 0.015 | kNN: 0.094 | Accuracy: 0.775 | (22.11 s)
    MI Estimate: -0.007 | kNN: 0.062 | Accuracy: 0.744 | (23.32 s)
    MI Estimate: 0.019 | kNN: 0.115 | Accuracy: 0.777 | (24.52 s)
    MI Estimate: 0.002 | kNN: 0.101 | Accuracy: 0.775 | (25.74 s)
    MI Estimate: 0.025 | kNN: -0.003 | Accuracy: 0.787 | (26.95 s)
    MI Estimate: -0.002 | kNN: 0.044 | Accuracy: 0.787 | (28.15 s)
    MI Estimate: 0.014 | kNN: 0.021 | Accuracy: 0.760 | (29.36 s)
    MI Estimate: 0.006 | kNN: 0.075 | Accuracy: 0.754 | (30.57 s)
    MI Estimate: 0.023 | kNN: 0.016 | Accuracy: 0.773 | (31.78 s)
    MI Estimate: 0.017 | kNN: 0.088 | Accuracy: 0.754 | (32.99 s)
    MI Estimate: 0.013 | kNN: 0.022 | Accuracy: 0.775 | (34.20 s)
    MI Estimate: 0.032 | kNN: 0.067 | Accuracy: 0.799 | (35.41 s)
    MI Estimate: 0.001 | kNN: 0.026 | Accuracy: 0.754 | (36.61 s)
    MI Estimate: 0.009 | kNN: 0.068 | Accuracy: 0.777 | (37.82 s)
    MI Estimate: 0.010 | kNN: 0.030 | Accuracy: 0.785 | (39.03 s)
    MI Estimate: 0.016 | kNN: 0.055 | Accuracy: 0.779 | (40.23 s)
    MI Estimate: 0.018 | kNN: 0.082 | Accuracy: 0.789 | (41.44 s)
    MI Estimate: 0.009 | kNN: 0.053 | Accuracy: 0.768 | (42.65 s)
    MI Estimate: -0.001 | kNN: 0.047 | Accuracy: 0.816 | (43.86 s)
    MI Estimate: 0.013 | kNN: 0.050 | Accuracy: 0.777 | (45.07 s)
    MI Estimate: 0.017 | kNN: 0.095 | Accuracy: 0.762 | (46.27 s)
    MI Estimate: 0.012 | kNN: 0.070 | Accuracy: 0.801 | (47.48 s)
    MI Estimate: 0.012 | kNN: 0.039 | Accuracy: 0.791 | (48.68 s)
    MI Estimate: -0.003 | kNN: 0.030 | Accuracy: 0.805 | (49.88 s)
    MI Estimate: 0.022 | kNN: 0.059 | Accuracy: 0.773 | (51.09 s)
    MI Estimate: -0.008 | kNN: 0.091 | Accuracy: 0.797 | (52.30 s)
    MI Estimate: 0.004 | kNN: 0.046 | Accuracy: 0.799 | (53.51 s)
    MI Estimate: 0.025 | kNN: 0.029 | Accuracy: 0.811 | (54.70 s)
    MI Estimate: 0.013 | kNN: 0.053 | Accuracy: 0.736 | (55.92 s)
    MI Estimate: 0.002 | kNN: 0.013 | Accuracy: 0.771 | (57.12 s)
    MI Estimate: -0.004 | kNN: -0.011 | Accuracy: 0.811 | (58.33 s)
    MI Estimate: -0.000 | kNN: 0.026 | Accuracy: 0.746 | (59.54 s)
    MI Estimate: 0.004 | kNN: 0.035 | Accuracy: 0.814 | (60.76 s)
    MI Estimate: 0.011 | kNN: 0.080 | Accuracy: 0.777 | (61.97 s)
    MI Estimate: 0.004 | kNN: 0.025 | Accuracy: 0.773 | (63.18 s)
    MI Estimate: 0.006 | kNN: 0.046 | Accuracy: 0.764 | (64.39 s)
    MI Estimate: -0.001 | kNN: 0.028 | Accuracy: 0.775 | (65.59 s)
    MI Estimate: 0.013 | kNN: 0.013 | Accuracy: 0.801 | (66.80 s)
    MI Estimate: 0.006 | kNN: 0.078 | Accuracy: 0.756 | (68.01 s)
    MI Estimate: 0.018 | kNN: -0.004 | Accuracy: 0.771 | (69.24 s)
    MI Estimate: 0.009 | kNN: 0.108 | Accuracy: 0.758 | (70.45 s)
    MI Estimate: 0.017 | kNN: 0.072 | Accuracy: 0.766 | (71.66 s)
    MI Estimate: 0.011 | kNN: 0.016 | Accuracy: 0.771 | (72.87 s)
    MI Estimate: 0.010 | kNN: 0.044 | Accuracy: 0.762 | (74.08 s)
    MI Estimate: 0.024 | kNN: 0.083 | Accuracy: 0.779 | (75.28 s)
    MI Estimate: -0.008 | kNN: 0.076 | Accuracy: 0.768 | (76.49 s)
    MI Estimate: 0.007 | kNN: 0.078 | Accuracy: 0.793 | (77.69 s)
    MI Estimate: 0.012 | kNN: 0.064 | Accuracy: 0.795 | (78.90 s)
    MI Estimate: 0.009 | kNN: 0.043 | Accuracy: 0.795 | (80.10 s)
    MI Estimate: 0.009 | kNN: 0.035 | Accuracy: 0.791 | (81.32 s)
    MI Estimate: 0.026 | kNN: 0.062 | Accuracy: 0.770 | (82.53 s)
    MI Estimate: 0.022 | kNN: 0.083 | Accuracy: 0.777 | (83.74 s)
    MI Estimate: 0.016 | kNN: 0.041 | Accuracy: 0.766 | (84.95 s)
    MI Estimate: 0.010 | kNN: 0.050 | Accuracy: 0.770 | (86.17 s)
    MI Estimate: 0.027 | kNN: 0.091 | Accuracy: 0.771 | (87.38 s)
    MI Estimate: 0.022 | kNN: 0.055 | Accuracy: 0.787 | (88.59 s)
    MI Estimate: 0.018 | kNN: 0.085 | Accuracy: 0.781 | (89.80 s)
    MI Estimate: 0.004 | kNN: 0.069 | Accuracy: 0.785 | (91.01 s)
    MI Estimate: 0.023 | kNN: 0.059 | Accuracy: 0.787 | (92.22 s)
    MI Estimate: 0.018 | kNN: 0.065 | Accuracy: 0.795 | (93.43 s)
    MI Estimate: 0.004 | kNN: 0.072 | Accuracy: 0.789 | (94.65 s)
    MI Estimate: 0.028 | kNN: 0.041 | Accuracy: 0.797 | (95.87 s)
    MI Estimate: 0.014 | kNN: 0.078 | Accuracy: 0.768 | (97.07 s)
    Training Complete. Time elapsed: 98.156 s

```python
plot_decision_surface(sess_p, model_p, title='Decision Surface with KL penalty, $\lambda=1$')
```

### Convergence

Look at joint loss, classifier loss, classifier performance.

```python
# No penalty
trials = 25
metric_runs = list()
for _ in range(trials):
    print('Trial: {}'.format(_))
    sess, model, metrics = train(args)
    metric_runs.append(metrics)
    sess.close()
metrics_df = pd.concat(metric_runs)
```

```python
sess.close()
args.penalty = True
args.lambda_MI = 1.0
metric_runs = list()
for _ in range(trials):
    print('Trial: {}'.format(_))
    sess_p, model_p, metrics_p = train(args)
    metric_runs.append(metrics_p)
    sess_p.close()
metrics_p_df = pd.concat(metric_runs)
```

```python
sess_p.close()
args.penalty = True
args.lambda_MI = 1.0
args.KL_loss = True
metric_runs = list()
for _ in range(trials):
    print('Trial: {}'.format(_))
    sess_p, model_p, metrics_p = train(args)
    metric_runs.append(metrics_p)
    sess_p.close()
metrics_p_kl_df = pd.concat(metric_runs)
```

```python
sess_p.close()
args.penalty = True
args.lambda_MI = 10.0
args.KL_loss = True
metric_runs = list()
for _ in range(trials):
    print('Trial: {}'.format(_))
    sess_p, model_p, metrics_p = train(args)
    metric_runs.append(metrics_p)
    sess_p.close()
metrics_p10_df = pd.concat(metric_runs)
```

```python
sess_p.close()
args.penalty = True
args.lambda_MI = 10.0
args.KL_loss = False
metric_runs = list()
for _ in range(trials):
    print('Trial: {}'.format(_))
    sess_p, model_p, metrics_p = train(args)
    metric_runs.append(metrics_p)
    sess_p.close()
metrics_p10_df_kl = pd.concat(metric_runs)
```

```python
# Convergence plots
smooth = lambda df: df.rolling(window=2, center=True).mean().dropna()
for col in metrics_df.columns:
    plt.figure()
    sns.lineplot(x=metrics_df.index, y=col, data=metrics_df, label='Unpenalized')
    sns.lineplot(x=metrics_p_df.index, y=col, data=metrics_p_df, label='JS-pen')
    sns.lineplot(x=metrics_p_kl_df.index, y=col, data=metrics_p_kl_df, label='KL-pen')
    plt.title(col)
    plt.legend(loc='best')
    plt.xlabel('Iterations')
```

```python

```


# Introduction to seismic modelling - 2D acoustic propagator

In this notebook we will highlight various aspects of seismic inversion based on Devito operators. Here we aim to highlight the core ideas behind seismic modelling, where we create a numerical model that captures the processes involved in a seismic survey. This forward model will then form the basis for further 'additional' tutorials (located in `additional_material`) on the implementation of inversion processes using Devito operators.

## Goals

- Introduce the seismic `Model` class
- Introduce 'sources' and 'receivers' (sparse operations)
- Build an acoustic forward propagator

## Modelling workflow

The core process we are aiming to model is a seismic survey, which consists of two main components:

- **Source** - A source is positioned at a single or a few physical locations where artificial pressure is injected into the domain we want to model. In the case of land survey, it is usually dynamite blowing up at a given location, or a vibroseis (a vibrating engine generating continuous sound waves). For a marine survey, the source is an air gun sending a bubble of compressed air into the water that will expand and generate a seismic wave.
- **Receiver** - A set of microphones or hydrophones are used to measure the resulting wave and create a set of measurements called a *Shot Record*. These measurements are recorded at multiple locations, and usually at the surface of the domain or at the bottom of the ocean in some marine cases.

In order to create a numerical model of a seismic survey, we need to solve the wave equation and implement source and receiver interpolation to inject the source and record the seismic wave at sparse point locations in the grid.

## The acoustic seismic wave equation
The acoustic wave equation for the square slowness $m$, defined as $m=\frac{1}{c^2}$, where $c$ is the speed of sound in the given physical media, and a source $q$ is given by:

\begin{cases}
 &m \frac{\partial^2 u(\mathbf{x},t)}{\partial t^2} - \nabla^2 u(\mathbf{x},t) = q \ \text{in } \Omega \\
 &u(\mathbf{x},0) = 0 \\
 &\frac{\partial u(\mathbf{x},t)}{\partial t}|_{t=0} = 0
\end{cases}

with the zero initial conditions to guarantee unicity of the solution.
The boundary conditions are Dirichlet conditions:
\begin{equation}
 u(\mathbf{x},t)|_{\delta\Omega} = 0
\end{equation}

where $\delta\Omega$ is the surface of the boundary of the model $\Omega$.

# Finite domains

The last piece of the puzzle is the computational limitation. In the field, the seismic wave propagates in every direction to an ''infinite'' distance. However, solving the wave equation in a mathematically/discrete infinite domain is not feasible. In order to compensate, Absorbing Boundary Conditions (ABC) or Perfectly Matched Layers (PML) are required to mimic an infinite domain. These two methods allow us to approximate an unbounded media by damping and absorbing the waves at the limit of the domain to avoid reflections.

The simplest of these methods is the absorbing damping mask. (Note that we will explore the implementation of a more sophisticated boundary condition in the following notebook).The core idea is to extend the physical domain and to add a *sponge mask* in this extension that will absorb the incident waves. The acoustic wave equation with this damping mask can be rewritten as:

\begin{cases}
 &m \frac{\partial^2 u(\mathbf{x},t)}{\partial t^2} - \nabla^2 u(\mathbf{x},t) + \eta \frac{\partial u(\mathbf{x},t)}{\partial t}=q  \ \text{in } \Omega \\
 &u(\mathbf{x},0) = 0 \\
 &\frac{\partial u(\mathbf{x},t)}{\partial t}|_{t=0} = 0
\end{cases}

where $\eta$ is the damping mask which is equal to $0$ inside the physical domain and increasing inside the sponge layer. Multiple choice of profile can be chosen for $\eta$ from, e.g., linear to exponential.

# Seismic modelling with devito

We describe here a step by step setup of seismic modelling with Devito in a simple 2D case. We will create a physical model of our domain and define a single source along with a set of receivers. But first, we initialize some basic utilities.

```python
import numpy as np
%matplotlib inline
```

## Define the physical problem

The first step is to define the physical model:

- What are the physical dimensions of interest
- What is the velocity profile of this physical domain

We will create a simple velocity model here by hand for demonstration purposes. This model essentially consists of two layers, each with a different velocity: $1.5km/s$ in the top layer and $2.5km/s$ in the bottom layer. We will use this simple model a lot in the following tutorials, so we will rely on a utility function to create it again later.

```python
#NBVAL_IGNORE_OUTPUT
from examples.seismic import Model, plot_velocity

# Define a physical size
shape = (101, 101)  # Number of grid point (nx, nz)
spacing = (10., 10.)  # Grid spacing in m. The domain size is now 1km by 1km
origin = (0., 0.)  # What is the location of the top left corner. This is necessary to define
# the absolute location of the source and receivers

# Define a velocity profile. The velocity is in km/s
v = np.empty(shape, dtype=np.float32)
v[:, :51] = 1.5
v[:, 51:] = 2.5

# With the velocity and model size defined, we can create the seismic model that
# encapsulates this properties. We also define the size of the absorbing layer as 10 grid points
model = Model(vp=v, origin=origin, shape=shape, spacing=spacing,
              space_order=2, nbl=10, bcs="damp")

plot_velocity(model)
```

# Acquisition geometry

To fully define our problem setup we also need to define the source that injects the wave to the model and the set of receiver locations at which to the wavefield is sampled. The source time signature will be modelled using a 'Ricker' wavelet whose form is defined as

\begin{equation}
  q(t) = (1-2\pi^2 f_0^2 (t - \frac{1}{f_0})^2 )e^{- \pi^2 f_0^2 (t - \frac{1}{f_0})}.
\end{equation}

To include such a source signature in our model we first need to define the time duration of our model and the timestep size, which is dictated by the CFL condition and our grid spacing. Luckily, our `Model` utility provides us with the critical timestep size, so we can fully discretize our model time axis as an array:

```python
from examples.seismic import TimeAxis

t0 = 0.  # Simulation starts a t=0
tn = 1000.  # Simulation last 1 second (1000 ms)
dt = model.critical_dt  # Time step from model grid spacing

time_range = TimeAxis(start=t0, stop=tn, step=dt)
```

Recall that we can view the documentation for the `TimeAxis` object (or any other object via):

```python
print(time_range.__doc__)
```

        Data object to store the TimeAxis. Exactly three of the four key arguments
        must be prescribed. Because of remainder values it is not possible to create
        a TimeAxis that exactly adhears to the inputs therefore start, stop, step
        and num values should be taken from the TimeAxis object rather than relying
        upon the input values.

        The four possible cases are:
        start is None: start = step*(1 - num) + stop
        step is None: step = (stop - start)/(num - 1)
        num is None: num = ceil((stop - start + step)/step);
                     because of remainder stop = step*(num - 1) + start
        stop is None: stop = step*(num - 1) + start

        Parameters
        ----------
        start : float, optional
            Start of time axis.
        step : float, optional
            Time interval.
        num : int, optional
            Number of values (Note: this is the number of intervals + 1).
            Stop value is reset to correct for remainder.
        stop : float, optional
            End time.

(Note: Values of the time steps are stored in `time_range.time_values`.)

We position the source at a depth of $20m$ and in the center of the $x$ axis ($\mathbf{x}_{src}=(500m, 20m)$), with a peak wavelet frequency of $10Hz$.

```python
#NBVAL_IGNORE_OUTPUT
from examples.seismic import RickerSource

f0 = 0.010  # Source peak frequency is 10Hz (0.010 kHz)
src = RickerSource(name='src', grid=model.grid, f0=f0,
                   npoint=1, time_range=time_range)

# Set source coordinates
src.coordinates.data[0, 0] = np.array(model.domain_size[0]) * .5
src.coordinates.data[0, 1] = 20.  # Depth is 20m

# We can plot the time signature to see the wavelet via:
src.show()
```

Similarly to our source object, we can now define our receiver geometry as a symbol of type `Receiver`. It is worth noting here that both utility classes, `RickerSource` and `Receiver` are thin wrappers around Devito's `SparseTimeFunction` type, which encapsulates sparse point data and allows us to inject and interpolate values into and out of the computational grid. As we have already seen, both types provide a `.coordinates` property to define the position within the domain of all points encapsulated by that symbol.

In this example we will position receivers at the same depth as the source, every $10m$ along the x axis. The `rec.data` property will be initialized, but left empty, as we will compute the receiver readings during the simulation.

```python
#NBVAL_IGNORE_OUTPUT
from examples.seismic import Receiver

# Create symbol for 101 receivers
rec = Receiver(name='rec', grid=model.grid, npoint=101, time_range=time_range)

# Prescribe even spacing for receivers along the x-axis
rec.coordinates.data[:, 0] = np.linspace(0, model.domain_size[0], num=101)
rec.coordinates.data[:, 1] = 20.  # Depth is 20m

# We can now show the source and receivers within our domain:
# Red dot: Source location
# Green dots: Receiver locations (every 4th point)
plot_velocity(model, source=src.coordinates.data,
              receiver=rec.coordinates.data[::4, :])
```

# Finite-difference discretization

The finite-difference approximation is derived from Taylor expansions of the continuous field after removing the error term.

## Time discretization

We only consider the second order time discretization for now. From the Taylor expansion, the second order discrete approximation of the second order time derivative is:
\begin{equation}
\begin{aligned}
 \frac{\partial^2 u(\mathbf{x},t)}{\partial t^2} = \frac{\mathbf{u}(\mathbf{x},\mathbf{t+\Delta t}) - 2 \mathbf{u}(\mathbf{x},\mathbf{t}) + \mathbf{u}(\mathbf{x},\mathbf{t-\Delta t})}{\mathbf{\Delta t}^2} + O(\mathbf{\Delta t}^2).
\end{aligned}
\end{equation}

where $\mathbf{u}$ is the discrete wavefield, $\mathbf{\Delta t}$ is the discrete
time-step (distance between two consecutive discrete time points) and $O(\mathbf{\Delta
  t}^2)$ is the temporal discretization error term. The discretized approximation of the
second order time derivative is then given by dropping the error term. This derivative is represented in Devito by `u.dt2` where u is a `TimeFunction` object.

## Spatial discretization

We define the discrete Laplacian as the sum of the second order spatial
derivatives in the three dimensions:
\begin{equation}
\begin{aligned}
\Delta \mathbf{u}(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{t})= \sum_{j=1}^{j=\frac{k}{2}} \Bigg[\alpha_j \Bigg(&
\mathbf{u}(\mathbf{x+jdx},\mathbf{y},\mathbf{z},\mathbf{t})+\mathbf{u}(\mathbf{x-jdx},\mathbf{y},\mathbf{z},\mathbf{t}) + \\
&\mathbf{u}(\mathbf{x},\mathbf{y+jdy},\mathbf{z},\mathbf{t})+\mathbf{u}(\mathbf{x},\mathbf{y-jdy},\mathbf{z}\mathbf{t}) + \\
&\mathbf{u}(\mathbf{x},\mathbf{y},\mathbf{z+jdz},\mathbf{t})+\mathbf{u}(\mathbf{x},\mathbf{y},\mathbf{z-jdz},\mathbf{t})\Bigg) \Bigg] + \\
&3\alpha_0 \mathbf{u}(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{t}).
\end{aligned}
\end{equation}

where $\alpha_j$ are the appropriate finite difference weights for the approximation. This derivative is represented in Devito by `u.laplace`.

## Wave equation

With the space and time discretization defined, we can fully discretize the wave-equation with the combination of time and space discretizations and obtain the following second order in time and $k^{th}$ order in space discrete stencil to update one grid point at position $\mathbf{x}, \mathbf{y},\mathbf{z}$ at time $\mathbf{t}$, i.e.
\begin{equation}
\begin{aligned}
\mathbf{u}(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{t+\Delta t}) = &2\mathbf{u}(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{t}) - \mathbf{u}(\mathbf{x},\mathbf{y}, \mathbf{z},\mathbf{t-\Delta t}) +\\
& \frac{\mathbf{\Delta t}^2}{\mathbf{m(\mathbf{x},\mathbf{y},\mathbf{z})}} \Big(\Delta \mathbf{u}(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{t}) + \mathbf{q}(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{t}) \Big).
\end{aligned}
\end{equation}

```python
# In order to represent the wavefield u and the square slowness we need symbolic objects
# corresponding to time-space-varying field (u, TimeFunction) and
# space-varying field (m, Function)
from devito import TimeFunction

# Define the wavefield with the size of the model and the time dimension
u = TimeFunction(name="u", grid=model.grid, time_order=2, space_order=2)

# We can now write the PDE
pde = model.m * u.dt2 - u.laplace + model.damp * u.dt

# The PDE representation is as on paper
pde
```

$\displaystyle \operatorname{damp}{\left(x,y \right)} \frac{\partial}{\partial t} u{\left(t,x,y \right)} - \frac{\partial^{2}}{\partial x^{2}} u{\left(t,x,y \right)} - \frac{\partial^{2}}{\partial y^{2}} u{\left(t,x,y \right)} + \frac{\frac{\partial^{2}}{\partial t^{2}} u{\left(t,x,y \right)}}{\operatorname{vp}^{2}{\left(x,y \right)}}$

```python
# This discrete PDE can be solved in a time-marching way updating u(t+dt) from the previous time step
# Devito as a shortcut for u(t+dt) which is u.forward. We can then rewrite the PDE as
# a time marching updating equation known as a stencil using customized SymPy functions
from devito import Eq, solve

stencil = Eq(u.forward, solve(pde, u.forward))
```

# Source injection and receiver interpolation

With a numerical scheme to solve the homogenous wave equation, we need to add the source to introduce seismic waves and to implement the measurement operator, and interpolation operator. This operation is linked to the discrete scheme and needs to be done at the proper time step. The semi-discretized in time wave equation with a source reads:

\begin{equation}
\begin{aligned}
\mathbf{u}(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{t+\Delta t}) = &2\mathbf{u}(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{t}) - \mathbf{u}(\mathbf{x},\mathbf{y}, \mathbf{z},\mathbf{t-\Delta t}) +\\
& \frac{\mathbf{\Delta t}^2}{\mathbf{m(\mathbf{x},\mathbf{y},\mathbf{z})}} \Big(\Delta \mathbf{u}(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{t}) + \mathbf{q}(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{t}) \Big).
\end{aligned}
\end{equation}

It shows that in order to update $\mathbf{u}$ to the time $\mathbf{t+\Delta t}$ we must inject the value of the source term $\mathbf{q}$ at time $\mathbf{t}$. In Devito, this corresponds to the update of $u$ at index $t+1$ (t = time implicitly) with the source of time $t$.

```python
# Finally we define the source injection and receiver read function to generate the corresponding code
src_term = src.inject(field=u.forward, expr=src * dt**2 / model.m)

# Create interpolation expression for receivers
rec_term = rec.interpolate(expr=u.forward)
```

# Devito operator and solve
After constructing all the necessary expressions for updating the wavefield, injecting the source term and interpolating onto the receiver points, we can now create the Devito operator that will generate the C code at runtime. When creating the operator, Devito's two optimization engines will log which performance optimizations have been performed:
* **DSE:** The Devito Symbolics Engine will attempt to reduce the number of operations required by the kernel.
* **DLE:** The Devito Loop Engine will perform various loop-level optimizations to improve runtime performance.

**Note**: The argument `subs=model.spacing_map` causes the operator to substitute values for our current grid spacing into the expressions before code generation. This reduces the number of floating point operations executed by the kernel by pre-evaluating certain coefficients.

```python
#NBVAL_IGNORE_OUTPUT
from devito import Operator

op = Operator([stencil] + src_term + rec_term, subs=model.spacing_map)
```

Now we can execute the create operator for a number of timesteps. We specify the number of timesteps to compute with the keyword `time` and the timestep size with `dt`.

```python
#NBVAL_IGNORE_OUTPUT
op(time=time_range.num-1, dt=model.critical_dt)
```

    Operator `Kernel` run in 0.01 s

    PerformanceSummary([(PerfKey(name='section0', rank=None),
                         PerfEntry(time=0.0026959999999999888, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
                        (PerfKey(name='section1', rank=None),
                         PerfEntry(time=2.5000000000000008e-05, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
                        (PerfKey(name='section2', rank=None),
                         PerfEntry(time=0.00042800000000000276, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])

After running our operator kernel, the data associated with the receiver symbol `rec.data` has now been populated due to the interpolation expression we inserted into the operator. This allows us the visualize the shot record:

```python
#NBVAL_IGNORE_OUTPUT
from examples.seismic import plot_shotrecord

plot_shotrecord(rec.data, model, t0, tn)
```

# Running Devito programs with MPI

MPI is immediatley supported for any Devito program. Simply type, e.g.:

```
DEVITO_MPI=1 mpirun <mpi args> python my_program.py
```

## Exercise: Run script with MPI

In the folder `additional_material` there is a python script named `acoustic_script.py` which essentially the python code introduced above but with the plotting instructions removed.

As a quick exercise, let us run this script with `mpirun` and satisfy ourselves that it returns the same results when run over 1 and 2 ranks:
- Before running with mpi we first need to install `mpi4py`. To do so, run the two following commands:
    - `sudo apt install mpich -y`
    - Then, in the `Devito` directory (`~/devito`) `pip install -r requirements-optional.txt --user`
- In the `additional_material` folder within the `transform2020` directory first run the script in serial: `python acoustic_script.py`.
- Then run over two mpi processes (`DEVITO_MPI=1 mpirun -np 2 python acoustic_script.py`) and check the resultant `norm` is the same. (Note 1: When running over two mpi processes each rank will print the norm to screen. Note 2: If executing via the jupyterhub terminal please do not run over 'too' many mpi processes as this may cause issues for other users.)