Copy of Final project_code_PINN_continous time model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 6013,
     "status": "ok",
     "timestamp": 1730632317551,
     "user": {
      "displayName": "Shucheng Li",
      "userId": "03239258419160329131"
     },
     "user_tz": -480
    },
    "id": "bZeFdPhvuMM3",
    "outputId": "6fc96892-6c7d-4b91-83b3-784af3a19db0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: pyDOE in /opt/anaconda3/envs/RAG/lib/python3.11/site-packages (0.3.8)\n",
      "Requirement already satisfied: numpy in /opt/anaconda3/envs/RAG/lib/python3.11/site-packages (from pyDOE) (1.26.4)\n",
      "Requirement already satisfied: scipy in /opt/anaconda3/envs/RAG/lib/python3.11/site-packages (from pyDOE) (1.14.1)\n"
     ]
    }
   ],
   "source": [
    "!pip install pyDOE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "RkMTccdsrWq7"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import scipy.io\n",
    "from scipy.interpolate import griddata\n",
    "from pyDOE import lhs\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "class PhysicsInformedNN:\n",
    "    def __init__(self, x0, u0, v0, tb, X_f, layers, lb, ub):\n",
    "        self.lb = lb\n",
    "        self.ub = ub\n",
    "\n",
    "        self.x0 = torch.tensor(x0, dtype=torch.float32, requires_grad=True)\n",
    "        self.t0 = torch.tensor(np.zeros_like(x0), dtype=torch.float32, requires_grad=True)\n",
    "        self.u0 = torch.tensor(u0, dtype=torch.float32)\n",
    "        self.v0 = torch.tensor(v0, dtype=torch.float32)\n",
    "\n",
    "        self.x_lb = torch.tensor(np.zeros_like(tb) + lb[0], dtype=torch.float32, requires_grad=True)\n",
    "        self.t_lb = torch.tensor(tb, dtype=torch.float32, requires_grad=True)\n",
    "\n",
    "        self.x_ub = torch.tensor(np.zeros_like(tb) + ub[0], dtype=torch.float32, requires_grad=True)\n",
    "        self.t_ub = torch.tensor(tb, dtype=torch.float32, requires_grad=True)\n",
    "\n",
    "        self.x_f = torch.tensor(X_f[:, 0:1], dtype=torch.float32, requires_grad=True)\n",
    "        self.t_f = torch.tensor(X_f[:, 1:2], dtype=torch.float32, requires_grad=True)\n",
    "\n",
    "        # Initialize the neural network\n",
    "        self.model = self.build_model(layers)\n",
    "        self.optimizer = torch.optim.Adam(self.model.parameters())\n",
    "\n",
    "    def build_model(self, layers):\n",
    "        model = []\n",
    "        num_layers = len(layers)\n",
    "        for i in range(num_layers - 2):\n",
    "            model.append(nn.Linear(layers[i], layers[i+1]))\n",
    "            model.append(nn.Tanh())\n",
    "        model.append(nn.Linear(layers[-2], layers[-1]))\n",
    "        return nn.Sequential(*model)\n",
    "\n",
    "    def forward_uv(self, x, t):\n",
    "        X = torch.cat([x, t], dim=1)\n",
    "        uv = self.model(X)\n",
    "        u = uv[:, 0:1]\n",
    "        v = uv[:, 1:2]\n",
    "        u_x = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True)[0]\n",
    "        v_x = torch.autograd.grad(v, x, grad_outputs=torch.ones_like(v), create_graph=True)[0]\n",
    "        return u, v, u_x, v_x\n",
    "\n",
    "    def net_f_uv(self, x, t):\n",
    "        u, v, u_x, v_x = self.forward_uv(x, t)\n",
    "        u_t = torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u), create_graph=True)[0]\n",
    "        v_t = torch.autograd.grad(v, t, grad_outputs=torch.ones_like(v), create_graph=True)[0]\n",
    "        u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), create_graph=True)[0]\n",
    "        v_xx = torch.autograd.grad(v_x, x, grad_outputs=torch.ones_like(v_x), create_graph=True)[0]\n",
    "        f_u = u_t + 0.5 * v_xx + (u**2 + v**2) * v\n",
    "        f_v = v_t - 0.5 * u_xx - (u**2 + v**2) * u\n",
    "        return f_u, f_v\n",
    "\n",
    "    def loss_func(self):\n",
    "        u0_pred, v0_pred, _, _ = self.forward_uv(self.x0, self.t0)\n",
    "        u_lb_pred, v_lb_pred, u_x_lb_pred, v_x_lb_pred = self.forward_uv(self.x_lb, self.t_lb)\n",
    "        u_ub_pred, v_ub_pred, u_x_ub_pred, v_x_ub_pred = self.forward_uv(self.x_ub, self.t_ub)\n",
    "        f_u_pred, f_v_pred = self.net_f_uv(self.x_f, self.t_f)\n",
    "\n",
    "        loss = torch.mean((self.u0 - u0_pred) ** 2) + \\\n",
    "               torch.mean((self.v0 - v0_pred) ** 2) + \\\n",
    "               torch.mean((u_lb_pred - u_ub_pred) ** 2) + \\\n",
    "               torch.mean((v_lb_pred - v_ub_pred) ** 2) + \\\n",
    "               torch.mean((u_x_lb_pred - u_x_ub_pred) ** 2) + \\\n",
    "               torch.mean((v_x_lb_pred - v_x_ub_pred) ** 2) + \\\n",
    "               torch.mean(f_u_pred ** 2) + \\\n",
    "               torch.mean(f_v_pred ** 2)\n",
    "        return loss\n",
    "\n",
    "    def train(self, n_iter):\n",
    "        for it in range(n_iter):\n",
    "            self.optimizer.zero_grad()\n",
    "            loss = self.loss_func()\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "            if it % 10 == 0:\n",
    "                print(f'Iter: {it}, Loss: {loss.item():.3e}')\n",
    "\n",
    "    def predict(self, X_star):\n",
    "        x_star = torch.tensor(X_star[:, 0:1], dtype=torch.float32)\n",
    "        t_star = torch.tensor(X_star[:, 1:2], dtype=torch.float32)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            u_pred, v_pred, _, _ = self.forward_uv(x_star, t_star)\n",
    "            f_u_pred, f_v_pred = self.net_f_uv(x_star, t_star)\n",
    "\n",
    "        return u_pred.numpy(), v_pred.numpy(), f_u_pred.numpy(), f_v_pred.numpy()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "WTJ8_Z85se3P"
   },
   "outputs": [],
   "source": [
    "# Load data from NLS.mat\n",
    "data = scipy.io.loadmat('NLS.mat')\n",
    "t = data['tt'].flatten()[:, None]       # time points\n",
    "x = data['x'].flatten()[:, None]        # spatial points\n",
    "Exact = data['uu']                      # complex solution\n",
    "\n",
    "# Separate real and imaginary parts of the solution\n",
    "Exact_u = np.real(Exact)\n",
    "Exact_v = np.imag(Exact)\n",
    "Exact_h = np.sqrt(Exact_u**2 + Exact_v**2)\n",
    "\n",
    "# Define bounds, layer configuration, and number of samples\n",
    "lb = np.array([-5.0, 0.0])\n",
    "ub = np.array([5.0, np.pi/2])\n",
    "layers = [2, 100, 100, 100, 100, 2]\n",
    "N0 = 50\n",
    "N_b = 50\n",
    "N_f = 20000\n",
    "\n",
    "# Sampling initial conditions\n",
    "idx_x = np.random.choice(x.shape[0], N0, replace=False)\n",
    "x0 = x[idx_x, :]\n",
    "u0 = Exact_u[idx_x, 0:1]\n",
    "v0 = Exact_v[idx_x, 0:1]\n",
    "\n",
    "# Sampling boundary conditions\n",
    "idx_t = np.random.choice(t.shape[0], N_b, replace=False)\n",
    "tb = t[idx_t, :]\n",
    "\n",
    "# Collocation points\n",
    "X_f = lb + (ub - lb) * lhs(2, N_f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "aEAL_eeWtiHW",
    "outputId": "34fcb982-dbed-44d8-9f37-3177fbbeb114"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iter: 0, Loss: 1.040e+00\n",
      "Iter: 10, Loss: 3.152e-01\n",
      "Iter: 20, Loss: 2.466e-01\n",
      "Iter: 30, Loss: 1.674e-01\n",
      "Iter: 40, Loss: 1.394e-01\n",
      "Iter: 50, Loss: 1.281e-01\n",
      "Iter: 60, Loss: 1.170e-01\n",
      "Iter: 70, Loss: 1.121e-01\n",
      "Iter: 80, Loss: 1.068e-01\n",
      "Iter: 90, Loss: 1.026e-01\n",
      "Iter: 100, Loss: 9.679e-02\n",
      "Iter: 110, Loss: 9.182e-02\n",
      "Iter: 120, Loss: 8.464e-02\n",
      "Iter: 130, Loss: 8.438e-02\n",
      "Iter: 140, Loss: 8.049e-02\n",
      "Iter: 150, Loss: 7.429e-02\n",
      "Iter: 160, Loss: 7.083e-02\n",
      "Iter: 170, Loss: 9.098e-02\n",
      "Iter: 180, Loss: 7.496e-02\n",
      "Iter: 190, Loss: 7.073e-02\n",
      "Iter: 200, Loss: 6.840e-02\n",
      "Iter: 210, Loss: 6.681e-02\n",
      "Iter: 220, Loss: 6.570e-02\n",
      "Iter: 230, Loss: 6.505e-02\n",
      "Iter: 240, Loss: 6.433e-02\n",
      "Iter: 250, Loss: 6.417e-02\n",
      "Iter: 260, Loss: 7.145e-02\n",
      "Iter: 270, Loss: 6.276e-02\n",
      "Iter: 280, Loss: 6.219e-02\n",
      "Iter: 290, Loss: 6.149e-02\n",
      "Iter: 300, Loss: 6.088e-02\n",
      "Iter: 310, Loss: 6.037e-02\n",
      "Iter: 320, Loss: 5.971e-02\n",
      "Iter: 330, Loss: 5.917e-02\n",
      "Iter: 340, Loss: 5.989e-02\n",
      "Iter: 350, Loss: 6.396e-02\n",
      "Iter: 360, Loss: 6.155e-02\n",
      "Iter: 370, Loss: 5.846e-02\n",
      "Iter: 380, Loss: 5.820e-02\n",
      "Iter: 390, Loss: 5.727e-02\n",
      "Iter: 400, Loss: 5.669e-02\n",
      "Iter: 410, Loss: 5.622e-02\n",
      "Iter: 420, Loss: 5.577e-02\n",
      "Iter: 430, Loss: 5.638e-02\n",
      "Iter: 440, Loss: 7.771e-02\n",
      "Iter: 450, Loss: 5.969e-02\n",
      "Iter: 460, Loss: 5.498e-02\n",
      "Iter: 470, Loss: 5.452e-02\n",
      "Iter: 480, Loss: 5.371e-02\n",
      "Iter: 490, Loss: 5.319e-02\n",
      "Iter: 500, Loss: 5.265e-02\n",
      "Iter: 510, Loss: 5.218e-02\n",
      "Iter: 520, Loss: 5.171e-02\n",
      "Iter: 530, Loss: 5.125e-02\n",
      "Iter: 540, Loss: 5.079e-02\n",
      "Iter: 550, Loss: 5.039e-02\n",
      "Iter: 560, Loss: 1.028e-01\n",
      "Iter: 570, Loss: 5.380e-02\n",
      "Iter: 580, Loss: 5.041e-02\n",
      "Iter: 590, Loss: 4.985e-02\n",
      "Iter: 600, Loss: 4.945e-02\n",
      "Iter: 610, Loss: 4.911e-02\n",
      "Iter: 620, Loss: 4.864e-02\n",
      "Iter: 630, Loss: 4.817e-02\n",
      "Iter: 640, Loss: 4.775e-02\n",
      "Iter: 650, Loss: 4.735e-02\n",
      "Iter: 660, Loss: 4.695e-02\n",
      "Iter: 670, Loss: 4.657e-02\n",
      "Iter: 680, Loss: 5.610e-02\n",
      "Iter: 690, Loss: 4.985e-02\n",
      "Iter: 700, Loss: 5.226e-02\n",
      "Iter: 710, Loss: 4.592e-02\n",
      "Iter: 720, Loss: 4.597e-02\n",
      "Iter: 730, Loss: 4.498e-02\n",
      "Iter: 740, Loss: 4.470e-02\n",
      "Iter: 750, Loss: 4.427e-02\n",
      "Iter: 760, Loss: 4.785e-02\n",
      "Iter: 770, Loss: 7.732e-02\n",
      "Iter: 780, Loss: 4.637e-02\n",
      "Iter: 790, Loss: 4.654e-02\n",
      "Iter: 800, Loss: 4.568e-02\n",
      "Iter: 810, Loss: 4.395e-02\n",
      "Iter: 820, Loss: 4.342e-02\n",
      "Iter: 830, Loss: 4.306e-02\n",
      "Iter: 840, Loss: 4.276e-02\n",
      "Iter: 850, Loss: 4.242e-02\n",
      "Iter: 860, Loss: 4.211e-02\n",
      "Iter: 870, Loss: 4.181e-02\n",
      "Iter: 880, Loss: 4.150e-02\n",
      "Iter: 890, Loss: 4.120e-02\n",
      "Iter: 900, Loss: 4.091e-02\n",
      "Iter: 910, Loss: 4.061e-02\n",
      "Iter: 920, Loss: 4.032e-02\n",
      "Iter: 930, Loss: 4.036e-02\n",
      "Iter: 940, Loss: 5.643e-02\n",
      "Iter: 950, Loss: 4.128e-02\n",
      "Iter: 960, Loss: 4.240e-02\n",
      "Iter: 970, Loss: 4.156e-02\n",
      "Iter: 980, Loss: 4.058e-02\n",
      "Iter: 990, Loss: 3.974e-02\n",
      "Iter: 1000, Loss: 3.949e-02\n",
      "Iter: 1010, Loss: 3.912e-02\n",
      "Iter: 1020, Loss: 3.884e-02\n",
      "Iter: 1030, Loss: 3.856e-02\n",
      "Iter: 1040, Loss: 3.829e-02\n",
      "Iter: 1050, Loss: 3.803e-02\n",
      "Iter: 1060, Loss: 3.778e-02\n",
      "Iter: 1070, Loss: 8.358e-02\n",
      "Iter: 1080, Loss: 4.214e-02\n",
      "Iter: 1090, Loss: 3.891e-02\n",
      "Iter: 1100, Loss: 3.752e-02\n",
      "Iter: 1110, Loss: 3.735e-02\n",
      "Iter: 1120, Loss: 3.724e-02\n",
      "Iter: 1130, Loss: 3.680e-02\n",
      "Iter: 1140, Loss: 3.656e-02\n",
      "Iter: 1150, Loss: 3.633e-02\n",
      "Iter: 1160, Loss: 3.671e-02\n",
      "Iter: 1170, Loss: 4.777e-02\n",
      "Iter: 1180, Loss: 4.444e-02\n",
      "Iter: 1190, Loss: 3.792e-02\n",
      "Iter: 1200, Loss: 3.860e-02\n",
      "Iter: 1210, Loss: 3.643e-02\n",
      "Iter: 1220, Loss: 3.638e-02\n",
      "Iter: 1230, Loss: 3.594e-02\n",
      "Iter: 1240, Loss: 3.562e-02\n",
      "Iter: 1250, Loss: 3.537e-02\n",
      "Iter: 1260, Loss: 3.514e-02\n",
      "Iter: 1270, Loss: 3.491e-02\n",
      "Iter: 1280, Loss: 3.469e-02\n",
      "Iter: 1290, Loss: 3.447e-02\n",
      "Iter: 1300, Loss: 1.082e-01\n",
      "Iter: 1310, Loss: 5.711e-02\n",
      "Iter: 1320, Loss: 5.102e-02\n",
      "Iter: 1330, Loss: 3.684e-02\n",
      "Iter: 1340, Loss: 3.633e-02\n",
      "Iter: 1350, Loss: 3.612e-02\n",
      "Iter: 1360, Loss: 3.576e-02\n",
      "Iter: 1370, Loss: 3.504e-02\n",
      "Iter: 1380, Loss: 3.476e-02\n",
      "Iter: 1390, Loss: 3.444e-02\n",
      "Iter: 1400, Loss: 3.420e-02\n",
      "Iter: 1410, Loss: 3.398e-02\n",
      "Iter: 1420, Loss: 3.377e-02\n",
      "Iter: 1430, Loss: 3.356e-02\n",
      "Iter: 1440, Loss: 3.337e-02\n",
      "Iter: 1450, Loss: 3.618e-02\n",
      "Iter: 1460, Loss: 5.335e-02\n",
      "Iter: 1470, Loss: 3.585e-02\n",
      "Iter: 1480, Loss: 3.476e-02\n",
      "Iter: 1490, Loss: 3.353e-02\n",
      "Iter: 1500, Loss: 3.288e-02\n",
      "Iter: 1510, Loss: 3.249e-02\n",
      "Iter: 1520, Loss: 3.222e-02\n",
      "Iter: 1530, Loss: 3.200e-02\n",
      "Iter: 1540, Loss: 3.181e-02\n",
      "Iter: 1550, Loss: 3.162e-02\n",
      "Iter: 1560, Loss: 3.143e-02\n",
      "Iter: 1570, Loss: 3.129e-02\n",
      "Iter: 1580, Loss: 7.208e-02\n",
      "Iter: 1590, Loss: 4.022e-02\n",
      "Iter: 1600, Loss: 3.352e-02\n",
      "Iter: 1610, Loss: 3.501e-02\n",
      "Iter: 1620, Loss: 3.149e-02\n",
      "Iter: 1630, Loss: 3.152e-02\n",
      "Iter: 1640, Loss: 3.110e-02\n",
      "Iter: 1650, Loss: 3.078e-02\n",
      "Iter: 1660, Loss: 3.063e-02\n",
      "Iter: 1670, Loss: 3.040e-02\n",
      "Iter: 1680, Loss: 3.020e-02\n",
      "Iter: 1690, Loss: 3.001e-02\n",
      "Iter: 1700, Loss: 2.982e-02\n",
      "Iter: 1710, Loss: 2.963e-02\n",
      "Iter: 1720, Loss: 2.947e-02\n",
      "Iter: 1730, Loss: 7.487e-02\n",
      "Iter: 1740, Loss: 1.109e-01\n",
      "Iter: 1750, Loss: 6.318e-02\n",
      "Iter: 1760, Loss: 3.579e-02\n",
      "Iter: 1770, Loss: 4.052e-02\n",
      "Iter: 1780, Loss: 3.326e-02\n",
      "Iter: 1790, Loss: 3.303e-02\n",
      "Iter: 1800, Loss: 3.201e-02\n",
      "Iter: 1810, Loss: 3.135e-02\n",
      "Iter: 1820, Loss: 3.097e-02\n",
      "Iter: 1830, Loss: 3.065e-02\n",
      "Iter: 1840, Loss: 3.037e-02\n",
      "Iter: 1850, Loss: 3.012e-02\n",
      "Iter: 1860, Loss: 2.989e-02\n",
      "Iter: 1870, Loss: 2.968e-02\n",
      "Iter: 1880, Loss: 2.947e-02\n",
      "Iter: 1890, Loss: 2.927e-02\n",
      "Iter: 1900, Loss: 2.908e-02\n",
      "Iter: 1910, Loss: 2.890e-02\n",
      "Iter: 1920, Loss: 2.872e-02\n",
      "Iter: 1930, Loss: 2.854e-02\n",
      "Iter: 1940, Loss: 2.837e-02\n",
      "Iter: 1950, Loss: 2.820e-02\n",
      "Iter: 1960, Loss: 2.803e-02\n",
      "Iter: 1970, Loss: 2.786e-02\n",
      "Iter: 1980, Loss: 2.769e-02\n",
      "Iter: 1990, Loss: 2.753e-02\n",
      "Iter: 2000, Loss: 2.736e-02\n",
      "Iter: 2010, Loss: 2.720e-02\n",
      "Iter: 2020, Loss: 2.703e-02\n",
      "Iter: 2030, Loss: 2.687e-02\n",
      "Iter: 2040, Loss: 2.676e-02\n",
      "Iter: 2050, Loss: 2.283e-01\n",
      "Iter: 2060, Loss: 3.945e-02\n",
      "Iter: 2070, Loss: 3.233e-02\n",
      "Iter: 2080, Loss: 2.991e-02\n",
      "Iter: 2090, Loss: 2.903e-02\n",
      "Iter: 2100, Loss: 2.756e-02\n",
      "Iter: 2110, Loss: 2.739e-02\n",
      "Iter: 2120, Loss: 2.702e-02\n",
      "Iter: 2130, Loss: 2.677e-02\n",
      "Iter: 2140, Loss: 2.657e-02\n",
      "Iter: 2150, Loss: 2.638e-02\n",
      "Iter: 2160, Loss: 2.620e-02\n",
      "Iter: 2170, Loss: 2.603e-02\n",
      "Iter: 2180, Loss: 2.587e-02\n",
      "Iter: 2190, Loss: 2.571e-02\n",
      "Iter: 2200, Loss: 2.555e-02\n",
      "Iter: 2210, Loss: 2.540e-02\n",
      "Iter: 2220, Loss: 2.525e-02\n",
      "Iter: 2230, Loss: 2.566e-02\n",
      "Iter: 2240, Loss: 3.731e-02\n",
      "Iter: 2250, Loss: 3.268e-02\n",
      "Iter: 2260, Loss: 2.878e-02\n",
      "Iter: 2270, Loss: 2.702e-02\n",
      "Iter: 2280, Loss: 2.539e-02\n",
      "Iter: 2290, Loss: 2.508e-02\n",
      "Iter: 2300, Loss: 2.488e-02\n",
      "Iter: 2310, Loss: 2.465e-02\n",
      "Iter: 2320, Loss: 2.448e-02\n",
      "Iter: 2330, Loss: 2.433e-02\n",
      "Iter: 2340, Loss: 2.418e-02\n",
      "Iter: 2350, Loss: 2.405e-02\n",
      "Iter: 2360, Loss: 2.653e-02\n",
      "Iter: 2370, Loss: 3.501e-02\n",
      "Iter: 2380, Loss: 3.440e-02\n",
      "Iter: 2390, Loss: 2.512e-02\n",
      "Iter: 2400, Loss: 2.529e-02\n",
      "Iter: 2410, Loss: 2.395e-02\n",
      "Iter: 2420, Loss: 2.386e-02\n",
      "Iter: 2430, Loss: 2.360e-02\n",
      "Iter: 2440, Loss: 2.340e-02\n",
      "Iter: 2450, Loss: 2.327e-02\n",
      "Iter: 2460, Loss: 2.324e-02\n",
      "Iter: 2470, Loss: 4.384e-02\n",
      "Iter: 2480, Loss: 2.542e-02\n",
      "Iter: 2490, Loss: 3.157e-02\n",
      "Iter: 2500, Loss: 2.375e-02\n",
      "Iter: 2510, Loss: 2.407e-02\n",
      "Iter: 2520, Loss: 2.320e-02\n",
      "Iter: 2530, Loss: 3.564e-02\n",
      "Iter: 2540, Loss: 3.246e-02\n",
      "Iter: 2550, Loss: 2.266e-02\n",
      "Iter: 2560, Loss: 2.332e-02\n",
      "Iter: 2570, Loss: 2.231e-02\n",
      "Iter: 2580, Loss: 2.220e-02\n",
      "Iter: 2590, Loss: 2.217e-02\n",
      "Iter: 2600, Loss: 2.288e-02\n",
      "Iter: 2610, Loss: 5.461e-02\n",
      "Iter: 2620, Loss: 3.052e-02\n",
      "Iter: 2630, Loss: 2.488e-02\n",
      "Iter: 2640, Loss: 2.288e-02\n",
      "Iter: 2650, Loss: 2.187e-02\n",
      "Iter: 2660, Loss: 2.165e-02\n",
      "Iter: 2670, Loss: 2.143e-02\n",
      "Iter: 2680, Loss: 2.133e-02\n",
      "Iter: 2690, Loss: 2.185e-02\n",
      "Iter: 2700, Loss: 3.561e-02\n",
      "Iter: 2710, Loss: 3.439e-02\n",
      "Iter: 2720, Loss: 2.416e-02\n",
      "Iter: 2730, Loss: 2.168e-02\n",
      "Iter: 2740, Loss: 2.141e-02\n",
      "Iter: 2750, Loss: 2.122e-02\n",
      "Iter: 2760, Loss: 2.091e-02\n",
      "Iter: 2770, Loss: 4.831e-02\n",
      "Iter: 2780, Loss: 2.196e-02\n",
      "Iter: 2790, Loss: 2.761e-02\n",
      "Iter: 2800, Loss: 2.182e-02\n",
      "Iter: 2810, Loss: 2.153e-02\n",
      "Iter: 2820, Loss: 2.078e-02\n",
      "Iter: 2830, Loss: 2.070e-02\n",
      "Iter: 2840, Loss: 2.051e-02\n",
      "Iter: 2850, Loss: 2.036e-02\n",
      "Iter: 2860, Loss: 2.025e-02\n",
      "Iter: 2870, Loss: 2.013e-02\n",
      "Iter: 2880, Loss: 2.002e-02\n",
      "Iter: 2890, Loss: 1.993e-02\n",
      "Iter: 2900, Loss: 3.132e-02\n",
      "Iter: 2910, Loss: 3.533e-02\n",
      "Iter: 2920, Loss: 2.241e-02\n",
      "Iter: 2930, Loss: 2.165e-02\n",
      "Iter: 2940, Loss: 2.075e-02\n",
      "Iter: 2950, Loss: 1.999e-02\n",
      "Iter: 2960, Loss: 2.011e-02\n",
      "Iter: 2970, Loss: 3.279e-02\n",
      "Iter: 2980, Loss: 2.980e-02\n",
      "Iter: 2990, Loss: 1.982e-02\n",
      "Iter: 3000, Loss: 2.083e-02\n",
      "Iter: 3010, Loss: 1.945e-02\n",
      "Iter: 3020, Loss: 1.927e-02\n",
      "Iter: 3030, Loss: 1.929e-02\n",
      "Iter: 3040, Loss: 1.913e-02\n",
      "Iter: 3050, Loss: 2.610e-02\n",
      "Iter: 3060, Loss: 3.022e-02\n",
      "Iter: 3070, Loss: 2.360e-02\n",
      "Iter: 3080, Loss: 2.128e-02\n",
      "Iter: 3090, Loss: 1.975e-02\n",
      "Iter: 3100, Loss: 1.915e-02\n",
      "Iter: 3110, Loss: 1.879e-02\n",
      "Iter: 3120, Loss: 1.865e-02\n",
      "Iter: 3130, Loss: 1.896e-02\n",
      "Iter: 3140, Loss: 4.374e-02\n",
      "Iter: 3150, Loss: 2.336e-02\n",
      "Iter: 3160, Loss: 1.893e-02\n",
      "Iter: 3170, Loss: 1.948e-02\n",
      "Iter: 3180, Loss: 1.854e-02\n",
      "Iter: 3190, Loss: 1.821e-02\n",
      "Iter: 3200, Loss: 1.817e-02\n",
      "Iter: 3210, Loss: 3.401e-02\n",
      "Iter: 3220, Loss: 3.752e-02\n",
      "Iter: 3230, Loss: 2.240e-02\n",
      "Iter: 3240, Loss: 1.872e-02\n",
      "Iter: 3250, Loss: 1.831e-02\n",
      "Iter: 3260, Loss: 1.807e-02\n",
      "Iter: 3270, Loss: 1.782e-02\n",
      "Iter: 3280, Loss: 1.768e-02\n",
      "Iter: 3290, Loss: 1.757e-02\n",
      "Iter: 3300, Loss: 1.747e-02\n",
      "Iter: 3310, Loss: 1.739e-02\n",
      "Iter: 3320, Loss: 2.827e-02\n",
      "Iter: 3330, Loss: 3.581e-02\n",
      "Iter: 3340, Loss: 2.336e-02\n",
      "Iter: 3350, Loss: 1.769e-02\n",
      "Iter: 3360, Loss: 1.825e-02\n",
      "Iter: 3370, Loss: 2.006e-02\n",
      "Iter: 3380, Loss: 2.230e-02\n",
      "Iter: 3390, Loss: 2.102e-02\n",
      "Iter: 3400, Loss: 1.860e-02\n",
      "Iter: 3410, Loss: 1.704e-02\n",
      "Iter: 3420, Loss: 1.705e-02\n",
      "Iter: 3430, Loss: 1.704e-02\n",
      "Iter: 3440, Loss: 1.762e-02\n",
      "Iter: 3450, Loss: 4.880e-02\n",
      "Iter: 3460, Loss: 1.965e-02\n",
      "Iter: 3470, Loss: 1.941e-02\n",
      "Iter: 3480, Loss: 1.676e-02\n",
      "Iter: 3490, Loss: 1.685e-02\n",
      "Iter: 3500, Loss: 1.641e-02\n",
      "Iter: 3510, Loss: 1.670e-02\n",
      "Iter: 3520, Loss: 4.457e-02\n",
      "Iter: 3530, Loss: 1.686e-02\n",
      "Iter: 3540, Loss: 1.644e-02\n",
      "Iter: 3550, Loss: 1.619e-02\n",
      "Iter: 3560, Loss: 1.626e-02\n",
      "Iter: 3570, Loss: 1.782e-02\n",
      "Iter: 3580, Loss: 4.117e-02\n",
      "Iter: 3590, Loss: 2.075e-02\n",
      "Iter: 3600, Loss: 1.629e-02\n",
      "Iter: 3610, Loss: 1.635e-02\n",
      "Iter: 3620, Loss: 1.596e-02\n",
      "Iter: 3630, Loss: 1.579e-02\n",
      "Iter: 3640, Loss: 1.561e-02\n",
      "Iter: 3650, Loss: 1.549e-02\n",
      "Iter: 3660, Loss: 1.537e-02\n",
      "Iter: 3670, Loss: 1.626e-02\n",
      "Iter: 3680, Loss: 3.228e-02\n",
      "Iter: 3690, Loss: 2.180e-02\n",
      "Iter: 3700, Loss: 1.786e-02\n",
      "Iter: 3710, Loss: 1.637e-02\n",
      "Iter: 3720, Loss: 1.606e-02\n",
      "Iter: 3730, Loss: 1.520e-02\n",
      "Iter: 3740, Loss: 1.508e-02\n",
      "Iter: 3750, Loss: 1.500e-02\n",
      "Iter: 3760, Loss: 1.534e-02\n",
      "Iter: 3770, Loss: 5.168e-02\n",
      "Iter: 3780, Loss: 1.604e-02\n",
      "Iter: 3790, Loss: 1.536e-02\n",
      "Iter: 3800, Loss: 1.491e-02\n",
      "Iter: 3810, Loss: 1.474e-02\n",
      "Iter: 3820, Loss: 1.482e-02\n",
      "Iter: 3830, Loss: 1.694e-02\n",
      "Iter: 3840, Loss: 2.788e-02\n",
      "Iter: 3850, Loss: 1.721e-02\n",
      "Iter: 3860, Loss: 1.501e-02\n",
      "Iter: 3870, Loss: 1.511e-02\n",
      "Iter: 3880, Loss: 1.435e-02\n",
      "Iter: 3890, Loss: 1.429e-02\n",
      "Iter: 3900, Loss: 2.680e-02\n",
      "Iter: 3910, Loss: 2.935e-02\n",
      "Iter: 3920, Loss: 1.449e-02\n",
      "Iter: 3930, Loss: 1.573e-02\n",
      "Iter: 3940, Loss: 1.402e-02\n",
      "Iter: 3950, Loss: 1.422e-02\n",
      "Iter: 3960, Loss: 2.547e-02\n",
      "Iter: 3970, Loss: 3.024e-02\n",
      "Iter: 3980, Loss: 1.752e-02\n",
      "Iter: 3990, Loss: 1.539e-02\n",
      "Iter: 4000, Loss: 1.427e-02\n",
      "Iter: 4010, Loss: 1.405e-02\n",
      "Iter: 4020, Loss: 3.034e-02\n",
      "Iter: 4030, Loss: 2.791e-02\n",
      "Iter: 4040, Loss: 1.549e-02\n",
      "Iter: 4050, Loss: 1.374e-02\n",
      "Iter: 4060, Loss: 1.391e-02\n",
      "Iter: 4070, Loss: 1.349e-02\n",
      "Iter: 4080, Loss: 1.341e-02\n",
      "Iter: 4090, Loss: 1.329e-02\n",
      "Iter: 4100, Loss: 1.318e-02\n",
      "Iter: 4110, Loss: 1.312e-02\n",
      "Iter: 4120, Loss: 1.382e-02\n",
      "Iter: 4130, Loss: 6.648e-02\n",
      "Iter: 4140, Loss: 1.738e-02\n",
      "Iter: 4150, Loss: 1.472e-02\n",
      "Iter: 4160, Loss: 1.381e-02\n",
      "Iter: 4170, Loss: 1.316e-02\n",
      "Iter: 4180, Loss: 1.298e-02\n",
      "Iter: 4190, Loss: 2.413e-02\n",
      "Iter: 4200, Loss: 3.076e-02\n",
      "Iter: 4210, Loss: 1.770e-02\n",
      "Iter: 4220, Loss: 1.460e-02\n",
      "Iter: 4230, Loss: 1.347e-02\n",
      "Iter: 4240, Loss: 1.293e-02\n",
      "Iter: 4250, Loss: 1.276e-02\n",
      "Iter: 4260, Loss: 1.258e-02\n",
      "Iter: 4270, Loss: 1.265e-02\n",
      "Iter: 4280, Loss: 4.153e-02\n",
      "Iter: 4290, Loss: 1.374e-02\n",
      "Iter: 4300, Loss: 1.637e-02\n",
      "Iter: 4310, Loss: 1.348e-02\n",
      "Iter: 4320, Loss: 1.242e-02\n",
      "Iter: 4330, Loss: 1.251e-02\n",
      "Iter: 4340, Loss: 1.439e-02\n",
      "Iter: 4350, Loss: 2.919e-02\n",
      "Iter: 4360, Loss: 2.306e-02\n",
      "Iter: 4370, Loss: 1.635e-02\n",
      "Iter: 4380, Loss: 1.383e-02\n",
      "Iter: 4390, Loss: 1.274e-02\n",
      "Iter: 4400, Loss: 1.213e-02\n",
      "Iter: 4410, Loss: 1.210e-02\n",
      "Iter: 4420, Loss: 1.196e-02\n",
      "Iter: 4430, Loss: 1.201e-02\n",
      "Iter: 4440, Loss: 5.857e-02\n",
      "Iter: 4450, Loss: 2.069e-02\n",
      "Iter: 4460, Loss: 1.415e-02\n",
      "Iter: 4470, Loss: 1.429e-02\n",
      "Iter: 4480, Loss: 1.244e-02\n",
      "Iter: 4490, Loss: 1.199e-02\n",
      "Iter: 4500, Loss: 1.195e-02\n",
      "Iter: 4510, Loss: 1.181e-02\n",
      "Iter: 4520, Loss: 1.172e-02\n",
      "Iter: 4530, Loss: 1.166e-02\n",
      "Iter: 4540, Loss: 1.162e-02\n",
      "Iter: 4550, Loss: 1.632e-02\n",
      "Iter: 4560, Loss: 1.511e-02\n",
      "Iter: 4570, Loss: 1.399e-02\n",
      "Iter: 4580, Loss: 1.477e-02\n",
      "Iter: 4590, Loss: 1.300e-02\n",
      "Iter: 4600, Loss: 1.200e-02\n",
      "Iter: 4610, Loss: 1.167e-02\n",
      "Iter: 4620, Loss: 1.150e-02\n",
      "Iter: 4630, Loss: 1.138e-02\n",
      "Iter: 4640, Loss: 1.133e-02\n",
      "Iter: 4650, Loss: 1.126e-02\n",
      "Iter: 4660, Loss: 1.121e-02\n",
      "Iter: 4670, Loss: 1.116e-02\n",
      "Iter: 4680, Loss: 1.315e-02\n",
      "Iter: 4690, Loss: 1.159e-02\n",
      "Iter: 4700, Loss: 2.015e-02\n",
      "Iter: 4710, Loss: 1.413e-02\n",
      "Iter: 4720, Loss: 1.207e-02\n",
      "Iter: 4730, Loss: 1.145e-02\n",
      "Iter: 4740, Loss: 1.146e-02\n",
      "Iter: 4750, Loss: 4.586e-02\n",
      "Iter: 4760, Loss: 2.002e-02\n",
      "Iter: 4770, Loss: 1.391e-02\n",
      "Iter: 4780, Loss: 1.144e-02\n",
      "Iter: 4790, Loss: 1.105e-02\n",
      "Iter: 4800, Loss: 1.097e-02\n",
      "Iter: 4810, Loss: 1.096e-02\n",
      "Iter: 4820, Loss: 1.186e-02\n",
      "Iter: 4830, Loss: 4.959e-02\n",
      "Iter: 4840, Loss: 1.727e-02\n",
      "Iter: 4850, Loss: 1.319e-02\n",
      "Iter: 4860, Loss: 1.183e-02\n",
      "Iter: 4870, Loss: 1.131e-02\n",
      "Iter: 4880, Loss: 1.091e-02\n",
      "Iter: 4890, Loss: 1.074e-02\n",
      "Iter: 4900, Loss: 1.065e-02\n",
      "Iter: 4910, Loss: 1.059e-02\n",
      "Iter: 4920, Loss: 1.054e-02\n",
      "Iter: 4930, Loss: 1.049e-02\n",
      "Iter: 4940, Loss: 1.044e-02\n",
      "Iter: 4950, Loss: 1.069e-02\n",
      "Iter: 4960, Loss: 7.999e-02\n",
      "Iter: 4970, Loss: 2.566e-02\n",
      "Iter: 4980, Loss: 1.201e-02\n",
      "Iter: 4990, Loss: 1.089e-02\n",
      "Training time: 3100.6974 seconds\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "all the input arrays must have same number of dimensions, but the array at index 0 has 1 dimension(s) and the array at index 1 has 2 dimension(s)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[6], line 9\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mTraining time: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00melapsed\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m.4f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m seconds\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m      8\u001b[0m \u001b[38;5;66;03m# Prediction\u001b[39;00m\n\u001b[0;32m----> 9\u001b[0m X_star \u001b[38;5;241m=\u001b[39m \u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mhstack\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrepeat\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mlen\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mt\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtile\u001b[49m\u001b[43m(\u001b[49m\u001b[43mt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mlen\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m2\u001b[39m)\n\u001b[1;32m     10\u001b[0m u_pred, v_pred, f_u_pred, f_v_pred \u001b[38;5;241m=\u001b[39m model\u001b[38;5;241m.\u001b[39mpredict(X_star)\n\u001b[1;32m     12\u001b[0m \u001b[38;5;66;03m# Compute and display errors\u001b[39;00m\n",
      "File \u001b[0;32m/opt/anaconda3/envs/RAG/lib/python3.11/site-packages/numpy/core/shape_base.py:357\u001b[0m, in \u001b[0;36mhstack\u001b[0;34m(tup, dtype, casting)\u001b[0m\n\u001b[1;32m    355\u001b[0m \u001b[38;5;66;03m# As a special case, dimension 0 of 1-dimensional arrays is \"horizontal\"\u001b[39;00m\n\u001b[1;32m    356\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m arrs \u001b[38;5;129;01mand\u001b[39;00m arrs[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mndim \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[0;32m--> 357\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_nx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconcatenate\u001b[49m\u001b[43m(\u001b[49m\u001b[43marrs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdtype\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdtype\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcasting\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcasting\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    358\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    359\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _nx\u001b[38;5;241m.\u001b[39mconcatenate(arrs, \u001b[38;5;241m1\u001b[39m, dtype\u001b[38;5;241m=\u001b[39mdtype, casting\u001b[38;5;241m=\u001b[39mcasting)\n",
      "\u001b[0;31mValueError\u001b[0m: all the input arrays must have same number of dimensions, but the array at index 0 has 1 dimension(s) and the array at index 1 has 2 dimension(s)"
     ]
    }
   ],
   "source": [
    "# Initialize and train model\n",
    "model = PhysicsInformedNN(x0, u0, v0, tb, X_f, layers, lb, ub)\n",
    "start_time = time.time()\n",
    "model.train(5000)\n",
    "elapsed = time.time() - start_time\n",
    "print(f'Training time: {elapsed:.4f} seconds')\n",
    "\n",
    "# Prediction\n",
    "X_star = np.hstack((np.repeat(x, len(t)), np.tile(t, len(x)))).reshape(-1, 2)\n",
    "u_pred, v_pred, f_u_pred, f_v_pred = model.predict(X_star)\n",
    "\n",
    "# Compute and display errors\n",
    "u_star = Exact_u.T.flatten()[:, None]\n",
    "v_star = Exact_v.T.flatten()[:, None]\n",
    "error_u = np.linalg.norm(u_star - u_pred, 2) / np.linalg.norm(u_star, 2)\n",
    "error_v = np.linalg.norm(v_star - v_pred, 2) / np.linalg.norm(v_star, 2)\n",
    "print(f'Error u: {error_u:.3e}')\n",
    "print(f'Error v: {error_v:.3e}')"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyPe5GD6YGH8hKYT2+AAz+NE",
   "provenance": [
    {
     "file_id": "1zEm5lknsdwGCUCSt7b4fUXqxeFngHE6P",
     "timestamp": 1730633483601
    }
   ]
  },
  "kernelspec": {
   "display_name": "RAG",
   "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}