Inverse Burger PINN reproducing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "social-compensation",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import tensor\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torch.autograd import Variable, grad\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits import mplot3d\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "import scipy\n",
    "import scipy.io as io\n",
    "from pyDOE import lhs\n",
    "\n",
    "from sklearn.metrics import mean_squared_error, mean_absolute_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "consolidated-payroll",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = io.loadmat('data/burgers_shock.mat')\n",
    "\n",
    "N = 2000\n",
    "\n",
    "t = data['t'].flatten()[:,None]\n",
    "x = data['x'].flatten()[:,None]\n",
    "Exact = np.real(data['usol']).T\n",
    "\n",
    "X, T = np.meshgrid(x,t)\n",
    "\n",
    "X_star = np.hstack((X.flatten()[:,None], T.flatten()[:,None]))\n",
    "u_star = Exact.flatten()[:,None]              \n",
    "\n",
    "# Doman bounds\n",
    "lb = X_star.min(0)\n",
    "ub = X_star.max(0)\n",
    "\n",
    "idx = np.random.choice(X_star.shape[0], N, replace=False)\n",
    "X_u_train = X_star[idx, :]\n",
    "u_train = u_star[idx,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "annoying-musician",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Network(nn.Module):\n",
    "    def __init__(self, model):\n",
    "        super(Network, self).__init__()\n",
    "        self.model = model\n",
    "        self.model.apply(self.xavier_init)\n",
    "        self.lambda_1 = torch.nn.Parameter(torch.tensor([0.0]))\n",
    "        self.lambda_2 = torch.nn.Parameter(torch.tensor([-6.0]))\n",
    "        \n",
    "    def xavier_init(self, m):\n",
    "        if type(m) == nn.Linear:\n",
    "            torch.nn.init.xavier_uniform_(m.weight)\n",
    "            m.bias.data.fill_(0.01)\n",
    "        \n",
    "    def forward(self, x, t):\n",
    "        return self.model(torch.cat([x, t], dim=1))\n",
    "    \n",
    "    def loss(self, x, t, y_input):\n",
    "        lambda_1 = self.lambda_1\n",
    "        lambda_2 = torch.exp(self.lambda_2)\n",
    "        \n",
    "        uf = self.forward(x, t)\n",
    "        \n",
    "        # PDE Loss calculation\n",
    "        u_t = self.gradients(uf, t)[0]\n",
    "        u_x = self.gradients(uf, x)[0]\n",
    "        u_xx = self.gradients(u_x, x)[0]\n",
    "        l_eq = (u_t + lambda_1*uf*u_x - lambda_2*u_xx)\n",
    "        l_eq = (l_eq**2).mean()\n",
    "        \n",
    "        # Loss on the boundary condition\n",
    "        mse = F.mse_loss(uf, y_input, reduction='mean')\n",
    "        \n",
    "        return l_eq + mse\n",
    "    \n",
    "    def gradients(self, func, x):\n",
    "        return grad(func, x, create_graph=True, retain_graph=True, grad_outputs=torch.ones(func.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "thorough-leave",
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_nodes = 50\n",
    "\n",
    "model = nn.Sequential(nn.Linear(2, hidden_nodes), \n",
    "                        nn.Tanh(), \n",
    "                        nn.Linear(hidden_nodes, hidden_nodes),\n",
    "                        nn.Tanh(), \n",
    "                        nn.Linear(hidden_nodes, hidden_nodes),\n",
    "                        nn.Tanh(), \n",
    "                        nn.Linear(hidden_nodes, hidden_nodes),\n",
    "                        nn.Tanh(),\n",
    "                        nn.Linear(hidden_nodes, 1))\n",
    "\n",
    "network = Network(model=model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "activated-bahrain",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_u_train = tensor(X_u_train).float().requires_grad_(True)\n",
    "u_train = tensor(u_train).float().requires_grad_(True)\n",
    "\n",
    "X_star = tensor(X_star).float().requires_grad_(True)\n",
    "u_star = tensor(u_star).float().requires_grad_(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "victorian-workstation",
   "metadata": {},
   "outputs": [],
   "source": [
    "# optimizer = torch.optim.Adam(network.parameters(), lr=3e-4)  # metaopt also has .parameters()\n",
    "optimizer = torch.optim.LBFGS(network.parameters(), lr=5e-2, max_iter=50, max_eval=50, line_search_fn='strong_wolfe')\n",
    "\n",
    "best_train_loss = 1e6; epochs = 600\n",
    "network.train()\n",
    "\n",
    "weights_path = './saved_path_inverse_burger/pinn.pth'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "horizontal-allah",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  0.04405350238084793\n",
      "Epoch 100:  2.3927972506498918e-05\n",
      "Epoch 200:  1.0691312127164565e-05\n",
      "Epoch 300:  1.0691312127164565e-05\n",
      "Epoch 400:  1.0691312127164565e-05\n",
      "Epoch 500:  1.0691312127164565e-05\n"
     ]
    }
   ],
   "source": [
    "for i in range(epochs):    \n",
    "    ### Add the closure function to calculate the gradient. For LBFGS.\n",
    "    def closure():\n",
    "        if torch.is_grad_enabled():\n",
    "            optimizer.zero_grad()\n",
    "        l = network.loss(X_u_train[:, 0:1], X_u_train[:, 1:2], u_train)\n",
    "        if l.requires_grad:\n",
    "            l.backward()\n",
    "        return l\n",
    "\n",
    "    optimizer.step(closure)\n",
    "\n",
    "    # calculate the loss again for monitoring\n",
    "    l = closure()\n",
    "\n",
    "    if i > 400 and float(l.item()) < best_train_loss:\n",
    "        torch.save(network.state_dict(), weights_path)\n",
    "        best_train_loss = float(l.item())\n",
    "\n",
    "    if (i % 100) == 0:\n",
    "        print(\"Epoch {}: \".format(i), l.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "italic-mining",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### Loading the best weights ###\n",
    "network.load_state_dict(torch.load(weights_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "military-manor",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Network(\n",
       "  (model): Sequential(\n",
       "    (0): Linear(in_features=2, out_features=50, bias=True)\n",
       "    (1): Tanh()\n",
       "    (2): Linear(in_features=50, out_features=50, bias=True)\n",
       "    (3): Tanh()\n",
       "    (4): Linear(in_features=50, out_features=50, bias=True)\n",
       "    (5): Tanh()\n",
       "    (6): Linear(in_features=50, out_features=50, bias=True)\n",
       "    (7): Tanh()\n",
       "    (8): Linear(in_features=50, out_features=1, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "network.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "facial-brown",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.1516401767730713, 2.289382591936315)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nu = 0.01 / np.pi\n",
    "\n",
    "error_lambda_1 = np.abs(network.lambda_1.detach().item() - 1.0)*100\n",
    "error_lambda_2 = np.abs(torch.exp(network.lambda_2).detach().item() - nu) / nu * 100\n",
    "\n",
    "error_lambda_1, error_lambda_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "opened-omaha",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1.0, 0.9984835982322693)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1.0, network.lambda_1.detach().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "blessed-seminar",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.003183098861837907, 0.003255972173064947)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nu, torch.exp(network.lambda_2).detach().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "public-nudist",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}