{
 "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",
    "[Previous Notebook](Start_Here.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]\n",
    "[3](Part3.ipynb)\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](Part3.ipynb)\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Steady State Flow Using Neural Networks - Part 2 \n",
    "\n",
    "**Contents of this notebook:**\n",
    "\n",
    "- [Approaching the Problem](#Approaching-the-Problem)\n",
    "- [Data and Task](#Data-and-Task)\n",
    "- [Model and Loss](#Model-and-Loss) \n",
    "- [Training and Evaluation](#Training-and-Evaluation) \n",
    "- [Building our First Model](#Building-our-First-Model)\n",
    "\n",
    "\n",
    "**By the End of this Notebook you will:**\n",
    "\n",
    "- Understand the working pipeline of DL Network.\n",
    "- Building our first model using Fully Connected networks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Approaching the Problem\n",
    "\n",
    "####  As mentioned in the Introduction to Deep Learning notebook we will follow the same steps in this notebook as well:\n",
    " \n",
    "- Data\n",
    "- Tasks\n",
    "- Model\n",
    "- Loss\n",
    "- Learning\n",
    "- Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data and Task \n",
    "\n",
    "Like we saw in the previous notebook, our input and output data are as follows :\n",
    "\n",
    "<img src=\"images/flow_example.png\">\n",
    "\n",
    "This Simulated Flow lines were calculated using Lattice Boltzmann method ([Mechsys](http://mechsys.nongnu.org/)). \n",
    "\n",
    "Let us import our dataset and see some of the input output pairs : "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing Necessary Libaries \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",
    "# reload(data_utils) # you need to execute this in case you modify the plotting scripts in data_utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input dataset is stored as a `TFRecords` file which is suitable for storing large datasets which cannot be loaded into memory and Tensorflow takes batches of data to optimize our data pipelining process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up our dataset\n",
    "dataset = tf.data.TFRecordDataset('data/train.tfrecords')\n",
    "# Transform binary data into image arrays\n",
    "dataset = dataset.map(data_utils.parse_flow_data) \n",
    "\n",
    "batched_dataset = dataset.batch(32, drop_remainder=True)\n",
    "\n",
    "# Create an iterator for reading a batch of input and output data\n",
    "iterator = iter(batched_dataset)\n",
    "boundary, vflow = next(iterator)\n",
    "\n",
    "print('Input shape:', boundary.shape.as_list())\n",
    "print('Output shape:', vflow.shape.as_list())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have set up a batch of 32 images with 128x256 resolution. The input data has only 1 channel (last dimension), which describes the boundary. The output data has 2 channels the x and y velocity component of the flow.\n",
    "\n",
    "Let us now display some of the training examples. Feel free to change `plot_idx`, and try plotting into a single figure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_idx = 10 # set it between 0-31\n",
    "\n",
    "data_utils.plot_flow_data(boundary[plot_idx,:,:,:], vflow[plot_idx,:,:,:])\n",
    "\n",
    "# You can plot the input and output data into a single figure\n",
    "data_utils.plot_flow_data(boundary[plot_idx,:,:,:], vflow[plot_idx,:,:,:], single_plot=True)\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# How a Boundary is defined ?  \n",
    "\n",
    "# Running counter over the Input \n",
    "from collections import Counter\n",
    "c = np.array(boundary[plot_idx,:,:,:]).flatten()\n",
    "c = Counter(c)\n",
    "c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the Input data, we can notice the boundary is defined by 1's and 0's, which is also reported to us by the counter above.\n",
    "\n",
    "Let us now understand how the flow lines are described."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2,figsize=(16, 4))\n",
    "\n",
    "ax1.hist(np.array(vflow[plot_idx,:,:,0]).flatten(), density=True, bins=30)\n",
    "ax1.set_ylabel('Count')\n",
    "ax1.set_xlabel('Value')\n",
    "ax1.set_title(\"X Values\")\n",
    "\n",
    "\n",
    "ax2.hist(np.array(vflow[plot_idx,:,:,1]).flatten(), density=True, bins=30)\n",
    "ax2.set_ylabel('Count')\n",
    "ax2.set_xlabel('Value')\n",
    "ax2.set_title(\"Y Values\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the X and Y Values, we can observe that each pixel denotes the velocity vector corresponding to that pixel :\n",
    "\n",
    "Positive meaning $+x$ or $+y $ direction and negative meaning the $-x$ or $-y $ direction respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Plotting Heatmap of the 2-Channels\n",
    "import matplotlib.pyplot as plt\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2,figsize=(14, 14))\n",
    "\n",
    "# Plotting X vector and Y vector separately.\n",
    "\n",
    "# X Vector \n",
    "\n",
    "ax1.imshow(vflow[plot_idx,:,:,0], cmap='hot', interpolation='nearest')\n",
    "ax1.set_title(\"X Vector\")\n",
    "\n",
    "# Y Vector\n",
    "\n",
    "ax2.imshow(-vflow[plot_idx,:,:,1], cmap='hot', interpolation='nearest')\n",
    "ax2.set_title(\"Y Vector\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we understand our input and output data, our task is now to predict the velocity vectors of both the $x$ and $y$ channels from our model.\n",
    "\n",
    "Let us now split the dataset into Training, Test and Validation Data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model and Loss\n",
    "\n",
    "We will be building the following Models and benchmarking them as we proceed :\n",
    "\n",
    "* Simple Fully Connected Networks \n",
    "    - *3 Layer Network*\n",
    "    - *5 Layer Network*\n",
    "* Convolution Neural Networks \n",
    "    - *Binary Boundary*\n",
    "    - *Signed Distance Function*\n",
    "* Advanced Networks\n",
    "    - *Gated Residual Network*\n",
    "    - *Non-Gated Residual Network*\n",
    "    \n",
    "    \n",
    "There are a variety of functions also discussed in the *Introduction to Deep Learning Notebook*. In this *Task*, we will be using Squared Error Loss.\n",
    "\n",
    "$$ \\mathrm{Loss}(\\hat{v}_x, \\hat{v}_y, v_x, v_y) = \\sum_{i=1}^{n_\\mathrm{batch}} \n",
    "  \\sum_{x=1,y=1}^{nh,nw} \\left(\\left(v_x^i(x,y) - \\hat{v}_x^i(x,y)\\right)^2 + \n",
    "   \\left(v_y^i(x,y) - \\hat{v}_y^i(x,y)\\right)^2\\right) . $$\n",
    "  \n",
    "This can be implemented easily using [tf.nn.l2_loss](https://www.tensorflow.org/api_docs/python/tf/nn/l2_loss)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training and Evaluation\n",
    "\n",
    "* Epochs: We will be using 25 - 30 epochs. Readers are free to change and play with the values. \n",
    "* Activation Function: We will start with ReLu and improving upon same\n",
    "* Optimizer: We will be using Adam Optimizer with Learning Rate of 0.0001\n",
    "* Test Set: We will be using a set of 28 boundary conditions as part of the out Test set.\n",
    "\n",
    "Now we have an idea on how to proceed. Let us start building our First Fully Connected Model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building our First Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's import our data and divide them into training, test and validation sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "dataset_size = 3000   # 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",
    "    # 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": "markdown",
   "metadata": {},
   "source": [
    "### Loss function \n",
    "\n",
    "Now let us define  Loss Function using [tf.nn.l2_loss](https://www.tensorflow.org/api_docs/python/tf/nn/l2_loss)"
   ]
  },
  {
   "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 square 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",
    "3 - Layer Fully Connected Network : \n",
    "\n",
    "* Input Layer of Size ( 128 * 256 * 1 )\n",
    "* Hidden Layer of Size ( 256 ) \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 layer\n",
    "    x = layers.Flatten()(input)\n",
    "    x = layers.Dense(256, activation='relu')(x)\n",
    "    \n",
    "    \n",
    "    ### Define output layer and reshape it to nh x nw x 2. \n",
    "    ### (Note that the extra batch dimension is handled automatically by Keras)\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": "markdown",
   "metadata": {},
   "source": [
    "Let us define the input and parameters of our model : "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Inputs and Outputs\n",
    "input = tf.keras.Input(shape=(128,256,1))\n",
    "output = fully_connected(input)\n",
    "# Use Keras Functional API to Create our Model and Define our Optimizer and Loss Function\n",
    "fc_model = tf.keras.Model(inputs = input, outputs=output)\n",
    "fc_model.compile(tf.keras.optimizers.Adam(0.0001), loss=loss_image)\n",
    "fc_model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let us now train our Model : "
   ]
  },
  {
   "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)"
   ]
  },
  {
   "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": [
    "In the upcoming notebook let us define a 5 Layer fully connected network and train it.\n",
    "\n",
    "## Important:\n",
    "<mark>Shutdown the kernel before clicking on “Next Notebook” to free up the GPU memory.</mark>\n",
    "\n",
    "\n",
    "## Licensing\n",
    "This material is released by NVIDIA Corporation under the Creative Commons Attribution 4.0 International (CC BY 4.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Previous Notebook](Start_Here.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]\n",
    "[3](Part3.ipynb)\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](Part3.ipynb)\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)"
   ]
  }
 ],
 "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
}
