{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&ensp;\n",
    "[Home Page](../Start_Here.ipynb)\n",
    "\n",
    "\n",
    "[Previous Notebook](Part2.ipynb)\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "[1](Start_Here.ipynb)\n",
    "[2](Part2.ipynb)\n",
    "[3]\n",
    "[4](Part4.ipynb)\n",
    "[5](Competition.ipynb)\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "[Next Notebook](Part4.ipynb)\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Steady State Flow Using Neural Networks - Part 3\n",
    "\n",
    "**Contents of the this notebook:**\n",
    "\n",
    "- [Improving Our Fully Connected Model](#Building-our-First-Model)\n",
    "- [Building a Convolutional Model](#Convolutional-model)\n",
    "- [Input Data Manipulation](#Input-Data-Manipulation)\n",
    "\n",
    "**By the End of this notebook participants will:**\n",
    "\n",
    "- Benchmark three different models and their performance\n",
    "- Understand how input data manipulation can help in building a better model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Improving our Fully Connected Network\n",
    "\n",
    "Let us import libraries, dataset and define the Loss function "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing Necessary Libraries \n",
    "from __future__ import print_function\n",
    "\n",
    "import sys\n",
    "sys.path.append('/workspace/python/source_code')\n",
    "\n",
    "import numpy as np\n",
    "import utils.data_utils as data_utils\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import layers\n",
    "import tensorflow.keras.backend as K\n",
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import time\n",
    "\n",
    "import importlib\n",
    "\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "# reload(data_utils) # you need to execute this in case you modify the plotting scripts in data_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "dataset_size = 1500   # Number of elements in the train.tfrecords\n",
    "validation_size = 256 # Number of elements to use for validation\n",
    "\n",
    "# derive some quantities\n",
    "train_size = dataset_size - validation_size\n",
    "train_batches = int(train_size / batch_size)\n",
    "validation_batches= int(validation_size / batch_size)\n",
    "\n",
    "test_size = 28\n",
    "test_batches = int(test_size/batch_size)\n",
    "print('Number of batches in train/validation/test dataset:', train_batches, '/', validation_batches, '/', test_batches)\n",
    "\n",
    "def init_datasets():\n",
    "    dataset = tf.data.TFRecordDataset('data/train.tfrecords')\n",
    "    dataset = dataset.take(dataset_size)\n",
    "    # Transform binary data into image arrays\n",
    "    dataset = dataset.map(data_utils.parse_flow_data)\n",
    "    \n",
    "    training_dataset = dataset.skip(validation_size).shuffle(buffer_size=512)\n",
    "    training_dataset = training_dataset.batch(batch_size, drop_remainder=True)\n",
    "    training_dataset = training_dataset.repeat()\n",
    "\n",
    "    validation_dataset = dataset.take(validation_size).batch(batch_size, drop_remainder=True)\n",
    "    validation_dataset = validation_dataset.repeat()\n",
    "\n",
    "    # Read test dataset\n",
    "    test_dataset = tf.data.TFRecordDataset('data/test.tfrecords')\n",
    "    test_dataset = test_dataset.map(data_utils.parse_flow_data) # Transform binary data into image arrays\n",
    "    test_dataset = test_dataset.batch(batch_size, drop_remainder = True).repeat()\n",
    " \n",
    "    return training_dataset, validation_dataset, test_dataset\n",
    "\n",
    "training_dataset, validation_dataset, test_dataset = init_datasets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an iterator for reading a batch of input and output data\n",
    "iterator = iter(training_dataset)\n",
    "boundary, vflow = next(iterator)\n",
    "\n",
    "print('input shape:', boundary.shape.as_list())\n",
    "print('output shape:', vflow.shape.as_list())\n",
    "\n",
    "plot_idx = 3 # set it between 0 and batch_size\n",
    "\n",
    "data_utils.plot_flow_data(boundary[plot_idx,:,:,:], vflow[plot_idx,:,:,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_image(vflow_hat, vflow):\n",
    "    ''' Defines the loss for the predicted flow.\n",
    "    \n",
    "    Arguments:\n",
    "    vflow_hat -- predicted flow, shape (?, nh, nw, 2)\n",
    "    vflow   -- target flow from the simulation, shape (?, nh, nw, 2)\n",
    "    \n",
    "    Returns: the L2 loss\n",
    "    '''\n",
    "    ### define the squure error loss (~ 1 line of code)\n",
    "    loss = tf.nn.l2_loss(vflow_hat - vflow)\n",
    "    ###\n",
    "                         \n",
    "    # Add a scalar to tensorboard\n",
    "    tf.summary.scalar('loss', loss)\n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model \n",
    "\n",
    "5 - Layer Fully Connected Network : \n",
    "\n",
    "* Input Layer of Size ( 128 * 256 * 1 )\n",
    "* Hidden Layer of Size ( 1024 ) \n",
    "* Hidden Layer of Size ( 1024 ) \n",
    "* Hidden Layer of Size ( 1024 ) \n",
    "* Output Layer of Size ( 128 * 256 * 2 ) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fully_connected(input):\n",
    "    # Arguments:\n",
    "    # input -- input layer for the network, expected shape (?,nh,nw,1)\n",
    "    # Returns -- predicted flow (?, nh, nw, 2)\n",
    "    \n",
    "    nh = K.int_shape(input)[1]\n",
    "    nw = K.int_shape(input)[2]\n",
    "    \n",
    "    # define the hidden layers\n",
    "    x = layers.Flatten()(input)\n",
    "    \n",
    "     \n",
    "    ### Add three dense hidden layers with 1024 hidden units each\n",
    "    x = layers.Dense(1024, activation='relu')(x)\n",
    "    x = layers.Dense(1024, activation='relu')(x)\n",
    "    x = layers.Dense(1024, activation='relu')(x)\n",
    "    ##\n",
    "   \n",
    "    x = layers.Dense(nh*nw*2, activation='relu')(x)\n",
    "    output = layers.Reshape((nh,nw,2))(x)\n",
    "    ###\n",
    "    \n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input = tf.keras.Input(shape=(128,256,1))\n",
    "output = fully_connected(input)\n",
    "\n",
    "\n",
    "### Define a new keras model with the above input and output, and compile it with Adam optimizer\n",
    "fc_model = tf.keras.Model(inputs = input, outputs=output)\n",
    "fc_model.compile(tf.keras.optimizers.Adam(0.0001), loss=loss_image)\n",
    "###\n",
    "fc_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = fc_model.fit(training_dataset, epochs=30, steps_per_epoch=train_batches,\n",
    "          validation_data=validation_dataset, validation_steps=validation_batches, \n",
    "            callbacks=[tf.keras.callbacks.TensorBoard(log_dir='/tmp/fc3')])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let us Plot the train History\n",
    "data_utils.plot_keras_loss(history)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test\n",
    "\n",
    "We will evaluate the model on the test dataset, and plot some of the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "test_loss = fc_model.evaluate(test_dataset, steps=3)\n",
    "print('The loss over the test dataset', test_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "x, vxy = data_utils.load_test_data(1) # you can try different numbers between 1 and 28\n",
    "vxy_hat = fc_model.predict(x)\n",
    "data_utils.plot_test_result(x, vxy, vxy_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot a vertical slice of the velocity field for better comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_idx=120 # x coordinate for the slice\n",
    "\n",
    "vx = np.squeeze(vxy[0,:,x_idx,0])                 # test velocity fields\n",
    "vy = np.squeeze(vxy[0,:,x_idx,1])\n",
    "\n",
    "vx_hat = np.squeeze(vxy_hat[0,:,x_idx,0])         # predicted velocity field\n",
    "vy_hat = np.squeeze(vxy_hat[0,:,x_idx,1])\n",
    "\n",
    "fig = plt.figure(figsize=(16,6))\n",
    "\n",
    "# plot the x component of the velocity\n",
    "ax = fig.add_subplot(121)\n",
    "ax.plot(vx, label='ground truth')\n",
    "ax.plot(vx_hat, label='predicted')\n",
    "ax.legend()\n",
    "ax.set_xlabel('y')\n",
    "ax.set_ylabel('vx')\n",
    "\n",
    "# plot the y component of the velocity\n",
    "ax = fig.add_subplot(122)\n",
    "ax.plot(vy, label='ground truth')\n",
    "ax.plot(vy_hat, label='predicted')\n",
    "ax.legend()\n",
    "ax.set_xlabel('y')\n",
    "ax.set_ylabel('vy')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have bbserved the decrease in the loss function but it is still not sufficient for practical purposes. Let us now define a Convolution Model and see how it performs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convolutional model\n",
    "We will re-create the network from [Convolutional Neural Networks for Steady Flow Approximation](https://www.autodeskresearch.com/publications/convolutional-neural-networks-steady-flow-approximation). Here is an illustration from the paper: <img src='images/convnet.png' width='800px'>\n",
    "\n",
    "The number of filters and the kernel size are shown below for the conv/deconv operations. The dimension of the feature maps is indicated below in the boxes. The strides are the same as the kernel sizes.\n",
    "\n",
    "The direct connection between the input and the output layers is just a multiplication that zeros the flow inside the objects.\n",
    "\n",
    "To Learn about Convolutional Neural Networks and how they work, visit [Convolution Neural Network Notebook](../Intro_to_DL/CNN's.ipynb)\n",
    "\n",
    "### Model\n",
    "\n",
    "We define the encoding/decoding part separately, and then we combine them.\n",
    "\n",
    "We will set the parameters for [conv2d](https://keras.io/layers/convolutional/#conv2d), and add the fully connected layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv(input):\n",
    "    # Define layers to calculate the convolution and FC part of the network\n",
    "    # Arguments:\n",
    "    # input -- (?, nh, nw, nc)\n",
    "    # Returns: (? 1,1,1024)\n",
    "    \n",
    "    \n",
    "    ### Set the number of filters for the first convolutional layer\n",
    "    x = layers.Conv2D(128, (16,16), strides=(16,16), padding='same', name='conv1', activation='relu')(input)\n",
    "    \n",
    "    \n",
    "    ### Set the number of filters and kernel size for the second convolutional layer \n",
    "    x = layers.Conv2D(512, (4,4), strides=(4,4), padding='same', name='conv2', activation='relu')(x)\n",
    "    ###\n",
    "    \n",
    "    x = layers.Flatten()(x)\n",
    "    \n",
    "    \n",
    "    ### Add a denslayer with ReLU activation\n",
    "    x = layers.Dense(1024, activation='relu')(x)\n",
    "    ###\n",
    "    \n",
    "    # Reshape the output as 1x1 image with 1024 channels:\n",
    "    x = layers.Reshape((1,1,1024))(x)\n",
    "    \n",
    "    return(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define one of the decoding branch using [Conv2DTranspose](https://keras.io/layers/convolutional/#conv2dtranspose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def deconv(input, suffix):\n",
    "    # Define layers that perform the deconvolution steps\n",
    "    # Arguments:\n",
    "    # input -- (?, 1, 1, 1024)\n",
    "    # suffix -- name_suffix\n",
    "    # Returns -- (?,128,256,1)\n",
    "    x = layers.Conv2DTranspose(512, (8,8), strides=(8,8), activation='relu', name=\"deconv1_\"+suffix)(input)\n",
    "    \n",
    "    \n",
    "    ### Add the 2nd and 3rd Conv2DTranspose layers\n",
    "    x = layers.Conv2DTranspose(256, (8,4), strides=(8,4), activation='relu', name=\"deconv2_\"+suffix)(x)\n",
    "    x = layers.Conv2DTranspose(32, (2,2), strides=(2,2), activation='relu', name=\"deconv3_\"+suffix)(x)\n",
    "    ###\n",
    "    \n",
    "    x = layers.Conv2DTranspose(1, (2,2), strides=(2,2), activation='relu', name=\"deconv4_\"+suffix)(x)\n",
    "    x = layers.Permute((2,1,3))(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_deconv(input):\n",
    "    # Combine the convolution / deconvolution steps\n",
    "    x = conv(input)\n",
    "    \n",
    "    vx = deconv(x, \"vx\")\n",
    "    \n",
    "    # create a mask to zero out flow at (and inside) the boundary \n",
    "    vx = layers.Lambda(lambda x: x[0]*(1-x[1]), name='mask_vx')([vx, input])\n",
    "    \n",
    "     \n",
    "    ### Add decoder for vy\n",
    "    vy = deconv(x, \"vy\")\n",
    "    ### \n",
    "    \n",
    "    vy = layers.Lambda(lambda x: x[0]*(1-x[1]), name='mask_vy')([vy, input])\n",
    "    \n",
    "    output = layers.concatenate([vx, vy], axis=3)\n",
    "    \n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compile the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K.clear_session()\n",
    "\n",
    "# we need to re-init the dacaset because of Clearing our session\n",
    "training_dataset, validation_dataset, test_dataset = init_datasets()\n",
    "\n",
    "input = tf.keras.Input(shape=(128,256,1), name=\"boundary\")\n",
    "output = conv_deconv(input)\n",
    "conv_model = tf.keras.Model(inputs = input, outputs=output)\n",
    "conv_model.compile(tf.keras.optimizers.Adam(0.0001), loss=loss_image)\n",
    "conv_model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#But Training our model from scratch will take a long time\n",
    "#So we will load a partially trained model to speedup the process \n",
    "conv_model.load_weights(\"conv_model.h5\")\n",
    "\n",
    "history = conv_model.fit(training_dataset, epochs=5, steps_per_epoch=train_batches,\n",
    "          validation_data=validation_dataset, validation_steps=validation_batches, \n",
    "            callbacks=[tf.keras.callbacks.TensorBoard(log_dir='/tmp/conv')])\n",
    "\n",
    "data_utils.plot_keras_loss(history)\n",
    "# not much improvement after 20 epochs, takes 25sec/epoch on v100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test\n",
    "\n",
    "We will evaluate the model on the test dataset, and plot some of the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loss = conv_model.evaluate(test_dataset, steps=3)\n",
    "print('The loss over the test dataset', test_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = data_utils.load_test_data(1) # You can try different values between 1 and 28\n",
    "y_hat = conv_model.predict(x)\n",
    "data_utils.plot_test_result(x, y, y_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Input Data Manipulation\n",
    "\n",
    "### Use signed distance function as the input feature\n",
    "\n",
    "To improve the performance of the model, we will use a different encoding of the input data. Instead of giving 0s and 1s, we calculate the [signed distance function (SDF)](https://en.wikipedia.org/wiki/Signed_distance_function) of the input data.\n",
    "\n",
    "Let $B$ denote the set of points inside solid objects, and $\\partial B$ its boundary. We define $d(\\vec{r},\\partial B)$ as the distance between point $\\vec{r}$ and the boundary $B$.\n",
    "\n",
    "$$ d(\\vec{r}, \\partial B) = \\min_{\\vec{x} \\in \\partial B} | \\vec{r} - \\vec{x}|.$$\n",
    "\n",
    "The signed distance function is defined as\n",
    "\n",
    "$$\\mathrm{SDF}(\\vec{r}) = \\begin{cases}\n",
    "  -d(\\vec{r}, \\partial B) & \\mbox{ if } \\vec{r} \\in B \\\\\n",
    "   d(\\vec{r}, \\partial B)&  \\mbox{ if } \\vec{r} \\notin B\n",
    "\\end{cases}$$\n",
    "\n",
    "\n",
    "For every point in the grid, the SDF tells the distance to the closest boundary point. The plot below illustrates the SDF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_utils.plot_sdf(x[:,:,:], plot_boundary=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following functions will create two new input files where the SDF is added as the second channel of the input data. Let it run for a minute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_utils.create_sdf_file('train')\n",
    "data_utils.create_sdf_file('test')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load our new dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "dataset_size = 1500   # Number of elements in the train.tfrecords\n",
    "validation_size = 256 # Number of elements to use for validation\n",
    "\n",
    "# derive some quantities\n",
    "train_size = dataset_size - validation_size\n",
    "train_batches = int(train_size / batch_size)\n",
    "validation_batches= int(validation_size / batch_size)\n",
    "\n",
    "test_size = 28\n",
    "test_batches = int(test_size/batch_size)\n",
    "\n",
    "def init_sdf_datasets():\n",
    "    # Set up a dataset\n",
    "    sdf_dataset = tf.data.TFRecordDataset('data/train_sdf.tfrecords')\n",
    "    sdf_dataset = sdf_dataset.take(dataset_size)\n",
    "    # Transform binary data into image arrays\n",
    "    sdf_dataset = sdf_dataset.map(data_utils.parse_sdf_flow_data) \n",
    "\n",
    "    sdf_training_dataset = sdf_dataset.skip(validation_size).shuffle(buffer_size=512)\n",
    "    sdf_training_dataset = sdf_training_dataset.batch(batch_size, drop_remainder=True)\n",
    "    sdf_training_dataset = sdf_training_dataset.repeat()\n",
    "\n",
    "    sdf_validation_dataset = sdf_dataset.take(validation_size).batch(batch_size, drop_remainder=True)\n",
    "    sdf_validation_dataset = sdf_validation_dataset.repeat()\n",
    "\n",
    "    # Read test dataset\n",
    "    sdf_test_dataset = tf.data.TFRecordDataset('data/test_sdf.tfrecords')\n",
    "    sdf_test_dataset = sdf_test_dataset.map(data_utils.parse_sdf_flow_data) # Transform binary data into image arrays\n",
    "    sdf_test_dataset = sdf_test_dataset.batch(batch_size, drop_remainder = True).repeat()\n",
    "\n",
    "    print('Number of batches in train/validation/test dataset:', train_batches, '/', validation_batches, '/', test_batches)\n",
    "    return sdf_training_dataset,sdf_validation_dataset,sdf_test_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the the SDF for the training data : "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sdf_training_dataset, sdf_validation_dataset, sdf_test_dataset = init_sdf_datasets()\n",
    "# Create an iterator for reading a batch of input and output data\n",
    "iterator = iter(sdf_training_dataset)\n",
    "boundary, vflow = next(iterator)\n",
    "boundary, vflow = next(iterator)\n",
    "\n",
    "print('input shape:', boundary.shape.as_list())\n",
    "print('output shape:', vflow.shape.as_list())\n",
    "\n",
    "plot_idx = 2 # set it between 0 and batch_size\n",
    "\n",
    "data_utils.plot_sdf(boundary[plot_idx,:,:,0],boundary[plot_idx,:,:,1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can notice above that now our input data has 2 Channel holding boundary data and signed distance function respectively. \n",
    "\n",
    "We will use the boundary condition as a mask to remove flow lines in those areas.\n",
    "So, let's modify the the convolutional network model to use the new SDF input feature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_deconv_sdf(input):\n",
    "    # Combine the convolution / deconvolution steps\n",
    "    boundary = layers.Lambda(lambda x : x[:,:,:,0:1], name=\"boundary_slice\")(input)\n",
    "    sdf = layers.Lambda(lambda x : x[:,:,:,1:2], name=\"sdf_slice\")(input)\n",
    "    \n",
    "    \n",
    "    ### Calculate the encoding using the SDF\n",
    "    x = conv(sdf)\n",
    "    ###\n",
    "    \n",
    "    vx = deconv(x, \"vx\")\n",
    "    \n",
    "    # create a mask to zero out flow at (and inside) the boundary \n",
    "    vx = layers.Lambda(lambda x: x[0]*(1-x[1]), name='mask_vx')([vx, boundary])\n",
    "    \n",
    "    vy = deconv(x, \"vy\")\n",
    "    vy = layers.Lambda(lambda x: x[0]*(1-x[1]), name='mask_vy')([vy, boundary])\n",
    "    \n",
    "    output = layers.concatenate([vx, vy], axis=3)\n",
    "    \n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K.clear_session()\n",
    "\n",
    "# we need to re-init the dacaset because of Clearing our session\n",
    "sdf_training_dataset, sdf_validation_dataset, sdf_test_dataset = init_sdf_datasets()\n",
    "\n",
    "# Define Input Outputs and Train the Model\n",
    "input = tf.keras.Input(shape=(128,256,2), name=\"boundary\")\n",
    "output = conv_deconv_sdf(input)\n",
    "conv_sdf_model = tf.keras.Model(inputs = input, outputs=output)\n",
    "conv_sdf_model.compile(tf.keras.optimizers.Adam(0.0001), loss=loss_image)\n",
    "conv_sdf_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#But Training our model from scratch will take a long time\n",
    "#So we will load a partially trained model to speedup the process \n",
    "conv_sdf_model.load_weights(\"conv_sdf_model.h5\")\n",
    "\n",
    "history = conv_sdf_model.fit(sdf_training_dataset, epochs=5, steps_per_epoch=train_batches,\n",
    "          validation_data=sdf_validation_dataset, validation_steps=validation_batches)\n",
    "\n",
    "# Plot Training data\n",
    "data_utils.plot_keras_loss(history)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loss = conv_sdf_model.evaluate(sdf_test_dataset, steps=3)\n",
    "print('The loss over the test dataset', test_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = data_utils.load_test_data(1) # You can try different values between 1 and 28\n",
    "sdf = np.reshape(data_utils.calc_sdf(x[0,:,:,0]),(1,x.shape[1], x.shape[2],1))\n",
    "x = np.concatenate((x, sdf), axis=3)\n",
    "y_hat = conv_sdf_model.predict(x)\n",
    "data_utils.plot_test_result(x[:,:,:,0:1], y, y_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We find the Signed distance function performed better than the Boundary defined input. Let us understand why this is the case:\n",
    "\n",
    "From the Research paper : \n",
    "\n",
    "```\n",
    "Geometry can be represented in multiple ways, such as boundaries and geometric parameters. However, those               representations are not effective for neural networks since the vectors' semantic meaning varies.\n",
    "```\n",
    "\n",
    "```\n",
    "The values of SDF on the sampled Cartesian grid not only provide local geometry details but also contain additional     information of the global geometry structure.\n",
    "```\n",
    "\n",
    "To put the above in simple words, when our Convolution neural networks learn, we have seen that it also convolutes over an area of the set kernel size where it takes the considerations of the neighbouring pixels and not just a single pixel, this makes the signed distance function a rightful choice as it assigns values to all the pixels in the input image.\n",
    "\n",
    "In the upcoming notebook, let us introduce some advance networks and train them.\n",
    "\n",
    "## Important:\n",
    "<mark>Shutdown the kernel before clicking on “Next Notebook” to free up the GPU memory.</mark>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "[Previous Notebook](Part2.ipynb)\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "[1](Start_Here.ipynb)\n",
    "[2](Part2.ipynb)\n",
    "[3]\n",
    "[4](Part4.ipynb)\n",
    "[5](Competition.ipynb)\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "[Next Notebook](Part4.ipynb)\n",
    "\n",
    "\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&ensp;\n",
    "[Home Page](../Start_Here.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
