# TMM for the TE And TM polarization only

One thing which isn't super clear to me is the fact that in scattering matrices, if you consider off-normal incidence, the incoming ray has two ports, the reflected ray has two ports and likewise for the transmission side.

The Maxwell's equations, typically formulated, is a 4x4 system of equations, but these don't actually correspond to the ports, do they

### Significance of the transfer matrix
The transfer matrix T is just a transfer function telling me how to map input amplitudes to output amplitudes. In dealing with homogeneous structures, the only inputs are plane waves and the only outputs are also plane waves.

\begin{equation}
    \begin{bmatrix}
        T_{11} & T_{12}\\
        T_{21} & T_{22}
    \end{bmatrix}
\end{equation}

The transfer matrix, in effect, works the same as the scattering matrix (assuming the modes are sorted, which they trivially will be for the TE and TM case (since we only have a 2-by-2 matrix)

```python
import sys
sys.path.append('../')
from TMM_functions import PQ_matrices as pq
from TMM_functions import scatter_matrices as sm
from TMM_functions import redheffer_star as rs
from TMM_functions import generate_initial_conditions as ic
import numpy as np
import matplotlib.pyplot as plt
from numpy.linalg import cond
import cmath;
from scipy import linalg as LA
from numpy.linalg import solve as bslash
import time
from convolution_matrices.convmat1D import *
# from RCWA_1D_functions.grating_fft.grating_conv import *

def nonHermitianEigenSorter(eigenvalues):
    N = len(eigenvalues);
    sorted_indices=[];
    sorted_eigs = [];
    for i in range(N):
        eig = eigenvalues[i];
        if(np.real(eig)>0 and np.imag(eig) == 0):
            sorted_indices.append(i); sorted_eigs.append(eig);
        elif(np.real(eig)==0 and np.imag(eig) > 0):
            sorted_indices.append(i); sorted_eigs.append(eig);
        elif(np.real(eig)>0 and abs(np.imag(eig)) > 0):
            sorted_indices.append(i); sorted_eigs.append(eig);
    return sorted_eigs, sorted_indices;

# Moharam et. al Formulation for stable and efficient implementation for RCWA
plt.close("all")
'''
'''
L0 = 1e-6;
e0 = 8.854e-12;
mu0 = 4*np.pi*1e-8;
c0 = 1/np.sqrt(e0*mu0)

n_ridge = 3.48;              # ridge
n_groove = 3.48;                # groove (unit-less)
lattice_constant = 0.7;  # SI units

## simulation parameters
theta = (0)*np.pi/180;
phi = (0)*np.pi/180;
e_r = 1;
m_r = 1;
e_t = 1;
m_t = 1;
n_i = np.sqrt(e_r);

normal_vector= [0,0,-1]

# epsilon_tensor = np.array([[12,1,0],[0,2,0],[1,0,12]]);
# epsilon_tensor = np.array([[12,10,0],[10,12,0],[0,0,2]]);
# epsilon_tensor = np.array([[12,5,0],[5,12,0],[0,0,12]]);
# epsilon_tensor = np.array([[12,1,0],[0,2,0],[1,0,12]]);

epsilon_tensor = np.array([[12,0,10],[0,12,0],[10,0,12]]);
## if the exy eyx term is equal to exx eyy, then there is a degeneracy which you have to account for.
epsilon_tensor = np.array([[12,0,0],[0,12,5],[0,10,12]]);

mu_tensor = np.array([[1,0,0],[0,1,0],[0,0,1]])
I = np.matrix(np.eye(2))

## specify the initial kx and ky
k0=1
kx = k0 *np.sin(theta)*np.cos(phi);

```

```python

```

```python
## =================== specify initial conditions ================== ##
# c+ and c- are now one dimensional.
cref = np.array([1,1]); #incident light
```

```python
## vacuum
a12 = -1;
a21 = kx**2/1 -1
A0 = np.matrix([[0, a12],
                  [a21, 0]]);
print(A0)
lam0, W0 = np.linalg.eig(A0); #this is W
print(W0)
```

    [[ 0. -1.]
     [-1.  0.]]
    [[ 0.70710678  0.70710678]
     [-0.70710678  0.70710678]]

## DOING TMM for the TE and TM polarization using Scattering Matrices
Now, we only have one transverse E and H component. The first thing to note is how are they related? (should be with a dz factor)

### Important Questions to Answer
Why do we need an eigenmode decomposition in the general case (especially since analytically, we don't need to do that)

```python
ref = [];
trans = []
wvlen_scan = np.linspace(0.5, 4, 100);
thickness = 0.5; # 1 layer
E0 = 1;
for wvlen in wvlen_scan:
    k0 = 2*np.pi/wvlen;
    omega = c0*k0;
    kx = np.sin(theta)*np.cos(phi);
    # we will build the system by rows?
    a11 = -1j*(kx*(epsilon_tensor[2,0]/epsilon_tensor[2,2]))
    a12 = -epsilon_tensor[0,0] + epsilon_tensor[0,2]*epsilon_tensor[2,0]/epsilon_tensor[2,2]
    a21 = kx**2/epsilon_tensor[2,2] -1
    a22 = 1j*kx*epsilon_tensor[2,0]/epsilon_tensor[2,2];

    A = np.matrix([[a11, a12],
                  [a21, a22]]);
    #print(np.linalg.cond(A))
    eigenvals, eigenmodes = np.linalg.eig(A); #this is W
    rounded_eigenvals = np.round(eigenvals, 3)
    propagation_phase = np.diag(np.exp(1j*eigenvals*k0*thickness)); #we shouldn't have to scale by 1j...theoretically
    print(propagation_phase); #one is huge, and one is small....
    #print(propagation_phase.shape)
#     plt.plot(np.real(rounded_eigenvals), np.imag(rounded_eigenvals), 'o')
#     plt.show()

    ## need Wref and Wtran
    kz0 = np.sqrt(omega**2-kx**2)

    W1 = eigenmodes;
    ## compute transfer matrix;
    T = np.linalg.inv(W0)@W1@propagation_phase;
    ## process the coefficients
    #print(T)
    ## convert ctran into a transmission
    # what about reflection? (typically, we derived these from the scattering matrix)
    Er = T[0,0]*E0
    Et = T[0,1]*E0; #only one out of plane component so we don't need to recover any additional components

    r = np.linalg.norm(Er);
    t = np.linalg.norm(Et);
    ref.append(r);
    trans.append(t);

ref = np.array(ref);
trans = np.array(trans);
```

    [[-0.97466986+0.22364851j  0.        +0.j        ]
     [ 0.        +0.j         -0.97466986-0.22364851j]]
    [[0.09198012+0.99576084j 0.        +0.j        ]
     [0.        +0.j         0.09198012-0.99576084j]]
    [[0.97602508+0.21765809j 0.        +0.j        ]
     [0.        +0.j         0.97602508-0.21765809j]]
    [[0.62712299-0.77892025j 0.        +0.j        ]
     [0.        +0.j         0.62712299+0.77892025j]]
    [[-0.30684896-0.95175822j  0.        +0.j        ]
     [ 0.        +0.j         -0.30684896+0.95175822j]]
    [[-0.93138898-0.36402552j  0.        +0.j        ]
     [ 0.        +0.j         -0.93138898+0.36402552j]]
    [[-0.91073385+0.41299378j  0.        +0.j        ]
     [ 0.        +0.j         -0.91073385-0.41299378j]]
    [[-0.40981348+0.91216934j  0.        +0.j        ]
     [ 0.        +0.j         -0.40981348-0.91216934j]]
    [[0.23310795+0.97245087j 0.        +0.j        ]
     [0.        +0.j         0.23310795-0.97245087j]]
    [[0.74194828+0.67045712j 0.        +0.j        ]
     [0.        +0.j         0.74194828-0.67045712j]]
    [[0.98314026+0.18285304j 0.        +0.j        ]
     [0.        +0.j         0.98314026-0.18285304j]]
    [[0.94821604-0.31762611j 0.        +0.j        ]
     [0.        +0.j         0.94821604+0.31762611j]]
    [[0.70274915-0.71143773j 0.        +0.j        ]
     [0.        +0.j         0.70274915+0.71143773j]]
    [[0.33861532-0.9409249j 0.        +0.j       ]
     [0.        +0.j        0.33861532+0.9409249j]]
    [[-0.05750374-0.99834529j  0.        +0.j        ]
     [ 0.        +0.j         -0.05750374+0.99834529j]]
    [[-0.41946905-0.90776964j  0.        +0.j        ]
     [ 0.        +0.j         -0.41946905+0.90776964j]]
    [[-0.70560894-0.70860146j  0.        +0.j        ]
     [ 0.        +0.j         -0.70560894+0.70860146j]]
    [[-0.89623153-0.44358656j  0.        +0.j        ]
     [ 0.        +0.j         -0.89623153+0.44358656j]]
    [[-0.9884577-0.1514971j  0.       +0.j       ]
     [ 0.       +0.j        -0.9884577+0.1514971j]]
    [[-0.99064733+0.13644734j  0.        +0.j        ]
     [ 0.        +0.j         -0.99064733-0.13644734j]]
    [[-0.91755676+0.39760481j  0.        +0.j        ]
     [ 0.        +0.j         -0.91755676-0.39760481j]]
    [[-0.78663643+0.61741649j  0.        +0.j        ]
     [ 0.        +0.j         -0.78663643-0.61741649j]]
    [[-0.61547563+0.78815592j  0.        +0.j        ]
     [ 0.        +0.j         -0.61547563-0.78815592j]]
    [[-0.42021862+0.9074229j  0.        +0.j       ]
     [ 0.        +0.j        -0.42021862-0.9074229j]]
    [[-0.21472138+0.97667534j  0.        +0.j        ]
     [ 0.        +0.j         -0.21472138-0.97667534j]]
    [[-0.01022857+0.99994769j  0.        +0.j        ]
     [ 0.        +0.j         -0.01022857-0.99994769j]]
    [[0.18461179+0.98281152j 0.        +0.j        ]
     [0.        +0.j         0.18461179-0.98281152j]]
    [[0.36353457+0.93158071j 0.        +0.j        ]
     [0.        +0.j         0.36353457-0.93158071j]]
    [[0.52234311+0.85273541j 0.        +0.j        ]
     [0.        +0.j         0.52234311-0.85273541j]]
    [[0.6585598+0.7525284j 0.       +0.j       ]
     [0.       +0.j        0.6585598-0.7525284j]]
    [[0.77108246+0.6367353j 0.        +0.j       ]
     [0.        +0.j        0.77108246-0.6367353j]]
    [[0.85986971+0.51051355j 0.        +0.j        ]
     [0.        +0.j         0.85986971-0.51051355j]]
    [[0.92566666+0.3783401j 0.        +0.j       ]
     [0.        +0.j        0.92566666-0.3783401j]]
    [[0.96977425+0.24400389j 0.        +0.j        ]
     [0.        +0.j         0.96977425-0.24400389j]]
    [[0.9938612+0.11063417j 0.       +0.j        ]
     [0.       +0.j         0.9938612-0.11063417j]]
    [[0.99981472-0.0192491j 0.        +0.j       ]
     [0.        +0.j        0.99981472+0.0192491j]]
    [[0.9896252-0.14367314j 0.       +0.j        ]
     [0.       +0.j         0.9896252+0.14367314j]]
    [[0.96529946-0.26114546j 0.        +0.j        ]
     [0.        +0.j         0.96529946+0.26114546j]]
    [[0.92879771-0.37058712j 0.        +0.j        ]
     [0.        +0.j         0.92879771+0.37058712j]]
    [[0.88198941-0.47126922j 0.        +0.j        ]
     [0.        +0.j         0.88198941+0.47126922j]]
    [[0.82662426-0.56275424j 0.        +0.j        ]
     [0.        +0.j         0.82662426+0.56275424j]]
    [[0.76431473-0.64484338j 0.        +0.j        ]
     [0.        +0.j         0.76431473+0.64484338j]]
    [[0.69652747-0.71753012j 0.        +0.j        ]
     [0.        +0.j         0.69652747+0.71753012j]]
    [[0.62458118-0.78095989j 0.        +0.j        ]
     [0.        +0.j         0.62458118+0.78095989j]]
    [[0.54964923-0.83539555j 0.        +0.j        ]
     [0.        +0.j         0.54964923+0.83539555j]]
    [[0.47276556-0.88118825j 0.        +0.j        ]
     [0.        +0.j         0.47276556+0.88118825j]]
    [[0.39483266-0.91875305j 0.        +0.j        ]
     [0.        +0.j         0.39483266+0.91875305j]]
    [[0.31663098-0.9485488j 0.        +0.j       ]
     [0.        +0.j        0.31663098+0.9485488j]]
    [[0.23882886-0.97106168j 0.        +0.j        ]
     [0.        +0.j         0.23882886+0.97106168j]]
    [[0.16199277-0.98679194j 0.        +0.j        ]
     [0.        +0.j         0.16199277+0.98679194j]]
    [[0.0865974-0.99624339j 0.       +0.j        ]
     [0.       +0.j         0.0865974+0.99624339j]]
    [[0.01303531-0.99991504j 0.        +0.j        ]
     [0.        +0.j         0.01303531+0.99991504j]]
    [[-0.05837382-0.99829479j  0.        +0.j        ]
     [ 0.        +0.j         -0.05837382+0.99829479j]]
    [[-0.12737468-0.99185467j  0.        +0.j        ]
     [ 0.        +0.j         -0.12737468+0.99185467j]]
    [[-0.19376838-0.9810473j  0.        +0.j       ]
     [ 0.        +0.j        -0.19376838+0.9810473j]]
    [[-0.25740524-0.96630355j  0.        +0.j        ]
     [ 0.        +0.j         -0.25740524+0.96630355j]]
    [[-0.31817822-0.94803092j  0.        +0.j        ]
     [ 0.        +0.j         -0.31817822+0.94803092j]]
    [[-0.37601703-0.92661275j  0.        +0.j        ]
     [ 0.        +0.j         -0.37601703+0.92661275j]]
    [[-0.43088285-0.90240787j  0.        +0.j        ]
     [ 0.        +0.j         -0.43088285+0.90240787j]]
    [[-0.48276358-0.87575072j  0.        +0.j        ]
     [ 0.        +0.j         -0.48276358+0.87575072j]]
    [[-0.53166971-0.84695178j  0.        +0.j        ]
     [ 0.        +0.j         -0.53166971+0.84695178j]]
    [[-0.57763063-0.81629827j  0.        +0.j        ]
     [ 0.        +0.j         -0.57763063+0.81629827j]]
    [[-0.62069136-0.784055j  0.        +0.j      ]
     [ 0.        +0.j       -0.62069136+0.784055j]]
    [[-0.66090976-0.75046538j  0.        +0.j        ]
     [ 0.        +0.j         -0.66090976+0.75046538j]]
    [[-0.69835405-0.71575249j  0.        +0.j        ]
     [ 0.        +0.j         -0.69835405+0.71575249j]]
    [[-0.73310063-0.68012018j  0.        +0.j        ]
     [ 0.        +0.j         -0.73310063+0.68012018j]]
    [[-0.76523229-0.64375425j  0.        +0.j        ]
     [ 0.        +0.j         -0.76523229+0.64375425j]]
    [[-0.79483656-0.60682357j  0.        +0.j        ]
     [ 0.        +0.j         -0.79483656+0.60682357j]]
    [[-0.82200437-0.56948118j  0.        +0.j        ]
     [ 0.        +0.j         -0.82200437+0.56948118j]]
    [[-0.84682887-0.53186546j  0.        +0.j        ]
     [ 0.        +0.j         -0.84682887+0.53186546j]]
    [[-0.86940444-0.49410113j  0.        +0.j        ]
     [ 0.        +0.j         -0.86940444+0.49410113j]]
    [[-0.88982586-0.45630026j  0.        +0.j        ]
     [ 0.        +0.j         -0.88982586+0.45630026j]]
    [[-0.90818762-0.41856331j  0.        +0.j        ]
     [ 0.        +0.j         -0.90818762+0.41856331j]]
    [[-0.9245833-0.38097995j  0.       +0.j        ]
     [ 0.       +0.j         -0.9245833+0.38097995j]]
    [[-0.93910512-0.34362998j  0.        +0.j        ]
     [ 0.        +0.j         -0.93910512+0.34362998j]]
    [[-0.95184356-0.30658416j  0.        +0.j        ]
     [ 0.        +0.j         -0.95184356+0.30658416j]]
    [[-0.96288698-0.2699049j  0.        +0.j       ]
     [ 0.        +0.j        -0.96288698+0.2699049j]]
    [[-0.97232148-0.23364703j  0.        +0.j        ]
     [ 0.        +0.j         -0.97232148+0.23364703j]]
    [[-0.9802306-0.19785844j  0.       +0.j        ]
     [ 0.       +0.j         -0.9802306+0.19785844j]]
    [[-0.98669525-0.16258068j  0.        +0.j        ]
     [ 0.        +0.j         -0.98669525+0.16258068j]]
    [[-0.99179357-0.12784955j  0.        +0.j        ]
     [ 0.        +0.j         -0.99179357+0.12784955j]]
    [[-0.99560089-0.0936956j  0.        +0.j       ]
     [ 0.        +0.j        -0.99560089+0.0936956j]]
    [[-0.99818967-0.06014464j  0.        +0.j        ]
     [ 0.        +0.j         -0.99818967+0.06014464j]]
    [[-0.99962952-0.02721817j  0.        +0.j        ]
     [ 0.        +0.j         -0.99962952+0.02721817j]]
    [[-0.99998717+0.00506622j  0.        +0.j        ]
     [ 0.        +0.j         -0.99998717-0.00506622j]]
    [[-0.99932653+0.03669444j  0.        +0.j        ]
     [ 0.        +0.j         -0.99932653-0.03669444j]]
    [[-0.99770874+0.06765556j  0.        +0.j        ]
     [ 0.        +0.j         -0.99770874-0.06765556j]]
    [[-0.99519217+0.09794154j  0.        +0.j        ]
     [ 0.        +0.j         -0.99519217-0.09794154j]]
    [[-0.99183254+0.12754687j  0.        +0.j        ]
     [ 0.        +0.j         -0.99183254-0.12754687j]]
    [[-0.98768298+0.15646833j  0.        +0.j        ]
     [ 0.        +0.j         -0.98768298-0.15646833j]]
    [[-0.98279405+0.18470478j  0.        +0.j        ]
     [ 0.        +0.j         -0.98279405-0.18470478j]]
    [[-0.9772139+0.21225689j  0.       +0.j        ]
     [ 0.       +0.j         -0.9772139-0.21225689j]]
    [[-0.97098831+0.23912695j  0.        +0.j        ]
     [ 0.        +0.j         -0.97098831-0.23912695j]]
    [[-0.96416077+0.26531871j  0.        +0.j        ]
     [ 0.        +0.j         -0.96416077-0.26531871j]]
    [[-0.95677257+0.29083716j  0.        +0.j        ]
     [ 0.        +0.j         -0.95677257-0.29083716j]]
    [[-0.9488629+0.31568845j  0.       +0.j        ]
     [ 0.       +0.j         -0.9488629-0.31568845j]]
    [[-0.94046893+0.33987969j  0.        +0.j        ]
     [ 0.        +0.j         -0.94046893-0.33987969j]]
    [[-0.93162587+0.36341882j  0.        +0.j        ]
     [ 0.        +0.j         -0.93162587-0.36341882j]]
    [[-0.9223671+0.38631456j  0.       +0.j        ]
     [ 0.       +0.j         -0.9223671-0.38631456j]]
    [[-0.9127242+0.40857623j  0.       +0.j        ]
     [ 0.       +0.j         -0.9127242-0.40857623j]]

```python
plt.figure();
plt.plot(ref);
plt.plot(trans)
```

```python
propagation_phase
```

    array([[-0.97466986+0.22364851j,  0.        +0.j        ],
           [ 0.        +0.j        , -0.97466986-0.22364851j]])

```python
print(eigenvals)
```

    [ 3.46410162-0.j -3.46410162+0.j]

```python
print(T)
```

## Check with Analytic Form

The analytic form of TMM appears to be analytic and stable (for TE and TM, but why is that?)
This analytic form is also derived very simply by matching coefficients it appears. There is no matrix inversion or analysis of modes or anything in the analytic forms...

```python
er = 12;
mur = 1;

def transfer_matrix(k,L, er, mu): ## this is the TE polarization
    Zc = 1/k; #np.sqrt(mu/er);
    M = np.array([[np.cos(k*L), 1*Zc*np.sin(k*L)],\
                  [(-1/Zc)*np.sin(k*L), np.cos(k*L)]], dtype = 'complex')
    return M;

def RT(M, kL, kR):
    '''
    kL: k of incident
    kR: k of transmission medium
    '''
    denominator = (-M[1,0]+kL*kR*M[0,1])+1j*(kL*M[1,1]+kR*M[0,0])
    r = ((M[1,0]+ kL*kR*M[0,1]) + 1j*(kL*M[1,1] - kR*M[0,0]))/denominator;

    t = 2j*kL*np.exp(-1j*kR*L)/denominator;
    return r,t;

L = 0.5; #layer thickness
wvlen_scan = np.linspace(0.5, 2, 301)
spectra = [];
for wvlen in wvlen_scan:
    k = 2*np.pi/wvlen;
    #in the medium
    kprime = np.sqrt(er)*k;
    M1 = transfer_matrix(kprime,L, er, mur)
    M2 = transfer_matrix(2*k, L, 4, 1);
    M = M1@M2;
    #assume a structure with air for reflection and transmission methods
    r, t = RT(M, k, k)

    R = abs(r)**2; T = abs(t)**2;
    spectra.append([R,T])
spectra = np.array(spectra)

plt.plot(wvlen_scan, spectra)
plt.show()

```

## Check with the Impedance Transformation
FROM EE236, the impedance transform starts at the transmission layer and then moves back up to the reflection region

```python
def impedanceTransformer(Z_0, Z0, kz, z):
    return Z_0 * (Z0 - 1j*Z_0*np.tan(kz*z))/(Z_0 - 1j*Z0*np.tan(kz*z));

```

```python
## single layer
Z_0 = np.sqrt(mur/er);
Z0 = 1;

spectra = [];
L= 0.5;
for wvlen in wvlen_scan:
    kz = 2*np.pi/wvlen;
    Gamma_ref = (impedanceTransformer(Z_0, Z0, kz, L) - Z0)/(impedanceTransformer(Z_0, Z0, kz, L) + Z0);
    spectra.append(abs(Gamma_ref)**2);

spectra = np.array(spectra);
plt.plot(wvlen_scan, spectra);
plt.plot(wvlen_scan, 1-spectra);
plt.show();
```

```python

```


## Contents

* amortized variational inference (VAEs)
* improving amortized variational inference using KL annleaing
* improving amortized variational inference using constraint optimization
* avoiding latent space distribution matching using GANs

```python
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math

import tensorflow as tf
import tensorflow_datasets as tfds

import numpy as np

# Plotting library.
from matplotlib import pyplot as plt
import seaborn as sns

# Tensorflow probability utilities
import tensorflow_probability as tfp

tfd = tfp.distributions

import warnings

```

```python
sns.set(rc={"lines.linewidth": 2.8}, font_scale=2)
sns.set_style("whitegrid")
```

```python
# Don't forget to select GPU runtime environment in Runtime -> Change runtime type
device_name = tf.test.gpu_device_name()
if device_name != '/device:GPU:0':
  raise SystemError('GPU device not found')
print('Found GPU at: {}'.format(device_name))
```

    Found GPU at: /device:GPU:0

```python
def gallery(array, ncols=10, rescale=False):
    """Data visualization code."""
    if rescale:
      array = (array + 1.) / 2
    nindex, height, width, intensity = array.shape
    nrows = nindex//ncols
    assert nindex == nrows*ncols
    # want result.shape = (height*nrows, width*ncols, intensity)
    result = (array.reshape(nrows, ncols, height, width, intensity)
              .swapaxes(1,2)
              .reshape(height*nrows, width*ncols, intensity))
    return result
```

```python
def show_digits(axis, digits, title=''):
  axis.axis('off')
  ncols = int(np.sqrt(digits.shape[0]))
  axis.imshow(gallery(digits, ncols=ncols).squeeze(axis=2),
                 cmap='gray')
  axis.set_title(title, fontsize=15)
```

```python
def show_latent_interpolations(generator, prior, session):
  a = np.linspace(0.0, 1.0, BATCH_SIZE)
  a = np.expand_dims(a, axis=1)

  first_latents = prior.sample()[0]
  second_latents = prior.sample()[0]

  # To ensure that the interpolation is still likely under the Gaussian prior,
  # we use Gaussian interpolation - rather than linear interpolation.
  interpolations =  np.sqrt(a) * first_latents + np.sqrt(1 - a) * second_latents

  ncols = int(np.sqrt(BATCH_SIZE))
  samples_from_interpolations = generator(interpolations)
  samples_from_interpolations_np = sess.run(samples_from_interpolations)
  plt.gray()
  axis = plt.gca()
  show_digits(
     axis, samples_from_interpolations_np, title='Latent space interpolations')
```

```python
BATCH_SIZE = 64
NUM_LATENTS = 10
TRAINING_STEPS = 10000
NUM_SVI_UPDATES = 100
```

```python
DECODER_VARIABLE_SCOPE = "decoder"
```

```python
tf.reset_default_graph()
```

## Get the Data

We will use the MNIST dataset. Luckly, TensorFlow comes with a simple way to load it.

```python
warnings.filterwarnings('ignore')
mnist = tf.contrib.learn.datasets.load_dataset("mnist")
print(mnist.train.images.shape)
print(type(mnist.train.images))
```

    Extracting MNIST-data/train-images-idx3-ubyte.gz
    Extracting MNIST-data/train-labels-idx1-ubyte.gz
    Extracting MNIST-data/t10k-images-idx3-ubyte.gz
    Extracting MNIST-data/t10k-labels-idx1-ubyte.gz
    (55000, 784)
    <class 'numpy.ndarray'>

```python
print(mnist.train.images.shape)
print(type(mnist.train.images))
```

    (55000, 784)
    <class 'numpy.ndarray'>

### Transform the Data from Numpy Arrays to in Graph Tensors.

This allows us to use TensorFlow datasets, which ensure that a new batch from the data is being fed at each session.run. This means that we do not need to use feed_dicts to feed data to each session.

```python
def make_tf_data_batch(np_data, shuffle=True):
  # Reshape the data to image size.
  images = np_data.reshape((-1, 28, 28, 1))

  # Create the TF dataset.
  dataset = tf.data.Dataset.from_tensor_slices(images)

  # Shuffle and repeat the dataset for training.
  # This is required because we want to do multiple passes through the entire
  # dataset when training.
  if shuffle:
    dataset = dataset.shuffle(100000).repeat()

  # Batch the data and return the data batch.
  batch_data = dataset.batch(BATCH_SIZE)
  one_shot_iterator = tf.compat.v1.data.make_one_shot_iterator(batch_data)
  data_batch = one_shot_iterator.get_next()

  return data_batch
```

```python
real_data = make_tf_data_batch(mnist.train.images)
print(real_data.shape)
```

    (?, 28, 28, 1)

# Amortized variational inference (VAEs)

Instead of learning one set of posterior variables per data point, we can use function approximation to learn the distributional variables. Specifically, the posterior parameters for $x_i$ will be the output of a *learned* function $f_\theta(x_i)$, where $\theta$ are parameters shared across all data points. Can you think of why this is useful?

Objective - maximize:
\begin{equation}
 \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x) \log p_\theta(x|z)}  - \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z))
