## Wave equation

#### Problem Setup

The general solution is given by:
$u(x,t) = F(x-ct) + G(x+ct)$ with F, G some functions.

Take $F(x) = x^2$ and $G(x) = \sin(x)$ and $c=1$.

Thus: $u(x,t) = (x-t)^2 + \sin(x + t)$.

Set $f = 0$.

Consider $u$ to be a Gaussian process:

$u \sim \mathcal{GP}(0, k_{uu}(x_i, x_j; \tilde{\theta}))$ with the hyperparameters $\tilde{\theta} = \{\theta, l_x, l_t\}$.

And the linear operator:

$\mathcal{L}_x^c = \frac{d^2}{dt^2} \cdot - c \frac{d^2}{dx^2} \cdot$

so that

$\mathcal{L}_x^c u = f$

Problem at hand: Estimate $c$ (should be $c = 1$ in the end).

#### Step 1: Simulate data

```python
import time
import numpy as np
import sympy as sp
from scipy.optimize import minimize
import matplotlib.pyplot as plt
import warnings
```

$x \in [0, 1]^n, \; t \in [0,1]^n$

```python
def get_simulated_data(n = 20):
    t = np.random.rand(n)
    x = np.random.rand(n)
    y_u = np.multiply(x-t, x-t) + np.sin(x+t)
    y_f = 0*x
    return(x, t, y_u, y_f)

(x, t, y_u, y_f) = get_simulated_data()
```

#### Step 2: Evaluate kernels

1) $k_{uu}(y_i, y_j; \tilde{\theta}) = \theta exp(-\frac{1}{2l_x}(x_i-x_j)^2 - \frac{1}{2l_t}(t_i-t_j)^2)$, where $y_i = (x_i, t_i)$, $y_j = (x_j, t_j)$.

```python
x_i, x_j, t_i, t_j, theta, l_x, l_t, c = sp.symbols('x_i x_j t_i t_j theta l_x l_t c')
kuu_sym = theta*sp.exp(-1/(2*l_x)*((x_i - x_j)**2) - 1/(2*l_t)*((t_i - t_j)**2))
kuu_fn = sp.lambdify((x_i, x_j, t_i, t_j, theta, l_x, l_t), kuu_sym, "numpy")
def kuu(x, t, theta, l_x, l_t):
    k = np.zeros((x.size, x.size))
    for i in range(x.size):
        for j in range(x.size):
            k[i,j] = kuu_fn(x[i], x[j], t[i], t[j], theta, l_x, l_t)
    return k
```

2) $k_{ff}(y_i,y_j; \tilde{\theta}, c)
= \mathcal{L}_{y_i}^c \mathcal{L}_{y_j}^c k_{uu}(y_i, y_j; \tilde{\theta}) \\
= \frac{d^4}{dt_i^2 dt_j^2}k_{uu} - c\frac{d^4}{dt_i^2 dx_j^2}k_{uu} - c\frac{d^4}{dx_i^2 dt_j^2}k_{uu} + c^2\frac{d^4}{dx_i^2 dx_j^2}k_{uu}$

```python
kff_sym = sp.diff(kuu_sym, t_i, t_i, t_j, t_j) \
        - c*sp.diff(kuu_sym, t_i, t_i, x_j, x_j) \
        - c*sp.diff(kuu_sym, x_i, x_i, t_j, t_j) \
        + c**2*sp.diff(kuu_sym, x_i, x_i, x_j, x_j)
kff_fn = sp.lambdify((x_i, x_j, t_i, t_j, theta, l_x, l_t, c), kff_sym, "numpy")
def kff(x, t, theta, l_x, l_t, c):
    k = np.zeros((x.size, x.size))
    for i in range(x.size):
        for j in range(x.size):
            k[i,j] = kff_fn(x[i], x[j], t[i], t[j], theta, l_x, l_t, c)
    return k
```

3) $k_{fu}(y_i,y_j;\tilde{\theta}, c)
= \mathcal{L}_{\tilde{x}_i}^c k_{uu}(y_i, y_j; \tilde{\theta})
= \frac{d^2}{dt_i^2}k_{uu} - c\frac{d^2}{dx_i^2}k_{uu}$

```python
kfu_sym = sp.diff(kuu_sym, t_i, t_i) - c*sp.diff(kuu_sym, x_i, x_i)
kfu_fn = sp.lambdify((x_i, x_j, t_i, t_j, theta, l_x, l_t, c), kfu_sym, "numpy")
def kfu(x, t, theta, l_x, l_t, c):
    k = np.zeros((x.size, x.size))
    for i in range(x.size):
        for j in range(x.size):
            k[i,j] = kfu_fn(x[i], x[j], t[i], t[j], theta, l_x, l_t, c)
    return k
```

4) $k_{uf}(y_i, y_j; \tilde{\theta}, c)$ is given by the transpose of $k_{fu}(y_i, y_j; \tilde{\theta}, c)$.

```python
def kuf(x, t, theta, l_x, l_t, c):
    return kfu(x, t, theta, l_x, l_t, c).T
```

#### Steps 3 and 4: Compute NLML and optimize the hyperparameters

```python
def nlml(params, x, t, y1, y2, s):
    params = np.exp(params)
    K = np.block([
        [kuu(x, t, params[0], params[1], params[2]) + s*np.identity(x.size), kuf(x, t, params[0], params[1], params[2], params[3])],
        [kfu(x, t, params[0], params[1], params[2], params[3]), kff(x, t, params[0], params[1], params[2], params[3]) + s*np.identity(x.size)]
    ])
    y = np.concatenate((y1, y2))
    val = 0.5*(np.log(abs(np.linalg.det(K))) + np.mat(y) * np.linalg.inv(K) * np.mat(y).T)
    return val.item(0)
```

```python
def minimize_restarts(x,t,y_u,y_f,n=10):
    nlml_wp = lambda params: nlml(params, x, t, y_u, y_f, 1e-7)
    all_results = []
    for it in range(0,n):
        all_results.append(minimize(nlml_wp, np.random.rand(4), method="Nelder-Mead"))
    filtered_results = [m for m in all_results if 0==m.status]
    return min(filtered_results, key = lambda x: x.fun)
```

```python
m = minimize_restarts(x, t, y_u, y_f, 5)
np.exp(m.x[3]) # This is the optimized value for our parameter c
```

    0.9978007101765757

#### Step 5: Plotting the behavior for varied parameters

The logarithms of the optimal hyperparameters are given by (arranged in $[\theta, l_x, l_t, c]$):

```python
m.x
```

    array([ 6.67851868e+00,  2.71167786e+00,  2.66415024e+00, -2.20171181e-03])

We want to plot the behavior of the nlml-function around the minimizer:

```python
lin0 = np.linspace(5, 9, 200) # Set to 200
lin1 = np.linspace(0, 5, 200)
lin3 = np.linspace(0, 0.1, 200)
res0 = [nlml((q, m.x[1], m.x[2], m.x[3]), x, t, y_u, y_f, 1e-7) for q in lin0]
res1 = [nlml((m.x[0], q, m.x[2], m.x[3]), x, t, y_u, y_f, 1e-7) for q in lin1]
res2 = [nlml((m.x[0], m.x[1], q, m.x[3]), x, t, y_u, y_f, 1e-7) for q in lin1]
res3 = [nlml((m.x[0], m.x[1], m.x[2], q), x, t, y_u, y_f, 1e-7) for q in lin3]

def show_1(lin0, lin1, lin3, res0, res1, res2, res3):
    f, (ax1, ax2) = plt.subplots(ncols=2, nrows=2, figsize=(13,7))
    f.suptitle("Local behavior of the nlml around the optimum")
    ax1[0].plot(lin0, res0)
    ax1[0].set(xlabel= r"Value of $\theta$", ylabel= "nlml")
    ax1[1].plot(lin1, res1)
    ax1[1].set(xlabel= r"Value of $l_x$", ylabel= "nlml")
    ax2[0].plot(lin1, res2)
    ax2[0].set(xlabel= r"Value of $l_t$", ylabel= "nlml")
    ax2[1].plot(lin3, res3)
    ax2[1].set(xlabel= r"Value of c", ylabel= "nlml")
    plt.show()
```

```python
show_1(lin0, lin1, lin3, res0, res1, res2, res3);
```

```python
lin = np.linspace(0, 10, 50)
res = [nlml((q, m.x[1], m.x[2], m.x[3]), x, t, y_u, y_f, 1e-7) for q in lin]
plt.plot(lin, res);
```

```python
lin = np.linspace(0, 10, 50)
res = [nlml((m.x[0], q, m.x[2], m.x[3]), x, t, y_u, y_f, 1e-7) for q in lin]
plt.plot(lin, res);
```

```python
lin = np.linspace(0, 10, 50)
res = [nlml((m.x[0], m.x[1], q, m.x[3]), x, t, y_u, y_f, 1e-7) for q in lin]
plt.plot(lin, res);
```

