PINNs_NavierStokes_HFM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7758f28f",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "    <h1> Tutorial 2 </h1> \n",
    "    <h2> Physics Informed Neural Networks Part 4</h2>\n",
    "    <h2> Navier Stokes Hidden Fluid Mechanics Example </h2>\n",
    "</div>    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cf2aeb8",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "This notebook is based on two papers: *[Physics-Informed Neural Networks:  A Deep Learning Framework for Solving Forward and Inverse Problems Involving Nonlinear Partial Differential Equations](https://www.sciencedirect.com/science/article/pii/S0021999118307125)* and *[Hidden Physics Models:  Machine Learning of Nonlinear Partial Differential Equations](https://www.sciencedirect.com/science/article/pii/S0021999117309014)* with the help of  Fergus Shone and Michael Macraild.\n",
    "\n",
    "These tutorials will go through solving Partial Differential Equations using Physics Informed Neural Networks, focusing on the Burgers Equation and a more complex example using the Navier Stokes Equations.\n",
    "\n",
    "**This introduction section is replicated in all PINN tutorial notebooks (please skip if you've already been through).** "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ccff06d",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "If you have not already then in your repository directory please run the following code. Via the terminal (mac or linux) or gitbash (windows)\n",
    "    \n",
    "```bash\n",
    "git submodule init\n",
    "git submodule update --init --recursive\n",
    "```\n",
    "\n",
    "**If this does not work please clone the [PINNs](https://github.com/maziarraissi/PINNs) repository into your Physics_Informed_Neural_Networks folder.**\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10361a00",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "<h1>Physics Informed Neural Networks</h1>\n",
    "\n",
    "For a typical Neural Network using algorithms like gradient descent to look for a hypothesis, data is the only guide. However, if the data are noisy or sparse and we already have governing physical models, we can use the knowledge we already know to optimize and inform the algorithms. This can be done via [feature engineering](https://en.wikipedia.org/wiki/Feature_engineering) or by adding a physical inconsistency term to the loss function.\n",
    "\n",
    "\n",
    "<a href=\"https://towardsdatascience.com/physics-guided-neural-networks-pgnns-8fe9dbad9414\">\n",
    "<img src=\"https://miro.medium.com/max/700/1*uM2Qh4PFQLWLLI_KHbgaVw.png\">\n",
    "</a>   \n",
    " \n",
    " \n",
    "## The very basics\n",
    "\n",
    "If you are new to neural networks there is a [toy neural network python code example](https://github.com/cemac/LIFD_ENV_ML_NOTEBOOKS/tree/main/ToyNeuralNetwork) included in the [LIFD ENV ML Notebooks Repository]( https://github.com/cemac/LIFD_ENV_ML_NOTEBOOKS). There we build a two-layer neural network to illustrate the fundamentals of how Neural Networks work, and give equivalent code using the Python machine learning library [TensorFlow](https://keras.io/).\n",
    "\n",
    "    \n",
    "## Recommended reading\n",
    "    \n",
    "The in-depth theory behind neural networks will not be covered here as this tutorial is focusing on application of machine learning methods. If you wish to learn more, here are some great starting points.  \n",
    " \n",
    "\n",
    "* [All you need to know on Neural networks](https://towardsdatascience.com/nns-aynk-c34efe37f15a) \n",
    "* [Introduction to Neural Networks](https://victorzhou.com/blog/intro-to-neural-networks/)\n",
    "* [Physics Guided Neural Networks](https://towardsdatascience.com/physics-guided-neural-networks-pgnns-8fe9dbad9414)\n",
    "* [Maziar Rassi's Physics informed GitHub web Page](https://maziarraissi.github.io/PINNs/)\n",
    "\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbf6505b",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Hidden Fluid Mechanics #\n",
    "\n",
    "In this notebook, we will utilise a more advanced implementation of PINNs, taken from Maziar Raissi's paper Hidden Fluid Mechanics. In many fluid flow scenarios, direct measurement of variables such as velocity and pressure is not possible. However, we may have access to measurements of some passive scalar field (such as smoke concentration in wind tunnel testing). This work aims to use PINNs to uncover hidden velocity and pressure fields for flow problems, utilising data drawn only from measurements of some passive scalar, $c(t,x,y)$. This scalar is governed by the transport equation:\n",
    "\n",
    "\\begin{equation}\n",
    "c_t + u c_x + v c_y = \\text{Pec}^{-1} \\left(c_{xx} + c_{yy}\\right)\n",
    "\\end{equation}\n",
    "\n",
    "where $\\mathrm{Pec}$ is the Peclet number, and $u, v$ are the two velocity components. Then, the Navier-Stokes equations are given by:\n",
    "\n",
    "\\begin{equation}    \n",
    "u_t + uu_x + vu_y = -p_x +\\mathrm{Re}^{-1}(u_{xx} + u_{yy}),\n",
    "\\\\\n",
    "v_t + uv_x + vv_y = -p_y +\\mathrm{Re}^{-1}(v_{xx} + v_{yy}),\n",
    "\\\\\n",
    "u_x + v_y = 0\n",
    "\\end{equation}\n",
    "\n",
    "for pressure $p$ and Reynolds number $\\mathrm{Re}$. Note that the only constraint on the velocity and pressure predictions is that they satisfy the above equations.\n",
    "\n",
    "The structure of the neural network is seen below:\n",
    "\n",
    "![](../images/Network.png)\n",
    "\n",
    "The network has four inputs, as expected, and six outputs, namely the three velocity components, the pressure, the concentration, c, and one final variable, d. d(t,x,y,z) is an 'auxiliary variable', defined to be the complement of c (i.e. d = 1 - c), and is governed by the same transport equation as c. Its inclusion improves prediction accuracy, and helps in detecting boundary locations.\n",
    "\n",
    "The data for this notebook was generated by Raissi et al. using a spectral element solver called NekTar. The Navier-Stokes and transport equations are numerically approximated to a high degree of accuracy.\n",
    "\n",
    "The fluid problem at hand is 2D channel flow over an obstacle. A crude diagram of the flow domain can be seen below:\n",
    "![](../images/hfminfo2.png)\n",
    "\n",
    "It may be useful to read the original paper alongside this notebook, which can be found at https://arxiv.org/pdf/1808.04327.pdf. The source code can be accessed via this Github repository: https://github.com/maziarraissi/HFM.\n",
    "\n",
    "This script was orignally written by Maziar Raissi et al. in their PINNs repository, before being modified by Michael MacRaild and Fergus Shone. All figures are from https://arxiv.org/pdf/1808.04327.pdf.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e242716",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "<h1> Python </h1>\n",
    "\n",
    "    \n",
    "## TensorFlow \n",
    "    \n",
    "There are many machine learning Python libraries available. [TensorFlow](https://www.tensorflow.org/) is one such library. If you have GPUs on your machine, TensorFlow will automatically use them and run the code even faster!\n",
    "\n",
    "## Further Reading\n",
    "\n",
    "* [Running Jupyter Notebooks](https://jupyter.readthedocs.io/en/latest/running.html#running)\n",
    "* [TensorFlow optimizers](https://www.tutorialspoint.com/tensorflow/tensorflow_optimizers.htm)\n",
    "\n",
    "</div>\n",
    "    \n",
    "<hr>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a747281",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ffffcc; padding: 10px;\">\n",
    "    \n",
    "<h1> Requirements </h1>\n",
    "\n",
    "These notebooks should run with the following requirements satisfied.\n",
    "\n",
    "<h2> Python Packages: </h2>\n",
    "\n",
    "* Python 3\n",
    "* TensorFlow > 2\n",
    "* numpy \n",
    "* matplotlib\n",
    "* scipy\n",
    "\n",
    "<h2> Data Requirements</h2>\n",
    "    \n",
    "This notebook refers to some data included in the GitHub repository, imported via the git submodules command mentioned in the installation instructions.\n",
    "    \n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0cb13b6",
   "metadata": {},
   "source": [
    "**Contents:**\n",
    "\n",
    "1. [1D Heat Equation Non ML Example](PINNs_1DHeatEquations_nonML.ipynb)\n",
    "2. [1D Equation PINN Example](PINNs_1DHeatEquationExample.ipynb)\n",
    "3. [Navier-Stokes PINNs discovery of PDE’s](PINNs_NavierStokes_example.ipynb)\n",
    "3. **[Navier-Stokes PINNs Hidden Fluid Mechanics](PINNs_Navier_Stokes_HFM.ipynb)**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6db561a7",
   "metadata": {},
   "source": [
    "<hr>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "700a7b60",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "Load in all required modules (including some auxiliary code) and turn off warnings. Make sure Keras session is clear.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9bff273",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For readability: disable warnings\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da61a6e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import requests\n",
    "import zipfile\n",
    "sys.path.insert(0, 'PINNs/Utilities/')\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.io\n",
    "from scipy.interpolate import griddata\n",
    "from itertools import product, combinations\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from mpl_toolkits.mplot3d.art3d import Poly3DCollection\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "import matplotlib.gridspec as gridspec\n",
    "from time import time\n",
    "import scipy.sparse as sp\n",
    "import scipy.sparse.linalg as la"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f9a8de",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Define PINN Functions\n",
    "    \n",
    "In this section we define a number of functions to be called later in the script, such as error metrics, the general neural network class etc. In the original github repository, these functions are defined in a separate script called 'utilities.py', which is called in each example script. To keep this notebook contained, we include them here instead.\n",
    "\n",
    "The basic functionality of this code is the same as the previous examples, but some of the syntax is updated to streamline the code. \n",
    " \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f87e7414",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate tensorflow session and initialise variables\n",
    "def tf_session():\n",
    "    # tf session\n",
    "    config = tf.compat.v1.ConfigProto(allow_soft_placement=True,\n",
    "                            log_device_placement=True)\n",
    "    config.gpu_options.force_gpu_compatible = True\n",
    "    sess = tf.compat.v1.Session(config=config)\n",
    "    \n",
    "    # init\n",
    "    init = tf.compat.v1.global_variables_initializer()\n",
    "    sess.run(init)\n",
    "    \n",
    "    return sess\n",
    "\n",
    "# produce relative error between prediction and ground truth\n",
    "def relative_error(pred, exact):\n",
    "    if type(pred) is np.ndarray:\n",
    "        return np.sqrt(np.mean(np.square(pred - exact))/np.mean(np.square(exact - np.mean(exact))))\n",
    "    return tf.sqrt(tf.reduce_mean(tf.square(pred - exact))/tf.reduce_mean(tf.square(exact - tf.reduce_mean(exact))))\n",
    "\n",
    "# produce MSE of predicition and ground truth\n",
    "def mean_squared_error(pred, exact):\n",
    "    if type(pred) is np.ndarray:\n",
    "        return np.mean(np.square(pred - exact))\n",
    "    return tf.reduce_mean(tf.square(pred - exact))\n",
    "\n",
    "# generate gradient of Y w.r.t. x using automatic differentiation\n",
    "def fwd_gradients(Y, x):\n",
    "    dummy = tf.ones_like(Y)\n",
    "    G = tf.compat.v1.gradients(Y, x, grad_ys=dummy, colocate_gradients_with_ops=True)[0]\n",
    "    Y_x = tf.compat.v1.gradients(G, dummy, colocate_gradients_with_ops=True)[0]\n",
    "    return Y_x\n",
    "\n",
    "# calculate 2D velocity gradients\n",
    "def Gradient_Velocity_2D(u, v, x, y):\n",
    "    \n",
    "    Y = tf.concat([u, v], 1)\n",
    "    \n",
    "    Y_x = fwd_gradients(Y, x)\n",
    "    Y_y = fwd_gradients(Y, y)\n",
    "    \n",
    "    u_x = Y_x[:,0:1]\n",
    "    v_x = Y_x[:,1:2]\n",
    "    \n",
    "    u_y = Y_y[:,0:1]\n",
    "    v_y = Y_y[:,1:2]\n",
    "    \n",
    "    return [u_x, v_x, u_y, v_y]\n",
    "\n",
    "# basic neural network class. This function constructs the network architecture, \n",
    "# normalises the inputs using the mean and standard deviation and specifies the\n",
    "# activation function, which is chosen here to be sigmoid.\n",
    "class neural_net(object):\n",
    "    def __init__(self, *inputs, layers):\n",
    "        \n",
    "        self.layers = layers\n",
    "        self.num_layers = len(self.layers)\n",
    "        \n",
    "        if len(inputs) == 0:\n",
    "            in_dim = self.layers[0]\n",
    "            self.X_mean = np.zeros([1, in_dim])\n",
    "            self.X_std = np.ones([1, in_dim])\n",
    "        else:\n",
    "            X = np.concatenate(inputs, 1)\n",
    "            self.X_mean = X.mean(0, keepdims=True)\n",
    "            self.X_std = X.std(0, keepdims=True)\n",
    "        \n",
    "        self.weights = []\n",
    "        self.biases = []\n",
    "        self.gammas = []\n",
    "        \n",
    "        for l in range(0,self.num_layers-1):\n",
    "            in_dim = self.layers[l]\n",
    "            out_dim = self.layers[l+1]\n",
    "            W = np.random.normal(size=[in_dim, out_dim])\n",
    "            b = np.zeros([1, out_dim])\n",
    "            g = np.ones([1, out_dim])\n",
    "            # tensorflow variables\n",
    "            self.weights.append(tf.Variable(W, dtype=tf.float32, trainable=True))\n",
    "            self.biases.append(tf.Variable(b, dtype=tf.float32, trainable=True))\n",
    "            self.gammas.append(tf.Variable(g, dtype=tf.float32, trainable=True))\n",
    "            \n",
    "    def __call__(self, *inputs):\n",
    "                \n",
    "        H = (tf.concat(inputs, 1) - self.X_mean)/self.X_std\n",
    "    \n",
    "        for l in range(0, self.num_layers-1):\n",
    "            W = self.weights[l]\n",
    "            b = self.biases[l]\n",
    "            g = self.gammas[l]\n",
    "            # weight normalization\n",
    "            V = W/tf.norm(W, axis = 0, keepdims=True)\n",
    "            # matrix multiplication\n",
    "            H = tf.matmul(H, V)\n",
    "            # add bias\n",
    "            H = g*H + b\n",
    "            # activation\n",
    "            if l < self.num_layers-2:\n",
    "                H = H*tf.sigmoid(H)\n",
    "                \n",
    "        Y = tf.split(H, num_or_size_splits=H.shape[1], axis=1)\n",
    "    \n",
    "        return Y\n",
    "\n",
    "# produce the necessary output gradients and then form the Navier-Stokes residuals for 2D flow\n",
    "def Navier_Stokes_2D(c, u, v, p, t, x, y, Pec, Rey):\n",
    "    \n",
    "    Y = tf.concat([c, u, v, p], 1)\n",
    "    \n",
    "    Y_t = fwd_gradients(Y, t)\n",
    "    Y_x = fwd_gradients(Y, x)\n",
    "    Y_y = fwd_gradients(Y, y)\n",
    "    Y_xx = fwd_gradients(Y_x, x)\n",
    "    Y_yy = fwd_gradients(Y_y, y)\n",
    "    \n",
    "    c = Y[:,0:1]\n",
    "    u = Y[:,1:2]\n",
    "    v = Y[:,2:3]\n",
    "    p = Y[:,3:4]\n",
    "    \n",
    "    c_t = Y_t[:,0:1]\n",
    "    u_t = Y_t[:,1:2]\n",
    "    v_t = Y_t[:,2:3]\n",
    "    \n",
    "    c_x = Y_x[:,0:1]\n",
    "    u_x = Y_x[:,1:2]\n",
    "    v_x = Y_x[:,2:3]\n",
    "    p_x = Y_x[:,3:4]\n",
    "    \n",
    "    c_y = Y_y[:,0:1]\n",
    "    u_y = Y_y[:,1:2]\n",
    "    v_y = Y_y[:,2:3]\n",
    "    p_y = Y_y[:,3:4]\n",
    "    \n",
    "    c_xx = Y_xx[:,0:1]\n",
    "    u_xx = Y_xx[:,1:2]\n",
    "    v_xx = Y_xx[:,2:3]\n",
    "    \n",
    "    c_yy = Y_yy[:,0:1]\n",
    "    u_yy = Y_yy[:,1:2]\n",
    "    v_yy = Y_yy[:,2:3]\n",
    "    \n",
    "    e1 = c_t + (u*c_x + v*c_y) - (1.0/Pec)*(c_xx + c_yy)\n",
    "    e2 = u_t + (u*u_x + v*u_y) + p_x - (1.0/Rey)*(u_xx + u_yy) \n",
    "    e3 = v_t + (u*v_x + v*v_y) + p_y - (1.0/Rey)*(v_xx + v_yy)\n",
    "    e4 = u_x + v_y\n",
    "    \n",
    "    return e1, e2, e3, e4\n",
    "\n",
    "# calculate the strain rate for generating results later in the script\n",
    "def Strain_Rate_2D(u, v, x, y):\n",
    "    \n",
    "    [u_x, v_x, u_y, v_y] = Gradient_Velocity_2D(u, v, x, y)\n",
    "    \n",
    "    eps11dot = u_x\n",
    "    eps12dot = 0.5*(v_x + u_y)\n",
    "    eps22dot = v_y\n",
    "    \n",
    "    return [eps11dot, eps12dot, eps22dot]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08530832",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Define PINN (HFM) Class\n",
    "    \n",
    "Here the main PINN class is defined. It is shorter than the PINN class in the previous notebooks as many of the functions that were previously included inside are defined externally in the previous section, and called within this class. There is also a change in the way the network is trained, with the implementation of batch training (or mini-batch training). This is useful when very large amounts of data are used, as it divides the entire training set into subsets called batches, avoiding issues with exceeding data allocation. Also, when implemented correctly, batch training can be more efficient. It is worth reading up on batch training in general, as it is a useful machine learning approach. \n",
    "   \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a61a3a37",
   "metadata": {},
   "outputs": [],
   "source": [
    "class HFM(object):\n",
    "    # notational conventions\n",
    "    # _tf: placeholders for input/output data and points used to regress the equations\n",
    "    # _pred: output of neural network\n",
    "    # _eqns: points used to regress the equations\n",
    "    # _data: input-output data\n",
    "    # _star: preditions\n",
    "\n",
    "    def __init__(self, t_data, x_data, y_data, c_data,\n",
    "                       t_eqns, x_eqns, y_eqns,\n",
    "                       layers, batch_size):\n",
    "        \n",
    "        # specs\n",
    "        self.layers = layers\n",
    "        self.batch_size = batch_size\n",
    "        \n",
    "        # flow properties\n",
    "        self.Pec = tf.Variable(15.0, dtype=tf.float32, trainable = True)\n",
    "        self.Rey = tf.Variable(5.0, dtype=tf.float32, trainable = True)\n",
    "                \n",
    "        # data\n",
    "        [self.t_data, self.x_data, self.y_data, self.c_data] = [t_data, x_data, y_data, c_data]\n",
    "        [self.t_eqns, self.x_eqns, self.y_eqns] = [t_eqns, x_eqns, y_eqns]\n",
    "        \n",
    "        # placeholders\n",
    "        [self.t_data_tf, self.x_data_tf, self.y_data_tf, self.c_data_tf] = [tf.compat.v1.placeholder(tf.float32, shape=[None, 1]) for _ in range(4)]\n",
    "        [self.t_eqns_tf, self.x_eqns_tf, self.y_eqns_tf] = [tf.compat.v1.placeholder(tf.float32, shape=[None, 1]) for _ in range(3)]\n",
    "        \n",
    "        # physics \"uninformed\" neural networks\n",
    "        self.net_cuvp = neural_net(self.t_data, self.x_data, self.y_data, layers = self.layers)\n",
    "        \n",
    "        [self.c_data_pred,\n",
    "         self.u_data_pred,\n",
    "         self.v_data_pred,\n",
    "         self.p_data_pred] = self.net_cuvp(self.t_data_tf,\n",
    "                                           self.x_data_tf,\n",
    "                                           self.y_data_tf)\n",
    "         \n",
    "        # physics \"informed\" neural networks\n",
    "        [self.c_eqns_pred,\n",
    "         self.u_eqns_pred,\n",
    "         self.v_eqns_pred,\n",
    "         self.p_eqns_pred] = self.net_cuvp(self.t_eqns_tf,\n",
    "                                           self.x_eqns_tf,\n",
    "                                           self.y_eqns_tf)\n",
    "        \n",
    "        [self.e1_eqns_pred,\n",
    "         self.e2_eqns_pred,\n",
    "         self.e3_eqns_pred,\n",
    "         self.e4_eqns_pred] = Navier_Stokes_2D(self.c_eqns_pred,\n",
    "                                               self.u_eqns_pred,\n",
    "                                               self.v_eqns_pred,\n",
    "                                               self.p_eqns_pred,\n",
    "                                               self.t_eqns_tf,\n",
    "                                               self.x_eqns_tf,\n",
    "                                               self.y_eqns_tf,\n",
    "                                               self.Pec,\n",
    "                                               self.Rey)\n",
    "        \n",
    "        [self.eps11dot_eqns_pred,\n",
    "         self.eps12dot_eqns_pred,\n",
    "         self.eps22dot_eqns_pred] = Strain_Rate_2D(self.u_eqns_pred,\n",
    "                                                   self.v_eqns_pred,\n",
    "                                                   self.x_eqns_tf,\n",
    "                                                   self.y_eqns_tf)\n",
    "        \n",
    "        # loss\n",
    "        self.loss = mean_squared_error(self.c_data_pred, self.c_data_tf) + \\\n",
    "                    mean_squared_error(self.e1_eqns_pred, 0.0) + \\\n",
    "                    mean_squared_error(self.e2_eqns_pred, 0.0) + \\\n",
    "                    mean_squared_error(self.e3_eqns_pred, 0.0) + \\\n",
    "                    mean_squared_error(self.e4_eqns_pred, 0.0)\n",
    "        \n",
    "        # optimizers\n",
    "        self.learning_rate = tf.compat.v1.placeholder(tf.float32, shape=[])\n",
    "        self.optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate = self.learning_rate)\n",
    "        self.train_op = self.optimizer.minimize(self.loss)\n",
    "        \n",
    "        self.sess = tf_session()\n",
    "        \n",
    "        # Add ops to save and restore all the variables.\n",
    "        self.netSaveDir = \"modelckpts/HFM/\"\n",
    "        self.saver = tf.compat.v1.train.Saver()\n",
    "\n",
    "    def train(self, total_time, learning_rate):\n",
    "        \n",
    "        N_data = self.t_data.shape[0]\n",
    "        N_eqns = self.t_eqns.shape[0]\n",
    "        \n",
    "        start_time = time()\n",
    "        running_time = 0\n",
    "        it = 0\n",
    "        while running_time < total_time:\n",
    "            \n",
    "            idx_data = np.random.choice(N_data, min(self.batch_size, N_data))\n",
    "            idx_eqns = np.random.choice(N_eqns, self.batch_size)\n",
    "            \n",
    "            (t_data_batch,\n",
    "             x_data_batch,\n",
    "             y_data_batch,\n",
    "             c_data_batch) = (self.t_data[idx_data,:],\n",
    "                              self.x_data[idx_data,:],\n",
    "                              self.y_data[idx_data,:],\n",
    "                              self.c_data[idx_data,:])\n",
    "\n",
    "            (t_eqns_batch,\n",
    "             x_eqns_batch,\n",
    "             y_eqns_batch) = (self.t_eqns[idx_eqns,:],\n",
    "                              self.x_eqns[idx_eqns,:],\n",
    "                              self.y_eqns[idx_eqns,:])\n",
    "\n",
    "\n",
    "            tf_dict = {self.t_data_tf: t_data_batch,\n",
    "                       self.x_data_tf: x_data_batch,\n",
    "                       self.y_data_tf: y_data_batch,\n",
    "                       self.c_data_tf: c_data_batch,\n",
    "                       self.t_eqns_tf: t_eqns_batch,\n",
    "                       self.x_eqns_tf: x_eqns_batch,\n",
    "                       self.y_eqns_tf: y_eqns_batch,\n",
    "                       self.learning_rate: learning_rate}\n",
    "            \n",
    "            self.sess.run([self.train_op], tf_dict)\n",
    "            \n",
    "            # Print\n",
    "            if it % 10 == 0:\n",
    "                elapsed = time() - start_time\n",
    "                running_time += elapsed/3600.0\n",
    "                [loss_value,\n",
    "                 Pec_value,\n",
    "                 Rey_value,\n",
    "                 learning_rate_value] = self.sess.run([self.loss,\n",
    "                                                       self.Pec,\n",
    "                                                       self.Rey,\n",
    "                                                       self.learning_rate], tf_dict)\n",
    "                print('It: %d, Loss: %.3e, Pec: %.3f, Rey: %.3f, Time: %.2fs, Running Time: %.2fh, Learning Rate: %.1e'\n",
    "                      %(it, loss_value, Pec_value, Rey_value, elapsed, running_time, learning_rate_value))\n",
    "                sys.stdout.flush()\n",
    "                start_time = time()\n",
    "            it += 1\n",
    "\n",
    "            if it % 1000 == 0:\n",
    "                save_path = self.saver.save(self.sess, self.netSaveDir + 'model_at_iter%s.ckpt'%(it))\n",
    "                print('Model saved in path: %s' % save_path)\n",
    "        save_path = self.saver.save(self.sess, self.netSaveDir + 'model.ckpt')\n",
    "        print('Model saved in path: %s' % save_path)\n",
    "    \n",
    "    def predict(self, t_star, x_star, y_star):\n",
    "        tf_dict = {self.t_data_tf: t_star, self.x_data_tf: x_star, self.y_data_tf: y_star}\n",
    "        \n",
    "        c_star = self.sess.run(self.c_data_pred, tf_dict)\n",
    "        u_star = self.sess.run(self.u_data_pred, tf_dict)\n",
    "        v_star = self.sess.run(self.v_data_pred, tf_dict)\n",
    "        p_star = self.sess.run(self.p_data_pred, tf_dict)\n",
    "        \n",
    "        return c_star, u_star, v_star, p_star\n",
    "    \n",
    "    def predict_eps_dot(self, t_star, x_star, y_star):\n",
    "                \n",
    "        tf_dict = {self.t_eqns_tf: t_star, self.x_eqns_tf: x_star, self.y_eqns_tf: y_star}\n",
    "        \n",
    "        eps11dot_star = self.sess.run(self.eps11dot_eqns_pred, tf_dict)\n",
    "        eps12dot_star = self.sess.run(self.eps12dot_eqns_pred, tf_dict)\n",
    "        eps22dot_star = self.sess.run(self.eps22dot_eqns_pred, tf_dict)\n",
    "        \n",
    "        return eps11dot_star, eps12dot_star, eps22dot_star"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0ec214f",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Main Code\n",
    "    \n",
    "In this section we run each of our previously defined functions and classes to train the network and generate results. We specify the batch size, network width and depth, process the training and test data, train the network and output error results. Note that the 'train' function is commented out, so if you would like to re-train the network feel free to uncomment that line (although be warned: it will overwrite the previously trained weights). \n",
    "\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2fa6010",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "    \n",
    "# If you have not already downloaded the data required the following cell will download the additional data for you \n",
    "    \n",
    "may take a minute or two depending on your internet speed\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aafe7d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "if [ ! -e Data/Stenosis2D.mat ];\n",
    "    then\n",
    "    echo 'grabbing addtional data'\n",
    "    wget https://gitlab.com/CEMACHELEN1/lifd_ml_large_file_store/-/archive/main/lifd_ml_large_file_store-main.zip\n",
    "    mv lifd_ml_large_file_store-main.zip Data\n",
    "    cd Data\n",
    "    unzip lifd_ml_large_file_store-main.zip\n",
    "    rm -f lifd_ml_large_file_store-main.zip\n",
    "    cd ../\n",
    "    mv Data/lifd_ml_large_file_store-main/Physics_Informed_Neural_Networks/Data/Stenosis2D.mat Data/\n",
    "    rm -rf Data/lifd_ml_large_file_store-main\n",
    "    ls Data\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16acb23e-60d5-4fe7-83e6-5330ce71fd45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if the file 'Data/Stenosis2D.mat' exists\n",
    "if not os.path.exists('Data/Stenosis2D.mat'):\n",
    "    print(\"Grabbing additional data...\")\n",
    "    \n",
    "    # Download the ZIP file\n",
    "    url = \"https://gitlab.com/CEMACHELEN1/lifd_ml_large_file_store/-/archive/main/lifd_ml_large_file_store-main.zip\"\n",
    "    response = requests.get(url)\n",
    "    \n",
    "    # Save the ZIP file\n",
    "    with open(\"lifd_ml_large_file_store-main.zip\", \"wb\") as zip_file:\n",
    "        zip_file.write(response.content)\n",
    "    \n",
    "    # Extract the contents\n",
    "    with zipfile.ZipFile(\"lifd_ml_large_file_store-main.zip\", \"r\") as zip_ref:\n",
    "        zip_ref.extractall(\"Data\")\n",
    "    \n",
    "    # Clean up\n",
    "    os.remove(\"lifd_ml_large_file_store-main.zip\")\n",
    "    \n",
    "    # Move the required file\n",
    "    os.rename(\"Data/lifd_ml_large_file_store-main/Physics_Informed_Neural_Networks/Data/Stenosis2D.mat\", \"Data/Stenosis2D.mat\")\n",
    "    \n",
    "    # List the contents of the 'Data' directory\n",
    "    print(\"Contents of 'Data':\")\n",
    "    for item in os.listdir(\"Data\"):\n",
    "        print(item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebf784d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This line of code is required to prevent some tensorflow errors arrising from the\n",
    "# inclusion of some tensorflw v 1 code \n",
    "tf.compat.v1.disable_eager_execution()\n",
    "\n",
    "batch_size = 10000\n",
    "\n",
    "layers = [3] + 10*[4*50] + [4]\n",
    "\n",
    "# Load Data\n",
    "data = scipy.io.loadmat('Data/Stenosis2D.mat')\n",
    "\n",
    "t_star = data['t_star'] # T x 1\n",
    "x_star = data['x_star'] # N x 1\n",
    "y_star = data['y_star'] # N x 1\n",
    "\n",
    "T = t_star.shape[0]\n",
    "N = x_star.shape[0]\n",
    "\n",
    "U_star = data['U_star'] # N x T\n",
    "V_star = data['V_star'] # N x T\n",
    "P_star = data['P_star'] # N x T\n",
    "C_star = data['C_star'] # N x T    \n",
    "\n",
    "# Rearrange Data \n",
    "T_star = np.tile(t_star, (1,N)).T # N x T\n",
    "X_star = np.tile(x_star, (1,T)) # N x T\n",
    "Y_star = np.tile(y_star, (1,T)) # N x T    \n",
    "\n",
    "######################################################################\n",
    "######################## Noiseless Data ##############################\n",
    "######################################################################\n",
    "\n",
    "T_data = T # int(sys.argv[1])\n",
    "N_data = N # int(sys.argv[2])\n",
    "idx_t = np.concatenate([np.array([0]), np.random.choice(T-2, T_data-2, replace=False)+1, np.array([T-1])] )\n",
    "idx_x = np.random.choice(N, N_data, replace=False)\n",
    "t_data = T_star[:, idx_t][idx_x,:].flatten()[:,None]\n",
    "x_data = X_star[:, idx_t][idx_x,:].flatten()[:,None]\n",
    "y_data = Y_star[:, idx_t][idx_x,:].flatten()[:,None]\n",
    "c_data = C_star[:, idx_t][idx_x,:].flatten()[:,None]\n",
    "\n",
    "T_eqns = T\n",
    "N_eqns = N\n",
    "idx_t = np.concatenate([np.array([0]), np.random.choice(T-2, T_eqns-2, replace=False)+1, np.array([T-1])] )\n",
    "idx_x = np.random.choice(N, N_eqns, replace=False)\n",
    "t_eqns = T_star[:, idx_t][idx_x,:].flatten()[:,None]\n",
    "x_eqns = X_star[:, idx_t][idx_x,:].flatten()[:,None]\n",
    "y_eqns = Y_star[:, idx_t][idx_x,:].flatten()[:,None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3accbf2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Training\n",
    "model = HFM(t_data, x_data, y_data, c_data,\n",
    "            t_eqns, x_eqns, y_eqns,\n",
    "            layers, batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52598941",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "  \n",
    "# Loading Pre trained model option \n",
    "    \n",
    "If the training time is too slow you can skip the following line and load in a pretrained model instead set `loadweights = True` in the next cell. You can play around with different number of iterations to see the effects e.g. setting `saver.restore(sess, netSaveDir + 'model_at_iter3000.ckpt')`\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1a889ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.train(total_time = 40, learning_rate=1e-3) # UNCOMMENT THIS LINE IF YOU WOULD LIKE TO RE-TRAIN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb659075",
   "metadata": {},
   "outputs": [],
   "source": [
    "loadweights = True\n",
    "if loadweights:\n",
    "    print(\"loading pre trained model\")\n",
    "    netSaveDir = \"modelckpts/HFM/\"\n",
    "    saver = tf.compat.v1.train.Saver()\n",
    "    saver.restore(model.sess, netSaveDir + 'model_at_iter3000.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "970ca6fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "Shear = np.zeros((300,t_star.shape[0]))\n",
    "\n",
    "for snap in range(0,t_star.shape[0]):\n",
    "\n",
    "    x1_shear = np.linspace(15,25,100)[:,None]\n",
    "    x2_shear = np.linspace(25,35,100)[:,None]\n",
    "    x3_shear = np.linspace(35,55,100)[:,None]\n",
    "\n",
    "    x_shear = np.concatenate([x1_shear,x2_shear,x3_shear], axis=0)\n",
    "\n",
    "    y1_shear = 0.0*x1_shear\n",
    "    y2_shear = np.sqrt(25.0 - (x2_shear - 30.0)**2)\n",
    "    y3_shear = 0.0*x3_shear\n",
    "\n",
    "    y_shear = np.concatenate([y1_shear,y2_shear,y3_shear], axis=0)\n",
    "\n",
    "    t_shear = T_star[0,snap] + 0.0*x_shear\n",
    "\n",
    "    eps11_dot_shear, eps12_dot_shear, eps22_dot_shear = model.predict_eps_dot(t_shear, x_shear, y_shear)\n",
    "\n",
    "    nx1_shear = 0.0*x1_shear\n",
    "    nx2_shear = 6.0 - x2_shear/5.0\n",
    "    nx3_shear = 0.0*x3_shear\n",
    "\n",
    "    nx_shear = np.concatenate([nx1_shear,nx2_shear,nx3_shear], axis=0)\n",
    "\n",
    "    ny1_shear = -1.0 + 0.0*y1_shear\n",
    "    ny2_shear = -y2_shear/5.0\n",
    "    ny3_shear = -1.0 + 0.0*y3_shear\n",
    "\n",
    "    ny_shear = np.concatenate([ny1_shear,ny2_shear,ny3_shear], axis=0)\n",
    "\n",
    "    shear_x = 2.0*(1.0/5.0)*(eps11_dot_shear*nx_shear + eps12_dot_shear*ny_shear)\n",
    "    shear_y = 2.0*(1.0/5.0)*(eps12_dot_shear*nx_shear + eps22_dot_shear*ny_shear)\n",
    "\n",
    "    shear = np.sqrt(shear_x**2 + shear_y**2)\n",
    "\n",
    "    Shear[:,snap] = shear.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "026b84c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test Data\n",
    "snap = np.array([55])\n",
    "t_test = T_star[:,snap]\n",
    "x_test = X_star[:,snap]\n",
    "y_test = Y_star[:,snap]\n",
    "\n",
    "c_test = C_star[:,snap]\n",
    "u_test = U_star[:,snap]\n",
    "v_test = V_star[:,snap]\n",
    "p_test = P_star[:,snap]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c270c3fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prediction\n",
    "c_pred, u_pred, v_pred, p_pred = model.predict(t_test, x_test, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9350e429",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Error\n",
    "error_c = relative_error(c_pred, c_test)\n",
    "error_u = relative_error(u_pred, u_test)\n",
    "error_v = relative_error(v_pred, v_test)\n",
    "error_p = relative_error(p_pred - np.mean(p_pred), p_test - np.mean(p_test))\n",
    "\n",
    "print('Error c: %e' % (error_c))\n",
    "print('Error u: %e' % (error_u))\n",
    "print('Error v: %e' % (error_v))\n",
    "print('Error p: %e' % (error_p))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45273a7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "################# Save Data ###########################\n",
    "\n",
    "C_pred = 0*C_star\n",
    "U_pred = 0*U_star\n",
    "V_pred = 0*V_star\n",
    "P_pred = 0*P_star"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23fc6d78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the relative error of each variable per time-step\n",
    "for snap in range(0,t_star.shape[0]):\n",
    "    t_test = T_star[:,snap:snap+1]\n",
    "    x_test = X_star[:,snap:snap+1]\n",
    "    y_test = Y_star[:,snap:snap+1]\n",
    "\n",
    "    c_test = C_star[:,snap:snap+1]\n",
    "    u_test = U_star[:,snap:snap+1]\n",
    "    v_test = V_star[:,snap:snap+1]\n",
    "    p_test = P_star[:,snap:snap+1]\n",
    "\n",
    "    # Prediction\n",
    "    c_pred, u_pred, v_pred, p_pred = model.predict(t_test, x_test, y_test)\n",
    "\n",
    "    C_pred[:,snap:snap+1] = c_pred\n",
    "    U_pred[:,snap:snap+1] = u_pred\n",
    "    V_pred[:,snap:snap+1] = v_pred\n",
    "    P_pred[:,snap:snap+1] = p_pred\n",
    "\n",
    "    # Error\n",
    "    error_c = relative_error(c_pred, c_test)\n",
    "    error_u = relative_error(u_pred, u_test)\n",
    "    error_v = relative_error(v_pred, v_test)\n",
    "    error_p = relative_error(p_pred - np.mean(p_pred), p_test - np.mean(p_test))\n",
    "    print('Time step: ', snap)\n",
    "    print('Error c: %e' % (error_c))\n",
    "    print('Error u: %e' % (error_u))\n",
    "    print('Error v: %e' % (error_v))\n",
    "    print('Error p: %e' % (error_p))\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a5425f0",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "    \n",
    "# Plots\n",
    "    \n",
    "Since the original code included no Python-based plotting functionality, the plots have been taken from the original [paper](https://www.sciencedirect.com/science/article/pii/S0021999117309014). Here we see, at one time step, the concentration fields c and d. The inverted nature of these two variables can be seen. \n",
    "\n",
    "![](../images/Plots_c-d.png)\n",
    "\n",
    "In the plots below, we see comparisons between ground truth and prediction for each output variable at one time step. We see that the PINN is capable of replicating each variable with a good degree of accuracy, although it is worth noting that the pressure prediction can only be accurate up to a constant, as it only appears in the Navier-Stokes equation in derivative form. This discrepency can be seen in the colourbar alongside the pressure plots.\n",
    "\n",
    "![](../images/HFMplot2.png)\n",
    "\n",
    "Finally, plots for ground truth and predicted wall shear stress (WSS) can be found below. Here, plotted is the WSS magnitude on the lower boundary of the domain at every time step. What is most impressive here is that we have not imposed any boundary conditions on the lower wall and we still obtain very accurate predictions based only on information from within the flow.\n",
    "\n",
    "![](../images/Plots_wss.png)\n",
    "    \n",
    "    \n",
    "</div>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}