## Comparison between Lambert Beers Law and Transfer Matrix absorption and their effects in a two temperature model
In this file the absorption calculated with an exponential decay, according to Lambert Beers law and the local absorption evaluated, using the Transfer Matrix Method are being compared within the framework of a two temperature model calculation originally proposed by Anisimov et al. but more recently used by J. Hohlfeld et al. and D. Schick et al. among others.
The sample case under consideration is :

| Layer       	    | length    	|  refractive index 	|
|-------------------|--------------	|----------------------	|
| Platinum       	| 3nm layer    	| n = 1.0433 + i3.0855 	|
| Cobalt         	| 15nm layer   	| n = 1.0454 + i3.2169 	|
| Chromium       	| 5nm layer    	| n = 2.0150 + i2.8488 	|
| Magnesium Oxide 	| inf- layer 	| n = 1.7660            |

The procedure of this session is to:

* Run a local absorption computation to obtain the local absorption profile of the layer (This is just done for reference, in fact the local absorption module is also implemented in the NTMpy package.)

* Run different simulations with the NTMpy, changing the input source from a Lambert Beer description to a Transfer Matrix description  of the absorption.

```python
import sys
sys.path.append("C:/Users/lUKAS/Documents/UDCM/Objects")
import TMM_abs as atmm #Import absorption file (The algorithm here is implemented in the NTMpy package)
from NTMpy import NTMpy as ntm
from matplotlib import pyplot as plt
import numpy as np
import numericalunits as u #Numerical units to show physcal dimensions. (In principle not required)
u.reset_units('SI')
```

```python
#Define the complex refractive index for every layer and the length of every layer
#layers:  Air   Pt          Co                  Cr          MgO
n_list = [1, 1.0433+3.0855j,1.0454+3.2169j,2.0150+2.8488j,1.766]
n_Pt = n_list[1]; n_Co = n_list[2]; n_Cr = n_list[3]
d_list = [np.inf,3,15,5 ,np.inf] #in nm
th0    = np.pi/4 #in rad. (0 is perpendicular to the surface)
lam0   = 400     # wavelength in vacuum
pol    = 'p'     #polarization
```