\end{equation}

For more information, see:
  * https://arxiv.org/abs/1312.6114

```python
def standard_decoder(z):
  with tf.variable_scope(DECODER_VARIABLE_SCOPE, reuse=tf.AUTO_REUSE):
    h = tf.layers.dense(z, 7 * 7 * 64, activation=tf.nn.relu)
    h = tf.reshape(h, shape=[BATCH_SIZE, 7, 7, 64])
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=1,
          kernel_size=5,
          strides=2,
          activation=None,  # Do not activate the last layer.
          padding='same')(h)
    return tf.distributions.Bernoulli(h)
```

```python
def big_decoder(z):
  with tf.variable_scope(DECODER_VARIABLE_SCOPE, reuse=tf.AUTO_REUSE):
    h = tf.layers.dense(z, 4000, activation=tf.nn.relu)
    h = tf.layers.dense(h, 4000, activation=tf.nn.relu)
    h = tf.layers.dense(h, 4000, activation=tf.nn.relu)
    h = tf.layers.dense(h, 7 * 7 * 64, activation=tf.nn.relu)
    h = tf.reshape(h, shape=[BATCH_SIZE, 7, 7, 64])
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=1,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=1,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=1,
          kernel_size=5,
          strides=2,
          activation=None,  # Do not activate the last layer.
          padding='same')(h)
    return tf.distributions.Bernoulli(h)
```

