Inverse Burger PINN experiments (The automatic approach).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "782d2e5b-eef0-43bd-b1f7-5aaf37921e6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\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\n",
    "\n",
    "from pysr import pysr, best, best_callable\n",
    "\n",
    "### For automatic PDE recovery ###\n",
    "from gplearn.genetic import SymbolicRegressor\n",
    "import sympy\n",
    "from sympy import *\n",
    "from sympy.parsing.sympy_parser import parse_expr\n",
    "from sympy.core import evaluate\n",
    "from utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2c0802d9-3751-43ac-b2f4-f739a7ab2bab",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = io.loadmat('data/burgers_shock.mat')\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",
    "N = 2000\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": "a222e088-1130-4be8-a41b-7ced0b82ea46",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded from ./saved_path_inverse_burger/equations_from_pysr.pkl\n",
      "-0.6860516*x0*x1 + 0.0020575821*x2\n"
     ]
    }
   ],
   "source": [
    "choice = 'pysr'\n",
    "\n",
    "if choice == 'gplearn':\n",
    "    ### Using a gplearn model ###\n",
    "    program = pickle_load('./data/gp_symreg.pkl')._program\n",
    "elif choice == 'pysr':\n",
    "    ### Using a pysr model ###\n",
    "    program = best(pickle_load('./saved_path_inverse_burger/equations_from_pysr.pkl'))\n",
    "\n",
    "expr, vars = build_exp(program)\n",
    "expr = expr.subs(Integer(-1), Float(-1.0, precision=53))\n",
    "# expr = manipulate_expr(expr)\n",
    "print(expr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "34c157d8-8190-4419-8544-47229aee323b",
   "metadata": {},
   "outputs": [],
   "source": [
    "constants = list(expr.atoms(Number))\n",
    "\n",
    "# Replace all positive coeffs with Sympy symbols \n",
    "with evaluate(False):\n",
    "    for i in range(len(constants)):\n",
    "        var = Symbol('C'+str(i))\n",
    "        expr = expr.subs(abs(constants[i]), var)\n",
    "variables = list(expr.atoms(Symbol))\n",
    "\n",
    "# Use the log form for the parameters that are less than 0.005\n",
    "# Have to store the indices which require torch.exp\n",
    "exp_indices = []\n",
    "for i in range(len(constants)):\n",
    "    constants[i] = abs(constants[i])\n",
    "    if constants[i] < 0.005:\n",
    "        constants[i] = log(constants[i])\n",
    "        exp_indices.append(i)\n",
    "\n",
    "pde_params = nn.Parameter(data=torch.tensor(constants, dtype=torch.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "23fd1aac-7584-4b05-b88e-7c6dd4fce6fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[C0, C1, x0, x1, x2]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### Sorting Symbols ###\n",
    "values = [string2int(str(e)) for e in variables]\n",
    "tmp = []\n",
    "for i in range(len(values)):\n",
    "    tmp.append((values[i], variables[i]))\n",
    "tmp.sort()\n",
    "variables = [e[1] for e in tmp]\n",
    "del tmp\n",
    "\n",
    "variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4a544034-e266-438f-b8a5-42369dfd0156",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-- loss_fn --\n",
      "C0*x2 + (-C1)*x0*x1\n",
      "-- pde_params --\n",
      "Parameter containing:\n",
      "tensor([-6.1862,  0.6861], requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "fn = lambdify(tuple(variables), expr)\n",
    "\n",
    "print('-- loss_fn --')\n",
    "print(expr)\n",
    "\n",
    "print('-- pde_params --')\n",
    "print(pde_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f665c3e4-d95d-40e1-81a2-235749b87b49",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Network(nn.Module):\n",
    "    def __init__(self, model, pde_parameters, loss_fn, exp_indices):\n",
    "        super(Network, self).__init__()\n",
    "        self.model = model\n",
    "        self.model.apply(self.xavier_init)\n",
    "        self.callable_loss_fn = loss_fn\n",
    "        self.pde_parameters = pde_parameters\n",
    "        self.exp_indices = exp_indices\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, update_network_params=True, update_pde_params=True):\n",
    "        total_loss = 0.0\n",
    "        uf = self.forward(x, t)\n",
    "        if update_network_params:\n",
    "            mse_loss = self.mean_squared(uf - y_input)\n",
    "            total_loss = total_loss + mse_loss\n",
    "        \n",
    "        if update_pde_params:\n",
    "            tmp_params = torch.zeros(self.pde_parameters.shape)\n",
    "            for idx in range(self.pde_parameters.shape[0]):\n",
    "                if idx in self.exp_indices: \n",
    "                    tmp_params[idx] = torch.exp(self.pde_parameters[idx])\n",
    "                else: \n",
    "                    tmp_params[idx] = self.pde_parameters[idx]\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",
    "            \n",
    "            l_eq = self.mean_squared(u_t-self.callable_loss_fn(tmp_params[0], tmp_params[1], uf, u_x, u_xx))\n",
    "\n",
    "            total_loss = total_loss + l_eq\n",
    "            \n",
    "        return total_loss\n",
    "    \n",
    "    def get_gradients_dict(self, x, t):\n",
    "        self.eval()\n",
    "        \n",
    "        uf = self.forward(x, t)\n",
    "        ß\n",
    "        ### PDE Loss calculation ###\n",
    "        # first-order derivatives\n",
    "        u_t = self.gradients(uf, t)[0]\n",
    "        u_x = self.gradients(uf, x)[0]\n",
    "        # Homo second-order derivatives\n",
    "        u_tt = self.gradients(u_t,t)[0]\n",
    "        u_xx = self.gradients(u_x, x)[0]\n",
    "        # Hetero second-order derivatives\n",
    "        u_xt = self.gradients(u_t, x)[0]\n",
    "        u_tx = self.gradients(u_x, t)[0]\n",
    "        \n",
    "        return {'uf':uf, 'u_x':u_x, 'u_xx':u_xx}, u_t\n",
    "    \n",
    "    def mean_squared(self, inn_tensor):\n",
    "        return (inn_tensor**2).mean()\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": 8,
   "id": "1b81ffa8-8d5d-4cbb-88c2-57703af3930e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained = False\n",
    "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, pde_parameters=pde_params, loss_fn=fn, exp_indices=exp_indices)\n",
    "if pretrained: network.load_state_dict(torch.load(\"./saved_path_inverse_burger/nn_without_physical_reg.pth\"), strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0b054684-f93d-499e-a386-bade84793145",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_u_train = torch.tensor(X_u_train).float().requires_grad_(True)\n",
    "u_train = torch.tensor(u_train).float().requires_grad_(True)\n",
    "\n",
    "X_star = torch.tensor(X_star).float().requires_grad_(True)\n",
    "u_star = torch.tensor(u_star).float().requires_grad_(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9c23c389-85ab-4619-a125-a07c71a9a586",
   "metadata": {},
   "outputs": [],
   "source": [
    "# optimizer = torch.optim.Adam(network.parameters(), lr=5e-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",
    "epochs = 500\n",
    "# weights_path = './saved_path_inverse_burger/frozen_pinn.pth'\n",
    "# weights_path = './saved_path_inverse_burger/nn_with_physical_reg_from_symreg.pth'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "54cbc0cf-6fa6-426e-9224-c36537f18879",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  0.049615528434515\n",
      "Epoch 100:  4.703639933723025e-05\n",
      "Epoch 200:  3.808536712313071e-05\n",
      "Epoch 300:  3.808536712313071e-05\n",
      "Epoch 400:  3.808536712313071e-05\n"
     ]
    }
   ],
   "source": [
    "network.train(); best_train_loss = 1e6\n",
    "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, update_network_params=True, update_pde_params=True)\n",
    "        if l.requires_grad:\n",
    "            l.backward()\n",
    "        return l\n",
    "\n",
    "    optimizer.step(closure)\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": 12,
   "id": "e2b23e59-35fa-4805-ba2e-21a34b44cfe0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.0031686916882638083, 0.991890013217926]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "est_lambdas = [e.detach().item() for e in network.pde_parameters]\n",
    "for idx in exp_indices:\n",
    "    est_lambdas[idx] = np.exp(est_lambdas[idx])\n",
    "est_lambdas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "793e8371-c880-4657-96be-b5ef370b1c52",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Loading the best weights ###\n",
    "# network.load_state_dict(torch.load(weights_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "82bbb2ad-86f3-48b0-a809-622bae347af4",
   "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": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "network.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "91a3d2cc-1d0e-424d-b66f-65d8225d31b7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.4526147065938181, 0.8109986782073975)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### pls chk this code before actual evaluation ###\n",
    "nu = 0.01 / np.pi\n",
    "grounds = [nu, 1.0]\n",
    "\n",
    "error_lambda_1 = np.abs(est_lambdas[0] - grounds[0]) / nu * 100\n",
    "error_lambda_2 = np.abs(est_lambdas[1] - grounds[1])*100\n",
    "\n",
    "error_lambda_1, error_lambda_2 # (0.3306746482849121, 1.30792548173677)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33c210ef-4a35-4326-94ce-c623e04914d6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "891596cd-723e-4384-a0b8-0db8492c0f9f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca1db110-772c-4251-a14f-cd1e002041f0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "286d944b-d18c-4df5-886d-fd4b9c776fe7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6399819a-445c-4e7d-b740-fe09db76da93",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52b9ebdf-7554-4b7e-8897-577a87ab7f84",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "29001626-bf09-4b39-87f4-e24d52c100d6",
   "metadata": {},
   "source": [
    "### Symbolic regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68d0204a-3c97-4298-b43d-bcddf6d2cdfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "grads_dict, target = network.get_gradients_dict(X_u_train[:, 0:1], X_u_train[:, 1:2])\n",
    "index2features = grads_dict.keys()\n",
    "print(index2features)\n",
    "\n",
    "G = torch.cat(list(grads_dict.values()), dim=1).detach().numpy()\n",
    "target = torch.squeeze(target).detach().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4192994-11be-4cfa-ae75-92d06a263908",
   "metadata": {},
   "outputs": [],
   "source": [
    "equations = pysr(G, target, niterations=100, binary_operators=[\"plus\", \"sub\", \"mult\"], unary_operators=[], batching=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e0eefc8-7168-454f-8755-625b0818a12c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(best(equations))\n",
    "fn = best_callable(equations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03bd241e-ffc1-48a3-8235-5b9a9ca67989",
   "metadata": {},
   "outputs": [],
   "source": [
    "uf = grads_dict['uf']\n",
    "u_x = grads_dict['u_x']\n",
    "u_xx = grads_dict['u_xx']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ad9d5e5-9f3f-4c9b-825e-7d52b6825912",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exreacted equation (for further fine-tuning)\n",
    "# u_t + 0.6860763*uf*u_x - 0.0020577204*u_xx"
   ]
  }
 ],
 "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
}