Based on the suggestions made by [Steven J. Byrnes](https://arxiv.org/abs/1603.02720) the total absorption, reflection and transmission $A$, $R$, $T$ are calculated based on the idea of transfer matrices, where the refractive index, the incident angle and the length of every layer are taken into consideration. The parameters $A,R,T$ are normalized with respect to the incident power.

The results obtained have been compared with [Steven J. Byrnes- Sample 4](https://github.com/sbyrnes321/tmm/blob/master/examples.ipynb) and to results obtained by calculations, executed via COMSOL Multiphysics, where we find that the simulation done in 1D here matches almost exactly with a 2D simulation done with COMSOL Multiphysics where each layer attributes were extended homogeneously on to one more degree of freedom. Hence our 1D model is a good approximation even to a higher dimensional case.

```python
plotpoints                  = 500
[M,M_n,t,r,T,R,A,theta]     = atmm.TM(th0,lam0,n_list,d_list,pol)
[absorp,grid]               = atmm.absorption(th0,lam0,n_list,d_list,pol,plotpoints)
```

```python
print(f"Total transmission at the end {T:.2}\n"\
      f"Total reflectance at the first layer {R:.2}\n"\
      f"Total Absortpin at the entire material {A:.2}")
```

    Total transmission at the end 0.14
    Total reflectance at the first layer 0.51
    Total Absortpin at the entire material 0.36

```python
#--> Result obtained by calling the function absorption()
plt.figure()
plt.suptitle('Local absorbtion profile of multi layer thin film', fontsize=12)
plt.title(r"$\lambda=400$nm, $\theta_0=45^°$, polarization =$p$, total Absorbtion = {:.0f} %"\
          .format(np.round(A,2)*100),fontsize=10)
plt.xlabel("Distance from surface (nm)");
plt.ylabel(r"Absorbtion per unit incident power $\left(\frac{1}{nm}\right)$")
plt.plot(grid,absorp)
plt.show()
```

Now we are considering the two-temperature model, where the equation is given below, i.e. two coupled differential equations. Note that the source term $S(x,t)$ is the part which we are going to consider with the two different approaches. Since the deposited energy and therefor also the heating depends mostly on where and how much energy is deposited.
In addition, one has to mention that we are considering zero flux boundary condition on both sides and an initial temperature of 300K along the entire material.

That is
\begin{align}
    \begin{cases}
        \partial_x(\varphi(0,t))= \partial_x(\varphi(L,t)) &= 0 &\text{ Neumann boundary condition}\\
        \varphi(x,0) &= 300 K &\text{ initial condition}
    \end{cases}
\end{align}

However, before starting the simulation the variables, given in the coupled heat diffusion equation below have to be expressed.

\begin{align}
    \begin{cases}
        C_i^E(\varphi^E(x,t))\cdot\rho_i\cdot\partial_t\varphi^E(x,t) &= \partial_x\left(k^E_i(\varphi^E_i(x,t))\cdot \partial_x\varphi^E_i(x,t)\right) + G_i\cdot(\varphi^L_i(x,t)-\varphi^E_i(x,t)) + S(x,t) \\
        C_i^L(\varphi^L(x,t))\cdot\rho_i\cdot\partial_t\varphi^L(x,t) &= \partial_x\left(k^L_i(\varphi^L_i(x,t))\cdot \partial_x\varphi^L_i(x,t)\right) + G_i\cdot(\varphi^E_i(x,t)-\varphi^L_i(x,t))
    \end{cases}
\end{align}

Where $S(x,t)$ is the source applied on the material, $C_i^{E,L}$ and $k_i^{E,L}$ are heat capacity and conductivity for every $i$´th layer with respect to the electron/lattice system, as a generic function of the temperature $\varphi^{E,L}(x,t)$ of the system. The constants $\rho_i$ and $G_i$ describe the density and the electron-lattice coupling for every layer.

```python
#Platinum
length_Pt   = 3*u.nm #Length of the Material
k_el_Pt     = 73*u.W/(u.m*u.K);#Heat conductivity
rho_Pt      = 1e3*21*u.kg/(u.m**3)#Density
C_el_Pt     = lambda Te: (740*u.J/(u.m**3*u.K**2))/(1e3*21*u.kg/(u.m**3)) *Te #Electron heat capacity
C_lat_Pt    = 2.8e6*u.J/(u.m**3*u.K**2)/rho_Pt#Lattice heat capacity
G_Pt        = 1e16*25*u.W/(u.m**3*u.K)        #Lattice electron coupling constant
#Cobalt
length_Co   = 15*u.nm;
k_el_Co     = 100*u.W/(u.m*u.K);
rho_Co      = 1e3*8.86*u.kg/(u.m**3)
C_el_Co     = lambda Te: (704*u.J/(u.m**3*u.K**2))/(1e3*8.86*u.kg/(u.m**3)) *Te
C_lat_Co    = 4e6*u.J/(u.m**3*u.K**2)/rho_Co
G_Co        = 1e16*93*u.W/(u.m**3*u.K)
#Chromium
length_Cr   = 5*u.nm;
k_el_Cr     = 95*u.W/(u.m*u.K);
rho_Cr      = 1e3*7.15*u.kg/(u.m**3)
C_el_Cr     = lambda Te: (194*u.J/(u.m**3*u.K**2))/(1e3*7.15*u.kg/(u.m**3)) *Te
C_lat_Cr    = 3.3e6*u.J/(u.m**3*u.K**2)/rho_Cr
G_Cr        = 1e16*42*u.W/(u.m**3*u.K)

```

```python
# Source
s                           = ntm.source()
s.spaceprofile              = "LB"# We are going to consider the Lambert Beer case first
s.timeprofile               =  "Gaussian"
s.FWHM                      = 0.1*u.ps
s.fluence                   = 1*u.mJ/u.cm**2
s.t0                        = 0.5*u.ps #Peake of Gaussian
s.lambda_vac                = 400 #in nm
```

Note that the Lambert Beer "law" is a heuristic one, which can be ralated to the refractive index,by
\begin{align}
    \alpha &= \frac{4\pi}{\lambda_0}\cdot \operatorname{Im}(n) \\ \nonumber
     I(z)  &= I_0e^{-\alpha z}
\end{align}
Where the optical penetration depth $\delta_p$ = $\frac{1}{\alpha}$, $\lambda_0$ is the wavelength in vacuum and as indicated $n$ is the refractive index.

```python
#Two temperatures are considered, electron and lattice
sim = ntm.simulation(2,s)
#add parameters for all layers and both systems
#lengt,refractive_index, conductivity [electron, lattice], heatCapacity [electron, lattice], density, linear Coupling
sim.addLayer(length_Pt,n_Pt,[k_el_Pt,k_el_Pt],[C_el_Pt,C_lat_Pt],rho_Pt,G_Pt) #Platinum
sim.addLayer(length_Co,n_Co,[k_el_Co,k_el_Co],[C_el_Co,C_lat_Co],rho_Co,G_Co) #Cobalt
sim.addLayer(length_Cr,n_Cr,[k_el_Cr,k_el_Cr],[C_el_Cr,C_lat_Cr],rho_Cr,G_Cr) #Chromium
sim.final_time = 5*u.ps
```

```python
v = ntm.visual(sim)
#output of v.source is the full matrix of the source(x,t)
so               = v.source()
[tt,avTemp]      = v.average()#Averaged in space
print(avTemp.shape)
```

In the first picture we can see a surface and in the second a contour plot of the source that has been applied to the system. We see that the optical penetration depth of Platinum is very similar to the one of Cobalt and Chromium, that is their imaginary part of the refractive index is very much alike, which is why the transition appears between the different layers appears almost to be smooth.

What we can do now, in order to make our calculation of the deposited energy in the system more accurate, is to multiply the amplitude of the source times $A$, that is the amount of energy absorbed. Now we are considering some results obtained by the transfer matrix formalism, but not the entire local absorption profile.
Changing the parameters and recalculation leads to:

```python
oldfluence     = s.fluence
#Here the coefficient of the absorbed amount has been added
s.fluence      = oldfluence*A
#Pass on the source with the modified incident fluence and simulate again
sim            = ntm.simulation(2,s)
sim.addLayer(length_Pt,n_Pt,[k_el_Pt,k_el_Pt],[C_el_Pt,C_lat_Pt],rho_Pt,G_Pt) #Platinum
sim.addLayer(length_Co,n_Co,[k_el_Co,k_el_Co],[C_el_Co,C_lat_Co],rho_Co,G_Co) #Cobalt
sim.addLayer(length_Cr,n_Cr,[k_el_Cr,k_el_Cr],[C_el_Cr,C_lat_Cr],rho_Cr,G_Cr) #Chromium
sim.final_time = 5*u.ps
#Visualize result
v              = ntm.visual(sim)
so             = v.source()
[tt,avT]       = v.average()
```

Comparing this output now to what we we were seeing before, we can clearly see that the temperature is way lower than before. This is since we are considering that not all the incident power is deposited in the material but some of it is back reflected.
Note however, that the shape of the pulse heating the material is the same!

Let us look at the full heat map via a contour plot to also see which areas in space get heated up.

```python
#Input is a string corresponding to the system under consideration
v.contour('1') #Electron system
```

```python
v.contour('2') #Lattice system
```

Now we will switch from the Lambert Beer law, which takes the optical penetration depth of every material into consideration to the Transfer matrix method to calculate the local absorption at every grid point, as shown above.

The only thing that has to be changed in the code is the spaceprofile type (from `"LB"` to `"TMM"`).
Some additional input parameters also have to be given, i.e. `polarization = "s"` or `"p"`.
And the incident angle `s.theta_in = number` a float number.

```python
#Transfer matrix method is now considered
s                           = ntm.source()
s.spaceprofile              =  "TMM" #Change to Transfer Matrix Method
s.FWHM                      = 0.1*u.ps
s.fluence                   = 1*u.mJ/u.cm**2
s.t0                        = 0.5*u.ps
s.lambda_vac                = lam0
#provide additional input for the TM method, as indicated above
s.theta_in                  = th0
s.polarization              = pol

sim = ntm.simulation(2,s)
sim.addLayer(length_Pt,n_Pt,[k_el_Pt,k_el_Pt],[C_el_Pt,C_lat_Pt],rho_Pt,G_Pt) #Platinum
sim.addLayer(length_Co,n_Co,[k_el_Co,k_el_Co],[C_el_Co,C_lat_Co],rho_Co,G_Co) #Cobalt
sim.addLayer(length_Cr,n_Cr,[k_el_Cr,k_el_Cr],[C_el_Cr,C_lat_Cr],rho_Cr,G_Cr) #Chromium
sim.final_time = 5*u.ps

v                = ntm.visual(sim)
so               = v.source()
[tt,avTemp] = v.average()
```

Since now we also take reflections within the material into consideration, the average temperature is higher. Most of all we see that the shape of the source is different.

We can see lots of absorption on and close to the surface of the material. On the other hand, we can also see an absorption peak at the very end. This can be explained through the appearing reflections within the layer, that occure on the Chromium- MgO- edge (recall our stack structure from the table above).

The dynamics can also be depicted in a not averaged form, that is by showing the contour plots of time and space dynamics for each system, lattice and electrons, separately.

Comparing the contour plots with respect to Lambert Beer and TMM heating, one could not only see a difference in the average temperature but also which points in space have been over/ underestimated.

```python
#Input is a string corresponding to the system under consideration
v.contour('1') #Electron system
```

```python
v.contour('2') #lattice system
```

```python
#Same plot as in the beginning of the session just integrated as a method in the visual class
[T,R,A,absorption,xflat,grid] = v.localAbsorption()
```

Summarizing we find the TMM matrix the best method and recommend it to the user if they are interested in finding a precise description of local absorption of energy, which is in the zero-flux boundary case the most important factor of heating a material. Also, different incident angles and polarization can be considered, which makes this representation a more realistic approach to physical experiments.

However, in order to get quick results and for pedagogical reasons/ consistency checks on can also use the Lambert Beer spacial profile. Keep in mind though that one most likely overshoots the energy deposit.

Since for some configuration the time step for integrating the coupled differential equation must be chosen very small, the amount of produced data can be huge. Therefore we integrated an algorithm, that uses _block averaging_, i.e. statistics in order to compress the obtained data. We can, in this case, lower the amount of data obtained from almost 2 million-time-steps (in the case above) to only a few hundred data points, still giving us a reliable result.

```python
blocks = 120 #number of points in the time grid
#Apply the blocking method to both, electron and lattice system
[blockvecE,blocktt, errorE , sigmaE] = v.blocking(avTemp[0],tt,blocks)
[blockvecL,blocktt, errorL , sigmaL] = v.blocking(avTemp[1],tt,blocks)
#Depict the data
plt.figure()
plt.suptitle(f"Average temperature blocked in time with {blocks} blocks",fontsize = 12)
plt.title(r"$err(T) = \frac{1}{\sqrt{n_b}}\sum_{i=1}^{n_b}\frac{(<T>_t-\bar{T}_i)^2}{n_b-1}$",fontsize = 9)
plt.xlabel("Time in ps"); plt.ylabel("Temperature in K")
plt.grid()
plt.errorbar(blocktt,blockvecE,yerr=errorE,ms = 3,fmt='r-o',linewidth=2,label = "Electron")
plt.errorbar(blocktt,blockvecL,yerr=errorL,ms = 3,fmt='k-o',linewidth=2,label = "Lattice")
plt.legend()
plt.show()
```


# Autoregressive flows and RealNVP

```python
from IPython.display import Image
```

This reading contains an overview of normalising flows, and introduces two popular normalising flow models: masked autoregressive flow (MAF) and RealNVP.

You'll also learn about the considerations in different architectures and the tradeoff between computational complexity and learning power.

## Introduction

Before any theory, we'll discuss an example of how normalizing flows work. Suppose you have a standard normal distribution (mean 0, variance 1). It has a single mode at 0, so, even after scaling and shifting, can't be fit well to data with two modes. However, you've seen how bijectors applied to distributions can create other distributions. A natural question is then: can we create a bimodal distribution (one with two modes) from a bijector applied to a standard normal distribution? It turns out that this is possible with the `Softsign` bijector. This is a differentiable approximation to the sign function (1 if $x$ is nonnegative, -1 if $x$ is negative). Passing a standard normal distribution through this bijector transforms the probability distribution as follows:

```python
# Run this cell to download and view an example transformed distribution

!wget -q -O normal_to_bimodal.png --no-check-certificate "https://docs.google.com/uc?export=download&id=1H8hY33aWjmpoa_mqAKQ7Xma2yfR6vr2E"
Image("normal_to_bimodal.png", width=1000)
```

As you can see, the bijector created a bimodal distribution from a standard normal one! This is just one from a huge class of possible bijectors available in TensorFlow Probability. Furthermore, since you can chain them together, it's possible to create very complicated bijectors that change standard distributions (e.g. a normal) to very complicated ones. This is how a normalizing flow works: it creates a complicated distribution by applying a bijector to a simple, well-understood and computationally implemented distribution (such as a Gaussian). In this reading, you'll learn how this works and see some implementations from previous research.

## Normalizing flows

### The one-dimensional case

The main idea of normalizing flows is to create a random variable $X$ (with complicated distribution $P$) by applying a bijector $f$ to a random variable $Z$ (with a simple distribution). For example, suppose $Z \sim N(0, 1)$ has a standard normal distribution. The goal is to find a bijector $f$ so that $ X = f(Z) \sim P $ for some target distribution $P$. Under this transformation, we can calculate the log-density using the change of variables equation:
$$ \log p(x) = \log p(z) - \log \left| \frac{\partial f}{\partial z}(z) \right| $$
where $z = f^{-1}(x)$.

Finding an $f$ that changes the distribution as required is not trivial: if the target distribution $P$ is very complex, a simple $f$ (such as a scale or shift) won't do the trick. However, we know that composing bijectors with one another creates more bijectors. Hence, one approach is to combine multiple simple bijectors $f_k$ to create a more complicated $f$:
$$ f = f_K \circ f_{K-1} \circ \ldots \circ f_1. $$

This series, where a base distribution is transformed by a series of bijectors after each other, is called a *normalizing flow*:
$$ z_0 = z$$
$$ z_k = f_k(z_{k-1}) \quad k=1, \ldots, K. $$
$$ x = z_K $$
Furthermore, the log-probability can be calculated by summing the contributions from each of the bijectors:
$$ \log p(x) = \log p(z) - \sum_{k=1}^K \log \left| \frac{\partial f_k}{\partial z_{k-1}}(z_{k-1}) \right| $$

This, however, still doesn't answer the question of how to construct the $f_k$. Usually, this is done by giving each $f_k$ some simple functional form, such as a scale and shift followed by a simple nonlinearity such as a sigmoid or ReLU. Each $f_k$ will have some parameters (such as the scale and shift values), and these can be learned via standard methods such as maximum likelihood estimation given some training data.

### The higher-dimensional case

The results above generalise straightforwardly to higher dimensions. Suppose that $\mathbf{z} \sim N(0, \mathbf{I})$ is distributed according to a multivariate unit Gaussian. The normalizing flow is then
$$ \mathbf{z}_0 = \mathbf{z}$$
$$ \mathbf{z}_k = \mathbf{f}_k(\mathbf{z}_{k-1}) \quad k=1, \ldots, K. $$
The log-probability involves the determinant of the transformation, as you'll remember from an earlier reading:
$$ \log p(\mathbf{x}) = \log p(\mathbf{z}) - \sum_{k=1}^K \log \left(\left| \det \left( \frac{\partial \mathbf{f}_k}{\partial \mathbf{z}_{k-1}}(\mathbf{z}_{k-1}) \right) \right|\right) $$
where we use the shorthand notation $\frac{\partial \mathbf{a}}{\partial \mathbf{b}}$ for the matrix with components $\frac{\partial \mathbf{a}_i}{\partial \mathbf{b}_j}$, where $i$ and $j$ index the components of $\mathbf{a}$ and $\mathbf{b}$ respectively.

Let's see an example of this from an [early research paper](https://arxiv.org/abs/1505.05770). In the figure below, the left column is the density of the target distribution $P$, and the right columns are the normalizing flow approximations with $K$=2, 8 and 32 bijectors (each with a simple form and some trainable parameters).

```python
# Run this cell to download and view example normalising flow approximations

!wget -q -O normalising_flow.png --no-check-certificate "https://docs.google.com/uc?export=download&id=16t_qyKcx3PblwB2ZRo4D0PDSxPmEugvh"
Image("normalising_flow.png", width=600)
```

As you can see, the approximation improves as the number of bijectors in the flow increases.

The reason this is useful is that it allows us to "learn" a complex distribution from data and then manipulate it. For example, to draw a new sample from the learned distribution, simply draw $\mathbf{z}$ from a standard unit Gaussian and transform it to the correct space using the normalizing flow (the series of bijectors).

**Note**: Throughout this reading, we use the index $k$ to refer to the bijector in the normalizing flow and indices $i$ and $j$ to refer to dimensions of the probability space (from 1 to $D$). From here on, for clarity, we consider a normalizing flow formed of only one bijector (with $K=1$), so that we may drop the indices $k$. The equation becomes $\mathbf{x} = \mathbf{f}(\mathbf{z})$. The reason for doing this is that we now use indices to refer to components of the vectors (e.g. $\mathbf{x} = [x_1, \ldots, x_D]^T)$ where $D$ is the dimensionality. For normalizing flows with $K>1$, the results apply for each $k$.

### Computational concerns

The above theory provides, in principle, a framework to learn and manipulate complex distributions by building them up from a simple one. There is one key difficulty, however, when going to a practical implementation. This comes from the need to calculate the determinant $\left| \det \left( \frac{\partial \mathbf{f}}{\partial \mathbf{z}} \right) \right|$ to determine the density of the transformed variable $\mathbf{x}$. The computational cost (number of operations) to calculate a determinant for a general matrix with $D$ dimensions scales as $\mathcal{O}(D^3)$. This makes general normalizing flow density calculations intractable, and some simplifications, as outlined below, are required.

### Autoregressive flows

For some matrices, calculating a determinant is easy. For example, for a lower or upper triangular matrix, the determinant is the product of the diagonal elements, of which there are $D$, meaning the determinant calculation scales linearly. Hence, to attain a linear scaling of the determinant in the number of dimensions, it is enough to enforce that $\frac{\partial f_i}{\partial z_j} = 0$ whenever $j > i$. In other words, the component $f_i$ depends only on $z_1, \ldots z_i$.

Autoregressive models can be reinterpreted as normalising flows that fulfil this requirement. These are models that model the joint density $p(\mathbf{x})$ as the product of conditionals $\prod_i p(x_i \mid \mathbf{x}_{1:i-1})$. For example, the conditionals could be parameterised as Gaussians:

$$
\begin{align}
p(x_i \mid\mathbf{x}_{1:i-1}) &= \mathcal{N}(x_i\mid\mu_i, \exp(\sigma_i)^2),\\
\text{where}\qquad \mu_i &= f_{\mu_i}(\mathbf{x}_{1:i-1})\\
\text{and}\qquad \sigma_i &= f_{\sigma_i}(\mathbf{x}_{1:i-1}).
\end{align}
$$

In the above equations, the mean and standard deviations of each conditional distribution are computed using (parameterised) functions of all previous variables. The above can alternatively be written as:

$$ x_i = \mu_i(\mathbf{x}_{1:i-1}) + \exp(\sigma_i(\mathbf{x}_{1:i-1})) z_i \quad \quad i=1, \ldots, D$$

where $z_i \sim N(0, 1)$ is sampled from a unit Gaussian. This last equation shows how the autoregressive model can be viewed as a transformation $f$ from the random variables $\mathbf{z}\in\mathbb{R}^D$ to the data $\mathbf{x}\in\mathbb{R}^D$.

This is an example of an *autoregressive* process where $x_i$ depends only on the components of $\mathbf{z}$ that are lower than or equal to $i$ but not any of the higher ones. The dependence on lower dimensions of $\mathbf{z}$ happens indirectly through the $x_i$ dependence in the $f_{\mu_i}$ and $f_{\sigma_i}$.

## Implementations

### Masked Autoregressive Flow (MAF)

An implementation of the above autoregressive flow appears in the following paper:
- George Papamakarios, Theo Pavlakou, Iain Murray (2017). [Masked Autoregressive Flow for Density Estimation](http://papers.nips.cc/paper/6828-masked-autoregressive-flow-for-density-estimation.pdf). In *Advances in Neural Information Processing Systems*, 2017.

Here, the authors use the above equations, using a masked autoencoder for distribution estimation ([MADE](http://proceedings.mlr.press/v37/germain15.pdf)) to implement the functions $f_{\mu_i}$ and $f_{\sigma_i}$. For clarity, let's see how $\mathbf{x}$ is sampled. This is done as follows:
1. $x_1 = f_{\mu_1} + \exp(f_{\sigma_1})z_1$ for $z_1 \sim N(0, 1)$
2. $x_2 = f_{\mu_2}(x_1) + \exp(f_{\sigma_2}(x_1))z_2$ for $z_2 \sim N(0, 1)$
2. $x_3 = f_{\mu_3}(x_1, x_2) + \exp(f_{\sigma_3}(x_1, x_2))z_3$ for $z_3 \sim N(0, 1)$

and so on. For the $f_{\mu_i}$ and $f_{\sigma_i}$, they use the same MADE network across the $i$, but mask the weights so that $x_i$ depends on $x_j$ for all $j<i$ but not any others. By re-using the same network, weights can be shared and the total number of parameters is significantly lower.

A note on computational complexity: determining $\mathbf{x}$ from $\mathbf{z}$ is relatively slow, since this must be done sequentially: first $x_1$, then $x_2$, and so on up to $x_D$. However, determining $\mathbf{z}$ from $\mathbf{x}$ is fast: each of the above equations can be solved for $z_i$ at the same time:
$$ z_i = \frac{x_i - f_{\mu_i}}{\exp(f_{\sigma_i})} \quad \quad i=0, \ldots, D-1$$
Hence, the *forward* pass through the bijector (sampling $\mathbf{x}$) is relatively slow, but the *inverse* pass (determining $\mathbf{z}$), which is used in the likelihood calculations used to train the model, is fast.

### Inverse Autoregressive Flow (IAF)

The inverse autoregressive flow reverses the dependencies to make the forward pass parallelisable but the inverse pass sequential. Details can be found in the following paper:
- Diederik Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, Max Welling (2016). [Improved Variational Inference with Inverse Autoregressive Flow](http://papers.nips.cc/paper/6581-improved-variational-inference-with-inverse-autoregressive-flow.pdf). In *Advances in Neural Information Processing Systems*, 2016.

It uses the same equations:

$$ x_i = \mu_i + \exp(\sigma_i) z_i \quad \quad i=1, \ldots, D$$

but has the scale and shift functions depend on the $z_i$ instead of the $x_i$:

$$ \mu_i = f_{\mu_i}(z_1, \ldots, z_{i-1}) \quad \quad \sigma_i = f_{\sigma_i}(z_1, \ldots, z_{i-1}).$$

Note that now the forward equation (determining $\mathbf{x}$ from $\mathbf{z}$) can be parallelised, but the reverse transformations require determining $z_1$, followed by $z_2$, etc. and must hence be solved in sequence.

### Real-NVP and NICE

A further simplification of these approaches can be found in these papers:
- Laurent Dinh, Jascha Sohl-Dickstein, Samy Bengio (2016). [Density estimation using Real NVP](https://arxiv.org/abs/1605.08803).
- Laurent Dinh, David Krueger, Yoshua Bengio (2014). [NICE: Non-linear Independent Components Estimation](https://arxiv.org/abs/1410.8516).

The first uses a reduced versions of the above equations, for some chosen $L$:

$$
\begin{align}
x_i &= z_i \qquad &i = 1, \ldots, d \\
x_i &= \mu_i + \exp(\sigma_i) z_i \qquad &i = d+1, \ldots D
\end{align}
$$

where

$$
\begin{align}
\mu_i &= f_{\mu_i}(z_1, \ldots, z_{d})\\
\sigma_i &= f_{\sigma_i}(z_1, \ldots, z_{d})
\end{align}
$$

Hence, nothing happens for the first $d$ dimensions, but the $z_i$ values across these values transform the $x_i$ values for remaining $D-d$. Note that, in this case, both the forward and backward pass of the flow can be done fully in parallel. The second paper is even simpler, and omits the scale term altogether.

There is, of course, a catch: such a simple form means the flow typically needs a lot of bijections (a high $K$ value) to be able to describe complicated distributions. Furthermore, the dimensions that are transformed ($D-L$ in total) and not transformed ($L$ in total) must be permuted in the different bijections: otherwise the first $L$ dimensions of $\mathbf{z}$ are never changed throughout the whole normalizing flow, which greatly limits the expressive power. You'll be creating such a normalizing flow in this week's programming assignment.

### Further reading and resources

Besides the papers cited above, there are two great blog posts that explain that material as well:
- [Normalizing Flows](http://akosiorek.github.io/ml/2018/04/03/norm_flows.html) by Adam Kosiorek
- [Normalizing Flows Tutorial](https://blog.evjang.com/2018/01/nf1.html) by Eric Jang.

Both of these offer slightly more detail than we have space for here. They also have some great visuals. Happy reading and have fun implementing these ideas in the next few lessons!

Week3/Autoregressive flows and RealNVP.ipynb


# sympy Laplace Transforms for solving linear ODEs
> We've found that the Laplace transform  utility in sympy version up until now (v1.9) doesn't do what we need for solving linear ODEs.  So we made some needed improvements and included the code in the [MATH280 Package](https://github.com/ejbarth/MATH280).
- toc: true
- comments: true
- categories: [MATH280]

```python
from sympy import *
import MATH280

x = Function("x")
t,s = symbols("t s",real=True,  nonnegative=true )
```

    MATH280

## Current sympy version 1.9 and hoping for a fix soon

A famously intense pen-and-paper technique from elementary ODEs class:  Laplace Transforms.
Sympy v1.9 has Laplace capabilities built-in, but with maddening shortcomings.   This may be addressed in an upcomming new version.  I see in  [the pre-release notes for sympy 1.10](https://github.com/sympy/sympy/wiki/release-notes-for-1.10) "The Laplace transform is now rule-based, can transform a far wider range of expressions, and works better with undefined functions, including the ability to transform differential equations. "

Let's work with a 2nd order, linear, constant-coefficient, nonhomogeous IVP:
$$ \ddot{x}+3\dot{x}+4x = \sin(2t), \;\;\;\;\;\; x(0)=1, \dot{x}(0)=0$$

```python
ode = Eq(x(t).diff(t,2)+3*x(t).diff(t)+4*x(t) , sin(2*t) )
ode
```

$\displaystyle 4 x{\left(t \right)} + 3 \frac{d}{d t} x{\left(t \right)} + \frac{d^{2}}{d t^{2}} x{\left(t \right)} = \sin{\left(2 t \right)}$

## laplace_transform() in sympy 1.9

First let's see how the built-in sympy `laplace_transform()` handles that equation:

```python
laplace_transform(ode,t,s)
```

Look at that!  **'Equality' object has no attribute**.  `laplace_transform()` would prefer that we assume the expression we enter is equal to 0:

```python
Lx=laplace_transform(x(t).diff(t,2)+3*x(t).diff(t)+4*x(t) - sin(2*t),t,s)
Lx
```

    (4*LaplaceTransform(x(t), t, s) + 3*LaplaceTransform(Derivative(x(t), t), t, s) + LaplaceTransform(Derivative(x(t), (t, 2)), t, s) - 2/(s**2 + 4),
     0,
     True)

We see the output above is a tuple that includes some conditional statements at the end.  To hide that and see just the transform itself, use the option `noconds=True`

```python
Lx=laplace_transform(x(t).diff(t,2)+3*x(t).diff(t)+4*x(t) - sin(2*t),t,s,noconds=True)
Lx
```

$\displaystyle 4 \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) + 3 \mathcal{L}_{t}\left[\frac{d}{d t} x{\left(t \right)}\right]\left(s\right) + \mathcal{L}_{t}\left[\frac{d^{2}}{d t^{2}} x{\left(t \right)}\right]\left(s\right) - \frac{2}{s^{2} + 4}$

### Laplace Transform and Derivatives

Notice in that output another hassle:  `laplace_transform()` ignores the single most useful property of the Laplace Transform:
$$ {\cal L}\{x'(t)\} = s{\cal L}\{x(t)\} - x(0), \mbox{ and } {\cal L}\{x''(t)\} = s^2{\cal L}\{x(t)\} - sx(0)- x'(0).$$

## laplace() in MATH280

We've reworked that and included an alternative `laplace()` in the [MATH280 Module](https://github.com/ejbarth/MATH280).

### Solving an equation with Laplace Transforms in four steps:

#### 1. take the transform of everything

```python
L = MATH280.laplace(ode,t,s)
L
```

$\displaystyle s^{2} \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) + 3 s \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) - s x{\left(0 \right)} + 4 \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) - 3 x{\left(0 \right)} - \left. \frac{d}{d t} x{\left(t \right)} \right|_{\substack{ t=0 }} = \frac{2}{s^{2} + 4}$

#### 2. plug in the initial conditions

The second step in solving the equation is to plug in the initial conditions:

```python
L0=L.subs(x(0),1).subs(Subs(Derivative(x(t), t), t, 0),0)
L0
```

$\displaystyle s^{2} \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) + 3 s \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) - s + 4 \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) - 3 = \frac{2}{s^{2} + 4}$

#### 3. solve for the lapace transform of the solution function

The third step is to solve the resulting equation for the symbol ${\cal L}\{x(t)\}$

```python
Lx=solve(L0,LaplaceTransform(x(t),t,s))
Lx
```

    [(s**3 + 3*s**2 + 4*s + 14)/(s**4 + 3*s**3 + 8*s**2 + 12*s + 16)]

#### 4. look up the laplace transform to determine the solution

The fourth and final step is to "look up" that complicated expression in the variable $s$ to determine the function of $t$ with that Laplace transform. The built-in sympy funtion `inverse_laplace_transform()` works fine for that, with the slight annoyance that it doesn't align perfectly with the list format of output from `solve()`.  So I've included in  MATH280 a little wrapper called `laplaceInv()` that extracts the zeroth entry from that list. **I've noticed that this step can run really slowly.**

```python
sol = MATH280.laplaceInv(Lx,s,t)
sol
```

$\displaystyle - \frac{\cos{\left(2 t \right)} \theta\left(t\right)}{6} + \frac{\sqrt{7} e^{- \frac{3 t}{2}} \sin{\left(\frac{\sqrt{7} t}{2} \right)} \theta\left(t\right)}{2} + \frac{7 e^{- \frac{3 t}{2}} \cos{\left(\frac{\sqrt{7} t}{2} \right)} \theta\left(t\right)}{6}$

Something to notice is that every term as a $\theta(t)$.  That's the unit step function:  $\theta(t)=1$ if $t\gt 0$ and otherwise zero.  Seems to be enforcing the assumption that our solution only makes sense for nonnegative time $t$.

We can plot the solution with sympy plot:

```python
plot(sol,(t,0,20))
```

And just in case there was any doubt, we notice that `dsolve()` produces the same solution.

```python
dsolve(ode,x(t),ics={x(0):1, x(t).diff(t).subs(t,0):0})
```

$\displaystyle x{\left(t \right)} = \left(\frac{\sqrt{7} \sin{\left(\frac{\sqrt{7} t}{2} \right)}}{2} + \frac{7 \cos{\left(\frac{\sqrt{7} t}{2} \right)}}{6}\right) e^{- \frac{3 t}{2}} - \frac{\cos{\left(2 t \right)}}{6}$

## Discontinuous Forcing Functions

We haven't yet explored what I think is *the* reason to consider Laplace Transforms in the first place:  discontinuous forcing functions.

### A first-order model with a square wave switch function

Suppose we model a capacitor with an external voltage source that switches on at $t=2$ and off at $t=5$.
The equation, in some idealized units, could be
$$\dot{x} - x = \theta(t-2) - \theta(t-5), \;\;\;\;\;  x(0)=0 $$
where $\theta(t-a)$ is the Heaviside unit step function that turns on at $t=a$.

```python
capmodel = Eq(x(t).diff(t)+x(t), Heaviside(t-2)-Heaviside(t-5))
capmodel
```

$\displaystyle x{\left(t \right)} + \frac{d}{d t} x{\left(t \right)} = - \theta\left(t - 5\right) + \theta\left(t - 2\right)$

We'll complete the four steps as in the previous example:

```python
L=MATH280.laplace(capmodel,t,s)
L
```

$\displaystyle s \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) + \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) - x{\left(0 \right)} = \frac{e^{- 2 s}}{s} - \frac{e^{- 5 s}}{s}$

```python
L0=L.subs(x(0),0)
L0
```

$\displaystyle s \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) + \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) = \frac{e^{- 2 s}}{s} - \frac{e^{- 5 s}}{s}$

```python
Ls=solve(L0,LaplaceTransform(x(t),t,s))
Ls
```

    [(exp(3*s) - 1)*exp(-5*s)/(s*(s + 1))]

```python
capsol=MATH280.laplaceInv(Ls,s,t)
capsol
```

$\displaystyle - \theta\left(t - 5\right) + \theta\left(t - 2\right) + e^{5} e^{- t} \theta\left(t - 5\right) - e^{2} e^{- t} \theta\left(t - 2\right)$

The functional form of the solution looks a little opaque, so lets make a plot of the solution together with the right hand side function:

```python
plot(capsol,
     Heaviside(t-2)-Heaviside(t-5),
    (t,0,15),)
```

### A second-order example with a delta function

Here's an example from the undergraduate ODE textbook by Blanchard, Devaney and Hall:
$$\ddot{x} +2\dot{x}+3x= \delta(t-1)+3\delta(t-4), \;\;\;\;\; x(0)=0, \dot{x}(0)=0$$

```python
ode2 = Eq(x(t).diff(t,2)+2*x(t).diff(t)+3*x(t),DiracDelta(t-1)+3*DiracDelta(t-4))
ode2
```

$\displaystyle 3 x{\left(t \right)} + 2 \frac{d}{d t} x{\left(t \right)} + \frac{d^{2}}{d t^{2}} x{\left(t \right)} = 3 \delta\left(t - 4\right) + \delta\left(t - 1\right)$

```python
L=MATH280.laplace(ode2,t,s)
L
```

$\displaystyle s^{2} \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) + 2 s \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) - s x{\left(0 \right)} + 3 \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) - 2 x{\left(0 \right)} - \left. \frac{d}{d t} x{\left(t \right)} \right|_{\substack{ t=0 }} = e^{- s} + 3 e^{- 4 s}$