```python
ENCODER_VARIABLE_SCOPE = 'encoder'
```

```python
def encoder(x):
  with tf.variable_scope(ENCODER_VARIABLE_SCOPE, reuse=tf.AUTO_REUSE):

    h = tf.layers.Conv2D(
          filters=8,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(x)
    h = tf.layers.Conv2D(
          filters=16,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2D(
          filters=32,
          kernel_size=5,
          strides=1,
          activation=tf.nn.relu,
          padding='same')(h)

    out_shape = 1
    for s in h.shape.as_list()[1:]:
      out_shape*= s

    h = tf.reshape(h, shape=[BATCH_SIZE, out_shape])
    mean = tf.layers.dense(h, NUM_LATENTS, activation=None)
    scale = tf.layers.dense(h, NUM_LATENTS, activation=None)
    return multi_normal(loc=mean, log_scale=scale)
```

```python
def multi_normal(loc, log_scale):
  # We model the latent variables as independent
  return tfd.Independent(
      distribution=tfd.Normal(loc=loc, scale=tf.exp(log_scale)),
      reinterpreted_batch_ndims=1)
```

```python
def make_prior():
  # Zero mean, unit variance prior.
  prior_mean = tf.zeros(shape=(BATCH_SIZE, NUM_LATENTS), dtype=tf.float32)
  prior_log_scale = tf.zeros(shape=(BATCH_SIZE, NUM_LATENTS), dtype=tf.float32)

  return multi_normal(prior_mean, prior_log_scale)
```

```python
def bound_terms(data_batch, variational_posterior, decoder_fn):

  # !!COMPLETE>> Draw a sample from the variational posterior
  #
  # latent_sample = ??

  # !!COMPLETE>> Compute the logprob from the decoder model evaluated at latent_sample
  #
  # all_log_probs = ??

  # Since we used broadcasting, log_prob of a Bernoulli will return a log_prob for
  # each dimension.
  all_log_probs.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

  # Reduce sum over the data dimensions
  likelihood_term = tf.reduce_sum(all_log_probs, axis=[1, 2, 3])

  # Reduce mean over the batch dimensions
  likelihood_term = tf.reduce_mean(likelihood_term)

  # !!COMPLETE>> Compute the KL divergence term
  # kl_term = ??

  # Reduce over the batch dimension.
  kl_term.shape.assert_is_compatible_with([BATCH_SIZE])
  kl_term = tf.reduce_mean(kl_term)

  return likelihood_term, kl_term
```

```python
real_data = make_tf_data_batch(mnist.train.images)
print(real_data.shape)
```

    (?, 28, 28, 1)

```python
prior = make_prior()
```

```python
# Note: we no longer have to use a variable to store the data.
# We will perform one encoder update per decoder update, so it is OK for the
# data batch to be refreshed at each run.
variational_posterior = encoder(real_data)
decoder = standard_decoder
```

```python
# !!COMPLETE>> Compute the variational bound
# train_elbo = ??

loss =  - train_elbo
```

```python
samples = decoder(prior.sample()).mean()
samples.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

reconstructions = decoder(variational_posterior.sample()).mean()
```

```python
# We now perform joint optimization on the encoder and decoder variables.
optimizer = tf.train.AdamOptimizer(0.001, beta1=0.9, beta2=0.9)
decoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=DECODER_VARIABLE_SCOPE)
encoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=ENCODER_VARIABLE_SCOPE)
autoencoder_vars = encoder_vars + decoder_vars
update_op = optimizer.minimize(loss, var_list=autoencoder_vars)
```

## Training

```python
sess = tf.Session()

# Initialize all variables
# sess.run(tf.initialize_all_variables())
sess.run(tf.global_variables_initializer())
```

```python
# %hide_pyerr  # - uncomment to interrupt training without a stacktrace
losses = []
kls = []
likelihood_terms = []

for i in range(TRAINING_STEPS):
  sess.run(update_op)

  if i % 100 == 0:
    iteration_loss, iteration_likelihood, iteration_kl  = sess.run(
          [loss, likelihood_term, kl_term])
    print('Iteration {}. Loss {}. KL {}'.format(
          i, iteration_loss, iteration_kl))
    losses.append(iteration_loss)
    kls.append(iteration_kl)
    likelihood_terms.append(iteration_likelihood)
```

    Iteration 0. Loss 539.196533203125. KL 0.02204948477447033
    Iteration 100. Loss 198.71803283691406. KL 4.2688117027282715
    Iteration 200. Loss 177.23019409179688. KL 10.226387023925781
    Iteration 300. Loss 132.45806884765625. KL 14.541561126708984
    Iteration 400. Loss 125.61270141601562. KL 16.109647750854492
    Iteration 500. Loss 116.99410247802734. KL 16.198226928710938
    Iteration 600. Loss 117.24712371826172. KL 18.95182228088379
    Iteration 700. Loss 134.2333984375. KL 17.74468421936035
    Iteration 800. Loss 120.09989166259766. KL 17.780799865722656
    Iteration 900. Loss 112.07304382324219. KL 17.73116683959961
    Iteration 1000. Loss 115.76332092285156. KL 18.482267379760742
    Iteration 1100. Loss 110.12649536132812. KL 17.3060302734375
    Iteration 1200. Loss 112.2479476928711. KL 17.51970672607422
    Iteration 1300. Loss 111.29019927978516. KL 18.701194763183594
    Iteration 1400. Loss 116.5347900390625. KL 17.516788482666016
    Iteration 1500. Loss 114.019287109375. KL 18.3753719329834
    Iteration 1600. Loss 111.26945495605469. KL 18.403675079345703
    Iteration 1700. Loss 106.58517456054688. KL 16.488300323486328
    Iteration 1800. Loss 112.21721649169922. KL 18.761083602905273
    Iteration 1900. Loss 111.12541198730469. KL 18.653030395507812
    Iteration 2000. Loss 106.52731323242188. KL 19.023784637451172
    Iteration 2100. Loss 104.4305419921875. KL 17.932323455810547
    Iteration 2200. Loss 109.64266204833984. KL 18.720788955688477
    Iteration 2300. Loss 114.87144470214844. KL 19.40646743774414
    Iteration 2400. Loss 106.15460205078125. KL 18.692188262939453
    Iteration 2500. Loss 112.76716613769531. KL 19.469356536865234
    Iteration 2600. Loss 103.92877197265625. KL 18.354677200317383
    Iteration 2700. Loss 111.34396362304688. KL 18.478622436523438
    Iteration 2800. Loss 108.03636169433594. KL 17.9420166015625
    Iteration 2900. Loss 99.22251892089844. KL 18.88661003112793
    Iteration 3000. Loss 108.01915740966797. KL 19.20111846923828
    Iteration 3100. Loss 105.20238494873047. KL 19.60022735595703
    Iteration 3200. Loss 105.30938720703125. KL 17.72589111328125
    Iteration 3300. Loss 97.84925079345703. KL 17.209144592285156
    Iteration 3400. Loss 108.55044555664062. KL 18.621307373046875
    Iteration 3500. Loss 106.50871276855469. KL 19.582073211669922
    Iteration 3600. Loss 104.05218505859375. KL 18.530948638916016
    Iteration 3700. Loss 107.37323760986328. KL 18.88671112060547
    Iteration 3800. Loss 107.99214935302734. KL 18.779151916503906
    Iteration 3900. Loss 110.96134948730469. KL 19.390949249267578
    Iteration 4000. Loss 96.7219009399414. KL 19.22564125061035
    Iteration 4100. Loss 105.692138671875. KL 18.31119155883789
    Iteration 4200. Loss 106.78297424316406. KL 18.4332275390625
    Iteration 4300. Loss 103.08122253417969. KL 18.773365020751953
    Iteration 4400. Loss 103.89761352539062. KL 19.114713668823242
    Iteration 4500. Loss 101.34099578857422. KL 18.082468032836914
    Iteration 4600. Loss 106.24536895751953. KL 18.91722869873047
    Iteration 4700. Loss 105.0051498413086. KL 19.090538024902344
    Iteration 4800. Loss 111.025634765625. KL 18.013797760009766
    Iteration 4900. Loss 98.26335906982422. KL 18.462234497070312
    Iteration 5000. Loss 99.50084686279297. KL 17.637489318847656
    Iteration 5100. Loss 110.61246490478516. KL 19.32965087890625
    Iteration 5200. Loss 103.47390747070312. KL 19.181732177734375
    Iteration 5300. Loss 105.43096923828125. KL 19.317718505859375
    Iteration 5400. Loss 110.90916442871094. KL 19.179569244384766
    Iteration 5500. Loss 104.08894348144531. KL 19.338199615478516
    Iteration 5600. Loss 112.50633239746094. KL 19.25715446472168
    Iteration 5700. Loss 101.44236755371094. KL 18.742740631103516
    Iteration 5800. Loss 107.12516021728516. KL 18.28661346435547
    Iteration 5900. Loss 104.33573150634766. KL 19.145286560058594
    Iteration 6000. Loss 102.23605346679688. KL 18.68785858154297
    Iteration 6100. Loss 107.91573333740234. KL 20.11248016357422
    Iteration 6200. Loss 101.89085388183594. KL 19.257585525512695
    Iteration 6300. Loss 107.33724212646484. KL 19.3143310546875
    Iteration 6400. Loss 105.61495208740234. KL 19.286752700805664
    Iteration 6500. Loss 102.5097427368164. KL 19.22541046142578
    Iteration 6600. Loss 102.29480743408203. KL 18.95319366455078
    Iteration 6700. Loss 100.42646789550781. KL 19.44255256652832
    Iteration 6800. Loss 102.34622192382812. KL 19.113353729248047
    Iteration 6900. Loss 104.66024780273438. KL 19.521411895751953
    Iteration 7000. Loss 103.05218505859375. KL 19.736644744873047
    Iteration 7100. Loss 103.5003890991211. KL 19.10436248779297
    Iteration 7200. Loss 109.19738006591797. KL 19.215492248535156
    Iteration 7300. Loss 100.52944946289062. KL 18.11672592163086
    Iteration 7400. Loss 103.93933868408203. KL 18.80729103088379
    Iteration 7500. Loss 100.49058532714844. KL 19.344200134277344
    Iteration 7600. Loss 104.36833190917969. KL 19.871599197387695
    Iteration 7700. Loss 103.64473724365234. KL 18.923250198364258
    Iteration 7800. Loss 104.57835388183594. KL 19.684059143066406
    Iteration 7900. Loss 102.67532348632812. KL 19.903072357177734
    Iteration 8000. Loss 101.18742370605469. KL 18.677898406982422
    Iteration 8100. Loss 106.8114013671875. KL 18.89575958251953
    Iteration 8200. Loss 105.57196044921875. KL 18.670217514038086
    Iteration 8300. Loss 103.58242797851562. KL 19.73116111755371
    Iteration 8400. Loss 96.03123474121094. KL 19.177391052246094
    Iteration 8500. Loss 94.8234634399414. KL 19.494247436523438
    Iteration 8600. Loss 99.21393585205078. KL 18.731712341308594
    Iteration 8700. Loss 105.25141906738281. KL 19.716747283935547
    Iteration 8800. Loss 99.65660095214844. KL 18.591737747192383
    Iteration 8900. Loss 107.36187744140625. KL 19.888652801513672
    Iteration 9000. Loss 98.2471923828125. KL 20.029090881347656
    Iteration 9100. Loss 102.62771606445312. KL 19.227155685424805
    Iteration 9200. Loss 99.278076171875. KL 18.250972747802734
    Iteration 9300. Loss 94.84703826904297. KL 17.79936981201172
    Iteration 9400. Loss 101.58551025390625. KL 19.67104721069336
    Iteration 9500. Loss 109.45127868652344. KL 19.814083099365234
    Iteration 9600. Loss 105.96746063232422. KL 19.2408447265625
    Iteration 9700. Loss 102.26542663574219. KL 19.273231506347656
    Iteration 9800. Loss 104.65763854980469. KL 19.263690948486328
    Iteration 9900. Loss 95.91481018066406. KL 18.94552230834961

