PINNs_NavierStokes_example.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 3</h2>\n",
    "    <h2> PINN Navier Stokes 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": "0c2d1f44",
   "metadata": {},
   "source": [
    "<hr>\n",
    "\n",
    "\n",
    "<div style=\"background-color: #e6ccff; padding: 10px;\">\n",
    "    \n",
    "<h1> Machine Learning Theory </h1>\n",
    "<a href=\"https://victorzhou.com/series/neural-networks-from-scratch/\">\n",
    "<img src=\"https://victorzhou.com/media/nn-series/network.svg\">\n",
    "</a>\n",
    "\n",
    "    \n",
    "## Physics informed Neural Networks\n",
    "\n",
    "Neural networks work by using lots of data to calculate weights and biases which minimise the loss function, enabling them to act as universal function approximators. Because they rely purely on observed relationships in the data, these networks lose their robustness when data is limited. However, by using known physical laws or empirically validated relationships, the solutions from neural networks can be sufficiently constrained by disregarding unrealistic solutions.\n",
    "    \n",
    "A Physics Informed Neural Network considers a parameterized and nonlinear partial differential equation in the general form\n",
    "\n",
    "\n",
    "\n",
    "    \n",
    "\\begin{align}\n",
    "     u_t + \\mathcal{N}[u; \\lambda] &= 0, && x \\in \\Omega, t \\in [0,T],\\\\\n",
    "\\end{align}\n",
    "    \n",
    "\n",
    "\n",
    "where $\\mathcal{u(t,x)}$ denotes the hidden solution, $\\mathcal{N}$ is a nonlinear differential operator acting on $u$, $\\mathcal{\\lambda}$ and $\\Omega$ is a subset of $\\mathbb{R}^D$ (the prescribed data). This set up encapsulates a wide range of problems such as diffusion processes, conservation laws,  advection-diffusion-reaction  systems,  and  kinetic  equations and conservation laws.\n",
    "\n",
    "Here we will go though this for the 1D Heat equation and for the Navier Stokes equations.\n",
    "\n",
    "\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 Heat Equation PINN Example](PINNs_1DEquationExample.ipynb)\n",
    "3. **[Navier-Stokes PINNs discovery of PDE’s](PINNs_Navier_Stokes_example.ipynb)**\n",
    "4. [Navier-Stokes PINNs Hidden Fluid Mechanics](PINNs_NavierStokes_HFM.ipynb)\n"
   ]
  },
  {
   "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. \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",
    "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",
    "import time\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": "37a3b261",
   "metadata": {},
   "source": [
    "<hr>\n",
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "<h1> Navier-Stokes inverse data driven discovery of PDE’s </h1>\n",
    "\n",
    "The Navier-Stokes equations describe the physics of many phenomena of scientific and engineering interest. They may be used to model the weather, ocean currents, water flow in a pipe and air flow around a wing. The Navier-Stokes equations in their full and simplified forms help with the design of aircraft and cars, the study of blood flow, the design of power stations, the analysis of the dispersion of pollutants, and many other applications. Let us consider the Navier-Stokes equations in two dimensions (2D) given explicitly by\n",
    "\n",
    "\\begin{equation}    \n",
    "u_t + \\lambda_1 (u u_x + v u_y) = -p_x + \\lambda_2(u_{xx} + u_{yy}),\\\\\n",
    "v_t + \\lambda_1 (u v_x + v v_y) = -p_y + \\lambda_2(v_{xx} + v_{yy}),\n",
    "\\end{equation}\n",
    "   \n",
    "where $u(t, x, y)$ denotes the $x$-component of the velocity field, $v(t, x, y)$ the $y$-component, and $p(t, x, y)$ the pressure. Here, $\\lambda = (\\lambda_1, \\lambda_2)$ are the unknown parameters. Solutions to the Navier-Stokes equations are sought in the set of divergence-free functions; i.e.,\n",
    "\n",
    "\\begin{equation} \n",
    "u_x + v_y = 0.\n",
    "\\end{equation}\n",
    "       \n",
    "This extra equation is the continuity equation for incompressible fluids that describes the conservation of mass of the fluid. We make the assumption that\n",
    "\n",
    "\\begin{equation}    \n",
    "u = \\psi_y,\\ \\ \\ v = -\\psi_x,\n",
    "\\end{equation}\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbf6505b",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "\n",
    "for some latent function $\\psi(t,x,y)$. Under this assumption, the continuity equation will be automatically satisfied. Given noisy measurements\n",
    "\n",
    "\\begin{equation}\n",
    "\\{t^i, x^i, y^i, u^i, v^i\\}_{i=1}^{N}\n",
    "\\end{equation}\n",
    "    \n",
    "of the velocity field, we are interested in learning the parameters $\\lambda$ as well as the pressure $p(t,x,y)$. We define $f(t,x,y)$ and $g(t,x,y)$ to be given by\n",
    "\n",
    "\\begin{equation}\n",
    "\\begin{array}{c}\n",
    "f := u_t + \\lambda_1 (u u_x + v u_y) + p_x - \\lambda_2(u_{xx} + u_{yy}),\\\\\n",
    "g := v_t + \\lambda_1 (u v_x + v v_y) + p_y - \\lambda_2(v_{xx} + v_{yy}),\n",
    "\\end{array}\n",
    "\\end{equation}\n",
    "\n",
    "and proceed by jointly approximating \n",
    "\n",
    "\\begin{equation}\n",
    "\\begin{bmatrix}\n",
    "\\psi(t,x,y) & p(t,x,y)\n",
    "\\end{bmatrix}\n",
    "\\end{equation}\n",
    "    \n",
    "using a single neural network with two outputs. This prior assumption results into a [physics informed neural network](https://arxiv.org/abs/1711.10566) \n",
    "    \n",
    "\\begin{equation}\n",
    "\\begin{bmatrix}\n",
    "f(t,x,y) & g(t,x,y)\n",
    "\\end{bmatrix}.\n",
    "\\end{equation}\n",
    "    \n",
    "The parameters $\\lambda$ of the Navier-Stokes operator as well as the parameters of the neural networks \n",
    "\n",
    "\\begin{equation}\n",
    "\\begin{bmatrix}\n",
    "\\psi(t,x,y) & p(t,x,y)\n",
    "\\end{bmatrix}\n",
    "\\end{equation}\n",
    "and \n",
    "    \n",
    "\\begin{equation}\n",
    "\\begin{bmatrix}\n",
    "f(t,x,y) & g(t,x,y)\n",
    "\\end{bmatrix}\n",
    "\\end{equation}\n",
    "    \n",
    "can be trained by minimizing the mean squared error loss\n",
    "\n",
    "\\begin{equation}\n",
    "\\begin{array}{rl}\n",
    "MSE :=& \\frac{1}{N}\\sum_{i=1}^{N} \\left(|u(t^i,x^i,y^i) - u^i|^2 + |v(t^i,x^i,y^i) - v^i|^2\\right) \\\\\n",
    "    +& \\frac{1}{N}\\sum_{i=1}^{N} \\left(|f(t^i,x^i,y^i)|^2 + |g(t^i,x^i,y^i)|^2\\right).\n",
    "\\end{array}\n",
    "\\end{equation}\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02ad8792",
   "metadata": {},
   "outputs": [],
   "source": [
    "def xavier_init( size):\n",
    "    in_dim = size[0]\n",
    "    out_dim = size[1]        \n",
    "    xavier_stddev = np.sqrt(2/(in_dim + out_dim))\n",
    "    return tf.Variable(tf.random.truncated_normal([in_dim, out_dim], stddev=xavier_stddev), dtype=tf.float32)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebcf2227",
   "metadata": {},
   "outputs": [],
   "source": [
    "def initialize_NN( layers):        \n",
    "    weights = []\n",
    "    biases = []\n",
    "    num_layers = len(layers) \n",
    "    for l in range(0,num_layers-1):\n",
    "        W = xavier_init(size=[layers[l], layers[l+1]])\n",
    "        b = tf.Variable(tf.zeros([1,layers[l+1]], dtype=tf.float32), dtype=tf.float32)\n",
    "        weights.append(W)\n",
    "        biases.append(b)        \n",
    "    return weights, biases\n",
    "      "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f9a8de",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Initialise the neural network \n",
    "    \n",
    "`init` is called passing in the training data `x_train`, `y_train`, `t_train`, `u_train` and  `v_train` with information about the neural network layers.\n",
    "    \n",
    "# Extract vars\n",
    "    \n",
    "`init` reformats some of the data and outputs model features that we need to pass into the training function `train`.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f64a43d",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "# Advanced \n",
    "    \n",
    "    \n",
    "Once you have run through the notebook once you may wish to alter the optimizer used in the `init()` function to see the large effect optimizer choice may have.\n",
    "    \n",
    "We've highlighted in the comments a number of possible optimizers to use from the [tf.compat.v1.train](https://www.tensorflow.org/api_docs/python/tf/compat/v1/train) module. \n",
    "*This method was chosen to limit TensorFlow version modifications required from the original source code.*\n",
    "    \n",
    "You can learn more about different optimizers [here](https://towardsdatascience.com/optimizers-for-training-neural-network-59450d71caf6).\n",
    "    \n",
    "</div>\n",
    "\n",
    "# Init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f87e7414",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init(x, y, t, u, v, layers):\n",
    "        # This line of code is required to prevent some TensorFlow errors arrising from the\n",
    "        # inclusion of some TensorFlow v1 code \n",
    "        tf.compat.v1.disable_eager_execution()\n",
    "        X = np.concatenate([x, y, t], 1)\n",
    "        # lb and ub denote lower and upper bounds on the inputs to the network\n",
    "        # these bounds are used to normalise the network variables\n",
    "        lb = X.min(0)\n",
    "        ub = X.max(0)\n",
    "                \n",
    "        X = X\n",
    "        \n",
    "        x = X[:,0:1]\n",
    "        y = X[:,1:2]\n",
    "        t = X[:,2:3]\n",
    "        \n",
    "        u = u\n",
    "        v = v\n",
    "        \n",
    "        layers = layers\n",
    "        \n",
    "        # Initialize NN\n",
    "        weights, biases = initialize_NN(layers)        \n",
    "        \n",
    "        # Initialize parameters\n",
    "        lambda_1 = tf.Variable([0.0], dtype=tf.float32)\n",
    "        lambda_2 = tf.Variable([0.0], dtype=tf.float32)\n",
    "        \n",
    "        \n",
    "        \n",
    "        # tf placeholders and graph\n",
    "        ## This converts the data into a Tensorflow format\n",
    "        sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(allow_soft_placement=True,\n",
    "                                                     log_device_placement=True))\n",
    "\n",
    "        x_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, x.shape[1]])\n",
    "        y_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, y.shape[1]])\n",
    "        t_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, t.shape[1]])\n",
    "\n",
    "        u_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, u.shape[1]])\n",
    "        v_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, v.shape[1]])\n",
    "        u_pred, v_pred, p_pred, f_u_pred, f_v_pred = net_NS(x_tf, y_tf, t_tf,lambda_1, lambda_2, weights, biases,lb, ub)\n",
    "\n",
    "        loss = tf.reduce_sum(tf.square(u_tf - u_pred)) + \\\n",
    "            tf.reduce_sum(tf.square(v_tf - v_pred)) + \\\n",
    "            tf.reduce_sum(tf.square(f_u_pred)) + \\\n",
    "            tf.reduce_sum(tf.square(f_v_pred))\n",
    "\n",
    "        \n",
    "        ##############################################################################################\n",
    "        #                                                                                            #\n",
    "        #  The optimizer is something that can be tuned to different requirements.                   #\n",
    "        #  We have not investigated using different optimizers. The orignal code uses L-BFGS-B which # \n",
    "        #  is not TensorFlow 2 compatible.                                                           #\n",
    "        #                                                                                            #\n",
    "        #  SELECT OPTIMIZER BY UNCOMMENTING one of the below lines AND RERUNNING CODE.               #\n",
    "        #  You can also edit the learning rate to see the effect of that.                            #\n",
    "        #                                                                                            #\n",
    "        ##############################################################################################\n",
    "        \n",
    "        learning_rate = 0.001\n",
    "        optimizer = tf.compat.v1.train.MomentumOptimizer(learning_rate, 0.9)\n",
    "        # optimizer = tf.compat.v1.train.AdagradOptimizer(learning_rate) # 8 %\n",
    "        # optimizer = tf.compat.v1.train.ProximalGradientDescentOptimizer(learning_rate)  \n",
    "        # optimizer = tf.compat.v1.train.GradientDescentOptimizer(learning_rate) \n",
    "        # optimizer = tf.compat.v1.train.AdadeltaOptimizer(learning_rate) # yeilds poor results\n",
    "        # ptimizer = tf.compat.v1.train.FtrlOptimizer(learning_rate) \n",
    "        # optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate)\n",
    "        \n",
    "        # LEAVE THESE OPIMISERS ALONE\n",
    "        optimizer_Adam = tf.compat.v1.train.AdamOptimizer()\n",
    "        train_op_Adam = optimizer_Adam.minimize(loss)                    \n",
    "\n",
    "        init = tf.compat.v1.global_variables_initializer()\n",
    "        sess.run(init)\n",
    "        xvars=[X,lb,ub,x,y,t,u,v]\n",
    "        NNvars=[layers, weights, biases, lambda_1, lambda_2]\n",
    "        tfvars=[sess, x_tf,y_tf, t_tf ,u_tf,v_tf]\n",
    "        preds=[u_pred,v_pred, p_pred, f_u_pred, f_v_pred]\n",
    "        optvars=[loss, optimizer,optimizer_Adam,train_op_Adam]\n",
    "        \n",
    "        \n",
    "        return xvars,NNvars,tfvars,preds,optvars"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afb966da",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "`neural_net()` constructs the network Y where X is a matrix containing the input and output coordinates, i.e. x,t,u and X is normalised so that all values lie between -1 and 1 (this improves training)\n",
    "\n",
    "`net_NS()` is where the PDE is encoded\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a61a3a37",
   "metadata": {},
   "outputs": [],
   "source": [
    "def neural_net(X, weights, biases,lb, ub):\n",
    "    \n",
    "    num_layers = len(weights) + 1\n",
    "\n",
    "    H = 2.0*(X - lb)/(ub - lb) - 1.0\n",
    "    for l in range(0,num_layers-2):\n",
    "        W = weights[l]\n",
    "        b = biases[l]\n",
    "        H = tf.tanh(tf.add(tf.matmul(H, W), b))\n",
    "    W = weights[-1]\n",
    "    b = biases[-1]\n",
    "    Y = tf.add(tf.matmul(H, W), b)\n",
    "    return Y\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b43d1d84",
   "metadata": {},
   "outputs": [],
   "source": [
    "def net_NS( x, y, t,lambda_1, lambda_2, weights, biases,lb,ub):\n",
    "    \n",
    "    psi_and_p = neural_net(tf.concat([x,y,t], 1), weights, biases, lb,ub)\n",
    "    psi = psi_and_p[:,0:1]\n",
    "    p = psi_and_p[:,1:2]\n",
    "\n",
    "    u = tf.gradients(psi, y)[0]\n",
    "    v = -tf.gradients(psi, x)[0]  \n",
    "\n",
    "    u_t = tf.gradients(u, t)[0]\n",
    "    u_x = tf.gradients(u, x)[0]\n",
    "    u_y = tf.gradients(u, y)[0]\n",
    "    u_xx = tf.gradients(u_x, x)[0]\n",
    "    u_yy = tf.gradients(u_y, y)[0]\n",
    "\n",
    "    v_t = tf.gradients(v, t)[0]\n",
    "    v_x = tf.gradients(v, x)[0]\n",
    "    v_y = tf.gradients(v, y)[0]\n",
    "    v_xx = tf.gradients(v_x, x)[0]\n",
    "    v_yy = tf.gradients(v_y, y)[0]\n",
    "\n",
    "    p_x = tf.gradients(p, x)[0]\n",
    "    p_y = tf.gradients(p, y)[0]\n",
    "\n",
    "    f_u = u_t + lambda_1*(u*u_x + v*u_y) + p_x - lambda_2*(u_xx + u_yy) \n",
    "    f_v = v_t + lambda_1*(u*v_x + v*v_y) + p_y - lambda_2*(v_xx + v_yy)\n",
    "\n",
    "    return u, v, p, f_u, f_v\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08530832",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Load data and set input parameters \n",
    "   \n",
    "</div>\n",
    "\n",
    "\n",
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "Once you have run through the notebook once you may wish to alter any the following:\n",
    "    \n",
    "- number of data training points `N_train`,\n",
    "- number of layers in the network `layers`,\n",
    "- number of neurons per layer `layers`,\n",
    "    \n",
    "to see the impact on the results.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43d3a8bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_train = 5000\n",
    "# structure of network: \n",
    "# 8 fully connected layers with 20 nodes per layer\n",
    "layers = [3, 20, 20, 20, 20, 20, 20, 20, 20, 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "993a55e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Data\n",
    "data = scipy.io.loadmat('PINNs/main/Data/cylinder_nektar_wake.mat')\n",
    "\n",
    "U_star = data['U_star'] # N x 2 x T\n",
    "P_star = data['p_star'] # N x T\n",
    "t_star = data['t'] # T x 1\n",
    "X_star = data['X_star'] # N x 2\n",
    "\n",
    "N = X_star.shape[0]\n",
    "T = t_star.shape[0]\n",
    "\n",
    "# Rearrange Data \n",
    "XX = np.tile(X_star[:,0:1], (1,T)) # N x T\n",
    "YY = np.tile(X_star[:,1:2], (1,T)) # N x T\n",
    "TT = np.tile(t_star, (1,N)).T # N x T\n",
    "\n",
    "UU = U_star[:,0,:] # N x T\n",
    "VV = U_star[:,1,:] # N x T\n",
    "PP = P_star # N x T\n",
    "\n",
    "x = XX.flatten()[:,None] # NT x 1\n",
    "y = YY.flatten()[:,None] # NT x 1\n",
    "t = TT.flatten()[:,None] # NT x 1\n",
    "\n",
    "u = UU.flatten()[:,None] # NT x 1\n",
    "v = VV.flatten()[:,None] # NT x 1\n",
    "p = PP.flatten()[:,None] # NT x 1\n",
    "\n",
    "######################################################################\n",
    "######################## Noiseles Data ###############################\n",
    "######################################################################\n",
    "# Training Data    \n",
    "## We downsample the boundary data to leave N_train randomly distributed points\n",
    "## This makes the training more difficult - \n",
    "## if we used all the points then there is not much for the network to do!\n",
    "idx = np.random.choice(N*T, N_train, replace=False)\n",
    "x_train = x[idx,:]\n",
    "y_train = y[idx,:]\n",
    "t_train = t[idx,:]\n",
    "u_train = u[idx,:]\n",
    "v_train = v[idx,:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09952e84",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "If this fails you may need to restarted the notebook with a flag:\n",
    "```bash\n",
    "jupyter notebook --NotebookApp.iopub_data_rate_limit=1.0e10\n",
    "\n",
    "```\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "345d3bde",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Initalize the neural network \n",
    "    \n",
    "`init` is called passing in the training data `x_train`,  `y_train`, `u_train` and `v_train` with information about the neural network layers. The bound information `lb` `ub` is included in the `init()` function.\n",
    "    \n",
    "# Extract vars\n",
    "    \n",
    "`init` reformats some of the data and outputs model features that we need to pass into the training function `train`.\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba0d6968",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "xvars, NNvars, tfvars, preds, optvars = init(x_train, y_train, t_train, u_train, v_train, layers)\n",
    "X, lb, ub, x, y, t, u, v = xvars\n",
    "layers, weights, biases, lambda_1, lambda_2 = NNvars\n",
    "sess, x_tf,y_tf, t_tf ,u_tf,v_tf = tfvars\n",
    "u_pred,v_pred, p_pred, f_u_pred, f_v_pred = preds\n",
    "loss, optimizer, optimizer_Adam, train_op_Adam = optvars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08433368",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(sess, nIter,x_tf, y_tf, t_tf,u_tf, v_tf,x, y, t,u, v, loss, train_op_Adam, optimizer, path): \n",
    "    # Add ops to save and restore all the variables.\n",
    "    # path = \"modelckpts/NS/\"\n",
    "    netSaveDir = path\n",
    "    saver = tf.compat.v1.train.Saver()\n",
    "    tf_dict = {x_tf: x, y_tf: y, t_tf: t,\n",
    "             u_tf: u, v_tf: v}\n",
    "\n",
    "    start_time = time()\n",
    "    for it in range(nIter):\n",
    "        sess.run(train_op_Adam, tf_dict)\n",
    "\n",
    "        # Print\n",
    "        if it % 50 == 0:\n",
    "            elapsed = time() - start_time\n",
    "            loss_value = sess.run(loss, tf_dict)\n",
    "            lambda_1_value = sess.run(lambda_1)\n",
    "            lambda_2_value = sess.run(lambda_2)\n",
    "            print('It: %d, Loss: %.3e, l1: %.3f, l2: %.5f, Time: %.2f' % \n",
    "                (it, loss_value, lambda_1_value, lambda_2_value, elapsed))\n",
    "            start_time = time()\n",
    "        \n",
    "        if it % 2000 == 0:\n",
    "            save_path = saver.save(sess, netSaveDir + 'model_at_iter%s.ckpt'%(it))\n",
    "            print('Model saved in path: %s' % save_path)\n",
    "\n",
    "    print('Model saved in path: %s' % save_path)\n",
    "    return optimizer.minimize(loss)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96976c5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict( x_star, y_star, t_star, u_pred, v_pred, p_pred):\n",
    "  \n",
    "    tf_dict = {x_tf: x_star, y_tf: y_star, t_tf: t_star}\n",
    "\n",
    "    u_star = sess.run(u_pred, tf_dict)\n",
    "    v_star = sess.run(v_pred, tf_dict)\n",
    "    p_star = sess.run(p_pred, tf_dict)\n",
    "\n",
    "    return u_star, v_star, p_star\n",
    "\n",
    "def plot_solution(X_star, u_star, index):\n",
    "  \n",
    "    lb = X_star.min(0)\n",
    "    ub = X_star.max(0)\n",
    "    nn = 200\n",
    "    x = np.linspace(lb[0], ub[0], nn)\n",
    "    y = np.linspace(lb[1], ub[1], nn)\n",
    "    X, Y = np.meshgrid(x,y)\n",
    "\n",
    "    U_star = griddata(X_star, u_star.flatten(), (X, Y), method='cubic')\n",
    "\n",
    "    plt.figure(index)\n",
    "    plt.pcolor(X,Y,U_star, cmap = 'jet')\n",
    "    plt.colorbar()\n",
    "  \n",
    "  \n",
    "def axisEqual3D(ax):\n",
    "    extents = np.array([getattr(ax, 'get_{}lim'.format(dim))() for dim in 'xyz'])\n",
    "    sz = extents[:,1] - extents[:,0]\n",
    "    centers = np.mean(extents, axis=1)\n",
    "    maxsize = max(abs(sz))\n",
    "    r = maxsize/4\n",
    "    for ctr, dim in zip(centers, 'xyz'):\n",
    "        getattr(ax, 'set_{}lim'.format(dim))(ctr - r, ctr + r)\n",
    "      \n",
    "      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af33bf85",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set random seeds\n",
    "np.random.seed(1234)\n",
    "tf.random.set_seed(1234)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3accacd4",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "**Training might take a long time depending on value of Train_iterations**\n",
    "\n",
    "If you set Train_iterations too low the end results will be garbage. 20000 was used to achieve excellent results in the original papers but this value is too high to run on a laptop. \n",
    "\n",
    "* If you are using a machine with GPUs please set `Train_iterations=20000` to achieve the best results\n",
    "* If you are using a well spec'ed laptop/computer and can leave this setting `Train_iterations=10000` should suffice (may take a while)\n",
    "* If you are using a low spec'ed laptop/computer or cannont leave the code running `Train_interations=5000` is the reccomended values (high errors will remain)\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baf8859c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training\n",
    "Train_iterations = 20000"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05197413",
   "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_iter15000.ckpt')`.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed5e0da8",
   "metadata": {},
   "outputs": [],
   "source": [
    "train(sess, Train_iterations,x_tf, y_tf, t_tf, u_tf, v_tf,x, y, t,u_train, v_train, loss, train_op_Adam, optimizer_Adam, \"modelckpts/NS/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "676ddc60",
   "metadata": {},
   "outputs": [],
   "source": [
    "loadweights = False\n",
    "if loadweights:\n",
    "    print(\"loading pre trained model\")\n",
    "    netSaveDir = \"modelckpts/NS/\"\n",
    "    saver = tf.compat.v1.train.Saver()\n",
    "    saver.restore(sess, netSaveDir + 'model_at_iter19000.ckpt')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "899a1899",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Use trained model to predict from data sample\n",
    "    \n",
    "`predict` will predict `u`, `v` and `p` using the trained model.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c33cdc7",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "The `predict` function has an option `load=False` set by default. Alter this to `load=True` if you wish to load the previously trained model.\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a551e020",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test Data\n",
    "snap = np.array([100])\n",
    "x_star = X_star[:,0:1]\n",
    "y_star = X_star[:,1:2]\n",
    "t_star = TT[:,snap]\n",
    "\n",
    "u_star = U_star[:,0,snap]\n",
    "v_star = U_star[:,1,snap]\n",
    "p_star = P_star[:,snap]\n",
    "\n",
    "# Prediction\n",
    "u_pred,v_pred, p_pred, f_u_pred, f_v_pred = preds\n",
    "#\n",
    "# EDIT BELOW TO LOAD TRAINED MODEL --> load=True\n",
    "#\n",
    "u_pred, v_pred, p_pred = predict(x_star, y_star, t_star, u_pred, v_pred, p_pred)\n",
    "lambda_1_value = sess.run(lambda_1)\n",
    "lambda_2_value = sess.run(lambda_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0ec214f",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Calculate Errors\n",
    "    \n",
    "If you have set the number of training iterations large enough the errors should be small.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23fc6d78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Error\n",
    "error_u = np.linalg.norm(u_star-u_pred,2)/np.linalg.norm(u_star,2)\n",
    "error_v = np.linalg.norm(v_star-v_pred,2)/np.linalg.norm(v_star,2)\n",
    "error_p = np.linalg.norm(p_star-p_pred,2)/np.linalg.norm(p_star,2)\n",
    "\n",
    "error_lambda_1 = np.abs(lambda_1_value - 1.0)*100\n",
    "error_lambda_2 = np.abs(lambda_2_value - 0.01)/0.01 * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02f0d716",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Error u: %e' % (error_u))    \n",
    "print('Error v: %e' % (error_v))    \n",
    "print('Error p: %e' % (error_p))    \n",
    "print('Error l1: %.5f%%' % (error_lambda_1))                             \n",
    "print('Error l2: %.5f%%' % (error_lambda_2))        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34ad7ed3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Predict for plotting\n",
    "lb = X_star.min(0)\n",
    "ub = X_star.max(0)\n",
    "nn = 200\n",
    "x = np.linspace(lb[0], ub[0], nn)\n",
    "y = np.linspace(lb[1], ub[1], nn)\n",
    "X, Y = np.meshgrid(x,y)\n",
    "\n",
    "UU_star = griddata(X_star, u_pred.flatten(), (X, Y), method='cubic')\n",
    "VV_star = griddata(X_star, v_pred.flatten(), (X, Y), method='cubic')\n",
    "PP_star = griddata(X_star, p_pred.flatten(), (X, Y), method='cubic')\n",
    "P_exact = griddata(X_star, p_star.flatten(), (X, Y), method='cubic')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b88ff3e5",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Using Noisy Data\n",
    "    \n",
    "We're now going to repeat the previous steps but include some noise in our data to see the effect of that on our results.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16f8f591",
   "metadata": {},
   "outputs": [],
   "source": [
    "######################################################################\n",
    "########################### Noisy Data ###############################\n",
    "######################################################################\n",
    "noise = 0.01        \n",
    "u_train = u_train + noise*np.std(u_train)*np.random.randn(u_train.shape[0], u_train.shape[1])\n",
    "v_train = v_train + noise*np.std(v_train)*np.random.randn(v_train.shape[0], v_train.shape[1])    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb88c53f",
   "metadata": {},
   "outputs": [],
   "source": [
    "xvars, NNvars, tfvars, preds, optvars = init(x_train, y_train, t_train, u_train, v_train, layers)\n",
    "X, lb, ub, x, y, t, u, v = xvars\n",
    "layers, weights, biases, lambda_1, lambda_2 = NNvars\n",
    "sess, x_tf, y_tf, t_tf ,u_tf,v_tf = tfvars\n",
    "u_pred, v_pred, p_pred, f_u_pred, f_v_pred = preds\n",
    "loss, optimizer, optimizer_Adam, train_op_Adam = optvars"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33d97748",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "**Training might take a while depending on value of Train_iterations**\n",
    "\n",
    "If you set Train_iterations too low the end results will be garbage. 20000 was used to achieve excellent results. \n",
    "\n",
    "* If you are using a machine with [GPUs](https://towardsdatascience.com/what-is-a-gpu-and-do-you-need-one-in-deep-learning-718b9597aa0d) please set `Train_iterations` to 20000 and this will run in a few mins\n",
    "* If you are using a well spec'ed laptop/computer then setting `Train_iterations=10000` but it will take a little while\n",
    "* If you are using a low spec'ed laptop/computer or cannont leave the code running `Train_iterations=5000` is the reccomended value (this solution may not be accurate)\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0d21566",
   "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_iter15000.ckpt')`.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0a26045",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training\n",
    "train(sess, Train_iterations, x_tf, y_tf, t_tf, u_tf, v_tf, x, y, t, u_train, v_train, loss, train_op_Adam, optimizer_Adam,\"modelckpts/NSn/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b36d3e6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "loadweights = False\n",
    "if loadweights:\n",
    "    print(\"loading pre trained model\")\n",
    "    netSaveDir = \"modelckpts/NSn/\"\n",
    "    saver = tf.compat.v1.train.Saver()\n",
    "    saver.restore(sess, netSaveDir + 'model_at_iter18000.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97f1938c",
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_1_value_noisy = sess.run(lambda_1)\n",
    "lambda_2_value_noisy = sess.run(lambda_2)\n",
    "\n",
    "error_lambda_1_noisy = np.abs(lambda_1_value_noisy - 1.0)*100\n",
    "error_lambda_2_noisy = np.abs(lambda_2_value_noisy - 0.01)/0.01 * 100\n",
    "\n",
    "print('Error l1: %.5f%%' % (error_lambda_1_noisy))                             \n",
    "print('Error l2: %.5f%%' % (error_lambda_2_noisy))     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c255a86a",
   "metadata": {},
   "outputs": [],
   "source": [
    "######################################################################\n",
    "############################# Plotting ###############################\n",
    "######################################################################    \n",
    "# Load Data\n",
    "data_vort = scipy.io.loadmat('PINNs/main/Data/cylinder_nektar_t0_vorticity.mat')\n",
    "\n",
    "x_vort = data_vort['x'] \n",
    "y_vort = data_vort['y'] \n",
    "w_vort = data_vort['w'] \n",
    "modes = np.ndarray.item(data_vort['modes'])\n",
    "nel = np.ndarray.item(data_vort['nel'])    \n",
    "\n",
    "xx_vort = np.reshape(x_vort, (modes+1,modes+1,nel), order = 'F')\n",
    "yy_vort = np.reshape(y_vort, (modes+1,modes+1,nel), order = 'F')\n",
    "ww_vort = np.reshape(w_vort, (modes+1,modes+1,nel), order = 'F')\n",
    "\n",
    "box_lb = np.array([1.0, -2.0])\n",
    "box_ub = np.array([8.0, 2.0])\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.axis('off')\n",
    "plt.figure(figsize=(16, 8))\n",
    "####### Row 0: Vorticity ##################    \n",
    "gs0 = gridspec.GridSpec(1, 2)\n",
    "gs0.update(top=1-0.06, bottom=1-2/4 + 0.12, left=0.0, right=1.0, wspace=0)\n",
    "ax = plt.subplot(gs0[:, :])\n",
    "\n",
    "for i in range(0, nel):\n",
    "    h = ax.pcolormesh(xx_vort[:,:,i], yy_vort[:,:,i], ww_vort[:,:,i], cmap='seismic',shading='gouraud',  vmin=-3, vmax=3) \n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "fig.colorbar(h, cax=cax)\n",
    "\n",
    "ax.plot([box_lb[0],box_lb[0]],[box_lb[1],box_ub[1]],'k',linewidth = 1)\n",
    "ax.plot([box_ub[0],box_ub[0]],[box_lb[1],box_ub[1]],'k',linewidth = 1)\n",
    "ax.plot([box_lb[0],box_ub[0]],[box_lb[1],box_lb[1]],'k',linewidth = 1)\n",
    "ax.plot([box_lb[0],box_ub[0]],[box_ub[1],box_ub[1]],'k',linewidth = 1)\n",
    "\n",
    "ax.set_aspect('equal', 'box')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$y$')\n",
    "ax.set_title('Vorticity', fontsize = 10)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d794243e",
   "metadata": {},
   "outputs": [],
   "source": [
    "####### Row 1: Training data ##################\n",
    "########      u(t,x,y)     ###################  \n",
    "plt.figure(figsize=(20, 8))\n",
    "gs1 = gridspec.GridSpec(1, 2)\n",
    "gs1.update(top=1-2/4, bottom=0.0, left=0.01, right=0.99, wspace=0)\n",
    "ax = plt.subplot(gs1[:, 0],  projection='3d')\n",
    "ax.axis('off')\n",
    "\n",
    "r1 = [x_star.min(), x_star.max()]\n",
    "r2 = [data['t'].min(), data['t'].max()]       \n",
    "r3 = [y_star.min(), y_star.max()]\n",
    "\n",
    "for s, e in combinations(np.array(list(product(r1,r2,r3))), 2):\n",
    "    if np.sum(np.abs(s-e)) == r1[1]-r1[0] or np.sum(np.abs(s-e)) == r2[1]-r2[0] or np.sum(np.abs(s-e)) == r3[1]-r3[0]:\n",
    "        ax.plot3D(*zip(s,e), color=\"k\", linewidth = 0.5)   \n",
    "\n",
    "ax.scatter(x_train, t_train, y_train, s = 0.1)\n",
    "# Predict for plotting\n",
    "lb = X_star.min(0)\n",
    "ub = X_star.max(0)\n",
    "nn = 200\n",
    "x = np.linspace(lb[0], ub[0], nn)\n",
    "y = np.linspace(lb[1], ub[1], nn)\n",
    "X, Y = np.meshgrid(x,y)\n",
    "ax.contourf(X,UU_star,Y, zdir = 'y', offset = t_star.mean(), cmap='rainbow', alpha = 0.8)\n",
    "\n",
    "ax.text(x_star.mean(), data['t'].min() - 1, y_star.min() - 1, '$x$')\n",
    "ax.text(x_star.max()+1, data['t'].mean(), y_star.min() - 1, '$t$')\n",
    "ax.text(x_star.min()-1, data['t'].min() - 0.5, y_star.mean(), '$y$')\n",
    "ax.text(x_star.min()-3, data['t'].mean(), y_star.max() + 1, '$u(t,x,y)$')    \n",
    "ax.set_xlim3d(r1)\n",
    "ax.set_ylim3d(r2)\n",
    "ax.set_zlim3d(r3)\n",
    "axisEqual3D(ax)\n",
    "\n",
    "########      v(t,x,y)     ###################        \n",
    "ax = plt.subplot(gs1[:, 1],  projection='3d')\n",
    "ax.axis('off')\n",
    "\n",
    "r1 = [x_star.min(), x_star.max()]\n",
    "r2 = [data['t'].min(), data['t'].max()]       \n",
    "r3 = [y_star.min(), y_star.max()]\n",
    "\n",
    "for s, e in combinations(np.array(list(product(r1,r2,r3))), 2):\n",
    "    if np.sum(np.abs(s-e)) == r1[1]-r1[0] or np.sum(np.abs(s-e)) == r2[1]-r2[0] or np.sum(np.abs(s-e)) == r3[1]-r3[0]:\n",
    "        ax.plot3D(*zip(s,e), color=\"k\", linewidth = 0.5)   \n",
    "\n",
    "ax.scatter(x_train, t_train, y_train, s = 0.1)\n",
    "ax.contourf(X,VV_star,Y, zdir = 'y', offset = t_star.mean(), cmap='rainbow', alpha = 0.8)\n",
    "\n",
    "ax.text(x_star.mean(), data['t'].min() - 1, y_star.min() - 1, '$x$')\n",
    "ax.text(x_star.max()+1, data['t'].mean(), y_star.min() - 1, '$t$')\n",
    "ax.text(x_star.min()-1, data['t'].min() - 0.5, y_star.mean(), '$y$')\n",
    "ax.text(x_star.min()-3, data['t'].mean(), y_star.max() + 1, '$v(t,x,y)$')    \n",
    "ax.set_xlim3d(r1)\n",
    "ax.set_ylim3d(r2)\n",
    "ax.set_zlim3d(r3)\n",
    "axisEqual3D(ax)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15b0cf51",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "plt.figure(figsize=(20, 8))\n",
    "ax.axis('off')\n",
    "\n",
    "######## Row 2: Pressure #######################\n",
    "########      Predicted p(t,x,y)     ########### \n",
    "gs2 = gridspec.GridSpec(1, 2)\n",
    "gs2.update(top=1, bottom=1-1/2, left=0.1, right=0.9, wspace=0.5)\n",
    "ax = plt.subplot(gs2[:, 0])\n",
    "h = ax.imshow(PP_star, interpolation='nearest', cmap='rainbow', \n",
    "            extent=[x_star.min(), x_star.max(), y_star.min(), y_star.max()], \n",
    "            origin='lower', aspect='auto')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "\n",
    "fig.colorbar(h, cax=cax)\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$y$')\n",
    "ax.set_aspect('equal', 'box')\n",
    "ax.set_title('Predicted pressure', fontsize = 10)\n",
    "\n",
    "########     Exact p(t,x,y)     ########### \n",
    "ax = plt.subplot(gs2[:, 1])\n",
    "h = ax.imshow(P_exact, interpolation='nearest', cmap='rainbow', \n",
    "            extent=[x_star.min(), x_star.max(), y_star.min(), y_star.max()], \n",
    "            origin='lower', aspect='auto')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "\n",
    "fig.colorbar(h, cax=cax)\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$y$')\n",
    "ax.set_aspect('equal', 'box')\n",
    "ax.set_title('Exact pressure', fontsize = 12)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18f381f2",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "    \n",
    "Predicted versus exact instantaneous pressure field at a representative time instant. By definition, the pressure can be recovered up to a constant, hence justifying the different magnitude between the two plots. This remarkable qualitative agreement highlights the ability of physics-informed neural networks to identify the entire pressure field, despite the fact that no data on the pressure are used during model training. \n",
    "\n",
    "**NB** train must be set to approx 10000 to achieve the desired results.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eace04b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "######## Row 3: Table #######################\n",
    "gs3 = gridspec.GridSpec(1, 2)\n",
    "gs3.update(top=1-1/2, bottom=0.0, left=0.0, right=1.0, wspace=0)\n",
    "ax = plt.subplot(gs3[:, :])\n",
    "ax.axis('off')\n",
    "plt.rc('text', usetex=False)\n",
    "s=''\n",
    "s = s + \"Correct PDE \\n \"\n",
    "s = s + \"$u_t + (u u_x + v u_y) = -p_x + 0.01 (u_{xx} + u_{yy})$\"\n",
    "s = s + \"\\n\"\n",
    "s = s + \"$v_t + (u v_x + v v_y) = -p_y + 0.01 (v_{xx} + v_{yy})$\"\n",
    "s = s + \"\\n \\n \\n\"\n",
    "s = s + r'Identified PDE (clean data) '\n",
    "s = s + \"\\n\"\n",
    "s = s + '$u_t + %.3f (u u_x + v u_y) = -p_x + %.5f (u_{xx} + u_{yy})$' % (lambda_1_value, lambda_2_value)\n",
    "s = s + \"\\n\"\n",
    "s = s + '$v_t + %.3f (u v_x + v v_y) = -p_y + %.5f (v_{xx} + v_{yy})$' % (lambda_1_value, lambda_2_value)\n",
    "s = s + \"\\n\\n \\n\"\n",
    "\n",
    "s = s + r'Identified PDE (1% noise) & '\n",
    "s = s + \"\\n\"\n",
    "s = s + '$u_t + %.3f (u u_x + v u_y) = -p_x + %.5f (u_{xx} + u_{yy})$' % (lambda_1_value_noisy, lambda_2_value_noisy)\n",
    "s = s + \"\\n\"\n",
    "s = s + '$v_t + %.3f (u v_x + v v_y) = -p_y + %.5f (v_{xx} + v_{yy})$' % (lambda_1_value_noisy, lambda_2_value_noisy)\n",
    "s = s + \"\\n\"\n",
    "\n",
    "plt.rc('font',family='serif')\n",
    "\n",
    "plt.rc('font',size=16)\n",
    "ax.text(0,0,s,fontsize=18)\n",
    "\n",
    "# savefig('./figures/NavierStokes_prediction') "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a5425f0",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\"> \n",
    "    \n",
    "If you have not been able to run enough training iterations the figures produced running 10000 iterations can be found:\n",
    "    \n",
    "* [Solution with network trained over 10000 iterations](figures/PINNS_NS_10000_PDE.png)\n",
    "* [Figure comparing predicted vs exact with network trained over 10000 iterations](figures/PINNS_NS_10000_predict_vs_exact.png)\n",
    "\n",
    "**Further Work**\n",
    "\n",
    "Congratulations, you have now trained your another physics-informed neural network!\n",
    "\n",
    "This network contains a number of hyper-parameters that could be tuned to give better results. Various hyper-parameters include:\n",
    "- number of data training points `N_train`\n",
    "- number of `layers` in the network\n",
    "- number of neurons per layer\n",
    "- optimisation \n",
    "\n",
    "It is also possible to use different sampling techniques for training data points. We randomly select $N_u$ data points, but alternative methods could be choosing only boundary points or choosing more points near the $t=0$ boundary.\n",
    "\n",
    "Return [here](#init) to alter optimization method used.\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68081d83",
   "metadata": {},
   "source": [
    "<hr>\n",
    "\n",
    "<div style=\"background-color: #e6ccff; padding: 10px;\">\n",
    "\n",
    "## Next steps\n",
    "\n",
    "Now we've demonstrated using PINNs for more complex equations we can take a brief look at Hidden Fluid Mechanics (*this final notebook is beyond the scope of these tutorials but provided to give a brief example of the methodology*).\n",
    "    \n",
    "[Navier-Stokes PINNs Hidden Fluid Mechanics](PINNs_NavierStokes_HFM.ipynb)\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
}