PINNS_navier-stokes_equation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "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",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "This notebook is based on two papers: *[Physics-Informed Neural Networks:  A Deep LearningFramework for Solving Forward and Inverse ProblemsInvolving Nonlinear Partial Differential Equations](https://www.sciencedirect.com/science/article/pii/S0021999118307125)* and *[Hidden Physics Models:  Machine Learning of NonlinearPartial 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 Neuaral Networks focusing on the Burgers Equation and a more complex example using the Navier Stokes Equation\n",
    "\n",
    "**This introduction section is replicated in all PINN tutorial notebooks (please skip if you've already been through)** "
   ]
  },
  {
   "cell_type": "markdown",
   "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 algorithims like gradient descent to look for a hypothesis, data is the only guide, however if the data is noisy or sparse and we already have governing physical models we can use the knowledge we already know to optamize and inform the algoithms. This can be done via [feature enginnering]() or by adding a physicall inconsistency term to the loss function.\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 know nothing about 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). Creating a 2 layer neural network to illustrate the fundamentals of how Neural Networks work and the equivlent 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",
    "* [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",
    "<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 from data alone to minimise the loss function enabling them to act as universal fuction approximators. However these loose their robustness when data is limited. However by using know physical laws or empirical validated relationships the solutions from neural networks can be sufficiently constrianed by disregardins no realistic solutions.\n",
    "    \n",
    "A Physics Informed Nueral Network considers a parameterized and nonlinear partial differential equation in the genral form;\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",
    "where $\\mathcal{u(t,x)}$ denores 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 perscribed data). This set up an encapuslate 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 Navier stokes equations\n",
    "\n",
    "\n",
    "</div>    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "<h1> Python </h1>\n",
    "\n",
    "## Pytorch\n",
    "    \n",
    "There are many machine learning python libraries available, [Pytorch](https://pytorch.org/) a is one such library the original code for this tutorial is in Tensorflow, if you are interested in the difference please take a look [here](https://github.com/cemac/LIFD_Physics_Informed_Neural_Networks). If you have GPUs on the machine you are using Pytorch will automatically use them and run the code even faster! This will work automatically with google co-lab if using your own machine please ensure that the GPU enabled pytorch is installed \n",
    "\n",
    "## Further Reading\n",
    "\n",
    "* [Running Jupyter Notebooks](https://jupyter.readthedocs.io/en/latest/running.html#running)\n",
    "* [PyTorch optimizers](https://nbviewer.org/github/bentrevett/a-tour-of-pytorch-optimizers/blob/main/a-tour-of-pytorch-optimizers.ipynb)\n",
    "\n",
    "\n",
    "\n",
    "</div>\n",
    "    \n",
    "<hr>\n",
    "\n",
    "<div style=\"background-color: #f4b85d; 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",
    "* PyTorch\n",
    "* numpy \n",
    "* matplotlib\n",
    "* scipy\n",
    "\n",
    "<h2> Data Requirements</h2>\n",
    "    \n",
    "This notebook referes to some data included in the git hub repositroy\n",
    "    \n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "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",
    "    * [1D Heat Equation Forwards](#1D-Heat-Equation-Forwards)\n",
    "    * [1D Heat Equation Inverse](#1D-Heat-Equation-Inverse)\n",
    "3. **[Navier-Stokes PINNs discovery of PDE’s](PINNs_NavierStokes_example.ipynb)**\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "Load in all required modules (includig some auxillary code) and turn off warnings. \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For readability: disable warnings\n",
    "import warnings\n",
    "import os\n",
    "warnings.filterwarnings('ignore')\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import scipy.io\n",
    "from scipy.interpolate import griddata\n",
    "import time\n",
    "import matplotlib.pyplot as plt\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",
    "import scipy.sparse as sp\n",
    "import scipy.sparse.linalg as la\n",
    "from pyDOE import lhs\n",
    "import requests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "This code defines a neural network with Xavier initialization in PyTorch which is a custom layer initialization using Xavier initialization, which is designed to keep the scale of the gradients roughly the same in all layers.\n",
    "\n",
    "`NeuralNet()` constructs the network U where X is a matrix containing the input and output coordinates, i.e. x,t,u\n",
    "and X is normalised so that all values lie between -1 and 1, this improves training. Applies the layers sequentially with the tanh activation function, except for the last layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# Xavier Initialization for PyTorch\n",
    "def xavier_init(size):\n",
    "    \"\"\"\n",
    "    Initialize the weights of a layer using Xavier initialization.\n",
    "\n",
    "    Parameters:\n",
    "    size (tuple): A tuple containing the dimensions of the layer (in_dim, out_dim).\n",
    "\n",
    "    Returns:\n",
    "    torch.Tensor: A tensor with the initialized weights.\n",
    "    \"\"\"\n",
    "    in_dim, out_dim = size\n",
    "    xavier_stddev = np.sqrt(2 / (in_dim + out_dim))\n",
    "    return torch.randn(in_dim, out_dim) * xavier_stddev\n",
    "\n",
    "# Neural Network Initialization\n",
    "class NeuralNet(nn.Module):\n",
    "    \"\"\"\n",
    "    A simple feedforward neural network with tanh activation functions.\n",
    "\n",
    "    Attributes:\n",
    "    layers (nn.ModuleList): A list of linear layers.\n",
    "    lb (torch.Tensor): Lower bound for input normalization.\n",
    "    ub (torch.Tensor): Upper bound for input normalization.\n",
    "    \"\"\"\n",
    "    def __init__(self, layers):\n",
    "        \"\"\"\n",
    "        Initialize the neural network.\n",
    "\n",
    "        Parameters:\n",
    "        layers (list): A list containing the number of neurons in each layer.\n",
    "        \"\"\"\n",
    "        super(NeuralNet, self).__init__()\n",
    "        self.layers = nn.ModuleList()\n",
    "        for i in range(len(layers) - 1):\n",
    "            self.layers.append(nn.Linear(layers[i], layers[i+1]))\n",
    "        self.lb = None\n",
    "        self.ub = None\n",
    "\n",
    "    def forward(self, X):\n",
    "        \"\"\"\n",
    "        Forward pass through the network.\n",
    "\n",
    "        Parameters:\n",
    "        X (torch.Tensor): The input tensor.\n",
    "\n",
    "        Returns:\n",
    "        torch.Tensor: The output of the network.\n",
    "        \"\"\"\n",
    "        H = 2.0 * (X - self.lb) / (self.ub - self.lb) - 1.0\n",
    "        for l in range(len(self.layers) - 1):\n",
    "            H = torch.tanh(self.layers[l](H))\n",
    "        return self.layers[-1](H)"
   ]
  },
  {
   "cell_type": "markdown",
   "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",
    "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 searched 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>\n",
    "\n",
    "<hr>\n",
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "<h1> Navier-Stokes inverse data driven discovery of PDE’s </h1>\n",
    "\n",
    "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 searched 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>\n",
    "\n",
    "<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": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "This code defines a Physics-Informed Neural Network (PINN) model for solving the Navier-Stokes equations. Below are the key components:\n",
    "\n",
    "#### PINN Class\n",
    "- **Class Name:** `PINN`\n",
    "- **Purpose:** Defines a PINN model to predict fluid dynamics using the Navier-Stokes equations.\n",
    "- **Attributes:**\n",
    "  - `model` (NeuralNet): Neural network model for the PINN.\n",
    "  - `x`, `y`, `t` (torch.Tensor): Input coordinates and time as PyTorch tensors.\n",
    "  - `u`, `v` (torch.Tensor): Velocity components as PyTorch tensors.\n",
    "  - `lambda_1`, `lambda_2` (torch.Tensor): Learnable parameters for the Navier-Stokes equations.\n",
    "  - `optimizer` (torch.optim.Adam): Optimizer for training the model.\n",
    "- **Methods:**\n",
    "  - `__init__(self, x, y, t, u, v, layers)`: Initializes the PINN model with input data and network layers.\n",
    "  - `predict(self, x_star, y_star, t_star)`: Predicts the velocity components and pressure at given coordinates and time.\n",
    "  - `compute_loss(self)`: Computes the loss function for training the model.\n",
    "  - `net_NS(self)`: Defines the Navier-Stokes equation model and computes necessary gradients.\n",
    "  - `train(self, nIter)`: Trains the model for a specified number of iterations.\n",
    "\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "### Autograd in PyTorch\n",
    "\n",
    "This is a more complex model than out 1D heat equation so we need to consider our computation of gradients a bit more carefully. Autograd is PyTorch's automatic differentiation engine that powers neural network training. It provides automatic computation of gradients for tensor operations, which is essential for backpropagation.\n",
    "\n",
    "- PyTorch records operations on tensors that have `requires_grad=True` to create a computation graph.\n",
    "- When `.backward()` is called on a tensor, PyTorch traverses this graph to compute gradients for all tensors involved in the computation.\n",
    "\n",
    "- **Computation Graph**:\n",
    "  - The computation graph is a directed acyclic graph where nodes represent operations and edges represent tensors.\n",
    "  - By default, PyTorch frees the computation graph after the backward pass to save memory.\n",
    "\n",
    "- **`retain_graph=True`**:\n",
    "  - By default, PyTorch frees the computation graph after the backward pass. If you need to perform multiple backward passes on the same graph (e.g., higher-order gradients), you must retain the graph.\n",
    "  - Setting `retain_graph=True` ensures that the graph is not freed, allowing for subsequent backward passes.\n",
    "  - Here the gradients of 'u','v'.'p' with respect to 'x','y','t' are computed multiple times. Retaining the graph ensures that it is available for subsequent computations \n",
    "\n",
    "- **`create_graph=True`**:\n",
    "  - When computing higher-order derivatives, you need to create a new computation graph during the backward pass.\n",
    "  - Setting `create_graph=True` allows PyTorch to construct a new graph during the backward pass, enabling the computation of higher-order gradients.\n",
    "  - Example:\n",
    "    ```python\n",
    "    grad = torch.autograd.grad(outputs=y, inputs=x, create_graph=True)\n",
    "    ```\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PINN Model\n",
    "class PINN:\n",
    "    def __init__(self, x, y, t, u, v, layers):\n",
    "        # Set up model and lambda parameters\n",
    "        self.model = NeuralNet(layers)\n",
    "        self.model.lb = torch.from_numpy(x.min(0)).float()\n",
    "        self.model.ub = torch.from_numpy(x.max(0)).float()\n",
    "        \n",
    "        # Convert input data to PyTorch tensors\n",
    "        self.x = torch.tensor(x, requires_grad=True).float()\n",
    "        self.y = torch.tensor(y, requires_grad=True).float()\n",
    "        self.t = torch.tensor(t, requires_grad=True).float()\n",
    "        self.u = torch.tensor(u).float()\n",
    "        self.v = torch.tensor(v).float()\n",
    "        \n",
    "        # Initialize lambda_1 and lambda_2 as learnable parameters\n",
    "        self.lambda_1 = torch.tensor([0.0], requires_grad=True, dtype=torch.float32)\n",
    "        self.lambda_2 = torch.tensor([0.0], requires_grad=True, dtype=torch.float32)\n",
    "        \n",
    "        # Optimizer\n",
    "        self.optimizer = optim.Adam(\n",
    "            list(self.model.parameters()) + [self.lambda_1, self.lambda_2],\n",
    "            lr=0.001\n",
    "        )\n",
    "\n",
    "    # Function to predict values\n",
    "    def predict(self, x_star, y_star, t_star):\n",
    "        X_star = torch.cat([x_star, y_star, t_star], dim=1)\n",
    "        psi_and_p = self.model(X_star)\n",
    "        psi = psi_and_p[:, 0:1]\n",
    "        p = psi_and_p[:, 1:2]\n",
    "        \n",
    "        u = torch.autograd.grad(psi, y_star, torch.ones_like(psi), retain_graph=True, create_graph=True)[0]\n",
    "        v = -torch.autograd.grad(psi, x_star, torch.ones_like(psi), retain_graph=True, create_graph=True)[0]\n",
    "        \n",
    "        return u, v, p\n",
    "\n",
    "    def compute_loss(self):\n",
    "        u_pred, v_pred, p_pred, f_u_pred, f_v_pred = self.net_NS()\n",
    "        loss = torch.mean((self.u - u_pred)**2) + \\\n",
    "               torch.mean((self.v - v_pred)**2) + \\\n",
    "               torch.mean(f_u_pred**2) + \\\n",
    "               torch.mean(f_v_pred**2)\n",
    "        return loss\n",
    "\n",
    "    # Navier-Stokes equation model\n",
    "    def net_NS(self):\n",
    "        X = torch.cat([self.x, self.y, self.t], dim=1)\n",
    "        psi_and_p = self.model(X)\n",
    "        psi = psi_and_p[:, 0:1]\n",
    "        p = psi_and_p[:, 1:2]\n",
    "        \n",
    "        # Compute gradients\n",
    "        u = torch.autograd.grad(psi, self.y, torch.ones_like(psi), retain_graph=True, create_graph=True)[0]\n",
    "        v = -torch.autograd.grad(psi, self.x, torch.ones_like(psi), retain_graph=True, create_graph=True)[0]\n",
    "        \n",
    "        # Additional gradients for Navier-Stokes\n",
    "        u_t = torch.autograd.grad(u, self.t, torch.ones_like(u), retain_graph=True, create_graph=True)[0]\n",
    "        u_x = torch.autograd.grad(u, self.x, torch.ones_like(u), retain_graph=True, create_graph=True)[0]\n",
    "        u_y = torch.autograd.grad(u, self.y, torch.ones_like(u), retain_graph=True, create_graph=True)[0]\n",
    "        u_xx = torch.autograd.grad(u_x, self.x, torch.ones_like(u_x), retain_graph=True, create_graph=True)[0]\n",
    "        u_yy = torch.autograd.grad(u_y, self.y, torch.ones_like(u_y), retain_graph=True, create_graph=True)[0]\n",
    "        \n",
    "        v_t = torch.autograd.grad(v, self.t, torch.ones_like(v), retain_graph=True, create_graph=True)[0]\n",
    "        v_x = torch.autograd.grad(v, self.x, torch.ones_like(v), retain_graph=True, create_graph=True)[0]\n",
    "        v_y = torch.autograd.grad(v, self.y, torch.ones_like(v), retain_graph=True, create_graph=True)[0]\n",
    "        v_xx = torch.autograd.grad(v_x, self.x, torch.ones_like(v_x), retain_graph=True, create_graph=True)[0]\n",
    "        v_yy = torch.autograd.grad(v_y, self.y, torch.ones_like(v_y), retain_graph=True, create_graph=True)[0]\n",
    "        \n",
    "        p_x = torch.autograd.grad(p, self.x, torch.ones_like(p), retain_graph=True, create_graph=True)[0]\n",
    "        p_y = torch.autograd.grad(p, self.y, torch.ones_like(p), retain_graph=True, create_graph=True)[0]\n",
    "        \n",
    "        f_u = u_t + self.lambda_1 * (u * u_x + v * u_y) + p_x - self.lambda_2 * (u_xx + u_yy)\n",
    "        f_v = v_t + self.lambda_1 * (u * v_x + v * v_y) + p_y - self.lambda_2 * (v_xx + v_yy)\n",
    "        \n",
    "        return u, v, p, f_u, f_v\n",
    "\n",
    "    def train(self, nIter):\n",
    "        for it in range(nIter):\n",
    "            # Zero the gradients\n",
    "            self.optimizer.zero_grad()\n",
    "            # Compute the loss\n",
    "            loss = self.compute_loss()\n",
    "            # Backpropagation\n",
    "            loss.backward()\n",
    "            # Update the model parameters\n",
    "            self.optimizer.step()\n",
    "            # Print the training progress every 50 iterations\n",
    "            if it % 50 == 0:\n",
    "                print(f\"Iter: {it}, Loss: {loss.item():.3e}, Lambda1: {self.lambda_1.item():.3f}, Lambda2: {self.lambda_2.item():.5f}\")\n",
    "\n",
    "    def save(self, path):\n",
    "        torch.save({\n",
    "            'model_state_dict': self.model.state_dict(),\n",
    "            'lambda_1': self.lambda_1,\n",
    "            'lambda_2': self.lambda_2\n",
    "        }, path)\n",
    "\n",
    "    def load(self, path):\n",
    "        if not os.path.exists(path):\n",
    "            raise FileNotFoundError(f\"No such file: '{path}'\")\n",
    "        \n",
    "        checkpoint = torch.load(path)\n",
    "        self.model.load_state_dict(checkpoint['model_state_dict'])\n",
    "        self.lambda_1 = torch.tensor(checkpoint['lambda_1'], requires_grad=True, dtype=torch.float32)\n",
    "        self.lambda_2 = torch.tensor(checkpoint['lambda_2'], requires_grad=True, dtype=torch.float32)\n",
    "   \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "### Optimizer Choice in the `PINN` Class\n",
    "\n",
    "The `PINN` class uses the Adam optimizer for training the model. Below is an expanded explanation of the optimizer choice:\n",
    "\n",
    "\n",
    "**Adam Optimizer**: \n",
    "  - The Adam optimizer is an adaptive learning rate optimization algorithm designed for training deep learning models. It combines the advantages of two other popular optimizers: AdaGrad and RMSProp.\n",
    "  - Adam computes individual adaptive learning rates for different parameters from estimates of first and second moments of the gradients.\n",
    "  - It is well-suited for problems with large datasets and/or high-dimensional parameter spaces.\n",
    "\n",
    "#### Optimizer Initialization\n",
    "- **Optimizer Used**: `torch.optim.Adam`\n",
    "- **Learning Rate**: `0.001` which is a common default value for the Adam optimizer. This value controls the step size at each iteration while moving towards a minimum of the loss function.\n",
    "- **Parameters Optimized**: \n",
    "  - Model parameters (weights and biases of the neural network)\n",
    "  - Learnable parameters for the Navier stokes equation `lambda_1` and `lambda_2`\n",
    "\n",
    "\n",
    "#### Optimizer Choice alternatives\n",
    "\n",
    "`LBFGS`: Limited-memory Broyden–Fletcher–Goldfarb–Shanno (quasi-Newton)\n",
    "\n",
    "`Adagrad`: Adaptive Subgradient [link to paper for more info](https://jmlr.org/papers/volume12/duchi11a/duchi11a.pdf)\n",
    "\n",
    "`SGD`: stochastic gradient descent (optionally with momentum\n",
    "\n",
    "`Adadelta`: Adaptive learning rate  (poor in this case) [link to paper for more info](https://arxiv.org/abs/1212.5701)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Load data and set input parameters \n",
    "      \n",
    "A feedforward neural network of the following structure is assumed:\n",
    "- the input is scaled elementwise to lie in the interval $[-1, 1]$,\n",
    "- followed by 8 fully connected layers each containing 20 neurons and each followed by a hyperbolic tangent activation function,\n",
    "- two fully connected output layer.\n",
    "\n",
    "This setting results in a network with a first hidden layer: $3 \\cdot 20 + 20 = 80$; $8$ intermediate layers: each $20 \\cdot 20 + 20 = 540$; output layer: $20 \\cdot 2 + 2 = 42$).\n",
    "\n",
    "1. Input Layer:\n",
    "\n",
    "Number of Neurons: 3\n",
    "This layer takes in the input features. In the context of the Navier-Stokes equations, these could be spatial and temporal coordinates (e.g., x, y, t).\n",
    "\n",
    "2. Hidden Layers:\n",
    "\n",
    "Number of Layers: 8\n",
    "Number of Neurons per Layer: 20\n",
    "Each hidden layer consists of 20 neurons. These layers apply transformations to the input data through learned weights and biases, followed by an activation function (typically ReLU or tanh).\n",
    "\n",
    "3. Output Layer:\n",
    "\n",
    "\n",
    "Number of Neurons: 2\n",
    "This layer produces the final output of the network. In the context of the Navier-Stokes equations, the outputs could represent the velocity components (e.g., u, v) or other relevant quantities.\n",
    "\n",
    "    \n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize and Train the Model\n",
    "layers = [3, 20, 20, 20, 20, 20, 20, 20, 20, 2]\n",
    "data = scipy.io.loadmat('Data/cylinder_nektar_wake.mat')\n",
    "\n",
    "U_star = data['U_star'] # velocity coordinates (N x 2 x T)\n",
    "P_star = data['p_star'] # pressure coordinates  (N x T)\n",
    "t_star = data['t'] # temporal coordinates  (T x 1)\n",
    "X_star = data['X_star'] # Spatial coordinates (N x 2)\n",
    "\n",
    "\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",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "#### 1. Data Preparation\n",
    "\n",
    "- **Input Data**: Prepare the input data by downsampling and organizing it into training sets.\n",
    "  - `x_train`, `y_train`, `t_train`: Input features (spatial and temporal coordinates).\n",
    "  - `u_train`, `v_train`: Target outputs (velocity components)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "######################################################################\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",
    "N_train=5000\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert training data to tensor format and initialize the model\n",
    "pinn = PINN(x_train, y_train, t_train, u_train, v_train, layers)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "## Training\n",
    "\n",
    "\n",
    "Now we've intialised our model we can call our training funtion, all we need to pass in is `nIter` which is the number of training interations, our training function outlined in our PINN class, for each iteration the training function:\n",
    "\n",
    "1. zeros the gradients, \n",
    "2. computes a mean squared errors between the predicted and actual velocity components (`u` and `v`), as well as the residuals of the Navier-Stokes equations (`f_u` and `f_v`).\n",
    "3. Performs back propagation This computes the gradients of the loss with respect to the model parameters, including the learnable parameters `lambda_1` and `lambda_2`.\n",
    "4. Optimizer step: The optimizer updates the model parameters based on the computed gradients \n",
    "5. Every 50 iterations, the function prints the current iteration number, the loss value, and the values of lambda_1 and lambda_2. You expect all values to change each iteration, if they not changing or the loss is very high etc it can indicate something has gone awry before completing the training.\n",
    "\n",
    "### Loss Function in the `compute_loss` Method\n",
    "\n",
    "The `compute_loss` method calculates the loss for training the Physics-Informed Neural Network (PINN) model. The loss function is composed of two main parts:\n",
    "\n",
    "1. **Data Loss**: This part measures the difference between the predicted and actual velocity components (`u` and `v`). It ensures that the model's predictions are close to the observed data.\n",
    "2. **Physics Loss**: This part enforces the physical constraints of the Navier-Stokes equations. It ensures that the predicted velocity components (`u` and `v`) satisfy the Navier-Stokes equations.\n",
    "\n",
    "#### Detailed Breakdown\n",
    "\n",
    "- **Predicted Values**: The method first calls `net_NS()` to obtain the predicted velocity components (`u_pred` and `v_pred`), pressure (`p_pred`), and the residuals of the Navier-Stokes equations (`f_u_pred` and `f_v_pred`).\n",
    "\n",
    "- **Data Loss**: \n",
    "  - `torch.mean((self.u - u_pred)**2)`: Mean squared error between the actual and predicted `u` velocity component.\n",
    "  - `torch.mean((self.v - v_pred)**2)`: Mean squared error between the actual and predicted `v` velocity component.\n",
    "\n",
    "- **Physics Loss**:\n",
    "  - `torch.mean(f_u_pred**2)`: Mean squared error of the residuals of the Navier-Stokes equation for the `u` component.\n",
    "  - `torch.mean(f_v_pred**2)`: Mean squared error of the residuals of the Navier-Stokes equation for the `v` component.\n",
    "\n",
    "- **Total Loss**: The total loss is the sum of the data loss and the physics loss. This combined loss ensures that the model not only fits the observed data but also adheres to the underlying physical laws.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #D3D3D3; padding: 10px;\">\n",
    "\n",
    "if you wish to load a pre-trained model you will need to copy this code into a cell. Available trained models are on hugging face [huggingface.co/CEMAC/NS_PINN/tree/main](https://huggingface.co/CEMAC/NS_PINN/tree/main)\n",
    "\n",
    "````python\n",
    "\n",
    "url = 'https://huggingface.co/CEMAC/NS_PINN/resolve/main/pinn_model_niter_20k.pth'\n",
    "model_path = 'pinn_model_niter_20k.pth'\n",
    "\n",
    "if not os.path.isfile(model_path):\n",
    "    response = requests.get(url, stream=True)\n",
    "    if response.status_code == 200:\n",
    "        with open(model_path, 'wb') as f:\n",
    "            for chunk in response.iter_content(chunk_size=1024):\n",
    "                if chunk:  # Filter out keep-alive chunks\n",
    "                    f.write(chunk)\n",
    "        print(\"File downloaded successfully.\")\n",
    "    else:\n",
    "        print(f\"Failed to download file. HTTP status code: {response.status_code}\")\n",
    "else:\n",
    "    print(\"File already exists.\")\n",
    "pinn.load(model_path)\n",
    "\n",
    "````\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# Train the model\n",
    "nIter = 20000\n",
    "pinn.train(nIter)\n",
    "\n",
    "# Save the model\n",
    "model_path = 'pinn_model_noise_niter_'+str(nIter)+'.pth'\n",
    "pinn.save(model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "## Prediction with Test Data\n",
    "\n",
    "Now we've trained our Physics-Informed Neural Network (PINN) for the Navier-Stokes equations we can use some test data to \n",
    "\n",
    "1. Test Data Preparation:\n",
    "\n",
    "* `snap` is an array containing the snapshot index for which predictions are to be made.\n",
    "* `x_star`, `y_star`, and `t_star` are the spatial and temporal coordinates of the test data.\n",
    "* `u_star`, `v_star`, and `p_star` are the true velocity components and pressure at the test data points.\n",
    "\n",
    "2. Prediction:\n",
    "\n",
    "The `predict` method of the `pinn` object is called with the test data coordinates to obtain the predicted velocity components (`u_pred`, `v_pred`) and pressure (`p_pred`).\n",
    "\n",
    "3. Evaluation:\n",
    "\n",
    "Compute the error between the predicted and true values to evaluate the model's performance.\n",
    "Visualize the predictions and compare them with the true values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "\n",
    "\n",
    "u_pred, v_pred, p_pred = pinn.predict(torch.tensor(x_star, dtype=torch.float32, requires_grad=True), \n",
    "                                      torch.tensor(y_star, dtype=torch.float32, requires_grad=True), \n",
    "                                      torch.tensor(t_star, dtype=torch.float32, requires_grad=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_1 = pinn.lambda_1\n",
    "lambda_2 = pinn.lambda_2\n",
    "\n",
    "# Error\n",
    "error_u = np.linalg.norm(u_star-u_pred.detach().numpy(),2)/np.linalg.norm(u_star,2)\n",
    "error_v = np.linalg.norm(v_star-v_pred.detach().numpy(),2)/np.linalg.norm(v_star,2)\n",
    "error_p = np.linalg.norm(p_star-p_pred.detach().numpy(),2)/np.linalg.norm(p_star,2)\n",
    "\n",
    "error_lambda_1 = np.abs(lambda_1.detach().numpy() - 1.0)*100\n",
    "error_lambda_2 = np.abs(lambda_2.detach().numpy() - 0.01)/0.01 * 100\n",
    "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))        \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",
    "\n",
    "UU_star = griddata(X_star, u_pred.detach().numpy().flatten(), (X, Y), method='cubic')\n",
    "VV_star = griddata(X_star, v_pred.detach().numpy().flatten(), (X, Y), method='cubic')\n",
    "PP_star = griddata(X_star, p_pred.detach().numpy().flatten(), (X, Y), method='cubic')\n",
    "P_exact = griddata(X_star, p_star.flatten(), (X, Y), method='cubic')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "outline some plotting tools\n",
    "\n",
    "`plot_solution` is a helper function that grids the data and plots the velocity field`\n",
    "\n",
    "`axisEqual3D` function is designed to set equal scaling for all three axes (x, y, z) in a 3D plot created using Matplotlib. This ensures that the units are equally scaled across all axes, which is important for accurately representing 3D data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def plot_solution(X_star, u_star, index):\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",
    "    plt.show()\n",
    "\n",
    "# Axis equalizer function remains unchanged\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "######################################################################\n",
    "############################# Plotting ###############################\n",
    "######################################################################    \n",
    "# Load Data\n",
    "data_vort = scipy.io.loadmat('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,
   "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",
    "try:\n",
    "    x_star=x_star.detach().numpy()\n",
    "    y_star=y_star.detach().numpy()\n",
    "except:\n",
    "    pass\n",
    "\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
   "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": "markdown",
   "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,
   "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",
    "# Convert training data to tensor format and initialize the model\n",
    "pinn = PINN(x_train, y_train, t_train, u_train, v_train, layers)\n",
    "\n",
    "# Train the model\n",
    "nIter = 20000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #D3D3D3; padding: 10px;\">\n",
    "\n",
    "if you wish to load a pre-trained model you will need to copy this code into a cell. Available trained models are on hugging face [huggingface.co/CEMAC/NS_PINN/tree/main](https://huggingface.co/CEMAC/NS_PINN/tree/main)\n",
    "\n",
    "````python\n",
    "\n",
    "url = 'https://huggingface.co/CEMAC/NS_PINN/resolve/main/pinn_model_noise_niter_20k.pth'\n",
    "model_path = 'pinn_model_noise_niter_20k.pth'\n",
    "\n",
    "if not os.path.isfile(model_path):\n",
    "    response = requests.get(url, stream=True)\n",
    "    if response.status_code == 200:\n",
    "        with open(model_path, 'wb') as f:\n",
    "            for chunk in response.iter_content(chunk_size=1024):\n",
    "                if chunk:  # Filter out keep-alive chunks\n",
    "                    f.write(chunk)\n",
    "        print(\"File downloaded successfully.\")\n",
    "    else:\n",
    "        print(f\"Failed to download file. HTTP status code: {response.status_code}\")\n",
    "else:\n",
    "    print(\"File already exists.\")\n",
    "pinn.load(model_path)\n",
    "\n",
    "````\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "pinn.train(nIter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "\n",
    "#\n",
    "# EDIT BELOW TO LOAD TRAINED MODEL --> load=True\n",
    "#\n",
    "u_pred, v_pred, p_pred = pinn.predict(torch.tensor(x_star, dtype=torch.float32, requires_grad=True), \n",
    "                                      torch.tensor(y_star, dtype=torch.float32, requires_grad=True), \n",
    "                                      torch.tensor(t_star, dtype=torch.float32, requires_grad=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_1_noisey = pinn.lambda_1\n",
    "lambda_2_noisey = pinn.lambda_2\n",
    "\n",
    "# Error\n",
    "error_u = np.linalg.norm(u_star-u_pred.detach().numpy(),2)/np.linalg.norm(u_star,2)\n",
    "error_v = np.linalg.norm(v_star-v_pred.detach().numpy(),2)/np.linalg.norm(v_star,2)\n",
    "error_p = np.linalg.norm(p_star-p_pred.detach().numpy(),2)/np.linalg.norm(p_star,2)\n",
    "\n",
    "error_lambda_1_noisey = np.abs(lambda_1_noisey.detach().numpy() - 1.0)*100\n",
    "error_lambda_2_noisey = np.abs(lambda_2_noisey.detach().numpy() - 0.01)/0.01 * 100\n",
    "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_noisey))                             \n",
    "print('Error l2: %.5f%%' % (error_lambda_2_noisey))        \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",
    "\n",
    "UU_star = griddata(X_star, u_pred.detach().numpy().flatten(), (X, Y), method='cubic')\n",
    "VV_star = griddata(X_star, v_pred.detach().numpy().flatten(), (X, Y), method='cubic')\n",
    "PP_star = griddata(X_star, p_pred.detach().numpy().flatten(), (X, Y), method='cubic')\n",
    "P_exact = griddata(X_star, p_star.flatten(), (X, Y), method='cubic')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "try:\n",
    "    x_star=x_star.detach().numpy()\n",
    "    y_star=y_star.detach().numpy()\n",
    "except:\n",
    "    pass\n",
    "\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
   "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": "markdown",
   "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 20000 iterations can be found:\n",
    "    \n",
    "* [Solution with network trained over 20000 iterations](figures/PINNS_NS_20000_PDE.png)\n",
    "* [Figure comparing predicted vs exact with network trained over 20000 iterations](figures/PINNS_NS_20000_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",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}