```python
Ls=L.subs(x(0),0).subs(Subs(Derivative(x(t), t), t, 0),0)
Ls
```

$\displaystyle s^{2} \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) + 2 s \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) + 3 \mathcal{L}_{t}\left[x{\left(t \right)}\right]\left(s\right) = e^{- s} + 3 e^{- 4 s}$

```python
Lx=solve(Ls,LaplaceTransform(x(t),t,s))
Lx
```

    [(exp(3*s) + 3)*exp(-4*s)/(s**2 + 2*s + 3)]

```python
sol2 = MATH280.laplaceInv(Lx,s,t)
sol2
```

$\displaystyle \frac{3 \sqrt{2} e^{4} e^{- t} \sin{\left(\sqrt{2} t \right)} \cos{\left(4 \sqrt{2} \right)} \theta\left(t - 4\right)}{2} + \frac{\sqrt{2} e e^{- t} \sin{\left(\sqrt{2} t - \sqrt{2} \right)} \theta\left(t - 1\right)}{2} - \frac{3 \sqrt{2} e^{4} e^{- t} \sin{\left(4 \sqrt{2} \right)} \cos{\left(\sqrt{2} t \right)} \theta\left(t - 4\right)}{2}$

```python
plot(sol2,(t,0,20))
```

We can see in the plot above how the impulses at $t=1$ and $t=4$ energize the system.