### Visualize the loss in time

```python
fig, axes = plt.subplots(1, 3, figsize=(3*8,5))

axes[0].plot(losses, label='Negative ELBO')
axes[0].set_title('Time', fontsize=15)
axes[0].legend()

axes[1].plot(kls, label='KL')
axes[1].set_title('Time', fontsize=15)
axes[1].legend()

axes[2].plot(likelihood_terms, label='Likelihood Term')
axes[2].set_title('Time', fontsize=15)
axes[2].legend()
```

### Generate samples and latent interpolations

```python
real_data_vals, final_samples_vals, data_reconstructions_vals = sess.run(
      [real_data, samples, reconstructions])
```

```python
fig, axes = plt.subplots(1, 3, figsize=(3*4,4))

show_digits(axes[0], real_data_vals, 'Data')
show_digits(axes[1], data_reconstructions_vals, 'Reconstructions')
show_digits(axes[2], final_samples_vals, 'Samples')
```

```python
show_latent_interpolations(lambda x: decoder(x).mean(), prior, sess)
```

## Questions about amortized variational inference

* What do you observe about sample quality and reconstruction quality?
* What do you observe about the ELBO and KL term?
* Which gradient estimation method is used to compute the gradients with respect to the encoder parameters?

# KL annealing

Objective - maximize:
\begin{equation}
 \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x) \log p_\theta(x|z)}  - \alpha \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z))
\end{equation}

Where $\alpha$ changes during training, to weigh in the KL term more.

```python
tf.reset_default_graph()
```

```python
real_data = make_tf_data_batch(mnist.train.images)
prior = make_prior()
decoder = standard_decoder
```

```python
kl_coefficient = tf.Variable(
    tf.zeros(shape=(1), dtype=tf.float32),
    trainable=False,
    name='kl_coeff')

kl_step = 1.0 / TRAINING_STEPS
update_kl_coeff = tf.assign(kl_coefficient, kl_coefficient + kl_step)
```

```python
variational_posterior = encoder(real_data)

# !!COMPLETE>> Compute the variational bound with rescaled KL coefficient
# train_elbo = ??
#

loss =  - train_elbo
```

```python
samples = decoder(prior.sample()).mean()
samples.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

reconstructions = decoder(variational_posterior.sample()).mean()
```

```python
# We now perform joint optimization on the encoder and decoder variables.
optimizer = tf.train.AdamOptimizer(0.001, beta1=0.9, beta2=0.9)
decoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=DECODER_VARIABLE_SCOPE)
encoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=ENCODER_VARIABLE_SCOPE)
autoencoder_vars = encoder_vars + decoder_vars
variables_update_op = optimizer.minimize(loss, var_list=autoencoder_vars)

# Ensure that a variable update is followed by an update in the KL coefficient.
with tf.control_dependencies([variables_update_op]):
  update_op = tf.identity(update_kl_coeff)
```

## Training

```python
sess = tf.Session()

# Initialize all variables
sess.run(tf.initialize_all_variables())
```

```python
# %hide_pyerr  # - uncomment to interrupt training without a stacktrace
losses = []
kls = []
likelihood_terms = []

for i in range(TRAINING_STEPS):
  sess.run(update_op)

  if i % 100 == 0:
    iteration_loss, iteration_likelihood, iteration_kl  = sess.run(
          [loss, likelihood_term, kl_term])
    print('Iteration {}. Loss {}. KL {}'.format(
          i, iteration_loss, iteration_kl))
    losses.append(iteration_loss)
    kls.append(iteration_kl)
    likelihood_terms.append(iteration_likelihood)
```

    Iteration 0. Loss [539.4792]. KL 0.04104074090719223
    Iteration 100. Loss [194.9967]. KL 44.48965835571289
    Iteration 200. Loss [168.335]. KL 52.840965270996094
    Iteration 300. Loss [131.64783]. KL 62.42009735107422
    Iteration 400. Loss [117.0551]. KL 51.43534469604492
    Iteration 500. Loss [100.880455]. KL 47.36773681640625
    Iteration 600. Loss [98.15953]. KL 41.8781623840332
    Iteration 700. Loss [97.19895]. KL 41.61931610107422
    Iteration 800. Loss [89.82359]. KL 37.182716369628906
    Iteration 900. Loss [91.434395]. KL 37.03062057495117
    Iteration 1000. Loss [97.83667]. KL 35.20429992675781
    Iteration 1100. Loss [89.1202]. KL 33.190818786621094
    Iteration 1200. Loss [87.98577]. KL 33.383323669433594
    Iteration 1300. Loss [88.122856]. KL 33.20022964477539
    Iteration 1400. Loss [87.82192]. KL 30.42446517944336
    Iteration 1500. Loss [89.218185]. KL 31.81494903564453
    Iteration 1600. Loss [88.70567]. KL 31.473731994628906
    Iteration 1700. Loss [89.785675]. KL 31.230899810791016
    Iteration 1800. Loss [95.66893]. KL 30.48082160949707
    Iteration 1900. Loss [94.65039]. KL 30.083908081054688
    Iteration 2000. Loss [91.38539]. KL 29.352718353271484
    Iteration 2100. Loss [92.5235]. KL 28.619653701782227
    Iteration 2200. Loss [87.64369]. KL 28.069053649902344
    Iteration 2300. Loss [93.1667]. KL 28.68340301513672
    Iteration 2400. Loss [93.84164]. KL 28.539382934570312
    Iteration 2500. Loss [90.91083]. KL 28.491025924682617
    Iteration 2600. Loss [91.34407]. KL 27.924732208251953
    Iteration 2700. Loss [91.4208]. KL 27.621349334716797
    Iteration 2800. Loss [93.23716]. KL 27.289260864257812
    Iteration 2900. Loss [86.89399]. KL 26.907886505126953
    Iteration 3000. Loss [90.4547]. KL 26.34688377380371
    Iteration 3100. Loss [92.44128]. KL 27.99745750427246
    Iteration 3200. Loss [86.27027]. KL 26.08523941040039
    Iteration 3300. Loss [96.30439]. KL 26.65250587463379
    Iteration 3400. Loss [94.76277]. KL 25.680145263671875
    Iteration 3500. Loss [89.287125]. KL 26.495065689086914
    Iteration 3600. Loss [93.47805]. KL 25.70299530029297
    Iteration 3700. Loss [96.46541]. KL 26.557954788208008
    Iteration 3800. Loss [88.54142]. KL 25.336076736450195
    Iteration 3900. Loss [95.405624]. KL 26.10382080078125
    Iteration 4000. Loss [92.5049]. KL 24.51131820678711
    Iteration 4100. Loss [90.58173]. KL 25.58831787109375
    Iteration 4200. Loss [96.25146]. KL 25.72899627685547
    Iteration 4300. Loss [97.52775]. KL 25.512435913085938
    Iteration 4400. Loss [93.37403]. KL 24.511112213134766
    Iteration 4500. Loss [92.61507]. KL 24.15967559814453
    Iteration 4600. Loss [98.13136]. KL 24.7989444732666
    Iteration 4700. Loss [90.45614]. KL 24.81452751159668
    Iteration 4800. Loss [97.08648]. KL 23.383087158203125
    Iteration 4900. Loss [97.60376]. KL 24.295560836791992
    Iteration 5000. Loss [88.023735]. KL 24.104740142822266
    Iteration 5100. Loss [90.13156]. KL 24.34149169921875
    Iteration 5200. Loss [88.34305]. KL 23.70738410949707
    Iteration 5300. Loss [92.43582]. KL 23.96887969970703
    Iteration 5400. Loss [97.05117]. KL 24.261497497558594
    Iteration 5500. Loss [89.448296]. KL 23.773601531982422
    Iteration 5600. Loss [91.999756]. KL 23.84142303466797
    Iteration 5700. Loss [102.921165]. KL 23.941131591796875
    Iteration 5800. Loss [99.04912]. KL 23.20785903930664
    Iteration 5900. Loss [100.200645]. KL 24.855825424194336
    Iteration 6000. Loss [89.12278]. KL 23.080841064453125
    Iteration 6100. Loss [94.78172]. KL 22.66788101196289
    Iteration 6200. Loss [98.34577]. KL 22.407012939453125
    Iteration 6300. Loss [93.268585]. KL 22.565532684326172
    Iteration 6400. Loss [94.34651]. KL 22.229541778564453
    Iteration 6500. Loss [93.3301]. KL 22.404951095581055
    Iteration 6600. Loss [92.20555]. KL 22.457979202270508
    Iteration 6700. Loss [93.61355]. KL 22.577659606933594
    Iteration 6800. Loss [92.897255]. KL 21.75212287902832
    Iteration 6900. Loss [96.808044]. KL 22.82540512084961
    Iteration 7000. Loss [95.47898]. KL 22.13536834716797
    Iteration 7100. Loss [99.639275]. KL 22.168386459350586
    Iteration 7200. Loss [91.112015]. KL 21.828296661376953
    Iteration 7300. Loss [96.96475]. KL 22.717498779296875
    Iteration 7400. Loss [103.48729]. KL 22.02381134033203
    Iteration 7500. Loss [100.304855]. KL 21.41842269897461
    Iteration 7600. Loss [93.50812]. KL 21.563304901123047
    Iteration 7700. Loss [100.20209]. KL 21.966285705566406
    Iteration 7800. Loss [94.06003]. KL 20.694252014160156
    Iteration 7900. Loss [99.726]. KL 21.32509994506836
    Iteration 8000. Loss [102.53663]. KL 21.641376495361328
    Iteration 8100. Loss [95.9506]. KL 20.099872589111328
    Iteration 8200. Loss [99.36583]. KL 22.242008209228516
    Iteration 8300. Loss [98.1301]. KL 20.949371337890625
    Iteration 8400. Loss [99.98479]. KL 20.578990936279297
    Iteration 8500. Loss [101.882706]. KL 21.06544303894043
    Iteration 8600. Loss [97.40351]. KL 20.493473052978516
    Iteration 8700. Loss [90.48268]. KL 20.477787017822266
    Iteration 8800. Loss [100.02028]. KL 20.4517879486084
    Iteration 8900. Loss [99.91766]. KL 20.730117797851562
    Iteration 9000. Loss [95.83577]. KL 19.801795959472656
    Iteration 9100. Loss [101.65898]. KL 19.877809524536133
    Iteration 9200. Loss [97.11158]. KL 20.092958450317383
    Iteration 9300. Loss [100.94263]. KL 19.882980346679688
    Iteration 9400. Loss [94.284546]. KL 19.719219207763672
    Iteration 9500. Loss [95.98587]. KL 19.380163192749023
    Iteration 9600. Loss [99.67559]. KL 19.859458923339844
    Iteration 9700. Loss [108.477135]. KL 20.849803924560547
    Iteration 9800. Loss [100.90942]. KL 20.312353134155273
    Iteration 9900. Loss [99.55182]. KL 19.708248138427734