```python
lin = np.linspace(-1, 1, 50)
res = [nlml((m.x[0], m.x[1], m.x[2], q), x, t, y_u, y_f, 1e-7) for q in lin]
plt.plot(lin, res);
```

#### Step 6: Analysis of the error

In this section we want to analyze the error of our algorithm using two different ways and plot its time complexity.

```python
res = np.zeros((5,25))
timing = np.zeros((5,25))

# Needed for L2-Norm-calculation: in columns 25 vectors for five runs
X = np.zeros((25, 25, 5))
T = X

warnings.filterwarnings("ignore")
for k in range(5):
    for n in range(25):
        start_time = time.time()
        (x, t, y_u, y_f) = get_simulated_data(n)

        # Storing the x and t-values
        for j in range(n):
            X[j][n][k] = x[j]
            T[j][n][k] = t[j]

        m = minimize(nlml, np.random.rand(4), args=(x, t, y_u, y_f, 1e-7), method="Nelder-Mead")
        res[k][n] = np.exp(m.x[3])
        timing[k][n] = time.time() - start_time
```

**1. Plotting the error in our estimate for c:**

The error is given by $| c_{estimate} - c_{true} |$.

```python
lin = np.linspace(8, res.shape[1] - 1, res.shape[1] - 8)
ones = np.ones(res.shape[1])
est = np.repeat(0.043, len(lin))

def show_2(lin, ones, est):
    f, (ax1, ax2) = plt.subplots(ncols=2, nrows=2, figsize=(13,7))
    to_del = np.linspace(0, 7, 8)

    a0 = np.delete(np.abs(res[0,:] - ones), to_del)
    ax1[0].plot(lin, a0)
    ax1[0].plot(lin, est, linestyle='dashed', color='green')
    ax1[0].set(xlabel= r"Number of data points", ylabel= "Error")

    a1 = np.delete(np.abs(res[1,:] - ones), to_del)
    ax1[1].plot(lin, a1)
    ax1[1].plot(lin, est, linestyle='dashed', color='green')
    ax1[1].set(xlabel= r"Number of data points", ylabel= "Error")

    a2 = np.delete(np.abs(res[2,:] - ones), to_del)
    ax2[0].plot(lin, a2)
    ax2[0].plot(lin, est, linestyle='dashed', color='green')
    ax2[0].set(xlabel= r"Number of data points", ylabel= "Error")

    a3 = np.delete(np.abs(res[3,:] - ones), to_del)
    ax2[1].plot(lin, a3)
    ax2[1].plot(lin, est, linestyle='dashed', color='green')
    ax2[1].set(xlabel= r"Number of data points", ylabel= "Error")

    plt.show()
```

```python
show_2(lin, ones, est)
```

We ran the algorithm five times and plotted the respective outcomes in different colors:

```python
lin = np.linspace(8, res.shape[1] - 1, res.shape[1] - 8)

def show_3(lin, ones, res):
    plt.figure(figsize=(5,3))

    for i in range(res.shape[0]):
        to_del = np.linspace(0, 7, 8)
        a_i = np.delete(np.abs(res[i,:] - ones), to_del)
        plt.plot(lin, a_i)
        plt.suptitle('Error in our estimate for c')
        plt.ylabel('Error')
        plt.xlabel('Number of data points')

    est1 = np.repeat(0.041, len(lin))
    plt.plot(lin, est1, color='blue', linestyle='dashed')
    plt.show();
```

```python
show_3(lin, ones, res)
```

We see that for n sufficiently large (in this case $n \geq 10$), we can assume the error to be bounded by 0.041.

**2. Plotting the error between the solution and the approximative solution:**

```python
lin = np.linspace(8, res.shape[1] - 1, res.shape[1] - 8)
ones = np.ones(res.shape[1])
diff = np.ndarray(res.shape[1])

def show_4(lin, ones, res, diff):
    plt.figure(figsize=(5,3))

    to_del = np.linspace(0, 7, 8)
    for i in range(res.shape[0]):
        for j in range(res.shape[1]):
            diff[j] = np.linalg.norm((X[:,j,i] - np.sqrt(res[i,j])*T[:,j,i])**2 + \
                    np.sin(X[:,j,i]+np.sqrt(res[i,j])*T[:,j,i])-(X[:,j,i]-T[:,j,i])**2 - \
                    np.sin(X[:,j,i]+T[:,j,i]))
        diff_i = np.delete(diff, to_del)
        plt.suptitle('$L^2$-error in our estimate for c')
        plt.plot(lin, diff_i)
        plt.ylabel('Error')
        plt.xlabel('Number of data points')

    est = np.repeat(0.015, len(lin))
    plt.plot(lin, est, color='blue', linestyle='dashed')
    plt.show()
```

```python
show_4(lin, ones, res, diff)
```

The $L^2$-error is in our case bounded by 0.015 for $n \geq 10$.

**3. Plotting the execution time:**

```python
lin = np.linspace(1, timing.shape[1], timing.shape[1])

for i in range(timing.shape[0]):
    plt.plot(lin, timing[i,:])
    plt.ylabel('Execution time in seconds')
    plt.xlabel('Number of data points')
    plt.show()
```

```python
lin = np.linspace(1, timing.shape[1], timing.shape[1])

def show_5(lin, timing):
    plt.figure(figsize=(5,3))

    for i in range(timing.shape[0]):
        plt.suptitle('Execution time of our algorithm')
        plt.plot(lin, timing[i,:])
        plt.ylabel('Seconds')
        plt.xlabel('Number of data points')

    est = lin**(1.33)
    plt.plot(lin, est, color='blue', linestyle='dashed')
    plt.show()
```

```python
show_5(lin, timing)
```

Curiously, the time complexity seems to be around $\mathcal{O}(n^{4/3})$ (blue-dashed line).

Assuming an equal amount of function evaluations in the Nelder-Mead algorithm for different values of n,
we would have been expecting a time complexity of $\mathcal{O}(n^3)$, due to the computation of the inverse of an $n\times n$-matrix in every evaluation of $\textit{nlml}$. This could probably be seen with larger values of n.


# Programming Assignment

## RealNVP for the LSUN bedroom dataset

### Instructions

In this notebook, you will develop the RealNVP normalising flow architecture from scratch, including the affine coupling layers, checkerboard and channel-wise masking, and combining into a multiscale architecture. You will train the normalising flow on a subset of the LSUN bedroom dataset.

Some code cells are provided for you in the notebook. You should avoid editing provided code, and make sure to execute the cells in order to avoid unexpected errors. Some cells begin with the line:

`#### GRADED CELL ####`

Don't move or edit this first line - this is what the automatic grader looks for to recognise graded cells. These cells require you to write your own code to complete them, and are automatically graded when you submit the notebook. Don't edit the function name or signature provided in these cells, otherwise the automatic grader might not function properly.

### How to submit

Complete all the tasks you are asked for in the worksheet. When you have finished and are happy with your code, press the **Submit Assignment** button at the top of this notebook.

### Let's get started!

We'll start running some imports, and loading the dataset. Do not edit the existing imports in the following cell. If you would like to make further Tensorflow imports, you should add them here.

```python
#### PACKAGE IMPORTS ####

# Run this cell first to import all required packages. Do not make any imports elsewhere in the notebook

import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np
import matplotlib.pyplot as plt
from IPython.display import Image

from tensorflow.keras import Model, Input
from tensorflow.keras.layers import Conv2D, BatchNormalization
from tensorflow.keras.optimizers import Adam

tfd = tfp.distributions
tfb = tfp.bijectors

# If you would like to make further imports from tensorflow, add them here
from tensorflow.keras import layers
from tensorflow.keras.regularizers import l2

```

#### The LSUN Bedroom Dataset

