{
 "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": 2,
   "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",
    "\n",
    "np.random.seed(1234)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": 4,
   "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 = lambda x: 2.0/(2.0 - np.sqrt(2) * np.cos(x)) - 1.0\n",
    "v0 = lambda x: np.zeros_like(x)\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": 5,
   "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_weight.pt')\n",
    "    def forward(self, x):\n",
    "        out = self.layers(x)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PhysicsInformedNN():\n",
    "    def __init__(self, period,  X_IC, u_IC, v_IC, X_BC, u_BC, v_BC, X_PDE, layers,n_batches, batch_size):\n",
    "        # IC data point\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",
    "        # BC data point\n",
    "        self.t_BC = torch.tensor(X_BC[:, 0:1]).float().to(device)\n",
    "        self.x_BC = torch.tensor(X_BC[:, 1:2]).float().to(device)\n",
    "        self.u_BC = torch.tensor(u_BC).float().to(device)\n",
    "        self.v_BC = torch.tensor(v_BC).float().to(device) \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",
    "        # equation related parameters\n",
    "        \n",
    "        # layers to build Neural Net\n",
    "        self.layers = layers\n",
    "        # deep neural networks\n",
    "        self.n_batches = n_batches\n",
    "        self.batch_size = batch_size\n",
    "\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",
    "        \n",
    "        \n",
    "        self.iter = 0\n",
    "    # evaluater neural network, no transformation\n",
    "    def NN_eval(self, t, x): \n",
    "        NN = self.dnn(torch.cat([t, x], dim = 1)) \n",
    "        u = NN[:, 0][:, None]\n",
    "        v = NN[:, 1][:, None]\n",
    "        return u, v\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",
    "        # compute PBC loss\n",
    "        uBC_pred_left, vBC_pred_left = self.NN_eval(self.t_BC, self.x_BC)\n",
    "        uBC_pred_right, vBC_pred_right = self.NN_eval(self.t_BC, self.x_BC+self.period)\n",
    "        loss_BC = torch.mean(( uBC_pred_left - uBC_pred_right)**2.0) + torch.mean((vBC_pred_left - vBC_pred_right)**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_BC + 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, PBC: %10.4e, PDE: %10.4e' % (loss_IC.item(), loss_BC.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",
    "            for batch_id in range(self.n_batches):\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 PBC loss\n",
    "                uBC_pred_left, vBC_pred_left = self.NN_eval(self.t_BC, self.x_BC)\n",
    "                uBC_pred_right, vBC_pred_right = self.NN_eval(self.t_BC, self.x_BC+self.period)\n",
    "                loss_BC = torch.mean(( uBC_pred_left - uBC_pred_right)**2.0) + torch.mean((vBC_pred_left - vBC_pred_right)**2.0)\n",
    "        \n",
    "       \n",
    "        # compute PDE loss\n",
    "                pde1_pred, pde2_pred = self.pde_eval(self.t_PDE[batch_id*self.batch_size:(batch_id + 1)*self.batch_size, :], self.x_PDE[batch_id*self.batch_size:(batch_id + 1)*self.batch_size, :])\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_BC + 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, PBC: %10.4e, PDE: %10.4e' % (loss_IC.item(), loss_BC.item(), loss_PDE.item()))\n",
    "                start_time = end_time\n",
    "                # change the learning rate\n",
    "                self.scheduler.step()\n",
    "        # Using the second-order L-BFGS optimizer\n",
    "#         print('Starting with L-BFGS')\n",
    "        self.start_time = time.time()\n",
    "        self.optimizer_LBFGS.step(self.loss_func)                \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": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting with Adam\n",
      "Iter  1000, Total: 3.5598e-02, Time: 264.67 secs\n",
      "IC: 1.7181e-02, PBC: 9.6240e-05, PDE: 1.8321e-02\n",
      "Iter  2000, Total: 1.4366e-02, Time: 281.42 secs\n",
      "IC: 5.0651e-03, PBC: 5.1091e-05, PDE: 9.2498e-03\n",
      "Iter  3000, Total: 2.8111e-03, Time: 275.13 secs\n",
      "IC: 1.0386e-03, PBC: 1.4255e-05, PDE: 1.7582e-03\n",
      "Iter  4000, Total: 6.3112e-02, Time: 274.74 secs\n",
      "IC: 2.9820e-02, PBC: 8.8919e-05, PDE: 3.3203e-02\n",
      "Iter  5000, Total: 1.9902e-02, Time: 275.18 secs\n",
      "IC: 7.0726e-03, PBC: 4.3678e-05, PDE: 1.2785e-02\n",
      "Iter  6000, Total: 4.7788e-03, Time: 277.45 secs\n",
      "IC: 1.3388e-03, PBC: 2.4256e-05, PDE: 3.4158e-03\n",
      "Iter  7000, Total: 7.2314e-02, Time: 279.85 secs\n",
      "IC: 3.5892e-02, PBC: 3.8060e-04, PDE: 3.6041e-02\n",
      "Iter  8000, Total: 1.5937e-02, Time: 284.91 secs\n",
      "IC: 5.8825e-03, PBC: 2.1803e-05, PDE: 1.0033e-02\n",
      "Iter  9000, Total: 4.2635e-03, Time: 282.43 secs\n",
      "IC: 1.3494e-03, PBC: 2.5041e-05, PDE: 2.8890e-03\n",
      "Iter 10000, Total: 1.0830e-01, Time: 276.21 secs\n",
      "IC: 4.8878e-02, PBC: 9.4497e-04, PDE: 5.8473e-02\n",
      "Iter 11000, Total: 1.8552e-02, Time: 270.61 secs\n",
      "IC: 7.1221e-03, PBC: 9.5714e-05, PDE: 1.1334e-02\n",
      "Iter 12000, Total: 2.5608e-03, Time: 255.25 secs\n",
      "IC: 8.3132e-04, PBC: 1.6257e-05, PDE: 1.7132e-03\n",
      "Iter 13000, Total: 7.7537e-02, Time: 253.71 secs\n",
      "IC: 3.3615e-02, PBC: 1.8346e-04, PDE: 4.3738e-02\n",
      "Iter 14000, Total: 2.0531e-02, Time: 254.91 secs\n",
      "IC: 8.4553e-03, PBC: 1.4067e-04, PDE: 1.1935e-02\n",
      "Iter 15000, Total: 4.0686e-03, Time: 294.31 secs\n",
      "IC: 1.4578e-03, PBC: 3.9676e-05, PDE: 2.5711e-03\n",
      "Iter 16000, Total: 9.4922e-02, Time: 296.03 secs\n",
      "IC: 6.2893e-02, PBC: 7.6809e-05, PDE: 3.1953e-02\n",
      "Iter 17000, Total: 1.8989e-02, Time: 296.86 secs\n",
      "IC: 9.1964e-03, PBC: 6.8900e-05, PDE: 9.7238e-03\n",
      "Iter 18000, Total: 4.5286e-03, Time: 256.88 secs\n",
      "IC: 1.7501e-03, PBC: 7.3138e-05, PDE: 2.7054e-03\n",
      "Iter 19000, Total: 1.2444e-03, Time: 223.36 secs\n",
      "IC: 3.6777e-04, PBC: 5.7268e-06, PDE: 8.7086e-04\n",
      "Iter 20000, Total: 3.1516e-02, Time: 210.71 secs\n",
      "IC: 1.5665e-02, PBC: 3.6313e-05, PDE: 1.5814e-02\n",
      "Iter 21000, Total: 6.9697e-03, Time: 210.71 secs\n",
      "IC: 3.2389e-03, PBC: 7.0565e-05, PDE: 3.6602e-03\n",
      "Iter 22000, Total: 2.0249e-03, Time: 210.15 secs\n",
      "IC: 7.0689e-04, PBC: 1.3439e-05, PDE: 1.3046e-03\n",
      "Iter 23000, Total: 4.3990e-04, Time: 208.32 secs\n",
      "IC: 1.0606e-04, PBC: 2.0743e-06, PDE: 3.3176e-04\n",
      "Iter 24000, Total: 3.5292e-02, Time: 207.20 secs\n",
      "IC: 1.6749e-02, PBC: 4.2776e-04, PDE: 1.8114e-02\n",
      "Iter 25000, Total: 7.7015e-03, Time: 210.77 secs\n",
      "IC: 3.3909e-03, PBC: 1.4566e-04, PDE: 4.1650e-03\n",
      "Iter 26000, Total: 1.9704e-03, Time: 209.67 secs\n",
      "IC: 7.1669e-04, PBC: 1.6417e-05, PDE: 1.2373e-03\n",
      "Iter 27000, Total: 4.9596e-04, Time: 208.81 secs\n",
      "IC: 1.4869e-04, PBC: 3.0688e-06, PDE: 3.4421e-04\n",
      "Iter 28000, Total: 2.3499e-02, Time: 210.40 secs\n",
      "IC: 1.0646e-02, PBC: 7.4798e-05, PDE: 1.2778e-02\n",
      "Iter 29000, Total: 5.5573e-03, Time: 211.65 secs\n",
      "IC: 2.3137e-03, PBC: 5.1306e-05, PDE: 3.1923e-03\n",
      "Iter 30000, Total: 1.2712e-03, Time: 240.15 secs\n",
      "IC: 4.2629e-04, PBC: 1.1917e-05, PDE: 8.3304e-04\n",
      "Iter 31000, Total: 3.2233e-04, Time: 465.63 secs\n",
      "IC: 7.5923e-05, PBC: 6.6767e-06, PDE: 2.3973e-04\n",
      "Iter 32000, Total: 1.6300e-02, Time: 462.00 secs\n",
      "IC: 7.5334e-03, PBC: 5.4099e-05, PDE: 8.7129e-03\n",
      "Iter 33000, Total: 4.4673e-03, Time: 459.32 secs\n",
      "IC: 1.7873e-03, PBC: 2.7060e-05, PDE: 2.6529e-03\n",
      "Iter 34000, Total: 1.3564e-02, Time: 448.43 secs\n",
      "IC: 6.5900e-03, PBC: 8.4748e-05, PDE: 6.8889e-03\n",
      "Iter 35000, Total: 3.1164e-03, Time: 449.23 secs\n",
      "IC: 1.3810e-03, PBC: 1.5009e-05, PDE: 1.7204e-03\n",
      "Iter 36000, Total: 1.0111e-03, Time: 450.90 secs\n",
      "IC: 3.6678e-04, PBC: 3.0214e-06, PDE: 6.4126e-04\n",
      "Iter 37000, Total: 2.7752e-04, Time: 446.23 secs\n",
      "IC: 7.3539e-05, PBC: 1.6833e-06, PDE: 2.0229e-04\n",
      "Iter 38000, Total: 6.3205e-03, Time: 447.97 secs\n",
      "IC: 2.9227e-03, PBC: 3.9986e-05, PDE: 3.3578e-03\n",
      "Iter 39000, Total: 1.3982e-03, Time: 507.14 secs\n",
      "IC: 6.2415e-04, PBC: 8.1554e-06, PDE: 7.6591e-04\n",
      "Iter 40000, Total: 4.0927e-04, Time: 542.50 secs\n",
      "IC: 1.4654e-04, PBC: 2.6725e-06, PDE: 2.6006e-04\n",
      "Iter 41000, Total: 1.5971e-04, Time: 539.19 secs\n",
      "IC: 4.5099e-05, PBC: 9.5401e-07, PDE: 1.1366e-04\n",
      "Iter 42000, Total: 9.2265e-03, Time: 541.42 secs\n",
      "IC: 4.5408e-03, PBC: 1.7478e-04, PDE: 4.5109e-03\n",
      "Iter 43000, Total: 1.8625e-03, Time: 543.68 secs\n",
      "IC: 8.9474e-04, PBC: 1.0180e-05, PDE: 9.5761e-04\n",
      "Iter 44000, Total: 5.2655e-04, Time: 538.56 secs\n",
      "IC: 2.1366e-04, PBC: 1.0842e-06, PDE: 3.1180e-04\n",
      "Iter 45000, Total: 1.6131e-04, Time: 540.18 secs\n",
      "IC: 4.9985e-05, PBC: 7.9934e-07, PDE: 1.1052e-04\n",
      "Iter 46000, Total: 7.2270e-05, Time: 222.69 secs\n",
      "IC: 1.9299e-05, PBC: 4.6012e-07, PDE: 5.2510e-05\n",
      "Iter 47000, Total: 4.7634e-03, Time: 202.07 secs\n",
      "IC: 2.2171e-03, PBC: 2.5223e-05, PDE: 2.5212e-03\n",
      "Iter 48000, Total: 1.3327e-03, Time: 201.46 secs\n",
      "IC: 5.6237e-04, PBC: 6.0681e-06, PDE: 7.6422e-04\n",
      "Iter 49000, Total: 3.1168e-04, Time: 202.05 secs\n",
      "IC: 1.0561e-04, PBC: 1.1798e-06, PDE: 2.0488e-04\n",
      "Iter 50000, Total: 1.0252e-04, Time: 203.43 secs\n",
      "IC: 2.2805e-05, PBC: 9.9087e-07, PDE: 7.8726e-05\n"
     ]
    }
   ],
   "source": [
    "layers = [2, 32, 32, 32, 32, 32, 32, 32, 2]\n",
    "N_IC = 128\n",
    "ptsIC=np.load('ptsIC.npy')\n",
    "ptsBC=np.load('ptsBC.npy')\n",
    "ptsPDE=np.load('random_data.npy')\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",
    "n_batches = 5\n",
    "batch_size = 3280\n",
    "uBC = np.load('uBC.npy')\n",
    "vBC = np.load('vBC.npy')\n",
    "model = PhysicsInformedNN(period, ptsIC, uIC, vIC, ptsBC, uBC, vBC, ptsPDE, layers,n_batches, batch_size)\n",
    "model.train(50000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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": 9,
   "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": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 128256 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()\n",
    "fig.savefig('NLSu_1D_BL.png',format='png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 138276 characters]",
      "text/plain": [
       "<Figure size 1800x600 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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()\n",
    "fig.savefig('NLSv_1D_BL.png',format='png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "l2: 0.06604281628685388\n",
      "l1: 0.05646342899668441\n",
      "linf: 0.0729687755219959\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": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "l2: 0.10538184768270127\n",
      "l1: 0.11081683524797876\n",
      "linf: 0.11354052674612172\n"
     ]
    }
   ],
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
