PINNs_1DHeatEquationExample.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fc2b8b36",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "    <h1> Tutorial 2 </h1> \n",
    "    <h2> Physics Informed Neural Networks Part 2</h2>\n",
    "    <h2> 1D Heat Equation PINNs Equation Example </h2>\n",
    "</div>    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dc7a8ba",
   "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 1D Heat 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).** \n",
    "\n",
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "If you have not already then please run the following code in the LIFD_ENV_ML_NOTEBOOKS directory via the terminal (mac or linux)  or git bash (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",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a32bc58b",
   "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": "45be2930",
   "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": "f82d6a60",
   "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": "85906b8e",
   "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": "19e9738e",
   "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",
    "4. [Navier-Stokes PINNs Hidden Fluid Mechanics](PINNs_NavierStokes_HFM.ipynb)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab403585",
   "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": "301ac7de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For readability: disable warnings\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a356009b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\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\n",
    "import subprocess\n",
    "subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"pyDOE\"])\n",
    "from pyDOE import lhs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "109fdc1b",
   "metadata": {},
   "source": [
    "<hr>\n",
    "\n",
    "# Solving 1D heat equations via Neural Networks\n",
    "\n",
    "# 1D Heat Equation Forwards\n",
    "\n",
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "**Model Problem: 1D Heat Equation**\n",
    "\n",
    "We begin by describing the first model problem - the one-dimensional heat equation. \n",
    "\n",
    "The heat equation is the prototypical parabolic partial differential equation and can be applied to modelling the diffusion of heat through a given region, hence its name. Read more about the heat equation here: https://en.wikipedia.org/wiki/Heat_equation.\n",
    "\n",
    "In 1D, the heat equation can be written as:\n",
    "\n",
    "\\begin{equation}\n",
    "\\frac{\\partial u}{\\partial t} = k \\frac{\\partial^2 u}{\\partial x^2 },\n",
    "\\end{equation}\n",
    "    \n",
    "where $k$ is a material parameter called the coefficient of thermal diffusivity.\n",
    "\n",
    "This equation can be solved using numerical methods, such as finite differences or finite elements. For this notebook, we have solved the above equation numerically on a domain of $x \\in [0,1]$ and $t \\in [0, 0.25]$. Solving this equation numerically gives us a spatiotemporal domain $(x,t)$ and corresponding values of the solution $u$.\n",
    "\n",
    "</div>\n",
    "    \n",
    "<div style=\"background-color: #e6ccff; padding: 10px;\">\n",
    "\n",
    "    \n",
    "Here we will describe the architecture of the PINN we use to solve the 1D heat equation in this notebook. \n",
    "![PINNS.png](PINNS.png)\n",
    "    \n",
    "    \n",
    "Net U in the above diagram approximates a function that maps from $(x,t) \\mapsto u$. $\\sigma$ represents the biases and weights for the each neuron of the network. These $\\sigma$ values are the network parameters that are updated after each iteration. AD means Automatic Differentiation - this is the chain rule-based differentiation procedure that allows for differentiation of network outputs with respect to its inputs, e.g. differentiating $u$ with respect to $x$, or calculating $\\frac{\\partial u}{\\partial x}$. The I node in the AD section represents the identity operation, i.e. keeping $u$ fixed without applying any differentiation. \n",
    "\n",
    "After the automatic differentiation part of the network, we have two separate loss function components - the data loss and the PDE loss. The data loss term is calculated by finding the difference between the network outputs/predictions $u$ and the ground truth values of $u$, which could come from simulation or experiment. The data loss term enforces the network outputs to match known data points, which are represented by the pink box labelled \"Data\". The PDE loss term is where we add the \"physics-informed\" part of the network. Using automatic differentiation, we are able to calculate derivatives of our network outputs, and so we are able to construct a loss function that enforces the network to match the PDE that is known to govern the system. In this case, the PDE loss term is defined as:\n",
    " \n",
    "\\begin{equation}\n",
    "f = \\frac{\\partial u}{\\partial t} - k \\frac{\\partial^2 u}{\\partial x^2 },\n",
    "\\end{equation}\n",
    "    \n",
    "where f is the residual of the 1D heat equation. By demanding that $f$ is minimised as our network train, we ensure that the network outputs obey the underlying PDE that governs the system. We then calculate the total loss of the system as a sum of the data loss and the PDE loss.\n",
    "\n",
    "The loss is calculated after each pass through the network and when it is above a certain tolerance, the weights and biases are updated using a gradient descent step. When the loss falls below the tolerance the network is trained. In inference mode, we can then input a fine mesh of spatiotemporal coordinates and the network will find the solution at each of these points.\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc3575fe",
   "metadata": {},
   "source": [
    "<hr>\n",
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "**$u(x,t)$** can then be defined below as the function `net_u` and the physics informed neural network **$f(x,t)$** is outlined in function `net_f`\n",
    "\n",
    "`neural_net()` 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)\n",
    "\n",
    "`net_u()` constructs a network that takes input x,t and outputs the solution u    \n",
    "    \n",
    "`net_f()`  the f network is where the PDE is encoded:\n",
    "    \n",
    "1. we read in the value of k first so that it can be included in the equations \n",
    "2. then we evaluate u for the X_f input coordinates (collocation points)\n",
    "3. then we use tensorflow differentiation to calculate the derivatives of the solution\n",
    "4. finally we encode the PDE in residual form, as f->0, u_t = k*u_xx, which is the governing eq\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9b297dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def neural_net(X, weights, biases,lb,ub):\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",
    "def net_u(x_tf, t_tf, weights, biases, lb, ub):  \n",
    "    u = neural_net(tf.concat([x_tf,t_tf],1), weights, biases, lb, ub)\n",
    "    return u\n",
    "\n",
    "\n",
    "def net_f(x_tf, t_tf, weights, biases,lb, ub,k): \n",
    "    u = net_u(x_tf,t_tf,weights, biases, lb, ub)\n",
    "    u_t = tf.gradients(u, t_tf)[0]\n",
    "    u_x = tf.gradients(u, x_tf)[0]\n",
    "    u_xx = tf.gradients(u_x, x_tf)[0]\n",
    "    f = k*u_xx\n",
    "    return f"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "082afce2",
   "metadata": {},
   "source": [
    "<hr>\n",
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "    \n",
    "###  Intialise everying \n",
    "    \n",
    "The `init` function will take our gridded data X and U, and build our neural networks from the functions defined above ready to train the model.\n",
    "\n",
    "Variables to be defined here:\n",
    "\n",
    "`X_u`: Input coordinates, e.g. spatial and temporal coordinates.\n",
    "\n",
    "`u`: Output corresponding to each input coordinate. \n",
    "\n",
    "`X_f`: Collocation points at which the governing equations are satisfied. These coordinates will have the same format as the X_u coordinates, e.g. $(x,t)$.\n",
    "\n",
    "layers: specify the structure of the u network.\n",
    "\n",
    "`lb`: Vector containing the lower bound of all of the coordinate variables, e.g. $x_{min}$, $t_{min}$.\n",
    "\n",
    "`ub`: Vector containing the upper bound of all of the coordinate variables, e.g. $x_{max}$, $t_{max}$.\n",
    "\n",
    "`k`: This is the constant material parameter for this specific problem. For this problem, the heat equation, $k$ represents thermal diffusivity.\n",
    "    \n",
    "\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b00beade",
   "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": "f470f47e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## the initialisation type can be played with, however we kept it as in the original code\n",
    "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",
    "\n",
    "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",
    "\n",
    "def init(X, u, layers, lb, ub, k):\n",
    "        # This line of code is required to prevent some TensorFlow errors arising from the\n",
    "        # inclusion of some Tensorflow v1 code\n",
    "        tf.compat.v1.disable_eager_execution()\n",
    "        \n",
    "        \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",
    "       \n",
    "    \n",
    "        ## the first two columns of X contain x_u and t_u\n",
    "        x = X[:,0:1]\n",
    "        t = X[:,1:2]\n",
    "        \n",
    "        \n",
    "        layers = layers\n",
    "        \n",
    "        # Initialize NN\n",
    "        weights, biases = initialize_NN(layers)  \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",
    "        t_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, t.shape[1]])\n",
    "        u_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, u.shape[1]])\n",
    "        \n",
    "        # u predictions take the input coordinates x_u and t_u and calculate the corresponding u value\n",
    "        u_pred = net_u(x_tf, t_tf,weights, biases, lb, ub)\n",
    "        # f predictions take the input coordinates x_f and t_f and calculate the corresponding f value\n",
    "        f_pred = net_f(x_tf, t_tf,weights, biases,lb, ub,k)\n",
    "        \n",
    "        ## the loss function is defined by the sum of the prediction loss: u-u_pred \n",
    "        ## and the PDE loss: f_pred\n",
    "        ## the PDE loss only requires one term, as it is a residual and tends to zero\n",
    "        loss_PDE = tf.reduce_mean(tf.square(f_pred))\n",
    "        loss_data = tf.reduce_mean(tf.square(u_tf - u_pred))\n",
    "        loss = loss_PDE + 5*loss_data\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) # yields poor results\n",
    "        # ptimizer = tf.compat.v1.train.FtrlOptimizer(learning_rate) \n",
    "         \n",
    "        \n",
    "        \n",
    "        \n",
    "        # LEAVE THESE OPIMIZERS 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, t, u]\n",
    "        NNvars = [layers, weights, biases]\n",
    "        tfvars = [sess, x_tf, t_tf ,u_tf]\n",
    "        preds = [u_pred, f_pred]\n",
    "        optvars = [loss, optimizer,optimizer_Adam,train_op_Adam]\n",
    "        return xvars, NNvars, tfvars, preds, optvars"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "295fbd70",
   "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",
    "- one fully connected output layer.\n",
    "\n",
    "This setting results in a network with a first hidden layer: $2 \\cdot 20 + 20 = 60$ parameters; $7$ intermediate layers: each $20 \\cdot 20 + 20 = 540$; output layer: $20 \\cdot 1 + 1 = 21$).\n",
    "    \n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bf4360e",
   "metadata": {},
   "source": [
    " <div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "    \n",
    "# Number of collocation points \n",
    "    \n",
    "`2000` collocation points is the default setting for this example. This can be increased to improve results at cost of computational speed. The original work set this `N_u=10000` running on GPUs in a few minutes.\n",
    "    \n",
    "    \n",
    "The network takes in data in coordinate pairs: $(x,t) \\mapsto u$.     \n",
    "</div>\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_u`\n",
    "- number of collocation training points `N_f`\n",
    "- number of layers in the network `layers`\n",
    "- number of neurons per layer `layers`\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb011bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 1\n",
    "N_u = 100 #100 # number of data points\n",
    "N_f = 2000 # Collocation points \n",
    "# structure of network: two inputs (x,t) and one output u\n",
    "# 8 fully connected layers with 20 nodes per layer\n",
    "layers = [2, 20, 20, 20, 20, 20, 20, 20, 20, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a13e702a",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = scipy.io.loadmat(\"Data/heatEquation_data.mat\")\n",
    "t = data['t'].flatten()[:,None] # read in t and flatten into column vector\n",
    "x = data['x'].flatten()[:,None] # read in x and flatten into column vector\n",
    " # Exact represents the exact solution to the problem, from the Matlab script provided\n",
    "Exact = np.real(data['usol']).T # Exact has structure of nx times nt\n",
    "\n",
    "\n",
    "print(\"usol shape = \", Exact.shape)\n",
    "\n",
    "# We need to find all the x,t coordinate pairs in the domain\n",
    "X, T = np.meshgrid(x,t)\n",
    "\n",
    "# Flatten the coordinate grid into pairs of x,t coordinates\n",
    "X_star = np.hstack((X.flatten()[:,None], T.flatten()[:,None])) # coordinates x,t\n",
    "u_star = Exact.flatten()[:,None]   # corresponding solution value with each coordinate            \n",
    "\n",
    "\n",
    "print(\"X has shape \", X.shape, \", X_star has shape \", X_star.shape)\n",
    "    \n",
    "# Doman bounds (-1,1)\n",
    "lb = X_star.min(0)\n",
    "ub = X_star.max(0)  \n",
    "\n",
    "print(\"Lower bounds of x,t: \", lb)\n",
    "print(\"Upper bounds of x,t: \", ub)\n",
    "\n",
    "## train using internal points\n",
    "X_u_train = X_star\n",
    "u_train = u_star\n",
    "\n",
    "## Generate collocation points using Latin Hypercube sampling within the bounds of the spationtemporal coordinates\n",
    "# Generate N_f x,t coordinates within range of upper and lower bounds\n",
    "X_f_train = lb + (ub-lb)*lhs(2, N_f) # the 2 denotes the number of coordinates we have - x,t \n",
    "\n",
    "## In addition, we add the X_u_train coordinats from the boundaries to the X_f coordinate set\n",
    "X_f_train = np.vstack((X_f_train, X_u_train)) # stack up all training x,t coordinates for u and f \n",
    "\n",
    "## We downsample the boundary data to leave N_u 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(X_star.shape[0], N_u, replace=False)\n",
    "X_u_train = X_star[idx,:]\n",
    "u_train = u_star[idx,:]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d358bde2",
   "metadata": {},
   "outputs": [],
   "source": [
    "## making a plot to show the distribution of training data\n",
    "plt.scatter(X_f_train[:,1], X_f_train[:,0], marker='x', color='red',alpha=0.1)\n",
    "plt.scatter(X_u_train[:,1], X_u_train[:,0], marker='x', color='black')\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('t')\n",
    "plt.title('Data points and collocation points (red crosses)')\n",
    "plt.legend(['Collocation Points', 'Data Points'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ee085c0",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Initalise the neural network \n",
    "    \n",
    "`init` is called passing in the training data `X_u_train` and `u_train` with information about the neural network layers and bounds `lb` `ub`.\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": "4f419efd",
   "metadata": {},
   "outputs": [],
   "source": [
    "xvars, NNvars, tfvars, preds, optvars = init(X_u_train, u_train, layers, lb, ub,k)\n",
    "X, lb, ub, x, t, u = xvars\n",
    "layers, weights, biases=NNvars\n",
    "sess, x_tf, t_tf ,u_tf=tfvars\n",
    "u_pred, f_pred,=preds\n",
    "loss, optimizer,optimizer_Adam,train_op_Adam = optvars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce0ba9ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(sess, nIter,x_tf, t_tf, u_tf,x, t,u_train, loss, train_op_Adam, optimizer_Adam): \n",
    "    tf_dict = {x_tf: x,  t_tf: t, u_tf: u}\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",
    "\n",
    "            print('It: %d, Loss: %.3e, Time: %.2f' % \n",
    "                (it, loss_value, elapsed))\n",
    "            start_time = time()\n",
    "\n",
    "    \n",
    "    optimizer.minimize(loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "416fa291",
   "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. 50000 was used to achieve excellent results. \n",
    "\n",
    "* If you are using a machine with GPUs please set `Train_iterations` to 50000 and this will run quickly\n",
    "* If you are using a well spec'ed laptop/computer and can leave this setting `Train_iterations=50000` but it will take upto 10 mins\n",
    "* If you are using a low spec'ed laptop/computer or cannont leave the code running `Train_iterations=20000` is the reccomended value (this solution may not be accurate)\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1678fa0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training\n",
    "Train_iterations=50000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b706205",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train(sess, Train_iterations,x_tf, t_tf, u_tf,x, t,u_train, loss, train_op_Adam, optimizer_Adam)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc849d67",
   "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` using the trained model.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13fba79a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(sess, x_star,u_star, u_pred, f_pred):\n",
    "    tf_dict = {x_tf: x_star, t_tf: u_star}\n",
    "    u_star = sess.run(u_pred, tf_dict)\n",
    "    f_star = sess.run(f_pred, tf_dict)\n",
    "        \n",
    "    \n",
    "\n",
    "    return u_star, f_star"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b086ac50",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "u_pred, f_pred=preds\n",
    "u_pred, f_pred = predict(sess,X_star[:,0:1],X_star[:,1:2], u_pred, f_pred)\n",
    "\n",
    "error_u = np.linalg.norm(u_star-u_pred,2)/np.linalg.norm(u_star,2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c2f300b",
   "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",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fae4a8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"f_pred mean = \", np.mean(f_pred))\n",
    "print('Error u: %e' % (error_u))\n",
    "print('Percent error u: ',  100*error_u)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5c9d3c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set grid values back to full data set size for plotting\n",
    "\n",
    "t = data['t'].flatten()[:,None]\n",
    "x = data['x'].flatten()[:,None]\n",
    "X, T = np.meshgrid(x,t) \n",
    "\n",
    "\n",
    "U_pred = griddata(X_star, u_pred.flatten(), (X,T), method='cubic')\n",
    "Error = np.abs(Exact - U_pred)\n",
    "percentError = 100*np.divide(Error, Exact)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a040f967",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "# Plot Exact and Predicted $(u,t)$\n",
    "    \n",
    "\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de46f6ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(15,15))\n",
    "ax.axis('off')\n",
    "\n",
    "print(\"--------- Errors ---------\")\n",
    "print('Percent error u: ',  100*error_u)\n",
    "print(\"--------------------------\")\n",
    "\n",
    "\n",
    "####### Row 0: u(t,x) ##################\n",
    "gs0 = gridspec.GridSpec(3, 2)\n",
    "gs0.update(top=1-0.06, bottom=1-1/3, left=0.15, right=0.85, wspace=0, hspace=1)\n",
    "\n",
    "\n",
    "########## Prediction ##################\n",
    "ax = plt.subplot(gs0[0, :])\n",
    "h = ax.imshow(U_pred.T, interpolation='nearest', cmap='rainbow',\n",
    "              extent=[t.min(), t.max(), x.min(), x.max()],\n",
    "              origin='lower', aspect='auto', vmin = 0, vmax = 1)\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(X_u_train[:,1], X_u_train[:,0], 'kx', label = 'Data (%d points)' % (u_train.shape[0]), markersize = 4, clip_on = False)\n",
    "\n",
    "line = np.linspace(x.min(), x.max(), 2)[:,None]\n",
    "ax.plot(t[2]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[5]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[10]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "\n",
    "ax.set_xlabel('$t$')\n",
    "ax.set_ylabel('$x$')\n",
    "ax.legend(frameon=False, loc = 'best')\n",
    "ax.set_title('$u(t,x) - Prediction$', fontsize = 10)\n",
    "\n",
    "########## Exact ##################\n",
    "ax = plt.subplot(gs0[1, :])\n",
    "i = ax.imshow(Exact.T, interpolation='nearest', cmap='rainbow',\n",
    "              extent=[t.min(), t.max(), x.min(), x.max()],\n",
    "              origin='lower', aspect='auto', vmin = 0, vmax = 1)\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "fig.colorbar(i, cax=cax)\n",
    "\n",
    "ax.plot(X_u_train[:,1], X_u_train[:,0], 'kx', label = 'Data (%d points)' % (u_train.shape[0]), markersize = 4, clip_on = False)\n",
    "\n",
    "line = np.linspace(x.min(), x.max(), 2)[:,None]\n",
    "ax.plot(t[2]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[5]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[10]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "\n",
    "ax.set_xlabel('$t$')\n",
    "ax.set_ylabel('$x$')\n",
    "ax.legend(frameon=False, loc = 'best')\n",
    "ax.set_title('$u(t,x)$ - Exact', fontsize = 10)\n",
    "\n",
    "########## Error ##################\n",
    "ax = plt.subplot(gs0[2, :])\n",
    "j = ax.imshow(percentError.T, interpolation='nearest', cmap='rainbow',\n",
    "              extent=[t.min(), t.max(), x.min(), x.max()],\n",
    "              origin='lower', aspect='auto', vmin = 0, vmax = 10)\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "fig.colorbar(j, cax=cax)\n",
    "\n",
    "ax.plot(X_u_train[:,1], X_u_train[:,0], 'kx', label = 'Data (%d points)' % (u_train.shape[0]), markersize = 4, clip_on = False)\n",
    "\n",
    "line = np.linspace(x.min(), x.max(), 2)[:,None]\n",
    "ax.plot(t[2]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[5]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[10]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "\n",
    "ax.set_xlabel('$t$')\n",
    "ax.set_ylabel('$x$')\n",
    "ax.legend(frameon=False, loc = 'best')\n",
    "ax.set_title('$u(t,x)$ - Percent Error', fontsize = 10)\n",
    "\n",
    "####### Row 1: u(t,x) slices ##################\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15,15))\n",
    "ax.axis('off')\n",
    "gs1 = gridspec.GridSpec(1, 3)\n",
    "gs1.update(top=1-1/3, bottom=0, left=0.1, right=0.9, wspace=0.5)\n",
    "\n",
    "ax = plt.subplot(gs1[:, 0])\n",
    "ax.plot(x,Exact[2,:], 'b-', linewidth = 2, label = 'Exact')\n",
    "ax.plot(x,U_pred[2,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$u(t,x)$')\n",
    "ax.set_title('$t = ' + str(t[2,0]) + '$', fontsize = 10)\n",
    "ax.axis('square')\n",
    "ax.set_xlim([0,1])\n",
    "ax.set_ylim([0,1])\n",
    "\n",
    "ax = plt.subplot(gs1[:, 1])\n",
    "ax.plot(x,Exact[5,:], 'b-', linewidth = 2, label = 'Exact')\n",
    "ax.plot(x,U_pred[5,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$u(t,x)$')\n",
    "ax.axis('square')\n",
    "ax.set_xlim([0,1])\n",
    "ax.set_ylim([0,1])\n",
    "ax.set_title('$t = ' + str(t[5,0]) + '$', fontsize = 10)\n",
    "ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.35), ncol=5, frameon=False)\n",
    "\n",
    "ax = plt.subplot(gs1[:, 2])\n",
    "ax.plot(x,Exact[10,:], 'b-', linewidth = 2, label = 'Exact')\n",
    "ax.plot(x,U_pred[10,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$u(t,x)$')\n",
    "ax.axis('square')\n",
    "ax.set_xlim([0,1])\n",
    "ax.set_ylim([0,1])\n",
    "ax.set_title('$t = ' + str(t[10,0]) + '$', fontsize = 10);\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "075549af",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "**Results**\n",
    "\n",
    "Above are the results of the PINN. The error for recreating the full solution field is $\\approx 10 \\%$, despite using only $N_u = 100$ data points. This shows the power of PINNs to learn from sparse measurements by augmenting the available observational data with knowledge of the underlying physics (i.e. governing equations).\n",
    "\n",
    "The three colourmaps show the PINN prediction, the exact solution from the numerical method and the relative error between these two fields. We can see that the errors are largest near $t=0$ and $x=0$, but that overall the agreement is very good.\n",
    "\n",
    "On the colourmap, we can see three vertical white lines, which show the location in time of the three profile plots of $u$ against $x$. The three heat profiles at these times are plotted against the exact solution found using numerical methods. The profiles can be seen to be in very good agreement, but show worse agreement.\n",
    "\n",
    "**Further Work**\n",
    "\n",
    "Congratulations, you have now trained your first 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_u\n",
    "- number of collocation training points N_f\n",
    "- number of layers in the network\n",
    "- number of neurons per layer\n",
    "- weightings for the data and PDE loss terms in the loss function (currently we use loss = loss_PDE + 5*loss_data)\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. Choosing boundary points for training could help to reduce the errors seen in these regions.\n",
    "\n",
    "Feel free to try out some of these changes if you like!\n",
    "\n",
    "There are 3 subsequent PINNs notebooks to follow, which look at the inverse heat equation problem and two Navier-Stokes fluid flow problems. \n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ffaec48",
   "metadata": {},
   "source": [
    "# 1D Heat Equation Inverse\n",
    "\n",
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "Remembering that in 1D, the heat equation can be written as:\n",
    "\n",
    "\\begin{equation}\n",
    "\\frac{\\partial u}{\\partial t} = k \\frac{\\partial^2 u}{\\partial x^2 },\n",
    "\\end{equation}\n",
    "\n",
    "where $k$ is a material parameter called the coefficient of thermal diffusivitiy. For this notebook, we have solved the above equation numerically on a domain of $x \\in [0,1]$ and $t \\in [0, 0.25]$. Solving this equation numerically gives us a spatiomtemporal domain $(x,t)$ and corresponding values of the solution $u$.\n",
    "\n",
    "**Now we will let $k$ be an unknown input parameter in the PINN**. In reality, we know the value of $k$, as we set it when solving the system numerically, but for the sake of this example let's imagine we do not know the value of $k$ when we come to use the PINN. This corresponds to real-world problems where we may have observational data, knowledge of the governing equations, but little information for some input parameters for the system.\n",
    "\n",
    "The network architecture for this example is the same as for the previous [example](#1D-Heat-Equation-Forwards). The only difference is that this time we do not know the value for $k$, and so in each training iteration we do not only updates the network weights and biases, but also the value of $k$. Through training, the network will then optimise the value of $k$ such that it fits with the observed data.\n",
    "    \n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37b7826c",
   "metadata": {},
   "outputs": [],
   "source": [
    "## this is the k value used to generate the data\n",
    "## we use this to compare to the value found by the PINN\n",
    "k_exact = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0914331",
   "metadata": {},
   "source": [
    "<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_u`\n",
    "- number of collocation training points `N_f`\n",
    "- number of layers in the network `layers`\n",
    "- number of neurons per layer `layers`\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20519b79",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_u = 100 #100 # number of data points\n",
    "N_f = 2000 # Collocation points \n",
    "# structure of network: two inputs (x,t) and one output u\n",
    "# 8 fully connected layers with 20 nodes per layer\n",
    "layers = [2, 20, 20, 20, 20, 20, 20, 20, 20, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11107983",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This code duplicated from above in case you have been playing with parameters\n",
    "data = scipy.io.loadmat(\"Data/heatEquation_data.mat\")\n",
    "t = data['t'].flatten()[:,None] # read in t and flatten into column vector\n",
    "x = data['x'].flatten()[:,None] # read in x and flatten into column vector\n",
    " # Exact represents the exact solution to the problem, from the Matlab script provided\n",
    "Exact = np.real(data['usol']).T # Exact has structure of nx times nt\n",
    "\n",
    "\n",
    "# print(\"t = \", t.transpose())\n",
    "# print(\"x = \", x.transpose())\n",
    "print(\"usol shape = \", Exact.shape)\n",
    "\n",
    "# We need to find all the x,t coordinate pairs in the domain\n",
    "X, T = np.meshgrid(x,t)\n",
    "\n",
    "# Flatten the coordinate grid into pairs of x,t coordinates\n",
    "X_star = np.hstack((X.flatten()[:,None], T.flatten()[:,None])) # coordinates x,t\n",
    "u_star = Exact.flatten()[:,None]   # corresponding solution value with each coordinate            \n",
    "\n",
    "\n",
    "print(\"X has shape \", X.shape, \", X_star has shape \", X_star.shape)\n",
    "    \n",
    "# Doman bounds (-1,1)\n",
    "lb = X_star.min(0)\n",
    "ub = X_star.max(0)  \n",
    "\n",
    "print(\"Lower bounds of x,t: \", lb)\n",
    "print(\"Upper bounds of x,t: \", ub)\n",
    "\n",
    "## train using internal points\n",
    "X_u_train = X_star\n",
    "u_train = u_star\n",
    "\n",
    "## Generate collocation points using Latin Hypercube sampling within the bounds of the spationtemporal coordinates\n",
    "# Generate N_f x,t coordinates within range of upper and lower bounds\n",
    "X_f_train = lb + (ub-lb)*lhs(2, N_f) # the 2 denotes the number of coordinates we have - x,t \n",
    "\n",
    "## In addition, we add the X_u_train coordinats from the boundaries to the X_f coordinate set\n",
    "X_f_train = np.vstack((X_f_train, X_u_train)) # stack up all training x,t coordinates for u and f \n",
    "\n",
    "## We downsample the boundary data to leave N_u 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(X_star.shape[0], N_u, replace=False)\n",
    "X_u_train = X_star[idx,:]\n",
    "u_train = u_star[idx,:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb252d1d",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "Now we will use all the same functions as before except we will modify `k` and the train function to handle a changing `k` value.\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90247095",
   "metadata": {},
   "outputs": [],
   "source": [
    "# k as tensorflow variable with inital value set\n",
    "k = tf.Variable([1.0], dtype=tf.float32)\n",
    "xvars, NNvars, tfvars, preds, optvars = init(X_u_train, u_train, layers, lb, ub,k)\n",
    "X, lb, ub, x, t, u = xvars\n",
    "layers, weights, biases = NNvars\n",
    "sess, x_tf, t_tf ,u_tf = tfvars\n",
    "u_pred, f_pred,=preds\n",
    "loss, optimizer,optimizer_Adam,train_op_Adam = optvars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "352ed46d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(sess, nIter,x_tf, t_tf, u_tf,x, t,u_train, loss, train_op_Adam, optimizer_Adam): \n",
    "    tf_dict = {x_tf: x,  t_tf: t, u_tf: u}\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",
    "            k_value = sess.run(k)\n",
    "            print('It: %d, Loss: %.3e, k: %.3f, Time: %.2f' % \n",
    "                          (it, loss_value, k_value, elapsed))\n",
    "            start_time = time()\n",
    "\n",
    "    optimizer_Adam.minimize(loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b0583a5",
   "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. 50000 was used to achieve excellent results. \n",
    "\n",
    "* If you are using a machine with GPUs please set `Train_iterations` to 50000 and this will run quickly\n",
    "* If you are using a well spec'ed laptop/computer and can leave this setting `Train_iterations=50000` but it will take upto 10 mins\n",
    "* If you are using a low spec'ed laptop/computer or cannont leave the code running `Train_iterations=20000` is the reccomended value (this solution may not be accurate)\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "994e8ec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training\n",
    "Train_iterations=50000\n",
    "train(sess, Train_iterations,x_tf, t_tf, u_tf,x, t,u_train, loss, train_op_Adam, optimizer_Adam)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aceb088b",
   "metadata": {},
   "outputs": [],
   "source": [
    "u_pred, f_pred=preds\n",
    "u_pred, f_pred = predict(sess, X_star[:,0:1], X_star[:,1:2], u_pred, f_pred)\n",
    "\n",
    "error_u = np.linalg.norm(u_star - u_pred, 2)/np.linalg.norm(u_star, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de55ec7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"f_pred mean = \", np.mean(f_pred))\n",
    "print('Error u: %e' % (error_u))\n",
    "print('Percent error u: ',  100*error_u)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "281887c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set grid values back to full data set size for plotting\n",
    "\n",
    "t = data['t'].flatten()[:,None]\n",
    "x = data['x'].flatten()[:,None]\n",
    "X, T = np.meshgrid(x,t) \n",
    "\n",
    "\n",
    "U_pred = griddata(X_star, u_pred.flatten(), (X,T), method='cubic')\n",
    "Error = np.abs(Exact - U_pred)\n",
    "percentError = 100*np.divide(Error, Exact)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6efd0892",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(15,15))\n",
    "ax.axis('off')\n",
    "\n",
    "print(\"--------- Errors ---------\")\n",
    "print('Percent error u: ',  100*error_u)\n",
    "print(\"--------------------------\")\n",
    "\n",
    "\n",
    "####### Row 0: u(t,x) ##################\n",
    "gs0 = gridspec.GridSpec(3, 2)\n",
    "gs0.update(top=1-0.06, bottom=1-1/3, left=0.15, right=0.85, wspace=0, hspace=1)\n",
    "\n",
    "\n",
    "########## Prediction ##################\n",
    "ax = plt.subplot(gs0[0, :])\n",
    "h = ax.imshow(U_pred.T, interpolation='nearest', cmap='rainbow',\n",
    "              extent=[t.min(), t.max(), x.min(), x.max()],\n",
    "              origin='lower', aspect='auto', vmin = 0, vmax = 1)\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(X_u_train[:,1], X_u_train[:,0], 'kx', label = 'Data (%d points)' % (u_train.shape[0]), markersize = 4, clip_on = False)\n",
    "\n",
    "line = np.linspace(x.min(), x.max(), 2)[:,None]\n",
    "ax.plot(t[2]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[5]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[10]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "\n",
    "ax.set_xlabel('$t$')\n",
    "ax.set_ylabel('$x$')\n",
    "ax.legend(frameon=False, loc = 'best')\n",
    "ax.set_title('$u(t,x) - Prediction$', fontsize = 10)\n",
    "\n",
    "########## Exact ##################\n",
    "ax = plt.subplot(gs0[1, :])\n",
    "i = ax.imshow(Exact.T, interpolation='nearest', cmap='rainbow',\n",
    "              extent=[t.min(), t.max(), x.min(), x.max()],\n",
    "              origin='lower', aspect='auto', vmin = 0, vmax = 1)\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "fig.colorbar(i, cax=cax)\n",
    "\n",
    "ax.plot(X_u_train[:,1], X_u_train[:,0], 'kx', label = 'Data (%d points)' % (u_train.shape[0]), markersize = 4, clip_on = False)\n",
    "\n",
    "line = np.linspace(x.min(), x.max(), 2)[:,None]\n",
    "ax.plot(t[2]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[5]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[10]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "\n",
    "ax.set_xlabel('$t$')\n",
    "ax.set_ylabel('$x$')\n",
    "ax.legend(frameon=False, loc = 'best')\n",
    "ax.set_title('$u(t,x)$ - Exact', fontsize = 10)\n",
    "\n",
    "########## Error ##################\n",
    "ax = plt.subplot(gs0[2, :])\n",
    "j = ax.imshow(percentError.T, interpolation='nearest', cmap='rainbow',\n",
    "              extent=[t.min(), t.max(), x.min(), x.max()],\n",
    "              origin='lower', aspect='auto', vmin = 0, vmax = 10)\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "fig.colorbar(j, cax=cax)\n",
    "\n",
    "ax.plot(X_u_train[:,1], X_u_train[:,0], 'kx', label = 'Data (%d points)' % (u_train.shape[0]), markersize = 4, clip_on = False)\n",
    "\n",
    "line = np.linspace(x.min(), x.max(), 2)[:,None]\n",
    "ax.plot(t[2]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[5]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "ax.plot(t[10]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "\n",
    "ax.set_xlabel('$t$')\n",
    "ax.set_ylabel('$x$')\n",
    "ax.legend(frameon=False, loc = 'best')\n",
    "ax.set_title('$u(t,x)$ - Percent Error', fontsize = 10)\n",
    "\n",
    "####### Row 1: u(t,x) slices ##################\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15,15))\n",
    "ax.axis('off')\n",
    "gs1 = gridspec.GridSpec(1, 3)\n",
    "gs1.update(top=1-1/3, bottom=0, left=0.1, right=0.9, wspace=0.5)\n",
    "\n",
    "ax = plt.subplot(gs1[:, 0])\n",
    "ax.plot(x,Exact[2,:], 'b-', linewidth = 2, label = 'Exact')\n",
    "ax.plot(x,U_pred[2,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$u(t,x)$')\n",
    "ax.set_title('$t = ' + str(t[2,0]) + '$', fontsize = 10)\n",
    "ax.axis('square')\n",
    "ax.set_xlim([0,1])\n",
    "ax.set_ylim([0,1])\n",
    "\n",
    "ax = plt.subplot(gs1[:, 1])\n",
    "ax.plot(x,Exact[5,:], 'b-', linewidth = 2, label = 'Exact')\n",
    "ax.plot(x,U_pred[5,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$u(t,x)$')\n",
    "ax.axis('square')\n",
    "ax.set_xlim([0,1])\n",
    "ax.set_ylim([0,1])\n",
    "ax.set_title('$t = ' + str(t[5,0]) + '$', fontsize = 10)\n",
    "ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.35), ncol=5, frameon=False)\n",
    "\n",
    "ax = plt.subplot(gs1[:, 2])\n",
    "ax.plot(x,Exact[10,:], 'b-', linewidth = 2, label = 'Exact')\n",
    "ax.plot(x,U_pred[10,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$u(t,x)$')\n",
    "ax.axis('square')\n",
    "ax.set_xlim([0,1])\n",
    "ax.set_ylim([0,1])\n",
    "ax.set_title('$t = ' + str(t[10,0]) + '$', fontsize = 10);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42ca918c",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "\n",
    "**Results**\n",
    "\n",
    "Above are the results of the PINN. The error for recreating the full solution field is $\\approx 5 \\%$, despite using only $N_u = 100$ data points. This shows the power of PINNs to learn solution fields from sparse measurements, even when some of the input parameters are unknown.\n",
    "\n",
    "The three colourmaps show the PINN prediction, the exact solution from the numerical method and the relative error between these two fields. We can see that the errors are largest near $t=0$ and $x=0$, but that overall the agreement is very good.\n",
    "\n",
    "On the colourmap, we can see three vertical white lines, which show the location in time of the three profile plots of $u$ against $x$. The three heat profiles at these times are plotted against the exact solution found using numerical methods. The profiles can be seen to be in very good agreement, but show worse agreement.\n",
    "\n",
    "**Further Work**\n",
    "\n",
    "Congratulations, you have now trained your second 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_u,\n",
    "- number of collocation training points N_f,\n",
    "- number of layers in the network,\n",
    "- number of neurons per layer,\n",
    "- weightings for the data and PDE loss terms in the loss function (currently we use loss = loss_PDE + 5*loss_data),\n",
    "- initialisation value for k,\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](#1D-Heat-Equation) to try out some of these changes if you like, or [here](#init) to alter the optimization method used.\n",
    " \n",
    "    \n",
    " </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8ca4e16",
   "metadata": {},
   "source": [
    "<hr>\n",
    "\n",
    "\n",
    "<div style=\"background-color: #e6ccff; padding: 10px;\">\n",
    "    \n",
    "# Next Steps\n",
    "    \n",
    "    \n",
    "Next we move on to a more complex example using the Navier Stokes Equation in the notebook linked below\n",
    "  \n",
    "[Navier-Stokes PINNs discovery of PDE’s](PINNs_NavierStokes_example.ipynb)\n",
    "    \n",
    " <hr>\n",
    "\n",
    "**Contact Details** \n",
    "\n",
    "For any questions, corrections or comments, please contact either:\n",
    "\n",
    "Michael MacRaild - scmm@leeds.ac.uk\n",
    "\n",
    "Fergus Shone - mm16f2s@leeds.ac.uk\n",
    "    \n",
    "</div>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}