In this assignment, you will use a subset of the [LSUN dataset](https://www.yf.io/p/lsun). This is a large-scale image dataset with 10 scene and 20 object categories. A subset of the LSUN bedroom dataset has been provided, and has already been downsampled and preprocessed into smaller, fixed-size images.

* F. Yu, A. Seff, Y. Zhang, S. Song, T. Funkhouser and J. Xia. "LSUN: Construction of a Large-scale Image Dataset using Deep Learning with Humans in the Loop". [arXiv:1506.03365](https://arxiv.org/abs/1506.03365), 10 Jun 2015

Your goal is to develop the RealNVP normalising flow architecture using bijector subclassing, and use it to train a generative model of the LSUN bedroom data subset. For full details on the RealNVP model, refer to the original paper:

* L. Dinh, J. Sohl-Dickstein and S. Bengio. "Density estimation using Real NVP". [arXiv:1605.08803](https://arxiv.org/abs/1605.08803), 27 Feb 2017.

#### Import the data

The dataset required for this project can be downloaded from the following link:

https://drive.google.com/file/d/1scbDZrn5pkRjF_CeZp66uHVQC9o1gIsg/view?usp=sharing

You should upload this file to Drive for use in this Colab notebook. It is recommended to unzip it on Drive, which can be done using the `zipfile` package:

>```
import zipfile
with zipfile.ZipFile("/path/to/lsun_bedroom.zip","r") as zip_ref:
   zip_ref.extractall('lsun_bedroom_data')
```

```python
# Run this cell to connect to your Drive folder

from google.colab import drive
drive.mount('/content/gdrive')
```

    Mounted at /content/gdrive

```python
# unpack data

import zipfile
with zipfile.ZipFile("/content/gdrive/MyDrive/Datasets/lsun_bedroom.zip","r") as zip_ref:
  zip_ref.extractall('lsun_bedroom_data')
```

#### Load the dataset

The following functions will be useful for loading and preprocessing the dataset. The subset you will use for this assignment consists of 10,000 training images, 1000 validation images and 1000 test images.

The images have been downsampled to 32 x 32 x 3 in order to simplify the training process.

```python
# Functions for loading and preprocessing the images

def load_image(filepath):
    raw_img = tf.io.read_file(filepath)
    img_tensor_int = tf.image.decode_jpeg(raw_img, channels=3)
    img_tensor_flt = tf.image.convert_image_dtype(img_tensor_int, tf.float32)
    img_tensor_flt = tf.image.resize(img_tensor_flt, [32, 32])
    img_tensor_flt = tf.image.random_flip_left_right(img_tensor_flt)
    return img_tensor_flt, img_tensor_flt

def load_dataset(split):
    train_list_ds = tf.data.Dataset.list_files('/content/lsun_bedroom_data/{}/*.jpg'.format(split), shuffle=False)
    train_ds = train_list_ds.map(load_image)
    return train_ds
```

```python
# Load the training, validation and testing datasets splits

train_ds = load_dataset('train')
val_ds = load_dataset('val')
test_ds = load_dataset('test')
```

```python
# Shuffle the datasets

shuffle_buffer_size = 1000
train_ds = train_ds.shuffle(shuffle_buffer_size)
val_ds = val_ds.shuffle(shuffle_buffer_size)
test_ds = test_ds.shuffle(shuffle_buffer_size)
```

```python
# Display a few examples

n_img = 4
f, axs = plt.subplots(n_img, n_img, figsize=(14, 14))

for k, image in enumerate(train_ds.take(n_img**2)):
    i = k // n_img
    j = k % n_img
    axs[i, j].imshow(image[0])
    axs[i, j].axis('off')
f.subplots_adjust(wspace=0.01, hspace=0.03)
```

```python
# Batch the Dataset objects

batch_size = 64
train_ds = train_ds.batch(batch_size)
val_ds = val_ds.batch(batch_size)
test_ds = test_ds.batch(batch_size)
```

### Affine coupling layer

We will begin the development of the RealNVP architecture with the core bijector that is called the _affine coupling layer_. This bijector can be described as follows: suppose that $x$ is a $D$-dimensional input, and let $d<D$. Then the output $y$ of the affine coupling layer is given by the following equations:

$$
\begin{align}
y_{1:d} &= x_{1:d} \tag{1}\\
    y_{d+1:D} &= x_{d+1:D}\odot \exp(s(x_{1:d})) + t(x_{1:d}), \tag{2}
\end{align}
$$

where $s$ and $t$ are functions from $\mathbb{R}^d\rightarrow\mathbb{R}^{D-d}$, and define the log-scale and shift operations on the vector $x_{d+1:D}$ respectively.

The log of the Jacobian determinant for this layer is given by $\sum_{j}s(x_{1:d})_j$.

The inverse operation can be easily computed as

$$
\begin{align}
x_{1:d} &= y_{1:d}\tag{3}\\
x_{d+1:D} &= \left(y_{d+1:D} - t(y_{1:d})\right)\odot \exp(-s(y_{1:d})),\tag{4}
\end{align}
$$

In practice, we will implement equations $(1)$ and $(2)$ using a binary mask $b$:

$$
\begin{align}
\text{Forward pass:}\qquad y &= b\odot x + (1-b)\odot\left(x\odot\exp(s(b\odot x)) + t(b\odot x)\right),\tag{5}\\
\text{Inverse pass:}\qquad x &= b\odot y + (1-b)\odot\left(y - t(b\odot x) \odot\exp( -s(b\odot x))\right).\tag{6}
\end{align}
$$

Our inputs $x$ will be a batch of 3-dimensional Tensors with `height`, `width` and `channels` dimensions. As in the original architecture, we will use both spatial 'checkerboard' masks and channel-wise masks:

```python
# Run this cell to download and view a figure to illustrate the checkerboard and binary masks

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

Figure 1. Spatial checkerboard mask (left) and channel-wise mask (right). From the original paper.

#### Custom model for log-scale and shift

You should now create a custom model for the shift and log-scale parameters that are used in the affine coupling layer bijector. We will use a convolutional residual network, with two residual blocks and a final convolutional layer. Using the functional API, build the model according to the following specifications:

* The function takes the `input_shape` and `filters` as arguments
* The model should use the `input_shape` in the function argument to set the shape in the Input layer (call this layer `h0`).
* The first hidden layer should be a Conv2D layer with number of filters set by the `filters` argument, and a ReLU activation
* The second hidden layer should be a BatchNormalization layer
* The third hidden layer should be a Conv2D layer with the same number of filters as the input `h0` to the model, and a ReLU activation
* The fourth hidden layer should be a BatchNormalization layer
* The fifth hidden layer should be the sum of the fourth hidden layer output and the inputs `h0`. Call this layer `h1`
* The sixth hidden layer should be a Conv2D layer with filters set by the `filters` argument, and a ReLU activation
* The seventh hidden layer should be a BatchNormalization layer
* The eighth hidden layer should be a Conv2D layer with the same number of filters as `h1` (and `h0`), and a ReLU activation
* The ninth hidden layer should be a BatchNormalization layer
* The tenth hidden layer should be the sum of the ninth hidden layer output and `h1`
* The eleventh hidden layer should be a Conv2D layer with the number of filters equal to twice the number of channels of the model input, and a linear activation. Call this layer `h2`
* The twelfth hidden layer should split `h2` into two equal-sized Tensors along the final channel axis. These two Tensors are the shift and log-scale Tensors, and should each have the same shape as the model input
* The final layer should then apply the `tanh` nonlinearity to the log_scale Tensor. The outputs to the model should then be the list of Tensors `[shift, log_scale]`

All Conv2D layers should use a 3x3 kernel size, `"SAME"` padding and an $l2$ kernel regularizer with regularisation coefficient of `5e-5`.

_Hint: use_ `tf.split` _with arguments_ `num_or_size_splits=2, axis=-1` _to create the output Tensors_.

In total, the network should have 14 layers (including the `Input` layer).

```python
#### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def get_conv_resnet(input_shape, filters):
    """
    This function should build a CNN ResNet model according to the above specification,
    using the functional API. The function takes input_shape as an argument, which should be
    used to specify the shape in the Input layer, as well as a filters argument, which
    should be used to specify the number of filters in (some of) the convolutional layers.
    Your function should return the model.
    """
    h0 = layers.Input(shape=input_shape)
    h = layers.Conv2D(filters=filters, kernel_size=(3,3), padding="same", kernel_regularizer=l2(5e-5), activation="relu")(h0)
    h = layers.BatchNormalization()(h)

    h = layers.Conv2D(filters=input_shape[-1], kernel_size=(3,3), padding="same", kernel_regularizer=l2(5e-5), activation="relu")(h)
    h = layers.BatchNormalization()(h)

    h1 = layers.Add()([h0, h])
    h = layers.Conv2D(filters=filters, kernel_size=(3,3), padding="same", kernel_regularizer=l2(5e-5), activation="relu")(h1)
    h = layers.BatchNormalization()(h)

    h = layers.Conv2D(filters=input_shape[-1], kernel_size=(3,3), padding="same", kernel_regularizer=l2(5e-5), activation="relu")(h)
    h = layers.BatchNormalization()(h)

    h = layers.Add()([h1, h])
    h2 = layers.Conv2D(filters=2*input_shape[-1], kernel_size=(3,3), padding="same", kernel_regularizer=l2(5e-5), activation="linear")(h)

    shift, log_scale = layers.Lambda(lambda t: tf.split(t, num_or_size_splits=2, axis=-1))(h2)

    log_scale = layers.Activation(activation="tanh")(log_scale)

    model = Model(inputs=h0, outputs=[shift, log_scale])
    return model

```

```python
# Test your function and print the model summary

conv_resnet = get_conv_resnet((32, 32, 3), 32)
conv_resnet.summary()
```

    Model: "model"
    __________________________________________________________________________________________________
    Layer (type)                    Output Shape         Param #     Connected to
    ==================================================================================================
    input_1 (InputLayer)            [(None, 32, 32, 3)]  0
    __________________________________________________________________________________________________
    conv2d (Conv2D)                 (None, 32, 32, 32)   896         input_1[0][0]
    __________________________________________________________________________________________________
    batch_normalization (BatchNorma (None, 32, 32, 32)   128         conv2d[0][0]
    __________________________________________________________________________________________________
    conv2d_1 (Conv2D)               (None, 32, 32, 3)    867         batch_normalization[0][0]
    __________________________________________________________________________________________________
    batch_normalization_1 (BatchNor (None, 32, 32, 3)    12          conv2d_1[0][0]
    __________________________________________________________________________________________________
    add (Add)                       (None, 32, 32, 3)    0           input_1[0][0]
                                                                     batch_normalization_1[0][0]
    __________________________________________________________________________________________________
    conv2d_2 (Conv2D)               (None, 32, 32, 32)   896         add[0][0]
    __________________________________________________________________________________________________
    batch_normalization_2 (BatchNor (None, 32, 32, 32)   128         conv2d_2[0][0]
    __________________________________________________________________________________________________
    conv2d_3 (Conv2D)               (None, 32, 32, 3)    867         batch_normalization_2[0][0]
    __________________________________________________________________________________________________
    batch_normalization_3 (BatchNor (None, 32, 32, 3)    12          conv2d_3[0][0]
    __________________________________________________________________________________________________
    add_1 (Add)                     (None, 32, 32, 3)    0           add[0][0]
                                                                     batch_normalization_3[0][0]
    __________________________________________________________________________________________________
    conv2d_4 (Conv2D)               (None, 32, 32, 6)    168         add_1[0][0]
    __________________________________________________________________________________________________
    lambda (Lambda)                 [(None, 32, 32, 3),  0           conv2d_4[0][0]
    __________________________________________________________________________________________________
    activation (Activation)         (None, 32, 32, 3)    0           lambda[0][1]
    ==================================================================================================
    Total params: 3,974
    Trainable params: 3,834
    Non-trainable params: 140
    __________________________________________________________________________________________________

You can also inspect your model architecture graphically by running the following cell. It should look something like the following:

```python
# Run this cell to download and view an example model plot

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

```python
# Plot the model graph

tf.keras.utils.plot_model(conv_resnet, show_layer_names=False, rankdir='LR')
```

```python
# Check the output shapes are as expected

print(conv_resnet(tf.random.normal((1, 32, 32, 3)))[0].shape)
print(conv_resnet(tf.random.normal((1, 32, 32, 3)))[1].shape)
```

    (1, 32, 32, 3)
    (1, 32, 32, 3)

#### Binary masks

Now that you have a shift and log-scale model built, we will now implement the affine coupling layer. We will first need functions to create the binary masks $b$ as described above. The following function creates the spatial 'checkerboard' mask.

It takes a rank-2 `shape` as input, which correspond to the `height` and `width` dimensions, as well as an `orientation` argument (an integer equal to `0` or `1`) that determines which way round the zeros and ones are entered into the Tensor.

```python
# Function to create the checkerboard mask

def checkerboard_binary_mask(shape, orientation=0):
    height, width = shape[0], shape[1]
    height_range = tf.range(height)
    width_range = tf.range(width)
    height_odd_inx = tf.cast(tf.math.mod(height_range, 2), dtype=tf.bool)
    width_odd_inx = tf.cast(tf.math.mod(width_range, 2), dtype=tf.bool)
    odd_rows = tf.tile(tf.expand_dims(height_odd_inx, -1), [1, width])
    odd_cols = tf.tile(tf.expand_dims(width_odd_inx, 0), [height, 1])
    checkerboard_mask = tf.math.logical_xor(odd_rows, odd_cols)
    if orientation == 1:
        checkerboard_mask = tf.math.logical_not(checkerboard_mask)
    return tf.cast(tf.expand_dims(checkerboard_mask, -1), tf.float32)
```

This function creates a rank-3 Tensor to mask the `height`, `width` and `channels` dimensions of the input. We can take a look at this checkerboard mask for some example inputs below. In order to make the Tensors easier to inspect, we will squeeze out the single channel dimension (which is always 1 for this mask).

```python
# Run the checkerboard_binary_mask function to see an example
# NB: we squeeze the shape for easier viewing. The full shape is (4, 4, 1)

tf.squeeze(checkerboard_binary_mask((4, 4), orientation=0))
```

    <tf.Tensor: shape=(4, 4), dtype=float32, numpy=
    array([[0., 1., 0., 1.],
           [1., 0., 1., 0.],
           [0., 1., 0., 1.],
           [1., 0., 1., 0.]], dtype=float32)>

```python
# The `orientation` should be 0 or 1, and determines which way round the binary entries are

tf.squeeze(checkerboard_binary_mask((4, 4), orientation=1))
```

    <tf.Tensor: shape=(4, 4), dtype=float32, numpy=
    array([[1., 0., 1., 0.],
           [0., 1., 0., 1.],
           [1., 0., 1., 0.],
           [0., 1., 0., 1.]], dtype=float32)>

You should now complete the following function to create a channel-wise mask. This function takes a single integer `num_channels` as an input, as well as an `orientation` argument, similar to above. You can assume that the `num_channels` integer is even.

The function should return a rank-3 Tensor with singleton entries for `height` and `width`. In the channel axis, the first `num_channels // 2` entries should be zero (for `orientation=0`) and the final `num_channels // 2` entries should be one (for `orientation=0`). The zeros and ones should be reversed for `orientation=1`. The `dtype` of the returned Tensor should be `tf.float32`.

```python
#### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def channel_binary_mask(num_channels, orientation=0):
    """
    This function takes an integer num_channels and orientation (0 or 1) as
    arguments. It should create a channel-wise binary mask with
    dtype=tf.float32, according to the above specification.
    The function should then return the binary mask.
    """
    if orientation == 0:
        return tf.concat([tf.zeros((1,1, num_channels//2), dtype=tf.float32),
                          tf.ones((1,1, num_channels - num_channels//2), dtype=tf.float32)], axis=-1)

    return tf.concat([tf.ones((1,1, num_channels//2), dtype=tf.float32),
                      tf.zeros((1,1, num_channels - num_channels//2), dtype=tf.float32)], axis=-1)

```

```python
# Run your function to see an example channel-wise binary mask

channel_binary_mask(6, orientation=0)
```

    <tf.Tensor: shape=(1, 1, 6), dtype=float32, numpy=array([[[0., 0., 0., 1., 1., 1.]]], dtype=float32)>

```python
#### GRADED CELL ####

# Complete the following functions.
# Make sure to not change the function names or arguments.

def forward(x, b, shift_and_log_scale_fn):
    """
    This function takes the input Tensor x, binary mask b and callable
    shift_and_log_scale_fn as arguments.
    This function should implement the forward transformation in equation (5)
    and return the output Tensor y, which will have the same shape as x
    """
    shift, log_scale = shift_and_log_scale_fn(b * x)
    return b * x + (1 - b) * (x * tf.math.exp(log_scale) + shift)

def inverse(y, b, shift_and_log_scale_fn):
    """
    This function takes the input Tensor x, binary mask b and callable
    shift_and_log_scale_fn as arguments.
    This function should implement the forward transformation in equation (5)
    and return the output Tensor y, which will have the same shape as x
    """
    shift, log_scale = shift_and_log_scale_fn(b * y)
    return b * y + (1 - b) * ((y - shift) * tf.math.exp(-log_scale))

```

The new bijector class also requires the `log_det_jacobian` methods to be implemented. Recall that the log of the Jacobian determinant of the forward transformation is given by $\sum_{j}s(x_{1:d})_j$, where $s$ is the log-scale function of the affine coupling layer.

You should now complete the following functions to define the `forward_log_det_jacobian` and `inverse_log_det_jacobian` methods of the affine coupling layer bijector.

* Both functions `forward_log_det_jacobian` and `inverse_log_det_jacobian` takes an input Tensor `x` (or `y`), a rank-3 binary mask `b`, and the `shift_and_log_scale_fn` callable
* These arguments are the same as the description for the `forward` and `inverse` functions
* The `forward_log_det_jacobian` function should implement the log of the Jacobian determinant for the transformation $(5)$
* The `inverse_log_det_jacobian` function should implement the log of the Jacobian determinant for the transformation $(6)$
* Both functions should reduce sum over the last three axes of the input Tensor (`height`, `width` and `channels`)

```python
#### GRADED CELL ####

# Complete the following functions.
# Make sure to not change the function names or arguments.

def forward_log_det_jacobian(x, b, shift_and_log_scale_fn):
    """
    This function takes the input Tensor x, binary mask b and callable
    shift_and_log_scale_fn as arguments.
    This function should compute and return the log of the Jacobian determinant
    of the forward transformation in equation (5)
    """
    _, log_scale = shift_and_log_scale_fn(b * x)
    return tf.reduce_sum((1 - b) * log_scale, axis=[-1,-2,-3])

def inverse_log_det_jacobian(y, b, shift_and_log_scale_fn):
    """
    This function takes the input Tensor y, binary mask b and callable
    shift_and_log_scale_fn as arguments.
    This function should compute and return the log of the Jacobian determinant
    of the forward transformation in equation (6)
    """
    _s_log_scale = shift_and_log_scale_fn(b * y)
    return -tf.reduce_sum((1 - b) * log_scale, axis=[-1,-2,-3])

```

You are now ready to create the coupling layer bijector, using bijector subclassing. You should complete the class below to define the `AffineCouplingLayer`.

* You should complete the initialiser `__init__`, and the internal class method `_get_mask`
* The `_forward`, `_inverse`, `_forward_log_det_jacobian` and `_inverse_log_det_jacobian` methods are completed for you using the functions you have written above. Do not modify these methods
* The initialiser takes the `shift_and_log_scale_fn` callable, `mask_type` string (either `"checkerboard"` or `"channel"`, `orientation` (integer, either `0` or `1`) as required arguments, and allows for extra keyword arguments
  * The required arguments should be set as class attributes in the initialiser (note that the `shift_and_log_scale_fn` attribute is being used in the `_forward`, `_inverse`, `_forward_log_det_jacobian` and `_inverse_log_det_jacobian` methods)
  * The initialiser should call the base class initialiser, and pass in any extra keyword arguments
  * The class should have a required number of event dimensions equal to 3
* The internal method `_get_mask` takes a `shape` as an argument, which is the shape of an input Tensor
  * This method should use the `checkerboard_binary_mask` and `channel_binary_mask` functions above, as well as the `mask_type` and `orientation` arguments passed to the initialiser to compute and return the required binary mask
  * This method is used in each of the `_forward`, `_inverse`, `_forward_log_det_jacobian` and `_inverse_log_det_jacobian` methods

```python
#### GRADED CELL ####

# Complete the following class.
# Make sure to not change the class or method names or arguments.

class AffineCouplingLayer(tfb.Bijector):
    """
    Class to implement the affine coupling layer.
    Complete the __init__ and _get_mask methods according to the instructions above.
    """

    def __init__(self, shift_and_log_scale_fn, mask_type, orientation, **kwargs):
        """
        The class initialiser takes the shift_and_log_scale_fn callable, mask_type,
        orientation and possibly extra keywords arguments. It should call the
        base class initialiser, passing any extra keyword arguments along.
        It should also set the required arguments as class attributes.
        """
        super(AffineCouplingLayer, self).__init__(forward_min_event_ndims=3, **kwargs)
        self.shift_and_log_scale_fn = shift_and_log_scale_fn
        self.mask_type = mask_type
        self.orientation = orientation

    def _get_mask(self, shape):
        """
        This internal method should use the binary mask functions above to compute
        and return the binary mask, according to the arguments passed in to the
        initialiser.
        """
        if self.mask_type == "channel":
            return channel_binary_mask(shape[-1], self.orientation)
        return checkerboard_binary_mask(shape[1:], self.orientation)

    def _forward(self, x):
        b = self._get_mask(x.shape)
        return forward(x, b, self.shift_and_log_scale_fn)

    def _inverse(self, y):
        b = self._get_mask(y.shape)
        return inverse(y, b, self.shift_and_log_scale_fn)

    def _forward_log_det_jacobian(self, x):
        b = self._get_mask(x.shape)
        return forward_log_det_jacobian(x, b, self.shift_and_log_scale_fn)

    def _inverse_log_det_jacobian(self, y):
        b = self._get_mask(y.shape)
        return inverse_log_det_jacobian(y, b, self.shift_and_log_scale_fn)
```

```python
# Test your function by creating an instance of the AffineCouplingLayer class

affine_coupling_layer = AffineCouplingLayer(conv_resnet, 'channel', orientation=1,
                                            name='affine_coupling_layer')
```

```python
# The following should return a Tensor of the same shape as the input

affine_coupling_layer.forward(tf.random.normal((16, 32, 32, 3))).shape
```

    TensorShape([16, 32, 32, 3])

```python
# The following should compute a log_det_jacobian for each event in the batch

affine_coupling_layer.forward_log_det_jacobian(tf.random.normal((16, 32, 32, 3)), event_ndims=3).shape
```

    TensorShape([16])

#### Combining the affine coupling layers

In the affine coupling layer, part of the input remains unchanged in the transformation $(5)$. In order to allow transformation of all of the input, several coupling layers are composed, with the orientation of the mask being reversed in subsequent layers.

```python
# Run this cell to download and view a sketch of the affine coupling layers

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

Figure 2. RealNVP alternates the orientation of masks from one affine coupling layer to the next. From the original paper.

Our model design will be similar to the original architecture; we will compose three affine coupling layers with checkerboard masking, followed by a batch normalization bijector (`tfb.BatchNormalization` is a built-in bijector), followed by a squeezing operation, followed by three more affine coupling layers with channel-wise masking and a final batch normalization bijector.

The squeezing operation divides the spatial dimensions into 2x2 squares, and reshapes a Tensor of shape `(H, W, C)` into a Tensor of shape `(H // 2, W // 2, 4 * C)` as shown in Figure 1.

The squeezing operation is also a bijective operation, and has been provided for you in the class below.

```python
# Bijector class for the squeezing operation

class Squeeze(tfb.Bijector):

    def __init__(self, name='Squeeze', **kwargs):
        super(Squeeze, self).__init__(forward_min_event_ndims=3, is_constant_jacobian=True,
                                      name=name, **kwargs)

    def _forward(self, x):
        input_shape = x.shape
        height, width, channels = input_shape[-3:]
        y = tfb.Reshape((height // 2, 2, width // 2, 2, channels), event_shape_in=(height, width, channels))(x)
        y = tfb.Transpose(perm=[0, 2, 1, 3, 4])(y)
        y = tfb.Reshape((height // 2, width // 2, 4 * channels),
                        event_shape_in=(height // 2, width // 2, 2, 2, channels))(y)
        return y

    def _inverse(self, y):
        input_shape = y.shape
        height, width, channels = input_shape[-3:]
        x = tfb.Reshape((height, width, 2, 2, channels // 4), event_shape_in=(height, width, channels))(y)
        x = tfb.Transpose(perm=[0, 2, 1, 3, 4])(x)
        x = tfb.Reshape((2 * height, 2 * width, channels // 4),
                        event_shape_in=(height, 2, width, 2, channels // 4))(x)
        return x

    def _forward_log_det_jacobian(self, x):
        return tf.constant(0., x.dtype)

    def _inverse_log_det_jacobian(self, y):
        return tf.constant(0., y.dtype)

    def _forward_event_shape_tensor(self, input_shape):
        height, width, channels = input_shape[-3], input_shape[-2], input_shape[-1]
        return height // 2, width // 2, 4 * channels

    def _inverse_event_shape_tensor(self, output_shape):
        height, width, channels = output_shape[-3], output_shape[-2], output_shape[-1]
        return height * 2, width * 2, channels // 4
```

You can see the effect of the squeezing operation on some example inputs in the cells below. In the forward transformation, each spatial dimension is halved, whilst the channel dimension is multiplied by 4. The opposite happens in the inverse transformation.

```python
# Test the Squeeze bijector

squeeze = Squeeze()
squeeze(tf.ones((10, 32, 32, 3))).shape
```

    TensorShape([10, 16, 16, 12])

```python
# Test the inverse operation

squeeze.inverse(tf.ones((10, 4, 4, 96))).shape
```

    TensorShape([10, 8, 8, 24])

We can now construct a block of coupling layers according to the architecture described above. You should complete the following function to chain together the bijectors that we have constructed, to form a bijector that performs the following operations in the forward transformation:

* Three `AffineCouplingLayer` bijectors with `"checkerboard"` masking with orientations `0, 1, 0` respectively
* A `BatchNormalization` bijector
* A `Squeeze` bijector
* Three more `AffineCouplingLayer` bijectors with `"channel"` masking with orientations `0, 1, 0` respectively
* Another `BatchNormalization` bijector

The function takes the following arguments:
* `shift_and_log_scale_fns`: a list or tuple of six conv_resnet models
  * The first three models in this list are used in the three coupling layers with checkerboard masking
  * The last three models in this list are used in the three coupling layers with channel masking
* `squeeze`: an instance of the `Squeeze` bijector

_NB: at this point, we would like to point out that we are following the exposition in the original paper, and think of the forward transformation as acting on the input image. Note that this is in contrast to the convention of using the forward transformation for sampling, and the inverse transformation for computing log probs._

```python
#### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def realnvp_block(shift_and_log_scale_fns, squeeze):
    """
    This function takes a list or tuple of six conv_resnet models, and an
    instance of the Squeeze bijector.
    The function should construct the chain of bijectors described above,
    using the conv_resnet models in the coupling layers.
    The function should then return the chained bijector.
    """
    block = [AffineCouplingLayer(shift_and_log_scale_fns[0], 'checkerboard', orientation=0),
             AffineCouplingLayer(shift_and_log_scale_fns[1], 'checkerboard', orientation=1),
             AffineCouplingLayer(shift_and_log_scale_fns[2], 'checkerboard', orientation=0),
             tfb.BatchNormalization(),
             squeeze,
             AffineCouplingLayer(shift_and_log_scale_fns[3], 'channel', orientation=0),
             AffineCouplingLayer(shift_and_log_scale_fns[4], 'channel', orientation=1),
             AffineCouplingLayer(shift_and_log_scale_fns[5], 'channel', orientation=0),
             tfb.BatchNormalization()
            ]

    return tfb.Chain(list(reversed(block)))

```

```python
# Run your function to create an instance of the bijector

checkerboard_fns = []
for _ in range(3):
    checkerboard_fns.append(get_conv_resnet((32, 32, 3), 512))
channel_fns = []
for _ in range(3):
    channel_fns.append(get_conv_resnet((16, 16, 12), 512))

block = realnvp_block(checkerboard_fns + channel_fns, squeeze)
```

```python
# Test the bijector on a dummy input

block.forward(tf.random.normal((10, 32, 32, 3))).shape
```

    TensorShape([10, 16, 16, 12])

#### Multiscale architecture

The final component of the RealNVP is the multiscale architecture. The squeeze operation reduces the spatial dimensions but increases the channel dimensions. After one of the blocks of coupling-squeeze-coupling that you have implemented above, half of the dimensions are factored out as latent variables, while the other half is further processed through subsequent layers. This results in latent variables that represent different scales of features in the model.

```python
# Run this cell to download and view a sketch of the multiscale architecture

!wget -q -O multiscale.png --no-check-certificate "https://docs.google.com/uc?export=download&id=19Sc6PKbc8Bi2DoyupHZxHvB3m6tw-lki"
Image("multiscale.png", width=700)
```

Figure 3. RealNVP creates latent variables at different scales by factoring out half of the dimensions at each scale. From the original paper.

The final scale does not use the squeezing operation, and instead applies four affine coupling layers with alternating checkerboard masks.

The multiscale architecture for two latent variable scales is implemented for you in the following bijector.

```python
# Bijector to implement the multiscale architecture

class RealNVPMultiScale(tfb.Bijector):

    def __init__(self, **kwargs):
        super(RealNVPMultiScale, self).__init__(forward_min_event_ndims=3, **kwargs)

        # First level
        shape1 = (32, 32, 3)  # Input shape
        shape2 = (16, 16, 12)  # Shape after the squeeze operation
        shape3 = (16, 16, 6)  # Shape after factoring out the latent variable
        self.conv_resnet1 = get_conv_resnet(shape1, 64)
        self.conv_resnet2 = get_conv_resnet(shape1, 64)
        self.conv_resnet3 = get_conv_resnet(shape1, 64)
        self.conv_resnet4 = get_conv_resnet(shape2, 128)
        self.conv_resnet5 = get_conv_resnet(shape2, 128)
        self.conv_resnet6 = get_conv_resnet(shape2, 128)
        self.squeeze = Squeeze()
        self.block1 = realnvp_block([self.conv_resnet1, self.conv_resnet2,
                                    self.conv_resnet3, self.conv_resnet4,
                                    self.conv_resnet5, self.conv_resnet6], self.squeeze)

        # Second level
        self.conv_resnet7 = get_conv_resnet(shape3, 128)
        self.conv_resnet8 = get_conv_resnet(shape3, 128)
        self.conv_resnet9 = get_conv_resnet(shape3, 128)
        self.conv_resnet10 = get_conv_resnet(shape3, 128)
        self.coupling_layer1 = AffineCouplingLayer(self.conv_resnet7, 'checkerboard', 0)
        self.coupling_layer2 = AffineCouplingLayer(self.conv_resnet8, 'checkerboard', 1)
        self.coupling_layer3 = AffineCouplingLayer(self.conv_resnet9, 'checkerboard', 0)
        self.coupling_layer4 = AffineCouplingLayer(self.conv_resnet10, 'checkerboard', 1)
        self.block2 = tfb.Chain([self.coupling_layer4, self.coupling_layer3,
                                 self.coupling_layer2, self.coupling_layer1])

    def _forward(self, x):
        h1 = self.block1.forward(x)
        z1, h2 = tf.split(h1, 2, axis=-1)
        z2 = self.block2.forward(h2)
        return tf.concat([z1, z2], axis=-1)

    def _inverse(self, y):
        z1, z2 = tf.split(y, 2, axis=-1)
        h2 = self.block2.inverse(z2)
        h1 = tf.concat([z1, h2], axis=-1)
        return self.block1.inverse(h1)

    def _forward_log_det_jacobian(self, x):
        log_det1 = self.block1.forward_log_det_jacobian(x, event_ndims=3)
        h1 = self.block1.forward(x)
        _, h2 = tf.split(h1, 2, axis=-1)
        log_det2 = self.block2.forward_log_det_jacobian(h2, event_ndims=3)
        return log_det1 + log_det2

    def _inverse_log_det_jacobian(self, y):
        z1, z2 = tf.split(y, 2, axis=-1)
        h2 = self.block2.inverse(z2)
        log_det2 = self.block2.inverse_log_det_jacobian(z2, event_ndims=3)
        h1 = tf.concat([z1, h2], axis=-1)
        log_det1 = self.block1.inverse_log_det_jacobian(h1, event_ndims=3)
        return log_det1 + log_det2

    def _forward_event_shape_tensor(self, input_shape):
        height, width, channels = input_shape[-3], input_shape[-2], input_shape[-1]
        return height // 4, width // 4, 16 * channels

    def _inverse_event_shape_tensor(self, output_shape):
        height, width, channels = output_shape[-3], output_shape[-2], output_shape[-1]
        return 4 * height, 4 * width, channels // 16
```

```python
# Create an instance of the multiscale architecture

multiscale_bijector = RealNVPMultiScale()
```

#### Data preprocessing bijector

We will also preprocess the image data before sending it through the RealNVP model. To do this, for a Tensor $x$ of pixel values in $[0, 1]^D$, we transform $x$ according to the following:

$$
T(x) = \text{logit}\left(\alpha + (1 - 2\alpha)x\right),\tag{7}
$$

where $\alpha$ is a parameter, and the logit function is the inverse of the sigmoid function, and is given by

$$
\text{logit}(p) = \log (p) - \log (1 - p).
$$

You should now complete the following function to construct this bijector from in-built bijectors from the bijectors module.

* The function takes the parameter `alpha` as an input, which you can assume to take a small positive value ($\ll0.5$)
* The function should construct and return a bijector that computes $(7)$ in the forward pass

```python
#### GRADED CELL ####

# Complete the following function.
# Make sure to not change the function name or arguments.

def get_preprocess_bijector(alpha):
    """
    This function should create a chained bijector that computes the
    transformation T in equation (7) above.
    This can be computed using in-built bijectors from the bijectors module.
    Your function should then return the chained bijector.
    """
    return  tfb.Chain([tfb.Invert(tfb.Sigmoid()),
                       tfb.Shift(alpha),
                       tfb.Scale(1 - 2 * alpha)])

```

```python
# Create an instance of the preprocess bijector

preprocess = get_preprocess_bijector(0.05)
```

#### Train the RealNVP model

Finally, we will use our RealNVP model to train

We will use the following model class to help with the training process.

```python
# Helper class for training

class RealNVPModel(Model):

    def __init__(self, **kwargs):
        super(RealNVPModel, self).__init__(**kwargs)
        self.preprocess = get_preprocess_bijector(0.05)
        self.realnvp_multiscale = RealNVPMultiScale()
        self.bijector = tfb.Chain([self.realnvp_multiscale, self.preprocess])

    def build(self, input_shape):
        output_shape = self.bijector(tf.expand_dims(tf.zeros(input_shape[1:]), axis=0)).shape
        self.base = tfd.Independent(tfd.Normal(loc=tf.zeros(output_shape[1:]), scale=1.),
                                    reinterpreted_batch_ndims=3)
        self._bijector_variables = (
            list(self.bijector.variables))
        self.flow = tfd.TransformedDistribution(
            distribution=self.base,
            bijector=tfb.Invert(self.bijector),
        )
        super(RealNVPModel, self).build(input_shape)

    def call(self, inputs, training=None, **kwargs):
        return self.flow

    def sample(self, batch_size):
        sample = self.base.sample(batch_size)
        return self.bijector.inverse(sample)
```

```python
# Create an instance of the RealNVPModel class

realnvp_model = RealNVPModel()
realnvp_model.build((1, 32, 32, 3))
```

```python
# Compute the number of variables in the model

print("Total trainable variables:")
print(sum([np.prod(v.shape) for v in realnvp_model.trainable_variables]))
```

    Total trainable variables:
    315180

Note that the model's `call` method returns the `TransformedDistribution` object. Also, we have set up our datasets to return the input image twice as a 2-tuple. This is so we can train our model with negative log-likelihood as normal.

```python
# Define the negative log-likelihood loss function

def nll(y_true, y_pred):
    return -y_pred.log_prob(y_true)
```

It is recommended to use the GPU accelerator hardware on Colab to train this model, as it can take some time to train. Note that it is not required to train the model in order to pass this assignment. For optimal results, a larger model should be trained for longer.

```python
# Compile and train the model

optimizer = Adam()
realnvp_model.compile(loss=nll, optimizer=Adam())
realnvp_model.fit(train_ds, validation_data=val_ds, epochs=10)
```

    Epoch 1/10
    938/938 [==============================] - 387s 392ms/step - loss: 213.6944 - val_loss: -3272.2449
    Epoch 2/10
    938/938 [==============================] - 363s 387ms/step - loss: -4292.1904 - val_loss: -5019.3613
    Epoch 3/10
    938/938 [==============================] - 364s 388ms/step - loss: -5596.1196 - val_loss: -5903.5815
    Epoch 4/10
    938/938 [==============================] - 363s 387ms/step - loss: -6373.9087 - val_loss: -6692.9321
    Epoch 5/10
    938/938 [==============================] - 364s 387ms/step - loss: -6852.5996 - val_loss: -7120.7026
    Epoch 6/10
    938/938 [==============================] - 363s 387ms/step - loss: -7192.7466 - val_loss: -7292.8345
    Epoch 7/10
    938/938 [==============================] - 364s 388ms/step - loss: -7274.7817 - val_loss: -7627.1870
    Epoch 8/10
    938/938 [==============================] - 364s 388ms/step - loss: -7667.7085 - val_loss: -7847.4912
    Epoch 9/10
    938/938 [==============================] - 363s 387ms/step - loss: -7834.7275 - val_loss: -7753.6006
    Epoch 10/10
    938/938 [==============================] - 363s 387ms/step - loss: -7931.5469 - val_loss: -8150.2329

    <keras.callbacks.History at 0x7f7e35489890>

```python
# Evaluate the model

realnvp_model.evaluate(test_ds)
```

    157/157 [==============================] - 25s 151ms/step - loss: -8143.5366

    -8143.53662109375

#### Generate some samples

```python
# Sample from the model

samples = realnvp_model.sample(8).numpy()
```

    /usr/local/lib/python3.7/dist-packages/keras/engine/base_layer.py:2215: UserWarning: `layer.apply` is deprecated and will be removed in a future version. Please use `layer.__call__` method instead.
      warnings.warn('`layer.apply` is deprecated and '

```python
# Display the samples

n_img = 8
f, axs = plt.subplots(2, n_img // 2, figsize=(14, 7))

for k, image in enumerate(samples):
    i = k % 2
    j = k // 2
    axs[i, j].imshow(image[0])
    axs[i, j].axis('off')
f.subplots_adjust(wspace=0.01, hspace=0.03)
```

Congratulations on completing this programming assignment! In the next week of the course we will look at the variational autoencoder.


# Variability in the Arm Endpoint Stiffness

In this notebook, we will calculate the feasible endpoint stiffness of a
simplified arm model for an arbitrary movement. The calculation of the feasible
muscle forces and the generation of the movement is presented in
feasible_muscle_forces.ipynb. The steps are as follows:

1. Generate a movement using task space projection
2. Calculate the feasible muscle forces that satisfy the movement
3. Calculate the feasible endpoint stiffness

```python
# notebook general configuration

%load_ext autoreload
%autoreload 2

# imports and utilities
import numpy as np
import sympy as sp
from IPython.display import display, Image
sp.interactive.printing.init_printing()

import logging
logging.basicConfig(level=logging.INFO)

# plot
%matplotlib inline
from matplotlib.pyplot import *
rcParams['figure.figsize'] = (10.0, 6.0)

# utility for displaying intermediate results
enable_display  = True
def disp(*statement):
    if (enable_display):
        display(*statement)
```

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

## Step 1: Task Space Inverse Dynamics Controller

The task space position ($x_t$) is given as a function of the generalized
coordinates ($q$)

\begin{equation}\label{equ:task-position}
x_t = g(q), x_t \in \Re^{d}, q \in \Re^{n}, d \leq n
\end{equation}

The first and second derivatives with respect to time (the dot notation depicts
a derivative with respect to time) are given by

\begin{equation}\label{equ:task-joint-vel}
  \dot{x}_t = J_t(q) \dot{q}, \; J_t(q) =
  \begin{bmatrix}
  \frac{\partial g_1}{\partial q_1} & \cdots & \frac{\partial g_1}{\partial q_n} \\
      \vdots & \ddots & \vdots \\
      \frac{\partial g_d}{\partial q_1} & \cdots & \frac{\partial g_d}{\partial q_n}
  \end{bmatrix}
  \in \Re^{d\times n}
\end{equation}

\begin{equation}\label{equ:task-joint-acc}
  \ddot{x}_t = \dot{J}_t\dot{q} + J_t\ddot{q}
\end{equation}

The task Jacobian defines a dual relation between motion and force
quantities. The virtual work principle can be used to establish the link between
task and join space forces (augmented by the null space)

\begin{equation}\label{equ:joint-task-forces-vw}
  \begin{aligned}
    \tau^T \delta q &= f_t^T \delta x_t \\
    \tau^T \delta q &= f_t^T J_t \delta q \\
    \tau &= J_t^T f_t + N_{J_t} \tau_0, \; N_{J_t} = (I - J_t^T \bar{J}_t^T)
  \end{aligned}
\end{equation}

where $N_{J_t} \in \Re^{n \times n}$ represents the right null space of $J_t$
and $\bar{J}_t$ the generalized inverse. Let the joint space equations of motion
(EoMs) have the following form

\begin{equation}\label{equ:eom-joint-space}
  \begin{gathered}
    M(q) \ddot{q} + f(q, \dot{q}) =  \tau  \\
    f(q, \dot{q}) = \tau_g(q) + \tau_c(q, \dot{q}) + \tau_{o}(q, \dot{q})
  \end{gathered}
\end{equation}

where $M \in \Re^{n \times n}$ denotes the symmetric, positive definite joint
space inertia mass matrix, $n$ the number of DoFs of the model and ${q, \dot{q},
\ddot{q}} \in \Re^{n}$ the joint space generalized coordinates and their
derivatives with respect to time. The term $f \in \Re^{n}$ is the sum of all
joint space forces, $\tau_g \in \Re^{n}$ is the gravity, $\tau_c \in \Re^{n}$
the Coriolis and centrifugal and $\tau_{o} \in \Re^{n}$ other generalized
forces. Term $\tau \in \Re^{n}$ denotes a vector of applied generalized forces
that actuate the model.

We can project the joint space EoMs in the task space by multiplying both sides
from the left with $J_t M^{-1}$

\begin{equation}\label{equ:eom-task-space}
  \begin{gathered}
    J_t M^{-1}M \ddot{q} + J_t M^{-1}f = J_t M^{-1}\tau \\
    \ddot{x}_t - \dot{J}_t\dot{q} + J_t M^{-1}f = J_t M^{-1} (J^T_t f_t + N_{J_t} \tau_0) \\
    \Lambda_t(\ddot{x}_t + b_t) + \bar{J}_t^T f = f_t
  \end{gathered}
\end{equation}

where $\Lambda_t=(J_tM^{-1}J_t^T)^{-1} \in \Re^{d \times d}$ represents the task
space inertia mass matrix, $b_t = - \dot{J}_t\dot{q}$ the task bias term and
$\bar{J}_t^T = \Lambda_m RM^{-1} \in \Re^{d \times n}$ the generalized inverse
transpose of $J_t$ that is used to project joint space quantities in the task
space. Note that $\bar{J}_t^T N_{J_t} \tau_0 = 0$.

The planning will be performed in task space in combination with a Proportional
Derivative (PD) tracking scheme

\begin{equation}\label{equ:pd-controller}
  \ddot{x}_t = \ddot{x}_d + k_p (x_d - x_t) + k_d (\dot{x}_d - x_t)
\end{equation}

where $x_d, \dot{x}_d, \ddot{x}_d$ are the desired position, velocity and
acceleration of the task and $k_p = 50, k_d = 5$ the tracking gains.

The desired task goal is derived from a smooth sigmoid function that produces
bell-shaped velocity profiles in any direction around the initial position of
the end effector

\begin{equation}\label{equ:sigmoid}
  \begin{gathered}
    x_d(t) = [x_{t,0}(0) + a (tanh(b (t - t_0 - 1)) + 1) / 2, x_{t,1}(0)]^T, \;
    \dot{x}_d(t) = \frac{d x_d(t)}{dt}, \; \ddot{x}_d(t) = \frac{d \dot{x}_d(t)}{dt} \\
        x_d^{'} = H_z(\gamma) x_d, \; \dot{x}_d^{'} = H_z(\gamma) \dot{x}_d,
    \; \ddot{x}_d^{'} = H_z(\gamma) \ddot{x}_d
  \end{gathered}
\end{equation}

where $x_{t, 0}$, $x_{t, 1}$ represent the $2D$ components of $x_t$, $a = 0.3$,
$b = 4$ and $t_0 = 0$. Different directions of movement are achieved by
transforming the goals with $H_z(\gamma)$, which defines a rotation around
the $z$-axis of an angle $\gamma$.

```python
# import necessary modules
from model import ArmModel
from projection import TaskSpace
from controller import TaskSpaceController
from simulation import Simulation
```

```python
# construct model gravity disabled to improve execution time during numerical
# integration note that if enabled, different PD gains are required to track the
# movement accurately
model = ArmModel(use_gravity=0, use_coordinate_limits=1, use_viscosity=1)
model.pre_substitute_parameters()
```

```python
# simulation parameters
t_end = 2.0
angle = np.pi  # direction of movement
fig_name = 'results/feasible_stiffness/feasible_forces_ts180'

# define the end effector position in terms of q's
end_effector = sp.Matrix(model.ee)
disp('x_t = ', end_effector)

# task space controller
task = TaskSpace(model, end_effector)
controller = TaskSpaceController(model, task, angle=angle)

# numerical integration
simulation = Simulation(model, controller)
simulation.integrate(t_end)

# plot simulation results
fig, ax = subplots(2, 3, figsize=(15, 10))
simulation.plot_simulation(ax[0])
controller.reporter.plot_task_space_data(ax[1])
fig.tight_layout()
fig.savefig(fig_name + '.pdf', format='pdf', dpi=300)
fig.savefig(fig_name + '.eps', format='eps', dpi=300)
```

## Step 2: Calculation of the Feasible Muscle Force Space

The feasible muscle forces are calculated below. Initially, the moment arm and
maximum muscle force quantities are computed for each instance of the
movement. Then the following inequality is formed assuming a linear muscle model

\begin{equation}\label{equ:linear-muscle-null-space-inequality}
  \begin{gathered}
    f_m = f_{max} \circ a_m = f_m^{\parallel} +
    N_{R} f_{m0},\; 0 \preceq a_m \preceq 1
    \rightarrow \\
    \begin{bmatrix}
      - N_{R} \\
      \hdashline
      N_{R}
    \end{bmatrix}
    f_{m0} \preceq
    \begin{bmatrix}
      f_m^{\parallel} \\
      \hdashline
      f_{max} - f_m^{\parallel}
    \end{bmatrix} \\
    Z f_{m0} \preceq \beta
  \end{gathered}
\end{equation}

where $a_m \in \Re^{m}$ represents a vector of muscle activations, $f_{max} \in
\Re^{m}$ a vector specifying the maximum muscle forces, $\circ$ the Hadamard
(elementwise) product, $f_m^{\parallel}$ the particular muscle force solution
that satisfies the action, $N_{R}$ the moment arm null space and $f_{m0}$ the
null space forces.

The next step is to sample the inequality $Z f_{m0} \leq \beta$. This is the
bottleneck of the analysis. The *convex_bounded_vertex_enumeration* uses the
lsr method, which is a vertex enumeration algorithm for finding the vertices
of a polytope in $O(v m^3)$.

```python
# import necessary modules
from analysis import FeasibleMuscleSetAnalysis

# initialize feasible muscle force analysis
feasible_muscle_set = FeasibleMuscleSetAnalysis(model, controller.reporter)
```

## Step 3: Calculate the Feasible Task Stiffness

In the following section, we will introduce a method for calculating the
feasible muscle forces that satisfy the motion and the physiological muscle
constraints. As the muscles are the main actors of the system, it is important
to examine the effect of muscle redundancy on the calculation of limbs'
stiffness.

The muscle stiffness is defined as

\begin{equation}\label{equ:muscle-stiffness}
  K_m = \frac{\partial f_m}{\partial l_{m}},\;  K_m \in \Re^{m \times m}
\end{equation}

where $f_m \in \Re^{m}$ represents the muscle forces, $l_{m} \in \Re^{m}$ the
musculotendon lengths and $m$ the number of muscles. The joint stiffness is
defined as

\begin{equation}\label{equ:joint-stiffness}
  K_j = \frac{\partial \tau}{\partial q},\; K_j \in \Re^{n \times n}
\end{equation}

where $\tau \in \Re^{n}$, $q \in \Re^{n}$ are the generalized forces and
coordinates, respectively and $n$ the DoFs of the system. Finally, the task
stiffness is defined as

\begin{equation}\label{equ:task-stiffness}
  K_t = \frac{\partial f_t}{\partial x_t},\; K_t \in \Re^{d \times d}
\end{equation}

where $f_t \in \Re^{d}$ denotes the forces, $x_t \in \Re^{d}$ the positions and
$d$ the DoFs of the task.

The derivation starts with a model for computing the muscle stiffness matrix
$K_m$. The two most adopted approaches are to either use the force-length
characteristics of the muscle model or to approximate it using the definition of
the short range stiffness, where the latter is shown to explain most of the
variance in the experimental measurements. The short range stiffness is
proportional to the force developed by the muscle ($f_m$)

\begin{equation}\label{equ:short-range-stiffness}
  k_{s} = \gamma \frac{f_m}{l_m^o}
\end{equation}

where $\gamma = 23.4$ is an experimentally determined constant and $l_m^o$ the
optimal muscle length. This definition will be used to populate the diagonal
elements of the muscle stiffness matrix, whereas inter-muscle coupling
(non-diagonal elements) will be assumed zero since it is difficult to measure
and model in practice.

The joint stiffness is related to the muscle stiffness through the following
relationship

\begin{equation}\label{equ:joint-muscle-stiffness}
    K_j = -\frac{\partial R^T}{\partial q} \bullet_2 f_m - R^T K_m R
\end{equation}

where the first term captures the varying effect of the muscle moment arm ($R
\in \Re^{m \times n}$), while the second term maps the muscle space stiffness to
joint space. The notation $\bullet_2$ denotes a product of a rank-3 tensor
($\frac{\partial R^T}{\partial q} \in \Re^{n \times m \times n}$, a 3D matrix)
and a rank-1 tensor ($f_m \in \Re^{m}$, a vector), where the index $2$ specifies
that the tensor dimensional reduction (by summation) is performed across the
second dimension, resulting in a reduced rank-2 tensor of dimensions $n \times
n$.

In a similar manner, the task stiffness is related to the muscle stiffness
through the following relationship

\begin{equation}\label{equ:task-muscle-stiffness}
  K_t  = -J_t^{+T} \left(\frac{\partial J_t^T}{\partial q} \bullet_2
    f_t + \frac{\partial R^T}{\partial q} \bullet_2 f_m + R^T
    K_m R\right) J_t^{+}
\end{equation}

where the task Jacobian matrix ($J_t \in \Re^{d \times n}$) describes the
mapping from joint to task space ($\Re^{n} \rightarrow \Re^{d}$), $+$ stands for
the Moore-Penrose pseudoinverse and $+T$ the transposed pseudoinverse operator.

Algorithm for calculating the feasible joint stiffness:

**Step 1:** Calculate the feasible muscle forces $f_m^{\oplus}$ that satisfy the
task and the physiological muscle constraints

**Step 2:** Calculate the muscle stiffness matrix $K_m$ using the short range
stiffness model

\begin{equation*}\label{equ:short-range-stiffness-2}
  k_s = \gamma \frac{f_m}{l_m^o},\; \gamma = 23.4
\end{equation*}

**Step 3:** Calculate the task $K_t$ and joint $K_j$ stiffness

\begin{equation*}
  \begin{gathered}
    K_j = -\frac{\partial R^T}{\partial q} \bullet_2 f_m - R^T K_m R \\
    K_t  = -J_t^{+T} \left(\frac{\partial J_t^T}{\partial q} \bullet_2
    f_t + \frac{\partial R^T}{\partial q} \bullet_2 f_m + R^T
    K_m R\right) J_t^{+}
  \end{gathered}
\end{equation*}

```python
# import necessary modules
from analysis import StiffnessAnalysis
from util import calculate_stiffness_properties

base_name = 'results/feasible_stiffness/feasible_stiffness_ts180_'

# initialize stiffness analysis
stiffness_analysis = StiffnessAnalysis(model, task, controller.reporter,
                                       feasible_muscle_set)

# calculate feasible stiffness
calculate_stiffness_properties(stiffness_analysis, base_name, 0, t_end, 0.2, 500)
```

```python
Image(url=base_name + 'anim.gif')
```

The left diagram shows the feasible major and minor axes of the endpoint
stiffness using scaled ($\text{scaling} = 0.0006$) ellipses (ellipses are
omitted for visibility reasons). The ellipse is a common way to visualize the
task stiffness, where the major axis (red) of the ellipse is oriented along the
maximum stiffness and the area is proportional to the determinant of $K_t$,
conveying the stiffness amplitude. The stiffness capacity (area) is increased in
the last pose, since the arm has already reached its final position and muscle
forces are not needed for it to execute any further motion. The second diagram
(middle) depicts the distribution of ellipse parameters (area and orientation
$\phi$). Finally, the rightmost box plot shows the feasible joint stiffness
distribution at three distinct time instants. Experimental measurements have
showed that the orientation of stiffness ellipses varies in a range of about
$30^{\circ}$. While our simulation results confirm this, they also reveal a
tendency of fixation towards specific directions for higher stiffness
amplitudes. The large variation of feasible stiffness verifies that this type of
analysis conveys important findings that complement experimental observations.