```python
fig, axes = plt.subplots(1, 3, figsize=(3*8,5))

axes[0].plot(losses, label='Negative ELBO')
axes[0].set_title('Time', fontsize=15)
axes[0].legend()

axes[1].plot(kls, label='KL')
axes[1].set_title('Time', fontsize=15)
axes[1].legend()

axes[2].plot(likelihood_terms, label='Likelihood Term')
axes[2].set_title('Time', fontsize=15)
axes[2].legend()
```

```python
real_data_vals, final_samples_vals, data_reconstructions_vals = sess.run(
      [real_data, samples, reconstructions])
```

```python
fig, axes = plt.subplots(1, 3, figsize=(3*4,4))

show_digits(axes[0], real_data_vals, 'Data')
show_digits(axes[1], data_reconstructions_vals, 'Reconstructions')
show_digits(axes[2], final_samples_vals, 'Samples')
```

```python
show_latent_interpolations(lambda x: decoder(x).mean(), prior, sess)
```

## Questions about KL annealing

* What do you observe about the KL behaviour throughout training as opposed to amortized variational inference without any KL annealing?

# Constrained optimization

Instead of using KL annealing, constrained optimization can be used to automatically tuned the relative weight of the likelihood and kl terms. This removes the need to manually create an optimization schedule, which can be problem specific.

The objective now becomes:

\begin{equation}
  \text{minimize } \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z)) \text{ such that }  \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x) \log p_\theta(x|z)} > \alpha
\end{equation}

This can be solved using the use of Lagrange multipliers. The objective becomes:

\begin{equation}
  \text{minimize }  \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z)) + \lambda  (\mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\alpha - \log p_\theta(x|z)))
\end{equation}

The difference compared to the KL annealing is that:

   * $\lambda$ is a learned parameter - it will be learned using stochastic gradient descent, like the network parameters.  The difference is that the lagrangian has to solve a maximization problem. You can see this intuitively: the graadient with respect to $\lambda$ in the objective above is $\mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\alpha - \log p_\theta(x|z))$. If $ \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\alpha - \log p_\theta(x|z))> 0$, the constraint is not being satisfied, so the value of the lagrangian needs to increase. This will be done by doing gradient ascent, instead of gradient descent. Note that for $\lambda$ to be a valid lagranian in a minimization problem, it has to be positive.
   * The practicioner has to specify the hyperparameter $\alpha$, which determines the reoncstruction quality of the model.
   * the coefficient is in front of the likelihood term, not the KL term. This is mainly for convenience, as it is easier to specify the hyperparameter $\alpha$ for the likelihood (reconstruction loss).

For more assumptions made by this method, see the Karush–Kuhn–Tucker conditions.

For more information, see:
  * http://bayesiandeeplearning.org/2018/papers/33.pdf

```python
tf.reset_default_graph()
```

```python
real_data = make_tf_data_batch(mnist.train.images)
prior = make_prior()
decoder = standard_decoder
```

```python
# Unlike in the KL annealing case, we learn the coefficient.
# We initialize the lagrangian such that after the softplus the coefficient is
# approximately 1.
lagrangian_var = tf.Variable(
    0.55 * tf.ones(shape=(1), dtype=tf.float32),
    trainable=True,
    name='lagrangian')

# Ensure that the lagrangian is positive and has stable dynamics.
lagrangian = tf.nn.softplus(lagrangian_var)
```

```python
# How good do we want the reconstruction loss to be?
# We can look at previous runs to get an idea what a reasonable value would be.
reconstruction_target = - 90
```

```python
variational_posterior = encoder(real_data)
likelihood_term, kl_term = bound_terms(
    real_data, variational_posterior, decoder)

# Note: now the elbo is a different quanitity to what we optimize.
# !!COMPLETE>> Implement the variational bound
# train_elbo = ??

# !!COMPLETE>> Implement the lagrangian loss function that is optimised.
# loss =  ??
#

```

```python
samples = decoder(prior.sample()).mean()
samples.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

reconstructions = decoder(variational_posterior.sample()).mean()
```

```python
tf.trainable_variables()
```

    [<tf.Variable 'lagrangian:0' shape=(1,) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d/kernel:0' shape=(5, 5, 1, 8) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d/bias:0' shape=(8,) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d_1/kernel:0' shape=(5, 5, 8, 16) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d_1/bias:0' shape=(16,) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d_2/kernel:0' shape=(5, 5, 16, 32) dtype=float32_ref>,
     <tf.Variable 'encoder/conv2d_2/bias:0' shape=(32,) dtype=float32_ref>,
     <tf.Variable 'encoder/dense/kernel:0' shape=(1568, 10) dtype=float32_ref>,
     <tf.Variable 'encoder/dense/bias:0' shape=(10,) dtype=float32_ref>,
     <tf.Variable 'encoder/dense_1/kernel:0' shape=(1568, 10) dtype=float32_ref>,
     <tf.Variable 'encoder/dense_1/bias:0' shape=(10,) dtype=float32_ref>,
     <tf.Variable 'decoder/dense/kernel:0' shape=(10, 3136) dtype=float32_ref>,
     <tf.Variable 'decoder/dense/bias:0' shape=(3136,) dtype=float32_ref>,
     <tf.Variable 'decoder/conv2d_transpose/kernel:0' shape=(5, 5, 32, 64) dtype=float32_ref>,
     <tf.Variable 'decoder/conv2d_transpose/bias:0' shape=(32,) dtype=float32_ref>,
     <tf.Variable 'decoder/conv2d_transpose_1/kernel:0' shape=(5, 5, 1, 32) dtype=float32_ref>,
     <tf.Variable 'decoder/conv2d_transpose_1/bias:0' shape=(1,) dtype=float32_ref>]

```python
lagrangian_optimizer = tf.train.GradientDescentOptimizer(0.001)

# We now perform joint optimization on the encoder and decoder variables.
optimizer = tf.train.AdamOptimizer(0.001, beta1=0.9, beta2=0.9)
decoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=DECODER_VARIABLE_SCOPE)
encoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=ENCODER_VARIABLE_SCOPE)
autoencoder_vars = encoder_vars + decoder_vars
autoencoder_variables_update_op = optimizer.minimize(
      loss, var_list=autoencoder_vars)

# Ensure that a variable update is followed by an update to the Lagrangian.
with tf.control_dependencies([autoencoder_variables_update_op]):
  # Ensure that the lagrangian solves a maximization problem instead of a
  # minimization problem by changing the sign of the loss function.
  update_op = lagrangian_optimizer.minimize(- loss, var_list=[lagrangian_var])
```

## Training

```python
sess = tf.Session()

# Initialize all variables
sess.run(tf.initialize_all_variables())
```

