{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PINN Solution of the Nonlinear Shroedinger PDE\n",
    "\n",
    "This PyTorch code demonstrates the application of physically-informed neural networks (PINN) in the solution of a well-known Nonlinear Shroedinger PDE with periodic boundary condition\n",
    "\\begin{aligned}\n",
    "  &\\mathbf{u}_t = i * \\mathbf{u}_{xx} + i * |\\mathbf{u}|^2 * \\mathbf{u} \\in [0, T]\\times[-L, L]\\\\\n",
    "  &\\mathbf{u}(0, x) = \\mathbf{u}_0(x), \\quad \\forall x \\in [-L, L] \\\\\n",
    "  &\\mathbf{u}(t, -L) = \\mathbf{u}(t, L), \\quad \\forall t \\in [0, T]\n",
    "\\end{aligned}\n",
    "where $i^2 = -1$, and $[-L, L]$ covers one full period, i.e. $T = 2L$.  If we let $\\mathbf{u} = u + i * v$, then we obtain two PDEs\n",
    "\\begin{aligned}\n",
    "u_t &= -v_{xx} - (u^2 + v^2) * v \\\\\n",
    "v_t &= u_{xx} + (u^2 + v^2) * u\n",
    "\\end{aligned}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Libraries and Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.optim\n",
    "import torch.optim.lr_scheduler as lr_scheduler\n",
    "from itertools import chain\n",
    "from collections import OrderedDict\n",
    "from pyDOE import lhs\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.io\n",
    "from scipy.interpolate import griddata\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "import matplotlib.gridspec as gridspec\n",
    "import time\n",
    "# set the random seed\n",
    "np.random.seed(1234)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working on mps\n"
     ]
    }
   ],
   "source": [
    "if torch.backends.mps.is_available():\n",
    "    device = torch.device('mps')\n",
    "elif torch.cuda.is_available():\n",
    "    device = torch.device('cuda')\n",
    "else:\n",
    "    device = torch.device('cpu')\n",
    "#\n",
    "print(f\"Working on {device}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = np.pi\n",
    "xlo = -L\n",
    "xhi = L\n",
    "period = xhi - xlo\n",
    "tlo = 0.0\n",
    "thi = 2.0\n",
    "sqrt2 = torch.tensor(np.sqrt(2)).float().to(device)\n",
    "# taking A = 1 and B = 1\n",
    "u0_ten = lambda x: 2.0/(2.0 - sqrt2 * torch.cos(x)) - 1.0\n",
    "v0_ten = lambda x: torch.zeros_like(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Physics-informed Neural Networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the deep neural network\n",
    "class DNN(torch.nn.Module):\n",
    "    def __init__(self, layers):\n",
    "        super(DNN, self).__init__()\n",
    "        # parameters\n",
    "        self.depth = len(layers) - 1\n",
    "        # set up layer order dict\n",
    "        self.activation = torch.nn.Tanh\n",
    "        layer_list = list()\n",
    "        for i in range(self.depth - 1): \n",
    "            layer_list.append(\n",
    "                ('layer_%d' % i, torch.nn.Linear(layers[i], layers[i+1]))\n",
    "            )\n",
    "            layer_list.append(('activation_%d' % i, self.activation()))\n",
    "            \n",
    "        layer_list.append(\n",
    "            ('layer_%d' % (self.depth - 1), torch.nn.Linear(layers[-2], layers[-1]))\n",
    "        )\n",
    "        layerDict = OrderedDict(layer_list)\n",
    "        # deploy layers\n",
    "        self.layers = torch.nn.Sequential(layerDict)\n",
    "        self.layers[0].weight = torch.load('initial_weights_PBC.pt')\n",
    "    def forward(self, x):\n",
    "        out = self.layers(x)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PhysicsInformedNN():\n",
    "    def __init__(self, period, m,  X_IC, u_IC, v_IC, X_PDE, layers):\n",
    "        # IC data point\n",
    "        m_vec = np.expand_dims(np.arange(1, m + 1), axis = 0)\n",
    "        self.ms = torch.tensor(2.0 * np.pi/period * m_vec).float().to(device)\n",
    "        layers[0] = int(2 * m + 2)\n",
    "        self.t_IC = torch.tensor(X_IC[:, 0:1]).float().to(device)\n",
    "        self.x_IC = torch.tensor(X_IC[:, 1:2]).float().to(device)\n",
    "        self.u_IC = torch.tensor(u_IC).float().to(device)\n",
    "        self.v_IC = torch.tensor(v_IC).float().to(device)\n",
    " \n",
    "        \n",
    "        self.period = torch.tensor(period).float().to(device)\n",
    "        # PDE data, gradients will be computed on these points so requires_grad = True\n",
    "        self.t_PDE = torch.tensor(X_PDE[:, 0:1], requires_grad=True).float().to(device)\n",
    "        self.x_PDE = torch.tensor(X_PDE[:, 1:2], requires_grad=True).float().to(device)\n",
    "        \n",
    "        # layers to build Neural Net\n",
    "        self.layers = layers\n",
    "        # deep neural networks\n",
    "        self.dnn = DNN(layers).to(device)    \n",
    "        # prepare the optimizer\n",
    "        self.optimizer_Adam = torch.optim.Adam(self.dnn.parameters(), lr = 1e-3)\n",
    "        # add a learning rate scheduler\n",
    "        self.optimizer_LBFGS = torch.optim.LBFGS(\n",
    "            self.dnn.parameters(), \n",
    "            lr=1.0, \n",
    "            max_iter=10000, \n",
    "            max_eval=5000, \n",
    "            history_size=50,\n",
    "            tolerance_grad=1e-7, \n",
    "            tolerance_change=1.0 * np.finfo(float).eps,\n",
    "            line_search_fn=\"strong_wolfe\"       # can be \"strong_wolfe\"\n",
    "        )      \n",
    "        self.scheduler = lr_scheduler.ExponentialLR(self.optimizer_Adam, gamma=0.99)\n",
    "        self.iter = 0\n",
    "    # evaluater neural network, no transformation\n",
    "    def NN_eval(self, t, x):  \n",
    "        x_trans = torch.matmul(x, self.ms)\n",
    "        NN = self.dnn(torch.cat([t, torch.ones_like(x), torch.cos(x_trans), torch.sin(x_trans)], dim = 1))\n",
    "        u = NN[:, 0][:, None]\n",
    "        v = NN[:, 1][:, None]\n",
    "        unew =  u\n",
    "        vnew =  u\n",
    "        return unew,vnew\n",
    "    # compute the PDE\n",
    "    def pde_eval(self, t, x):\n",
    "        \"\"\" The pytorch autograd version of calculating residual \"\"\"\n",
    "        u, v = self.NN_eval(t, x)\n",
    "        # compute the derivatives for u\n",
    "        u_t    = torch.autograd.grad(u,     t, grad_outputs = torch.ones_like(u), retain_graph = True, create_graph=True)[0]\n",
    "        u_x    = torch.autograd.grad(u,     x, grad_outputs = torch.ones_like(u), retain_graph = True, create_graph=True)[0]\n",
    "        u_xx   = torch.autograd.grad(u_x,   x, grad_outputs = torch.ones_like(u), retain_graph = True, create_graph=True)[0]\n",
    "        v_t    = torch.autograd.grad(v,     t, grad_outputs = torch.ones_like(u), retain_graph = True, create_graph=True)[0]\n",
    "        v_x    = torch.autograd.grad(v,     x, grad_outputs = torch.ones_like(u), retain_graph = True, create_graph=True)[0]\n",
    "        v_xx   = torch.autograd.grad(v_x,   x, grad_outputs = torch.ones_like(u), retain_graph = True, create_graph=True)[0]   \n",
    "        # pre-compute the norm of \\mathbf{u}, so that it can be re-used\n",
    "        U_norm = torch.pow(u, 2.0) + torch.pow(v, 2.0)\n",
    "        Eq1    = u_t + v_xx + U_norm * v\n",
    "        Eq2    = v_t - u_xx - U_norm * u\n",
    "        return Eq1, Eq2\n",
    "    # compute the total loss for the second-order optimizer\n",
    "    def loss_func(self):\n",
    "        # reset the gradient\n",
    "        self.optimizer_LBFGS.zero_grad()\n",
    "        # compute IC loss\n",
    "        uIC_pred, vIC_pred = self.NN_eval(self.t_IC, self.x_IC)\n",
    "        loss_IC = torch.mean((self.u_IC - uIC_pred)**2.0) + torch.mean((self.v_IC - vIC_pred)**2.0)\n",
    "   \n",
    "       \n",
    "        # compute PDE loss\n",
    "        pde1_pred, pde2_pred = self.pde_eval(self.t_PDE, self.x_PDE)\n",
    "        loss_PDE = torch.mean(pde1_pred ** 2) + torch.mean(pde2_pred ** 2)     \n",
    "        # compute the total loss, it can be weighted\n",
    "        loss = loss_IC  + loss_PDE\n",
    "        # backward propagation\n",
    "        loss.backward()\n",
    "        # increase the iteration counter\n",
    "        self.iter += 1\n",
    "        # output\n",
    "        # output the progress\n",
    "        if self.iter % 1000 == 0:\n",
    "            end_time = time.time()\n",
    "            print('Iter %5d, Total: %10.4e, Time: %.2f secs' % (self.iter, loss.item(), end_time - self.start_time))\n",
    "            print('IC: %10.4e,  PDE: %10.4e' % (loss_IC.item(),  loss_PDE.item()))\n",
    "            self.start_time = end_time\n",
    "        return loss\n",
    "    #\n",
    "    def train(self, nIter):\n",
    "        # start the timer\n",
    "        start_time = time.time()        \n",
    "        # start the training with Adam first\n",
    "        self.dnn.train()\n",
    "        print('Starting with Adam')\n",
    "        for epoch in range(nIter):\n",
    "            # compute IC loss\n",
    "            uIC_pred, vIC_pred = self.NN_eval(self.t_IC, self.x_IC)\n",
    "            loss_IC = torch.mean((self.u_IC - uIC_pred)**2.0) + torch.mean((self.v_IC - vIC_pred)**2.0)\n",
    "            # compute PDE loss\n",
    "            pde1_pred, pde2_pred = self.pde_eval(self.t_PDE, self.x_PDE)\n",
    "            loss_PDE = torch.mean(pde1_pred ** 2) + torch.mean(pde2_pred ** 2)   \n",
    "            # compute the total loss, it can be weighted\n",
    "            loss = loss_IC +  loss_PDE\n",
    "            # Backward and optimize\n",
    "            self.optimizer_Adam.zero_grad()\n",
    "            loss.backward()\n",
    "            self.optimizer_Adam.step() \n",
    "            # output the progress\n",
    "            if (epoch + 1) % 1000 == 0:\n",
    "                end_time = time.time()\n",
    "                print('Iter %5d, Total: %10.4e, Time: %.2f secs' % (epoch + 1, loss.item(), end_time - start_time))\n",
    "                print('IC: %10.4e,  PDE: %10.4e' % (loss_IC.item(), loss_PDE.item()))\n",
    "                start_time = end_time\n",
    "                # change the learning rate\n",
    "                self.scheduler.step()\n",
    "\n",
    "        self.start_time = time.time()\n",
    "        self.optimizer_LBFGS.step(self.loss_func)    \n",
    "\n",
    "        \n",
    "    def predict(self, X):\n",
    "        t = torch.tensor(X[:, 0:1], requires_grad=True).float().to(device)\n",
    "        x = torch.tensor(X[:, 1:2], requires_grad=True).float().to(device)\n",
    "        self.dnn.eval()\n",
    "        u, v = self.NN_eval(t, x)\n",
    "        u = u.detach().cpu().numpy()\n",
    "        v = v.detach().cpu().numpy()\n",
    "        return u,v\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting with Adam\n",
      "Iter  1000, Total: 1.2021e-01, Time: 63.18 secs\n",
      "IC: 1.0186e-01,  PDE: 1.8346e-02\n",
      "Iter  2000, Total: 8.9401e-02, Time: 109.19 secs\n",
      "IC: 7.6479e-02,  PDE: 1.2921e-02\n",
      "Iter  3000, Total: 1.1619e-01, Time: 111.46 secs\n",
      "IC: 1.0102e-01,  PDE: 1.5165e-02\n",
      "Iter  4000, Total: 9.8283e-02, Time: 110.35 secs\n",
      "IC: 8.4315e-02,  PDE: 1.3967e-02\n",
      "Iter  5000, Total: 8.7488e-02, Time: 111.27 secs\n",
      "IC: 7.3649e-02,  PDE: 1.3839e-02\n",
      "Iter  6000, Total: 6.9147e-02, Time: 111.91 secs\n",
      "IC: 5.6427e-02,  PDE: 1.2720e-02\n",
      "Iter  7000, Total: 5.6942e-02, Time: 112.02 secs\n",
      "IC: 4.4917e-02,  PDE: 1.2025e-02\n",
      "Iter  8000, Total: 4.5652e-02, Time: 111.51 secs\n",
      "IC: 3.3452e-02,  PDE: 1.2200e-02\n",
      "Iter  9000, Total: 2.9806e-02, Time: 110.31 secs\n",
      "IC: 2.2435e-02,  PDE: 7.3716e-03\n",
      "Iter 10000, Total: 2.7080e-02, Time: 110.49 secs\n",
      "IC: 2.1258e-02,  PDE: 5.8217e-03\n",
      "Iter 11000, Total: 5.3088e-02, Time: 110.63 secs\n",
      "IC: 4.2236e-02,  PDE: 1.0852e-02\n",
      "Iter 12000, Total: 3.0874e-02, Time: 109.92 secs\n",
      "IC: 2.1607e-02,  PDE: 9.2667e-03\n",
      "Iter 13000, Total: 2.0509e-02, Time: 110.14 secs\n",
      "IC: 1.3500e-02,  PDE: 7.0090e-03\n",
      "Iter 14000, Total: 1.4518e-02, Time: 108.09 secs\n",
      "IC: 8.4247e-03,  PDE: 6.0934e-03\n",
      "Iter 15000, Total: 1.1506e-02, Time: 108.79 secs\n",
      "IC: 6.3035e-03,  PDE: 5.2029e-03\n",
      "Iter 16000, Total: 1.0442e-02, Time: 110.75 secs\n",
      "IC: 5.4521e-03,  PDE: 4.9901e-03\n",
      "Iter 17000, Total: 8.2860e-03, Time: 109.74 secs\n",
      "IC: 4.9148e-03,  PDE: 3.3713e-03\n",
      "Iter 18000, Total: 7.0270e-03, Time: 108.90 secs\n",
      "IC: 3.9624e-03,  PDE: 3.0646e-03\n",
      "Iter 19000, Total: 5.7892e-03, Time: 109.51 secs\n",
      "IC: 3.1984e-03,  PDE: 2.5907e-03\n",
      "Iter 20000, Total: 5.9812e-03, Time: 110.23 secs\n",
      "IC: 3.1608e-03,  PDE: 2.8203e-03\n",
      "Iter 21000, Total: 3.4258e-03, Time: 111.23 secs\n",
      "IC: 1.8809e-03,  PDE: 1.5449e-03\n",
      "Iter 22000, Total: 3.1079e-03, Time: 110.79 secs\n",
      "IC: 1.7514e-03,  PDE: 1.3565e-03\n",
      "Iter 23000, Total: 3.9549e-03, Time: 111.32 secs\n",
      "IC: 1.9738e-03,  PDE: 1.9811e-03\n",
      "Iter 24000, Total: 2.5228e-03, Time: 109.99 secs\n",
      "IC: 1.4640e-03,  PDE: 1.0588e-03\n",
      "Iter 25000, Total: 2.9999e-03, Time: 108.83 secs\n",
      "IC: 1.2847e-03,  PDE: 1.7152e-03\n",
      "Iter 26000, Total: 2.9609e-03, Time: 108.84 secs\n",
      "IC: 1.4102e-03,  PDE: 1.5507e-03\n",
      "Iter 27000, Total: 3.6604e-02, Time: 109.36 secs\n",
      "IC: 3.0736e-02,  PDE: 5.8678e-03\n",
      "Iter 28000, Total: 4.7405e-03, Time: 109.35 secs\n",
      "IC: 1.2969e-03,  PDE: 3.4436e-03\n",
      "Iter 29000, Total: 1.9480e-03, Time: 109.11 secs\n",
      "IC: 1.0381e-03,  PDE: 9.0986e-04\n",
      "Iter 30000, Total: 1.7285e-03, Time: 108.72 secs\n",
      "IC: 9.8706e-04,  PDE: 7.4148e-04\n",
      "Iter 31000, Total: 2.9940e-03, Time: 109.12 secs\n",
      "IC: 1.2282e-03,  PDE: 1.7657e-03\n",
      "Iter 32000, Total: 4.1573e-03, Time: 109.07 secs\n",
      "IC: 2.7886e-03,  PDE: 1.3687e-03\n",
      "Iter 33000, Total: 1.5878e-03, Time: 109.01 secs\n",
      "IC: 9.1618e-04,  PDE: 6.7165e-04\n",
      "Iter 34000, Total: 2.1715e-03, Time: 109.52 secs\n",
      "IC: 9.8593e-04,  PDE: 1.1856e-03\n",
      "Iter 35000, Total: 1.3982e-03, Time: 109.53 secs\n",
      "IC: 8.4071e-04,  PDE: 5.5745e-04\n",
      "Iter 36000, Total: 1.3868e-03, Time: 108.56 secs\n",
      "IC: 8.1019e-04,  PDE: 5.7661e-04\n",
      "Iter 37000, Total: 1.6523e-03, Time: 109.29 secs\n",
      "IC: 8.4550e-04,  PDE: 8.0678e-04\n",
      "Iter 38000, Total: 1.2938e-03, Time: 109.85 secs\n",
      "IC: 7.7064e-04,  PDE: 5.2315e-04\n",
      "Iter 39000, Total: 1.3783e-03, Time: 109.04 secs\n",
      "IC: 7.9389e-04,  PDE: 5.8436e-04\n",
      "Iter 40000, Total: 1.0655e-03, Time: 109.66 secs\n",
      "IC: 7.3064e-04,  PDE: 3.3481e-04\n",
      "Iter 41000, Total: 1.1079e-03, Time: 108.55 secs\n",
      "IC: 7.4191e-04,  PDE: 3.6604e-04\n",
      "Iter 42000, Total: 1.0184e-03, Time: 108.36 secs\n",
      "IC: 7.1979e-04,  PDE: 2.9858e-04\n",
      "Iter 43000, Total: 1.0334e-03, Time: 96.81 secs\n",
      "IC: 7.2929e-04,  PDE: 3.0414e-04\n",
      "Iter 44000, Total: 1.0514e-03, Time: 77.29 secs\n",
      "IC: 7.0851e-04,  PDE: 3.4285e-04\n",
      "Iter 45000, Total: 1.6036e-03, Time: 75.19 secs\n",
      "IC: 7.9338e-04,  PDE: 8.1025e-04\n",
      "Iter 46000, Total: 9.5743e-04, Time: 75.85 secs\n",
      "IC: 7.0221e-04,  PDE: 2.5522e-04\n",
      "Iter 47000, Total: 9.5155e-04, Time: 76.43 secs\n",
      "IC: 7.0014e-04,  PDE: 2.5141e-04\n",
      "Iter 48000, Total: 9.3837e-04, Time: 76.39 secs\n",
      "IC: 6.9996e-04,  PDE: 2.3841e-04\n",
      "Iter 49000, Total: 9.4411e-04, Time: 69.20 secs\n",
      "IC: 7.0107e-04,  PDE: 2.4305e-04\n",
      "Iter 50000, Total: 9.3967e-04, Time: 63.88 secs\n",
      "IC: 6.9012e-04,  PDE: 2.4955e-04\n",
      "Iter  1000, Total: 7.5780e-04, Time: 112.27 secs\n",
      "IC: 5.5542e-04,  PDE: 2.0239e-04\n",
      "Iter  2000, Total: 6.2911e-04, Time: 110.96 secs\n",
      "IC: 4.7358e-04,  PDE: 1.5553e-04\n",
      "Iter  3000, Total: 5.3255e-04, Time: 109.64 secs\n",
      "IC: 3.9539e-04,  PDE: 1.3716e-04\n",
      "Iter  4000, Total: 4.4966e-04, Time: 102.10 secs\n",
      "IC: 3.2000e-04,  PDE: 1.2967e-04\n",
      "Iter  5000, Total: 3.5149e-04, Time: 95.71 secs\n",
      "IC: 2.4209e-04,  PDE: 1.0939e-04\n"
     ]
    }
   ],
   "source": [
    "layers = [2, 32, 32, 32, 32, 32, 32, 32, 2]\n",
    "N_IC = 96\n",
    "ptsIC=np.load('ptsIC.npy')\n",
    "ptsBC=np.load('ptsBC.npy')\n",
    "ptsPDE=np.load('random_data.npy')\n",
    "ptsPDE= np.vstack((ptsPDE,ptsBC))\n",
    "x_IC = np.expand_dims(np.linspace(xlo, xhi, N_IC), axis = 1)\n",
    "uIC = u0(x_IC)\n",
    "vIC = v0(x_IC)\n",
    "uBC = np.load('uBC.npy')\n",
    "vBC = np.load('vBC.npy')\n",
    "model = PhysicsInformedNN(period,  ptsIC, uIC, vIC,  ptsPDE, layers)\n",
    "model.train(50000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = np.linspace(tlo, thi, 201)\n",
    "x = np.linspace(xlo, xhi, 501)\n",
    "T, X = np.meshgrid(t, x)\n",
    "pts_flat = np.hstack((T.flatten()[:, None], X.flatten()[:, None]))\n",
    "u_pred, v_pred = model.predict(pts_flat)           \n",
    "u_pred = griddata(pts_flat, u_pred.flatten(), (T, X), method='cubic')\n",
    "v_pred = griddata(pts_flat, v_pred.flatten(), (T, X), method='cubic')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = scipy.io.loadmat('Data/NLS.mat')\n",
    "t = data['t'].flatten()[:,None]\n",
    "x2 = data['x'].flatten()[:,None]\n",
    "u_sol = np.real(data['Exact']).T\n",
    "v_sol = np.imag(data['Exact']).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 150016 characters]",
      "text/plain": [
       "<Figure size 1800x600 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "####### Row 1: u(t,x) slices ################## \n",
    "\n",
    "\"\"\" The aesthetic setting has changed. \"\"\"\n",
    "\n",
    "fig = plt.figure(figsize=(18, 6))\n",
    "#\n",
    "ax = plt.subplot(1, 3, 1)\n",
    "ax.plot(x2, u_sol[0,:], 'bo-', linewidth = 2, label = 'Exact')       \n",
    "ax.plot(x, u_pred[:, 0], 'rx--', linewidth = 2, label = 'u')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$u(t,x)$')    \n",
    "ax.set_title('$t = %.1f$' %(t[0]), fontsize = 24)\n",
    "\n",
    "\n",
    "plt.locator_params(axis = 'y', nbins = 5)\n",
    "plt.locator_params(axis = 'x', nbins = 5)\n",
    "for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +\n",
    "             ax.get_xticklabels() + ax.get_yticklabels()):\n",
    "    item.set_fontsize(24)\n",
    "\n",
    "ax = plt.subplot(1, 3, 2)\n",
    "ax.plot(x,u_sol[100,:], 'bo-', linewidth = 2, label = 'Exact')       \n",
    "ax.plot(x, u_pred[:, 100], 'rx--', linewidth = 2, label = 'u')\n",
    "\n",
    "ax.set_xlabel('x')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax.set_title('$t = %.1f$' %(t[100]), fontsize = 24)\n",
    "ax.legend(\n",
    "    loc='upper center', \n",
    "    bbox_to_anchor=(0.5, -0.15), \n",
    "    ncol=5, \n",
    "    frameon=False, \n",
    "    prop={'size': 24}\n",
    ")\n",
    "plt.locator_params(axis = 'y', nbins = 5)\n",
    "plt.locator_params(axis = 'x', nbins = 5)\n",
    "for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +\n",
    "             ax.get_xticklabels() + ax.get_yticklabels()):\n",
    "    item.set_fontsize(24)\n",
    "\n",
    "ax = plt.subplot(1, 3, 3)\n",
    "ax.plot(x,u_sol[200,:], 'bo-', linewidth = 2, label = 'Exact')       \n",
    "ax.plot(x, u_pred[:, 200], 'rx--', linewidth = 2, label = 'u')\n",
    "\n",
    "ax.set_xlabel('$x$')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax.set_title('$t = %.1f$' %(t[200]), fontsize = 24)\n",
    "plt.locator_params(axis = 'y', nbins = 5)\n",
    "plt.locator_params(axis = 'x', nbins = 5)\n",
    "for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +\n",
    "             ax.get_xticklabels() + ax.get_yticklabels()):\n",
    "    item.set_fontsize(24)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(18, 6))\n",
    "#\n",
    "ax = plt.subplot(1, 3, 1)\n",
    "ax.plot(x2, v_sol[0,:], 'bo-', linewidth = 2, label = 'Exact')       \n",
    "ax.plot(x, v_pred[:, 0], 'rx--', linewidth = 2, label = 'u')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$v(t,x)$')    \n",
    "ax.set_title('$t = %.1f$' %(t[0]), fontsize = 24)\n",
    "\n",
    "\n",
    "plt.locator_params(axis = 'y', nbins = 5)\n",
    "plt.locator_params(axis = 'x', nbins = 5)\n",
    "for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +\n",
    "             ax.get_xticklabels() + ax.get_yticklabels()):\n",
    "    item.set_fontsize(24)\n",
    "\n",
    "ax = plt.subplot(1, 3, 2)\n",
    "ax.plot(x,v_sol[100,:], 'bo-', linewidth = 2, label = 'Exact')       \n",
    "ax.plot(x, v_pred[:, 100], 'rx--', linewidth = 2, label = 'u')\n",
    "\n",
    "ax.set_xlabel('x')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax.set_title('$t = %.1f$' %(t[100]), fontsize = 24)\n",
    "ax.legend(\n",
    "    loc='upper center', \n",
    "    bbox_to_anchor=(0.5, -0.15), \n",
    "    ncol=5, \n",
    "    frameon=False, \n",
    "    prop={'size': 24}\n",
    ")\n",
    "plt.locator_params(axis = 'y', nbins = 5)\n",
    "plt.locator_params(axis = 'x', nbins = 5)\n",
    "for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +\n",
    "             ax.get_xticklabels() + ax.get_yticklabels()):\n",
    "    item.set_fontsize(24)\n",
    "\n",
    "ax = plt.subplot(1, 3, 3)\n",
    "ax.plot(x,v_sol[200,:], 'bo-', linewidth = 2, label = 'Exact')       \n",
    "ax.plot(x, v_pred[:, 200], 'rx--', linewidth = 2, label = 'u')\n",
    "\n",
    "ax.set_xlabel('$x$')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax.set_title('$t = %.1f$' %(t[200]), fontsize = 24)\n",
    "plt.locator_params(axis = 'y', nbins = 5)\n",
    "plt.locator_params(axis = 'x', nbins = 5)\n",
    "for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +\n",
    "             ax.get_xticklabels() + ax.get_yticklabels()):\n",
    "    item.set_fontsize(24)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def relative_error_l2(pred,exact):\n",
    "    error_l2 = np.sqrt(np.sum(np.power(pred - exact,2)))\n",
    "    relative = error_l2/np.sqrt(np.sum(np.power(exact,2)))\n",
    "    return relative\n",
    "def relative_error_l1(pred,exact):\n",
    "    error_l1 = np.sum(np.abs(pred-exact))\n",
    "    relative = error_l1/np.sum(np.abs(exact))\n",
    "    return relative\n",
    "def relative_error_linf(pred,exact):\n",
    "    error_linf = np.max(np.abs(pred-exact))\n",
    "    relative = error_linf/np.max(np.abs(exact))\n",
    "    return relative"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "l2: 1.004449529572271\n",
      "l1: 1.0024359215885241\n",
      "linf: 1.1825159622595356\n"
     ]
    }
   ],
   "source": [
    "print(f'l2: {relative_error_l2(u_pred.T,u_sol)}')\n",
    "print(f'l1: {relative_error_l1(u_pred.T,u_sol)}')\n",
    "print(f'linf: {relative_error_linf(u_pred.T,u_sol)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'l2: {relative_error_l2(v_pred.T,v_sol)}')\n",
    "print(f'l1: {relative_error_l1(v_pred.T,v_sol)}')\n",
    "print(f'linf: {relative_error_linf(v_pred.T,v_sol)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 133888 characters]",
      "text/plain": [
       "<Figure size 2400x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(24, 6))\n",
    "\n",
    "ax = plt.subplot(1, 3, 1)\n",
    "h = ax.imshow(u_sol.T, interpolation='nearest', cmap='rainbow', \n",
    "              extent=[tlo, thi, x.min(), x.max()], \n",
    "              origin='lower', aspect='auto')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.10)\n",
    "cbar = fig.colorbar(h, cax=cax)\n",
    "cbar.ax.tick_params(labelsize=24) \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax.set_xlabel('$t$', size=24)\n",
    "ax.set_ylabel('$x$', size=24)\n",
    "\n",
    "ax.set_title('$GS: u_{exact}(t,x)$', fontsize = 24) # font size doubled\n",
    "ax.tick_params(labelsize=24)\n",
    "\n",
    "ax = plt.subplot(1, 3, 2)\n",
    "h = ax.imshow(u_pred, interpolation='nearest', cmap='rainbow', \n",
    "              extent=[tlo, thi, x.min(), x.max()], \n",
    "              origin='lower', aspect='auto')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.10)\n",
    "cbar = fig.colorbar(h, cax=cax)\n",
    "cbar.ax.tick_params(labelsize=24) \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax.set_xlabel('$t$', size=24)\n",
    "\n",
    "\n",
    "ax.set_title('$NLS: u_{pred}(t,x)$', fontsize = 24) # font size doubled\n",
    "ax.tick_params(labelsize=24)\n",
    "\n",
    "ax = plt.subplot(1, 3, 3)\n",
    "h = ax.imshow(np.abs(u_pred-u_sol.T), interpolation='nearest', cmap='rainbow', \n",
    "              extent=[tlo, thi, x.min(), x.max()], \n",
    "              origin='lower', aspect='auto')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.10)\n",
    "cbar = fig.colorbar(h, cax=cax)\n",
    "cbar.ax.tick_params(labelsize=24) \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax.set_xlabel('$t$', size=24)\n",
    "\n",
    "\n",
    "ax.set_title('$NLS(Error): u(t,x)$', fontsize = 24) # font size doubled\n",
    "ax.tick_params(labelsize=24)\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(24, 6))\n",
    "\n",
    "ax = plt.subplot(1, 3, 1)\n",
    "h = ax.imshow(v_sol.T, interpolation='nearest', cmap='rainbow', \n",
    "              extent=[tlo, thi, x.min(), x.max()], \n",
    "              origin='lower', aspect='auto')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.10)\n",
    "cbar = fig.colorbar(h, cax=cax)\n",
    "cbar.ax.tick_params(labelsize=24) \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax.set_xlabel('$t$', size=24)\n",
    "ax.set_ylabel('$x$', size=24)\n",
    "\n",
    "ax.set_title('$NLS: v_{exact}(t,x)$', fontsize = 24) # font size doubled\n",
    "ax.tick_params(labelsize=24)\n",
    "\n",
    "ax = plt.subplot(1, 3, 2)\n",
    "h = ax.imshow(v_pred, interpolation='nearest', cmap='rainbow', \n",
    "              extent=[tlo, thi, x.min(), x.max()], \n",
    "              origin='lower', aspect='auto')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.10)\n",
    "cbar = fig.colorbar(h, cax=cax)\n",
    "cbar.ax.tick_params(labelsize=24) \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax.set_xlabel('$t$', size=24)\n",
    "\n",
    "ax.set_xticks([0 ,0.5 ,1])\n",
    "ax.set_yticks([-1 ,0 ,1])\n",
    "ax.set_title('$NLS: v_{pred}(t,x)$', fontsize = 24) # font size doubled\n",
    "ax.tick_params(labelsize=24)\n",
    "\n",
    "ax = plt.subplot(1, 3, 3)\n",
    "h = ax.imshow(np.abs(v_pred-v_sol.T), interpolation='nearest', cmap='rainbow', \n",
    "              extent=[tlo, thi, x.min(), x.max()], \n",
    "              origin='lower', aspect='auto')\n",
    "divider = make_axes_locatable(ax)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.10)\n",
    "cbar = fig.colorbar(h, cax=cax)\n",
    "cbar.ax.tick_params(labelsize=24) \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax.set_xlabel('$t$', size=24)\n",
    "\n",
    "ax.set_xticks([0 ,0.5 ,1])\n",
    "ax.set_yticks([-1 ,0 ,1])\n",
    "ax.set_title('$NLS(Error): v(t,x)$', fontsize = 24) # font size doubled\n",
    "ax.tick_params(labelsize=24)\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "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.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