```python
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
plt.style.use('classic')
%matplotlib inline
```

# Class 6: Pandas and Statsmodels

## Background

In their article "The New Kaldor Facts: Ideas, Institutions, Population, and Human Capital" in the January 2010 issue of the *American Economic Journal: Macroeconomics*, Charles Jones and Paul Romer document six empirical features about macroeconomics that are yet to be adequately explained by a general theory of economic growth (link to article: [https://www.jstor.org/stable/25760291](https://www.jstor.org/stable/25760291)). Among these facts is the observation that total factor productivity is positively correlated with a country's GDP per worker.

Recall the Cobb-Douglas production function with human capital:

\begin{align}
Y & = A K^{\alpha}\left(hL\right)^{1-\alpha}, \label{eqn:production}
\end{align}

where $Y$ denotes the production of final goods and services, $K$ is the stock of physical capital, $L$ is the labor force (either number of workers or number of worker hours), $h$ is human capital per worker, and $A$ is *total factor productivity* or TFP which represents all other factors that might affect production that are not captured by $K$, $L$, or $h$.

For most countries, we can measure $Y$, $K$, $L$ and $h$. Measurements may not be perfect; for example, we often associate $h$ with years of schooling in a country even though, strictly speaking, on-the-job experience and training also contribute to human capital accumulation. Regardless, the point is that the production function *implies* a value for TFP:

\begin{align}
A & = \frac{Y}{K^{\alpha}(hL)^{1-\alpha}}
\end{align}

TFP is like a residual in a linear regression: it is the part of production that is not explained by observable inputs for production. Remarkably, Jones and Romer show that TFP explains a lot of the variation in GDP per worker across countries. This fact is captured in Figure 1.

Figure 1. TFP and GDP per capita across countries  relative to the US. Source: Jones and Romer (2010)

In this Notebook, we will replicate Figure 1 and use a linear regression procedure to estimate the relationship between TFP and GDP per worker. The data required for the exercise is named `cross_country_production.csv` and is available here: https://github.com/letsgoexploring/econ126/raw/master/Data/Csv/cross_country_production.csv
The file contains production data in 2014 for 144 countries. In addition to the name and country code, the file contains production data for 144 countries in 2014:

| Name               | Description               | Units                       |
|--------------------|---------------------------|-----------------------------|
| `gdp`              | real GDP                  | Millions of 2001 US dollars |
| `labor`            | Average number of workers | Millions                    |
| `physical_capital` | Physical capital stock    | Millions of 2001 US dollars |
| `human_capital`    | Human capital per worker  | Index (unitless)            |

We can use the data to compute TFP for each country in the sample and then to replicate Figure 1.

## Replicate Figure from Jones and Romer (2010)

Import the production data. Add new columns containing GDP per worker, physical capital per worker, and TFP. Construct three scatter plots:
1. GDP per worker against physical capital per worker
2. GDP per worker against human capital per worker
3. TFP against GDP per worker

```python
# Import data into a DataFrame called 'df'. Use the leftmost column as the index

# Print the first five rows of the imported data

```

```python
# Add a column to df called 'gdp_pw' equal to GDP divided by labor divided by 1000 so units are in thousands of $

# Add a column to df called 'physical_capital_pw' equal to physical capital divided by labor divided by 1000 so units are in thousands of $

# Set alpha to 1/3

# Add a column to df called 'tfp' equal to TFP computed using the production function

# Sort the rows of df so that the GDP per worker column is in ascending order

```

```python
# Construct a scatter plot of real GDP per worker against human capital per worker with:
#    1. All countries as blue circles, size 75, opacity (alpha) 0.25, label: 'All countries'
#    2. USA as red square, size 75, opacity (alpha) 0.5, label='USA'
#    3. Add a legend in the upper left corner with scatterpoints=1
#    4. y-axis limits: [-10,250]
#    5. x-axis label: 'Human capital per worker\n(Index)'
#    6. y-axis label: 'Real GDP per worker\n(Thousands of 2001 $)'
#    7. Add a grid if you'd like

```

```python
# Construct a scatter plot of real GDP per worker against physical capital per worker with:
#    1. All countries as blue circles, size 75, opacity (alpha) 0.25, label: 'All countries'
#    2. USA as red square, size 75, opacity (alpha) 0.5, label='USA'
#    3. Add a legend in the upper left corner with scatterpoints=1
#    4. y-axis limits: [-10,250]
#    5. x-axis label: 'Physical capital per worker\n(Thousands of 2001 $)'
#    6. y-axis label: 'Real GDP per worker\n(Thousands of 2001 $)'
#    7. Add a grid if you'd like

```

```python
# Construct a scatter plot of TFP against real GDP per worker with:
#    1. All countries as blue circles, size 75, opacity (alpha) 0.25, label: 'All countries'
#    2. USA as red square, size 75, opacity (alpha) 0.5, label='USA'
#    3. Add a legend in the upper left corner with scatterpoints=1
#    4. x-axis limits: [-10,250]
#    4. y-axis limits: [-10,1600]
#    6. x-axis label: 'Real GDP capital per worker\n(Thousands of 2001 $)'
#    7. y-axis label: 'TFP (Unitless)'
#    8. Add a grid if you'd like

```

## Estimate Relationship Between TFP and GDP per Worker

The Statsmodels module is an excellent tool for a broad range of statistical applications. Learn more about Statsmodels here: https://www.statsmodels.org/stable/index.html. We can use the module to estimate the relationship between TFP and GDP per worker in our data.

Since the scatter plot reveals some curvature in the relationship between TFP and GDP per worker, let's specifiy the linear regression model with the variables in logs:

\begin{align}
\log \text{TFP}_i & = \beta_0 + \beta_1 \log \text{GDP}_i + \epsilon_i
\end{align}

Once we've obtained the estimated coefficients $\hat{\beta}_0$ and $\hat{\beta}_1$, we can compute predicted values for TFP $\widehat{TFP}_i$ using by setting $\epsilon_i=0$ and solving the regression equation for TFP:

\begin{align}
\widehat{TFP}_i & = \exp\left(\hat{\beta}_0 + \hat{\beta}_1 \log \text{GDP}_i\right)
\end{align}

```python
# Import the Statsmodels module

```

```python
# Create a variable called 'y' that will be to the dependent variable in the regression

# Create a variable called 'x' that will be to the independent variable in the regression

# Use the function sm.add_constant() to add a constant to x

# Print the first five rows of x

```

```python
# Create an OLS model using Statsmodels

# Fit the OLS model to data

# Print the summary2() method of results to print summary of regression results:

```

Estimated coefficients are stored in the `params` attribute of `results`. Residuals of the regression are stored in the attribute `resid`. And predicted values of the dependent variable are stored in `fittedvalues`.

```python
# Print estimated coefficients of linear regression

```

```python
# Print R-squared of the regression

```

**Question**

1. What is the interpretation of the coefficient on log GDP per worker?
2. What is the R-squared of the regression? What does it mean?

**Answer**

1.

2.

```python
# Print the first five rows of residuals.

```

```python
# Print the mean of the residuals of the regression.

```

```python
# Create a variable called 'yhat' equal to the fitted values of y:

# Print the first five rows of yhat

```

Recall that we can use estimated parameters can be used to compute out-of-sample predictions:

\begin{align}
\widehat{TFP}_i & = \exp\left(\hat{\beta}_0 + \hat{\beta}_1 \log \text{GDP}_i\right)
\end{align}

```python
# Create variable called 'gdp_values' that is an array from 0.1 to 300

# Create variable called 'beta0' equal to the constant of the regression

# Create variable called 'beta1' equal to the coefficient on log GDP per worker from the regression

# Create a variable called 'ahat' that is equal to the predicted values of TFP corresponding to gdp_values

```

```python
# Recreate the scatter plot of TFP against real GDP per worker with:
#    1. All features and settings from the previous plot
#    2. A plot of ahat against gdp_values; linewidth: 3, opacity: 0.5, label: 'OLS', color: magenta (c='m')

```