```python
# %hide_pyerr  # - uncomment to interrupt training without a stacktrace
losses = []
kls = []
likelihood_terms = []
lagrangian_values = []

for i in range(TRAINING_STEPS):
  sess.run(update_op)

  if i % 100 == 0:
    iteration_loss, iteration_likelihood, iteration_kl, lag_val  = sess.run(
          [loss, likelihood_term, kl_term, lagrangian])
    print('Iteration {}. Loss {}. KL {}. Lagrangian {}'.format(
          i, iteration_loss, iteration_kl, lag_val))
    losses.append(iteration_loss)
    kls.append(iteration_kl)
    likelihood_terms.append(iteration_likelihood)
    lagrangian_values.append(lag_val)
```

    Iteration 0. Loss [537.7765]. KL 0.0048736026510596275. Lagrangian [1.1971195]
    Iteration 100. Loss [1897.5077]. KL 26.191951751708984. Lagrangian [18.056318]
    Iteration 200. Loss [1980.6298]. KL 41.9063720703125. Lagrangian [27.26814]
    Iteration 300. Loss [1119.1221]. KL 52.965667724609375. Lagrangian [31.972078]
    Iteration 400. Loss [581.9747]. KL 59.35569763183594. Lagrangian [34.11227]
    Iteration 500. Loss [362.915]. KL 54.883079528808594. Lagrangian [35.39999]
    Iteration 600. Loss [317.58914]. KL 50.1129150390625. Lagrangian [36.20063]
    Iteration 700. Loss [402.21625]. KL 52.07322692871094. Lagrangian [36.832962]
    Iteration 800. Loss [44.509346]. KL 52.564788818359375. Lagrangian [37.237885]
    Iteration 900. Loss [145.89885]. KL 50.428306579589844. Lagrangian [37.375645]
    Iteration 1000. Loss [252.88945]. KL 48.64805603027344. Lagrangian [37.489204]
    Iteration 1100. Loss [86.101105]. KL 49.36549377441406. Lagrangian [37.484894]
    Iteration 1200. Loss [-187.32822]. KL 48.465309143066406. Lagrangian [37.415928]
    Iteration 1300. Loss [276.67462]. KL 47.247196197509766. Lagrangian [37.183254]
    Iteration 1400. Loss [44.90869]. KL 48.37846755981445. Lagrangian [36.87293]
    Iteration 1500. Loss [-240.8709]. KL 47.56217956542969. Lagrangian [36.602734]
    Iteration 1600. Loss [-65.4957]. KL 45.485939025878906. Lagrangian [36.239353]
    Iteration 1700. Loss [-161.15503]. KL 47.329833984375. Lagrangian [35.786396]
    Iteration 1800. Loss [-35.120934]. KL 45.93567657470703. Lagrangian [35.233425]
    Iteration 1900. Loss [-127.0646]. KL 44.629615783691406. Lagrangian [34.69096]
    Iteration 2000. Loss [-123.93813]. KL 44.614219665527344. Lagrangian [34.06938]
    Iteration 2100. Loss [86.43119]. KL 42.25973892211914. Lagrangian [33.48723]
    Iteration 2200. Loss [-196.51482]. KL 44.49226379394531. Lagrangian [32.8319]
    Iteration 2300. Loss [-6.7258034]. KL 43.77875518798828. Lagrangian [32.23462]
    Iteration 2400. Loss [-35.24203]. KL 41.98087692260742. Lagrangian [31.480381]
    Iteration 2500. Loss [19.750977]. KL 41.99729919433594. Lagrangian [30.742231]
    Iteration 2600. Loss [6.487568]. KL 41.64616775512695. Lagrangian [29.94041]
    Iteration 2700. Loss [-94.59319]. KL 40.834129333496094. Lagrangian [29.157299]
    Iteration 2800. Loss [-189.4112]. KL 42.273353576660156. Lagrangian [28.291048]
    Iteration 2900. Loss [-188.75612]. KL 42.7293586730957. Lagrangian [27.508793]
    Iteration 3000. Loss [-29.464428]. KL 40.41386795043945. Lagrangian [26.63594]
    Iteration 3100. Loss [-252.8194]. KL 42.25995635986328. Lagrangian [25.777622]
    Iteration 3200. Loss [-99.08545]. KL 40.733768463134766. Lagrangian [24.927343]
    Iteration 3300. Loss [-136.38887]. KL 41.26321792602539. Lagrangian [24.02633]
    Iteration 3400. Loss [-171.0202]. KL 40.27965545654297. Lagrangian [23.092607]
    Iteration 3500. Loss [-160.03505]. KL 40.045406341552734. Lagrangian [22.180038]
    Iteration 3600. Loss [-285.6041]. KL 39.8399658203125. Lagrangian [21.194817]
    Iteration 3700. Loss [-157.98004]. KL 37.2536506652832. Lagrangian [20.18167]
    Iteration 3800. Loss [-99.951935]. KL 39.57304382324219. Lagrangian [19.211649]
    Iteration 3900. Loss [-133.08795]. KL 37.35369110107422. Lagrangian [18.226156]
    Iteration 4000. Loss [-178.85277]. KL 38.57893371582031. Lagrangian [17.189882]
    Iteration 4100. Loss [-39.641556]. KL 37.0748291015625. Lagrangian [16.26722]
    Iteration 4200. Loss [-79.26456]. KL 35.941978454589844. Lagrangian [15.233492]
    Iteration 4300. Loss [-156.79112]. KL 37.362064361572266. Lagrangian [14.184778]
    Iteration 4400. Loss [-71.75191]. KL 36.911827087402344. Lagrangian [13.099477]
    Iteration 4500. Loss [-91.16696]. KL 35.89760208129883. Lagrangian [12.021125]
    Iteration 4600. Loss [-97.58989]. KL 33.92449951171875. Lagrangian [10.975303]
    Iteration 4700. Loss [-80.60414]. KL 33.63560104370117. Lagrangian [9.959117]
    Iteration 4800. Loss [-67.91059]. KL 33.325218200683594. Lagrangian [8.903631]
    Iteration 4900. Loss [-65.18633]. KL 33.25416564941406. Lagrangian [7.808928]
    Iteration 5000. Loss [-15.535587]. KL 31.095558166503906. Lagrangian [6.787668]
    Iteration 5100. Loss [-24.973703]. KL 31.360139846801758. Lagrangian [5.7431364]
    Iteration 5200. Loss [-18.540977]. KL 30.122482299804688. Lagrangian [4.6958375]
    Iteration 5300. Loss [-7.657318]. KL 28.350543975830078. Lagrangian [3.7137]
    Iteration 5400. Loss [0.45495033]. KL 27.149545669555664. Lagrangian [2.808532]
    Iteration 5500. Loss [10.220221]. KL 26.92540740966797. Lagrangian [2.0341866]
    Iteration 5600. Loss [0.73215485]. KL 23.452598571777344. Lagrangian [1.4687269]
    Iteration 5700. Loss [12.473395]. KL 21.205371856689453. Lagrangian [1.1070313]
    Iteration 5800. Loss [16.350008]. KL 20.87267303466797. Lagrangian [0.87566775]
    Iteration 5900. Loss [12.372095]. KL 19.502777099609375. Lagrangian [0.7397758]
    Iteration 6000. Loss [17.3273]. KL 17.203771591186523. Lagrangian [0.6608266]
    Iteration 6100. Loss [18.498072]. KL 17.817684173583984. Lagrangian [0.6036913]
    Iteration 6200. Loss [13.487437]. KL 16.91201400756836. Lagrangian [0.5713629]
    Iteration 6300. Loss [16.41008]. KL 17.004230499267578. Lagrangian [0.5440622]
    Iteration 6400. Loss [14.306404]. KL 17.139183044433594. Lagrangian [0.5397322]
    Iteration 6500. Loss [14.950502]. KL 16.598512649536133. Lagrangian [0.5288986]
    Iteration 6600. Loss [18.11353]. KL 16.513784408569336. Lagrangian [0.54235137]
    Iteration 6700. Loss [16.32279]. KL 15.744285583496094. Lagrangian [0.5353231]
    Iteration 6800. Loss [20.099459]. KL 16.940044403076172. Lagrangian [0.54523987]
    Iteration 6900. Loss [16.931604]. KL 16.437591552734375. Lagrangian [0.5344547]
    Iteration 7000. Loss [15.184028]. KL 15.469598770141602. Lagrangian [0.5300546]
    Iteration 7100. Loss [13.839041]. KL 15.779718399047852. Lagrangian [0.51980776]
    Iteration 7200. Loss [16.66923]. KL 15.932589530944824. Lagrangian [0.53124344]
    Iteration 7300. Loss [12.075546]. KL 15.629206657409668. Lagrangian [0.5210471]
    Iteration 7400. Loss [14.645985]. KL 15.711103439331055. Lagrangian [0.5296076]
    Iteration 7500. Loss [13.683577]. KL 15.71375846862793. Lagrangian [0.52480024]
    Iteration 7600. Loss [16.807558]. KL 16.17851448059082. Lagrangian [0.52315646]
    Iteration 7700. Loss [17.495083]. KL 17.071975708007812. Lagrangian [0.51749647]
    Iteration 7800. Loss [16.16618]. KL 15.620613098144531. Lagrangian [0.5190346]
    Iteration 7900. Loss [12.866259]. KL 15.346925735473633. Lagrangian [0.52493703]
    Iteration 8000. Loss [15.973631]. KL 15.430500030517578. Lagrangian [0.5158491]
    Iteration 8100. Loss [16.417778]. KL 15.713497161865234. Lagrangian [0.51503956]
    Iteration 8200. Loss [17.234903]. KL 15.84148120880127. Lagrangian [0.51808584]
    Iteration 8300. Loss [15.351126]. KL 15.346195220947266. Lagrangian [0.5144993]
    Iteration 8400. Loss [16.169954]. KL 15.392377853393555. Lagrangian [0.50217485]
    Iteration 8500. Loss [19.457584]. KL 15.516424179077148. Lagrangian [0.50745887]
    Iteration 8600. Loss [18.475319]. KL 15.140362739562988. Lagrangian [0.4975889]
    Iteration 8700. Loss [16.657322]. KL 15.216124534606934. Lagrangian [0.49648476]
    Iteration 8800. Loss [14.188514]. KL 14.20589828491211. Lagrangian [0.5018922]
    Iteration 8900. Loss [14.499733]. KL 15.683168411254883. Lagrangian [0.5002055]
    Iteration 9000. Loss [15.6828785]. KL 15.228012084960938. Lagrangian [0.5045635]
    Iteration 9100. Loss [12.713363]. KL 15.692811965942383. Lagrangian [0.5091483]
    Iteration 9200. Loss [16.391518]. KL 15.2518949508667. Lagrangian [0.5008657]
    Iteration 9300. Loss [19.191256]. KL 16.065940856933594. Lagrangian [0.5002946]
    Iteration 9400. Loss [11.031858]. KL 14.29248046875. Lagrangian [0.49477336]
    Iteration 9500. Loss [13.649986]. KL 14.990755081176758. Lagrangian [0.48877287]
    Iteration 9600. Loss [15.478353]. KL 15.449694633483887. Lagrangian [0.4823132]
    Iteration 9700. Loss [16.540375]. KL 15.143045425415039. Lagrangian [0.48222673]
    Iteration 9800. Loss [16.783539]. KL 14.739253044128418. Lagrangian [0.48703077]
    Iteration 9900. Loss [13.969278]. KL 15.331631660461426. Lagrangian [0.49168152]

```python
fig, axes = plt.subplots(2, 2, figsize=(2*8, 2* 5))

axes[0, 0].plot(losses, label='Negative ELBO')
axes[0, 0].set_title('Time', fontsize=15)
axes[0, 0].legend()

axes[0, 1].plot(kls, label='KL')
axes[0, 1].set_title('Time', fontsize=15)
axes[0, 1].legend()

axes[1, 0].plot(likelihood_terms, label='Likelihood Term')
axes[1, 0].set_title('Time', fontsize=15)
axes[1, 0].legend()

axes[1, 1].plot(lagrangian_values, label='Lagrangian Values')
axes[1, 1].set_title('Time', fontsize=15)
axes[1, 1].legend()
```

### Generate samples and latent interpolations

```python
real_data_vals, final_samples_vals, data_reconstructions_vals = sess.run(
      [real_data, samples, reconstructions])
```

```python
fig, axes = plt.subplots(1, 3, figsize=(3*4,4))

show_digits(axes[0], real_data_vals, 'Data')
show_digits(axes[1], data_reconstructions_vals, 'Reconstructions')
show_digits(axes[2], final_samples_vals, 'Samples')
```

```python
show_latent_interpolations(lambda x: decoder(x).mean(), prior, sess)
```

## Discussion about constrained optimization

* What do you observe about the behaviour of the likelihood and KL term throughout training? How is it different than in Amortized Variational inference with and without KL annealing?
* What do you notice about the behaviour of the lagrangian during training?
* How would you choose the hyperparameters of the constrained optimization problem?

Machine Learning Summer School 2019 (London, UK)/tutorials/vi/MLSS2019_VAEs_exercise.ipynb


## GeostatsPy: Bootstrap for Subsurface Data Analytics in Python

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

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

### PGE 383 Exercise: Bootstrap for Subsurface Data Analytics in Python

Here's a simple workflow, demonstration of bootstrap for subsurface modeling workflows. This should help you get started with building subsurface models that integrate uncertainty in the sample statistics.

#### Bootstrap

Uncertainty in the sample statistics
* one source of uncertainty is the paucity of data.
* do 200 or even less wells provide a precise (and accurate estimate) of the mean? standard deviation? skew? P13?

Would it be useful to know the uncertainty in these statistics due to limited sampling?
* what is the impact of uncertainty in the mean porosity e.g. 20%+/-2%?

**Bootstrap** is a method to assess the uncertainty in a sample statistic by repeated random sampling with replacement.

Assumptions
* sufficient, representative sampling, identical, idependent samples

Limitations
1. assumes the samples are representative
2. assumes stationarity
3. only accounts for uncertainty due to too few samples, e.g. no uncertainty due to changes away from data
4. does not account for boundary of area of interest
5. assumes the samples are independent
6. does not account for other local information sources

The Bootstrap Approach (Efron, 1982)

Statistical resampling procedure to calculate uncertainty in a calculated statistic from the data itself.
* Does this work?  Prove it to yourself, for uncertainty in the mean solution is standard error:

\begin{equation}
\sigma^2_\overline{x} = \frac{\sigma^2_s}{n}
\end{equation}

Extremely powerful - could calculate uncertainty in any statistic!  e.g. P13, skew etc.
* Would not be possible access general uncertainty in any statistic without bootstrap.
* Advanced forms account for spatial information and sampling strategy (game theory and Journel’s spatial bootstrap (1993).

Steps:

1. assemble a sample set, must be representative, reasonable to assume independence between samples

2. optional: build a cumulative distribution function (CDF)
    * may account for declustering weights, tail extrapolation
    * could use analogous data to support

3. For $\ell = 1, \ldots, L$ realizations, do the following:

    * For $i = \alpha, \ldots, n$ data, do the following:

        * Draw a random sample with replacement from the sample set or Monte Carlo simulate from the CDF (if available).

6. Calculate a realization of the sammary statistic of interest from the $n$ samples, e.g. $m^\ell$, $\sigma^2_{\ell}$. Return to 3 for another realization.

7. Compile and summarize the $L$ realizations of the statistic of interest.

This is a very powerful method.  Let's try it out.

#### Objective

In the PGE 383: Stochastic Subsurface Modeling class I want to provide hands-on experience with building subsurface modeling workflows. Python provides an excellent vehicle to accomplish this. I have coded a package called GeostatsPy with GSLIB: Geostatistical Library (Deutsch and Journel, 1998) functionality that provides basic building blocks for building subsurface modeling workflows.

The objective is to remove the hurdles of subsurface modeling workflow construction by providing building blocks and sufficient examples. This is not a coding class per se, but we need the ability to 'script' workflows working with numerical methods.

#### Getting Started

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

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

You will need to copy the data file to your working directory.  They are available here:

* Tabular data - sample_data_biased.csv at https://git.io/fh0CW

There are exampled below with these functions. You can go here to see a list of the available functions, https://git.io/fh4eX, other example workflows and source code.

```python
import geostatspy.GSLIB as GSLIB          # GSLIB utilies, visualization and wrapper
import geostatspy.geostats as geostats    # GSLIB methods convert to Python
```

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

```python
import numpy as np                        # ndarrys for gridded data
import pandas as pd                       # DataFrames for tabular data
import os                                 # set working directory, run executables
import matplotlib.pyplot as plt           # for plotting
from scipy import stats                   # summary statistics
import math                               # trig etc.
import scipy.signal as signal             # kernel for moving window calculation
import random
```

#### Set the working directory

I always like to do this so I don't lose files and to simplify subsequent read and writes (avoid including the full address each time).

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

#### Loading Tabular Data

Here's the command to load our comma delimited data file in to a Pandas' DataFrame object.

```python
df = pd.read_csv('sample_data_biased.csv') # load our data table (wrong name!)
```

Let's drop some samples so that we increase the variations in bootstrap samples for our demonstration below.

```python
df = df.sample(frac = 0.9)                  # extract 50 random samples to reduce the size of the dataset
print('Using ' + str(len(df)) + ' number of samples')
```

    Using 260 number of samples

Visualizing the DataFrame would be useful and we already learned about these methods in this demo (https://git.io/fNgRW).

We can preview the DataFrame by printing a slice or by utilizing the 'head' DataFrame member function (with a nice and clean format, see below). With the slice we could look at any subset of the data table and with the head command, add parameter 'n=13' to see the first 13 rows of the dataset.

```python
print(df.iloc[0:5,:])                   # display first 4 samples in the table as a preview
df.head(n=13)                           # we could also use this command for a table preview
```

           X    Y  Facies  Porosity        Perm
    209  530  789       1  0.154871   51.991055
    146  670   49       0  0.118435   11.076648
    196  500  699       1  0.121085    9.033307
    57   420  874       1  0.205473  954.697202
    76   840  949       1  0.141307   18.635578

      X
      Y
      Facies
      Porosity
      Perm

      209
      530
      789
      1
      0.154871
      51.991055

      146
      670
      49
      0
      0.118435
      11.076648

      196
      500
      699
      1
      0.121085
      9.033307

      57
      420
      874
      1
      0.205473
      954.697202

      76
      840
      949
      1
      0.141307
      18.635578

      273
      820
      919
      1
      0.104364
      2.766127

      90
      430
      809
      1
      0.181352
      212.452297

      197
      490
      419
      1
      0.098454
      2.881800

      53
      410
      884
      1
      0.218254
      1923.286133

      256
      410
      669
      1
      0.118083
      22.382541

      101
      740
      379
      1
      0.132215
      29.031980

      181
      640
      669
      0
      0.105467
      1.389270

      252
      200
      359
      0
      0.077625
      0.540486

#### Summary Statistics for Tabular Data

The table includes X and Y coordinates (meters), Facies 1 and 0 (1 is sandstone and 0 interbedded sand and mudstone), Porosity (fraction), and permeability as Perm (mDarcy).

There are a lot of efficient methods to calculate summary statistics from tabular data in DataFrames. The describe command provides count, mean, minimum, maximum, and quartiles all in a nice data table. We use transpose just to flip the table so that features are on the rows and the statistics are on the columns.

```python
df.describe().transpose()
```

      count
      mean
      std
      min
      25%
      50%
      75%
      max

      X
      260.0
      474.769231
      255.099997
      0.000000
      300.000000
      430.000000
      670.000000
      990.000000

      Y
      260.0
      530.242308
      304.354191
      9.000000
      269.000000
      544.000000
      839.000000
      999.000000

      Facies
      260.0
      0.807692
      0.394874
      0.000000
      1.000000
      1.000000
      1.000000
      1.000000

      Porosity
      260.0
      0.134177
      0.038388
      0.058548
      0.106235
      0.124065
      0.153088
      0.228790

      Perm
      260.0
      215.012711
      576.830397
      0.075819
      3.416283
      12.826386
      62.273419
      5308.842566

#### Visualizing Tabular Data with Location Maps

It is natural to set the x and y coordinate and feature ranges manually. e.g. do you want your color bar to go from 0.05887 to 0.24230 exactly? Also, let's pick a color map for display. I heard that plasma is known to be friendly to the color blind as the color and intensity vary together (hope I got that right, it was an interesting Twitter conversation started by Matt Hall from Agile if I recall correctly). We will assume a study area of 0 to 1,000m in x and y and omit any data outside this area.

```python
xmin = 0.0; xmax = 1000.0               # range of x values
ymin = 0.0; ymax = 1000.0               # range of y values
pormin = 0.05; pormax = 0.25;           # range of porosity values
nx = 100; ny = 100; csize = 10.0
cmap = plt.cm.plasma                    # color map
```

Let's try out locmap. This is a reimplementation of GSLIB's locmap program that uses matplotlib. I hope you find it simpler than matplotlib, if you want to get more advanced and build custom plots lock at the source. If you improve it, send me the new code. Any help is appreciated. To see the parameters, just type the command name:

```python
GSLIB.locmap
```

    <function geostatspy.GSLIB.locmap(df, xcol, ycol, vcol, xmin, xmax, ymin, ymax, vmin, vmax, title, xlabel, ylabel, vlabel, cmap, fig_name)>

Now we can populate the plotting parameters and visualize the porosity data.

```python
plt.subplot(111)
GSLIB.locmap_st(df,'X','Y','Porosity',xmin,xmax,ymin,ymax,pormin,pormax,'Well Data - Porosity','X(m)','Y(m)','Porosity (fraction)',cmap)
plt.subplots_adjust(left=0.0, bottom=0.0, right=1.0, top=1.2, wspace=0.2, hspace=0.2)
plt.show()
```

#### Declustering

Let's calculate some declustering weights. There is a demonstration on declustering here https://git.io/fhgJl if you need more information.

```python
wts, cell_sizes, dmeans = geostats.declus(df,'X','Y','Porosity',iminmax = 1, noff= 10, ncell=100,cmin=10,cmax=2000)
df['Wts'] = wts                            # add weights to the sample data DataFrame
df.head()                                  # preview to check the sample data DataFrame

def weighted_avg_and_std(values, weights): # function to calculate weighted mean and st. dev., from Eric O Lebigot, stack overflow,
    average = np.average(values, weights=weights)
    variance = np.average((values-average)**2, weights=weights)
    return (average, math.sqrt(variance))

sample_avg, sample_stdev = weighted_avg_and_std(df['Porosity'],df['Wts'])
print('Declustered mean = ' + str(round(sample_avg,3)) + ' and declustered standard deviation = ' + str(round(sample_stdev,3)))
```

    There are 260 data with:
       mean of      0.13417733651923078
       min and max  0.058547873 and 0.228790002
       standard dev 0.03831442273384987
    Declustered mean = 0.121 and declustered standard deviation = 0.032

##### A Couple of Bootstrap Realizations

We will attempt boostrap by-hand and manually loop over $L$ realizations and draw $n$ samples to calculate the summary statistics of interest, mean and variance. The choice function from the random package simplifies sampling with replacement from a set of samples with weights.

This command returns a ndarray with k samples with replacment from the 'Porosity' column of our DataFrame (df) accounting for the data weights in column 'Wts'.
```p
samples1 = random.choices(df['Porosity'].values, weights=df['Wts'].values, cum_weights=None, k=len(df))
```

It is instructive to look at a couple of these realizations from the original declustered data set.

```python
samples1 = random.choices(df['Porosity'].values, weights=df['Wts'].values, cum_weights=None, k=len(df))
samples2 = random.choices(df['Porosity'].values, weights=df['Wts'].values, cum_weights=None, k=len(df))

print('Bootstrap means, realization 1 = ' + str(np.average(samples1)) + ' and realization 2 = ' + str(np.average(samples2)))

plt.subplot(131)
GSLIB.hist_st(df['Porosity'],pormin,pormax,False,False,20,df['Wts'],'Porosity (fraction)','Declustered Porosity')

plt.subplot(132)
GSLIB.hist_st(samples1,pormin,pormax,False,False,20,None,'Bootstrap Sample - Realizaton 1','Bootstrap Porosity 1')

plt.subplot(133)
GSLIB.hist_st(samples2,pormin,pormax,False,False,20,None,'Bootstrap Sample - Realizaton 2','Bootstrap Porosity 2')

plt.subplots_adjust(left=0.0, bottom=0.0, right=3.0, top=1.2, wspace=0.2, hspace=0.2)
plt.show()
```

Note that the bootstrap distributions vary quite a bit from the original.

#### Summarizations Over Bootstrap Realizations

Let's make a loop to conduct $L$ resamples and calculate the average and standard deviation for each  ($m^\ell$, $\sigma^2_{\ell}$,  for $\ell = 0,\dots,L-1$). We then summarization over these $L$ realizations.

I did not find any built-in, concise functions to accomplish this, i.e. with a single line of code, so we are going to do it by hand.

To understand this code there are just a couple of Python concepts that you need to add to your Python arsenal.

1. declaring arrays - NumPy has a lot of great array (ndarray) functionality. There are build in functions to make a ndarray of any length (and dimension). This includes 'zeros', 'ones' and 'rand', so when we use this code:

```p
mean = np.zeros(L); stdev = np.zeros(L)
```

    we're making arrays of length $L$ pre-populated with zeros.

2. For Loops - when we are using the command below, we are instructing the computer to loop over all the indented code below the command for $l = 0,1,2,\ldots,L-1$ times. For each loop the $l$ variable increments, so we can use this to save each result to a different index in the arrays mean and stdev. Note, Python arrays index starting at 0 and stop at the length - 1.

```p
for l in range(0, L):
```

    we are running each bootstrap resampled realization, calculating the average and standard deviation and storing them in the arrays that we already declared.

```python
L = 1000                                 # set the number of realizations
mean = np.zeros(L); stdev = np.zeros(L)    # declare arrays to hold the realizations of the statistics
for l in range(0, L):                      # loop over realizations
    samples = random.choices(df['Porosity'].values, weights=df['Wts'].values, cum_weights=None, k=len(df))
    mean[l] = np.average(samples)
    stdev[l] = np.std(samples)

plt.subplot(121)
GSLIB.hist_st(mean,0.11,0.15,False,False,50,None,'Average Porosity (fraction)','Bootstrap Uncertainty in Porosity Average')

plt.subplot(122)
GSLIB.hist_st(stdev,0.015,0.045,False,False,50,None,'Standard Deviation Porosity (fraction)','Bootstrap Uncertainty in Porosity Standard Deviation')

plt.subplots_adjust(left=0.0, bottom=0.0, right=3.0, top=1.2, wspace=0.2, hspace=0.2)
plt.show()

print('Summary Statistics for Bootstrap Porosity Mean Realizations:')
print(stats.describe(mean))
print('P10 ' + str(round(np.percentile(mean,10),3)) + ', P50 ' + str(round(np.percentile(mean,50),3)) + ', P90 ' + str(round(np.percentile(mean,90),3)))

print('\nSummary Statistics for Bootstrap Porosity Standard Deviation Realizations:')
print(stats.describe(stdev))
print('P10 ' + str(round(np.percentile(stdev,10),3)) + ', P50 ' + str(round(np.percentile(stdev,50),3)) + ', P90 ' + str(round(np.percentile(stdev,90),3)))

```

#### Comments

This was a basic demonstration of bootstrap. Much more could be done, you could replace the statistics, average and standard deviation with any other statistics, for example P90, kurtosis, P13 etc. I have other demonstrations on the basics of working with DataFrames, ndarrays, univariate statistics, plotting data, declustering, data transformations, trend modeling and many other workflows available at https://github.com/GeostatsGuy/PythonNumericalDemos and https://github.com/GeostatsGuy/GeostatsPy.

I hope this was helpful,

*Michael*

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

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


## Introduction

In this project we have choosen to analyze a simple two period economic model with two different social security systems, more specifically the pay-as-you-go system and the fully funded system.
First, we import the relevant packages:

```python
import numpy as np
import matplotlib.pyplot as plt  # So that we can do our graphical analysis
plt.style.use('ggplot') # Choice of plotstyle.
import ipywidgets as widgets # So that we can create our interactive figure
from scipy import optimize # To be able to optimize
import sympy as sm # What we will use when solving for steady state
from sympy import symbols # For when we have to define our variables
from sympy.plotting import plot # So that we can plot our Solow Diagram

# Autoreload modules when code is run
%load_ext autoreload
%autoreload 2

# Importing the local module (modelproject.py)
import modelproject
```

    The autoreload extension is already loaded. To reload it, use:
      %reload_ext autoreload

>**Note**: Most of the following model is taken from the MacroIII exam Winter 2019

In this model we have:
* $K_t$ is capital
* $L_t$ is labor (growing with a constant rate of $n>0$)

Identical competitive firms maximize their profits using the Cobb-Douglas technology:
* $Y_t = K^{\alpha}_tL^{1-\alpha}_t$ with $\alpha\in (0,1)$

So, we have
$$ F(K_{t},L_{t})=K_{t}^{\alpha}L_{t}^{1-\alpha}\Rightarrow f(\tilde{k}_{t})=\tilde{k}_{t}^{\alpha} $$

Assume full depreciation ($\delta = 0$)

With these assumptions, profit maximization leads to the following:
$$1+r_t = \alpha k^{\alpha-1}_t$$
$$w_t = (1-\alpha)k^{\alpha}_t$$

Utility for the young individuals born in period t is:
$$ U_t = lnc_{1t} + \frac{1}{1+\rho} lnc_{2t+1} $$  with $\rho >-1.$

To start, the government runs a pay-as-you-go (PAYG) social security system where the young pay the contribution $d_t$ that is the fraction $\tau \in (0,1)$ of their wages.

Finally, we have the following budget constraints:
$$ c_{1t} + s_t = (1-\tau)w_t $$
$$ c_{2t+1} = s_t (1+r_{t+1})+(1+n)d_{t+1} $$

Now, we will find the capital-per-worker in steady state.
Under the PAYG system, we can derive the Euler equation:
$$ c_{2t+1} = \frac{1+r_{t+1}}{1+\rho}c_{1t} $$
Replacing $c_{1t}$ and $c_{2t+1}$ with their budget constraints yields the savings equation:
$$ s_t=\frac{1}{2+\rho}w_t(1-\tau)-\tau \frac{1+\rho}{2+\rho} \frac{1+n}{1+r_{t+1}}w_{t+1} $$
From this it is straightforward (using $k_{t+1}=\frac{s_t}{1+n}$ and the equilibrium conditions) finding the capital accumulation equation:
$$ k_{t+1} = \frac{1}{1+\frac{1+\rho}{2+\rho}\cdot\frac{(1-\alpha)}{\alpha}\cdot\tau}\cdot\left[\frac{(1-\alpha)(1-\tau)}{(1+n)(2+\rho)}\cdot{k_t}^\alpha\right] $$

In the following, we will look at the Analytical solution and the numerical soluation of this specific model for choosen parametervalues.

## Analytical solution

We want to show that the capital-per-worker in steady state is:
$$ \bar k =\left[\frac{1}{1+\frac{1+\rho}{2+\rho}\frac{(1-\alpha)}{\alpha}\tau}\frac{(1-\alpha)}{(1+n)}\frac{(1-\tau)}{(2+\rho)}\right]^\frac{1}{1-\alpha} $$

```python
# We can use sympy to analytically find the steady state expression
# First, we define all the symbols:
k = sm.symbols('k')
alpha = sm.symbols('alpha')
rho = sm.symbols('rho')
tau = sm.symbols('tau')
n = sm.symbols('n')
```

```python
# Next, we can define the capital accumulation equation and ensure that it is printing it out correctly (which we see it is)
ss = sm.Eq(k,((((1-alpha)*(1-tau))/((1+n)*(2+rho)))*k**alpha)/(1+((1+rho)/(2+rho))*((1-alpha)/alpha)*tau))
ss
```

$\displaystyle k = \frac{k^{\alpha} \left(1 - \alpha\right) \left(1 - \tau\right)}{\left(1 + \frac{\tau \left(1 - \alpha\right) \left(\rho + 1\right)}{\alpha \left(\rho + 2\right)}\right) \left(n + 1\right) \left(\rho + 2\right)}$

```python
# We solve for the Steady State:
kss = sm.solve(ss,k)[0]
sm.simplify(kss)
```

$\displaystyle \left(\frac{- \alpha n \rho \tau + \alpha n \rho - \alpha n \tau + 2 \alpha n - \alpha \rho \tau + \alpha \rho - \alpha \tau + 2 \alpha + n \rho \tau + n \tau + \rho \tau + \tau}{\alpha \left(\alpha \tau - \alpha - \tau + 1\right)}\right)^{\frac{1}{\alpha - 1}}$

>**Note**: As you can see from the output above, we find an expression that is rather long and complicated. We have tried alot and spent alot of time on it, but we can't seem to make it work for us so that we get the exact equation. It works for us later on, when we find the other Steady State.

```python
# We turn this steady state solution for capital-per-worker into a Python function to use below in the numerical solution
ss_func = sm.lambdify((n,tau,rho,alpha),kss)
```

## Numerical Solution
In this section we will parametrizise and find the numerical solution.

We will now define our objective function. More specifically, re-writing our steady state for the capital-per-worker (moving $k^*$ to the other side), we can get a **root-finding problem** which can be solved with the bisection method.
$$ 0= \frac{1}{(2+\rho)(1+n)}(1-\tau)(1-\alpha)f(\tilde{k}_{t}) - \frac{1+\rho}{2+\rho}\tau \frac{1-\alpha}{\alpha}{k}^{\ast}-{k}^{\ast} $$

```python
# We define some model parameters:
n = 0.01
tau = 0.1
rho = 0.3
alpha = 1/3

# We solve numerically for the steady state (see code in the local module modelproject.py):
solution = modelproject.solve_for_ss(n,tau,rho,alpha)

# We print the two solutions:
print(f'Analytical solution is: {ss_func(n,tau,rho,alpha):.3f}')
print(f'Numerical solution is: {solution.root:.3f}')
```

    Analytical solution is: 0.112
    Numerical solution is: 0.112

From these results we see that the steady state level of capital-per-worker is 0.112 under the pay-as-you-go social security system under these circumstances and parametersizes.

## Further Analysis

```python
# We can visualize the steady state in a Solow Diagram for parametervalues n=0.01, tau=0.1, rho=0.3 and alpha=1/3 to see how the equations develops over time:
x = symbols('k_(t)')
p1 = plot((0.26*x**0.33-0.115*x),(x, 0, 0.2), show=False, line_color='b') # Displaying the capital accumulation locus
p2 = plot(x,(x, 0, 0.2), show=False, line_color='black') # Displaying the 45degree line
p1.append(p2[0]) # Combining the two plots in one diagram
p1.show()
```

In this diagram we can find the Steady State value where the 45degree line crosses the capital line, which we see is 0.112 like we found in our solutions above.

```python
# We will now visualize the steady state value for different values of n,tau,rho
# This is to display how a change in the three variables change the Steady State value
def interactive_figure(n,tau,rho):
    fig = plt.figure()
    ax = fig.add_subplot(1,1,1)
    ax.plot(ss_func(n,tau,rho,alpha),'o',label='SS PAYG') # Plot of SS under PAYG
    ax.legend(loc='lower right',facecolor='white',frameon=True) # Box with name of plot down right
    ax.set_ylabel('$Capital$') # Label of y-axis
    ax.set_xlim([-0.5,0.5]) # The length of the x-axis
    ax.set_ylim([-0.2,0.4]) # The length of the y-axis
    plt.suptitle('Capital in Steady State', size = 20) # Giving the diagram a title

# We create the sliders for the variables n, tau and rho:
widgets.interact(interactive_figure,
    n=widgets.FloatSlider(description="n", min=0, max=3, step=0.05,),
    tau=widgets.FloatSlider(description="tau", min=0, max=0.5, step=0.05,),
    rho=widgets.FloatSlider(description="rho", min=0, max=0.5, step=0.05,)
);
```

    interactive(children=(FloatSlider(value=0.0, description='n', max=3.0, step=0.05), FloatSlider(value=0.0, desc…

We see that for higher n,tau and rho respectively, we get a lower steady state capital-per-worker under the PAYG social security system.
The higher n means a larger population and this hollows out capita-per-worker as the capital has to be spread out between more people.
A higher tau lowers capital in two ways:
1. There is less income when young
2. You recieve income when old.
Hence, there is less incentive to save when young.
Finally, a rise in rho means a decrease in the patience of the consumer. When consumers are less patient they save less and consume more now. Hence, less capital is accumulated.

## An Extension/Change of Our Baseline Model
Now we suppose that the government at time T decides to switch to the fully funded social security system where the young pay a contribution $d_T$ which amounts to a fraction $\tau\in (0,1)$ of their wages.
These contributions are paid out the next period together with the interest rate. The new budget constraints are:
$$ c_{1t}+s_t = (1-\tau)w_t $$
$$ c_{2t+1} = (s_t+\tau w_t)(1+r_{t+1})  ,    t\ge T $$

In this case we need to find a new steady state capital-per-worker.

Following the same procedure as above, we get the following capital accumulation equation:
$$ k_{t+1} = \frac{s_t+d_t}{1+n}=\frac{1}{(1+n)(2+\rho)}(1-\alpha)f{(\tilde{k_t})} $$

We want to find the following Steady State level of capital-per-worker:
$$ \bar k' = \left[\frac{(1-\alpha)}{(1+n)(2+\rho)}\right]^\frac{1}{1-\alpha}$$

```python
# Again, we can use sympy to analytically find the steady state expression
# We define all the symbols:
k1 = sm.symbols('k1')
alpha1 = sm.symbols('alpha1')
rho1 = sm.symbols('rho1')
tau1 = sm.symbols('tau1')
n1 = sm.symbols('n1')
```

```python
# We define the capital accumulation equation and solve for k to find the Steady State value:
ss1 = sm.Eq(k1,((1/((2+rho1)*(1+n1)))*(1-alpha1)*k1**alpha1))
kss1 = sm.solve(ss1,k1)[0]
kss1
```

$\displaystyle \left(\frac{1 - \alpha_{1}}{n_{1} \rho_{1} + 2 n_{1} + \rho_{1} + 2}\right)^{- \frac{1}{\alpha_{1} - 1}}$

By looking at this output and the Steady State capital-per-worker we wanted to find, it is easy to see that it is indeed the same and what we wanted.

```python
# We turn it into a Python function
ss1_func = sm.lambdify((n1,tau1,rho1,alpha1),kss1)
```

We define the objective function in this case so that we can find the numerical solution.
$$ 0 = \frac{1}{(1+n)(2+\rho)}(1-\alpha)f(\bar{k}_{t}) - {k}^{\ast}$$

```python
# In the following, we use a different method compared to the one before
# We define some model parameters:
n1 = 0.01
tau1 = 0.1
rho1 = 0.3
alpha1 = 1/3

# We solve numerically for the steady state:
f = lambda k1: k1**alpha1
obj_kss1 = lambda kss1: kss1 - ((1/((1+n1)*(2+rho1)))*(1-alpha1)*f(kss1)) # Our objective function
result1 = optimize.root_scalar(obj_kss1,bracket=[0.1,100],method='brentq') # Our method used for the optimization
print(f'The Steady State for k is:',result1.root)
```

    The Steady State for k is: 0.15374085667498896

So, the steady state level of capital-per-worker is 0.154 under the Fully Funded social security system under these circumstances and parametersizes.

Now, we can also find the steady state capital-per-worker in this case for different values of $\alpha_1$.

```python
alphas = [0,0.1,0.2,1/3]
for alpha1 in alphas:
    f = lambda k1: k1**alpha1
    obj_kss1 = lambda kss1: kss1 - ((1/((1+n1)*(2+rho1)))*(1-alpha1)*f(kss1))
    result1 = optimize.root_scalar(obj_kss1,bracket=[0.1,100],method='brentq')
    print(f'For alpha1 = {alpha1:.3f}, the Steady State for k is:',result1.root)
```

    For alpha1 = 0.000, the Steady State for k is: 0.43047783039173493
    For alpha1 = 0.100, the Steady State for k is: 0.3486879984656627
    For alpha1 = 0.200, the Steady State for k is: 0.26381590295580987
    For alpha1 = 0.333, the Steady State for k is: 0.15374085667498896

The closer alpha is to 1 the smaller are the diminishing returns and if every unit of capital produces less, at one point capital will stop accumulating.

```python
# We will now visualize the steady state value for both social security systems side by side
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
line, = ax.plot(ss_func(n,tau,rho,alpha),'o',label='SS PAYG', color='c') # Plotting the Steady State under PAYG
plt.annotate('PAYG', xy=(0, 0.112), xytext=(0.01, 0.12),arrowprops=dict(facecolor='black', shrink=0.1),) # Adding the arrow
ax1 = fig.add_subplot(1,1,1)
line1, = ax1.plot(ss1_func(n,tau,rho,alpha),'o',label='SS Fully Funded',color='m') # Plotting the Steady State under Fully Funded
plt.annotate('FF', xy=(0, 0.154), xytext=(0.015, 0.155),arrowprops=dict(facecolor='black', shrink=0.1),) # Adding the arrow
plt.ylim(0.1, 0.16)
ax.legend(loc='lower right',facecolor='white',frameon=True) # Adding the box down right
ax.set_ylabel('$Capital$') # Adding the label on the y-axis
ax.set_title('Capital in Steady State') # Adding a title to the diagram
```

## Conclusion

Comparing the two shows us that the capital-per-worker under the pay-as-you-go social security system is larger than what it is under the fully funded social security system for the same parametervalues in this model.
This makes sense as the Fully Funded social security system does not have an effect on the capital accumulation because Public savings exactly offset private savings. Under the PAYG system though, there is this offsetting effect from tau mentioned above.